[论文解读] BOOT: Data-free Distillation of Denoising Diffusion Models with Bootstrapping

Jiatao Gu, Shuangfei Zhai+3 more|arXiv (Cornell University)|Jun 8, 2023

Advanced Neuroimaging Techniques and Applications被引用 4

一句话总结

BOOT 提出了一种用于去噪扩散模型的数据无关知识蒸馏方法，通过连续采样步骤的自举（bootstrapping）训练单步学生模型，无需真实数据或大规模离线数据生成。该方法在实现与教师模型相当的图像质量的同时，使推理速度提升数个数量级，即使在大型文本到图像模型（如 Stable Diffusion 和 DeepFloyd IF）上也表现优异。

ABSTRACT

Diffusion models have demonstrated excellent potential for generating diverse images. However, their performance often suffers from slow generation due to iterative denoising. Knowledge distillation has been recently proposed as a remedy that can reduce the number of inference steps to one or a few without significant quality degradation. However, existing distillation methods either require significant amounts of offline computation for generating synthetic training data from the teacher model or need to perform expensive online learning with the help of real data. In this work, we present a novel technique called BOOT, that overcomes these limitations with an efficient data-free distillation algorithm. The core idea is to learn a time-conditioned model that predicts the output of a pre-trained diffusion model teacher given any time step. Such a model can be efficiently trained based on bootstrapping from two consecutive sampled steps. Furthermore, our method can be easily adapted to large-scale text-to-image diffusion models, which are challenging for conventional methods given the fact that the training sets are often large and difficult to access. We demonstrate the effectiveness of our approach on several benchmark datasets in the DDIM setting, achieving comparable generation quality while being orders of magnitude faster than the diffusion teacher. The text-to-image results show that the proposed approach is able to handle highly complex distributions, shedding light on more efficient generative modeling.

研究动机与目标

• 为解决扩散模型推理速度缓慢的问题，其需要在许多步骤中迭代去噪。
• 克服现有蒸馏方法的局限性，这些方法需要大规模合成数据生成或访问真实训练数据。
• 实现大规模文本到图像扩散模型的高效蒸馏，这些模型因隐私或规模限制而无法获取训练数据。
• 开发一种无需数据的高效训练范式，利用连续时间步预测的自举机制。
• 实现与完整扩散教师模型质量相当的高保真度单步生成，同时显著缩短推理时间。

提出的方法

• 该方法训练一个时间条件的学生模型，基于单一噪声输入，预测预训练扩散教师模型在所有时间步的输出。
• 它使用一种新颖的 Signal-ODE，源自 PF-ODE，以稳定训练并提升跨时间步的泛化能力。
• 通过训练模型基于 $\bm{x}_{t'}$（其中 $t' > t$）的知识来预测 $\bm{x}_t$，实现自举，利用同一扩散轨迹上的确定性映射。
• 在训练过程中定期应用边界损失，以提升采样质量和多样性，同时降低频率以保持效率。
• 该方法与无分类器引导（CFG）兼容，可有效蒸馏条件生成和文本到图像模型。
• 学生模型仅通过教师模型的前向传播进行端到端训练，无需真实数据，适用于隐私敏感或大规模模型。

实验结果

研究问题

• RQ1 能否在不访问真实训练数据或大规模合成数据生成的情况下，有效训练单步学生模型？
• RQ2 如何通过连续时间步预测的自举提升数据无关设置下的蒸馏效率与性能？
• RQ3 所提方法能否实现与完整扩散教师模型相当的高保真度图像生成，同时实现数个数量级的推理加速？
• RQ4 该方法在复杂、高容量的文本到图像扩散模型（如 Stable Diffusion 和 DeepFloyd IF）上表现如何？
• RQ5 Signal-ODE 和边界损失在稳定训练和提升采样多样性方面起到什么作用？

主要发现

• BOOT 在无条件及类别条件图像生成基准测试（包括 FFHQ、LSUN 和 ImageNet）上，生成质量与完整扩散教师模型相当。
• 该方法实现了单步推理，推理速度相比教师模型提升 50–1000×，具体取决于数据集和设置。
• 在文本到图像生成任务中，BOOT 有效蒸馏了 DeepFloyd IF 和 Stable Diffusion 模型，无需访问训练数据即可从多样化提示生成高质量图像。
• 蒸馏后的模型保持了高多样性和保真度，通过在教师和学生模型间使用相同初始噪声进行定性对比得以验证。
• 每 4 个训练步骤应用一次边界损失，可提升采样质量和多样性，且训练成本增加有限。
• 该方法计算效率高，仅需 8 张 A100 GPU 和 3–7 天训练时间（适用于最大 1B 参数的模型），且除文本提示外无额外数据加载开销。