🌀

フローマッチングの数学的原理とその応用

に公開

はじめに

本記事は、近年の生成AI分野で注目を集めるフローマッチング(Flow Matching, FM)の理論的基盤を、数学的な側面から解説することを目的とします。フローマッチングは、拡散モデル(Diffusion Models)の一般化と見なすことができ、より柔軟なパス設計やシミュレーションフリーな学習を実現することで、生成モデルの新たな可能性を切り拓いています。

僕自信が実験をやってみたい組み合わせにである、事前学習済みの拡散モデル(DDPM)をフローマッチングの枠組みでファインチューニングするアプローチは、両者の長所を組み合わせるという疑問を答えたいです。本記事では、まずフローマッチングを支える数学的な道具立てを整理し、次にその学習アルゴリズムの核心に迫ります。最後に、拡散モデルとの関係性や応用について少しだけ考察し、両方の関係の深堀は次の記事に書くつもりだ。今回は基礎をメインに勉強ノートをまとめてたものになります、僕は数学や基礎を学ぶよりコードを書く方が好きなんだけど自分の仮設を他人に説明する時には結構困ってしまうので実装の裏側に数学の勉強をやってみました。フローマッチングの実装は比較的に簡単ので今回は省くとします。

この記事の知識は【入門】フローマッチングのエッセンスScaling Rectified Flow Transformers for High-Resolution Image SynthesisからまとめたものでGeminiくんと解読した内容になります。特にこの「【入門】フローマッチングのエッセンス」は本当に分かりやすくて時間があればぜひ読んでみてください。


第0章: フローマッチングの基礎的な実装

フローマッチングはDDPMより実装は非常に単純なので数学の裏側を学ぶ前に実際に学習の時に使われてる方式とコードを見た方が良いと思って、説明の前に事前に知って欲しい知識だ。

FluxやStable Diffusion 3のような、より高品質なオープンモデルはこの方式で潜在変数をノイズに混ぜる。x_t = (1-t)x_0 + t*x_1 という単純な線形補間パスを仮定する。このとき、条件付きベクトル場は u_t(x_t | x_1) = x_1 - x_0 となる。v_\thetaは、x_t、時刻t、テキスト条件cを入力として、このベクトル場(速度)を予測するように学習される。損失関数は以下の通りである。損失は普通のMSEで予測された速度と本物の速度の距離を小さくするように学習させる。

L_{FM} = \mathbb{E}_{t, x_0, x_1, c} [ ||(x_1 - x_0) - v_\theta((1-t)x_0 + t*x_1, t, c)||^2 ]

Pytorchの実装は以下になります。論文や数学の基礎ではUniformからサンプリングされる前提なんだけど学習率を上がるために時刻t はLogit-Normal分布からサンプリングを行う。中間ステップ t=0.5 はノイズと元画像の融合が両方の信号を同じ割合で持つのでモデルの予測が難しい。

# 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 の考え方から設定をしました。無駄のない二つのサンプルを繋がることが出来るし、個人的に一番良いポイントは正解ラベルである速度u_t(x_t | x_1) = x_1 - x_0は時刻tに依存しないことだ。

DDPMは時刻 t と共にノイズの量を調整するのでモデルの予測量はいつも同じじゃないので複数のタスクを同時に学ばせてることになるということは色んな論文で書かれている。だからフローマッチングは少しでもその問題を解決出来ることを期待します。

サンプリングの時は学習された速度を従ってランダムにサンプルされたノイズを x_1(画像)に変換させる、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による決定論的変換

フローマッチングは、単純なソース分布 p_0(例:標準正規分布 \mathcal{N}(0, I))に従うサンプル X_0 を、複雑なターゲット分布 p_1(データ分布、例えば猫の画像)に従うサンプル X₁ へと変換するプロセスをモデル化します。この変換は、決定論的な 常微分方程式(Ordinary Differential Equation, ODE) によって記述されます。

\frac{d(X_t)}{dt} = u_t(X_t)

この式は、「ある時刻 t における点 X_t の瞬間的な変化(速度)は、ベクトル場 u_t(X_t) によって与えられる」ことを意味します。 ここで、u_t(x)速度場(Velocity Field) と呼ばれ、時空の各点 (t, x) における流れの向きと速さを定義する「ルールブック」に相当します。

