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

发表时间: 2026-06 · arXiv:2606.12370 (Qwen Team, Alibaba)

作者/机构: 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 (Qwen Team, Alibaba Inc.)

A1 主要贡献

本文介绍了一种名为 Bebop 的方法，系统性地研究了多词元预测（Multi-Token Prediction, MTP）在大型语言模型（LLM）强化学习（RL）后训练阶段的应用，并为将 MTP 集成到大规模 RL 流程中提供了实用方案。RL 训练的计算成本高昂，其中推理展开（rollout）阶段是主要瓶颈。虽然 MTP 通过推测解码为加速展开提供了自然方案，但许多研究观察到 MTP 的接受率在 RL 训练期间会显著下降，导致加速效果有限。

本文的核心问题是解决 RL 训练中 MTP 接受率下降的瓶颈。研究目标是通过深入分析 MTP 接受率下降的根本原因，提出新的训练目标和策略，以在整个 RL 过程中保持高且稳定的 MTP 接受率，从而显著加速 RL 训练。

主要贡献如下：
* 揭示 MTP 接受率的熵约束：本文首次揭示了 MTP 接受率与目标模型的熵波动之间存在根本性的制约关系。实验表明，在 RL 阶段，MTP 接受率随着模型熵的增加呈现出清晰的负线性下降关系（如图 1a 所示）。与普遍认为的策略更新导致分布不匹配是主要原因的观点不同，本文通过分解分析证明，熵的变化是导致接受率下降的主导因素。
* 提出端到端 TV 损失：针对传统 MTP 训练目标（交叉熵或 KL 散度）在概率性拒绝采样（Rejection Sampling）设置下的次优性，本文提出了一种新颖的端到端总变差（Total Variation, TV）损失函数。该损失函数直接优化多步拒绝采样的接受率，能产生对熵变化不敏感的 MTP 模型，使接受率提升约 10%，在多种任务中达到高达 95% 的接受率，并带来高达 25% 的额外推理吞吐量增益。
* 提出 RL 的 MTP 适应策略：本文证明，通过在 RL 训练前进行一次轻量级的 MTP 训练（使用 e2e TV 损失），并结合拒绝采样，即可在整个 RL 过程中保持稳定的高接受率和加速效果。由策略更新引起的分布不匹配影响可以忽略不计，从而无需在 RL 期间进行昂贵的在线 MTP 更新。
* 广泛的实验验证：通过在 Qwen3.5、Qwen3.6 和 Qwen3.7 等模型上进行的大规模实验，涵盖数学推理、代码生成和智能体任务，验证了 Bebop 的有效性。实验结果表明，该方法在异步 RL 训练流程中实现了高达 1.8 倍的端到端加速。


图 1: (a) 在 RL 过程中，MTP 接受率随策略熵波动呈线性下降；使用我们新颖的 e2e TV 损失训练的 MTP，在拒绝采样下很大程度上消除了这种熵依赖性。每个点代表在不同大小的 Qwen3.5、3.6 和 3.7 模型在各种任务的训练运行中，一个 RL 步骤的平均熵和接受长度。(b) TV 训练的 MTP 与策略模型实现了明显更好的分布重叠，从而带来优越的接受率和加速效果。

A3 背景知识

2.1 多词元预测与推测解码

MTP 作为一种推测解码范式。作为一种有效的推测解码范式【19, Fast inference from transformers via speculative decoding, 2023, ICML】【4, Accelerating large language model decoding with speculative sampling, 2023, ICML】，多词元预测（MTP）通过为自回归 LLM 增加轻量级的草稿头（draft heads）来顺序预测未来的多个词元【13, Better & faster large language models via multi-token prediction, 2024, ICML】【7, Deepseek-v3 technical report, 2024, ArXiv preprint】【46, Qwen3 technical report, 2025, ArXiv preprint】。设 $p(·|x, y_{<t})$ 表示目标（主干）模型在位置 $t$ 的下一词元分布，而 $q(·|x, y_{<t})$ 表示草稿头的预测分布。在推理过程中，MTP 采用“先草稿后验证”的模式：$\gamma$ 个草稿头链式地提出候选词元 $\hat{y}_{t+1}, \dots, \hat{y}_{t+\gamma}$，其中每个头都将前一个头的隐藏状态作为输入；然后，这 $\gamma$ 个候选词元在一次前向传播中与目标模型进行验证。</p>

接受长度决定吞吐量。每个验证步骤接受的词元期望数量，我们称之为接受长度，直接决定了推理吞吐量。这个接受长度取决于验证期间使用的具体接受方法，详见下一节。

2.2 接受方法

在推测解码中，通常使用两种接受方法：仅目标采样（Target-Only Sampling）和拒绝采样（Rejection Sampling）。图 13 展示了代表性模型在每种方法下的接受率分布。

仅目标采样。在这种方法下，草稿词元被贪婪地选择为 $\hat{y} = \arg \max_y q(y)$，并以概率 $p(\hat{y})$ 被接受，仅使用目标模型的概率。单步接受率为：

如果被拒绝，输出词元将从残差分布 $p_{\text{resid}}(y) \propto p(y) \mathbf{1}[y \neq \hat{y}]$ 中重新采样，确保整体输出分布无偏。值得注意的是，对于接受率相对较低的草稿模型，仅目标采样可能比拒绝采样产生更高的吞吐量，因为其更简单的接受标准避免了缓存和计算草稿概率向量的开销。

拒绝采样。在拒绝采样下【19, Fast inference from transformers via speculative decoding, 2023, ICML】【4, Accelerating large language model decoding with speculative sampling, 2023, ICML】，草稿词元 $\hat{y} \sim q(\cdot)$ 以概率 $\min(1, p(\hat{y})/q(\hat{y}))$ 被接受。期望的单步接受率为：

其中 $d_{\text{TV}}(p, q) = \frac{1}{2} \sum_y |p(y) - q(y)|$ 是总变差距离（Total Variation distance）【20, Markov Chains and Mixing Times, 2017, American Mathematical Society】。该方法提供了无偏保证：无论草稿质量如何，输出分布都精确等于目标分布 $p$。

2.3 LLM 的强化学习

标准 RL 框架。我们考虑 LLM 的标准 RL 框架，其中策略 $\pi_\theta$（即 LLM）为提示 $x \sim D$ 生成轨迹 $y$，并接收标量奖励 $R(x, y)$。我们采用 GRPO（Shao et al., 2024）【37, Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024, ArXiv preprint】，该方法为每个提示从展开策略 $\pi_{\theta_{\text{old}}}$ 中采样一组 $G$ 个轨迹 $\{y_1, \dots, y_G\}$，并优化以下裁剪后的代理目标：

其中，$r_{i,t} = \pi_\theta(y_{i,t}|x, y_{i,<t}) / \pi_{\theta_{\text{old}}}(y_{i,t}|x, y_{i,<t})$ 是重要性采样比率，$\hat{A}_i = (R(x, y_i) - \mu_G) / \sigma_G$ 是组归一化的优势。</p>


图 2: 在 SWE-bench RL 训练期间，使用 Qwen3.5-3.6 Plus 的每步 MTP 接受率。每条线代表一次独立的 RL 运行。后面的 MTP 步骤表现出逐渐增大的退化：在训练过程中，步骤 1 下降 1.2%，步骤 2 下降 2.6%，步骤 3 下降 3.5%。

