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

发表时间: 2025-11 · arXiv:2511.13841 (UCSD WukLab et al., MLSys 2026)

Zelei Shao, Vikranth Srivatsa, Sanjana Srivastava, Qingyang Wu, Alpay Ariyak, Xiaoxia Wu, Ameen Patel, Jue Wang, Percy Liang, Tri Dao, Ce Zhang, Yiying Zhang, Ben Athiwaratkun, Chenfeng Xu, Junxiong Wang

A1 主要贡献

本文旨在解决大型语言模型(LLM)强化学习(RL)后训练中的一个关键效率瓶颈:rollout(轨迹生成)阶段耗时过长。研究发现,在现代LLM的RL训练中,rollout阶段占据了超过70%的总训练时间,远超反向传播和参数更新的开销。这一瓶颈主要由LLM解码的自回归特性、生成长度的增加以及为追求高精度而增大的样本量共同导致。

核心问题与洞察
现有用于加速LLM推理的系统(如服务系统)的优化技术无法直接有效地应用于RL训练的rollout阶段,原因如下:
1. 长尾延迟问题 (Insight-1):RL rollout要求一个批次内的所有样本全部完成推理后才能进入下一训练阶段。而服务系统优化的指标(如首个令牌时间TTFT和每输出令牌时间TPOT)会导致长序列耗时更长,从而在RL rollout中产生严重的“长尾”现象,即少数长序列决定了整个批次的完成时间。
2. 样本复用特性 (Insight-2):RL训练在每个迭代中会重复使用同一组样本(prompts),而服务系统通常假定每个用户请求都是独立的,因此未能利用这种请求重复出现的特性来加速。
3. 模型动态变化 (Insight-3):RL训练中,目标模型的权重是持续更新的,这使得在传统服务场景中有效的、基于固定模型的优化方法(如使用预训练的草稿模型)变得不再适用。

研究目标与创新点
为解决上述问题,本文提出了一套名为DAS (Distribution-Aware Speculative decoding)的系统与机器学习协同设计方法,该方法通过对传统推测解码(Speculative Decoding, SD)进行三方面关键改造,以适应RL rollout的独特性质,从而在不改变模型输出和训练质量的前提下,显著提升rollout效率。

主要创新点
1. 自适应、非参数化草稿模型:针对RL训练中目标模型动态变化的问题(Insight-3),DAS不使用需要重新训练的神经网络草稿模型,而是利用样本复用特性(Insight-2),基于近期rollout历史构建一个通过后缀树实现的非参数化草稿模型。该后缀树结构会动态剪枝,仅保留近期相关的历史rollout,从而能持续与当前策略模型保持一致,确保高接受率。
2. 感知长度分布的推测策略:针对长尾延迟问题(Insight-1),DAS提出了一种感知分布的推测解码方法。它不再对所有样本统一进行推测,而是优先将更多的“草稿预算”分配给那些预计生成长度更长、难度更高、主导rollout完成时间的任务。这种差异化的策略通过为每个问题构建和更新独立的后缀树来实现,从而能更好地捕捉领域特定模式,提高推测准确性。

核心贡献总结
- 提出了关于RL rollout阶段独特属性及其对系统设计影响的三个关键洞察。
- 设计并实现了完整的RL系统DAS,相比当前最先进的RL系统,可将rollout时间减少高达50%。
- 提出了一种专为RL rollout设计的、无需训练、自我演进的推测解码方法。
- 提出了一种感知分布的推测解码方法,有效减少了RL rollout中的长尾延迟。

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

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

