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

1 · 动机:为什么 RL rollout 配 EAGLE 很难

1.1 历史脉络:从 serving 到 rollout 的迁移

Speculative decoding 起家于 LLM serving 场景:固定权重的 target,小巧 draft model 提议多 token,target 一次 verify。EAGLE / Medusa / SpecInfer 这条线在 serving 上做到 2–3× wall-clock 加速。

2024 起,RL post-training (RLHF / GRPO / DAPO) 占总训练时间的比例急速上升:reasoning 模型 trajectory 从 ~512 token 拉到 16k–32k,DeepSeek-R1 进一步把 long-CoT 推到几万 token。论文里的硬数据:rollout 占总训练时间 > 70%。这是 Amdahl 视角下"该被攻击的部分"。

把 spec decoding 直接搬过来?——三个 insight 决定不能简单照搬:

• Insight-1 (long tail): Serving 优化 TTFT/TPOT,RL rollout 必须等 batch 里最长那条跑完才能进下一步。短序列先完成,effective batch size 塌缩,只剩几条 straggler 决定 makespan。
• Insight-2 (reuse): Serving 假设每个 request 是新的。RL 每个 epoch 把同一批 prompt 重新跑一遍,trajectory 之间有大量结构与词汇复用——这是免费的 prior。
• Insight-3 (drift): Serving 里 target 模型权重是冻的。RL 里每个 learner step 之后 target 都变,任何预训练好的 draft 都会变 stale。

1.2 别的方案为什么不够 (alternatives table)

1.3 为什么这事不平凡

表面上"用历史 trajectory 当 draft 来源"听起来一句话能讲清,但要在生产 RL 系统里落地,有三个非平凡点:

1. 策略漂移 vs. 历史质量。一开始的 trajectory 反映早期 policy,后期已经"过期"。直接用全部历史 → 接受率掉。需要某种 sliding window + recency weighting。
2. 跨 problem 干扰。同一个 domain 的不同问题,token-level 模式不互通。一条全局 suffix tree 在统计上稀释信号,在系统上还吃 query 时间。需要 per-problem sharding。
3. 一刀切的 budget 是错的。短序列对 makespan 没贡献,给它发 draft tokens 反而增加 verification overhead;长序列才是 makespan 的 bottleneck。需要根据预测长度动态分配 budget。

这三点叠加起来才是 DAS 的论文价值。如果只换数据结构(suffix tree),会被 SuffixDecoding (Oliaro 2025) 截胡;如果只做 budget 分配,会被 SpecServe / Liu 2024 截胡。这三个 component 的耦合+针对 RL 训练循环的工程化才是新的。

2 · 背景速查 (术语 / 数据结构)

Suffix tree 30 秒回顾

给定字符串 S = "banana$",suffix tree 把 所有后缀 banana$, anana$, nana$, ana$, na$, a$, $ 压成一棵 trie,公共前缀合并到同一节点上。从 root 到任意叶子的路径就是某个 suffix。Ukkonen 1995 给出在线、单次扫描、O(n) 时间构造的算法——这正好契合 RL 里"每 epoch 进来一批新 trajectory,要 incremental 维护"的场景。

查询:给定当前 decode 已生成的 prefix q(长度 m),从 root 沿 q 指引的边下降,落到的节点的子树就是所有以 q 结尾的历史延续。在子树里取频次最高 / 路径最长的分支,就是 draft proposal。查询时间 O(m)。

3 · 方法 A — Suffix Tree Drafter

3.1 核心思路

放弃训练 draft model,改用一个对最近 N 个 epoch 的 trajectory 增量维护的 suffix tree。每次 target 即将解码下一段时:

1. 取当前 decode prefix 末尾 m 个 token (m 一般 8–32) 作为 query;
2. 在 suffix tree 里 longest-match,定位到一个内部节点;
3. 从该节点出发沿"最高频 child"路径走 d 步 (draft length),拿到 d 个 draft tokens;
4. target 一次 forward 把 (q, draft_1..d) 全 verify,接受最长 prefix。

