擴散模型=流匹配？谷歌DeepMind部落格深度詳解這種驚人的等價性

机器之心發表於2024-12-13
模型谷歌

擴散模型和流匹配實際上是同一個概念的兩種不同表達方式嗎？

從表面上看，這兩種方法似乎各有側重：擴散模型專注於透過迭代的方式逐步去除噪聲，將資料還原成清晰的樣本。

而流匹配則側重於構建可逆變換系統，目標是學習如何將簡單的基礎分佈精確地對映到真實資料分佈。

因為流匹配的公式很簡單，並且生成樣本的路徑很直接，最近越來越受研究者們的歡迎，於是很多人都在問：

「到底是擴散模型好呢？還是流匹配好？」

現在，這個困擾已得到解答。Google DeepMind 的研究團隊發現，原來擴散模型和流匹配就像一枚硬幣的兩面，本質上是等價的 (尤其是在流匹配採用高斯分佈作為基礎分佈時)，只是不同的模型設定會導致不同的網路輸出和取樣方案。

這無疑是個好訊息，意味著這兩種框架下的方法可以靈活搭配，發揮組合技了。比如在訓練完一個流匹配模型後，不必再侷限於傳統的確定性取樣方法，完全可以引入隨機取樣策略。

圖片
連結：https://diffusionflow.github.io

在這篇部落格的開頭，作者們寫道：「我們的目標是幫助大家能夠自如地交替使用這兩種方法，同時在調整演算法時擁有真正的自由度 —— 方法的名稱並不重要，重要的是理解其本質。」

擴散模型與流匹配

擴散模型主要分為前向過程和反向兩個階段。

前向過程用公式表示為：
圖片
其中 z_t 是在時間點 t 時的帶噪聲資料，x 代表原始資料，ε 代表隨機噪聲，a_t 和 σ_t 是控制噪聲新增程度的引數。

若滿足圖片，稱為「方差保持」，意味著在每個時間步驟中，噪聲的方差保持不變或接近不變。

DDIM 取樣器的反向過程用公式表示為：
圖片其中，圖片
而在流匹配中，前向過程視為資料 x 和噪聲項 ε 之間的線性插值：
圖片
取樣

人們普遍認為，這兩個框架在生成樣本的方式上有所不同：流匹配取樣是確定性的，具有直線路徑，而擴散模型取樣是隨機性的，具有曲線路徑。下面文章將澄清這一誤解：首先關注更簡單的確定性取樣，稍後再討論隨機情況。

假設你想使用訓練好的降噪器模型將隨機噪聲轉換為資料點。可以先回想一下 DDIM 的更新 圖片，有趣的是，重新排列項可以用以下公式來表達，這裡涉及幾組網路輸出和重新引數化：
圖片
圖片
我們再回到公式（4）中的流匹配更新，和上述方程看起來很相似。如果在最後一行將網路輸出設為 圖片，並令 圖片，可以得到 圖片、圖片 這樣我們就恢復了流匹配更新！更準確地說，流匹配更新可以被視為重引數化取樣常微分方程（ODE）的尤拉積分：
圖片
對於 DDIM 取樣器而言，普遍存在以下結論：DDIM 取樣器對於應用於噪聲排程 α_t、σ_t 的線性縮放是不變的，因為縮放不會影響 圖片 和 圖片，這對於其他取樣器來說並不成立，例如機率流 ODE 的尤拉取樣器。

為了驗證上述結論，本文展示了使用幾種不同的噪聲排程得到的結果，每種排程都遵循流匹配排程，並具有不同的縮放因子。如下圖，隨意調整滑塊，在最左側，縮放因子是 1，這正是流匹配排程，而在最右側，縮放因子是 圖片。可以觀察到 DDIM（以及流匹配取樣器）總是給出相同的最終資料樣本，無論排程的縮放如何。對於機率流 ODE 的尤拉取樣器，縮放確實會產生真正的差異：可以看到路徑和最終樣本都發生了變化。
圖片
看到這裡，需要思考一下。人們常說流匹配會產生直線路徑，但在上圖中，其取樣軌跡看起來是彎曲的。

