擴散模型原理 B14 · Reflow 與本章小結
The Principles of Diffusion Models · 第 5 章 §5.4–5.5

Reflow:換掉配對,把彎曲的流拉直

B13 的結尾留下一個問題:每條訓練用的條件路徑都是直線,生成時的邊際軌跡卻是彎的,少步數的 Euler 法就會出錯。問題出在配對:獨立配對下,同一個 \(\mathbf x_t\) 可能來自方向完全不同的端點對,邊際速度只能取平均。Rectified flow 的解法是:用學到的 ODE 把每個來源點送到它的終點,再用這個「有依賴關係的配對」重新訓練。這個步驟叫 Rectify,反覆做就是 Reflow。它不會增加運送成本,也會讓軌跡越來越直。但直不等於最佳:直線的配對不一定是最佳運送。本頁最後是第 5 章的結語。

書的範圍:第 5 章 §5.4(選讀,含 §5.4.1–5.4.3)、§5.5 閱讀時間:約 40 分鐘

讀完這頁你應該能夠

  1. 說明為什麼獨立配對會讓邊際軌跡彎曲,以及 Rectify 怎麼用學到的 ODE 造出有依賴關係的配對。
  2. 說出 Rectify 為什麼保持邊際分佈,以及 Reflow 的兩個步驟(重新擬合、產生新配對)。
  3. 用 Jensen 不等式證明 Rectify 不增加凸成本,並用伸縮和證明 Reflow 的直度上界。
  4. 用旋轉的例子說明:直線配對不一定是最佳運送配對。
  5. 說出第 5 章和前兩個觀點(變分、score-based)共有的結構:連續時間的生成過程,加上條件化技巧。

前置知識

  • B13:典型仿射插值 \((1-t)\mathbf x_0+t\mathbf x_1\)、條件路徑直不代表邊際軌跡直、一步 Euler 在獨立配對下會塌縮。
  • B12:CFM 的最小化解是條件速度的後驗平均;「每個條件速度生成它的條件路徑 ⟹ 邊際速度生成邊際路徑」。
  • Jensen 不等式:凸函數 \(c\) 滿足 \(c(\mathbb E[\mathbf u])\le\mathbb E[c(\mathbf u)]\);條件版本 \(c(\mathbb E[\mathbf u\mid\mathbf x])\le\mathbb E[c(\mathbf u)\mid\mathbf x]\)。
名詞。「配對」(coupling)是一個聯合分佈 \(\pi(\mathbf x_0,\mathbf x_1)\),它的兩個邊際分別是 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\)。獨立配對是 \(p_{\text{src}}(\mathbf x_0)p_{\text{tgt}}(\mathbf x_1)\);確定性配對是 \(p_{\text{src}}(\mathbf x_0)\,\delta(\mathbf x_1-\mathbf T(\mathbf x_0))\),每個來源點只配一個終點。最佳運送(OT)的正式定義在第 7 章(B18)。

書把 §5.4 標為選讀。「書 §」標出原文位置。灰框「轉述者補充」是書裡沒有的內容;藍框「推導」展開書中省略的步驟。

01典型的仿射流書 §5.4

給定兩個端點分佈 \(p_0=p_{\text{src}}\)、\(p_1=p_{\text{tgt}}\),flow matching(Lipman et al. 2022)和 rectified flow(RF;Liu et al. 2022)定義條件路徑時最自然、最常用的選擇是線性插值 \(a_t=1-t\)、\(b_t=t\),得到

$$\mathbf x_t=(1-t)\mathbf x_0+t\mathbf x_1,\qquad\mathbf x_0\sim p_{\text{src}},\ \mathbf x_1\sim p_{\text{tgt}}.$$

在這個選擇下,訓練目標化簡為

$$\mathbb E_{t\sim\mathcal U[0,1]}\,\mathbb E_{\mathbf x_0,\mathbf x_1}\Big[\big\Vert\mathbf v_{\bm\phi}(\mathbf x_t,t)-(\mathbf x_1-\mathbf x_0)\big\Vert_2^2\Big].$$

這個線性流有幾個吸引人的性質。特別是,它允許一個叫 Reflow 的反覆改良方法:在保持邊際的同時,逐步把分佈之間的路徑拉直。

02Rectify:從亂配到有結構的配對書 §5.4.1

為什麼要換配對?

在獨立配對 \(\pi(\mathbf x_0,\mathbf x_1)=p_{\text{src}}(\mathbf x_0)p_{\text{tgt}}(\mathbf x_1)\) 下,每條條件插值 \((1-t)\mathbf x_0+t\mathbf x_1\) 都是直線。但如 B13 第 6 節強調的,對應的邊際速度是把許多不相關的端點對的方向平均起來,所以誘導出的邊際 ODE 軌跡不一定是直的。這種彎曲的邊際動態,只用少數幾步數值積分時比較難算準。