3.2 SVG: 一棵小的 suffix tree

假设过去两条 rollout 拼起来是 S = "the cat sat on the mat$" (只为示意,实际是 BPE token 序列)。下图把它的 (示意性) suffix tree 画出来,并展示 query "the " 如何 longest-match:

3.3 Worked example: 一个 token 的命运

假设当前 epoch t,policy π_t 正在为问题 P3 生成第 1273 个 token。decode prefix 末尾 16 token 是 "...therefore the answer is "。

1. 查询:取末 8 token "the answer is " 走 P3 的 suffix tree。
2. 命中:P3 在过去 16 个 epoch 共 256 条 trajectory 里 (sliding window),有 47 条以 "the answer is " 开头的延续;最频繁的一条是 "\\boxed{42}\\n\\nLet me verify"。
3. 提议 d=8:draft = "\\boxed", "{", "4", "2", "}", "\\n", "\\n", "Let"。
4. Verify:target 一次 forward,前 5 个被接受,第 6 个 "\\n" 处 logits argmax 是空格——拒绝。
5. 结果:1 次 forward 推进了 5 token (而不是 1),省 4 次 forward。Tree 也用刚刚生成的这条 trajectory online update。

3.4 反向论证:不用 suffix tree 会怎样

• 用 suffix array (SA):更省空间但动态更新需要 O(n) 重建。Figure 5 报 update 时间差 > 3 个数量级 (sub-ms vs. seconds);RL 每 iteration 都有新 trajectory → SA 不可用。
• 用 hash table of n-grams:查询是 O(1) 但只能单一 n;suffix tree 支持任意长度longest-match 一次完成。
• 用 EAGLE 风格 draft head:每次 policy update 都要 fine-tune draft;额外 GPU 时间 + 工程复杂度。

4 · 方法 B — Per-problem sharding & sliding window

4.1 为什么 per-problem 而不是全局一棵树

论文 Figure 6 的结论:per-problem suffix tree 比 global 一棵树 接受率更高且 查询更便宜。原因:

• 统计上:问题 P1 的解题 pattern 与 P2 几乎不重合,把 P1 的 token 塞进 P2 的索引只是噪声。
• 系统上:全局树体量大,query 要走更深、更多碰撞;per-problem 树各自小,locality 好。

4.2 SVG: per-problem shard + sliding window

4.3 实证细节

Figure 6 显示三种 history 配置的平均接受 token 数:

对小模型禁用 prefix trie 的 routing,因为 routing 的 CPU 开销已经接近 draft 收益本身。

5 · 方法 C — Length-aware speculation budget

5.1 直觉

RL batch 里的 trajectory 长度分布严重偏斜(Figure 1):前 100 步后 effective batch size 已经塌缩到 <30%,只剩几条 long tail 在跑。这意味着:

• 给短序列发 draft → 提议被频繁拒绝 + 增加 verification token → 反而变慢;
• 给长序列发 draft → 它本来就要决定 makespan,每节省一个 forward pass 都直接砍 wall-clock。

所以"per-request 平均分 budget"是错的。budget 应该向 long-tail 倾斜。

5.2 SVG: budget allocation 概念图

5.3 Length class policy (运行时)

因为 generation 长度在线不可精确预测(同一 prompt 不同 epoch 长度 spread 巨大,见 Figure 9),论文用一个 hierarchical heuristic:

1. Init from history:对每个 prompt,从历史 trajectory 长度分布算一个 {Long, Medium, Short} 多数类作为初始预测。
2. Online update:解码到第 l 个 token 时,根据 P(class | l, init) 把当前 request 重新分类。比如初始预测 Medium,但已经生成了 8k 还没结束,就上调为 Long,加大 draft budget。
3. Class → budget 表是 offline tune 出来的(典型: Long=32, Medium=8, Short=0)。