我らはこの p_1 をしらない、サンプリングは出来るそれは学習のデータセット。なのでフローマッチングは p₁ という未知分布から新しいサンプリングを生成する手法である、生成AIと同様に最終的に p₁ を近似したい。

このODEを初期値 X_0 のもとで t=0 から t=1 まで解く(積分する)ことで得られる軌跡がフロー(Flow) \psi_t(X_0) です。

  • 速度場 u_t(x): 流れの「ルール」そのもの。ニューラルネットワークで学習する対象です。
  • フロー \psi_t(X_0): ルール u_t に従って初期点 X_0 から移動した結果としての「軌跡」。X_t = \psi_t(X_0) となります。

ODEは瞬間のルールであり、その解であるフローが具体的な軌跡を描きます。速度場 u_t とフロー \psi_t は1対1に対応しており、片方が決まればもう一方も一意に定まります。

フローを直感的にノベルと「速度を積分すると変位になる」という定義もできます。

\int_{0}^{1} \frac{dX_t}{dt} dt = \int_{0}^{1} u_t(X_t) dt

左辺を積分すると、微積分学の基本定理により、単に X_1 - X_0 となります。

X_1 - X_0 = \int_{0}^{1} u_t(X_t) dt

これを移項すると、最終的なゴール X1 を得るための理想的な式が得られます。

X_1 = X_0 + \int_{0}^{1} u_t(X_t) dt

X0 というノイズに、t=0 から t=1 までの全行程にわたる速度 u_t(X_t) の影響をすべて足し合わせれば(積分すれば)、最終的な画像 X_1 が得られる、ということを意味しています。

この方式はeulerサンプラーを使用する推論の時は積分の部分を小さいステップhをニューラルネットワークの出力である u_\theta(X_t) に掛け算すると数値積分で近似をするのだ。

速度場とフローの等価性

フローマッチングにおける速度場 u(t,x)(流れのルール)とフロー \psi_t(x)(流れの軌跡)は、表裏一体の関係にあり、どちらか一方を決めればもう一方も一意に定まります。

1. 速度場からフローへ

これは直感的な関係です。流れのルールである速度場 u(t,x) が与えられれば、そのルールに従って初期点 x₀ が時間と共にどう動くかを計算(ODEを解く)することで、その軌跡であるフロー \psi_t(x_0) が一意に決まります。

2. フローから速度場へ

逆に、全ての点の軌跡であるフロー \psi_t(x) が分かっている場合、それを生み出す瞬間のルール(速度場)u(t,x) も計算できます。これは以下のステップで理解できます。

  1. 問いの設定: 我々が知りたいのは「時刻 t に、位置 x にいる点の速度は何か?」ということです。これが u(t,x) にあたります。

  2. 出身地を特定する: まず、その位置 x にいる点が、元々はどの初期点 x_0 から来たのかを特定します。これはフローの逆写像 \psi_t^{-1} を使って、x_0 = \psi_t^{-1}(x) と求められます。

  3. その出身者の速度を計算する: 次に、その「出身地 x_0」からスタートした点が、時刻 t において持つべき速度を計算します。これは、フロー \psi_t(x_0) を時間 t で微分した \dot{\psi}_t(x_0) で与えられます。

  4. 組み合わせる: 最後に、ステップ2で求めた x_0 をステップ3の式に代入することで、目的の速度場が得られます。

u(t,x) = \dot{\psi}_t( \psi_t^{-1}(x) )

この式は、「時刻 t、位置 x での速度とは、その点 x の出身地 \psi_t^{-1}(x) から来た点が、時刻 t に持つべき速度 \dot{\psi}_t のことである」と解釈できます。

結論

このように、速度場とフローは1対1で対応する等価な関係です。この理論的な保証があるため、フローマッチングでは、ニューラルネットワークで直接モデル化しやすい速度場 u(t,x) を学習の対象としてパラメータ化するアプローチを取ります。

1.2 確率分布の変換:プッシュフォワード