RL 训练流程与瓶颈。LLM 的 RL 训练通常在一个包含三个阶段的循环中进行：(1) 展开，使用当前策略在推理引擎中生成轨迹，可能涉及多轮沙箱或工具交互；(2) 奖励，使用奖励模型或验证器评估这些生成的轨迹；(3) 更新，在训练引擎中使用策略梯度方法优化策略。异步 RL 或部分展开框架被普遍采用，以减轻展开过程中长尾轨迹造成的“气泡”开销【11, Areal: A large-scale asynchronous reinforcement learning system for language reasoning, 2025, ArXiv preprint】【42, Reinforcement learning optimization for large-scale learning: An efficient and user-friendly scaling library, 2025, ArXiv preprint】【40, Slime: An llm post-training framework for rl scaling, 2025, Github】【31, Seer: Online context learning for fast synchronous llm reinforcement learning, 2025, ArXiv preprint】【28, Forge: Scalable agent rl framework and algorithm, 2026, MiniMax News】。尽管有异步设计，展开阶段仍然是主要的计算瓶颈。虽然 MTP 提供了一个强大的加速范式来减轻这一负担，但其在 RL 环境中的直接应用暴露出独特的性能差距，需要进一步优化。

2.4 RL 训练期间 MTP 的退化

MTP 接受率下降现象。在 RL 训练期间，MTP 接受率在各个预测步骤中显著下降。如图 2 所示，越靠后的步骤下降幅度越大。每步接受率的下降范围从步骤 1 的 1.2% 到步骤 3 的 3.5%。

退化原因分析。近期工作【28, Forge: Scalable agent rl framework and algorithm, 2026, MiniMax News】【6, Respec: Towards optimizing speculative decoding in reinforcement learning systems, 2026, MLSys】【21, Power up speculative decoding in reinforcement learning, 2025, Github】主要将这种退化归因于分布不匹配。具体来说，由于主干权重更新导致草稿头落后，静态的草稿预测 $q = q_\phi(\cdot | x, y_{<t})$ 与演变的目标分布 $p = \pi_\theta(\cdot | x, y_{<t})$ 之间出现了差距。虽然这种不匹配确实存在，但我们认为这一观点是不完整的。我们识别出目标模型在 RL 训练期间的熵 $H(p)$ 变化是另一个根本驱动因素。这些熵变固有地改变了可实现的接受界限，而与草稿的准确性无关。这两个因素通过多步接受结构复合作用：<br /> 1. 单步退化：由于草稿-目标分布之间的持续分歧，TV 距离 $d_{\text{TV}}(p, q)$ 增大，导致每词元接受率 $\alpha_i$ 不断下降。
2. 多步复合：对于 $\gamma$ 步 MTP，期望接受长度涉及每步接受率的乘积，因此退化会乘法式地复合：$E[L] = \sum_{j=1}^{\gamma} \prod_{i=1}^{j} \alpha_i$。

关键洞察：熵是主导因素。至关重要的是，我们在 §3 和图 3 中的分解分析挑战了传统以不匹配为中心的观点。我们证明，熵驱动的成分实际上主导了 RL 训练期间接受率的波动。分布不匹配的成分相对较小。这一关键洞察重塑了我们对 MTP 退化的理解，并直接激发了我们后续的优化策略。

3. 目标熵对 MTP 接受率的约束

本节分析目标模型的熵如何从根本上约束 MTP 接受率，这解释了 RL 训练中由熵变化驱动的接受率退化。这也进一步激发了我们在 §4 中的训练目标。

3.1 公式化

熵的定义。考虑生成过程中的一个固定位置 $t$。设 $p \in \Delta^{|V|}$ 表示目标模型的下一词元分布， $q \in \Delta^{|V|}$ 表示草稿头的分布，其中 $V$ 是词汇表。我们定义目标熵为：

它衡量目标模型预测的不确定性。低熵表示一个自信、尖锐的分布，而高熵表示一个分散的分布。我们旨在理解 $H(p)$ 如何约束在公式 (1) 和 (2) 中定义的接受率 $\alpha_{\text{TO}}$ 和 $\alpha_{\text{RS}}$。

3.2 使用仅目标采样的 MTP

高低熵下的表现。在仅目标采样下，接受率取决于草稿的贪婪预测 $\hat{y} = \arg \max_y q(y)$ 与目标的高概率区域的对齐程度。当目标熵 $H(p)$ 较低（即 $p$ 在少数几个词元上呈尖峰）时，即使是中等准确的草稿模型也能通过将概率质量放在主导词元上实现高接受率。相反，当 $H(p)$ 较高时，目标分布散布在许多词元上，降低了 $\max_y p(y)$ 并增加了排名错误的概率。

命题 1 (仅目标采样下熵依赖的接受率)。对于一个训练良好的草稿模型，$\alpha_{\text{TO}} = \max_y p(y)$，这是 $H(p)$ 的一个单调递减函数，其下界为 $\exp(-H(p))$，并且经验上可以很好地近似为线性关系（图 1a）：

其中 $a_{\text{TO}}, b_{\text{TO}}$ 为正常数。不完美草稿下的排名错误会使斜率更陡峭，但仍保持线性关系（§D.2）。

证明简述。当草稿正确识别目标的前 1 个词元时（$\arg \max q = \arg \max p$），接受率简化为 $\alpha_{\text{TO}} = \max_y p(y)$。通过对凹函数对数应用詹森不等式，$\log(\max_y p(y)) \ge -H(p)$，即 $\max_y p(y) \ge \exp(-H(p))$。将 $\alpha_{\text{TO}} = f(H)$ 写成某个平滑递减函数 $f$ 的形式，并在参考熵 $\bar{H}$ 附近进行一阶泰勒展开：

由于 $f$ 是递减的，$b_{\text{TO}} = -f'(\bar{H}) > 0$。完整推导包括不完美草稿的修正，见 §D.2。□ 这种线性关系在不同模型大小、任务和训练阶段都非常稳健，如图 1a 所示。

3.3 使用拒绝采样的 MTP

接受率与 TV 距离。在拒绝采样下，接受率等于 $p$ 和 $q$ 之间的 TV 重叠（公式 (2)）。我们可以使用恒等式 $|a - b| = a + b - 2 \min(a, b)$ 和概率归一化来分解 TV 距离：

因此，最大化接受率等同于最小化 TV 距离：

结果是，接受率不再直接受策略熵的限制。然而，经验结果显示，切换到拒绝采样后，与熵的联系仍然存在。在进一步调查中，我们发现在使用 CE/KL 训练的草稿模型下，即使是微小的每词元不匹配，当 $p$ 具有高熵时也会累积，导致更大的 TV 距离。这促使我们对训练目标如何影响这种关系进行更深入的分析。

命题 2 (CE/KL 训练的拒绝采样下熵依赖的接受率)。在使用 CE/KL 训练的草稿模型下，拒绝采样接受率满足：

其中 $a_{\text{RS}}, b_{\text{RS}}$ 为正常数，其中 $b_{\text{RS}}$ 与 $b_{\text{TO}}$ 相当，但经验上略陡（§D.3, 图 8）。

证明简述。CE/KL 梯度 $q_j - p_j$ 产生均匀的每词元不匹配 $|\eta_v| \lesssim \sigma$。由于有效支持集大小约为 $|S_{\text{eff}}| \approx \exp(H(p))$，TV 距离累积为 $d_{\text{TV}} \approx \sigma \exp(H(p))$，从而得到 $\alpha_{\text{RS}} \approx 1 - \frac{\sigma}{2} \exp(H(p))$。在操作熵范围内对指数进行线性化，即可得到所述形式。详见 §D.3。□ 因此，在使用 CE/KL 训练的 MTP 下，拒绝采样和仅目标采样都对熵的变化敏感。由于策略熵在 RL 训练期间波动显著，这种敏感性从根本上限制了可实现的加速效果。

