突破熵界限:通过MTP与拒绝采样加速RL训练 Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling
提出端到端TV损失与拒绝采样方案,将RL训练加速1.8倍
前置知识
Multi-Token Prediction (MTP)
MTP是一种推测解码技术,通过在主模型上添加轻量级draft heads来并行预测多个未来token。例如,gamma等于3的MTP会同时预测接下来的3个token,然后用主模型一次性验证,将原本需要3次前向传播的操作减少到1次。这是当前LLM推理加速的主流方案之一,已被DeepSeek、Qwen等主流模型采用。
MTP是本文的核心加速手段,理解其工作原理(draft-then-verify范式)和接受率机制是读懂本文的基础。
Rejection Sampling (拒绝采样)
拒绝采样是MTP验证阶段的一种策略。对于draft token,以概率min(1, p(y)/q(y))接受,其中p是目标模型分布,q是draft分布。接受率等于分布重叠sum_v min(p(v), q(v)),等于1减去Total Variation距离。与target-only sampling(贪婪采样)不同,拒绝采样利用完整的分布重叠来决定接受,对熵的变化更鲁棒。
本文的核心发现之一就是拒绝采样相比传统的target-only sampling能够大幅减轻熵对接受率的影响,这是Bebop方案的关键组件。
Total Variation Distance (TV距离)
TV距离是衡量两个概率分布差异的指标,定义为二分之一乘以对所有token v的绝对差求和。在推测解码中,拒绝采样的接受率恰好等于1减去TV距离,因此直接优化TV距离等价于优化接受率。本文提出的TV损失就是基于这一观察,直接最小化draft和target分布之间的TV距离。
TV损失是本文的核心技术创新,它直接优化决定接受率的量,而不是像CE/KL损失那样间接优化。
GRPO (Group Relative Policy Optimization)
GRPO是一种RL算法,从当前策略采样一组G个轨迹,然后优化裁剪代理目标。其中r是重要性采样比,A是组归一化优势。这是现代LLM RL训练的主流算法之一,通过组内归一化降低方差,提高训练稳定性。
本文的实验主要在GRPO框架下进行,理解RL训练流程(rollout到reward到update)和熵在探索中的作用有助于理解为什么RL阶段熵会波动。
Entropy (熵)
熵H(p)等于对所有token v的p(v)乘以log p(v)求和并取负,衡量分布的不确定性。低熵表示模型预测很自信(分布集中在少数token上),高熵表示模型预测很不确定(概率分散在许多token上)。在RL训练中,为了鼓励探索,策略模型通常保持较高的熵,甚至熵会逐渐上升。本文的核心发现是MTP接受率与熵之间存在负线性关系,约等于a减去b乘以H(p)。
熵是本文分析MTP接受率下降的根本原因。理解熵如何影响接受率是理解Bebop方案为何有效的关键。
研究动机
现有的MTP加速方案在RL训练中接受率会显著下降。具体来说,在SFT阶段MTP接受率可以稳定在百分之九十以上,但进入RL阶段后,接受率会逐步下降,第1步下降约1.2%,第2步下降约2.6%,第3步下降约3.5%。这导致MTP在RL训练中的加速效果大打折扣。更严重的是,在某些agent任务中,rollout阶段可以长达128K token、200轮对话,如果MTP接受率下降,整个RL训练pipeline的效率会受到严重影响。例如,Qwen3.5-3.6 Plus在SWE-bench RL训练中,如果MTP接受率持续下降,整个异步RL训练的端到端加速比可能从预期的1.8倍降至1.2倍以下。
本文的目标是本文的目标是在整个RL训练过程中保持MTP的高接受率,从而实现稳定的加速效果。具体来说,希望达到:第一,接受率在RL训练中保持稳定,不随策略熵的变化而显著波动;第二,接受率提升到百分之九十五以上(gamma等于3时接近理论极限);第三,消除了在RL阶段进行MTP在线更新的需要,降低内存和计算开销;第四,端到端RL训练加速比达到1.5到1.8倍。
与已有工作不同的是,现有工作主要从分布失配的角度分析MTP接受率下降,认为是因为策略模型权重更新导致frozen的draft heads与target模型分布不一致。然而本文通过分解分析发现,分布失配的影响实际上是次要的,主要驱动因素是策略模型熵的波动。更独特的是,本文发现使用拒绝采样替代target-only sampling可以大幅减轻熵的影响,并进一步提出直接优化TV距离的训练目标,这是与现有工作的本质区别。
核心方法
Bebop方案的核心思想有三点:第一,使用拒绝采样替代target-only sampling,使得接受率由分布重叠决定而非被max y p(y)限制;第二,提出端到端TV损失,直接优化多步拒绝采样的接受率;第三,在RL训练前进行一次轻量级的MTP预训练(使用TV损失),然后在整个RL过程中保持draft heads冻结,无需在线更新。直觉上,传统CE/KL训练会产生均匀的per-token失配,当熵高时(有效支持大小约等于exp(H(p))),这些小失配会累积导致接受率下降。而TV损失的梯度与q成正比,只关注高概率token,产生概率比例失配,这使得TV距离小于等于delta除以2,与熵无关。
核心创新点有三:第一,揭示并利用了拒绝采样与target-only sampling在接受率对熵的敏感性上的根本差异,前者接受率是分布重叠,后者接受率被max y p(y)限制,后者随熵增加直接下降,前者对熵变化更鲁棒;第二,提出端到端TV损失,直接优化决定拒绝采样接受率的TV距离,而不是像CE/KL那样间接优化。TV损失的梯度有界(小于等于1),且只作用于q不为零的token,产生概率比例失配结构;第三,通过分解分析证明,在拒绝采样下,策略更新导致的分布失配对接受率的影响可以忽略,因此无需在RL阶段进行昂贵的MTP在线更新。
方法步骤详情
方法步骤的完整描述如下:第一,在SFT阶段训练MTP heads,使用端到端TV损失,这个损失直接优化期望接受长度,自然地给早期步骤更高权重;第二,在RL rollouts中使用拒绝采样,draft token从q分布采样,以概率min(1, p(y)/q(y))接受。实现时使用fused backward kernel降低内存峰值;第三,在RL训练过程中冻结MTP heads,无需更新draft参数,保持预训练阶段的TV优化效果。如果必须进行MTP co-training(如target-only sampling场景),使用独立的学习率和梯度范数归一化,避免干扰主模型的RL优化;第四,可选地,如果RL探索导致策略熵超出SFT数据覆盖范围,可以在RL阶段继续使用TV损失更新MTP heads以扩展覆盖范围。
技术新颖性
技术新颖性体现在:第一,首次提出直接优化TV距离作为MTP训练目标,与传统的CE/KL损失本质不同。CE/KL梯度是q减p,对所有token均匀施加优化压力;TV梯度是负q乘以指示函数减S,只作用于q不为零的token,自动忽略长尾,产生概率比例失配;第二,首次系统分析了MTP接受率在RL训练中的两个驱动因素(熵和失配),并通过分解分析证明熵是主导因素,失配可以忽略;第三,首次证明在拒绝采样下,预RL训练足以维持整个RL过程的稳定接受率,消除了对在线MTP更新的需求;第四,提出了端到端多步TV损失,捕获多步验证的乘法结构,这与现有工作的固定位置权重方案本质不同。
实验结果
核心发现包括:第一,MTP接受率与策略熵存在明确的负线性关系,约等于a减去b乘以H(p),这在不同模型、任务、训练阶段都成立。CE/KL训练的斜率b约为负1.68;第二,拒绝采样相比target-only sampling对熵的变化更鲁棒,在CE/KL训练下两者斜率相似,但拒绝采样基线接受率更高;第三,TV训练将熵-接受率斜率降低百分之九十五以上(从负1.68降至负0.06),大幅解耦了接受率与熵的关系;第四,在SFT阶段,e2e TV损失相比CE基线在推理任务上提升3.3到8.0%接受率,在Agent任务上提升6.7%,在OOD的MT-Bench上提升2.3%。Qwen3.7-Plus在Agent任务上达到98.6%接受率;第五,在RL训练中,RS w/ TV方案保持稳定接受长度,而target-only持续下降。Reasoning RL中RS w/ TV实现1.5到1.8倍延迟降低,Agent RL中rollout阶段加速达2.4倍;第六,RL阶段更新MTP权重收益有限,从RS w/ TV检查点开始,继续用CE损失更新会使接受率退化到RS w/ CE基线,用TV损失更新则几乎没有额外提升;第七,接受率提升与吞吐量提升呈线性关系(r等于0.81),每提升10%接受率可带来约25%额外吞吐量增益。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| SFT Rejection Sampling (Math) | Acceptance Rate | 78.0% (e2e TV) | 75.0% (CE) | +3.0% |
| SFT Rejection Sampling (Code) | Acceptance Rate | 74.6% (e2e TV) | 71.3% (CE) | +3.3% |
| SFT Rejection Sampling (SWE) | Acceptance Rate | 83.1% (e2e TV) | 75.1% (CE) | +8.0% |
| SFT Rejection Sampling (Agent) | Acceptance Rate | 97.0% (e2e TV) | 90.3% (CE) | +6.7% |
| SFT Rejection Sampling (MT-Bench OOD) | Acceptance Rate | 67.6% (e2e TV) | 65.3% (CE) | +2.3% |
| RL Training Latency (Reasoning) | Speedup Ratio | 1.5-1.8× | 1.0× (no MTP) | +50-80% |
| RL Training Latency (Agent) | Rollout Speedup | 2.4× | 1.0× (no MTP) | +140% |
| End-to-End Async RL | Acceleration | 1.8× | 1.0× (no MTP) | +80% |
局限与改进
局限性包括:第一,熵-接受率关系的理论分析依赖于建模假设(均匀失配 vs 概率比例失配),这些假设由梯度结构启发驱动而非严格证明,tightening这些假设仍是开放问题;第二,TV训练的熵不变性是分布条件的:它只在SFT训练数据覆盖的熵范围内成立。如果RL探索驱动策略熵显著超出这个范围,draft head会遇到分布外目标分布,失配比delta不再有界,恢复到类似CE/KL训练的熵-接受率依赖;第三,全词表TV损失计算在大型词表上内存峰值较高,虽然使用fused kernel缓解了这个问题,但对于更大规模的词表可能仍需优化;第四,在极端高熵场景(如tau大于1.2的温度采样),拒绝采样也会出现接受率下降,虽然比target-only更温和但仍然存在;第五,实验主要在Qwen系列模型上进行,虽然包括不同规模(35A3B到Max),但在其他架构(如transformer变体)上的泛化性有待验证。
独立分析的弱点
独立分析的弱点包括:第一,对分布外熵的脆弱性,当RL探索导致策略熵超出SFT数据范围时,TV训练的熵不变性会失效。改进方向是在RL阶段进行TV损失的在线更新,或使用更强大的数据增强策略覆盖更广的熵范围;第二,内存开销,全词表TV损失计算在大型词表(如128K词汇)上内存峰值较高。改进方向是开发更高效的top-K TV近似算法,或者使用分层或聚类词汇结构;第三,温度敏感性,虽然拒绝采样比target-only对温度更鲁棒,但在极高温度(tau大于1.5)下接受率仍会下降。改进方向是开发温度自适应的MTP策略,或结合温度缩放的draft训练;第四,长序列生成中的位置依赖性,在序列后期(特别是在长思维链推理中),接受率可能下降。改进方向是开发自适应gamma策略,根据局部熵动态调整draft长度;第五,对模型规模的依赖,虽然Qwen3.7模型能达到95%接受率,但较小模型(如27B)在相同训练下接受率较低(约70-80%)。改进方向是针对小模型设计更高效的draft架构或训练策略。
未来方向
未来研究方向包括:第一,自适应MTP策略,根据局部熵估计动态调整draft长度gamma,在低熵区域使用更长的draft,在高熵区域使用更短的draft,进一步优化吞吐量;第二,温度感知训练,在SFT阶段使用多种温度进行TV损失训练,使draft head对更广的熵范围鲁棒;第三,分层TV损失,将词汇分组,在每组内计算TV距离,然后加权组合,在保持接受率优化的同时降低内存开销;第四,跨任务迁移学习,探索在不同任务间共享或迁移MTP heads,减少每个任务的训练成本;第五,与其他加速方案的结合,研究Bebop与KV cache压缩、量化、早退等其他加速技术的协同效应;第六,在更多RL算法中的应用,验证Bebop在PPO、DPO等其他RL算法中的有效性;第七,理论收紧,严格证明概率比例失配假设下的熵不变性,或找到更紧的TV距离上界。
复现评估
复现评估:作者已在GitHub上开源实现。论文提供了详细的实验配置:使用Qwen3.5、3.6、3.7系列模型,在Megatron框架上训练,全局batch size为256,序列长度256K,学习率3.5乘以10的负5次方(SFT)或1到2乘以10的负6次方(RL)。SFT训练1个epoch,RL训练200步左右。实验覆盖了多种任务:数学推理、代码生成、Agent任务等。计算资源方面,Qwen3.6-35A3B的SFT训练需要数百GPU小时,RL训练也需要相当的资源。对于资源受限的研究者,复现完整实验可能需要访问大规模集群。但核心的TV损失训练和拒绝采样验证可以在较小的模型上快速验证。代码基于SGLang和veRL框架,这两个都是开源的,降低了复现门槛。总体而言,复现难度中等偏高,主要是由于需要大规模计算资源,但代码和配置的可用性良好。
论文图表
图(a)展示了MTP接受率与策略熵之间的线性负相关关系,每个点代表一个RL训练步的平均熵和接受长度。不同颜色的点代表不同模型规模和任务。灰色点是传统CE训练的结果,显示接受率随熵增加线性下降;蓝色点是TV训练的结果,显示接受率几乎不随熵变化。图(b)展示了draft分布与target分布的重叠情况,TV训练的draft与target的重叠明显大于CE训练的draft。
这张图是论文的核心发现,直观展示了:第一,熵是MTP接受率下降的主导因素;第二,TV训练能够解耦接受率与熵的关系。这是理解Bebop方案为何有效的关键证据。
该图展示了SWE-bench RL训练过程中三个MTP步骤的接受率变化趋势。横轴是训练步数(0到350),纵轴是接受率。三条曲线分别代表MTP步骤1(蓝)、2(绿)、3(橙)。所有曲线都呈下降趋势,步骤1从约96%降至94.8%(下降1.2%),步骤2从94%降至91.4%(下降2.6%),步骤3从92%降至88.5%(下降3.5%)。
这张图展示了MTP在RL训练中接受率退化的具体幅度,量化了问题的严重性。这为后续分析提供了背景,也是论文要解决的核心问题的直观展示。