Logit-Contribution Scoring Identifies Non-Literal Retrieval Heads¶

会议: NeurIPS 2026
arXiv: 2607.01002
代码: https://github.com/aryopg/locos
领域: 可解释性
关键词: 大语言模型, 可解释性, 检索头, 注意力机制, 长文本理解

一句话总结¶

提出写入感知的检索头检测方法 LOCOS，通过将注意力头 OV 电路输出投影到答案词解嵌入方向并进行单步空间对比，首次精准定位执行非字面语义综合检索的关键注意力头。

研究背景与动机¶

在大语言模型的长上下文推理中，模型从海量输入中提取信息主要依赖一类稀疏的注意力头，即检索头。然而在真实长文本问答场景中，上下文检索往往不是字面意义上的逐字复制：用户的提问与目标段落可能完全没有词面重叠，模型必须定位相关语义片段、解析其深层含义，并通过语义综合生成答案。这种非字面检索（Non-Literal Retrieval）是模型实现长文本泛化与推理的核心支柱。

然而，现有的检索头识别方法（如基于 NIAH 的字面词匹配或加权注意力累加）存在根本性的结构盲区。它们在检测准则上仅依据注意力权重矩阵，即只观测注意力头在源序列上的读取位置（QK 电路），并强制要求被注意到的 token 与最终生成的 token 保持字面一致。但在非字面检索中，注意力头可能高度聚焦于背景线索短语，而其输出-值电路（OV 电路）却向残差流写入指向语义推理答案的更新方向。读取位置与写入内容在非字面场景下发生解耦，导致已有注意力探测器从原理上漏掉了所有非字面检索头。

本文切入的角度是：注意力头是一个完整的“读取-写入”双阶段通路，决定检索有效性的关键在于该头最终向正确答案词的 Logit 空间实际注入了多少有效贡献。核心 idea：提出基于 Logit 贡献评分的写入感知检测方法 LOCOS，通过在单步解码中计算 OV 电路输出在目标答案词解嵌入向量上的标量投影，并引入针尖与非针尖区域的长度归一化空间对比，在单次前向传播中精准捕获非字面检索头。

方法详解¶

整体框架¶

LOCOS 旨在克服传统基于纯注意力分布评分无法感知输出值变换的弊端。对于给定的长文本探测样本，模型在解码目标答案词时，LOCOS 直接切入每个注意力头（Layer $l$, Head $h$）在各个源位置的输出向量，度量其通过解嵌入矩阵映射到答案词 Logit 的瞬时贡献。整体流程包含位置级 Logit 贡献投影、针尖与背景的长度归一化空间对比，以及跨测试样本的稳健池化聚合三个核心阶段。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["长文本探测输入与目标答案词"] --> B["位置级 Logit 贡献投影<br/>QK 权重结合 OV 向量投影到解嵌入空间"]
    B --> C["长度归一化空间对比<br/>消除全序列背景基线与词频先验偏差"]
    C --> D["单步跨样本稳健池化<br/>均值聚合与无截断全局检索头排序"]
    D --> E["输出关键非字面检索头集合"]

关键设计¶

1. 位置级 Logit 贡献投影：打通 OV 写入电路与答案词空间 以往方法仅观察注意力权重 $\alpha_{t,j}^{(l,h)}$，无法感知不同源位置的值向量经投影矩阵 $W_O$ 变换后所携带的语义信息。LOCOS 结合注意力读取权重与输出变换向量，将位置 $j$ 的输出向量 $\mathbf{o}_{t,j}^{(l,h)} = \alpha_{t,j}^{(l,h)} W_O^{(l,h)} \mathbf{v}_{t,j}^{(l,h)}$ 直接与目标答案词 $y_t$ 在解嵌入矩阵中的行向量 $\mathbf{u}_{y_t}$ 做内积，计算该位置对答案词 Logit 的直接标量贡献： $$\phi_{t,j}^{(l,h)} = \mathbf{u}_{y_t}^\top \mathbf{o}_{t,j}^{(l,h)} = \alpha_{t,j}^{(l,h)} \cdot \mathbf{u}_{y_t}^\top W_O^{(l,h)} \mathbf{v}_{t,j}^{(l,h)}$$ 该设计将注意力头拆解为读取与写入两部分。即使某个头对线索位置仅分配中等大小的注意力权重，只要其 OV 回路将该特征成功旋转并投射到正确答案的词表方向，$\phi_{t,j}^{(l,h)}$ 便会展现出极大的正向标量值；相反，盲目分配高注意力但输出正交于答案方向的头则获得接近 0 的评分。

