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

百家 作者:新智元 2022-10-20 13:46:37



  新智元报道  

编辑:Aeneas David
【新智元导读】斯坦福、谷歌大脑新作:无需分类器,两步蒸馏,将扩散模型采样速度提升256倍。

最近,无分类器的指导扩散模型(classifier-free guided diffusion models)在高分辨率图像生成方面非常有效,并且已经被广泛用于大规模扩散框架,包括DALL-E 2、GLIDE和Imagen。


然而,无分类器指导扩散模型的一个缺点是它们在推理时的计算成本很高。因为它们需要评估两个扩散模型——一个类别条件模型(class-conditional model) 和一个无条件模型(unconditional model),而且需要评估数百次。


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

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

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

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

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

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

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

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

扩散模型的背景

通过来自数据分布的样本x,噪声调度函数研究者通过最小化加权均方差来训练了具有参数θ的扩散模型
其中是信噪比,是预先指定的加权函数。
一旦训练了扩散模型,就可以使用离散时间DDIM采样器从模型中采样。
具体来说,DDIM采样器从 z1 ∼ N (0,I)开始,更新如下
其中,N是采样步骤的总数。使用,会生成最终样本。
无分类器指导是一种有效的方法,可以显著提高条件扩散模型的样本质量,已经广泛应用于包括GLIDE,DALL·E 2和Imagen。
它引入了一个指导权重参数来衡量样本的质量和多样性。为了生成样本,无分类器指导在每个更新步骤都会使用作为预测模型,来评估条件扩散模型和联合训练的
由于每次采样更新都需要评估两个扩散模型,因此使用无分类器指导进行采样通常很昂贵。
为了解决这个问题,研究者使用了渐进式蒸馏(progressive distillation)  ,这是一种通过重复蒸馏提高扩散模型采样速度的方法。
在以前,这种方法不能直接被直接用在引导模型的蒸馏上,也不能在确定性DDIM采样器以外的采样器上使用。而在这篇论文中,研究者解决了这些问题。

蒸馏无分类器的指导扩散模型

他们的办法是,将无分类器的指导扩散模型进行蒸馏。
对于一个训练有素的教师引导模型,他们采取了两个步骤。

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


其中
为了结合指导权重w,研究者引入了w条件模型,其中w作为学生模型的输入。为了更好地捕捉特征,他们将傅里叶嵌入应用w,然后用Kingma等人使用的时间步长的方式,把它合并到扩散模型的主干中。
由于初始化在性能中起着关键作用,研究者初始化学生模型时,使用的是与教师条件模型相同的参数(除了新引入的与w-conditioning相关的参数)。
第二步,研究者设想了一个离散的时间步长场景,并且通过每次将采样步数减半,逐步将学习模型从第⼀步蒸馏成具有可学习参数η2、步⻓更少的学⽣模型
其中,N表⽰采样步骤的数量,对于,研究者开始训练学生模型,让它用一步来匹配教师模型的两步DDIM采样的输出(例如:从t/N到t - 0.5/N,从t - 0.5/N到t - 1/N)。
将教师模型中的2N个步骤蒸馏成学生模型中的N个步骤以后,我们可以将新的N-step学生模型作为新的教师模型,然后重复同样的过程,将教师模型蒸馏成N/2-step的学生模型。在每⼀步,研究者都会⽤教师模型的参数来初始化学⽣模型。
N-step的确定性和随机采样
⼀旦模型被训练出来,对于,研究者就可以通过DDIM更新规则来执行采样。研究者注意到,对于蒸馏模型,这个采样过程在给定初始化的情况下是确定的。
另外,研究者也可以进行N步的随机采样。使用两倍于原始步长的确定性采样步骤( 即与N/2-step确定性采样器相同),然后使用原始步长进行一次随机步回(即用噪声扰动)。
,当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/

公众号 关注网络尖刀微信公众号
随时掌握互联网精彩
赞助链接