巴黎综合理工学院提出Di[M]O！掩码扩散模型的一步生成革命！-CSDN博客

本文链接：https://blog.csdn.net/DataSourceAI/article/details/146462359

🌐 社群导航

🔗 点击加入➡️【AIGC/LLM/MLLM/3D/自动驾驶】技术交流群

🔗 点击加入➡️【2-3月CCF投稿】交流群

最新论文解读系列

论文名：Di[M]O: Distilling Masked Diffusion Models into One-step Generator

论文链接：https://arxiv.org/pdf/2503.15457

开源代码：https://yuanzhi-zhu.github.io/DiMO/

导读

最近，掩码扩散模型（MDMs）已成为生成式建模领域的一个突出框架，在广泛的任务中表现出强大的性能。与按顺序生成数据的自回归模型不同，MDMs能够更快地生成离散数据。此外，与连续扩散模型相比，MDMs在多模态框架内进行视觉建模方面具有特别的优势，因为它们利用统一的词汇表无缝集成视觉和文本内容。它们的多功能性使其在图像生成、文本生成、视频生成、多模态建模、蛋白质设计、音频合成和运动生成等方面得到了成功应用。此外，已经提出了几种算法来指导MDMs的生成过程，进一步增强了它们的灵活性和与用户定义目标的一致性。

简介

掩码扩散模型（MDMs）已成为一种强大的生成式建模技术。尽管它们取得了显著的成果，但通常存在多步推理速度慢的问题。在本文中，我们提出了Di[M]O，这是一种将掩码扩散模型蒸馏为一步生成器的新方法。Di 解决了两个关键挑战：（1）使用中间步骤信息进行一步生成的难处理性，我们通过令牌级分布匹配来解决这个问题，该方法借助辅助模型，通过“策略内框架”优化模型输出的对数几率；（2）初始分布缺乏熵，我们通过一种令牌初始化策略来解决这个问题，该策略在保持与教师训练分布相似性的同时注入随机性。我们展示了Di[M]O在类别条件和文本条件图像生成方面的有效性，在大幅减少推理时间的同时，令人印象深刻地实现了与多步教师模型输出相媲美的性能。据我们所知，我们是首个成功实现掩码扩散模型一步蒸馏的团队，也是首个将离散蒸馏应用于文本到图像生成的团队，为高效生成式建模开辟了新途径。

方法与模型

我们提出了将 [M]DM 蒸馏为一步生成器（Distilling [M]DM into One-step generator，Di[M]O）的新方法，用于将多步 MDM 教师模型蒸馏为一步生成器。图3展示了的概述。

图3. Di[M]O流程。我们的方法将一个代价高昂的多步MDM教师模型提炼为一个单步生成器。给定使用我们提出的令牌初始化策略采样得到的，单步生成器（学生模型）生成对数概率，从中采样得到图像令牌序列。然后通过前向掩码扩散过程处理这些令牌以获得中间状态。对于每个中间状态，我们交替更新单步生成器和辅助模型：单步生成器通过最小化令牌级别的条件散度进行优化，而辅助模型使用交叉熵损失进行训练，以对生成的令牌的分布进行建模，并形成更新的梯度。教师模型在训练期间保持冻结。

1. 一步在线策略蒸馏

我们的目标是训练学生模型，将初始随机输入分布转换为与教师模型的多步生成分布极为相似的输出。受近期关于策略内蒸馏 [1, 2, 5, 35, 37, 84, 93] 的研究启发，我们提议首先对源自学生模型自身预测的中间状态进行采样，然后使教师模型的条件预测分布与学生模型的条件预测分布相匹配。我们将此过程记为，其中是散度。通过对所有可能的中间状态强制执行这种匹配，我们有效地确保了学生模型的整体生成分布与教师模型的生成分布一致。整体蒸馏目标可表述如下：

