← 返回 2026-06-11

突破熵界限:通过MTP与拒绝采样加速RL训练 Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling

Yucheng Li, Huiqiang Jiang, Yang Xu, Jianxin Yang, Yi Zhang, Yizhong Cao, Yuhao Shen, Fan Zhou, Rui Men, Jianwei Zhang, An Yang, Bowen Yu, Bo Zheng, Fei Huang, Junyang Lin, Dayiheng Liu, Jingren Zhou 📅 2026-06-10 👍 21 2026-07-13 08:37
LLM训练加速 多令牌预测 强化学习 推测解码

提出端到端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损失,捕获多步验证的乘法结构,这与现有工作的固定位置权重方案本质不同。

Decomposition of acceptance length changes during RL training. Delta alpha (total, gray) is decomposed into an entropy-driven component Delta alpha entropy = b乘以(Ht minus H0) (orange) and a draft-target mismatch component Delta alpha mismatch (green). Under target-only sampling, both entropy increase and growing mismatch contribute to acceptance degradation. Under rejection sampling with CE loss, the degradation is almost entirely entropy-driven, with mismatch remaining near zero. RS with TV loss shows near-zero change across all components, confirming the stability of TV-trained drafts.
Figure 3: Decomposition of acceptance length changes during RL training. Delta alpha (total, gray) is decomposed into an entropy-driven component Delta alpha entropy = b乘以(Ht minus H0) (orange) and a draft-target mismatch component Delta alpha mismatch (green). Under target-only sampling, both entropy increase and growing mismatch contribute to acceptance degradation. Under rejection sampling with CE loss, the degradation is almost entirely entropy-driven, with mismatch remaining near zero. RS with TV loss shows near-zero change across all components, confirming the stability of TV-trained drafts.
Accept length under target-only sampling with CE loss vs. TV loss during SFT training. Acceptance rates are nearly identical (less than 0.3% difference) across all tasks, confirming that target-only acceptance depends on the target distribution rather than the draft's distributional shape.
Figure 5: Accept length under target-only sampling with CE loss vs. TV loss during SFT training. Acceptance rates are nearly identical (less than 0.3% difference) across all tasks, confirming that target-only acceptance depends on the target distribution rather than the draft's distributional shape.
Entropy loss vs. accept length across three RL workloads in Qwen3.6-Plus and Qwen3.7-Max. Each point represents one training step; the line shows the linear fit. TO and RS w/ CE exhibit a strong negative correlation (slope approx minus 1.68), while RS w/ TV remains nearly flat (slope approx minus 0.06), confirming that TV training decouples acceptance from entropy.
Figure 8: Entropy loss vs. accept length across three RL workloads in Qwen3.6-Plus and Qwen3.7-Max. Each point represents one training step; the line shows the linear fit. TO and RS w/ CE exhibit a strong negative correlation (slope approx minus 1.68), while RS w/ TV remains nearly flat (slope approx minus 0.06), confirming that TV training decouples acceptance from entropy.
(a) Entropy gap Delta H vs. DKL(q double vertical bar p) across models and tasks. (b) Entropy gap correlates negatively with RS acceptance rate (r equals minus 0.54). (c) KL divergence shows no such correlation (r equals 0.13), indicating that entropy gap, rather than KL, is the relevant predictor of RS acceptance.
Figure 10: (a) Entropy gap Delta H vs. DKL(q double vertical bar p) across models and tasks. (b) Entropy gap correlates negatively with RS acceptance rate (r equals minus 0.54). (c) KL divergence shows no such correlation (r equals 0.13), indicating that entropy gap, rather than KL, is the relevant predictor of RS acceptance.
RS decision boundary across models. Nearly all model-task combinations fall in the RS-better region, confirming that rejection sampling is beneficial for virtually all practical MTP deployments.
Figure 13: RS decision boundary across models. Nearly all model-task combinations fall in the RS-better region, confirming that rejection sampling is beneficial for virtually all practical MTP deployments.

实验结果

核心发现包括:第一,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%额外吞吐量增益。

Gradient comparison across training objectives. C denotes a global constant (S for TV, DKL(q double vertical bar p) for reverse KL). See derivations.
Table 1: Gradient comparison across training objectives. C denotes a global constant (S for TV, DKL(q double vertical bar p) for reverse KL). See derivations.
MTP acceptance rate (%) under rejection sampling across tasks and training objectives under gamma equals 3 on Qwen3.5-35A3B. All results are measured at convergence. Delta denotes improvement over CE loss baseline.
Table 2: MTP acceptance rate (%) under rejection sampling across tasks and training objectives under gamma equals 3 on Qwen3.5-35A3B. All results are measured at convergence. Delta denotes improvement over CE loss baseline.
MTP acceptance rate (%) under rejection sampling across tasks and training objectives under gamma equals 3 on different models. Qwen3.7 models are trained with e2e TV loss; all others are trained with CE loss.
Table 3: MTP acceptance rate (%) under rejection sampling across tasks and training objectives under gamma equals 3 on different models. Qwen3.7 models are trained with e2e TV loss; all others are trained with CE loss.
CE loss (solid) vs. TV loss (dashed) during SFT training. TV loss consistently achieves higher acceptance rates across all MTP steps, with especially pronounced gains on agentic tasks.
Figure 4: CE loss (solid) vs. TV loss (dashed) during SFT training. TV loss consistently achieves higher acceptance rates across all MTP steps, with especially pronounced gains on agentic tasks.
Accept length during RL training across different workloads in Qwen3.6-Plus and Qwen3.7-Max. Rejection sampling with TV loss (RS w/ TV) consistently maintains higher accept lengths compared to target-only (TO) and rejection sampling with CE loss (RS w/ CE).
Figure 6: Accept length during RL training across different workloads in Qwen3.6-Plus and Qwen3.7-Max. Rejection sampling with TV loss (RS w/ TV) consistently maintains higher accept lengths compared to target-only (TO) and rejection sampling with CE loss (RS w/ CE).
Training latency comparison during RL using Qwen3.6-35A3B and Qwen3.6-Plus. MTP with rejection sampling (RS w/ TV) substantially reduces per-step latency compared to training without MTP (w/o MTP) and target-only sampling (TO).
Figure 7: Training latency comparison during RL using Qwen3.6-35A3B and Qwen3.6-Plus. MTP with rejection sampling (RS w/ TV) substantially reduces per-step latency compared to training without MTP (w/o MTP) and target-only sampling (TO).
(a) Accept length during RL training with and without MTP weight updates. Updating MTP weights with CE loss causes the acceptance rate to converge toward the corresponding non-updated baseline, while target-only sampling with CE loss updates can even degrade acceptance due to distribution mismatch. (b) Accept rate delta (RS minus No-RS) vs. throughput speedup ratio (RS over No-RS) across 8 models and 3 tasks (r equals 0.81). Higher acceptance rate gains from rejection sampling translate directly to greater throughput improvements.
Figure 9: (a) Accept length during RL training with and without MTP weight updates. Updating MTP weights with CE loss causes the acceptance rate to converge toward the corresponding non-updated baseline, while target-only sampling with CE loss updates can even degrade acceptance due to distribution mismatch. (b) Accept rate delta (RS minus No-RS) vs. throughput speedup ratio (RS over No-RS) across 8 models and 3 tasks (r equals 0.81). Higher acceptance rate gains from rejection sampling translate directly to greater throughput improvements.
查看结构化数据
任务指标本文基线提升
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框架,这两个都是开源的,降低了复现门槛。总体而言,复现难度中等偏高,主要是由于需要大规模计算资源,但代码和配置的可用性良好。