Rectification 透過改變端點的配對處理這個問題:用一條學好的 ODE 把每個來源樣本送到模型產生的終點,得到一個有依賴關係的配對,再用它定義新的條件插值。

用有依賴關係的配對來 rectify

不用任意的配對,而是把一個預先訓練好的 diffusion model \(\mathbf v_{\bm\phi^\times}(\cdot,t)\) 當成 PF-ODE 的漂移,確定性地運送每個來源點。從 \(\mathbf z(0)=\mathbf z_0\sim p_{\text{src}}\) 出發,積分

$$\frac{d\mathbf z(t)}{dt}=\mathbf v_{\bm\phi^\times}(\mathbf z(t),t),\qquad t\in[0,1],$$

得到 \(\hat{\mathbf z}_1:=\mathbf z(1)\),它落在預訓練模型學到的資料空間附近。得到的配對 \((\mathbf z_0,\hat{\mathbf z}_1)\) 是一個有依賴關係的配對:它沿著一條有結構、由模型引導的路徑,而不是任意的插值。這個想法自然地推廣到 \(\mathbf x_t=a_t\mathbf x_0+b_t\mathbf x_1\) 形式的仿射參考路徑。

書中的演算法:Rectify 需要:參考路徑 \(\{\mathbf x_t\}_{t\in[0,1]}\)(例如 \(\mathbf x_t=a_t\mathbf x_0+b_t\mathbf x_1\))。
  1. 預先訓練。在選定的路徑上擬合 \(\mathbf v_{\bm\phi^\times}\): $$\bm\phi^\times\in\arg\min_{\bm\phi}\ \mathbb E_{t,\mathbf x_0,\mathbf x_1}\Big[\Big\Vert\mathbf v_{\bm\phi}(\mathbf x_t,t)-\frac{d\mathbf x_t}{dt}\Big\Vert_2^2\Big].$$
  2. Rectify。抽 \(\mathbf z_0\sim p_{\text{src}}\),積分 \(\frac{d\mathbf z(t)}{dt}=\mathbf v_{\bm\phi^\times}(\mathbf z(t),t)\)、\(\mathbf z(0)=\mathbf z_0\),得到 \(\hat{\mathbf z}_1=\mathbf z(1)\) 和軌跡 \(\{\mathbf z(t)\}_{t\in[0,1]}\)。
輸出:有依賴關係(一致)的配對 \((\mathbf z_0,\hat{\mathbf z}_1)\),或整條軌跡。

為什麼可行:保持邊際的結構

令 \(\bm\Phi_{0\to t}\) 是由上面那條(以預訓練的 \(\mathbf v_{\bm\phi^\times}\) 定義的)ODE 生成的 flow map,則 \(\mathbf z(t)=\bm\Phi_{0\to t}(\mathbf z_0)\)、\(\hat{\mathbf z}_1=\bm\Phi_{0\to1}(\mathbf z_0)\)。Rectify 把每個來源點和它的流終點配在一起,得到確定性的聯合分佈

$$\pi_{\text{Rectify}}(\mathbf z_0,\mathbf z_1)=p_{\text{src}}(\mathbf z_0)\,\delta\big(\mathbf z_1-\bm\Phi_{0\to1}(\mathbf z_0)\big).$$

立刻有兩個結果:

  • 來源邊際保持不變:\(\int\pi_{\text{Rectify}}(\mathbf z_0,\mathbf z_1)\,d\mathbf z_1=p_{\text{src}}(\mathbf z_0)\)。
  • 沿流的 pushforward:\((\bm\Phi_{0\to t})_\#p_{\text{src}}=\mathrm{Law}(\mathbf z(t))\),也就是時間 \(t\) 的分佈是 \(p_{\text{src}}\) 經 \(\bm\Phi_{0\to t}\) 的 pushforward。

如果 \(\mathbf v_{\bm\phi^\times}\) 等於給定參考路徑 \(\mathbf x_t\) 的理想漂移,那麼所有中間的邊際都相同:

$$\mathrm{Law}(\mathbf z(t))=\mathrm{Law}(\mathbf x_t)\quad\text{對所有 }t\in[0,1],\qquad(\bm\Phi_{0\to1})_\#p_{\text{src}}=p_{\text{tgt}}.$$
為什麼「理想漂移」保證每個時刻的邊際都相同

理想漂移就是 CFM 的最小化解 \(\mathbf v^*(\mathbf x,t)=\mathbb E[\tfrac{d\mathbf x_t}{dt}\mid\mathbf x_t=\mathbf x]\)。B12 第 5 節的命題說:只要每條條件路徑由它的條件速度生成(直線插值顯然是),這個後驗平均的速度就生成邊際路徑 \(\mathrm{Law}(\mathbf x_t)\)。所以從 \(p_{\text{src}}\) 出發解這條 ODE,每個時刻的分佈都等於 \(\mathrm{Law}(\mathbf x_t)\);特別是 \(t=1\) 時等於 \(p_{\text{tgt}}\)。因此 \(\pi_{\text{Rectify}}\) 的第二個邊際也是 \(p_{\text{tgt}}\):它是一個合法的配對,只是每個 \(\mathbf z_0\) 只配一個終點。

小結。Rectification 把獨立的端點配對換成由學到的流誘導的配對。反覆做 Reflow,可以得到越來越容易用少數幾步數值積分算準的邊際運送。對典型路徑,反覆套用 Rectify(即 Reflow)會進一步拉直軌跡,而且不增加運送成本,讓訓練更容易。

03Reflow:反覆把流拉直書 §5.4.2

為什麼要 Reflow?獨立配對常常在 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\) 之間誘導出不規則、蜿蜒的 ODE 軌跡,增加模擬時的離散化誤差和變異。這引出一個自然的問題:

