学习重点：使用因果证据集监督稀疏注意力路由

# 使用因果证据集监督稀疏注意力路由

来源：https：//arxiv.org/html/2607.21692

## 学习本质：使用因果证据集监督稀疏注意力路由

###### 摘要

稀疏注意力通过允许每个查询仅读取输入的选择部分来降低长上下文的成本。这些选择器通常通过蒸馏密集教师的注意力模式来训练，假设注意力能够揭示教师实际使用的上下文。我们在检索任务上测试这一假设，其中每个答案的证据是确切已知的。通过屏蔽部分上下文并测量答案是否变化，我们发现注意力和因果依赖经常不一致，蒸馏选择器继承了这种不匹配。教师会注意它们已学会忽略的过时事实，并且即使它们依赖于相同的证据，其注意力也可能因训练运行而异。在一个两步引用任务中，答案处的注意力跳过了中间步骤，因为它在前向传递中已被更早解析：一个基于注意力训练的选择器达到了41%的准确率，而同一个基于因果证据训练的选择器则达到了99%，与教师持平。这些证据集不需要标注：仅通过屏蔽从冻结教师中恢复，训练出的选择器可达到相同准确率。我们在预训练模型中也发现了同样的冲突：Qwen2.5-3B在58%的冲突事实示例中对过时事实的关注多于当前事实，尽管回答正确；而Gemma-2-9B在仅限两个相关句子时，准确率从56%提升到99%。注意力显示模型看向哪里，但不一定显示其答案依赖什么；在我们测试的所有场景中，作为训练目标，这种依赖匹配或优于注意力。

## 1.引言

