跳转至

Free Draft-and-Verification: Toward Lossless Parallel Decoding for Diffusion Large Language Models

会议: NeurIPS 2026
arXiv: 2510.00294
领域: 代码智能
关键词: 扩散大语言模型, 并行解码, 投机解码, 无损加速, 块扩散

一句话总结

针对扩散大语言模型(DLLM)单步单 token 解码速度慢且传统并行解码导致质量严重下降的矛盾,FreeDave 提出一种免训练、免额外模型的自验证并行解码机制,理论证明能以最少前向步数严格还原标准解码轨迹,在数学推理与代码生成上实现最高 2.83 倍无损推理加速。

研究背景与动机

扩散大语言模型(Diffusion Large Language Models, DLLMs)凭借全序列双向自注意力机制,能够同时建模文本序列中所有位置的边际分布,在克服自回归模型的“逆转诅咒”(Reversal Curse)以及受限数据学习方面展现出巨大潜力。从原理上看,DLLM 天生具备一步预测多个 token 的并行解码能力;然而,在实际应用中,要保证顶尖的生成质量,现有模型通常必须退化为与序列长度等步数的一步一 token“静态解码”(Static Decoding)。缺乏自回归因果注意力的特性也导致其与标准 KV Cache 不兼容,推理吞吐量受限于频繁刷新缓存的高昂代价。

为了提升吞吐量,近期研究提出了基于置信度阈值的并行解码(Threshold-based Parallel Decoding),即在每步将所有预测置信度超过阈值的掩码位置同时解掩码。然而大量实证表明,这种启发式并行解码极易引入次优解码路径,在复杂推理和代码补全任务中引发严重的性能退化,甚至在较大块尺寸(如 block size 32)下因模型过度自信而崩溃为整段重复 token。自回归模型中成熟的投机解码(Speculative Decoding)虽然通过“草稿-验证”保证了输出分布一致性,但直接迁移至 DLLM 需要额外训练或引入较小的草稿模型,并引入独立的草稿生成和验证阶段,草稿与目标模型分布不对齐时反而带来负优化。

这构成了扩散大模型高效推理的核心矛盾:启发式并行解码速度快但有损严重,静态单步解码质量高但速度缓慢且不实用。本文的切入点在于重新审视 DLLM 的数学本质——预训练扩散模型本身就可以在单次批量前向中对自身历史前瞻步的预测进行自验证。核心 idea:将投机解码的“草稿-验证”机制无缝内化至扩散大模型自身的单步前向中,利用静态重掩码调度器在前一步预测分布上前瞻生成多步候选序列,并在下一步以批处理前向实现自我验证与分布复用,做到完全无需额外草稿模型且理论保证无损的最优并行解码。

方法详解

整体框架

FreeDave 旨在不改变模型参数、无需外挂草稿模型的前提下,以最少的模型前向传播调用次数严格复现静态解码序列。在每个时间步,DLLM 前向输出当前掩码序列对应的预测分布;静态重掩码调度器根据此分布一次性前瞻 \(d\) 步,采样出解除不同数量掩码的候选草稿批次;在紧接着的下一个时间步,DLLM 以 Batch 形式并行对这些候选输入进行前向验证,并直接复用被接受候选的输出分布作为下一轮迭代的基底。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["当前时间步输入序列 x_ti 与预测分布 x~0"] --> B["前瞻候选生成<br/>静态重掩码调度器根据 x~0 并行构建 d 个草稿序列"]
    B --> C["自验证批量前向<br/>DLLM 对所有草稿候选执行单次并行前向推理"]
    C --> D["单步目标匹配验证<br/>检验草稿序列与对应一步目标序列是否严格精确匹配"]
    D --> E["贪心最长前缀接受与分布复用<br/>接受前 m 步转移并直接复用第 m+1 候选分布推进时间步"]

关键设计

1. 理论最优的可行路径搜索:基于自验证的无损并行跳步

