← 返回 2026-03-24

基于快照的广义离散扩散模型 Generalized Discrete Diffusion from Snapshots

Oussama Zekri, Théo Uscidda, Nicolas Boullé, Anna Korba 📅 2026-03-22 👍 11 2026-07-13 08:36
扩散模型 注意力机制 生成式模型 离散扩散 语言建模
arXiv HuggingFace PDF

统一框架支持任意离散扩散，首次大规模超越自回归

前置知识

连续时间马尔可夫链

连续时间马尔可夫链是一种随机过程，其状态随时间连续变化，且未来状态仅依赖于当前状态。它由率矩阵Q刻画，其中Q(i,j)表示从状态i跳转到状态j的瞬时速率。对角线元素Q(i,i)保证概率质量守恒。演化遵循Kolmogorov前向方程dq_t/dt = Q_t q_t，转移概率为q_t(x_t|x_0) = K_t(., x_0)，其中K_t是时间序指数。本文用CTMC描述离散扩散的前向噪声过程。

CTMC是理解GDDS方法论的核心。任何率矩阵Q_t都可以通过统一框架建模，这是GDDS支持任意噪声过程的数学基础。

ELBO证据下界

证据下界是通过变分推断得到的对数似然下界。对于潜在变量模型，ELBO为log p_theta(x_0)大于等于期望log p_theta(x_0, omega)减去期望log q(omega|x_0)。在扩散模型中，潜在变量omega通常是完整的噪声路径。本文改用快照s = (x_t, t)作为潜在变量，得到更简化的目标函数：积分从0到1，期望负log mu_theta(x_t, t)[x_0] dt。

ELBO是扩散模型的标准训练目标。本文推导的基于快照的ELBO使训练对任意噪声过程都高效可行，这是GDDS训练的核心创新。

均匀化Uniformization

均匀化是一种精确表示CTMC转移的技术，通过泊松过程分解跳转。给定率矩阵Q_t = f(t)(F_t - I)，在时间t内的跳转数N_t服从强度f_bar(t)的泊松分布。条件在N_t = k时，跳转时间为有序的随机时间序列，每次跳转按F_{T_k}采样。转移矩阵为K_t = alpha_t I + (1-alpha_t) Pi_t。这使GDDS能够对任意率矩阵进行精确噪声采样，只需列访问而无需计算矩阵指数。

均匀化是实现高效前向噪声的关键技术。它使GDDS支持任意噪声过程，且在大词汇量（如50257）下仍然计算可行。

自回归模型

自回归模型通过条件概率链式法则对序列建模：p(x_1, ..., x_n)等于乘积p(x_i | x_<i)。训练时预测下一个token，生成时按顺序采样。GPT系列是典型代表。AR模型的优势是精确似然计算，缺点是生成速度慢且无法并行解码。本文将扩散模型与AR在相同计算预算下对比，首次证明扩散模型在此规模上能超越AR基线。

AR是语言建模的主导范式。本文在OWT上用262B token重新训练AR作为基线，比较结果证明GDDS的性能优势，这对评估扩散模型进展至关重要。

掩码扩散模型MDM

掩码扩散通过逐步用特殊token 替换原始token来添加噪声。前向核为q_t(x_t|x_0) = alpha_t delta_{x_t=x_0} + (1-alpha_t) delta_{x_t=}，其中alpha_t在0到1之间递减。逆向过程训练网络从部分掩码上下文预测被掩码的原始token。BERT等模型使用的MLM掩码语言建模是特例。MDM是离散扩散的主要方法之一，本文将其纳入GDDS框架并在相同设置下复现对比。

MDM是重要的离散扩散基线。本文证明GDDS Absorb在Text8上达到1.16 BPC，而重新训练的MDM为1.58 BPC，验证了GDDS框架的优越性。

研究动机

