擴散模型=流匹配?GoogleDeepMind博客深度詳解這種驚人的等價性

機器之心報導

機器之心編輯部

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

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

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

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

到底是擴散模型好呢?還是流匹配好?

現在,這個困擾已得到解答。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.