← 返回 2026-02-04

在离散扩散模型中平衡理解与生成能力 Balancing Understanding and Generation in Discrete Diffusion Models

Yue Liu, Yuzhong Zhao, Zheyong Xie, Qixiang Ye, Jianbin Jiao, Yao Hu, Shaosheng Cao, Yunfan Liu 📅 2026-02-01 👍 17 2026-07-13 08:35
噪声核设计 图像生成 大模型扩展 离散扩散模型 语言模型

通过平稳噪声核统一掩码与均匀噪声扩散,实现理解-生成帕累托前沿突破

前置知识

离散去噪扩散概率模型(D3PM)

D3PM 是将连续扩散模型推广到离散状态空间的框架。它定义了两个马尔可夫过程:前向过程通过行随机矩阵 $Q_{t|s}$ 逐步将干净数据 $x_0$(以 one-hot 向量表示)腐蚀为含噪状态 $z_t$,反向过程则通过神经网络 $f_θ$ 预测干净数据分布 $\tilde{x}_0$ 来迭代恢复原始数据。训练目标是变分下界(VLB/ELBO),在连续时间极限下简化为 KL 散度的加权积分。D3PM 的核心灵活性在于可以自由选择转移矩阵的结构,从而衍生出不同的噪声范式。

XDLM 的理论框架直接建立在 D3PM 之上,理解 D3PM 的前向-反向过程和训练目标是理解 XDLM 如何统一不同范式的基础。

掩码扩散语言模型(MDLM)

MDLM 采用吸收态转移(absorbing-state transition)作为噪声核,即以特殊 [MASK] token 逐步替换原始 token。其噪声核矩阵 $M$ 是一个吸收矩阵,将所有状态映射到 [MASK] 状态。前向过程 $Q_{t|s} = \alpha_{t|s}I + \beta_{t|s}M$ 随时间推进逐步增加掩码比例。训练目标简化为仅在被掩码位置上的交叉熵损失。MDLM 在似然建模和零样本泛化方面表现优异,但少步生成质量较差,因为模型只能从 [MASK] 状态向外填充,无法修正已生成的 token。

MDLM 是 XDLM 的一个重要基线和特例($k=0$ 时退化为 MDLM),理解其优势(语义理解)和劣势(少步生成差)是理解 XDLM 动机的关键。

均匀噪声扩散语言模型(UDLM)

UDLM 的噪声核使用均匀分布 $J/N$($J$ 为全 1 矩阵,$N$ 为词表大小),即前向过程以概率将任何 token 替换为词表中随机选择的 token。这种噪声结构使得模型在推理时可以从任意 token 状态互相转移,而不仅仅是从 [MASK] 填充。UDLM 在少步生成(如 4-32 步)时质量显著优于 MDLM,但在多步推理和零样本泛化方面落后于 MDLM,因为它缺乏 MDLM 那种结构性的吸收态先验。

UDLM 是 XDLM 的另一个特例($k=1$ 时退化为 UDLM),其少步生成优势和理解能力劣势构成了 XDLM 要解决的核心矛盾。

平稳噪声核(Stationary Noise Kernel)

平稳噪声核是指噪声核矩阵 $K$ 在所有时间步上保持不变,即 $K$ 不依赖于时间 $t$。这与 GIDD 等方法中使用的时间非齐次(time-inhomogeneous)过程形成对比,后者在每个时间步上使用不同的噪声分布。平稳性使得增量噪声的结构与边际噪声的结构一致,从而允许训练目标被高效分解为静态常数和动态调度的乘积,避免了 GIDD 中需要在每个时间步重新计算噪声分布的开销。

这是 XDLM 的核心设计原则,平稳性使得数学推导和高效实现成为可能,是区别于 GIDD 的关键所在。

Classifier-Free Guidance(CFG)

CFG 是条件生成中常用的推理时增强技术。它同时训练条件模型和无条件模型,在采样时将两者的输出按比例混合:$\hat{\epsilon} = \epsilon_\theta(z_t, c) + s \cdot (\epsilon_\theta(z_t, c) - \epsilon_\theta(z_t, \emptyset))$,其中 $s$ 为引导尺度。CFG 可以显著提升生成质量和条件一致性,但在图像生成中不同模型对 CFG 的敏感度不同,这影响了不同方法在有无 CFG 设置下的相对表现。