2. 长度归一化空间对比：消除全序列背景基线与词频先验偏差 若直接将全序列的 Logit 贡献相加，模型自带的词频先验（Token Frequency Priors）或非特定语义背景会在长文本中积累虚高数值，掩盖针尖（Needle）位置的真实检索作用。LOCOS 摒弃了跨生成步的传统时间对比，创新性地提出在单次解码步内进行针尖区域与非针尖区域（Off-needle）的空间对比。设针尖区间为 $[s_\tau, e_\tau)$，序列总有效长度为 $N_t$，分别定义针尖贡献总和 $\Phi_t^{(l,h),+}$ 与经长度缩放的背景贡献 $\Phi_t^{(l,h),-}$： $$\Phi_t^{(l,h),+} = \sum_{j=s_\tau}^{e_\tau - 1} \phi_{t,j}^{(l,h)}, \qquad \Phi_t^{(l,h),-} = \frac{e_\tau - s_\tau}{N_t - (e_\tau - s_\tau)} \sum_{j \notin [s_\tau, e_\tau)} \phi_{t,j}^{(l,h)}$$ 缩放因子 $(e_\tau - s_\tau) / [N_t - (e_\tau - s_\tau)]$ 将非针尖区域的贡献严格校准到与针尖相同的等效跨度长度。这一差值对比能够彻底消除与位置无关的全局基础贡献，精准提取出仅由目标针尖触发的局域增量。

3. 单步跨样本稳健池化：无截断保留正负对比信号 为了获得稳定的头级评分，LOCOS 在所有通过正确性初筛（ROUGE-1 recall $> 0.5$）的测试样本 $\mathcal{D}_{\text{pass}}$ 和有效答案生成步 $\mathcal{A}^\tau$ 上进行平均池化，生成全局评分： $$S_{l,h} = \frac{1}{\sum_{\tau \in \mathcal{D}_{\text{pass}}} |\mathcal{A}^\tau|} \sum_{\tau \in \mathcal{D}_{\text{pass}}} \sum_{t \in \mathcal{A}^\tau} \Big(\Phi_t^{(l,h),+} - \Phi_t^{(l,h),-}\Big)$$ 关键的是，LOCOS 严格禁止在零处进行人工截断（Unclamped）。正分显著的头代表真正从针尖提炼答案的检索头；而严重负分的头则代表答案 Logit 主要源自非针尖的通用上下文先验或参数化记忆。保留负值不仅为后续分析提供了天然的下界对照组，还保证了排名的线性无偏性。

实验关键数据¶

主实验¶

在代表性非字面长文本检索基准 NoLiMa（800 篇留出测试集）上，通过对评选出的 Top-$k$ 注意力头进行均值消融（Mean-Ablation，用校准均值向量 $\bar{\mathbf{q}}$ 替换查询向量），验证各方法所选头的因果重要性。在六种开源模型规模下，LOCOS 所选头被消融后均导致性能急剧崩塌。

模型架构	评估指标	LOCOS 消融 (k=50)	Wu/NIAH 消融 (k=50)	原始 Baseline	性能降幅 (vs Baseline)
Qwen3-8B	ROUGE-L	0.000	0.292	0.401	-100.0%
Qwen3-14B	ROUGE-L	0.002	0.285	0.418	-99.5%
Qwen3-32B	ROUGE-L	0.077	0.274	0.395	-80.5%
Gemma-3-12B	ROUGE-L	0.000	0.230	0.420	-100.0%
Gemma-3-27B	ROUGE-L	0.015	0.260	0.410	-96.3%
OLMo-3.1-32B	ROUGE-L	0.001	0.250	0.380	-99.7%

消融实验¶

为验证 OV 电路投影与空间对比各自的独立贡献，以及排除“高 Logit 投影头必然敏感”的循环论证假设，论文在 Qwen3-8B 和 Qwen3-32B 上对比了完整 LOCOS、纯注意力空间对比变体（$\alpha$-spatial）、倒数 Bottom-$k$ 头（强负分头）以及随机头消融：

