PaperlayerView exact revision on arXiv

Source-linked AI summary

Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads

Aryo Pradipta Gema, Beatrice Alex, Pasquale Minervini

arXiv:2607.01002v1cs.CLcs.AIcs.LG

TL;DR

Existing retrieval-head detectors focus on where attention reads and therefore miss non-literal retrieval implemented through what heads write. LOCOS scores answer-directed OV-circuit contributions with needle-versus-off-needle contrast, and its top-ranked heads cause stronger, retrieval-specific degradation across long-context evaluations.

  • Problem

    Existing detectors use literal-copy criteria based on attention patterns, missing heads that synthesize answers from attended meaning through their OV circuits.

  • Method

    LOCOS projects each head’s weighted OV output onto the correct answer-token unembedding direction and contrasts needle with off-needle positions in one forward pass.

  • Results

    Across three model families, LOCOS-selected head ablations collapse NoLiMa ROUGE-L at lower head counts than attention-based baselines; on Qwen3-8B, 50 heads reduce ROUGE-L from 0.401 to 0.000 versus 0.292 for the strongest baseline.

  • Takeaways & Limitations

    The selected heads are retrieval-specific: the same ablation preserves parametric recall and arithmetic reasoning while degrading long-context retrieval performance.

  • Takeaways & Limitations

    LOCOS may miss diffuse contextual integration when distractors are answer-related, and its reported ablation magnitudes are not established for mixture-of-experts, encoder–decoder, or state-space-hybrid architectures.

Abstract

from arXiv · show

In long-context use, large language models frequently synthesize answers from the meaning of a relevant context span rather than literally copy-pasting them. Identifying which attention heads perform this synthesis matters for interpreting long-context model behavior. Yet existing detectors miss these heads by construction: they reward heads whose attended token matches the generated token, a literal-copy criterion that captures where a head reads but not what it writes through its output-value (OV) circuit, the very mechanism that carries non-literal retrieval. We introduce Logit-Contribution Scoring (LOCOS), a write-aware detector that scores each head by the projection of its OV-circuit output onto the answer-token unembedding direction, contrasting needle and off-needle source positions in a single forward pass. Across three model families (Qwen3, Gemma-3, OLMo-3.1), mean-ablating the top LOCOS heads on the NoLiMa non-literal retrieval benchmark collapses ROUGE-L at lower head counts than prior attention-based detections; on Qwen3-8B, ablating 50 heads drives ROUGE-L from 0.401 to 0.000 while the strongest baseline still retains 0.292. The selected heads are retrieval-specific: parametric recall and arithmetic reasoning stay at baseline under the same ablation. On Qwen3-8B, the same ablation also drops MuSiQue from 0.55 to 0.08 and BABI-Long from 0.62 to 0.20, while a random-heads control stays within 0.05 of baseline.

1 Introduction

Non-literal retrieval requires models to synthesize answers from relevant context rather than copy matching tokens. LOCOS addresses this gap by measuring what attention heads write through their OV circuits, and its selected heads produce stronger retrieval-specific ablation effects.

  • Motivation: Existing detectors observe attention patterns, capturing where heads read but not what their OV circuits propagate.This misses non-literal retrieval, where the answer may not lexically overlap with the attended span.
  • Method: LOCOS measures each head’s contribution to the correct answer token by projecting its weighted OV output onto the answer unembedding direction.The method contrasts needle and off-needle positions within a single forward pass.
  • Motivation: Figure 1 contrasts literal retrieval of “Eiffel Tower” with non-literal synthesis of “Yuki” from the same context.The example illustrates why lexical copying is insufficient for non-literal retrieval.
  • Results: Ablating the top LOCOS heads on NoLiMa produces steeper ROUGE-L degradation than evaluated baselines across six configurations and three model families.On Qwen3-8B, top-50 ablation reduces ROUGE-L from 0.401 to 0.000, while the strongest attention-based baseline retains 0.292.