其中和是一个加权函数。由于生成器现在只能将作为输入并直接生成，我们应用前向扩散过程，从其输出中得到作为伪中间状态。此散度是针对所有可能的和计算的，以涵盖每一个可能的中间状态。

词元级别的分布匹配。公式中的条件散度可以进一步分解，并表示为给定时每个掩码词元的平均散度，定义如下：

其中是中掩码词元的数量。与连续扩散模型蒸馏 [60, 66, 119] 不同，在连续扩散模型蒸馏中，散度是在模型的完全生成潜变量上计算的，而我们提出的 Di[M]O 则通过公式中的分解在词元级别计算散度。

2. 损失梯度的近似

直接优化公式是不可行的，因为的精确形式未知。与先前的工作 [41, 77, 125] 类似，我们试图近似训练目标的以下梯度（见附录A.1）：

其中，对数几率是学生模型的直接输出。通过这种方式，我们可以将损失梯度视为两项的乘积：散度相对于模型对数几率的梯度（橙色部分），以及模型输出相对于学生模型参数的雅可比矩阵（蓝色部分）。通过分别近似这两项，我们可以解决直接优化原始目标的难处理性问题：

对的近似：鉴于式(6)中突出显示的散度梯度项可视为教师模型输出和未知学生模型输出的泛函，我们采用中的策略，并引入一个辅助模型来近似未知的。辅助模型作为一个预测器进行训练，使用单步生成器生成的数据，并通过式(3)中的MDM训练损失进行优化，以得到。
对的近似：模型雅可比项的近似必须满足一个关键要求：它需要提供梯度信息以更新模型参数。考虑到这一约束，我们提出在一致性属性[102]的假设下，用单步输出替代难以处理的隐式项来进行近似。

这些近似得到以下损失梯度：

为了有效地从教师模型转移先验知识，我们用教师模型的预训练权重初始化生成器和辅助模型。

广义杰弗里散度。与依赖反向KL（RKL）散度的连续扩散提炼相关工作[60, 66, 119]不同，所提出的Di[M]O并不局限于这种选择。相反，它可以利用其他令牌级别的散度度量，因为我们在式(7)中没有对散度的形式施加限制。这是一个特别有趣的特性，因为已知RKL会表现出不期望的模式搜索行为[18]。

具体而言，我们提议探索广义杰弗里散度的使用，它被定义为前向KL（FKL）散度和RKL散度的线性组合，并研究其在模式搜索之外的效果（见附录F）：

其中为超参数。这些散度的显式梯度在附录A.2中给出。

3. 令牌初始化策略

与初始分布通常为标准高斯分布的连续扩散模型不同，MDM（掩码扩散模型，Masked Diffusion Model）的初始状态由一系列确定性的相同 [M] 标记组成。这种固定的初始化方式会导致一步生成时出现模式崩溃，因为生成器只能输出具有有限样本多样性的固定对数几率。因此，我们探索了三种可能的标记序列初始化策略，包括固定初始化策略，具体如下（见图 3 左图）：

全掩码标记：遵循标准的 MDM 初始化，我们仅用 [M] 标记来初始化序列。
全随机标记：为了使熵最大化，我们选择具有随机值的图像标记来填充初始标记序列。
混合策略：通过固定初始掩码比例，并将标记中的其余替换为随机图像标记，将前两种策略结合起来。

在我们的工作中，我们采用混合策略，这基于两个关键假设，并得到了图 5 中消融分析的实证支持：（1）为了防止模式崩溃，一步生成器的初始化应包含一定程度的随机性；（2）由于一步生成器继承自教师模型，其目标是从部分掩码输入中预测所有标记，因此初始化必须保留类似的模式，即包含一些掩码标记，以避免输入分布不匹配。

此外，受先前自回归蒸馏工作 [53] 的启发，该工作提出用随机高斯噪声替换所有初始标记嵌入，因此我们在当前方案中添加高斯扰动，以进一步将随机性从离散空间扩展到实数空间。具体来说，在随机选择个标记和个随机图像标记后，我们使用保持方差的方案 [38,101] ，以固定的噪声水平用高斯噪声对所有标记嵌入 e 进行扰动，如图 3 所示。