書中的問題能不能學到一種配對,它誘導的運送路徑更接近兩個分佈之間的直線,同時仍然保持正確的邊際?

這就是 Reflow 的動機:反覆套用 Rectify 更新配對,讓接下來的流越來越容易積分。

核心想法:用 Rectify 遞迴地拉直。從乘積配對 \(\pi^{(0)}:=p_{\text{src}}(\mathbf x_0)p_{\text{tgt}}(\mathbf x_1)\) 上的典型插值 \(\mathbf x_t=(1-t)\mathbf x_0+t\mathbf x_1\) 開始。套用 Rectify 把獨立配對換成由學到的 ODE 誘導的依賴配對 \((\mathbf z_0,\hat{\mathbf z}_1)\)。重複這個程序(即 Reflow)傾向於產生更直、更容易用少數幾步算準的軌跡。精確的拉直保證在第 4 節。

Reflow 程序。每一輪做兩步:

  • 重新擬合流:用目前配對的樣本訓練一個新的速度場: $$\bm\phi_{k+1}=\arg\min_{\bm\phi}\mathcal L\big(\bm\phi\,\big|\,\pi^{(k)}\big),\qquad\mathcal L\big(\bm\phi\,\big|\,\pi^{(k)}\big):=\mathbb E_{t,(\mathbf z_0^{(k)},\hat{\mathbf z}_1^{(k)})\sim\pi^{(k)}}\Big[\big\Vert\mathbf v_{\bm\phi}(\mathbf z_t,t)-(\hat{\mathbf z}_1^{(k)}-\mathbf z_0^{(k)})\big\Vert^2\Big],$$ 其中 \(\mathbf z_t=(1-t)\mathbf z_0^{(k)}+t\hat{\mathbf z}_1^{(k)}\)。
  • 產生新配對:從新的來源樣本 \(\mathbf z_0^{(k+1)}\sim p_{\text{src}}\) 出發解學到的 ODE:\(\hat{\mathbf z}_1^{(k+1)}\leftarrow\mathbf z_0^{(k+1)}+\int_0^1\mathbf v_{\bm\phi_{k+1}}(\mathbf z(t),t)\,dt\),並定義更新後的配對 \(\pi^{(k+1)}(\mathbf z_0,\mathbf z_1):=p_{\text{src}}(\mathbf z_0)\,\delta\big(\mathbf z_1-\hat{\mathbf z}_1^{(k+1)}\big)\)。

換句話說,Reflow 可以看成反覆套用 Rectify 運算子,產生一串逐步改良的配對:

$$\pi^{(k+1)}=\texttt{Rectify}\big(\pi^{(k)}\big),$$

流和配對一起演化,得到越來越穩定的運送路徑。

第一輪和之後各輪的差別

第一輪(\(k=0\))的訓練資料是獨立抽的 \((\mathbf x_0,\mathbf x_1)\),就是普通的 FM/RF 訓練;學到的流叫 1-rectified flow。之後各輪的訓練資料是「上一個流的輸入與輸出」,需要先用上一個模型生成一大批配對(Demo A 用 2,000 對);用它訓練出的流叫 2-rectified flow,依此類推。代價是每一輪都要額外生成資料並重新訓練,而且模型誤差會一輪輪累積:生成的 \(\hat{\mathbf z}_1\) 只是近似地服從 \(p_{\text{tgt}}\)。

Demo A1-rectified flow 與 2-rectified flow轉述者補充