论文在图像生成实验中分别报告了无 CFG 和有 CFG(scale=2.0)的结果,两种设置下各模型的排名和最优 $k$ 值不同,理解 CFG 对正确解读实验结果至关重要。

研究动机

在离散扩散模型领域,MDLM 和 UDLM 代表两种截然不同的噪声范式,各自在特定维度上表现优异但存在显著短板。MDLM 采用吸收态掩码机制,在似然建模和零样本泛化方面表现突出(如在 OWT 验证集上 PPL 为 23.321),但在少步生成场景下表现极差——在 OWT 上仅 8 步采样时 PPL 高达 711.382,远差于 UDLM 的 183.991。这在图像生成中尤为明显:MDLM 在 ImageNet-1K 上 4 步生成的 FID 为 80.752,而 UDLM 仅为 49.861,差距达 17.6%。然而 UDLM 的优势在多步推理时逆转,其在 1024 步时 OWT PPL 反而恶化到 106.666(MDLM 为 41.545),且零样本平均 PPL 达 59.574,比 MDLM 高出近 6 个点。这种互补性意味着现有方法都无法在理解和生成两个维度上同时达到令人满意的性能,形成了一个明显的性能权衡困境。

本文的目标是本文的目标是打破 MDLM 和 UDLM 之间的帕累托权衡,设计一个统一的离散扩散框架,使得单一模型既能继承 MDLM 的强语义理解和零样本泛化能力,又能继承 UDLM 的少步高效生成质量。具体而言,作者希望通过理论推导证明两种范式可以被统一到同一个数学框架下,并通过一个可控的插值参数 $k$ 在两种噪声模式之间实现平滑过渡。同时,该框架需要在计算效率上具有实用性,能够扩展到大词表场景和 8B 参数量的大语言模型。

与已有工作不同的是,本文的独特切入角度在于对噪声核施加平稳性约束。与最相关的竞品 GIDD 不同,GIDD 使用时间非齐次过程,在每个时间步上动态混合掩码和均匀噪声分布,这导致计算开销巨大(需要在每个时间步重新计算噪声分布)。XDLM 则强制噪声核 $K$ 在所有时间步上保持不变,这一设计选择使得增量噪声的结构与边际噪声一致,从而允许训练目标被分解为静态常数和动态调度的乘积。更重要的是,作者证明了在平稳性约束下,MDLM 和 UDLM 是同一框架的两个极端特例,这为理论统一提供了严格的数学基础。此外,后验概率的标量化重构进一步消除了大词表场景下的内存瓶颈,使得框架具备实际可用性。

核心方法

XDLM 的核心思想可以用一个简单的类比来理解:想象你在黑板上写字,MDLM 的方式是先用橡皮擦掉所有字(变成空白),然后一个字一个字地重新写;UDLM 的方式则是把黑板上的每个字随机替换成别的字,然后逐步修正。两者各有优劣——前者结构清晰但灵活度低,后者灵活但缺乏结构约束。XDLM 的做法是在每一步同时使用两种策略:以比例 $k$ 进行均匀随机替换,以比例 $\mu = 1-k$ 进行掩码操作。从技术路线看,XDLM 将前向过程定义为 $Q_{t|s} = \alpha_{t|s}I + \beta_{t|s}K$,其中噪声核 $K$ 被构造为 $K = \frac{k}{N}J + \mu M$($J$ 为全 1 矩阵,$M$ 为掩码吸收矩阵)。通过将后验概率和 KL 散度重写为标量形式,避免了直接对大矩阵的运算,从而实现了内存高效的训练和采样。

XDLM 最核心的创新是噪声核的平稳性设计及其带来的理论统一性。具体而言,作者引入了一个混合噪声核 $K = \frac{k}{N}J + \mu M$,其中 $k$ 是均匀噪声权重,$\mu = 1-k$ 是掩码噪声权重。这个看似简单的构造有两个深远意义:第一,它满足瞬时混合(instantaneous mixing)条件,即无论起始状态如何,一步转移后都收敛到同一个目标分布 $\pi = ku + \mu m$($u$ 为均匀分布,$m$ 为掩码态),这使得噪声核与最优传输路径(Flow Matching)对齐;第二,它建立了严格的特例关系——当 $k=0$ 时退化为 MDLM,当 $k=1$ 时退化为 UDLM。这与 GIDD 的本质区别在于:GIDD 的混合过程是时间依赖的(time-inhomogeneous),噪声特征随扩散时间线变化;而 XDLM 的噪声核是时间不变的(stationary),将噪声特性与时间调度解耦,从而允许训练目标中的静态部分和动态部分分离计算。

