フローマッチングの数学的原理とその応用
はじめに
本記事は、近年の生成AI分野で注目を集めるフローマッチング(Flow Matching, FM)の理論的基盤を、数学的な側面から解説することを目的とします。フローマッチングは、拡散モデル(Diffusion Models)の一般化と見なすことができ、より柔軟なパス設計やシミュレーションフリーな学習を実現することで、生成モデルの新たな可能性を切り拓いています。
僕自信が実験をやってみたい組み合わせにである、事前学習済みの拡散モデル(DDPM)をフローマッチングの枠組みでファインチューニングするアプローチは、両者の長所を組み合わせるという疑問を答えたいです。本記事では、まずフローマッチングを支える数学的な道具立てを整理し、次にその学習アルゴリズムの核心に迫ります。最後に、拡散モデルとの関係性や応用について少しだけ考察し、両方の関係の深堀は次の記事に書くつもりだ。今回は基礎をメインに勉強ノートをまとめてたものになります、僕は数学や基礎を学ぶよりコードを書く方が好きなんだけど自分の仮設を他人に説明する時には結構困ってしまうので実装の裏側に数学の勉強をやってみました。フローマッチングの実装は比較的に簡単ので今回は省くとします。
この記事の知識は【入門】フローマッチングのエッセンスとScaling Rectified Flow Transformers for High-Resolution Image SynthesisからまとめたものでGeminiくんと解読した内容になります。特にこの「【入門】フローマッチングのエッセンス」は本当に分かりやすくて時間があればぜひ読んでみてください。
第0章: フローマッチングの基礎的な実装
フローマッチングはDDPMより実装は非常に単純なので数学の裏側を学ぶ前に実際に学習の時に使われてる方式とコードを見た方が良いと思って、説明の前に事前に知って欲しい知識だ。
FluxやStable Diffusion 3のような、より高品質なオープンモデルはこの方式で潜在変数をノイズに混ぜる。
Pytorchの実装は以下になります。論文や数学の基礎ではUniformからサンプリングされる前提なんだけど学習率を上がるために時刻
# x1はデータのサンプルで、潜在変数である
nt = torch.randn((batch_size,), device=x1.device)
t = torch.sigmoid(nt)
# t の形を潜在変数に合わせる: [B] -> [B, 1, 1, 1]
t = t.view([batch_size, *([1] * len(x1.shape[1:]))])
# これはノイズでp0からサンプリングされたもの
x0 = torch.randn_like(x1)
xt = (1 - t) * x0 + t * x1
# cond_inputはテキストや画像の条件付きだ
v_pred = model(xt, t, encoder_hidden_states=cond_input)
loss = F.mse_loss((x1 - x0), v_pred,)
見ての通り学習自体はシンプルでノイズの追加は元画像とノイズを混ぜるように分かりやすい。このxtの方式は 最適輸送(Optimal Transport, OT の考え方から設定をしました。無駄のない二つのサンプルを繋がることが出来るし、個人的に一番良いポイントは正解ラベルである速度
DDPMは時刻
サンプリングの時は学習された速度を従ってランダムにサンプルされたノイズを Eulerサンプラーのコードは以下になります。
@torch.no_grad()
def sample_euler(
model_wrapper: torch.nn.Module,
x: torch.Tensor,
steps: int=30,
):
# 0から1を小さいステップに分割する
sigmas = torch.linspace(0, 1, steps + 1)
for i in range(len(sigmas)-1):
sigma_cur = sigmas[i]
sigma_next = sigmas[i+1]
dt = sigma_next - sigma_cur
v = model(x, sigma_cur)
# 速度に沿って小さいステップで潜在変数を更新する
x = x + v * dt
return x
実装を理解した後に数学の説明はもっと簡単になると思います。
第1章: フローモデルの数学的基礎
フローマッチングの「フロー(流れ)」を理解するには、まず個々のデータ点がどのように連続的に変換されるか、そしてその変換が確率分布全体にどう影響するかを数学的に記述する必要があります。
1.1 フローと速度場:ODEによる決定論的変換
フローマッチングは、単純なソース分布
この式は、「ある時刻
我らはこの
このODEを初期値
-
速度場
: 流れの「ルール」そのもの。ニューラルネットワークで学習する対象です。u_t(x) -
フロー
: ルール\psi_t(X_0) に従って初期点u_t から移動した結果としての「軌跡」。X_0 となります。X_t = \psi_t(X_0)
ODEは瞬間のルールであり、その解であるフローが具体的な軌跡を描きます。速度場
フローを直感的にノベルと「速度を積分すると変位になる」という定義もできます。
左辺を積分すると、微積分学の基本定理により、単に
これを移項すると、最終的なゴール
この方式はeulerサンプラーを使用する推論の時は積分の部分を小さいステップ
速度場とフローの等価性
フローマッチングにおける速度場
1. 速度場からフローへ
これは直感的な関係です。流れのルールである速度場
2. フローから速度場へ
逆に、全ての点の軌跡であるフロー
-
問いの設定: 我々が知りたいのは「時刻
に、位置t にいる点の速度は何か?」ということです。これがx にあたります。u(t,x) -
出身地を特定する: まず、その位置
にいる点が、元々はどの初期点x から来たのかを特定します。これはフローの逆写像x_0 を使って、\psi_t^{-1} と求められます。x_0 = \psi_t^{-1}(x) -
その出身者の速度を計算する: 次に、その「出身地
」からスタートした点が、時刻x_0 において持つべき速度を計算します。これは、フローt を時間\psi_t(x_0) で微分したt で与えられます。\dot{\psi}_t(x_0) -
組み合わせる: 最後に、ステップ2で求めた
をステップ3の式に代入することで、目的の速度場が得られます。x_0
この式は、「時刻
結論
このように、速度場とフローは1対1で対応する等価な関係です。この理論的な保証があるため、フローマッチングでは、ニューラルネットワークで直接モデル化しやすい速度場
1.2 確率分布の変換:プッシュフォワード
フロー
プッシュフォワード測度は元分布p0とp_t(x)を繋がる方式です。確率の「保存則」と「密度の変化」という2つのアイデアから成り立っています。変数変換の公式により、時刻
ここで、
1.3 確率パスと連続の方程式
ソース分布
この「地図」の時間変化と、それを実現するための「ルール」である速度場
直感的なアナロジーを使用すればもっと分かりやすいと思います。「部屋の中の人の流れ」、部屋の中に大勢の人がいるとしましょう。
-
: 時刻p_t(x) 、場所t における「人の密度」。時刻x のランダム変数t の確率密度だ。x -
: 場所u_t(x) での「人の動く速さと向き」。x -
: 場所\frac{\partial p_t(x)}{\partial t} での「密度の時間変化」。その場所は混雑してきているか、それとも空いてきているか。x -
: 「人の流れの勢い(確率流束)」。密度が高い場所で人が速く動けば、勢いは大きくなります。p_t(x)u_t(x) -
(発散): ある一点から「流れが湧き出しているか、吸い込まれているか」の度合い。\nabla_x \cdot (\dots)
連続の方程式を書き換えると
- もしある地点から人がどんどん出ていけば(発散が正)、その地点の密度は時間と共に減少します
。(\frac{\partial p_t}{\partial t} < 0) - もしある地点に人がどんどん流れ込んでくれば(発散が負、つまり収束)、その地点の密度は時間と共に増加します
。(\frac{\partial p_t}{\partial t} > 0)
「人の流れ」のアナロジーを用いて見事に解説されている通り、この式は「ある地点の確率密度の時間変化は、その地点への確率の流れ込み(あるいは流れ出し)の正味の量に等しい」ことを意味します。
この方程式の存在が極めて重要です。なぜなら、 我々が扱いやすい確率パス
1.4 瞬間的変数変換と尤度計算
生成モデルの性能を測る重要な指標の一つに「尤度(Likelihood)」があります。これは、モデルが学習データ(本物の画像など)をどれだけもっともらしく生成できるかの確率を示します。フローモデルの理論的な美しさの一つは、この尤度を正確に計算できる点にあります。
最終的な対数尤度
この式が示しているのは単一の粒子
-
発散が正 (
): その点から確率の流れが「湧き出し」ている状態。密度は時間と共に減少します。\nabla_x \cdot u > 0 -
発散が負 (
): その点に確率の流れが「吸い込まれ」ている状態。密度は時間と共に増加します。\nabla_x \cdot u < 0
「会計の基本原則」として考える方がもっと自然だと思います。最終的な値 - 初期の値 = トータルの変化量、これを対数確率密度に当てはめる
- トータルの変化量 =
\int_{0}^{1} (\text{瞬間の変化率}) dt - トータルの変化量 =
\int_{0}^{1} (\frac{d}{dt} \log p_t) dt
「瞬間の変化率」の正体は「負の発散」でした。
\frac{d}{dt} \log p_t = -\nabla_x \cdot u
これを上の式に代入します。
トータルの変化量 =\int_{0}^{1} (-\nabla_x \cdot u) dt
全てを合体させるとステップ2の式に、ステップ4の結果を代入します。
$$
\log p_1(x_1) - \log p_0(x_0) = \int_{0}^{1} (-\nabla_x \cdot u) dt
$$
この「瞬間の法則」を
この式は、フローモデルが理論上、正確な尤度計算能力を持つことの根拠となります。 しかし、従来の連続正規化フロー(CNF)では、学習のたびにこの式の右辺、特に高次元データでの「発散の計算」と積分を実行するための「ODEシミュレーション」が膨大な計算コストとなり、大きなボトルネックとなっていました。
フローマッチングは、まさにこの課題を克服するために生まれました。フローマッチングはこのコストのかかる尤度計算を学習プロセスから完全に切り離し、速度場そのものをより直接的かつ効率的に学習するアプローチを取ります。
第2章: フローマッチングの学習原理
前節で見たように、尤度を直接最大化する学習は困難を伴います。そこでフローマッチングは、より巧妙な方法で、望ましい速度場
2.1 条件付き確率パスの導入
我々が本当に知りたいのは、全てのノイズとデータの組み合わせを平均した、計算不可能な周辺確率パス
そこでフローマッチングでは、戦略を転換します。学習の時は答え(
- 目的地を決める: データセットから本物の画像
を一枚ランダムに選びます。x1 - 出発点を決める: ノイズ
をランダムに生成します。x0 - 中間地点を作る:
のように、出発点と目的地を使って中間地点X_t = (1-t)X_0 + tX_1 を作ります。X_t はx_t とx_0 に条件付けられる。x_1 - 正解ベクトルを計算する:
とX_t を使って、単純な式で「正解」の条件付き速度x_1 (例:u_t(X_t|x_1) )を計算します。X_1 - X_0
ニューラルネットワーク
学習時には 「目的地X1という答えをカンニング」 することで、あらゆる地点から目的地へ向かうための「正しい方向感覚」をモデルに徹底的に叩き込むのです。我らが計算出来ない
要約すると「ある特定の目的地
周辺確率パス
目的地である訓練サンプル
-
t=0 (スタート地点):
p_0(x|x_1) = p_0(x) - スタート時点では、目的地
の情報はなく、全てのパスはソース分布x₁ (例:ガウスノイズ)から出発します。p₀
- スタート時点では、目的地
-
t=1 (ゴール地点):
p_1(x|x_1) = \delta(x_1) - ゴール時点では、全てのパスは目的地である
の一点に収束している必要があります(x₁ はデルタ関数)。\delta
- ゴール時点では、全てのパスは目的地である
この条件付きパスを導入することで、学習の「正解ラベル」となるターゲット速度場が、次に示すように簡単に手に入るようになります。
2.2 生成速度場の導出
設計した条件付き確率パス
最も代表的で強力な**線形補間パス(Optimal Transport, OT)**を例に見てみましょう。このパスでは、時刻
このとき、目的地
-
フローの時間微分を求める: このパスは、
からX_0 への条件付きフローX_t と見なせます。その時間微分(速度)は、単純に\psi_t(X_0|x_1) = (1-t)X_0 + t x_1 となります。\dot{\psi}_t = x_1 - X_0 -
モデルの入力変数で表現する: しかし、モデル
が学習時に入力するのは、その瞬間の位置u_\theta であり、出発点X_t を知ることはできません。そこで、X_0 という速度を、x_1 - X_0 を使って表現し直す必要があります。X_t
の式をX_t = (1-t)X_0 + t x_1 について解くと、X_0
となります。X_0 = \frac{X_t - t x_1}{1-t} -
条件付き速度場を導出する: ステップ1の
X₀にステップ2の式を代入すると、最終的な条件付き速度場が得られます。
この式は非常に直感的です。
-
: 現在地x_1 - X_t から目的地X_t へ向かう方向ベクトルを示しています。x_1 -
: 時間の進行度を表す係数です。ゴールが近づく(\frac{1}{1-t} が1に近づく)ほど分母が0に近づき、速度が急激に増大します。これは、残されたわずかな時間で目的地に正確にたどり着くために「急加速」する必要があることを意味しています。t
このようにして、我々はニューラルネットワークが学習すべき「正解」の速度ベクトルを、訓練サンプル
2.3 周辺化トリック:単純な「条件付き」学習から複雑な「周辺」学習へ
前節では、特定の目的地
モデルが最終的に学習すべき真の目標は、全てのノイズとデータの組み合わせを平均した、いわば「流れの平均的な振る舞い」である周辺速度場
周辺速度場:真の学習目標
ある時空
より厳密には、この平均は事後分布
この
立ちはだかる壁:計算不可能性
しかし、この
となり、分母に計算不可能な周辺確率パス
フローマッチングの妙技:周辺化トリック
ここで、フローマッチングの最も巧妙かつ強力な理論的帰結が登場します。それは、この計算不可能な周辺速度場
計算が簡単な条件付き速度場
このトリックを支えるのが、原著で示されている以下の定理です。
定理(周辺化トリック):
もし、設計した条件付き速度場が、対応する条件付き確率パス u_t(x|z) を正しく生成するならば、それらを平均して得られる周辺速度場 p_t(x|z) は、周辺確率パス u_t(x) を正しく生成する。 p_t(x)
この定理は、我々が行う学習の正当性を保証してくれます。個々の単純なルール(条件付き速度場)がそれぞれの単純な経路(条件付きパス)を正しく生成するように学習させれば、それらのルールを脳内で(暗黙的に)平均した結果であるモデルの出力は、全ての経路を束ねた複雑な流れ(周辺パス)を正しく生成できるようになる、というわけです。
勿論、条件付きだからこそデータ量に依存するので非常に大きいなデータセットは必要になります。Stable Diffusion 3は1Bくらいの画像で事前学習されたらしくて、LuminaImagen2.0も100M以上のサンプルを含まれるデータセットを使ったらしい。
結論:シミュレーションフリー学習の理論的根幹
周辺化トリックにより、我々は扱いやすい「条件付き」の世界でモデルを訓練するだけで、本来の目的である複雑で計算不可能な「周辺」の世界のダイナミクスを捉えることができます。
これにより、学習プロセスにおいて
2.4 フローマッチング損失:学習の羅針盤
前節までで、計算が困難な「周辺速度場」の代わりに、扱いやすい「条件付き速度場」を導出できることを見ました。本節では、この条件付き速度場を「正解ラベル」として、ニューラルネットワーク
理想的な損失と現実的な損失
我々が本当に最小化したい理想的な損失は、モデルの予測
ここで
しかし、前述の通り
この損失関数では、期待値
勾配の等価性:CFMがFMを達成する理由
ここで最も重要な理論的保証が登場します。それは、 「FM損失とCFM損失のパラメータ
この
フローマッチングの学習を実用的なものにするための核心的な結果です。この定理が意味するのは、計算が簡単な
この「勾配の等価性」こそが、「周辺化トリック」をアルゴリズムとして実現する数学的なエンジンです。我々は、扱いやすい条件付き量をターゲットにして回帰学習を行うだけで、目的の複雑な周辺量を学習できることが、これにより保証されるのです。
第3章: フローマッチングの拡張性と柔軟性
フローマッチングのフレームワークは、その設計の柔軟性により、様々なタスクやデータ形式へと拡張することが可能です。本章では、条件付き生成への応用や、パス設計の多様な選択肢、そして拡散モデルとの深い関係性について掘り下げます。
3.1 条件付き生成への応用
これまでの議論は、主にデータ分布
この場合、学習データは
CFM損失も同様に拡張され、データサンプル
サンプリング時には、所望の条件
3.2 条件付けの多様な選択肢
CFM損失を構築する際の「条件」の選び方には、いくつかのバリエーションがあり、タスクに応じて最適なものを選択できます(図3.3参照)。
-
ターゲット条件付け (
): これまで主に見てきたケースです。ノイズからデータを生成する一般的な生成タスクに適しています。Z = X_1 -
ソース条件付け (
): ソースサンプルZ = X_0 を条件とし、そこからターゲットX_0 へ向かうフローを学習します。これは、ある画像X_1 を別のスタイルX_0 へと変換するような、画像編集やスタイル変換タスクに応用できます。ですがQwen-Imageみたいな最近の画像編集モデルは普通のターゲット条件付けで学習させられた、編集をしたい画像は条件X_1 、テキストと同じ扱いで入力をします。y -
ペア条件付け (
): ソースとターゲットのペアを固定し、その間を補間するフローを学習します。この場合、速度場はZ = (X_0, X_1) には陽に依存せず、X_t の関数となります。これは「確率的補間(Stochastic Interpolant)」という関連研究と密接な繋がりがあります。X_0, X_1, t
重要なのは、これらの条件付けの方法が適切に設計されていれば、「周辺化トリック」によって、最終的に学習される周辺速度場
3.3 最適輸送との接続:最も効率的な経路の探求
確率パスの設計には無限の自由度がありますが、どのようなパスが「良い」パスなのでしょうか。その一つの理想的な指針を与えるのが、 最適輸送(Optimal Transport, OT) の理論です。
OTは、ある分布を別の分布へと輸送する際に、総コスト(例えば、移動距離やエネルギー)が最小となるような「最も効率的な輸送計画」を見つける数学理論です。これをフローマッチングに当てはめると、「サンプルがノイズからデータへと、できるだけまっすぐ、無駄なく移動する」ようなフローを学習することに対応します。
そして、理論的な解析により、2.2節で紹介した線形補間パス
マルチタスクの問題を完全に解決出来る期待がないですが、せめて和らげることが出来るだろうか。実際にノイズに近いステップは最終的に生成される画像の影響は他のステップより高いので学習の時はUniformからサンプリングではなくてはLogit-Normal分布からサンプリングを行う。
このOTとの関連性は、フローマッチングの最もシンプルで基本的な損失関数
3.4 拡散モデルとの架け橋:アフィンフローとガウスパス
線形補間パスは、より一般的なアフィン条件付きフローの特殊なケースと見なすことができます。最適輸送はこのフローの例の一つだ。
ここで
特に重要なのは、ソース分布
この形式は、拡散モデル(DDPM)で用いられるノイズ付加過程と非常によく似た形をしています。実際に、
この発見は、フローマッチングと拡散モデルの間に深い理論的な繋がりがあることを示唆しています。拡散モデルのスコアマッチング学習は、特定のガウスパスを仮定した フローマッチングの特殊な場合 として捉えることができます。一方で、フローマッチングは勾配(スコア)である必要のない 、より一般的なベクトル場を学習できるため、表現力の高い変換をモデル化できる可能性があります。
この関係性は、DDPMで事前学習したモデルの知識を、フローマッチングの枠組みでファインチューニングするアプローチの理論的な妥当性を強く支持するものです。DDPMが学習した「ノイズを除去する能力」は、FMにおける「ノイズからデータへ向かう速度場」の優れた初期値として機能する可能性が高いと考えられます。
第3章のまとめ:フローマッチングの動作原理と展望
本記事を通じて、フローマッチングの数学的基盤からその学習アルゴリズムの核心までを概観しました。
- 理論的基盤: フローマッチングは、ODEで記述される決定論的なフローを用いて分布間を変換します。確率パスと速度場は「連続の方程式」で結びついており、これが学習の理論的正当性を保証します。
- 学習の核心: 計算不可能な周辺速度場を直接学習する代わりに、計算可能な条件付き速度場をターゲットとするCFM損失を用います。周辺化トリックと勾配の等価性により、この単純な回帰学習が、真の目標である周辺速度場の学習を達成することを保証します。
- 柔軟性と拡張性: 条件付けの仕方や確率パスの設計に高い自由度があり、条件付き生成や画像編集など多様なタスクに適用可能です。特に、最適輸送に動機付けられた線形パスは、シンプルかつ強力な選択肢です。
- 拡散モデルとの関係: アフィンフローとガウスパスの概念を通じて、フローマッチングが拡散モデルを包含する、より一般的なフレームワークであることが示唆されます。
Discussion