One-sentence Summary

Researchers from the Institute of Information Engineering, Chinese Academy of Sciences, the University of Chinese Academy of Sciences, and Tencent demonstrate that multimodal retrieval heads remain sparse and largely shared across audio-visual modalities, and that training these heads with a ranking objective improves Qwen3-VL-8B's page Recall@1 on MMDocIR by 6.03 points and document retrieval across three model families.

Key Contributions

• A paired analysis of retrieval heads across audio-visual evidence modalities and 18 language, vision-language, and omni-modal checkpoints shows that retrieval heads remain sparse in the new modality, are largely shared across modalities, and are preserved from backbones in omni-modal models.
• Ablation experiments provide causal evidence that retrieval heads are necessary for audio-visual retrieval, and that a shared cross-modal subset making up fewer than 2% of all heads also supports broader multimodal reasoning such as document understanding, spatial perception, and visual mathematics.
• A ranking-based optimization method trains retrieval-head attention to favor relevant pages over distractors, improving Qwen3-VL-8B page Recall@1 on MMDocIR by 6.03 points and improving document retrieval, including layout retrieval without layout supervision, across three model families.

Introduction

Multimodal large language models must often locate a small piece of relevant evidence inside long inputs that span many modalities, making evidence localization a central challenge. Prior work in language and vision-language models identified sparse retrieval heads that transfer relevant context to the output, but it remained unclear whether this mechanism generalizes to newer modalities such as audio and whether it could be used to improve models rather than only analyze them. The authors address this by studying paired audio-visual evidence across 18 omni-modal checkpoints, showing that retrieval heads are sparse, shared across modalities, preserved from backbones, and causally important for multimodal reasoning. They then use these findings to optimize retrieval-head attention with a ranking objective, improving multimodal document retrieval across multiple model families.

Dataset

The authors use three evidence-annotated benchmark families for unified retrieval-head detection.

• Text benchmarks:

  • Source: evidence-annotated subsets of LongMemEval and Natural Questions, processed by QRHead.
  • Each question includes annotated evidence needed to answer it.

• Interleaved image-text benchmark:

  • Source: MM-NIAH.
  • It places text or natural-image needles inside long interleaved contexts.
  • Evidence annotations are available for each question.

• Audio-visual benchmark:

  • Source: newly constructed paired AV-NIAH examples built on background videos from MLVU.
  • A six-digit number is inserted into a background video in two modality conditions.
  • Audio condition: synthesized speech reads the digits while the video track stays unchanged.
  • Visual condition: the number appears in a centered banner for one second while the audio track stays unchanged.
  • Both conditions share the same source video, answer, insertion depth, and the question “What is the secret number?”
  • The speech or banner time span is mapped to the corresponding audio or video tokens, forming the evidence span G_i.

• How the data are used:

  • The data are used to analyze 18 checkpoints spanning language, vision-language, unified, and omni-modal models.
  • Evaluation targets the language decoder that produces text, including the Thinker of Qwen2.5-Omni and the understanding expert of BAGEL.
  • Several omni-modal checkpoints are adapted from language or vision-language backbones in the same set, which supports retention analysis.

• Notes on scale and processing:

  • Subset sizes, filtering rules, training split, and mixture ratios are not reported in this excerpt.
  • No separate cropping strategy is described.
  • The described data are used for detection and evaluation rather than model training.

Method

Unified Retrieval-Head Detection and Guided Ranking

The authors propose a unified method for identifying retrieval heads, comparing them across models, and using their attention patterns to train document ranking. The core idea is that a retrieval head should direct more attention from question tokens to evidence tokens than to the rest of the input.

For example $i$i, let $Q_i$Qi​ denote the question-token positions and $G_i$Gi​ the evidence-token positions. The raw evidence attention of head $(l,h)$(l,h) is

where $a_{lh}^{(i)}(q,t)$alh(i)​(q,t) is the attention from position $q$q to position $t$t in head $h$h of layer $l$l. Part of this attention reflects positional or content preferences that do not depend on the question. To remove those preferences, the authors recompute the same statistic with an uninformative question on the same context, denoted $u_{lh}^{(i,\varnothing)}$ulh(i,∅)​. The calibrated retrieval score is then

Heads are ranked by $s_{lh}$slh​. Unless otherwise noted, retrieval heads are defined as the top 5% of heads in this ranking. All attention values are computed under the model’s native visibility constraints.

Comparison and intervention protocol

Two types of comparison structure the analysis. Within a model, changing only the evidence modality tests whether retrieval heads are shared across modalities. Across related checkpoints, holding the text task fixed tests whether heads are retained under modality adaptation. In both cases, Spearman’s $\rho$ρ over all heads measures agreement of the full ranking, while overlap among top-ranked sets measures agreement among the strongest heads.

Attention scores identify candidate heads, but the authors also test whether the model actually relies on them. A head set $S$S is ablated by zeroing the outputs of its heads before the attention output projection. As a control, the same number of heads is removed from each layer, drawn at random from outside the retrieval heads, across three seeds. Let $t$t be the task score after ablating $S$S, and let $\bar{r}$rˉ be the mean score after matched random ablations. The effect of $S$S is its additional drop:

