AtManRL: Towards Faithful Reasoning via Differentiable Attention Saliency

📄 arXiv: 2604.16158v1 📥 PDF

作者: Max Henning Höth, Kristian Kersting, Björn Deiseroth, Letitia Parcalabescu

分类: cs.CL, cs.AI, cs.LG

发布日期: 2026-04-17

备注: 14 pages, 8 figures, 1 table


💡 一句话要点

AtManRL:通过可微注意力显著性实现语言模型的可信推理

🎯 匹配领域: 支柱一:机器人控制 (Robot Control) 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 可解释性AI 思维链推理 强化学习 注意力机制 显著性分析

📋 核心要点

  1. 现有大语言模型的CoT推理缺乏可信度,推理过程可能只是伴随答案,而非真正影响答案。
  2. AtManRL通过可微注意力操纵,学习识别CoT中关键tokens,并以此作为奖励信号,引导模型生成更具影响力的推理过程。
  3. 实验表明,AtManRL能有效识别关键推理tokens,并训练出更透明的推理模型,提升模型的可解释性。

📝 摘要(中文)

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

🔬 方法详解

问题定义:现有的大语言模型虽然能够进行思维链推理,但其推理过程的可信度难以保证。模型生成的推理步骤可能与最终答案没有实际的因果关系,仅仅是表面上的“解释”,缺乏真正的可解释性。因此,如何确保推理过程能够真正影响模型的最终预测,是当前面临的一个重要挑战。

核心思路:AtManRL的核心思路是通过强化学习来训练一个注意力掩码,该掩码能够识别思维链中对产生正确答案至关重要的tokens。通过操纵注意力权重,模型能够学习到哪些tokens对于推理过程是真正重要的,从而生成更具影响力的推理过程。这种方法鼓励模型关注关键信息,避免冗余或无关的推理步骤。

技术框架:AtManRL的整体框架基于GRPO(Gradient Ratio Policy Optimization)。首先,模型生成一个思维链推理过程。然后,通过一个可微的注意力掩码来调整注意力权重,突出显示对最终答案有重要影响的tokens。接着,根据模型的答案是否正确以及注意力掩码所识别的tokens的重要性,计算奖励信号。奖励信号包括基于结果的奖励(答案正确则奖励高)和基于显著性的奖励(关键tokens被关注则奖励高)。最后,使用GRPO算法来优化模型,使其既能生成正确的答案,又能生成可信的推理过程。

关键创新:AtManRL的关键创新在于引入了可微注意力操纵和显著性奖励信号。通过可微的注意力掩码,模型能够学习到哪些tokens对于推理过程是真正重要的。显著性奖励信号则鼓励模型关注这些关键tokens,从而生成更具影响力的推理过程。这种方法与传统的基于结果的强化学习方法不同,它不仅关注答案的正确性,还关注推理过程的可解释性。

关键设计:AtManRL的关键设计包括:1) 使用加性注意力掩码来调整注意力权重,该掩码可以学习到每个token的重要性;2) 设计显著性奖励信号,该信号基于注意力掩码所识别的tokens的重要性来计算;3) 将显著性奖励与基于结果的奖励集成到GRPO框架中,以共同优化正确性和可解释性。具体而言,显著性奖励的计算方式是基于注意力权重的梯度,梯度越大,表示该token对最终答案的影响越大,因此奖励越高。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,AtManRL在GSM8K和MMLU数据集上取得了显著的性能提升。具体来说,AtManRL能够有效地识别思维链中对产生正确答案至关重要的tokens,并且能够训练出更透明的推理模型。与基线模型相比,AtManRL在可解释性方面有显著提升,同时保持了较高的准确率。

🎯 应用场景

AtManRL具有广泛的应用前景,例如可以应用于教育领域,帮助学生理解解题思路;应用于金融领域,解释投资决策的依据;应用于医疗领域,辅助医生进行诊断和治疗。通过提高模型的可解释性,AtManRL能够增强人们对人工智能系统的信任,促进人工智能技术的广泛应用。

📄 摘要(原文)

Large language models (LLMs) increasingly rely on chain-of-thought (CoT) reasoning to solve complex tasks. Yet ensuring that the reasoning trace both contributes to and faithfully reflects the processes underlying the model's final answer, rather than merely accompanying it, remains challenging. We introduce AtManRL, a method that leverages differentiable attention manipulation to learn more faithful reasoning through reinforcement learning. By training an additive attention mask that identifies tokens in the CoT crucial for producing correct answers, we derive a saliency reward signal that encourages the model to generate reasoning traces that genuinely influence its final predictions. We integrate this saliency reward with outcome-based rewards within the GRPO framework to jointly optimize for correctness and interpretability. Experiments on GSM8K and MMLU with Llama-3.2-3B-Instruct demonstrate that our approach can identify influential reasoning tokens and enable training more transparent reasoning models.