A2 方法细节

4. 为 RL 训练优化 MTP

如上所述，由于熵的限制，MTP 接受率在 RL 训练期间会显著下降。本节我们开发了新颖的端到端 TV 损失来应对这一挑战。

4.1 TV 损失：直接优化接受率

动机。传统的 MTP 训练最小化目标分布和草稿分布之间的交叉熵（CE）损失或 KL 散度。然而，拒绝采样的接受率由 TV 距离（公式 8）决定，而非 KL 散度。根据 Pinsker 不等式，$d_{\text{TV}}(p, q) \le \sqrt{D_{\text{KL}}(p\|q)/2}$，KL 散度只提供了一个间接的上界，最小化它并不能有效地最小化 TV 距离。这促使我们直接优化 TV 距离作为 MTP 的训练目标。

TV 损失。我们提出直接最小化 TV 距离：

其中 $p$ 被视为常数（从计算图中分离），梯度只通过 $q$ 传播。

梯度分析。设草稿头输出 logits $z \in \mathbb{R}^{|V|}$，其中 $q_j = \text{softmax}(z)_j$。TV 损失关于 $z_j$ 的梯度为：

命题 3 (有界梯度)。TV 损失的梯度是有界的：对所有 $j$，$\left| \frac{\partial \mathcal{L}_{\text{TV}}}{\partial z_j} \right| \le 1$。

证明。由于 $q_j \in [0, 1]$ 且 $|\mathbf{1}[q_j \le p_j] - S| \le 1$（因为指示函数和 $S \in [0, 1]$），我们有 $\left| \frac{\partial \mathcal{L}_{\text{TV}}}{\partial z_j} \right| = q_j \cdot |\mathbf{1}[q_j \le p_j] - S| \le 1$。□ 这种有界梯度属性确保了训练的稳定性，与 KL 散度形成对比，后者的梯度 $\frac{\partial D_{\text{KL}}}{\partial z_j} = q_j - p_j$ 在 $q$ 和 $p$ 显著不同时可能出现大幅值。

直观解释。TV 损失梯度在拒绝采样机制方面具有自然的解释：
* 对于 $q_j \le p_j$ 的词元（会被接受的词元）：梯度增加 logit，鼓励草稿分配更多概率质量。
* 对于 $q_j > p_j$ 的词元（会被拒绝的词元）：梯度减小 logit，抑制过分自信的预测。
* 对于 $q_j \approx 0$ 的词元（不相关的词元）：梯度自动地 $\approx 0$（因为它与 $q_j$ 成正比），避免了在词汇表长尾部分浪费优化精力。
这种选择性的梯度行为与 KL 散度形成对比，后者对所有词元都施加梯度，无论它们与接受决策的相关性如何。

表 1: 不同训练目标的梯度比较。C 表示一个全局常数（TV 为 S，反向 KL 为 DKL(q∥p)）。推导见附录 A-C。

CE、KL 和 TV 梯度的比较。表 1 总结了三种训练目标的梯度结构。关键区别在于梯度是否与 $q_j$ 成正比：CE 损失产生均匀的每词元不匹配 $(q_j - p_j)$，将优化精力均匀分布在整个词汇表上，包括不相关的低概率词元。相反，反向 KL 和 TV 损失都表现出与 $q_j$ 成正比的梯度，并具有自然的长尾抑制功能，将更新集中在草稿已经分配了不可忽略概率质量的词元上。然而，尽管有这个共同属性，反向 KL 相对于 CE 在接受率上的改进微乎其微（§6），因为它的“零强制”（zero-forcing）行为允许草稿丢弃 $p$ 的模式，并且其不对称惩罚全局性地驱动 $q \le p$——这两者都减少了 TV 重叠 $\sum_v \min(p, q)$（详见 §C 的详细分析）。TV 损失通过直接优化与接受率相关的量，并产生一个与概率成比例的不匹配，从而将接受率与目标熵解耦，避免了这些陷阱。

4.2 端到端多步 TV 损失

多步接受长度。对于 $\gamma$ 步 MTP，期望的接受长度是：

其中 $\alpha_i = 1 - d_{\text{TV}}(p_i, q_i)$ 是第 $i$ 步的每步接受率。直接优化平均每步 TV 距离 $\frac{1}{\gamma} \sum_{i=1}^{\gamma} d_{\text{TV}}(p_i, q_i)$ 并没有考虑到多步接受的乘法结构。

端到端 TV 损失。因此，我们提出了端到端（e2e）TV 损失：

这个损失直接优化归一化的期望接受长度，自然地给予较早的步骤更大的权重（因为它们出现在更多的乘积项中），并捕捉了多步验证的复合效应。这可以看作是一种动态的步进式加权方案：由于 $\alpha_i$ 取决于当前的草稿质量，每个位置的有效权重会随着训练的进行而自动调整，将重点转移到当前限制接受率的步骤上。这与先前使用固定位置相关权重的工作形成对比，例如依赖于头的损失权重【2, Medusa: Simple LLM inference acceleration framework with multiple decoding heads, 2024, ICML】【23, EAGLE-3: Scaling up inference acceleration of large language models via training-time test, 2026, NeurIPS】，指数衰减的块位置权重【5, Dflash: Block diffusion for flash speculative decoding, 2026, ArXiv preprint】，对被拒绝位置的固定衰减【18, Draft-opd: On-policy distillation for speculative draft models, 2026, ArXiv preprint】，或在 CE 基础上对每个位置加权【43, D-pace: Dynamic position-aware cross-entropy for parallel speculative drafting, 2026, ArXiv preprint】。

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

TV 损失的优越性分析。引入 TV 损失后，我们现在分析为什么它在 RL 的背景下从根本上优于 CE/KL 训练，因为在 RL 中目标熵是持续变化的。公式 (5) 和 (9) 中的线性关系描述了用 CE/KL 损失训练的草稿模型；我们表明，训练目标的选择从根本上改变了熵-接受率关系。完整推导见 §D；这里我们陈述主要结果。

Pinsker 不等式和 KL-TV 差距。根据 Pinsker 不等式：

$D_{\text{KL}}/2$ 只提供了 $d_{\text{TV}}$ 的一个上界，KL 优化在最小化 TV 距离方面效率低下：最小化 KL 散度并不能有效地最小化 TV 距离，而 TV 距离是直接决定拒绝采样接受率的量。

CE/KL 训练：均匀不匹配。KL 散度梯度 $\frac{\partial D_{\text{KL}}}{\partial z_j} = q_j - p_j$ 施加的优化压力与绝对差异 $|q_j - p_j|$ 成正比，而与 $p_j$ 相对于其他词元的大小无关。在一个容量有限的草稿模型下，这会产生近似均匀的每词元不匹配：对于一个常数 $\sigma$，$|q^*(v) - p(v)| \lesssim \sigma$。如命题 2 所示，这种均匀不匹配会在有效支持集 $|S_{\text{eff}}| \approx \exp(H(p))$ 上累积，产生一个依赖于熵的接受率。

TV 训练：概率成比例的不匹配。TV 损失梯度（公式 11）与 $q_j$ 成正比，将优化集中在高概率词元上，并自动忽略长尾。在一个容量有限的草稿模型下，每个词元获得的优化资源与其概率 $q_j \approx p_j$ 成正比，因此每词元不匹配也与 $p(v)$ 成比例，而不是保持在一个统一的水平。这会产生概率成比例的不匹配：对于一个常数 $\delta$，$|q^*(v) - p(v)| \lesssim \delta \cdot p(v)$（详见 §D.4 的详细推导）。