6 · 公式推导:最优 draft budget

6.1 Per-pass latency 模型

profiling 显示 (DeepSeek-R1-Distill-7B, Figure 8):

c_base 是 per-pass 固定开销 (kernel launch、DRAM→SRAM 搬权重);c_tok 是每 token 边际成本。Mean relative error ≈ 12%。

6.2 接受 token 的饱和模型

给 request i 提议 p_i 个 draft token,接受数:

k_i ∈ (0,1] 是 capacity factor (drafter 与 target 的最大 alignment);α_i 是接受效率。p_i → ∞ 时 A_i → k_il_i(永远不会 100% 因为 drafter ≠ target)。

6.3 总目标

剩余需 forward 的 token 数: l_i(1 − k_i + k_ie^{−α_ip_i/l_i})。Batch 的 forward pass 数由最长那条决定:

总延迟:

6.4 闭式解

把约束 l_i(1 − k_i + k_ie^{−α_ip_i/l_i}) ≤ N_fwd 在最优处取等:

6.5 物理直觉

三个观察(直接对应论文 §4.2.2):

1. p_i^* ∝ l_i:越长的 request 应分越大 budget(因为它本就是 makespan 的瓶颈)。
2. l_i ≤ N_fwd ⇒ p_i^* = 0:反正它早完成,不必 speculate。
3. c_base ≫ c_tok:per-pass 主导(典型小 batch 场景)→ 应激进减 N_fwd,即提议更多 token;反之大 batch 时 c_tok 主导,应保守。

6.6 数值敏感度小表

取 c_base=20ms, c_tok=0.5ms, k=0.6, α=0.4:

7 · Worked example: 一个 prompt 的完整命运

设 RL 训练第 t 步,problem P3 = "证明欧拉公式":

1. 历史信号:过去 16 epoch 它的 generation 长度均值 11k、最大 18k → init class = Long → draft budget per round = 32。
2. Drafter 状态:P3 的 suffix tree 已吸收最近 16 epoch × 16 sample = 256 条 trajectory,共 ~2.8M token,树节点 ~600k。
3. Decode 启动:π_t 开始生成。每步当前 prefix 末 16 token 进 P3 tree:longest-match → 沿最高频路径取 32 token draft → target 一次 batched verify。
4. 典型 round:32 提议被接受 18 个 (k≈0.6),N_fwd(此 request) 从 11000 降到 ~6500,直接砍掉 40% forward pass。
5. Online update:本 round 真正生成的 18 个新 token incrementally insert 到 P3 tree (Ukkonen O(18))。
6. 跨步:已生成 9k 还没 EOS → online length classifier 把 class 维持 Long;若初始误判为 Medium,这里会 escalate。
7. 本步结束:P3 trajectory 总长 13k,贡献最大 makespan。Tree 把 epoch t-16 的某条 trajectory 移出窗口、把 epoch t 的 16 条 sample 全部加进。
8. 下一步:π_t+1 ← π_t(GRPO 更新),但 P3 tree 几乎不变(因为 policy 漂移是渐变的,recent window 仍 representative)→ 接受率保持。

8 · 实验关键结果

8.1 Math RL (DeepSeek-R1-Distill-Qwen-7B, DSR-sub)

8.2 Code RL (Qwen3-8B, DeepCoder)

Math 收益 (-50%) 比 code (-25%) 大,因为 math reasoning 的 trajectory 更长 + 模板化 (\boxed, Let me verify, Therefore...) 更多,suffix tree 命中率高。

8.3 关键 ablation

• 无 length-aware budget (Unlimited):把 draft budget 拉到无限大,反而比 DAS 慢 15%——verification cost 把收益吃掉。论文这条直接证实了 length-aware 的价值。
• Sequence length 8k vs 16k:8k 仍然 >30% 加速 → 不依赖超长 trajectory。
• Batch 32 vs 16:fractional speedup 类似 → 不依赖特定 batching regime。