2 Background

Literal retrieval heads can copy answer tokens, but NoLiMa requires recovering answers from meaning without lexical overlap. Attention heads therefore need both correct source selection and answer-aligned writing through their OV circuits.

  • Read-and-write circuits: An attention head’s QK circuit determines where it reads, while its OV circuit maps source values into what it writes.A retrieval head requires both stages to select relevant positions and produce an answer-aligned update.
  • Prior detection: Existing retrieval-head detection rewards attention to the needle and matching between the attended token and generated token.This is a literal-copy criterion inherited from literal needle-in-a-haystack tasks.
  • Literal versus non-literal retrieval: Literal NIAH places an answer-bearing needle in distractor context, whereas NoLiMa requires semantic recovery without lexical overlap.This breaks the equivalence between attending to the needle and copying the answer token.
  • Read-and-write circuits: Figure 2 shows that high attention can still yield near-zero answer contribution, whereas moderate attention can yield a large contribution when the OV output aligns with the answer.This is why attention-only methods can select the wrong head for non-literal retrieval.
  • Non-literal retrieval: Non-literal retrieval can attend to a semantically relevant phrase and write an answer-aligned direction even when no attended token matches the answer.For example, attending to “Paris” can support producing “France.”

3 Logit-Contribution Scoring

LOCOS detects retrieval heads by scoring answer-directed OV writes rather than attention alone. It uses needle-versus-off-needle spatial contrast and aggregates these contributions across correct answer steps.

  • Core idea: LOCOS scores each head by what it writes toward the answer instead of where it allocates attention.The design targets non-literal retrieval because attended content is transformed through the OV circuit before becoming an answer.
  • Per-position contribution: Per-position logit contribution measures the weighted OV output’s projection onto the correct answer-token unembedding direction.A head receives little contribution when its OV output is orthogonal to the answer, even under strong attention.
  • Spatial contrast: Spatial contrast compares needle and length-normalized off-needle contributions within one decoding step.It isolates position-specific writes and can cancel uniform contributors such as position-independent token-frequency priors.
  • Aggregation: The final score averages contributions over trial and answer-step pairs, weighting each answer step equally.Aggregation pools answer steps from trials that pass the correctness filter without per-trial normalization.

4 Experiments

Across six model configurations, LOCOS-ranked head ablations most strongly disrupt NoLiMa non-literal retrieval, while controls show that the effect is selective and retrieval-focused. The experiments further examine OV-specific scoring, layer distribution, retrieval specificity, literal retrieval, and downstream transfer.

  • Ablation Comparison Across Scoring Methods: Across six model configurations, top-ranked LOCOS heads produce steeper NoLiMa ROUGE-L degradation than evaluated baselines.The benchmark spans Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B), and OLMo-3.1 (32B).
  • Ablation Comparison Across Scoring Methods: 0.000 ROUGE-L remains on Qwen3-8B after ablating 50 LOCOS heads, versus 0.292 for the strongest attention-based baseline.At k=5, LOCOS already reaches 0.321, while Wu/NIAH remains at 0.406.
  • Isolating the OV Contribution: Under identical spatial-contrast aggregation, LOCOS is stronger on Qwen3-8B, Qwen3-32B, and Gemma-3-12B, comparable on Qwen3-14B and OLMo-3.1-32B, and weaker at large k on Gemma-3-27B.The comparison isolates the per-position observable: OV projection for LOCOS versus attention for the control.
  • Bottom-k Control: Spatial Source Matters: Bottom-k heads leave NoLiMa ROUGE-L near the random baseline despite large absolute answer-logit contributions, countering a circularity explanation.Bottom-k heads contribute predominantly from off-needle positions, whereas top-k ablations produce steep degradation.
  • Layer Distribution: LOCOS scores concentrate in late layers for Qwen3 and Gemma-3-27B, but the layer pattern is family-dependent rather than universal.Tuned-lens and causal-attribution checks support the late-layer concentration on the examined models; top-k heads are interpreted collectively.
  • Retrieval Specificity: LOCOS achieves the highest dissociation-score peak in every model configuration while parametric accuracy remains near baseline under the corresponding ablations.The dissociation score measures relative retrieval degradation against relative parametric-accuracy degradation.
  • Retrieval Specificity: Top-k LOCOS ablation damages NoLiMa more than NIAH, although literal retrieval also declines in five of six models.The widening NoLiMa–NIAH gap indicates that the selected set includes both shared retrieval heads and heads more aligned with non-literal retrieval.
  • Downstream Long-Context Evaluation: On Qwen3-8B, top-50 LOCOS ablation lowers MuSiQue from 0.55 to 0.08 and BABILong from 0.62 to 0.20.Across twelve model–benchmark cells, LOCOS is most damaging in six, while random-head ablation stays within 0.05 of baseline in eleven.

