Elamur Research⌕
☰
home / papers / 2026 / From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models

From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models

Paper · 2026 · Youmi Ma and Naoaki Okazaki

Summary

Mechanistic interpretability work has found "retrieval heads": a small set of attention heads that copy information from earlier in the context. Until now, nobody had shown that this finding could make models better. The authors propose RetMask, which turns the finding into a training signal. First, they detect retrieval heads with a Needle-In-A-Haystack scoring procedure. Next, they build an ablated copy of the model by zeroing the output-projection columns of those heads. Both the original and the ablated model then answer prompts from existing instruction datasets. The original model's answer is treated as preferred and the ablated model's answer as rejected, and the original model is trained on these pairs with Direct Preference Optimization. No human labels or LLM judge are needed.

Across four models from three families, RetMask consistently improves HELMET long-context scores at 8K to 128K tokens, while math, coding and general-knowledge performance stays intact. Gains are larger when retrieval ability is concentrated in fewer heads. The work shows that interpretability insights can directly guide model improvement.

Key points

  • Models: Llama-3.1-8B-Instruct, Qwen3-8B and Olmo-3-7B (Instruct and Think).
    • Retrieval-head thresholds were τ = 0.1 for Llama-3.1 and τ = 0.05 for the others.
    • This gave 40 masked heads in Llama-3.1 and 79 in Qwen3.
  • Main result on Llama-3.1 at 128K: the HELMET average rose from 46.40 to 48.68 (+2.28).
    • Generation with citations improved by about 70% (3.09 → 5.25).
    • Passage re-ranking improved by about 32% (13.73 → 18.16).
  • Baselines: RetMask beat several alternative DPO setups, including smaller-model contrast, random masking and masking non-retrieval heads.
  • Training prompts: drawn from LMSYS-Chat-1M, WildChat and Guru-RL-92K.
    • Training mostly uses self-synthesis, where the model generates its own preference pairs.
    • Code is available at github.com/YoumiMa/RetMask.
  • Effect on retrieval scores: after training, Llama-3.1's mean retrieval score grew from 0.017 to 0.020.
    • The masked heads gained about 0.051 on average, while the other heads barely changed.
    • Qwen3, whose retrieval scores are less sparse, improved only modestly (+5%).

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

🎓Institute of Science TokyoTokyo · Japan10 works · 20 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksInformation Systems · 1,437 works

🧭 Topics 3Research subjects and languages the work is about.Language models & pretraining · 3,320 worksInformation retrieval & search · 473 worksInterpretability · 130 works

Cite

Cite as Youmi Ma and Naoaki Okazaki. 2026. From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.1380

@inproceedings{ma2026interpretability,
  title = {{From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models}},
  author = {Youmi Ma and Naoaki Okazaki},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.1380},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.1380}
}