Transformer通过注意力读取其上下文：每个新标记对每个更早的标记进行评分，并取其加权混合。成本随上下文长度的平方增长，而在长上下文中，大部分成本花在了从未影响输出的内容上。稀疏注意力通过让每个查询只读取上下文的一小部分子集来降低成本。这个子集必须被选择。目前最强的选择器通过蒸馏（Ahmad & Yun，2026 (https：//arxiv.org/html/2607.21692#bib.bib2)）训练：一个小型网络，即 *选择器*，学习复制完整密集模型（即 *教师*）的注意力，以便在推理时模型只读取教师本应关注的内容。这种设计假设模型的注意力能够识别其需要的上下文。

本文测试了这一假设。测试需要一个场景，其中模型所需的上下文可以独立于其注意力而被知晓，因此我们构建了一个。我们的任务是合成型键值检索问题，其中上下文是一系列固定宽度的 *块*，每个块包含一个 *记录*，后跟一个查询。一条记录是一个存储的事实：一个键和其下存储的值，写成一行简短的标记。通过构造，我们知道哪些记录决定了答案。对于某个示例，如果将该模型的注意力限制在这些块和查询上，而其答案保持不变，则一组块是 *充分的*。充分性通过干预来验证，即在每一层屏蔽进入某个块的注意力，同时保持标记位置不变。监督目标是 *因果证据集*（第2.3节 (https：//arxiv.org/html/2607.21692#S2.SS3)）：当存在任何这样的块时，每个单独维持答案的块；否则，每个屏蔽后会改变答案的块。这个标签并不总是一个最小的充分集。在这些任务上，我们从多个随机种子训练密集教师，冻结它们，并比较选择器（我们称之为 *路由器*）的两种监督目标：教师的注意力权重，以及因果证据集。部署路由器会屏蔽其未选择的每个块，因此模型仅从保留的块中回答。由此产生的准确率即为它的 *路由准确率*。

最清晰的结果来自于一个需要跟随引用来回答的任务。一条记录可能持有一个指向另一条记录的指针来代替值，因此被查询的记录开始一条由两个指针组成的链，最终指向一个值，证据就是那个三块链。我们将此任务称为 *多跳* 检索，以区别于 *单记录* 任务，其中一条记录决定答案。一个密集教师以99%的准确率解决了该任务。但是指针步骤在更早的层和更早的位置已被解析，到产生答案时，注意力在不到10%的示例上覆盖了整个链。路由器继承了这种遗漏。使用相同架构、数据和预算，一个在标注链上监督的路由器以99%的路由准确率完成任务，而同一个通过蒸馏教师注意力获得的路由器仅达到41%。逐元素召回率解释了这种差距。它在第一个指针上为1.00，在最终值上为0.99，但在中间块上为0.00。注意力显示了答案被读取的位置，而非计算的位置，模仿它会剪掉中间的计算。从外部观察无法得知这何时发生。同样的蒸馏在不同训练运行中在0.41和0.98之间波动，而一个拒绝捷径的任务变体又使其正常工作。注意力与其说是一个错误的目标，不如说是一个不经干预就无法验证的目标，而干预本身就是更好的标签。

单记录任务以较温和的形式得出了相同的结论。在一个键被多次写入且只有最后一次写入重要的任务上，每个教师，在每个种子上，都在每一个示例上花费注意力于过时的写入。注意力在种子间一致，并且一致地包含了不重要的证据。它也很分散。在答案位置覆盖90%的注意力质量大约需要32个块中的14个，而因果证据大约占据1.5个。通过相同的干预恢复的因果证据集，在种子间的一致性远高于注意力（Jaccard指数0.79到0.90，而大小匹配的注意力集为0.38到0.46）。种植回答记录的冗余副本会破坏单块干预，因为屏蔽一个副本不会改变任何东西，因此没有单个块看起来是必要的。可互换的充分集家族在这种情况下仍然存在，并且随着冗余度的增加，它在种子间的一致性更好，从一个副本时的0.81增加到四个副本时的0.90。

这些测量对路由很重要，因为监督是一种选择。我们训练具有相同架构和输入的路由器，一个模仿块级注意力分布（模仿路由器），一个在因果证据集上训练（因果路由器），每个都以硬掩码的形式部署在冻结的教师上，预算为10%的块。在分布内，两种监督不相上下，都紧跟密集教师；在训练长度四倍的情况下，因果路由器在唯一检索上的每个种子都领先9个准确率点，经过学习率扫描和基线三倍训练后依然如此。在教师的注意力和标注证据基本一致的地方，比如上述的最近写入任务，两种监督不分伯仲，这正是测量所预测的模式。

本文做出了四项贡献。

- • 一个带有已知证据的测试平台，其中的干预表明注意力和依赖性会分离：注意力覆盖过时的写入，在答案处错过跳链，并在证据固定的种子间变化。
- • 匹配的路由实验，其中影响来自标签而非损失。因果监督在多跳任务上达到0.99，而注意力蒸馏为0.41，后者从未选择链的中间块。路由器还必须足够表达：平均池化在任何监督下都无法处理链。
- • 一个无需标注的因果证据集估计器。在其标签上训练的路由器与基于标注训练的路由器相匹配，而所有单信号替代方案在测量预测的失败之处均告失败。
- • 在三个系列的冻结预训练模型上观察到同样的冲突和路由排序，包括冲突事实修复、精确标签恢复，以及在反事实SQuAD上保持的排序。

先前的工作要么模仿注意力，要么质疑它。但没有一项将其作为路由目标与已知证据进行比较。学习型选择器模仿教师注意力（Ahmad & Yun，2026 (https：//arxiv.org/html/2607.21692#bib.bib2)；Roy et al.，2021 (https：//arxiv.org/html/2607.21692#bib.bib18)）。关于端到端稀疏训练的工作表明，学习到的门控几乎不比随机的好，因为模型的其余部分与掩码共同适应（Aquino-Michaels，2026 (https：//arxiv.org/html/2607.21692#bib.bib3)）；这就是为什么我们首先密集训练教师然后将其冻结。电路发现方法事后恢复训练模型的因果子图（Bhaskar et al.，2024 (https：//arxiv.org/html/2607.21692#bib.bib7)），但它们不会产生一个在注意力计算之前运行的、每个输入一个的预测器。其他相关工作在第9节 (https：//arxiv.org/html/2607.21692#S9)中讨论。

## 2. 设置

### 2.1 任务

每个示例是一个由 \(n\) 个长度为 \(w\) 的块组成的序列，后跟一个短查询，除非另有说明，否则 \(n=32\) 且 \(w=8\)。每个块包含一条记录，WRITE ff mm vv，并填充了填充符：一个两标记的键（族 ff，成员 mm）和一个单标记值 vv。两标记的键提供了分级干扰项相似度，因为干扰项可以共享族标记而不共享键。查询是 QUERY ff mm vv，模型根据在位置 vv 处的预测来评分。表1 (https：//arxiv.org/html/2607.21692#S2.T1) 列出了这些任务。

表1：任务。充分集通过构造可知，并通过干预验证（第2.2节 (https：//arxiv.org/html/2607.21692#S2.SS2)）。在最近写入任务中，较早的写入带有不同的值，因此用一个值回答是错误的。在重复证据中，每个副本单独就是充分的，因此目标是可互换集的集合；我们称一个示例的最小充分集族为其 *族*。在多跳中，指针记录持有一个第二个键来代替值。聚合要求统计命名族中的记录数量，因此每个标记的块都是必要的，只有完整的标记集才是充分的。这是预期的边界情况。

两个协议细节对复现很重要：值直接在查询键处预测，因为键和值之间的答案标记阻止了在这种规模下形成检索回路；训练序列携带八个查询；评估总是单查询，是训练格式的精确前缀（附录B (https：//arxiv.org/html/2607.21692#A2)；Arora et al.，2023 (https：//arxiv.org/html/2607.21692#bib.bib4)）。

### 2.2 教师、干预和合理性检查

教师是仅解码器Transformer，6层、8头、模型宽度384、10.7M参数、旋转位置嵌入，使用AdamW在任务混合上训练至高准确率（跨种子每任务0.90到0.99；最终数据见附录G (https：//arxiv.org/html/2607.21692#A7)）。所有教师随后被冻结。训练期间不施加稀疏性，因此监督标签不能被吸收到权重中（Aquino-Michaels，2026 (https：//arxiv.org/html/2607.21692#bib.bib3)）。

屏蔽一个块会向其每个层上的位置注意力对数添加 \(-10^4\)，对块外的每个查询使用与因果掩码相同的填充值；标记位置不变，因此干预移除的是信息，而非位置结构。限制到一个集意味着屏蔽其补集。两个方向都进行了检查：对于解决了的示例，仅保留标注的充分集可保持答案；在唯一检索示例上，屏蔽标注块会改变答案的概率为98%，而屏蔽一个随机其他块改变答案的概率为0%。屏蔽一个块并用填充符物理覆盖其记录，在200个检查的检索示例中100%产生相同的答案。在唯一检索上，答案位置注意力最高的块在99%的解决示例上是标注的块，因此至少在一个块重要且没有干扰项共享键的情况下，注意力和充分性是一致的。下面的分歧不是读取机制故障的产物。

### 2.3 标签

三个对象需要单独命名，因为我们的结果显示它们会分离。*任务证据* 是通过构造决定答案的标注记录。一个块是 *必要的*，如果屏蔽它会改变全上下文答案。一个集合是 *充分的*，如果限制到该集合能保留答案，如上定义。在解决了的单记录示例上，这三者一致。冗余破坏了必要性（第4节 (https：//arxiv.org/html/2607.21692#S4)），而限制引起的回退可能会将充分集缩小到任务证据之下（附录D (https：//arxiv.org/html/2607.21692#A4)），因此每次测量都会说明它使用哪个对象。对于每个示例和每个冻结教师，我们计算三种块标签。注意力标签采纳答案位置注意力质量最大的 \(k\) 个块，在层上池化，在头上取最大值或平均值，加上一个自适应变体，该变体取覆盖上下文块中该质量90%的最小集合。消融排序采纳 \(k\) 个块，其中单独屏蔽每个块最能降低答案对数概率。我们将 *因果证据集* 定义为两种情况。如果某个单独的块本身就能保留答案，则该集合是每个能做到这一点的块，从而在冗余下恢复所有副本。如果没有单个块充分，则该集合是每个屏蔽后改变答案的块，从而在链下恢复必要的支持。这个名称是中性的，因为成员在一个场景中是充分的替代品，在另一个场景中是必要的块。该集合没有固定大小。干预决定每个示例包含多少个块。因此，与注意力的比较使用自适应注意力变体，该变体也选择自己的大小。

## 3. 注意力包含不重要的证据

在最近写入任务中，一个键被多次写入，只有最后一次写入是正确的，自适应注意力集在100%的示例上包含至少一个过时的写入，对于五个种子中的每一个都是如此。这种行为是系统性的，而非噪音。注意力检索所有匹配查询键的内容，并将仲裁留给后续计算。在固定的 \(k=3\) 预算下，这里的注意力标签在种子间的一致性大约与因果证据标签一样好（Jaccard 0.76）。稳定性本身并不能证明一个目标，因为其稳定集包含不重要的证据。

注意力也很分散。取决于种子，覆盖90%答案位置质量需要32个块中的13到15个；因果证据集在单记录任务上平均为1.4到1.6个块。一个从注意力蒸馏出来的选择器被教导的目标比证据所需的大九倍。

## 4. 因果证据集是稳定的

表2 (https：//arxiv.org/html/2607.21692#S4.T2) 报告了跨种子标签一致性。两个规律性突出。首先，在单记录任务上，因果证据集的一致性大约是大小匹配注意力的两倍：不同的训练运行以不同的注意力和相同的证据解决了这些任务。其次，冗余区分了估计器。种植 \(m\) 个副本时，屏蔽任一单个副本不再改变答案：单个因果块的翻转率从一个副本时的0.91下降到两个副本时的0.06和四个副本时的0.001，消融排序标签也随之下降（0.25到0.11）。族标签则反向移动，从0.81到0.90，因为所有副本单独都是充分的并且都被恢复。当多个答案同样好时，模型使用哪一个不可识别，但可用答案的集合是可识别的，并且它是稳定的对象。Reco