Retrieval Heads in Transformers

• Retrieval heads are a distinct subset of attention heads in Transformers that extract task-relevant evidence from long input contexts.
• They are identified through varied paradigms such as copy-paste, query-focused selection, and contrastive ranking, each influencing model behavior and performance.
• Their use enhances long-context processing efficiency by reducing memory demands and accelerating decoding in systems like Llama-2 and Llama-3.

Searching arXiv for relevant papers on retrieval heads and related variants. arxiv_search query: "retrieval heads long-context attention heads"

Retrieval heads are a sparse, functionally specialized subset of Transformer attention heads that are disproportionately responsible for retrieving task-relevant information from long input context. In the earliest mechanistic formulation, they were defined through a copy-consistent pattern on Needle-in-a-Haystack-style tasks: when the model generates an answer token, a retrieval head places its highest attention on the matching token occurrence inside the answer-bearing span in the prompt. Subsequent work broadened the concept from literal copy-paste retrieval to query-focused evidence selection, end-to-end functional importance under cache restriction, non-literal retrieval mediated by the OV circuit, multilingual retrieval-to-generation transitions, and multimodal evidence routing over visual tokens (Wu et al., 2024, Gema et al., 1 Jul 2026, Li et al., 26 May 2026).

1. Canonical definition and mechanistic origin

Wu et al. introduced retrieval heads as attention heads that “redirect information from the input to the output” during long-context factual generation. In their operationalization, if $w$ is the current generated token, $r_q$ is the inserted answer span, and $j=\arg\max(\mathbf a)$ is the maximally attended input position for one head, then a retrieval event occurs when $w \in a$, $c_j = w$, and $j \in r_q$. The head-level score is

$$\text{RetrievalScore}(h) = \frac{|T_h \cap a|}{|a|}.$$

Heads with $\text{RetrievalScore}(h) > 0.1$ are treated as retrieval heads (Wu et al., 2024).

That paper established the canonical empirical profile of retrieval heads as universal, sparse, intrinsic, dynamically activated, and causal. Across the explored long-context-capable models, about 3% to 6% of attention heads had retrieval score $> 0.1$, less than 5% of heads were retrieval in the stronger practical sense, roughly 45% to 73% of heads had retrieval score exactly 0, and roughly 25% to 52% had only weak copy-paste behavior. Retrieval-head layouts were stable within model families, with Pearson correlation greater than 0.8 between a base model and its long-context or chat-adapted variant, but less than 0.1 across families. In Llama-2-7B-80K, 12 strongest retrieval heads had activation frequency 1; in Yi-6B-200K, 36 did. Masking the top retrieval heads sharply degraded Needle-in-a-Haystack, extractive QA, and chain-of-thought reasoning, whereas masking random non-retrieval heads had much smaller effect, grounding the later literature’s use of retrieval heads as a mechanistic intervention target (Wu et al., 2024).

2. Conceptual broadening beyond literal copy-paste

Later work retained the intuition that only a minority of heads perform long-range retrieval, but relaxed the original copy-based criterion. DuoAttention distinguished retrieval heads from streaming heads by asking which heads significantly alter model outputs when restricted to recent tokens and attention sinks. During identification, each head $(i,j)$ is assigned a gate $r_q$0 in the convex mixture

$r_q$1

and heads above threshold retain full attention at inference. This recasts retrieval heads as an end-to-end functional subset rather than a purely descriptive attention-map phenomenon (Xiao et al., 2024).

RazorAttention preserved the retrieve/process split but framed retrieval heads more narrowly as the small subset of heads that can “essentially pay attention to all input tokens” and support long-range recall. For RoPE models it linked retrieval heads to induction heads and echo heads, selecting the top 14% by induction score and top 1% by echo score; it reported that in RoPE models only a small fraction—around 15%—support long-range retrieval (Tang et al., 2024).