命题 4 (TV 训练下熵依赖性的降低)。当每词元不匹配满足 $|q^*(v) - p(v)| \lesssim \delta \cdot p(v)$ 时，TV 距离的界限与熵无关：

从而得到 $\alpha_{\text{RS}}^{\text{TV}} \ge 1 - \delta/2$。实际上，草稿头容量有限，所以 $\delta$ 可能表现出弱的熵依赖性 $\delta = \delta(H)$，但经验上，与 CE/KL 训练相比，熵-接受率斜率降低了超过 95%（图 8）。

证明简述。TV 梯度与 $q_j$ 成正比（公式 11），所以每个词元的优化资源与其概率成比例，产生 $|q^*(v) - p(v)| \lesssim \delta \cdot p(v)$（§D.4）。求和：$d_{\text{TV}} = \frac{1}{2} \sum_v |q^* - p| \le \frac{\delta}{2} \sum_v p(v) = \frac{\delta}{2}$，这个结果与熵无关，因为 $\sum_v p(v) = 1$。□ 这一分析解释了为什么经验上 TV 训练的草稿模型在不同的目标熵下能实现更稳定的接受率，而 CE/KL 训练的模型则表现出强烈的负相关性（图 8）。

5. RL 的 MTP 适应策略

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

5.1 分解：RL 中的熵与不匹配

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

其中 $b$ 是从每个实验早期阶段估计的熵-接受率斜率，$H_0$ 是初始熵，$\Delta \alpha_t = \alpha_t - \alpha_0$ 是在步骤 $t$ 的总接受率变化。第一项捕捉了仅由熵变化引起的接受率变化（假设草稿-目标关系固定），而残差项则捕捉了由于主干权重更新导致的草稿-目标不匹配增长的影响。

分解结果分析。如图 3 所示：(1) 在仅目标采样下，熵增加和不匹配增长都导致接受率下降，因为贪婪的草稿预测与不断演变的目标越来越不一致。(2) 在使用 CE 损失的拒绝采样下，退化几乎完全由熵驱动（$\Delta\alpha_{\text{mismatch}} \approx 0$），表明 RL 权重更新对草稿-目标 TV 重叠没有显著影响。(3) 在使用 TV 损失的拒绝采样下，所有分量都观察到接近零的变化，证实了 TV 训练的草稿对熵变化和权重更新都具有鲁棒性。

5.2 RL 前的适应已足够

关键实践洞察。分解分析得出了一个关键的实践洞察：由于在拒绝采样下，RL 权重更新引起的草稿-目标不匹配可以忽略不计，因此在 RL 期间更新 MTP 头是不必要的。


图 3: RL 训练期间接受长度变化的分解。$\Delta\alpha$ (总变化，灰色) 被分解为熵驱动分量 $\Delta\alpha_{\text{entropy}} = b \cdot (H_t - H_0)$ (橙色) 和草稿-目标不匹配分量 $\Delta\alpha_{\text{mismatch}}$ (绿色)。在仅目标采样下，熵增加和不匹配增长都导致接受率下降。在使用 CE 损失的拒绝采样下，退化几乎完全由熵驱动，不匹配接近于零。使用 TV 损失的 RS 在所有分量上都显示接近零的变化，证实了 TV 训练草稿的稳定性。

一次性预训练。在 RL 开始前的 SFT 阶段，应用一次性的 TV 损失预训练适应，就足以产生在整个 RL 训练期间保持高接受率的草稿模型（图 6）。这消除了维护 MTP 优化器状态的内存开销和 RL 期间 MTP 梯度更新的计算成本。

在线更新的负面影响。经验上，如图 9a 所示，从一个训练良好的 TV 检查点开始，在 RL 期间继续更新 MTP 权重并没有带来显著改善。更糟糕的是，在 RL 期间使用 CE 损失进行更新会导致接受率退化到 RS w/ CE 的基线水平，因为 CE 损失会使草稿分布更平滑，从而侵蚀 TV 训练带来的增益（§7.2）。

5.3 MTP 和主干的交叉训练

联合训练策略。当需要在 RL 期间进行 MTP 协同训练时（例如，对于不匹配不可忽略的仅目标采样），我们发现使用独立的学习率和独立的梯度范数归一化进行联合训练提供了最佳的权衡。主干的梯度不受 MTP 损失的影响（它只流经草稿头），确保 MTP 训练不会干扰主干的 RL 优化。

A4 实验环境

  • 模型架构:

    • 主要使用 Qwen3.5-35A3B、Qwen3.6-35A3B、Qwen3.6-Plus 和 Qwen3.7-Plus/Max 等一系列 Qwen 模型。
    • MTP 模块包含 5 个预测头，实验中主要验证 γ=3（一次验证 4 个词元）的情况。

  • 数据集:

    • SFT 阶段: 使用混合的 RFT（Reinforcement from Teacher）数据，以及领域特定的数据，包括代码、智能体和推理任务数据。MT-Bench【51, Judging llm-as-a-judge with mt-bench and chatbot arena, 2023, NeurIPS】用作分布外评估。
    • RL 阶段:
      • 推理 RL: 包含数学推理、代码推理和指令遵循的长思维链任务，评估基准为 HMMT25【9, Beyond benchmarks: Matharena as an evaluation platform for mathematics with llms, 2026, ArXiv preprint】、AIME25【49, American invitational mathematics examination (aime) 2025, 2025】 和 LiveCodeBench【15, Livecodebench: Holistic and contamination free evaluation of large language models for code, 2025, ICLR】。最大生成长度 64K。
      • SWE RL: 多轮代码编辑任务，涉及思考、工具调用和执行。评估基准为 SWE-Verified【16, Swe-bench: Can language models resolve real-world github issues?, 2024, ICLR】。最大生成长度 128K，最多 200 轮。

  • 硬件配置: 实验在大型 GPU 集群上进行，但未明确指出具体型号和数量。

  • 软件配置:
    • 代码实现: 基于 Megatron-LM【39, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, ArXiv preprint】进行 SFT 训练。
    • RL 框架: 使用 SGLang【52, Sglang: Efficient execution of structured language model programs, 2024, NeurIPS】作为 rollout 引擎，并构建在 veRL【38, Hybridflow: A flexible and efficient rlhf framework, 2024, ArXiv preprint】之上的异步 RL 框架。
    • RL 算法: GRPO【37, Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024, ArXiv preprint】。

A4 实验结果

6.1 多步 MTP 训练提升接受率