フロー \psi_t は個々のデータ点を動かしますが、我々の真の関心事は確率分布全体の変換です。ソース分布 p₀ に従う確率変数 X₀ がフローによって X_t = \psi_t(X_0) に変換されるとき、X_t が従う分布 p_t は、p_0\psi_t によって「押し出した」ものと考えられます。これをプッシュフォワード測度と呼び、p_t = (\psi_t)_{\#} p_0 と表記します。

プッシュフォワード測度は元分布p0とp_t(x)を繋がる方式です。確率の「保存則」と「密度の変化」という2つのアイデアから成り立っています。変数変換の公式により、時刻 t における確率密度 p_t(x) は以下のように計算されます。

p_t(x) = p_0(\psi_t^{-1}(x)) *|\det J_{\psi_t^{-1}}(x)|

ここで、|\det J_{\psi_t^{-1}}(x)| はフローの逆写像 \psi_t^{-1} のヤコビ行列式(ヤコビアン)の絶対値です。 これは、フローによる微小体積の変化率を補正する項です。これは確率の総和が1に保たれる「確率保存則」のもとで、空間の伸縮によって密度がどう変化するかを記述しています。この方式は中間確率p_t(x)とフローを繋がる。

1.3 確率パスと連続の方程式

ソース分布 p_0 とターゲット分布 p_1 を滑らかに結ぶ、時間的に連続な分布の連なり p_t(x)確率パス と呼びます。 これは、生成プロセスにおける「中間目標の地図」と考えることができます。

この「地図」の時間変化と、それを実現するための「ルール」である速度場 u_t は、無関係ではありません。両者は 連続の方程式(Continuity Equation) という物理法則(質量保存則)によって固く結びついています。

\frac{\partial p_t(x)}{\partial t} + \nabla_x \cdot (p_t(x) u_t(x)) = 0

直感的なアナロジーを使用すればもっと分かりやすいと思います。「部屋の中の人の流れ」、部屋の中に大勢の人がいるとしましょう。

  • 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} = - \nabla_x \cdot (p_t u_t) となります。これは、「ある場所の人の密度が増加する速さ (\frac{\partial p_t}{\partial t}) は、その場所に流れ込んでくる人の正味の量 ( - \nabla_x \cdot (p_t u_t)) に等しい。」自然な質量保存則を述べています。

  • もしある地点から人がどんどん出ていけば(発散が正)、その地点の密度は時間と共に減少します (\frac{\partial p_t}{\partial t} < 0)
  • もしある地点に人がどんどん流れ込んでくれば(発散が負、つまり収束)、その地点の密度は時間と共に増加します (\frac{\partial p_t}{\partial t} > 0)

「人の流れ」のアナロジーを用いて見事に解説されている通り、この式は「ある地点の確率密度の時間変化は、その地点への確率の流れ込み(あるいは流れ出し)の正味の量に等しい」ことを意味します。

この方程式の存在が極めて重要です。なぜなら、 我々が扱いやすい確率パス p_t を設計すれば、それを生成するターゲット速度場 u_t が理論的に必ず存在する ことが保証されるからです。 フローマッチングは、この u_t をニューラルネットワークで直接近似することを目指します。

1.4 瞬間的変数変換と尤度計算

生成モデルの性能を測る重要な指標の一つに「尤度(Likelihood)」があります。これは、モデルが学習データ(本物の画像など)をどれだけもっともらしく生成できるかの確率を示します。フローモデルの理論的な美しさの一つは、この尤度を正確に計算できる点にあります。

最終的な対数尤度 \log p_1(x_1) は、初期の対数尤度 \log p_0(x_0) に、プロセス全体の「変化量」を足し合わせることで得られます。この「瞬間の変化率」を捉えるのが、 瞬間的変数変換(Instantaneous Change of Variables) の公式です。

\frac{d}{dt} \log(p_t(\psi_t(x_0))) = -\nabla_x \cdot u(t, \psi_t(x_0))

この式が示しているのは単一の粒子 x_0 の軌跡 \psi_t(x_0) を追いかけたときの、その場所における確率密度の時間変化です。 この変化率は、その点における速度場 u_t発散(divergence) \nabla_x \cdot u という、空間的な性質によって決まります。

  • 発散が正 (\nabla_x \cdot u > 0): その点から確率の流れが「湧き出し」ている状態。密度は時間と共に減少します。
  • 発散が負 (\nabla_x \cdot u < 0): その点に確率の流れが「吸い込まれ」ている状態。密度は時間と共に増加します。

