Skip to content

feat(models): Make MimiModel encoding padding-aware to ensure batch-to-individual consistency #43378

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
harshaljanjani wants to merge 5 commits into
base: main
Choose a base branch
from
Open

feat(models): Make MimiModel encoding padding-aware to ensure batch-to-individual consistency #43378

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4f05c5d
fix: Make MimiModel encoding padding aware for batch-individual consi…
harshaljanjani Jan 20, 2026
274f406
Merge branch 'main' into fix/mimi-batch-correctness
harshaljanjani Feb 10, 2026
191d904
fix: Batched encoding with true batch-parallel padding
harshaljanjani Feb 10, 2026
641f08f
Merge branch 'main' into fix/mimi-batch-correctness
harshaljanjani Feb 23, 2026
f7122e3
Merge branch 'main' into fix/mimi-batch-correctness
harshaljanjani Mar 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 36 additions & 5 deletions src/transformers/models/mimi/modeling_mimi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,12 +487,21 @@ def __init__(self, config: MimiConfig):
conv_layer = self.get_submodule(layername)
setattr(conv_layer, "layer_idx", layer_idx)

def forward(self, hidden_states, padding_cache=None):
def forward(self, hidden_states, padding_cache=None, output_lengths=None):
for layer in self.layers:
if isinstance(layer, (MimiConv1d, MimiResnetBlock)):
hidden_states = layer(hidden_states, padding_cache=padding_cache)
else:
hidden_states = layer(hidden_states)
# zero out positions after valid lengths so that garbage from conv bias
# does not leak into boundary positions at later strided convolutions.
if output_lengths is not None:
if isinstance(layer, MimiConv1d):
output_lengths = layer._get_output_length(output_lengths)
time_mask = torch.arange(
hidden_states.shape[-1], device=hidden_states.device
) < output_lengths.unsqueeze(1)
hidden_states = hidden_states * time_mask.unsqueeze(1)
return hidden_states


Expand Down Expand Up @@ -1483,12 +1492,26 @@ def _encode_frame(
Encodes the given input using the underlying VQVAE. The padding mask is required to compute the correct scale.
"""

# TODO: @eustlb, let's make the encoder support padding_mask so that batched inputs are supported.
embeddings = self.encoder(input_values, padding_cache=padding_cache)
input_lengths = None

@harshaljanjani harshaljanjani Feb 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leaving some reasoning here, I took a deeper look and wrote the trace script on main, they diverged at transformer position 12 (2.60e-02), downsample position 6 (0.33). The root cause is that the conv bias produces non-zero garbage at padded positions, and then later the strided convolutions at the boundary mix it into valid outputs (pytorch/audio#2242 documents the identical issue with wav2vec2).

Change adapts wav2vec2 patterns:
→ Copied modeling_wav2vec2.py#L680-L683 (zero padded tokens), adapted to run inside MimiEncoder.forward() after every layer using a time mask.
→ Copied modeling_wav2vec2.py#L1006-L1025 (compute output lengths after convs), adapted it to iterate over _mimiconv1d_layer_names and call _get_output_length.
→ Copied modeling_wav2vec2.py#L1027-L1045 (build attention mask from lengths), adapted it to use torch.arange(...) < encoder_output_lengths.
→ Also, just noting this down since it's Mimi-specific, the downsample uses pad_mode="replicate", so garbage positions should contain the last valid embedding (gather + torch.where) to match individual encoding behavior.

All ops are batched (no per-sample loops) now, happy to make further changes if needed :)

if padding_mask is not None and padding_cache is None:
padding_mask_2d = padding_mask.any(dim=1) if padding_mask.dim() == 3 else padding_mask
input_lengths = padding_mask_2d.sum(dim=-1)
embeddings = self.encoder(input_values, padding_cache=padding_cache, output_lengths=input_lengths)
attention_mask = None
encoder_output_lengths = None
if input_lengths is not None:
encoder_output_lengths = input_lengths
for layer_name in self.encoder._mimiconv1d_layer_names:
encoder_output_lengths = self.encoder.get_submodule(layer_name)._get_output_length(
encoder_output_lengths
)
attention_mask = torch.arange(embeddings.shape[-1], device=embeddings.device).unsqueeze(
0
) < encoder_output_lengths.unsqueeze(1)

# TODO: @eustlb, convert the padding mask to attention mask.
encoder_outputs = self.encoder_transformer(
embeddings.transpose(1, 2),
attention_mask=attention_mask,
past_key_values=past_key_values,
use_cache=use_streaming,
return_dict=return_dict,
Expand All @@ -1498,10 +1521,18 @@ def _encode_frame(
elif len(encoder_outputs) > 1:
past_key_values = encoder_outputs[1]
embeddings = encoder_outputs[0].transpose(1, 2)
embeddings = self.downsample(embeddings, padding_cache=padding_cache)

if encoder_output_lengths is not None:
last_valid_idx = (encoder_output_lengths - 1).clamp(min=0)
last_valid_emb = embeddings.gather(2, last_valid_idx.view(-1, 1, 1).expand(-1, embeddings.shape[1], 1))
garbage_mask = torch.arange(embeddings.shape[-1], device=embeddings.device).unsqueeze(
0
) >= encoder_output_lengths.unsqueeze(1)
embeddings = torch.where(garbage_mask.unsqueeze(1), last_valid_emb, embeddings)
embeddings = self.downsample(embeddings, padding_cache=padding_cache)
codes = self.quantizer.encode(embeddings, num_quantizers)
codes = codes.transpose(0, 1)

return codes, past_key_values, padding_cache

def get_encoded_length(self, input_length: torch.LongTensor) -> torch.LongTensor:
Expand Down