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

稳定扩散采样速度提高一倍!只需10到25步的扩散模型采样算法

时间:2023-03-21 16:26:06 科技观察

要说今年AI领域最具影响力的进展,爆发式的AI映射绝对是其中之一。设计师只需要输入图片的文字描述,AI就可以生成一张高质量的高分辨率图片。目前StabilityAI的开源模型StableDiffusion应用最为广泛。该模型一经开源,便在社区内引起广泛讨论。然而,使用扩散模型的最大问题是其极慢的采样速度。模型采样需要从一张纯噪声的图片开始,一步一步去噪,最终得到清晰的图片。在这个过程中,模型必须连续计算至少50到100步才能获得更高质量的图片,这导致生成一张图片所需的时间是其他深度生成模型的50到100倍,极大地限制了模型的部署和部署。降落。为了加快扩散模型的采样速度,许多研究者从硬件优化的角度入手。例如,Google使用JAX语言在TPU上编译和运行模型。OneFlow团队[1]使用自研编译器让StableDiffusion实现“一秒输出”。”。这些方法都是基于50步采样算法PNDM[2],随着步数的减少,这个算法急剧下降。就在几天前,这个记录又被打破了!官方Demo[3]更新了StableDiffusion显示8张图片的采样时间从8秒缩短到4秒!快了一倍!OneFlow团队基于自主研发的深度学习编译技术,成功缩短了以往的“一秒图片”在不降低采样效果的情况下,将“半秒图像输出”转化为“半秒图像输出”!在GPU上得到一张高清图片只需不到0.5秒!相关工作已发表在[1]中。其实核心驱动力这些作品中的一部分来自于清华大学朱军教授领导的TSAIL团队提出的DPM-Solver,一种专为扩散模型设计的高效求解器:该算法不需要任何额外的训练,适用于Discrete-time和continuous-时间扩散离子模型几乎可以在20到25步内收敛,仅需10到15步就可以获得非常高质量的采样。在StableDiffusion上,25-stepDPM-Solver可以获得比50-stepPNDM更好的采样质量,所以采样速度直接翻倍!项目链接:DPM-Solver:AFastODESolverforDiffusionProbabilisticModelSamplinginAround10Steps:https://arxiv.org/abs/2206.00927(NeurIPS2022Oral)DPM-Solver++:FastSolverforGuidedSamplingofDiffusionProbabilisticModels:https://arxiv.org/abs/2211.01095项目开源代码:https://github.com/LuChengTHU/dpm-solver项目在线演示:https://huggingface.co/spaces/LuChengTHU/dpmsolver_sdm扩散定义模型采用采样的方式,扩散模型通过定义一个不断加噪声的正向过程,逐渐将图片变成高斯噪声,然后通过定义一个反向过程,逐渐将高斯噪声降噪为清晰的图片,得到采样:在采样过程中,根据是否加入附加噪声,扩散模型可分为两类:一类是扩散随机微分方程模型(DiffusionSDE),一类是扩散常微分方程(DiffusionODE)。两个模型的训练目标函数相同,通过最小化与噪声的均方误差来训练“噪声预测网络”:基于DiffusionSDE的采样过程可以看作是一个离散化的随机微分方程如下:并在[4]中证明,DDPM[5]是上述SDE的一阶离散化。基于DiffusionODE的采样过程可以看作是以下常微分方程的离散化:并且在[6]中证明了DDIM[7]是上述ODE的一阶离散化。但是,这些一阶离散化方法收敛速度很慢,扩散模型的采样通常需要100到1000次串行计算才能得到高质量的图片。通常,为了加快扩散模型的采样速度,研究人员往往通过使用高阶求解器对扩散ODE进行加速,例如经典的Runge-Kutta方法(RK45),因为ODE不会带来额外的随机性,离散化步骤尺寸可以选的相对大一些。给出时间s的解后,Runge-Kutta方法基于如下离散化:这种离散化将DiffusionODE视为一个整体的黑盒,并丢失了ODE的已知信息。收敛。DPM-Solver:为扩散模型设计的求解器DPM-Solver是基于DiffusionODE的半线性结构。通过精确解析地计算ODE中的线性项,我们可以得到:剩余的积分项是关于时间的复数积分。但是DPM-Solver的提出者发现通过改变log-SNR(对数信噪比)可以得到一个非常简单的积分形式:剩下的积分是关于噪声预测模型的指数积分(指数加权积分)。通过对噪声预测模型进行泰勒展开,我们可以得到积分的估计:这个估计有两项:一项是全导数部分(向量),另一项是系数部分(标量)。DPM-Solver的另一个核心贡献是可以通过分部积分解析计算系数:而剩下的全导数部分可以用传统ODE求解器的数值方法来近似(没有任何导数运算):基于对于上面4分,DPM-Solver已经尽可能准确地计算了所有已知项,只对神经网络部分进行了近似,从而最小化了离散化误差:另外,根据这个推导,我们可以得到DDIM本质上是第一个——DPM-Solver的阶形式,这也可以解释为什么DDIM在步数较少的情况下仍然可以取得很好的加速效果:在实验中,DPM-Solver获得了远超其他采样算法的加速效果,仅在15-20steps几乎可以收敛:而且论文中的量化结果表明,DPM-Solver引入的额外计算量完全可以忽略不计,即加速度ef对步数的影响与加速对时间的影响成正比——因此,基于DPM-Solver的25步,Stable-Diffusion模型的采样速度直接翻倍!例如,下图展示了随着步数的变化,不同采样算法对Stable-Diffusion的影响。可以看出DPM-Solver在10到15步就可以得到非常高质量的采样:使用DPM-SolverDPM-Solver使用起来非常简单,要么基于作者官方提供的代码,要么使用主流的Diffusers库.例如,基于作者提供的官方代码(https://github.com/LuChengTHU/dpm-solver),只需要3行:官方代码支持全部4种扩散模型:并支持无条件采样、分类器引导andclassifier-freeguidance:基于Diffusers库的DPM-Solver也很简单,只需要定义一个scheduler即可:另外,作者团队还提供了在线demo:https://huggingface.co/spaces/LuChengTHU/dpmsolver_sdm如下图是15步的例子。可以看出图像质量已经很高了:相信基于DPM-Solver,扩散模型的采样速度将不再是瓶颈。关于作者DPM-Solver论文的第一作者是清华大学TSAIL团队的陆诚博士。他还在知乎上关于扩散模型的讨论中写了一篇介绍扩散模型原理的介绍。目前已有2000+点赞:https://www.zhihu.com/question/536012286/answer/2533146567清华大学TSAIL团队长期致力于贝叶斯机器学习理论与算法研究,并且是世界上最早研究深度概率生成模型的团队之一,首先在贝叶斯模型、高效算法和概率编程库方面取得了系统和深入的研究成果。团队另一位博士生鲍帆提出了Analytic-DPM[8][9],为扩散模型的最优均值和方差给出了一种简单而令人惊讶的解析形式,并获得了ICLR2022杰出论文奖。在概率编程方面,机器之心早在2017年就报道了团队发布的“算算”深度概率编程库(https://zhusuan.readthedocs.io/en/latest/)[10],即世界上最早的。深度概率模型的编程库之一。此外,值得一提的是,扩散概率模型的两位核心作者宋洋和宋家明,本科时均在朱军教授指导下进行科研训练,后赴斯坦福大学深造博士学位。论文的共同作者周宇豪、陈健飞、李崇轩也是TSAIL课题组培养的优秀博士生。周宇豪是学生,陈健飞和李崇轩分别在清华大学计算机系和人民大学高瓴人工智能学院任教。