← 返回 2026-01-27

扩散中的扩散:在半自回归扩散中恢复全局一致性 Diffusion In Diffusion: Reclaiming Global Coherence in Semi-Autoregressive Diffusion

Linrui Ma, Yufei Cui, Kai Han, Yunhe Wang 📅 2026-01-20 👍 7 2026-07-13 08:35
半自回归生成 多阶段生成 文本生成 离散扩散模型 语言模型

通过多阶段草稿-修订框架恢复块扩散模型的全局一致性

前置知识

离散扩散语言模型

离散扩散模型是一类基于扩散过程的生成模型,将传统连续高斯扩散推广到离散的token空间。其前向过程通过掩码(masking)或转移矩阵逐步将干净文本转化为噪声(特殊[MASK]标记),反向过程则训练神经网络从噪声中逐步恢复原始文本。关键代表包括MDLM和SEDD等模型,它们通过连续时间步的离散扩散过程实现了与自回归模型可比的语言建模困惑度。这类模型的核心优势在于全局双向注意力机制,能够在生成时同时利用左侧和右侧上下文,实现真正的全局规划和文本填充。

本文的核心贡献是在保持离散扩散模型全局能力的同时引入高效的块结构,因此理解离散扩散的基本框架是理解本文方法的前提。

块扩散模型(Block Diffusion / BD3-LM)

块扩散模型是一种半自回归架构,将整个序列划分为固定大小的块(block),块间采用自回归顺序生成,块内则并行进行双向扩散去噪。这种设计允许使用KV Cache进行推理加速,将推理复杂度从全局扩散的$O(L^2)$降低到线性级别。BD3-LM(Block Discrete Denoising Diffusion Model)是该范式的代表工作,通过将每个块视为一个'元token',在保持局部双向生成能力的同时利用自回归先验提升性能。

本文提出的Diffusion in Diffusion方法直接基于BD3-LM框架进行改进,理解块扩散的工作机制和局限性是理解本文动机的关键。

KV Cache(键值缓存)

KV Cache是现代Transformer语言模型中广泛使用的推理加速技术。在自回归生成过程中,每个新token的生成需要对所有之前的token计算注意力,KV Cache将之前计算过的Key和Value矩阵缓存起来,避免重复计算,使每步推理的计算量从$O(n^2)$降低到$O(n)$。全局扩散模型由于不存在因果顺序,无法利用KV Cache,这是其推理效率低下的根本原因。

块扩散模型的核心优势之一就是能在块间使用KV Cache,而全局扩散模型无法使用。本文试图在保持全局能力的同时不牺牲这一效率优势,理解KV Cache的作用是理解这种权衡的基础。

生成困惑度(Generative Perplexity, Gen PPL)

生成困惑度是评估语言模型生成质量的核心指标,通过一个预训练的外部模型(如GPT-2 Large)来评估模型生成文本的流畅性和连贯性。与在测试集上计算的静态困惑度不同,Gen PPL评估的是模型无条件采样生成的实际文本质量,更能反映模型的开放域生成能力。较低的Gen PPL表示生成文本更流畅、更连贯、更接近自然语言分布。

本文的核心实验结果以Gen PPL作为主要评估指标,模型在OpenWebText上将Gen PPL从25.7降至21.9,理解该指标是理解本文性能提升幅度的基础。

掩码比率(Masking Ratio)与重掩码(Remasking)

掩码比率$\gamma$控制扩散过程中将多少比例的token替换为[MASK]标记。在前向扩散中,较大的掩码比率意味着更多的token被噪声化;在生成过程中,掩码比率影响生成的多样性和质量。重掩码机制是本文的创新之一,指在多阶段生成的中间阶段,将已生成的某些token重新转化为[MASK]状态,让模型在更大上下文窗口下重新生成这些token,从而修正早期生成中的错误。

本文提出的快照置信度重掩码策略是核心创新之一,理解掩码和重掩码机制是理解该策略工作原理的基础。

研究动机

当前离散扩散语言模型研究分为两大范式:全局扩散模型和块扩散模型。全局扩散模型拥有强大的全局规划能力和完整的双向上下文理解,但其$O(L^2)$的推理复杂度和无法使用KV Cache的特性导致推理极慢,严重限制了长序列生成的应用场景。块扩散模型(如BD3-LM)通过引入自回归先验实现了块间KV Cache加速和线性推理复杂度,但其宏观层面的自回归范式带来了'短视'问题——块的生成严格依赖之前块的输出,无法回头修改已生成的内容,本质上退化为局部模型。论文指出,已有块扩散研究在序列长度$L=1024$时Gen PPL为25.7,而自回归模型仅为14.1,两者存在显著性能差距。这种速度-质量-全局一致性的'不可能三角'使得现有方法难以在保持推理效率的同时获得全局规划能力。

