ATTNPO: 用于高效推理的注意力引导过程监督

# AttnPO: 用于高效推理的注意力引导过程监督

来源：https://arxiv.org/html/2602.09953

Shuaiyi Nie1,2, Siyu Ding3‡, Wenyuan Zhang1,2, Linhao Yu4, Tianmeng Yang3, Yao Chen1,2, Weichong Yin3, Yu Sun3, Hua Wu3, Tingwen Liu1,2†

1信息工程研究所，中国科学院
2网络安全学院，中国科学院大学
3百度公司
4天津大学

{nieshuaiyi, liutingwen}@iie.ac.cn

###### 摘要

使用强化学习和可验证奖励(RLVR)训练的大规模推理模型在复杂推理任务上表现出色，但常常过度思考，产生冗余推理而不能带来性能提升。现有的轨迹级长度惩罚通常无法有效缩短推理长度且会降低准确率，因为它们统一对待所有推理步骤，缺乏细粒度信号来区分冗余与必要性。与此同时，过程监督方法通常计算开销大，且面临不准确的信用分配问题。为了解决这些问题，我们提出了 AttnPO，一个低开销的过程监督强化学习框架，利用模型内在的注意力信号进行步级信用分配。我们首先识别一组特殊的注意力头，它们自然地聚焦于关键步骤，同时抑制冗余步骤。通过利用这些头的注意力分数，我们采用两个子策略来通过阻止冗余步骤的产生来缓解过度思考，同时通过减少对关键步骤的惩罚来保留准确率。实验结果表明 AttnPO 在 9 个基准上大幅减少推理长度并显著改进性能¹¹代码链接：https://github.com/NieSYsc20/AttnPO。

## 1 引言

大规模推理模型(LRM)的最近进展已在复杂推理任务上展现出非凡的性能。这一成功主要由结果监督的强化学习方法驱动，以 GRPO 为代表，它通过使用可验证的结果奖励来估计群体优势，从而简化了 PPO 中的批评模型。LRM 的一个特点是能够生成采用自我反思和多路径探索等高级策略的深思熟虑的长思维链(CoT)。然而，长 CoT 会导致过度思考：LRM 无差别地应用冗长推理，在琐碎操作上浪费计算。

在结果监督强化学习中集成长度惩罚(图 1(a))被广泛采用来缓解过度思考。核心思想是给既正确又简洁的回应分配更高的奖励。然而，这种粗粒度反馈对所有推理步骤分配统一的信用，无法提供区分必要和冗余步骤的判别学习信号。

因此，最近的研究聚焦于过程监督来降低分配给冗余步骤的信用权重。如图 1 所示，基于样本的方法(图 1(b))通过蒙特卡洛采样最终解决方案来估计中间步骤的边际性能增益，而基于模型的方法(图 1(c))训练一个奖励模型来确定第一个正确答案的位置，仅将后续步骤视为冗余。然而，这些方法面临两个关键限制：(1)高开销，需要额外的采样或奖励模型；(2)不准确的信用分配，无法提供独立的步级信号，可能导致错误的信用分配。

因此出现了一个关键问题：我们能否以几乎零的额外资源成本，仅依赖模型的内在信号来实现细粒度的步级监督，以区分哪些步骤是关键的或冗余的？

在本研究中，我们深入研究了模型的内在注意力机制，并发现了一个显著的现象：在最终答案生成期间(在<|answer|>标记之后)，一个注意力头的子集自然地聚焦于关键步骤，通过向它们分配更高的注意力权重来抑制冗余步骤。我们称这些为**关键焦点头(KFH)**。基于这一发现，我们的实验进一步揭示：(1)位置分布，KFH 主要位于中间到后期的层；(2)训练稳定性，KFH 的行为在带有长度惩罚的强化学习训练下变化甚微；(3)稀疏充分性，少量 KFH 足以进行区分，随着更多头的添加，性能饱和迅速。

基于这些见解，我们提出了 **AttnPO**(注意力引导策略优化)，一个新颖的强化学习框架，通过利用 KFH 的注意力分数实现步级监督，以可忽略不计的训练开销实现优异性能。具体来说，AttnPO 通过两个策略对正确的响应进行逐步缩放结果级优势：(1)对于具有正结果优势的正确响应，AttnPO 衰减分配给冗余步骤的正优势，防止过度鼓励它们并缓解过度思考；(2)对于具有负结果优势的正确响应，AttnPO 衰减分配给关键步骤的负优势，避免对有效推理的过度惩罚并缓解性能下降。

我们的贡献有三个方面：(1)我们首次揭示了 LRM 中 KFH 的存在，它们自然地聚焦于关键推理步骤并抑制冗余步骤；(2)基于这一见解，我们提出了 AttnPO，一个低开销的过程监督强化学习框架，利用内在注意力信号进行步级信用分配；(3)实证结果表明 AttnPO 显著缩短推理长度的同时改进准确率。例如，在 DeepSeek-R1-Distill-Qwen-1.5B 上，AttnPO 在六个数学基准上平均获得 +7.3 个百分点的准确率提升，同时推理长度减少 60%。

## 2 预备知识

##### LRM 的两阶段生成

给定输入问题 q，LRM M 生成响应 o=(T,F)，分为两个由特殊<|answer|>标记显式分隔的阶段。首先，模型产生深思熟虑的思考过程 T，其中可能包含冗余推理结构，如不必要的反思或验证。然后，它通过总结 T 并仅保留解决问题所需的关键步骤来输出简洁的最终解决方案 F。

##### 步骤分割

思考过程 T 可以分解为有意义的中间步骤，其边界通常由特定的特殊短语标记，这些短语捕捉推理中的关键点，如暂停、反思或方向转变。因此我们使用特殊短语来分割思考过程，这些短语分为三类：

