并非所有Token都同等重要：通过强化学习中的Token重要性实现高效LLM推理

# 并非所有令牌都很重要：通过强化学习中的令牌显著性实现高效LLM推理

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

Hanbing Liu1 Lang Cao2[1] Yuanyi Ren3[1] Mengyu Zhou4 Haoyu Dong5 Xiaojun Ma5 Shi Han5 Dongmei Zhang5

1清华大学
2伊利诺伊大学厄巴纳-香槟分校
3北京大学
4阿里巴巴通义大模型应用团队
5微软

[email protected], [email protected], [email protected]

微软实习期间完成的工作。
隶属于清华大学深圳国际研究生院泛在数据使能深圳市重点实验室。
通讯作者。[email protected]。

## 摘要

大型语言模型（LLM）展现出强大的推理能力，但经常产生不必要的冗长解释，降低效率。尽管强化学习（RL）已被用于改进推理，但大多数方法专注于准确性，并依赖于忽视个别令牌不同贡献的统一长度奖励，通常会损害正确性。我们通过令牌显著性的视角重新审视RL中的长度优化。观察到许多思维链（CoT）令牌对最终答案的贡献很小，我们引入了显著性感知的长度奖励，选择性地惩罚不重要的令牌，减少冗余同时保留关键推理。我们还提出了一种动态长度奖励，鼓励在训练早期进行更详细的推理，并随着学习进行逐步转向简洁性。将这些组件集成到标准策略优化中得到了一个框架，既提高了推理效率，也提高了准确性。在多个基准上的实验展示了响应长度的大幅减少，同时保留或改进了正确性，突出了为高效LLM推理建模令牌显著性的重要性。代码已发布在https://github.com/microsoft/Bingo。

## 1 引言

大型语言模型（LLM）在从算术问题求解到常识推理的各种任务中展现了令人印象深刻的推理能力。最近工作的一个关键观察是，足够大的模型可以展现出新兴的推理能力，例如思维链（CoT）推理，无需显式监督。尽管取得了这些成功，但一个主要挑战仍然存在：LLM通常生成不必要的冗长或冗余推理轨迹，导致计算成本、冗余和延迟方面的低效。因此，改进LLM的推理效率已成为一个重要的研究方向。

该领域的先前工作可以大致分为两类：监督微调（SFT）方法和强化学习（RL）方法。基于SFT的方法专注于构建压缩的推理轨迹并训练模型模仿它们。虽然这些方法可能有效，但它们依赖于高质量的压缩监督，这很难获得，并且通常缺乏跨不同任务的可泛化性。基于RL的方法通常引入基于长度的奖励，惩罚过长的响应以鼓励简洁性。然而，RL方法中此类奖励或惩罚的设计仍未被充分探索，通常过于简化。例如，O1-Pruner对所有样本应用统一惩罚，假设每个响应都应该缩短。这个假设经常导致性能下降，因为并非所有推理轨迹同样冗长——有些需要更详细的步骤才能得出正确答案。

为了解决这个问题，其他工作提出了更选择性的惩罚策略，将惩罚条件与样本正确性或估计难度联系起来。这些方法通常对简单问题分配更强的惩罚，对更具挑战的情况分配较弱的惩罚。然而，准确估计问题难度仍然是一个基本挑战，未解决的困难问题通常导致不必要的冗长响应，进一步破坏推理效率。

尽管兴趣不断增长，当前长度奖励的设计仍然有限，因为它们通常无法充分促进简洁推理同时保留答案准确性。例如，先前工作在很大程度上忽视了令牌级别贡献对推理整体效率的影响。

在这项工作中，我们从令牌显著性这一新颖视角来处理这个问题。我们的动机来自于LLM中观察到的令牌冗余，其中思维链（CoT）推理中的许多令牌对最终答案的贡献很小。我们主张并非所有令牌对高效推理都同样重要——许多是不重要的，如冗余短语或不必要的中间步骤，可以删除而不会降低性能。现有奖励设计通常忽视了这种区别。相反，我们引入了显著性感知的长度奖励，选择性地仅惩罚那些对最终答案没有有意义贡献的不重要令牌，同时保留关键推理步骤。

