AtManRL: 通过可微分注意力显著性实现忠实推理

# AtManRL：通过可微分注意力显著性实现忠实推理
来源：https://arxiv.org/html/2604.16158
Max Henning Höth Aleph Alpha Research Lab1141

&Kristian Kersting TU Darmstadt Hessian\.AI Lab1141

&Björn Deiseroth Aleph Alpha Research Lab1141

&Letitia Parcalabescu Aleph Alpha Research Lab1141

###### 摘要

大型语言模型(LLMs)越来越依赖思维链(CoT)推理来解决复杂任务。然而，确保推理轨迹既对模型最终答案有贡献，又能忠实反映答案背后的推理过程（而不仅仅是伴随答案出现）仍然具有挑战性。我们提出AtManRL，一种利用可微分注意力操纵通过强化学习学习更忠实推理的方法。通过训练一个加性注意力掩码来识别CoT中对产生正确答案至关重要的token，我们推导出一个显著性奖励信号，鼓励模型生成真正影响其最终预测的推理轨迹。我们将此显著性奖励与GRPO框架中基于结果的奖励相结合，以共同优化正确性和可解释性。在GSM8K和MMLU上使用Llama-3.2-3B-Instruct进行的实验表明，我们的方法可以识别有影响力的推理token，并训练出更透明的推理模型。

## 1 引言