本节评估不同损失目标在 SFT 阶段对 MTP 接受率的影响，比较了 CE 损失、KL 损失、反向 KL 损失、TV 损失和 e2e TV 损失。

  • 拒绝采样接受率：如表 2 所示，在 Qwen3.5-35A3B 上，与 CE 基线相比，e2e TV 损失在所有任务上都显著提升了拒绝采样的接受率。在分布内任务（数学、代码、智能体、SWE）上提升 3-8%，在分布外任务 MT-Bench 上提升 2.3%。特别是在智能体任务上，接受率从 90.3% 提升至 97.0%。
  • TV 损失的优势：图 4 显示，在 Qwen3.6 和 Qwen3.7 模型上，TV 损失始终优于 CE 损失。CE 损失在训练中导致第一步接受率持续下降，而 TV 损失则保持稳定。e2e TV 损失的优势在后续 MTP 步骤中更为明显，在第 3 步比 CE 损失高出约 5%。智能体任务受益最大，提升高达 8%。
  • 仅目标采样接受率：如图 5 所示，在仅目标采样下，所有训练目标的接受率几乎相同（差异 < 0.3%），这证实了该方法的接受率主要取决于目标分布本身，而不是草稿分布的形状。
  • 吞吐量：如图 9b 所示，接受率的提升与吞吐量增益大致呈线性关系。使用 e2e TV 训练的 Qwen3.7 Plus 在所有数据集上始终优于使用 CE 损失训练的 Qwen3.6 Plus。
  • 模型规模效应：如表 3 所示，MTP 接受率随着模型规模的增大而提高。Qwen3.7 模型（使用 e2e TV 损失）的接受率最高可达 95%，尤其是在智能体任务上。

表 2: 在 Qwen3.5-35A3B 上，γ=3 时不同任务和训练目标下的 MTP 拒绝采样接受率 (%)。所有结果均为收敛时测量。∆ 表示相比 CE 损失基线的提升。


(b) 在智能体和混合任务上的接受长度 (Hybrid, Agent, Long-Horizon, SWE-Bench)。

图 4: SFT 训练期间 CE 损失 (实线) vs. TV 损失 (虚线)。TV 损失在所有 MTP 步骤中始终获得更高的接受率，在智能体任务上增益尤其显著。
图 5: SFT 训练期间使用 CE 损失 vs. TV 损失的仅目标采样接受长度。所有任务的接受率几乎相同 (<0.3% 差异)，证实了仅目标采样接受率取决于目标分布而非草稿的分布形状。

表 3: 在不同模型上，γ=3 时不同任务和训练目标下的 MTP 拒绝采样接受率 (%)。Qwen3.7 模型使用 e2e TV 损失训练；其他所有模型使用 CE 损失训练。

6.2 TV 损失稳定 RL 训练中的 MTP 加速

本节在 RL 场景下验证 Bebop 的有效性，涵盖推理 RL 和 SWE RL 两种工作负载。

  • 接受长度稳定性：图 6 显示，在 RL 训练中，Bebop（使用拒绝采样和 TV 损失）即使在策略保持高熵的情况下，也能维持稳定或提升的接受长度。相比之下，仅目标采样和使用 CE 损失的拒绝采样则表现出接受率下降。这种优势在 SWE 等高熵任务中尤为明显。
  • 训练延迟降低：图 7 显示，与不使用 MTP 的训练相比，使用拒绝采样的 MTP 将每步 RL 训练延迟降低了 1.5-1.8 倍。在智能体 RL 中，rollout 阶段的加速比高达 2.4 倍。
  • 熵-接受率关系验证：图 1a 和图 8 证实了 §3 中建立的线性熵-接受率关系。TV 损失训练显著降低了熵-接受率的斜率（超过 95%，如从 -1.68 降至 -0.06），并提高了截距。这证实 TV 损失通过更好地对齐分布和将接受率与目标熵解耦，实现了稳定的 MTP 加速。


图 6: 在 Qwen3.6-Plus 和 Qwen3.7-Max 的不同工作负载下，RL 训练期间的接受长度。使用 TV 损失的拒绝采样 (RS w/ TV) 始终比仅目标采样 (TO) 和使用 CE 损失的拒绝采样 (RS w/ CE) 保持更高的接受长度。


图 7: 在 Qwen3.6-35A3B 和 Qwen3.6-Plus 上进行 RL 期间的训练延迟比较。与不使用 MTP (w/o MTP) 和仅目标采样 (TO) 相比，使用拒绝采样的 MTP (RS w/ TV) 显著降低了每步延迟。


图 8: 在 Qwen3.6-Plus 和 Qwen3.7-Max 的三个 RL 工作负载中，熵损失与接受长度的关系。每个点代表一个训练步骤；线条显示线性拟合。TO 和 RS w/ CE 表现出强烈的负相关（斜率 ≈ -1.68），而 RS w/ TV 几乎保持平坦（斜率 ≈ -0.06），证实了 TV 训练将接受率与熵解耦。

6.3 RL 期间更新 MTP 权重的收益

本节研究在 RL 阶段在线更新 MTP 权重的必要性。

  • 在线更新无显著收益：如图 9a 所示，对于经过良好 SFT 训练的模型，在 RL 期间继续更新 MTP 参数并没有带来显著的接受率提升。接受率曲线紧密跟随不更新的基线。
  • 错误更新导致性能下降：从 RS w/ TV 检查点开始，如果在 RL 期间使用 CE 损失进行在线更新，接受率会逐渐退化到 RS w/ CE 的水平。这表明 CE 损失会侵蚀 TV 训练带来的分布优势。对于仅目标采样，使用 CE 损失更新甚至可能因分布不匹配而导致接受率下降。
  • 结论：在 RL 前进行一次充分的 MTP 训练（使用 TV 损失和拒绝采样）是足够且高效的策略，无需在 RL 期间进行在线更新。


图 9: (a) RL 训练期间有无 MTP 权重更新的接受长度。使用 CE 损失更新 MTP 权重导致接受率收敛到相应的不更新基线，而使用 CE 损失更新的仅目标采样甚至可能因分布不匹配而降低接受率。(b) 接受率增量 (RS - No-RS) vs. 吞吐量加速比 (RS / No-RS)，涵盖 8 个模型和 3 个任务 (r = 0.81)。拒绝采样带来的更高接受率增益直接转化为更大的吞吐量提升。

A7 补充细节

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

TV 损失对分布形状的影响。我们分析了 TV 损失与 CE/KL 训练相比如何影响草稿分布的熵。TV 损失产生的草稿分布熵更接近目标熵（但略高），表明草稿变得更“尖锐”，与目标的峰值预测更对齐。相反，CE/KL 训练倾向于产生更平滑的草稿分布，将概率质量分散在整个词汇表中，这对于拒绝采样是次优的，因为重叠 $\sum_v \min(p(v), q(v))$ 是通过匹配目标形状来最大化的。

尖锐化效应的来源。这种尖锐化效应源于 TV 损失梯度的选择性行为（公式 11）：它将优化精力集中在决策边界附近（$q_j \approx p_j$）的词元上，同时忽略不相关的低概率词元。图 10 说明了草稿-目标熵差距与 KL 距离之间的关系。具有训练良好 MTP 头的模型在草稿和目标分布之间表现出较小的熵差距，同时具有较大的 KL 距离（也见图 1b）。


图 10: (a) 熵差距 ∆H vs. DKL(q∥p) 在不同模型和任务上的关系。(b) 熵差距与 RS 接受率呈负相关 (r = -0.54)。(c) KL 散度没有这样的相关性 (r = 0.13)，表明熵差距而非 KL 是 RS 接受率的相关预测指标。

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

不同损失下的度量演变。图 11 显示了在 RL 期间使用不同损失更新 MTP 权重时，各种 MTP 指标如何演变。与 CE 损失相比，TV 损失产生的草稿熵更接近目标模型，但 KL 距离更大。此外，因为 TV 损失产生更尖锐的草稿分布，相应的 $\alpha_{p>q}$ 较低，而 $\alpha_{q>p}$ 较高。

损失切换的影响。当在 RL 期间使用不同的损失进行 MTP 权重更新时，MTP 指标会朝着该损失特有的模式转变。例如，对于 RS w/ TV + CE 损失，草稿熵在训练过程中逐渐增加。


