Stable Diffusion采样速度翻倍!仅需10到25步的扩散模型采样算法

本文 机器之心编辑部 授权

清华大学计算机系朱军教授带领的 TSAIL 团队提出 DPM-Solver(NeurIPS 2022 Oral,约前 1.7%)和 DPM-Solver++,将扩散模型的快速采样算法提升到了极致:无需额外训练,仅需 10 到 25 步就可以获得极高质量的采样。

要说 AI 领域今年影响力最大的进展,爆火的 AI 作图绝对是其中之一。设计者只需要输入对图片的文字描述,就可以由 AI 生成一张质量极高的高分辨率图片。目前,使用范围最广的当属 StabilityAI 的开源模型 Stable Diffusion,模型一经开源就在社区引起了广泛的讨论。

然而,扩散模型在使用上最大的问题就是其极慢的采样速度。模型采样需要从纯噪声图片出发,一步一步不断地去噪,最终得到清晰的图片。在这个过程中,模型必须串行地计算至少 50 到 100 步才可以获得较高质量的图片,这导致生成一张图片需要的时间是其它深度生成模型的 50 到 100 倍,极大地限制了模型的部署和落地。

为了加速扩散模型的采样,许多研究者从硬件优化的角度出发,例如 Google 使用 JAX 语言将模型编译运行在 TPU 上,OneFlow 团队 [1] 使用自研编译器将 Stable Diffusion 做到了“一秒出图”。这些方法都基于 50 步的采样算法 PNDM[2],该算法在步数减少时采样效果会急剧下降。

就在几天前,这一纪录又被刷新了!Stable Diffusion 的官方 Demo[3]更新显示,采样 8 张图片的时间从原来的 8 秒钟直接被缩短至了 4 秒钟!快了整整一倍!

7d9048e7734be9377a6d2515a6c6df57.png

而基于自研深度学习编译器技术的 OneFlow 团队更是在不降低采样效果的前提下,成功将之前的 “一秒出图” 缩短到了 “半秒出图”!在 GPU 上仅仅使用不到 0.5 秒就可以获得一张高清的图片!相关工作已经发布在[1] 中。

事实上,这些工作的核心驱动力都来自于清华大学朱军教授带领的 TSAIL 团队所提出的DPM-Solver,一种针对于扩散模型特殊设计的高效求解器:该算法无需任何额外训练,同时适用于离散时间与连续时间的扩散模型,可以在 20 到 25 步内几乎收敛,并且只用 10 到 15 步也能获得非常高质量的采样。在 Stable Diffusion 上,25 步的 DPM-Solver 就可以获得优于 50 步 PNDM 的采样质量,因此采样速度直接翻倍!

项目链接:

  • DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps:https://arxiv.org/abs/2206.00927(NeurIPS 2022 Oral)

  • DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models:https://arxiv.org/abs/2211.01095

  • 项目开源代码:https://github.com/LuChengTHU/dpm-solver

  • 项目在线 Demo:https://huggingface.co/spaces/LuChengTHU/dpmsolver_sdm

扩散模型的定义与采样方法

扩散模型通过定义一个不断加噪声的前向过程来将图片逐步变为高斯噪声,再通过定义了一个逆向过程将高斯噪声逐步去噪变为清晰图片以得到采样:

616ecfc2e19db33e556d57dce1e12002.png

在采样过程中,根据是否添加额外的噪声,可以将扩散模型分为两类:一类是扩散随机微分方程模型(Diffusion SDE),另一类是扩散常微分方程(Diffusion ODE)。两种模型的训练目标函数都一样,通过最小化与噪声的均方误差来训练一个“噪声预测网络”:

26b013d258f35bd816ce6c8c3051fd10.png

基于 Diffusion SDE 的采样过程可以视为离散化如下随机微分方程:

370aa49a47b6ad263213602add6c134d.png

并且 [4] 中证明,DDPM[5] 是对上述 SDE 的一阶离散化。

而基于 Diffusion ODE 的采样过程可以视为离散化如下常微分方程:

59247f58a47cae3d7092bcdafaf19a61.png

并且 [6] 中证明,DDIM[7]是对上述 ODE 的一阶离散化。

