Technical Summary: Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads

1. Problem Statement

Large Language Models (LLMs) frequently synthesize answers from the semantic meaning of a relevant context span rather than literally copying tokens. While existing mechanistic interpretability work has identified "retrieval heads" responsible for context access, current detection methods are fundamentally limited to literal retrieval.

Existing detectors (e.g., Wu et al., 2025; Fu et al., 2025) rely on attention-weight observables. They reward heads where the attended source token matches the generated answer token. This "token-matching" criterion successfully identifies heads that copy-paste information (literal retrieval) but fails to detect heads that perform non-literal retrieval. In non-literal scenarios, a model may attend to a source phrase (e.g., "Eiffel Tower") but synthesize a different answer token (e.g., "Yuki") via its output-value (OV) circuit. Attention-based scoring captures where a head reads but ignores what it writes, missing the critical mechanism by which non-literal synthesis occurs.

2. Methodology: Logit-Contribution Scoring (LOCOS)

The authors introduce Logit-Contribution Scoring (LOCOS), a "write-aware" detector that evaluates attention heads based on their contribution to the final answer logit, rather than their attention distribution.

Core Mechanism

LOCOS decomposes an attention head's operation into two circuits:

1. QK Circuit: Determines where the head reads (attention weights $\alpha$).
2. OV Circuit: Determines what the head writes to the residual stream (output projection $W_O v$).

The method scores a head $(l, h)$ by projecting its OV-circuit output onto the unembedding direction of the correct answer token ($u_{y_t}$).

The Three-Step Procedure

1. Per-Position Logit Contribution:
   For a decoding step $t$ generating token $y_t$, the contribution of a source position $j$ through head $(l, h)$ is calculated as:
   $$\phi^{(l,h)}_{t,j} = u_{y_t}^\top o^{(l,h)}_{t,j} = \alpha^{(l,h)}_{t,j} \cdot u_{y_t}^\top W^{(l,h)}_O v^{(l,h)}_{t,j}$$
   This scalar $\phi$ captures the full pipeline: a head can have high attention ($\alpha$) but near-zero contribution if its OV output is orthogonal to the answer direction. Conversely, a head with moderate attention can have high contribution if its OV circuit transforms the attended representation into an answer-aligned vector.

2. Spatial Contrast:
   To isolate retrieval-specific signals, LOCOS compares contributions from needle positions (the relevant context span $[s_\tau, e_\tau)$) against off-needle positions.
   $$\Phi^{(l,h),+}_t = \sum_{j \in \text{needle}} \phi^{(l,h)}_{t,j}$$
   $$\Phi^{(l,h),-}_t = \frac{e_\tau - s_\tau}{N_t - (e_\tau - s_\tau)} \sum_{j \notin \text{needle}} \phi^{(l,h)}_{t,j}$$
   The rescaling factor ensures the off-needle sum is comparable to the needle sum by length. The score for a single step is the difference: $\Phi^{(l,h),+}_t - \Phi^{(l,h),-}_t$. This spatial contrast (needle vs. off-needle within a single step) replaces the temporal contrast (answer vs. non-answer steps) used in prior work.

3. Aggregation:
   The final score $S_{l,h}$ is the mean of these spatial contrasts across all decoding steps where the model generates a correct answer, aggregated over all successful trials in a benchmark. Negative scores indicate heads whose contributions to the answer logit originate predominantly from off-needle positions.

3. Key Contributions

• Write-Aware Detection: The paper establishes that identifying non-literal retrieval requires measuring the OV circuit's output, not just attention weights.
• LOCOS Algorithm: A single-forward-pass method that scores heads by their projected contribution to the answer logit, contrasting needle and off-needle sources.
• Causal Validation: Through mean-ablation experiments, the authors demonstrate that LOCOS identifies heads that are causally critical for non-literal retrieval, whereas attention-based baselines fail to do so.

4. Experimental Results

The method was evaluated on the NoLiMa benchmark (a non-literal needle-in-a-haystack task) across three model families: Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B), and OLMo-3.1 (32B).

Ablation Performance (NoLiMa)

• Superior Degradation: Mean-ablating the top-$k$ LOCOS heads causes a steeper decline in ROUGE-L compared to all attention-based baselines (Wu/NIAH-scored, Wu/NoLiMa-scored, and random heads).
• Qwen3-8B Case Study: Ablating the top 50 LOCOS heads collapses ROUGE-L from 0.401 to 0.000. In contrast, the strongest attention-based baseline (Wu/NIAH) retains a ROUGE-L of 0.292 at the same depth.
• Specificity: Ablating LOCOS heads severely degrades contextual retrieval while leaving parametric recall (e.g., city-country associations) and arithmetic reasoning intact. This confirms the identified heads are retrieval-specific, not generically important for model output.

Literal vs. Non-Literal Specificity

• Ablating top-$k$ LOCOS heads degrades both non-literal (NoLiMa) and literal (NIAH) retrieval.
• However, the drop is significantly steeper on NoLiMa. This indicates that while LOCOS captures heads relevant to both regimes, it uniquely identifies the subset of heads responsible for non-literal synthesis that token-matching detectors miss.

Downstream Transfer

• On downstream long-context benchmarks (MuSiQue and BABILong), ablating top-50 LOCOS heads significantly degrades performance, particularly on the Qwen3 family (e.g., MuSiQue accuracy on Qwen3-8B drops from 0.55 to 0.08).
• The ranking of LOCOS vs. attention-based baselines varies by benchmark, suggesting that multi-hop QA tasks recruit both literal and non-literal retrieval circuits.

Layer Distribution

• LOCOS scores are heavily concentrated in late layers for the Qwen3 family and Gemma-3-27B.
• The authors note this concentration may be influenced by the "direct-path assumption" (heads closer to the output have a more direct impact on logits), but causal patching experiments suggest the late-layer pattern is not merely an artifact.

5. Significance and Claims

The paper claims that LOCOS provides the first mechanistic identification of non-literal retrieval heads.

• Limitation of Prior Work: Existing detectors are "blind" to non-literal retrieval because they reward heads where the attended token matches the generated token. This misses heads that attend to a source and synthesize a different answer.
• Mechanistic Insight: By measuring the projection of the OV circuit onto the answer direction, LOCOS captures the "write" operation essential for synthesis.
• Practical Utility: The identified heads are causally critical for long-context reasoning. This finding has implications for:
  • Interpretability: Providing a more accurate map of how LLMs access and synthesize context.
  • Efficiency: Enabling more targeted KV cache compression by preserving only the specific heads responsible for retrieval, rather than uniform or attention-based heuristics.
  • Hallucination Mitigation: Offering a mechanism to identify and intervene on heads that fail to retrieve relevant context correctly.

The authors remain modest regarding generalization, noting that their findings are specific to the tested architectures (standard decoder-only transformers) and the NoLiMa benchmark. They do not claim that LOCOS identifies all retrieval mechanisms, but rather that it successfully isolates the subset missed by token-matching criteria.