RL rollout中的长尾瓶颈。RL后训练以迭代方式运行,其中actor生成一组固定的轨迹,learner则更新策略。在实践中,像VeRL和OpenRLHF这样的系统倾向于使用数据并行(DP)的rollout worker来扩展解码吞吐量,仅在训练模型无法装入内存时才使用张量并行(TP)。解码开始时并行度最高,但随着短序列的完成,有效批次大小会逐渐“崩溃”,最终由少数长响应(即“掉队者”)决定了整个步骤的完成时间(makespan)——这是在LLM系统中也观察到的典型长尾行为。这导致GPU利用率低下,因为在掉队者继续解码时,其他worker处于空闲状态。多项RL/LLM研究报告称,尽管采用了复杂的批处理技术,解码阶段仍主导了墙钟时间,并且加速器利用率不足。图1展示了一个代表性设置下有效批次大小的剖析:大约100个解码步骤后,并行度急剧下降,证实了RL训练环境中存在长尾运行时瓶颈。这启发我们使用推测解码(speculative decoding)【Leviathan et al., Fast inference from transformers via speculative decoding, 2023b】:通过起草多个token并并行验证它们,我们可以缩短掉队者的运行时间并回收空闲的GPU时间,从而在不改变模型输出的情况下实现显著的端到端加速。虽然推测解码通常在小批量、对延迟敏感的服务场景中进行评估,但RL训练在解码过程中表现出有效的批次崩溃——短序列首先完成,而掉队者主导了步骤。这暴露了推测可以利用的空闲计算能力。
图1. 在没有DAS的情况下,rollout期间有效批次大小的崩溃。测量于DeepSeek-distilled 7B模型(DeepScaleR提示)。随着解码的进行,短序列首先完成,有效批次大小缩小,留下少数长序列的“掉队者”来决定步骤的完成时间。使用我们的方法,我们既可以减少总延迟,又可以减轻长尾掉队者的影响。

RL Rollout与近期历史的高度相似性。在不同epoch之间,针对相同提示(prompt)生成的轨迹表现出明显的词汇和结构上的重用。在图2中,我们观察到与近期轨迹的相似度较高,而与久远轨迹的相似度较低,这与策略漂移(policy drift)的现象一致。这种“新近度偏见”(recency bias)意味着,一个基于历史索引的、无模型的草稿生成器(例如,后缀数组【Manber & Myers, Suffix arrays: a new method for on-line string searches, 1993】或后缀树【Ukkonen, On-line construction of suffix trees, 1995】)可以挖掘近期的延续(continuations)来提出高质量的草稿,即使策略在不断演进。
图2. (左) 使用N-gram计算的每个迭代的内容相似性(重用率)。(右) Qwen2.5-7B-Instruct在不同epoch间的成对相似度。靠近对角线的块状结构表明,rollout与最近epoch的rollout最为相似,且相似度随着时间距离的增加而衰减。这反映了策略漂移:随着策略的不断更新,较早的生成对于当前行为的预测性会降低。

A2 方法细节

我们引入了一个名为DAS(distribution-aware speculative decoding)的感知分布的推测解码框架,这是一种通过调整推测解码来加速rollout的RL后训练方法。图3提供了该系统的概览。DAS建立在使用基于偏好或可验证性奖励来优化策略的RL流水线之上,它维护一个基于历史索引的、非参数化的草稿生成器,该生成器通过近期的rollout不断刷新,以在策略演进过程中保持高接受率。然后,感知分布的推测解码组件决定为每个请求分配多少推测预算,优先考虑那些生成时间长、延迟高的问题。通过将自适应的预算分配与一个在线的、与分布对齐的推测器(以在线后缀树构建的方式增量更新)相结合,该系统在不修改奖励循环的情况下降低了rollout延迟。我们在4.1节中阐述了选择自适应、非参数化草稿生成方法的原因,并在4.2节中描述了如何分配草稿预算以减少长尾序列的运行时间。

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

4.1 感知分布的草稿生成器

4.1.1 参数化的静态草稿生成器

EAGLE系列技术的局限性。EAGLE【Li et al., EAGLE: Speculative sampling requires rethinking feature uncertainty, 2024a】、EAGLE-2【Li et al., EAGLE-2: Faster inference of language models with dynamic draft trees, 2024b】和EAGLE-3【Li et al., EAGLE-3: Scaling up inference acceleration of large language models via training-time test, 2025】是当前最先进的基于推测解码的推理时服务技术。一个轻量级的EAGLE头通过外推内部特征来构建一个令牌草稿树,然后目标模型并行验证该树,从而在推理期间实现显著的延迟降低。特别是EAGLE-2,它依赖于一个经过良好校准的草稿生成器,其置信度能紧密预测令牌的接受情况,并据此动态地生成草稿树。然而,在RL训练中,策略是非平稳的:模型权重在每次learner更新后都会改变,因此这种校准会迅速失效。图2表明,从早期epoch模型收集的轨迹所导致的接受率将显著低于从后期检查点收集的轨迹。在实践中,EAGLE要么必须忍受接受率下降(从而导致加速效果降低),要么必须在整个训练过程中反复重新训练/调整其头部及其建树阈值——这给本已由rollout主导的阶段增加了计算和工程开销。

