Beat the Long Tail: Distribution-Aware Speculative Decoding for RL Training

作者/机构: Zelei Shao * 1 2, Vikranth Srivatsa * 2 3, Sanjana Srivastava 2, Qingyang Wu 2, Alpay Ariyak 2, Xiaoxia Wu 2, Ameen Patel 4, Jue Wang 2, Percy Liang 2, Tri Dao 2, Ce Zhang 2, Yiying Zhang 3, Ben Athiwaratkun 2, Chenfeng Xu 2, Junxiong Wang 2

A1 主要贡献

强化学习(RL)后训练已成为对齐大型语言模型(LLM)的关键,但其效率日益受到“rollout”阶段的限制,在该阶段中,长轨迹是逐个 token 生成的。本文首先指出了一个主要瓶颈——rollout 长度的长尾分布,其中一小部分长生成占据了大部分的墙上时钟时间。同时,本文也发现了一个互补的机会——历史 rollout 的可用性揭示了在训练周期中稳定的提示级别模式。基于这些观察,本文提出了 DAS,一个分布感知的推测解码框架,它在不改变模型输出的情况下加速 RL rollout。

DAS 框架整合了两个核心思想:
1. 自适应、非参数化的草稿模型(drafter):该模型利用一个增量维护的后缀树,从近期的 rollout 中构建而成。由于在 RL 训练期间,目标模型不断变化,预训练的草稿模型会失效。DAS 利用训练样本的重复使用特性,基于后缀树进行推测解码,并通过动态修剪后缀结构来与当前策略保持一致,只保留近期和相关的历史 rollout,以提高局部性和可靠性。
2. 长度感知的推测策略:该策略将更激进的草稿预算分配给主导完成时间(makespan)的长轨迹。系统不是维护一个单一的全局后缀树,而是为每个问题构建和更新独立的树,这能更好地捕捉领域特定的模式,并提高异构任务中推测的准确性。

本文在 VeRL 和 vLLM 的基础上实现了 DAS,并在数学和代码推理任务上进行了评估。实验结果表明,DAS 将 rollout 时间减少了高达 50%,同时保持了完全相同的训练曲线,证明了分布感知的推测解码可以在不损害学习质量的情况下显著加速 RL 后训练。

本文的主要贡献如下:
* 提出了关于 RL rollout 阶段独特属性及其系统级影响的三个关键见解。
* 实现了 DAS,一个完整的 RL 系统,相比于当前最先进的 RL 系统,可将 rollout 时间减少高达 50%。
* 提出了一种专为 RL rollout 设计的、无需训练、自我演进的推测解码方法。
* 提出了一种分布感知的推测解码方法,能有效减少 RL rollout 中的尾部延迟。

A3 背景知识/关键观察/设计原则

RL 训练中的 Rollout

本节我们介绍 RL 工作负载的独有特性。

RL rollout 中的长尾瓶颈。强化学习后训练以迭代方式运行,其中 actor 生成一组固定的轨迹,而 learner 则更新策略。在实践中,像 VeRL 和 OpenRLHF 这样的系统倾向于使用数据并行(DP)的 rollout worker 来扩展解码吞吐量,只有当训练模型无法装入内存时才使用张量并行(TP)。解码开始时并行度是满的,但随着短序列的完成,有效批次大小会“坍塌”,只剩下少数长的响应(掉队者)决定了整个步骤的完成时间——这是在 LLM 系统中也观察到的经典长尾行为。这导致 GPU 利用率低下,因为当掉队者仍在解码时,其他 worker 处于空闲状态;多项 RL/LLM 研究报告称,尽管采用了复杂的批处理技术,解码阶段仍主导了墙上时钟时间,并且加速器利用率不足。这促使我们使用推测解码:通过起草多个 token 并行验证它们,我们可以缩短掉队者的运行时间,并利用空闲的 GPU 时间,从而在不改变模型输出的情况下实现显著的端到端加速。虽然推测解码通常在小批量、对延迟敏感的服务场景中进行评估,但 RL 训练在解码过程中表现出有效的批次坍塌——短序列先完成,而掉队者主导了整个步骤。这暴露了推测可以利用的空闲容量。


图 1. 在没有 DAS 的情况下,rollout 期间有效批大小的坍塌。测量于 DeepSeek-distilled 7B 模型 (Guo et al., 2025) (DeepScaleR (Luo et al., 2025b) 提示)。随着解码的进行,短序列首先完成,有效批大小缩小,留下几个长的掉队者来决定步骤的完成时间。使用我们的方法,我们既可以减少总延迟,又可以减轻长尾掉队者的影响。

