Organizations: Institute of Information Engineering, Chinese Academy of Sciences · School of Cyber Security, University of Chinese Academy of Sciences · Tencent
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.
Figures & tables
Group
Checkpoints
Input
Language
Qwen2.5-3B/7B, Qwen3-4B/8B
T
Vision-language
Qwen2.5-VL-3B/7B, Qwen3-VL-4B/8B, MiniCPM-V-4.5
T, I, V
Unified
BAGEL-7B
T, I
Omni-modal
Qwen2.5-Omni-3B/7B, MiniCPM-o-2.6/4.5, SALMONN 2+ 3B/7B, SALMONN Pro 4B/8B
T, I, A, V
Table 1: Checkpoints used for mechanism analysis. Inputs: text (T), image (I), video (V), audio (A). Measurements use the language decoder.
Figure 2: Retrieval heads are sparse in omni-modal models. (a,b) Calibrated scores ( Equation 2 ) under audio evidence for two representative models, normalized to [0,1] per model; rows are layers and columns are heads. (c) Share of heads in three normalized-score bins for eight omni-modal models under audio (A) and visual (V) evidence.
Figure 3: Spoken and displayed evidence recruit overlapping heads. (a) Retrieval heads (top 5%) detected under audio and visual evidence in each model: heads selected in both conditions (purple, with count over set size) and heads selected in only one. Independent random sets would share fewer than three heads. (b) Spearman correlation of all head scores between the two conditions. Gray bars give the range obtained when one of the three detection source–number pairs is left out; they show sensitivity to the detection material rather than a confidence interval.
Figure 4: Retrieval structure is retained across modality adaptation. Spearman correlation of head scores between a backbone and a related modality-adapted checkpoint on two text retrieval tasks. Shading marks ρ≥0.95 .
Figure 5: Retrieval-head ablation disrupts audio-visual retrieval. Exact match (EM, %) on held-out AV-NIAH examples without ablation, after ablating an equal number of layer-matched random heads, and after ablating the top-5% retrieval heads, for audio (a,b) and visual (c,d) evidence with 100-second and 300-second backgrounds (about 8K and 24K input tokens).
Figure 6: The ablation effect is robust to how the head set is chosen. Additional EM drop beyond matched random ablation ( Equation 3 ), with 95% intervals, for three head sets: (a) the top-20 heads of the tested modality, (b) the top-5% heads, and (c) the AV-shared heads, which appear in both the audio and visual top-20 sets. Bars cover both evidence modalities and background lengths (input tokens in parentheses).
Qwen2.5-Omni-7B
SALMONN Pro 4B
Score (%)
Additional drop E (pp)
Score (%)
Additional drop E (pp)
Task
Base
AV-shared
Image-ret.
Base
AV-shared
Image-ret.
MMLongBench-Doc
24.1
5.8
5.3
39.6
10.0
2.5
CV-Bench
75.5
16.3
5.7
80.5
14.8
6.8
MMStar
67.0
12.2
2.8
63.0
5.3
0.3
CountBenchQA
83.5
9.6
3.4
87.4
2.2
− 5.2
Table 2: AV-shared heads contribute to downstream multimodal tasks. Base is the unmodified task score (%); E=rˉ−t is the additional drop beyond matched random ablation (percentage points; larger means greater damage). AV-shared heads belong to both the audio and visual top-20 sets; Image-ret. heads are selected by natural-image retrieval. Each target set contains 14 heads in Qwen2.5-Omni-7B and 10 in SALMONN Pro 4B, with a separate size- and layer-matched random control. Bold denotes Holm-adjusted p<0.05 : AV-shared effects are significant on 4/8 tasks per model, versus at most 1/8 for Image-ret.
Figure 7: Functional head selection matters at every head-set size and cutoff. Recall of Qwen3-VL-8B when k retrieval heads (orange) or k random heads (blue) define the ranking score, before (dashed) and after (solid) training; shading marks the training gain. Top: Page R@1/3/5/10. Bottom: Layout R@1/5/10.
Page
Layout
Method
Macro
Micro
Macro
Micro
Baseline results
DPR
27.20
26.91
19.24
19.09
ColBERT
45.93
45.43
31.57
31.26
DSE docmatix
50.41
49.98
29.04
29.55
ColPali
55.27
54.58
33.28
33.06
Table 3: MMDocIR Recall@1 (%). Upper block: Baseline Results. Lower block: Qwen3-VL-8B attention ranking with k=16 heads under one shared configuration, so rows are directly comparable.
Page R@1
Layout R@1
Model
Base
Trained
Gain
Base
Trained
Gain
Qwen3-VL-8B
56.45
62.48
+6.03
47.14
48.60
+1.46
Qwen2.5-VL-7B
52.41
58.10
+5.69
31.79
32.60
+0.81
Gemma3-12B
44.75
46.92
+2.17
23.88
26.04
+2.16
Table 4: Ranking training across model families. MMDocIR Micro R@1 (%) with retrieval heads before and after page-level training; gains in points. Candidate windows differ across families, so rows compare training effects within a model.
Figure 8: Ranking training has little effect on out-of-domain performance. Scores on four out-of-domain benchmarks before and after the main training: MMMU, MMMU-Pro (standard), MathVista, and MMLongBench-Doc. Every change is within one point, and several model–benchmark pairs even improve slightly (green).