ReSpec: Towards Optimizing Speculative Decoding in Reinforcement Learning Systems

1 · 动机:naïve SD 在 RL 里为什么不工作

1.1 历史脉络

Speculative decoding 在 serving 系统里已经被打磨得非常成熟:SGLang / TensorRT-LLM / OpenAI predicted-outputs / AWS SageMaker 都默认带 SD。EAGLE-3 把 lossless SD 的 acceptance length 推到 ~4–6 tokens,1.5–2× 的加速基本是即插即用的。

把这套搬进 RL post-training 听上去很顺:RL 一个 step 里 generation 占 75–86% 时间(Table 1),减半就是端到端减半。GRPO/DAPO 这类 group sampling 又把每个 prompt 复制成 8/16 个 rollout,token 量再翻几倍。看起来 SD 是天造地设的优化。

但作者把 EAGLE-3 直接挂在 VeRL+SGLang 上一跑,事情立刻不对劲(Figure 5):Qwen2.5-7B 在 math 数据集上,带 EAGLE-3 跑 200 步,reward 不仅没涨,反而从 0.3 跌到 0.2;不带 SD 的 baseline 同期能涨到 0.4。SD 不是没加速,而是把 actor 训坏了。

1.2 别的方案为什么不够

在 ReSpec 之前,RL × SD 这条线已经有几个方向,但都各有缺口:

1.3 为什么这事不平凡

把上面的故事聚成三个 research gap(论文里 G1/G2/G3,这是文章的骨架):

三个 gap 中,G3 是最反直觉的:经典 SD 的 acceptance test (Eq.1) 数学上保证每个 token 的 marginal 分布等于 target——很多人因此以为 SD 是 lossless 的、不影响训练。但 RL 的 return 是非线性 trajectory function:第 5 个 token 不一样,后面 100 个 token 的奖励可能差出一个数量级。Eq.3 给出的 multi-token 方差是 (1+δ)^|T| − 1,指数膨胀。drafter 偏爱的"看起来合理"的短语会反复出现在 replay buffer 里(论文那个例子:drafter 老说 "the key to success",而 actor 真正高 reward 的续写是 "the key of life"),于是 gradient 就被这些 lower-reward continuation 反复"投票"。这是 distribution bias,不是 variance,根本无法靠多采样消除。

另外两个 opportunity 让方案有抓手:① generation 工作流是高度倾斜的(Figure 6:active batch 在一个 rollout window 内从 32 衰减到 1),意味着 SD 配置不能 static,得跟 active batch 走;② SD 的 verification 步骤已经天然在算 target logits,这正是 KD 想要的 soft target——白送的监督信号,只要顺手存下来就能拿来 distill。

2 · 背景速查

SD acceptance test (one-shot 版)

其中 q 是 drafter,p 是 target。被拒时从 residual r(x) = max(0, p−q) / Σ max(0, p−q) 重采样。marginal 分布无偏是 Leviathan'23 的核心定理。

SD cost model (Eq.4)

C_q: drafter 单 token 成本;C_p: target 单 token 成本;α: 并行验证效率(0<α≤1);r: 接受率。当 batch 已经把 GPU 喂饱时 α→1 但 C_q 反而成了纯开销 → SD 退化为净亏损。

3 · G1 解法:Adaptive Speculative Decoding Server

核心观察(Figure 7 那张热力图):同一份 SD 配置 (s=5, t=8, n=32) 在 batch=2 时 1.46× speedup,在 batch=32 时变成 0.76×(亏 24%)。所以"挑一个最好的 (s,t,n) 用一辈子"这条路根本走不通。

ReSpec 的方案分两层:

• Solver (offline):训练前在目标硬件上跑一轮 profiling,枚举 (s,t,n) × batch size 的吞吐表,拟合一个 lookup model。这一步 lightweight,只做一次。
• Scheduler (runtime):每个 decoding step 看当前 active batch,查表选最优配置;如果连 SD 本身都不划算(预测 speedup < 1),直接降级到纯 target decode。

难点在状态切换:从 spec 切回 non-spec,或反过来,KV cache 状态不一样。ReSpec 用一个 per-request 的 flag,把 decoding 抽象成两个状态:

• Non-spec → Spec:复用 prefill 接口,把当前 decode batch "promote" 成 spec-enabled batch,不改 decoding kernel
• Spec → Non-spec:直接丢掉 speculative metadata(候选树等),继续普通 decode

Worked example

假设 H100 上 Qwen2.5-7B + EAGLE-3,profiler 测到下表(简化版):

追踪一个 prompt 的生命周期:rollout 开始时 batch=32 → Scheduler 选 non-spec(虽然 EAGLE-3 已经加载,但不调用 draft path);跑了 200 token 后 28 个 sequence 已经 EOS,active=4 → 切到 spec,先用一次 prefill-style forward 把 KV cache 补到当前位置(这是论文里 "Extend Once" 那一步,Figure 9b 中间的灰条),之后正常 spec decode。这样大批和长尾两段都吃到甜头。

