--- title: Retrieval Heads in Transformers url: https://www.emergentmind.com/topics/retrieval-heads type: topic --- # Retrieval Heads in Transformers 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 [2404.15574][2607.01002][2605.27243]. ## 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 [2404.15574]. 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 [2404.15574]. ## 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 \(\alpha_{i,j}\in[0,1]\) in the convex mixture \[ attn_{i,j} = \alpha_{i,j}\cdot full\_attn + (1-\alpha_{i,j})\cdot streaming\_attn, \] 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 [2410.10819]. 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 [2407.15891]. 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: \[ \mathrm{QRscore}_h(q,d_i) =\frac{1}{|q|}\sum_{t_q\in q}\sum_{t_d \in d_i}A_h^{ t_q\rightarrow t_d}. \] 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 [2506.09944]. 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: \[ \phi^{(l,h)}_{t,j} = \mathbf{u}_{y_t}^\top \mathbf{o}^{(l,h)}_{t,j}, \] 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 [2607.01002]. ## 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 | [2404.15574] | | Induction/echo profiling | Top induction-score and echo-score heads on repeated-token probes | [2407.15891] | | Functional cache sensitivity | Output deviation under sink+recent restriction with sparse gates \(\alpha_{i,j}\) | [2410.10819] | | Query-focused retrieval | Query-to-gold-evidence attention mass | [2506.09944] | | Contrastive reranking heads | Positive-vs-negative attention discrimination | [2510.02219] | | Write-aware non-literal retrieval | OV-circuit logit contribution from needle vs off-needle | [2607.01002] | | Multimodal evidence routing | Question-to-text/image evidence attention with null-question calibration | [2605.27243] | Optimization-based identification in DuoAttention freezes model weights and trains only the gates \(\alpha_{i,j}\) with a distillation-plus-sparsity objective, \[ \mathcal{L} = \mathcal{L}_{\text{distill}} + \lambda \mathcal{L}_{\text{reg}}, \] using synthetic passkey supervision because ordinary language modeling underprovides long-range retrieval signal. The paper reports \(\lambda=0.05\), about **2,000 steps**, **AdamW**, and **8 A100 GPUs** [2410.10819]. RazorAttention uses a training-free synthetic repeated-token probe: generate \(K=2500\) 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 \(L_h=\max(4000, N/5)\) and **4** sink tokens [2407.15891]. Contrastive selection appears again in CoRe heads for listwise reranking. There, a head score is computed from positive versus hard-negative document attention, \[ S_{CoRe}(h) = \frac{\exp(s_{pos}^h / t)} {\exp(s_{pos}^h / t) + \sum_i \exp(s_{neg,i}^h / t)}, \] and the top **8 heads**, constituting **less than 1%** of all heads, are used for reranking [2510.02219]. In multilingual settings, retrieval heads remain attention-defined, but the paper on Retrieval-Transition Heads introduces a cross-lingual score \[ RTS_h = \sum_{t=1}^{N} \mathcal{A}^t \] based on source-target token alignment \(f:\mathcal N_{src}\to\mathcal N_{tgt}\), separating same-language retrieval from latent-to-target-language transition behavior [2602.22453]. In multimodal settings, the detection method is again attention-based but operates over annotated text spans or visual-token spans: \[ S_h(x) = \frac{1}{|q|} \sum_{t_q \in q} \sum_{g_i \in G_x} \sum_{t=s_i}^{e_i} A^h_{t_q \rightarrow t}, \] with null-question calibration to remove query-independent biases [2605.27243]. ## 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 [2410.10819]. RazorAttention adopts the same headwise asymmetry but compresses non-retrieval heads using local windows, sink tokens, and a **compensation token**: \[ \hat{\mathbf k} = \frac{1}{N_d}\sum_{m\in \mathcal D}\mathbf k_m, \qquad \hat{\mathbf v} = \frac{1}{N_d}\sum_{m\in \mathcal D}\mathbf v_m. \] Using \(L_h=\max(4000,N/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** [2407.15891]. 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** [2508.12407]. 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 [2602.11374]. ## 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: \[ p(x_t\mid x_{