方法步骤详情

XDLM 的方法可以分为以下几个关键步骤:(1) 前向过程定义:使用 $Q_{t|s} = \alpha_{t|s}I + \beta_{t|s}K$ 逐步腐蚀数据,其中 $\alpha_{t|s}$ 和 $\beta_{t|s}$ 是满足 $\alpha_{t|s} + \beta_{t|s} = 1$ 的标量调度。噪声核 $K$ 在所有时间步上保持恒定。(2) 标量原语定义:引入噪声率函数 $r(e) = \frac{k}{N} + \mu \delta_{e,m}$ 和前向扩散映射 $f_t(x, e) = \alpha_t p_{x,e} + \beta_t r(e)$,将矩阵运算转化为标量运算。(3) 标量后验推导:证明后验概率可以写为 $q(z_s = e | z_t, x) = \frac{f_s(x, e) f_{t|s}(e, z_t)}{f_t(x, z_t)}$,避免了显式构造 $N \times N$ 矩阵。(4) 标量 KL 散度:将训练损失中的 KL 散度项重写为 $D_{KL} = \frac{\beta_{t|s}\alpha_s r(z_t)}{f_t(x, z_t)} h_t(x, z_t, \tilde{x}_0)$,其中 $h_t$ 是包含对数差异项的辅助函数。(5) 渐近极限处理:当 $s \to t$ 时,利用 L'Hôpital 法则将 $h_t$ 简化,避免数值不稳定。(6) 训练目标:在连续时间极限下,优化 $L_{vb} = \mathbb{E}_{t,x_0,z_t} \left[ \frac{-\alpha_t' r(z_t)}{f_t(x, z_t)} \cdot h_t(x, z_t, \tilde{x}_0) \right]$。(7) 采样时使用 Gumbel 噪声的祖先采样(ancestral sampling),并可选地引入重掩码(re-masking)策略进行 token 修正。

技术新颖性

XDLM 的技术新颖性体现在三个层面。首先,在理论层面,它是第一个严格证明 MDLM 和 UDLM 是同一统一框架特殊情形的工作。虽然 D3PM 早期就观察到 BERT 可以视为混合腐败的一步特例,但从未从噪声核的角度建立这种统一关系。GIDD 也尝试了插值,但依赖的是时间非齐次过程,缺乏 XDLM 的理论简洁性。其次,在计算层面,标量化重构是关键贡献。直接计算后验概率需要对 $N \times N$ 矩阵进行运算($N$ 为词表大小),在大词表场景下内存消耗巨大。XDLM 通过辅助函数 $r(e)$、$f_t(x,e)$ 将问题转化为 $O(N)$ 的标量运算,使前向吞吐量达到 396,398 tokens/s,是 GIDD(199,516 tokens/s)的近两倍,内存仅需 31.4 GB(UDLM 59.7 GB,GIDD 40.9 GB)。第三,在实际效果层面,XDLM 展现出独特的训练动态——与 MDLM 快速收敛但早饱和不同,XDLM 持续改进,在后期超越 MDLM,暗示其具有更好的长期扩展潜力。

XDLM 噪声核结构与理解-生成权衡图
Figure 1: XDLM 噪声核结构与理解-生成权衡图
ImageNet-1K 8 步采样动态可视化
Figure 7: ImageNet-1K 8 步采样动态可视化

实验结果

