论文

突破熵界:通过带拒绝采样的 MTP 加速 RL 训练

突破熵界:通过带拒绝采样的 MTP 加速 RL 训练

强化学习(RL)已成为现代大语言模型的关键组成部分,但 rollout 阶段仍然是 RL 训练流程的主要瓶颈。虽然多令牌预测(MTP)通过投机解码自然提供了加速 rollout 的解决方案,但许多研究发现 MTP 接受率在 RL 训练期间显著下降,导致加速性能有限。 为了解决这一瓶颈,我们提出了 Bebop,对 LLM 后训练中的 MTP 进行系统性研究,并提供将 MTP 集成到大规模 RL 流水线中的实用方案。首先,我们揭示了 MTP 接受率根本上受模型熵波动的限制,该波动与 RL 阶段熵的上升呈明显的负线性关系。其次,我们表明与贪婪草稿采样相比,概率拒绝采样在很大程度上减轻了 RL 中熵引入的干扰。我们进一步发现传统的 MTP 训练目标(交叉熵或 KL)在此设置下是次优的,因此我们提出了一种新颖的端到端 TV 损失,直接优化多步拒绝采样接受率,在数学推理、代码生成和代理任务上实现了约 10% 的接受率提升,最高达到 95% 的接受率和 25% 的额外推理吞吐量增益。第三,我们测试了 RL 期间的各种在线 MTP 训练策略,并表明带有端到端 TV 损失和拒绝采样的预 RL MTP 训练在整个 RL 过程中保持一致的接受率和加速,无需昂贵的在线 MTP 更新。 我们提供了大量实验和分析来验证我们的发现。实验结果表明,我们的方法在 Qwen3.5、Qwen3.6 和 Qwen3.7 模型的异步 RL 训练中实现了高达 1.8 倍的端到端加速。

论文精读

TL;DR 揭示 MTP 接受率随 RL 熵增线性下降的规律,通过拒绝采样与端到端 TV 损失直接优化多步接受率,在 RL 全程保持高加速比,实现高达 1.8 倍异步训练提速。

问题

问题背景

大语言模型(LLM)的后训练阶段,特别是强化学习(RL)微调,已成为对齐与能力提升的关键。但RL训练的 rollout 阶段需生成大量样本,通常采用自回归逐 token 解码,耗时极大,成为整个训练流程的吞吐瓶颈。

现有方法局限

多token预测(MTP) 结合推测解码是实现加速的常用方案:草稿模型一次预测多个未来token,由目标模型并行验证。然而,现有研究发现,在RL训练过程中,MTP的接受率会严重退化,导致加速效果大打折扣。传统MTP训练使用交叉熵(CE)或KL散度作为损失,且验证阶段多采用贪心采样。这种方式存在两个根本局限:

  • 熵增敏感:RL训练使模型熵升高,而贪心采样下的接受率与模型熵呈近似负线性关系,熵增导致接受率快速下降。
  • 训练目标与验证脱节:CE/KL损失未直接优化接收率,尤其在概率拒绝采样验证下,会造成分布失配,使草稿分布过于均匀或错误集中,进一步压低接受率。

为什么这个问题难/重要

挑战在于:RL训练中模型分布持续变化,MTP模型需要动态适应以维持高接受率;同时,高熵环境中验证步骤的误差放大效应更显著,传统逐token损失难以捕获多步联合接受概率。业界亟需一种能在整个RL周期内稳定保持高加速比的MTP方案,以降低大模型训练成本、缩短实验迭代时间。

行业类比

类似自动驾驶实时规划中,预测模型需在车辆动态行为分布不断变化时保持高准确率,否则累积误差会导致控制失效。MTP在RL训练中的加速同样需要鲁棒的分布匹配,以应对模型策略的持续漂移。