「会計の基本原則」として考える方がもっと自然だと思います。最終的な値 - 初期の値 = トータルの変化量、これを対数確率密度に当てはめる \log p_1(x_1) - \log p_0(x_0) = \text{log pのトータルの変化量} 。数学の基本定理(微積分学の基本定理)により、「トータルの変化量」は「瞬間の変化率」を最初から最後まで積分することで得られます。

  1. トータルの変化量 = \int_{0}^{1} (\text{瞬間の変化率}) dt
  2. トータルの変化量 = \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
    $$

この「瞬間の法則」を t=0 から t=1 まで積分することで、プロセス全体の対数尤度の変化量を計算でき、最終的なターゲット尤度の式が導かれます。

\log p_1(x_1) = \log p_0(x_0) - \int_{0}^{1} \nabla_x \cdot u(t, \psi_t(x_0)) dt

この式は、フローモデルが理論上、正確な尤度計算能力を持つことの根拠となります。 しかし、従来の連続正規化フロー(CNF)では、学習のたびにこの式の右辺、特に高次元データでの「発散の計算」と積分を実行するための「ODEシミュレーション」が膨大な計算コストとなり、大きなボトルネックとなっていました。

フローマッチングは、まさにこの課題を克服するために生まれました。フローマッチングはこのコストのかかる尤度計算を学習プロセスから完全に切り離し、速度場そのものをより直接的かつ効率的に学習するアプローチを取ります。

第2章: フローマッチングの学習原理

前節で見たように、尤度を直接最大化する学習は困難を伴います。そこでフローマッチングは、より巧妙な方法で、望ましい速度場 u_t(x) をニューラルネットワークに学習させます。その鍵となるのが「条件付け」という考え方です。

2.1 条件付き確率パスの導入

我々が本当に知りたいのは、全てのノイズとデータの組み合わせを平均した、計算不可能な周辺確率パス p_t(x) と、それを生成する周辺速度場 u_t(x) です。 しかし、これらを直接扱うのは、「霧のようにぼやけた目標」を追いかけるようなもので、極めて困難です。

そこでフローマッチングでは、戦略を転換します。学習の時は答え(X1)を見ながら「ベクトルの向き」(速度)を覚える、処理は以下のように行います。

  1. 目的地を決める: データセットから本物の画像 x1 を一枚ランダムに選びます。
  2. 出発点を決める: ノイズ x0 をランダムに生成します。
  3. 中間地点を作る: X_t = (1-t)X_0 + tX_1 のように、出発点と目的地を使って中間地点 X_t を作ります。x_tx_0x_1に条件付けられる。
  4. 正解ベクトルを計算する: X_tx_1 を使って、単純な式で「正解」の条件付き速度 u_t(X_t|x_1)(例:X_1 - X_0)を計算します。

ニューラルネットワーク u_t^\theta に中間地点 Xt を見せて、「君がこの地点で予測すべき速度ベクトルは、この『正解ベクトル』だよ」と教えます。これがCFM損失 D(u_t(X_t|X_1), u_t^\theta(X_t)) の意味です。

学習時には 「目的地X1という答えをカンニング」 することで、あらゆる地点から目的地へ向かうための「正しい方向感覚」をモデルに徹底的に叩き込むのです。我らが計算出来ない u_t(x)より、簡単に計算できる条件付き速度 u_t(X_t|x_1) を代わりに学習やります。

要約すると「ある特定の目的地 x₁ に向かう、一本の明確な経路」だけを考えるのです。これが**条件付き確率パス p_t(x|x_1) ** の導入です。

周辺確率パス p_t(x) は、これらの無数の条件付きパスを、ターゲット分布 q(x₁) で平均(周辺化)したものと見なすことができます。

p_t(x) = \int p_t(x|x_1) q(x_1) dx_1

目的地である訓練サンプル x₁ を一つ固定することで、計算が困難だった問題が、扱いやすい問題へと変わります。この条件付き確率パスは、以下の境界条件を満たすように設計されます。

  1. t=0 (スタート地点): p_0(x|x_1) = p_0(x)
    • スタート時点では、目的地 x₁ の情報はなく、全てのパスはソース分布 p₀(例:ガウスノイズ)から出発します。
  2. t=1 (ゴール地点): p_1(x|x_1) = \delta(x_1)
    • ゴール時点では、全てのパスは目的地である x₁ の一点に収束している必要があります(\delta はデルタ関数)。