我们也观察到有效处理困难问题对高效推理至关重要。先前工作表明鼓励扩展的CoT推理可以通过启用更深的探索来改进性能，这可能有助于解决更难的问题。因此，直观地使用长度作为困难问题的激励是合理的。然而，LLM应该不仅准确而且简洁地解决困难问题。应用静态长度激励可能导致不必要的冗长响应，这仍可能无法产生正确答案。为了解决这个问题，我们整合了在训练过程中自适应的动态长度奖励。此奖励应用于错误样本中的重要令牌，以平衡探索和效率。具体来说，它在早期训练阶段鼓励更长的推理以促进探索，并逐步在后期阶段转向惩罚过度长度以促进简洁性。

基于这些见解，我们引入了**Bingo**（**B**oosting **E**fficient **R**eason**IN**G in Policy Optimization，策略优化中的高效推理增强），一个RL框架，将我们提出的两种奖励机制整合到标准RL算法（如近端策略优化（PPO））中。这使得推理准确性和效率的联合优化成为可能。在多样化推理基准上的广泛实验表明，Bingo通过减少冗余计算同时保持或改进准确性，一致地优于强基线。

总结来说，本文作出了以下关键贡献：

- **令牌显著性洞察**。我们在策略优化中引入令牌显著性的概念，区分推理轨迹中的重要和不重要令牌。这一洞察激发了我们的显著性感知长度奖励，它显式地惩罚不信息性令牌同时保留关键推理内容，实现更有针对性和有效的长度控制。

- **动态长度控制**。我们提出了一种动态长度奖励策略，在训练过程中调整奖励信号——在早期阶段鼓励更长的推理以促进探索，并随着模型收敛逐步促进简洁性。

- **效率导向的RL框架**。我们开发了Bingo，一个新的强化学习框架，集成了两种奖励策略。在多个推理基准上的广泛实验以及全面的分析展示了其有效性。

## 2 相关工作

**大型语言模型的强化学习**。强化学习（RL）已成为将大型语言模型（LLM）与人类偏好对齐的强大范例。在来自人类反馈的强化学习（RLHF）中，近端策略优化（PPO）算法与人类偏好数据一起被用于训练奖励模型，该模型引导LLM的微调。基于PPO，随后的工作如GRPO和REINFORCE++提出了改进的变体以解决其局限性。除了对齐外，RL在改进LLM的推理能力方面也显示了前景。早期研究表明奖励引导的训练可以增强多步推理性能。最近，DeepSeek-R1表明大规模RL可以在广泛的任务中大幅提升推理能力，指向了一个有前景的未来工作方向。

**大型语言模型的高效推理**。最近的进展使语言模型能够通过推理时技术（如思维链提示和后训练）进行强大的推理。更近期的工作已转向优化准确性和效率。一些方法在推理时改进效率，如令牌预算感知推理或提示策略，如"无思考推理"和思维链草稿。其他方法通过监督微调（SFT）应用后训练优化，包括TokenSkip、TwT、LightThinker和C3oT。这些SFT方法主要构造包含关键信息的高质量压缩推理路径，并在其上训练模型。

同时，基于RL的方法通常通过将长度控制或惩罚融入奖励函数来改进效率。例如，O1-Pruner使用离线长度奖励与样本比较对比平均长度。Kimi k1.5仅对正确样本应用在线惩罚。基于先前基于RL的方法，我们推进长度奖励设计，使LLM能够平衡推理准确性与计算效率。

![图2：Bingo框架的示意图。给定生成的CoT轨迹，LLM首先区分重要和不重要令牌。然后根据令牌类型和样本正确性计算动态长度奖励。在训练的早期探索阶段（k(t)≥β），奖励鼓励扩展推理...]