A Broader Look into MLLM's Retrieval Heads

A Broader Look into MLLM's Retrieval Heads. By Jia Wang and four others, from the Institute of Information Engineering of the Chinese Academy of Sciences, the University of Chinese Academy of Sciences, and one other institution. Published on the twenty-third of October.

Finding the relevant evidence inside a long input is a core capability of multimodal large language models. Earlier studies trace that capability to retrieval heads: a sparse set of attention heads that carry information from the context to the output. Those analyses, though, were confined to a limited set of modalities and the models built for them, and how retrieval heads could be used to improve a model was left largely open. The question here is whether the mechanism grows along with the model's input modalities. The answer reported is yes: when a new modality is added, retrieval stays sparse, it largely reuses the same heads, and those heads matter well beyond retrieval.

Audio-visual input serves as the newly added modality. The authors build paired inputs in which the same answer is either spoken aloud or shown on screen, so the evidence content is held fixed and only its modality changes. Under that control, retrieval heads remain sparse in the new modality and are shared to a large extent across modalities. Across eighteen checkpoints, spanning language, vision-language and omni-modal models, the omni-modal models largely preserve the retrieval heads of the backbones they were built on.

Ablation confirms that the heads are doing the work. Removing them eliminates nearly all audio-visual retrieval, so the new modality is leaning on the same narrow machinery as the old ones. The heads shared across modalities are a very small group, making up fewer than two percent of all heads. Yet ablating them does not only hurt retrieval. It also degrades document understanding, spatial perception and visual mathematics, tasks that are not framed as retrieval at all.

Those results turn the heads from an object of study into a lever. Building on the findings, the authors treat retrieval heads as something to optimize directly, training their attention with a ranking objective. The mechanism found by analysis becomes the target of training. That raises the page Recall@1 of Qwen3-VL-8B on MMDocIR by six point zero three points, and improves document retrieval across three model families.

The modality extension studied is audio-visual input specifically, set against language and vision-language baselines, over eighteen checkpoints. The cross-modal claim rests on paired inputs where content is fixed and only the presentation varies, spoken versus on screen. The headline gain from optimizing the heads is reported for one model on one benchmark, with the broader improvement described across three model families for document retrieval.

A small group of attention heads, under two percent of the total, turns out to be shared across modalities and inherited from the backbone. The authors use that group in two ways: as a lens for reading what a multimodal model does, and as a direct target for making it better.

A Broader Look into MLLM's Retrieval Heads. By Jia Wang and four others, from the Institute of Information Engineering of the Chinese Academy of Sciences, the University of Chinese Academy of Sciences, and one other institution. Published on the twenty-third of October. This is the long version: the whole paper rather than the abstract.

A multimodal model handed a long input often needs only a sliver of it. The answer to a question may sit in one second of a long video, surrounded by everything else. In language models, earlier work traced this ability to retrieval heads: a small set of attention heads that carry information from the context to the output, and whose removal sharply cuts long-context retrieval. That account was later extended to images. This study asks whether the mechanism stretches to audio and video, and whether it can be used to make a model better rather than only to describe it.

Two ideas carry the whole study. The first is how a head earns the name. Each example has a context, a question, and an annotated stretch of evidence. For every head, the authors measure how much attention flows from the question tokens to the evidence tokens, averaged over the question. Some of that attention has nothing to do with the question: a head may simply favour certain positions or content. So they rerun the same context with an uninformative question, subtract that baseline, and average over examples. Heads are ranked by this calibrated score, and the top five percent are called retrieval heads.

The second idea is the paired test. A six-digit number is planted in a background video. In the audio condition a synthesized voice reads the digits and the picture is untouched. In the visual condition the number appears in a centred banner for one second and the soundtrack is untouched. The video, the answer, the insertion point and the question, what is the secret number, are all shared. Only the modality of the evidence changes.

The analysis covers eighteen checkpoints: language models, vision-language models, one unified model, and eight omni-modal models that take text, image, audio and video. Several of the omni-modal ones are adapted from backbones in the same set, which lets the authors ask whether heads survive adaptation.

Scores only nominate heads. Ablation tests whether the model needs them. A head set is removed by zeroing its outputs before the attention output projection. The control removes the same number of heads in each layer, drawn at random from outside the retrieval heads. The effect of a set is how much further the task score falls than under that matched random removal, always on examples held out from head selection.

The second half turns the finding into a method on MMDocIR, a benchmark for retrieving the right page, or the right region of a page, from a long document. Each candidate page is scored by summing, over a fixed set of sixteen selected heads, the attention from the question to that page's visual tokens: the same evidence attention as before. Within each example the scores are rescaled between their minimum and maximum, so the logits have a fixed range. A ranking loss then pushes the relevant page above one hard negative from the same document. The loss is defined only on the selected heads, but its gradients update low-rank adapters on the attention projections. At inference, the uninformative-question baseline is subtracted once more.

The analysis lands four findings. Sparsity: in all eight omni-modal models, fewer than one in ten heads accounts for half the total positive score, under both spoken and displayed evidence. Sharing: between sixty-nine and eighty-three percent of the top heads are selected under both conditions, where random sets would share fewer than three, and rank correlation over all heads runs from zero point eight eight to zero point nine two. Retention: the SALMONN models are nearly identical to their vision-language backbones. Necessity: in Qwen2.5-Omni-7B, audio exact match falls from ninety-eight point three percent to zero when the retrieval heads go, while matched random removal leaves it at ninety-six point nine.

The most telling result is downstream. The heads shared by audio and visual evidence number fourteen in one model and ten in the other, fewer than two percent of all heads. Removing them hurts all eight benchmarks beyond random removal, most on visual mathematics and spatial perception, and the damage exceeds that from an equally sized set of image-retrieval heads in fifteen of sixteen model and task pairs.

For ranking, random heads give a page Recall at one of thirty-eight percent, untrained retrieval heads fifty-six point four five, and trained ones sixty-two point four eight: a gain of six point zero three points, which holds in all ten document domains. The training data have page labels and no layout labels at all, yet layout retrieval also improves, by about one and a half points. Random heads gain from training too, but stay far below retrieval heads at every head-set size.

The paper has no limitations section, so these come from what the evidence reports. The audio-visual analysis rests on one kind of target, a spoken or displayed six-digit number. The authors themselves find that a different target, natural-image evidence, gives lower correlations, so the sharing result is about matched content rather than modality in general. The downstream ablations cover two models, and the shared heads were significant on four of eight tasks per model. The ranking gains are smaller outside the Qwen family, about two points for Gemma3-12B, and layout gains are modest. Out-of-domain scores moved by less than a point either way: little harm, but no broad benefit.

The idea to carry away is that interpretability can hand you a training target. Retrieval heads turn out to be a compact, shared piece of machinery: the same few heads find a number whether it is heard or seen, they survive the move from a backbone to an omni-modal model, and about a dozen of them matter well beyond the task that found them. Once you know where evidence gets located, you can train exactly that attention, and the choice of heads still matters after training.