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

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

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

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

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

就在幾天前，這一紀錄又被刷新了！Stable Diffusion 的官方 Demo更新顯示，采樣 8 張圖片的時間從原來的 8 秒鐘直接被縮短至了 4 秒鐘！快了整整一倍！

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

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

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

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

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

并且證明，DDPM 是對上述 SDE 的一階離散化。

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

并且證明，DDIM是對上述 ODE 的一階離散化。

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

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

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

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

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

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

該估計中存在兩項：一項是全導數(shù)部分（向量），另一項是系數(shù)部分（標量）。DPM-Solver 的另一個核心貢獻是，該系數(shù)可以通過分部積分被解析地計算：

而剩余的全導數(shù)部分則可以通過傳統(tǒng) ODE 求解器的數(shù)值方法來近似估計（無需任何求導運算）：

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

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

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

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

使用 DPM-Solver

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

官方代碼對 4 種擴散模型都進行了支持：

并且同時支持 unconditional sampling、classifier guidance 和 classifier-free guidance：

而基于 Diffusers 庫的 DPM-Solver 同樣很簡單，只需要定義 scheduler 即可：

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

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