集成复杂性与非平稳环境的挑战。此外,EAGLE的树形草稿生成和并行验证内核是为高吞吐量推理服务而设计的。在一个actor-learner RL流水线中,检查点是随时间演变的,将这些组件集成并持续重新校准,远比其预期的仅推理用例要复杂得多。由于这些原因,尽管EAGLE在推理方面取得了很好的效果,但在非平稳的RL训练机制中难以高效扩展。持续更新神经模型的高昂开销促使我们采用一种基于文本索引的无模型方法,特别是后缀树【Ukkonen, On-line construction of suffix trees, 1995】和后缀数组【Manber & Myers, Suffix arrays: a new method for on-line string searches, 1993】。

4.1.2 自适应非参数化草稿生成器

基于后缀树的自适应草稿生成器设计。基于上述动机,我们探索了一种自适应的非参数化草稿生成器。我们利用一个基于近期rollout的滑动窗口构建的后缀草稿生成器,通过维护一个关于过去生成内容的紧凑后缀树。在rollout时,我们在树中定位当前上下文与树的最长匹配(对于长度为m的上下文,自上而下的查找时间为O(m),m是我们查找的查询/模式的长度——例如,当前的解码前缀),然后沿着匹配路径提出一个多令牌的草稿。目标模型在一个或几个批处理步骤中验证这个草稿。验证后,我们在线更新树,这样推测器就能持续适应不断演变的策略。这种设计捕捉了训练各epoch中反复出现的模式,同时避免了训练或维护一个独立的、基于训练的神经草稿生成器的需要。重要的是,后缀树的构建和在线后缀树的更新通过Ukkonen算法【Ukkonen, On-line construction of suffix trees, 1995】以线性时间运行,这使其非常适合增量式地接收新路径。

与参数化草稿生成器的性能对比。图4比较了在RL训练中,一个参数化的静态草稿生成器(EAGLE)与我们的非参数化自适应草稿生成器。我们绘制了每轮验证平均接受的令牌数。尽管EAGLE的接受曲线大致平坦,但我们的草稿生成器的接受率随着训练的进行而持续提高,因为它会根据近期的rollout不断更新。
PROTECTED_IMAGE_26____PROTECTED_IMAGE_27

后缀树与后缀数组的比较。除了后缀树,一个自然的想法是利用后缀数组(SA),其灵感来自Infinigram【Liu et al., Infini-gram: Scaling unbounded n-gram language models to a trillion tokens, 2024a】,作为一种空间效率更高的替代方案。一个标准的SA支持通过二分查找在$O(m \log n)$时间内搜索子串,其中m是模式长度,n是语料库中的令牌数【Manber & Myers, Suffix arrays: a new method for on-line string searches, 1993】。通过增加一个LCP数组【Kasai et al., Linear-time longest-common-prefix computation in suffix arrays and its applications, 2001】可以减少比较次数,而使用增强后缀数组(ESA)【Abouelhoda et al., Replacing suffix trees with enhanced suffix arrays, 2004】则能通过有效模拟后缀树遍历,以更好的常数和缓存局部性实现$O(m)$时间的模式匹配。然而,SA和ESA本质上是静态的:动态更新通常需要(部分的)$O(n)$重建,这在RL训练中是不可取的,因为每个迭代都会有新的轨迹产生。我们的在线后缀树索引则以适度的额外空间换取了快速的增量更新和$O(m)$的最长匹配查询,这与非平稳的策略环境天然契合。我们在图5中比较了后缀树与后缀数组的推测时间和更新成本。后缀树在两个指标上都表现出卓越的性能:推测时间快2-20倍,而更新成本则显示出巨大优势,保持在亚毫秒级别,相比之下后缀数组的重建时间则不断攀升。更新性能上超过三个数量级的差异证实了后缀数组尽管空间效率高,但在在线RL训练中是不切实际的,因为每个迭代都必须快速更新新的轨迹。

全局后缀树的局限性。一个单一的、不断增长的、覆盖所有过去rollout的全局索引在统计和系统两个层面都存在问题。在统计上,策略漂移导致较早的续写与模型当前的条件分布对齐不佳,这会降低推测解码的接受率——而这恰恰是决定加速效果的关键量(见图2)。此外,尽管RL中的所有问题可能属于同一领域,但它们的多样性意味着一个问题的模式很少能可靠地转移到另一个问题上。与此一致,图6显示,启用一个全局树所带来的增益小于为每个问题维护一个树,并且由于树更大而增加了额外的开销。
图6. (左) 每轮验证平均接受的令牌数 vs. 训练步骤。我们比较了三种用于构建草稿生成器的历史记录:全局+请求、问题+请求、仅问题;问题范围的历史记录在接受率上超过了全局记录。(右) 推测解码延迟(毫秒/令牌);全局+请求由于查询和维护一个大型全局索引的成本而始终较慢,而问题范围的分片查询成本更低,因此延迟也更低。

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

