这篇论文告诉你怎样用 LOCOS 找出模型里真正干活的注意力头,在长上下文检索任务上效果拔群,比旧方法准很多。
论文提出 LOCOS(Logit-Contribution Scoring),一种写感知的注意力头检测方法。在 Qwen3、Gemma-3、OLMo-3.1 三个模型家族上,平均消融 LOCOS 选出的 top 头后,在 NoLiMa 非字面检索基准上 ROUGE-L 下降比之前基于注意力的检测方法更快。以 Qwen3-8B 为例,消融 50 个头使 ROUGE-L 从 0.401 降至 0.000,而最强基线仍保留 0.292。同一消融在 MuSiQue 上从 0.55 降至 0.08,在 BABI-Long 上从 0.62 降至 0.20,而随机头消融偏差在 0.05 以内。所选头是检索专用的,参数回忆和算术推理保持在基线水平。
Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads
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.