Query-focused work then shifted attention from output-token copying to query-conditioned evidence localization. QRHead scores a head by the attention mass from query tokens into gold evidence: $r_q$2 This reframing was motivated by long-context QA and re-ranking, where the crucial behavior is selecting relevant context for a query rather than copying the final answer verbatim (Zhang et al., 11 Jun 2025).

The strongest conceptual break came from LOCOS, which argued that standard detectors miss non-literal retrieval heads by construction because they reward heads whose attended token matches the generated token. LOCOS scores each head by the projection of its OV-circuit output onto the answer-token unembedding direction: $r_q$3 and contrasts needle versus off-needle source positions. In this view, a retrieval head must not only read from the relevant span but also write an answer-aligned update into the residual stream; many important retrieval heads are therefore invisible to attention-only literal-copy detectors (Gema et al., 1 Jul 2026).

3. Identification paradigms

The retrieval-head literature now contains several distinct identification regimes, each tied to a different notion of what retrieval means.

Identification regime	Signal used	Representative paper
Copy-paste retrieval	Argmax attention on matching answer token in needle span	(Wu et al., 2024)
Induction/echo profiling	Top induction-score and echo-score heads on repeated-token probes	(Tang et al., 2024)
Functional cache sensitivity	Output deviation under sink+recent restriction with sparse gates $r_q$4	(Xiao et al., 2024)
Query-focused retrieval	Query-to-gold-evidence attention mass	(Zhang et al., 11 Jun 2025)
Contrastive reranking heads	Positive-vs-negative attention discrimination	(Tran et al., 2 Oct 2025)
Write-aware non-literal retrieval	OV-circuit logit contribution from needle vs off-needle	(Gema et al., 1 Jul 2026)
Multimodal evidence routing	Question-to-text/image evidence attention with null-question calibration	(Li et al., 26 May 2026)

Optimization-based identification in DuoAttention freezes model weights and trains only the gates $r_q$5 with a distillation-plus-sparsity objective,

$r_q$6

using synthetic passkey supervision because ordinary language modeling underprovides long-range retrieval signal. The paper reports $r_q$7, about 2,000 steps, AdamW, and 8 A100 GPUs (Xiao et al., 2024).

RazorAttention uses a training-free synthetic repeated-token probe: generate $r_q$8 random tokens, repeat them 4 times, compute per-head induction and echo scores from attention weights, and select the top 14% induction heads plus top 1% echo heads. For non-retrieval heads it retains a local window $r_q$9 and 4 sink tokens (Tang et al., 2024).

Contrastive selection appears again in CoRe heads for listwise reranking. There, a head score is computed from positive versus hard-negative document attention,

$j=\arg\max(\mathbf a)$0

and the top 8 heads, constituting less than 1% of all heads, are used for reranking (Tran et al., 2 Oct 2025).

In multilingual settings, retrieval heads remain attention-defined, but the paper on Retrieval-Transition Heads introduces a cross-lingual score

$j=\arg\max(\mathbf a)$1

based on source-target token alignment $j=\arg\max(\mathbf a)$2, separating same-language retrieval from latent-to-target-language transition behavior (Patel et al., 25 Feb 2026). In multimodal settings, the detection method is again attention-based but operates over annotated text spans or visual-token spans: $j=\arg\max(\mathbf a)$3 with null-question calibration to remove query-independent biases (Li et al., 26 May 2026).

4. Retrieval heads as a systems primitive

A major systems consequence of the retrieval-head hypothesis is that full-context KV retention is not uniformly necessary across heads. DuoAttention makes this explicit: retrieval heads keep a full KV cache, while streaming heads use a constant-size cache consisting of sink tokens and a recent window. The paper reports near-full long-context performance with only 25% retrieval heads on MHA Llama-2-7B-32K and 50% retrieval heads on GQA Llama-3-8B-1048K. Quantitatively, it reports up to 2.55× memory reduction for MHA and 1.67× for GQA, up to 2.18× decoding speedup for MHA and 1.50× for GQA, up to 1.73× prefilling speedup for MHA and 1.63× for GQA, and up to 2.38× prefilling memory reduction for MHA and 1.53× for GQA. With 8-bit weights and 4-bit KV cache quantization, DuoAttention enables 3.30 million tokens on a single A100-80GB for Llama-3-8B, a 6.4× capacity increase over naïve BF16 full attention (Xiao et al., 2024).

