Rectified Flow

Rectified flowは、単純な分布（通常はガウスノイズ）からデータ分布へサンプルを運ぶ速度場を学習する生成モデリング手法です。その際、できるだけ直線的な経路をたどります。学習は単純な回帰で行います。ノイズサンプルとデータサンプルをペアにし、その直線上の点を選び、ニューラルネットワークに直線の方向を予測させます。生成時には、ニューラルODEと同じ仕組みである常微分方程式（ODE）ソルバーを使って、学習済み速度を積分します。経路が直線的であれば、少数の大きなステップで追跡できます。テキサス大学オースティン校のXingchao Liu、Chengyue Gong、Qiang Liuは、2022年9月にFlow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowでこの手法を発表しました。また、MetaのYaron Lipmanらによる密接に関連するフローマッチングも同時期に発表されました。

Rectified flowの学習方法#

ノイズx0とデータx1に対して、点x_t = (1 - t) * x0 + t * x1は、一定速度x1 - x0で、t = 0のノイズからt = 1のデータへ移動します。ネットワークv(x_t, t)は、ランダムなペアと時刻を用い、二乗誤差の損失関数によってその速度を予測するよう学習します。任意の点を通る直線は複数あるため、ネットワークはそれらの平均方向を学習します。ネットワークが定めるODE経路は互いに交差せず、異なるデータ点へ向かう直線が出会う場所で曲がります。PyTorchによる1次元のrectified flowでは、ガウスノイズを-2と+2の2つのクラスターに変換するよう学習します。

import torch
from torch import nn

torch.manual_seed(0)
velocity = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(velocity.parameters(), lr=1e-3)

for step in range(3000):
    x0 = torch.randn(256, 1)  # noise
    x1 = 2.0 * torch.sign(torch.randn(256, 1)) + 0.1 * torch.randn(256, 1)  # data: two modes at -2 and +2
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # a point on the straight line from noise to data
    loss = ((velocity(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the line's direction
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate by Euler integration of dx/dt = v(x, t), with different step counts
for steps in [1, 2, 4, 32]:
    x = torch.randn(10000, 1)
    with torch.no_grad():
        for i in range(steps):
            x = x + velocity(torch.cat([x, torch.full_like(x, i / steps)], dim=1)) / steps
    near_mode = ((x.abs() - 2).abs() < 0.3).float().mean()
    print(f"{steps:>2} Euler steps: {near_mode:.0%} of samples within 0.3 of a mode")

オイラー法のステップが1回の場合、モード付近に到達するサンプルは0%です。t = 0では、すべてのノイズ点からの平均方向が、2つのクラスターの中間にあるデータ平均0へ向かいます。ステップを2回にすると、サンプルの39%がモードから0.3以内に入り、4回では81%、32回では96%になります。これは、学習された経路がクラスターに向かって曲がり、その曲線をたどるには小さなステップが必要だからです。

Reflow：経路の直線化#

論文で再帰的整流と呼ばれるReflowは、こうした曲がりを取り除きます。学習済みフローでノイズとサンプルのペアを生成し、そのペアを使って新しいフローを学習します。ランダムなペアの代わりに、このペアを使います。各ペアは最初のフローの経路ですでに結ばれているため、ペア間の直線が衝突することはほとんどありません。また論文では、新しい結合によって輸送コストが増加することはないと証明しています。Reflowを繰り返すと経路はほぼ直線になり、著者らはオイラー法の1ステップで高品質な画像生成と画像変換ができることを報告しました。

同じチームは知識蒸留をStable Diffusion 1.5に適用し、InstaFlow（2023年）を開発しました。1ステップモデルはMS COCO 2017-5kでFID 23.3を達成し、漸進的蒸留の37.2を上回りました。画像生成にかかる時間は0.09秒で、学習にはA100 GPUを199日分使用しました。

テキストから画像を生成するモデルにおけるRectified flow#

Stability AIのStable Diffusion 3論文『Scaling Rectified Flow Transformers for High-Resolution Image Synthesis』（2024年3月）では、複数の学習方式を比較し、学習時刻をlogit-normal分布からサンプリングするrectified flowを採用しました。この分布は、速度予測が最も難しい経路の中間部分に学習を集中させます。このモデルは、潜在空間で動作するマルチモーダルな拡散Transformer（DiT）と、この目的関数を組み合わせています。元Stability AIの研究者が設立したBlack Forest Labsは、FLUX.1 devモデルを120億パラメーターのrectified flow Transformerと説明しています。また、蒸留版のFLUX.1 [schnell]は1～4ステップで生成します。これらのテキストから画像を生成するモデルは、ビジョン学習用の合成データの一般的な供給源です。Ultralytics Platformでは、類似画像を生成する機能をデータセットのキュレーションに組み込んでいます。データセット内の画像を1枚指定すると、新しいシーンの説明を作成し、最大16種類のバリエーションを生成して、採用した画像をラベルなしでデータセットに追加し、アノテーションに利用できる状態にします。

Rectified flowと関連用語の比較#

• フローマッチング：選択した確率経路に沿った速度を回帰することで、ODE生成器を学習する一般的なフレームワークです。独立したノイズとデータの間を直線で結ぶrectified flowは、最も広く使われる経路であり、両者の名称はしばしば同じ意味で使われます。
• 拡散モデル：スケジュールに沿ってノイズを加えます。確率フロー経路は曲線を描き、決定論的サンプラーでは通常20～50ステップかかります。
• 正規化フロー：学習済みの可逆層を固定された順序で積み重ね、ノイズをデータに変換します。厳密な最尤法で学習します。
• 整合性モデル：曲線状の軌道上の任意の点から終点まで直線的に移動する方法を学習するため、ネットワークの呼び出し1回でサンプルを生成できます。