清华大学计算机系朱军教授带领的 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 秒钟!快了整整一倍!
而基于自研深度学习编译器技术的 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
扩散模型的定义与采样方法
扩散模型通过定义一个不断加噪声的前向过程来将图片逐步变为高斯噪声,再通过定义了一个逆向过程将高斯噪声逐步去噪变为清晰图片以得到采样:
在采样过程中,根据是否添加额外的噪声,可以将扩散模型分为两类:一类是扩散随机微分方程模型(Diffusion SDE),另一类是扩散常微分方程(Diffusion ODE)。两种模型的训练目标函数都一样,通过最小化与噪声的均方误差来训练一个“噪声预测网络”:
基于 Diffusion SDE 的采样过程可以视为离散化如下随机微分方程:
并且 [4] 中证明,DDPM[5] 是对上述 SDE 的一阶离散化。
而基于 Diffusion ODE 的采样过程可以视为离散化如下常微分方程:
并且 [6] 中证明,DDIM[7]是对上述 ODE 的一阶离散化。
然而,这些一阶的离散化方法收敛速度极慢,扩散模型的采样通常需要 100 到 1000 次串行计算才可以得到高质量的图片。通常情况下,为了加速扩散模型的采样,研究者往往通过对 Diffusion ODE 使用高阶求解器来进行加速,例如经典的 Runge-Kutta 方法(RK45),这是因为 ODE 不会带来额外的随机性,离散化步长可以相对选取得更大一些。在给定 s 时刻的解后,Runge-Kutta 方法基于离散化如下积分:
这样的离散化将 Diffusion ODE 整体看做一个黑盒,损失了 ODE 的已知信息,在小于 50 步的情况下就难以收敛了。
DPM-Solver:专为扩散模型设计的求解器
DPM-Solver 基于 Diffusion ODE 的半线性(semi-linear)结构,通过精确且解析地计算 ODE 中的线性项,我们可以得到:
剩余的积分项是一个关于时间的复杂的积分。然而,DPM-Solver 的提出者发现,该积分可以通过对 log-SNR(对数信噪比)做换元后得到一个非常简单的形式:
剩余的积分是一个关于噪声预测模型的指数积分(exponentially weighted integral)。通过对噪声预测模型做泰勒展开,我们可以得到该积分的一个估计:
该估计中存在两项:一项是全导数部分(向量),另一项是系数部分(标量)。DPM-Solver 的另一个核心贡献是,该系数可以通过分部积分被解析地计算:
而剩余的全导数部分则可以通过传统 ODE 求解器的数值方法来近似估计(无需任何求导运算):
基于以上 4 点,DPM-Solver 做到了尽可能地准确计算所有已知项,只对神经网络部分做近似,因此最大程度地减小了离散化误差:
此外,基于该推导,我们可以得到 DDIM 本质上是 DPM-Solver 的一阶形式,这也能解释为什么 DDIM 在步数较少时依然可以获得很好的加速效果: