SpecCTMC:扩散语言模型的精确推测采样

SpecCTMC 用无需训练的推测采样加速扩散语言模型的精确连续时间采样,评估次数降 1.73-4.38 倍。
扩散语言模型通过并行生成令牌加速文本生成,但独立采样会破坏令牌依赖、改变输出分布;精确连续时间采样虽能保真,却需为每个令牌顺序执行一次模型评估。本文介绍 SpecCTMC,一种无需训练的推测采样器:它利用模型此前已计算的预测连续提议多个令牌更新,并行评估中间状态,再通过精确的接受-校正程序保证结果严格服从原始连续时间马尔可夫链的联合分布,且不需要独立草稿模型。作者证明其保持事件时间、生成顺序与令牌取值的联合分布,并在 LLaDA-8B-Instruct 和 Dream-v0-Instruct-7B 的五个基准上把顺序评估次数减少 1.73 至 4.38 倍,同时保持参考分布。
扩散语言模型(Diffusion Language Model)通过并行生成令牌来加速文本生成,但独立采样多个令牌会破坏依赖关系,改变输出分布。精确连续时间采样虽能避免这种近似,却通常需要为每个令牌执行一次顺序模型评估。本文介绍 SpecCTMC,一种无需训练即可减少这种顺序瓶颈的推测采样器,它不依赖独立的草稿模型,在保持原始分布的同时显著减少模型评估次数。