图 11: 在 RL 训练期间使用不同 MTP 损失目标时 MTP 指标的演变。TV 损失产生的草稿熵更接近目标，但 KL 距离更大，$\alpha_{p>q}$ 更低，$\alpha_{q>p}$ 更高。在 RL 期间切换 MTP 训练损失会导致指标向新损失的特征模式转变。

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

两种方法的敏感性差异。尽管 §5.1 的分析表明 RL 期间模型更新的幅度相对较小，但仅目标采样和拒绝采样在对 RL 策略更新引起的排名变化敏感性方面存在重要区别。

仅目标采样的脆弱性。仅目标采样的接受依赖于草稿词元是否落在目标模型的高概率区域（如 top-k）。这是一个离散标准。当 RL 梯度步骤导致 top-1 词元改变时，即使概率变化很小（例如，$p(v_1)$ 从 0.31 降到 0.29，而 $p(v_2)$ 从 0.29 升到 0.31），仍然偏爱旧 top-1 的草稿模型会经历从接受到拒绝的不连续跳跃。

拒绝采样的平滑退化。在拒绝采样下，接受率 $\alpha_{\text{RS}} = \sum_v \min(p(v), q(v))$ 是两个分布的连续函数。相同的排名变化对 TV 重叠产生的影响可以忽略不计，因为 $\min(p(v_1), q(v_1)) + \min(p(v_2), q(v_2))$ 对小的概率交换几乎不敏感。

高熵放大了脆弱性差距。当目标熵高时，多个词元具有相似的概率，使得 RL 更新下排名变化更频繁。这对仅目标采样的影响不成比例，因为每次排名翻转都可能导致离散的接受失败。尽管存在这种质的差异，我们经验上观察到仅目标采样和拒绝采样的熵-接受率斜率相似（$b_{\text{TO}} \approx b_{\text{RS}}$；见 §3），这表明仅目标采样的离散脆弱性被同样影响 CE/KL 训练下拒绝采样的累积 TV 距离增长所抵消。

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

温度对熵和接受率的影响。采样温度 $\tau$ 直接影响目标模型的熵：$H(p_\tau) = H(\text{softmax}(z/\tau))$ 随 $\tau$ 单调增加。结合 §3 中建立的线性熵-接受率关系，这意味着更高的温度会导致更低的 MTP 接受率。

实验验证。图 12a 证实了这一点：拒绝采样在不同温度下保持相对稳定的接受长度，而仅目标采样随着温度的升高而急剧下降。这对 RL 训练具有实际意义，因为通常使用更高的温度来鼓励探索。我们的分析为理解通过温度缩放进行探索的吞吐量成本提供了一个定量框架。

7.5 拒绝采样决策边界

决策条件。当 $d_{\text{TV}}(p, q) < 1 - p(\hat{y})$ 且 $\hat{y} = \arg \max_y q(y)$ 时，拒绝采样的性能优于仅目标采样（见 §E）。这个决策边界提供了一个简单的诊断方法：如果草稿-目标 TV 距离小于目标模型在草稿 top-1 词元之外的概率质量，则首选 RS。

实验验证。图 13 在八个具有原生训练 MTP 头的模型中可视化了这一边界，涵盖了三类任务。几乎所有的模型-任务组合（24 个中的 23 个）都稳固地落在 RS-better 区域，证实了对于原生 MTP 模型，拒绝采样始终优于仅目标采样。这证实了启用拒绝采样对于几乎所有实际的 MTP 部署都是有益的。


图 12: (a) 平均接受长度作为采样温度的函数。拒绝采样保持相对稳定的接受长度，而仅目标采样在较高温度下急剧下降。(b) MTP 接受率与输出长度（8 个模型平均）。RS 在所有生成位置上都保持对仅目标采样的稳定优势。


图 13: 不同模型的 RS 决策边界（见 §7.5）。几乎所有的模型-任务组合都落在 RS-better 区域，证实了拒绝采样对几乎所有实际的 MTP 部署都是有益的。

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

位置依赖的接受率。如图 12b 所示，我们观察到 MTP 接受率随生成序列中的位置系统性地变化。在早期位置（靠近提示），目标模型倾向于具有较低的熵（更可预测的延续），导致较高的接受率。随着生成的进行，尤其是在具有长思维链的推理任务中，熵可能增加，接受率可能下降。这种位置依赖的接受模式表明，自适应 MTP 策略——根据估计的局部熵调整草稿长度 $\gamma$——可能进一步提高吞吐量。

7.7 智能体 RL 与“气泡”问题

智能体 RL 的特点。如图 14a 所示，在智能体 RL 设置中（例如，SWE-bench【16, Swe-bench: Can language models resolve real-world github issues?, 2024, ICLR】），模型生成长的、多轮的交互，涉及工具调用、代码执行和迭代改进。这些设置表现出特别长的生成长度和可变的熵剖面，产生周期性的接受率波动，并随着生成的进行而趋于增加。

MTP 在智能体设置中的优势。MTP 在智能体设置中特别有益，原因有二：(1) 长生成包含丰富的结构化输出——如样板代码、工具调用格式和重复模式——这些都高度可预测，在这些段落中产生高接受率；(2) 多轮交互和长尾生成降低了有效的运行批量大小，在这种情况下，MTP 的延迟优势被放大，因为推理引擎在远离计算饱和的状态下运行。实际上，我们的实验表明，智能体工作负载从我们提出的 TV 损失训练中获得了最大的接受率改进（5%）。


图 14: (a) 智能体 RL 期间的接受长度。平均接受长度稳定在约 3.7，而最小-最大范围揭示了跨步骤的周期性波动。(b) 不同 top-K 截断值下的 MTP 损失曲线。较小的 K 导致明显的损失尖峰和训练不稳定，而即使 K = 20,000 也显示出比全词汇表 TV 损失更慢的收敛速度。

7.8 Top-K TV 近似的不稳定性

全词汇表 TV 损失的挑战与解决方案。计算全词汇表的 TV 损失会在大词汇表上产生高峰值内存。为了解决这个问题，我们采用了一个融合的反向传播核，减少了中间激活的大小（见 §F）。我们还尝试通过 top-K 截断来近似 TV 损失，以进一步减少峰值内存。

Top-K 近似的问题。然而，即使 K = 20,000，我们也观察到损失收敛速度略有减慢和性能下降。较小的 K 值会导致明显的损失尖峰，如图 14b 所示。最终，我们采用了融合的全词汇表 TV 损失，而不是 top-K 近似。

A5 结论

本文提出了 Bebop，一项关于在大型语言模型强化学习背景下应用多词元预测（MTP）的系统性研究。我们的分析揭示了三个关键发现：
1. 在仅目标采样和拒绝采样下，MTP 的接受率都受到目标模型熵的线性制约。
2. Bebop 的端到端 TV 损失直接优化多步拒绝采样的接受率，相比传统的 CE/KL 目标，带来了约 10% 的接受率提升，最高可达 95% 的接受率，以及高达 25% 的额外推理吞吐量增益。
3. 在 RL 训练前进行轻量级的 TV 损失和拒绝采样适应，足以在整个 RL 训练过程中保持高 MTP 接受率，从而无需昂贵的在线 MTP 更新。

通过对 Qwen3.5、3.6 和 3.7 模型的广泛实验证明，Bebop 在异步 RL 流程中实现了高达 1.8 倍的端到端加速。