この条件付きパスを導入することで、学習の「正解ラベル」となるターゲット速度場が、次に示すように簡単に手に入るようになります。

2.2 生成速度場の導出

設計した条件付き確率パス p_t(x|x_1) を生成する 条件付き速度場 u_t(x|x_1) は、どうすれば求められるでしょうか。これもまた、条件 x₁ を固定した上で「連続の方程式」を満たす必要があります。 幸いなことに、多くの実用的なパス設計では、この条件付き速度場を解析的に導出できます。

最も代表的で強力な**線形補間パス(Optimal Transport, OT)**を例に見てみましょう。このパスでは、時刻 t のサンプル X_t は、ソースサンプル X_0 とターゲットサンプル X_1 の線形結合で定義されます。実際にSOTAモデルは使用している確率パスでもある。

X_t = (1-t)X_0 + tX_1

このとき、目的地 x_1 に向かうための条件付き速度場 u_t(X_t|x_1) は、以下のように導出できます。

  1. フローの時間微分を求める: このパスは、X_0 から X_t への条件付きフロー \psi_t(X_0|x_1) = (1-t)X_0 + t x_1 と見なせます。その時間微分(速度)は、単純に \dot{\psi}_t = x_1 - X_0 となります。

  2. モデルの入力変数で表現する: しかし、モデル 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} となります。

  3. 条件付き速度場を導出する: ステップ1の X₀ にステップ2の式を代入すると、最終的な条件付き速度場が得られます。

u_t(X_t|x_1) = x_1 - \frac{X_t - t x_1}{1-t} = \frac{x_1 - X_t}{1-t}

この式は非常に直感的です。

  • x_1 - X_t: 現在地 X_t から目的地 x_1 へ向かう方向ベクトルを示しています。
  • \frac{1}{1-t}: 時間の進行度を表す係数です。ゴールが近づく(t が1に近づく)ほど分母が0に近づき、速度が急激に増大します。これは、残されたわずかな時間で目的地に正確にたどり着くために「急加速」する必要があることを意味しています。

このようにして、我々はニューラルネットワークが学習すべき「正解」の速度ベクトルを、訓練サンプル x_1 とノイズ X_0 から簡単に計算できるようになりました。これが、シミュレーションフリーな学習を可能にするフローマッチングの核心です。

2.3 周辺化トリック:単純な「条件付き」学習から複雑な「周辺」学習へ

前節では、特定の目的地 x_1 を仮定することで、計算可能な条件付き速度場 u_t(x|x_1) を導出しました。しかし、これはあくまで個別の目的地へ向かうための一時的な「正解ラベル」に過ぎません。生成時、モデルは特定の目的地を知らない状態で、もっともらしい流れを自ら作り出さなければなりません。

モデルが最終的に学習すべき真の目標は、全てのノイズとデータの組み合わせを平均した、いわば「流れの平均的な振る舞い」である周辺速度場 u_t(x) です。本節では、なぜ単純な「条件付き」の学習で、この複雑な「周辺」の目標を達成できるのか、その鍵となる 周辺化トリック(Marginalization Trick) について解説します。

周辺速度場:真の学習目標

ある時空 (t, x) における周辺速度場 u_t(x) とは、直感的には、その点 x に到達しうる全ての目的地 x_1 の可能性を考慮し、それぞれの目的地へ向かうための条件付き速度 u_t(x|x_1) を平均したものです。

より厳密には、この平均は事後分布 p_{1|t}(x_1|x) で重み付けされます。事後分布 p_{1|t}(x_1|x) とは、「時刻 t に点 x にいるという観測のもとで、本来の目的地が x_1 であった確率」を意味します。これを数式で表すと以下のようになります。

u_t(x) = \mathbb{E}_{X_1 \sim p_{1|t}(\cdot|x)}[u_t(x|X_1)] = \int u_t(x|x_1) p_{1|t}(x_1|x) dx_1

この u_t(x) こそが、ニューラルネットワークが最終的に近似すべき、普遍的な「ルールブック」なのです。

立ちはだかる壁:計算不可能性

しかし、この u_t(x) を直接計算して学習のターゲットにしようとすると、大きな壁にぶつかります。事後分布 p_{1|t}(x_1|x) をベイズの定理で展開すると、

p_{1|t}(x_1|x) = \frac{p_t(x|x_1) q(x_1)}{p_t(x)}