二維。來源是左邊兩團 \(\mathcal N((-2.5,\pm2),0.4^2\mathbf I)\),目標是右邊兩團 \(\mathcal N((2.5,\pm2),0.4^2\mathbf I)\),各占一半。獨立配對下,一半的直線會從上團連到下團、在中間交叉。這個例子的 1-rectified flow(獨立配對上的邊際速度)可以精確算出(給定兩端各自的團時,\((\mathbf x_t,\mathbf x_1-\mathbf x_0)\) 是聯合高斯)。按「訓練 2-rectified flow」:用精確的 1-rectified flow 把 2,000 個來源點送到終點(RK4 60 步),得到新的配對,再用 Reflow 損失訓練一個兩層、每層 48 個神經元的 MLP(2,500 步)。

來源點從上團出發的軌跡從下團出發的軌跡Euler N 步的終點

瀏覽器實測(和 Node 原型一致;2-rectified flow 的訓練約 5 秒)。把「顯示」依序切成四個選項看:(1) 獨立配對的直線一半從上團連到下團,在中間交叉;(2) 1-rectified flow 不會交叉(ODE 軌跡不能相交),於是上半的軌跡往中間擠、在 \(x\approx0\) 附近幾乎碰到下半的軌跡再彎回去;(3) 把 1-rectified flow 的起點和終點直接連起來,得到的新配對全是水平的直線,上團只配上團;(4) 用這個配對訓練出的 2-rectified flow,軌跡就是這些直線。表裡的數字:

  • 直度從 5.367 降到 0.0004:2-rectified flow 幾乎是等速直線。
  • 成本從獨立配對的 33.66 降到 25.000,2-rectified 是 24.995:第 4 節命題說的「不增加成本」。這個例子裡第一次 Rectify 就把上團配上團、下團配下團,成本已經是平移距離的平方 \(5^2=25\);第二次 Rectify 幾乎不再改變配對,只把路徑拉直。
  • 一步 Euler:1-rectified flow 的終點全部落在 \((2.5,0)\),兩團中間的空白處(把「終點用 Euler」設成 1 就看得到那個孤零零的紫點),平均誤差 2.017;這是 B13 第 6 節的塌縮:在 \(t=0\) 時 \(\mathbf v^*(\mathbf x,0)=\mathbb E[\mathbf x_1]-\mathbf x\)。2-rectified flow 的一步 Euler 誤差只有 0.025,兩步 0.013。

這個例子刻意選得很對稱,所以一輪 Reflow 就幾乎完美。實際的影像資料上,模型誤差和複雜的配對都會讓效果打折扣,書也只保證直度的上界(第 4 節)。

04Reflow 的性質書 §5.4.3

兩個關鍵的理論性質讓 Reflow 有用:它降低運送成本,也拉直軌跡。

I. Reflow 從不增加運送成本

令 \(c(\mathbf y)\) 是凸的成本函數(例如 \(p\ge1\) 時的 \(\Vert\mathbf y\Vert_2^p\))。每一步 Rectify 形成的新配對 \((\mathbf z_0,\hat{\mathbf z}_1)\),成本都不比原本的差:

書中的命題:Rectify 可能降低運送成本 假設速度場是理想的 \(\mathbf v^*=\mathbf v_{\bm\phi^\times}\),則 $$\mathbb E\big[c(\hat{\mathbf z}_1-\mathbf z_0)\big]\le\mathbb E\big[c(\mathbf x_1-\mathbf x_0)\big].$$ 書:由 Jensen 不等式得到,完整推導見 Liu et al. (2022)。
證明:兩次 Jensen

沿著 rectified flow 的軌跡,\(\hat{\mathbf z}_1-\mathbf z_0=\int_0^1\mathbf v^*(\mathbf z(t),t)\,dt\)。

第一次 Jensen(對時間平均)。\(dt\) 在 \([0,1]\) 上是機率測度,\(c\) 凸:

$$c\Big(\int_0^1\mathbf v^*(\mathbf z(t),t)\,dt\Big)\le\int_0^1c\big(\mathbf v^*(\mathbf z(t),t)\big)\,dt.$$

換成參考路徑的分佈。取期望、交換積分;由第 2 節,\(\mathrm{Law}(\mathbf z(t))=\mathrm{Law}(\mathbf x_t)\),所以 \(\mathbb E\,c(\mathbf v^*(\mathbf z(t),t))=\mathbb E\,c(\mathbf v^*(\mathbf x_t,t))\)。