现有离散扩散模型存在两个核心瓶颈。首先，前向噪声过程对离散空间中的邻域概念如语言中的语义邻近性视而不见。大多数方法采用简单的token级污染规则，包括掩码或均匀替换，这种盲目噪声忽略了词汇之间的语义结构关系。例如将bank替换为river或money比替换为随机词汇更有意义，但现有方法无法利用这种先验。其次，均值参数化将去噪器同时用于控制链何时跳转和跳转到何处，这种耦合在LLM尺度下变得日益受限。在Text8数据集上，重新训练的AR基准达到1.35 BPC，MDM为1.58 BPC，UDLM为1.67 BPC，都显著落后于GDDS Absorb的1.16 BPC。

本文的目标是本文目标是提出一个真正通用的离散扩散框架，能够支持任意马尔可夫噪声过程，同时保持计算可扩展性。具体而言，GDDS旨在建立数学框架覆盖所有现有离散扩散方法，并允许语义感知的结构化噪声。开发高效的前向噪声算法，仅需列访问率矩阵即可精确采样，适用于大词汇量如GPT-2的50257。设计简洁的ELBO训练目标，直接从快照训练标准生成架构，避免完整路径的复杂依赖。

与已有工作不同的是，本文的独特切入点是将潜在变量从完整的噪声路径omega简化为单个快照s = (x_t, t)。这背后的洞察是强大的模型主要在快照噪声上下文上训练，为什么变分潜在变量必须是整个路径。这种选择与Li & He的观点一致，即去噪模型应该通过均值参数化预测干净量而非噪声量。关键权衡在于丢弃完整路径导致内在信息损失IPG大于等于0，但可以换取更好的校准，从而产生更强的生成模型。在OWT上，GDDS Uniform验证困惑度达到10.97，相比重新训练的UDLM的36.82提升70%以上，证明这种简化是有效的。

核心方法

GDDS的整体思路基于插值离散扩散框架。给定混合率alpha_t从1递减到0和列随机混合矩阵Pi_t，转移算子定义为K_t = alpha_t I + (1-alpha_t) Pi_t。这涵盖了掩码扩散和均匀扩散等现有方法。更一般的Pi_t选择允许结构和token相关的污染机制，如语义感知的SIK。逆向过程采用jump-states参数化，直接学习跳转目标j_theta(x_t, t)，而跳转强度由固定时间表r[x](t)控制。训练使用基于快照的ELBO，在快照上做交叉熵。网络是标准Transformer，时间嵌入通过AdaLN注入。推断时使用祖先采样，从完全噪声开始逐步逆向采样。

核心创新点在于三个方面。第一是广义插值公式K_t = alpha_t I + (1-alpha_t) Pi_t，数学上等价于任意CTMC的解，但更灵活且易于计算。第二是均匀化采样，通过泊松过程N_t和跳转核F_t精确采样x_t，只需列访问Q_t，使大词汇量下的精确采样可行。第三是jump-states参数化与快照ELBO，将传统均值参数化的耦合问题解耦，网络只预测跳转目标，跳转时间由固定调度控制，ELBO简化为单个交叉熵。这与传统方法本质不同，SEDD的ELBO包含线性得分项和交叉熵两项，均值参数化则包含theta相关的反向速率项。

方法步骤详情

GDDS的完整方法流程分为前向噪声和逆向学习。前向噪声给定清洁序列x_0和目标时间t，首先计算混合率alpha_t = exp(-f_bar(t))，然后采样泊松跳转数N_t服从Poisson(f_bar(t))分布。如果N_t等于0则x_t等于x_0，否则采样跳转时间0小于T_1小于...小于T_{N_t}小于等于t，对每个跳转执行z_k从F_{T_k}(., z_{k-1})采样，最终返回x_t等于z_{N_t}。序列级可并行执行所有位置。逆向学习训练时采样x_0从数据分布和t从均匀分布，用上述算法得到x_t，然后最小化负log mu_theta(x_t, t)[x_0]。网络是标准Transformer，时间嵌入通过AdaLN注入。推断时使用祖先采样，给定时间网格从x_{t_K}采样开始，逐步用插值贝叶斯核逆向采样。

技术新颖性