滑动窗口选择树。在不同的epoch和模型中,我们观察到针对相同提示的生成内容之间存在显著的相似性(图2)。时间上相近的轨迹倾向于产生更长的被接受的草稿,而来自更早epoch的轨迹则产生较短的接受长度,这表明随着策略的漂移,匹配质量会下降。由于策略随时间演变,历史久远的生成变得预测性较差。因此,我们从一个近期轨迹的滑动窗口构建草稿生成器,并在每次迭代中刷新索引。窗口大小控制着一个偏差-稳定性(bias–stability)的权衡:较短的窗口能快速适应,但提供的匹配较少;而较长的窗口能提高覆盖率,但有陈旧的风险。在实践中,我们将窗口更新速率与优化器的步长尺度联系起来(例如,较大的参数更新意味着较短的窗口),并对源自较早epoch的匹配应用轻微的权重衰减。
图7. (左) 不同历史窗口大小下,每轮验证平均接受的令牌数 vs. 训练步骤。较大的窗口(如16、32、all)能提供更高的接受率,因为它们提供了更多的匹配续写,而更高的接受率已知能直接转化为更少的目标模型前向传播和更低的解码成本。(右) 每令牌推测解码延迟 vs. 训练步骤。“window all”显示出最高的延迟,因为查询和维护一个完整的全局历史记录更昂贵,并且包含了陈旧的轨迹,因此在实践中,中等大小的窗口(16或32)在接受率和延迟之间取得了比“window all”更好的平衡。

4.2 感知长度的推测策略

4.2.1 Rollout延迟估计

线性延迟模型。我们通过首先对每次前向传播的延迟进行建模,然后将其在所有前向传播中累加,来估计端到端的rollout延迟。通过性能剖析,我们发现一个简单的线性模型能够捕捉主要行为(平均相对误差约12%):
公式 Latency_fwd = c_base + c_tok * n_toks
其中,$n_{toks}$是处理的令牌数(已接受+推测的)。$c_{base}$项捕捉了每次传播的开销,如移动参数/激活值(全局→共享内存)、内核启动和临时分配。$c_{tok}$项代表了GPU/CPU上每个令牌的平均计算成本。
图8. DeepSeek-R1-Distill-Qwen-7B上解码延迟与令牌数的关系。显示出清晰的线性关系。

总rollout延迟。总rollout延迟为 $Latency_{total} = C + c_{base}N_{fwd} + c_{tok}N_{toks}$,其中C包括非前向传播的开销,如输入准备、调度和输出格式化。因此,rollout延迟可分解为三部分:(1)非前向传播开销$C$,(2)基础(每次传播)成本 $c_{base}N_{fwd}$,以及(3)依赖于令牌的成本 $c_{tok}N_{toks}$。为了减少延迟,我们应该最小化前向传播的次数 $N_{fwd}$ 和解码的总令牌数 $N_{toks}$(推测和非推测的)。在推测解码中存在一个权衡:增加推测令牌可以减少 $N_{fwd}$,但提出太多令牌会引入额外的系统开销(见图8)。因此,一个最优策略必须平衡推测令牌的预算,以最大化整体的rollout加速效果。这个公式也解释了为什么长生成主导了总延迟:它们不仅产生了更高的依赖于令牌的成本,还决定了整个批次所需的前向传播次数,从而放大了基础成本部分。

4.2.2 最优推测令牌预算

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

总接受令牌数遵循一个饱和形式:
公式 A_i(p_i) = k_i l_i (1 - e^(-alpha_i p_i / l_i))
其中,当 $p_i \to \infty$ 时,$A_i(p_i) \to k_i l_i$,这反映了两个模型之间固有的不匹配限制。剩余待生成的令牌数为:
公式 R_i = l_i - A_i(p_i)
完成所有请求所需的总前向传播次数为:
公式 N_fwd = max_i R_i
相应的rollout延迟模型为:
公式 Latency_total = C + c_base N_fwd + c_tok * sum(p_i)
约束条件为 $p_i \ge 0$ 以及任何系统层面对推测扩展的限制。