第二次 Jensen(條件期望)。\(\mathbf v^*(\mathbf x_t,t)=\mathbb E[\mathbf x_1-\mathbf x_0\mid\mathbf x_t]\),所以 \(c(\mathbf v^*(\mathbf x_t,t))\le\mathbb E[c(\mathbf x_1-\mathbf x_0)\mid\mathbf x_t]\)。再取期望:\(\mathbb E\,c(\mathbf v^*(\mathbf x_t,t))\le\mathbb E\,c(\mathbf x_1-\mathbf x_0)\),對每個 \(t\) 都成立。合起來:

$$\mathbb E\,c(\hat{\mathbf z}_1-\mathbf z_0)\le\int_0^1\mathbb E\,c\big(\mathbf v^*(\mathbf x_t,t)\big)\,dt\le\int_0^1\mathbb E\,c(\mathbf x_1-\mathbf x_0)\,dt=\mathbb E\,c(\mathbf x_1-\mathbf x_0).\qquad\blacksquare$$

兩個不等號各有意義:第一個在軌跡是等速直線時取等號;第二個在給定 \(\mathbf x_t\) 時端點方向已經確定(沒有交叉需要平均)時取等號。

遞迴地套用這個結果,就知道 Reflow 過程不會增加運送成本。

II. Reflow 拉直路徑

遞迴地套用 Reflow 會促使軌跡越來越直,雖然直度的量度不保證每一輪都單調改善。為了量化,對一條路徑 \(\mathbf Y=\{\mathbf y_t\}_{t\in[0,1]}\) 定義直度泛函

$$\mathcal S(\mathbf Y):=\int_0^1\mathbb E\Big[\Big\Vert\mathbf y_1-\mathbf y_0-\frac{d\mathbf y_t}{dt}\Big\Vert_2^2\Big]\,dt.$$

若 \(\mathcal S(\mathbf Y)=0\),\(\mathbf Y\) 恰好是直線(而且以等速走)。

書中的命題:Reflow 拉直隨機路徑 對 rectified 路徑 \(\mathbf Z^{(k)}\), $$\min_{k\in\{0,\dots,K\}}\mathcal S\big(\mathbf Z^{(k)}\big)\le\frac{\mathbb E\big[\Vert\mathbf x_1-\mathbf x_0\Vert^2\big]}{K}.$$ 書:見 Liu et al. (2022) 的 Theorem 3.7。
證明的骨架(轉述者整理的伸縮和論證)

記第 \(k\) 輪的配對為 \((\mathbf x_0^{(k)},\mathbf x_1^{(k)})\)(\(k=0\) 是獨立配對),用它訓練出的理想流的軌跡為 \(\mathbf Z^{(k+1)}\),它的終點就是下一輪的配對。

(a) 直度 = 動能 − 位移平方。因為 \(\int_0^1\frac{d\mathbf z_t}{dt}dt=\mathbf z_1-\mathbf z_0\),展開平方:

$$\mathcal S(\mathbf Z)=\mathbb E\Vert\mathbf z_1-\mathbf z_0\Vert^2-2\,\mathbb E\Big\langle\mathbf z_1-\mathbf z_0,\int_0^1\tfrac{d\mathbf z_t}{dt}dt\Big\rangle+\int_0^1\mathbb E\Big\Vert\tfrac{d\mathbf z_t}{dt}\Big\Vert^2dt=\int_0^1\mathbb E\Vert\mathbf v(\mathbf z_t,t)\Vert^2dt-\mathbb E\Vert\mathbf z_1-\mathbf z_0\Vert^2.$$

(b) 動能不超過上一輪的位移平方。和上一個證明一樣:邊際保持 \(+\) 條件 Jensen(\(\Vert\mathbb E[\mathbf u\mid\mathbf x]\Vert^2\le\mathbb E[\Vert\mathbf u\Vert^2\mid\mathbf x]\))給出 \(\int_0^1\mathbb E\Vert\mathbf v^{(k+1)}\Vert^2dt\le\mathbb E\Vert\mathbf x_1^{(k)}-\mathbf x_0^{(k)}\Vert^2\)。

(c) 伸縮和。合起來,\(\mathcal S(\mathbf Z^{(k+1)})\le\mathbb E\Vert\mathbf x_1^{(k)}-\mathbf x_0^{(k)}\Vert^2-\mathbb E\Vert\mathbf x_1^{(k+1)}-\mathbf x_0^{(k+1)}\Vert^2\)。對 \(k=0,\dots,K-1\) 加起來,右邊逐項相消:

$$\sum_{k=1}^{K}\mathcal S\big(\mathbf Z^{(k)}\big)\le\mathbb E\Vert\mathbf x_1-\mathbf x_0\Vert^2-\mathbb E\Vert\mathbf x_1^{(K)}-\mathbf x_0^{(K)}\Vert^2\le\mathbb E\Vert\mathbf x_1-\mathbf x_0\Vert^2.$$