XDLM 在语言和图像两个模态上均展示了优越的性能平衡。在零样本语言理解方面,XDLM 在七个外部基准上的平均 PPL 为 54.110,与 MDLM(53.650)和 GIDD(53.384)相当,显著优于 UDLM(59.574),差距达 5.4 个点。在 OWT 验证集上,XDLM($k=0.1$)的 PPL 为 24.097,虽略高于 MDLM(23.321)和 GIDD(23.136),但远优于 UDLM(25.937)。在语言生成质量方面,XDLM 在少步和多步两种体制下都表现稳健:在 OWT 上 8 步采样时 PPL 为 301.640,大幅优于 MDLM(711.382),接近 UDLM(183.991);在 1024 步时 PPL 为 52.609,保持了高质量。在 ImageNet-1K 条件生成中,XDLM 在无 CFG 设置下 16 步 FID 为 25.774,优于 UDLM(26.242)并大幅领先 MDLM(28.785);在 4 步少步生成中 FID 为 54.085,远优于 MDLM(80.752)。在 CFG(scale=2.0)设置下,XDLM 在 4 步和 8 步分别取得最佳 FID(13.550 和 8.956),16 步 FID 为 8.625,优于 UDLM(8.980)。最令人瞩目的是大模型扩展实验:在 LLaDA-8B 上仅用 600 步持续预训练,LLaDA-XDLM 在 MBPP 代码生成任务上取得 15.0 分,相比原始 LLaDA 的 6.8 分提升了 120%,将不可编译的失败样本从 429 个降至 304 个。训练动态分析揭示了性能交叉现象:MDLM 和 GIDD 在训练初期占优但很快饱和,而 XDLM 持续改进,在 700k 步后超越或追平 MDLM。

OWT 验证集困惑度对比
Table 1: OWT 验证集困惑度对比
零样本跨数据集性能对比
Table 2: 零样本跨数据集性能对比
ImageNet-1K 无 CFG 条件生成性能
Table 3: ImageNet-1K 无 CFG 条件生成性能
ImageNet-1K 有 CFG(scale=2.0)条件生成性能
Table 4: ImageNet-1K 有 CFG(scale=2.0)条件生成性能
LLaDA-XDLM 在推理和代码生成基准上的评估
Table 16: LLaDA-XDLM 在推理和代码生成基准上的评估
计算效率对比
Table 17: 计算效率对比
语言生成质量对比
Figure 2: 语言生成质量对比
LLaDA-8B 持续预训练评估
Figure 3: LLaDA-8B 持续预训练评估
训练动态中的性能交叉现象
Figure 4: 训练动态中的性能交叉现象
ImageNet-1K 无 CFG 定性对比
Figure 5: ImageNet-1K 无 CFG 定性对比
ImageNet-1K 有 CFG(scale=2.0)定性对比
Figure 6: ImageNet-1K 有 CFG(scale=2.0)定性对比
查看结构化数据
任务指标本文基线提升
零样本语言理解(7个基准平均) PPL (↓) 54.110 MDLM 53.650 / GIDD 53.384 / UDLM 59.574 比 UDLM 低 5.4 个点,与 MDLM 仅差 0.46
OWT 验证集似然 PPL (↓) 24.097 MDLM 23.321 / GIDD 23.136 / UDLM 25.937 比 UDLM 低 1.84 个点
ImageNet-1K 4步无CFG生成 FID (↓) 54.085 MDLM 80.752 / GIDD 86.842 / UDLM 49.861 比 MDLM 降低 26.7(33%)
ImageNet-1K 16步无CFG生成 FID (↓) 25.774 MDLM 28.785 / GIDD 35.403 / UDLM 26.242 最佳,比 UDLM 低 0.47
ImageNet-1K 4步CFG=2.0生成 FID (↓) 13.550 MDLM 33.468 / GIDD 41.000 / UDLM 14.055 最佳,比 UDLM 低 0.51
ImageNet-1K 16步CFG=2.0生成 FID (↓) 8.625 MDLM 6.725 / GIDD 7.076 / UDLM 8.980 比 UDLM 低 0.36
LLaDA-8B MBPP代码生成(32步) Score (↑) 15.00 LLaDA 6.80 / LLaDA-MDLM 4.40 比原始 LLaDA 提升 120%
LLaDA-8B GSM8K(32步) Score (↑) 29.26 LLaDA 24.49 / LLaDA-MDLM 24.41 比原始 LLaDA 提升 19.5%

局限与改进

