清華大學計算機系朱軍教授帶領的 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 在步數較少時依然可以獲得很好的加速效果：
在實驗中，DPM-Solver 獲得了遠超其它取樣演算法的加速效果，僅僅在 15-20 步就幾乎可以收斂：
並且在論文中定量的結果顯示，DPM-Solver 引入的額外計算量完全可以忽略，即對於步數的加速效果直接正比於時間上的加速效果——因此，基於 25 步的 DPM-Solver，Stable-Diffusion 模型的取樣速度直接翻倍！例如，下圖展示了不同取樣演算法在 Stable-Diffusion 上隨著步數變化的效果，可見 DPM-Solver 在 10 到 15 步就可以獲得非常高質量的取樣：
使用 DPM-Solver
DPM-Solver 的使用非常簡單，既可以基於作者提供的官方程式碼，也可以使用主流的 Diffusers 庫。例如，基於作者提供的官方程式碼（https://github.com/LuChengTHU/dpm-solver），只需要 3 行：
官方程式碼對 4 種擴散模型都進行了支持：
並且同時支持 unconditional sampling、classifier guidance 和 classifier-free guidance：
而基於 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]，為擴散模型的最優均值和方差給出了簡單、令人吃驚的解析形式，獲得 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
