当前位置: 首页 > 科技观察

Stanford-GoogleBrain:双重蒸馏,引导扩散模型采样提速256倍!

时间:2023-03-13 11:33:28 科技观察

最近,无分类器引导扩散模型(classifier-freeguideddiffusionmodels)在高分辨率图像生成中非常有效,并被广泛应用于大规模扩散框架,包括DALL-E2、GLIDE和Imagen.然而,无分类器引导扩散模型的一个缺点是它们在推理时的计算量很大。因为他们需要评估两个扩散模型——一个类条件模型和一个无条件模型——数百次。为了解决这个问题,斯坦福大学和谷歌大脑的学者提出了使用两步蒸馏的方法来提高非分类器引导扩散模型的采样效率。论文地址:https://arxiv.org/abs/2210.03142如何将非分类器引导扩散模型细化为快速采样模型?首先,对于没有分类器指导的预训练模型,我们首先学习一个单一的模型来匹配条件和非条件模型的组合输出。随后,研究人员逐渐将这个模型提炼为采样步骤更少的扩散模型。可以看出,在ImageNet64x64和CIFAR-10上,该方法能够生成视觉上与原始模型相当的图像。只需4个采样步骤,即可获得与原始模型相当的FID/IS分数,同时采样速度提高了256倍。可以看出,通过改变指导权重w,研究人员提炼出的模型能够在样本多样性和质量之间进行权衡。只需一个采样步骤,即可获得视觉上令人愉悦的结果。扩散模型背景使用来自数据分布的样本x,噪声调度函数研究人员通过最小化加权均方误差来训练具有参数θ的扩散模型。其中是信噪比,是预先指定的加权函数。训练扩散模型后,可以使用离散时间DDIM采样器从模型中采样。具体来说,DDIM采样器从z1~N(0,I)开始,更新如下,其中N是采样步数的总数。使用,生成最终样本。Classifier-freeguidance是一种可以显着提高条件扩散模型样本质量的有效方法,已被广泛应用,包括GLIDE、DALL·E2和Imagen。它引入了一个指导权重参数来衡量样本的质量和多样性。为了生成样本,在每个更新步骤中使用无分类器指导作为预测模型来评估条件扩散模型和联合训练模型。使用无分类器指导进行采样通常很昂贵,因为每次采样更新都需要评估两个扩散模型。为了解决这个问题,研究人员使用了渐进式蒸馏,这是一种通过重复蒸馏来提高扩散模型采样速度的方法。以前,这种方法不能直接用于引导模型的蒸馏,也不能用于确定性DDIM采样器以外的采样器。在本文中,研究人员解决了这些问题。在没有分类器的情况下提取监督扩散模型他们的方法是在没有分类器的情况下提取监督扩散模型。对于训练有素的教师指导模型,他们采取两个步骤。在第一步中,研究人员引入了一个具有可学习参数η1的连续时间学生模型,以匹配教师模型在任何时间步长t∈[0,1]的输出。在指定了他们感兴趣的一系列教学优势后,他们使用以下目标优化了学生模型。in.为了结合指导权重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步学生模型作为新的教师模型,然后重复同样的过程将教师模型提炼成N/2-学生模型步骤模型。在每个步骤中,研究人员使用教师模型的参数初始化学生模型。N步确定性和随机抽样模型训练完成后,研究人员可以通过DDIM更新规则进行抽样。研究人员指出,对于蒸馏模型,这种采样过程在给定初始化的情况下是确定性的。此外,研究人员还可以进行N步随机抽样。使用两倍于原始步长的确定性采样步长(即与N/2步确定性采样器相同),然后使用原始步长随机后退一步(即受噪声干扰)。,当t>1/N时,可以使用以下更新规则——其中,.当t=1/N时,研究人员使用派生自的确定性更新公式。值得注意的是,我们注意到执行随机抽样需要在与确定性抽样器略有不同的时间步长上评估模型,并且需要对边缘情况的训练算法进行小的修改。其他的蒸馏方法还有一种直接将progressivedistillation应用于bootstrapmodel的方法,即沿用teachermodel的结构,直接将studentmodel蒸馏成联合训练的conditional和unconditionalmodel。研究人员尝试后发现,这种方法效果不佳。实验和结论模型实验在两个标准数据集上进行:ImageNet(64*64)和CIFAR10。实验中探索了引导权重w的不同范围,并且观察到所有范围都具有可比性,因此[wmin,wmax]=[0,4]用于实验。第一步和第二步模型使用SNR损失进行训练。基线标准包括DDPM祖先采样和DDIM采样。为了更好地理解如何合并引导权重w,使用固定w值训练的模型作为参考。为了公平比较,实验对所有方法使用相同的预训练教师模型。使用U-Net(Ronnebergeretal.,2015)架构作为基线,并使用相同的U-Net主干,引入结构嵌入w作为两步学生模型。上图展示了所有方法在ImageNet64x64上的表现。其中D和S分别代表确定性和随机采样器。在实验中,以引导区间w∈[0,4]为条件的模型训练与具有固定w值的模型训练的性能相当。在更少的步骤中,我们的方法明显优于DDIM基线性能,并且在8到16步时几乎达到教师模型的性能水平。ImageNet64x64SamplingQualityAssessedbyFIDandISScoreCIFAR-10SamplingQualityAssessedbyFIDandISScore我们还提取了教师模型的编码过程,并对风格迁移进行了实验。具体来说,为了在两个域A和B之间执行风格转换,域A中的图像使用在域A上训练的扩散模型进行编码,然后使用在域B上训练的扩散模型进行解码。由于编码过程可以理解为采样DDIM的过程相反,我们提取无分类器指导的编码器和解码器,并与DDIM编码器和解码器进行比较,如上图所示。我们还探讨了信道强度w变化对性能的影响。总之,我们提出了一个引导扩散模型的蒸馏,以及一个从蒸馏模型中抽样的随机采样器。根据经验,我们的方法仅需一步即可实现视觉上的高体验采样,并且仅需8到16步即可获得与教师的FID/IS分数相当的分数。