PaperlayerView exact revision on arXiv

Source-linked AI summary

A Broader Look into MLLM's Retrieval Heads

Jia Wang, Yuqing Li, Jiangnan Li, Zheng Lin, Mo Yu

arXiv:2610.32411v1cs.CL

TL;DR

Prior retrieval-head analyses cover a limited range of modalities and models, and using these insights to improve multimodal capabilities remains underexplored. This paper analyzes paired audio-visual evidence across model checkpoints and trains retrieval-head attention with a ranking objective. It finds sparse, shared heads that support retrieval and reports improved document retrieval from their optimization.

  • Problem

    Prior retrieval-head analyses cover a limited set of modalities and models, while how to leverage retrieval heads to improve multimodal capabilities remains less explored.

    Source: §1 · p. 2

  • Method

    The paper compares paired audio-visual inputs with fixed evidence content across 18 checkpoints, then trains selected retrieval-head attention with a ranking objective for document retrieval.

    Source: §1 · p. 2, §3.2 · p. 3

  • Results

    Retrieval heads are sparse, largely shared across modalities and retained from omni-modal backbones; ranking training raises Qwen3-VL-8B's page Recall@1 on MMDocIR by 6.03 points over the untrained scorer.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Ablations show that retrieval heads are necessary for audio-visual retrieval, and the cross-modally shared subset also supports broader multimodal reasoning.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1. Introduction

The introduction asks whether retrieval heads extend across modalities and model adaptation, and whether this mechanism can be used to improve multimodal capabilities. It studies paired audio-visual evidence and reports shared, functionally important heads alongside a ranking-based optimization method.

1. Introduction

  • The study compares 18 language, vision-language, and omni-modal checkpoints using paired audio-visual inputs and related text and image-text retrieval tasks.The paired video inputs place the same answer in speech or on-screen text while holding the background, question, and answer fixed.Source: §1 · p. 2
  • Retrieval heads are sparse, largely shared across spoken and displayed evidence, and largely retained when omni-modal models adapt from language or vision-language backbones.The analysis uses paired inputs that hold the answer and background fixed while changing whether the evidence is spoken or shown.Source: §1 · p. 2
  • Ablating retrieval heads severely disrupts audio-visual retrieval, while ablating the cross-modally shared subset—fewer than 2% of all heads—substantially impairs broader multimodal reasoning.The reported broader tasks include document understanding, spatial perception, and visual mathematics.Source: §1 · p. 2
  • Ranking training raises Qwen3-VL-8B's page Recall@1 on MMDocIR by 6.03 points over the untrained scorer and improves document retrieval across three model families.It also improves layout retrieval without layout supervision.Source: §1 · p. 2
  • The paper contributes a paired cross-modal analysis, evidence that retrieval heads support audio-visual retrieval and broader reasoning, and a ranking method that improves document retrieval.Its motivating gap is that prior analyses cover limited modalities and models, while using retrieval-head insights to improve multimodal capabilities remains less explored.Source: §1 · p. 2

2. Related work

Prior work characterizes retrieval heads in language models and extends functional-head analyses to vision-language models. Using these foundations, this paper extends retrieval-mechanism analysis and explores improving multimodal evidence use.

Functional attention heads

  • Language-model studies identify sparse retrieval heads and develop methods for analyzing them, while vision-language work finds specialized heads for visual evidence retrieval and related functions.The literature also includes analyses of dynamic retrieval heads and causal interventions connecting selected heads to evidence localization and downstream reasoning.Source: §2 · p. 2

Functional head utilization

  • Although retrieval-head use has informed language-model methods, exploiting these mechanisms to improve multimodal capabilities remains relatively underexplored.The paper frames its work as extending retrieval-mechanism analysis and exploring their potential to improve multimodal evidence use.Source: §2 · pp. 2–3
  • Prior language-model applications include attention-based retrieval, efficient inference, and interventions intended to improve evidence use.Emerging multimodal applications suggest similar potential, but broader use of these mechanisms remains less explored.Source: §2 · p. 2

3. Unified retrieval-head detection

The detection method ranks heads by question-to-evidence attention, using annotated evidence spans and a calibration against attention that does not depend on the question. It is applied across text, image-text, and paired audio-visual tasks, with head agreement and controlled ablations used to assess sharing, retention, and necessity.

3.1. Formulation of Detection Method

  • The method scores heads by attention from question tokens to annotated evidence tokens, calibrating for positional or content preferences with an uninformative question on the same context.Retrieval heads are ordinarily the top 5% of heads by this ranking, and attention is computed under each model's native visibility constraints.Source: §3.1 · p. 3

3.2. Detection data and models

  • Detection uses evidence-annotated text and image-text benchmarks plus paired AV-NIAH videos that keep the six-digit answer, source video, insertion depth, and question fixed while varying speech versus banner evidence.The analysis covers 18 checkpoints across language, vision-language, unified, and omni-modal models, measuring the language decoder.Source: §3.2 · p. 3, Table 1 · p. 3

3.3. Comparison and intervention protocol

  • Head sharing and retention are assessed through ranking agreement and top-set overlap, while necessity is tested by comparing ablations with matched random head removals.Ablations zero selected head outputs before the attention output projection; tests use held-out examples and paired cluster-bootstrap tests with Holm correction.Source: §3.3 · p. 3