核心洞察

  • RL 阶段模型熵的上升给多令牌预测(MTP)的接受率带来了明确的瓶颈:**接受率与熵成反线性关系**,这解释了为何许多现有方法在高熵 RL 训练中失效。
  • 传统的交叉熵或 KL 散度训练目标在拒绝采样下是次优的,而提出的 **端到端 TV 损失直接优化多步拒绝采样的接受率**,从梯度分析上证明其能更精准地缩小草稿分布与目标分布的 mismatch,在数学推理、代码生成和智能体任务中实现高达 95% 的接受率和最高 25% 的额外推理吞吐量。

方法

输入与背景

针对 LLM 强化学习(RL)训练的 rollout 瓶颈,Bebop 以**多 token 预测(MTP)**作为推测解码的加速手段。输入为 RL 训练中需要大量采样的生成序列,MTP 通过一次前向预测多个 token 的 draft 分布,再经目标模型验证来加速推理。然而,随着 RL 训练推进,模型的 熵(entropy) 逐渐升高,导致 MTP 的接受率显著下降,加速效果受限。

关键模块

熵约束分析:Bebop 首先从理论上揭示了 MTP 接受率与目标模型熵之间的负线性关系。RL 中模型分布愈发不确定(熵增),greedy draft sampling 下的分布错配急剧放大,接受率自然塌缩。

拒绝采样增强:为缓解熵的干扰,Bebop 引入**概率拒绝采样(rejection sampling)**代替常用的 target-only(greedy)验证。在拒绝采样下,draft token 按概率被接受或回退至目标分布,使得分布错配的影响被平滑化,接受率随熵的退化更平缓。

TV 损失直接优化接受率:传统 MTP 训练使用交叉熵(CE)或 KL 散度作为目标,但这些损失在 RL 场景下对提高拒绝采样接受率并非最优。Bebop 提出端到端 TV 损失,直接最大化多步拒绝采样的期望接受率。TV 损失的梯度能促使 draft 分布更尖锐地贴合目标分布的高概率区域,且自身具有有界性,训练更稳定。相比 CE/KL 造成均匀或无序的分布错配,TV 训练引致的概率比例错配(probability-proportional mismatch)更有利于高接受率。

预 RL 适应策略:实验表明,在 RL 开始前使用 TV 损失与拒绝采样完成 MTP 训练(pre-RL adaptation),即可在整个 RL 过程中维持高接受率,无需昂贵的在线 MTP 更新。此外,交替冻结 backbone 与 MTP 的交叉训练可进一步提升模型质量。

输出

最终产出在数学推理、代码生成、智能体任务上均能保持高达 95% 接受率的 MTP 加速方案,在 Qwen 系列模型的异步 RL 训练中实现最高 1.8 倍端到端加速。

差异点

Bebop 区别于仅用 CE/KL 损失训练 MTP 并依赖 greedy 验证的方案,其通过熵感知的拒绝采样和直接优化接受率的 TV 损失,系统性解决了 RL 训练中 MTP 退化问题。

实验

实验设计

论文在 Qwen3.5 / 3.6 / 3.7 三个模型规模上,针对数学推理、代码生成与智能体三类任务进行 RL 训练,系统评估不同 MTP 训练目标(CE / KL / TV loss)与接受采样方法(目标采样 / 拒绝采样)对加速效果的影响。实验测量多步 MTP 的接受率、推理吞吐量与端到端训练加速比,并对比在线与离线 MTP 更新策略。

关键发现

  1. MTP 接受率与模型熵呈 负线性关系,RL 阶段熵上升导致接受率退化。
  2. 拒绝采样 较目标采样受熵扰动更小,接受率更高且更稳定。
  3. 提出的 TV loss 直接优化多步拒绝采样接受率,相比 CE/KL loss 提升约 10% 绝对接受率,最高达 95%,吞吐量增益 25%
  4. 预 RL 阶段用 TV loss + 拒绝采样训练 MTP,即可在整个 RL 过程中维持稳定的接受率与加速比,无需昂贵的在线 MTP 更新。
  5. 在异步 RL 训练中取得最高 1.8 倍 端到端加速。

