Aryo Pradipta Gema, Beatrice Alex, Pasquale Minervini
4 min
Abstract
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.
Alex: And when they tested this — when they removed the heads LOCOS identified as important — what happened?
Sam: The model's ability to answer questions collapsed much faster than when they used older methods to select which heads to remove. That's the key validation. It shows LOCOS is finding the heads that genuinely matter, not just the ones that look busy.
Alex: Did removing those heads affect anything else, like basic arithmetic or general knowledge?
Sam: The effect was specific to retrieval. The model's ability to do arithmetic or recall facts it already knew stayed intact. That specificity is meaningful — it suggests LOCOS is isolating a distinct piece of machinery, not just disrupting the model broadly.
Alex: You mentioned a benchmark called NoLiMa. What makes that test particularly useful here?
Sam: NoLiMa is designed to test exactly the cases where copy-pasting fails. The answer to a question doesn't share any words with the relevant part of the source text — so the model genuinely has to reason rather than pattern-match. Older methods that rely on token-matching struggle badly on these tasks. LOCOS, because it tracks what the model writes rather than just where it looks, successfully identifies the heads doing that underlying reasoning.
Alex: So the older tools were essentially blind to an entire category of thinking the model was doing.
Sam: That's a fair way to put it. And the practical implications go beyond just understanding the model better. If you can precisely identify which components handle retrieval, you can start to optimize how models use memory, or improve how you diagnose failures — figuring out why a model got something wrong, rather than just observing that it did.
Alex: That shift from observation to diagnosis feels like it matters quite a bit for building more reliable systems.
Sam: It does. There's a meaningful difference between knowing a model gave a wrong answer and knowing which part of its reasoning broke down. LOCOS is a step toward that second kind of understanding — and that's what makes it worth paying attention to.
Alex: Thanks for listening to ResearchPod.