\(K\) 個非負數的和不超過 \(\mathbb E\Vert\mathbf x_1-\mathbf x_0\Vert^2\),最小的那個就不超過它的 \(1/K\)。這是對 \(k=1,\dots,K\) 的 rectified 流取最小值,比書上對 \(k\in\{0,\dots,K\}\) 取最小值的寫法更強,所以蘊含書的命題。它也說明了為什麼「不保證每一輪都單調改善」:上界只約束總和,不約束個別項的順序。

書中的圖:Reflow 的示意
書中的圖(§5.4.3):Reflow 的示意。三格分別是 (a) 1-rectified flow、(b) 2-rectified flow、(c) 3-rectified flow。中間一圈紅色的點是來源(八個小團排成一圈),外面一圈紫色的點是目標(八個團排成較大的一圈);軌跡依方向交替著藍色和綠色。(a) 的軌跡在中心附近彎曲、互相擠壓,有些繞到別的方向;(b) 和 (c) 幾乎都是從內圈直直往外的直線束。書的圖說:路徑隨著 Rectify 程序越來越直。圖片出處:改編自 Liu et al. (2022)。

線性插值核的 FM 或 RF 寫法,加上 Reflow 程序,提供了更簡單的訓練目標,也提供了改良隨機配對的實用方法。理論細節見 Liu et al. (2022) 的兩篇論文。

III. 和最佳運送的關聯

最後要指出:直線配對在最佳運送(OT)的意義下不一定是最佳的。這牽涉第 7 章才介紹的術語(B18),不熟 OT 的讀者可以先看那一節。

二次成本的最佳運送有一個標誌性的特徵:粒子沿直線走。在 \(\mathbf x_0\) 的粒子經由 \(\mathbf x_t=(1-t)\mathbf x_0+t\,\mathbf T(\mathbf x_0)\) 移到 \(\mathbf T(\mathbf x_0)\),\(\mathbf T\) 是最佳運送映射。但並不是每個產生這種直線路徑的映射 \(\mathbf S\)(\(\mathbf x_t=(1-t)\mathbf x_0+t\,\mathbf S(\mathbf x_0)\))都是最佳的:只有當 \(\mathbf S\) 最小化 Monge 成本 \(\mathbb E\Vert\mathbf x_0-\mathbf S(\mathbf x_0)\Vert^2\) 時,它才給出最佳的流。所以直線路徑是必要條件,但不是充分條件;最佳性還取決於正確的端點映射 \(\mathbf T\)。

書中的例子:直線配對不一定是最佳 令 \(p_{\text{src}}=p_{\text{tgt}}=\mathcal N(\mathbf 0,\mathbf I)\)。對成本 \(c(\mathbf x,\mathbf y)=\Vert\mathbf x-\mathbf y\Vert^p\)(\(p>0\)),\(c\)-最佳配對是恆等配對 \(\pi^*\),即 \((\mathbf x,\mathbf x)\) 的分佈,\(\mathbf x\sim p_{\text{src}}\)。

現在考慮配對 \(\pi_{\mathbf A}\):\((\mathbf x,\mathbf A\mathbf x)\) 的分佈,\(\mathbf x\sim p_{\text{src}}\),\(\mathbf A\) 是旋轉矩陣,滿足 \(\mathbf A^\top\mathbf A=\mathbf I\)、\(\det\mathbf A=1\)、\(\mathbf A\ne\mathbf I\),而且 \(-1\) 不是它的特徵值。那麼 \(\pi_{\mathbf A}\) 是 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\) 之間合法的配對,對應 \(\mathbf x\) 和 \(\mathbf A\mathbf x\) 之間的直線路徑,但對任何 Hessian 可逆、二次可微的嚴格凸成本 \(c\),它都不是 \(c\)-最佳的。次佳來自旋轉這個變換。如第 7 章(B20)會討論的,即使去掉旋轉,也不一定得到最佳配對。

例子裡的三個條件各在做什麼

為什麼是合法配對。\(\mathcal N(\mathbf 0,\mathbf I)\) 在正交變換下不變:\(\mathbf A\mathbf x\sim\mathcal N(\mathbf 0,\mathbf A\mathbf A^\top)=\mathcal N(\mathbf 0,\mathbf I)\),所以兩個邊際都對。

為什麼不是最佳。以二次成本為例,\(\mathbb E\Vert\mathbf x-\mathbf A\mathbf x\Vert^2=\mathbb E[\mathbf x^\top(2\mathbf I-\mathbf A-\mathbf A^\top)\mathbf x]=2D-2\operatorname{tr}\mathbf A>0\)(\(\mathbf A\ne\mathbf I\) 時 \(\operatorname{tr}\mathbf A<D\)),恆等配對的成本卻是 0。二維、轉角 \(\theta\) 時是 \(4(1-\cos\theta)\)。