消融配置	Qwen3-8B ROUGE-L (k=50)	Qwen3-32B ROUGE-L (k=50)	配置说明与因果分析
原始 Baseline	0.401	0.395	未进行任何消融的正常生成性能
随机头对照 (Random)	0.385	0.388	随机采样 50 个头，检索性能几乎完全保持
Bottom-50 消融	0.398	0.390	消融非针尖贡献极大的头，完全不影响检索能力
纯注意力空间对比 ($\alpha$-spatial)	0.148	0.346	仅利用空间对比但丢弃 OV 投影，消融损伤显著放缓
完整 LOCOS (Top-50)	0.000	0.077	兼顾 OV 写入与空间对比，检索功能彻底瓦解

关键发现¶

OV 投影是定位非字面检索的核心杠杆：在相同空间对比聚合策略下，单纯依赖注意力权重 $\alpha$ 在 Qwen3-32B 上消融 50 个头仅使 ROUGE-L 从 0.395 微跌至 0.346，而加入 OV 投影的 LOCOS 则直接击穿至 0.077，证明非字面场景下值向量的语义变换才是因果命脉。
检索功能与参数化记忆高度解耦：在分离度分析（Dissociation Score）中，消融 LOCOS Top-$k$ 头导致 NoLiMa 检索大幅暴跌的同时，城市-国家常识问答、PopQA 参数化召回以及基础算术推理（Addition/Subtraction）准确率与未消融基线完全持平，表明所选头专门行使上下文检索职责而非通用输出调节。
下游真实长文本任务具备强迁移性：在未经探测调优的多跳推理基准 MuSiQue 和实体追踪任务 BABILong 上，消融 Qwen3-8B 的 Top-50 LOCOS 头使 MuSiQue 准确率从 0.55 断崖式下跌至 0.08，BABILong 准确率从 0.62 暴跌至 0.20，证实其捕获的是通用的深层长文本理解电路。

亮点与洞察¶

将机械可解释性推进至“读取-写入统一视图”：传统注意力归因长期被束缚在“看哪里即代表用哪里”的误区中，LOCOS 严格证明了注意力和 OV 输出在非字面场景的显式解耦，通过线性解嵌入投影在单次前向中同时完成了读写闭环。
空间对比替代时间对比的严谨优雅：传统方法依赖多步生成中的时间步对比，对多步均持续激活的稳定检索头完全失效；LOCOS 借助针尖与背景的长度归一化空间差分，使得在单个解码步内即可纯净剥离全序列常数偏置与高频词干扰。
对 GQA 架构 KV Cache 压缩提供全新指导粒度：论文发现检索活跃度在模型后部层呈现家族聚类，且在 KV Group 维度分布均匀，这为按 (Layer, KV-Group) 实施非均匀细粒度缓存压缩策略提供了强有力的机械依据。

局限与展望¶

依赖直接路径线性近似假设：LOCOS 计算标量贡献时采用 Direct-Path 假设，忽略了深层 LayerNorm 和后续 MLP 层的复杂非线性调制。尽管 Tuned Lens 实验表明主要后期层结论依然稳健，但早期层可能存在被下游 MLP 变换掩盖的隐式检索头。
对富含竞争干扰项的语境敏感：空间对比假设非针尖区域为中性背景；若超长文本中充斥大量与答案语义紧密相关的高度干扰性伪线索，$\Phi^-$ 会异常升高，可能低估执行弥散性全局上下文整合的注意力头。
动态长度自适应机制：在超大规模文本中，不同段落密度的信息衰减速度不同，固定空间差分权重在极端非均匀文本流中可能需要动态温度调节。
待向跨模态与复杂架构拓展：当前验证主要覆盖 Dense 与 GQA 结构，尚未推广至 MoE 动态路由、状态空间混合架构（如 Mamba-Transformer）及图文多模态输入场景。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从 OV 写入电路线索揭示非字面检索头的存在，打破传统纯注意力检测盲区。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 大模型家族 6 个尺寸，深入开展了因果均值消融、反向控制、任务解耦与下游迁移评测。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，理论证明完备（将已有方法严谨归为特例），图表逻辑清晰。
价值: ⭐⭐⭐⭐⭐ 为长文本大模型机制解释、幻觉根因分析及下一代高效 KV 缓存压缩提供了坚实基石。