SPOT:基于提交顺序预测的扩散语言模型投机解码加速 On the SPOT: Enhancing Speculative Decoding for Diffusion Language Models via Commit-Order Prediction
0.62M小模型学习DLM提交顺序构建草稿树,共享噪声验证实现最高3.4倍无损加速
前置知识
块扩散语言模型(Block Diffusion LM)
用离散扩散生成文本的模型:序列被分成块,块内用双向注意力对掩码位置去噪(一次前向可并行给出所有掩码位置的分布),块间保持因果。每步解码选择置信度最高的掩码位置填入一个 token,整块填满后再进入下一块,SDAR、LLaDA、Dream 是其代表。与自回归模型逐 token 从左到右生成不同,它每次前向都知道所有空格的候选,但只正式提交一个。
本文的草稿树、块状态、提交顺序等核心概念全部建立在块内双向、块间因果的结构上;不理解它就无法理解为什么位置顺序会成为自回归模型所没有的新预测维度。
投机解码(Speculative Decoding)
无损加速大模型推理的标准范式:廉价草稿模型快速提出若干未来 token,目标模型在单次前向中并行验证这些提议,接受与自身分布一致的前缀、拒绝其余,保证输出与无草稿解码完全一致(经典方案用拒绝采样做分布校正)。实际加速比取决于草稿被接受的平均长度与验证开销之比。
本文把该范式迁移到扩散模型,但验证对象从 token 序列变成块状态转移,且需要为任意采样温度重新设计精确验证方案,这是全文两大贡献之一的出发点。
Gumbel–Max 技巧与随机耦合
Gumbel–Max 技巧指出:对类别分布 $p$,取 $\arg\max_v(\log p_v + G_v)$($G_v$ 为独立标准 Gumbel 噪声)得到的样本恰服从 $p$,是温度采样的可复现实现。耦合指让两个分布的采样共享同一随机数:Daliri 等人证明共享 Gumbel 噪声下两方输出一致的概率至少为 $(1-\delta)/(1+\delta)$,其中 $\delta$ 为两分布的总变差距离。
SPOT 在 $T>0$ 时的精确验证完全依赖按绝对位置共享的 Gumbel 噪声表:条件于噪声,目标转移变成确定性函数,状态相等验证重新成立。
动态阈值解码(Dynamic Threshold Decoding)
Fast-dLLM 提出的有损加速规则:每次前向把所有置信度 $c_s=\max_v p_{\theta,s}(v)$ 超过阈值 $\tau$ 的掩码位置一次性全部提交(若没有位置过线则提交最高置信的一个),大幅减少前向次数但改变解码轨迹。本文主配置取 $\tau=0.9$。
SPOT-Sets 的目标就是在这个更快的解码器之上再做投机:草稿节点从单个位置变成整组提交位置及其 token,理解阈值规则才能理解提交集合草稿的含义。
总变差距离(Total Variation Distance)
离散分布 $p,q$ 的总变差 $D_{TV}(p,q)=\tfrac12\sum_z|p(z)-q(z)|$,取值 0 到 1,度量分布差异大小。它决定耦合采样的一致概率上界 $1-D_{TV}$,也决定拒绝采样的最大接受率。本文中草稿所用的陈旧条件分布与目标新鲜条件分布之间的 $\delta_T=D_{TV}$ 越小,共享噪声验证的接受率越高。
它是理解为什么耦合验证随温度升高会退化、以及为什么草稿质量只影响速度而不影响正确性的数学钥匙。
研究动机
扩散语言模型(DLM)承诺并行生成,但实践中存在未解决的速度–质量权衡:精确解码每次前向只提交一个掩码位置,SDAR-8B 生成一个 256 token 的样本约需 270 次目标前向;更快的多 token 方案(如 Fast-dLLM 的动态阈值解码、置信度 top-k)改变解码轨迹、有损。投机解码是自回归模型的标准无损加速,但无法直接迁移到 DLM:块内注意力是双向的,提交一个位置会改变块内所有剩余掩码位置的表示与输出分布,目标模型必须同时决定写什么 token 和下一个写哪里。近期工作(SSD 的置信度自链、Spiffy 的离线校准草稿图、MRP 的残差 logit 外推、BlockSpec 式置信度树)改用状态视角——提出假设未来状态再验证——使验证成为可能,但留下两个未解问题:草稿器到底应该预测什么信息才能超前多个转移?以及当目标模型随机采样而非贪心解码时,如何对整棵草稿树做精确验证?
本文的目标是本文要把投机解码无损地搬到块扩散语言模型上,并同时覆盖贪心(T=0)与任意温度采样:用一个极轻量的草稿器(0.62M 参数)预测目标模型的提交顺序、递归构建候选块状态树,由目标模型单次批量前向验证,使输出严格满足 $X^{\text{Spec}}_{T,D}\stackrel{d}{=}X^{\text{DF}}_T$。量化目标是在 SDAR-8B/4B/1.7B、LLaDA2.0-mini、LLaDA-8B、Dream-7B 共六个 DLM 上,相对无草稿解码的 NFE 缩减与吞吐提升全面超过校准式(Spiffy)与免训练(SSD、BlockSpec 式)基线,并在动态阈值解码之上进一步加速。最终实现 NFE 缩减最高 6.8 倍、吞吐最高 3.4 倍,且草稿器检查点可以跨目标零再训练迁移,SPOT-Sets 把阈值解码从平均 5.0 倍推到 7.9 倍。
与已有工作不同的是,本文的独特切入点是一个此前被忽视的实证分解:在投机解码关心的视野内,token 几乎是已知的,位置顺序才是瓶颈。陈旧分布的 argmax 在 92%(首块)及不低于 90%(其后各块)的提交上与最终 token 一致,每位置保留 8 个陈旧候选即可覆盖 89% 的八步连续提交;而上一轮最高置信位置下一步仍被提交的概率到第 9 块跌至 3%。活体先知实验(Figure 2a、Table 7)进一步证明真位置比真 token 值钱:位置先知把 NFE 缩减 5.4–8.3 倍,token 先知只有 3.0–6.5 倍。据此作者把草稿器的全部学习容量押在提交顺序上,token 直接复用陈旧分布;验证端则把 Daliri 等人的免通信 Gumbel 耦合引入扩散解码,用共享噪声把采样温度下的树验证化归为状态相等判断,填补了任意温度精确树验证的空白。
核心方法
直觉上,DLM 解码像填 crossword:大部分空格该填什么早已写在草稿纸上(陈旧分布),没人知道的是下一笔落在哪个格子。SPOT 的技术路线分三步。第一步草稿:一个 0.62M 参数的三层 transformer 编码器($d_{model}=128$、8 头)只读目标模型的陈旧输出分布——每个掩码位置的 top-32 概率(原始值与对数)、8 个汇总统计、2 个路径特征(已提交位置集合与深度)——对给定假设状态 $y$ 输出下一提交位置分布 $q_\phi(s\mid y)$。第二步建树:每步把预测位置与陈旧 token 假设(argmax,且当 top-1 概率低于 0.6 时附加次优候选作为对冲)配对生成子节点,按前缀质量 $m(z)=m(y)\,q_\phi(s\mid y)\,\tilde p_{t,s}(v)$ 递减做层批最佳优先搜索(DDTree),把预算 $D$ 分给质量最高的候选状态。第三步验证:目标模型对根与全部草稿状态做一次 $(1+D)L$ 位置的批量前向,逐边做状态相等检验,接受唯一与自身轨迹一致的前缀路径;$T>0$ 时通过共享 Gumbel 噪声让验证依然精确。
核心创新有两点。其一是学习提交顺序、复用陈旧 token:与 SSD(按置信度排序的自链)、Spiffy(离线校准的静态草稿图)、BlockSpec(陈旧置信度树)不同,SPOT 用可训练模型预测目标轨迹中的下一个提交位置,并在每次假设提交后做前缀条件化的重新预测——因为块内双向注意力意味着每次提交都会改写后续位置的分布,顺序必须逐状态重估。这正是作者证明的信息瓶颈:该模型在保留轨迹上以 75% 准确率预测下一次提交(陈旧置信度排序仅 22%),而 token 已被陈旧 argmax 以 92% 以上命中率覆盖。其二是耦合噪声验证:$T>0$ 时独立采样会让状态相等验证误杀合法草稿,SPOT 让草稿器与目标共享按绝对输出位置索引的 Gumbel 噪声表 $G_{a,v}$,条件于 $G$ 的目标转移 $\Phi_G$ 成为确定性函数,状态相等验证重归精确(Proposition 2),对噪声边缘化后严格恢复目标采样分布。这使树形批量验证在任意温度下无损,且不需要经典拒绝采样所要求的顺序提议接口。
方法步骤详情
训练(Algorithm 1):在目标模型自身的无草稿轨迹上,保留每轮的陈旧条件分布 $\tilde p_t$ 与后续提交序列 $(s^\star_t,s^\star_{t+1},\dots)$;对每个前缀深度 $d$,重放前 $d$ 个提交构造假设状态 $y_{t,d}$(陈旧分布保持不变),以下一个真实提交 $s^\star_{t+d}$ 为监督,最小化 $\mathcal{L}_{order}(\phi)=-\mathbb{E}\log q_\phi(s^\star\mid y)$。SDAR-8B 草稿器用 CodeAlpaca 约 2.1 万条轨迹、Adam 学习率 $3\times10^{-4}$ 训 10 轮,单卡不到 2 GPU 小时,保留集下一次提交准确率 0.75(深度 0 为 0.77、深度 7 为 0.72)。推理(Algorithm 2):从当前状态 $x_t$ 与陈旧分布 $\tilde p_t$ 出发,对每个活跃前缀 $y$ 用 $q_\phi(\cdot\mid y)$ 给剩余位置打分,与 token 假设集 $H_t(s)$ 配对生成子节点并按公式 (3) 累积前缀质量;DDTree 层批扩展、保留质量最高的 $D$ 个节点。验证:一次前向同时返回根与每个草稿副本处目标自身的下一提交($T=0$ 为 $\Phi_0$,$T>0$ 为固定 $G$ 后的 $\Phi_G$),接受相等的边,否则回退提交目标自身结果。SPOT-γ 变体额外用前缀质量下限 $\gamma$(SDAR-8B 上 0.02)丢弃低质量候选,让实际验证预算随草稿置信度自适应。
技术新颖性
新颖性可沿四个维度定位(Table 6)。第一,学习顺序是独有的:Spiffy 的图是离线校准的静态结构,SSD 与 BlockSpec 用陈旧置信度现场排序,MRP 学习残差 logit 外推但不学顺序、且只在贪心下验证;SPOT 是唯一用训练模型(且仅 0.62M)预测提交顺序的方法。第二,多步树加单次批量验证:预算通过分支而非链条生效(换成链 NFE 增加 7.0–11.7),深块每轮接受 token 数可达 14.5(Spiffy 图被深度 6 封顶在约 5)。第三,任意温度下的精确性:greedy 匹配随温度漂移(T=1.5 时 SPOT 输出准确率 0.15–0.16 对模型真实的 0.11),独立采样匹配在 T=2 几乎不接受(257.5 NFE 对无草稿 267.1),只有耦合噪声在所有温度保持分布一致且与基线共享同一验证器使对比只关乎草稿质量。第四,可迁移性:草稿器不读隐状态、token 嵌入或文本,只读 top-32 概率分布,因此检查点可跨六个 DLM(含不同分词器的 MoE 模型)零重训练使用。
实验结果
主结果(Table 2,SDAR-8B b32,n=350):T=0 时 SPOT 相对无草稿解码在 D8 用 3.2–4.7 倍更少前向(如 MBPP 从 269.2 降到 56.7 NFE),D20 达 4.1–6.8 倍(MBPP 39.5),吞吐为无草稿的 2.5–3.4 倍;对 Spiffy 校准图在所有任务-预算组合少 16–36% 前向、吞吐 1.04–1.41 倍。基线随预算崩塌:从 D8 到 D20 自由链吞吐从 2.3–2.6 倍跌至 1.4–1.7 倍、置信度树跌至 0.86–1.10 倍(MBPP 上比无草稿还慢),因为它们每轮只接受 1.1–2.9 个 token 而 SPOT 接受 2.5–7.5 个。T=1(耦合噪声验证)结论不变:比 Spiffy 少 16–35% 前向、吞吐 1.08–1.37 倍,所有结果分布保持(准确率与无草稿差小于 0.04);MRP 报告的 1.17–1.40 倍加速在相同目标上被 SPOT 的 2.5–3.8 倍(D8)大幅超越。SPOT-Sets 叠加在阈值解码(τ=0.9)上:阈值解码单独 3.3–6.1 倍(平均 5.0 倍),SPOT-Sets-γ 达 5.2–9.9 倍(平均 7.9 倍),比阈值解码少 47–48% 前向。迁移(Table 3/15):SPOT 检查点零再训练跨六模型,在 LLaDA-8B 与 Dream-7B 把 256 次前向砍到 55–88 与 51–80 次;SPOT-Sets 只需 99 条提示重拟合输入层(0.27M/5.1M 参数)即保持在原生效 8% 以内。消融(Table 19):树比链在 D20 省 11.7 NFE;质量地板 γ 用最多 4 次额外 NFE 把 D20 吞吐从 3.19 倍提到 4.13 倍。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 文本生成(SDAR-8B,T=0,长度256) | 每样本目标前向次数 NFE | D8 为 56.7–84.6,D20 为 39.5–66.0(缩减最高 6.8 倍) | 无草稿解码 269.2–271.8 | NFE 减少 3.2–6.8 倍 |
| 吞吐量(SDAR-8B,T=0,D8) | tok/s 相对无草稿解码倍数 | 2.50–3.38 倍(SPOT-γ 最高 3.42–3.48 倍) | Spiffy 2.29–2.64 倍;SSD 链 2.28–2.63 倍 | 相对最强基线 Spiffy 提升 1.04–1.41 倍 |
| T=1 采样(六模型验证) | NFE 相对 Spiffy / 分布保持 | 少 16–35% 前向,吞吐 1.08–1.37 倍,准确率偏差 <0.04 | Spiffy 校准图(耦合噪声验证下) | 全部任务-预算组合领先且分布无损 |
| 阈值解码加速(SDAR-8B,τ=0.9,T=0) | 相对标准解码的吞吐倍数 | SPOT-Sets-γ 5.2–9.9 倍(平均 7.9 倍) | 阈值解码单独 3.3–6.1 倍(平均 5.0 倍) | 平均提升 59%,exact 变体 7.6 倍 |
| 下一提交位置预测 | 保留轨迹上的下一提交准确率 | 75%(深度 0 达 77%) | 陈旧置信度排序 22% | 提升 53 个百分点 |
| 跨模型零迁移(LLaDA-8B / Dream-7B) | 每样本 NFE(近似前缀缓存下) | LLaDA-8B 55–88,Dream-7B 51–80 | 无草稿 256 次前向 | 前向减少约 3–4.7 倍 |
局限与改进
作者承认的限制:收益在验证昂贵或接受率低时缩水——宽块下一个草稿要付出 $(1+D)L$ 个位置的验证成本,双向注意力下精确的紧凑树打包仍未解决;低接受任务上优势收窄,SDAR-1.7B 的 MATH500 D20 上 Spiffy 反超 0.08–0.09 倍,靠 SPOT-γ 才夺回;T≥2 时 T=0 训练的顺序模型被噪声感知排序超越(D8 上 114.8 对 187.2 NFE);SPOT-Sets 依赖目标隐状态与嵌入,换目标需短校准,其 5.1M 草稿占一轮时间约 6%,且主结果使用的 product max-commit 规则会改变输出(exact 变体为 7.6 倍)。我补充的观察:固定 256 生成长度把答案后的结束符也计入加速比——四任务答案实际仅 55–251 token,按答案截断后收益缩水至约 1.04–2.54 倍(Table 18);理论精确性依赖批量副本不改变数学条件分布的理想化假设,浮点批形状效应在 LLaDA2.0-mini 的 MoE 路由上已导致阈值解码输出不可复现;与 MRP 的对比是跨实现、跨块的间接比较,只比相对增益。
独立分析的弱点
第一,验证成本的结构性浪费:每个草稿是整个块的独立副本,D=20 时一次前向要做 $21\times32$ 个位置,被拒绝的草稿浪费全部算力,这是基线随预算崩塌、SPOT 吞吐在 D12–D16 见顶的根源;改进方向是研究双向注意力下可证明精确的紧凑树打包,或在副本间共享计算。第二,训练分布依赖:顺序模型在 T=0 轨迹上训练,T≥2 时被噪声感知排序反超;改进方向是把共享噪声作为输入特征做温度条件化训练或在线微调。第三,全注意力目标失守:LLaDA-8B、Dream-7B 没有精确前缀缓存,只能搭配近似缓存测量迁移上限,失去无损标签;改进方向是设计对缓存刷新友好的草稿布局。第四,SPOT-Sets 的 product max-commit 为速度牺牲一致性:与阈值解码输出仅 20–54% 完全一致(Table 26),部署上必须明示;改进方向是把严格性做成运行时开关并给出质量–速度帕累托前沿。第五,预算分配仍有空间:错误分析显示 D20 时 MBPP 上 32% 的首次拒绝发生在真实位置排名第一的情况下(Figure 7),是树预算耗尽而非模型能力不足,自适应预算或更好的节点分配可直接兑现这些被浪费的接受。
未来方向
作者方向:解决宽块下的精确树打包;把耦合验证与温度自适应草稿训练结合(覆盖 T≥2 的噪声感知排序);把 SPOT-Sets 每目标 99 条提示的短校准自动化。基于成果可延伸的方向:其一,把顺序模型升级为读隐状态的更强版本——SPOT-Sets 已证明隐状态有价值(token 命中率 87/73/62% 对陈旧的 83/60/44%),可在 0.62M 与 5.1M 之间系统扫描容量–迁移权衡;其二,预算 $D$、质量地板 $\gamma$ 的在线自适应控制,与吞吐最优预算 D12–D16 的目标依赖性(Figure 5)配合做运行时调参;其三,把位置+内容的双轴投机推广到多模态扩散模型或图像扩散模型的空间位置解码;其四,耦合噪声思想可扩展到 beam search、树搜索解码的精确加速,以及Speculative Decoding 以外的序列决策验证;其五,错误归因显示 62–81% 首次拒绝源于位置错过,把位置预测与置信度校准联合训练有望进一步压低 NFE。
复现评估
复现条件很好。代码与草稿检查点以 MIT 协议开源(github.com/Eliezer318/SPOT-dLLM,HuggingFace 镜像 huggingface.co/EliezerS/SPOT-dLLM),附复现指令。训练成本极低:0.62M 顺序模型单卡不到 2 GPU 小时(约 2.1 万条 CodeAlpaca 轨迹),SPOT-Sets 草稿约 30 分钟(约 100 万样本,ShareGPT+UltraChat);大规模吞吐实验用两台 8 卡 RTX PRO 6000 Blackwell 机器、每卡单任务,但这是吞吐测量的要求而非方法本身的要求。评测数据全部公开(GSM8K/MBPP/HumanEval/MATH500 标准切分,n=350/164,长度 256),基线用 Spiffy 官方校准流程,SSD/BlockSpec 式基线在同一解码器与验证器下复现。评估协议极为细致:种子、机器归属、NFE 与吞吐口径、5% 平局带、种子方差(T=0 吞吐比中位极差 0.05)均有披露。风险点是批量浮点效应(MoE 目标上阈值解码输出本身依赖批形状)与跨实现基线对比。总体复现难度:低到中,单卡即可复现草稿训练与小规模验证。
论文图表
(a) 给出位置先知、token 先知与联合先知的吞吐上界:位置先知在所有任务都高于 token 先知,完美草稿可达 6.2–7.0 倍;(b) 显示每位置保留 top-$t$ 陈旧候选时未来 $K$ 步提交 token 全被覆盖的概率,top-8 已覆盖约 89% 的八步连续提交。
这是全文动机的实证核心,直接证明顺序是瓶颈而 token 已知,是方法设计的依据。
瀑布图展示组件逐级叠加后的平均吞吐:集合草稿+树 6.72 倍,+自适应地板 6.85 倍,+跨块草稿 7.28 倍,+草稿刷新 7.54 倍,+product max-commit 7.90 倍;两个块端设计贡献最大(各 +0.43 与 +0.26)。
量化 SPOT-Sets 每个设计选择的边际贡献,说明跨块草稿与草稿刷新为何是关键工程创新。
左图显示草稿器的前缀质量分与实际接受率几乎校准(最高质量 bin 接受率 96%,最低 bin 3%);右图按块分解首次拒绝原因:早块以位置错过为主(占 62–81%),深块以预算耗尽为主(MBPP D20 达 32%,最后一块 90%)。
提供改进方向的诊断证据,解释 NFE 为何随预算持续改善、质量地板 γ 为何能换吞吐。