8.4 与 EAGLE 的接受率曲线

论文 Figure 4: EAGLE 静态 drafter 的 avg accepted tokens 在训练全过程基本持平(因为它没法跟随 policy 漂移);DAS 的 nonparametric drafter 随训练步骤上升——因为 sliding window 里的 trajectory 越来越像当前 policy。这是"online drafter"的核心价值。

9 · 与同类工作对比

9.1 DAS 的"工程优雅"在哪

1. Drafter 没参数。不用第二个 GPU model,显存占用 ~MB 级 CPU 内存。
2. 不改 reward loop。VeRL/OpenRLHF 不用动,只在 vLLM rollout call 上挂一个 hook。
3. Lossless。verification 是严格 rejection sampling,输出分布与无 SD 的 baseline 完全一致——SPEC-RL 的 lenience 路线就放弃了这条。
4. O(m) 查询 + O(Δn) 更新。所有数据结构开销在 CPU,target 永远是 GPU bottleneck。
5. 可与 EAGLE 正交叠加。理论上 DAS 的 budget 模型可以套在 neural drafter 之上(只是论文没做这组实验)。

个人 take:DAS 是 spec-decoding-for-RL 这条线里最容易在已有生产栈上"今晚就上"的方案。NVIDIA 那篇要训 EAGLE-3 + 还得做 capture/replay;FastGRPO 还要管 draft model 显存;DAS 只要一个 suffix tree 库 + 一个 budget 调度器。

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

10.1 局限

• 第一个 epoch 冷启动:tree 是空的,接受率近 0。论文未细谈 warm-up 策略——可能直接用 prompt 自身做 PLD-style fallback。
• 问题数极大时 sharding 开销:10k+ problems 时 per-problem 树的内存 + 路由代价上升。论文最大实验 ~1209 个 prompt。
• 非模板化 domain:open-ended dialogue / 创作类 RL 的 trajectory 没有 \boxed 这类 motif,suffix tree 命中率会显著低,DAS 收益预计 < 10%。
• Length classifier 是 offline tune 的:新 domain 上需要重新校准 P(class|l, init) 表。
• 对 batch 极小或极大 regime 的覆盖不足:大 batch (≥256) 下 c_tok 主导,文中公式预测的 budget 会变小,但实测未给出。
• 仅在 ≤8B 规模实测:14B 提到但 70B+ 未验证。policy 漂移更慢的大模型可能让更长的 sliding window 也仍然 representative,反而对 DAS 利好。

10.2 个人 take

• 这篇论文的 ML 贡献其实有限,真正的贡献是问题工程化的清晰度:三个 insight + 三个 component 一一对应。
• "per-problem suffix tree + sliding window"这个组合是 SuffixDecoding 在 RL 场景的自然延伸,但把它跟 length-aware budget 耦合是新的、且实验上证明必要(unlimited budget 反而慢 15%)。
• 与 ReSpec / NeMo-RL 比,DAS 没有 EAGLE-3 那种"neural drafter 跟着 target 一起进化"的优雅,但它换来的是部署上的极简——这种 trade-off 在生产 RL 里非常值。

10.3 待验证问题

1. cold-start 阶段(第 1–2 epoch)实测加速率是多少?有没有 fallback 到 PLD?
2. 把 sliding window 大小自适应到 KL(π_t‖π_t-N) 上能不能进一步提收益?
3. 多模态 / agentic RL(tool 调用穿插)里 trajectory 不连续,suffix tree 怎么 segment?
4. 把 DAS budget 模型套在 EAGLE-3 drafter 上,长 trajectory 上能不能拿到比 DAS+suffix 更高的 k?
5. k_i 和 α_i 在线如何估计?Paper 说"profiling"但没给在线估计的算法。
6. Off-policy / async RL 下,policy lag 大,sliding window 还成立吗?