局限性。我们对熵-接受率关系的理论分析依赖于建模假设（均匀 vs. 概率成比例的不匹配），这些假设是基于梯度结构启发式地提出的，而非形式化证明；完善这些假设仍是一个悬而未决的问题。此外，TV 训练保证的熵不变性是分布条件的：它在 SFT 训练数据覆盖的熵范围内成立，但当 RL 探索将策略熵显著推到该范围之外时，草稿头会遇到分布外的目标分布，此时不匹配比率 $\delta$ 不再有界，从而恢复了与 CE/KL 训练相当的熵-接受率依赖性。在这种情况下，建议在 RL 期间进行带 TV 损失的 MTP 协同训练，以将草稿头的有效覆盖范围扩展到新的熵域。

A6 附录

A. TV 损失梯度推导

TV 损失梯度推导。我们提供 TV 损失梯度（公式 11）的完整推导。设草稿头输出 logits $z \in \mathbb{R}^{|V|}$，其中 $q_j = \text{softmax}(z)_j = \frac{e^{z_j}}{\sum_k e^{z_k}}$。目标模型概率 $p$ 被视为常数（分离）。TV 损失为：

对 $z_j$ 的梯度为：

由于 $p$ 是常数，$\min(p_v, q_v)$ 对 $q_v$ 的次梯度为：

使用链式法则和 softmax 的雅可比矩阵 $\frac{\partial q_v}{\partial z_j} = q_v(\delta_{vj} - q_j)$：

因此：

其中 $S = \sum_v \mathbf{1}[q_v \le p_v] \cdot q_v \in [0, 1]$。

有界性。由于 $q_j \in [0, 1]$ 且 $|\mathbf{1}[q_j \le p_j] - S| \le 1$：

B. 与前向 KL 散度梯度的比较

前向 KL 散度梯度。作为比较，前向 KL 散度 $D_{\text{KL}}(p\|q) = \sum_v p_v \log \frac{p_v}{q_v}$ 对 $z_j$ 的梯度为：

与 TV 损失梯度的主要区别：
1. 前向 KL 梯度对每个 $q_j \neq p_j$ 的词元都施加非零作用力，包括概率可忽略的词元。TV 梯度与 $q_j$ 成正比，因此自动忽略低概率词元。
2. 前向 KL 梯度不区分在拒绝采样下会被接受还是拒绝的词元。TV 梯度通过指示函数 $\mathbf{1}[q_j \le p_j]$ 明确包含了这种区别。
3. 当 $q_j \gg p_j$（草稿过分自信）时，前向 KL 梯度可能很大。TV 梯度则以 $q_j$ 为界。

C. 反向 KL 散度分析

反向 KL 散度。前面的分析集中于前向 KL 散度 $D_{\text{KL}}(p\|q)$，它在常数范围内等同于 CE 损失。一个自然的问题是，反向 KL 散度 $D_{\text{KL}}(q\|p) = \sum_v q_v \log \frac{q_v}{p_v}$ 是否是更好的拒绝采样训练目标。

梯度推导。反向 KL 散度对草稿 logits $z_j$ 的梯度为：

梯度结构比较。表 1 总结了三种梯度结构。反向 KL 梯度与 TV 梯度共享了理想的 $q_j$ 比例性，这意味着低概率词元自动接收到可忽略的优化压力。这表明反向 KL 应该比前向 KL 产生更与 $q_j$ 成比例的不匹配，从而表现出比前向 KL 更弱的熵-接受率耦合。

为何反向 KL 仍然次优。尽管梯度结构有所改进，但反向 KL 在最大化拒绝采样接受率方面仍然次优，原因有三：
1. 零强制行为。反向 KL 不惩罚 $q(v) \to 0$ 的情况，即使 $p(v) > 0$，因为 $\lim_{q\to 0} q \log(q/p) = 0$。这种“模式寻求”特性允许草稿丢弃 $p$ 的模式，直接放弃了这些词元上的重叠 $\min(p(v), q(v))$，从而降低了接受率。
2. 不对称的过高/过低估计惩罚。拒绝采样的接受率取决于 $\sum_v \min(p(v), q(v))$，它对称地惩罚过高估计（$q > p$）和过低估计（$q < p$）。反向 KL 施加了不对称的惩罚：过高估计（$q_j > p_j$）比过低估计招致更强的梯度。这驱使草稿在大多数词元上趋向于 $q(v) \le p(v)$，虽然这确保了单个词元的接受概率 $\min(1, p/q) = 1$，但减少了这些词元的采样概率，导致次优的总重叠。
3. 间接优化目标。与前向 KL 一样，反向 KL 并不直接优化 $d_{\text{TV}}(p, q)$。反向 KL 梯度中的对数比率 $\log(q_j/p_j)$ 提供了一个软的、非线性的信号，而 TV 梯度的指示函数 $\mathbf{1}[q_j \le p_j]$ 提供了一个与拒绝采样决策边界对齐的硬的、直接的信号。

总结。就优化拒绝采样接受率的适用性而言：反向 KL 通过更好的容量分配（梯度 $\propto q_j$）改进了前向 KL，但由于其零强制行为和不对称惩罚结构而仍然次优。TV 损失直接优化了我们关心的量，并避免了这两种失败模式。

D. 不同训练目标下的熵-接受率关系

D.1 设置和符号

设置。考虑一个固定位置 $t$。$p \in \Delta^{|V|}$ 是目标分布，$q \in \Delta^{|V|}$ 是草稿分布。草稿模型容量有限，其与 $p$ 的不匹配结构取决于训练目标。
两种不匹配结构：
* 均匀不匹配 (CE/KL 训练): $q^*(v) = p(v) + \eta_v$，其中 $|\eta_v| \lesssim \sigma$ 且 $\sum_v \eta_v = 0$。
* 概率成比例不匹配 (TV 训练): $|q^*(v) - p(v)| \lesssim \delta \cdot p(v)$。

D.2 仅目标采样

接受率。贪婪选择草稿词元 $\hat{y} = \arg \max_y q(y)$，接受率为：

完美草稿情况。当草稿正确识别 top-1 词元时，接受率为 $\alpha_{\text{TO}} = \max_y p(y)$。

此值是 $H(p)$ 的单调递减函数，下界为 $\exp(-H(p))$。

线性化。对 $\alpha_{\text{TO}} = f(H(p))$ 在平均熵 $\bar{H}$ 附近进行一阶泰勒展开：

由于 $f$ 递减，$b_{\text{TO}} = -f'(\bar{H}) > 0$，得到线性关系：

不完美草稿修正。在高熵下，排名错误更频繁，导致接受率下降，从而加强了负斜率。线性近似仍然成立，但斜率可能更陡：

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

接受率。$\alpha_{\text{RS}} = 1 - d_{\text{TV}}(p, q)$。
梯度结构。$D_{\text{KL}}(p\|q)$ 的梯度为 $q_j - p_j$，导致近似均匀的每词元不匹配 $q^*_{\text{CE}}(v) = p(v) + \eta_v$，其中 $|\eta_v| \lesssim \sigma$。

TV 距离推导。在均匀不匹配下：

主要贡献来自有效支持集 $S_\tau(p)$，其大小 $|S_\tau(p)| \approx \exp(H(p))$。

其中，在有效支持集上，不匹配累积为：

因此，接受率为：

线性近似。对指数项进行泰勒展开得到：

代入后得到线性关系，其中 $a_{\text{RS}}$ 和 $b_{\text{RS}}$ 为正常数。这解释了 CE/KL 训练下熵与接受率的负相关性。

CE/KL 次优性根源。问题在于容量分配。KL 梯度将有限的模型容量大致均匀地分配给所有词元。高熵分布将概率质量分散到更多词元上，导致累积的每词元残差形成更大的 TV 距离，即使 KL 散度本身也在被最小化。