本文的目标是本文的核心目标是在不牺牲半自回归扩散模型推理效率优势的前提下,将全局双向上下文能力重新引入块扩散模型框架。具体而言,作者希望设计一种新的生成框架,使得模型既能利用小块进行快速草稿生成(保留KV Cache加速),又能利用大块甚至全序列长度的双向感受野进行全局修订,最终实现:(1)在OpenWebText数据集上建立离散扩散模型的新SOTA;(2)大幅缩小与自回归模型的Gen PPL差距;(3)在仅使用基线模型26%微调预算的条件下完成上述目标,证明方法的数据效率。

与已有工作不同的是,现有研究通常将全局规划能力和推理效率视为不可调和的对立面:更强的全局能力必然牺牲推理速度,更好的自回归先验和更快的推理必然牺牲全局视角。本文的独特切入点在于提出'扩散中的扩散'这一多阶段框架,打破了这种二元对立。核心洞察是:人类写作过程本身就是'快速起草+全局修订'的两阶段过程,而现有的块扩散模型只有'起草'阶段,缺失了'修订'阶段。因此,本文不是在现有框架内做权衡优化,而是引入了一个全新的生成范式——通过渐进式增加块大小来扩大双向感受野,通过快照置信度机制识别需要修订的token,在保持块扩散框架效率优势的同时恢复全局修订能力。这种方法本质上是在块扩散框架内嵌套了一个全局扩散过程,形成了'扩散中的扩散'结构。

核心方法

本文提出的方法称为DIFFUSION IN DIFFUSION(也称为结构化块扩散),其整体思路受人类写作过程启发:先快速写出草稿,再进行全局修订以确保连贯性。技术路线采用多阶段渐进式生成:第一阶段使用小块(如$B^{(1)}=4$)进行快速的半自回归扩散生成,专注于局部语法和内容,产出'草稿'序列;后续阶段逐步增大块大小($B^{(1)} < B^{(2)} < \cdots < B^{(K)} = L$),利用更大的双向感受野对前一阶段的输出进行全局修订。在阶段之间,通过快照置信度重掩码机制识别最需要修改的token,将其重置为[MASK]状态,然后在更大上下文窗口下重新生成。当最终阶段的块大小等于序列长度$L$时,该阶段退化为全序列掩码扩散过程(类似MDLM),但拥有来自第一阶段的结构化先验和部分初始化。这种'草稿-修订'范式本质上是在高效块扩散框架内嵌套全局扩散过程,实现了速度与质量的统一。

本文的核心创新点有三个,与已有方法存在本质区别。第一是多阶段结构化扩散框架:不同于标准块扩散模型只进行一次单一块大小的生成,本文通过渐进式增加块大小实现'先局部后全局'的生成策略,这与Speculative Decoding等方法追求推理加速的目标不同,本文追求的是生成质量的提升。第二是快照置信度重掩码(Snapshot Confidence Remasking):不同于事后置信度评估(对生成序列重新做forward pass获取likelihood),本文在每个token生成时记录其被采样瞬间的模型预测概率$p_\theta(\hat{x}_i^{(k-1)} | \mathbf{x}_t^b, \mathbf{x}_{<b})$作为置信度信号。这种'快照'方式捕捉了生成过程中的动态不确定性,而事后评估容易产生'过度自信'问题——模型倾向于肯定自己之前的幻觉输出,无法提供有意义的修订指导。第三是混合尺度训练策略(Mix-Scale Training):不同于标准训练固定单一块大小,本文将块大小作为随机变量从双模分布中采样,90%使用小块训练草稿能力,10%使用大块训练修订能力,使模型同时具备两种能力。

方法步骤详情

完整的多阶段生成流程如下:首先将序列初始化为全[MASK]状态$\mathbf{x} \leftarrow \mathbf{m}^L$,置信度轨迹$S \leftarrow \mathbf{0}^L$。对于每个阶段$k=1,\ldots,K$:(1)如果$k>1$,执行阶段间重掩码:根据置信度轨迹$S$升序排列,选择置信度最低的$\lfloor\gamma^{(k)} L\rfloor$个token的索引$I$,将这些位置重置为[MASK]状态;(2)初始化KV Cache为空;(3)将序列划分为$B^{(k)}$个块,对每个块$b=1,\ldots,B^{(k)}$执行:使用块扩散采样生成当前块$x_b$及其置信度$S_b$,同时获取该块的Key和Value缓存$K_b, V_b$,将生成结果合并回序列,将KV缓存累积。在第一阶段,块大小较小(如$B^{(1)}=4$),模型快速逐块生成草稿;在后续阶段,块大小增大,模型在保留高置信度token的同时,对低置信度token进行全局双向修订。最终阶段当$B^{(K)}=L$时,整个序列作为一个块进行全序列扩散。整个过程的计算开销约为基线方法的1.1-3倍NFEs(函数评估次数),可通过掩码比率$\gamma$灵活控制。

