Elamur Research⌕
☰
home / papers / 2024 / Retrieval Head Mechanistically Explains Long-Context Factuality

Retrieval Head Mechanistically Explains Long-Context Factuality

Paper · 2024 · Wenhao Wu, Yizhong Wang, Guangxuan Xiao et al.

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 WashingtonSeattle · United States263 works · 180,581 citations🎓University of EdinburghEdinburgh · United Kingdom255 works · 68,771 citations🎓Massachusetts Institute of TechnologyCambridge · United States159 works · 65,157 citations🎓Peking UniversityBeijing · China100 works · 35,618 citations

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 9,086 worksArtificial Intelligence · 8,889 worksArts and Humanities · 9,296 worksLanguage and Linguistics · 9,206 worksComputer Vision and Pattern Recognition · 1,282 works

🧭 Topics 1Research subjects and languages the work is about.Information retrieval & search · 145 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. DOI 10.48550/arxiv.2404.15574 · arXiv 2404.15574

@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}
}

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. ↗
  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. 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
  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, 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
  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. 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. 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
  9. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. find
  10. Jiatao Gu, Zhengdong Lu, Hang Li, and Victor O.K. Li. Incorporating copying mechanism in sequence-to-sequence learning. 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. 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
  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
  14. Greg Kamradt. Needle in a haystack - pressure testing llms. gkamradt/LLMTest_NeedleInAHaystack, 2023. 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
  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
  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
  18. Mistral. Model card for mistral-7b-instruct-v0.2, April 2024. URL 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
  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. 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
  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
  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. 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
  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.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wu-2024-retrieval-head-mechanistically-explains-long/