现有启发式并行解码之所以掉点,根本原因在于其强行选取的解码路径偏离了预定义时间离散化诱导的“神谕路径”(Oracle Path)。作者在离散吸收态扩散框架下,形式化定义了解码路径、可行路径与最优路径:若一段跨越多步的并行转移决策集 \(\mathcal{I}_{t_a \rightarrow t_b}\) 能够严格等价于神谕路径上对应逐步决策集的并集 \(\bigcup_j \mathcal{I}_{t_j \rightarrow t_{j+1}}\),则该路径称为可行路径(Feasible Path)。所有可行路径生成的最终序列与标准一步一 token 静态解码完全一致。

为了在未知可行空间中搜寻最优可行步长,作者设计了验证器函数 \(h(n, d)\): $\(h(n,d) = \max\left\{i \in \{1,\dots,d-1\} \;\middle|\; \mathcal{I}_{t_n \rightarrow t_{n+i+1}} = \mathcal{I}_{t_n \rightarrow t_{n+i}} \cup \mathcal{I}_{t_{n+i} \rightarrow t_{n+i+1}}\right\}\)$ 理论定理证明,只要搜索预算 \(d\) 大于最优可行路径单次跳跃的最大步数 \(\|\mathbf{p}^\star\|_\infty\)(序列总长度 \(L\) 即为其天然上界),基于自验证的贪心最长匹配策略即可精确命中全局最少前向次数的最优路径,在数学上严格保证了解码过程的无损性。

2. 零额外开销的自验证流水线:草稿生成与验证的分布复用

自回归投机解码需要专门的轻量级模型生成草稿并串行验证,带来了双重推理开销。FreeDave 揭示出 DLLM 天生具备“免费自验证”能力:在时间步 \(t_i\),模型单次前向输出全序列的边际预测概率 \(\tilde{x}_0 = f_\theta(x_{t_i}, t_i)\)。由于重掩码调度器 \(g\) 通常是不带参数的规则算法(如贪心选择置信度最高的掩码位置),系统可以在 CPU/GPU 内存中无损计算出前瞻 \(1\) 到 \(d\) 步的全部候选序列 \(\mathbf{X}_{draft} = [x_{draft, t_{i+1}}, \dots, x_{draft, t_{i+d}}]\)。

在下一步,DLLM 将这些候选打包成大小为 \(d\) 的 Batch 一次性传入网络,得到所有候选的新预测分布 \(\tilde{\mathbf{X}}_{draft, 0}\)。调度器对每个候选再推演一步得到目标序列 \(\mathbf{X}_{target}\),只要 \(x_{draft, t_{i+k+1}} = x_{target, t_{i+k}}\),即表明第 \(k\) 步草稿与逐步演化的标准路径吻合。当验证通过最长前缀 \(m\) 步时,时间步直接跳跃至 \(t_{i+m+1}\),而第 \(m+1\) 个候选在前向中计算出的预测分布被直接用作下一轮迭代的输入。因此,草稿生成与验证完全是标准静态解码前向过程中的“副产品”(Byproduct),未引入任何额外的串行步骤或外置模型调用。

3. 面向块扩散的急迫模式(Eager Mode):突破块边界约束的前瞻预测

为了加速双向注意力推理,主流开源 DLLM 广泛采用块扩散(Block Diffusion)架构,即块间自回归、块内掩码扩散,配合半自回归 KV 缓存。然而,当标准并行解码接近当前块尾部时,若剩余掩码数少于草稿步数 \(d\),算法将被迫缩减 \(d\),导致块边界处的并行度严重受限。

FreeDave 提出了专为块扩散定制的急迫模式(Eager Mode,记为 FreeDave+)。当当前块的未掩码 token 数量不足 \(d\) 时,FreeDave+ 允许重掩码调度器跨越块边界,提前向未来的后续数据块提交草稿候选与自验证请求,同时在块优先序上严格保持从左至右的因果依赖。论文在理论上证明,跨块急迫验证等价于一种“延迟缓存更新”的静态解码变体,生成后续 token 时所使用的已计算 KV 状态比传统静态解码更具信息量,不仅打破了 \(d \le B\) 的硬性限制(例如在块大小 \(B=4\) 下成功使用 \(d=8\)),还在消除块边界瓶颈的同时维持了至少不亚于标准静态解码的输出质量。