5 Related Work

Prior retrieval-head detectors rely on attention-based observables, while LOCOS builds on QK/OV circuit analysis and projects per-head writes into unembedding space.

  • Attention weights alone have contested explanatory faithfulness, motivating task-grounded attribution methods such as LOCOS.
  • LOCOS extends QK/OV circuit analysis by measuring answer-aligned output writes at per-head, per-position granularity.
  • Attention-based detectors identify retrieval heads using where they attend, including token matching, semantic attention mass, and retrieval-specific cache criteria.

6 Conclusion

LOCOS detects retrieval-relevant heads through answer-logit contributions rather than read-location matching. Across evaluated configurations, its selected heads strongly affect non-literal retrieval while remaining specific to retrieval behavior, subject to scope limitations.

  • LOCOS scores OV-circuit output along the answer-token unembedding direction and contrasts needle with off-needle positions in one forward pass.
  • Across six configurations, ablating top-k LOCOS heads collapses NoLiMa ROUGE-L at lower k than every evaluated attention-based baseline.
  • Parametric recall and arithmetic reasoning remain at baseline under the same ablation, supporting retrieval-specific head selection.
  • LOCOS captures heads in both literal and non-literal retrieval, with its advantage being identification of the non-literal subset missed by token matching.
  • The reported architecture coverage excludes mixture-of-experts routing, encoder–decoder stacks, and state-space hybrids, so transfer of ablation effects requires verification.

A Datasets, Models, and Licenses

The paper documents its external assets, versions, and licensing terms, and states that model weights and dataset content are not redistributed.

  • Tab. 1 lists every external asset used, together with its version and license.
  • All assets are used within their respective license terms, and the paper does not redistribute model weights or dataset content.
  • Gemma Terms of Use permit research and commercial use subject to a prohibited-use policy and attribution requirement.

B Experimental Setup Details

The experiments use documented hyperparameters, calibrated mean-ablation vectors, greedy ROUGE-L evaluation, bootstrap intervals, and specified hardware and per-experiment costs.

  • Mean-ablation calibration averages post-Q-projection queries within each trial, then averages equally across 50 trials.
  • During ablation, calibrated query vectors replace post-Q-projection queries before RoPE, while value and output projections remain unchanged.
  • The resulting attention distribution becomes content-independent because attention logits depend only on keys and position-dependent RoPE rotations.
  • ROUGE-L uses longest common subsequence normalized by gold-answer length, with greedy decoding at temperature = 0 and a maximum of 50 new tokens.
  • Each head score receives a 95% bootstrap confidence interval from 1,000 resamples of ROUGE-passing trials.
  • The Wu et al. token-matching score requires attention inside the needle and a matching attended token, undercounting retrieval heads on NoLiMa.
  • Experiments ran on 2× NVIDIA H100 80GB, with tensor parallelism varying by model size.
  • Detection uses one forward pass per probing trial, while ablation evaluation adds one pass per trial and ablation depth plus a 50-trial calibration pass.

