A Broader Look into MLLM's Retrieval Heads
Abstract
Locating relevant evidence in long inputs is a core capability of Multimodal Large Language Models. Prior studies attribute this capability to retrieval heads, a sparse set of attention heads that transfer information from the context to the output. However, these analyses are confined to a limited set of modalities and the corresponding models. Moreover, how retrieval heads can be leveraged to enhance the model calls for further exploration. In this paper, we ask whether this mechanism expands together with the input modalities of the model. Taking audio-visual input as a newly added modality, we construct paired inputs in which the same answer is either spoken or shown on screen. With the evidence content fixed, retrieval heads remain sparse in the new modality and are shared to a large extent across modalities. Across 18 language, vision-language, and omni-modal checkpoints, omni-modal models largely preserve the retrieval heads of their backbones. Ablating these heads eliminates nearly all audio-visual retrieval. The heads shared across modalities make up fewer than 2\% of all heads, yet ablating them also degrades document understanding, spatial perception, and visual mathematics. Building on these findings, we turn retrieval heads into a target of optimization by training their attention with a ranking objective. This raises the page Recall@1 of Qwen3-VL-8B on MMDocIR by 6.03 points and improves document retrieval across three model families. Functional attention heads can thus serve both as a lens on multimodal models and as a target for improving them.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper investigates the role of retrieval heads in omni-modal models, demonstrating their sparsity, cross-modal sharing, and causal importance, and proposes a ranking-based optimization method that improves multimodal document retrieval. By extending mechanistic interpretability to audio-visual inputs and linking it to practical model improvement, the study provides a valuable bridge between understanding and optimizing multimodal evidence localization.
Methodology
The paper proposes a unified framework for analyzing "retrieval heads" in Multimodal Large Language Models (MLLMs), extending prior work from text and vision-language models to omni-modal (audio-visual) settings. The methodology involves constructing paired audio-visual inputs (AV-NIAH) where the same evidence is presented either as speech or on-screen text, allowing for controlled comparison of attention mechanisms. The authors define a calibrated retrieval score based on attention from question tokens to evidence tokens, subtracting a baseline. They then perform causal ablation studies to verify the necessity of these heads. Furthermore, they introduce a novel optimization technique where the attention of identified retrieval heads is fine-tuned using a ranking loss to improve document retrieval performance. This dual approach of mechanistic analysis and targeted optimization is methodologically sound and well-structured.
Results & evaluation
The experimental evaluation is extensive, covering 18 checkpoints across language, vision-language, and omni-modal models. The authors demonstrate that retrieval heads are sparse, shared across modalities (audio/visual), and preserved during model adaptation. Ablation studies show that removing these heads significantly degrades performance on audio-visual retrieval and broader multimodal tasks. The optimization component shows a 6.03 point improvement in Recall@1 on MMDocIR for Qwen3-VL-8B, with consistent gains across other model families. The use of paired inputs and controlled ablations provides strong causal evidence for the findings.
Reproducibility
The paper provides detailed descriptions of the detection method, ablation protocol, and training procedure. Specific hyperparameters such as the number of heads selected ($k=16$) and the use of LoRA adapters are mentioned. However, specific implementation details for the AV-NIAH dataset construction (e.g., exact video sources, speech synthesis parameters) are somewhat brief, though references to prior work are provided. The code availability is not explicitly stated in the text, which is a minor gap for full reproducibility.
Limitations
The analysis is primarily focused on a specific set of models (Qwen, SALMONN, Gemma) and may not generalize to all MLLM architectures. The optimization method is tested only on document retrieval tasks, and its applicability to other multimodal reasoning tasks is not explored. The paper relies on the assumption that attention scores directly correlate with functional importance, which, while supported by ablation, is a simplification of the complex dynamics within transformer layers.
Broader impact
This work contributes to the understanding of how MLLMs process multimodal information, offering insights into the shared mechanisms across modalities. The proposed optimization technique provides a practical method for enhancing retrieval capabilities in multimodal models, which could be applicable in various downstream tasks requiring precise evidence localization. The findings on the preservation of retrieval heads during modality adaptation have implications for model design and fine-tuning strategies. The paper investigates the role of retrieval heads in omni-modal models, demonstrating their sparsity, cross-modal sharing, and causal importance, and proposes a ranking-based optimization method that improves multimodal document retrieval. By extending mechanistic interpretability to audio-visual inputs and linking it to practical model improvement, the study provides a valuable bridge between understanding and optimizing multimodal evidence localization.
Author of this paper? Send a correction.