GDDS的技术新颖性体现在多个维度。数学上它是最初的完全通用且计算高效的离散扩散框架。广义插值公式理论上可表示任意率矩阵，而均匀化使大词汇量下的精确采样可行。训练上快照ELBO比路径目标更简洁，与标准架构自然对齐，避免了双流机制等复杂设计。实验上这是首次在同等计算预算下，离散扩散模型在语言建模上超越AR基线。在OWT上GDDS Gauss达到7.65困惑度，低于AR的20.49。语义感知噪声SIK是另一个新颖点，利用GPT-2嵌入计算高斯核，使噪声过程尊重语义相似性。这带来显著的零样本迁移增益，在PTB上GDDS Gauss困惑度53.65相比AR的147.90降低64%。

Overview of GDDS. A clean sequence x0 is first noised exactly by the forward CTMC at a sampled time t in 0 to 1, yielding a snapshot sequence (xt, t). The mean parametrization is then used as a denoiser: given the snapshot, the model predicts the clean-token posterior directly from (xt, t), so training is performed on snapshots rather than through a full path-wise objective.
Figure 2: Overview of GDDS. A clean sequence x0 is first noised exactly by the forward CTMC at a sampled time t in 0 to 1, yielding a snapshot sequence (xt, t). The mean parametrization is then used as a denoiser: given the snapshot, the model predicts the clean-token posterior directly from (xt, t), so training is performed on snapshots rather than through a full path-wise objective.
Snapshot vs. path-wise training. The forward process corrupts the clean sequence My name is David. The blue path shows the beginning of the noising trajectory omega of one tracked position position 4. Path-wise objectives condition on the entire trajectory omega, whereas our GDDS snapshot objective uses only one random-time observation s = (x_{t_star}, t_star).
Figure 3: Snapshot vs. path-wise training. The forward process corrupts the clean sequence My name is David. The blue path shows the beginning of the noising trajectory omega of one tracked position position 4. Path-wise objectives condition on the entire trajectory omega, whereas our GDDS snapshot objective uses only one random-time observation s = (x_{t_star}, t_star).

实验结果

实验结果在三个层面验证了GDDS的有效性。在Text8字符级建模上训练1百万步，GDDS Absorb达到1.16 BPC，相比重新训练的AR的1.35提升14%，相比MDM的1.58提升27%，相比UDLM的1.67提升31%。这是首次离散扩散在此任务上超越AR。在OWT词元级建模上训练50万步262B token，GDDS Gauss验证困惑度7.65相比AR的20.49降低63%，GDDS Uniform的10.97和Absorb的8.98也显著优于对应基线UDLM的36.82和MDM的31.03。零样本迁移到7个下游数据集包括PTB、Wikitext103、LM1B、Lambada、AG News、Pubmed、Arxiv时，GDDS Gauss在所有数据集上取得最佳结果。生成质量多样性权衡上在64解码步时，GDDS Absorb达到与MDM在1024步时相同的熵约5.70但Gen-PPL更低约125相比150，采样效率提升16倍。Distinct-n指标显示GDDS Absorb与MDM相当优于UDLM。