D.4 使用 TV 训练的拒绝采样

梯度结构。TV 损失的梯度 $\partial\mathcal{L}_{\text{TV}}z = -q_j[\mathbf{1}[q_j \le p_j] - S]$ 与 $q_j$ 成正比。这意味着高概率词元得到强优化，低概率词元被忽略。

TV 梯度作为自校正机制。TV 训练驱动概率比率 $r_j = q_j/p_j \to 1$。当 $r_j < 1$ 时，梯度增加 $z_j$；当 $r_j > 1$ 时，梯度减小 $z_j$。这种校正的幅度与 $q_j$ 成正比，确保有效支持集上的对数比率误差 $| \log r_j |$ 有界。

假设：有界对数比率误差。我们假设在 TV 训练下，对于有效支持集中的所有 $j$，存在一个常数 $\epsilon$，使得：

这导致概率成比例的不匹配：
其中 $\delta = e^\epsilon - 1$。

TV 距离推导。在概率成比例不匹配下：

这个界限与 $H(p)$ 无关，因此接受率：

这证明了命题 4。在实践中，由于容量有限，$\delta$ 可能有微弱的熵依赖性，但相比 CE/KL，熵-接受率相关性大大减弱。

E. 拒绝采样决策边界推导

决策边界推导。我们推导拒绝采样（RS）何时比仅目标采样（TO）具有更高的接受率。RS 的接受率为 $\alpha_{\text{RS}} = 1 - d_{\text{TV}}(p, q)$，TO 的接受率为 $\alpha_{\text{TO}} = p(\hat{y})$。RS 优于 TO 的条件是 $\alpha_{\text{RS}} > \alpha_{\text{TO}}$：

这意味着，只要草稿-目标 TV 距离小于目标在草稿贪婪预测之外的概率质量，RS 就更优。

F. 融合的 TV 损失核

融合核实现。我们提供了融合 TV 损失实现的伪代码。前向传播（算法 1）在一个核中计算每词元 TV 损失和辅助量 $S$。反向传播（算法 2）计算对草稿 logits 的梯度。两个核都以大小为 BLOCK_V 的块迭代词汇表，以限制寄存器和共享内存使用，从而实现全词汇表 TV 损失计算而无需物化 softmax 输出。

算法 1 TV 损失前向核（每词元位置）
输入: 草稿 logits z ∈ R|V|, 目标 log-probs log p ∈ R|V|
输出: TV 损失 ℓ, 辅助标量 S
1: // 步骤 1: 数值稳定的 softmax 分母
2: m ← max_v z_v             // 全局 logit 最大值
3: D ← Σ_v exp(z_v - m)      // exp-sum
4: // 步骤 2: 分块累加重叠和 S
5: overlap ← 0; S ← 0
6: for v_start = 0 to |V| step BLOCK_V do
7:   v ← [v_start, ..., v_start + BLOCK_V-1]
8:   q ← exp(z[v] - m) / D    // 草稿概率
9:   p ← exp(log p[v])        // 目标概率
10:  overlap += Σ min(q, p)
11:  S += Σ q · 1[q ≤ p]
12: end for
13: ℓ ← clamp(1 - overlap, 0, τ_max) // τ_max: 可选的裁剪
14: return ℓ, S
算法 2 TV 损失反向核
输入: 草稿 logits z, 目标 log-probs log p, 缓存的 (m, D, S, gout)
输出: 梯度 ∇ℓ ∈ R|V|
1: for v_start = 0 to |V| step BLOCK_V do
2:   v ← [v_start, ..., v_start + BLOCK_V-1]
3:   q ← exp(z[v] - m) / D
4:   p ← exp(log p[v])
5:   ∇z[v] ← q · (S - 1 + 1[q > p]) · gout
6: end for
7: return ∇z

G. 拒绝采样推理实现

实现概述。在生产推理引擎中实现基于 MTP 的拒绝采样需要修改草稿和验证阶段。与仅目标采样不同，拒绝采样需要：(1) 从草稿分布 $q$ 中采样草稿词元，(2) 缓存草稿概率用于验证，(3) 在验证期间计算接受比率 $\min(1, p(\hat{y})/q(\hat{y}))$。我们描述了两种实现策略。

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

SGLang 实现。第一种方法在 SGLang 中实现，直接使用多项式采样从草稿分布中采样词元。
* 草稿阶段：通过多项式采样从 $q(\cdot)$ 中采样 $\hat{y}$，并缓存完整的草稿概率向量 $q \in \mathbb{R}^{|V|}$。
* 验证阶段：使用融合的 Triton 核处理每个请求。该核执行两个阶段：
1. 顺序接受：对每个草稿步骤 $i$，以概率 $\min(1, p_i(\hat{y}_i)/q_i(\hat{y}_i))$ 接受 $\hat{y}_i$，直到第一次拒绝。
2. 残差重采样：如果草稿被拒绝，或所有草稿都被接受，从残差分布中采样下一个词元。这通过两遍 CDF 反演实现。

G.2 Gumbel-Max 技巧 (vLLM)

vLLM 实现。第二种方法在 vLLM 中实现，利用 Gumbel-Max 技巧避免了显式的 CDF 反演。

  • 草稿阶段：使用 Gumbel-Max 技巧采样词元：$v^* = \arg \max_v [\log q(v)/\tau + G_v]$，其中 $G_v$ 是 Gumbel 噪声。缓存经过温度缩放的草稿 logits。
  • 验证阶段：分为两个核：
    1. 接受核：一个顺序 Triton 核迭代草稿步骤，计算接受概率并记录第一个被拒绝的步骤索引。
    2. 残差 logits 核：一个并行的 Triton 核在 logit 空间计算残差分布。然后使用与草稿阶段相同的 Gumbel-Max 采样从该残差分布中抽取重采样的词元。

算法 4 链式拒绝采样验证 (Gumbel-Max / vLLM)
输入: 草稿词元 $\hat{y}_1, \dots, \hat{y}_\gamma$; 草稿 logits $z_{q_1}, \dots, z_{q_\gamma} \in \mathbb{R}^{|V|}$; 目标概率 $p_1, \dots, p_\gamma \in \mathbb{R}^{|V|}$; 目标 logits $z_{p_\gamma}$
输出: 接受的词元数 $n$; 位置 $n+1$ 的输出词元 $y^*$
1: // 核 1: 顺序接受
2: $n \leftarrow \gamma$
3: for $i = 1$ to $\gamma$ do
4: $q_i(\hat{y}_i) \leftarrow \text{softmax}(z_{q_i})_{\hat{y}_i}$
5: $u \leftarrow \text{tl.rand}(\text{seed}, i)$
6: if $u \cdot q_i(\hat{y}_i) \ge p_i(\hat{y}_i)$ then
7: $n \leftarrow i - 1$; break
8: end if
9: end for
10: // 核 2: 残差 logits
11: if $n < \gamma$ then
12: $z_{\text{resid}}(v) \leftarrow \log \max(0, p_{n+1}(v) - q_{n+1}(v))$ for all $v$
13: else
14: $z_{\text{resid}}(v) \leftarrow z_{p_\gamma}(v)$ for all $v$
15: end if
16: // Gumbel-Max 重采样
17: $G_v \sim \text{Gumbel}(0, 1)$ for all $v$
18: $y^* \leftarrow \arg \max_v (z_{\text{resid}}(v) + G_v)$
19: return $n, y^*$