RL Rollout 与近期历史的高度相似性。在不同的训练周期(epoch)中,针对相同提示生成的轨迹表现出明显的词汇和结构复用。在图 2 中,我们观察到与近期轨迹的相似度较高,而与较远轨迹的相似度较低,这与策略漂移(policy drift)的现象一致。这种“新近偏好”(recency bias)意味着,一个基于历史索引的、无模型的草稿模型(例如,后缀数组或后缀树)可以挖掘近期的延续序列来提出高质量的草稿,即使策略在不断演变。


图 2. (左) 使用 N-gram 计算每个迭代的内容相似度以衡量复用率。(右) Qwen2.5-7B-Instruct 在不同训练周期之间的成对相似度。靠近对角线的块状结构表明,rollout 与最近周期的 rollout 最为相似,并且相似度随着时间距离的增加而衰减。这反映了策略漂移:随着策略的不断更新,较早的生成对当前行为的预测性变差。

A2 方法细节

分布感知的推测解码框架

我们引入了一个分布感知的推测解码框架(DAS),这是一种通过调整推测解码来加速 rollout 的 RL 后训练方法。图 3 提供了该系统的概览。DAS 建立在通过基于偏好或可验证性奖励来优化策略的 RL 流程之上,它维护一个基于历史索引的、非参数化的草稿模型,该模型会根据近期的 rollout 持续刷新,以在策略演变过程中保持高接受率。然后,分布感知的推测解码组件决定为每个请求分配多少推测预算,优先考虑那些长的、高延迟的问题。通过将自适应的预算分配与一个在线的、与分布对齐的推测器(以在线后缀树构建的方式增量更新)相结合,该系统在不修改奖励循环的情况下降低了 rollout 延迟。


图 3. 我们在 RL 训练中 rollout 加速框架的概览。(左) 长度感知的草稿预算分配。我们根据近期的 rollout 估计每个问题的长度,并相应地分配草稿预算:预计会很长或很难的问题会被分配更激进的推测预算,而简单的问题则几乎不进行推测。该策略在一个滑动窗口内的近期轨迹上更新,因此它会随着训练中策略的变化而自适应。(右) 分布感知的、自我演进的推测解码。对于每个问题分片,我们维护一个后缀树推测器,该推测器根据最新的 rollout 增量更新。在解码时,推测器从高频的后缀匹配中提取多-token 草稿,目标模型并行验证它们;被接受的 token 以更低的 rollout 延迟推进生成。

4.1 分布感知的草稿模型

4.1.1 参数化的静态草稿模型

EAGLE 系列方法是当前最先进的推测解码推理时服务技术。EAGLE【索引241.15077, EAGLE: Speculative sampling requires rethinking feature uncertainty, 2024a, arXiv】、EAGLE-2【索引2406.16858, EAGLE-2: Faster inference of language models with dynamic draft trees, 2024b, arXiv】和 EAGLE-3【索引2503.01840, EAGLE-3: Scaling up inference acceleration of large language models via training-time test, 2025, arXiv】通过一个轻量级的 EAGLE 头来推断内部特征,构建一个 token 草稿树,然后由目标模型并行验证,从而在推理过程中显著降低延迟。特别是 EAGLE-2,它依赖于一个经过良好校准的草稿模型,其置信度能紧密预测 token 的接受情况,并据此动态地生成草稿树。

然而,在 RL 训练中,策略是非平稳的。模型权重在每次 learner 更新后都会改变,因此这种校准会迅速漂移。图 2 表明,从早期 epoch 模型收集的轨迹导致的接受率会显著低于从后期检查点收集的轨迹。在实践中,EAGLE 必须要么容忍接受率下降(从而降低加速效果),要么在整个训练过程中反复重新训练或重新调整其头部和建树阈值——这给本已由 rollout 主导的阶段增加了计算和工程开销。

此外,EAGLE 的树状草稿生成和并行验证核心是为高吞吐量的推理服务而设计的。在一个 actor-learner 的 RL 流程中,检查点随时间演变,集成并持续重新校准这些组件比其预期的仅推理用例要复杂得多。由于这些原因,尽管 EAGLE 在推理方面取得了优异的成果,但它很难在非平稳的 RL 训练机制中高效扩展。