となり、分母に計算不可能な周辺確率パス p_t(x) が現れてしまいます。このため、u_t(x) を直接計算することは事実上不可能であり、学習のターゲットとして利用することはできません。

フローマッチングの妙技:周辺化トリック

ここで、フローマッチングの最も巧妙かつ強力な理論的帰結が登場します。それは、この計算不可能な周辺速度場 u_t(x) を、学習ターゲットにする必要はない、という驚くべき事実です。

計算が簡単な条件付き速度場 u_t(x|x_1) をターゲットとして回帰学習を行うだけで、その結果として得られるモデルは、真の周辺速度場 u_t(x) を正しく学習することが理論的に保証されています。これが「周辺化トリック」です。

このトリックを支えるのが、原著で示されている以下の定理です。

定理(周辺化トリック):
もし、設計した条件付き速度場 u_t(x|z) が、対応する条件付き確率パス p_t(x|z) を正しく生成するならば、それらを平均して得られる周辺速度場 u_t(x) は、周辺確率パス p_t(x) を正しく生成する。

この定理は、我々が行う学習の正当性を保証してくれます。個々の単純なルール(条件付き速度場)がそれぞれの単純な経路(条件付きパス)を正しく生成するように学習させれば、それらのルールを脳内で(暗黙的に)平均した結果であるモデルの出力は、全ての経路を束ねた複雑な流れ(周辺パス)を正しく生成できるようになる、というわけです。

勿論、条件付きだからこそデータ量に依存するので非常に大きいなデータセットは必要になります。Stable Diffusion 3は1Bくらいの画像で事前学習されたらしくて、LuminaImagen2.0も100M以上のサンプルを含まれるデータセットを使ったらしい。

結論:シミュレーションフリー学習の理論的根幹

周辺化トリックにより、我々は扱いやすい「条件付き」の世界でモデルを訓練するだけで、本来の目的である複雑で計算不可能な「周辺」の世界のダイナミクスを捉えることができます。

これにより、学習プロセスにおいて p_t(x)u_t(x) の直接計算やシミュレーションを一切行う必要がなくなります。これこそが、フローマッチングが「シミュレーションフリー」な学習を実現し、高い効率性とスケーラビリティを達成できる理論的な心臓部なのです。

2.4 フローマッチング損失:学習の羅針盤

前節までで、計算が困難な「周辺速度場」の代わりに、扱いやすい「条件付き速度場」を導出できることを見ました。本節では、この条件付き速度場を「正解ラベル」として、ニューラルネットワーク u_t^\theta(x) を訓練するための損失関数、すなわちフローマッチング損失を定式化します。

理想的な損失と現実的な損失

我々が本当に最小化したい理想的な損失は、モデルの予測 u_t^\theta(X_t) と、真の(しかし計算不可能な)周辺速度場 u_t(X_t) との差です。これはフローマッチング(FM)損失と呼ばれます。

L_{FM}(\theta) = \mathbb{E}[ D(u_t(X_t), u_t^\theta(X_t)) ]

ここで D(u, v) は2つのベクトル uv の「距離」を測る尺度で、一般的にはBregmanダイバージェンスが用いられます。Bregmanダイバージェンスは、関数とその線形近似との差として定義される非対称な尺度です。最も一般的な例は、二乗誤差(MSE) で、これは D(u, v) = \frac{1}{2} ||u - v||^2 に相当します。説明は少し難しいかもしれません、ですが実装は普通のMSEでモデルの予測と真の条件付き速度の差を小さくするように学習させる。

しかし、前述の通り u_t(X_t) は直接計算できないため、このFM損失を学習に用いることはできません。そこで、我々は計算可能な条件付き速度場 u_t(X_t|X_1) をターゲットとする、実践的な条件付きフローマッチング(CFM)損失を導入します。

L_{CFM}(\theta) = \mathbb{E}[ D(u_t(X_t|X_1), u_t^\theta(X_t)) ]

この損失関数では、期待値 \mathbb{E} は、時刻 t、ソースサンプル X_0、ターゲットサンプル X_1 をそれぞれの分布からランダムにサンプリングすることで計算されます。X_tX_0X_1 から(例えば線形補間で)生成され、u_t(X_t|X_1) は解析的に計算できるため、この損失は容易に評価できます。

