斯坦福/谷歌大脑：两次蒸馏，引导扩散模型采样提速256倍！

---

  新智元报道  

为了解决这个问题，斯坦福大学和谷歌大脑的学者提出使用两步蒸馏(two-step distillation)的方法来提升无分类器指导扩散模型的采样效率。

论文地址：https://arxiv.org/abs/2210.03142

如何将无分类器指导扩散模型提炼成快速采样的模型？

首先，对于一个预先训练好的无分类器指导模型，研究者首先学习了一个单一的模型，来匹配条件模型和无条件模型的组合输出。

随后，研究者逐步将这个模型蒸馏成一个采样步骤更少的扩散模型。

可以看到，在ImageNet 64x64和CIFAR-10上，这种方法能够在视觉上生成与原始模型相当的图像。

只需4个采样步骤，就能获得与原始模型相当的FID/IS分数，而采样速度却高达256倍。

可以看到，通过改变指导权重w，研究者蒸馏的模型能够在样本多样性和质量之间进行权衡。而且只用一个取样步骤，就能获得视觉上愉悦的结果。

通过来自数据分布的样本x，噪声调度函数研究者通过最小化加权均方差来训练了具有参数θ的扩散模型。

其中是信噪比，和是预先指定的加权函数。

一旦训练了扩散模型，就可以使用离散时间DDIM采样器从模型中采样。

其中，N是采样步骤的总数。使用，会生成最终样本。

它引入了一个指导权重参数来衡量样本的质量和多样性。为了生成样本，无分类器指导在每个更新步骤都会使用作为预测模型，来评估条件扩散模型和联合训练的。

对于一个训练有素的教师引导模型，他们采取了两个步骤。

第一步，研究者引入了一个连续时间的学生模型，它具有可学习的参数η1，来匹配教师模型在任意时间步长t ∈ [0, 1] 的输出。指定一系列他们有兴趣的指导强度后，他们使用以下目标来优化学生模型。

其中。

第二步，研究者设想了一个离散的时间步长场景，并且通过每次将采样步数减半，逐步将学习模型从第⼀步蒸馏成具有可学习参数η2、步⻓更少的学⽣模型。

其中，N表⽰采样步骤的数量，对于和，研究者开始训练学生模型，让它用一步来匹配教师模型的两步DDIM采样的输出（例如：从t/N到t - 0.5/N，从t - 0.5/N到t - 1/N）。

⼀旦模型被训练出来，对于，研究者就可以通过DDIM更新规则来执行采样。研究者注意到，对于蒸馏模型，这个采样过程在给定初始化的情况下是确定的。

，当t > 1/N时，可用以下的更新规则——

其中，。

当t=1/N时，研究者使用确定性更新公式，从得出。

模型实验在两个标准数据集上进行：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

关注公众号：拾黑（shiheibook）了解更多

[广告]赞助链接：

四季很好，只要有你，文娱排行榜：https://www.yaopaiming.com/
让资讯触达的更精准有趣：https://www.0xu.cn/