接續上一篇: 讀 Flow Matching 前要先理解的東西 (建議先閱讀)
Flow matching 模型在時間 $t=0$ 的時候從常態分佈出發 $p_0(x)=\mathcal{N}(0,I)$, 隨著時間變化其 pdf, 例如時間 $t$ 時的 pdf 變化成為 $p_t(x)$, 直到時間 $t=1$ 時希望變成接近目標分佈 $q(x)$, 即希望 $p_1(x)\approx q(x)$.
概念是頭尾 pdf 確定後, 中間這無限多種可能的 $p_t(x)$ 變化經過作者的巧妙設定, 讓學 vector field $u_t(x)$ 變的可能! (不學 pdf 而是學 vector field)
結果就是用 NN 學到的 $u_t(x)$ 可以讓 pdf 從開頭的常態分佈一路變化到最後的資料目標分佈 $q(x)$.
Vector Field 決定了 Probability Path
注意到這一路變化的過程要滿足 continuity equation (質量守恆)
$$\begin{align}
\frac{\partial p_t(x)}{\partial t}
+
\nabla\cdot(p_t(x) u_t(x))=0
\end{align}$$ 其中 $u_t(x)$ 表示時間 $t$ 時在位置 $x$ 的 vector field.
從 continuity equation 出發, 可推導出 Instantaneous Change of Variables Theorem [1]:
$$\frac{d\log p_t(x)}{dt}=-\text{div}(u_t(x))$$ 因此得到 log-likelihood, 就是積分起來的結果: $$\log p_{\color{orange}{t}}(x)=\log p_0(z)+\int_{\hat{t}=0}^{\color{orange}{t}}-\text{div}(u_{\hat{t}}(x))d\hat{t}$$ 因此可以發現, 如果有每個時間 $t$ 的 vector field $u_t(x)$, 透過上式就可以得到對應的 pdf $p_t(x)$.
[觀念]: Vector field $u_t(x)$ 決定了 probability path $p_t(x)$.
Flow Matching Loss
既然 $u_t(x)$ 決定了 $p_t(x)$, 讓 NN 直接學 $u_t(x)$ 就好了, 因此目標函式:
$$\mathcal{L}_{FM}(\theta)=\mathbb{E}_{t,x\sim {\color{orange}{p_t(x)}}}\|v_t(x)-{\color{orange}{u_t(x)}}\|^2, \qquad (\star)$$ 其中 $v_t(x)$ 是我們的 NN, 具有參數 $\theta$.
$t\sim U[0,1]$ 表示對任意時間 $t$ 都要去學 vector field.
注意到我們根本不知道真實的 ${\color{orange}{u_t(x)}}$ 和 ${\color{orange}{p_t(x)}}$, 所以 $\mathcal{L}_{FM}$ 該怎麼計算?
High level picture 的想法是, 我們定義出頭尾分佈, 頭的分佈就用標準常態分佈, 尾的分佈從 training data 設計出來. 中間變化的分佈論文經過特殊設計.
設計好後, 就能找出 ${\color{orange}{u_t(x)}}$ 和 ${\color{orange}{p_t(x)}}$ 了, 找出來後應該就有機會算 $\mathcal{L}_{FM}$ 了.
以下詳細說明.
頭尾的分佈長相
我們先來想像一下怎麼從 $\mathcal{N}(0,I)$ 一路轉換成目標分佈 $q(x)$.
假設我們有 $N$ 筆 training data $x_1=\{x_1^1,x_1^2,...,x_1^N\}$.
對第 $n$ 筆資料 $x_1^n$, 設定 $t=0$ 時條件機率分佈:
$$\begin{align}
p_0(x|x_1^n)=p_0(x)=\mathcal{N}(0,I)
\end{align}$$ 也就是初始分佈與資料無關, 全部都是標準常態分佈
接著設定 $t=1$ 時條件機率分佈:
$$\begin{align}
p_1(x|x_1^n)=\mathcal{N}\left(x_1^n,\sigma_{min}^2I\right)
\end{align}$$ 並假設每一筆訓練資料機率相同, $q(x_1^n)=1/N,\forall n=1,2,...,N$. 則
$$p_1(x)=\sum_{n=1}^N p_1(x|x_1^n)q(x_1^n)\\
=\frac{1}{N}\sum_{n=1}^N p_1(x|x_1^n)=\frac{1}{N}\sum_{n=1}^N \mathcal{N}\left(x_1^n,\sigma_{min}^2I\right)$$ 這是個 Gaussin Mixture Model (GMM) 分佈
理論上 $N\rightarrow\infty,\sigma_{min}\rightarrow0$ 時 $p_1(x)$ 就會等於目標分佈 $q(x)$. 或我們寫成連續積分形式
$$\begin{align}
p_1(x)=\int p_1(x|x_1)q(x_1)dx_1
\end{align}$$ 其中 $p_1(x|x_1)=\mathcal{N}(x_1,\sigma_{min}^2I)$, $q(x_1)$ 是目標分佈, 也就是從手上的 training data 採樣就可以.
事實上 $q(x_1)$ 可以設定的很彈性, 不一定要從 training data 採樣. 更廣泛來說下式數學上本就成立:
$$p_1(x)=\int p_1(x|{\color{orange}{z}})q({\color{orange}{z}})d{\color{orange}{z}}$$ $z$ 可以是很彈性的變量, 而 $q(z)$ 是它的 distribution.
現在我們先想成 $z$ 就是 training data $x_1$ 即可. $q(z)$ 就是目標分佈 $q(x_1)$, 即手上的 data 分佈
後面會說明 $z$ 和 $q(z)$ 的其他選擇.
任何時間 $t$ 都可以這樣寫: $$\begin{align}
p_{\color{orange}{t}}(x)=\int p_{\color{orange}{t}}(x|x_1)q(x_1)dx_1
\end{align}$$ 因此我們設定好了式 (2) 的頭 $p_0(x)$ 和式 (4) 的尾 $p_1(x)$ 的分佈
中間分佈的設定
由式 (2) 和 (4) 我們設定好了頭尾分佈, 那中間的 pdf $p_t(x)$ 怎麼設定?
或是因為 $u_t(x)$ 決定了 $p_t(x)$, 我們改問 $u_t(x)$ 怎麼設定?
可以想像有無窮多種 $u_t(x)$ 的設定方法 (因為只要滿足頭尾 pdf 即可)
但還是要有個變化規則, 這樣 NN 才知道要學什麼.
$u_t(x)$ 的規則是 (注意到這是人為定義的):
$$\begin{align}
u_t(x) \triangleq \int u_t(x|x_1)\frac{p_t(x|x_1)q(x_1)}{p_t(x)}dx_1
\end{align}$$ 怕大家看到這裡忘記: $x_1$ 是訓練資料 $\{x_1^1,x_1^2,...,x_1^N\}$, 然後 $q(x_1)$ 表示目標分佈 (資料分佈).
其中特別說明 $u_t(x|x_1)$ 與 $p_t(x|x_1)$ 是一對的, 也就是說 $u_t(x|x_1)$ 產生 $p_t(x|x_1)$ 因此需滿足質量守恆 continuity equation (1). 這裡我們先假設已經有這樣一對的東西.
因為 $u_t(x)$ 是人為定義, 我們需要驗證它是合法的 vector field, i.e. 能產生 $p_t(x)$ (5), 即要滿足 continuity equation (1). 證明放在 Appendix.
因此得到 $u_t(x)$ 和 $p_t(x)$ 是一對的, 式 (6) 定義的 $u_t(x)$ 是 well-defined.
接著說明 $u_t(x|x_1)$ 和 $p_t(x|x_1)$ 的長相, 其中 $p_t(x|x_1)$ 論文這樣定義:
$$\begin{align}
p_t(x|x_1)\triangleq\mathcal{N}\left(
\mu_t(x_1),\sigma_t(x_1)^2I
\right)
\end{align}$$ Mean 和 covariance matrix 都是 $x_1$ 的函數.
由於要滿足頭尾 pdf 的設定, i.e. 式 (2) $p_0(x|x_1)=\mathcal{N}(0,I)$ 和式 (3) $p_1(x|x_1)=\mathcal{N}(x_1,\sigma_{min}^2I)$.
所以 $\mu_t,\sigma_t$ 需滿足:
$$\mu_0(x_1)=0,\quad \sigma_0(x_1)=1 \\
\mu_1(x_1)=x_1,\quad \sigma_1(x_1)=\sigma_{min}$$ 這裡有個問題, 什麼長相的 $u_t(x|x_1)$ 會產生式 (7) $p_t(x|x_1)$ 的定義?
論文的 Theorem 3 證明了存在唯一 $u_t(x|x_1)$ 產生式 (7) $p_t(x|x_1)$:
$$\begin{align}
u_t(x|x_1)=\frac{\sigma_t'(x_1)}{\sigma_t(x_1)}(x-\mu_t(x_1))+\mu_t'(x_1)
\end{align}$$ 證明放在 Appendix
這結果實在太方便了, 基本上只要把 $\mu_t$ 和 $\sigma_t$ 定義好, $u_t(x|x_1)$ 有 closed form solution!
講到現在可能有點亂了, 整理一下目前的故事線:
定義了條件機率 $p_t(x|x_1)$ 的長相後 (7), 我們可以得到存在唯一的條件向量場 $u_t(x|x_1)$ 長相 (8)
有了 $p_t(x|x_1)$ 和 $u_t(x|x_1)$, 我們就能得出向量場 ${\color{orange}{u_t(x)}}$ (6) 和 ${\color{orange}{p_t(x)}}$ (5).
Flow matching loss $\mathcal{L}_{FM}(\theta)$ 就能計算了!
$$\begin{align}
\mathcal{L}_{FM}(\theta)=\mathbb{E}_{t,x\sim {\color{orange}{p_t(x)}}}\|v_t(x)-{\color{orange}{u_t(x)}}\|^2
\end{align}$$ 真的能算… 吧? 是吧?
好像還是不行, (5) 和 (6) 這樣的積分形式很難算.
Conditional Flow Matching Loss
論文的再一個重要發現為把 $\mathcal{L}_{FM}$ 轉為等價且實際可計算的 loss, a.k.a. Cnditional Flow Matching $\mathcal{L}_{CFM}$:
$$\begin{align}
\mathcal{L}_{CFM}(\theta)=\mathbb{E}_{t,q(x_1),{\color{orange}{p_t(x|x_1)}}}\|v_t(x)-{\color{orange}{u_t(x|x_1)}}\|^2
\end{align}$$ 由於 ${\color{orange}{p_t(x|x_1)}}$ 和 ${\color{orange}{u_t(x|x_1)}}$ 很容易計算 (式 (7), (8)), 因此 $\mathcal{L}_{CFM}$ 很容易算
論文證明 $\nabla_\theta\mathcal{L}_{CFM}(\theta)=\nabla_\theta\mathcal{L}_{FM}(\theta)$! 證明放在 Appendix.
Gradient 一樣, 因此找出來的最佳解 $\theta$ 也一樣.
有意思的是 $\mathcal{L}_{CFM}$ 用的 ground truth 是 conditional vector field $u_t(x|x_1)$, 但是這樣學出來的 $v_t(x)$ 卻是等價於去逼近原本的 (unconditional) vector field $u_t(x)$! 重點是 $\mathcal{L}_{CFM}$ 又容易計算.
Conditional Probability 和 Conditional Vector Field 的選擇
條件分佈 $p_t(x|x_1)$ 如下設定後會有唯一的條件向量場 $u_t(x|x_1)$:
$$p_t(x|x_1)\triangleq\mathcal{N}\left(
\mu_t(x_1),\sigma_t(x_1)^2I
\right) \\
u_t(x|x_1)=\frac{\sigma_t'(x_1)}{\sigma_t(x_1)}(x-\mu_t(x_1))+\mu_t'(x_1)$$ $\sigma_t’$ 表示 $\sigma_t$ 對 $t$ 偏微分, $\mu_t’$ 表示 $\mu$ 對 $t$ 偏微分, 其中
$$\mu_0(x_1)=0,\quad \sigma_0(x_1)=1 \\
\mu_1(x_1)=x_1,\quad \sigma_1(x_1)=\sigma_{min}$$ 這樣的設定能讓我們方便計算 $\mathcal{L}_{CFM}$ (等價於 $\mathcal{L}_{FM}$).
同時保證 NN 學出來的 $v_t(x)$ 其對應的 $p_t(x)$ 能滿足頭是標準常態分佈, 尾是目標分佈.
因此我們可以設計不同的 $\mu_t$ 和 $\sigma_t$, 就可以有不同的中間 pdf $p_t$ 演變方式.
Conditional flows from FM (Lipman et al.)
選擇 conditional probability path $p_t(x|x_1)$ 的 $\mu_t,\sigma_t$ 為:
$$\mu_t(x)=tx_1 \\
\sigma_t(x)=1-(1-\sigma_{min})t$$ 套 (8) 得到具有唯一的 conditional vector field $u_t(x|x_1)$:
$$u_t(x|x_1)=\frac{x_1-(1-\sigma_{min})x}{1-(1-\sigma_{min})t}$$ 讀者可自行對 $\mu_t$ 和 $\sigma_t$ 對 $t$ 偏微分驗證.
I-CFM (Independent Coupling)
選擇 conditional probability path $p_t(x|z)$ 的 $\mu_t,\sigma_t$ 為:
$$\mu_t(x)=tx_1+(1-t)x_0 \\
\sigma_t(x)=\sigma_{min}$$ 其中 $z\triangleq(x_0,x_1)$ 且 $x_0$ 從開頭分佈 $p_0$ 去採樣, $x_1$ 從目標分佈 $q$ 去採樣, 且 $x_0$ 和 $x_1$ 獨立採樣:
$$q(z)\triangleq p_0(x_0)p_1(x_1)=\mathcal{N}(0,I)q(x_1)$$ 套 (8) 導致具有唯一的 conditional vector field $u_t(x|x_1)$:
$$u_t(x|x_1)=x_1-x_0$$ 詳細可參考論文 [2]
OT-CFM (Optimal transport)
選擇 conditional probability path $p_t(x|z)$ 的 $\mu_t,\sigma_t$ 為:
$$\mu_t(x)=tx_1+(1-t)x_0 \\
\sigma_t(x)=\sigma_{min}$$ $z\triangleq(x_0,x_1)$ 的 $x_0$ 和 $x_1$ 不再是獨立採樣, 而是根據 2-Wasserstein distance, 即最小搬運 cost, 詳細參考論文 [2], 論文裡給出這三種設定差別:
Left: Conditional flows from FM (Lipman et al., 2023), I-CFM (§3.2.2), and OT-CFM (§3.2.3).
Right: Learned flows (green) from moons (blue) to 8gaussians (black) using I-CFM (centre-right) and OT-CFM (far right).
圖示很精闢
Inference and Sampling
我們花了大把的篇幅來說明怎麼學 vector field $u_t(x)$.
如同 “讀 Flow Matching 前要先理解的東西” 說的 vector field 我們想成是速度的話, 理論上時間 $t$ 時 $p_t(x)$ 的任一點 $x_t$ 都能夠得出演變的路徑.
$$\frac{d}{dt}x_t=u_t(x_t)$$ 這個路徑我們叫 flow.
因此採樣就是從 $p_0(x)=\mathcal{N}(0,I)$ 出發得到一個點 $x_0$, 離散的估計 $\Delta t$ 時間後的位置為:
$$x_{\Delta t}=x_0+\frac{d}{dt}x_0\cdot \Delta t \\
=x_0+u_0(x_0)\cdot \Delta t$$ 所以時間 $t$ 時, $t+\Delta t$ 的位置為:
$$x_{t+\Delta t}=x_t+\frac{d}{dt}x_t\cdot \Delta t \\
=x_t+u_t(x_t)\cdot \Delta t$$ 當 $\Delta t$ 每次都走很小步的時候就會比較準, 但要花很多次 iteration.
這種 Naive 的方法為 Euler method.
實際上有更多好的選擇, 請參考 [4] Numerical Methods for Ordinary Differential Equations
Toy Example Codes
這篇文章 “Flow Matching: Matching flows instead of scores” [3] 有給出 Conditional flows from FM (Lipman) 和 I-CFM 的 sample codes. 完整 source code 在 [6].
算是一個很清楚的展示, 容易擴展.
Flow matching 理論雖然複雜精美, 實作上卻十分簡潔! 精彩!
Appendix
證明直接擷取自論文 [5], 為了完整性做紀錄而已
Vector field 決定了 probability path
證明 $u_t(x)$ 能產生 $p_t(x)$
證明存在唯一 $u_t(x|x_1)$ 產生式 (7) $p_t(x|x_1)$
先說明一個觀念, vector field $u_t(x)$ 我們可以想成速度, 即位置對時間的微分, 有了速度就可以知道下一個時間點的位置
因此隨著時間變化的位置稱為 flow $\phi_t(x)$, 則 $d\phi_t(x)/dt=u_t(\phi_t(x))$. 下面證明的 $\psi_t$ 也是 flow 只是用的是條件 vector field $u_t(x|x_1)$.
Conditional FM 與 FM Losses 等價
證明 $\nabla_\theta\mathcal{L}_{CFM}(\theta)=\nabla_\theta\mathcal{L}_{FM}(\theta)$.
Reference
- Neural Ordinary Differential Equations, Appendix A: Proof of the Instantaneous Change of Variables Theorem
- Improving and generalizing flow-based generative models with minibatch optimal transport [arxiv]
- Flow Matching: Matching flows instead of scores
- Numerical Methods for Ordinary Differential Equations
- Flow Matching for Generative Modeling [arxiv]
- Toy Example Codes from [here]