Command Palette
Search for a command to run...
面向快速图像合成的改进分布匹配蒸馏方法
面向快速图像合成的改进分布匹配蒸馏方法
Tianwei Yin Michaël Gharbi Taesung Park Richard Zhang Eli Shechtman Frédo Durand William T. Freeman
摘要
近期研究在将昂贵的扩散模型蒸馏为高效的单步生成器方面展现出潜力。其中,分布匹配蒸馏(DMD)能够生成在分布上与其教师模型相匹配的单步生成器,即蒸馏过程并不强制要求与教师模型的采样轨迹一一对应。然而,为确保实际训练中的稳定性,DMD 需要借助一个额外的回归损失,该损失通过教师模型使用确定性采样器经过大量步数生成的大规模噪声–图像对来计算。这不仅在大规模文本到图像合成中计算开销巨大,还限制了学生模型的质量,使其过度依赖于教师模型原有的采样路径。我们提出了 DMD2,一套旨在突破这一限制并改进 DMD 训练的技术。首先,我们消除了回归损失及昂贵数据集构建的需求。我们表明,由此产生的不稳定性源于“虚假”判别器未能以足够精度估计生成样本的分布,并提出了一种双时间尺度更新规则作为解决方案。其次,我们将 GAN 损失集成到蒸馏过程中,对生成样本与真实图像进行判别。这使我们能够在真实数据上训练学生模型,从而减轻教师模型对“真实”分数的不完美估计,进而提升生成质量。第三,我们引入了一种新的训练流程,使学生模型能够进行多步采样,并通过在训练期间模拟推理时的生成器样本,解决了先前工作中训练与推理输入不匹配的问题。综合来看,我们的改进为单步图像生成设立了新的基准:在 ImageNet-64×64 上取得了 1.28 的 FID 分数,在零样本 COCO 2014 上取得了 8.35 的 FID 分数,在推理成本降低 500 倍的情况下超越了原始教师模型。此外,我们展示了该方法能够通过蒸馏 SDXL 生成百万像素级图像,在少步方法中展现出卓越的视觉质量,并超越了教师模型。我们已开源代码和预训练模型。
一句话总结
DMD2 由 MIT 和 Adobe Research 的研究人员提出,通过双时间尺度 critic 更新消除回归损失,加入 GAN 损失以利用真实数据训练,并支持多步采样,在一轮生成中达到当前最佳水平,ImageNet-64×64 上 FID 为 1.28,零样本 COCO 2014 上 FID 为 8.35,超越教师模型,同时推理成本降低 500×。
核心贡献
- 该方法消除了 DMD 中的回归损失以及昂贵的数据集构建,并识别出虚假 critic 分布估计不准确的问题,提出双时间尺度更新规则作为补救措施,从而解决了由此带来的不稳定性。
- 集成了 GAN 损失,用于区分生成样本与真实图像,支持在真实数据上训练,减轻教师模型真实分数估计的不完善,并提升视觉质量。
- 新的训练流程在学生模型中引入多步采样,通过在训练期间模拟推理时的生成器样本来解决训练与推理输入不匹配的问题,在 ImageNet-64×64 上 FID 达到 1.28,在零样本 COCO 2014 上 FID 达到 8.35,创下新基准。
引言
扩散模型能够生成高质量的视觉内容,但需要大量迭代去噪步骤,导致高分辨率合成速度慢且成本高。先前的蒸馏方法通过训练少量步骤的学生模型来加速采样,但往往会降低质量。分布匹配蒸馏(DMD)曾在一轮生成中取得当时最佳结果,但仍依赖回归损失,这需要预先计算数百万个噪声-图像对,且将学生模型质量限制在教师模型水平。作者提出 DMD2,一种改进的分布匹配方法,完全消除回归损失,同时通过双时间尺度更新规则保持训练稳定性。进一步集成 GAN 目标,在分布层面操作,超越教师模型的质量,并设计反向模拟以支持多步生成,避免以往方法中常见的域不匹配问题。最终得到一个快速、可扩展的生成器,在一轮和少步图像合成中创下新基准。
方法
作者提出一种改进的分布匹配蒸馏框架,将昂贵的多步扩散模型蒸馏为高效的小步数生成器。
首先,移除标准 DMD 中使用的回归损失,该损失需要昂贵的预计算噪声-图像对。为在没有该正则化的情况下稳定训练,采用双时间尺度更新规则。具体来说,虚假分数估计器的更新频率高于生成器,以确保其准确跟踪生成器非平稳的输出分布。
为超越教师模型的质量,作者集成了 GAN 目标。参考框架图:
该架构包含一个小步数生成器 Gθ 用于生成虚假图像。这些图像被扩散,并与扩散后的真实图像通过冻结的真实分数函数(教师)和动态训练的虚假分数函数进行比较。分布匹配目标的梯度由这些分数之间的差异计算得出。此外,在虚假扩散去噪器的瓶颈层上添加了一个判别器分支。该判别器训练用于区分真实和虚假扩散样本,而生成器则最小化标准的非饱和 GAN 损失。
对于更大的模型,作者将方法扩展以支持使用固定时间步调度进行多步生成。在推理过程中,模型在去噪和注入噪声之间交替进行。然而,在噪声真实图像上训练多步生成器会导致与推理时不匹配,因为推理时的输入来自前一步生成器的输出。如下图所示:
作者通过在训练期间模拟推理时的反向过程来解决这一问题。不使用真实图像的前向扩散,而是通过运行当前学生生成器若干步来生成噪声合成图像。然后生成器对这些模拟图像进行去噪,确保训练与测试域之间的一致性。
最后,训练过程在两个主要步骤之间交替进行:使用分布匹配梯度和 GAN 损失优化生成器,以及使用虚假数据上的去噪分数匹配目标和 GAN 分类损失优化虚假分数估计器及其附加的判别器。
实验
实验在类条件 ImageNet-64×64 和零样本文本到图像合成 COCO 上评估 DMD2,使用 FID、CLIP 分数和人类偏好研究。蒸馏后的 SDXL 和 SDv1.5 模型仅需 1–4 个采样步数即可在图像质量和提示词一致性上匹配或超越教师扩散模型,同时大幅降低计算量。消融实验结果表明,移除原始 DMD 回归损失并采用双时间尺度更新规则及附加的 GAN 项,可防止训练不稳定和过饱和输出,并且集成的分布匹配与 GAN 框架相比任一单独组件都能实现更高的真实感和文本一致性。
在 ImageNet-64×64 上,单步 StyleGAN-XL 在先前的一轮方法中取得最低 FID(1.52),而多步扩散模型如 RIN 需要 1000 步才达到 1.23。现有蒸馏技术表现不佳,FID 分数均高于 15。所提出的方法优于所有蒸馏基线,甚至超越教师扩散模型的 ODE 采样器。StyleGAN-XL 在所有单步基线中取得最佳 FID(1.52),但仍落后于多步扩散模型如 RIN(1.23)和 ADM(2.07)。先前的蒸馏方法(渐进式蒸馏、BOOT)FID 高于 15,而所提出的一轮模型显著优于它们以及教师扩散模型。
DMD2 在 COCO 2014 上以单次前向传播实现强大的零样本文本到图像生成,FID 超越 SDXL 教师及其他加速方法,同时保持有竞争力的 CLIP 分数。该一轮模型在 FID 上比教师模型提升 4.91 分,用户研究证实其在使用 25 倍更少步数的情况下,在 24% 的样本中图像质量优于教师模型。DMD2 单次前向传播在 FID(19.01)上取得所有方法中最低值,相较 SDXL 教师提升 4.91 分。DMD2 的 CLIP 分数(0.336)与 SDXL-Turbo(0.337)几乎相同,表明文本-图像对齐能力相当。在一项用户研究中,DMD2 四步生成器在 24% 的样本中图像质量优于其教师,同时使用 25 倍更少的前向传播。
从 DMD 中移除回归损失使 ImageNet FID 从 2.62 升至 3.48,表明训练不稳定。双时间尺度更新规则(TTUR)在无回归损失的情况下稳定训练,将 FID 恢复到 2.61,与原始 DMD 持平。在 TTUR 之上添加 GAN 损失将 FID 进一步改善至 1.51,比仅用 TTUR 的变体提升 1.1 分,是所有消融中的最佳结果。省略回归损失使 FID 从 2.62 增至 3.48,反映不稳定。TTUR 无回归损失将 FID 恢复到 2.61,与基线 DMD 持平。结合 TTUR 与 GAN 损失取得最佳 FID 1.51,优于在原始 DMD 上添加 GAN(2.56)和单独添加 TTUR(2.52)。
从 DMD2 中移除 GAN 损失导致 FID 急剧上升和图像过饱和。纯 GAN 变体(无分布匹配)获得最低 FID,但 patch 级和文本对齐分数最差,而省略反向模拟会降低 patch FID。完整的 DMD2 模型取得最佳 patch FID 和并列最高 CLIP,提供保真度与文本对齐的良好平衡。没有 GAN 项时,FID 恶化至 26.90,生成图像变得过饱和和过度平滑。消融分布匹配(纯 GAN)产生最低 FID(13.77),但 Patch FID 最差(27.96)且 CLIP 最低(0.307),表明文本对齐和 patch 级质量较差。移除反向模拟使 Patch FID 增至 24.21,显示其对图像保真度的重要性。完整 DMD2 取得最佳 Patch FID(20.86)和并列最高 CLIP(0.332),同时保持有竞争力的 FID(19.32)。
在 ImageNet-64×64 上,所提出的一轮模型显著优于先前的蒸馏方法,甚至超越教师扩散 ODE 采样器,尽管多步扩散模型仍实现最佳质量。对于零样本 COCO 上的文本到图像任务,DMD2 单次前向传播在 FID 上相比 SDXL 教师取得大幅提升,并在用户研究中受到偏好,同时使用 25 倍更少步数。消融研究揭示,移除回归损失导致训练不稳定,可通过 TTUR 恢复,而添加 GAN 损失进一步提升性能。在 DMD2 中,GAN 损失防止过饱和,而结合分布匹配、反向模拟和 GAN 的完整模型在 patch 级保真度和文本对齐之间提供最佳平衡。