Improved Distribution Matching Distillation for Fast Image Synthesis

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

本论文旨在解决分布匹配蒸馏（DMD）在大规模文本到图像合成中的高成本问题，同时提高生成质量。

DMD2是一种改进的DMD方法，通过消除回归损失和集成GAN损失来提高训练稳定性和生成质量，同时还解决了多步采样中的训练-推理输入不匹配问题。

DMD2在ImageNet-64x64和COCO 2014数据集上取得了1.28和8.35的FID分数，与原始教师模型相比，在推理成本减少500倍的情况下，实现了更高的生成质量。此外，DMD2还可以生成分辨率更高的图像，并且开源了代码。

相关研究包括GAN、VAE、PixelCNN等在图像生成领域的方法，以及其他分布匹配蒸馏的改进方法，如FD、FD-GAN等。