大数跨境
0
0

采样提速256倍,蒸馏扩散模型生成图像质量媲美教师模型,只需4步

采样提速256倍,蒸馏扩散模型生成图像质量媲美教师模型,只需4步 极市平台
2022-10-11
0
↑ 点击蓝字 关注极市平台

作者丨杜伟、陈萍
来源丨机器之心
编辑丨极市平台

极市导读

 

斯坦福大学联合谷歌大脑使用「两步蒸馏方法」提升无分类器指导的采样效率,在生成样本质量和采样速度上都有非常亮眼的表现。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

去噪扩散概率模型(DDPM)在图像生成、音频合成、分子生成和似然估计领域都已经实现了 SOTA 性能。同时无分类器(classifier-free)指导进一步提升了扩散模型的样本质量,并已被广泛应用在包括 GLIDE、DALL·E 2 和 Imagen 在内的大规模扩散模型框架中。

然而,无分类器指导的一大关键局限是它的采样效率低下,需要对两个扩散模型评估数百次才能生成一个样本。这一局限阻碍了无分类指导模型在真实世界设置中的应用。尽管已经针对扩散模型提出了蒸馏方法,但目前这些方法不适用无分类器指导扩散模型。

为了解决这一问题,近日斯坦福大学和谷歌大脑的研究者在论文《On Distillation of Guided Diffusion Models》中提出使用两步蒸馏(two-step distillation)方法来提升无分类器指导的采样效率。

在第一步中,他们引入单一学生模型来匹配两个教师扩散模型的组合输出;在第二步中,他们利用提出的方法逐渐地将从第一步学得的模型蒸馏为更少步骤的模型。

利用提出的方法,单个蒸馏模型能够处理各种不同的指导强度,从而高效地对样本质量和多样性进行权衡。此外为了从他们的模型中采样,研究者考虑了文献中已有的确定性采样器,并进一步提出了随机采样过程。

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

研究者在 ImageNet 64x64 和 CIFAR-10 上进行了实验,结果表明提出的蒸馏模型只需 4 步就能生成在视觉上与教师模型媲美的样本,并且在更广泛的指导强度上只需 8 到 16 步就能实现与教师模型媲美的 FID/IS 分数,具体如下图 1 所示。

此外,在 ImageNet 64x64 上的其他实验结果也表明了,研究者提出的框架在风格迁移应用中也表现良好。

方法介绍

接下来本文讨论了蒸馏无分类器指导扩散模型的方法 (distilling a classifier-free guided diffusion model) 。给定一个训练好的指导模型, 即教师模型 之后本文分两步完成。

第一步引入一个连续时间学生模型 , 该模型具有可学习参数 , 以匹配教师模型在任意时间步 处的输出。给定一个优化范围 [w_min, , 对学生模型进行优化:

其中, 。为了合并指导权重 , 本文引入了一个 条件模 型, 其中 作为学生模型的输入。为了更好地捕捉特征, 本文还对 应用傅里叶嵌入。此外, 由于初始化在模型性能中起着关键作用, 因此本文初始化学生模型的参数与教师模型相同。

在第二步中, 本文将离散时间步(discrete time-step)考虑在内, 并逐步将第一步中的蒸馏模型 转化为步数较短的学生模型 , 其可学习参数为 2, 每次采样步数减半。设 为采样步数, 给定 max , 然后根据 Salimans & Ho 等人 提出的方法训练学生模型。在将教师模型中的 步蒸馏为学生模型中的 步之后, 之后使用 步学生模型作为新的教师模型, 这个过程不断重复, 直到将教师模型蒸馏为 步学生模型。

步可确定性和随机采样:一旦模型 训练完成,给定一个指定的 [w_min, w_max], 然后使用 DDIM 更新规则执行采样。

实际上, 本文也可以执行 步随机采样, 使用两倍于原始步长的确定性采样步骤, 然后使用原始步 长向后执行一个随机步骤。对于 , 当 时, 本文使用以下更新规则

实验

实验评估了蒸馏方法的性能,本文主要关注模型在 ImageNet 64x64 和 CIFAR-10 上的结果。他们探索了指导权重的不同范围,并观察到所有范围都具有可比性,因此实验采用 [w_min, w_max] = [0, 4]。图 2 和表 1 报告了在 ImageNet 64x64 上所有方法的性能。

本文还进行了如下实验。具体来说,为了在两个域 A 和 B 之间执行风格迁移,本文使用在域 A 上训练的扩散模型对来自域 A 的图像进行编码,然后使用在域 B 上训练的扩散模型进行解码。由于编码过程可以理解为反向 DDIM 采样过程,本文在无分类器指导下对编码器和解码器进行蒸馏,并与下图 3 中的 DDIM 编码器和解码器进行比较。

本文还探讨了如何修改指导强度 w 以影响性能,如下图 4 所示。

公众号后台回复“开学”获取CVPR、ICCV、VALSE等资源合集~

极市干货
算法竞赛:国际赛事证书,220G数据集开放下载!ACCV2022国际细粒度图像分析挑战赛开赛!
技术综述BEV 学术界和工业界方案、优化方法与tricks综述PyTorch下的可视化工具(网络结构/训练过程可视化)
极视角动态:极视角与华为联合发布基于昇腾AI的「AICE赋能行业解决方案」算法误报怎么办?自训练工具使得算法迭代效率提升50%!

CV技术社群邀请函 #



△长按添加极市小助手
添加极市小助手微信(ID : cvmart2)

备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳)


即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群


极市&深大CV技术交流群已创建,欢迎深大校友加入,在群内自由交流学术心得,分享学术讯息,共建良好的技术交流氛围。

点击阅读原文进入CV社区

收获更多技术干货

【声明】内容源于网络
0
0
极市平台
为计算机视觉开发者提供全流程算法开发训练平台,以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。
内容 8155
粉丝 0
极市平台 为计算机视觉开发者提供全流程算法开发训练平台,以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。
总阅读919
粉丝0
内容8.2k