算法概述。我们现在在算法 1 中总结我们的方法，其中生成器和辅助模型进行迭代更新。在每次迭代中，生成器从初始状态一步生成标记，并使用式（7）中的损失梯度进行优化，而辅助模型则以这些生成的标记为目标进行训练。

算法 1 Di[M]O 蒸馏

实验与结果

教师模型。我们在类别条件图像生成和文本到图像生成任务上进行了广泛的蒸馏实验。对于类别条件生成，我们采用性能最佳且公开可用的类别条件 MDM，即 MaskGit [7] 作为教师模型。对于文本到图像生成，我们采用最近的 Meissonic [4] 作为教师模型，并在蒸馏过程中使用 LAION - Aesthetics - 6+ 提示数据集 [14] 进行提示。

评估指标。我们用于比较 ImageNet 结果的指标是弗雷歇 inception 距离（Fréchet Inception Distance，FID）[34] 和 inception 得分（Inception Score，IS）[88]。我们还使用精度（Prec.）和召回率（Rec.）[48]、密度（Den.）和覆盖率（Cov.）[65] 来进一步评估生成图像的保真度和多样性。对于文本到图像生成，我们遵循文献 [4] 并测量 HPSv2 [112] 和 Geneval [26] 得分。有关这两个任务的更多指标和结果，请参阅附录 D。在我们对 ImageNet 的实验中，我们使用张生成图像进行消融实验，使用张生成图像进行基准测试，并使用 Clean - FID [73] 将它们与 ImageNet 验证集中的张图像进行比较。

实验设置。两个教师模型均采用无分类器引导（Classifier-Free Guidance，CFG）[36]。对于MaskGit教师模型，我们在蒸馏过程中保持2的CFG系数，而对于Meissonic教师模型，我们使用4的CFG系数。我们的生成器始终在有条件的情况下进行训练，在推理过程中无需使用CFG，进一步提高了采样效率。MaskGit的所有消融实验均以64的批量大小进行30000次迭代训练。我们使用不同的温度对最终的检查点进行评估，以获得最佳的指标结果。MaskGit实验和Meissonic实验的学习率分别为和。更多详细信息请参见附录C。

1. 类别条件图像生成

在表1中，我们展示了我们的与各种加速方法的定量性能比较，这些加速方法包括高阶采样器 - 梯形法[81]、其他蒸馏技术以及教师模型。与需要32步才能达到7.1的FID（弗雷歇 inception 距离）的高阶采样器相比，我们的方法在保持有竞争力的性能的同时，将步数显著减少到仅1步，FID为6.91，IS（ inception 得分）为214.0。与蒸馏方法相比，我们的方法与di4c [31]取得了相当的结果，di4c使用4步，FID为6.79，IS为209.2。我们的方法在一步内实现了6.91的FID和更高的214.0的IS。此外，与不同步数的教师模型（MaskGit [7]）相比，我们的方法在性能上与教师模型在16步时达到的6.60的FID非常接近，而我们的方法仅需一步即可实现。这些比较凸显了我们的方法在匹配多步教师模型性能的同时，显著降低计算成本的有效性。更多细节请参考附录C。

表1. 类别条件ImageNet - 256上的定量结果。* 表示从原始图表中估算出的数值。

2. 消融实验

在本节中，我们对我们方法的关键超参数进行消融实验，以验证我们的设计选择：（1）初始掩码比率；（2）杰弗里系数；（3）高斯扰动强度：

初始掩码比率。我们通过分析初始掩码比率的影响来开始我们的消融实验，该比率是我们算法中的一个关键因素。如图5a所示，我们的结果证实了4.3节中的假设，即两个极端值和都不起作用。具体来说，设置会导致模式崩溃，而非常低的值会导致训练不稳定，如图5a子图中的断裂曲线所示。最优选择似乎是，它实现了最低的FID。可视化演示见图4。