在下面的互動式圖表中，我們可以透過滑塊更改右側資料分佈的方差。
圖片
不過，在像影像這樣的真實資料集上找到這樣的直線路徑要複雜得多。但結論仍然是相同的：最優的積分方法取決於資料分佈。

我們可以從確定性取樣中得到的兩個重要結論：

  1. 取樣器的等價性：DDIM 與流匹配取樣器等價，並且對噪聲排程的線性縮放不變。
  2. 對直線性的誤解：流匹配排程僅在模型預測單個點時才是直線。

訓練

對於擴散模型，學習模型是透過最小化加權均方誤差（MSE）損失來完成的：
圖片
流匹配也符合上述訓練目標：
圖片
網路應該輸出什麼

下面總結了文獻中提出的幾個網路輸出，包括擴散模型使用的幾個版本和流匹配使用的其中一個版本。
圖片
然而，在實踐中，模型的輸出可能會產生非常大的影響。例如，基於相似的原因，圖片 在低噪聲水平下是有問題的，因為 圖片 沒有資訊量，並且錯誤在 圖片 中被放大了。

因此，一種啟發式方法是選擇一個網路輸出，它是 圖片、圖片的組合，這適用於 圖片 和流匹配向量場 圖片

如何選擇加權函式

加權函式是損失函式中最重要的部分，它平衡了影像、影片和音訊等資料中高頻和低頻分量的重要性。這一點至關重要，因為這些訊號中的某些高頻分量是人類無法感知的。如果透過加權情況來檢視損失函式，可以得出以下結果：
圖片
即公式 (7) 中的條件流匹配目標與擴散模型中常用的設定相同。下面繪製了文獻中常用的幾個加權函式。
圖片
流匹配加權（也稱為 v-MSE + 餘弦排程加權）會隨著 λ 的增加而呈指數下降。該團隊在實驗中發現了另一個有趣的聯絡：Stable Diffusion 3 加權 [9]（這是流匹配的一種重新加權版本）與擴散模型中流行的 EDM 加權 [10] 非常相似。

如何選擇訓練噪聲排程？

最後討論訓練噪聲排程，因為在以下意義上，它對訓練的重要程度最低：

1. 訓練損失不會隨訓練噪聲排程變化。具體來說，損失函式可以重寫為 圖片 它只與端點（λ_max, λ_min）有關，但與中間的排程 λ_t 無關。在實踐中，應該選擇合適的 λ_max, λ_min，使得兩端分別足夠接近乾淨資料和高斯噪聲。λ_t 可能仍然會影響訓練損失的蒙特卡洛估計量的方差。一些文獻中提出了一些啟發式方法來在訓練過程中自動調整噪聲排程。這篇博文有一個很好的總結：https://sander.ai/2024/06/14/noise-schedules.html#adaptive

2. 類似於取樣噪聲排程，訓練噪聲排程不會隨線性擴充套件（linear scaling）而變化，因為人們可以輕鬆地將線性擴充套件應用於 z_t，並在網路輸入處進行 unscaling 以獲得等價性。噪聲排程的關鍵定義屬性是對數訊雜比 λ_t。

3. 人們可以根據不同的啟發式方法為訓練和取樣選擇完全不同的噪聲排程：對於訓練，最好有一個噪聲排程來最小化蒙特卡洛估計量的方差；而對於取樣，噪聲排程與 ODE / SDE 取樣軌跡的離散化誤差和模型曲率更相關。

總結

下面給出了訓練擴散模型 / 流匹配的一些要點：

  1. 加權中的等價性：加權函式對於訓練很重要，它平衡了感知資料不同頻率分量的重要性。流匹配加權與常用的擴散訓練加權方法相同。
  2. 訓練噪聲排程的不重要性：噪聲排程對訓練目標的重要性要小得多，但會影響訓練效率。
  3. 網路輸出的差異：流匹配提出的網路輸出是新的，它很好地平衡了 圖片

