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

• URL: http://arxiv.org/abs/2601.11020v2
• Date: Thu, 22 Jan 2026 11:02:26 GMT
• Title: From Interpretability to Performance: Optimizing Retrieval Heads for Long-Context Language Models
• Authors: Youmi Ma, Naoaki Okazaki,
• Abstract summary: This work investigates whether retrieval heads can be leveraged to enhance the long-context capabilities of LLMs.<n>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.<n> Experiments across three model families reveal that the effectiveness depends on retrieval head organization.
• Score: 19.62954865335739
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Explicit Multi-head Attention for Inter-head Interaction in Large Language Models [70.96854312026319]
  Multi-head Explicit Attention (MEA) is a simple yet effective attention variant that explicitly models cross-head interaction.<n>MEA shows strong robustness in pretraining, which allows the use of larger learning rates that lead to faster convergence.<n>This enables a practical key-value cache compression strategy that reduces KV-cache memory usage by 50% with negligible performance loss.
  arXiv  Detail & Related papers  (2026-01-27T13:45:03Z)
• Retrieval Heads are Dynamic [101.60087217027949]
  Recent studies have identified "retrieval heads" in Large Language Models (LLMs)<n>In this paper, we investigate retrieval heads from a dynamic perspective.
  arXiv  Detail & Related papers  (2026-01-07T02:29:24Z)
• Reinforcement Learning Improves Traversal of Hierarchical Knowledge in LLMs [7.424730923663806]
  We show that RL-enhanced models consistently outperform their base and supervised fine-tuned (SFT) counterparts on pure knowledge recall tasks.<n>We hypothesize these gains stem not from newly acquired data, but from improved procedural skills in navigating and searching existing knowledge hierarchies within the model parameters.
  arXiv  Detail & Related papers  (2025-11-08T08:56:29Z)
• Learning Compact Representations of LLM Abilities via Item Response Theory [35.74367665390977]
  We explore how to learn compact representations of large language models (LLMs)<n>We frame this problem as estimating the probability that a given model will correctly answer a specific query.<n>To learn these parameters jointly, we introduce a Mixture-of-Experts (MoE) network that couples model- and query-level embeddings.
  arXiv  Detail & Related papers  (2025-10-01T12:55:34Z)
• Structured Thinking Matters: Improving LLMs Generalization in Causal Inference Tasks [0.7988085110283119]
  Recent results from the Corr2Cause dataset benchmark reveal that state-of-the-art LLMs only marginally outperform random baselines.<n>We provide the model with the capability to structure its thinking by guiding the model to build a structured knowledge graph.<n> Experiments on the test subset of the Corr2Cause dataset benchmark with Qwen3-32B model (reasoning model) show substantial gains over standard direct prompting methods.
  arXiv  Detail & Related papers  (2025-05-23T15:37:40Z)
• Exploring Training and Inference Scaling Laws in Generative Retrieval [50.82554729023865]
  Generative retrieval reformulates retrieval as an autoregressive generation task, where large language models generate target documents directly from a query.<n>We systematically investigate training and inference scaling laws in generative retrieval, exploring how model size, training data scale, and inference-time compute jointly influence performance.
  arXiv  Detail & Related papers  (2025-03-24T17:59:03Z)
• DeepRAG: Thinking to Retrieve Step by Step for Large Language Models [92.87532210660456]
  We propose DeepRAG, a framework that models retrieval-augmented reasoning as a Markov Decision Process (MDP)<n>By iteratively decomposing queries, DeepRAG dynamically determines whether to retrieve external knowledge or rely on parametric reasoning at each step.<n> Experiments show that DeepRAG improves retrieval efficiency and boosts answer accuracy by 26.4%, demonstrating its effectiveness in enhancing retrieval-augmented reasoning.
  arXiv  Detail & Related papers  (2025-02-03T08:22:45Z)
• SEAL: Scaling to Emphasize Attention for Long-Context Retrieval [8.805524738976075]
  We introduce a novel approach called Scaling to Emphasize Attention for Long-context retrieval (SEAL)<n>We observe that specific attention heads are closely tied to long-context retrieval, showing positive or negative correlation with retrieval scores.<n>We propose a learning-based mechanism that leverages generated data to emphasize these heads.
  arXiv  Detail & Related papers  (2025-01-25T14:09:39Z)
• Learning from Neighbors: Category Extrapolation for Long-Tail Learning [62.30734737735273]
  We offer a novel perspective on long-tail learning, inspired by an observation: datasets with finer granularity tend to be less affected by data imbalance.<n>We introduce open-set auxiliary classes that are visually similar to existing ones, aiming to enhance representation learning for both head and tail classes.<n>To prevent the overwhelming presence of auxiliary classes from disrupting training, we introduce a neighbor-silencing loss.
  arXiv  Detail & Related papers  (2024-10-21T13:06:21Z)
• Prospector Heads: Generalized Feature Attribution for Large Models & Data [82.02696069543454]
  We introduce prospector heads, an efficient and interpretable alternative to explanation-based attribution methods. We demonstrate how prospector heads enable improved interpretation and discovery of class-specific patterns in input data.
  arXiv  Detail & Related papers  (2024-02-18T23:01:28Z)
• Augmenting Unsupervised Reinforcement Learning with Self-Reference [63.68018737038331]
  Humans possess the ability to draw on past experiences explicitly when learning new tasks. We propose the Self-Reference (SR) approach, an add-on module explicitly designed to leverage historical information. Our approach achieves state-of-the-art results in terms of Interquartile Mean (IQM) performance and Optimality Gap reduction on the Unsupervised Reinforcement Learning Benchmark.
  arXiv  Detail & Related papers  (2023-11-16T09:07:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.