持续更新神经模型的高昂开销促使我们采用一种基于文本索引的无模型方法,特别是后缀树【索引Ukkonen, 1995, On-line construction of suffix trees】和后缀数组【索引Manber & Myers, 1993, Suffix arrays: a new method for on-line string searches】。

4.1.2 自适应非参数化草稿模型

基于上述动机,我们探索了一种自适应的非参数化草稿模型。我们利用一个基于后缀的草稿模型,该模型通过在一个滑动窗口内的近期 rollout 上维护一个紧凑的后缀树来构建。在 rollout 时,我们定位当前上下文与树之间的最长匹配(对于长度为 m 的上下文,自上而下查找的时间复杂度为 O(m),其中 m 是我们查找的查询/模式的长度,例如当前的解码前缀),然后沿着匹配的路径提出一个多-token 的草稿。目标模型在一个或几个批处理步骤中验证这个草稿。验证后,我们在线更新树,这样推测器就能持续适应不断演变的策略。

这种设计捕捉了训练周期中反复出现的模式,同时避免了训练或维护一个单独的基于训练的神经草稿模型。重要的是,后缀树的构建和在线后缀树的更新通过 Ukkonen 算法【索引Ukkonen, 1995, On-line construction of suffix trees】以线性时间运行,这使其非常适合增量地接收新的路径。


图 4. RL 训练中每轮验证平均接受的 token 数。我们比较了一个静态学习的草稿模型 EAGLE【索引2406.16858, EAGLE-2: Faster inference of language models with dynamic draft trees, 2024b, arXiv】和我们的免训练草稿模型。虽然静态草稿模型的接受率保持平稳,但我们的非参数化草稿模型会根据近期的 rollout 进行更新,并跟踪不断演变的策略,从而随着时间的推移产生更高的接受长度。更高的接受长度意味着每个生成的 token 需要更少的目标模型前向传播,从而降低 rollout 延迟。

后缀树与后缀数组的比较。除了后缀树,受 Infinigram【索引2401.17377, Infini-gram: Scaling unbounded n-gram language models to a trillion tokens, 2024a, arXiv】的启发,使用后缀数组(SA)作为一种更节省空间的选择是很自然的想法。一个标准的 SA 支持通过二分搜索在 O(m log n) 时间内查找长度为 m 的模式的子串,其中 n 是语料库中的 token 数量。通过增加一个 LCP 数组【索引Kasai et al., 2001, Linear-time longest-common-prefix computation in suffix arrays and its applications】可以减少比较次数,而使用增强型后缀数组(ESA)【索引Abouelhoda et al., 2004, Replacing suffix trees with enhanced suffix arrays】则能通过有效模拟后缀树遍历,以更好的常数和缓存局部性在 O(m) 时间内实现模式匹配。然而,SA 和 ESA 本质上是静态的:动态更新通常需要(部分)O(n) 的重建,这在 RL 训练中是不理想的,因为每个迭代都会有新的轨迹产生。

我们的在线后缀树索引则通过牺牲适度的额外空间来换取快速的增量更新和 O(m) 的最长匹配查询,这与非平稳的策略环境天然契合。我们在图 5 中比较了后缀树与后缀数组的推测时间和更新成本。后缀树在这两个指标上都表现出优越的性能:推测时间快了 2-20 倍,而更新成本则显示出巨大优势,保持在亚毫秒级别,相比之下,后缀数组的重建时间则不断攀升。这种超过三个数量级的更新性能差异证实了后缀数组尽管空间效率高,但在在线 RL 训练中是不切实际的,因为每个迭代都必须快速更新新的轨迹。


图 5. 后缀树和后缀数组数据结构的性能比较。左图:不同语料库大小下的推测时间。右图:插入100个 token 的更新时间(对数尺度)。

全局后缀树的问题。一个单一的、不断增长的、覆盖所有过去 rollout 的全局索引在统计和系统两个层面都会失败。在统计上,策略漂移导致较早的延续序列与模型当前的条件分布对齐不佳,这会降低推测解码的接受率——而这正是决定加速效果的关键量(见图 2)。此外,尽管 RL 中的所有问题可能都属于同一领域,但它们的多样性意味着来自一个问题的模式很少能可靠地迁移到另一个问题。与此一致,图 6 显示,使用全局树比为每个问题维护一棵树带来的增益要小,但由于树更大,反而增加了额外的开销。