论文的局限性主要体现在以下几个方面。首先,作者明确承认 XDLM 尚未从头开始大规模预训练,所有大模型实验都是在已有 LLaDA-8B(一个预训练好的 MDLM)上进行持续预训练,这意味着 XDLM 在大规模预训练中的潜力尚未被充分验证。其次,性能交叉现象——即 UDLM 和 XDLM 在大采样步数下似乎能超越 MDLM——虽然被观察到但未得到充分理论解释,作者也承认这是未来研究方向。第三,XDLM 在无 CFG 的 16 步 ImageNet-1K 生成中 FID 为 25.774,虽优于 UDLM(26.242),但差距较小(0.47 点),而 MDLM 在有 CFG 的 16 步设置下仍取得最佳 FID(6.725),说明在特定条件下掩码范式仍具优势。第四,从计算效率角度看,虽然 XDLM 大幅优于 GIDD 和 UDLM,但与纯 MDLM 相比仍有差距——前向吞吐量为 396,398 tokens/s(MDLM 为 424,294),采样吞吐量为 7,108 tokens/s(MDLM 为 8,789)。最后,XDLM 的最优混合比 $k$ 需要根据任务和采样步数进行调优,缺乏自适应机制。

独立分析的弱点

从独立分析的角度,XDLM 存在以下几个可以改进的弱点。第一,混合比 $k$ 的选择目前是全局固定的超参数,但最优 $k$ 依赖于采样步数和任务类型——少步生成偏好较大的 $k$(如 0.5-0.9),多步推理偏好较小的 $k$(如 0.1),图像和文本的最优值也不同。改进方向是设计自适应的、时间步依赖的调度策略,让模型在采样过程中动态调整 $k$,例如在早期利用较大 $k$ 快速建立全局结构,在后期切换到较小 $k$ 精细打磨。第二,重掩码(re-masking)策略虽然在定性分析中展示了有趣的纠错能力(如将 coffin、trumpets 等随机初始 token 退回 [MASK]),但论文中缺乏对这一机制的系统性消融实验,不清楚重掩码的具体贡献有多大。第三,标量 KL 散度推导中涉及对全词表的求和项 $\frac{k\beta_s}{N\alpha_s} \sum_{e \in V} \log \frac{f_s(x, e)}{f_s(\tilde{x}_0, e)}$,虽然理论上 $O(N)$,但在实际超大词表(如 10 万以上)场景下仍可能成为瓶颈,可以考虑采样近似或解析消去。第四,论文在 LLaDA 实验中仅使用了 600 步持续预训练和 10B token 数据,实验规模偏小,缺乏更充分的大规模验证。

未来方向

基于论文成果,未来研究可以从以下几个方向展开。作者提出的最直接方向是从头大规模预训练 XDLM,这将允许更全面地探索模型的涌现性质,尤其是在长上下文理解和复杂推理任务上的表现。其次,深入研究性能交叉现象的理论基础,理解为什么 UDLM 和 XDLM 在大采样步数体制下能超越 MDLM,这可能揭示离散扩散模型训练动力学的深层规律。第三,为 XDLM 开发特定领域的采样策略,论文中已展示了文本和图像生成的不同最优配置,但缺乏针对代码生成、多模态推理等特定任务的定制优化。第四,探索文本和视觉模态在单一统一模型中的交互与平衡,目前的实验分别在语言和图像上进行,多模态联合训练的效果未知。第五,开发后训练方案(post-training schemas)和推理加速技术,如蒸馏、缓存机制等,这对实际部署至关重要。第六,将平稳噪声核的设计思想扩展到其他离散结构化数据,如蛋白质序列、分子图、音乐符号等。

复现评估

论文的复现条件总体较好。代码已开源(论文中提供了 GitHub 链接),这大幅降低了复现门槛。训练配置方面,所有实验使用 8xH800 GPU 集群,语言模型实验使用 AdamW 优化器($\beta_1=0.9$, $\beta_2=0.999$),学习率 $3 \times 10^{-4}$,batch size 512,EMA rate 0.9999,这些超参数与 MDLM/UDLM 前作一致,说明实验设置是公平的。LM1B 和 OWT 是常用公开数据集,ImageNet-1K 也是标准基准。LLaDA-8B 持续预训练仅需 600 步,计算成本可控。然而,完全复现所有实验仍需要较大算力——仅训练 OWT 模型就需要 8xH800 GPU 上约 1M 步的训练。标量化重构的实现细节在论文附录中有完整推导,且有代码辅助,复现难度中等偏低。对于希望在自己项目中应用 XDLM 的研究者,最实用的方案是在已有 MDLM 模型上进行类似 LLaDA-XDLM 的持续微调,仅需少量额外训练即可获得显著收益。