---
title: "Retrieval Head Mechanistically Explains Long-Context Factuality"
description: "The authors propose that a small group of attention heads, which they call retrieval heads, performs a… 25 references, cited by 89 papers here."
authors: ["Wenhao Wu", "Yizhong Wang", "Guangxuan Xiao", "Hao Peng", "Yao Fu"]
year: 2024
venue: "arXiv:2404.15574"
source: "https://arxiv.org/abs/2404.15574"
doi: "10.48550/arxiv.2404.15574"
arxiv: "2404.15574"
licence: "CC BY 4.0 (arXiv)"
licence_url: "http://creativecommons.org/licenses/by/4.0/"
summary_licence: "Elamur Cosmic License"
canonical: "https://research.elamur.ai/free/en/papers/wu-2024-retrieval-head-mechanistically-explains-long/"
---
[Elamur Research](https://md.research.elamur.ai/free/en/) · [this page in HTML](https://research.elamur.ai/free/en/papers/wu-2024-retrieval-head-mechanistically-explains-long/)

[home](https://md.research.elamur.ai/free/en/) / [papers](https://md.research.elamur.ai/free/en/papers/) / [2024](https://md.research.elamur.ai/free/en/papers/2024/page/3/) / Retrieval Head Mechanistically Explains Long-Context Factuality

# Retrieval Head Mechanistically Explains Long-Context Factuality

**Authors:** [Wenhao Wu](https://md.research.elamur.ai/free/en/authors/wenhao-wu/), [Yizhong Wang](https://md.research.elamur.ai/free/en/authors/yizhong-wang/), [Guangxuan Xiao](https://md.research.elamur.ai/free/en/authors/guangxuan-xiao/), [Hao Peng](https://md.research.elamur.ai/free/en/authors/hao-peng/) and [Yao Fu](https://md.research.elamur.ai/free/en/authors/yao-fu/) · **Year:** 2024 · **Venue:** arXiv:2404.15574 · **Licence:** [CC BY 4.0 (arXiv)](http://creativecommons.org/licenses/by/4.0/)

**Links:** [Publisher's page](https://arxiv.org/abs/2404.15574) · [PDF](https://arxiv.org/pdf/2404.15574) · [DOI 10.48550/arxiv.2404.15574](https://doi.org/10.48550/arxiv.2404.15574) · [arXiv 2404.15574](https://arxiv.org/abs/2404.15574) · [OpenAlex](https://openalex.org/W4395482887)

**Citations:** 1 citations (OpenAlex, 2026-10-14) · [89](https://md.research.elamur.ai/free/en/papers/wu-2024-retrieval-head-mechanistically-explains-long/cited-by/) in our library, of the 141,840 papers whose bibliographies we read

## Summary
The paper investigates how long-context transformer LLMs locate and reuse facts placed anywhere in a long input. The authors propose that a small group of attention heads, which they call retrieval heads, performs a conditional copy-paste from the context into the generated output. To find these heads, they define a retrieval score based on Needle-in-a-Haystack trials. A head scores a hit when its most-attended input token lies inside the inserted needle and matches the token currently being generated. The score is a token-level recall over the needle, averaged across many context lengths and needle positions. The study covers LLaMA, Yi, Qwen and Mistral/Mixtral, including base, chat, length-extended and MoE variants. Retrieval heads appear in every model that can retrieve reliably. They are rare, they already exist after short-context pretraining, and they are causally required: masking them makes the model hallucinate fluent but wrong answers and hurts chain-of-thought reasoning. These results tie factuality to a specific circuit and inform hallucination research and KV cache compression.

## Key points
- Models studied: Llama-2-7B (plus 80K and 13B-64K extensions), Mistral-7B-v0.2 / Instruct, Mixtral-8x7B-v0.1, Yi-6B / Yi-6B-200K / Yi-34B-200K, and Qwen1.5-14B / Chat.
- Detection protocol: 20 context lengths between 1K and 50K tokens, each with 10 needle depths, giving about 600 tests per model. A head counts as a retrieval head if its average score exceeds 0.1.
- Fewer than 5% of heads reach a retrieval score above 0.5. Continued pretraining for longer context, chat alignment and sparse upcycling all reuse the base model's heads.
- In LLaMA 2 7B 80K, masking the top 20 retrieval heads lowers Needle-in-a-Haystack accuracy to 63.6, while masking 20 random heads leaves it at 94.7. With masked retrieval heads, CoT results on GSM8K, MMLU and MuSiQue drop, whereas direct answering is barely affected.
- The authors argue that full attention over the whole KV cache is needed for retrieval. As an example, Mistral v0.1, which uses sliding-window attention, fails the needle test, while v0.2 with full attention passes. They suggest that pruning the cache of non-retrieval heads could reduce memory use, since 100K tokens need more than 50GB for LLaMA 2 7B.

*Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.*

## Institutions

- [University of Washington](https://md.research.elamur.ai/free/en/universities/university-of-washington/) — University, Seattle, United States · 263 works · 180,581 citations
- [University of Edinburgh](https://md.research.elamur.ai/free/en/universities/university-of-edinburgh/) — University, Edinburgh, United Kingdom · 255 works · 68,771 citations
- [Massachusetts Institute of Technology](https://md.research.elamur.ai/free/en/universities/massachusetts-institute-of-technology/) — University, Cambridge, United States · 159 works · 65,157 citations
- [Peking University](https://md.research.elamur.ai/free/en/universities/peking-university/) — University, Beijing, China · 100 works · 35,618 citations

📚 Shelves (5): [🖥️ Computer Science](https://md.research.elamur.ai/free/en/shelves/computer-science/) · 9,085 works · [🤖 Artificial Intelligence](https://md.research.elamur.ai/free/en/shelves/computer-science/artificial-intelligence/) · 8,888 works · [🎭 Arts and Humanities](https://md.research.elamur.ai/free/en/shelves/arts-and-humanities/) · 9,296 works · [🗣️ Language and Linguistics](https://md.research.elamur.ai/free/en/shelves/arts-and-humanities/language-and-linguistics/) · 9,206 works · [👁️ Computer Vision and Pattern Recognition](https://md.research.elamur.ai/free/en/shelves/computer-science/computer-vision-and-pattern-recognition/) · 1,282 works

🧭 Topics (1): [🔍 Information retrieval & search](https://md.research.elamur.ai/free/en/topics/information-retrieval/) · 146 works

## Cite

Cite as: Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng and Yao Fu. 2024. Retrieval Head Mechanistically Explains Long-Context Factuality. arXiv:2404.15574.

The paper itself, not this page. BibTeX from the library's record of the work.

```bibtex
@misc{wu2024retrieval,
  title = {{Retrieval Head Mechanistically Explains Long-Context Factuality}},
  author = {Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng and Yao Fu},
  year = {2024},
  doi = {10.48550/arxiv.2404.15574},
  eprint = {2404.15574},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2404.15574}
}
```

- [← previous PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference](https://md.research.elamur.ai/free/en/papers/yang-2024-pyramidinfer-pyramid-kv-cache-compression/)
- [next → Graph Chain-of-Thought: Augmenting Large Language Models by Reasoning on Graphs](https://md.research.elamur.ai/free/en/papers/jin-2024-graph-chain-thought-augmenting-large/)
- [↑ back to papers of 2024](https://md.research.elamur.ai/free/en/papers/2024/page/3/)

## References

25 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

1. Anthropic. Model card and evaluations for claude models, July 2023. URL [https://www.anthropic.com/product](https://www.anthropic.com/product). [↗](https://www.anthropic.com/product)
2. Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. [Qwen technical report](https://md.research.elamur.ai/free/en/papers/bai-2023-qwen-technical-report/). arXiv preprint arXiv:2309.16609, 2023.
3. Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah… [find](https://scholar.google.com/scholar?q=Trenton+Bricken%2C+Adly+Templeton%2C+Joshua+Batson%2C+Brian+Chen%2C+Adam+Jermyn%2C+Tom+Conerly%2C+Nick+Turner%2C+Cem+Anil%2C+Carson+Denison%2C+Amanda+Askell%2C+Robert+Lasenby%2C+Yifan+Wu%2C+Shauna+Kravec%2C+Nicholas+Schiefer%2C+)
4. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. [Training verifiers to solve math word problems](https://md.research.elamur.ai/free/en/papers/cobbe-2021-training-verifiers-solve-math-word/), 2021.
5. Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.19427, 2024. [find](https://scholar.google.com/scholar?q=Soham+De%2C+Samuel+L+Smith%2C+Anushan+Fernando%2C+Aleksandar+Botev%2C+George+Cristian%2DMuraru%2C+Albert+Gu%2C+Ruba+Haroun%2C+Leonard+Berrada%2C+Yutian+Chen%2C+Srivatsan+Srinivasan%2C+et+al%2E+Griffin%3A+Mixing+gated+linear+re)
6. Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. [Data engineering for scaling language models to 128k context](https://md.research.elamur.ai/free/en/papers/fu-2024-data-engineering-scaling-language-models/). arXiv preprint arXiv:2402.10171, 2024.
7. Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. [Model tells you what to discard: Adaptive kv cache compression for llms](https://md.research.elamur.ai/free/en/papers/ge-2023-model-tells-discard-adaptive-kv/). arXiv preprint arXiv:2310.01801, 2023.
8. Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint arXiv:2012.14913, 2020. [find](https://scholar.google.com/scholar?q=Mor+Geva%2C+Roei+Schuster%2C+Jonathan+Berant%2C+and+Omer+Levy%2E+Transformer+feed%2Dforward+layers+are+key%2Dvalue+memories%2E+arXiv+preprint+arXiv%3A2012%2E14913%2C+2020%2E)
9. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. [find](https://scholar.google.com/scholar?q=Albert+Gu+and+Tri+Dao%2E+Mamba%3A+Linear%2Dtime+sequence+modeling+with+selective+state+spaces%2E+arXiv+preprint+arXiv%3A2312%2E00752%2C+2023%2E)
10. Jiatao Gu, Zhengdong Lu, Hang Li, and Victor O.K. Li. [Incorporating copying mechanism in sequence-to-sequence learning](https://md.research.elamur.ai/free/en/papers/gu-2016-incorporating-copying-mechanism-sequence-sequence/). In Katrin Erk and Noah A. Smith, editors, Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1631–1640, Berlin,…
11. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. [Measuring massive multitask language understanding](https://md.research.elamur.ai/free/en/papers/hendrycks-2020-measuring-massive-multitask-language-understanding/). In International Conference on Learning Representations, 2020.
12. Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023. [find](https://scholar.google.com/scholar?q=Albert+Q+Jiang%2C+Alexandre+Sablayrolles%2C+Arthur+Mensch%2C+Chris+Bamford%2C+Devendra+Singh+Chaplot%2C+Diego+de+las+Casas%2C+Florian+Bressand%2C+Gianna+Lengyel%2C+Guillaume+Lample%2C+Lucile+Saulnier%2C+et+al%2E+Mistral+7b)
13. Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024. [find](https://scholar.google.com/scholar?q=Albert+Q+Jiang%2C+Alexandre+Sablayrolles%2C+Antoine+Roux%2C+Arthur+Mensch%2C+Blanche+Savary%2C+Chris+Bamford%2C+Devendra+Singh+Chaplot%2C+Diego+de+las+Casas%2C+Emma+Bou+Hanna%2C+Florian+Bressand%2C+et+al%2E+Mixtral+of+expe)
14. Greg Kamradt. Needle in a haystack - pressure testing llms. gkamradt/LLMTest_NeedleInAHaystack, 2023. [https://github.com/](https://github.com/) [↗](https://github.com/)
15. Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao. Gear: An efficient kv cache compression recipefor near-lossless generative inference of llm. arXiv preprint arXiv:2403.05527, 2024. — [arXiv:2403.05527](https://arxiv.org/abs/2403.05527) [find](https://scholar.google.com/scholar?q=Hao+Kang%2C+Qingru+Zhang%2C+Souvik+Kundu%2C+Geonhwa+Jeong%2C+Zaoxing+Liu%2C+Tushar+Krishna%2C+and+Tuo+Zhao%2E+Gear%3A+An+efficient+kv+cache+compression+recipefor+near%2Dlossless+generative+inference+of+llm%2E+arXiv+prepr)
16. Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05055, 2022. [find](https://scholar.google.com/scholar?q=Aran+Komatsuzaki%2C+Joan+Puigcerver%2C+James+Lee%2DThorp%2C+Carlos+Riquelme+Ruiz%2C+Basil+Mustafa%2C+Joshua+Ainslie%2C+Yi+Tay%2C+Mostafa+Dehghani%2C+and+Neil+Houlsby%2E+Sparse+upcycling%3A+Training+mixture%2Dof%2Dexperts+from+)
17. Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. In search of needles in a 10m haystack: Recurrent memory finds what llms miss. arXiv preprint arXiv:2402.10790, 2024. — [arXiv:2402.10790](https://arxiv.org/abs/2402.10790) [find](https://scholar.google.com/scholar?q=Yuri+Kuratov%2C+Aydar+Bulatov%2C+Petr+Anokhin%2C+Dmitry+Sorokin%2C+Artyom+Sorokin%2C+and+Mikhail+Burtsev%2E+In+search+of+needles+in+a+10m+haystack%3A+Recurrent+memory+finds+what+llms+miss%2E+arXiv+preprint+arXiv%3A2402)
18. Mistral. Model card for mistral-7b-instruct-v0.2, April 2024. URL [https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2). [↗](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2)
19. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022. [find](https://scholar.google.com/scholar?q=Catherine+Olsson%2C+Nelson+Elhage%2C+Neel+Nanda%2C+Nicholas+Joseph%2C+Nova+DasSarma%2C+Tom+Henighan%2C+Ben+Mann%2C+Amanda+Askell%2C+Yuntao+Bai%2C+Anna+Chen%2C+et+al%2E+In%2Dcontext+learning+and+induction+heads%2E+arXiv+preprin)
20. Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://md.research.elamur.ai/free/en/papers/gemini-team-2024-gemini-1-5-unlocking-multimodal/). arXiv preprint arXiv:2403.05530,…
21. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. [find](https://scholar.google.com/scholar?q=Hugo+Touvron%2C+Louis+Martin%2C+Kevin+Stone%2C+Peter+Albert%2C+Amjad+Almahairi%2C+Yasmine+Babaei%2C+Nikolay+Bashlykov%2C+Soumya+Batra%2C+Prajjwal+Bhargava%2C+Shruti+Bhosale%2C+et+al%2E+Llama+2%3A+Open+foundation+and+fine%2Dtun)
22. Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020. [find](https://scholar.google.com/scholar?q=Sinong+Wang%2C+Belinda+Z+Li%2C+Madian+Khabsa%2C+Han+Fang%2C+and+Hao+Ma%2E+Linformer%3A+Self%2Dattention+with+linear+complexity%2E+arXiv+preprint+arXiv%3A2006%2E04768%2C+2020%2E)
23. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. [Chain-of-thought prompting elicits reasoning in large language models](https://md.research.elamur.ai/free/en/papers/wei-2022-chain-thought-prompting-elicits-reasoning/). In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Info…
24. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023. [find](https://scholar.google.com/scholar?q=Guangxuan+Xiao%2C+Yuandong+Tian%2C+Beidi+Chen%2C+Song+Han%2C+and+Mike+Lewis%2E+Efficient+streaming+language+models+with+attention+sinks%2E+arXiv+preprint+arXiv%3A2309%2E17453%2C+2023%2E)
25. Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, et al. Yi: Open foundation models by 01. ai. arXiv preprint arXiv:2403.[04652, 2024](https://md.research.elamur.ai/free/en/papers/ai-2024-yi-open-foundation-models-01/).

*Updated 24 October 2026*