更深入地理解取樣器

這一節將更詳細地介紹各種不同的取樣器。

迴流運算元

流匹配中的迴流（Reflow）運算是使用直線將噪聲與資料點連線起來。透過基於噪聲執行一個確定性的取樣器，可以得到這些 (資料，噪聲) 對。然後，可以訓練模型，使之可以根據給定噪聲直接預測資料，而無需取樣。在擴散技術的相關文獻中，這同樣的方法是最早的蒸餾技術之一。

確定性取樣器與隨機取樣器

此前已經討論了擴散模型或流匹配的確定性取樣器。另一種方法是使用隨機取樣器，例如 DDPM 取樣器。

執行一個從 λ_t 到 λ_t+Δλ 的 DDPM 取樣步驟完全等價於執行一個到 λ_t+2Δλ 的 DDIM 取樣步驟，然後透過執行前向擴散重新噪聲化到 λ_t+Δλ。也就是說，透過前向擴散重新噪聲化恰好逆轉了 DDIM 所取得的一半進展。為了理解這一點，讓我們看一個 2D 示例。從相同的高斯分佈混合開始，我們可以執行一個小的 DDIM 取樣步驟，左圖帶有更新反轉的符號，右圖則是一個小的前向擴散步驟：
圖片
對於單個樣本而言，這些更新的行為完全不同：反轉的 DDIM 更新始終將每個樣本推離分佈模式，而擴散更新完全是隨機的。但是，在彙總所有樣本時，更新後得到的分佈是相同的。因此，如果執行 DDIM 取樣步驟（不反轉符號），然後執行前向擴散步驟，則整體分佈與更新之前的分佈保持不變。

透過重新加噪來撤消的 DDIM 步驟的比例是一個超引數，並且可以自由選擇（即不必一定是 DDIM 步驟的一半）。這個超引數在《Elucidating the design space of diffusion-based generative models》中被稱為 level of churn，可譯為「攪動水平」。有趣的是，將攪動新增到取樣器的效果是：減少取樣過程早期做出的模型預測對最終樣本的影響，並增加對後續預測的權重。如下圖所示：
圖片
在這裡，我們使用餘弦噪聲排程以及 圖片 預測將不同取樣器都執行了 100 個取樣步驟。忽略非線性相互作用，取樣器產生的最終樣本可以寫成取樣過程中做出的預測和高斯噪聲 e 的加權和：圖片
這些預測的權重 h_t 顯示在 y 軸上，而 x 軸上顯示不同的擴散時間 t。DDIM 會在此設定下對 圖片 預測賦予相等的權重，而 DDPM 則更注重在取樣結束時所做的預測。另請參閱《Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models》以瞭解 圖片中這些權重的解析表示式。

SDE 和 ODE 視角

前面，我們已經觀察到擴散模型和流匹配演算法之間的等價性。下面將使用 ODE 和 SDE 來形式化地描述正向過程和取樣的等價性，以實現理論上的完整性。

擴散模型

擴散模型的前向過程涉及到隨時間推移逐漸破壞一個資料，而該過程可使用以下隨機微分方程（SDE）來描述：
圖片
其中 dz 是無窮小的高斯（即布朗運動）。f_t 和 g_t 決定了噪聲排程。其生成過程由前向過程的逆過程給出，其公式為：
圖片
其中 ∇log p_t 是前向過程的分數。

請注意，這裡引入了一個附加引數 η_t，它控制的是推理時的隨機性。這與之前介紹的攪動（churn）引數有關。當離散化後向過程時，如果 η_t=0，則是恢復 DDIM；如果 η_t=1，則是恢復 DDPM。

流匹配