图 6. (左) 平均每轮验证接受的 token 数与训练步数的关系。我们比较了三种构建草稿模型的历史记录:全局+请求、问题+请求和仅问题;问题范围内的历史记录在接受率上超过了全局记录。(右) 推测解码延迟 (ms/token);由于查询和维护一个大型全局索引的成本较高,全局+请求的方式始终较慢,而问题范围的分片查询成本更低,因此延迟也更低。

每个请求的后缀树。基于以上分析,我们采用每个请求一个后缀树,并辅以一个轻量级的请求前缀树(trie)进行路由。这种按问题设计的主要优势在于,当模型倾向于重复某个模式时,前缀树可以快速识别该模式并路由到相应的后缀树,因为它是由先前的生成构建的。然而,前缀路由的好处取决于工作负载和模型。对于较小的模型,前缀路由的额外 CPU 开销可能会超过其收益。在这种情况下,我们禁用请求前缀树,直接查询按问题的后缀树。图 6 显示,使用请求前缀树可能会带来更高的接受率,但会产生更多的推测时间。

滑动窗口选择树。在不同的训练周期和模型中,我们观察到针对相同提示的生成之间存在显著的相似性(图 2)。时间上邻近的轨迹倾向于产生更长的被接受草稿,而来自更早周期的轨迹则产生较短的接受长度,这表明随着策略的漂移,匹配质量会下降。

由于策略随时间演变,历史久远的生成变得不那么具有预测性。因此,我们从一个滑动窗口内的近期轨迹中构建草稿模型,并在每次迭代中刷新索引。窗口大小控制着一个偏差-稳定性权衡:较短的窗口能快速适应,但提供的匹配较少;而较长的窗口能提高覆盖率,但有陈旧的风险。在实践中,我们将窗口更新速率与优化器的步长尺度联系起来(例如,较大的参数更新意味着较短的窗口),并对源自较早周期的匹配应用轻微的降权。


图 7. (左) 平均每轮验证接受的 token 数与训练步数的关系,比较了用于构建草稿模型的不同历史窗口大小。较大的窗口(如 16, 32, all)能提供更高的接受率,因为它们提供了更多的匹配延续,而更高的接受率已知能直接转化为更少的目标模型前向传播和更低的解码成本。(右) 每个 token 的推测解码延迟与训练步数的关系。window all 显示出最高的延迟,因为查询和维护一个完整的全局历史记录更昂贵,并且包含了陈旧的轨迹,因此在实践中,中等大小的窗口(16 或 32)在接受率和延迟之间取得了比 window all 更好的平衡。

4.2 长度感知的推测策略

4.2.1 Rollout 延迟估计

我们通过首先建模每次前向传播的延迟,然后将其累加到所有前向传播中来估计端到端的 rollout 延迟。通过性能分析,我们发现一个简单的线性模型能够捕捉主要行为(平均相对误差 ≈ 12%):


图 8. 在 DeepSeek-R1-Distill-Qwen-7B 上的解码延迟与 token 数的关系图。它显示了清晰的线性关系。

$$t_{\text{fwd}} = c_{\text{base}} + c_{\text{tok}} n_{\text{toks}},$$

其中 $n_{toks}$ 是处理的 token 数量(接受的 + 推测的)。$c_{base}$ 项捕获了每次传播的开销,例如移动参数/激活值(全局 → 共享内存)、内核启动和临时分配。$c_{tok}$ 项表示 GPU/CPU 上每个 token 的平均计算成本。

总的 rollout 延迟是 $T_{rollout} = C + c_{base}N_{fwd} + c_{tok}N_{toks}$,其中 C 包括非前向传播的开销,如输入准备、调度和输出格式化。

因此,rollout 延迟分解为三个部分:(1)非前向传播开销 C,(2)基础(每次传播)成本 $c_{base}N_{fwd}$,以及(3)依赖于 token 的成本 $c_{tok}N_{toks}$。为了减少延迟,我们应该同时最小化前向传播的次数 $N_{fwd}$ 和解码的 token 总数 $N_{toks}$(推测的和非推测的)。在推测解码中存在一个权衡:增加推测的 token 可以减少 $N_{fwd}$,但提出过多的 token 会引入额外的系统开销(见图 8)。因此,一个最优策略必须平衡推测-token 预算以最大化整体的 rollout 加速。