- **困惑短语**(例如"等等")表示模型进入反思、验证或分支阶段
- **进展短语**(例如"首先")表示当前推理的继续
- **总结短语**(例如"因此")总结前面的推理

为了避免步骤过短，我们自适应地合并短于阈值的步骤。在分割和合并期间，我们将以困惑短语开头的步骤从前面的上下文中分离出来，因为它们经常触发冗余的验证或反思。分离它们使得这种冗余更容易被检测和惩罚。详见附录 A.1。

## 3 关键焦点头

遵循第 2 节描述的两阶段生成范式，LRM 将通常包含冗余推理的深思熟虑的思考过程转变为简洁的最终解决方案，这需要进行非平凡的选择和汇聚关键信息。在 Transformer 架构中，注意力是信息选择的主要机制，前面的工作表明不同的注意力头专门化于不同的功能。基于这一观察，我们假设**在最终答案生成期间，存在一个注意力头的特定子集，它们有选择性地聚焦于关键推理步骤，同时忽视冗余步骤，从而汇总和整合思考过程**。

为了测试这个假设，我们检查了不同注意力头在关键和冗余步骤上的注意力分布。

### 3.1 探测设置

##### 探测数据构建

为了分析不同的注意力头如何注意关键步骤与冗余步骤，我们需要一个带有可靠步级注解的探测数据集。因此我们从 DeepScaler 中采样 300 个问题，这些问题对 DeepSeek-R1-Distill-Qwen-1.5B 来说不困难(在八个样本上平均准确率 ≥ 0.875)。这个限制很重要，因为在更难的问题上，LRM 通常表现出探索性推理。这种推理可能看起来冗余，但仍然有助于正确的解决方案，使得可靠的步骤注解困难，从而影响头部识别。

对于每个问题，我们采样八个响应，选择中等长度的正确响应，并按照第 A.1 节分割其思考过程。我们然后使用三个最先进的 LLM 对每个推理步骤进行注解，仅保留三个模型都同意标签的步骤，从而减少注解噪声。因此，对于每个思考过程 T，我们获得一组关键步骤 E_s 和一组冗余步骤 R_s。详情见 A.2.1。

##### 基于注意力的步骤分数

对于每层 l 的每个头 h，我们为每个推理步骤 s_k 计算基于注意力的步骤分数 S^{l,h}_{s_k}，其中 k 索引步骤索引，基于在最终答案生成期间使用的注意力权重。具体来说，令 a^{l,h}_{m→n} 表示最终解决方案 F 中标记 m 到步骤 s_k 中标记 n 的注意力权重。我们定义步骤分数 S^{l,h}_{s_k} 为：

$$\mathcal{S}^{l,h}_{s_k} = \frac{1}{|s_k|}\sum_{m \in \mathcal{F}}\sum_{n \in s_k} a^{l,h}_{m \to n}$$

这个分数衡量步骤 s_k 中每个标记从最终解决方案接收的平均注意力。

##### 评估指标：步骤排序准确率

我们评估每个注意力头区分关键和冗余步骤的能力有多好。对于每个思考过程 T，我们形成所有关键-冗余步骤对 P_T = {(s^+, s^-) | s^+ ∈ E_s, s^- ∈ R_s}。对于层 l 的头 h，我们定义步骤排序准确率(SRA)为：

$$\mathrm{SRA}_{l,h} = \frac{|\{(s^+, s^-) \in \mathcal{P}_\mathcal{T} | \mathcal{S}^{l,h}_{s^+} > \mathcal{S}^{l,h}_{s^-}\}|}{|\mathcal{P}_\mathcal{T}|}$$

SRA 衡量关键步骤接收比冗余步骤更高分数的步骤对的比例。我们通过在探测数据集中的所有示例上平均来报告最终 SRA。

### 3.2 探测分析

我们在 DeepSeek-R1-Distill-Qwen-1.5B/7B 上进行实验，做出以下观察；更详细的结果在附录 A.2.2：

- **在最终解决方案生成期间，某些注意力头有选择性地聚焦于关键推理步骤同时忽视冗余步骤，主要位于中间到后期的层。** 我们将这些专门的头称为关键焦点头。如图 2(a,b) 所示，在两个评估的 LRM 上，少量头达到 SRA 超过 0.9，最好的头在 1.5B 模型上达到 95%，在 7B 模型上达到 96%。相比之下，大多数头的性能在随机基线处或以下(SRA ≈ 0.5)。

- **从多个注意力头集成步骤分数 S_{s_k} 仅提供有限的收益。** 在图 2(c,d) 中，在 SRA 的 top-k 选择和贪心 SRA 最大化选择下，当组合少量头时性能略有改进，但随着添加更多头迅速饱和或下降。

- **在强化学习训练下带有长度惩罚，KFH 的行为变化甚微，表明它们的功能角色是稳健的。** 我们使用式 5 中描述的奖励函数训练模型，并在步骤 500 和 1000 处的检查点重新运行相同的探测过程。由此产生的 SRA 分布在检查点间显示高度一致，在所有情况下皮尔逊相关系数超过 0.85。

- **即使在探测中使用非困难问题，KFH 仍表现出对具有挑战性问题的一定程度的泛化。** 我们在具有挑战性的 AIME24 数据集上采样响应，并从 SRA 最高的头获得步骤分数。对于每个具有正确最终答案的思考过程 T，我们移除底部 30%(低分)或顶部 30%(高分)的步骤，然后附加提示"我认为我已经完成了思考。现在逐步给出最终解决方案。"来生成最终解决方案。如图 2(e,f) 所示，移除高分...