Research map / Architectures and efficiency / Attention and transformers

Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes LOCOS, a write-aware retrieval-head detector that scores heads by OV-circuit logit contribution to the answer token with needle-vs-off-needle spatial contrast, and validates it via mean-ablation on NoLiMa and downstream long-context tasks.

Problem

Existing retrieval-head detectors rely on attention-based token matching and miss heads that perform non-literal retrieval (synthesizing answers from the meaning of a context span) because they ignore what the head writes through its OV circuit.

Approach

Score each head by the projection of its weighted OV-circuit output onto the correct answer's unembedding vector, contrasting needle and length-normalized off-needle source positions within a single decoding step (spatial contrast), so the detector captures what heads write rather than where they read.

Key results

• On Qwen3-8B, ablating top-50 LOCOS heads drives NoLiMa ROUGE-L from 0.401 to 0.000, while Wu/NIAH-scored ablation retains 0.292.
• On Qwen3-8B, top-50 LOCOS ablation drops MuSiQue accuracy from 0.55 to 0.08 and BABILong from 0.62 to 0.20.
• Bottom-k ablation and random-head ablation leave NoLiMa ROUGE-L near baseline, and parametric recall and arithmetic remain stable under top-k ablation.

Surprising findings

• The attention-only spatial-contrast control is more damaging than LOCOS on Gemma-3-27B at k in {20,50}.
• The top-10 head sets from LOCOS and the Wu-style attention score overlap only marginally (2/10 on Qwen3-8B, 3/10 on Gemma-3-12B).

Limitations

• Top-k ablation also degrades literal NIAH retrieval, so the detected set is not strictly non-literal-specific.
• Evaluated only on dense transformers from three families; MoE, encoder-decoder, and state-space models not covered.
• Late-layer concentration may partly reflect the direct-path metric; tuned-lens check only at head level for Gemma-3-27B.

Experimental setup

Datasets: NoLiMa (onehop) NIAH PopQA MuSiQue BABILong

Benchmarks: NoLiMa NIAH PopQA MuSiQue BABILong

Metrics: ROUGE-L accuracy dissociation score

Baselines: Random heads Wu et al. retrieval score (NIAH-scored) Wu et al. retrieval score (NoLiMa-scored) attention-only spatial-contrast control

Builds on: retrieval heads induction heads direct logit attribution NoLiMa

Closest papers by meaning

• Retrieval Head Mechanistically Explains Long-Context Factuality (ICLR 2025 · similarity 0.90)
• The Atlas of In-Context Learning: How Attention Heads Shape In-Context Retrieval Augmentation (NeurIPS 2025 · similarity 0.88)
• Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism (ICML 2025 · similarity 0.87)
• Retrieval from Within: An Intrinsic Capability of Attention-Based Models (NeurIPS 2026 · similarity 0.87)
• NoLiMa: Long-Context Evaluation Beyond Literal Matching (ICML 2025 · similarity 0.87)
• Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning (ICLR 2025 · similarity 0.87)