[论文解读] On Distillation of Guided Diffusion Models

Chenlin Meng, Rombach, Robin+5 more|arXiv (Cornell University)|Oct 6, 2022

Advanced Neuroimaging Techniques and Applications被引用 8

一句话总结

本文提出一种两阶段蒸馏框架，以加速无分类器引导扩散模型，在仅1至4个去噪步骤内实现高保真图像生成。首先将条件与无条件教师模型的联合输出蒸馏到单一学生模型中，随后逐步将该模型蒸馏为更少的推理步数。该方法在保持ImageNet与LAION基准上竞争力的FID与IS分数的同时，实现比标准扩散模型快达256倍的采样速度。

ABSTRACT

Classifier-free guided diffusion models have recently been shown to be highly effective at high-resolution image generation, and they have been widely used in large-scale diffusion frameworks including DALLE-2, Stable Diffusion and Imagen. However, a downside of classifier-free guided diffusion models is that they are computationally expensive at inference time since they require evaluating two diffusion models, a class-conditional model and an unconditional model, tens to hundreds of times. To deal with this limitation, we propose an approach to distilling classifier-free guided diffusion models into models that are fast to sample from: Given a pre-trained classifier-free guided model, we first learn a single model to match the output of the combined conditional and unconditional models, and then we progressively distill that model to a diffusion model that requires much fewer sampling steps. For standard diffusion models trained on the pixel-space, our approach is able to generate images visually comparable to that of the original model using as few as 4 sampling steps on ImageNet 64x64 and CIFAR-10, achieving FID/IS scores comparable to that of the original model while being up to 256 times faster to sample from. For diffusion models trained on the latent-space (e.g., Stable Diffusion), our approach is able to generate high-fidelity images using as few as 1 to 4 denoising steps, accelerating inference by at least 10-fold compared to existing methods on ImageNet 256x256 and LAION datasets. We further demonstrate the effectiveness of our approach on text-guided image editing and inpainting, where our distilled model is able to generate high-quality results using as few as 2-4 denoising steps.

研究动机与目标

• 为解决无分类器引导扩散模型的高计算成本问题，此类模型在推理时需对两个独立模型进行数十至数百次评估。
• 开发一种蒸馏方法，使蒸馏后的模型在不损失图像质量或引导灵活性的前提下实现快速采样。
• 将蒸馏技术扩展至像素空间与潜在空间扩散模型，包括SOTA模型如Stable Diffusion。
• 在极少数采样步数下实现高效的文本到图像生成、图像编辑与图像修复。

提出的方法

• 引入一个单一学生模型，学习匹配教师模型中条件与无条件扩散模型的联合输出。
• 采用两阶段蒸馏流程：首先训练学生模型以模仿教师模型的输出，随后逐步将模型蒸馏为所需更少采样步数的版本。
• 利用现有的确定性采样方法，并提出一种随机采样过程，以在推理时使用蒸馏后的模型。
• 将该蒸馏框架应用于像素空间模型（如在CIFAR-10和ImageNet 64×64上）以及潜在空间模型（如在ImageNet 256×256和LAION 512×512上的Stable Diffusion）。
• 训练蒸馏后的模型以支持一系列引导权重，从而在样本多样性与质量之间实现可调的权衡。
• 在蒸馏过程中使用加权MSE损失函数，使学生模型的预测结果与教师模型在不同引导强度下的输出保持一致。

实验结果

研究问题

• RQ1 蒸馏能否有效应用于需要从两个模型进行推理的无分类器引导扩散模型？
• RQ2 单一蒸馏模型是否能在仅1至4个去噪步骤内实现高质量图像生成，同时保持引导灵活性？
• RQ3 所提出的两阶段蒸馏方法是否在FID与IS分数上与原始模型相当或更优？
• RQ4 该蒸馏模型是否能泛化至多样化的扩散模型架构，包括像素空间与潜在空间变体？
• RQ5 该方法是否能在文本到图像生成、图像编辑与图像修复等下游任务中，以极少采样步数取得成功应用？

主要发现

• 在仅使用4个采样步骤的情况下，蒸馏模型在ImageNet 64×64与CIFAR-10上达到与教师模型相当的FID分数，推理速度最高提升256倍。
• 在Stable Diffusion等潜在空间模型上，该方法将采样步数减少至1–4步，在ImageNet 256×256与LAION 512×512上实现比现有方法快至少10倍的推理速度。
• 该蒸馏模型在不同引导权重（w = 0 至 w = 4）下均保持高视觉质量，有效实现多样性与保真度之间的权衡。
• 在文本引导的图像编辑与图像修复任务中，该方法仅使用2–4个去噪步骤即可生成高质量结果，展现出在编辑任务中的强大性能。
• 与确定性采样相比，随机采样过程进一步提升了低步数下的样本质量。
• 该方法是首个在像素空间与潜在空间无分类器引导扩散模型中均实现有效蒸馏的工作，在各类基准上均表现出一致的性能提升。