Research map / Large language models / Language, safety and interpretability
Retrieval Head Mechanistically Explains Long-Context Factuality
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Identifies sparse, universal "retrieval heads" in long-context LLMs via a copy-paste retrieval score on needle-in-a-haystack tests, and shows by masking them that they causally drive factual retrieval and chain-of-thought reasoning.
Problem
How transformer LLMs retrieve relevant information from arbitrary positions in long contexts, and which internal components are responsible for factual retrieval versus hallucination.
Approach
Defines a retrieval score per attention head as the frequency with which the head's most-attended token is a needle token being copied to the output during decoding, finds that a small set (<5%) of heads has high scores across model families, and shows that masking these heads causes retrieval failure and hallucination while masking random heads does not.
Key results
- Retrieval heads make up only about 3% to 6% of attention heads across the studied models, and masking the top 50 (about 5%) drops needle-test performance below 50.
- Masking retrieval heads reduces extractive QA F1 by 9.2% and 23.1%, while masking random non-retrieval heads has no significant impact.
- Base and variant models (continual pretraining, chat, upcycling) show retrieval-score correlation above 0.8, while cross-family correlation is below 0.1.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Mask retrieval heads (ablation) | Mistral-7B-Instruct-v0.2 · Needle-in-a-Haystack | retrieval performance | Hurt |
| Mask random non-retrieval heads (ablation) | Mistral-7B-Instruct-v0.2 · Needle-in-a-Haystack | retrieval performance | No clear effect |
| Mask retrieval heads (ablation) | Mistral-7B-Instruct-v0.2 · ExtractQA | F1: 9.2%; 23.1% | Hurt |
| Mask random non-retrieval heads (ablation) | Mistral-7B-Instruct-v0.2 · ExtractQA | F1 | No clear effect |
| Mask retrieval heads (ablation) | Mistral-7B-Instruct-v0.2 · MMLU, MuSiQue, and GSM8K with chain-of-thought | performance | Hurt |
| Mask random heads (ablation) | Mistral-7B-Instruct-v0.2 · MMLU, MuSiQue, and GSM8K with answer-only prompting | performance | No clear effect |
| Mask retrieval heads (ablation) | Mistral-7B-Instruct-v0.2 · MMLU, MuSiQue, and GSM8K with answer-only prompting | performance | No clear effect |
| Use full attention instead of sliding-window attention | Mistral · Needle-in-a-Haystack | retrieval capability | Helped |
Surprising findings
- Retrieval heads already exist in short-context pretrained base models and persist unchanged through long-context extension, chat tuning, and sparse upcycling.
- Masking retrieval heads hurts chain-of-thought reasoning substantially but barely affects answer-only prompting on MMLU and GSM8K.
Limitations
- KV cache compression implications are speculative and not evaluated.
- Retrieval-CoT relationship is only preliminarily studied, with in-depth analysis left to future work.
- Detection uses a copy-paste heuristic with a threshold of 0.1 chosen by the authors.
Experimental setup
Datasets: synthetic needle-in-a-haystack samples up-to-date news extractive QA (GPT-4 generated) MMLU MuSiQue GSM8K
Benchmarks: Needle-in-a-Haystack extractive QA (news) MMLU MuSiQue GSM8K
Metrics: retrieval score activation frequency needle retrieval accuracy F1 accuracy Pearson correlation
Baselines: random non-retrieval head masking
Builds on: induction heads needle-in-a-haystack CopyNet
Closest papers by meaning
- The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation (NeurIPS 2025 · similarity 0.90)
- Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads (NeurIPS 2026 · similarity 0.90)
- Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism (ICML 2025 · similarity 0.89)
- Retrieval from Within: An Intrinsic Capability of Attention-Based Models (NeurIPS 2026 · similarity 0.88)
- Longer Context, Deeper Thinking: Uncovering the Role of Long-Context Ability in Reasoning (NeurIPS 2025 · similarity 0.87)
- Understanding Synthetic Context Extension via Retrieval Heads (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.