D Code and Data Availability

The release provides detection, ablation, evaluation, environment, and sampling materials, while documenting intended uses and important deployment boundaries.

  • The repository includes LOCOS detection scripts, a vLLM-based ablation driver, evaluation scripts, exact commands, pinned dependencies, and random seeds.
  • NoLiMa inputs are generated from the public release, while parametric control sets are distributed as JSON files.
  • LOCOS is a diagnostic tool for interpretability and retrieval-signal applications, not a new model or a modification of model weights.
  • Publication could support KV-cache compression, retrieval-aware decoding, and targeted hallucination-mitigation interventions.
  • The method requires white-box access and could theoretically enable adversarial retrieval-suppression inputs, although the authors assess the marginal publication risk as low.
  • Evaluation covers only English benchmarks and controls, so multilingual deployments should re-validate LOCOS-derived policies on target languages.

G Worked Example

The worked example shows how LOCOS detects a head whose moderate attention to a non-answer token becomes answer-relevant through its OV circuit, then aggregates this contrast across trials.

  • At needle position 205, attention α = 0.08 accompanies a per-position contribution of 1.3 toward the “Yuki” logit.
  • The OV circuit transforms “Kiasma” into an output aligned with “Yuki,” illustrating non-literal retrieval missed by attention-only scoring.
  • The 12 needle tokens contribute 4.7 to the answer logit, versus 0.23 from a comparable off-needle span.
  • After 200 passing trials and 247 answer steps, the head scores S16,1 = 3.8 with 95% bootstrap CI [3.1, 4.5] and consistency 0.94.
  • LOCOS requires architecture inputs from attention weights, value vectors, output projections, and unembedding matrices, with mechanical adaptations for GQA and multimodal model paths.
  • The ablation implementation uses masked query manipulation and trades continuous batching for an O(n^2) KV cache and single-request decoding.
  • Across models, LOCOS scores are right-skewed, with a small positive tail and strictly negative bottom-50 heads.

K KV-Group × Layer View of Logit-Contribution Scores

The KV-group × layer analysis places top LOCOS cells differently across model families and frames the view as a descriptive aid for compression rather than a standalone localization of retrieval circuitry.

  • GQA shares keys and values across query heads, making the KV group—not the individual head—the relevant cache and spatial-analysis unit.
  • Qwen3 places every top-10 cell in upper layers across all three scales, while Gemma-3-12B and OLMo-3.1-32B span broader layer ranges.
  • The top-10 cells are summarized at KV-group granularity, with distinct-group counts compared against uniform-allocation expectations.
  • The layer pattern is detector- and metric-conditional because direct-path projections may be systematically larger near the unembedding.
  • Bottom-50 heads have strictly negative LOCOS scores, identifying off-needle-dominant contributions rather than high retrieval relevance.
  • The figure is descriptive of LOCOS scores, not by itself a localization of retrieval circuitry; causal interpretation comes from ablations and specificity controls.
  • Bottom-k ablation leaves retrieval and parametric performance near baseline, supporting the specificity of top-k effects.
  • The heatmap encodes mean LOCOS by layer and KV group, with red boxes marking each model’s top-10 cells and spatial position—not absolute color magnitude—carrying the layer interpretation.

L Six-Model Versions of Main-Text Ablation Figures