為什麼要求 \(-1\) 不是特徵值(轉述者的解讀)。兩條直線 \((1-t)\mathbf x+t\mathbf A\mathbf x\) 和 \((1-t)\mathbf y+t\mathbf A\mathbf y\) 在某個 \(t\in(0,1)\) 相交,等價於 \(\big((1-t)\mathbf I+t\mathbf A\big)(\mathbf x-\mathbf y)=\mathbf 0\) 有非零解,也就是 \(\mathbf A\) 有特徵值 \(-(1-t)/t\),一個負實數。旋轉矩陣的特徵值都在單位圓上,唯一可能的負實數是 \(-1\)(\(t=\tfrac12\))。所以這個條件保證:任兩條直線路徑在同一時刻不會在同一點(畫在平面上的線段仍可能交叉,只是兩個粒子經過交點的時間不同)。這樣一來,給定 \((\mathbf x_t,t)\) 就唯一決定了它在哪條直線上,Reflow 的條件期望不需要平均任何東西,學到的流恰好沿著這些直線走,Rectify 會把 \(\pi_{\mathbf A}\) 原封不動地送回來。也就是說,\(\pi_{\mathbf A}\) 是 Reflow 的不動點:Reflow 停在這裡,不會再往最佳配對前進。Demo B 可以看到 \(\theta=180^\circ\) 時所有直線都在 \(t=\tfrac12\) 穿過原點。

和 OT 的關聯在第 7 章 §7.5.2(B20)會繼續探討。

Demo B直,但不是最佳:旋轉配對轉述者補充

二維 \(p_{\text{src}}=p_{\text{tgt}}=\mathcal N(\mathbf 0,\mathbf I)\),配對 \((\mathbf x,\mathbf A_\theta\mathbf x)\)。畫出 60 條直線路徑(從點到點,深淺表示時間),以及 \(t\) 時刻的位置。拖動 \(\theta\):任何 \(0<\theta<180^\circ\) 都是合法的直線配對,而且任兩個粒子不會在同一時刻撞在一起(藍點永遠分開),卻不是最佳的;\(\theta=180^\circ\) 時所有粒子在 \(t=\tfrac12\) 同時經過原點。

直線路徑起點 x時間 t 的位置

預設 \(\theta=60^\circ\):2,000 個樣本的二次成本是 1.992,公式 \(4(1-\cos60^\circ)=2\);恆等配對(什麼都不動)的成本是 0,所以這個直線配對遠非最佳。終點的樣本變異數是 (0.99, 1.01),確認 \(\mathbf A\mathbf x\) 仍服從 \(\mathcal N(\mathbf 0,\mathbf I)\)。另一個值得注意的地方:中途的分佈是 \(\mathcal N\big(\mathbf 0,(1-2t(1-t)(1-\cos\theta))\mathbf I\big)\),\(\theta=60^\circ\)、\(t=\tfrac12\) 時變異數只剩 0.75,因為弦比弧短,直線路徑會切過圓的內部。把 \(\theta\) 拉到 \(180^\circ\)、\(t=0.5\):所有藍點擠在原點,變異數是 0,這正是書要求「\(-1\) 不是特徵值」所排除的情形。

05本章結語書 §5.5

這一章闡明了 diffusion model 第三個、也是最後一個基礎觀點:根植於確定性流的原理。探索從 normalizing flow(NF)開始:它利用變數變換公式,學一個簡單先驗和資料分佈之間精確、可逆的映射(B11)。接著這個概念演進成連續時間的 neural ODE:一個學到的速度場支配整個轉換。但這個做法有很大的缺點:訓練迴圈裡需要昂貴的 ODE 模擬。

現代的 Flow Matching(FM)框架是這個挑戰優雅而有效率的解答(B12–B13)。它預先定義一條機率路徑 \(\{p_t\}_t\) 和一個滿足連續方程的對應速度場,為 ODE 流訂下明確的目標。關鍵的是,就像在變分和 score-based 觀點裡看到的,FM 使用了強大的條件化技巧,把「匹配算不出來的邊際速度場」這個問題,轉成對一個已知條件速度的簡單、好處理的回歸,讓訓練完全不需要模擬。這個觀點把 diffusion model 本身重新詮釋成:學一個把高斯先驗運送到資料分佈的確定性流的特例。

加上 flow 的觀點,diffusion 建模的三根概念支柱就全部介紹完了。一路走來浮現一個值得注意的模式:每個框架雖然各自起源於 VAE、EBM 或 NF,都收斂到一個連續時間的生成過程,並且都依賴條件化策略讓學習變得可行。

