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

arXiv: https://arxiv.org/abs/2601.11020 · HTML: https://mlsanity.com/paper/2601.11020/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge and representations inside LLMs, one of 13 directions in Mechanistic interpretability.
- Builds on 24 of its 40 references that are in the map.
- First posted on arXiv 16 Jan 2026.

## Placed in

- Mechanistic interpretability / [Knowledge and representations inside LLMs](https://mlsanity.com/mechinterp/knowledge-and-representations-inside-llms/)

## Builds on (most cited first, 24 of 24)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [WildChat: 1M ChatGPT Interaction Logs in the Wild](https://mlsanity.com/paper/2405.01470/) (2024)
- [LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset](https://mlsanity.com/paper/2309.11998/) (2023)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)
- [How to Train Long-Context Language Models (Effectively)](https://mlsanity.com/paper/2410.02660/) (2024)
- [Model Editing Harms General Abilities of Large Language Models: Regularization to the Rescue](https://mlsanity.com/paper/2401.04700/) (2024)
- [On the Multilingual Ability of Decoder-based Pre-trained Language Models: Finding and Controlling Language-Specific Neurons](https://mlsanity.com/paper/2404.02431/) (2024)
- [LongReward: Improving Long-context Large Language Models with AI Feedback](https://mlsanity.com/paper/2410.21252/) (2024)
- [Query-Focused Retrieval Heads Improve Long-Context Reasoning and Re-ranking](https://mlsanity.com/paper/2506.09944/) (2025)
- [Repetition Neurons: How Do Language Models Produce Repetitions?](https://mlsanity.com/paper/2410.13497/) (2024)
- [Improving Contextual Faithfulness of Large Language Models via Retrieval Heads-Induced Optimization](https://mlsanity.com/paper/2501.13573/) (2025)
- [Language-specific Neurons Do Not Facilitate Cross-Lingual Transfer](https://mlsanity.com/paper/2503.17456/) (2025)
- [LongAttn: Selecting Long-context Training Data via Token-level Attention](https://mlsanity.com/paper/2502.16860/) (2025)

## Built on by (most cited first, 4 of 4)

- [Retrieval Heads Meet Vision: Uncovering How VLMs Locate and Extract Visual Information](https://mlsanity.com/paper/2608.27417/) (2026)
- [Neuron Level Analysis of Large Language Model in Legal Domain Reasoning](https://mlsanity.com/paper/2606.15884/) (2026)
- [Does RoPE Prevent or Degrade Retrieval Heads? A Mechanistic Analysis Across Model Families](https://mlsanity.com/paper/2606.21249/) (2026)
- [Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads](https://mlsanity.com/paper/2607.01002/) (2026)

## Abstract (arXiv)

Advances in mechanistic interpretability have identified special attention heads, known as retrieval heads, that are responsible for retrieving information from the context. However, the role of these retrieval heads in improving model performance remains unexplored. This work investigates whether retrieval heads can be leveraged to enhance the long-context capabilities of LLMs. Specifically, we propose RetMask, a method that generates training signals by contrasting normal model outputs with those from an ablated variant in which the retrieval heads are masked. This mechanism-based approach achieves substantial improvements: +2.28 points on HELMET at 128K for Llama-3.1, with +70% gains on generation with citation and +32% on passage re-ranking, while preserving performance on general tasks. Experiments across four models in three families demonstrate that RetMask consistently improves long-context performance, where gains correlate with the sparsity of the retrieval score distribution: models with sparser distributions, where retrieval capabilities are concentrated in a small set of heads, respond more strongly, while those with less sparse distributions show more modest gains. These results validate the functional role of retrieval heads and show that mechanistic insights can be transformed into performance enhancements.