RazorAttention adopts the same headwise asymmetry but compresses non-retrieval heads using local windows, sink tokens, and a compensation token: $j=\arg\max(\mathbf a)$4 Using $j=\arg\max(\mathbf a)$5 and retrieval-head protection, it reports over 70% KV-cache reduction, or about 3.125x compression, while remaining compatible with FlashAttention and introducing negligible overhead (Tang et al., 2024).

ZigzagAttention preserves the retrieval/streaming distinction but argues that mixing both head types within one layer creates extra tensor-access and indexing overhead. It therefore enforces layerwise-exclusive retrieval or streaming assignment. At 50% budget on long-context LLaMA-3-8B, it reports 38.44 average LongBench versus 39.45 for DuoAttention and 39.78 for the full model, while achieving lower decoding latency and reporting up to 37% acceleration in 1k context length (Liu et al., 17 Aug 2025).

Retrieval-aware distillation pushes the systems argument further by preserving only retrieval-critical heads when converting a Transformer into a Transformer–SSM hybrid. There, ablating each head on a synthetic KV-retrieval task ranks retrieval importance, and preserving just 10 heads—about 2% of the 512 heads in Llama-3.2-1B—recovers 95.0 retrieval-heavy coverage, versus 49.2 with 0 heads and 100.0 for the teacher. The resulting hybrid is reported as 5–6× more memory-efficient than comparable hybrids, and the SSM state can be reduced by up to 8× once retrieval heads are preserved (Bick et al., 11 Feb 2026).

5. Optimization, decoding, and ranking applications

Once retrieval heads are treated as a mechanistic bottleneck rather than a descriptive artifact, they become direct targets for training and inference-time control. DeCoRe starts from the hypothesis that masking retrieval heads induces hallucination, then contrasts the base model against a retrieval-head-masked model during decoding: $j=\arg\max(\mathbf a)$6 On Llama3-8B-Instruct, the entropy-controlled variant improves XSum factKB from 47.61 to 66.10, MemoTrap Macro Acc from 65.86 to 74.14, NQ-Open EM from 69.68 to 70.66, and NQ-Swap EM from 60.62 to 66.08 (Gema et al., 2024).

RHIO uses retrieval-head masking to synthesize realistic unfaithful LFQA outputs, then jointly trains faithful and unfaithful conditional modes with control tokens and performs self-induced contrastive decoding. In its pilot study on CLAPNQ, masking 100 retrieval heads drops FaithScore from 80.14 to 35.85 for Llama-2-7B-Chat, from 80.09 to 41.09 for Llama-2-13B-Chat, and from 86.36 to 51.49 for Llama-2-70B-Chat, whereas masking 100 random non-retrieval heads leaves scores near baseline. On GroundBench, RHIO reaches 82.35 for 7B and 83.77 for 13B, exceeding GPT-4o’s 82.33 average faithfulness score (Huang et al., 23 Jan 2025).

RetMask similarly turns retrieval-head ablation into a preference-learning signal. It identifies retrieval heads via Needle-in-a-Haystack, zeros their contribution through $j=\arg\max(\mathbf a)$7, generates full-model and ablated-model outputs, and trains with DPO to prefer the former. On Llama-3.1-8B-Instruct, it reports +2.28 HELMET points at 128K, with about +70% on generation with citation and +32% on passage re-ranking, while preserving general-task performance (Ma et al., 16 Jan 2026).