思维链(CoT)提示(Wei et al., 2022 (https://arxiv.org/html/2604.16158#bib.bib6))、监督学习和强化学习(RL)方法(Yang et al., 2025 (https://arxiv.org/html/2604.16158#bib.bib23); OpenAI et al., 2024 (https://arxiv.org/html/2604.16158#bib.bib24); Guo et al., 2025 (https://arxiv.org/html/2604.16158#bib.bib25))通过引出推理轨迹，提升了大型语言模型(LLMs)的推理能力。通过在最终答案之前生成中间推理步骤，模型在复杂任务上往往能达到更高的准确性。

然而，推理轨迹的存在并不能保证模型实际利用它来得出答案。因此，一个核心问题是：*生成的CoT是否对模型的最终预测产生因果影响并具有解释力，还是仅仅作为一种风格化的附带产物？* 这个问题关乎*忠实性*的概念，即解释是否反映了模型的真实决策过程(Jacovi and Goldberg, 2020 (https://arxiv.org/html/2604.16158#bib.bib22))。不忠实的推理轨迹可能看起来*合理*且逻辑连贯，而模型却通过绕过所述推理的捷径得出正确答案(Agarwal et al., 2024 (https://arxiv.org/html/2604.16158#bib.bib29))。先前的工作表明，LLMs可以生成听起来合理的CoT解释，但这些解释与驱动其预测的机制并不一致(Turpin et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib26); Lanham et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib20); Barez et al., 2025 (https://arxiv.org/html/2604.16158#bib.bib27))。

为探究这一差距，我们区分了*显著性*和*忠实性*。我们将显著性定义为单个推理token对最终答案logits的可测量因果贡献。忠实性则有更高的要求，即推理轨迹必须准确反映产生答案的潜在推理过程。因此，显著性构成了忠实性的必要但不充分条件。确保推理轨迹的*显著性*（定义为推理token对最终预测的可测量影响）可以防止CoT退化为冗长但关联性弱的叙述。如果没有这样的约束，推理轨迹可能成为事后合理化解释，而非模型计算过程的可解释证据。

在此区分的指导下，我们提出AtManRL来强制执行推理轨迹的显著性，该方法使用强化学习显式训练模型生成显著的推理轨迹。我们的方法构建在AtMan(Deiseroth et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib5))之上，这是一种通过预定义掩码对注意力权重进行针对性修改的注意力操纵技术。先前的工作将AtMan用于事后可解释性，而我们则把注意力操纵掩码视为一个可学习的、可微分的对象。这使得我们能够：(i) *高效地*识别推理轨迹中哪些token对最终答案真正有影响力，(ii) 从这些贡献中推导出基于显著性的奖励信号，以及(iii) 将此信号融入强化学习，以鼓励生成显著的推理步骤，同时抑制无关或弱相关的解释性内容。

总体而言，我们的贡献如下：(1) 我们引入了一个*显著性奖励*，该奖励源自优化一个可微分注意力掩码，用以识别CoT中的显著token。(2) 我们将此显著性奖励与GRPO框架中基于结果的奖励相结合，以*共同优化正确性和推理质量*（就显著性而言）。(3) 我们在GSM8K和MMLU上使用Llama-3.2-3B-Instruct评估了我们的方法，结果表明我们可以在保持准确性的同时*减少无关推理*。

## 2 相关工作

**CoT / 推理轨迹**。CoT提示(Wei et al., 2022 (https://arxiv.org/html/2604.16158#bib.bib6))和诸如GRPO(Shao et al., 2024 (https://arxiv.org/html/2604.16158#bib.bib19))等RL方法鼓励LLMs生成推理轨迹。RL通过优化基于结果的奖励来提升推理性能。然而，结果奖励关注于答案正确性，并不强制要求推理轨迹对最终预测产生因果影响。相比之下，我们显式地奖励CoT token与答案之间的因果依赖关系。

**推理轨迹忠实性**。模型解释的忠实性已在可解释性研究中被广泛探讨。有工作证明并论证了CoT解释可能不忠实，模型有时会通过与其陈述逻辑相矛盾的推理得出正确答案(Turpin et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib26); Lanham et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib20); Parcalabescu and Frank, 2024 (https://arxiv.org/html/2604.16158#bib.bib28); Barez et al., 2025 (https://arxiv.org/html/2604.16158#bib.bib27))。过程奖励模型使用外部监督将奖励分配给中间推理步骤(Lightman et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib21))，从而提升了CoT的合理性。然而，合理性反映了与外部评估者的一致性，而非与模型内部计算的对齐。忠实性则反映了模型内因果产生答案的机制。因此，在我们的方法中，我们为每个样本学习一个注意力掩码，以验证每个token的因果影响力。

**关键推理Token**。有工作表明，单个CoT token（称为*关键token*）可能对LLM的输出产生巨大影响(Lin et al., 2025 (https://arxiv.org/html/2604.16158#bib.bib2))。Vassoyan et al. (2025 (https://arxiv.org/html/2604.16158#bib.bib8))鼓励对这些token进行探索，以提高RL微调效率。Yan et al. (2024 (https://arxiv.org/html/2604.16158#bib.bib7))干预注意力权重，以减轻对少样本示例中误导性token的过度依赖。与这些事后分析或操纵推理token的方法不同，我们使用可微分注意力操纵来学习token级别的显著性，并将其融入RL训练中。

**注意力操纵**。AtMan(Deiseroth et al., 2023 (https://arxiv.org/html/2604.16158#bib.bib5))引入了内存高效的注意力操纵技术用于Transformer可解释性，能够实现对单个token注意力的针对性抑制，以评估其影响力。我们将AtMan构建为一个可微分的注意力掩码，并通过SGD针对正确答案进行优化，以识别显著的推理token。

## 3 用于忠实推理的可微分注意力操纵

接下来，我们介绍我们的方法AtManRL，该方法通过将AtMan构建为可微分的注意力掩码来训练模型产生显著的推理轨迹。具体来说，我们将 (1) 回顾AtMan，(2) 描述我们如何学习该掩码，(3) 从优化后的掩码推导出一个显著性度量，最后 (4) 在训练中将显著性作为RL奖励整合进来。

### 3.1 背景：AtMan 注意力操纵

首先，我们回顾Deiseroth et al. (2023 (https://arxiv.org/html/2604.16158#bib.bib5))中引入的加性AtMan注意力操纵。在标准Transformer中，注意力输出计算为：O = softmax(H) · V，其中·表示矩阵乘法，预Softmax注意力分数由 H = Q K^T / √d 给出。这里，Q, K, V ∈ ℝ^{h×s×d} 分别表示查询、键和值张量，h为注意力头数，s为序列长度，d为头部维度。AtMan通过一个加性掩码 H^{AtMan} ∈ ℝ^{s×s} 操纵预Softmax分数 H：

H = Q · K^T / √d + H^{AtMan}   (1)

在Softmax之前应用掩码 H^{AtMan} 确保了Softmax之后产生的注意力分数仍然总和为1。此外，与XAI中的其他扰动方法（例如Shapley值）不同，这不会导致输入分布或位置嵌入的偏移，而是精细地操纵模型对每个单独token的注意力。掩码正值增加对选定token的注意力，而负值则抑制其影响。对于自回归模型，我们额外应用一个下三角因果掩码T，并计算 H_M = H ∘ T，其中∘表示Hadamard乘积。Deiseroth et al. (2023 (https://arxiv.org/html/2604.16158#bib.bib5))使用H^{AtMan}通过将一个固定的负值（视为超参数）分配给H^{AtMan}的对应列来抑制对单个token的注意力，从而分析每个独立token对LLM输出logits的影响。

### 3.2 训练 H^{AtMan} 掩码以测量显著性

由于掩码是加性地进入预Softmax注意力分数的，因此它保持完全可微分。我们将H^{AtMan}限制为仅作用于推理轨迹（CoT）内的token，而不修改对提示token或最终答案token的注意力。提示是固定的且不受模型控制，因此不构成奖励塑形的目标。相反，我们要求最终答案因果地依赖于推理轨迹。如果推理轨迹是显著的，扰动其注意力应会影响正确答案的概率。

我们将所有与CoT相关的掩码条目初始化为一个负常数 c = -0.4。此初始化统一抑制对推理token的注意力，并产生更平坦的Softmax后分布。从这种抑制状态出发，我们优化掩码以恢复正确答案的概率。

具体来说，为了训练掩码，我们在教师强制下最小化预测答案token y_{1:N}的logits的交叉熵损失，如图2 (https://arxiv.org/html/2604.16158#S3.F2)所示：

L_{mask} = -1/N ∑_{n=1}^{N} log P(y_n | c_{1:T}, y_{1:n-1}, H^{AtMan})   (2)

其中c_{1:T}表示CoT token。这意味着所有值为0的掩码会导致损失为0。在此阶段，掩码是唯一可训练的对象，以识别在抑制推理的情况下仍能保持答案似然的注意力配置。

我们在固定步数后停止优化掩码。我们将掩码除以初始化常数进行归一化：\hat{H}^{AtMan} = H^{AtMan} / c，并计算下三角（因果）区域上的平均归一化掩码值：

R_{Faithfulness}(a_i) = 1/|I_v| ∑_{w ∈ I_v} \hat{H}^{AtMan}_{w, v},   I_v = {w ∈ {1, ..., n} | w ≥ v}   (3)

这个量作为我们的显著性度量和rollout a_i的奖励。直观上，它衡量了推理token必须被重新启用多强的程度才能保持正确答案的概率。

参阅图注
图 1: 我们使用负值初始化加性注意力掩码 H^{AtMan} 以抑制对CoT token的注意力。然后我们优化该掩码200步以恢复正确答案概率。
参阅图注
图 2: 为了识别非显著token，我们为每个rollout优化掩码，目标是恢复原始标签的对数概率。

### 3.3 通过强化学习(RL)优化显著性

对于RL，我们将显著性奖励与标准的结果奖励相结合：

R_{Outcome}(a_i) = { 0, 如果 i=j; -1, 否则 }   (4)

其中j表示真实答案，a_i表示rollout预测。

因此，*总奖励*为：R_{total}(a_i) = R_{Outcome}(a_i) + R_{Faithfulness}(a_i)。遵循GRPO(Shao et al., 2024 (https://arxiv.org/html/2604.16158#bib.bib19))，我们计算组归一化奖励 \hat{R}_{total} = 1/N ∑_{i=1}^{N} R_{total}(a_i)，并定义优势 A(a_i) = R_{total}(a_i) - \hat{R}_{total}。然后我们使用裁剪后的GRPO目标更新策略：

L_{GRPO}(θ) = -1/N ∑_{i=1}^{N} min( (π_θ(a_i|q) / π_{θ_old}(a_i|q)) A(a_i), clip(π_θ(a_i|q) / π_{θ_old}(a_i|q), 1-ε, 1+ε) A(a_i) )   (5)

### 3.4 实现细节

**掩码优化**：我们使用AdamW (Loshchilov and Hutter, 2019 (https://arxiv.org/html/2604.16158#bib.bib15))，学习率为1e-3，betas为0.6和0.9999，权重衰减为0.05。我们训练200个梯度步来更新H^{AtMan}。

**值缩放**：在将H^{AtMan}加到注意力分数之前，我们将其限制在0的上限，以防止对每个token应用正值，因为我们只想检测非显著token。然后我们将其乘以因子10以加快收敛。

**RL训练细节**：我们使用GRPO微调模型8个epoch，每个查询对应8个rollout，最大生成长度为1024个token。对于每次更新，我们计算每批8个查询的显著性奖励。我们每批次执行两次梯度传递，使用大小为2的微批次重新计算策略的对数概率，用于裁剪后的RL目标。我们使用固定的1×10^{-6}学习率和ε=0.2的标准裁剪参数。我们在48个NVIDIA A100 GPU上进行了所有实验。

表 1：基线与AtManRL的对比