Questions about FlexiCodec eval & local attention

[jee019]

Hi,

I have two questions regarding the official FlexiCodec checkpoint and the merging transformer implementation.

1. Gap between paper results and official checkpoint evaluation

I first evaluated the official DualCodec checkpoint, and the results were either almost identical to the paper or slightly worse, which seems reasonable.

However, when I evaluated the official FlexiCodec checkpoint using the same evaluation code, the gap from the paper results seemed noticeably larger. In particular, at threshold = 0.867, I obtained WER(RVQ1) = 4.43.

For evaluation, I used the LibriSpeech test-clean subset filtered to 4–10 seconds (1,237 utterances).

My measured results are as follows:

Model	Threshold	WER (RVQ1) ↓	WER (RVQ1:8) ↓	PESQ ↑	MCD ↓	SIM ↑	UTMOS ↑	Avg Frame Rate (Hz)
FlexiCodec (Paper)	1.000	2.76	2.23	3.35	2.76	0.85	4.22	12.5
	0.910	2.98	2.28	3.03	3.10	0.78	4.21	8.3
	0.867	4.15	2.53	2.76	3.42	0.71	4.18	6.25
FlexiCodec (Official ckpt)	1.000	2.89	2.27	3.42	2.97	0.86	4.15	-
	0.910	3.04	2.34	3.13	3.27	0.80	4.16	8.39
	0.867	4.43	2.63	2.88	3.52	0.73	4.13	6.36

Could you help clarify why this gap occurs?
Is there any additional inference configuration or evaluation setup that should be used for the official FlexiCodec checkpoint?

2. Question about local attention in the merging transformer

In the FlexiCodec paper, the merging module is described as using a Transformer with local attention, and it says:

“Each token in the local attention transformer can attend to ℓ_k = 8 tokens left and right.”

I initially thought this was the reason why the inference config uses:

transformer_context_frames: 16

However, when I debugged the code, it seemed that in the Transformer, self.causal = False, so local attention did not appear to behave the way I expected.

Also, in the training config, I saw:

transformer_context_frames: 38 # 3-second context

So I wanted to ask whether I may have misunderstood how the local attention is implemented.

Could you clarify:

• how the local attention is actually applied in the merging transformer
• why the training config uses transformer_context_frames: 38

Thank you very much.

[jiaqili3]

Hi, thanks for asking.

Is there any additional inference configuration or evaluation setup that should be used for the official FlexiCodec checkpoint?

No there isn't. The configuration you are using is correct, and as in your evaluation you see some metrics higher and some lower, that is normal because we retrained for the released checkpoint. If you want to report the flexicodec performance in your paper, you can safely use your test results.

self.causal = False, so local attention did not appear to behave the way I expected.

That's a great observation. We apply the windowing + noncausal transformer hoping it would attend only inside the window. You pointed out in the underlying code it only works in causal mode. I think you are correct and we missed it out.

About the transformer_context_frames setting, we initially used the 38 setting and used 16 when training the released model. Both configs are ok. 38 may give slightly better performance because it has a larger window

[jiaqili3]

I just checked the documentation of the Mimi Transformer, it says

class StreamingMultiheadAttention(nn.Module):
    """Similar to `nn.MultiheadAttention` but with support for streaming, causal evaluation.

    Args:
        embed_dim (int): Dimension to project to.
        num_heads (int): Number of heads.
        causal (bool): Causal mask applied automatically.
        context (int, optional): Number of time steps the attention can access to.
            When causal, can access `context` time steps into the past, **and when non causal,
            can access `context // 2` steps in the past, and the same in the future.**
        rope (`RotaryEmbedding`, optional): Rope embedding to use.
        weights_per_step (int): use different weights per time step. If non zero, should correspond to the
            number of possible time steps.
        device (torch.device, optional): Device on which to initialize.
        dtype (torch.dtype, optional): dtype to use.
    """

so according to it the windowing works for both causal and noncausal mode.

[jee019]

Thanks for checking. I saw that docstring too, but I’m still a bit confused by the actual forward() implementation.

It looks like context is only used inside the if self.causal: branch:
https://github.com/AmphionTeam/FlexiCodec/blob/main/flexicodec/model_blocks/mimi/transformer.py#L403

if self.causal:
    ...
    if self.context is not None:
        attn_bias = attn_bias & (delta < self.context)
else:
    attn_bias = None

So when causal=False, attn_bias=None is passed to scaled_dot_product_attention, which seems to result in full bidirectional/global attention rather than windowed local attention.

Could you confirm whether this means that transformer_context_frames is effectively not used when causal=False, and therefore this setting has no effect in the current non-causal QueryTokenAggregator?

[jiaqili3]

Thanks for checking. I saw that docstring too, but I’m still a bit confused by the actual forward() implementation.

It looks like context is only used inside the if self.causal: branch: https://github.com/AmphionTeam/FlexiCodec/blob/main/flexicodec/model_blocks/mimi/transformer.py#L403

if self.causal:
    ...
    if self.context is not None:
        attn_bias = attn_bias & (delta < self.context)
else:
    attn_bias = None

So when causal=False, attn_bias=None is passed to scaled_dot_product_attention, which seems to result in full bidirectional/global attention rather than windowed local attention.

Could you confirm whether this means that transformer_context_frames is effectively not used when causal=False, and therefore this setting has no effect in the current non-causal QueryTokenAggregator?

After checking, I think it's indeed the case that the docstring does not match the model behavior. Since I directly used Mimi's transformer, I believe the problem is also present in the upstream code. To fix, I suggest changing line 412 of .../model_blocks/mimi/transformer.py into:

if self.context is not None:
      # Symmetric windowed attention: each token attends to context // 2
      # tokens in each direction.
      pos_k = pos_k.view(1, -1)
      pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(
          -1, 1
      )
      delta = (pos_q - pos_k).abs()
      attn_bias = (pos_k >= 0) & (delta <= self.context // 2)
  else:
      attn_bias = None

i don't know whether it will change the behavior of the released model. Please update if you have more comment.