并行生成的分布困境
掩码扩散语言模型(Masked Diffusion Language Model)的核心优势在于并行解码:一次前向传播可以同时预测多个位置的令牌。然而,这种并行性伴随着统计代价。当多个令牌被独立采样时,它们之间的依赖关系被忽略,导致生成结果的联合分布偏离模型原本定义的分布。
这种偏差并非无关紧要。在需要严格概率保证的场景中——如可控文本生成、科学计算或概率编程——分布保真度至关重要。传统的解决方案是采用精确的连续时间采样,它模拟连续时间马尔可夫链(CTMC)的演化过程,确保每个令牌的生成顺序、时间和取值都严格遵循底层链的联合分布。但代价也很明显:精确采样通常需要为每个令牌顺序执行一次模型评估,完全抵消了并行生成的速度优势。
连续时间马尔可夫链(CTMC) 是描述状态在连续时间上随机跳转的概率模型。在掩码扩散语言模型中,CTMC 将文本生成建模为:每个令牌位置从“掩码”状态跳转到具体词元的过程,跳转时刻和跳转目标由速率矩阵决定。精确采样要求严格模拟这一跳转过程,因此每次状态转移都需要基于当前完整序列重新计算模型输出,导致评估次数与令牌数成正比。并行采样则把多个位置的跳转视为独立事件,相当于把联合速率因子化为边缘速率的乘积,从而引入因子分解误差。理解 CTMC 的速率分解结构,是理解 SpecCTMC 为何必须做接受-校正的钥匙:只有校正步骤才能把独立提议“还原”为联合链的真实转移概率。
SpecCTMC 的三步机制
SpecCTMC 的设计目标是在不训练额外草稿模型的前提下打破这一顺序瓶颈。它的核心思想是利用模型已经计算出的预测结果来提议后续的令牌更新,从而将多个步骤并行化。
具体而言,SpecCTMC 的运行分为三个阶段。首先,它基于之前步骤中已计算的模型预测,连续提议若干个令牌更新。这些提议并非随机猜测,而是来自模型自身的输出分布,因此具有较高的接受潜力。其次,它并行评估这些提议的中间状态,而不是逐个顺序评估。最后,它应用一个精确的接受-校正(acceptance-and-correction)程序,确保最终采样的结果严格服从原始连续时间马尔可夫链的联合分布。
这一过程的关键在于,并行评估和校正步骤都不会引入时间离散化误差或额外的因子分解误差。与传统的推测解码(speculative decoding)不同,SpecCTMC 不需要一个单独的、较小的草稿模型来生成候选序列;它直接利用主模型自身的预测来提议更新,因此是一种无需训练的采样加速方案。
推测解码(speculative decoding) 原本是自回归语言模型中的加速技术:用一个小的草稿模型快速生成若干候选令牌,再由大模型并行验证并接受或拒绝,从而在保持目标分布不变的前提下减少大模型的前向次数。SpecCTMC 借用了“提议—验证—校正”的骨架,但关键差异在于草稿来源:它不训练独立草稿模型,而是复用主扩散模型在先前步骤中已经算出的预测分布作为提议分布。这消除了草稿模型与目标模型之间的分布不匹配问题,也省去了额外的训练和显存开销。同时,由于扩散模型的预测天然带有时间维度的信息,提议可以沿 CTMC 的事件时间轴连续展开,而非像自回归推测那样只能逐 token 向右扩展。
理论保证与实验验证
SpecCTMC 的作者证明,该采样器保持了底层连续时间马尔可夫链的事件时间、生成顺序和令牌取值的联合分布。这意味着,从统计意义上讲,SpecCTMC 生成的样本与直接连续时间采样生成的样本不可区分。这一性质对于需要精确分布保证的应用尤为重要。
实验在 LLaDA-8B-Instruct 和 Dream-v0-Instruct-7B 两个模型上进行,覆盖五个基准测试。结果显示,与直接连续时间采样相比,SpecCTMC 将顺序模型评估次数减少了 1.73 至 4.38 倍,同时保持了参考采样分布。这一加速幅度在无需额外训练的情况下相当可观,表明利用模型自身预测进行推测采样是一条有效的路径。
接受-校正(acceptance-and-correction)程序 是保证分布精确性的核心机制。其基本流程是:从提议分布中采样一个候选更新,以一定概率接受它,若拒绝则从校正分布中重新采样。接受概率的设计使得“提议—接受/拒绝—重采样”的复合转移概率恰好等于目标 CTMC 的真实转移概率,这一构造类似 Metropolis-Hastings 的接受准则。SpecCTMC 的贡献在于把该准则推广到连续时间、多令牌并行的场景,并证明校正后的样本在事件时间、生成顺序和令牌取值三个维度上都与直接 CTMC 采样同分布。换句话说,加速不来自近似,而来自把“串行评估”改写成“并行提议加精确拒绝”,从而在数学上不损失任何分布保真度。
对扩散语言模型生态的意义
SpecCTMC 解决了一个长期困扰扩散语言模型的矛盾:并行生成带来的速度优势与精确采样所需的顺序评估之间的冲突。通过消除这一瓶颈,它使得扩散语言模型在需要严格概率保证的场景中更具实用性。
有意思的是,SpecCTMC 的无需训练特性降低了采用门槛。研究者和开发者可以直接将其应用于现有的扩散语言模型,而无需重新训练或微调。代码已在 GitHub 上开源,进一步促进了可复现性和后续研究。
从更广的视角看,SpecCTMC 代表了扩散语言模型推理优化中的一个方向:不改变模型本身,而是通过更聪明的采样算法来释放并行潜力。随着扩散语言模型在规模和质量上的持续进步,这类精确且高效的采样方法将成为其落地应用的关键基础设施。
局限与未来方向
SpecCTMC 当前的方法仍有一些限制。它依赖于模型预测的质量来生成高接受率的提议;如果模型在某些上下文下预测不确定,接受率可能下降,加速效果也会减弱。此外,实验仅在两个模型和五个基准上验证,更广泛的模型架构和任务类型上的表现仍有待探索。
未来的工作可以探索如何将 SpecCTMC 与更高效的模型评估策略结合,例如缓存中间状态或动态调整提议长度。同时,将这一方法扩展到其他类型的扩散模型(如连续扩散模型)也是一个有趣的方向。
总体而言,SpecCTMC 为扩散语言模型的精确采样提供了一种实用且理论严谨的加速方案,为并行文本生成的分布保真度问题给出了新的答案。
LLaDA 与 Dream 是当前最具代表性的开源掩码扩散语言模型。LLaDA-8B 采用大规模掩码扩散预训练,在指令跟随任务上接近同规模自回归模型;Dream-v0-Instruct-7B 则探索了扩散模型在对话场景中的指令微调路径。两者都通过迭代式掩码预测生成文本,天然适合作为 SpecCTMC 的测试平台。选择这两个模型也说明该方法的适用面:它不依赖特定架构,只要模型能输出每个位置的令牌分布并支持中间状态评估即可。不过,不同模型的预测置信度分布差异较大,会直接影响提议接受率,这也是作者将接受率与预测质量挂钩、并呼吁在更广模型上验证的原因。
相关推荐
科技前沿Caro揭秘:LBJ如何用伪装骗过种族隔离派
罗伯特·卡罗在LBJ传记中记录前参议员塔尔梅奇的采访:约翰逊初入参议院时刻意迎合南方保守派,塔尔梅奇称未察觉其民权倾向。掌权后约翰逊推动1964年民权法案与1965年投票权法案,揭示政治伪装与改革时机的关系。

Commodore 64经典赛车游戏迎来Driver AI
探索Commodore 64经典赛车游戏《Stunt Car Racer》中的Driver AI实验,了解在资源受限的8位平台上实现自动驾驶AI的技术挑战与启示。

黄仁勋若被替换,开源AI会否生变?Reddit激辩
Reddit网友激辩:如果黄仁勋被替换,英伟达会反对开源AI吗?讨论指出英伟达支持开源AI源于商业利益,需拓宽GPU客户群,而非一人决策。