流匹配中 x 和 ε 之間的插值可以用以下常微分方程（ODE）描述：
圖片
假設該插值為 圖片
其生成過程只是在時間上反轉這個 ODE，並將 u_t 替換為其對 z_t 的條件期望。這是隨機插值（stochastic interpolants）的一個特例 —— 在這種情況下，它可以泛化成 SDE：
圖片
其中 ε_t 控制著推理時的隨機性。

兩個框架的等價性

這兩個框架都分別由三個超引數定義：擴散的三個引數是 f_t、g_t、η_t，而流匹配的三個引數是 α_t、σ_t、ε_t。透過從一組超引數推導得到另一組超引數，可以顯示這兩組超引數的等價性。從擴散到流匹配：
圖片
從流匹配到擴散：
圖片
總之，除了訓練考慮和取樣器選擇之外，擴散和高斯流匹配沒有根本區別。

結語

讀到這裡，想必你已經理解了擴散模型和高斯流匹配的等價性。不過，文中重點介紹的是流匹配為該領域帶來的兩個新模型規範：

  • 網路輸出：流匹配提出了一種網路輸出的向量場引數化方案，並且其不同於擴散文獻中使用的方案。當使用高階取樣器時，網路輸出可能會有所不同。它也可能影響訓練動態。
  • 取樣噪聲排程：流匹配利用了簡單的取樣噪聲排程 α_t = 1-t 和 σ_t = t，並且更新規則與 DDIM 相同。

該團隊最後表示：「如果能透過實證方式研究這兩個模型規範在不同的真實應用中的重要性，那一定會很有趣。我們將此留給了未來的工作。」

參考文獻

  1. Flow matching for generative modeling
  2. Lipman, Y., Chen, R.T., Ben-Hamu, H., Nickel, M. and Le, M., 2022. arXiv preprint arXiv:2210.02747.
  3. Flow straight and fast: Learning to generate and transfer data with rectified flow
  4. Liu, X., Gong, C. and Liu, Q., 2022. arXiv preprint arXiv:2209.03003.
  5. Building normalizing flows with stochastic interpolants
  6. Albergo, M.S. and Vanden-Eijnden, E., 2022. arXiv preprint arXiv:2209.15571.
  7. Stochastic interpolants: A unifying framework for flows and diffusions
  8. Albergo, M.S., Boffi, N.M. and Vanden-Eijnden, E., 2023. arXiv preprint arXiv:2303.08797.
  9. Denoising diffusion implicit models
  10. Song, J., Meng, C. and Ermon, S., 2020. arXiv preprint arXiv:2010.02502.
  11. Score-based generative modeling through stochastic differential equations
  12. Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S. and Poole, B., 2020. arXiv preprint arXiv:2011.13456.
  13. Understanding diffusion objectives as the elbo with simple data augmentation
  14. Kingma, D. and Gao, R., 2024. Advances in Neural Information Processing Systems, Vol 36.
  15. Diffusion is spectral autoregression [HTML]
  16. Dieleman, S., 2024.
  17. Scaling rectified flow transformers for high-resolution image synthesis
  18. Esser, P., Kulal, S., Blattmann, A., Entezari, R., Muller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F. and others,, 2024. Forty-first International Conference on Machine Learning.
  19. Elucidating the design space of diffusion-based generative models
  20. Karras, T., Aittala, M., Aila, T. and Laine, S., 2022. Advances in neural information processing systems, Vol 35, pp. 26565—26577.
  21. Knowledge distillation in iterative generative models for improved sampling speed [PDF]
  22. Luhman, E. and Luhman, T., 2021. arXiv preprint arXiv:2101.02388.
  23. Denoising diffusion probabilistic models
  24. Ho, J., Jain, A. and Abbeel, P., 2020. Advances in neural information processing systems, Vol 33, pp. 6840—6851.
  25. Progressive Distillation for Fast Sampling of Diffusion Models
  26. Salimans, T. and Ho, J., 2022. International Conference on Learning Representations.
  27. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models
  28. Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C. and Zhu, J., 2022. arXiv preprint arXiv:2211.01095.

相關文章