最优化求解。使用 $N_{fwd}$ 表示推测后有效的前向传播次数,我们可以将优化问题重构为:
公式 min_{N_fwd, p_i} C + c_base N_fwd + c_tok * sum(p_i)
受限于:
公式 l_i - k_i l_i (1 - e^(-alpha_i p_i / l_i)) <= N_fwd
在最优情况下,对于活跃的请求,约束是紧的:
公式 l_i - k_i l_i (1 - e^(-alpha_i p_i / l_i)) = N_fwd
这给出了推测预算的闭式解:
公式 p_i = - (l_i / alpha_i) * ln(1 - (1 - N_fwd / l_i) / k_i)
将式(7)代入式(5),得到一个单变量目标函数:
公式 Latency_total = C + c_base N_fwd - (c_tok / alpha_i) * sum(l_i * ln(1 - (1 - N_fwd / l_i) / k_i))
对 $N_{fwd}$ 求导,得到最优性条件:
公式 c_base + c_tok * sum(1 / (alpha_i * (k_i - 1 + N_fwd / l_i))) = 0

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

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

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

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

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

A4 实验环境

实验设置
- 基础框架:我们的系统在SOTA RL训练框架VeRL【Sheng et al., Hybridflow: A flexible and efficient rlhf framework, 2024】之上实现,该框架实现了核心的生成、奖励和训练循环。我们还利用了vLLM【Kwon et al., Efficient memory management for large language model serving with pagedattention, 2023】。
- 任务与数据集
- 数学推理 (Math RL):遵循One-Shot-RLVR【Wang et al., Reinforcement learning for reasoning in large language models with one training example, 2025】的训练流程,在DSR-sub数据集【Wang et al., 2025】上进行,该数据集包含来自DeepScaleR【Luo et al., Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl, 2025b】的1209个竞赛级数学问题。
- 代码生成 (Code RL):采用类似于DeepCoder【Luo et al., Deepcoder: A fully open-source 14b coder at o3-mini level, 2025a】的RL设置,奖励由生成代码的单元测试通过/失败来确定。

A4 实验结果

所有报告的时间均为实际墙钟rollout步骤时间,包括批处理、调度和验证开销。

5.1 数学推理 RL 实验

5.2 代码生成 RL 实验

5.3 消融研究

A5 结论

强化学习已成为大型语言模型后训练的核心,但在当前规模下,大部分墙钟时间成本来自rollout阶段。生成长轨迹速度缓慢,且每批次中少数最差的样本决定了整个步骤的时间。推测解码是一种行之有效的降低生成延迟的方法,它通过一个快速的草稿生成器提出多个令牌,并让主模型并行验证它们,而不会改变其输出分布。我们通过以下方式将推测解码应用于RL:(i)用一个从近期rollout中持续刷新的、非参数化的自适应草稿生成器取代固定的神经草稿生成器;(ii)引入一个感知分布的草稿预算策略,该策略将更多的推测预算分配给那些主导rollout时间的更难、更长的问题。我们在数学推理和代码RL设置中验证了这种方法,在这两种情况下,rollout都由可验证的信号(解决方案正确性或单元测试通过)来奖励。在两种情况下,我们的方法相对于无推测的基线,将rollout墙钟时间减少了超过30%,同时匹配了基线的奖励曲线,这表明自适应推测解码可以在不降低策略质量的情况下加速RL训练。

A6 附录

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

每轮接受动态。在推测解码期间,每个请求$r_i$会经历几轮草稿生成。设$d_{i,k}$为第k轮中提出的推测令牌数,而$a_{i,k}$为该轮的接受率。经验轨迹(例如,DYSPEC, DEEPSEEKR1)显示,接受率随着轮次的深度呈指数衰减:
公式 a_i,k = e^(-beta_i * k)
其中$\beta_i$量化了草稿模型与主模型不匹配程度的增长速度。

跨轮次聚合。经过$K_i$轮后,接受的总令牌数为:
公式 A_i = sum_{k=1}^{K_i} d_{i,k} a_{i,k}
假设每轮的提议长度$d_i$固定(即$p_i = K_i d_i$)。代入$K_i = p_i/d_i$得到:
公式 A_i = d_i * (1 - e^(-beta_i * K_i)) / (1 - e^(-beta_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$可实现的最大接受令牌比例。当$p_i \ll l_i$时,模型简化为线性区域$A_i \approx \alpha_i p_i$;当$p_i \gg l_i/\alpha_i$时,它饱和于$A_i \to l_i$。