Flow Matching

フローマッチングは、学習した速度場に従ってランダムノイズをデータに変換する生成モデルの学習手法です。ニューラルネットワークは、時間0から1の任意の時点と任意の点について、その点が進むべき方向と速度を学習します。サンプルの生成では、ガウスノイズから開始し、常微分方程式（ODE）ソルバーでその速度を時間1まで積分します。Meta AIとWeizmann Institute of ScienceのYaron Lipmanらが、ICLR 2023でFlow Matching for Generative Modelingを発表しました。また、同時期の2本の論文は、整流フローと確率的補間という名称で、ほぼ同じ目的関数に到達しました。フローマッチングは学習時にODEをシミュレーションせずに連続正規化フローを学習でき、拡散モデルを特殊ケースとして含みます。また、Stable Diffusion 3、FLUX.1、MetaのMovie Genをはじめ、近年の音声モデルやロボティクスモデルの多くで学習目的関数として使われています。

フローマッチングの仕組み#

学習中に生成は行いません。各ステップでは、データサンプルx1、標準ガウス分布からノイズサンプルx0、0から1までの時刻tを抽出し、両者を結ぶ直線上の点xt = (1 - t) * x0 + t * x1を設定します。その直線上の点は一定速度x1 - x0で移動するため、ネットワークv(xt, t)は単純な平均二乗誤差でx1 - x0を予測するよう学習されます。ノイズとデータの各ペアが定める経路は1つだけですが、任意の点を通る経路は多数あり、その点でのノイズからデータへのフロー全体の速度は、それらの平均です。この論文は、単純なペアごとの目標値に対する回帰が、計算困難な平均値に対する回帰と同じ勾配を与えることを証明しています。そのため学習は、ODEシミュレーションを必要としない通常の回帰損失関数になります。

生成時には、通常はオイラー法または高次ソルバーによるステップを使って、t = 0からt = 1までdx/dt = v(x, t)を積分します。このPyTorchコードは、2つの2D点群で小さな速度ネットワークを学習させた後、オイラー法のステップ数を1、2、8としてサンプリングします。

import torch
from torch import nn

torch.manual_seed(0)
# Velocity network v(x, t) for 2D points; the data are two clusters at (2, 2) and (-2, -2)
net = nn.Sequential(nn.Linear(3, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
centers = torch.tensor([[2.0, 2.0], [-2.0, -2.0]])

for step in range(2000):
    x1 = centers[torch.randint(0, 2, (256,))] + 0.2 * torch.randn(256, 2)  # data samples
    x0 = torch.randn(256, 2)  # noise samples
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # point on the straight path from noise to data
    loss = ((net(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the path's velocity
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate: start from noise and follow the learned velocity field with Euler steps
for steps in (1, 2, 8):
    x = torch.randn(1000, 2)
    with torch.no_grad():
        for i in range(steps):
            x = x + net(torch.cat([x, torch.full((1000, 1), i / steps)], dim=1)) / steps
    near = (torch.cdist(x, centers).min(dim=1).values < 0.6).float().mean()
    print(f"{steps}-step sampling: {near:.0%} of samples land near a cluster")

出力は1-step sampling: 0% of samples land near a clusterとなり、2ステップでは47%、8ステップでは94%になります。t = 0では、すべてのノイズ点の期待される到達先がデータの平均値で同じになるため、1ステップではすべてのサンプルが2つのクラスタ間の何もない領域に送られます。各学習経路は直線ですが、近くから始まった経路が異なるクラスタに到達し、それらを平均した学習済みフローは曲線になります。そのため、少ないステップで動作する生成器では、モデル自身のノイズとサンプルのペアで再学習してフローを直線化するリフローや、一貫性モデルへの蒸留などの手法が追加されます。

フローマッチングと関連手法の比較#

• 拡散モデル： 拡散学習では、ノイズスケジュールに沿ってデータを破損させ、追加されたノイズをネットワークが予測するよう学習させ、その過程を逆転させてサンプリングします。フローマッチングの論文は、その目的関数が、拡散モデルで使われる経路を含む任意のガウス確率経路に対応し、拡散モデルをより安定して学習できることを示しました。直線経路は、それとは別のより単純な選択肢です。両方の手法で、拡散トランスフォーマー（DiT）、分類器なしガイダンス、ODEサンプラーなどのアーキテクチャが使われます。
• 整流フロー： 同じ直線補間と速度損失を使うため、線形経路では2つの名称はほぼ同義です。整流フローの論文では、軌跡をさらに直線化するリフロー手順が追加されています。
• 正規化フロー： 離散正規化フローは可逆層を重ね、厳密な尤度で学習します。ニューラルODEの一種である連続正規化フローは、ODEで変換を定義しますが、最尤法で学習され、各学習ステップでODEを解く必要がありました。フローマッチングは同じ種類のODEモデルを回帰で学習し、その計算を必要としません。
• オプティカルフロー： 名称が似ていますが、関連はありません。オプティカルフローは動画フレーム間におけるピクセルの見かけ上の動きであり、フローマッチングはデータの可能空間内でサンプルを移動させます。

フローマッチングの用途#

• 画像生成： Stability AIは整流フロートランスフォーマーを使ってStable Diffusion 3を学習させ、中間的なノイズレベルを重視するロジット正規分布から学習時刻をサンプリングしました。Black Forest LabsのFLUX.1は、パラメーター数120億の整流フロートランスフォーマーです。こうしたテキストから画像を生成するモデルの画像は、YOLO26検出モデルの学習用データセットに希少な物体、天候、視点を加える合成データとして活用されます。
• 動画生成： MetaのMovie Genは、フローマッチングで学習されたパラメーター数300億のトランスフォーマーで、最大16秒、毎秒16フレームの1080p動画を生成します。
• 音声： MetaのVoiceboxは、50,000時間を超える音声データで非自己回帰型フローマッチングモデルを学習しました。ゼロショットのテキスト読み上げでは、単語誤り率が1.9%となり、VALL-Eの5.9%を下回る一方、最大20倍高速に動作しました。その後、F5-TTSはフローマッチングと拡散トランスフォーマーを組み合わせました。
• ロボット制御： Physical Intelligenceのπ0は、パラメーター数30億の視覚言語モデルに、パラメーター数3億のフローマッチング「アクションエキスパート」を追加します。積分ステップ10回で未来のアクションを50個ずつチャンクとして生成し、最大50 Hzでロボットを制御します。