Logit-Contribution Scoring 识别非字面检索头

# Logit-Contribution 评分识别非字面检索头  
来源：https://arxiv.org/html/2607.01002  

Aryo Pradipta Gema^Q^, Beatrice Alex^K^, Pasquale Minervini^Q,V^  
^Q^爱丁堡大学  
^K^赫瑞-瓦特大学  
^V^Miniml.AI  
{aryo.gema, p.minervini}@ed.ac.uk  
[email protected]  

###### 摘要  
在长上下文应用中，大语言模型通常根据相关上下文片段的意义综合生成答案，而非直接逐字复制。识别哪些注意力头执行这种综合操作，对解释长上下文模型行为至关重要。然而，现有检测方法在构造上遗漏了这些头：它们奖励注意力头关注的 token 与生成的 token 相匹配的情况，这是一种逐字复制标准，只捕捉了头在“哪里”读取（通过其输出值（OV）电路），但未捕捉其“写入”了什么，而 OV 电路正是执行非字面检索的机制。我们引入 Logit-Contribution 评分（LOCOS），这是一种关注“写入”的检测方法，通过将每个头的 OV 电路输出投影到答案 token 的解除嵌入方向上进行评分，并在单次前向传递中对比“针”（needle）和“非针”（off-needle）源位置。在三个模型族（Qwen3、Gemma-3、OLMo-3.1）上，对 NoLiMa 非字面检索基准的保留集进行平均消融实验，移除 top LOCOS 头后，在比先前基于注意力的检测方法更少的头数量上，ROUGE-L 性能崩溃：在 Qwen3-8B 上，消融 50 个头使 ROUGE-L 从 0.401 降至 0.000，而最强基线仍保持 0.292。所选头具有检索特异性：在相同消融下，参数召回和算术推理仍保持基线水平。在 Qwen3-8B 上，相同消融也使 MuSiQue 从 0.55 降至 0.08，BABILong 从 0.62 降至 0.20，而随机头对照组保持在基线 ±0.05 范围内。  
[![未加说明图片](https://arxiv.org/html/2607.01002v1/GitHub-Mark.png) locos](https://github.com/aryopg/locos)  
[![未加说明图片](https://arxiv.org/html/2607.01002v1/hf-logo.png) locos-results](https://huggingface.co/datasets/aryopg/locos-results)  

## 1 引言  
大语言模型（LLM）从其输入上下文中检索信息（而非依赖记忆参数知识）的能力，取决于一组稀疏的注意力头，称为*检索头*（Wuet al., 2025 (https://arxiv.org/html/2607.01002#bib.bib17)），这建立在早期关于归纳头的机制研究之上（Elhageet al., 2021 (https://arxiv.org/html/2607.01002#bib.bib24); Olssonet al., 2022 (https://arxiv.org/html/2607.01002#bib.bib25)）。然而，在实际应用中，上下文检索很少是逐字复制：用户的问题可能与相关段落没有词汇重叠，模型必须识别相关片段，解析其含义，并从中综合生成答案（如 Fig. 1 (https://arxiv.org/html/2607.01002#S1.F1) 所示）。然而，所有现有的识别方法，无论是通过 token 匹配启发式（Wuet al., 2025 (https://arxiv.org/html/2607.01002#bib.bib17)）还是加权注意力累积（Fuet al., 2025 (https://arxiv.org/html/2607.01002#bib.bib19); Linet al., 2025 (https://arxiv.org/html/2607.01002#bib.bib18)），都在逐字复制任务上评估头，并共享一个共同的可观测变量：每个头在源位置上的注意力模式。这个可观测变量捕捉了头在“哪里”分配注意力，但没有捕捉头通过其输出值（OV）电路传播“什么信息”，而这正是执行非字面检索的机制。两个具有相同注意力模式但不同 OV 电路的头，可以从相同位置传播完全不同的信息。对于字面检索（注意力 token 就是答案 token），注意力和 OV 输出是平凡对齐的，基于注意力的评分有效。在非字面检索中，一个头可能关注“埃菲尔铁塔”，同时向残差流写入“Yuki”方向。基于注意力的评分看到的是读取位置（“埃菲尔铁塔”）；OV 电路决定写入内容（“Yuki”），两者不必一致。这种区别在实际中很重要，而以往的检索头检测器都没有做出这种区分。  
参考图注  
Figure 1: 非字面检索需要综合。相同的上下文以不同方式回答两个问题：字面问题需要直接从“针”中读取“埃菲尔铁塔”，而非字面问题必须在综合上下文后产生“Yuki”。  
我们的方法 Logit-Contribution 评分（LOCOS），衡量每个注意力头在解除嵌入空间中对正确答案 token 的贡献（参见 Fig. 2 (https://arxiv.org/html/2607.01002#S3.F2)）。对于每个头在每个源位置，该方法计算头的加权 OV 电路输出在正确答案解除嵌入向量上的标量投影。聚合使用*空间对比*：针位置的 logit 贡献与长度归一化的非针位置贡献在单次解码步骤中进行比较。该方法每个探测试验只需一次前向传递。在三个模型族（即 Qwen3 (8B, 14B, 32B) (Team,2025b (https://arxiv.org/html/2607.01002#bib.bib33))、Gemma-3 (12B, 27B) (Team,2025a (https://arxiv.org/html/2607.01002#bib.bib34)) 和 OLMo-3.1 (32B) (Olmoet al., 2025 (https://arxiv.org/html/2607.01002#bib.bib35))）的六种配置上，对 NoLiMa 非字面检索基准保留集（Modarressiet al., 2025 (https://arxiv.org/html/2607.01002#bib.bib20)）进行的消融实验验证了该方法：平均消融排名靠前的 LOCOS 头比所有评估的基线产生更陡峭的 ROUGE-L 退化曲线（§ 4.2 (https://arxiv.org/html/2607.01002#S4.SS2)）。在 Qwen3-8B 上，消融前 50 个 LOCOS 头使 ROUGE-L 从 0.401 降至 0.000，而最强基于注意力的基线在同一深度保持 0.292。控制实验证实了检索特异性：在相同消融下，参数召回和算术推理保持完整（§ 4.6 (https://arxiv.org/html/2607.01002#S4.SS6)）。相同消融也降低了下游长上下文性能，在 Qwen3 族上最显著（§ 4.8 (https://arxiv.org/html/2607.01002#S4.SS8)）：在 Qwen3-8B 上，消融前 50 个 LOCOS 头使 MuSiQue 准确率从 0.55 降至 0.08，BABILong 从 0.62 降至 0.20；迁移在 Qwen3 族上最一致，而在 Gemma-3 和 OLMo-3.1 上，排名与基于注意力的基线相比具有基准依赖性（§ 4.8 (https://arxiv.org/html/2607.01002#S4.SS8)）。  

## 2 背景  
**符号说明**。考虑一个具有 L 层、每层 H 个注意力头、头维度 d_h 和模型维度 d = H·d_h 的 Transformer (Vaswaniet al., 2017 (https://arxiv.org/html/2607.01002#bib.bib1))。头 (l, h) 以权重 α_{t,j}^{(l,h)} 关注源位置，读取值向量 v_{t,j}^{(l,h)} ∈ R^{d_h}，并通过其输出投影 W_O^{(l,h)} ∈ R^{d×d_h} 向残差流写入。头 (l, h) 从源位置 j 输出的每位置输出为：  
o_{t,j}^{(l,h)} = α_{t,j}^{(l,h)} · W_O^{(l,h)} v_{t,j}^{(l,h)} ∈ R^d. (1)  
解除嵌入矩阵 W_U ∈ R^{|V|×d} 将残差流映射到 logits；u_y ∈ R^d 是其第 y 行。  
**作为读写电路的注意力头**。一个注意力头分解为 QK 电路（决定 α_{t,j}^{(l,h)}，即头在何处读取）和 OV 电路（通过 W_O 将每个源值映射到残差流，即头写入什么）(Elhageet al., 2021 (https://arxiv.org/html/2607.01002#bib.bib24))。头 (l, h) 在步骤 t 的完整输出为 Σ_j o_{t,j}^{(l,h)}，下一个 token 的 logits 来自通过 W_U 映射最终层残差流。头只有在两个阶段一致时才对检索有用：QK 电路必须选择正确的源位置，且 OV 电路必须写入与答案对齐的更新。  
**归纳头和检索头**。归纳头实现逐字复制模式 [A][B]...[A] → [B] (Olssonet al., 2022 (https://arxiv.org/html/2607.01002#bib.bib25))：它们对先前匹配的位置放置高 α_{t,j}^{(l,h)}，并写入一个投影到 u_{y_t} 上的输出，因为关注的 token 就是下一个 token。由 Wu 等人 (2025 (https://arxiv.org/html/2607.01002#bib.bib17)) 在字面“大海捞针”（NIAH）提示上识别的检索头，将此图景推广到长上下文事实性。其检测程序奖励那些 argmax 注意力落在针内且关注的 token 与生成 token 匹配的头，即逐字复制标准。  
**字面与非字面检索**。在 NIAH 设置中 (Kamradt,2023 (https://arxiv.org/html/2607.01002#bib.bib2))，一个“针”是插入较长干扰上下文（“草堆”）中的短答案承载片段；我们将其索引为 [s_τ, e_τ) 用于试验 τ，并将所有其他源位置称为“非针”。在字面 NIAH 中，答案 token 出现在针内，因此大的针注意力 α_{t,j}^{(l,h)} 通常与大的 logit 相关写入 u_{y_t}^⊤ o_{t,j}^{(l,h)} 一致。NoLiMa 打破了这种等价性 (Modarressiet al., [2025 (https://arxiv.org/html/2607.01002#bib.bib20))：答案必须从针的含义中恢复，且可能与针没有词汇重叠。一个头可以通过关注 [s_τ, e_τ) 内的语义相关短语并写入与答案对齐的方向来检索非字面信息，即使没有关注的 token 匹配 y_t。这是 LOCOS 针对的场景。  

## 3 Logit-Contribution 评分  
参考图注  
Figure 2: 一个注意力头有两个电路：它在何处读取（QK）以及它写入什么（OV）。Logit-Contribution 评分使用 OV 电路来识别*非字面检索头*。  
(a) 头每位置输出的解剖：QK 电路产生注意力权重 α_{t,j}；OV 电路产生 W_O v_j。基于注意力的方法只测量 α。Logit-Contribution 评分（LOCOS）测量 φ = u_{y_t}^⊤ (α · W_O v_j)，捕捉完整流程。  
(b) 对非字面检索的后果：两个头从“埃菲尔铁塔”读取以回答“Yuki”。头 A 关注强（α=0.30），但其 OV 输出与答案方向正交（φ≈0）。头 B 关注中等（α=0.08）但向答案写入（φ=1.3）。基于注意力的方法选择头 A；LOCOS 选择头 B。  

我们引入 LOCOS，它根据每个头向答案“写入”什么（而非分配注意力到何处）进行评分，因为非字面检索通过 OV 电路转换关注内容后才成为答案。我们定义三步过程：  
**每位置 Logit 贡献**。考虑一个探测试验 τ，具有针跨度 [s_τ, e_τ)，嵌入在 N_τ 个 token 的上下文中。设 A^τ 表示模型生成正确答案 token 的解码步骤集合（通过匹配标记化的金标准答案识别），y_t ∈ V 是在步骤 t 的正确 token，N_t 是在步骤 t 可用的键位置总数。源位置 j 通过头 (l, h) 对 y_t 的 logit 的贡献为：  
φ_{t,j}^{(l,h)} = u_{y_t}^⊤ o_{t,j}^{(l,h)} = α_{t,j}^{(l,h)} · u_{y_t}^⊤ W_O^{(l,h)} v_{t,j}^{(l,h)} ∈ R. (2)  
标量 φ_{t,j}^{(l,h)} 同时依赖于头在何处读取（通过 α_{t,j}^{(l,h)}）以及它提取了什么（通过 W_O^{(l,h)} v_{t,j}^{(l,h)}）。一个强烈关注某位置但其 OV 输出正交于 u_{y_t} 的头获得 φ≈0，尽管注意力高；相反，执行非字面检索（例如，关注“巴黎”以产生“法国”）的头获得大的 φ，因为其 OV 电路将关注表示转换为与答案对齐的输出。  
**空间对比**。对于试验 τ 中答案步骤 t 处的每个头 (l, h)，我们定义来自针和非针位置的 logit 贡献：  
Φ_t^{(l,h),+} = Σ_{j=s_τ}^{e_τ-1} φ_{t,j}^{(l,h)},  
Φ_t^{(l,h),-} = (e_τ - s_τ) / (N_t - (e_τ - s_τ)) · Σ_{j∉[s_τ, e_τ)} φ_{t,j}^{(l,h)}, (3)  
其中重新缩放因子 (e_τ - s_τ) / (N_t - (e_τ - s_τ)) 使 Φ_t^{(l,h),+} 和 Φ_t^{(l,h),-} 具有可比性：两者都表示长度为 (e_τ - s_τ) 的区域的对数贡献。对比 Φ_t^{(l,h),+} - Φ_t^{(l,h),-} 是*空间性*的：它在单个解码步骤内比较针和非针位置，而非基于注意力的方法使用的*时间性*对比（答案步骤与非答案步骤）。空间对比从单个答案步骤产生分数，识别那些持续关注针但仅从针位置写入答案相关内容的主，并抵消均匀贡献者（如 token 频率先验，其 φ_{t,j}^{(l,h)} 大但与位置无关）。  
**聚合**。我们通过正确性过滤器（ROUGE-1 召回 > ρ，默认 ρ=0.5）的所有试验中的所有答案步骤进行池化，不做每试验归一化。设 D_pass ⊆ {1,...,T} 表示通过试验的集合。最终分数为：  
S_{l,h} = (1 / Σ_{τ∈D_pass} |A^τ|) · Σ_{τ∈D_pass} Σ_{t∈A^τ} (Φ_t^{(l,h),+} - Φ_t^{(l,h),-}). (4)  
分数 S_{l,h} 是所有（试验，答案步骤）对上的均值，每个答案步骤权重相等。与基于注意力的方法不同，我们不对 S_{l,h} 进行零值截断：负值表示头对正确答案的 logit 贡献主要来自非针位置，这指向针特定检索以外的因素（例如，参数或上下文关联）。  
在 [附录 G (https://arxiv.org/html/2607.01002#A7) 中报告了一个带有每试验一致性诊断的示例。  
LOCOS 在 OV 电路不贡献位置相关信号时，将基于注意力的评分作为特殊情况恢复。¹¹Appx.O (https://arxiv.