Improved Distribution Matching Distillation for Fast Image Synthesis

最近的方法表明，将扩散模型提炼为高效的单步生成器是有前途的。其中，分布匹配蒸馏（DMD）生成与其教师在分布上相匹配的单步生成器，而不强制要求其与教师的采样轨迹一一对应。然而，为了确保稳定的训练，DMD需要使用由教师使用确定性采样器的多个步骤生成的大量噪声图像对计算额外的回归损失。这对于大规模的文本到图像合成来说是昂贵的，并限制了学生的质量，使其与教师的原始采样路径过于紧密相连。我们引入了DMD2，一组技术来解决这个问题并改进DMD的训练。首先，我们消除了回归损失和昂贵数据集构建的需要。我们表明，由于虚假评论家无法准确估计生成样本的分布，导致了结果的不稳定性，并提出了一个两个时间尺度的更新规则作为补救措施。其次，我们将GAN损失集成到蒸馏过程中，区分生成的样本和真实图像。这使我们能够在真实数据上训练学生模型，减轻教师模型对真实得分估计的不完美，并提高质量。最后，我们修改了训练过程，以实现多步采样。我们通过在训练时间模拟推理时间生成器样本来解决此设置中的训练推理输入不匹配问题。总的来说，我们的改进在单步图像生成方面设置了新的基准，ImageNet-64x64的FID得分为1.28，零样本COCO 2014的FID得分为8.35，尽管推理成本降低了500倍，但超过了原始教师。此外，我们展示了我们的方法可以通过蒸馏SDXL生成百万像素的图像，展示了在少步方法中的卓越视觉质量。

本文旨在解决Distribution Matching Distillation (DMD)在大规模文本到图像合成中的昂贵数据集构建和限制学生模型质量的问题，提出DMD2的改进方案。

DMD2方案采取两个策略，一是消除回归损失和昂贵数据集构建的需求，通过两个时间尺度的更新规则解决不准确的样本分布估计问题；二是将GAN损失集成到蒸馏过程中，区分生成样本和真实图像，并在真实数据上训练学生模型，提高质量。同时，修改训练过程以支持多步采样。

实验结果表明，DMD2在ImageNet-64x64和零样本COCO 2014上的FID得分分别为1.28和8.35，即使推理成本降低了500倍，也能超越原始教师。此外，DMD2还可以通过蒸馏SDXL生成百万像素的图像，显示出与少步方法相比出色的视觉质量。

与本文相关的研究包括：对蒸馏的改进，如Deep Mutual Learning和Variational Information Distillation；对图像生成的其他方法，如VAE和PixelCNN；以及对GAN的改进，如WGAN和StyleGAN。