一个完整示例

假设当前序列处理到第 \(i\) 步,当前块剩余 4 个掩码位置,配置草稿步数 \(d=4\): 1. 草稿生成:DLLM 输出当前分布 \(\tilde{x}_0\)。静态调度器按置信度递减顺序依次挑选出 4 个置信度最高的掩码位置及其预测 token,分别生成解除 1、2、3、4 个掩码的候选序列 \(\mathbf{X}_{draft} = [x^{(1)}, x^{(2)}, x^{(3)}, x^{(4)}]\)。 2. 批量自验证:将此 4 个序列组成一个 Batch 传入 DLLM 进行一次并行前向,获得对应的 4 组预测分布。调度器根据各候选的新分布推演下一步单 token 解掩码结果,获得校验目标序列 \(\mathbf{X}_{target} = [y^{(1)}, y^{(2)}, y^{(3)}, y^{(4)}]\)。 3. 前缀匹配:比对发现 \(x^{(2)} = y^{(1)}\) 且 \(x^{(3)} = y^{(2)}\),但 \(x^{(4)} \neq y^{(3)}\)(即第 3 步解掩码后,因上下文变化导致下一步最高置信度 token 与先前静态推测不符)。 4. 状态跃迁:成功匹配前缀长度 \(m=2\),算法一次性接受 3 个 token 的解掩码更新,时间步指针直接从 \(i\) 跳转到 \(i+3\);同时直接提取 \(x^{(3)}\) 在前向验证时产出的预测分布作为下一轮起点,无需重新计算。

损失函数 / 训练策略

FreeDave 属于完全的免训练(Training-free)与免模型修改(Model-free)算法,不需要在训练集上微调任何权重,也不引入任何辅助损失函数。在实际部署中,它直接作用于预训练离散掩码扩散语言模型(如 Dream-7B-Instruct、TraDo-4B-Instruct 及 TraDo-8B-Instruct)的推理执行流中。关键超参数为草稿步数 \(d\)(Dream 模型默认 \(d=4\),TraDo 模型配合 Eager Mode 默认 \(d=8\))。

实验关键数据

主实验

在单张 NVIDIA L40S GPU 上,评估各解码策略在数学推理(MATH500、GSM8K)与代码生成(MBPP、HumanEval)基准上的表现。对比基线包括一步一 token 的标准静态解码(Static)和基于阈值的启发式并行解码(Parallel)。评估指标包含准确率(Acc)、每秒生成 token 数(TPS)以及单次前向生成 token 数(Tokens Per Forward, TPF)。