References

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

  1. Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi, Lei Hou, Jie Tang, Yuxiao Dong, and Juanzi Li. 2024a. LongAlign: A recipe for long context alignment of large language models. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 1376–1395, Miami, Florida, USA.
  2. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. 2024b. LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Associat…
  3. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter,… — Language Models are Few-Shot Learners
  4. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, and 39 others. 202… — Evaluating Large Language Models Trained on Code
  5. Zhoujun Cheng, Shibo Hao, Tianyang Liu, Fan Zhou, Yutao Xie, Feng Yao, Yuexin Bian, Nilabjo Dey, Yonghao Zhuang, Yuheng Zha, Yi Gu, Kun Zhou, Yuqi Wang, Yuan Li, Richard Fan, Jianshu She, Chengqian Gao, Abulhair Saparov, Taylor W. Killian, and 5 others. 2025. Revisiting reinforcement learning for LL… find
  6. Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. 2022. Knowledge neurons in pretrained transformers. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (ACL), pages 8493–8502. find
  7. Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. 2025. How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (ACL), pages 7376–7399.
  8. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, and 542 o… — The Llama 3 Herd of Models
  9. Jia-Chen Gu, Hao-Xiang Xu, Jun-Yu Ma, Pan Lu, Zhen-Hua Ling, Kai-Wei Chang, and Nanyun Peng. 2024. Model editing harms general abilities of large language models: Regularization to the rescue. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 16…
  10. Tatsuya Hiraoka and Kentaro Inui. 2025. Repetition neurons: How do language models produce repetitions? In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL), pages 483–495. find
  11. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. 2024. RULER: What’s the real context size of your long-context language models? In First Conference on Language Modeling.
  12. Greg Kamradt. 2023. Needle in a haystack - pressure testing LLMs. find
  13. Takeshi Kojima, Itsuki Okimura, Yusuke Iwasawa, Hitomi Yanaka, and Yutaka Matsuo. 2024. On the multilingual ability of decoder-based pre-trained language models: Finding and controlling language-specific neurons. In Proceedings of the 2024 Conference of the North American Chapter of the Association…
  14. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. 2023. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles.
  15. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023. Let’s verify step by step. arXiv preprint arXiv:2305.20050.
  16. Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR). find
  17. Kevin Meng, David Bau, Alex J Andonian, and Yonatan Belinkov. 2022. Locating and editing factual associations in GPT. In Advances in Neural Information Processing Systems (NeurIPS). find
  18. Soumen Kumar Mondal, Sayambhu Sen, Abhishek Singhania, and Preethi Jyothi. 2025. Language-specific neurons do not facilitate cross-lingual transfer. In The Sixth Workshop on Insights from Negative Results in NLP, pages 46–62. find
  19. Team Olmo, Allyson Ettinger, Amanda Bertsch, Bailey Kuehl, David Graham, David Heineman, Dirk Groeneveld, Faeze Brahman, Finbarr Timbers, Hamish Ivison, Jacob Morrison, Jake Poznanski, Kyle Lo, Luca Soldaini, Matt Jordan, Mayee Chen, Michael Noukhovitch, Nathan Lambert, Pete Walsh, and 49 others. 20… — Olmo 3
  20. OpenAI. 2024. Learning to reason with LLMs. find
  21. Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. 2023. Direct preference optimization: Your language model is secretly a reward model. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS).
  22. David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. 2024. GPQA: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling (COLM).
  23. Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. 2024. Scaling LLM test-time compute optimally can be more effective than scaling model parameters. Preprint, arXiv:2408.03314.
  24. Tianyi Tang, Wenyang Luo, Haoyang Huang, Dongdong Zhang, Xiaolei Wang, Xin Zhao, Furu Wei, and Ji-Rong Wen. 2024. Language-specific neurons: The key to multilingual capabilities in large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (ACL)…
  25. Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, Louis Rouillard, Thomas Mesnard, Geoffrey Cideron, Jean bastien Grill, Sabela Ramos, Edouard Yvinec, Michelle Casbon, Etienne Pot, Ivo Penche… find
  26. Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, Tianle Li, Max Ku, Kai Wang, Alex Zhuang, Rongqi Fan, Xiang Yue, and Wenhu Chen. 2024. MMLU-pro: A more robust and challenging multi-task language understanding benchmark…
  27. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier-ric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, and 3 others. 2020. Transf…
  28. Longyun Wu, Dawei Zhu, Guangxiang Zhao, Zhuocheng Yu, Junfeng Ran, Xiangyu Wong, Lin Sun, and Sujian Li. 2025a. LongAttn: Selecting long-context training data via token-level attention. In Findings of the Association for Computational Linguistics: ACL 2025, pages 19367–19380. find
  29. Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. 2025b. Retrieval head mechanistically explains long-context factuality. In The Thirteenth International Conference on Learning Representations (ICLR).
  30. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, and 41 others. 2025. Qwen3 technical report. Preprint, arXiv:2505.09388.
  31. Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Izsak, Moshe Wasserblat, and Danqi Chen. 2025. HELMET: How to evaluate long-context models effectively and thoroughly. In The Thirteenth International Conference on Learning Representations (ICLR). find
  32. Jiajie Zhang, Zhongni Hou, Xin Lv, Shulin Cao, Zhenyu Hou, Yilin Niu, Lei Hou, Yuxiao Dong, Ling Feng, and Juanzi Li. 2025a. LongReward: Improving long-context large language models with AI feedback. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (ACL), pa…
  33. Wuwei Zhang, Fangcong Yin, Howard Yen, Danqi Chen, and Xi Ye. 2025b. Query-focused retrieval heads improve long-context reasoning and re-ranking. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 23802–23816. find
  34. Wenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie, Yejin Choi, and Yuntian Deng. 2024. Wildchat: 1m chatGPT interaction logs in the wild. In The Twelfth International Conference on Learning Representations (ICLR).
  35. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Tianle Li, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zhuohan Li, Zi Lin, Eric Xing, Joseph E. Gonzalez, Ion Stoica, and Hao Zhang. 2024. LMSYS-chat-1m: A large-scale real-world LLM conversation dataset. In The Twelfth International Conference on Learning…
  36. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. 2023. Judging LLM-as-a-judge with MT-bench and chatbot arena. In Thirty-seventh Conference on Neural Information Processing…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ma-2026-interpretability-performance-optimizing-retrieval-heads/