勾配の等価性:CFMがFMを達成する理由

ここで最も重要な理論的保証が登場します。それは、 「FM損失とCFM損失のパラメータ θ に関する勾配は等しい」 という驚くべき事実です。

L_{FM}(\theta) = L_{CFM}(\theta) + C

このCはパラメータ\thetaに依存しないので勾配は以下になります。

\nabla_\theta L_{FM}(\theta) = \nabla_\theta L_{CFM}(\theta)

フローマッチングの学習を実用的なものにするための核心的な結果です。この定理が意味するのは、計算が簡単な L_{CFM} の坂道を下っていく(最小化する)ことが、本来の目標である計算不可能な L_{FM} の坂道を下っていくことと、全く同じであるということです。

この「勾配の等価性」こそが、「周辺化トリック」をアルゴリズムとして実現する数学的なエンジンです。我々は、扱いやすい条件付き量をターゲットにして回帰学習を行うだけで、目的の複雑な周辺量を学習できることが、これにより保証されるのです。

第3章: フローマッチングの拡張性と柔軟性

フローマッチングのフレームワークは、その設計の柔軟性により、様々なタスクやデータ形式へと拡張することが可能です。本章では、条件付き生成への応用や、パス設計の多様な選択肢、そして拡散モデルとの深い関係性について掘り下げます。

3.1 条件付き生成への応用

これまでの議論は、主にデータ分布 q(x) 全体を学習する無条件生成を想定していました。しかし、フローマッチングの枠組みは、テキストや画像といった条件 y に基づいてサンプル x を生成する条件付き生成にも自然に拡張できます。

この場合、学習データは (x_1, y) のペアとなり、我々が学習するのは条件付きの確率パス p_t(x|y) を生成する速度場 u_t(x|y) です。ニューラルネットワークも、条件 y を追加の入力として受け取る u_t^\theta(x, y) へと変更します。

CFM損失も同様に拡張され、データサンプル x_1 とその条件 y の両方を用いてターゲットとなる条件付き速度場 u_t(X_t|x_1, y) を計算し、モデルの出力 u_t^\theta(X_t, y) との差を最小化するように学習を行います。

サンプリング時には、所望の条件 y をモデルに与えながらODEを解くことで、q(x|y) に従うサンプルを生成できます。これは、拡散モデルにおける分類器無しガイダンス(CFG)のような技術を、より一般的な形で実現するものと解釈できます。

3.2 条件付けの多様な選択肢

CFM損失を構築する際の「条件」の選び方には、いくつかのバリエーションがあり、タスクに応じて最適なものを選択できます(図3.3参照)。

  1. ターゲット条件付け (Z = X_1): これまで主に見てきたケースです。ノイズからデータを生成する一般的な生成タスクに適しています。
  2. ソース条件付け (Z = X_0): ソースサンプル X_0 を条件とし、そこからターゲット X_1 へ向かうフローを学習します。これは、ある画像 X_0 を別のスタイル X_1 へと変換するような、画像編集やスタイル変換タスクに応用できます。ですがQwen-Imageみたいな最近の画像編集モデルは普通のターゲット条件付けで学習させられた、編集をしたい画像は条件y、テキストと同じ扱いで入力をします。
  3. ペア条件付け (Z = (X_0, X_1)): ソースとターゲットのペアを固定し、その間を補間するフローを学習します。この場合、速度場は X_t には陽に依存せず、X_0, X_1, t の関数となります。これは「確率的補間(Stochastic Interpolant)」という関連研究と密接な繋がりがあります。

重要なのは、これらの条件付けの方法が適切に設計されていれば、「周辺化トリック」によって、最終的に学習される周辺速度場 u_t(x) は理論上同じものになるという点です。この柔軟性が、フローマッチングを多様な問題設定に適用可能にしています。

3.3 最適輸送との接続:最も効率的な経路の探求

確率パスの設計には無限の自由度がありますが、どのようなパスが「良い」パスなのでしょうか。その一つの理想的な指針を与えるのが、 最適輸送(Optimal Transport, OT) の理論です。

OTは、ある分布を別の分布へと輸送する際に、総コスト(例えば、移動距離やエネルギー)が最小となるような「最も効率的な輸送計画」を見つける数学理論です。これをフローマッチングに当てはめると、「サンプルがノイズからデータへと、できるだけまっすぐ、無駄なく移動する」ようなフローを学習することに対応します。