4 · G2 解法:Drafter Evolution via on-policy KD

EAGLE-3 在 serving 里只蒸馏一次就稳稳跑;在 RL 里 actor 每 step 在变,所以训练 100 步后,Figure 4 里 acceptance length 已经从 4.0 滑到 ~2.5。修复思路简单又巧妙:SD 的 verification 已经在算 target logits 了,顺手存下来当 KD 的 soft target,反过来更新 drafter。

具体流程(Algorithm 1 简化):

1. 每个 rollout sample 产生 (x, y, log p, log q, r) 五元组,其中 log p 是 verification 自动给出的 target logits
2. 这些五元组存进 replay buffer Q
3. 每 I 步训练一次 drafter,目标是 KL(softmax(log p) ‖ q_θ(·)),梯度只对当前 q_θ 反传(不对 log q 反传——log q 只是诊断)
4. 更新完的 drafter 推回 inference engine,下一个 RL step 用

Worked example: 一条 rollout 的生命周期

假设 prompt = "求 ∫₀¹ x²dx",actor 输出 100 个 token。其中第 17 个位置 SD verify 同时产出长度 V=151936(Qwen2.5 vocab) 的 logits 张量 log p_17 ∈ ℝ^V。存盘时只存 top-K=32(节省存储,论文默认这么做),其它位置 mask 掉。

句子级 r=0.85(答对了 1/3)。Loss = 0.85 · (0.32 + 0.18 + 0) = 0.425,反向更新到 drafter 的 token-mlp head。下一个 step,drafter 在"x²"位置的概率会被推高,acceptance length 期望提升 ~5%。

5 · G3 解法:Reward-Weighted KD 核心创新

这是 ReSpec 区别于其它 RL × SD 工作的真正新东西。前两个 gap 的解法多少有先例(adaptive SD 是 serving 里的常识,on-policy KD 是 DistillSpec 改 RL),但 G3 第一次被诊断为"SD 不仅没加速,还在偷换梯度",并且解法是把 RL 的 reward 信号注回到 SD 的 KD 损失里。

问题精确化

朴素 KD(标准 KL,所有样本平权)在 RL 里有个反馈环路:

1. drafter 学习全部 rollout(包括 reward=0 的失败 trajectory)
2. drafter 学到的"平均"分布偏向 actor 的 noisy 行为
3. 下一轮 SD 用这个 drafter 提议,acceptance test 在marginal上无偏,但在 multi-token 上把 actor 朝 drafter 的偏置方向"拉过去"
4. 更多 noisy rollout 进 buffer → drafter 更偏 → ...

论文给出的修正非常直接——给每个 sample 的 KD 损失加一个 reward 权重 w(r):

默认 w(r) = r(实际实现里 normalize + clip 防止极端值)。直觉:reward 高的 trajectory 是"actor 的优秀样本",让 drafter 重点学这类;reward 低的 trajectory 是 actor 的失败案例,drafter 不该去模仿。

Worked example: GRPO group of 5

同一道数学题 GRPO 采 5 个完成,reward = [1.0, 0.8, 0.3, 0.3, 0.0]。每条算自己的 KL_t = Σ_t KL(p_t ‖ q_t),假设都接近 0.4。

朴素 KD 的更新方向是 5 条样本的平均;reward-weighted 的更新方向 ≈ 0.42 · ∇top1 + 0.33 · ∇top2,实际把 drafter "拉"向 reward=1.0 那条 trajectory 的预测分布。failed sample 完全不传梯度。

消融对比 (Figure 10)

这张表是论文最重要的数据:平权 KD 比"完全不更新"还坏,因为它主动把 drafter 训成了"actor 的均值",而 actor 均值充满失败 trajectory。Reward-weighted 几乎完美追上 baseline 曲线。

6 · Async Update Overlap (G2 系统侧补丁)

drafter 也是个模型,反向传播要 GPU 时间。如果每个 RL step 都同步训一次 drafter,generation 必须等更新完成 → pipeline bubble。ReSpec 用两个机制摊薄这个开销:

• Replay buffer 累积:每 I 步才更新一次 drafter,且每次只用 buffer 的 1/I,把"频繁小更新"摊成"低频次但 batch 大"
• Async overlap:drafter 训练塞进 RL pipeline 的 idle slot——actor weights sync / reward inference 等阶段,GPU 有空,顺势跑 drafter forward+backward

worked example: I=1 时(每步更新),Qwen-7B reward 涨到 0.42;I=3 时降到 0.35;I=5 时降到 0.30 甚至低于 sync。所以默认 I=1 + async,既频繁又不阻塞。这是系统级 vs 算法级 freshness的小心 trade-off。