与基线对比深度解读

基线方案使用 CE 或 KL 训练 + 贪婪目标采样的 MTP,在 RL 后接受率急剧下滑,加速效果有限。Bebop 的 TV loss 训练 + 拒绝采样 从根本上改变了接受率-熵的退化曲线:

  • TV loss 使草稿分布更尖锐,概率质量更集中,对策略偏移具有自纠正能力;
  • 拒绝采样引入的长度无偏校正边界,抑制了高熵场景下的错误累积。
  • 因此,即使在 RL 策略剧烈变化时,接受率仍保持高位,无需在线微调,降低了整体训练成本,使多步预测成为大型 RL 流程中实用的加速手段。

行业影响

落地场景

Bebop 提供的加速方案直指大规模语言模型后训练的异步 RL 管线,可用于需要高频迭代的智能聊天机器人、代码生成助手、数学推理引擎以及多步 agent 任务。当模型需要通过 RL 对齐用户反馈(如 RLHF)或自我提升(如 RLAIF)时,rollout 阶段的推理解码耗时占主导,Bebop 通过优化 MTP 接受率将端到端训练加速至 1.8 倍,直接降低每次迭代的时间成本和算力消耗。

商业价值

在降本方面,训练加速意味着同样的推理预算可以支撑更多 RL 迭代,或更快完成模型版本迭代,显著减少 GPU 集群占用时间。在一次典型的大规模 RL 训练中(例如对 70B 模型进行数学推理对齐),三天内可节省数千 GPU 小时,直接转化为基础设施成本的下调。同时,更高吞吐的推理能力让线上服务的响应延迟更低,提升用户体验。对于提供 fine-tuning 平台的云服务商,可将 Bebop 作为一种加速选项嵌入其训练管线,形成差异化竞争力。

与现有工作流的接口

Bebop 的改造集中在 MTP 模块的训练损失函数和采样策略,可无缝接入已有异步 RL 框架(如 OpenRLHF、Ray)。具体集成方式是:

  1. 用预训练的 MTP draft 模型初始化,并将其冻结,无需在线更新权重,省去传统方案中昂贵的 MTP 协同训练开销。
  2. 将标准交叉熵/KL 损失替换为端到端 TV 损失,直接优化多步拒绝采样的接受率。
  3. 在 rollout 阶段采用概率拒绝采样替代贪心采样,使接受率对 RL 阶段的模型熵增长保持稳健。

真实场景举例

  • 代码助手平台:某编程工具需要根据用户反馈持续微调代码生成模型,RL 训练常需在数小时内完成。使用 Bebop 可将单次训练时间从 8 小时压缩至约 4.5 小时,每日可进行多轮迭代,加速 model drift 的适配。
  • 企业级 agent 系统:一个多步客服 agent 需要根据人工评估结果进行 RL 微调,rollout 链长且易出现“气泡”退化。Bebop 的 TV 损失与拒绝采样可稳定长链生成的接受率,维持 95% 以上,保障 agent 在线训练的经济性。

局限

  • TV损失计算中采用的Top-K截断近似可能导致训练不稳定,超参数K的选择对最终接受率影响较大,论文指出这种方法在对抗性训练场景下易出现梯度消失或爆炸,限制了方法的鲁棒性。
  • 方法假设在RL训练前完成MTP模块的预训练(pre-RL adaptation),尽管实验表明此策略足以应对策略更新带来的分布偏移,但当RL训练跨度长、模型能力变化剧烈时,预训练模块可能逐渐失效,而集成在线MTP更新又将引入高昂计算开销,此矛盾未完全解决。
  • 实验主要基于Qwen3系列模型和数学推理、代码生成、智能体任务,未覆盖更大规模模型(如>100B参数)及更开放的任务类型(如创意写作),在异构模型架构(如非Transformer)上的适用性也缺乏验证。
论文Yucheng Li2026-06-10原文

相关内容