这个公式也凸显了为什么长的生成主导了总延迟:它们不仅产生更高的依赖于 token 的成本,还决定了整个批次所需的前向传播次数,从而放大了基础成本部分。

4.2.2 最佳推测-token 预算

为了推导出一个好的推测-token 预算配置,考虑一个包含 n 个请求 $\{r_i\}_{i=1}^n$ 的批次,每个请求的特征如下:
* 目标生成长度 $l_i$,
* 接受效率参数 $\alpha_i > 0$,
* 总提议的 token 数量 $p_i$(包括推测和非推测的 token)。
* 草稿模型容量因子 $k_i \in (0, 1]$,表示请求 $r_i$ 可实现的最大接受 token 比例。

接受的 token 总数遵循一个饱和形式

$$A_i(p_i) = k_i l_i(1 - e^{-\alpha_i p_i / l_i}),$$

当 $p_i \rightarrow \infty$ 时,$A_i(p_i) \rightarrow k_i l_i$,这反映了两个模型之间固有的不匹配限制。剩余待生成的 token 数量为:

$$l_i - A_i(p_i) = l_i(1 - k_i + k_i e^{-\alpha_i p_i / l_i}),$$

完成所有请求所需的总前向传播次数为

相应的 rollout 延迟由以下模型描述

$$\begin{aligned} \begin{aligned} J(\mathbf{p}) & = c_{\text{base}} \cdot \max_{i} [l_{i}(1 - k_{i} + k_{i} e^{-\alpha_{i} p_{i} / l_{i}})] \\ & + c_{\text{tok}} \cdot \sum_{i=1}^{n} p_{i} + C, \end{aligned} \end{aligned}$$

约束条件为 $p_i \geq 0$ 以及任何系统层面对推测扩展的限制。

使用 $N_{fwd}$ 表示推测后有效的前向传播次数,我们可以将优化问题重新表述为:

$$\min_{\mathbf{p},\, N_{\mathrm{fwd}}} \quad c_{\mathrm{base}} N_{\mathrm{fwd}} + c_{\mathrm{tok}} \sum_{i} p_{i},$$ $$\text{s.t. } l_i(1-k_i+k_ie^{-\alpha_i p_i/l_i}) \leq N_{\text{fwd}}, \quad p_i \geq 0.$$

在最优状态下,对于活跃的请求,约束是紧的

$$l_{i}(1-k_{i}+k_{i}e^{-\alpha_{i}p_{i}/l_{i}})=N_{\mathrm{fwd},}$$

这给出了推测预算的闭式解

$$\begin{aligned} \begin{aligned} p_{i}^{*} & =-\frac{l_{i}}{\alpha_{i}} \ln \left(1-k_{i}\left(1-N_{\mathrm{fwd}} / l_{i}\right)\right), \text { for } N_{\mathrm{fwd}}<l_{i} ; \\ p_{i}^{*} & =0 \text { otherwise. } \end{aligned} \end{aligned}$$ <p>将方程式 (7) 代入方程式 (5) 得到一个单变量的目标函数

$$\begin{aligned} \begin{aligned} J(N_{\mathrm{fwd}}) = & \ c_{\mathrm{base}} N_{\mathrm{fwd}} + C \\ & + c_{\mathrm{tok}} \sum_{i:\, l_{i} > N_{\mathrm{fwd}}} \frac{l_{i}}{\alpha_{i}} \left[ -\ln(1 - k_{i}(1 - N_{\mathrm{fwd}}/l_{i})) \right] \end{aligned} \end{aligned}$$

对 $N_{fwd}$ 求导得到最优性条件

$$c_{\text {base}}-c_{\text {tok }} \sum_{i: l_{i}>N_{\text {fwd }}} \frac{k_{i}}{\alpha_{i}\left(1-k_{i}+k_{i} N_{\text {fwd }} / l_{i}\right)}=0.$$

观察。从这个公式中,我们得出三个关键观察:
1. 最优推测预算 $p_i^*$ 随着请求长度 $l_i$ 的增加而增加,长度相似的请求会获得相似的推测 token 预算。
2. 对于 $l_i \le N_{fwd}$ 的短生成,应跳过推测。
3. 容量因子 $k_i$ 限制了最大的推测增益。当 $k_i$ 很小(草稿模型较弱)时,$p^*$ 和可实现的加速都会减小,因为额外的推测 token 带来的回报递减。
4. 当 $c_{base} \gg c_{tok}$(基础成本主导的情况)时,最优策略优先减少前向传播的次数 $N_{fwd}$,这与在小批量 rollout 中的经验发现一致。