图4. ImageNet的可视化结果。使用不同的训练的生成器一步生成的图像与教师模型16步采样生成的图像进行比较。样本的类别标签从上到下分别为388、979和207。

图5. 在ImageNet上使用FID作为评估指标的消融实验。* 表示训练崩溃，超出了与其他结果的可比范围，我们在右上角的子图中以相同的x轴范围展示这些结果。

杰弗里系数。在此基础上，我们对广义杰弗里散度系数进行了消融实验（图5b）。我们的实验表明，降低通常会改善FID，这与先前工作[116]的观察结果一致。令人惊讶的是，即使对于的负值，我们的方法仍然有效，在时观察到最佳性能。

高斯扰动。最后，我们研究了高斯扰动对令牌嵌入的影响（图5c）。我们的结果表明，引入扰动可以进一步改善FID，进一步提高生成样本的质量。为了完整起见，附录图7中提供了相应的IS指标结果。

3. 文本到图像生成

我们通过在HPSv2 [112]和GenEval [26]基准测试上评估我们的蒸馏一步生成器，验证了在文本到图像生成中的有效性。

表2. HPS v2.0基准测试。分数收集自https://github.com/tgxs002/HPSv2。我们突出显示了最佳结果。

在表2中，我们展示了HPSv2 [112]基准测试，该测试用于评估人类对文本到图像（T2I）模型的偏好。我们的模型是从唯一开源的文本到图像MDM（Meis-sonic [4]）中提炼而来的，我们不仅将我们的一步生成器与教师模型进行了比较，还与近期的扩散模型[19, 21, 76, 83]及其提炼的一步生成器进行了比较。结果清楚地表明，我们的一步生成器与使用16到32步生成的教师模型相比具有有竞争力的性能，而当生成步数减少时，教师模型的性能会迅速下降（例如，48步生成的得分是28.83，而4步生成的得分降至24.66）。我们的一步生成器也优于其他基于连续扩散的一步生成器；然而，由于它们是从不同的教师模型中提炼而来的，因此直接比较存在局限性。

表3. GenEval基准测试。我们突出显示了最佳结果。

为了更好地评估我们的一步生成器的语义表达能力，表3展示了与扩散模型和我们的教师模型Meis-sonic [4]的比较。实验结果证实了我们之前的观察：我们的一步生成器在性能上可与使用16到32步的教师模型相媲美，而教师模型的性能会随着步数的减少而迅速下降。

局限性与未来工作。虽然我们的无数据蒸馏方法达到了教师模型级别的性能，但其应用目前仅限于我们当前的模型范围。未来，我们计划将我们的方法扩展到更强大的多模态扩散（MDM）教师模型，特别是用于图像和文本生成任务。虽然我们的方法无需数据，但引入真实数据来提升单步学生模型的性能，使其超越教师模型，可能会带来益处。最后，虽然我们利用广义杰弗里斯散度（Jeffreys divergence）来避免模式搜索行为，但我们打算探索更通用的 - 散度，以提高我们方法的灵活性和有效性。

图6. 从教师模型Meissonic中提炼的我们方法的定性结果。相应的提示词可在附录H中找到

总结

在这项工作中，我们提出了Di[M]O，这是一种新颖的方法，它利用词元级别的分布匹配将多步扩散模型（MDMs）的推理过程提炼为一步。具体而言，受策略内提炼（on-policy distillation）概念的启发，我们对基于从学生模型的一步生成中获得的伪中间状态的词元级别分布进行匹配。我们还对初始化策略和提炼目标的选择进行了广泛的实验，以提高的鲁棒性。我们的实验结果表明，我们提炼后的模型生成的图像质量与教师模型相当，而在推理过程中仅需1次采样步骤。这项工作展示了分布匹配方法在多步扩散模型提炼方面的强大能力，为探索离散数据高效生成的研究群体做出了贡献。