然而,这些一阶的离散化方法收敛速度极慢,扩散模型的采样通常需要 100 到 1000 次串行计算才可以得到高质量的图片。通常情况下,为了加速扩散模型的采样,研究者往往通过对 Diffusion ODE 使用高阶求解器来进行加速,例如经典的 Runge-Kutta 方法(RK45),这是因为 ODE 不会带来额外的随机性,离散化步长可以相对选取得更大一些。在给定 s 时刻的解后,Runge-Kutta 方法基于离散化如下积分:

65d669c4d746f6b065883cf262b3ed7a.png

这样的离散化将 Diffusion ODE 整体看做一个黑盒,损失了 ODE 的已知信息,在小于 50 步的情况下就难以收敛了。

DPM-Solver:专为扩散模型设计的求解器

DPM-Solver 基于 Diffusion ODE 的半线性(semi-linear)结构,通过精确且解析地计算 ODE 中的线性项,我们可以得到:

9fdf22e73060070330f30dbe3e8fbfe2.png

剩余的积分项是一个关于时间的复杂的积分。然而,DPM-Solver 的提出者发现,该积分可以通过对 log-SNR(对数信噪比)做换元后得到一个非常简单的形式:

5ad1711d954381d44afa26ab0419fecc.png

剩余的积分是一个关于噪声预测模型的指数积分(exponentially weighted integral)。通过对噪声预测模型做泰勒展开,我们可以得到该积分的一个估计:

1d0b904debbf3168b048ea19e60d5c8c.png

该估计中存在两项:一项是全导数部分(向量),另一项是系数部分(标量)。DPM-Solver 的另一个核心贡献是,该系数可以通过分部积分被解析地计算:

1571044f7a0e8e9362773b3cbdf94a1a.png

而剩余的全导数部分则可以通过传统 ODE 求解器的数值方法来近似估计(无需任何求导运算):

5cfbb459031dd4eccc7d065b4b273593.png

基于以上 4 点,DPM-Solver 做到了尽可能地准确计算所有已知项,只对神经网络部分做近似,因此最大程度地减小了离散化误差:

13b1ee7da0e6c82119d7c713389abbf2.png

此外,基于该推导,我们可以得到 DDIM 本质上是 DPM-Solver 的一阶形式,这也能解释为什么 DDIM 在步数较少时依然可以获得很好的加速效果:

a4e4e7e6cc333415f3047f1a13000d37.png

在实验中,DPM-Solver 获得了远超其它采样算法的加速效果,仅仅在 15-20 步就几乎可以收敛:

4787a1f3a47950e6ffb96a6d223c362a.png

并且在论文中定量的结果显示,DPM-Solver 引入的额外计算量完全可以忽略,即对于步数的加速效果直接正比于时间上的加速效果——因此,基于 25 步的 DPM-Solver,Stable-Diffusion 模型的采样速度直接翻倍!例如,下图展示了不同采样算法在 Stable-Diffusion 上随着步数变化的效果,可见 DPM-Solver 在 10 到 15 步就可以获得非常高质量的采样:

2cb8c233327586936bcdfeb8681a529d.png

使用 DPM-Solver

