MLLM retrieval heads: a few attention heads serve every modality

MLLM retrieval heads are a small set of attention heads that pull the answer out of a long input, and a new paper finds that the same few heads do the job whether the evidence arrives as text, an image, speech or video. The paper, "A Broader Look into MLLM's Retrieval Heads", checked 18 model checkpoints. It also turned the finding into a small training trick that raised page retrieval on a document benchmark by 6.03 points.

The result is useful because it suggests a model does not build a separate lookup mechanism for each new input type. That makes the heads a cheap thing to inspect, protect or train.

What a retrieval head is

A transformer, the architecture behind these models, has many attention heads in each layer. Each head decides which earlier tokens the model should look at while it writes the next one. Earlier work found that a few heads do the "find the needle in the long document" job. Those are retrieval heads: heads that carry evidence from the context to the answer.

The open question was what happens when a language model gains eyes and ears. Does it grow new heads for images and audio, or reuse the old ones? The authors studied audio-visual input, the hardest case, and compared it with text and images.

How the team found the heads

For each head, the authors measured how much attention flows from the question tokens to the tokens that hold the evidence. They subtracted what the same head does for an uninformative question, so heads that always attend to the same place do not score well. The top 5% of heads by this score count as retrieval heads.

For text they used evidence-annotated parts of LongMemEval and Natural Questions. For images mixed with text they used MM-NIAH. For audio and video they built paired tests. A six-digit secret number is hidden in a background video from MLVU, either spoken in the audio or shown as a one-second banner, and the model is asked "What is the secret number?"

They ran this on 18 checkpoints: Qwen2.5 and Qwen3 language models, Qwen2.5-VL, Qwen3-VL and MiniCPM-V-4.5 vision-language models, BAGEL-7B, and eight omni-modal models (Qwen2.5-Omni, MiniCPM-o, SALMONN 2+ and SALMONN Pro).

What they found

• The heads are sparse. In all eight omni-modal models, fewer than one in ten heads accounts for half of the total retrieval score, for both audio and visual evidence.
• The heads are shared across modalities. Between 69% and 83% of the top 5% heads are the same under the audio and the visual condition. Two random sets of that size would share fewer than three heads. The correlation between the two full rankings is 0.88 to 0.92 across the eight models.
• Adding modalities barely moves them. The SALMONN models rank heads almost the same way as the vision-language models they were built on (correlation of 0.97 or higher).
• The heads are causal. In Qwen2.5-Omni-7B, zeroing the top 5% of heads drops exact-match accuracy on the audio secret-number task from 98.3% to 0. Zeroing the same number of random heads, matched by layer, leaves it at 96.9%. The top 20 heads alone remove most audio retrieval.

The heads shared by audio and visual retrieval are very few: 14 in Qwen2.5-Omni-7B and 10 in SALMONN Pro 4B. Yet ablating them also hurts tasks that look unrelated to needle search. On Qwen2.5-Omni-7B, accuracy falls by 22.3 points on MathVista, 16.3 on CV-Bench and 12.2 on MMStar. On SALMONN Pro 4B, it falls by 18.8 points on MathVision and 14.8 on CV-Bench. Across 16 model and task pairs, removing the audio-visual heads hurt more than removing image-only retrieval heads in 15.

Training the heads to rank pages

The authors then used the heads as a retriever. For a document question, they sum the attention that the question tokens pay to each candidate page across 16 chosen heads, and rank pages by that score. No new retrieval model is needed.

They then trained only the attention of those heads, using LoRA adapters (small add-on weights) and a ranking loss with one relevant page and one hard negative from the same document. Training used the MMDocIR training set with page-level labels only.

On Qwen3-VL-8B, page Recall@1 (micro average) rose from 56.45 to 62.48, a gain of 6.03 points. Layout Recall@1 rose from 47.14 to 48.60. Random heads scored 38.00 on page recall, so the choice of heads matters. The paper lists published retrievers on the same table, such as ColPali at 54.58 and Col-Phi3 MMDocIR at 57.10, but says they were evaluated under different settings. The clean comparison is trained heads against untrained heads.

The gain carries over to other models, with smaller size. Qwen2.5-VL-7B gained 5.69 points on page recall and Gemma3-12B gained 2.17. On MMMU, MMMU-Pro, MathVista and MMLongBench-Doc, training changed scores by less than a point, so it did not visibly damage general ability.

What to be careful about

The audio-visual test is synthetic: a six-digit number in a fixed question format. It shows that these heads are necessary for that kind of lookup, not that they handle every kind of reasoning over video. The paper also says which heads get recruited depends on the retrieval target as well as the modality, and it has no separate limitations section. The page snippet we read did not state a code or data release, though the authors have a project page. The numbers come from the authors' own experiments and have not been independently reproduced.

Why it matters

For anyone building on long-context or multimodal models, the practical reading is that a few heads carry much of the lookup work, and that pruning, quantizing or fine-tuning them carelessly could break retrieval across modalities at once. It also gives teams a low-cost way to build a page retriever out of a model they already run. Watch whether the authors release code, and whether others reproduce the 6-point gain on models outside the Qwen, MiniCPM and Gemma families.