A more direct training intervention is MuDAF, which optimizes MDQA retrieval heads themselves through a head-level contrastive loss over question-side query vectors and passage-key representations. On Llama3.1-8B, Vanilla-SFT averages 37.8 on the reported long-context QA suite, whereas MuDAF reaches 50.5, with especially large gains on MDQA tasks: HotpotQA 46.8 → 69.6, 2WikiMultihopQA 50.5 → 66.2, and MuSiQue (LongBench) 28.9 → 48.2 (Liu et al., 19 Feb 2025).

Retrieval heads have also become a retrieval and reranking interface. QRRetriever uses QRHeads to score passages by accumulated query-to-passage attention mass; on Llama-3.1-8B-Instruct it improves LongMemEval from 46.5 under full context to 60.2 at Top-10 and CLIPPER from 31.3 to 44.4 at Top-5 (Zhang et al., 11 Jun 2025). QRRanker takes the next step and trains selected heads for listwise ranking, reaching 72.13 average Recall@10 across four datasets versus 67.59 for out-of-box QRHeads (Li et al., 12 Feb 2026). CoRe heads refine this further with positive-vs-negative contrastive head scoring; using 8 heads (<1%), CoRe-R reports BEIR average nDCG@10 of 52.1 on Llama-3.1-8B versus 50.2 for ICR and 51.3 for QR-R, and also shows that pruning the final 50% of layers preserves reranking accuracy while reducing memory by about 40% and latency by about 20% (Tran et al., 2 Oct 2025).

6. Multilingual and cross-lingual specialization

In multilingual models, retrieval heads are partly shared across languages but not fully language-invariant. Across English, German, Chinese, and Swahili, Qwen2.5-7B has 118 of 200 RH, or 59.0%, shared across all four languages, whereas Phi-3.5-3B has 94 of 211 RH, or 44.5%. Llama-3.1-8B shows especially strong bilingual English–Chinese overlap, with bilingual en–zh heads accounting for 19% of all RH, compared with 8.0% en–zh in Qwen-2.5-7B and 8.1% en–sw in Phi-3.5-3B (Patel et al., 25 Feb 2026).

The same work argues that standard retrieval heads are insufficient to explain multilingual generation, because reasoning may proceed in a latent conceptual space before being realized in a target language. It introduces Retrieval-Transition Heads (RTH), scored by source-target token alignment rather than same-language copying. RH and RTH have low overlap, low or negative cross-type rank correlation, and different layer profiles: in Qwen, prominent RTH lie mostly around layers 14–23, while language-specific RH concentrate later. Causally, masking top-$j=\arg\max(\mathbf a)$8 RTH hurts more than masking top-$j=\arg\max(\mathbf a)$9 RH on MMLU-ProX, MGSM, MLQA, and XQuAD. On MGSM, for example, Llama-3.1 drops by 54 points under RTH masking versus 21.8 under RH masking, and Qwen-2.5 drops by 17 points under RTH masking versus 5 under RH masking. Failure analysis further shows RH masking is dominated by retrieval failures, whereas RTH masking disproportionately causes Loss of Coherence, supporting the claim that RH retrieve while RTH mediate latent-to-target-language realization (Patel et al., 25 Feb 2026).

A plausible implication is that “retrieval head” is no longer a single invariant category once models must map retrieved content into a different output language. The multilingual literature therefore preserves the sparse-head premise while splitting retrieval from transition.

7. Multimodal, visual, and vision-language variants

The retrieval-head idea has also been extended to vision-LLMs, where the evidence-bearing substrate is often a visual token sequence rather than a text span. VERA identifies Visual Evidence Retrieval (VER) Heads as heads whose attention aligns with gold evidence patches in document-style images. In Qwen3-8B-VL, fewer than 1.65% of attention heads are classified as VER heads, Spearman rank correlations of per-head retrieval scores across datasets are consistently above 0.44, and masking the top-5 most frequent VER heads yields an average drop of -4.10 versus -0.45 for random masking and -0.66 for OCR-head masking. Leveraging these heads in the training-free VERA framework yields average relative improvements of 21.3% on Qwen3-VL-8B-Instruct and 20.1% on GLM-4.1V-Thinking across five benchmarks (Pei et al., 9 Feb 2026).

