Stable Diffusion采樣速度翻倍！僅需10到25步的擴(kuò)散模型采樣算法

清華大學(xué)計(jì)算機(jī)系朱軍教授帶領(lǐng)的 TSAIL 團(tuán)隊(duì)提出 DPM-Solver（NeurIPS 2022 Oral，約前 1.7%）和 DPM-Solver++，將擴(kuò)散模型的快速采樣算法提升到了極致：無(wú)需額外訓(xùn)練，僅需 10 到 25 步就可以獲得極高質(zhì)量的采樣。

要說(shuō) AI 領(lǐng)域今年影響力最大的進(jìn)展，爆火的 AI 作圖絕對(duì)是其中之一。設(shè)計(jì)者只需要輸入對(duì)圖片的文字描述，就可以由 AI 生成一張質(zhì)量極高的高分辨率圖片。目前，使用范圍最廣的當(dāng)屬 StabilityAI 的開(kāi)源模型 Stable Diffusion，模型一經(jīng)開(kāi)源就在社區(qū)引起了廣泛的討論。

然而，擴(kuò)散模型在使用上最大的問(wèn)題就是其極慢的采樣速度。模型采樣需要從純?cè)肼晥D片出發(fā)，一步一步不斷地去噪，最終得到清晰的圖片。在這個(gè)過(guò)程中，模型必須串行地計(jì)算至少 50 到 100 步才可以獲得較高質(zhì)量的圖片，這導(dǎo)致生成一張圖片需要的時(shí)間是其它深度生成模型的 50 到 100 倍，極大地限制了模型的部署和落地。

為了加速擴(kuò)散模型的采樣，許多研究者從硬件優(yōu)化的角度出發(fā)，例如 Google 使用 JAX 語(yǔ)言將模型編譯運(yùn)行在 TPU 上，OneFlow 團(tuán)隊(duì) [1] 使用自研編譯器將 Stable Diffusion 做到了“一秒出圖”。這些方法都基于 50 步的采樣算法 PNDM[2]，該算法在步數(shù)減少時(shí)采樣效果會(huì)急劇下降。

就在幾天前，這一紀(jì)錄又被刷新了！Stable Diffusion 的官方 Demo[3]更新顯示，采樣 8 張圖片的時(shí)間從原來(lái)的 8 秒鐘直接被縮短至了 4 秒鐘！快了整整一倍！

而基于自研深度學(xué)習(xí)編譯器技術(shù)的 OneFlow 團(tuán)隊(duì)更是在不降低采樣效果的前提下，成功將之前的 “一秒出圖” 縮短到了 “半秒出圖”！在 GPU 上僅僅使用不到 0.5 秒就可以獲得一張高清的圖片！相關(guān)工作已經(jīng)發(fā)布在[1] 中。

事實(shí)上，這些工作的核心驅(qū)動(dòng)力都來(lái)自于清華大學(xué)朱軍教授帶領(lǐng)的 TSAIL 團(tuán)隊(duì)所提出的DPM-Solver，一種針對(duì)于擴(kuò)散模型特殊設(shè)計(jì)的高效求解器：該算法無(wú)需任何額外訓(xùn)練，同時(shí)適用于離散時(shí)間與連續(xù)時(shí)間的擴(kuò)散模型，可以在 20 到 25 步內(nèi)幾乎收斂，并且只用 10 到 15 步也能獲得非常高質(zhì)量的采樣。在 Stable Diffusion 上，25 步的 DPM-Solver 就可以獲得優(yōu)于 50 步 PNDM 的采樣質(zhì)量，因此采樣速度直接翻倍！

項(xiàng)目鏈接：

• 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

• 項(xiàng)目開(kāi)源代碼：https://github.com/LuChengTHU/dpm-solver

• 項(xiàng)目在線 Demo：https://huggingface.co/spaces/LuChengTHU/dpmsolver_sdm

擴(kuò)散模型的定義與采樣方法