技术新颖性

本文的技术新颖性体现在多个层面。首先,'扩散中的扩散'框架本身是一个全新的生成范式,不同于Speculative Decoding(加速推理)、Levenshtein Transformer(迭代编辑)等方法,本文是在扩散模型框架内进行多阶段渐进式生成,保留了扩散模型的核心优势(双向注意力、灵活控制)。其次,快照置信度的概念是原创的,利用了扩散模型天然的不确定性信号——每个token在被'决定'时的模型预测概率,这比自回归模型中的token概率更能反映生成过程的动态质量。第三,混合尺度训练的双模分布设计(90%小块+10%大块)揭示了一个重要发现:少量的大块训练样本(仅10%)就足以解锁模型的全局修订能力,而不会损害其局部生成质量。实验表明,仅使用基线模型26%的微调预算(40K步vs 150K步)就实现了显著的性能提升,证明了该方法的极端数据效率。此外,方法提供了灵活的帕累托前沿——用户可以通过调整掩码比率$\gamma$在质量和效率之间进行权衡,这是单一阶段模型无法提供的。

DIFFUSION IN DIFFUSION方法概览
Figure 1: DIFFUSION IN DIFFUSION方法概览

实验结果

本文在OpenWebText数据集上的实验结果表明,Diffusion in Diffusion方法在多个维度上取得了显著突破。在生成质量方面,使用110M参数的Transformer骨干网络,在序列长度$L=1024$时,单阶段草稿模型的Gen PPL为27.4(略高于BD3-LM的25.7,因为仅用26%的微调预算),但经过第二阶段全局修订后,Gen PPL大幅降至21.9,实现了约20%的相对提升。在$L=2048$的更长序列上,方法达到Gen PPL 20.6,而基线BD3-LM的最佳结果为22.8,优势更加明显。在计算效率方面,方法展现了灵活的帕累托前沿:最小开销下(1.1K NFEs,仅比基线增加10%计算量)达到PPL 24.6,已超过标准BD3-LM的25.7;等计算量比较(1.5K NFEs)下,方法PPL为21.9而BD3-LM仅为25.0,证明性能提升源于框架设计而非单纯增加采样步数。消融实验进一步验证了:(1)全局上下文是必要的——第二阶段块大小$B^{(2)}$从4增大到1024时,困惑度单调下降;(2)掩码比率$\gamma$呈现U型最优曲线,在0.25-0.5之间效果最佳;(3)快照置信度重掩码策略显著优于随机掩码(PPL 30.26)和事后置信度(PPL 29.85),仅快照置信度能改善基线(PPL 21.85 vs 27.36);(4)混合尺度训练是必要的——仅用小块训练的模型在修订阶段PPL从27.95退化到31.97,双模混合训练($B \in \{4, 1024\}$)效果最佳。

