Flow Matching/zh
| Article | |
|---|---|
| Topic area | generative-models |
| Prerequisites | Diffusion Models, Optimal Transport, Neural Ordinary Differential Equations |
概述
Flow Matching 是一種用於 Continuous Normalizing Flows 的免模擬訓練框架,其中神經網絡直接回歸到一個目標的時間相關向量場,該場將一個簡單的先驗分布傳輸到數據分布。該方法由 Lipman、Chen、Ben-Hamu、Nickel 和 Le 於 2022 年提出,推廣並統一了若干早期方法,包括 Score Matching 與 Rectified Flow,並已成為圖像、視頻、音頻和分子生成建模的主流範式。與流模型的經典極大似然訓練相比,後者需要昂貴的軌跡模擬;與 Denoising Diffusion Probabilistic Models 相比,後者需要隨機微分方程的形式化,Flow Matching 提供了一個確定性且概念簡潔的替代方案:選取一條連接噪聲與數據的概率路徑,導出生成該路徑的向量場,並以均方誤差損失學習該向量場。
直覺
一個連續的歸一化流描述了分布空間中的一條曲線,通過常微分方程 (ODE) 將樣本從時間 $ t=0 $ 處的初始概率密度 $ p_0 $ 傳輸到時間 $ t=1 $ 處的目標密度 $ p_1 $。對於每個時刻 $ t \in [0,1] $,一個 Vector Field $ u_t(x) $ 給出了位置 $ x $ 處的瞬時速度;沿時間積分該向量場,便將樣本從先驗推送到數據分布。
核心挑戰在於,一般而言我們無法直接觀測到一個能將可處理的先驗傳輸到經驗數據分布的向量場。Flow Matching 繞開了這一難題,通過對單個數據點進行條件化、分段地構造該路徑。對於固定的數據樣本 $ x_1 $,可以容易地寫出從一個噪聲樣本到 $ x_1 $ 的光滑路徑,並讀出生成它的速度。在噪聲與數據的聯合採樣下對這些逐樣本的速度取平均,便得到將整個總體從先驗驅動到數據的無條件速度場。關鍵洞見是:將神經網絡回歸到條件速度,在期望意義下即可恢復無條件速度,從而完全無需計算邊緣密度。
概率路徑與向量場
一條概率路徑是按時間索引的密度族 $ \{p_t\}_{t \in [0,1]} $,其中 $ p_0 $ 是所選定的先驗(通常是一個標準的 Gaussian Distribution),$ p_1 $ 是數據分布。當 Continuity Equation 成立時,向量場 $ u_t $ 生成該路徑:
$ {\displaystyle \frac{\partial p_t(x)}{\partial t} + \nabla \cdot (p_t(x)\, u_t(x)) = 0.} $
等價地,從 $ p_0 $ 抽取並由 ODE $ dx/dt = u_t(x) $ 演化的樣本,在每個中間時刻都服從 $ p_t $。多種向量場可以生成同一條路徑,因此需要額外的結構(例如直線性,或關於某種傳輸代價的最優性)來挑出一個優選的向量場。
條件 Flow Matching
直接回歸 $ u_t $ 不可行,因為 $ u_t $ 依賴於未知的邊緣密度。條件 Flow Matching (CFM) 的目標通過對目標樣本 $ x_1 $ 進行條件化來解決此問題。對於選定的條件路徑 $ p_t(x \mid x_1) $(例如,一個均值在 $ t=0 $ 時為 $ 0 $、在 $ t=1 $ 時線性插值到 $ x_1 $ 的高斯分布)以及生成它的條件向量場 $ u_t(x \mid x_1) $,損失為
$ {\displaystyle \mathcal{L}_{\mathrm{CFM}}(\theta) = \mathbb{E}_{t,\, x_1,\, x \sim p_t(\cdot \mid x_1)}\!\left[\, \lVert v_\theta(t, x) - u_t(x \mid x_1) \rVert^2 \right],} $
其中 $ v_\theta $ 是所學習的向量場,$ t $ 在 $ [0,1] $ 上均勻採樣,$ x_1 $ 從數據中採樣。Lipman 等人證明,該目標關於 $ \theta $ 的梯度與對邊緣場 $ u_t $ 的回歸相同,儘管該邊緣場不可處理。關鍵的設計選擇是條件路徑;常用選擇包括保方差的高斯路徑、爆方差路徑,以及 Optimal Transport 位移線性插值 $ x_t = (1-t)\, x_0 + t\, x_1 $,後者產生異常簡潔的回歸目標 $ u_t(x \mid x_0, x_1) = x_1 - x_0 $。
訓練與推斷
訓練僅需採樣一個時刻、一個噪聲向量和一個數據點;以閉式計算條件速度;並最小化平方誤差。訓練期間無需模擬 ODE,無需輔助分數網絡,也無需跟蹤變分下界。小批量由相互獨立的三元組 $ (t, x_0, x_1) $ 組成,其中 $ x_0 $ 從先驗中抽取,$ x_1 $ 從數據集中抽取。
在推斷時,通過在從先驗中抽取的初始條件下,將所學到的 ODE $ dx/dt = v_\theta(t, x) $ 從 $ t=0 $ 積分到 $ t=1 $ 來生成樣本。可以使用任何黑盒 ODE 求解器;常見選擇包括自適應 Runge-Kutta 方法以及定步長 Euler 或 Heun's Method 積分器。由於使用線性(最優傳輸)插值訓練的 Flow Matching 往往產生近乎直線的軌跡,生成樣本通常只需少量求解器步驟,這與可能需要數十至數百步的擴散模型形成對比。
變體
Flow Matching 的若干變體會調整條件路徑、$ x_0 $ 與 $ x_1 $ 之間的耦合,或訓練過程:
- Rectified Flow(Liu 等,2022)使用與 OT-CFM 相同的線性插值進行訓練,然後在自身被拉直的軌跡上迭代地重新訓練模型,產生越來越直的流,從而支持單步或少步採樣。
- Stochastic Interpolants(Albergo 與 Vanden-Eijnden,2023)將該框架推廣至允許隨機動力學,將基於流和基於擴散的生成建模統一在一個插值形式之下。
- Optimal Transport Conditional Flow Matching(Tong 等,2023)用基於小批量的 Optimal Transport 耦合替代 $ x_0 $ 與 $ x_1 $ 的獨立耦合,從而銳化噪聲與數據之間的對齊並降低路徑曲率。
- Multisample Flow Matching(Pooladian 等,2023)提出了相關的批量耦合視角,並對相應估計量給出了理論分析。
- Riemannian Flow Matching 將該構造擴展到流形上的數據,用測地線插值代替歐氏插值,並採用流形感知的 ODE 積分器。
- Discrete Flow Matching 用連續時間馬爾可夫鏈取代 ODE,將該框架適配於分類數據。
與擴散模型的比較
擴散模型與 Flow Matching 緊密相關:兩者都學習從噪聲到數據的時間相關變換,並且都可以表述為對目標場的回歸問題。它們的差別在於過程與參數化的選擇。擴散模型通過隨機的正向與逆向過程加以表述,並學習 Score Function $ \nabla \log p_t(x) $;其訓練對應於 Flow Matching 家族中一種特定的保方差高斯路徑。Flow Matching 在 ODE 層面上是純確定性的,將路徑視為自由的設計選擇,並參數化速度而非分數。經驗上,OT 風格的 Flow Matching 產生更直的軌跡並支持更快的採樣,而擴散的隨機性在某些情形下可以提升樣本多樣性。基於分數的擴散採樣器可以被重新解釋為概率流 ODE 的 ODE 積分器,從而精確揭示了這兩族之間的數學橋梁。
局限性
Flow Matching 繼承了 Continuous Normalizing Flows 的常見困難:在軌跡彎曲或剛性時,推斷階段的 ODE 積分可能代價高昂;精確的對數似然計算需要 Hutchinson Trace Estimator 或昂貴的雅可比求值;而在高維流形上則可能需要精心選擇先驗,以免浪費建模容量。該框架還假設存在一條可處理的條件路徑,這在歐氏空間中是直接的,但在流形、圖或離散空間上則更加微妙。條件化、Classifier-Free Guidance 與無似然評估可以從擴散遷移到 Flow Matching,但有時需要謹慎調整,因為其底層對象是一個向量場而非分數。
應用
Flow Matching 已被應用於高解析度圖像生成,包括將 OT-CFM 擴展到數十億參數的文本到圖像模型;語音與音頻合成,其中直線軌跡支持實時生成;在 $ \mathrm{SE}(3) $ 流形上的蛋白質與分子結構生成;以及機器人領域的軌跡生成。許多近期的大規模生成系統都採用 rectified-flow 或 OT-CFM 訓練,因其簡潔性與少步推斷的特性。
參考文獻
Cite error: <ref> tag with name "lipman2022" defined in <references> has group attribute "" which does not appear in prior text.
Cite error: <ref> tag with name "liu2022" defined in <references> has group attribute "" which does not appear in prior text.
Cite error: <ref> tag with name "albergo2023" defined in <references> has group attribute "" which does not appear in prior text.
Cite error: <ref> tag with name "tong2023" defined in <references> has group attribute "" which does not appear in prior text.
Cite error: <ref> tag with name "pooladian2023" defined in <references> has group attribute "" which does not appear in prior text.