4. Retrieval heads in omni-modal models

The analysis tracks retrieval-head sparsity, cross-modal sharing, retention through modality adaptation, and causal contributions in audio-visual models. It finds that retrieval heads support audio-visual retrieval, while a small cross-modal subset also contributes to broader multimodal tasks.

4.1. Retrieval heads are sparse for audio evidence

  • Fewer than one in ten heads accounts for half of total positive retrieval score in every omni-modal model, for both spoken and displayed evidence.High scores concentrate in a few middle and later-middle-layer coordinates.Source: §4.1 · p. 4, Figure 2 · p. 4

4.2. Audio-visual evidence share retrieval heads

  • Between 69% and 83% of top-5% heads overlap across matched audio and visual evidence, while full head-score rankings correlate at ρ = 0.88–0.92 in all eight models.The paired conditions differ only in whether the answer is spoken or displayed; the agreement persists with 300-second backgrounds and when one detection video is left out.Source: §4.2 · p. 4, Figure 3 · p. 4

4.3. Retrieval heads are retained along modality adaptation

  • Every modality-adapted checkpoint remains strongly correlated with its backbone on text retrieval; SALMONN models are nearly identical to their vision-language backbones (ρ ≥0.97).Together with cross-modal sharing, this indicates that audio-visual adaptation largely reuses backbone retrieval heads for new evidence modalities.Source: Figure 4 · p. 4, §4.3 · p. 5

4.4. Causal contribution of retrieval heads

  • Ablating top-5% retrieval heads eliminates audio retrieval in all four models and removes most visual retrieval; in Qwen2.5-Omni-7B, audio EM falls from 98.3% to 0 versus 96.9% after matched random ablation.The top-20 and AV-shared sets also reduce retrieval beyond random ablation; AV-shared effects grow with context length.Source: Figure 5 · p. 5, Figure 6 · p. 5, §4.4 · p. 6
  • The ablation effect is not limited to the top-5% cutoff: top-20 heads remove most audio retrieval, and AV-shared heads also reduce audio retrieval beyond random ablation.In Qwen2.5-Omni-7B, AV-shared heads also lower visual retrieval, with effects that grow with context length.Source: Figure 6 · p. 5, §4.4 · p. 6
  • Ablating AV-shared heads lowers performance beyond random ablation across eight downstream benchmarks, exceeding image-retrieval-head effects in 15 of 16 model–task pairs.The benchmarks span document question answering, visual perception, multidisciplinary reasoning, and visual mathematics.Source: §4.4 · p. 6, Table 2 · p. 6

5. Retrieval-head-guided document ranking

The section turns retrieval-head attention into a document-ranking signal, training it to rank relevant pages and regions above distractors. The results test its gains across retrieval granularities, model families, and head-set choices.

5.1. Task and candidate scoring

  • For each question, candidate pages or regions are scored by summing selected retrieval-head attention over question positions and candidate visual tokens.MMDocIR candidates are either page screenshots or layout regions such as paragraphs, tables, and figures.Source: §5.1 · p. 6

5.2. Ranking objective

  • Ranking training selects retrieval heads once, normalizes candidate scores within each example, and trains on a relevant page against a same-document hard negative.The loss acts on selected heads while its gradients update LoRA adapters on attention projections.Source: §5.2 · pp. 6–7

5.4. Results

  • For every head-set size and recall cutoff, trained retrieval-head scorers improve, while random-head scorers remain far below them.The largest gains appear in page R@1, where the base model has the most room for improvement.Source: Figure 7 · p. 7, §5.4 · p. 8

5.3. Experimental setup

  • Training uses MMDocIR page-level relevance labels, then evaluates page and layout retrieval with Recall@K under both Micro and Macro averaging.Layout retrieval tests transfer to a finer granularity because training contains no layout-level data.Source: §5.3 · p. 7

5.4. Results

  • 6.03 Micro R@1 points: ranking training improves Qwen3-VL-8B page retrieval over the untrained retrieval-head scorer under identical settings.The gain holds across all ten document domains; layout retrieval also improves despite no layout-level training data.Source: §5.4 · p. 7, Table 3 · p. 8
  • Ranking training improves both MMDocIR granularities for all three model families, while out-of-domain benchmark changes remain within one point.Some model–benchmark pairs improve slightly; others change by less than one point.Source: §5.4 · p. 8, Figure 8 · p. 8, Table 4 · p. 8

6. Conclusion

Across 18 checkpoints, the paper finds sparse retrieval heads shared across evidence modalities and largely preserved by omni-modal models. It also shows that these heads support retrieval and broader multimodal tasks, and that optimizing their attention improves document retrieval.

  • Across 18 checkpoints, omni-modal models largely retain sparse retrieval heads from their backbones, with heads shared when evidence content is matched across modalities.Source: §6 · p. 8
  • Retrieval heads are necessary for audio-visual retrieval, and the small subset shared across modalities also contributes to broader multimodal tasks.Source: §6 · p. 8
  • Training selected-head attention with a ranking objective improves multimodal document retrieval across three model families.The paper presents functional attention heads as both an interpretability lens and an optimization target.Source: §6 · p. 8

Related papers

Loading 2610.32411v1…