超越单一令牌：通过离散MMD蒸馏离散扩散模型 Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD
提出D-MMD方法将离散扩散模型蒸馏为少步生成器，在文本和图像上超越教师
前置知识
离散扩散模型
离散扩散模型是一类用于生成离散数据（如文本token、图像像素值）的生成模型。它们通过逐步添加噪声将数据转换为均匀分布或掩码分布，然后学习反向去噪过程来生成数据。两种主要类型是均匀扩散（将token转换为均匀分布）和掩码扩散（将token转换为特殊掩码token）。与传统自回归模型不同，扩散模型可以并行处理token块，因此具有更好的加速器利用率。
本文的核心方法D-MMD专门针对离散扩散模型的蒸馏问题设计。理解离散扩散的前向过程（$z_t \sim \text{Cat}(\alpha_t x + (1-\alpha_t)\pi)$）和后验公式是理解D-MMD算法设计的基础。
模型蒸馏
模型蒸馏是一种将大型或复杂的教师模型的知识转移到小型或简化的学生模型的技术。在扩散模型背景下，蒸馏的目标是创建一个只需要少量采样步骤的生成器，同时保持生成质量。蒸馏方法可以分为确定性蒸馏（如DDIM、一致性模型）和随机性蒸馏（如MMD、分布匹配）。确定性蒸馏通常学习确定性的轨迹，而随机性蒸馏则学习匹配教师模型的分布。
本文的D-MMD属于随机性蒸馏方法，通过匹配教师模型的条件期望来训练学生生成器。理解蒸馏的基本概念有助于理解为什么D-MMD需要交替优化生成器和辅助模型，以及为什么它能够超越教师性能。
Moment Matching Distillation (MMD)
MMD是连续扩散模型蒸馏的一种成功方法，其核心思想是确保学生生成器的条件期望与教师模型的条件期望相匹配。MMD采用交替优化策略：生成器最小化教师模型的损失，同时最大化辅助模型的损失；辅助模型则学习匹配生成器的期望。数学上，MMD损失可以表示为 $\mathcal{L}_{MMD}(\eta) = \mathbb{E}_{g_\eta(z_t,s,z_s)}[w(s)\hat{x}_\eta(z_t)^\top \text{sg}(\hat{x}_\phi(z_s) - \hat{x}_\theta(z_s))]$。
D-MMD是MMD在离散域的推广。理解MMD的原理有助于理解D-MMD的设计动机。D-MMD的关键创新在于将MMD从连续空间推广到离散空间，使用交叉熵而不是欧几里得距离，并使用软概率向量而不是硬样本。
交叉熵损失
交叉熵是衡量两个概率分布之间差异的常用指标，在分类和生成任务中广泛使用。对于离散数据，交叉熵定义为 $\text{CE}(x|\hat{x}) = -\sum_c x_c \log \hat{x}_c$，其中$x$是真实标签的one-hot向量，$\hat{x}$是预测概率向量。在扩散模型中，交叉熵损失用于训练模型预测给定噪声状态下原始数据token的概率。权重函数$w(t)$可以根据时间步调整损失的重要性。
D-MMD将MMD中的欧几里得距离替换为交叉熵损失，这是从连续域到离散域的关键转换。D-MMD的生成器损失 $\mathcal{L}_{GEN}(\eta) = \text{CE}(\hat{x}_\eta|\hat{x}_\theta(z_s)) - \text{CE}(\hat{x}_\eta|\hat{x}_\phi(z_s))$ 和辅助模型损失都依赖于交叉熵。
GPT-2 Gradient Moment
这是本文提出的新评估指标，用于评估离散扩散模型的生成质量。其核心直觉是：如果参考模型（如GPT-2）在训练数据上收敛，那么它在真实数据上的损失梯度期望应该为零。因此，生成样本的质量可以通过其与真实数据梯度期望之间的距离来衡量。数学公式为 $\|\mathbb{E}_g[\nabla_\theta \log p^{\text{LLM}}_\theta(x)] - \mathbb{E}_q[\nabla_\theta \log p^{\text{LLM}}_\theta(x)]\|^2$，其中$g$是生成分布，$q$是数据分布。
本文提出的GPT-2 GM指标解决了传统生成困惑度指标的缺陷。生成困惑度容易被低温度采样'作弊'，而GPT-2 GM更加稳健。理解这个指标有助于评估D-MMD在文本生成任务上的性能。
研究动机
离散扩散模型面临严重的效率瓶颈。它们需要大量采样步骤才能生成高质量数据，因为每个token的建模假设了独立性，导致采样过程中的误差累积。虽然扩散语言模型相比因果LLM具有更好的加速器利用率（可以并行处理token块），但它们需要很多迭代才能收敛到合理的生成结果。例如，在CIFAR-10数据集上，均匀扩散教师模型需要1024步才能达到FID 7.5，掩码扩散教师模型需要512步才能达到FID 6.4。在Open Web Text数据集上，掩码教师模型需要256步才能达到GPT-2 GM 0.275。这种高计算成本和高FLOPs需求限制了离散扩散模型的实际应用。
本文的目标是本文的目标是将连续扩散域中成功的MMD（Moment Matching Distillation）方法推广到离散扩散模型，创建只需少量采样步骤的生成器。具体而言，作者希望开发一个能够蒸馏离散扩散过程的算法，使得蒸馏后的生成器在使用极少函数评估（NFEs）的情况下，其生成质量能够超越教师模型。这个算法应该同时支持不同的离散扩散过程（如掩码扩散和均匀扩散），并且在文本和图像生成任务上都能有效工作。
与已有工作不同的是，本文的独特切入角度在于识别出离散蒸馏方法的根本问题。现有的离散蒸馏方法如SDTT存在局限性：它们要么性能提升有限，要么会丢失模式。例如，SDTT无法完美地用一步模拟两个完美相关的硬币投掷结果——这是由于因子化输出分布的限制。Di4C方法通过扩展模型输出支持混合分布来部分解决这个问题，但这种方法的效果有限，因为需要指数级增长的混合数量来捕获所有token间的相关性。本文的D-MMD方法采取完全不同的思路：保持因子化输出分布不变，而是让生成器通过'坍缩'其因子化输出分布来匹配期望矩。换句话说，整个生成器本身变成了混合分布。
核心方法
D-MMD方法的核心思想是将连续MMD推广到更通用的扩散过程，特别是离散扩散。方法的直觉是通过交替优化策略，使生成器学习匹配教师模型的条件期望分布。技术路线上，D-MMD采用min-max框架：生成器目标是在教师模型下最小化损失，同时在辅助模型下最大化损失；辅助模型则学习最小化与生成器的损失，同时保持接近教师模型的分布。这种方法的关键创新在于对于离散数据，使用软概率向量而不是硬样本，并使用交叉熵损失代替欧几里得距离。
D-MMD的核心创新点在于将MMD从连续域成功推广到离散域，本质区别在于处理离散采样的方式。在连续MMD中，可以直接从生成器采样$x$并计算损失梯度。但在离散域，从分类分布$\text{Cat}(p=\hat{x}_\eta(z_t))$采样的硬token到参数$\eta$没有直接的梯度。D-MMD的解决方案是使用软概率向量$\hat{x}_\eta(z_t)$代替硬样本。这种设计使得D-MMD能够在对数概率空间中进行匹配，而不是在原始输出空间。另一个关键创新是D-MMD的损失函数设计：生成器损失$\mathcal{L}_{GEN}(\eta) = -\sum_c (\hat{x}_\eta)_c \log(\hat{x}_\theta(z_s)_c/\hat{x}_\phi(z_s)_c)$，这本质上是对$\hat{x}_\theta(z_s)$和$\hat{x}_\phi(z_s)$的对数概率差异进行更新。
方法步骤详情
D-MMD算法的完整训练流程如下。输入包括学生生成器$\hat{x}_\eta$、教师模型$\hat{x}_\theta$、辅助模型$\hat{x}_\phi$、训练步$i$、采样步数$k$、数据集$D$、权重函数$w(s)$和损失函数$\mathcal{L}(\cdot,\cdot,\cdot)$。算法首先采样$s, \delta_t \sim U(0,1), U(0,1/k)$，计算$t = \min(1, s+\delta_t)$。然后从数据集$D$采样数据点并扩散得到$z_t$。生成器生成概率向量$\hat{x}_\eta(z_t)$，并从中采样$x \sim \text{Cat}(p=\hat{x}_\eta(z_t))$。使用采样器（如后验$q(z_s|z_t,x)$）绘制样本$z_{s|x,z_t}$，并对$z_s$应用stop-gradient。如果训练步$i$是偶数，则优化生成器：最小化$\mathcal{L}_{GEN}(\eta) = \mathcal{L}_s(\hat{x}_\eta(z_t), \hat{x}_\theta(z_s), z_s) - \mathcal{L}_s(\hat{x}_\eta(z_t), \hat{x}_\phi(z_s), z_s)$。如果$i$是奇数，则优化辅助模型：最小化$\mathcal{L}_{AUX}(\phi) = \mathcal{L}_s(x, \hat{x}_\phi(z_s), z_s) + \mathcal{L}_s(\hat{x}_\theta(z_s), \hat{x}_\phi(z_s), z_s)$。对于掩码扩散，可以选择性地使用软目标$x \leftarrow \hat{x}_\eta(z_t)$。
技术新颖性
D-MMD的技术新颖性体现在多个方面。首先，它提出了一个更通用的MMD公式，可以应用于更一般的扩散过程，包括离散扩散。这个通用公式采用min-max框架：$\mathcal{L}_{D-MMD}(\eta) = \min_\eta \max_\phi \mathbb{E}_{g_\eta(z_t,x,s,z_s)}[\mathcal{L}_s(x, \hat{x}_\theta(z_s), z_s) - \mathcal{L}_s(x, \hat{x}_\phi(z_s), z_s) - \mathcal{L}_s(\hat{x}_\theta(z_s), \hat{x}_\phi(z_s), z_s)]$。其次，D-MMD针对离散域的特点，使用软概率向量和交叉熵损失，解决了离散梯度传播的问题。第三，D-MMD支持温度和top-p蒸馏，将教师模型的采样特性融入学生生成器。特别是对于top-p采样，D-MMD提出了动态降低logits的策略（$s_\theta(z_s) \leftarrow s_\theta(z_s) - (1-\text{mask}_{\text{top-}p}) \cdot \Delta$）来避免梯度爆炸。最后，D-MMD提出了GPT-2 Gradient Moment这一新的评估指标，解决了传统生成困惑度指标的缺陷。
实验结果
D-MMD在图像和文本生成任务上都展现了卓越的性能。在CIFAR-10图像生成任务上，Uniform D-MMD在32步达到FID 3.7，而Uniform教师需要1024步才能达到FID 7.5，性能提升约51%。Masked D-MMD的表现更为惊人：它在16步就超越了教师（FID 5.3 vs 教师512步的FID 6.7），在64步达到FID 3.5，远超教师512步的FID 6.7。在Open Web Text文本生成任务上，Masked D-MMD在16步就达到GPT-2 GM 0.236，而Masked教师需要256步才能达到GPT-2 GM 0.275。Uniform D-MMD同样表现出色，在8步达到GPT-2 GM 0.337，与Uniform教师256步的0.326相当。在块自回归扩散实验中，256块大小的模型上，16步D-MMD达到了与256步教师相同的性能（GPT-2 GM 0.225）。与现有方法相比，D-MMD明显优于Di4C（教师40步FID 8.0，Di4C混合20步FID 9.5，D-MMD Uniform 8步FID 5.0）和SDTT（教师64步GPT-2 GM 0.307，SDTT 64步0.293，D-MMD 16步0.236）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| CIFAR-10图像生成（Uniform扩散） | FID（越低越好） | 32步 FID 3.7 | 教师1024步 FID 7.5 | 51%性能提升，步数减少97% |
| CIFAR-10图像生成（Masked扩散） | FID（越低越好） | 16步 FID 5.3, 64步 FID 3.5 | 教师512步 FID 6.7 | 16步超越教师，64步提升48% |
| Open Web Text文本生成（Masked扩散） | GPT-2 GM（越低越好） | 16步 GM 0.236 | 教师256步 GM 0.275 | 14%性能提升，步数减少94% |
| Open Web Text文本生成（Uniform扩散） | GPT-2 GM（越低越好） | 8步 GM 0.337 | 教师256步 GM 0.326 | 相当性能，步数减少97% |
| CIFAR-10与Di4C对比 | FID（越低越好） | Uniform 8步 FID 5.0 | Di4C混合20步 FID 9.5 | 47%性能提升，步数减少60% |
| Open Web Text与SDTT对比 | GPT-2 GM（越低越好） | Masked 16步 GM 0.236 | SDTT 64步 GM 0.293 | 19%性能提升，步数减少75% |
局限与改进
作者承认D-MMD存在一些局限性。首先，对于非掩码扩散（如均匀扩散），训练辅助模型时必须使用硬样本$\hat{x}_\eta$，因为$z_s \sim q(z_s|x,z_t)$是与$x$和$z_t$一致的样本，而不是与$\hat{x}_\eta$一致。只有掩码扩散可以使用软目标，因为每维的掩码$z_s$不提供关于$x$的信息。其次，对抗优化的动力学可能很困难，往往依赖于特定的超参数设置。第三，学生模型超越教师性能的现象可能会在采样步数增加时出现性能退化，因为随着步数增加，性能会收敛到教师水平。我观察到的额外局限性是：D-MMD需要训练额外的辅助模型，增加了计算开销；对于非常长的序列，块自回归设置可能仍需要较大的块大小才能保持质量；GPT-2 GM指标虽然比生成困惑度更稳健，但仍然依赖于参考模型的选择和训练状态。
独立分析的弱点
D-MMD的几个具体弱点值得改进方向探索。第一，辅助模型的训练增加了计算开销，未来可以考虑设计更高效的辅助模型训练策略，如参数共享或渐进式训练。第二，对抗优化的不稳定性可能影响训练效果，可以探索更稳定的min-max优化算法，如使用梯度惩罚或正则化技术。第三，D-MMD在块自回归设置下的性能依赖于块大小选择，未来可以研究自适应块大小或动态块选择机制。第四，GPT-2 GM指标虽然改进了生成困惑度，但仍然有局限性，可以考虑开发更全面的评估指标，结合多个参考模型或人类评估。第五，D-MMD目前主要在无条件生成任务上验证，条件生成（如文本到图像、文本续写）的效果可能需要进一步研究。
未来方向
作者提出的未来方向包括探索D-MMD在其他离散数据类型上的应用，如音频、代码和分子生成。基于D-MMD的成果，可以延伸多个研究方向：1）将D-MMD与其他技术结合，如与自回归模型混合形成更强的混合生成模型；2）研究D-MMD在条件生成任务上的表现，如图像captioning、文本到图像生成等；3）开发D-MMD的变体，支持更复杂的采样策略，如分类器引导采样；4）研究D-MMD在多模态生成任务上的应用；5）探索D-MMD与强化学习的结合，用于序列决策任务；6）开发更高效的不需要辅助模型的蒸馏方法；7）研究D-MMD在大规模模型上的可扩展性；8）将GPT-2 GM指标推广到其他模态和任务，如图像、音频的梯度矩评估。
复现评估
从论文描述来看，D-MMD的复现具有中等难度。论文提供了详细的算法描述和超参数设置（如温度$\tau$、top-p值$p$、噪声条件参数$\Delta=2$等）。实验在标准数据集上进行（CIFAR-10、Open Web Text），这些数据集都是公开可获取的。论文没有明确说明代码是否开源，这是复现的一个潜在障碍。在算力需求方面，训练教师模型和蒸馏学生模型都需要相当的计算资源。例如，在CIFAR-10上训练需要生成50K样本进行FID评估，在Open Web Text上训练2%数据作为验证集。估计复现主要实验需要多GPU环境，训练时间可能在数天到数周，取决于具体配置。总的来说，如果有代码开源和足够的计算资源，D-MMD的复现应该是可行的，但需要一定的机器学习和扩散模型背景知识。
论文图表