The six-model appendix extends the ablation and mechanistic analyses, while showing that direct-path and tuned-lens interpretations have important nonlinear and model-specific limitations.

  • Six-model scope: The appendix reports full six-model versions of the ablation figures, while the main-text conclusions use all six configurations.
  • Non-literal retrieval: BABILong requires tracing an entity’s trajectory across interleaved narrative, making literal token copying insufficient.
  • Direct-path analysis: The direct-path derivation motivates the per-position score but characterizes what linear projection misses and how tuned lenses address part of the bias.
  • Direct-path analysis: A head’s contribution is approximately represented through its output and the final unembedding after replacing the final normalization Jacobian with the identity.
  • Direct-path analysis: The direct-path score is an approximation for the evaluated RMSNorm models, validated empirically with a tuned-lens variant.
  • Limitations: Final normalization, downstream attention re-mixing, and nonlinear MLP composition can distort or hide a head’s true causal contribution.
  • Tuned lens: A tuned lens absorbs normalization and downstream linear effects but cannot capture nonlinear composition, leaving some retrieval heads low-scoring under both linear probes.
  • Tuned lens: The tuned-lens variant preserves the Gemma-3-27B late-layer band and surfaces candidate layer-11 heads 26 and 27.

N.6 Beyond linear probes

LOCOS is compared with causal attribution and formalized as a write-aware score that separates where a head reads from what answer-aligned content it writes. The comparison supports both agreement in late-layer localization and meaningful divergence for non-literal retrieval.

  • Empirical comparison: Causal attribution also concentrates top-10 heads in upper layers on Qwen3-8B and Gemma-3-12B, corroborating LOCOS’s late-layer pattern.Causal-attribution peaks occur in layers 22–35 for Qwen3-8B and 35–45 for Gemma-3-12B; LOCOS occupies similar bands.
  • Empirical comparison: 2/10 and 3/10 top-10 heads overlap between causal attribution and LOCOS on Qwen3-8B and Gemma-3-12B, respectively.The small overlap is consistent with the scores capturing different aspects of retrieval-related computation.
  • Method divergence: LOCOS surfaces several Gemma-3-12B layer-0 heads that causal attribution does not flag, leaving two unresolved explanations.The possibilities are a direct-path artifact or distributed circuits whose parallel paths compensate during single-head intervention.
  • Score structure: LOCOS factorizes each per-position contribution into attention mass and answer-aligned write magnitude, while answer-orthogonal writes are irrelevant to the projected score.The factorization is ϕj = αj · ∥u∥2 · cj, so head selection can treat ϕj as proportional to αjcj.
  • Score structure: When answer-aligned write magnitude is position-independent, LOCOS reduces to length-normalized attention-mass contrast; non-literal retrieval violates this condition.For non-literal retrieval, cj differs between needle and off-needle positions, which attention-only scoring cannot observe.
  • Relation to prior scoring: Wu’s score falls from 0.97 on NIAH to 0.03 on NoLiMa because token matching fails while non-literal answer-aligned writing remains large.Under the literal-copy assumptions, Wu’s indicator is a hard-thresholded LOCOS special case; those assumptions fail on NoLiMa.

O.8 Summary

The paper’s summary connects its theoretical reductions to empirical predictions: LOCOS should resemble attention scoring for literal retrieval but diverge on NoLiMa, where its selected heads should be more causally critical. It also identifies practical implications and bounds the evidence to the evaluated setting.

  • Theoretical predictions: On literal NIAH, Wu/NIAH and LOCOS should rank heads similarly; on NoLiMa, LOCOS should credit answer-aligned writes despite token mismatch.These predictions follow from whether the OV write is position-independent or varies with needle versus off-needle content.
  • Theoretical predictions: LOCOS-only heads should produce stronger NoLiMa ablation effects than Wu-only heads under causal validation.The paper reports that the Wu/NIAH–NoLiMa score gap and Fig. 3 ablation separation are consistent with this prediction.
  • Implications: Attention-based KV-cache budgets may under-allocate cache to heads that write answer-aligned content rather than merely attending to answer-bearing positions.The paper proposes substituting LOCOS into head-aware KV-cache scoring as a direct test of this prediction.
  • Scope and validation: The abstract and introduction make three scoped claims: OV-based non-literal retrieval detection, causal validation across six configurations, and retrieval specificity.The checklist states that these claims are backed by experiments in §§ 4.2, 4.4, and 4.6.
Loading 2607.01002v1…