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

TL;DR

The paper investigates whether mechanistically identified retrieval heads can be leveraged to enhance long-context processing in large language models. It introduces RetMask, a retrieval-head masking approach that generates contrastive supervision by comparing outputs from the full model and an ablated variant, trained with Direct Preference Optimization. Results show strong gains for models with concentrated retrieval heads (notably Llama-3.1), modest improvements for models with distributed retrieval (Olmo-3), and robust improvements on long-context tasks such as generation with citations and passage re-ranking, while preserving general capabilities. The work provides mechanistic insights into when retrieval-head optimization is effective and demonstrates how interpretability findings can translate into practical performance gains, offering a diagnostic for model developers and a path toward better long-context LLMs.

Abstract

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 three model families reveal that the effectiveness depends on retrieval head organization: models with concentrated patterns of retrieval heads respond strongly, while those with distributed patterns show limited gains. This mechanistic relationship validates the function of retrieval heads and demonstrates that mechanistic insights can be transformed into performance enhancements.

Figures (5)

• Figure 1: An overview of this study. Given an LLM, we construct an ablated version by masking the retrieval heads wu2025retrieval. We then sample texts from both the normal and ablated models and conduct Direct Preference Optimization. The effectiveness of such a training depends on how the retrieval heads are organized within LLMs.
• Figure 2: Overview of Preliminaries (left) and RetMask (right). The example on the right is a real case extracted from the training data. We detect and mask retrieval heads for generating contrastive responses.
• Figure 3: Qwen3-8B's improvement on each task of HELMET when the input sequence length is 128K. RC stands for reasoning contents. Removing reasoning contents has a limited impact on RetMask.
• Figure 4: The retrieval score distribution of LLMs tested in this work. While attention heads in Llama-3.1-8B-Instruct and Qwen3-8B exhibit a concentrated pattern of retrieval capability, it is more distributed for Olmo-3-7B.
• Figure 5: The distribution of retrieval score before and after RetMask. We observe an increase in retrieval scores for both Llama-3.1-8B-Instruct and Qwen3-8B.