基准测试 模型 块大小 \(B\) 解码策略 准确率 Acc (%) TPS (速度倍率) TPF (前向效率倍率)
MATH500 Dream-7B-Instruct 32 Static 37.40 23.40 (1.00×) 0.94 (1.00×)
32 Parallel 0.00 (-37.40) - -
32 FreeDave 36.60 (-0.80) 35.43 (1.51×) 2.59 (2.76×)
TraDo-4B-Instruct 4 Static 74.40 24.14 (1.00×) 0.80 (1.00×)
4 Parallel 70.20 (-4.20) 53.37 (2.21×) 1.85 (2.31×)
4 FreeDave+ 73.80 (-0.60) 55.17 (2.29×) 1.96 (2.45×)
TraDo-8B-Instruct 4 Static 78.60 23.73 (1.00×) 0.80 (1.00×)
4 Parallel 74.00 (-4.60) 46.71 (1.97×) 1.73 (2.16×)
4 FreeDave+ 79.00 (+0.40) 48.27 (2.03×) 1.89 (2.36×)
GSM8K Dream-7B-Instruct 32 Static 71.72 22.01 (1.00×) 0.86 (1.00×)
32 Parallel 0.08 (-71.64) - -
32 FreeDave 71.34 (-0.38) 29.25 (1.33×) 2.20 (2.56×)
TraDo-4B-Instruct 4 Static 91.66 23.96 (1.00×) 0.80 (1.00×)
4 Parallel 89.08 (-2.58) 52.74 (2.20×) 1.88 (2.35×)
4 FreeDave+ 91.81 (+0.15) 53.40 (2.23×) 2.06 (2.58×)
TraDo-8B-Instruct 4 Static 93.03 23.01 (1.00×) 0.80 (1.00×)
4 Parallel 91.58 (-1.45) 40.03 (1.74×) 1.53 (1.92×)
4 FreeDave+ 92.95 (-0.08) 42.24 (1.84×) 1.85 (2.31×)
MBPP Dream-7B-Instruct 32 Static 47.80 19.20 (1.00×) 0.75 (1.00×)
32 Parallel 0.20 (-47.50) - -
32 FreeDave 48.20 (+0.40) 27.21 (1.42×) 2.02 (2.69×)
TraDo-4B-Instruct 4 Static 55.00 22.97 (1.00×) 0.80 (1.00×)
4 Parallel 49.00 (-6.00) 59.05 (2.57×) 1.95 (2.44×)
4 FreeDave+ 55.20 (+0.20) 54.05 (2.35×) 2.01 (2.51×)
TraDo-8B-Instruct 4 Static 63.20 22.64 (1.00×) 0.80 (1.00×)
4 Parallel 53.20 (-10.00) 51.57 (2.28×) 1.74 (2.18×)
4 FreeDave+ 63.20 (+0.00) 51.59 (2.28×) 2.00 (2.50×)
HumanEval Dream-7B-Instruct 32 Static 54.27 20.62 (1.00×) 0.81 (1.00×)
32 Parallel 0.00 (-54.27) - -
32 FreeDave 54.27 (+0.00) 30.77 (1.49×) 2.29 (2.83×)
TraDo-4B-Instruct 4 Static 70.73 24.48 (1.00×) 0.80 (1.00×)
4 Parallel 65.24 (-5.49) 59.76 (2.44×) 1.92 (2.40×)
4 FreeDave+ 70.73 (+0.00) 57.77 (2.36×) 2.06 (2.58×)
TraDo-8B-Instruct 4 Static 79.27 23.85 (1.00×) 0.80 (1.00×)
4 Parallel 76.22 (-3.05) 50.04 (2.10×) 1.68 (2.10×)
4 FreeDave+ 78.66 (-0.61) 48.03 (2.01×) 1.90 (2.38×)

消融实验

1. 理论无损性与非确定性归因分析(TraDo-4B-Instruct on HumanEval)

尽管理论证明 FreeDave 严格无损,但在真实硬件与框架下,实测准确率存在 \(\pm 1\%\) 以内的微小浮动。作者通过严密控制数值非确定性(固定随机种子、强制 cuDNN 确定性算法、禁用 SDPA/FlashAttention/Liger Kernel)以及算法配置(关闭 Eager Mode),测试了序列级与 token 级精确匹配率:

数值非确定性控制 算法控制 (禁用 Eager Mode) 序列级匹配率 \(R_{\text{sequence}}\) Token 级前缀匹配率 \(R_{\text{token}}\)
开启 开启 159 / 164 (96.95%) 325,993 / 336,127 (96.99%)
关闭 开启 157 / 164 (95.73%) 322,010 / 336,127 (95.80%)
关闭 关闭 (默认配置) 152 / 164 (92.68%) 312,248 / 336,127 (92.90%)

此外,当作者在严控环境下设置单草稿步数 \(d=1\),或强制拒绝所有草稿时,token 级匹配率均达到 100%。这严谨证实了非确定性完全源自批处理 GPU 算子的浮点非结合性与 Eager Mode 的微小状态差异,而非算法逻辑缺陷。

2. 草稿步数 \(d\) 的权衡分析(TraDo-4B on HumanEval)

在 \(d \in \{1, 2, 4, 8, 16, 32\}\) 的扫描中,当 \(d=1\) 时即退化为 Static 解码(TPS 24.48,TPF 0.80)。随着 \(d\) 增大,TPF 持续上升并呈现收敛态势(\(d=8\) 时 TPF 达 2.06,\(d=16\) 时增至约 2.25);然而因 Batch 增大带来的单次前向耗时与显存开销增长,TPS 在 \(d=8\) 时达到峰值拐点(57.77 TPS)。因此 \(d=8\) 为显存开销与吞吐率的最佳平衡点。

