跳到正文
北京时间
原文
HuggingFace Daily Papers(社区热门论文)·· 2026-06-10精选AI 评分73

Bebop:通过带拒绝采样的多token预测加速RL训练

Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling

AI 导读

Bebop系统研究多token预测(MTP)在大语言模型RL后训练中的加速瓶颈。发现MTP接受率与模型熵波动呈负线性关系;概率拒绝采样相比贪婪草稿采样能缓解熵干扰。提出端到端TV损失直接优化多步拒绝采样接受率,带来约10%提升,最高达95%接受率,额外推理吞吐提升25%。在Qwen3.5、Qwen3.6和Qwen3.7上,异步RL训练实现1.8倍端到端加速。预RL阶段的MTP训练结合e2e TV损失和拒绝采样可保持全程一致的加速,无需在线更新。

推荐理由

做RLHF的团队值得细看,MTP加拒绝采样直接把RL训练推到了1.8倍加速,还省去了在线更新MTP的成本。TV损失的设计让接受率提了10%,这套组合拳相当实用。

正文 · AI 翻译

曹宜忠、沈宇豪、范舟、阮睿、张建伟、杨安、余博文、郑波、黄飞、林俊阳、刘大义恒、周靖人

通义千问团队,阿里巴巴集团

摘要

强化学习已成为现代大语言模型的关键组成部分,然而 rollout 阶段仍然是 RL 训练流程中的主要瓶颈。尽管多 token 预测通过推测性解码为加速 rollout 提供了自然的解决方案,但许多研究观察到,在 RL 训练过程中,MTP 的接受率会显著下降,导致加速效果有限。为了解决这一瓶颈,我们提出了 Bebop,这是一项关于在 LLM 后训练中应用 MTP 的系统性研究,并提供了将 MTP 集成到大规模 RL 流程中的实用方案。首先,我们揭示了 MTP 接受率从根本上受到模型熵波动的限制,该波动与 RL 阶段熵值的上升呈现出清晰的负线性关系(§3)。其次,我们证明,与贪婪草稿采样相比,概率性拒绝采样在很大程度上缓解了 RL 中熵引入的干扰。我们进一步发现,传统的 MTP 训练目标(交叉熵或 KL 散度)在此类设置中并非最优,因此我们提出了一种新颖的端到端 TV 损失函数,该函数直接优化多步拒绝采样的接受率,从而提升接受率,在数学推理、代码生成和智能体任务中实现了高达 95% 的接受率以及高达 25% 的额外推理吞吐量增益(§4)。第三,我们测试了 RL 过程中的多种在线 MTP 训练策略,并表明,在 RL 之前使用端到端 TV 损失和拒绝采样进行 MTP 训练,能够在整个 RL 过程中保持一致的接受率和加速效果,从而消除了代价高昂的在线 MTP 更新的需求(§5)。我们提供了大量的实验和分析来验证我们的发现。实验结果表明,我们的方法在 Qwen3.5、Qwen3.6 和 Qwen3.7 模型的异步 RL 训练中实现了端到端的加速。

Refer to caption
(a) 熵值与接受长度
Refer to caption
(b) 草稿分布/目标分布
图 1:(a) 在强化学习中,多 token 预测(MTP)的接受率随策略熵的波动呈线性下降;采用我们提出的新型端到端全变差(TV)损失函数训练 MTP,可在拒绝采样下基本消除这种熵依赖性。每个数据点代表不同规模的 Qwen3.5、3.6 和 3.7 训练任务中,在某一强化学习步骤下的平均熵与接受长度。(b) 经全变差损失训练的 MTP 与策略模型实现了显著更好的分布重叠,从而获得更优的接受率和加速比。

1 引言

强化学习(RL)已成为现代大语言模型(LLM)训练的关键范式(OpenAI, 2026;Anthropic, 2026;Qwen Team, 2026b;DeepSeek-AI, 2026;GLM Team, 2026;Kimi Team, 2026;MiniMax, 2026a)。然而,LLM 的强化学习训练计算成本高昂,在单轮和多轮场景下,端到端时间主要被推理 rollout 所占据。尽管近期异步强化学习框架的进展(Fu et al., 2025;Wang et al., 2025;THUDM, 2025)可部分缓解长尾延迟问题,但 rollout 成本仍是强化学习训练中的主要瓶颈。多 token 预测(MTP)近期作为一种可扩展的推测解码范式,在加速 LLM 推理方面备受关注(DeepSeek-AI, 2024;Qwen Team, 2026a)。这自然引出一个问题:能否有效利用 MTP 来加速 LLM 的强化学习训练?

我们进行了大量实验,结果表明,在强化学习训练中直接使用 MTP 往往会导致接受率显著下降,从而使得加速效果有限。具体而言,强化学习过程中有两个因素可能影响 MTP 接受率:1)为了鼓励探索,策略模型通常会保持较大的熵——甚至呈现出逐渐上升的熵曲线,这使得预测草稿 token 变得更加困难,从而降低了接受率;2)策略模型的权重更新会导致策略模型与 MTP 模块(在强化学习训练中冻结)之间出现分布不匹配,这可能影响接受率。通过理论分析和实证分解(第 3 节),我们证明熵是导致接受率下降的主导因素,而策略更新引入的不匹配影响可以忽略不计(图 3)。为了应对熵边界挑战并确保 MTP 的加速效果,近期研究(Chen 等人,2026b;Li 等人,2025;MiniMax,2026b)提出了在强化学习过程中进行在线 MTP 训练以缓解这种退化,但这种方法会引入显著的内存和延迟开销,并且在许多强化学习任务中带来的改进有限。