擴(kuò)散模型通過(guò)定義一個(gè)不斷加噪聲的前向過(guò)程來(lái)將圖片逐步變?yōu)楦咚乖肼?，再通過(guò)定義了一個(gè)逆向過(guò)程將高斯噪聲逐步去噪變?yōu)榍逦鷪D片以得到采樣：

在采樣過(guò)程中，根據(jù)是否添加額外的噪聲，可以將擴(kuò)散模型分為兩類：一類是擴(kuò)散隨機(jī)微分方程模型（Diffusion SDE），另一類是擴(kuò)散常微分方程（Diffusion ODE）。兩種模型的訓(xùn)練目標(biāo)函數(shù)都一樣，通過(guò)最小化與噪聲的均方誤差來(lái)訓(xùn)練一個(gè)“噪聲預(yù)測(cè)網(wǎng)絡(luò)”：

基于 Diffusion SDE 的采樣過(guò)程可以視為離散化如下隨機(jī)微分方程：

并且 [4] 中證明，DDPM[5] 是對(duì)上述 SDE 的一階離散化。

而基于 Diffusion ODE 的采樣過(guò)程可以視為離散化如下常微分方程：

并且 [6] 中證明，DDIM[7]是對(duì)上述 ODE 的一階離散化。

然而，這些一階的離散化方法收斂速度極慢，擴(kuò)散模型的采樣通常需要 100 到 1000 次串行計(jì)算才可以得到高質(zhì)量的圖片。通常情況下，為了加速擴(kuò)散模型的采樣，研究者往往通過(guò)對(duì) Diffusion ODE 使用高階求解器來(lái)進(jìn)行加速，例如經(jīng)典的 Runge-Kutta 方法（RK45），這是因?yàn)?ODE 不會(huì)帶來(lái)額外的隨機(jī)性，離散化步長(zhǎng)可以相對(duì)選取得更大一些。在給定 s 時(shí)刻的解后，Runge-Kutta 方法基于離散化如下積分：

這樣的離散化將 Diffusion ODE 整體看做一個(gè)黑盒，損失了 ODE 的已知信息，在小于 50 步的情況下就難以收斂了。

DPM-Solver：專為擴(kuò)散模型設(shè)計(jì)的求解器

DPM-Solver 基于 Diffusion ODE 的半線性（semi-linear）結(jié)構(gòu)，通過(guò)精確且解析地計(jì)算 ODE 中的線性項(xiàng)，我們可以得到：

剩余的積分項(xiàng)是一個(gè)關(guān)于時(shí)間的復(fù)雜的積分。然而，DPM-Solver 的提出者發(fā)現(xiàn)，該積分可以通過(guò)對(duì) log-SNR（對(duì)數(shù)信噪比）做換元后得到一個(gè)非常簡(jiǎn)單的形式：

剩余的積分是一個(gè)關(guān)于噪聲預(yù)測(cè)模型的指數(shù)積分（exponentially weighted integral）。通過(guò)對(duì)噪聲預(yù)測(cè)模型做泰勒展開(kāi)，我們可以得到該積分的一個(gè)估計(jì)：

該估計(jì)中存在兩項(xiàng)：一項(xiàng)是全導(dǎo)數(shù)部分（向量），另一項(xiàng)是系數(shù)部分（標(biāo)量）。DPM-Solver 的另一個(gè)核心貢獻(xiàn)是，該系數(shù)可以通過(guò)分部積分被解析地計(jì)算：

而剩余的全導(dǎo)數(shù)部分則可以通過(guò)傳統(tǒng) ODE 求解器的數(shù)值方法來(lái)近似估計(jì)（無(wú)需任何求導(dǎo)運(yùn)算）：

基于以上 4 點(diǎn)，DPM-Solver 做到了盡可能地準(zhǔn)確計(jì)算所有已知項(xiàng)，只對(duì)神經(jīng)網(wǎng)絡(luò)部分做近似，因此最大程度地減小了離散化誤差：

