Causal attention for factorization blocks?

[songweii]

Thank you for your excellent work. I noticed that in the earlier version of the code, the AttnProjection module used causal attention (which was also explicitly described in the early version of the arxiv paper: https://arxiv.org/pdf/2502.20321v1).

To ensure compatibility with autoregressive generation, the factorization blocks are configured with causal attention.

However, it seems that in the current codebase it has been changed to standard bidirectional attention:

UniTok/models/vqvae.py

Line 41 in bb8012e

x = scaled_dot_product_attention(q, k, v)

May I ask what motivated this change, and whether there were any specific considerations behind it?

Or more directly: what differences would these two choices bring in practice? Specifically in terms of reconstruction, semantic accuracy, understanding capability, and autoregressive generation performance?

[machuofan]

Thanks for your interest in our work. In our original implementation, we used causal attention in the attention projection layers to ensure compatibility with AR inference:

UniTok/eval/liquid/model/quant.py

Line 431 in bb8012e

class CausalAttention(nn.Module):

However, we found that causal attention can slightly degrade understanding performance. As a result, we switched the tokenizer projection layer to use full attention.

For the MLLM stage, we reuse the tokenizer’s projection layer, then:

1. switch it back to causal attention, and
2. fine-tune the projection layer together with the MLLM.

This way, the projection layer is explicitly adapted to causal attention mode for AR usage, while the tokenizer benefits from full attention during training.