A distinct line of work argues that visual text extraction is not just multimodal retrieval under another name. “OCR heads” are defined through answer-token generation conditioned on evidence patch tokens, and are reported to involve 91.3% more heads on average than text-based retrieval for strongly active heads. Their Jaccard overlap with retrieval heads is low, especially for stronger retrieval-score bands, and their activation frequency closely aligns with OCR score, leading the paper to characterize them as less sparse, qualitatively distinct, and statically activated (Baek et al., 21 May 2025).

The broadest multimodal generalization is the 2026 study of multimodal retrieval heads (MMRetHeads) in long-context LVLMs. There, retrieval is defined as question-to-evidence attention over either text-token spans or visual-token spans, calibrated by a null question. Across six LVLMs, only 4.4–10.2% of heads account for 50% of positive calibrated retrieval-score mass. Masking the top-5% selected heads drops MMLongBench-Doc from 48.2% to 5.7% and SlideVQA from 71.2% to 8.9%, while random-head masking is far less damaging. The selected heads are partly shared across modalities—average text-vs-image top-5% overlap is 0.51, while text-vs-rendered-text overlap is 0.79—but image retrieval heads vary more with context length and haystack modality composition. Without further training, selected-head scoring on MMDocIR improves Recall@1 by 7.7/7.4 macro/micro points for page retrieval and 6.3/6.8 points for layout retrieval over the strongest reported baselines (Li et al., 26 May 2026).

Taken together, these multimodal results suggest that retrieval heads are best understood as a family of sparse evidence-routing mechanisms whose concrete instantiation depends on modality. Some variants are close to text-only retrieval; others, such as OCR heads and visual evidence retrieval heads, are functionally adjacent but not identical.

8. Limitations, controversies, and open questions

Despite the convergence on sparse, high-impact heads, the literature is explicit that retrieval heads are not a complete explanation of long-context competence. The original mechanistic work already noted that retrieval heads explain which part of the model seeks information from input tokens, but not fully how query/key features trigger them, how they interact with MLPs, or how they coordinate across layers (Wu et al., 2024). Synthetic-context-extension work sharpened this point by showing that retrieval heads are necessary but not totally sufficient: synthetic data can recover many of the same heads as real data and still underperform, implying upstream representation quality and downstream integration also matter (Zhao et al., 2024).

Identification itself remains contested. Attention-only detectors based on argmax token matching are simple and interpretable, but they can miss distributed retrieval, paraphrastic evidence use, and non-literal retrieval mediated by OV circuits (Gema et al., 1 Jul 2026). Conversely, end-to-end or write-aware methods rely on proxy tasks, hidden-state matching, or direct-path approximations, and can themselves be task-dependent. Several papers also emphasize model dependence: retrieval-head ratios differ between MHA and GQA, distributed retrieval organization can blunt optimization gains, and multilingual or multimodal settings introduce additional specialization (Xiao et al., 2024, Ma et al., 16 Jan 2026, Patel et al., 25 Feb 2026).

A stable consensus nevertheless emerges. Retrieval is not uniformly distributed across all heads. A small subset of heads repeatedly carries disproportionate responsibility for locating or routing long-range evidence, whether the evidence is a literal answer span, a paraphrased support passage, a target-language-aligned source token, or a visual patch. What varies across the literature is not the existence of this sparse subnetwork, but the operational definition of retrieval and the level at which it is measured: copying, query-to-evidence attention, output sensitivity under cache restriction, OV-aligned logit contribution, multilingual transition, or multimodal grounding.