SpecOPD:面向投机解码的分支并行在线策略蒸馏 SpecOPD: Branch-Parallel On-Policy Distillation for Speculative Decoding
用提案树并行化在线蒸馏的采样与打分,草稿验收长度最高提升41.45%。
前置知识
投机解码(Speculative Decoding, SD)
一种无损加速自回归解码的技术:轻量草稿模型一次提出多个候选 token,目标模型并行验证,被接受的连续前缀一次推进多步;若某位置被拒绝则回退并由目标模型重新采样该位置。由于验证可以批量进行,其加速比取决于“每轮验证平均接受多少 token”(验收长度 AL)与草稿成本之比。
本文的训练目标就是提高草稿模型的验收长度,全部设计都围绕 SD 的 draft-and-verify 循环展开,不理解这个循环就无法理解训练信号的来源。
并行草稿器(Parallel Drafter:DFlash / DSpark / DFlash V2)
区别于逐 token 自回归起草的传统草稿模型,并行草稿器把整块候选当作掩码 token 在一次前向中预测(块扩散式生成),显著降低提案成本。DFlash 将块扩散引入起草,DSpark 结合并行主干与轻量顺序建模及置信度调度验证,DFlash V2 是其升级版本。
SpecOPD 的监督改造正是建立在并行草稿训练的“锚点+掩码块”形式之上,论文训练并评测的三种草稿器全部属于这一类。
在线策略蒸馏(On-Policy Distillation, OPD)
蒸馏时学生模型在“自己采样出的序列”上接受教师(目标模型)的稠密逐 token 监督,而不是在教师生成的固定数据上做 SFT。由于训练分布与学生当前策略的分布一致,能缓解 off-policy 数据造成的分布偏移,从而突破 SFT 的验收长度平台期。
本文的核心问题是:SD 场景下的 OPD 通常要靠顺序的 draft-and-verify 采样来收集“自己采样的序列”,作者证明这一顺序依赖并非在线监督的内在属性,并将其完全并行化。
验收长度与 E2E-TV 目标
验收长度(AL)指每轮验证平均推进的 token 数(含 bonus token)。在给定历史 $h$ 下,精确拒绝采样的单 token 接受概率等于草稿与目标分布的重叠:$\alpha=\sum_v\min\{q_\theta(v|h),p(v|h)\}=1-\mathrm{TV}(q_\theta,p)$。E2E-TV(来自 Bebop)用前缀乘积 $\prod_{i\le j}\alpha_i$ 聚合,直接把验收长度期望作为可微代理损失。
SpecOPD 直接采用该代理做训练目标;理解它才能看懂实验中“训练损失下降但验收长度反而下降”这一关键现象及其消融。
树注意力(Tree Attention)
由 SpecInfer 引入、Medusa 采用的注意力掩码机制:把多个候选路径组成一棵树后,按父映射定义每个节点的严格祖先集合,掩码 $M_{u,v}$ 只允许节点看见祖先与自身,从而在单次批处理中同时打分树上全部节点,既共享公共前缀的 KV 计算,又不泄露兄弟分支或未来 token。
SpecOPD 用它让打包后的分支查询共享主干 KV,是实现全注意力目标上分支并行打分的核心机制;同一“前缀共享、分叉后隔离”原则还被推广到稀疏与线性注意力。
三类注意力机制:GQA、稀疏注意力(QSA)、线性注意力(GDN)
GQA 保留逐 token 的 KV 缓存,可用掩码任意控制可见性;QSA 等稀疏注意力用压缩索引器键做 top-k 选块再计算注意力;GDN 等线性门控 delta 网络把历史压缩进循环状态 $S_t$,逐 token 更新 $S_{b,j}=F(S_{b,j-1},z_{b,j})$,无法靠掩码恢复任意历史位置的状态。
论文的三个实验目标分别覆盖全注意力(GQA)、全-线性混合(GQA-GDN)和稀疏-线性混合(QSA-GDN),SpecOPD 必须为每类注意力设计不同的分支历史保持方案,这是方法的主要工程难点。
研究动机
投机解码通过目标模型并行验证草稿候选来加速生成,DFlash 等并行草稿器已把提案成本降到一次前向,但整体加速上限取决于验收长度(AL)。常见做法是用目标模型生成的固定回答做 SFT(序列级知识蒸馏),容易到达平台期;在线策略蒸馏(OPD)把训练转移到当前草稿策略采样的块上,能突破 SFT 平台,但 SD 定制的 OPD 需要“顺序的 draft-and-verify”采样:草稿提出一块、目标验证、被接受的前缀决定下一轮起点,整条响应轨迹的收集被完全串行化,训练只能等收集结束后再更新。作者在 8 卡上对 Qwen3.8-27B + DFlash V2 的实测执行画像(Figure 1a 上半)显示,Draft-OPD 在顺序 rollout(占 20.3%)、replay+训练(32.9%)与权重更新(46.3%)之间反复交替,rollout 阶段 GPU 利用率极低,平均记录 GPU 活动仅 26.37%。由此作者提出关键问题:在线策略训练必须继承 draft-and-verify 采样的顺序依赖吗?
本文的目标是本文目标是构建一个不依赖顺序全响应采样的 SD 定制 OPD 框架,让已训练好的并行草稿器突破验收长度平台期,同时保持稠密的 token 级目标监督。具体包括:将草稿采样与目标打分在从当前草稿策略于参考锚点出发的独立分支上并行化;把分支打包成提案树以共享目标模型对公共前缀的计算;设计能同时覆盖全注意力(GQA)、稀疏注意力(QSA)与线性注意力(GDN)三类异构目标的执行器;并采用直接以验收长度为目标的 E2E-TV 代理损失。最终在 8B/27B/125B 三个目标、DFlash/DSpark/DFlash V2 三种草稿架构和对话/代码/数学/智能体四类任务上验证,追求更高的平均 AL(最高 +41.45%)、更高的监督 token 吞吐(16.09×)和更高的实测生成加速(最高 3.56×)。
与已有工作不同的是,独特切入点是一个简单却未被前人利用的观察:对某个草稿块做在线监督,只需要它的参考前缀和自身已采样的历史,并不需要一条完整被验证的轨迹。传统并行草稿训练用参考 token 标签(off-policy 的 CE 监督),SpecOPD 用目标模型在草稿采样前缀上的分布替代这些标签,这是对标准并行草稿训练的一次“独特改造”。由于每个分支只依赖参考上下文加自身样本,不同锚点的分支彼此独立,联合分布 $Q_\theta(Y|x,A)=\prod_b\prod_i q_\theta(y_{b,i}|h_{b,i})$ 完全因子化,采样与打分无需任何跨分支顺序;同时分支共享参考主干,目标模型对公共前缀只需计算一次。这把“在线策略”(要求训练数据来自当前策略)与“并行计算”(要求消除轨迹内依赖)这对看似矛盾的性质统一起来:分支内自相关、分支间独立。
核心方法
直觉是把“草稿自己采样、目标在草稿历史上打分”组织成一棵树:参考序列是共享主干,每个锚点长出一条草稿分支;主干只算一次、分支间互不依赖,可并行推进。技术路线三步。第一步提案树构建:在参考回答上随机采样 $B$ 个锚点,草稿器在每个锚点并行采样 $W$ 个 token 的延续,条件历史为 $h_{b,i}=(x_{\le a_b},y_{b,<i})$,分支内自回归、分支间完全独立。第二步分支并行打分:目标模型只对主干做一次 prefill 并缓存表示,再把所有分支查询行打包批处理,按父节点读出 $p_{b,i}=\mathrm{softmax}(g_p(H_{\pi(y_{b,i})}))$;按注意力类型分别用树掩码(GQA)、分支一致索引器组(QSA)与索引状态回放(GDN)保证各分支历史正确。第三步用 E2E-TV 前缀乘积损失 $\mathcal{L}$ 直接优化验收长度,每步重采样分支。27B/DFlash V2 实测 GPU 活动从 26.37% 升至 77.23%,串行 rollout 变为 5.6% 并行起草加 60.0% 并行打分。
核心创新是“解耦”:把在线策略监督从完整响应 rollout 中解耦出来。Draft-OPD 等 SD 定制 OPD 中,每个新训练上下文依赖更早的验证结果(接受到哪里就从哪里继续),收集必然串行;SpecOPD 观察到草稿块的监督只需参考前缀加自身采样历史,于是用锚点把依赖切断:分支在参考主干上分叉后只消费自己的样本,分支间零依赖。与 off-policy SFT/CE 的本质区别在于监督发生在草稿自己生成的历史上——模型学到的是“在我的分布下目标会怎么判”,而非“在参考分布下贴合目标”,直接对准验收长度而非逐 token 交叉熵。此外它对三类异构注意力给出统一原则“前缀共享、分叉后各算各的”:GQA 用树掩码共享 KV;QSA 共享完整 4-token 索引器组、只重建跨锚点边界组;GDN 掩码无法隔离循环状态,改为从 FP32 检查点回放锚点状态再并行推进各分支。配合 E2E-TV 前缀乘积目标,构成首个免串行采样收集的 SD 定制 OPD 框架。
方法步骤详情
(1) 锚点采样:在含提示的参考序列 $x$ 的回答部分随机取 $B$ 个锚点,构造“锚 token+掩码”块。(2) 草稿分支采样:草稿器一次前向并行采样每分支 $W$ 个 token,条件历史为参考前缀加自身样本 $h_{b,i}$。(3) 树组装:按父映射 $\pi$(父为锚点或前序 token)与语义位置把分支挂上主干。(4) 主干预填充:GQA 缓存主干 KV;QSA 缓存完整 4-token 索引器组;GDN 每 $C=64$ chunk 存不可变 FP32 状态检查点。(5) 并行打分:GQA 用树掩码 $M_{u,v}$ 只放行祖先与自身、rotary 用语义位置;QSA 复用完整索引器组、仅重建跨界组键,top-k 选择、因果尾绕过;GDN 从检查点回放恢复锚点状态,再逐步 $S_{b,j}=F(S_{b,j-1},z_{b,j})$,因果卷积随祖先。(6) E2E-TV 训练:对逐位重叠 $\alpha_{b,i}=\sum_v\min\{q_\theta,p\}$ 求前缀乘积得损失,只对草稿概率求导,采样与目标分布视为常量。(7) 更新权重,下一轮重采样分支。
技术新颖性
新颖性体现在三个层面。问题层面:证明 SD 定制 OPD 的串行性不是在线监督的内在属性,而是 rollout 组织方式的产物,把训练数据收集从轨迹级并行化为分支级,这一重述本身即贡献。系统层面:已有树注意力工作(SpecInfer/Medusa)只服务全注意力验证,本文首次为稀疏 QSA 设计分支一致索引器组(共享完整组+仅重建跨界组,保证分叉后块选择与单独执行一致),为线性 GDN 设计索引状态回放(不可变 FP32 检查点+后缀回放+并行推进,避免重复全前缀执行),并在附录 D.3-D.4 给出数值对齐与确定性细节。目标层面:采用 E2E-TV 前缀乘积作为验收长度代理并配套每步重采样训练,实验(Figure 8)证明训练损失与验收长度可以脱钩——off-policy CE 损失下降但验收低于初始化,而 SpecOPD 在全部 7 个草稿位置提升前缀存活率。相比 Draft-OPD 的目标辅助 rollout 与错误状态回放、DistillSpec 的解码感知散度选择,SpecOPD 无需构造完整验证轨迹即可获得新鲜监督。
实验结果
验收长度(Table 1):DFlash(8B,$T=0$)AL 2.649→3.747(+41.45%),Terminal 1.620→2.940(+81.51%);DSpark 3.674→4.439。混合注意力收益温和:27B/125B($T=1$)+3.48%/+5.01%(3.621→3.802)。吞吐(Figure 6):27B 单并发 MT-Bench 达 3.56× AR,几何平均在并发 1/8/32 为 2.97×/2.23×/1.35×,SWE 在并发 32 慢于原版。训练效率(Figure 7):更新速率 0.077→0.381/s(4.95×)、22.94k→74.58k token/更新(3.25×),监督吞吐 1.77k→28.42k token/s(16.09×)。消融(Figure 8):同驻点下比 off-policy E2E-TV 高至多 15.71%、比 CE 高 6.96%;两基线损失降但验收低于初始化,SpecOPD 在 7 个草稿位置全部提升存活率,第 7 位 +4.33pp(CE -2.57pp,824 请求 bootstrap)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 8B DFlash 五任务平均(T=0) | 平均验收长度 Mean AL | 3.747 | 2.649(原 DFlash) | +41.45% |
| 8B DFlash Terminal 智能体(T=0) | 平均验收长度 Mean AL | 2.940 | 1.620 | +81.51% |
| 8B DSpark 五任务平均(T=0) | 平均验收长度 Mean AL | 4.439 | 3.674 | +20.8% |
| 27B DFlash V2 五任务平均(T=1) | 平均验收长度 Mean AL | 4.881 | 4.717 | +3.48% |
| 125B DFlash V2 五任务平均(T=1) | 平均验收长度 Mean AL | 3.802 | 3.621 | +5.01% |
| 27B DFlash V2 端到端生成(并发 1,MT-Bench) | 输出吞吐相对 AR 加速比 | 3.56× | AR 自回归解码 | 3.56× |
| 训练流水线(Qwen3.8-27B / DFlash V2,8 GPU) | 监督 token 吞吐 | 28.42k tokens/s | Draft-OPD 1.77k tokens/s | 16.09× |
| 消融:On-policy vs Off-policy(同驻点对比) | 五基准平均 AL 相对提升 | SpecOPD | off-policy E2E-TV / off-policy CE | 最高 +15.71% / +6.96% |
局限与改进
作者承认的局限:混合注意力目标上收益明显更小(27B 约 +3.48%、125B 约 +5.01%,远低于 8B 的 41.45%);27B 智能体任务几乎无变化且 $T=1$ Terminal 略降;125B 草稿器预训练缺少智能体轨迹,存在覆盖缺口;E2E-TV 前缀乘积对条件提案只是代理而非服务端验收的无偏估计(附录 C.2);训练吞吐测量排除了离线参考准备成本;SWE 在并发 32 下慢于原版 DFlash V2。我的补充观察:其一,在线范围有限——外层上下文仍来自参考序列,分支只在锚点后分叉,草稿器从未在完全由自己生成的长上下文上被监督,分布偏移只是被缓解而非消除;其二,锚点采样与每步重采样带来额外草稿前向开销,锚点数 $B$、分支长 $W$ 均为固定/随机超参,缺乏自适应机制;其三,SWE/Terminal 评测基于录制交互历史的续写效率,不能完全代表真实智能体在线 rollout 的加速;其四,与 off-policy 基线使用不同计算分配与梯度路径(作者亦承认),未给出严格等算力对比。
独立分析的弱点
(1) 参考依赖:分支锚点必须在参考回答上,管线要求目标离线重生成或数据自带回答,对缺乏高质量参考的新领域是硬约束;改进方向是周期性用上一版草稿策略自生成的整段轨迹替换参考主干,形成自举主干。(2) 混合注意力收益小:GQA-GDN 与 QSA-GDN 上仅 3-5%,可能是状态回放与索引器组重建的额外开销摊薄了并行收益;可尝试更大锚点批量、更长分支 $W$ 或异步化回放计算。(3) 块结构固定:$W$、锚点数量与位置均为随机/固定超参,未按草稿不确定度分配预算;可结合置信度调度(如 DSpark 的验证调度)做自适应锚点放置与分支长度分配。(4) 目标偏差:前缀乘积假设各位置独立接受,对条件提案有偏,极端情况下可能过估长前缀收益;可引入重要性加权或改用 Expected Decoding Rounds 类目标校正。(5) 高并发收益递减:并发 32 几何平均仅 1.35× 且 SWE 为负收益,说明验证成本随批规模上升后草稿质量增益被稀释,需研究大批并发下训练目标与验证开销的联合权衡。每个弱点都指向一个具体可操作的改进路径,而非泛泛而谈。
未来方向
作者明确提出:为 125B 草稿器补充智能体轨迹预训练以填补覆盖缺口(8B 上 +81.51% 的 Terminal 增益暗示该方向空间很大);附录 F 还讨论了与 RL rollout 及多模态生成的联系——分支并行采样加共享前缀打分的执行器天然适配 RL 的并行经验收集,也可迁移到扩散式多模态草稿训练。基于成果可延伸的方向:把索引状态回放推广到其他循环架构(Mamba/RWKV 变体)与 MoE 目标;将 E2E-TV 与 Expected Decoding Rounds 等验收目标组合或加权,修正条件提案下的代理偏差;在 FP8/量化目标上进一步验证(附录 A.5 已有 Flash-Next FP8 的初步正结果);研究锚点位置的自适应放置与课程式分支长度增长;以及利用 16.09× 的训练吞吐红利做每日级草稿持续更新,让草稿器紧贴目标模型的 RL/后训练迭代,形成推理-训练数据飞轮。
复现评估
论文首页标注了 Code 与 Blog 链接,但正文未确认代码与权重已公开,复现前需先核实。可复现要点:目标为公开的 Qwen3-8B、Qwen3.8-27B(Qwen3.8-Flash-Next-125B 为内部模型),125B 草稿器为内部预训练初始化,社区拿不到同等起点;训练数据 Open-PerfectBlend,规模 8 张高性能 GPU(27B 目标);MT-Bench/HumanEval/MATH-500 公开可测,SWE/Terminal 用录制交互历史(附录 A.1,公开程度未知);服务评测基于 SGLang(BF16、TP=1、FA3、CUDA Graphs)。实现难点在自定义 kernel:GDN 索引循环核(FP32 检查点+回放)、QSA 索引器组重建与批量 top-k、树掩码打包注意力,以及数值对齐与确定性(附录 D.3-D.4)。总体判断:对具备大模型训练基础设施的团队属中高难度,8 卡级算力起步;若 kernel 与 125B 初始化不开源,第三方只能较完整复现 8B 全注意力(GQA)部分。
论文图表
(a) 上下两半对比 Draft-OPD 与 SpecOPD 的流水线时间占比:前者在顺序 rollout(20.3%)、replay+训练(32.9%)与权重更新(46.3%)间交替,GPU 记录活动仅 26.37%;后者拆为输入准备(16.6%)、并行起草(5.6%)、并行打分(60.0%)与训练(17.6%),活动升至 77.23%。(b) 给出 4.95× 更新速率、3.25× 每更新监督 token、合计 16.09× 监督吞吐(1.77k→28.42k tokens/s)。
这是全文动机与训练效率结论的总览图:上半部分直观展示串行 rollout 造成 GPU 空转的问题,下半部分量化并行化后的收益,是理解为什么需要 SpecOPD 以及它带来多少训练加速的第一张图。