这些发现与直觉相符:一个请求产生的开销越大,就应该为它分配越多的计算精力。长的生成对总延迟的贡献不成比例,因此需要更激进的推测解码,而短的生成施加的开销最小,因此从推测中获益甚微。

4.2.3 通过运行时长度预测实现动态草稿预算


图 9. 每个点代表一个问题,其中 x 轴显示了 90 个周期的平均生成长度,y 轴显示了观察到的最大生成长度。广泛的分布和高的上限表明生成长度行为是高度动态的。

从上一节的步骤方程式 9 中,我们了解到准确的长度预测对于选择合适的预算草稿至关重要。然而,生成动态是高度随机的(见图 9),使得直接预测变得困难。因此,我们提出了一种结合了历史统计数据和运行时信号的分层启发式方法:
1. 长度类别策略。我们将请求划分为三个长度类别——长、中、短——每个类别对应一个推测预算。长请求使用较大(更激进)的推测预算,中请求使用中等预算,而短请求则禁用推测解码。
2. 从历史记录初始化。初始类别是根据与 r 相似的请求的历史分布选择的:$Init_r = \arg \max_{c \in \{\text{Long,Medium,Short}\}} \#\{ r' \sim r : r' \in c \}$,其中 $\#\{\cdot\}$ 计算与 r 相当的历史请求中属于类别 c 的数量。
3. 运行时更新。在生成过程中,我们根据观察到的部分长度 l 动态更新类别:$Class_r | l, Init_r = \arg \max_{c \in \{\text{Long,Medium,Short}\}} P(c | l, Init_r)$。我们从历史 rollout 统计数据中估计 $P(c | l, Init_r)$,以获得一个用于在线长度分类的实用先验。然后调整推测解码设置以匹配当前类别。

A4 实验环境

我们在 SOTA RL 训练框架 VeRL (Sheng et al., 2024) 上评估我们的系统,该框架实现了核心的生成、奖励和训练循环。

A4 实验结果

5.1 数学 RL

实验内容:在 DSR-sub 数据集上,使用 DeepSeek-R1-Distill-Qwen-7B 模型进行训练。配置为最大序列长度 16K,训练批大小 128,每个问题采样 16 次,在单台 8xH100 节点上运行(有效批大小为 256)。训练进行 30 步,采样温度 T=0.6。
实验结果:与 VeRL 基线相比,DAS 将总 rollout 时间减少了超过 50%。
分析结论:由于推测解码是一种无损加速方法,DAS 实现了与 VeRL 基线完全相同的奖励曲线,表明在显著加速的同时没有牺牲训练质量。


图 10. DeepSeek-R1-Distill-Qwen-7B 模型的训练曲线,比较了 VeRL 基线(橙色)和我们的 DAS 方法(绿色)。左图:每个训练步骤的生成时间(rollout 延迟)。DAS 相对于基线持续将生成时间减少了 50% 以上。右图:每个训练步骤的奖励。DAS 在所有步骤中都与基线奖励紧密匹配,表明尽管加速了,但训练质量没有损失。

5.2 代码 RL

实验内容:在 DeepCoder 任务上,使用 Qwen3-8B 模型进行训练。配置为每个问题采样 8 次,每个 GPU 的训练批大小为 32,最大序列长度 16K,采样温度 T=0.6。使用两台 8xH100 节点进行数据并行训练(有效批大小为 16)。
实验结果:DAS 将 rollout 时间减少了约 25%。
分析结论:DAS 在保持与基线相当的奖励的同时,显著减少了 rollout 时间。


图 11. Qwen3-8B 模型的训练曲线,比较了 VeRL 基线(橙色)和我们的 DAS 方法(绿色)。左图:每个训练步骤的生成时间(rollout 延迟)。DAS 在整个训练过程中实现了显著更低的生成时间。右图:每个训练步骤的奖励。DAS 在各个步骤中都与基线奖励紧密匹配,表明尽管加速了,但学习质量没有下降。

5.3 消融研究

分布感知的推测解码


