用块扩散草稿树加速投机解码 Accelerating Speculative Decoding with Block Diffusion Draft Trees
DDTree:从单次块扩散草稿的逐位分布构建最优草稿树,提升投机解码接受长度。
前置知识
投机解码(Speculative Decoding)
用一个小草稿模型先快速生成若干候选 token,再让大目标模型一次前向并行验证这些 token,接受与目标分布一致的部分,从而在不改变输出分布的前提下减少解码步数。基础论文 Leviathan et al. 和 Chen et al. 证明了其输出严格等价于自回归采样。
DDTree 是对投机解码流程的优化,需要先理解草稿-验证两阶段范式和接受长度这一核心评价指标,才能体会为什么提升草稿质量与多样性如此重要。
块扩散模型(Block Diffusion)
介于自回归与全双向扩散之间的生成范式:输入形如 [b, m, ..., m] 的部分掩码序列,模型在一次前向中并行预测接下来 L 个位置的 token,对每个位置输出边际分布 q_i(v),位置之间不互相条件化。Block Diffusion 提供了支持 KV cache 的去噪框架。
DDTree 的整篇方法都建立在单次块扩散前向产出 L 个边际分布这一特性上,边际之间的独立性是与自回归草稿器的本质区别,也是树构造目标的来源。
DFlash 块扩散草稿器
由 z-lab 提出的轻量级块扩散草稿器,条件于目标模型的多层特征(类似 EAGLE-3 的思路),能在一个前向中产出整个 draft block,性能优于 EAGLE-3。vanilla DFlash 每轮只验证单一贪心轨迹。
DDTree 是 DFlash 的直接扩展,论文中所有实验都基于 DFlash 检查点;理解 DFlash 已知能力强但每轮只走单条轨迹这一瓶颈,是把握 DDTree 改进动机的关键。
树注意力(Tree Attention)
由 SpecInfer (Miao et al.) 提出的高效验证机制:把多棵候选续写拼成一个前缀闭合的 token 树,序列化为张量喂给目标模型,通过 ancestor-only attention mask 让每个节点只能看到根、祖先和自己,从而一次前向评估所有分支的 logit。
DDTree 把 DFlash 的单轨迹替换为草稿树后,正是依赖树注意力来保持一次目标模型前向的验证成本不变;若没有树注意力,扩展到多分支会使验证成本线性膨胀,收益消失。
研究动机
自回归语言模型逐 token 解码,每生成一个 token 都要一次大模型前向,序列依赖使解码成为推理延迟的主要来源。投机解码用轻量草稿器先提议、目标模型并行验证,大幅降低步数,但端到端加速比取决于草稿质量。DFlash 用一个小型块扩散草稿器配合目标模型特征,实现一次前向产出整块草稿,在 Qwen3 系列上把速度推到 $5\times$–$7\times$ 自回归基线。然而 vanilla DFlash 每轮只验证单一轨迹(block 大小 16),即使单次块扩散前向天然地在每个未来位置都给出了完整的 token 分布 $\{q_i\}_{i=1}^L$,这些潜在的好候选被全部丢弃。从 Table 1 看,DFlash 在 AIME 2024 Qwen3-4B 上平均接受长度 $\tau=7.54$,远未达到块长度 16;在 SWE-bench Lite、MT-Bench 等长尾任务上 $\tau$ 更是只有 3–4。直觉上每轮多探索几条高概率候选会提高命中目标模型路径的概率,但朴素地多跑几条轨迹会线性增加验证成本,把延迟收益吃掉。
本文的目标是本文提出 DDTree(Diffusion Draft Tree),目标是在固定节点预算 $B$ 下,仅用一次块扩散草稿前向的逐位边际分布,自动挑选出 $B$ 个最有希望被目标模型接受的 token 前缀构成前缀闭合的草稿树,然后通过树注意力在一次目标模型前向里完成验证。具体而言:论文要在理论上证明所选的树在 drafter 的因子化代理目标下最优,并给出一个 $O(B\log B)$ 的最佳优先堆算法来高效构造它。终极目标是:在所有 3 个目标模型 × 10 个基准 × 2 个温度共 60 个评测点上,一致地超过 vanilla DFlash,典型速度提升再额外 +20%–+30%,并在温度 1.0 这种随机性更强的解码设置下同样取得稳定收益。
与已有工作不同的是，现有同类工作(Medusa、EAGLE-2/3、OPT-Tree、近期 DART)在两个层面与 DDTree 不同。第一,OPT-Tree、EAGLE-2 这类基于自回归草稿器的方法,每深入一层就需要一次额外的草稿前向,树越深成本越高;而 DDTree 利用块扩散单次前向给出 L 个边际,树的深度对草稿成本是零边际。第二,DART 同样基于单次并行 logits 构树,但依赖外部 N-gram trie 和 continuity-aware pruning 才能工作;DDTree 不需要任何外部统计,直接把 drafter 的因子化分布 $Q(y_{1:L}\mid c,b)=\prod_i q_i(y_i\mid c,b)$ 作为代理目标,从理论上给出最大化期望接受长度的最优树并通过最佳优先堆实现。换言之,DDTree 首次在块扩散设定下严格连接了单次前向信息与期望接受长度两个概念,既给出最优性证明,又给出可实现的 $O(B\log B)$ 算法,在简洁性、理论保证和工程友好性三个维度上同时形成差异。
核心方法
DDTree 沿用每轮一次块扩散草稿加一次目标模型树验证的两阶段框架,核心改动是:不再把 drafter 给出的 L 个逐位边际 {q_i} 折叠成单一轨迹,而是用它们在节点预算 B 内构造一棵前缀闭合的草稿树 T。直觉上,目标模型每轮接受多少 token 取决于它实际选出的路径有多深地落在草稿树里,因此应把预算投到 drafter 眼中概率最高的前缀上。技术上,论文先证明在 drafter 因子化代理分布 Q 下,期望接受长度 E_{Y~Q}[α_T(Y)] = Σ_{u∈T} q(u|c,b),其中 q(u) = ∏_i q_i(u_i);再证明按 q 降序取前 B 个前缀即得最优树 T_B;最后用最佳优先堆只展开 top-K(K = min(B,|V|))上的兄弟与第一孩子,按真实降序枚举这些前缀,时间复杂度 O(B log B)。得到 T_B 后用 tree attention 的 ancestor-only mask 一次前向给树上所有节点打分,再按目标模型解码规则走树,接受最长匹配前缀,把首个不匹配 token 作为下一轮 bonus。
DDTree 与已有树构造方法(EAGLE-2 动态树、DART 连续性剪枝树)的本质区别在于代理目标和信息源的统一:它把块扩散草稿的边际分布独立性这一内在结构直接转化为代理目标 $\sum_{u\in T}\prod_i q_i(u_i\mid c,b)$,而这个加性结构恰好让取概率最高的 B 个前缀成为严格最优,不需要任何外部统计或经验阈值,也不需要为树的不同深度付出额外的 drafter 前向。算法层面,用 log-prob 累加 $\sigma(\rho)=\sum_i \log q^{(\rho_i)}_i$ 把乘积变求和、放进 max-heap,既保证数值稳定又让堆上只出现兄弟与第一孩子两类扩展,实现既最优又极简。
方法步骤详情
DDTree 单轮解码包含四步。第一步(Drafter pass):取上下文 c 与上一轮 bonus b,把 [b, m, ..., m] 送入 DFlash,L 个位置的 logits,softmax 后得边际分布 q_i(v)。第二步(Tree construction,算法 1):取 K = min(B,|V|),每深度按 q_i 降序取 top-K token,初始化 max-heap 装入 ((1), σ((1))),循环 |T|<B 时弹出堆顶 ρ 加入草稿树,再压入兄弟与第一孩子,直到取出 B 节点。第三步(Verify):把草稿树在 bonus b 之后展平成 token id 序列,按深度分配 position id,构造 ancestor-only mask,做一次目标模型前向。第四步(Walk):按目标模型解码规则从根向下选 token,首处不匹配 token 成为下一轮 bonus,接受路径拼入输出,KV cache 裁剪。整轮 drafter 一次前向、target 一次前向,与 vanilla DFlash 算力账本相同。
技术新颖性
DDTree 技术新颖性可总结为三点。其一,代理目标的封闭形式:在块扩散设定下首次把 E[α_T] 写成 Σ_{u∈T} q(u) 的加性形式(Proposition 1),把树选择问题简化为排序问题,绕开 DART 依赖外部 N-gram 分数的负担。其二,最优性的严格证明:Proposition 2 证明 top-B 前缀构成的最优树在代理目标下严格最优;Lemma 1 把搜索空间压缩到 S_K 内;Proposition 3 证明 best-first 堆算法 1 恰好按真实降序枚举这些前缀,总成本 O(B log B)。其三,与 OPT-Tree 的关键区别:OPT-Tree 在自回归草稿器下每深一层就需要一次 drafter 前向;DDTree 借助块扩散的天然并行性,树的深度对草稿成本是零边际。作者在 Remark 1 中坦承代理目标仅在 drafter 因子化分布 Q 下最优,非真实目标分布 p,但只要 drafter 与 target 在每个位置上排序一致,top-B 前缀同时也是 target 视角下的好选择,这是经验上 60 个评测点都胜出的根据。
实验结果
论文在 Qwen3-4B、Qwen3-8B、Qwen3-Coder-30B-A3B-Instruct 三个目标模型上,跨 10 个数据集、两个温度(0.0 与 1.0)做端到端评测,DDTree 在 60 个组合上全部超过 vanilla DFlash。Figure 1 显示 T=0 下典型加速比:Qwen3-8B AIME 2024 从 DFlash 5.38× 提升到 7.35×,MATH-500 从 5.56× 到 7.52×,GSM8K 从 4.78× 到 6.75×。Table 1 显示加速比提升主要来自接受长度 τ 的同步增长,如 Qwen3-4B AIME T=0 τ 从 7.54 升到 10.37。温度 1.0 下 DDTree 同样稳定,Qwen3-4B AIME 2024 从 3.50× 升到 5.31×。Figure 3 揭示 trade-off:MATH-500 Qwen3-8B T=0 上,端到端速度在 B=256-512 达峰值,B=1024 验证成本反噬。Figure 4 印证:DDTree(B=512) 质量从 1-3 迁到 14-16。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME 2024 (Qwen3-4B, T=0) | 相对自回归速度 + 平均接受长度 | DFlash+DDTree: 7.27x 速度, tau=10.37 | DFlash: 5.56x 速度, tau=7.54 | 加速比 +1.71x (相对 +30.8%), 接受长度 +2.83 |
| MATH-500 (Qwen3-8B, T=0) | 相对自回归速度 + 平均接受长度 | DFlash+DDTree: 7.52x 速度, tau=10.73 | DFlash: 5.56x 速度, tau=7.79 | 加速比 +1.96x (相对 +35.3%), 接受长度 +2.94 |
| HumanEval (Qwen3-Coder-30B-A3B, T=0) | 相对自回归速度 + 平均接受长度 | DFlash+DDTree: 8.22x 速度, tau=10.72 | DFlash: 6.09x 速度, tau=8.02 | 加速比 +2.13x (相对 +35.0%), 接受长度 +2.70 |
| GSM8K (Qwen3-4B, T=1) | 相对自回归速度 + 平均接受长度 | DFlash+DDTree: 6.17x 速度, tau=8.87 | DFlash: 4.35x 速度, tau=5.99 | 加速比 +1.82x (相对 +41.8%), 接受长度 +2.88 |
| SWE-bench Lite (Qwen3-Coder-30B-A3B, T=0) | 相对自回归速度 + 平均接受长度 | DFlash+DDTree: 4.38x 速度, tau=5.71 | DFlash: 2.77x 速度, tau=3.61 | 加速比 +1.61x (相对 +58.1%), 接受长度 +2.10 |
| Alpaca (Qwen3-8B, T=0) | 相对自回归速度 + 平均接受长度 | DFlash+DDTree: 3.36x 速度, tau=5.09 | DFlash: 2.07x 速度, tau=3.12 | 加速比 +1.29x (相对 +62.3%), 接受长度 +1.97 |
局限与改进
作者在文中明确指出了几项局限。其一,代理目标基于 drafter 的因子化分布 $Q$,而非真实目标分布 $p$,因此 top-B 前缀最优仅在 drafter 视角下成立,当 drafter 与 target 排序不一致时会出现理论次优,可能在极冷僻词、长尾代码标识符等场景放大错误。其二,Figure 3 显示最优预算落在 256–512,但这个 sweet spot 强烈依赖硬件:在 8 卡 H200 上如此,换到 A100 或消费级 GPU 上需要重新标定(作者明确提醒 optimal budget can shift across hardware platforms)。其三,论文没有在更长上下文(>32K)上做实验,树注意力在长 prompt 上的 mask 开销、KV cache 压缩的 I/O 抖动,文中未给出数据。其四,DDTree 复用 DFlash 的 drafter,继承其对目标模型的依赖——换模型需要重新训练 DFlash 检查点;且当 drafter 边际分布熵极低时,树实际宽度会塌缩到接近 vanilla 行为,加速空间被压制。
独立分析的弱点
独立审视,DDTree 至少有以下可改进点。第一,代理目标只用到 drafter 的边际,没有利用 target 自身前几层隐藏信号;可让树构造阶段轻量地运行一个 target 浅层前向(只跑最后 1–2 层),用 target 视角的排序校正 drafter 的 top-B 选择,理论上让 $\alpha_T$ 在真实分布 $p$ 下也更优,代价是每轮固定算力增加。第二,树完全由概率最高的 B 个前缀构成,缺乏多样性约束,容易出现前 10 个前缀共享同一前 2 个 token 的退化形态;借鉴 Diverse Beam Search 或 Gumbel-top-k 引入多样性惩罚,可能让同样 B 预算覆盖更广潜在轨迹。第三,Figure 3 显示 B>512 后加速比回落,但当前算法是预算无关的,没有给出在线自适应选择 B 的策略。第四,论文用 8 卡 H200 做基准,token-level 加速比能否平移到 batch=1 的在线服务尚缺端到端延迟数据;代码基于 HF Transformers,对 vLLM、SGLang 等生产级引擎的兼容性文中未涉及。
未来方向
作者在结尾段给出了若干直接延伸方向。其一,把 DDTree 的单次前向边际到最优树框架推广到任何并行一次产出多位置分布的 drafter(Medusa 头、并行自回归适配器),观察代理目标是否依然保持加性结构;若成立,可在不重训 drafter 的情况下套用到多个现有系统。其二,联合训练 drafter 加树构造器,让 drafter 在训练时就知道边际会被切分成节点预算 B 的树,优化指标换成 $\mathbb{E}_{Y\sim p}[\alpha_{T_B}(Y)]$,预期能进一步压缩 drafter 与 target 的排序差距。其三,研究自适应预算 B 的控制器:基于上一轮接受长度、当前生成 token 的熵、KV cache 占用等信号,在线调整 B 取 16/64/256,在保证 SLA 延迟的前提下最大化吞吐。其四,实验集中在 4B–30B 区间,可向 70B+ 模型和长上下文(64K、128K)扩展,另外可探索 multi-block DDTree,即一次 drafter 前向产出 >1 个块的草稿,以匹配 EAGLE-3 类方法的远期 token 视野。
复现评估
DDTree 建立在已开源的 DFlash 检查点之上(https://huggingface.co/collections/z-lab/dflash),论文声明 Project Page 与 Code 链接,作者用 HF Transformers 实现核心代码,工程门槛对熟悉投机解码的研究者较低。算力方面,论文在 8 卡 H200 上做全量评测;若只需复现 Table 1 单点(如 Qwen3-4B AIME 2024),单卡 H100/A100-80G 即可跑通,完整 60 个配置复现成本估计在数百到上千 GPU-hour 级别。超参数方面,预算 B 需在目标硬件上重新搜索(论文网格 16, 32, 64, 128, 256, 512, 1024),需自行实现 drafter pass、树构造、tree attention 三个模块,作者给出算法伪代码与复杂度分析(Algorithm 1 + Proposition 3),实现上无明显黑盒。Benchmark 全部公开。复现难度中等,核心风险在权重较大、占 GPU 内存,以及 attention mask kernel 等工程细节。
论文图表