在本文中,我们提出了 Bebop(打破熵界以实现最优预测),并展示了使用概率拒绝采样(我们已在 https://github.com/sgl-project/sglang/pull/26312 发布了实现)替代常见的贪心目标模型独有采样(目标模型独有采样指投机解码的验证仅使用目标概率而不缓存草稿概率:它通过选取草稿 token 并在验证中使用),能够大幅缓解由策略熵波动(§3.3)驱动的接受率下降问题,并在接受率上带来显著提升。关键洞察在于,目标模型独有接受率从根本上受限于,该值会随熵升高而直接下降,而拒绝采样的接受率等于完整的分布重叠,因此对熵变化敏感度低得多。我们进一步发现,现有的多 token 预测(MTP)训练目标(如交叉熵或 KL 散度)对于拒绝采样而言并非最优:交叉熵/KL 仅间接改善决定拒绝采样接受率的分布重叠。这促使我们提出一种新颖的端到端全变差(TV)损失函数,该函数优化联合多步重叠,从而直接提升拒绝采样接受率。

Bebop 生成的 MTP 模型在整个强化学习训练过程中保持一致的接受率。这些接受率对熵变化基本保持不变。Bebop 仅通过在强化学习前使用端到端 TV 损失进行轻量级 MTP 训练阶段,并在 rollout 期间配合拒绝采样,即可实现这种稳定性,无需在强化学习期间进行 MTP 协同训练。

具体而言,我们做出了以下贡献:

  • •

    MTP 接受率的熵约束(§3)。我们展示了在强化学习训练中,MTP 接受率从根本上受到目标模型熵的约束,在不同任务和模型上均呈现出清晰的负线性关系。我们进一步展示了拒绝采样在强化学习中能大幅提升接受率,因为其接受率取决于策略-草稿重叠,且对熵变化敏感度较低。

  • •

    用于 MTP 训练的端到端全变分损失(§4)。我们发现,经过 CE/KL 训练的 MTP 在拒绝采样中会产生次优结果,因此引入了一种新颖的端到端全变分损失,直接优化多步拒绝采样的接受率。我们证明,端到端全变分损失能确保训练稳定,产生固有熵不变的 MTP,并在接受率上带来额外提升。

  • •

    用于强化学习的 MTP 适配策略(§5)。我们表明,通过使用端到端全变分损失和拒绝采样进行轻量级的强化学习前 MTP 训练,我们的 MTP 模块在整个强化学习训练过程中能提供一致的接受率。另一个因素,即由策略更新导致的策略-草稿不匹配,可以忽略不计,这消除了在强化学习期间进行昂贵的 MTP 在线训练的必要性。

  • •

    广泛的实证验证与分析(§6, §7)。通过在推理、编码及各类智能体任务上使用 Qwen3.5、3.6 和 3.7 模型进行大规模实验,我们验证了 Bebop,并提供了将 MTP 集成到强化学习流水线中的实用方案,实现了异步强化学习流水线的端到端加速。我们进一步分析了全变分损失如何塑造草稿分布、策略更新下拒绝采样的鲁棒性,以及温度和生成长度对接受率的影响。

2 预备知识

2.1 多 token 预测与推测解码

作为推测解码的一种有效范式(Leviathan 等人,2023;Chen 等人,2023),多 token 预测(MTP)通过轻量级的草稿头来增强自回归大语言模型,这些草稿头能够顺序预测多个未来 token(Gloeckle 等人,2024;DeepSeek-AI,2024;Yang 等人,2025)。设 表示目标(骨干)模型在位置 处的下一个 token 分布, 表示草稿头预测的分布。在推理过程中,MTP 采用“先草稿后验证”的范式运行:一系列草稿头顺序提出候选 token ,其中每个草稿头将前一个草稿头的隐藏状态作为输入;然后,这些候选 token 在单次前向传播中与目标模型进行验证。

每个验证步骤中预期被接受的模型 token 数量(我们称之为接受长度)直接决定了推理吞吐量。该接受长度取决于验证过程中使用的具体接受方法,详见下一节。

2.2 接受方法

在推测解码中,通常使用两种接受方法:仅目标模型采样和拒绝采样。图 13 展示了每种方法下代表性模型的接受率分布。

仅目标模型采样。

在仅目标模型采样下,草稿 token 被贪婪地选取并以概率 被接受,仅使用目标模型的概率。单步接受率为:

(1)

如果被拒绝,则输出 token 从残差分布 中重新采样,确保整体输出分布保持无偏。值得注意的是,对于接受率相对较低的草稿模型,仅目标模型采样可以比拒绝采样获得更高的吞吐量,因为更简单的接受标准避免了缓存和计算草稿概率向量的开销。

拒绝采样。

在拒绝采样(Leviathan 等人,2023;Chen 等人,2023)下,草稿 token 以概率 被接受。预期的单步接受率为:

(2)

其中 是总变差距离(Levin 和 Peres,2017)。该方法提供了无偏保证:无论草稿质量如何,输出分布都精确等于目标分布 。

2.3 大语言模型的强化学习

我们考虑大语言模型的标准强化学习框架,其中策略(即大语言模型)根据提示词生成轨迹并接收标量奖励 。我们采用 GRPO(Shao 等人,2024),该方法为每个提示词从 rollout 策略中采样一组轨迹,并优化裁剪后的替代目标函数:

(3)

其中 是重要性采样比率, 是组归一化优势值。

大语言模型的强化学习训练通常在一个三阶段循环中运行:(1) 展开阶段,使用当前策略在推理引擎中生成轨迹,可能涉及多轮沙箱或工具交互;(2) 奖励阶段,使用奖励模型或验证器评估生成的轨迹;(3) 更新阶段,在训练引擎中使用策略梯度方法优化策略。异步强化学习或部分展开框架通常被采用,以缓解展开过程中长尾轨迹带来的气泡开销(Fu 等人,2025;Wang 等人,2025;THUDM,2025;Qin 等人,2025;MiniMax,2026b)。尽管采用了异步设计,展开阶段仍然是主要的计算瓶颈。虽然多 token 预测(MTP)提供了一种强大的加速范式来缓解这一负担,但其在强化学习环境中的直接应用暴露出了独特的性能差距,需要进一步优化。

2.4 强化学习训练期间多 token 预测的性能退化

Refer to caption
图 2:使用 Qwen3.5-3.6 Plus 进行 SWE-bench 强化学习训练时,每步多 token 预测接受率。每条线代表一次独立的强化学习运行。较后的多 token 预测步骤表现出逐渐增大的退化:在训练过程中,步骤 1 下降 1.2%,步骤 2 下降 2.6%,步骤 3 下降 3.5%。

在强化学习训练期间,多 token 预测接受率在各个预测步骤上显著下降。如图 2 所示,较后的步骤经历了逐渐增大的下降。每步接受率的下降幅度从步骤 1 的 1.2% 到步骤 3 的 3.5%。

近期工作(MiniMax,2026b;Chen 等人,2026b;Li 等人,2025)主要将此退化归因于分布不匹配。具体来说,由于骨干网络权重更新导致草稿头落后,静态草稿预测与不断变化的目标分布之间出现了差距。虽然这种不匹配确实存在,但我们认为这一观点并不完整。我们识别出强化学习训练期间目标模型熵的变化是另一个根本驱动因素。无论草稿精度如何,这些熵的变化都会从根本上改变可实现的接受界限。这两个因素通过多步接受结构相互叠加:

  1. 1.

    单步退化:随着 TV 距离增大,逐 token 接受率持续下降,其根本原因在于草稿模型与目标模型分布之间的持续偏差。

  2. 2.

    多步复合效应:对于 - 步 MTP,期望接受长度涉及各步接受率的乘积,因此退化效应呈乘性复合:。

关键在于,我们在第 3 节及图 3 中进行的分解分析挑战了传统的以分布失配为中心的观点。我们证明,在 RL 训练过程中,熵驱动成分实际上主导了接受率的波动。分布失配成分则相对较小。这一关键洞察重塑了我们对 MTP 退化的理解,并直接启发了我们后续的优化策略。

3 MTP 接受率的目标熵约束

在本节中,我们分析目标模型的熵如何从根本上约束 MTP 接受率,这解释了 RL 训练期间由熵偏移导致的接受率退化现象。这进一步为我们第 4 节中的训练目标提供了动机。

3.1 公式化表述

考虑生成过程中的一个固定位置。设 表示目标模型的下一个 token 分布, 表示草稿头的分布,其中 为词表。我们将目标熵定义为:

(4)

该熵值衡量目标模型预测的不确定性。低熵表示置信度高、分布集中,而高熵则表示分布分散。

我们感兴趣的是理解 如何约束式 (1) 和式 (2) 中定义的可实现接受率 和 。

3.2 仅目标采样的 MTP

在仅目标采样条件下,接受率取决于草稿模型的贪心预测与目标模型高概率区域的匹配程度。当目标熵较低时(即 集中在少数 token 上),即使草稿模型精度一般,也能通过将概率质量集中在主导 token 上实现高接受率。反之,当 较高时,目标分布分散在多个 token 上,这会降低 并增加排序错误的概率。

命题 1(仅目标采样下与熵相关的接受率)

对于一个训练良好的草稿模型,该函数是 的单调递减函数,以 为下界,且经验上可很好地近似为线性(图 1a):

(5)

其中 为正的常数。在不完美草稿下,排序错误会加剧斜率,但线性关系保持不变(§D.2)。

证明思路。

当草稿正确识别出目标模型的前 1 个 token()时,接受率简化为 。对凹函数对数应用 Jensen 不等式,可得 ,即 。将 写为某个光滑递减函数,并在参考熵 附近进行一阶泰勒展开:

(6)

由于 是递减函数,因此 。包含不完美草稿修正的完整推导见 §D.2。∎

如图 1a 所示,这种线性关系在不同模型规模、任务和训练阶段中均表现出极强的鲁棒性。

3.3 带拒绝采样的 MTP

在拒绝采样下,接受率等于 与 之间的总变差重叠(式 (2))。我们可以利用恒等式和概率归一化来分解总变差距离:

(7)

因此,最大化接受率等价于最小化总变差距离:

(8)

由此,接受率不再直接受策略熵的约束。然而,经验结果表明,在切换为拒绝采样后,与熵的关联仍然存在。在进一步研究中,我们发现,在基于交叉熵/KL 散度训练的草稿模型下,当 具有高熵时,即使每个 token 的微小不匹配也会累积,导致更大的总变差距离。这促使我们更深入地分析训练目标如何影响这一关系,具体如下。

命题 2(交叉熵/KL 训练拒绝采样下的熵依赖接受率)。

在基于交叉熵/KL 训练的草稿模型下,拒绝采样的接受率满足:

(9)

其中 为正的常数, 与 相当,但经验上斜率略陡(§D.3,图 8)。

证明思路。

交叉熵/KL 梯度会产生均匀的逐 token 不匹配 。由于有效支撑集大小按 缩放,总变差距离累积为 ,从而得到 。在运行熵范围内对指数函数进行线性化,即可得到所述形式。详见 §D.3。∎

因此,在基于交叉熵/KL散度训练的MTP中,拒绝采样和仅目标采样都对熵的变化敏感。由于策略熵在强化学习训练期间波动显著,这种敏感性本质上限制了可实现的加速比。

4 针对强化学习训练优化MTP

如上所述,由于熵界的存在,MTP的接受率在强化学习训练期间可能会显著下降。在本节中,我们开发了新颖的端到端全变差损失来应对这一挑战。

4.1 全变差损失:直接优化接受率

动机。

传统的MTP训练最小化目标分布与草稿分布之间的交叉熵损失或KL散度。⁴⁴⁴ 在本文中,除非另有说明,“KL”均指前向KL散度。交叉熵和前向KL仅相差一个常数,且产生相同的梯度。我们在§C中单独分析反向KL散度。然而,拒绝采样的接受率由全变差距离决定(式(8)),而非KL散度。根据Pinsker不等式, ,因此KL仅提供一个间接上界,最小化KL并不能高效地最小化全变差距离。这促使我们将直接优化全变差距离作为MTP的训练目标。

全变差损失。

我们提出直接最小化全变差距离:

(10)

其中 被视为常数(从计算图中分离),梯度仅通过 传播。

梯度分析。

设草稿头输出logits为 ,且 。全变差损失关于 的梯度为:

(11)
命题3(有界梯度)。

全变差损失梯度是有界的:对于所有 成立。

证明。

由于 且 (因为指示函数 且 ),我们有 。∎

这种有界梯度属性确保了训练稳定性,这与KL散度形成对比——当 和 差异显著时,KL散度的梯度可能具有很大的幅值。

直观解释。

全变差损失梯度在拒绝采样机制方面有一个自然的解释:

  • •

    对于 的token(即会被接受的token):梯度增加logit,鼓励草稿分配更多概率质量。

  • •

    对于 的token(即会被拒绝的token):梯度减小logit,抑制过度自信的预测。

  • •

    对于(无关 token 的)token:梯度是自动的(因为它与 成正比),从而避免了对词汇表中长尾部分的无效优化投入。

这种选择性梯度行为与 KL 散度形成对比,后者对所有 token 都施加梯度,无论它们与接受决策的相关性如何。

CE、KL 和 TV 梯度的比较。

表 1 总结了三种训练目标的梯度结构。关键区别在于梯度是否与 成正比:CE 损失产生均匀的逐 token 失配(),将优化努力均匀分布在词汇表上,包括无关的低概率 token。相比之下,反向 KL 和 TV 损失都表现出与 成正比的梯度,具有自然的尾部抑制,将更新集中在草稿模型已分配不可忽略概率的 token 上。然而,尽管具有这一共同特性,反向 KL 在接受率上的提升相比 CE 可以忽略不计(§6),因为其零强制行为允许草稿模型丢弃 的模式,且其非对称惩罚驱动 全局变化——两者都减少了 TV 重叠(详见 §C 分析)。TV 损失通过直接优化与接受相关的量并产生与概率成正比的失配(将接受与目标熵解耦),避免了这些陷阱。

媒体内容 · 前往原文查看
表 1:各训练目标的梯度比较。 表示全局常数(TV 为 ,反向 KL 为 )。推导过程见 §A-§C。
属性 前向 KL 反向 KL TV 损失
梯度
? 否 是 是
尾部抑制 否 是 是

4.2 端到端多步 TV 损失

对于 步 MTP,期望接受长度为:

(12)

其中 是第 步的每步接受率。直接优化平均每步 TV 距离并未考虑多步接受的乘法结构。因此,我们提出端到端(e2e)TV 损失:

(13)

该损失函数直接优化归一化期望接受长度,自然地赋予较早步骤更高权重(因为它们出现在更多乘积项中),并捕捉多步验证的复合效应。这可以视为一种动态的逐步骤加权方案:由于该权重取决于当前草稿质量,每个位置的有效权重会随着训练进程自动调整,将重点转向当前限制接受率的步骤。这与先前使用固定位置相关权重的工作形成对比,例如基于头部位置的损失权重(Cai 等人,2024;Li 等人,2026)、指数衰减的块位置权重(Chen 等人,2026a)、对拒绝位置的固定衰减(Lei 等人,2026),或在交叉熵基础上使用逐位置权重(Wu 等人,2026)。

4.3 训练目标对熵-接受率关系的影响

在介绍了总变差损失之后,我们现在分析为何它在目标熵持续变化的强化学习场景中从根本上优于交叉熵/KL散度训练。公式(5)和(9)中的线性关系刻画了使用交叉熵/KL散度损失训练的草稿模型;我们证明训练目标的选择从根本上改变了熵-接受率关系。完整推导见§D节;此处我们陈述主要结论。

Pinsker不等式与KL–TV差距。

根据Pinsker不等式:

(14)

仅给出了 的上界,且KL散度优化在最小化TV距离方面效率低下地分配模型容量:最小化KL散度并不能高效地最小化TV距离,而TV距离是直接决定拒绝采样接受率的量。

交叉熵/KL散度训练:均匀失配。

KL散度梯度施加的优化压力与绝对差值 成正比,而与该差值相对于其他token的大小无关。在容量受限的草稿模型下,这会产生近似均匀的逐token失配:即 为常数。如命题2所示,这种均匀失配会在有效支持集 上累积,从而产生依赖于熵的接受率。

TV训练:概率比例失配。

TV 损失梯度(式 (11))与概率成正比,将优化集中在高概率 token 上,并自动忽略长尾部分。在容量受限的草稿模型下,每个 token 获得的优化资源与其概率成正比,因此每个 token 的错配也按比例缩放,而非保持均匀水平。这产生了概率比例错配:对于常数(详细推导见 §D.4)。

命题 4(TV 训练下熵依赖性的降低)。

当每个 token 的错配满足时,TV 距离的边界与熵无关:

(15)

得到。在实践中,草稿头的容量有限,因此可能表现出较弱的熵依赖性,但根据经验,与 CE/KL 训练相比,熵-接受率斜率降低了超过(图 8)。

证明梗概。

TV 梯度与成正比(式 (11)),因此每个 token 的优化资源随其概率缩放,产生(§D.4)。求和:,由于,该结果与熵无关。∎

该分析解释了经验观察结果:TV 训练的草稿模型在不同目标熵下实现了显著更稳定的接受率,而 CE/KL 训练的模型则表现出强烈的负相关性(图 8)。

5 面向 RL 的 MTP 适配策略

在 RL 流程中使用 MTP 的一个关键问题是,在 RL 训练期间是否需要在线更新 MTP 模块。我们通过一种分解分析来研究这一点,该分析将驱动接受率变化的两个因素分离开来。

5.1 分解:RL 中的熵与错配

利用 §3 中建立的线性熵-接受率关系,我们将 RL 训练期间接受长度的变化分解为:

(16)

其中是从每次实验早期阶段估计的熵-接受率斜率,是初始熵,是第步的总接受率变化。第一项捕捉了仅由熵变化引起的接受率变化(假设草稿-目标关系固定),而残差项则捕捉了由于主干网络权重更新导致的草稿-目标错配增长所带来的影响。

Refer to caption
图 3:RL 训练过程中接受长度变化的分解。(总计,灰色)被分解为熵驱动分量(橙色)和草稿-目标不匹配分量(绿色)。在仅目标采样下,熵增加和不断增长的不匹配共同导致了接受度的下降。在使用 CE 损失的拒绝采样下,性能下降几乎完全由熵驱动,不匹配度接近于零。使用 TV 损失的 RS 在所有分量上均显示出近乎为零的变化,证实了 TV 训练草稿的稳定性。

如图 3 所示:(1)在仅目标采样下,熵增加和不断增长的不匹配共同导致了接受度的下降,因为贪婪的草稿预测与不断变化的目标之间的对齐度越来越差。(2)在使用 CE 损失的拒绝采样下,性能下降几乎完全由熵驱动(),表明 RL 权重更新不会显著影响草稿-目标之间的 TV 重叠。(3)在使用 TV 损失的拒绝采样下,所有分量均观察到近乎为零的变化,证实了 TV 训练的草稿对熵变化和权重更新都具有鲁棒性。

5.2 预 RL 阶段的自适应已足够

分解分析引出了一个关键的实际见解:由于在拒绝采样下,RL 权重更新引起的草稿-目标不匹配可以忽略不计,因此在 RL 期间更新 MTP 头部是不必要的。在 RL 开始之前,于 SFT 阶段应用一次性的、使用 TV 损失的预 RL 自适应,就足以产生能够在整个 RL 训练过程中保持高接受率的草稿模型(图 6)。这消除了在 RL 期间维护 MTP 优化器状态的内存开销以及 MTP 梯度更新的计算成本。

经验上,如图 9a 所示,当从一个训练良好的 TV 检查点开始时,在 RL 期间继续更新 MTP 权重并不会带来显著的改进。更糟糕的是,在 RL 期间使用 CE 损失进行更新会导致接受率下降至 RS w/ CE 基线的水平,因为 CE 损失会使草稿分布变得更平滑,从而侵蚀了 TV 训练带来的收益(§7.2)。

5.3 MTP 与主干网络的交叉训练

当在强化学习阶段需要采用多 token 预测(MTP)联合训练时(例如,在仅目标采样中,不匹配问题不可忽略),我们发现,使用独立的学习率和独立的梯度范数归一化进行联合训练能提供最佳权衡。主干网络的梯度不受 MTP 损失的影响(该损失仅流经草稿头),从而确保 MTP 训练不会干扰主干网络的强化学习优化。

6 实验

我们通过三组实验验证了所提出的 Bebop 方法的有效性:(1)在监督微调阶段,不同多步 MTP 损失目标对接受率的影响;(2)在强化学习阶段,结合拒绝采样的端到端全变差(TV)损失在接受率、加速比和训练稳定性方面的优势;(3)在强化学习阶段更新 MTP 参数所带来的收益。

6.1 多步 MTP 训练提升接受率

我们首先评估了在监督微调阶段,不同损失目标如何影响 MTP 接受率。具体而言,我们比较了四种 MTP 训练目标:

  1. (1)

    CE 损失:草稿分布与目标分布之间的标准交叉熵;

  2. (2)

    KL 损失:KL 散度;

  3. (3)

    反向 KL 损失:反向 KL 散度(公式 (17));

  4. (4)

    TV 损失:单步 TV 距离(公式 (10));

  5. (5)

    端到端 TV 损失:端到端多步 TV 损失(公式 (13))。

我们在 Qwen3.5-35A3B(Qwen 团队,2026a)上使用混合 RFT 数据进行了主要实验。所有实验均采用恒定学习率,包含 3% 的预热步数,使用 Megatron(Shoeybi 等人,2019)训练 1 个 epoch,全局批大小为 256,序列长度为 256K。在多步 MTP 训练期间,我们在冻结大语言模型主干网络的情况下,对 5 个 MTP 步进行前向和反向传播。所有评估均使用(即目标模型每次验证 4 个 token)。我们进一步将实验扩展到 Qwen3.6-35A3B、Qwen3.6-Plus 和 Qwen3.7-Plus,在不同数据混合(包括领域特定数据(代码、智能体、推理)和混合 RFT 数据)上进行训练。吞吐量使用 SGLang 的 MTP 实现结合拒绝采样进行测量(实现细节见 §G)。

拒绝采样接受率。

表2报告了我们提出的端到端全变差损失(e2e TV loss)相较于交叉熵(CE)和KL散度基线在Qwen3.5-35A3B上的接受率提升情况。在所有任务中,端到端全变差损失在分布内任务(数学、代码、智能体、SWE)上持续将拒绝采样接受率提升3–8%,在分布外的MT-Bench(Zheng等人,2023)任务上提升幅度最高达2.3%。值得注意的是,在智能体任务上,CE基线已实现90.3%的高接受率,而端到端全变差损失进一步将其推高至97.0%,这一水平在强化学习训练和智能体推理中显著提升了展开效率。

除主要实验外,我们还在更广泛的模型和数据配置上进行了评估。如图4所示,我们在不同的数据混合上训练了Qwen3.6-35A3B、Qwen3.6-Plus和Qwen3.7-Plus,并追踪了训练过程中每一步的接受率。出现了几种模式。首先,CE损失在训练过程中导致第1步接受率出现显著且持续的下降,因为它将优化努力分散到整个词汇表上。相比之下,全变差损失则保持第1步接受率稳定或略有提升。其次,端到端全变差损失的优势在后来的MTP步骤中愈发明显:在第3步,全变差损失比CE损失高出约5%,而在第2步,差距为2.5–5%。第三,收益因任务而异:智能体任务受益最大,在Agent和SWE-Bench(Jimenez等人,2024)上提升幅度高达8%,而推理和对话任务的提升幅度为4–5%。最后,MTP接受率表现出很强的泛化能力。完全未使用智能体特定数据训练的模型在智能体任务上仍能达到约70%的接受率。具体而言,全变差损失在分布内领域上的改进幅度大于分布外任务。

媒体内容 · 前往原文查看
表2:在Qwen3.5-35A3B上,不同训练目标下各任务在拒绝采样中的MTP接受率(%)。所有结果均在收敛时测量。表示相较于CE损失基线的改进。
MTP损失 数学 代码 SWE 智能体 MTBench(分布外)
CE损失(基线)
KL损失
反向KL损失
全变差损失
端到端全变差损失(我们的方法)
Refer to caption
(a)推理和对话任务(数学、代码、MT-Bench)上的接受长度。
Refer to caption
(b) 在智能体任务与混合任务上的接受长度(混合任务、智能体任务、长周期任务、SWE-Bench)。
图 4:SFT 训练期间 CE 损失(实线)与 TV 损失(虚线)对比。在所有 MTP 步骤中,TV 损失始终获得更高的接受率,在智能体任务上提升尤为显著。

仅目标接受。

在仅目标采样下,所有训练目标的接受率几乎相同(差异为 0.3%),如图 5 所示。这符合预期:仅目标接受简化为草稿模型 top-1 排名是否正确的情况,其仅取决于目标分布,而非草稿模型的分布形态。相比之下,拒绝采样接受率取决于完整的分布重叠,而这正是 TV 损失发挥优势之处。这与我们在第 3 节中的分析一致。

Refer to caption
图 5:SFT 训练期间,使用 CE 损失与 TV 损失时,在仅目标采样下的接受长度。所有任务上的接受率几乎相同(差异为 0.3%),证实了仅目标接受取决于目标分布,而非草稿模型的分布形态。

吞吐量。

如图 9b 所示,接受率的提升大致线性地转化为吞吐量的增益。在所有数据集上,经过端到端 TV 损失训练的 Qwen3.7 Plus 始终优于经过 CE 损失训练的 Qwen3.6 Plus。这些增益有效加速了 RL 的 rollout 过程,在数十万 GPU 小时的规模下意义重大。

接受率随模型规模扩展。

如表 3 所示,经过多步 SFT 训练后,MTP 接受率随模型规模增大而持续提升。Qwen3.7 模型使用端到端 TV 损失训练,而 Qwen3.6 模型使用 CE 损失训练。接受率最高可达 95%,尤其在智能体任务上,表明草稿模型已几乎收敛至骨干模型。相反,随着模型规模减小,接受率出现不同程度的下降。

6.2 TV 损失在 RL 训练中稳定 MTP 加速效果

我们在 RL 设置下进行了大量实验,以证明 Bebop 的有效性。我们选取了两种涵盖不同生成模式的有代表性的工作负载。

  1. (1)

    推理强化学习:包括数学推理、代码推理和指令遵循在内的长链思维任务,最大生成长度为 64K 个模型 token。评估基准:HMMT25 (Dekoninck 等人,2026)、AIME25 (Zhang 和 Math-AI,2025) 以及 LiveCodeBench (Jain 等人,2025)。

  2. (2)

    SWE 强化学习:多轮代码编辑任务,每轮涉及思考、工具调用和工具执行,工具响应会追加到之前的上下文中。最大生成长度为 128K 个模型 token,最多 200 轮。评估基准:SWE-Verified (Jimenez 等人,2024)。

对于所有强化学习实验,我们使用 SGLang (Zheng 等人,2024) 作为基于 veRL (Sheng 等人,2024) 构建的异步强化学习框架中的推理引擎,学习率为 或 。

Refer to caption
(a) 推理强化学习。
Refer to caption
(b) SWE 强化学习。
Refer to caption
(c) Qwen-3.7 Max 中的 SWE 强化学习。
图 6:Qwen3.6-Plus 和 Qwen3.7-Max 在不同工作负载的强化学习训练过程中的接受长度。与仅目标采样(TO)和带 CE 损失的拒绝采样(RS w/ CE)相比,带 TV 损失的拒绝采样(RS w/ TV)始终能维持更高的接受长度。
媒体内容 · 前往原文查看
表 3:在不同模型上,各任务和训练目标下拒绝采样的 MTP 接受率(%)。Qwen3.7 模型使用端到端 TV 损失训练;其他所有模型均使用 CE 损失训练。
模型 数学 代码 混合 SWE 智能体 长周期 MTBench
Qwen3.7-Max
Qwen3.7-Plus
Qwen3.6-Plus
Qwen3.6-27B
Qwen3.6-35A3B
Refer to caption
(a) 推理强化学习。
Refer to caption
(b) SWE 强化学习。
Refer to caption
(c) 智能体强化学习。
图 7:使用 Qwen3.6-35A3B 和 Qwen3.6-Plus 进行强化学习时的训练延迟对比。与不使用 MTP(w/o MTP)和仅目标采样(TO)的训练相比,带拒绝采样的 MTP(RS w/ TV)显著降低了每步延迟。
Refer to caption
(a) 推理强化学习
Refer to caption
(b) SWE 强化学习
Refer to caption
(c) Qwen-3.7 Max 中的 SWE 强化学习
图 8:Qwen3.6-Plus 和 Qwen3.7-Max 在三种强化学习工作负载下的熵损失与接受长度关系。每个点代表一个训练步骤;线条表示线性拟合。TO 和 RS w/ CE 表现出强烈的负相关性(斜率 ),而 RS w/ TV 几乎保持平坦(斜率 ),证实了 TV 训练将接受率与熵解耦。

图 6 展示了强化学习训练过程中接受长度的变化趋势。通过拒绝采样和 TV 损失,Bebop 在整个训练过程中保持了稳定或持续提升的接受长度,即使策略模型保持了较高的熵值。在推理强化学习中,观察到的接受率提升主要是由训练过程中策略熵值的显著下降所驱动,而非仅仅依靠草稿对齐的改善。相比之下,SWE 工作负载的熵值略有上升,这使得它们成为检验训练目标鲁棒性的更直接测试:在此场景下,带 TV 损失的拒绝采样保持了稳定的接受长度,而仅以目标模型为目标的采样则持续退化。这种优势在 SWE 及其他高熵任务上最为显著,因为更高的接受长度直接转化为更快的 rollout 完成速度。此外,在更大的模型规模下(图 6c),带 TV 损失的拒绝采样表现出更强的熵不变趋势,并在整个强化学习训练过程中维持了高接受率,而仅以目标模型为目标的采样则显示出接受率的持续下降。

图 7 展示了相应的延迟改善。与未使用 MTP 的训练相比,采用拒绝采样的 MTP 实现了每一步强化学习训练延迟的降低,其中在智能体任务上的收益最大,其 rollout 阶段在智能体强化学习中实现了高达的加速。这些加速效果在所有工作负载上都是一致的,并在大规模部署中提供了显著的挂钟时间节省。

图 1a 和图 8 验证了第 3 节中建立的线性熵-接受度关系。值得注意的是,使用 TV 损失进行训练显著降低了熵-接受度斜率(降低了超过,例如从降至),并提高了截距。这证实了 TV 损失通过更好地使草稿分布与目标分布对齐,以及很大程度上将接受率与目标熵值解耦,从而提高了接受度,这与第 4.3 节分析的熵不变失配结构一致,进而在整个强化学习训练过程中实现了稳定的 MTP 加速收益。

6.3 在强化学习过程中更新 MTP 权重的优势

经过彻底的多阶段 SFT 训练后,模型已经实现了较高的接受率(例如,Qwen3.7-Max 的接受率超过 75%)。只要维持住接受率,MTP 加速收益就能在整个 RL 训练过程中得以保持。此外,第 4 节的分析以及图 1a 中的实验验证表明,采用 TV 损失的拒绝采样能够有效解耦熵与接受率之间的关系,从而在 RL 训练期间稳定接受率。为了进一步量化在 RL 训练期间更新 MTP 权重的收益,我们比较了以下训练配置:

  1. (1)

    RS w/ TV + TV 损失:从 RS w/ TV 检查点开始,并采用 TV 损失进行在线 MTP 训练;

  2. (2)

    RS w/ TV + CE 损失:从 RS w/ TV 检查点开始,并采用 CE 损失进行在线 MTP 训练;

  3. (3)

    TO + CE 损失:从 TO 检查点开始,并采用 CE 损失继续进行 MTP 训练。

Refer to caption
(a) 更新 MTP 权重时的接受长度。
Refer to caption
(b) 接受率增量与吞吐量比值的对比。
图 9:(a) 在 RL 训练期间,更新与不更新 MTP 权重时的接受长度。使用 CE 损失更新 MTP 权重会导致接受率向对应的未更新基线收敛,而仅使用 CE 损失更新目标采样甚至会因分布不匹配而降低接受率。(b) 在 8 个模型和 3 个任务上,接受率增量(RS 减去 No-RS)与吞吐量加速比(RS 除以 No-RS)的关系。拒绝采样带来的更高接受率增益直接转化为更大的吞吐量提升。

如图 9a 所示,随着采用 MTP 权重更新的 RL 训练推进,接受率会收敛至无权重更新时的对应基线水平。例如,尽管 RS w/ TV 因 TV 损失训练初始阶段获得了更高的接受率,但在 RL 训练期间使用 CE 损失更新 MTP 权重会导致接受率下降至 RS w/ CE 的水平。这种接受率的变化反映了草稿分布的改变:如 §7.2 所分析,CE 损失更新使 RS w/ TV 的草稿分布变得更加平滑,从而更接近 RS w/ CE 的分布。此外,对于已经充分训练的 MTP 权重,在 RL 阶段进一步更新参数不会带来显著改善,其接受率与未更新的基线高度吻合。在仅目标采样的情况下,使用 CE 损失更新甚至可能因草稿模型与目标模型之间的分布不匹配而导致接受率下降。

7 讨论

在本节中,我们将深入分析端到端 TV 损失与拒绝采样背后的机制,包括 TV 损失的分布效应、不同接受方法的鲁棒性对比,以及温度参数、生成长度和智能体工作负载对 MTP 接受率的影响分析。

7.1 TV 损失使草稿分布更尖锐

我们分析了与 CE/KL 训练相比,TV 损失如何影响草稿分布的熵。TV 损失产生的草稿分布其熵更接近目标熵(但略高),这表明草稿分布变得更加尖锐,且与目标模型峰值预测的对齐程度更高。相比之下,CE/KL 训练往往会产生更平滑的草稿分布,将概率质量分散在整个词表上,这对于拒绝采样而言并非最优——因为拒绝采样需要通过匹配目标分布的形状来最大化重叠区域。

这种锐化效应源于 TV 损失的梯度选择性行为(公式 (11)):它将优化重点集中在决策边界附近的模型 token 上,而忽略无关的低概率 token。图 10 展示了各模型中草稿与目标熵差与 KL 距离之间的关系。拥有训练良好的 MTP 头的模型,其草稿分布与目标分布之间的熵差更小,同时 KL 距离更大(另见图 1b)。

Refer to caption
(a) 熵差与 KL 散度。
Refer to caption
(b) 熵差与 RS 接受率。
Refer to caption
(c) KL 散度与 RS 接受率。
图 10:(a) 各模型与任务上的熵差对比。(b) 熵差与 RS 接受率呈负相关。(c) KL 散度未显示此类相关性,表明熵差(而非 KL)才是 RS 接受率的相关预测指标。

7.2 不同 MTP 训练损失导致不同的草稿分布模式

Refer to caption
图 11:在 RL 训练过程中,采用不同 MTP 损失目标时 MTP 指标的演变。TV 损失产生的草稿熵更接近目标,但 KL 距离更大, 更低, 更高。在 RL 过程中切换 MTP 训练损失会导致指标向新损失的特征模式偏移。

图 11 展示了在 RL 过程中使用不同损失更新 MTP 权重时,各项 MTP 指标的变化情况。与 CE 损失相比,TV 损失产生的草稿熵更接近目标模型,但 KL 距离更大。此外,由于 TV 损失产生更锐利的草稿分布,相应的 更低而 更高。当在 RL 过程中使用不同损失进行 MTP 权重更新时,MTP 指标会向该损失的特征模式偏移。例如,在使用 RS w/ TV + CE 损失时,草稿熵在训练过程中逐渐增加。

7.3 策略更新下接受方法的鲁棒性

尽管 §5.1 的分析表明 RL 过程中模型更新的幅度相对较小,但仅目标采样与拒绝采样在对 RL 策略更新引起的排序变化的敏感性上存在重要区别。

仅目标采样对排序变化较为脆弱。

仅目标模型接受依赖于草稿 token 是否落在目标模型的高概率区域(例如 top-)。这是一个离散标准:token 要么被接受,要么被拒绝。当强化学习梯度步骤导致 top-1 token 发生变化时,即使概率变化很小(例如,从 降至 ,而 从 升至 ),仍偏好旧 top-1 的草稿模型会经历从接受到拒绝的不连续跳变。

拒绝采样会平滑地退化。

在拒绝采样下,接受率是两个分布的连续函数。同样的排序变化对 TV 重叠度产生的影响可忽略不计,因为 对小的概率交换几乎不变。

高熵会放大脆弱性差距。

当目标熵较高时,多个 token 具有相似概率,使得排序变化在强化学习更新下更频繁。这会对仅目标模型采样产生不成比例的影响,因为每次排序翻转都可能导致离散的接受失败。尽管存在这种定性差异,我们凭经验观察到仅目标模型采样和拒绝采样具有相似的熵-接受率斜率(;参见第 3 节),这表明仅目标模型采样的离散脆弱性被累积的 TV 距离增长所抵消,而后者在 CE/KL 训练下同样影响拒绝采样。

7.4 温度与 MTP 接受率之间的相关性

采样温度直接影响目标模型的熵: 随 单调递增。结合第 3 节建立的线性熵-接受率关系,这意味着更高的温度会导致更低的 MTP 接受率。

Refer to caption
(a) 接受长度 vs. 温度。
Refer to caption
(b) 接受率 vs. 输出长度。
图 12:(a) 平均接受长度作为采样温度的函数。拒绝采样保持相对稳定的接受长度,而仅目标模型采样在较高温度下急剧退化。(b) MTP 接受率 vs. 输出长度(在 8 个模型上取平均)。RS 在所有生成位置上相对于仅目标模型采样保持稳定的优势。
Refer to caption
图 13:各模型的拒绝采样决策边界(见 §7.5)。几乎所有模型-任务组合都落在拒绝采样更优区域,证实拒绝采样对几乎所有实际 MTP 部署均有益处。

图 12a 证实了这一点:拒绝采样在不同温度下保持相对稳定的接受长度,而仅目标模型采样在温度升高时接受长度急剧下降。这对强化学习训练具有实际意义——训练中常使用较高温度来鼓励探索。我们的分析为理解通过温度缩放进行探索所带来的吞吐量成本提供了量化框架。

7.5 拒绝采样决策边界

当满足条件 且 时(见 §E),拒绝采样优于仅目标模型采样。该决策边界提供了一个简单的诊断方法:若草稿模型与目标模型之间的总变分距离小于目标模型在草稿模型 top-1 token 之外的概率质量,则应优先选择拒绝采样。

图 13 展示了八个原生训练 MTP 头的模型在该边界上的分布,涵盖三类任务。几乎所有模型-任务组合(24 组中的 23 组)都明确落在拒绝采样更优区域,证实对于原生 MTP 模型,拒绝采样始终优于仅目标模型采样。这确认了启用拒绝采样对几乎所有实际 MTP 部署均有益处。

7.6 生成长度与 MTP 接受率之间的相关性

如图 12b 所示,我们观察到 MTP 接受率随生成序列中的位置呈现系统性变化。在早期位置(靠近提示词处),目标模型通常具有较低的熵(即更可预测的延续),从而带来更高的接受率。随着生成过程推进,尤其是在具有长思维链的推理任务中,熵可能增加,接受率可能下降。这种位置依赖的接受模式表明,采用自适应 MTP 策略——根据估计的局部熵调整草稿长度——可进一步提升吞吐量。

7.7 智能体强化学习与气泡问题

如图 14a 所示,在智能体强化学习场景(例如 SWE-bench(Jimenez 等人,2024))中,模型会生成长篇幅、多轮次的交互,涉及工具调用、代码执行和迭代优化。这些场景呈现出特别长的生成长度和可变的熵分布,导致接受率出现周期性波动,且这种波动往往随着生成过程的推进而加剧。

Refer to caption
(a)智能体强化学习过程中的接受长度。
Refer to caption
(b)基于 top-k 截断的多 token 预测损失。
图 14:(a)智能体强化学习过程中的接受长度。平均接受长度稳定在 3.7,而最小-最大值范围则显示出跨步骤的周期性波动。(b)不同 top-k 截断值下的多 token 预测损失曲线。较小的 k 值会导致明显的损失尖峰和训练不稳定,而即使较大的 k 值,其收敛速度也比全词汇表的总变差损失更慢。

多 token 预测在智能体场景中尤其有益,原因有二:(1)长篇幅生成包含大量结构化输出——例如样板代码、工具调用格式和重复模式——这些内容高度可预测,从而在这些片段中获得高接受率;(2)多轮交互和长尾生成会降低有效运行批次大小,在这种状态下,由于推理引擎远未达到计算饱和,多 token 预测的延迟优势会被放大。事实上,我们的实验表明,通过我们提出的总变差损失训练,智能体工作负载实现了最大的接受率提升(5%)。

7.8 Top-K 总变差近似的稳定性问题

计算全词汇表的总变差损失在大词汇量下会产生较高的峰值内存。为解决此问题,我们采用了一种融合反向传播内核来减少中间激活的大小(参见 §F)。我们还尝试通过 top-k 截断来近似总变差损失,以进一步降低峰值内存。然而,即使使用较大的 k 值,我们也观察到损失收敛速度略有减慢和性能下降。较小的 k 值会导致明显的损失尖峰,如图 14b 所示。最终,我们采用了融合的全词汇表总变差损失,而非 top-k 近似方法。

8 相关工作

推测性解码。

推测解码(Leviathan 等人,2023;Chen 等人,2023)通过使用轻量级草稿模型提出多个 token,再由目标模型并行验证,从而加速自回归大语言模型的推理。目前已提出多种草稿架构,包括独立小模型(Miao 等人,2024;Shen 等人,2026)、早退头(Elhoushi 等人,2024)、辅助头(Cai 等人,2024;Li 等人,2024;2026)、MTP 头(DeepSeek-AI,2024;Gloeckle 等人,2024;Qwen Team,2026a)以及扩散模型(Chen 等人,2026a)。Bebop 专注于共享主干网络隐藏状态的 MTP 头,并分析其在强化学习训练动态下的行为。

大语言模型的强化学习。

强化学习已成为将大语言模型与人类偏好对齐(Schulman 等人,2017)以及增强推理和智能体能力(OpenAI,2026;DeepSeek-AI,2026;Qwen Team,2026b)的核心手段。现代后训练流程通常将强化学习拆分为展开(rollout)、奖励评估和策略更新三个阶段,而 GRPO(Shao 等人,2024)和 GSPO(Zheng 等人,2025)等算法则改进了优化目标本身。在系统层面,异步或部分展开框架通过将推理工作节点与训练工作节点解耦,减少了长尾轨迹带来的空闲时间(Fu 等人,2025;Wang 等人,2025;THUDM,2025;Qin 等人,2025)。然而,这些方法主要隐藏了长尾气泡,在长上下文、多轮交互和工具使用场景下,轨迹生成仍然是瓶颈。相关研究探讨了由训练-推理差异和策略陈旧性导致的强化学习不稳定性(Yao 等人,2025;Liu 等人,2025);近期的 MTP 方法则通过在线更新草稿头来解决草稿-目标不匹配问题(Chen 等人,2026b;Iso 等人,2026;MiniMax,2026b;Li 等人,2025)。然而,我们发现强化学习过程中接受率波动的主要驱动力是目标模型熵值的变化,而非草稿-目标不匹配,并且目标熵值与 MTP 接受长度呈线性关系,这一观察结果也被 Xiao 等人(2026)注意到。我们的工作具有互补性:它在不改变强化学习目标或调度器的情况下加速了展开过程,并指出熵值变化是导致 MTP 接受率下降的主导因素。

机器学习中的总变差距离。

TV 距离是用于比较概率分布的标准度量,已被应用于分布测试(Canonne, 2020)、生成式建模(Nowozin 等人, 2016)以及马尔可夫链的收敛性分析(Levin 和 Peres, 2017)。在投机解码中,拒绝采样的接受率等于分布重叠度,即(Leviathan 等人, 2023;Chen 等人, 2023)。这一关联催生了以接受率为导向的目标函数,包括用于直接优化投机解码接受率的 LK 损失(Samarin 等人, 2026)。然而,这些工作聚焦于固定目标模型的推理时投机解码。近期研究也开始探索在 OPD 中使用反向 KL 散度来优化学生模型(Lu 和 Lab, 2025;Lei 等人, 2026),但其训练目标与直接最大化拒绝采样接受率仍有本质差异。据我们所知,我们是首个提出将 TV 距离直接作为 MTP 头训练目标的研究,也是首个分析其在强化学习训练过程中行为的工作。

9 结论

我们提出了 Bebop,这是一项关于大语言模型强化学习背景下多 token 预测(MTP)的系统性研究。我们的分析揭示了三个关键发现:(1)在仅目标模型和拒绝采样两种场景下,MTP 接受率均受目标模型熵的线性约束;(2)Bebop 的端到端 TV 损失直接优化多步拒绝采样接受率,相比传统 CE/KL 目标函数,接受率最高可达 95%,推理吞吐量提升最高达 25%;(3)在强化学习训练前,使用 TV 损失和拒绝采样进行轻量级预适应,足以在整个强化学习训练过程中维持高 MTP 接受率,从而消除了对高成本的在线 MTP 更新的需求。基于 Qwen3.5、3.6 和 3.7 模型的大量实验表明,Bebop 在异步强化学习流水线中可实现端到端加速。

局限性。

我们对熵-接受率关系的理论分析依赖于建模假设(均匀失配与概率比例失配),这些假设是基于梯度结构的启发式动机,而非经过严格证明;收紧这些假设仍是一个开放性问题。此外,TV 训练所保证的熵不变性是分布条件性的:它仅在 SFT 训练数据所覆盖的熵范围内成立,但当 RL 探索将策略熵显著推至该范围之外时,草稿头会遇到分布外的目标分布,此时失配比不再有界,从而恢复出与 CE/KL 训练相当的熵-接受率依赖关系。在这种情况下,建议在 RL 期间使用 TV 损失进行 MTP 联合训练,以将草稿头的有效覆盖范围扩展到新的熵区间。

参考文献

  • Anthropic [2026] Anthropic. Claude fable 5 and claude mythos 5, 2026. URL https://www.anthropic.com/news/claude-fable-5-mythos-5.
  • Cai et al. [2024] Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple LLM inference acceleration framework with multiple decoding heads. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https://openreview.net/forum?id=PEpbUobfJv.
  • Canonne [2020] Clément L. Canonne. A survey on distribution testing: Your data is big. but is it blue? Theory of Computing, 9:1–100, 2020.
  • Chen et al. [2023] Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. In International Conference on Machine Learning, 2023.
  • Chen et al. [2026a] Jian Chen, Yesheng Liang, and Zhijian Liu. Dflash: Block diffusion for flash speculative decoding. ArXiv preprint, abs/2602.06036, 2026a. URL https://arxiv.org/abs/2602.06036.
  • Chen 等人 [2026b] Qiaoling Chen, Zijun Liu, Peng Sun, Shenggui Li, Guoteng Wang, Ziming Liu, Yonggang Wen, Siyuan Feng, 和 Tianwei Zhang。Respec:面向强化学习系统中推测解码的优化。第九届机器学习与系统会议,2026b。URL https://openreview.net/forum?id=HhDSxs7x2R。
  • DeepSeek-AI [2024] DeepSeek-AI。DeepSeek-V3 技术报告。ArXiv 预印本,abs/2412.19437,2024。URL https://arxiv.org/abs/2412.19437。
  • DeepSeek-AI [2026] DeepSeek-AI。DeepSeek-V4:迈向高效百万 token 上下文智能,2026。
  • Dekoninck 等人 [2026] Jasper Dekoninck, Nikola Jovanović, Tim Gehrunger, Kári Rögnvaldsson, Ivo Petrov, Chenhao Sun, 和 Martin Vechev。超越基准:Matharena 作为大语言模型数学评估平台。ArXiv 预印本,abs/2605.00674,2026。URL https://arxiv.org/abs/2605.00674。
  • Elhoushi 等人 [2024] Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed Aly, Beidi Chen, 和 Carole-Jean Wu。LayerSkip:实现早期退出推理与自推测解码。第62届计算语言学协会年会论文集(第1卷:长论文),泰国曼谷,2024年8月。计算语言学协会。URL https://aclanthology.org/2024.acl-long.681/。
  • Fu 等人 [2025] Wei Fu, Jiaxuan Gao, Xujie Shen, Chen Zhu, Zhiyu Mei, Chuyi He, Shusheng Xu, Guo Wei, Jun Mei, Jiashu Wang, Tongkai Yang, Binhang Yuan, 和 Yi Wu。Areal:面向语言推理的大规模异步强化学习系统。ArXiv 预印本,abs/2505.24298,2025。URL https://arxiv.org/abs/2505.24298。
  • GLM 团队 [2026] GLM 团队。GLM-5.1:面向长周期任务,2026。URL https://z.ai/blog/glm-5.1。访问日期:2026-04-07。
  • Gloeckle 等人 [2024] Fabian Gloeckle、Badr Youbi Idrissi、Baptiste Rozière、David Lopez-Paz 和 Gabriel Synnaeve。通过多 token 预测实现更好更快的大语言模型。载于第四十一届国际机器学习大会,ICML 2024,奥地利维也纳,2024年7月21-27日。OpenReview.net,2024年。URL https://openreview.net/forum?id=pEWAcejiU2。
  • Iso 等人 [2026] Hayate Iso、Tiyasa Mitra、Sudipta Mondal、Rasoul Shafipour、Venmugil Elango、Terry Kong、Yuki Huang、Seonjin Na、Izzy Putterman、Benjamin Chislett 等。通过系统集成的推测解码加速强化学习后训练 rollout。ArXiv 预印本,abs/2604.26779,2026年。URL https://arxiv.org/abs/2604.26779。
  • Jain 等人 [2025] Naman Jain、King Han、Alex Gu、Wen-Ding Li、Fanjia Yan、Tianjun Zhang、Sida Wang、Armando Solar-Lezama、Koushik Sen 和 Ion Stoica。Livecodebench:大语言模型在代码方面的全面且无污染评估。载于第十三届国际学习表征大会,2025年。URL https://openreview.net/forum?id=chfJJYC3iL。
  • Jimenez 等人 [2024] Carlos E. Jimenez、John Yang、Alexander Wettig、Shunyu Yao、Kexin Pei、Ofir Press 和 Karthik R. Narasimhan。SWE-bench:语言模型能解决真实的 GitHub 问题吗?载于第十二届国际学习表征大会,ICLR 2024,奥地利维也纳,2024年5月7-11日,2024年。URL https://openreview.net/forum?id=VTF8yNQM66。
  • Kimi 团队 [2026] Kimi 团队。Kimi k2.6:推进开源编程能力,2026年。URL https://www.kimi.com/blog/kimi-k2-6。访问日期:2026-04-07。
  • Lei 等人 [2026] Haodi Lei、Yafu Li、Haoran Zhang、Shunkai Zhang、Qianjia Cheng、Xiaoye Qu、Ganqu Cui、Bowen Zhou、Ning Ding、Yun Luo 和 Yu Cheng。Draft-OPD:面向推测草稿模型的在线策略蒸馏。ArXiv 预印本,abs/2605.29343,2026年。URL https://arxiv.org/abs/2605.29343。
  • Leviathan 等人 [2023] Yaniv Leviathan、Matan Kalman 和 Yossi Matias。通过推测解码实现 Transformer 的快速推理。收录于《国际机器学习大会,ICML 2023,2023 年 7 月 23-29 日,美国夏威夷火奴鲁鲁》,机器学习研究论文集第 202 卷,第 19274–19286 页。PMLR,2023 年。URL https://proceedings.mlr.press/v202/leviathan23a.html。
  • Levin 和 Peres [2017] David A. Levin 和 Yuval Peres。《马尔可夫链与混合时间》。美国数学学会,第 2 版,2017 年。
  • Li 等人 [2025] Jiajun Li、Yuzhen Zhou、Mao Cheng 和 Ruiguo Yang Yang。在强化学习中提升推测解码能力,2025 年。URL https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/slime/spec/readme-en.md。
  • Li 等人 [2024] Yuhui Li、Fangyun Wei、Chao Zhang 和 Hongyang Zhang。EAGLE:推测采样需要重新思考特征不确定性。收录于《第四十一届国际机器学习大会,ICML 2024,2024 年 7 月 21-27 日,奥地利维也纳》。OpenReview.net,2024 年。URL https://openreview.net/forum?id=1NdN7eXyb4。
  • Li 等人 [2026] Yuhui Li、Fangyun Wei、Chao Zhang 和 Hongyang Zhang。EAGLE-3:通过训练时测试扩展大语言模型的推理加速。收录于《第三十九届神经信息处理系统年度大会,2026 年》。URL https://openreview.net/forum?id=4exx1hUffq。
  • Liu 等人 [2025] Jiacai Liu、Yingru Li、Yuqian Fu、Jiawei Wang、Qian Liu 和 Yu Shen。当速度破坏稳定性:从训练-推理不匹配角度揭秘强化学习崩溃,2025 年。URL https://richardli.xyz/rl-collapse。
  • Lu 和 Lab [2025] Kevin Lu 和 Thinking Machines Lab。同策略蒸馏。Thinking Machines Lab:联结主义,2025 年。doi: 10.64434/tml.20251026。https://thinkingmachines.ai/blog/on-policy-distillation。
  • Miao 等人 [2024] Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, 等. Specinfer: 基于树的推测性推理与验证加速大语言模型服务. 收录于第29届ACM国际编程语言与操作系统架构支持会议论文集, 第3卷, 页码 932–949, 2024.
  • MiniMax [2026a] MiniMax. MiniMax M2.5: 为真实世界生产力而生. https://www.minimax.io/news/minimax-m25, 2026a.
  • MiniMax [2026b] MiniMax. Forge: 可扩展的智能体强化学习框架与算法. MiniMax 新闻, 2026b. URL https://www.minimax.io/news/forge-scalable-agent-rl-framework-and-algorithm. 访问日期: 2026-06-09.
  • Nowozin 等人 [2016] Sebastian Nowozin, Botond Cseke, 和 Ryota Tomioka. f-gan: 使用变分散度最小化训练生成式神经采样器. 收录于 NeurIPS, 2016, 巴塞罗那, 西班牙, 页码 271–279, 2016.
  • OpenAI [2026] OpenAI. GPT-5.5 系统卡, 2026. URL https://openai.com/index/gpt-5-5-system-card/.
  • Qin 等人 [2025] Ruoyu Qin, Weiran He, Weixiao Huang, Yangkun Zhang, Yikai Zhao, Bo Pang, Xinran Xu, Yingdi Shan, Yongwei Wu, 和 Mingxing Zhang. Seer: 用于快速同步大语言模型强化学习的在线上下文学习. ArXiv 预印本, abs/2511.14617, 2025. URL https://arxiv.org/abs/2511.14617.
  • Qwen 团队 [2026a] Qwen 团队. Qwen3.5: 迈向原生多模态智能体. https://qwen.ai/blog?id=qwen3.5, 2026a.
  • Qwen 团队 [2026b] Qwen 团队. Qwen3.7: 智能体前沿, 2026b. URL https://qwen.ai/blog?id=qwen3.7.
  • Samarin 等人 [2026] Alexander Samarin, Sergei Krutikov, Anton Shevtsov, Sergei Skvortsov, Filipp Fisin, 和 Alexander Golubev. Lk 损失: 用于推测性解码的直接接受率优化. ArXiv 预印本, abs/2602.23881, 2026. URL https://arxiv.org/abs/2602.23881.
  • Schulman 等人 [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, 和 Oleg Klimov. 近端策略优化算法. ArXiv 预印本, abs/1707.06347, 2017. URL https://arxiv.org/abs/1707.06347.
  • Shao 等人 [2024] Zhihong Shao、Peiyi Wang、Qihao Zhu、Runxin Xu、Junxiao Song、Xiao Bi、Haowei Zhang、Mingchuan Zhang、Y. K. Li、Y. Wu 和 Daya Guo。Deepseekmath:在开放语言模型中突破数学推理的极限。ArXiv 预印本,abs/2402.03300,2024 年。URL https://arxiv.org/abs/2402.03300。
  • Shen 等人 [2026] Yuhao Shen、Junyi Shen、Quan Kong、Tianyu Liu、Yao Lu 和 Cong Wang。SpecBranch:通过混合草稿和回滚感知分支并行实现推测解码。载于第十四届国际学习表征会议,2026 年。
  • Sheng 等人 [2024] Guangming Sheng、Chi Zhang、Zilingfeng Ye、Xibin Wu、Wang Zhang、Ru Zhang、Yanghua Peng、Haibin Lin 和 Chuan Wu。HybridFlow:一个灵活高效的 RLHF 框架。ArXiv 预印本,abs/2409.19256,2024 年。URL https://arxiv.org/abs/2409.19256。
  • Shoeybi 等人 [2019] Mohammad Shoeybi、Mostofa Patwary、Raul Puri、Patrick LeGresley、Jared Casper 和 Bryan Catanzaro。Megatron-LM:利用模型并行训练数十亿参数的语言模型。ArXiv 预印本,abs/1909.08053,2019 年。URL https://arxiv.org/abs/1909.08053。
  • THUDM [2025] THUDM。SLiMe:一个用于强化学习扩展的大语言模型后训练框架。https://github.com/THUDM/slime,2025 年。
  • Wang 等人 [2025] Weixun Wang、Shaopan Xiong、Gengru Chen、Wei Gao、Sheng Guo、Yancheng He、Ju Huang、Jiaheng Liu、Zhendong Li、Xiaoyang Li、Zichen Liu、Haizhou Zhao 等人。面向大规模学习的强化学习优化:一个高效且用户友好的扩展库。ArXiv 预印本,abs/2506.06122,2025 年。URL https://arxiv.org/abs/2506.06122。
  • Wu 等人 [2026] Tianyu Wu、Yu Yao、Zhenting Qi、Han Zheng、Zhuohan Wang、Haoran Ma、Lawrence Liao、Himabindu Lakkaraju、Ju Li 和 Yilun Du。D-PACE:面向并行推测草稿的动态位置感知交叉熵。ArXiv 预印本,abs/2605.18810,2026 年。URL https://arxiv.org/abs/2605.18810。
  • Xiao 等人 [2026] Bangjun Xiao、Tianyang Lu、Weiji Zhuang 等人。MiMo-V2-Flash 技术报告。ArXiv 预印本,abs/2601.02780,2026 年。URL https://arxiv.org/abs/2601.02780。
  • Yang 等人 [2025] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, 等. Qwen3 技术报告. ArXiv 预印本, abs/2505.09388, 2025. URL https://arxiv.org/abs/2505.09388.
  • Yao 等人 [2025] Feng Yao, Liyuan Liu, Dinghuai Zhang, Chengyu Dong, Jingbo Shang, 和 Jianfeng Gao. 你的高效强化学习框架悄然为你带来了离策略强化学习训练, 2025. URL https://fengyao.notion.site/off-policy-rl.
  • Zhang 和 Math-AI [2025] Yifan Zhang 和 Team Math-AI. 美国邀请赛数学考试 (AIME) 2025, 2025.
  • Zheng 等人 [2025] Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, 等. 分组序列策略优化. ArXiv 预印本, abs/2507.18071, 2025. URL https://arxiv.org/abs/2507.18071.
  • Zheng 等人 [2023] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, 和 Ion Stoica. 以 LLM 作为评判者的评判方法:基于 MT-Bench 和 Chatbot Arena. 收录于 NeurIPS, 2023, 2023.
  • Zheng 等人 [2024] Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark W. Barrett, 和 Ying Sheng. SGLang:结构化语言模型程序的高效执行. 收录于 NeurIPS, 2024, 2024.

附录 A TV 损失梯度的推导

我们提供 TV 损失梯度(公式 (11))的完整推导。

设草稿头输出 logits 为 。目标模型概率被视为常数(已分离)。TV 损失为:

关于 的梯度为:

由于 是常数, 关于 的次梯度为:

利用 softmax 雅可比矩阵 的链式法则:

因此:

有界性。

由于 且 :

附录 B 与前向 KL 散度梯度的比较

作为比较,前向 KL 散度关于 的梯度为:

与 TV 损失梯度的关键区别:

  1. 1.

    前向 KL 梯度对每个 的 token 施加非零力,包括概率可忽略的 token。TV 梯度与 成正比,因此会自动忽略低概率 token。

  2. 2.

    前向 KL 散度的梯度无法区分在拒绝采样下哪些 token 会被接受、哪些会被拒绝。而全变差(TV)梯度则通过指示函数明确纳入了这一区分。

  3. 3.

    前向 KL 散度的梯度在(草稿模型过度自信时)可能很大。全变差(TV)梯度则被 所界。

附录 C:反向 KL 散度分析

前述分析聚焦于前向 KL 散度,它在常数项上等价于交叉熵损失。一个自然的问题是:反向 KL 散度是否更适合作为拒绝采样的训练目标。

梯度推导。

反向 KL 散度相对于草稿模型 logits 的梯度为:

(17)

梯度结构对比。

表 1 总结了三种梯度结构。

反向 KL 梯度与全变差(TV)梯度共享理想的 -比例性,这意味着低概率 token 会自动受到可忽略的优化压力。这表明反向 KL 产生的失配程度应比前向 KL 的均匀失配更与 成比例,因此其熵-接受率耦合会比前向 KL 更弱。

为何反向 KL 仍非最优。

尽管梯度结构有所改进,但反向 KL 在最大化拒绝采样接受率方面仍非最优,原因有三:

  1. 1.

    零强制行为。反向 KL 在 时即使 也不施加惩罚,因为 。这种“模式寻找”特性允许草稿模型丢弃 中的某些模式,直接放弃这些 token 上的重叠部分,从而降低接受率。相比之下,前向 KL 是“零回避”的(当 且 时 ),强制要求完整的支持集覆盖。全变差(TV)损失既非零强制也非零回避:它选择性地将容量分配给边际重叠改进最大的 token。

  2. 2.

    非对称的高估/低估惩罚。拒绝采样的接受率取决于 ,它对高估()和低估()施加对称惩罚——两者都会使重叠区域减少 。反向 KL 散度施加的是非对称惩罚:高估(,因此 )产生的梯度远强于低估。这促使草稿模型在大多数 token 上趋向于 ,虽然能保证单个 token 的接受概率,但会降低这些 token 的采样概率,从而导致总重叠区域次优。

  3. 3.

    间接优化目标。与前向 KL 散度类似,反向 KL 散度并不直接优化 。反向 KL 散度梯度中的对数比值提供的是软性、非线性的信号,而总变差梯度中的指示函数则提供与拒绝采样决策边界对齐的硬性、直接信号。

总结。

就优化拒绝采样接受率的适用性而言:

反向 KL 散度通过更好的容量分配(梯度 )改进了前向 KL 散度,但由于其零强制行为和非对称惩罚结构,仍然不是最优的。总变差损失直接优化了目标量,并避免了这两种失败模式。

附录 D 不同训练目标下的熵与接受率关系

我们详细分析了在不同接受方法和训练目标下,目标模型的熵如何约束多 token 预测(MTP)接受率。

D.1 设定与符号

考虑生成过程中的一个固定位置。设 为目标模型的分布, 为草稿模型的分布。草稿模型由 参数化,其 logits 为 且 。由于模型容量有限,草稿模型通常无法完美匹配 ,而每个 token 的失配结构关键取决于训练目标。

我们将 在阈值 下的有效支撑集定义为 ,并回顾有效支撑集大小与熵通过困惑度相关联:。

在下面的分析中,我们根据训练目标考虑两种失配结构:

  • •

    均匀失配(CE/KL 训练):其中 且 , 在各 token 上近似均匀(理由见 §D.3)。

  • •

    概率比例失配(TV 训练):其中绝对误差随 token 概率缩放(基于 §D.4 中的容量分配假设推导得出)。

D.2 仅目标采样

在仅目标采样下,草稿 token 被贪婪地选择为并以概率接受,得到接受率:

(18)

完美草稿情况。

对于训练良好的草稿模型,其中(即草稿正确识别了目标的 top-1 token),接受率简化为:

(19)

与香农熵的关系。

量是 的单调递减函数:随着熵增加,分布扩散且最大概率降低。一个标准边界给出,因此接受率的下界为。

线性化。

由于是 的光滑单调递减函数,我们可以写成某个递减函数。在平均运行熵附近进行一阶泰勒展开:

(20)

由于递减,因此,从而得到:

(21)

下界提供了斜率的数量级估计:。经验上,该线性近似在模型规模、任务和训练阶段中均表现出显著的鲁棒性(图 1a)。

不完美草稿修正。

在均匀逐 token 失配的不完美草稿下,当目标前两个概率之间的差距满足时,会发生排序错误。高熵分布中 top token 之间的差距更小,使得排序错误更频繁。当排序错误发生时,接受率从降至,引入了额外的熵相关损失。两种效应共同强化了负斜率,因此线性近似仍然成立,但斜率可能更陡:

(22)

其中斜率经验上与相当(见 §6),尽管两者源于不同机制:由 的集中性和排序不稳定性驱动,而由逐 token TV 残差的累积驱动。

D.3 使用 CE/KL 训练的拒绝采样

拒绝采样接受率为(式 (2))。我们分析 CE/KL 训练如何通过其均匀逐 token 失配结构产生熵相关的接受率。

梯度结构。

关于 logits 的梯度为(见 §B)。梯度幅度由 与 之间的绝对差值决定,而非 本身的幅度。在基于梯度的优化中,每个 token 接收到的优化压力与 成正比,无论 还是 。这种均匀压力会产生近似均匀的逐 token 失配:当 时。

TV 距离推导。

在均匀逐 token 失配条件下:

(23)

该求和分解为有效支撑集及其补集两部分:

(24)

对于补集项,由于在有效支撑集之外 且 ,我们有 ,因此:

(25)

这是一个与 无关的小常数(因为 和 的大部分概率质量都集中在有效支撑集内)。因此,与熵相关的贡献来自有效支撑集,在该支撑集内失配在 层面完全实现。对于 且 :

(26)

因此:

(27)

线性近似。

在 在中等范围内变化(例如在 RL 训练期间)的情况下,指数函数可以通过围绕 的一阶泰勒展开进行线性化:

(28)

代入:

(29)

其中 和 为正的常数。这解释了在 CE/KL 训练下经验观察到的熵与接受率之间的线性负相关关系。

直观理解。

CE/KL 训练将优化资源均匀分配给所有 token。当 较低时, 集中在少数几个 token 上,草稿模型只需准确匹配这些 token——其余 token 上的加性误差对 的贡献可忽略不计。当 较高时, 分散到多个 token 上,均匀的加性误差累积成较大的 TV 距离。

为什么 CE/KL 训练对于拒绝采样是次优的。

Pinsker 不等式指出 ,关联了两种散度。然而,CE/KL 训练对于拒绝采样的次优性并非源于该界限本身的松散性,而在于 KL 梯度如何在词汇表上分配模型容量。

在均匀逐 token 失配条件下, 的二阶展开给出:

(30)
(31)

根据柯西–施瓦茨不等式,,这恢复了 Pinsker 界限 。当 时等号成立——即当 为均匀分布时该界限最紧。

根本问题在于容量分配。KL 梯度施加的优化压力与绝对差值成正比,将有限的模型容量大致均匀地分配给所有 token,包括那些目标概率可忽略的 token。在这种均匀分配下,每个 token 贡献一个均匀的失配量,由此产生的 TV 距离与有效支撑集中的 token 数量成比例:

(32)

高熵分布将概率质量分散到更多 token 上(),累积更多的逐 token 残差,从而形成更大的 TV 距离,即使 KL 散度也在被最小化。这就是为什么 CE/KL 训练的草稿模型表现出强烈的负熵-接受率相关性:KL 目标函数并不区分哪些 token 对拒绝采样的接受决策至关重要,哪些无关紧要。

D.4 基于 TV 训练的拒绝采样

梯度结构。

TV 损失相对于 logits 的梯度为(公式 (11))。

关键观察:梯度与 成正比。这意味着:

  • •

    高概率 token( 较大)会收到强烈的梯度信号,并被精确优化。

  • •

    低概率 token()收到的梯度近乎为零,因此优化器不会在它们身上浪费容量。

TV 梯度作为自校正机制。

定义概率比 。TV 梯度(公式 (11))作为一种自校正反馈,驱动 :

  • •

    当 (即 )时:指示函数 ,因此

    (33)

    梯度下降会增大 ,将 向上推,使其趋向 。

  • •

    当 (即 )时:指示函数 ,因此

    (34)

    梯度下降会减小 ,将 向下推,使其趋向 。

在这两种情况下,TV 训练都会驱动 ,即 。此外,由于对于典型词表大小,softmax 在局部满足 ,因此单步梯度更新会产生

(35)

其中 是学习率。校正幅度与 成正比:概率较大的 token 会收到更强的校正信号,确保在有效支撑集上 收敛到一个有界值 。尾部 token()收到的校正可忽略不计,但它们贡献的 TV 距离也可忽略不计。

这种自校正动态与 CE/KL 训练形成对比,后者的梯度将绝对差异均匀地推向零,而非将比率推向一。在有限容量下,CE/KL 均衡保持均匀性,这对应于有效支持中概率较小 token 的无限比率。

假设:有效支持上的 logit 比率误差有界。

上述自校正特性引出了以下假设。设 为在有限草稿容量下通过 TV 训练达到的解。由于式 (35) 中的校正幅度与有效支持中 token 的 成正比,这些 token 获得足够的梯度信号,将 驱动到有界区间内。

该假设在 log 比率空间 () 而非绝对空间 () 中表述,因为梯度下降作用于 logits ,且 softmax 满足 ,因此每个 logit 更新直接转化为 。由于 固定,:优化器的原生空间是 log 比率,因此均衡误差自然在 log 比率上有界。

我们假设:存在常数 ,使得对于所有 ,

(36)

尾部 token 可能具有较大的相对不确定性,但携带的概率质量可忽略不计,且贡献的 TV 距离可忽略。

推导失配界。

有界 logit 比率假设意味着

(37)

因此,对于有效支持中的每个 token,

(38)
(39)
(40)

令 ,我们得到

(41)

也就是说,在 TV 梯度容量分配所诱导的有界 logit 比率假设下,TV 训练产生概率比例失配 (),而非 CE/KL 训练的均匀失配 ()。在实践中,优化器动态(例如 Adam 的二阶矩归一化)可能会部分削弱原始的 比例性,因此比例失配应被视为一种建模近似,而非无条件定理。

TV 距离推导。

在具有常数 的概率比例失配下:

(42)

该界与 无关,得到:

(43)

这证明了命题 4。

实际考量。

上述分析假设 为常数,但在实际中,草稿头(draft head)的容量是有限的。当 增大时,有效支撑集(effective support)随之增长,要在更多 token 上保持均匀的相对精度,可能需要更大的模型容量。如果草稿头的容量不足, 可能会表现出较弱的熵依赖性,从而重新引入残余的(但已大幅衰减的)熵-接受率相关性。实验表明,在 TV 训练下,熵-接受率曲线的斜率相比 CE/KL 训练降低了超过 (例如, 对比 ),这证实了概率比例失配(probability-proportional mismatch)基本成立,但并非完美。

直观理解。

TV 训练将优化资源按每个 token 的概率成比例分配。当 较高且概率分布分散在多个 token 上时,每个 token 获得的优化努力相应减少,但它在 TV 距离中所占的权重也相应降低。这两个效应在很大程度上相互抵消,使得熵-接受率关系比 CE/KL 训练下弱得多。

附录 E 拒绝采样决策边界推导

我们推导出拒绝采样(rejection sampling)比仅使用目标采样(target-only sampling)获得更高接受率的条件。

接受率。

在仅使用目标采样下,接受率为 ,其中 是草稿模型(draft)的 top-1 token。在拒绝采样下,接受率为:

(44)

分解 。

利用恒等式 以及归一化条件 :

(45)
(46)
(47)

决策边界。

当 时,RS 优于仅使用目标采样:

(48)

这将两种接受方法的比较简化为一个简单的不等式:当草稿-目标 TVD(draft–target TVD)小于草稿贪婪预测之外的目标概率质量时,RS 更优。由于对于任何非退化分布,总存在一个足够对齐的草稿模型使得 RS 是有益的。

附录 F 融合 TV 损失核(Fused TV Loss Kernel)

我们提供了融合 TV 损失(TV loss)实现的伪代码。前向传播(算法 1)在单个内核启动中计算每个 token 的 TV 损失以及反向传播所需的辅助量。反向传播(算法 2)计算关于草稿 logits 的梯度。两个内核均以大小为 BLOCK_V 的 tile 为单位遍历词汇表,以限制寄存器和共享内存的使用,从而无需实例化 softmax 输出即可实现全词汇表的 TV 损失计算。

媒体内容 · 前往原文查看
算法 1 TV 损失前向内核(每个 token 位置)

1:草稿 logits,目标 log-probs

2:TV 损失,辅助标量

3:// 第 1 遍:数值稳定的 softmax 分母

4:全局 logit 最大值

5:exp 求和

6:// 第 2 遍:分块重叠与累加

7:;

8:对于 到 步长 BLOCK_V 执行

9:

10:草稿概率

11:目标概率

12:

13:

14:结束循环

15::可选裁剪

16:返回

媒体内容 · 前往原文查看
算法 2 TV 损失反向内核(每个 token 位置)

1:草稿 logits,目标 log-probs,缓存

2:梯度

3:对于 到 步长 BLOCK_V 执行

4:

5:

6:

7:

8:结束循环

9:返回

实现说明。

(1) 前向内核将 softmax 归一化与 TV 重叠计算融合在一起,避免了单独的 softmax 步骤。(2) 对于张量并行训练,在重叠计算之前通过跨 TP 秩的 all_reduce 计算 和 ;局部重叠值和 值在计算后同样进行规约。(3) 可选的 top-K 路径选择最大的草稿 logits,并仅在这些位置上计算 TV 和梯度,将内存从 降低到 ,且精度损失可忽略不计(因为尾部 token 的梯度)。

附录 G 拒绝采样推理实现

在生产推理引擎中为基于 MTP 的推测解码实现拒绝采样,需要修改草稿阶段和验证阶段。与仅基于目标采样的方法(该方法通过 选择草稿 token,并仅根据目标概率接受)不同,拒绝采样需要 (1) 从草稿分布中采样草稿 token(而非取 argmax),(2) 缓存草稿概率以供验证时使用,以及 (3) 在验证期间计算接受比率。我们描述了以下两种不同的实现策略。

G.1 多项草稿采样(SGLang)

第一种方法已在 SGLang 中实现(https://github.com/sgl-project/sglang/pull/26312),它直接使用多项采样从草稿分布中抽取草稿 token。

草稿阶段。

我们不通过 选择草稿 token,而是对草稿 logits 应用温度缩放,并通过多项采样进行采样。完整的草稿概率向量会与每个草稿 token 一起缓存,用于验证阶段。

验证阶段。

给定一条包含已缓存草稿概率的草稿 token 链,以及通过单次目标模型验证获得的目标概率,我们通过一个融合的 Triton 内核实现拒绝采样。该内核独立处理每个请求(每个请求对应一个 Triton 程序),并执行两个阶段:

  1. 1.

    顺序接受:对于每个草稿步骤 ,抽取 并在 时接受,即接受概率为 。在第一次拒绝时停止。

  2. 2.

    残差重采样:如果草稿 token 在步骤 被拒绝,或者所有草稿都被接受(奖励 token 情况),则从残差分布中采样下一个 token。对于步骤 的拒绝,残差分布为 ;对于奖励 token(全部接受),残差分布简化为 。内核通过词汇表上的两遍 CDF 逆变换计算该值:第一遍计算归一化常数 ,第二遍找到 token 使得累积和首次超过均匀随机变量 。

媒体内容 · 前往原文查看
算法 3 链式拒绝采样验证(多项 / SGLang)

1:草稿 token ;草稿概率 ;目标概率

2:已接受 token 计数 ;位置 的输出 token

3: 假设全部接受

4:对于 到 执行

5:

6: 如果 则拒绝

7: ;跳出循环

8: 结束条件

9:结束循环

10:// 通过两遍 CDF 逆变换进行残差重采样

11:如果 则在步骤 被拒绝

12: 对于所有

13:否则 奖励 token

14: 对于所有

15:结束条件

16: 第一遍:归一化

17:

18: 第二遍:CDF 逆变换

19:返回

内存开销。

主要开销在于缓存草稿概率向量:每个请求 ,其中 是 MTP 步骤数。

G.2 Gumbel-Max 技巧(vLLM)

第二种方法已在 vLLM 中实现(https://github.com/vllm-project/vllm/pull/35461),它利用 Gumbel-Max 技巧避免了残差重采样过程中的显式 CDF 逆变换。

草稿阶段。

草稿 token 使用 Gumbel-Max 技巧进行采样:对于每个词汇 token,计算 ,其中 是独立同分布的 Gumbel 噪声, 是采样温度。这等价于在温度缩放后从 中采样。经过温度缩放的草稿 logits(在添加 Gumbel 噪声之前)会被缓存起来用于验证。

验证阶段。

验证过程分为两个内核:

  1. 1.

    接受内核:一个顺序执行的 Triton 内核会遍历草稿步骤,根据缓存的目标概率和草稿概率计算 和 ,如果 成立则接受该 token,其中 是通过 tl.rand 生成的伪随机数,其种子由请求的随机种子和位置决定。该内核会记录第一个被拒绝步骤的索引。

  2. 2.

    残差 logits 内核:一个并行执行的 Triton 内核在 logit 空间中计算残差分布。对于在步骤 处被拒绝的情况:;对于奖励 token:(原始目标 logits)。然后使用与草稿阶段相同的 Gumbel-Max 采样方法,从该残差分布中重新采样得到最终的 token。

媒体内容 · 前往原文查看
算法 4 链式拒绝采样验证(Gumbel-Max / vLLM)

1:草稿 token ;草稿 logits ;目标概率 ;目标 logits

2:已接受 token 计数 ;位置 处的输出 token

3:// 内核 1:顺序接受

4:

5:for 到 do

6:

7:

8: if then

9: ;跳出循环

10: end if

11:end for

12:// 内核 2:残差 logits

13:if then

14: 对于所有

15:else

16: 对于所有

17:end if

18:// Gumbel-Max 重采样

19: 对于所有

20:

21:返回

来源:HuggingFace Daily Papers(社区热门论文) · arxiv.org