图 12. Qwen3-8B 策略在三种设置下每个训练步骤的 rollout 生成时间:VERL 基线(橙色)、DAS(无限制推测预算,蓝色)、DAS(带分布感知,绿色)。无限制预算变体(蓝色)允许草稿模型提出任意长的延续,这增加了验证成本,并使端到端增益相比分布感知版本(绿色)减少了 15%。

批次和序列长度
* 实验内容:在 Qwen3-8B 代码训练中,测试了 DAS 在不同操作轴上的鲁棒性。首先,将最大解码长度从 16k 调整到 8k。其次,将有效批大小从 32 调整到 16。
* 实验结果:在 8k 序列长度下,DAS 仍然实现了超过 30% 的端到端 rollout 加速。在较小的批大小下,也保持了相似的加速比例。
* 分析结论:DAS 的有效性不依赖于特定的批处理或序列长度,因为它通过减少顺序的目标模型前向传播次数来提速。这种不变性证明了方法的普适性。


图 13. Qwen3-8B 模型的训练曲线,比较了不同序列长度(8k)和批大小(16)下的 VeRL 基线(橙色)和 DAS(绿色)。左图:将最大生成长度从 16k 减少到 8k,仍然产生超过 30% 的端到端 rollout 加速,表明 DAS 持续加速了那些长的、高延迟的轨迹。右图:将有效批大小从 32 减少到 16,保留了相似的加速比例,显示 DAS 在不同批大小下仍然有效。我们省略了奖励图,因为它之前已经展示过,并且源于推测解码的分布保持特性。

A5 结论

强化学习已成为大型语言模型后训练的核心,但在当前规模下,大部分墙上时钟成本来自 rollout 阶段。生成长轨迹速度缓慢,并且每个批次中最慢的几个样本决定了整个步骤的时间。推测解码是一种行之有效的方法,通过让一个快速的草稿模型提出多个 token,并让主模型并行验证它们,从而在不改变其输出分布的情况下降低生成延迟。

本文通过以下方式将推测解码应用于 RL:
(i) 用一个从近期 rollout 中持续刷新的非参数化、自适应的草稿模型取代了固定的神经草稿模型。
(ii) 引入了一种分布感知的草稿预算策略,该策略将更多的推测预算分配给主导 rollout 时间的、更难、更长的问题。

我们在数学推理和代码 RL 场景中验证了该方法,这些场景中的 rollout 通过可验证的信号(解决方案正确性或单元测试通过)获得奖励。在这两种情况下,我们的方法相对于无推测的基线,将 rollout 的墙上时钟时间减少了 30% 以上,同时与基线的奖励曲线相匹配,这表明自适应推测解码可以在不降低策略质量的情况下加速 RL 训练。

A6 附录

C. 非线性接受模型的推导

每轮接受动态。在推测解码期间,每个请求 $r_i$ 会经历几轮草稿。令 $d_{i,k}$ 表示在第 k 轮中提出的推测 token 数量,而 $a_{i,k}$ 表示该轮的接受率。经验轨迹(例如,DYSPEC, DEEPSEEKR1)显示,接受率随着轮次的深度呈指数衰减:

$$a_{i,k} = a_{i,0}e^{-\beta_i(k-1)},$$

其中 $\beta_i$ 量化了草稿模型与主模型不匹配程度的增长速度。

跨轮次聚合。经过 $K_i$ 轮后,接受的 token 总数为:

$$A_i=\sum_{k=1}^{K_i} a_{i,k} d_{i,k}=a_{i,0} d_i \frac{1-e^{-\beta_i K_i}}{1-e^{-\beta_i}},$$

假设每轮的提议长度 $d_i$ 固定 ($p_i = K_i d_i$)。代入 $K_i = p_i / d_i$ 得到:

$$A_i(p_i) = \frac{a_{i,0}d_i}{1 - e^{-\beta_i}}(1 - e^{-\beta_i p_i / d_i}).$$

归一化与简化。为简单起见,我们重新参数化常数,以归一化形式表达相同的指数饱和行为:

$$A_i(p_i) = k_i l_i(1 - e^{-\alpha_i p_i / l_i}),$$

其中 $\alpha_i$ 编码了请求 i 的有效“草稿效率”,$k_i \in (0, 1]$ 表示请求 $r_i$ 可实现的最大接受 token 比例。当 $p_i \ll l_i$ 时,模型简化为线性区域 $A_i \approx \alpha_i p_i$;当 $p_i \gg l_i / \alpha_i$ 时,它饱和于 $A_i \rightarrow l_i$。