そして、理論的な解析により、2.2節で紹介した線形補間パス X_t = (1-t)X_0 + tX_1 は、この最適輸送の概念と深く関連していることが示されています。このパスは、CFM損失の上界を最小化するという意味で「最適」であり、学習を安定させ、サンプリング効率を向上させる効果が期待されます。そして前述の通り目標の条件付き速度は時刻tに依存しないので、直感的に全部のステップで同じ目標地点に向かうので速度予測が多目的(マルチタスク)学習にならないはずだ。

マルチタスクの問題を完全に解決出来る期待がないですが、せめて和らげることが出来るだろうか。実際にノイズに近いステップは最終的に生成される画像の影響は他のステップより高いので学習の時はUniformからサンプリングではなくてはLogit-Normal分布からサンプリングを行う。

このOTとの関連性は、フローマッチングの最もシンプルで基本的な損失関数 L_{CFM}(\theta) = \mathbb{E}[ ||(X_1 - X_0) - u_t^\theta((1-t)X_0 + tX_1)||^2 ] に強力な理論的裏付けを与えています。

3.4 拡散モデルとの架け橋:アフィンフローとガウスパス

線形補間パスは、より一般的なアフィン条件付きフローの特殊なケースと見なすことができます。最適輸送はこのフローの例の一つだ。

\psi_t(x_0|x_1) = \sigma_t x_0 + \alpha_t x_1

ここで \sigma_t\alpha_t は、t=0(\sigma_0, \alpha_0) = (1, 0)t=1(\sigma_1, \alpha_1) = (0, 1) となるように設計された時間依存のスカラー係数です。この一般化により、パスのダイナミクスをより柔軟に制御できます。

特に重要なのは、ソース分布 p_0 がガウス分布 \mathcal{N}(0, I) であるケースです。この場合、アフィンフローによって生成される条件付き確率パス p_t(x|x_1) もまた、ガウス分布となります。

p_t(x|x_1) = \mathcal{N}(x | \alpha_t x_1, \sigma_t^2 I)

この形式は、拡散モデル(DDPM)で用いられるノイズ付加過程と非常によく似た形をしています。実際に、\alpha_t\sigma_t のスケジュールを適切に選ぶことで、DDPMで標準的に用いられるVP(Variance Preserving)やVE(Variance Exploding)といったノイズスケジュールを再現できます。

この発見は、フローマッチングと拡散モデルの間に深い理論的な繋がりがあることを示唆しています。拡散モデルのスコアマッチング学習は、特定のガウスパスを仮定した フローマッチングの特殊な場合 として捉えることができます。一方で、フローマッチングは勾配(スコア)である必要のない 、より一般的なベクトル場を学習できるため、表現力の高い変換をモデル化できる可能性があります。

この関係性は、DDPMで事前学習したモデルの知識を、フローマッチングの枠組みでファインチューニングするアプローチの理論的な妥当性を強く支持するものです。DDPMが学習した「ノイズを除去する能力」は、FMにおける「ノイズからデータへ向かう速度場」の優れた初期値として機能する可能性が高いと考えられます。

第3章のまとめ:フローマッチングの動作原理と展望

本記事を通じて、フローマッチングの数学的基盤からその学習アルゴリズムの核心までを概観しました。

  • 理論的基盤: フローマッチングは、ODEで記述される決定論的なフローを用いて分布間を変換します。確率パスと速度場は「連続の方程式」で結びついており、これが学習の理論的正当性を保証します。
  • 学習の核心: 計算不可能な周辺速度場を直接学習する代わりに、計算可能な条件付き速度場をターゲットとするCFM損失を用います。周辺化トリック勾配の等価性により、この単純な回帰学習が、真の目標である周辺速度場の学習を達成することを保証します。
  • 柔軟性と拡張性: 条件付けの仕方や確率パスの設計に高い自由度があり、条件付き生成や画像編集など多様なタスクに適用可能です。特に、最適輸送に動機付けられた線形パスは、シンプルかつ強力な選択肢です。
  • 拡散モデルとの関係: アフィンフローとガウスパスの概念を通じて、フローマッチングが拡散モデルを包含する、より一般的なフレームワークであることが示唆されます。

Discussion