关键发现

  • 启发式并行解码在大块下完全失效:在 Dream-7B(block size 32)上,阈值并行解码在 MATH500 和 HumanEval 上准确率直接降为 0.00%,原因是模型在较长块中容易对全局产生过度自信,导致退化生成大量死循环重复 token;即使将块大小降至 4,并行解码在 HumanEval 上依然下跌 7.32%。相比之下,FreeDave 在所有块尺寸下均保持无损。
  • 代码任务对解码次优路径极度敏感:在 MBPP 上,TraDo-8B 使用启发式并行解码后准确率从 63.20% 暴跌至 53.20%(掉点达 10.00%),而 FreeDave+ 完美保持在 63.20%,并取得了 2.28 倍的 TPS 提升与 2.50 倍的 TPF 飞跃。
  • 前向效率突破:在 HumanEval 任务上,FreeDave 使得 Dream-7B 的 TPF 达到 2.29(2.83 倍加速),TraDo-4B 的 TPF 达到 2.06(2.58 倍加速),显著减少了昂贵的大模型推理前向计算次数。

亮点与洞察

  • DLLM 本身即是自验证器:不同于自回归模型必须依赖外部小模型来生成草稿,扩散模型的双向全局预测特性使得单次前向输出的边际分布就能直接产出前瞻候选,验证过程只需将候选批量送入自身,完成了概念上的优雅闭环。
  • 分布复用实现零冗余计算:被接受草稿的中间前向表征和预测分布无需丢弃,直接作为下一步解码的输入基础,彻底消除了推测解码中“草稿一次、验证一次、重跑一次”的算力浪费。
  • Eager Mode 打破块扩散瓶颈:巧妙地将未来块的草稿生成提前引入,突破了 \(d \le B\) 的硬性约束,让小块扩散模型(如 TraDo 的 \(B=4\))也能享受长步长并行验证的收益。

局限与展望

  • 显存与计算瓶颈的权衡:增大草稿步数 \(d\) 会线性增加批量前向的显存占用,在计算受限或显存受限的边缘设备上,较大的 \(d\) 可能导致 TPS 收益不及预期。
  • 缺乏多请求在线高并发评测:当前论文仅在单卡离线批处理场景下完成验证,未测试在 vLLM / SGLang 风格的高并发动态批处理及分布式 Serving 环境下的调度开销。
  • 更大规模 DLLM 的扩展性:实验覆盖了 4B-8B 规模的 DLLM,未来有待在数十亿至数百亿参数的更大规模离散扩散基座上验证更长前瞻步数的接受率与加速上限。

相关工作与启发

  • vs. 自回归投机解码(Speculative Decoding, Leviathan et al.):传统投机解码需要训练或加载一个较小且同分布的草稿模型,存在草稿与主模型不对齐导致的接受率骤降问题;FreeDave 实现了“单模型自草稿+自验证”,完全省去额外草稿模型。
  • vs. 启发式置信度并行解码(Fast-dLLM / Dimple):现有并行解码直接根据阈值同时解除多个 token 的掩码,放弃了对标准扩散去噪轨迹的保真,导致在代码和数学基准上掉点严重;FreeDave 通过严格的等价性验证机制,理论与实证均保证了解码路径的无损。
  • vs. 块扩散与自适应缓存(Block Diffusion / dLLM-Cache):此类工作主要从 KV Cache 稀疏化或半因果结构入手加速前向速度,与 FreeDave 处于正交维度;FreeDave 可以与这些缓存优化技术无缝叠加。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示并证明扩散大模型原生具备自验证并行解码能力,构思巧妙。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖三大主流开源 DLLM、四大权威数学/代码基准,并给出了深入的非确定性与算子级归因分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学形式化定义清晰严格,定理推导与实现细节高度自洽。
  • 价值: ⭐⭐⭐⭐⭐ 彻底攻克了扩散语言模型并行解码“加速即掉点”的核心阻碍,具有极高落地价值。