FlashAttention-Compatible Modality-Mutual Attention with Virtual Reordering

Source abstract

Multimodal large language models (MLLMs) process an interleaved image–text sequence with a single autoregressive backbone under causal attention. Since the causal mask prevents image tokens from attending to the question that follows them, their representations are computed without access to the query, which has been identified as a source of visual hallucination. Modality-mutual attention (MMA) mitigates this issue by unlocking the image-to-question region of the causal mask, and dual-order training (DOT) mitigates it by additionally training on the reversed text-first order. However, MMA turns the global causal mask into a sample-dependent dense mask. Consequently, attention cannot be executed via FlashAttention, which is the most efficient attention path, and falls back to a slower implementation. Furthermore, the positional information in the unlocked region does not match that of the reversed input order that motivates the unlock, that is, the positional semantics are inconsistent. To address both problems, we propose FlashAttention-compatible MMA with virtual reordering (FMMA-VR), which repairs the positional information of the unlocked region as if the two modalities were reordered, without physically reordering the sequence. The proposed method is implemented as a Triton-based region-aware FlashAttention-style kernel. On identically pretrained models, it accelerates the attention path by up to 3.59× over conventional MMA for long sequences during supervised fine-tuning (SFT). In post-SFT evaluations, it stays on par with conventional MMA across ten benchmark metrics under full fine-tuning. Under low-budget LoRA adaptation, where the positional shift of conventional MMA is not absorbed, FMMA-VR additionally improves the hallucination-sensitive metrics by large margins.