此外，基于該推導(dǎo)，我們可以得到 DDIM 本質(zhì)上是 DPM-Solver 的一階形式，這也能解釋為什么 DDIM 在步數(shù)較少時(shí)依然可以獲得很好的加速效果：

在實(shí)驗(yàn)中，DPM-Solver 獲得了遠(yuǎn)超其它采樣算法的加速效果，僅僅在 15-20 步就幾乎可以收斂：

并且在論文中定量的結(jié)果顯示，DPM-Solver 引入的額外計(jì)算量完全可以忽略，即對(duì)于步數(shù)的加速效果直接正比于時(shí)間上的加速效果——因此，基于 25 步的 DPM-Solver，Stable-Diffusion 模型的采樣速度直接翻倍！例如，下圖展示了不同采樣算法在 Stable-Diffusion 上隨著步數(shù)變化的效果，可見(jiàn) DPM-Solver 在 10 到 15 步就可以獲得非常高質(zhì)量的采樣：

使用 DPM-Solver

DPM-Solver 的使用非常簡(jiǎn)單，既可以基于作者提供的官方代碼，也可以使用主流的 Diffusers 庫(kù)。例如，基于作者提供的官方代碼（https://github.com/LuChengTHU/dpm-solver），只需要 3 行：

官方代碼對(duì) 4 種擴(kuò)散模型都進(jìn)行了支持：

并且同時(shí)支持 unconditional sampling、classifier guidance 和 classifier-free guidance：

而基于 Diffusers 庫(kù)的 DPM-Solver 同樣很簡(jiǎn)單，只需要定義 scheduler 即可：

此外，作者團(tuán)隊(duì)還提供了一個(gè)在線 Demo：https://huggingface.co/spaces/LuChengTHU/dpmsolver_sdm

下圖是 15 步的例子，可以看到圖像質(zhì)量已經(jīng)非常高：

相信基于 DPM-Solver，擴(kuò)散模型的采樣速度將不再是瓶頸。

關(guān)于作者

DPM-Solver 論文一作是來(lái)自清華大學(xué) TSAIL 團(tuán)隊(duì)的路橙博士，他在知乎上關(guān)于擴(kuò)散模型的討論中也寫了一篇關(guān)于擴(kuò)散模型原理的入門介紹，目前已有 2000 + 贊：https://www.zhihu.com/question/536012286/answer/2533146567

清華大學(xué) TSAIL 團(tuán)隊(duì)長(zhǎng)期致力于貝葉斯機(jī)器學(xué)習(xí)的理論和算法研究，是國(guó)際上最早研究深度概率生成模型的團(tuán)隊(duì)之一，在貝葉斯模型、高效算法和概率編程庫(kù)方面取得了系統(tǒng)深入的研究成果。團(tuán)隊(duì)另一位博士生鮑凡提出Analytic-DPM [8][9]，為擴(kuò)散模型的最優(yōu)均值和方差給出了簡(jiǎn)單、令人吃驚的解析形式，獲得 ICLR 2022 Outstanding Paper Award。在概率編程方面，機(jī)器之心早在 2017 年就報(bào)道了該團(tuán)隊(duì)發(fā)布的 “ZhuSuan” 深度概率編程庫(kù)（https://zhusuan.readthedocs.io/en/latest/）[10]，是國(guó)際上最早的面向深度概率模型的編程庫(kù)之一。另外，值得一提的是，擴(kuò)散概率模型的兩位核心作者宋飏和宋佳銘，本科時(shí)均在朱軍教授的指導(dǎo)下做科研訓(xùn)練，后來(lái)都去了斯坦福大學(xué)讀博士。論文的合作者周聿浩、陳鍵飛、李崇軒，也是TSAIL組培養(yǎng)的優(yōu)秀博士生，周聿浩為在讀，陳鍵飛和李崇軒分別在清華大學(xué)計(jì)算機(jī)系、人民大學(xué)高瓴人工智能學(xué)院任教。

參考資料：

[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