DPM-Solver 的使用非常简单,既可以基于作者提供的官方代码,也可以使用主流的 Diffusers 库。例如,基于作者提供的官方代码(https://github.com/LuChengTHU/dpm-solver),只需要 3 行:

076139f8e23d9d847fa7fc4b540fe6d8.png

官方代码对 4 种扩散模型都进行了支持:

66337aa5c61cf1021d894ab5bbae05ad.png

并且同时支持 unconditional sampling、classifier guidance 和 classifier-free guidance:

2a8b39693324584d8cea7129379d1f9c.png

而基于 Diffusers 库的 DPM-Solver 同样很简单,只需要定义 scheduler 即可:

7cbbb44d3e631db1350cc3e89640171b.png

此外,作者团队还提供了一个在线 Demo:https://huggingface.co/spaces/LuChengTHU/dpmsolver_sdm

下图是 15 步的例子,可以看到图像质量已经非常高:

579825626b474c67e492834248fde369.png

相信基于 DPM-Solver,扩散模型的采样速度将不再是瓶颈。

关于作者

DPM-Solver 论文一作是来自清华大学 TSAIL 团队的路橙博士,他在知乎上关于扩散模型的讨论中也写了一篇关于扩散模型原理的入门介绍,目前已有 2000 + 赞:https://www.zhihu.com/question/536012286/answer/2533146567

清华大学 TSAIL 团队长期致力于贝叶斯机器学习的理论和算法研究,是国际上最早研究深度概率生成模型的团队之一,在贝叶斯模型、高效算法和概率编程库方面取得了系统深入的研究成果。团队另一位博士生鲍凡提出Analytic-DPM [8][9],为扩散模型的最优均值和方差给出了简单、令人吃惊的解析形式,获得 ICLR 2022 Outstanding Paper Award。在概率编程方面,机器之心早在 2017 年就报道了该团队发布的 “ZhuSuan” 深度概率编程库(https://zhusuan.readthedocs.io/en/latest/)[10],是国际上最早的面向深度概率模型的编程库之一。另外,值得一提的是,扩散概率模型的两位核心作者宋飏和宋佳铭,本科时均在朱军教授的指导下做科研训练,后来都去了斯坦福大学读博士。论文的合作者周聿浩、陈键飞、李崇轩,也是TSAIL组培养的优秀博士生,周聿浩为在读,陈键飞和李崇轩分别在清华大学计算机系、人民大学高瓴人工智能学院任教。

参考资料:

[1] OneFlow 版本的 Stable-Diffusion:https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion

[2] Luping Liu et al., 2022, Pseudo Numerical Methods for Diffusion Models on Manifolds, https://arxiv.org/abs/2202.09778

[3] Stable-Diffusion 的官方 Demo:https://huggingface.co/spaces/runwayml/stable-diffusion-v1-5

[4] Yang Song et al., 2021, Score-Based Generative Modeling through Stochastic Differential Equations, https://arxiv.org/abs/2011.13456

[5] Jonathan Ho et al., 2020, Denoising Diffusion Probabilistic Models,  https://arxiv.org/abs/2006.11239

[6] Tim Salimans & Jonathan Ho, 2022, Progressive Distillation for Fast Sampling of Diffusion Models, https://arxiv.org/abs/2202.00512

[7] Jiaming Song et al., 2020, Denoising Diffusion Implicit Models, https://arxiv.org/abs/2010.02502

[8] Fan Bao, et al., Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models, https://arxiv.org/abs/2201.06503

[9] Fan Bao, et al. Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models, https://arxiv.org/abs/2206.07309v1

[10] https://www.jiqizhixin.com/articles/2017-09-20-7

猜您喜欢:

深入浅出stable diffusion:AI作画技术背后的潜在扩散模型论文解读

2749c9d65ce4e338347271a739afa97f.png 戳我,查看GAN的系列专辑~!

一顿午饭外卖,成为CV视觉的前沿弄潮儿!

最新最全100篇汇总!生成扩散模型Diffusion Models

ECCV2022 | 生成对抗网络GAN部分论文汇总

CVPR 2022 | 25+方向、最新50篇GAN论文

 ICCV 2021 | 35个主题GAN论文汇总

超110篇!CVPR 2021最全GAN论文梳理

超100篇!CVPR 2020最全GAN论文梳理

拆解组新的GAN:解耦表征MixNMatch

StarGAN第2版:多域多样性图像生成

附下载 | 《可解释的机器学习》中文版

附下载 |《TensorFlow 2.0 深度学习算法实战》

附下载 |《计算机视觉中的数学方法》分享

《基于深度学习的表面缺陷检测方法综述》

《零样本图像分类综述: 十年进展》

《基于深度神经网络的少样本学习综述》

《礼记·学记》有云:独学而无友,则孤陋而寡闻

欢迎加入 GAN/扩散模型 —交流微信群 !

扫描下面二维码,添加运营小妹好友,拉你进群。发送申请时,请备注,格式为:研究方向+地区+学校/公司+姓名。如 扩散模型+北京+北航+吴彦祖

f19feab940e530bf71c34308976278ea.jpeg

请备注格式:研究方向+地区+学校/公司+姓名

a977eefff3e88b16ca2843e74b2b8f0f.jpeg

点击 一顿午饭外卖,成为CV视觉的前沿弄潮儿!,领取优惠券,加入 AI生成创作与计算机视觉 知识星球!

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值