Bits Per Character on Text8. Baseline results reported from prior work. All models are trained for 1M steps. Best results per model family are in bold, while best results among retrained models are underlined.
Table 1: Bits Per Character on Text8. Baseline results reported from prior work. All models are trained for 1M steps. Best results per model family are in bold, while best results among retrained models are underlined.
OWT validation perplexity. Validation perplexity on OWT. Best results are in bold. Trained on the WebText dataset. Results taken from other papers are marked with symbols.
Table 2: OWT validation perplexity. Validation perplexity on OWT. Best results are in bold. Trained on the WebText dataset. Results taken from other papers are marked with symbols.
Zero-shot transfer perplexity. Zero-shot perplexity of OWT-trained models. Models are evaluated on validation splits of 7 downstream datasets without additional fine-tuning.
Table 3: Zero-shot transfer perplexity. Zero-shot perplexity of OWT-trained models. Models are evaluated on validation splits of 7 downstream datasets without additional fine-tuning.
Lexical diversity. Distinct-1-2-3 computed on 256 unconditional samples from OWT-trained models, measuring the fraction of unique n-grams among generated texts.
Table 4: Lexical diversity. Distinct-1-2-3 computed on 256 unconditional samples from OWT-trained models, measuring the fraction of unique n-grams among generated texts.
GDDS Training configuration
Table 5: GDDS Training configuration
Noising-time benchmark. Mean wall-clock latency in milliseconds for sampling xt from qt given x0 on batches of size 512 and sequence length 1024 equal to 524288 positions.
Table 6: Noising-time benchmark. Mean wall-clock latency in milliseconds for sampling xt from qt given x0 on batches of size 512 and sequence length 1024 equal to 524288 positions.
GDDS-SIK sampling ablation on OpenWebText. We report decoding budget K average sequence entropy and Gen-PPL for unconditional samples. Natural OWT text typically lies around entropy 5.60 to 5.70.
Table 12: GDDS-SIK sampling ablation on OpenWebText. We report decoding budget K average sequence entropy and Gen-PPL for unconditional samples. Natural OWT text typically lies around entropy 5.60 to 5.70.
Validation perplexity of the two-stream architecture. We train the two-stream architecture on Text8 and OWT under the same experimental setup as in Section C and report the BPC and validation perplexity.
Table 13: Validation perplexity of the two-stream architecture. We train the two-stream architecture on Text8 and OWT under the same experimental setup as in Section C and report the BPC and validation perplexity.
Zero-shot transfer of OWT-trained models. Zero-shot perplexity on three representative downstream validation sets from Table 3: PTB, LM1B, and Wikitext.
Figure 1: Zero-shot transfer of OWT-trained models. Zero-shot perplexity on three representative downstream validation sets from Table 3: PTB, LM1B, and Wikitext.
OWT training curves. Evolution of OWT validation perplexity during training for the retrained models reported in Table 2. This complements the final numbers in Table 2 by showing the full optimization trajectory, both axes are shown on logarithmic scales.
Figure 4: OWT training curves. Evolution of OWT validation perplexity during training for the retrained models reported in Table 2. This complements the final numbers in Table 2 by showing the full optimization trajectory, both axes are shown on logarithmic scales.
Generation quality-diversity tradeoff. Gen-PPL versus Entropy for K in 32 64 128 256 512 1024 decoding steps. Bubble radius increases with K. For reference, the AR baseline achieves Gen-PPL 56.82 at entropy 5.60.
Figure 5: Generation quality-diversity tradeoff. Gen-PPL versus Entropy for K in 32 64 128 256 512 1024 decoding steps. Bubble radius increases with K. For reference, the AR baseline achieves Gen-PPL 56.82 at entropy 5.60.
Training loss stability on Text8. Train loss curves for AR, MDM, GDDS Absorb, and Campbell two-stream training, using the same setup as in Section C.
Figure 6: Training loss stability on Text8. Train loss curves for AR, MDM, GDDS Absorb, and Campbell two-stream training, using the same setup as in Section C.
查看结构化数据
任务指标本文基线提升
Text8 BPC Bits Per Character越低越好 GDDS Absorb: 1.16 AR: 1.35, MDM: 1.58, UDLM: 1.67 较AR提升14%，较MDM提升27%
OWT Validation PPL Perplexity越低越好 GDDS Gauss: 7.65, GDDS Absorb: 8.98, GDDS Uniform: 10.97 AR: 20.49, MDM: 31.03, UDLM: 36.82 Gauss较AR降低63%，Absorb较MDM降低71%
Zero-shot PTB PPL Perplexity越低越好 GDDS Gauss: 53.65, GDDS Absorb: 103.04, GDDS Uniform: 115.12 AR: 147.90, MDM: 181.36, UDLM: 177.26 Gauss较AR降低64%，Absorb较MDM降低43%
Zero-shot Wikitext103 PPL Perplexity越低越好 GDDS Gauss: 34.56, GDDS Absorb: 46.49, GDDS Uniform: 42.63 AR: 42.91, MDM: 45.42, UDLM: 64.65 Gauss较AR降低19%，Absorb与MDM相当
Gen-PPL at K=64 Generative Perplexity越低越好 GDDS Absorb: 约125 MDM在K=1024时: 约150, UDLM: 约275 较MDM降低17%，采样步数少16倍

局限与改进

