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

作者：機器之心 2024-12-13 12:07:15

人工智能新聞

Google DeepMind 的研究團隊發現，原來擴散模型和流匹配就像一枚硬幣的兩面，本質上是等價的 (尤其是在流匹配采用高斯分布作為基礎分布時)，只是不同的模型設定會導致不同的網絡輸出和采樣方案。

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

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

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

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

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

現在，這個困擾已得到解答。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 的歐拉采樣器，縮放確實會產生真正的差異：可以看到路徑和最終樣本都發生了變化。

看到這里，需要思考一下。人們常說流匹配會產生直線路徑，但在上圖中，其采樣軌跡看起來是彎曲的。

在下面的交互式圖表中，我們可以通過滑塊更改右側數據分布的方差。

不過，在像圖像這樣的真實數據集上找到這樣的直線路徑要復雜得多。但結論仍然是相同的：最優的積分方法取決于數據分布。

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

采樣器的等價性：DDIM 與流匹配采樣器等價，并且對噪聲調度的線性縮放不變。
對直線性的誤解：流匹配調度僅在模型預測單個點時才是直線。

訓練

對于擴散模型，學習模型是通過最小化加權均方誤差（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 采樣軌跡的離散化誤差和模型曲率更相關。

總結

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

加權中的等價性：加權函數對于訓練很重要，它平衡了感知數據不同頻率分量的重要性。流匹配加權與常用的擴散訓練加權方法相同。
訓練噪聲調度的不重要性：噪聲調度對訓練目標的重要性要小得多，但會影響訓練效率。
網絡輸出的差異：流匹配提出的網絡輸出是新的，它很好地平衡了

更深入地理解采樣器

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

回流算子

流匹配中的回流（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》以了解中這些權重的解析表達式。