7 · 关键公式与数值敏感性

7.1 Multi-token 接受概率方差(Eq.3)

把 χ² divergence 近似为常数 δ,方差按序列长度指数增长。这给 G3 提供了理论:即使每个 token 单步 marginal 无偏,接受 5-token 块时方差是 (1+δ)⁵ − 1。

物理直觉:drafter 越 stale (δ↑),长 draft block (|T|↑) 的方差爆炸越快。这就是为什么 ReSpec 必须同时做 G2(让 δ 小)和 G3(用 reward 选好的 trajectory),仅其一不够——只解 G2 也会因为剩余 δ 在长 block 下放大;只解 G3 不让 drafter 更新,δ 会从 0.05 飙到 0.5,把 1.25 的方差变成 24.6。

7.2 SD cost (Eq.4) 在 RL 下的退化

r 是接受率(acceptance rate)。当 r → 0,cost → ∞,SD 比裸 decode 慢。把 G2 staleness 写进去:r 从 0.7 (新鲜) 降到 0.3 (stale 100 步) 时,假设 C_q=0.1 C_p,α=0.7:

8 · 实验关键结果

8.1 端到端 reward stability (Figure 12)

读法:naïve EAGLE-3 在 3B/7B 上明显伤 reward;ReSpec 不仅追平 baseline,有时还略高(可能 reward-weighted KD 起到了"选择性放大优秀样本"的隐性正则)。14B 上差距小,符合"大模型对 stale drafter 更鲁棒"的直觉。

8.2 端到端 speedup (Figure 13)

3B 的 4.5× max speedup 是瞬时峰值(应该在 long-tail 单 sequence 阶段);avg 1.84× 是更可靠的数。

8.3 消融 (Figure 14, Qwen-14B)

读法:贡献排序是 RW-KD ≫ Adaptive Server > Async overlap。RW-KD 是 main lifter——这跟我前面说的"它是真正的算法创新"一致。注意 RW-KD 既包含 G2 修复(让 SD 不退化)又包含 G3 修复(让 actor 不被带偏),所以 1.48× 里其实是"让 SD 在 RL 里能用"这个基本能力的全部。

9 · 与同类工作对比

定位:ReSpec 在三件事上都打了卡,而其它工作每件至少缺一件。如果说 NeMo-RL 是 "把 SD 装进 RL pipeline 的工程示范",SPEC-RL 是 "把 SD 用得更聪明",那 ReSpec 是 "在 RL 训练动态下重新设计 SD 算法"。Reward-weighted KD 是其它工作目前都没碰的角度。

10 · 局限 / 个人 take / 待验证问题

• 实验规模有限:只到 14B,没碰 MoE 或更大稠密模型。Qwen3-32B / DeepSeek-V3 671B 上 reward-weighted KD 是否仍有这么大边际收益未知。
• w(r)=r 太朴素:reward 在 GRPO 里通常 ∈ [0,1],但绝对值受 reward model scale 影响。换 task(代码 / agent / preference)后 clip 阈值需要重调,论文没给 sensitivity。
• Adaptive Server profiling 离线一次:模型一变(actor 蒸馏 / LoRA / quantize)profiling 表就需要重做。社区里需要一个轻量的 online profiler。
• 4.5× 是峰值,avg 1.5–1.8×;真正落到生产线的口径应该是 avg。Qwen-14B 才 1.50×,跟 NeMo-RL 报的 1.4×(8B sync)接近。
• Reward-weighted KD 与 RLHF 偏好奖励的耦合:如果 reward model 本身有偏(过拟合/sycophancy),drafter 会放大这个偏差;论文没讨论 reward model robustness 的 downstream 影响。
• EAGLE-3 限定:整套设计基本绑死 EAGLE-3。如果换成 Medusa / Lookahead / Cascade,接口和 KD 形式都要改。

待验证问题清单

1. 把 w(r) 换成 normalized advantage 或 quantile-rank,在 GRPO 里是否更稳?
2. Reward-weighted KD 对 reward hacking 是否敏感?如果 reward model 被 hack,drafter 会不会把 hack 放大?
3. 14B 上 RW-KD 增益缩小到 ~0,是否 ReSpec 在更大模型上的核心贡献会被稀释?
4. Async I=1 vs sync 在 64+ GPU 大集群上的 weight broadcast 成本如何?(论文是 16 GPU 测的)
5. Drafter staleness 与 actor learning rate 的耦合:大 LR 下 δ 增长更快,KD 更新频率是否需要随 LR 自动 scale?
6. 把 Adaptive Server 的 (s,t,n) 用 RL bandit 在线学,而不是 offline profile,是否更适应数据集 / horizon 漂移?

11 · Memory points