无条件生成的生成困惑度和函数评估次数
Table 1: 无条件生成的生成困惑度和函数评估次数
重掩码策略消融实验
Table 2: 重掩码策略消融实验
训练混合策略消融实验
Table 3: 训练混合策略消融实验
修订范围消融实验:不同块大小和修订比率下的生成困惑度
Figure 2: 修订范围消融实验:不同块大小和修订比率下的生成困惑度
查看结构化数据
任务指标本文基线提升
无条件文本生成(L=1024) 生成困惑度(Gen PPL)↓ 21.9(完整两阶段,1.5K NFEs) BD3-LM (L'=4): 25.7 Gen PPL降低3.8点(相对提升约15%),仅用26%微调预算
无条件文本生成(L=2048) 生成困惑度(Gen PPL)↓ 20.6(完整两阶段,3.0K NFEs) BD3-LM (L'=4): 23.6 Gen PPL降低3.0点(相对提升约13%)
等计算量对比(L=1024) 生成困惑度(Gen PPL)↓ 21.9(1.5K NFEs) BD3-LM (L'=4): 25.0(1.5K NFEs) 相同NFEs下Gen PPL降低3.1点
修订阶段消融(不同掩码策略) 生成困惑度(Gen PPL)↓ 21.85(快照置信度,$B^{(2)}=1024, \gamma=0.5$) Stage 1输出: 27.36 相对提升20.2%,随机掩码和事后置信度均无法改善基线

局限与改进

尽管本文取得了显著的实验结果,但仍存在多方面局限性。首先,实验规模仅限于110M参数的模型,未验证方法在更大规模(如1B、7B)模型上的效果,而当前语言模型研究的主流已转向数十亿参数规模。其次,评估仅在OpenWebText数据集上进行,未在更广泛的任务(如摘要、翻译、问答)上测试生成质量,Gen PPL虽然是标准指标但不能完全反映生成文本的实际可用性。第三,多阶段生成引入了额外的计算开销(1.1-3倍NFEs),虽然仍远低于全局扩散模型,但相比单阶段块扩散增加了复杂度。第四,方法需要仔细调节两个关键超参数——第二阶段块大小$B^{(2)}$和掩码比率$\gamma$,论文中的最优值($B^{(2)}=1024, \gamma=0.25-0.5$)可能不适用于不同序列长度或数据分布。第五,论文未讨论方法在条件生成(如文本填充、续写)场景下的表现,而这是实际应用中最常见的需求。此外,快照置信度机制要求在采样过程中记录每个token的预测概率,这增加了内存开销和实现复杂度。

独立分析的弱点

本文的一个主要弱点是实验规模有限。110M参数的模型虽然符合当前离散扩散模型研究的标准配置,但无法验证方法在更实际的大规模场景中的效果。改进方向包括在1B-7B规模的模型上进行验证,并与最新的大规模扩散语言模型(如MDLM-7B)进行对比。另一个弱点是评估维度单一——仅使用Gen PPL作为生成质量指标,缺乏人工评估、下游任务评估(如分类、摘要质量)和多样性指标(如distinct-n)。建议增加人类评估实验和多维度自动评估。此外,混合尺度训练中的双模分布参数($\lambda=0.1$)是经验值,缺乏理论分析或自动搜索机制,可以探索基于课程学习或自适应的块大小调度策略。快照置信度虽然优于事后置信度,但仍可能受到采样温度和解码策略的影响,可以探索更鲁棒的不确定性估计方法,如集成多个扩散步的预测或使用贝叶斯方法。最后,论文未充分讨论多阶段修订的收敛性——当阶段数$K>2$时,是否会出现过度修订导致质量下降的情况需要进一步分析。

未来方向

基于本文的框架,未来研究可以在多个方向展开。首先是规模化验证:将方法应用于更大规模的预训练扩散语言模型,特别是在代码生成、数学推理等需要全局规划能力的任务上评估。其次是扩展到条件生成场景:本文主要关注无条件生成,可以将'草稿-修订'范式应用于文本填充(infilling)、文本编辑、摘要生成等条件生成任务,利用全局修订能力提升条件一致性。第三是动态阶段调度:当前方法使用固定的阶段数和块大小序列,可以设计自适应的阶段调度策略,根据生成内容的复杂度动态决定何时需要修订以及修订的粒度。第四是与高效注意力机制的结合:可以将Flash Attention、线性注意力等高效注意力技术融入框架,进一步降低多阶段生成的计算开销。第五是探索更多重掩码策略:当前的快照置信度基于单步预测概率,可以探索基于信息论(如互信息)或梯度信息的token重要性评估方法。第六是与自回归模型的混合架构:探索将本文的全局修订能力引入到大规模自回归语言模型的后处理阶段,作为'质量增强器'使用。

复现评估

本文的复现可行性较高。论文明确说明了实验配置:使用110M参数的Transformer(12层、768隐藏维度),基于Arriola et al. (2025)的官方预训练检查点进行微调,在OpenWebText数据集上进行评估。训练细节清晰:使用40K梯度步(467B tokens)进行微调,采用方差减少技术(自适应学习掩码率范围参数$\beta, \omega$),所有超参数跟随原始实现。方法的核心代码改动相对有限——主要是多阶段生成循环、快照置信度记录逻辑和混合尺度训练的数据加载器。然而,完全复现仍面临一些挑战:(1)需要BD3-LM的预训练检查点,如果官方未开源则需要从头预训练(850K步);(2)快照置信度的记录需要修改采样过程,在去噪步骤中记录每个token被决定时的模型预测概率;(3)评估使用GPT-2 Large计算Gen PPL,需要确保评估代码和模型版本一致。总体而言,如果基座模型和数据可用,复现难度为中等,预计一个有经验的研究团队在2-4周内可以完成。