[Journal club] Flow Matching for Generative Modeling

Transcript

1. Yaron Lipman1,2 Ricky T. Q. Chen1 Heli Ben-Hamu2 Maximilian Nickel1 Matt Le1 1Meta AI (FAIR) 2Weizmann Institute of Science Flow Matching for Generative Modeling 2025 杉浦孔明研究室 妹尾 幸樹 LIPMAN, Yaron, et al. Flow Matching for Generative Modeling. In: 11th International Conference on Learning Representations, ICLR 2023. 2023. ICLR23
2. 概要 2 ▪ 背景 L 既存⼿法は訓練の安定性・効率性に課題 Continuous Normalizing Flow をスケーラブルかつ効率的に訓練したい ▪ 提案︓Flow Matching ▪ Continuous Normalizing Flow を simulation-free で スケーラブルかつ安定的・効率的に訓練するための理論 ▪ 結果 ▪ unconditional な画像⽣成タスクでは，全てのデータセット・尺度で Flow Matching w/ OT が最良 ▪ 効率的な学習・⾼速なサンプリング
3. 背景︓既存⼿法は訓練の安定性・効率性に課題 3 ▪ 拡散モデル (Diffusion Models, DM) J スケーラブルで⽐較的安定した訓練 L 単純な拡散過程の設計により，確率パスの空間が⾮常に限定的 L ⻑い訓練・推論時間 L 効率的なサンプリング⼿法の活⽤が必要 (e.g. DDIM [Song+, ICLR21]) ▪ Continuous Normalizing Flow (CNF) J 任意の確率パスを設計可能 (=⾼い設計⾃由度) L ⾼コストな常微分⽅程式 (ODE) シミュレーション L simulation-free な⼿法は勾配近似により訓練が不安定 L 計算コスト ⾼ ℒ!"# = 𝐿 𝒙$ + & %! %" 𝑓 𝒙% , 𝑡, 𝜃 𝑑𝑡 𝑓: ニューラルネット 𝐿: 損失関数 𝒙 ∈ ℝ𝒅 𝜃: パラメータ 𝒙": ノイズ 𝑡 : サンプリング時刻 CNFをスケーラブルかつ安定的・効率的に訓練したい
4. 関連研究 4 ⼿法 特徴 Neural ODE [Chen+, NeurIPS18] Normalizing Flow を連続時間で扱うCNFを提案 L ⾼コストなODEシミュレーション Moser Flow [Rozen+, NeurIPS21] 事前分布とデータ分布を線形補完することによって，CNFの simulation-free な訓練⽅法を提案 L データが⾼次元の場合に推定が困難な積分 DDPM [Ho+, NeurIPS20] score matching を⽤いて拡散モデルを訓練 L ⻑い訓練・推論時間 DDPM [Ho+, NeurIPS20] Neural ODE[Chen+, NeurIPS18]
5. ▪ 単純な事前分布 𝑝! (e.g. ガウス分布) を未知分布 𝑝" へと変形 ▪ 速度場 𝒗# によって連続時間に依存した微分同相写像 (フロー) 𝜙# を構成 ▪ フローは以下で定義される ODE で記述 ▪ フローによる変数変換 (⟹ “𝒗# が確率密度パス 𝑝# を⽣成する”※) ※ ⟺ 連続の[⽅程]式を満たす ▪ torchdiffeq [Chen+, 18] が速度場 𝒗% をニューラルネットで表現 前提︓Continuous Normalizing Flow (CNF) 5 𝑑 𝑑𝑡 𝜙# 𝒙 = 𝒗# 𝜙# 𝒙 , 𝜙! 𝒙 = 𝒙 𝑝# = [𝜙# ]∗ 𝑝! = 𝑝! 𝜙# %" 𝒙 det 𝜕𝜙# %" 𝜕𝒙 (𝒙) 𝜙: 0, 1 ×ℝ𝒅 → ℝ𝒅 𝑝: 0, 1 ×ℝ𝒅 → ℝ#" 𝒗: 0, 1 ×ℝ𝒅 → ℝ𝒅 𝑑 𝑑𝑡 𝑝# 𝒙 = −div 𝑝# 𝒙 𝒗# 𝒙 流⼊量 = 流出量 Outflow [Holderrieth, 25] Inflow
6. ▪ ガウス分布 𝑝! 𝒙 を未知分布 𝑞 𝒙" へ変形する ▪ 任意の時刻, 位置における速度場をニューラルネット 𝒗& 𝑡, 𝒙 で表現 ▪ Flow Matching (FM) Loss (Marginal VF 𝒖# 𝒙 を学習) J 𝒙" を中⼼とした鋭い分布は容易に表現可能 J 𝑝# 𝒙|𝒙" , 𝒖# 𝒙|𝒙" は⾃由に設計可能（後述） Ø 定理１ Marginal Vector Field (VF) を直接学習するのは困難 6 𝑝! 𝒙 = 𝒩 𝒙 | 0, 𝐼 , 𝒙"~𝑞 ℒ'( 𝜃 = 𝔼#,*! 𝒙 𝒗& 𝑡, 𝒙 − 𝒖# 𝒙 , 𝒖# 𝒙 = ∫ 𝒖# 𝒙|𝒙" *! 𝒙|𝒙" . 𝒙" *! 𝒙 𝑑𝒙" とすれば， 𝑝" 𝒙|𝒙" = 𝒩 𝒙 |𝒙" , 𝜎,𝐼 L 計算が困難 𝒖: 0, 1 ×ℝ𝒅 → ℝ𝒅 ( ※ 𝜎 は⼗分に⼩さい ) 𝒖# 𝒙|𝒙" が 𝑝# 𝒙|𝒙" を⽣成する ⟹ 𝒖# 𝒙 は 𝑝# 𝒙 を⽣成する [Lipman+, 24]
7. 定理１の証明︓連続の式と周辺化を⽤いた式変形 7 ▪ 𝒖# 𝒙|𝒙" が 𝑝# 𝒙|𝒙" を⽣成する ⟹ 𝒖# 𝒙 は 𝑝# 𝒙 を⽣成する 𝑑 𝑑𝑡 𝑝# 𝒙 = 𝑑 𝑑𝑡 @ 𝑝# 𝒙|𝒙" 𝑞 𝒙" 𝑑𝒙" = @ 𝑑 𝑑𝑡 𝑝# 𝒙|𝒙" 𝑞 𝒙" 𝑑𝒙" = @ −div 𝑝# 𝒙|𝒙" 𝒖# 𝒙|𝒙" 𝑞 𝒙" 𝑑𝒙" = −div @ 𝑝# 𝒙|𝒙" 𝒖# 𝒙|𝒙" 𝑞 𝒙" 𝑑𝒙" = −div 𝑝# 𝒙 @ 𝒖# 𝒙|𝒙" 𝑝# 𝒙|𝒙" 𝑞 𝒙" 𝑝# 𝒙 𝑑𝒙" = −div 𝑝# 𝒙 𝒖# 𝒙 = 𝒖# 𝒙 ∎ × 𝑝# 𝒙 𝑝# 𝒙 𝑑 𝑑𝑡 𝑝# 𝒙′ = −div 𝑝# 𝒙′ 𝒗# 𝒙′ 周辺化
8. CFM Loss を⽤いることで速度場の学習を実現 8 ▪ Conditional Flow Matching (CFM) Loss 𝓛𝐂𝑭𝑴を⽤いて訓練すれば，Marginal VF を学習することが出来る ▪ 証明 Ø 定理２ ∀𝒙 ∈ ℝ𝒅, ∀𝑡 ∈ 0, 1 , 𝑝# 𝒙 > 0 ⟹ ∇& ℒ'( 𝜃 = ∇& ℒ3'( 𝜃 ℒ3'( 𝜃 = 𝔼#,. 𝒙" ,*! 𝒙|𝒙" 𝒗& 𝑡, 𝒙 − 𝒖# 𝒙|𝒙" , 𝒗& 𝑡, 𝒙 − 𝒖% 𝒙 ' = 𝒗& 𝑡, 𝒙 ' − 2 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙 + 𝒖% 𝒙 ' 𝒗& 𝑡, 𝒙 − 𝒖% 𝒙|𝒙( ' = 𝒗& 𝑡, 𝒙 ' − 2 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙|𝒙( + 𝒖% 𝒙|𝒙( ' ∇& ℒ#) 𝜃 = ∇& ℒ!#) 𝜃 ⟺ ℒ#) 𝜃 = ℒ!#) 𝜃 + 𝑐𝑜𝑛𝑠𝑡 𝔼*# 𝒙 𝒗& 𝑡, 𝒙 ' = & 𝒗& 𝑡, 𝒙 ' 𝑝% 𝒙 𝑑𝒙 = & & 𝒗& 𝑡, 𝒙 ' 𝑝% 𝒙|𝒙( 𝑞 𝒙( 𝑑𝒙( 𝑑𝒙 = 𝔼, 𝒙" , *# 𝒙|𝒙" 𝒗& 𝑡, 𝒙 ' ①
9. 定理２の証明 (続き) : ℒ!" と ℒ456 の勾配は同じ 9 𝔼*# 𝒙 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙 = & 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙 , ∫ 𝒖% 𝒙|𝒙( 𝑝% 𝒙|𝒙( 𝑞 𝒙( 𝑑𝒙( 𝑝% 𝒙 𝑝% 𝒙 𝑑𝒙 = & 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙 , & 𝒖% 𝒙|𝒙( 𝑝% 𝒙|𝒙( 𝑞 𝒙( 𝑑𝒙( 𝑑𝒙 = & & 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙 , 𝒖% 𝒙|𝒙( 𝑝% 𝒙|𝒙( 𝑞 𝒙( 𝑑𝒙( 𝑑𝒙 = 𝔼, 𝒙" , *# 𝒙|𝒙" 𝒗& 𝑡, 𝒙 , 𝒖% 𝒙|𝒙( ② 𝒖% 𝒙 ', 𝒖% 𝒙|𝒙( ' は 𝜃 に⾮依存 ⟹ ③ ①, ②, ③ より, ℒ#) 𝜃 = ℒ!#) 𝜃 + 𝑐𝑜𝑛𝑠𝑡 ∎ 𝔼*# 𝒙 𝒖% 𝒙 ', 𝔼, 𝒙" , *# 𝒙|𝒙" 𝒖% 𝒙|𝒙( ' は定数
10. 確率パスの設計︓Gaussian Probability Path 10 ▪ ガウシアン条件付き確率パス (Gaussian conditional probability paths) ▪ 条件付きフロー 𝜓# 𝒙 は以下 ▪ 証明 𝑝# 𝒙|𝒙" = 𝒩 𝒙|𝜇# 𝒙" , 𝜎# 𝒙" ,𝐼 𝜇" 𝒙" = 𝒙", 𝜎" 𝒙" = 𝜎456 𝜓# 𝒙 = 𝜎# 𝒙" 𝒙 + 𝜇# 𝒙" 𝜇: 0, 1 ×ℝ𝒅 → ℝ𝒅 𝜎: 0, 1 ×ℝ → ℝ#" 𝜎$%& は⼗分に⼩さい Ø 定理３ 𝑝# 𝒙|𝒙" がガウシアン確率パス のとき，Conditional VF 𝒖# 𝒙|𝒙" は 𝒖# 𝒙|𝒙" = 7! # 𝒙" 7! 𝒙" 𝒙 − 𝜇# 𝒙" + 𝜇# 8 𝒙" である 𝒚 = 𝜓# 𝒙 = 𝜎# 𝒙" 𝒙 + 𝜇# 𝒙" ⟺ 𝒙 = 𝜓# %" 𝒚 = 𝒚%:! 𝒙" 7! 𝒙" より, 𝒚 を 𝑡 で微分して 𝒚8 = 𝒖# 𝒚|𝒙" = 𝜎𝒕 8 𝒙" 𝒙 + 𝜇𝒕 8 𝒙" = 𝜎𝒕 8 𝒙" 𝜎# 𝒙" 𝒚 − 𝜇# 𝒙" + 𝜇# 8 𝒙" 𝜓: 0, 1 ×ℝ𝒅 → ℝ𝒅 ∎ [Lipman+, 24]
11. 効率的な確率パス︓Optimal Transport Conditional VFs 11 ▪ Diffusion Conditional VFs ▪ Optimal Transport Conditional VFs 𝑝# 𝒙|𝒙" = 𝒩 𝒙|𝛼"%#𝒙", 1 − 𝛼"%# , 𝐼 𝒖# 𝒙|𝒙" = 𝛼"%# 8 1 − 𝛼"%# , 𝛼"%#𝒙 − 𝒙" 𝛼: ノイズスケジュール係数 拡散モデルと同様 𝜇# 𝒙" = 𝑡𝒙" , 𝜎# 𝒙" = 1 − 1 − 𝜎456 𝑡 𝒖# 𝒙|𝒙" = 𝒙" − 1 − 𝜎456 𝒙 1 − 1 − 𝜎456 𝑡 𝜓# 𝒙 = 1 − 1 − 𝜎456 𝑡 𝒙 + 𝑡𝒙" ℒ<=3'( 𝜃 = 𝔼#,. 𝒙" ,*! 𝒙|𝒙" 𝒗& 𝑡, 𝒙 − 𝒙" − 1 − 𝜎456 𝒙! , 𝒙"~𝑝" [Lipman+, 24] ⼀般的には0 L 遠回り J ⼀直線
12. 実験設定 12 ▪ Dataset ▪ CIFAR-10 ▪ ImageNet {32, 64, 128} ▪ Method ▪ OT path ▪ Diffusion path w/ FM ▪ Diffusion path w/ SM ▪ 学習環境・時間︓記載なし ▪ 評価指標 ▪ negative log-likelihood (NLL) ▪ Frechet Inception Distance (FID) ▪ number of function evaluations (NFE)
13. 定量的結果︓少ないサンプリングステップで⾼品質 13 ▪ 全てのデータセット・尺度で FM w/ OT が最良
14. 定量的結果︓FM は効率的に訓練・推論 14 ▪ FMは早く収束し，⾼速にサンプリング可能
15. 定性的結果︓早い段階でノイズを除去 15 J 線形にノイズが減少 L 最後の⽅で⼀気にノイズが減少
16. まとめ 16 ▪ 背景 L 既存⼿法は訓練の安定性・効率性に課題 Continuous Normalizing Flow をスケーラブルかつ効率的に訓練したい ▪ 提案︓Flow Matching ▪ Continuous Normalizing Flow を simulation-free で スケーラブルかつ安定的・効率的に訓練するための理論 ▪ 結果 ▪ unconditional な画像⽣成タスクでは，全てのデータセット・尺度で Flow Matching w/ OT が最良 ▪ 効率的な学習・⾼速なサンプリング