Flow Matching:不解 ODE,直接用回歸學速度場
B11 的 neural ODE 用最大概似訓練,每一步都得解 ODE。Flow Matching(FM)換一個思路:先指定一條從來源分佈到目標分佈的機率路徑,再用回歸直接學一個能實現這條路徑的速度場,訓練時完全不必模擬。關鍵和 DDPM、denoising score matching 一樣是條件化:邊際速度算不出來,但對每個訓練樣本有封閉形式的條件速度,而對條件速度做回歸,最佳解恰好是邊際速度。這一頁先從 Score SDE 抽出這個教訓,再建立一般的 FM 框架,最後看保證它成立的連續方程。
讀完這頁你應該能夠
- 說明 FM 和 neural ODE 的差別:前者用回歸學速度場、訓練不需要解 ODE(simulation-free)。
- 用 Score SDE 的 PF-ODE 說明「條件速度的後驗平均就是邊際速度」,並說出這和 DSM 是同一個條件化原則。
- 寫出 FM 的三個步驟:條件路徑與邊際、速度場與連續方程、條件流匹配(CFM)損失;證明 \(\mathcal L_{\text{FM}}=\mathcal L_{\text{CFM}}+C\),並說明為什麼 \(C\le0\)。
- 說明為什麼同一條機率路徑可以有很多不同的速度場(連續方程不唯一)。
- (選讀)說明質量守恆判準,以及為什麼「每個條件速度生成它的條件路徑」就保證邊際速度生成邊際路徑。
前置知識
- B11:neural ODE、它的最大概似訓練為什麼貴、本章翻轉的時間方向。
- B09:Fokker–Planck 方程、PF-ODE 的速度場、flow map 與 pushforward、同一個邊際有多種條件分佈。
- B06:DSM 與「條件 score 的後驗平均就是邊際 score」;B10:從 \(\alpha_t,\sigma_t\) 推出 \(f,g\) 的 Lemma。
- 連續方程:CS236 補充課 C;書的附錄 B(B34)。
「書 §」標出原文位置。灰框「轉述者補充」是書裡沒有的內容;藍框「推導」展開書中省略的步驟或附錄 D 的證明。書把 §5.2.4 標為選讀。
01Flow Matching 要解決什麼書 §5.2
Score SDE(第 4 章)和 neural ODE(B11)都把生成建模描述成一個連續的運送過程。Score SDE 透過隨機動態和對應的 PF-ODE 來寫這個運送,neural ODE 則直接參數化一條確定性的 ODE。在標準的生成設定裡,兩者都想把簡單的高斯先驗 \(p_{\text{prior}}\) 運送成資料分佈 \(p_{\text{data}}\)。
neural ODE 的一個實務限制是:基於概似的訓練需要反覆解 ODE。Flow Matching(FM;Lipman et al. 2022、2024;Tong et al.)避開這個代價:用一個回歸目標訓練神經網路直接預測速度場。所以只有生成時需要積分 ODE,訓練迴圈裡不需要,FM 的訓練是 simulation-free 的。
更一般地,FM 不限於「高斯到資料」的運送。它考慮一個來源分佈 \(p_{\text{src}}\) 和一個目標分佈 \(p_{\text{tgt}}\),兩者都假設容易取樣。標準的生成建模是特例
$$p_{\text{src}}=p_{\text{prior}},\qquad p_{\text{tgt}}=p_{\text{data}},$$\(p_{\text{prior}}\) 通常是高斯。
這一節採用 FM 的觀點。書的註腳:幾個相關的方法共享「用連續時間的流在兩端分佈之間運送」這個核心想法,只是寫法略有不同,包括 Flow Matching(Lipman et al. 2022;Neklyudov et al. 2023)、Rectified Flow(Liu et al. 2022;Heitz et al. 2023)和 Stochastic Interpolants(Albergo et al. 2023;Ma et al. 2024);書用 FM 當統稱。FM 的目標是學一個隨時間變化的速度場 \(\mathbf v_t(\mathbf x_t)\),讓它對應的 ODE 沿著一條指定的機率路徑 \(\{p_t\}_{t\in[0,1]}\) 運送樣本,滿足
$$p_0=p_{\text{src}},\qquad p_1=p_{\text{tgt}}.$$當 \(p_{\text{src}}\) 是高斯時,叫做 Gaussian Flow Matching。所以 FM 既是生成建模的 simulation-free 訓練方法,也是學習任意兩端分佈之間連續運送的一般框架。
02從 score-based 方法學到的事書 §5.2.1
用一個稍微不同但等價的寫法回顧 Score SDE(第 4 章),抽出啟發 FM 的關鍵洞見。這個分析揭露 diffusion model 怎麼隱含地學習機率流,並引出一個更直接的寫法。(這一節用 diffusion 時鐘:\(t=0\) 是資料。)
步驟 1:定義條件路徑與它的邊際密度
diffusion model 指定一族連續時間的密度 \(\{p_t\}_{t\in[0,1]}\),把 \(t=1\) 的簡單先驗 \(p_{\text{prior}}\)(例如高斯,當作來源)運送到 \(t=0\) 的目標資料分佈 \(p_{\text{data}}\):\(p_1(\mathbf x_1)=p_{\text{prior}}(\mathbf x_1)\)、\(p_0(\mathbf x_0)=p_{\text{data}}(\mathbf x_0)\)。這條路徑透過前向條件分佈隱含地定義
$$p_t(\mathbf x_t\mid\mathbf x_0)=\mathcal N(\mathbf x_t;\alpha_t\mathbf x_0,\sigma_t^2\mathbf I),\qquad\mathbf x_0\sim p_{\text{data}},$$它誘導出邊際密度 \(p_t(\mathbf x_t):=\int p_t(\mathbf x_t\mid\mathbf x_0)\,p_{\text{data}}(\mathbf x_0)\,d\mathbf x_0\)。條件高斯的變異數 \(\sigma_t^2\) 不斷增加,推動 \(p_t\) 往高斯先驗演化。
步驟 2:速度場
邊際密度 \(p_t\) 的時間演化由一個速度場 \(\mathbf v_t:\mathbb R^D\to\mathbb R^D\) 支配,它由 Fokker–Planck 方程推出(B09):
$$\mathbf v_t(\mathbf x):=f(t)\mathbf x-\frac12g^2(t)\nabla_{\mathbf x}\log p_t(\mathbf x),$$它透過 PF-ODE 定義一個確定性的粒子流:
$$\frac{d\mathbf x(t)}{dt}=\underbrace{f(t)\mathbf x(t)-\frac12g^2(t)\nabla_{\mathbf x}\log p_t\big(\mathbf x(t)\big)}_{\mathbf v_t(\mathbf x(t))}.$$這條 ODE 把起始的 \(\mathbf x(0)\sim p_{\text{data}}\) 往前運送,或把 \(\mathbf x(1)\sim p_{\text{prior}}\) 往回運送,讓 \(\mathbf x(t)\) 的演化邊際在每個 \(t\in[0,1]\) 都和 \(p_t\) 相符(見下面的「背後的規則」)。純量函數 \(f(t)\)、\(g(t)\) 由對應的前向 SDE 的係數決定,或等價地由條件路徑裡高斯核的參數 \(\alpha_t\)、\(\sigma_t\) 決定(B10 的 Lemma)。
步驟 3:透過條件化策略學習
PF-ODE 的速度依賴邊際 score \(\nabla_{\mathbf x_t}\log p_t(\mathbf x_t)\),它一般是未知的。所以 score-based 方法用神經網路 \(\mathbf s_{\bm\phi}(\mathbf x_t,t)\),透過理想的 score matching 目標學它:
$$\mathcal L_{\text{SM}}(\bm\phi)=\mathbb E_{t\sim\mathcal U[0,1],\,\mathbf x_t\sim p_t}\Big[\big\Vert\mathbf s_{\bm\phi}(\mathbf x_t,t)-\nabla_{\mathbf x_t}\log p_t(\mathbf x_t)\big\Vert^2\Big].$$學到 score 之後,PF-ODE 的速度就由它得到:\(\mathbf v_t(\mathbf x_t)=f(t)\mathbf x_t-\tfrac12g^2(t)\nabla_{\mathbf x_t}\log p_t(\mathbf x_t)\),其中 \(\mathbf s_{\bm\phi^\times}(\mathbf x,t)\approx\nabla_{\mathbf x}\log p_t(\mathbf x)\)。
困難在於:計算回歸目標時,邊際 score 本身拿不到。和第 3 章一樣,用條件化技巧克服它。條件分佈 \(p_t(\mathbf x_t\mid\mathbf x_0)\) 好處理,而它的條件 score 滿足(B06)
$$\nabla_{\mathbf x_t}\log p_t(\mathbf x_t)=\mathbb E_{\mathbf x_0\sim p(\cdot\mid\mathbf x_t)}\big[\nabla_{\mathbf x_t}\log p_t(\mathbf x_t\mid\mathbf x_0)\big].$$這啟發了 denoising score matching 目標
$$\mathcal L_{\text{DSM}}(\bm\phi):=\mathbb E_{t,\,\mathbf x_0\sim p_{\text{data}},\,\mathbf x_t\sim p_t(\cdot\mid\mathbf x_0)}\Big[\big\Vert\mathbf s_{\bm\phi}(\mathbf x_t,t)-\nabla_{\mathbf x_t}\log p_t(\mathbf x_t\mid\mathbf x_0)\big\Vert^2\Big].$$如 B06 的定理,\(\mathcal L_{\text{SM}}\) 和 \(\mathcal L_{\text{DSM}}\) 只差一個和 \(\bm\phi\) 無關的常數,所以有相同的母體最佳解 \(\mathbf s^*(\mathbf x_t,t)=\mathbb E_{\mathbf x_0\sim p(\cdot\mid\mathbf x_t)}[\nabla_{\mathbf x_t}\log p_t(\mathbf x_t\mid\mathbf x_0)]=\nabla_{\mathbf x_t}\log p_t(\mathbf x_t)\)。
重要的是,同一個條件化原則也可以直接寫在 PF-ODE 速度的層級。定義條件速度
$$\mathbf v_t(\mathbf x_t\mid\mathbf x_0):=f(t)\mathbf x_t-\frac12g^2(t)\nabla_{\mathbf x_t}\log p_t(\mathbf x_t\mid\mathbf x_0).$$用上面的條件 score 恆等式,立刻得到
「立刻」的那一步
\(f(t)\mathbf x_t\) 不依賴 \(\mathbf x_0\),取後驗平均不變;\(-\tfrac12g^2(t)\) 是常數,可以拿出期望外。所以
$$\mathbb E_{\mathbf x_0\mid\mathbf x_t}[\mathbf v_t(\mathbf x_t\mid\mathbf x_0)]=f(t)\mathbf x_t-\tfrac12g^2(t)\,\mathbb E_{\mathbf x_0\mid\mathbf x_t}[\nabla\log p_t(\mathbf x_t\mid\mathbf x_0)]=f(t)\mathbf x_t-\tfrac12g^2(t)\nabla\log p_t(\mathbf x_t)=\mathbf v_t(\mathbf x_t).$$也就是說,條件速度和條件 score 只差一個仿射變換,「取後驗平均」和仿射變換可以交換。
所以 score-based diffusion 透過好處理的條件目標學 score,再用它建構 PF-ODE 的速度。Flow Matching 會把同一個「條件 → 邊際」的原則直接套用在速度場上。
轉述者注:圖 (2)、(4) 的箭頭畫成從 \(\mathbf x_t\) 指向資料那一側,也就是生成時沿時間倒著走的方向。在這一節的 diffusion 時鐘裡,PF-ODE 往前(\(t\) 增加)的速度其實是遠離資料;看圖時把箭頭理解成「\(-\mathbf v_t\)」的方向即可,平均的道理完全相同。
背後的規則:Fokker–Planck 方程
邊際密度 \(p_t\) 依 Fokker–Planck 方程演化:
這條偏微分方程保證 PF-ODE 給出的密度和前向 SDE 的邊際相符。為了看出這點,回想 PF-ODE 的 flow map \(\bm\Psi_{s\to t}(\mathbf x_s)\)(B09),它把時間 \(s\) 的初始狀態 \(\mathbf x_s\) 直接帶到它在 \(t\) 的狀態。把 PF-ODE 從 \(t=1\) 倒著跑到 \(t=0\)、以 \(\mathbf x_1\sim p_{\text{prior}}\) 起始,透過 pushforward 公式得到隨時間變化的密度:
$$p_t^{\text{rev}}(\mathbf x)=\int\delta\big(\mathbf x-\bm\Psi_{1\to t}(\mathbf x_1)\big)\,p_{\text{prior}}(\mathbf x_1)\,d\mathbf x_1.$$Fokker–Planck 方程保證誘導出的密度路徑就是同一條演化中的密度:\(p_t^{\text{rev}}=p_t\)。特別是 \(p_0^{\text{rev}}=p_0=p_{\text{data}}\),在 \(t=0\) 還原了資料分佈。因為 ODE 的解映射是雙向的,也可以從 \(\mathbf x_0\sim p_{\text{data}}\) 出發、往前解 ODE,做平行的分析。
03Flow Matching 框架書 §5.2.2
第 2 節的分析揭露:diffusion model 之所以成功,是因為它在滿足邊界條件的同時,學了一個在分佈之間運送的速度場(具體地說是 score)。第 2 節的高斯條件路徑設計(變異數 \(\sigma_t^2\) 遞增)隱含地把一端錨定在高斯先驗,同時讓條件密度定義在整個空間上,使得 score 的梯度計算可行。
這一節介紹建立在這個洞見上的 FM 框架(上面那張四格圖同樣適用於 FM),並把它推廣到在兩個任意分佈 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\) 之間學習連續的流。
步驟 1:定義條件路徑與它的邊際密度
考慮 \(\mathbb R^D\) 上任意的來源與目標機率分佈 \(p_{\text{src}}\)、\(p_{\text{tgt}}\)。令(書的註腳:為了和 FM 文獻的標準記號一致,時間軸相對於前面的章節反過來,\(t=0\) 對應來源、\(t=1\) 對應目標)
$$p_0(\mathbf x)=p_{\text{src}}(\mathbf x),\qquad p_1(\mathbf x)=p_{\text{tgt}}(\mathbf x).$$FM 隱含地定義一族在兩端之間內插的中間密度 \(\{p_t\}_{t\in[0,1]}\)。每個邊際 \(p_t\) 都透過一個從已知分佈 \(\pi(\mathbf z)\) 抽出的潛在變數 \(\mathbf z\) 和一個條件分佈 \(p_t(\mathbf x_t\mid\mathbf z)\) 來表示:
$$p_t(\mathbf x_t)=\int p_t(\mathbf x_t\mid\mathbf z)\,\pi(\mathbf z)\,d\mathbf z,$$\((\pi(\mathbf z),\{p_t(\cdot\mid\mathbf z)\})\) 的選擇要滿足上面的邊界條件。書提醒:一般而言邊際密度 \(p_t\) 算不出來,因為要對 \(\pi(\mathbf z)\) 積分,而 \(\pi(\mathbf z)\) 和條件分佈 \(p_t(\mathbf x_t\mid\mathbf z)\) 都可能很複雜。儘管如此,以潛在變數 \(\mathbf z\) 為條件,讓 FM 能建模比第 2 節更廣泛的一類內插路徑。\(\mathbf z\) 的常見選擇有:
- 雙側條件化:\(\mathbf z=(\mathbf x_0,\mathbf x_1)\sim p_{\text{src}}(\mathbf x_0)\,p_{\text{tgt}}(\mathbf x_1)\),\(\pi\) 把來源和目標分佈耦合起來。這讓 FM 能定義任意分佈之間的運送。
- 單側條件化:\(\mathbf z=\mathbf x_0\) 或 \(\mathbf z=\mathbf x_1\)。來源分佈取成高斯時,特別會還原出類似 diffusion 的設定。
在所有情況下,條件路徑 \(p_t(\mathbf x_t\mid\mathbf z)\) 都應該能有效率地取樣,對應的條件速度也要好處理。在下一頁「先指定機率路徑」的建構裡,還會進一步使用有封閉形式的條件分佈。具體的建構在 §5.3(B13)。
雙側條件化的一個例子(B13 會正式推導)
最簡單的雙側選擇:\(\mathbf z=(\mathbf x_0,\mathbf x_1)\) 各自獨立抽樣,條件路徑取一個點 \(\mathbf x_t=(1-t)\mathbf x_0+t\mathbf x_1\)(直線內插)。對 \(t\) 微分,條件速度就是 \(\mathbf v_t(\mathbf x_t\mid\mathbf z)=\mathbf x_1-\mathbf x_0\):一個常數向量,完全不用解任何方程。兩個條件都滿足:\(\mathbf x_t\) 一行就抽得到,條件速度有封閉形式。Demo A 用的就是這個。
步驟 2:速度場
機率路徑 \(\{p_t\}_{t\in[0,1]}\) 規定了分佈該怎麼隨時間演化,但一般並沒有規定個別樣本該怎麼移動。換句話說,同一串邊際密度可以由不同的速度場實現。
書的註腳:關鍵的區別是,機率路徑規定的是分佈的演化,一般不規定實現它的粒子層級的運送。只有再加上額外的運送結構,速度才會被決定。例如,一個夠正則、可逆的 flow map \(\bm\Psi_{0\to t}\) 透過 \(\mathbf v_t=\partial_t\bm\Psi_{0\to t}\circ\bm\Psi_{0\to t}^{-1}\) 決定它的 Euler 速度;同樣地,一旦指定了前向 diffusion SDE,它的 PF-ODE 就選出特定的速度 \(\mathbf v_t(\mathbf x)=f(t)\mathbf x-\tfrac12g^2(t)\nabla_{\mathbf x}\log p_t(\mathbf x)\)。兩種情況都在邊際路徑 \(\{p_t\}\) 本身之外,加上了粒子層級的運送結構。
所以 FM 的目標是找一個速度場 \(\mathbf v_t(\mathbf x)\),讓它誘導的 ODE
$$\frac{d\mathbf x(t)}{dt}=\mathbf v_t(\mathbf x(t)),\qquad t\in[0,1],$$在每個時間 \(t\) 產生的邊際分佈都和指定的 \(p_t\) 相符。成立時,這條 ODE 可以把樣本從 \(\mathbf x(0)\sim p_{\text{src}}\) 往前運送到 \(p_{\text{tgt}}\),或從 \(\mathbf x(1)\sim p_{\text{tgt}}\) 倒著運送到 \(p_{\text{src}}\)(更正式的討論在第 5 節)。這個要求由連續方程表達(書的註腳:它是 Fokker–Planck 方程沒有擴散項的確定性類比):
任何滿足連續方程的速度場 \(\mathbf v_t\),配上初始條件 \(p_0=p_{\text{src}}\),都會生成一條邊際依指定路徑 \(\{p_t\}\) 演化的 ODE(在適當的正則性條件下)。所以解 ODE 就提供了 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\) 之間逐樣本的運送,同時匹配中間的分佈。
重要的是,連續方程約束了機率密度的演化,但一般不唯一決定速度場。不同的樣本軌跡可以產生完全相同的邊際密度演化。例如,在 \(p_t>0\) 的區域,若 \(\mathbf v_t\) 滿足連續方程,則對任何夠正則、散度為零(\(\nabla\cdot\tilde{\mathbf v}_t=0\))的場 \(\tilde{\mathbf v}_t\),
$$\mathbf v_t+\frac1{p_t}\tilde{\mathbf v}_t$$也滿足它。
為什麼加上 \(\tilde{\mathbf v}_t/p_t\) 不改變密度的演化
所以連續方程的每一項都不變。直覺:\(\tilde{\mathbf v}_t\) 是一股「流進多少就流出多少」的環流,它只讓粒子在等密度的方向上轉來轉去,不會讓任何地方的機率質量增加或減少。Demo B 用一個具體例子看這件事:平移的高斯,加不加旋轉,邊際完全相同,軌跡卻一個是直線、一個是螺旋。
因此 FM 要找的是某一個能實現指定機率路徑的速度場。對一般的來源與目標分佈,邊際密度 \(p_t\) 和這樣的速度場通常都沒有封閉形式。這引出接下來的條件化建構:用好處理的條件路徑和條件速度,學一個合法的邊際速度場。書在 §5.3.1(B13)考慮一個高斯到高斯的路徑配上特定的仿射 flow map,那時對應的速度場可以顯式算出。
步驟 3:透過條件化策略學習
選定機率路徑和對應的邊際速度場 \(\mathbf v_t\) 之後,FM 用神經網路 \(\mathbf v_{\bm\phi}\) 直接學這個速度場。理想的 FM 目標是
$$\mathcal L_{\text{FM}}(\bm\phi)=\mathbb E_{t,\,\mathbf x_t\sim p_t}\Big[\big\Vert\mathbf v_{\bm\phi}(\mathbf x_t,t)-\mathbf v_t(\mathbf x_t)\big\Vert^2\Big].$$這種參數化叫做 \(\mathbf v\)-prediction(速度預測),因為網路直接預測 ODE 的漂移。這和第 2 節的 score-based 建構不同,那裡先學 score,再用它建構 PF-ODE 的速度。
困難在於邊際速度 \(\mathbf v_t(\mathbf x_t)\) 一般算不出來。和 denoising score matching 一樣,用條件化策略克服:引入潛在變數 \(\mathbf z\sim\pi(\mathbf z)\),建構好處理的條件機率路徑 \(p_t(\mathbf x_t\mid\mathbf z)\) 和對應的條件速度場 \(\mathbf v_t(\mathbf x_t\mid\mathbf z)\)。如果每個條件速度場都生成它對應的條件機率路徑,那麼對 \(\mathbf z\) 取邊際就得到一個合法的邊際速度場:
也就是說,一般算不出來的邊際速度,是好處理的條件速度的條件平均。這直接類比第 2 節的條件 score 恆等式。這個觀察導出條件流匹配(conditional flow matching,CFM)目標:
這裡 \(\mathbf z\) 用來建構訓練樣本 \(\mathbf x_t\) 和條件回歸目標 \(\mathbf v_t(\mathbf x_t\mid\mathbf z)\),它不是神經網路 \(\mathbf v_{\bm\phi}(\mathbf x_t,t)\) 的額外輸入。重要的是,兩個目標只差一個和 \(\bm\phi\) 無關的常數 \(C\le0\):
$$\mathcal L_{\text{FM}}(\bm\phi)=\mathcal L_{\text{CFM}}(\bm\phi)+C,$$所以兩者有相同的母體最佳解 \(\mathbf v^*(\mathbf x_t,t)=\mathbf v_t(\mathbf x_t)\)。換句話說,雖然每個訓練樣本都用條件速度 \(\mathbf v_t(\mathbf x_t\mid\mathbf z)\) 當目標,對這些樣本平均就還原了想要的邊際速度場。
要讓 CFM 提供好處理、simulation-free 的訓練,實務上需要兩個條件:
- 不必解 ODE,就能輕鬆抽樣 \(\mathbf x_t\sim p_t(\cdot\mid\mathbf z)\);
- 條件速度 \(\mathbf v_t(\mathbf x_t\mid\mathbf z)\) 有好處理的式子,可以當回歸目標。
滿足這些條件的具體建構在 §5.3(B13)。所以就像 denoising score matching 透過好處理的條件 score 學到算不出來的邊際 score,Flow Matching 透過好處理的條件速度目標學到算不出來的邊際速度。
把上面的討論整理成定理:
證明(書說和 DSM 的論證完全相同;這裡寫出來)
固定 \(t\)。把兩個平方展開。\(\Vert\mathbf v_{\bm\phi}\Vert^2\) 項:只依賴 \(\mathbf x_t\),而先抽 \(\mathbf z\) 再抽 \(\mathbf x_t\mid\mathbf z\) 的 \(\mathbf x_t\) 恰好服從 \(p_t\),兩邊相同。交叉項:
$$\mathbb E_{\mathbf z,\mathbf x_t\mid\mathbf z}\big[\mathbf v_{\bm\phi}(\mathbf x_t)^\top\mathbf v_t(\mathbf x_t\mid\mathbf z)\big]=\mathbb E_{\mathbf x_t\sim p_t}\Big[\mathbf v_{\bm\phi}(\mathbf x_t)^\top\,\mathbb E_{\mathbf z\mid\mathbf x_t}[\mathbf v_t(\mathbf x_t\mid\mathbf z)]\Big]=\mathbb E_{\mathbf x_t\sim p_t}\big[\mathbf v_{\bm\phi}(\mathbf x_t)^\top\mathbf v_t(\mathbf x_t)\big],$$(把聯合分佈寫成 \(p_t(\mathbf x_t)\,p(\mathbf z\mid\mathbf x_t)\),再用邊際速度的定義),兩邊也相同。剩下的是目標的平方:
$$\mathcal L_{\text{CFM}}-\mathcal L_{\text{FM}}=\mathbb E\Vert\mathbf v_t(\mathbf x_t\mid\mathbf z)\Vert^2-\mathbb E\Vert\mathbf v_t(\mathbf x_t)\Vert^2=\mathbb E_{\mathbf x_t}\Big[\operatorname{tr}\operatorname{Cov}_{\mathbf z\mid\mathbf x_t}\big(\mathbf v_t(\mathbf x_t\mid\mathbf z)\big)\Big]\ge0,$$和 \(\bm\phi\) 無關。所以 \(C=-\mathbb E[\text{條件速度的後驗變異數}]\le0\),這就是書說 \(C\le0\) 的原因。最小化解:每個 \(\mathbf x_t\) 是一個平方誤差,最佳值是目標的條件平均 \(\mathbb E_{\mathbf z\mid\mathbf x_t}[\mathbf v_t(\mathbf x_t\mid\mathbf z)]=\mathbf v_t(\mathbf x_t)\)。
這是條件化技巧第三次給出好處理的訓練目標,如上圖所示。值得注意的是,變分、score-based 和 flow-based 這三種方法都反映了同一個背後的原則。
和第 2 節一樣(那裡條件密度和條件速度都要明確指定:\(p_t(\mathbf x_t\mid\mathbf x_0)=\mathcal N(\mathbf x_t;\alpha_t\mathbf x_0,\sigma_t^2\mathbf I)\)、\(\mathbf v_t(\mathbf x_t\mid\mathbf x_0)=f(t)\mathbf x_t-\tfrac12g^2(t)\nabla_{\mathbf x_t}\log p_t(\mathbf x_t\mid\mathbf x_0)\)),一般的條件流匹配框架也需要這兩個元件。但我們還沒在一般情況下建構條件密度 \(p_t(\mathbf x_t\mid\mathbf z)\) 或條件速度場 \(\mathbf v_t(\mathbf x_t\mid\mathbf z)\)。§5.3(B13)會介紹幾個常見的實例。
Demo A條件流匹配:雜訊很大的目標,平均出邊際速度轉述者補充
來源 \(p_{\text{src}}=\mathcal N(0,1)\),目標是 B08–B11 的 \(0.6\,\mathcal N(1.2,0.25^2)+0.4\,\mathcal N(-1,0.2^2)\),用上面補充框裡的雙側直線路徑:\(x_t=(1-t)x_0+tx_1\)、條件速度 \(x_1-x_0\),\(x_0\)、\(x_1\) 獨立抽樣。這個例子的邊際速度 \(v_t(x)=\mathbb E[x_1-x_0\mid x_t=x]\) 可以精確算出(給定混合成分時 \((x_t,x_1-x_0)\) 是聯合高斯)。按「訓練」:一個兩層、每層 48 個神經元的 MLP 只用 CFM 損失訓練 3,000 步,然後用 Euler 法從 \(t=0\) 解 ODE 到 \(t=1\) 生成樣本。
預設 \(t=0.5\) 時,點雲上下散得很開:同一個 \(x_t\) 對應的條件目標 \(x_1-x_0\) 有正有負,因為不知道它是從哪一對 \((x_0,x_1)\) 來的。灰色虛線是這團點在每個 \(x\) 的平均,也就是邊際速度。網路每一步只看到這些雜訊很大的單點目標,但最後學到的藍線在 \(p_t\) 有質量的範圍(底部灰色區域)貼在灰色虛線上,兩端沒有訓練樣本的地方才會分開;那些地方 ODE 的軌跡也幾乎不會經過。本機瀏覽器實測(訓練約 5.5 秒的純計算):
| 訓練步數 | 500 | 1000 | 1500 | 2000 | 2500 | 3000 |
|---|---|---|---|---|---|---|
| \(\mathcal L_{\text{FM}}\)(對精確 \(v_t\) 的誤差) | 0.0392 | 0.0195 | 0.0247 | 0.0339 | 0.0097 | 0.0061 |
這裡的 \(\mathcal L_{\text{FM}}\) 是用 \(t\in\{0.05,0.15,\dots,0.95\}\) 和精確密度做數值積分算的,訓練本身從頭到尾沒用到它。前 2,100 步的學習率是 \(3\times10^{-3}\),\(\mathcal L_{\text{FM}}\) 在 0.02–0.03 附近上下跳;學習率降到 \(7.5\times10^{-4}\) 後才降到 0.006。學不掉的部分 \(-C\)(條件速度的後驗變異數)是 1.428,比可以學的部分大兩百倍以上:CFM 損失的數值幾乎全是這一項,所以看 CFM 損失本身判斷不出網路學得好不好。生成 2,000 個樣本(Euler 100 步):學到的 \(v_{\bm\phi}\) 上團 59.4%、平均 0.303;精確的 \(v_t\) 上團 60.5%、平均 0.323;真值是 60% 和 0.320。精確速度場的小誤差來自 Euler 離散化和取樣。
04Diffusion model、一般的 FM 與 NODE 的比較書 §5.2.3
Diffusion model 與一般 FM 的比較
第 2 節的洞見引出保留同樣原則的擴展 FM 框架。為了突顯相似之處,整理成下表。這裡「一般的 FM」指雙側條件化、\(\mathbf x_0\sim p_{\text{src}}\) 和 \(\mathbf x_1\sim p_{\text{tgt}}\) 獨立抽樣的設定。
| 面向 | (score-based)diffusion model | 一般的 FM |
|---|---|---|
| 來源分佈 \(p_{\text{src}}\) | 高斯先驗 | 任意 |
| 目標分佈 \(p_{\text{tgt}}\) | 資料分佈 | 任意 |
| 潛在分佈 \(\pi(\mathbf z)\) | \(p_{\text{data}}\) | 見 §5.3(B13) |
| 條件分佈 \(p_t(\mathbf x_t\mid\mathbf z)\) | \(\mathcal N(\mathbf x_t;\alpha_t\mathbf x_0,\sigma_t^2\mathbf I)\) | 見 §5.3(B13) |
| 邊際分佈 \(p_t(\mathbf x_t)\) | \(\int p_t(\mathbf x_t\mid\mathbf x_0)p_{\text{data}}(\mathbf x_0)\,d\mathbf x_0\) | \(\int p_t(\mathbf x_t\mid\mathbf z)\pi(\mathbf z)\,d\mathbf z\) |
| 條件速度 \(\mathbf v_t(\mathbf x\mid\mathbf z)\) | \(f(t)\mathbf x-\tfrac12g^2(t)\nabla\log p_t(\mathbf x\mid\mathbf x_0)\) | 見 §5.3(B13) |
| 邊際速度 \(\mathbf v_t(\mathbf x)\) | \(f(t)\mathbf x-\tfrac12g^2(t)\nabla\log p_t(\mathbf x)\) | 條件速度的後驗平均(第 3 節) |
| 學習目標 | \(\mathcal L_{\text{SM}}=\mathcal L_{\text{DSM}}+C\) | \(\mathcal L_{\text{FM}}=\mathcal L_{\text{CFM}}+C\) |
| 背後的規則 | Fokker–Planck/連續方程 | |
書提醒:只要高斯時程滿足 B10 Lemma 裡和 diffusion SDE 相容的條件(SNR 不增),Gaussian FM 本質上就等價於對應的 diffusion 寫法(第 6 章,B15–B17 會詳談);所以除非特別說明,書不區分兩者。
和 NODE 的連結
FM 可以看成 NODE(B11)的 simulation-free 替代方案。CNF 在最大概似訓練時必須解 ODE,計算量很大;FM 則透過簡單的回歸損失直接回歸一個指定的速度場,繞過了這一點。關鍵的洞見是:一旦連接來源和目標分佈的邊際密度路徑固定下來,訓練時就不需要精確的模擬。
05(選讀)背後的規則書 §5.2.4
連續方程:質量守恆判準
和第 2 節的 PF-ODE 與 Fokker–Planck 分析類似,這裡給一個判準,檢驗 ODE 的流誘導出的密度路徑是否和指定的路徑 \(\{p_t\}_{t\in[0,1]}\) 一致。考慮描述粒子在隨時間變化的速度場 \(\mathbf v_t\) 下流動的 ODE \(\frac{d\mathbf x(t)}{dt}=\mathbf v_t(\mathbf x(t))\)。和前面一樣,這條 ODE 對任何 \(s,t\in[0,1]\) 定義了 flow map \(\bm\Psi_{s\to t}\),特別是把時間 0 的初始點 \(\mathbf x_0\sim p_{\text{src}}\) 運送到它在時間 \(t\) 的狀態。時間 \(t\) 誘導出的分佈是 pushforward
$$p_t^{\text{fwd}}(\mathbf x)=\int\delta\big(\mathbf x-\bm\Psi_{0\to t}(\mathbf x_0)\big)\,p_{\text{src}}(\mathbf x_0)\,d\mathbf x_0=:\bm\Psi_{0\to t}\#p_{\text{src}},$$所以只要 \(\mathbf x_0\sim p_{\text{src}}\),就有 \(\bm\Psi_{0\to t}(\mathbf x_0)\sim p_t^{\text{fwd}}\)。同樣地,也可以透過 \(\bm\Psi_{1\to0}(\mathbf x_1)\) 從 \(\mathbf x_1\sim p_{\text{tgt}}\) 倒著運送到 \(p_{\text{src}}\)。假設給定一條指定的密度路徑 \(\{p_t\}\),並建構了一個速度場 \(\{\mathbf v_t\}\) 來定義粒子流,自然的問題是:
一旦兩種密度演化對齊,就能靈活地用解 ODE 的方式在 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\) 之間運送樣本。一個有原則的檢驗方式就是連續方程,它刻畫了隨時間演化的密度的質量守恆:
證明(書的附錄 D.3,非正式的弱形式論證;嚴格版見 Villani 2008 的「質量守恆公式」)
預備:flow map 與流誘導的密度。對任何初始位置 \(\mathbf x_0\),flow map \(\bm\Psi_t\) 是 \(\frac{d}{dt}\bm\Psi_t(\mathbf x_0)=\mathbf v_t(\bm\Psi_t(\mathbf x_0))\)、\(\bm\Psi_0(\mathbf x_0)=\mathbf x_0\) 的唯一解,對 \(t\) 和 \(\mathbf x_0\) 都連續可微。流誘導的密度是初始密度的 pushforward:\(p_t^{\text{fwd}}(\mathbf x)=p_0(\bm\Psi_t^{-1}(\mathbf x))\,|\det\nabla\bm\Psi_t^{-1}(\mathbf x)|\)(B11 的變數變換公式)。
(⟹)\(p_t^{\text{fwd}}=p_t\) 推得連續方程。用弱形式:對任何有緊支撐的平滑測試函數 \(\varphi\),變數變換 \(\mathbf x=\bm\Psi_t(\mathbf y)\) 給出
$$\int p_t^{\text{fwd}}(\mathbf x)\varphi(\mathbf x)\,d\mathbf x=\int p_0(\mathbf y)\,\varphi(\bm\Psi_t(\mathbf y))\,d\mathbf y.$$兩邊對 \(t\) 微分。左邊是 \(\int\partial_tp_t^{\text{fwd}}\,\varphi\,d\mathbf x\)。右邊用連鎖律和 \(\partial_t\bm\Psi_t(\mathbf y)=\mathbf v_t(\bm\Psi_t(\mathbf y))\):\(\int p_0(\mathbf y)\nabla\varphi(\bm\Psi_t(\mathbf y))\cdot\mathbf v_t(\bm\Psi_t(\mathbf y))\,d\mathbf y\),再把變數換回 \(\mathbf x\),得到 \(\int p_t^{\text{fwd}}(\mathbf x)\nabla\varphi(\mathbf x)\cdot\mathbf v_t(\mathbf x)\,d\mathbf x\)。分部積分(\(\varphi\) 有緊支撐,邊界項為 0):\(=-\int\varphi\,\nabla\cdot(p_t^{\text{fwd}}\mathbf v_t)\,d\mathbf x\)。兩邊相等對任何 \(\varphi\) 都成立,所以
$$\partial_tp_t^{\text{fwd}}+\nabla\cdot(p_t^{\text{fwd}}\mathbf v_t)=0.$$再代入 \(p_t^{\text{fwd}}=p_t\) 就是要的式子。(書把這個方向標成「Sufficient Condition」,邏輯上它證明的是「連續方程是 \(p_t^{\text{fwd}}=p_t\) 的必要條件」;標籤的用詞不影響論證。)
(⟸)連續方程推得 \(p_t^{\text{fwd}}=p_t\)。上面已經證明 \(p_t^{\text{fwd}}\) 滿足同一條連續方程,而且 \(p_0^{\text{fwd}}=p_0\)(因為 \(\bm\Psi_0\) 是恆等映射)。連續方程可以寫成 \(\partial_tp+\mathbf v_t\cdot\nabla p+p\,\nabla\cdot\mathbf v_t=0\),是一條一階線性偏微分方程。若 \(\mathbf v_t\) 連續可微且全域 Lipschitz、\(p_t\) 夠平滑,特徵線法保證平滑解唯一。兩者滿足同一條方程與初始條件,所以 \(p_t=p_t^{\text{fwd}}\)。∎
從條件路徑到邊際路徑
如第 3 節,我們先定義條件機率路徑 \(p_t(\cdot\mid\mathbf z)\) 和對應的條件速度場 \(\mathbf v_t(\cdot\mid\mathbf z)\),再透過下式建構邊際速度場:
$$\mathbf v_t(\mathbf x)=\int\mathbf v_t(\mathbf x\mid\mathbf z)\,\frac{p_t(\mathbf x\mid\mathbf z)\,\pi(\mathbf z)}{p_t(\mathbf x)}\,d\mathbf z.$$但還需要確保這個邊際速度 \(\mathbf v_t\) 誘導的 ODE 流,它的密度路徑和指定的 \(p_t\) 一致。幸運的是,這個驗證完全可以在條件層級完成:如果每個條件速度場 \(\mathbf v_t(\cdot\mid\mathbf z)\) 都誘導出條件密度路徑 \(p_t(\cdot\mid\mathbf z)\),那麼得到的邊際速度 \(\mathbf v_t\) 也誘導出正確的邊際路徑。
證明(書:驗證 \((p_t,\mathbf v_t)\) 滿足連續方程;用「反推」的方式寫,順便看出邊際速度為什麼長這樣)
每個條件對 \((p_t(\cdot\mid\mathbf z),\mathbf v_t(\cdot\mid\mathbf z))\) 都滿足連續方程:\(\frac{d}{dt}p_t(\mathbf x\mid\mathbf z)=-\nabla\cdot\big(\mathbf v_t(\mathbf x\mid\mathbf z)\,p_t(\mathbf x\mid\mathbf z)\big)\)。我們要找一個 \(\mathbf v_t(\cdot)\) 讓 \(\frac{d}{dt}p_t(\mathbf x)=-\nabla\cdot(\mathbf v_t(\mathbf x)\,p_t(\mathbf x))\)。從 \(p_t\) 的定義出發:
$$\frac{d}{dt}p_t(\mathbf x)=\int\frac{d}{dt}p_t(\mathbf x\mid\mathbf z)\,\pi(\mathbf z)\,d\mathbf z=-\int\nabla\cdot\big(\mathbf v_t(\mathbf x\mid\mathbf z)p_t(\mathbf x\mid\mathbf z)\big)\pi(\mathbf z)\,d\mathbf z=-\nabla\cdot\Big(\int\mathbf v_t(\mathbf x\mid\mathbf z)p_t(\mathbf x\mid\mathbf z)\pi(\mathbf z)\,d\mathbf z\Big).$$和目標式比較,(差一個散度為零的項)\(\mathbf v_t(\mathbf x)p_t(\mathbf x)=\int\mathbf v_t(\mathbf x\mid\mathbf z)p_t(\mathbf x\mid\mathbf z)\pi(\mathbf z)\,d\mathbf z\),所以可以定義
$$\mathbf v_t(\mathbf x):=\int\mathbf v_t(\mathbf x\mid\mathbf z)\,\frac{p_t(\mathbf x\mid\mathbf z)}{p_t(\mathbf x)}\,\pi(\mathbf z)\,d\mathbf z,$$正是邊際速度的式子(\(\frac{p_t(\mathbf x\mid\mathbf z)\pi(\mathbf z)}{p_t(\mathbf x)}=p(\mathbf z\mid\mathbf x)\))。命題的證明本質上就是把這個論證倒過來:用這個定義代回去,\((p_t,\mathbf v_t)\) 滿足連續方程,再由質量守恆判準得到結論。∎
這個連結讓我們把可能算不出來的邊際速度場的建構,化約成定義比較簡單的條件速度場 \(\mathbf v_t(\cdot\mid\mathbf z)\),它們依建構本來就比較好處理。
Demo B同一串分佈,不同的流轉述者補充
二維的機率路徑:\(p_t=\mathcal N(\mathbf m(t),\mathbf I)\),中心 \(\mathbf m(t)=(4t-2,\,0)\) 從左平移到右。兩個速度場都實現這條路徑:
- 純平移:\(\mathbf v_t(\mathbf x)=\mathbf m'(t)=(4,0)\)。
- 平移加旋轉:\(\mathbf v_t(\mathbf x)=\mathbf m'(t)+\omega\,\mathbf J(\mathbf x-\mathbf m(t))\),\(\mathbf J=\begin{bmatrix}0&-1\\1&0\end{bmatrix}\)(逆時針轉 90°)。多出來的 \(\omega\mathbf J(\mathbf x-\mathbf m)\) 就是第 3 節的 \(\tilde{\mathbf v}_t/p_t\),\(\tilde{\mathbf v}_t=p_t\,\omega\mathbf J(\mathbf x-\mathbf m)\) 的散度為 0(\(p_t\) 是繞 \(\mathbf m\) 的圓對稱,而 \(\mathbf J(\mathbf x-\mathbf m)\) 永遠沿著等密度圓的切線方向)。
500 個粒子從 \(\mathcal N(\mathbf m(0),\mathbf I)\) 出發。拖動 \(t\),比較兩種流:分佈(灰色的等高線,也就是 \(p_t\))完全一樣,粒子的軌跡卻完全不同。
讀數裡兩組粒子的平均、變異數和共變異數在每個 \(t\) 都幾乎相同,而且和 \(p_t\) 的 \(\mathbf m(t)\)、\(\mathbf I\) 相符(和真值差約 0.1,是 500 個樣本的抽樣誤差,兩組共用同一批起點,所以誤差也一樣)。但同一個起點的位置不一樣:旋轉的流在時間 \(t\) 把相對位置轉了 \(\omega t\) 弧度,平均距離是 \(2|\sin(\omega t/2)|\,\mathbb E\Vert\mathbf u\Vert=2|\sin(\omega t/2)|\sqrt{\pi/2}\)。\(\omega=6\) 時:\(t=0.25\) 實測 1.75(公式 1.71)、\(t=0.5\) 實測 2.55(公式 2.50);\(t=1\) 時轉了 6 弧度,幾乎轉回一圈,距離只剩 0.36(公式 0.35)。\(\omega=0\) 時兩種流完全重合,距離是 0。這正是第 3 節說的:連續方程只看密度,看不到粒子怎麼「排隊」。FM 選的是條件速度的後驗平均那一個;B13 會看到,就算來源和目標都是高斯,不同的選擇也會給出不同的軌跡。