Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads

The Problem with Finding What Heads Really Do

When a large language model answers a question about information deep in a long context, it doesn't just copy words verbatim—it synthesizes meaning. A prompt might say "The CEO of Acme Corp is Alice" and ask "Who runs Acme?" The model generates "Alice" by understanding the context, not by mechanically selecting the token that appears in the passage. Identifying which attention heads perform this semantic synthesis has been impossible with existing tools, because those tools only look at what a head reads (which token it attends to), not what it writes (how it transforms that information). This blind spot leaves researchers unable to interpret long-context reasoning.

The authors' core insight: existing detectors reward heads when the attended token matches the generated answer—a literal-copy criterion that misses the actual mechanism of synthesis. The mechanism that matters is the OV circuit (output-value circuit)—the component that transforms attended information into contributions to the final logits. A head could attend to the right token but write garbage; or attend to irrelevant tokens but output semantically meaningful contributions.

How LOCOS Works

Logit-Contribution Scoring (LOCOS) flips the lens from read to write. For each head, the method:

1. Runs a single forward pass through the model on a long-context task
2. Scores each head by projecting its OV-circuit output onto the direction of the answer token in logit space (the "unembedding")
3. Contrasts "needle" positions (where the answer appears in context) against "off-needle" positions (irrelevant spans)

Heads scoring high have OV outputs that push the model toward generating the correct answer specifically when attending to needle positions. This captures heads that synthesize, not just retrieve verbatim.

Why It Matters and Where

On NoLiMa, a non-literal retrieval benchmark, ablating the top 50 LOCOS-identified heads on Qwen3-8B completely destroys performance (ROUGE-L drops from 0.401 to 0.000), while the best prior baseline still scores 0.292. This isn't a marginal improvement—it's the difference between zero and something. LOCOS identifies heads that are essential for the task.

Crucially, these heads are retrieval-specific. Ablating them doesn't harm parametric recall (facts memorized during training) or arithmetic reasoning, proving the method isolates retrieval behavior rather than general language competence. The effect generalizes: on Qwen3-8B, the same ablation also cripples MuSiQue (a multi-hop reasoning benchmark) and BABI-Long, while random head ablations degrade performance by only ~0.05.

The results hold across three model families (Qwen3, Gemma-3, OLMo-3.1), suggesting LOCOS captures a real and consistent phenomenon rather than a quirk of one architecture.

The practical payoff is twofold: mechanistic interpretability researchers can now locate the heads responsible for synthesis-based retrieval, and the method runs in a single forward pass, making it efficient enough for large models and long contexts.