斯坦福/谷歌大脑：两次蒸馏，引导扩散模型采样提速256倍！
新智元报道
新智元报道
【新智元导读】斯坦福、谷歌大脑新作：无需分类器，两步蒸馏，将扩散模型采样速度提升256倍。
然而，无分类器指导扩散模型的一个缺点是它们在推理时的计算成本很高。因为它们需要评估两个扩散模型——一个类别条件模型(class-conditional model) 和一个无条件模型(unconditional model)，而且需要评估数百次。
为了解决这个问题，斯坦福大学和谷歌大脑的学者提出使用两步蒸馏(two-step distillation)的方法来提升无分类器指导扩散模型的采样效率。
论文地址：https://arxiv.org/abs/2210.03142
如何将无分类器指导扩散模型提炼成快速采样的模型？
首先，对于一个预先训练好的无分类器指导模型，研究者首先学习了一个单一的模型，来匹配条件模型和无条件模型的组合输出。
随后，研究者逐步将这个模型蒸馏成一个采样步骤更少的扩散模型。
可以看到，在ImageNet 64x64和CIFAR-10上，这种方法能够在视觉上生成与原始模型相当的图像。
只需4个采样步骤，就能获得与原始模型相当的FID/IS分数，而采样速度却高达256倍。
可以看到，通过改变指导权重w，研究者蒸馏的模型能够在样本多样性和质量之间进行权衡。而且只用一个取样步骤，就能获得视觉上愉悦的结果。
蒸馏无分类器的指导扩散模型
第一步，研究者引入了一个连续时间的学生模型，它具有可学习的参数η1，来匹配教师模型在任意时间步长t ∈ [0, 1] 的输出。指定一系列他们有兴趣的指导强度
后，他们使用以下目标来优化学生模型。
模型实验在两个标准数据集上进行：ImageNet（64*64）和 CIFAR 10。
实验中探索了指导权重w的不同范围，并观察到所有的范围都有可比性，因此使用[wmin, wmax] = [0, 4]进行实验。使用信噪比损失训练第一步和第二步模型。
基线标准包括DDPM ancestral采样和DDIM采样。
为了更好地理解如何纳入指导权重w，使用一个固定的w值训练的模型作为参照。
为了进行公平比较，实验对所有的方法使用相同的预训练教师模型。使用U-Net（Ronneberger等人，2015）架构作为基线，并使用相同的U-Net主干，引入嵌入了w的结构作为两步学生模型。
上图为所有方法在ImageNet 64x64上的表现。其中D和S分别代表确定性和随机性采样器。
在实验中，以指导区间w∈[0, 4]为条件的模型训练，与w为固定值的模型训练表现相当。在步骤较少时，我们的方法明显优于DDIM基线性能，在8到16个步骤下基本达到教师模型的性能水平。
由FID和IS分数评估的ImageNet 64x64采样质量
由FID和IS评分评估的CIFAR-10采样质量
我们还对教师模型的编码过程进行蒸馏，并进行了风格转移的实验。具体来说，为了在两个领域A和B之间进行风格转换，用在领域A上训练的扩散模型对领域A的图像进行编码，然后用在领域B上训练的扩散模型进行解码。
由于编码过程可以理解为颠倒了的DDIM的采样过程，我们对具有无分类器指导的编码器和解码器都进行了蒸馏，并与DDIM编码器和解码器进行比较，如上图所示。我们还探讨了对引导强度w的改动对性能的影响。
总之，我们提出的引导扩散模型的蒸馏方法，以及一种随机采样器，从蒸馏后的模型中采样。从经验上看，我们的方法只用了一个步骤就能实现视觉上的高体验采样，只用8到16个步骤就能获得与教师相当的FID/IS分数。
https://twitter.com/chenlin_meng/status/1579384412068016128
https://www.reddit.com/r/MachineLearning/comments/y0iu5w/new_distilled_diffusion_models_research_can/
https://arxiv.org/abs/2210.03142
微信扫码关注该文公众号作者