S2D2:通过免训练自推测解码加速扩散语言模型 S2D2: Fast Decoding for Diffusion LLMs via Training-Free Self-Speculation
用同一模型的自回归模式作验证器,无需训练即加速扩散LLM解码
前置知识
块级自回归扩散解码 (Block-wise AR Diffusion Decoding)
块扩散模型(如 BD3/SDAR/LLaDA)将序列分成固定大小的块(block),块与块之间按自回归方式生成(左侧已完成块作为条件),块内则使用掩码扩散并行去噪。这种混合设计既保留了 KV 缓存的复用效率,又允许块内 token 并行生成。标准扩散去噪在每步对所有掩码位置产生候选 token 和置信度分数,通过置信度阈值或固定调度来决定哪些位置被揭开。块大小 B 是关键参数:B=1 退化为纯自回归,B 更大则并行度更高但序列级依赖更弱。
S2D2 的加速基础架构,理解块扩散如何在块间保持自回归、块内并行去噪是理解本文的前提
推测解码 (Speculative Decoding)
推测解码是一种加速自回归生成的技术:用一个轻量的草稿模型(drafter)一次生成多个 token,然后用目标模型(verifier)并行验证。验证时从左到右扫描,每个 token 以概率 min(1, q_i/p_i) 接受(q_i 是验证器概率,p_i 是草稿概率),在第一个拒绝位置从残差分布重采样。这个过程在数学上保证不改变目标模型的输出分布,同时通常能每步接受多个 token,从而减少总的前向传播次数。
S2D2 的核心算法框架,将推测解码从 AR 模型迁移到扩散模型
置信度阈值解码 (Confidence-Threshold Decoding)
这是扩散语言模型(如 LLaDA)中的标准解码策略:在每个去噪步,模型对所有掩码位置产生候选 token 和置信度分数(通常来自 softmax 概率),然后以固定阈值或动态阈值决定哪些位置被接受。问题在于:激进的阈值(低阈值)会揭开很多位置但质量下降,保守的阈值(高阈值)虽然质量好但需要更多去噪步。动态阈值试图自适应调整,但仍依赖草稿置信度本身做决定,缺乏序列级的一致性检验。
S2D2 要替代的基线方法,其缺陷直接构成动机
自回归度 (AR-ness) 诊断指标
AR-ness 衡量扩散模型生成的 token 与自回归从左到右解码的吻合程度。值为 1 表示完全左到右自回归解码。论文使用 DiffuCoder 引入的两种形式:局部 AR-ness(在局部窗口内衡量)和全局 AR-ness(在整个序列上衡量)。通过比较不同模型(SDAR vs LLaDA)和不同任务(GSM8K vs MBPP)的 AR-ness,论文发现模型在数学推理任务上的 AR 结构比代码生成更强,这与 AR 解码在不同任务上的相对竞争力一致。
论文用来分析扩散模型行为、解释 S2D2 为何有效的关键分析工具
2L 技巧 (2L Trick) 用于并行验证
对于位置对齐的扩散 LLM(如 SDAR、LLaDA),需要在一次前向传播中同时计算草稿 token 对其左侧上下文的条件概率。2L 技巧的做法是:将长度为 L 的草稿 span 复制一份,拼接成长度 2L 的序列,并构造特殊的注意力掩码矩阵 M_ver,使前 L 个位置能看到草稿 token 但不能看到自身,后 L 个位置(全掩码)使用标准因果注意力。这样一次前向传播就能得到所有位置的验证概率。对于右移模型(如 Dream、Fast-dLLM v2),标准因果掩码已经提供了验证视图。
S2D2 能在单次前向传播中获取所有位置验证分数的关键工程技巧
研究动机
块扩散语言模型虽然通过块内并行去噪提供了比纯自回归模型更快的解码潜力,但在实际需要的少步解码(few-step decoding)场景下,标准的置信度阈值解码存在严重缺陷。具体来说,以 LLaDA 和 SDAR 为代表的块扩散模型在生成时采用均值场、token 独立的参数化方式,这削弱了序列级依赖,当去噪步数减少时错误会迅速累积。以 SDAR-8B 在 GSM8K 上的实验为例,标准静态解码在 B=4 时只能达到 1.4× 加速,动态解码虽然达到 2.6× 加速但准确率从 89.3% 降到 89.3%(持平);而在更大的块大小(B=32)下,动态解码的准确率会从 89.3% 急剧下降到 83.0%,MBPP 上更是从 64.4% 暴跌到 43.6%。现有的解决方案要么需要额外训练(如 EDLM 引入 AR 能量模型和自归一化重要性采样),要么在推理时增加计算开销(如多候选重排序),无法在不增加训练成本的前提下实现高效加速。
本文的目标是本文的核心目标是回答一个速度优先的问题:能否在推理时利用自回归结构来加速块扩散模型的解码,同时不引入额外训练、不依赖辅助模型或架构修改、且保持块扩散的并行性优势?具体而言,作者希望设计一个即插即用(plug-and-play)的推理框架,能够直接应用于现有的预训练块扩散模型,在相同的计算预算下同时提升生成质量和解码速度,改善准确率-速度权衡的帕累托前沿。
与已有工作不同的是,S2D2 的独特切入角度源于一个关键观察:当块扩散模型的块大小减小到 1 时,它就退化为标准的自回归模型,可以充当验证器(verifier)。这意味着同一个预训练的块扩散模型天然具备两种模式——标准块扩散解码模式(作为草稿器)和块大小为 1 的自回归模式(作为验证器)。与现有方法的本质区别在于:EDLM 需要训练额外的 AR 能量模型并使用全局多候选重采样,ASSD 需要特定架构(XLNet 风格的任意子集 AR 建模),BlockSpec 和 SSD 依赖分层批处理——而 S2D2 完全是免训练的,仅利用模型本身的两种解码模式进行自推测。这种设计还将推测解码的接受-拒绝机制从单纯的加速工具重新解释为一种局部序列级修正机制,对应于随机、贪心的自回归残差能量校正。
核心方法
S2D2 的整体思路可以用一个直觉来概括:扩散模型解码就像一个快速但不太仔细的「草稿工」,它能并行提出很多 token 候选,但缺乏对序列整体一致性的把握;而同一个模型在块大小为 1 时变成了一个慢但可靠的「审核员」,它按严格的从左到右顺序检查每个 token。S2D2 让草稿工先干活,然后只在关键时刻请审核员来检查最关键的那一段。技术上,S2D2 在标准块扩散解码的每个去噪步中插入一个可选的自推测验证步骤:(1) 草稿前向传播产生候选 token 和置信度;(2) 找到第一个连续掩码 span C_t;(3) 通过轻量路由策略决定是否值得验证;(4) 如果验证,切换到块大小为 1 的自回归掩码模式,计算验证概率,从左到右执行推测接受;(5) 在第一个拒绝位置从残差分布重采样,然后回退到标准扩散解码。这种混合轨迹既有扩散的并行性,又有自回归的序列级纠正能力。
S2D2 的核心创新在于「自推测」(self-speculation)的概念:用同一模型的两种模式实现草稿-验证,而无需任何辅助模型或额外训练。与已有方法的本质区别有三点:第一,与 EDLM 不同,S2D2 不训练全局 AR 能量模型,而是在线地使用模型的块大小为 1 模式作为局部验证器,通过推测接受和残差重采样实现修正;第二,与推测解码的标准用法不同,S2D2 不追求完全恢复块大小为 1 的自回归解码,而是将 AR 验证作为混合扩散轨迹中的「局部序列级批评家」;第三,与置信度阈值解码不同,S2D2 的接受决策基于验证器归一化的概率比 min(1, q_i/p_i),而不仅仅是草稿置信度本身,这提供了一个更强的局部一致性检验。论文还从残差能量的角度给出了理论解释:接受概率 min(1, e^{-E_i}) 是单调递减的残差能量函数,其中 E_i = -log q_i + log p_i,因此低能量(AR 和扩散一致)的 token 更容易被接受,高能量(不一致)的 token 被拒绝并从修正分布重采样。
方法步骤详情
S2D2 的完整方法流程如下:(1)外层循环:遵循标准块级自回归解码框架(Algorithm 1),逐块生成,每块初始化为全掩码 [MASK]^B,利用已完成块的 KV 缓存。(2)草稿前向传播:在每个去噪步 t,对当前块执行草稿前向传播得到 logits ℓ,采样得到候选 token x̂ 和置信度 p,标识当前掩码位置集合 M_t。(3)找到第一个连续掩码 span C_t(从左到右第一个连续的掩码位置段)。(4)路由决策(Algorithm 4):使用轻量策略判断是否值得验证。支持四种策略:最小 span 策略(|C_t| ≥ τ_span 时验证)、分数阈值策略(计算期望接受前缀长度 K̂ 的分数 s ≥ τ_score 时验证)、滞后策略(使用 ON/OFF 状态和不同开关阈值避免震荡)、以及 UCB 上下文 bandit。期望接受前缀长度 K̂ = Σ_{k=1}^{L} Π_{i=1}^{k} α_i,其中 α_i 可用 margin-based(α_i = 1[m_i ≥ τ_margin])或 entropy-based(α_i = exp(-β·H̃_i))估计。(5)如果决定验证:将草稿 token 填入 C_t 位置,构造验证模式注意力掩码(2L 技巧或标准因果掩码),执行一次额外前向传播得到验证概率 q,从左到右逐位置以概率 min(1, q_i/p_i) 接受,在第一个拒绝位置从残差分布 (P_{ℓ_ver} - P_ℓ)^+ 重采样。(6)如果决定不验证:回退到标准置信度阈值解码。(7)更新块状态,进入下一个去噪步。
技术新颖性
S2D2 在技术上有多个新颖之处。首先,这是首个针对块扩散语言模型的免训练自推测解码方法,通过发现块大小为 1 时块扩散模型自然退化为自回归模型这一事实,巧妙地复用了同一预训练模型的两种模式。其次,验证模式的注意力掩码设计(2L 技巧)能够在单次前向传播中并行计算所有位置的验证概率,这是将推测解码从 AR 模型迁移到扩散模型的关键工程创新。第三,多种轻量路由策略的设计——特别是最小 span 策略的简洁有效性和滞后策略对模式震荡的抑制——使得验证的额外开销能够被有效控制。第四,论文将推测接受-拒绝机制与 EDLM 的残差能量理论建立了形式化联系:min(1, q_i/p_i) = min(1, e^{-E_i}),将推测解码解释为随机贪心的局部能量最小化,这一理论视角为理解自推测在扩散模型中的作用提供了新的分析框架。最后,部分因果注意力掩码的引入(Eq. 4)使得草稿轨迹更加 AR 化,进一步提升了验证的有效性。
实验结果
论文在三个主流块扩散模型家族(SDAR、Fast-dLLM v2、LLaDA2.1)的五个模型上进行了全面实验,覆盖 GSM8K(数学推理)、MBPP(代码生成)、HumanEval(代码生成)和 IFEval(指令跟随)四个基准。核心发现包括:(1)SDAR 家族上,S2D2 在所有模型规模(1.7B/4B/8B)和几乎所有任务上都优于动态置信度阈值基线。最具代表性的结果是 SDAR-1.7B config-B:S2D2 达到 4.7× 加速(相对 AR),比动态解码快 1.57×(4.7/3.1),同时平均准确率提升 4.5 个百分点(52.9% vs 48.4%)。SDAR-8B config-A 在保持 2.0-2.2× 加速的同时,平均准确率达到 72.6%,超过了 AR 基线的 71.7%。(2)Fast-dLLM v2 上,S2D2 在所有子块大小配置下都改善了准确率-速度前沿。在 SB=32 的最大并行配置下,S2D2 比动态解码快 1.07×(3.1× vs 2.9×),同时平均准确率提升 4.5 个百分点(52.6% vs 48.1%)。(3)LLaDA2.1-Mini 上,S2D2 与模型内置的 token 编辑自纠正机制互补。在保守设置(τ_mask=0.95, τ_edit=0.9)下,S2D2 在 2.2× 加速下达到 79.3% 准确率,比静态基线(0.5× 加速,79.2%)快 4.4× 且准确率略高。(4)AR-ness 诊断分析显示,模型在数学推理任务上的 AR 结构比代码生成更强(SDAR 在 MBPP 上 AR-ness 更高,LLaDA 则相反),这解释了 S2D2 在不同任务上的效果差异。(5)推理时间缩放趋势分析表明,S2D2 展现出比标准扩散解码更平坦的缩放曲线,尤其在大块大小下优势明显。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K (数学推理) | 准确率 (%) | SDAR-8B: 89.6% (config-A, 2.0×); SDAR-1.7B: 73.8% (config-B, 4.3×) | SDAR-8B AR: 89.3% (1.0×); SDAR-8B 动态解码: 89.3% (2.6×); SDAR-1.7B 动态解码: 77.2% (2.7×) | SDAR-8B config-A 准确率 +0.3 点同时 2.0× 加速; SDAR-1.7B config-B 4.3× 加速但准确率低于动态基线 |
| MBPP (代码生成) | 准确率 (%) | SDAR-8B: 62.0% (config-A, 2.1×); SDAR-1.7B: 44.4% (config-B, 2.9×) | SDAR-8B AR: 64.4% (1.0×); SDAR-8B 动态解码: 60.6% (2.1×); SDAR-1.7B 动态解码: 21.6% (1.3×) | SDAR-1.7B config-B 准确率从 21.6% 提升到 44.4% (+22.8 点),加速从 1.3× 到 2.9× |
| HumanEval (代码生成) | 准确率 (%) | SDAR-8B: 78.7% (config-A, 2.2×); SDAR-1.7B: 52.4% (config-B, 4.1×) | SDAR-8B AR: 75.6% (1.0×); SDAR-8B 动态解码: 78.0% (2.9×); SDAR-1.7B 动态解码: 51.8% (3.4×) | SDAR-8B config-A 准确率 +3.1 点但速度略慢于动态解码; SDAR-1.7B 保持准确率同时加速 |
| IFEval (指令跟随) | 准确率 (%) | SDAR-8B: 60.9% (config-B, 3.9×); SDAR-4B: 59.5% (config-A, 4.1×) | SDAR-8B AR: 57.7% (1.0×); SDAR-8B 动态解码: 54.0% (2.7×); SDAR-4B AR: 54.3% (1.0×) | SDAR-8B config-B 准确率 +6.9 点且加速 3.9×; SDAR-4B config-A 准确率 +5.2 点且加速 4.1× |
| LLaDA2.1-Mini 保守设置 (GSM8K) | 准确率 (%) | 90.8% (S2D2, 2.2× 加速) | AR: 90.8% (1.0×); 静态扩散: 91.0% (0.5×, 比 AR 慢); 动态扩散: 89.8% (2.1×) | 比静态基线快 4.4× 且准确率相当; 比动态基线快 1.05× 且准确率 +1.0 点 |
局限与改进
论文存在几个值得讨论的局限性。首先,S2D2 的成功前提是模型的块大小为 1 自回归模式是一个可靠的验证器。论文在附录中坦诚地展示了一个失败案例:当使用 D2F(扩散到 AR)变体的 LLaDA 模型时,其 AR 模式在 GSM8K 上仅达到 20% 准确率(远低于扩散解码的 50%),此时 S2D2 的残差重采样反而会降低准确率(从 50% 降到 33.3%),只有禁用残差重采样(仅拒绝不重采样)才能部分恢复。这说明 S2D2 对模型的 AR 条件质量有隐含要求。其次,验证需要一次额外的前向传播,在小 batch size 下可以被速度提升抵消,但在高 batch size 下收益递减——论文显示 SDAR-8B 在 batch size 64 时,S2D2 的吞吐量增益比 batch size 1 时更小,因为 2L 验证器构造增加了有效工作量。第三,S2D2 只验证第一个连续掩码 span 而非所有掩码位置,这意味着它无法捕获后续 span 中的不一致性。第四,路由策略的超参数(如 τ_span、τ_score、τ_on/off)需要在目标模型上进行搜索,虽然搜索成本远低于重新训练,但仍有调优开销。最后,论文的实验主要在数学推理和代码生成任务上,在长文本生成、对话、翻译等其他任务上的效果尚不明确。
独立分析的弱点
S2D2 存在几个可以改进的弱点:(1)仅验证第一个连续掩码 span:当前设计只对 C_t(第一个连续掩码段)执行验证,后续掩码位置完全依赖扩散解码。在长块(B=32 或更大)中,可能有多个不连续的掩码 span,仅验证第一个无法捕获后续位置的不一致性。改进方向可以是设计多 span 验证策略,或者在路由决策中考虑所有 span 的整体结构。(2)验证概率的 2L 构造在高 batch 下成为瓶颈:SDAR 使用 2L 技巧需要将序列长度翻倍,在高 batch size 下显存和计算开销显著增加。Fast-dLLM v2 因为使用右移架构不需要 2L,其批处理吞吐量增益更直接。改进方向可以是探索更高效的验证掩码设计或近似验证方法。(3)路由策略的启发式性质:四种路由策略都是启发式设计,虽然上下文 bandit 提供了一种自适应方式,但实验显示它并非最佳。改进方向可以是训练一个轻量的元学习路由器,或基于解码轨迹的在线学习策略。(4)接受概率估计器的选择影响下游性能:论文发现估计精度更高的 hard-margin 估计器不一定带来更好的解码准确率,反而 soft-entropy 估计器更稳定。这暗示估计器的校准性和鲁棒性比原始精度更重要,值得进一步研究。(5)缺乏对推理时间缩放的理论分析:论文展示了 S2D2 的推理时间缩放曲线更平坦,但没有给出理论解释为什么自推测验证能改善缩放行为。
未来方向
论文和其成果可以延伸出多个研究方向:(1)将 S2D2 与批处理优化的推测解码方法(如 SSD 的分层批处理)结合,充分利用两种加速范式的互补性。论文提到这一方向是正交的,可以组合使用。(2)探索将自推测思想扩展到非块扩散模型(如连续扩散 LM),虽然这需要新的验证模式设计。(3)研究更好的验证路由策略,特别是基于强化学习或元学习的自适应路由,使其能根据任务特性、解码进度和模型行为动态调整。(4)将残差能量校正的理论视角与扩散模型的训练目标联系起来,探索是否可以在训练阶段就优化模型使其更适合自推测加速。(5)将 S2D2 应用于更长序列生成和更多样化的任务(如对话、摘要、翻译),研究自推测在不同任务特性下的效果差异。(6)探索多步验证策略:当前只在每步验证一次,可以研究是否值得在某些关键步骤执行多轮草稿-验证迭代。(7)研究自推测与 LLaDA 等模型内置的 token 编辑机制的更深度协同,论文已显示两者互补,但最优组合策略尚未充分探索。
复现评估
论文在可复现性方面做得较好。代码已在 GitHub 开源(https://github.com/phymhan/S2D2)。实验覆盖了三个开源模型家族(SDAR 1.7B/4B/8B、Fast-dLLM v2、LLaDA2.1-Mini),这些模型权重都是公开的。论文使用标准基准(GSM8K、MBPP、HumanEval、IFEval)并在附录中详细说明了评估设置,包括提示模板、答案提取方法和评估后端(lm-eval-harness 或自定义脚本)。超参数搜索在附录 A.5 中有完整记录,搜索网格和选择策略都有说明。论文还在附录 A.6 中提供了 LLaDA2.1-Mini 的 5-seed 多随机种子结果以验证稳定性。硬件方面,所有实验在单张 NVIDIA H100 80GB 上进行,batch size 1/32/64 的吞吐量测试也有报告。主要的复现难度在于:路由策略的超参数需要在具体模型上重新搜索(虽然论文提供了推荐范围),以及 2L 验证掩码的实现需要注意不同架构(位置对齐 vs 右移)的差异。
论文图表