作者承认的主要局限性是语义感知核SIK的采样困难。虽然GDDS-SIK模型能达到很强的验证损失，在OWT上PPL约7.65，但祖先采样需要近似时间序指数，通过均匀化缓存的矩阵向量积实现。每个逆向步骤需要截断泊松级数并执行多次矩阵向量乘，比uniform和absorbing的直接采样慢得多。更重要的是近似误差会沿轨迹累积。另一个观察到的局限是快照ELBO的训练损失波动较大。虽然它更简洁，但每序列仅提供单一监督信号，导致训练曲线波动。Campbell估计器虽然损失更稳定，但在架构约束下明显表现不佳。这说明稳定性不一定转化为更好的似然性能。

独立分析的弱点

GDDS存在几个潜在弱点值得改进。首先语义感知核的采样是当前的主要瓶颈。祖先采样对SIK需要近似K_t和K_{t,s}，这在KNN实现中通过稀疏图实现，在KeOps实现中通过延迟块约简实现。SIK Gauss噪声时间约9毫秒相比Absorbing的0.09毫秒慢约100倍。第二个弱点是快照ELBO的高方差。虽然它更简洁，但每序列仅提供单一监督信号，导致训练曲线波动。第三个潜在弱点是时间调度的选择需要启发式设计。虽然固定跳转强度简化了训练，但最优调度可能因任务和数据而异。最后SIK的带宽参数和温度调度需要手动调节，缺乏自适应机制。改进方向包括开发自适应采样程序，为SIK设计直接去噪采样器，探索学习的时间表，自动调优SIK的带宽和温度。

未来方向

作者提出的未来方向包括设计更多样的语义感知核，基于词汇相似性强制有意义的噪声过程。更广泛的未来工作可以延伸到几个方向。一是自适应采样，结合置信度或不确定性的动态时间表，在高不确定性区域分配更多步，在低不确定性区域跳过。这符合近期证据表明自适应调度在扩散语言模型中可超越标准祖先解码。二是SIK专用采样器，开发不依赖精确前向转移算子的采样方法，使GDDS-SIK在训练和解码时都对K_t盲目。三是多模态扩展，GDDS框架可应用于其他离散数据如图分子，利用特定领域的相似性度量构建SIK。四是理论分析，深入研究快照与路径ELBO的信息校准权衡。五是大规模实验，在更大模型和数据集如GPT-3规模上验证GDDS的可扩展性。六是下游任务微调，探索GDDS预训练模型的微调策略。

复现评估

论文提供了良好的复现支持。代码和博客文章在项目页面公开。实验设置详细包括序列长度Text8为256OWT为1024、隐藏层大小768、注意力头数12、层数12、优化器AdamW学习率3.5乘10的负4次方、批大小512、bf16精度、梯度裁剪1.0、EMA衰减0.9999、最大步数Text8为1百万OWT为50万。硬件为单节点4块NVIDIA H100 GPU。数据集处理明确OWT使用GPT-2分词器，拼接文档并切分为1024序列。评估指标定义清晰包括BPC、PPL、Gen-PPL、序列熵、Distinct-n。基线复现说明完整，重新训练AR MDM UDLM使用相同训练配方，强调跨论文对比对实现细节敏感。SIK实现有两种包括KNN 64邻居和KeOps密集延迟计算，Table 6报告噪声时间基准。主要局限性是双流架构和SIK采样的复杂性需要额外实现努力。总体而言复现难度中等，需要理解CTMC理论和均匀化，但提供了足够细节和公开代码。

论文图表

figure_1..jpg
figure_1..jpg
figure_2..jpg
figure_2..jpg
figure_3..jpg
figure_3..jpg
figure_4..jpg
figure_4..jpg
figure_5..jpg
figure_5..jpg
figure_6..jpg
figure_6..jpg
table_1..jpg
table_1..jpg
table_11..jpg
table_11..jpg
table_12..jpg
table_12..jpg
table_13..jpg
table_13..jpg
table_2..jpg
table_2..jpg
table_3..jpg
table_3..jpg
table_4..jpg
table_4..jpg
table_5..jpg
table_5..jpg
table_6..jpg
table_6..jpg
table_7..jpg
table_7..jpg
table_8..jpg
table_8..jpg
table_9..jpg
table_9..jpg