A large $E$E means that removing $S$S is far more harmful than removing an equally sized, equally distributed set of other heads. All ablations use examples held out from head selection, and significance is assessed with paired cluster-bootstrap tests and Holm correction.

Retrieval-head-guided document ranking

The authors also test whether training retrieval-head attention improves document ranking. Given a question and candidates $\mathcal{D} = \{d_1, \ldots, d_m\}$D={d1​,…,dm​} from one document, the model ranks candidates by relevance. For a fixed retrieval-head set $\mathcal{H}_k$Hk​, question positions $Q$Q, and visual-token positions $T_j$Tj​ of candidate $d_j$dj​, the candidate score is

This applies the same attention statistic as Equation (1) over selected heads, treating each candidate as the evidence span.

Training proceeds in three steps. First, the head set $\mathcal{H}_k$Hk​ is selected once from the base model and kept fixed during training and inference. Second, raw attention scores are normalized within each example because their ranges are small and variable. The normalized score is

with $b = \operatorname{sg}(\min_j r_j)$b=sg(minj​rj​) and $\Delta = \operatorname{sg}(\max_j r_j - \min_j r_j)$Δ=sg(maxj​rj​−minj​rj​), where $\operatorname{sg}$sg stops gradients, $\gamma$γ sets the logit scale, and $\epsilon$ϵ guards against degenerate ranges. Detaching the endpoints fixes the logit range at $[0,\gamma]$[0,γ] while gradients still flow through every $r_j$rj​.

Third, the ranking objective is minimized. Each training example contains one relevant page $j^+$j+ and one hard negative from the same document, in random order. The loss is

The loss is defined only on the selected heads, but its gradients update LoRA adapters on the attention projections. At inference, candidates are ranked by $r_j$rj​ minus the score obtained with an uninformative question, matching the calibration used in Equation (2).

Experiment

The experiments first analyze retrieval heads in omni-modal models, showing that a sparse set of heads in middle and later layers handles spoken and displayed evidence, that matched audio-visual evidence recruits largely the same heads, and that these heads are inherited from language or vision-language backbones. Ablation studies confirm that removing these heads disrupts audio-visual retrieval and broader multimodal tasks, with a small cross-modal subset proving especially important. A document-ranking evaluation then uses retrieval-head guided scoring on MMDocIR, demonstrating that retrieval heads outperform random heads for page and layout retrieval, and that lightweight ranking training improves both granularities across multiple model families without layout supervision.

The mechanism analysis uses eighteen checkpoints spanning language, vision-language, unified, and omni-modal models. Input coverage ranges from text only in language models to text, image, video, and audio in omni-modal models. All measurements target the language decoder, and several omni-modal checkpoints are adapted from language or vision-language backbones to support retention analysis. Omni-modal checkpoints are the only group that covers all four input modalities: text, image, video, and audio. Vision-language checkpoints handle text, image, and video, while unified checkpoints cover text and image, and language checkpoints cover text only. Several omni-modal checkpoints are adapted from language or vision-language backbones, enabling retention comparisons across model types.

Ablating AV-shared heads produces additional performance drops beyond matched random ablation across reported downstream multimodal tasks and both models. The effect is generally stronger than ablating image-retrieval-selected heads, with significant damage on more tasks per model. This supports the conclusion that a small set of crossmodal heads contributes broadly to multimodal evidence use. AV-shared head ablation yields positive additional drops on every reported task for both models, indicating consistent damage beyond random ablation. The largest reported AV-shared drops occur on visual perception benchmarks such as CV-Bench, while image-retrieval heads produce smaller or even negative additional drops. AV-shared effects are statistically significant on four of eight tasks per model, compared with at most one of eight for image-retrieval heads.

Published baseline retrievers show a wide Recall@1 spread on page and layout retrieval, with ColPali and Col-Phi3 MMDocIR among the strongest. In the controlled Qwen3-VL-8B setting, retrieval-head ranking already outperforms random-head ranking before training. Fine-tuning improves page retrieval across all ten document domains and layout retrieval without layout supervision, and the trained retrieval-head scorer exceeds the published baselines though settings differ. Retrieval heads rank candidates better than random heads before training, and page retrieval improves after training in every document domain. Layout retrieval also improves without layout supervision, and the trained model exceeds published baselines even though evaluation settings differ.

Ranking training improves page and layout retrieval for all three model families. Page retrieval gains are largest for the Qwen models, and layout retrieval improves across all models despite lacking layout supervision. Both page and layout retrieval scores are higher after training for every model family. The two Qwen models show the largest page retrieval improvements. Gemma3-12B shows the smallest page gain but the largest layout gain among the evaluated models.

The experiments evaluate multimodal evidence use and retrieval across model families. A mechanism analysis spans eighteen language, vision-language, unified, and omni-modal checkpoints to examine language decoder retention across input modalities. Ablating a small set of AV-shared crossmodal heads causes consistent additional performance drops beyond random ablation, especially on visual perception tasks, indicating their broad role in multimodal integration. Retrieval experiments show that retrieval-head ranking outperforms random-head ranking before training, fine-tuning improves page and layout retrieval across domains and model families, and trained retrieval-head scorers exceed published baselines.