下一章會把這些平行的線索綜合成一個統一的框架:

  1. 正式證明變分、score-based 和 flow-based 觀點不只是類比,而是在根本上數學等價。
  2. 說明 Fokker–Planck 方程是支配三個觀點裡密度演化的共同規律,揭示它們只是描述同一個核心生成原理的不同鏡頭。

這個統一的視角會提供對現代 diffusion 典範完整而系統的理解。

轉述者整理:三個觀點在本系列的位置。
觀點起源連續時間的形式條件化技巧本系列
變分VAEDDPM 的前向/反向馬可夫鏈以 \(\mathbf x_0\) 為條件的反向核B03–B04
Score-basedEBMScore SDE、PF-ODEDSM:條件 scoreB05–B10
Flow-basedNFNeural ODE、Flow MatchingCFM:條件速度B11–B14

與 CS236 的對照

本頁CS236 複習頁
Rectified flow、少步數生成、蒸餾第 17 講
Normalizing flow、neural ODE(本章開頭的回顧)第 7 講、第 8 講
Euler 法的步數與誤差補充課 B
Jensen 不等式補充課 A

常見誤解

誤解:「Reflow 改變了生成的分佈」。理想情況下每一輪都保持兩端的邊際(第 2 節);改變的只是哪個來源點配到哪個終點。實務上模型誤差會一輪輪累積,所以 Reflow 的輪數通常很少。
誤解:「Reflow 的不動點就是最佳運送」。直線、而且粒子不會在同一時刻相撞的配對都是不動點,但旋轉配對就不是最佳(第 4 節、Demo B)。和 OT 的完整關係在第 7 章(B20)。
誤解:「直度上界保證每一輪都更直」。命題只約束 \(\mathcal S\) 的總和(或最小值);個別一輪不保證單調改善(第 4 節)。
誤解:「Rectify 需要重新設計路徑」。路徑仍是同一條線性插值;換掉的是端點的配對 \(\pi\)(第 2 節)。

自我測驗

1. 獨立配對下,為什麼邊際軌跡會彎?Rectify 怎麼處理?
同一個 \(\mathbf x_t\) 可能來自許多方向不同的端點對,邊際速度只能取平均,而這個平均隨 \((\mathbf x,t)\) 改變。Rectify 用學到的 ODE 把每個 \(\mathbf z_0\) 送到 \(\hat{\mathbf z}_1\),用這個確定性的配對重新訓練。
2. 為什麼 \(\pi_{\text{Rectify}}\) 仍然是 \(p_{\text{src}}\) 和 \(p_{\text{tgt}}\) 之間的合法配對?
第一個邊際是 \(p_{\text{src}}\)(依定義);理想漂移生成參考路徑的每個邊際,所以 \((\bm\Phi_{0\to1})_\#p_{\text{src}}=p_{\text{tgt}}\)。
3. 用兩次 Jensen 證明 \(\mathbb E\,c(\hat{\mathbf z}_1-\mathbf z_0)\le\mathbb E\,c(\mathbf x_1-\mathbf x_0)\)。
\(c(\int\mathbf v^*dt)\le\int c(\mathbf v^*)dt\);邊際保持讓 \(\mathbb E\,c(\mathbf v^*(\mathbf z(t),t))=\mathbb E\,c(\mathbf v^*(\mathbf x_t,t))\);條件 Jensen 給 \(c(\mathbb E[\mathbf x_1-\mathbf x_0\mid\mathbf x_t])\le\mathbb E[c(\mathbf x_1-\mathbf x_0)\mid\mathbf x_t]\)。
4. 直度泛函 \(\mathcal S\) 為什麼等於「動能 − 位移平方」?
展開 \(\Vert\mathbf z_1-\mathbf z_0-\dot{\mathbf z}_t\Vert^2\) 並用 \(\int_0^1\dot{\mathbf z}_tdt=\mathbf z_1-\mathbf z_0\):交叉項變成 \(-2\mathbb E\Vert\mathbf z_1-\mathbf z_0\Vert^2\)。
5. 在 \(\mathcal N(\mathbf 0,\mathbf I)\) 到自己的旋轉配對裡,二次成本是多少?為什麼它是 Reflow 的不動點?
\(2D-2\operatorname{tr}\mathbf A\),二維轉角 \(\theta\) 時是 \(4(1-\cos\theta)\)。\(-1\) 不是特徵值時,任兩條直線路徑不會在同一時刻經過同一點,給定 \((\mathbf x_t,t)\) 就知道在哪條線上,學到的流沿著這些直線走,Rectify 送回同一個配對。
6. 第 5 章結語說三個觀點共有哪兩個特徵?
都收斂到一個連續時間的生成過程;都靠條件化策略讓學習變得可行。