Flow Straight and Fast: Rectified Flowによるデータ生成と転送の学習

背景と学術的系譜

起源と学術的系譜

ある確率分布を別の確率分布へと変換する問題は、機械学習および統計学における基礎的な課題であり、一般に「輸送写像問題（transport mapping problem）」として知られている。歴史的に、この問題は最適輸送（Optimal Transport: OT）の分野から派生したものであり、分布間で質量を移動させるための最も効率的な手法を模索するものである。OTは厳密な数学的枠組みを提供する一方で、現代の画像生成やドメイン転送タスクに見られるような高次元空間においては、その解法が極めて困難であることが知られている。

従来のアプローチ、特にGenerative Adversarial Networks (GANs) やVariational Autoencoders (VAEs) といった生成モデルは、データ空間とLatent Space間の写像を学習することでこの問題の解決を試みてきた。しかし、これらのモデルは重大な課題を抱えている。GANは数値的な不安定性とモード崩壊（mode collapse）に悩まされ、VAEやその他の尤度ベースのモデルは、計算コストの高い複雑な推論手順を必要とすることが多い。より最近では、拡散モデル（diffusion models）や神経常微分方程式（neural ODEs）といった連続時間モデルが注目を集めている。これらは強力である一方、本質的に「無限ステップ」のプロセスであり、高コストなニューラルネットワークを繰り返し呼び出して複雑な微分方程式を解く必要があるため、リアルタイム応用や高速な推論には適さない。本論文の著者らは、これらの連続時間モデルの核心的な限界が、曲がった非直線的な軌道への依存にあることを見出した。この軌道は、正確にシミュレートするために多くの離散化ステップを必要とするのである。

直感的な専門用語

• Rectified Flow: 「高速道路を直線化する」プロセスと捉えるべきである。データ粒子が2つの分布間を曲がりくねった非効率な経路で移動するのではなく、最短の直線経路を辿るよう強制することで、移動を高速化し計算を容易にする手法である。
• Reflow: 初日に曲がりくねったルートを通った配送ドライバーが、交通状況を観察した後にルートを完璧な直線に「再流（reflow）」させる様子を想像されたい。前モデルが生成した経路を用いて反復的に学習することで、システムは自身の軌道を「直線化」し、極めて少ないステップ数で高品質な結果を得ることが可能となる。
• Coupling: 単なる「ペアリング計画」である。砂の山（分布 $\pi_0$）を特定の形状（分布 $\pi_1$）に移動させたい場合、Couplingとは個々の砂粒がどこへ行くべきかを指示する一連の命令セットである。
• Drift Force: ODEの文脈において、これはモデルの「ハンドル」に相当する。データ点が目的地に到達するために、任意の時刻 $t$ においてどの方向に移動すべきかを指示するニューラルネットワークである。
• Discretization Step: 動画の「フレームレート」と考えることができる。連続的な動きをシミュレートするために、それを小さな塊に分割する。ステップ数が多いほど滑らかだが低速なプロセスとなる。著者らは、非常に少ないステップ数（1ステップでさえも）で高品質な結果を得ることを目指している。

表記一覧

問題定義と制約

核心的な問題定式化とジレンマ

本論文は、高次元空間における2つの経験的なデータ分布 $\pi_0$ と $\pi_1$ 間の輸送写像を学習するという根本的な問題に取り組んでいる。これは、生成モデリング（例：ガウスノイズから画像への写像）やドメイン転送（例：あるスタイルから別のスタイルへの画像変換）など、様々な機械学習アプリケーションにおいて極めて重要なタスクである。

入力/現在の状態: 出発点は、$\mathbb{R}^d$ における2つの分布 $\pi_0$ および $\pi_1$ からの経験的な観測値（サンプル）である。この問題の重要な側面は、ペアとなる入出力データが欠如している点にある。すなわち、各サンプル $X_0 \sim \pi_0$ に対して、その「正しい」変換先や生成ターゲットであることが既知である $X_1 \sim \pi_1$ は存在しない。我々が持つのは、各分布からの独立したサンプルセットのみである。

出力/目標状態: 望ましい終着点は、無限データ極限において $Z_0 \sim \pi_0$ ならば $Z_1 := T(Z_0) \sim \pi_1$ となるような輸送写像 $T: \mathbb{R}^d \to \mathbb{R}^d$ を学習することである。より具体的には、本論文は $dZ_t = v(Z_t, t)dt$ という神経常微分方程式（ODE）モデルを学習し、可能な限り「直線的」な経路を辿って $\pi_0$ から $\pi_1$ へサンプルを輸送することを目指す。このODEは、新規データの生成やドメイン転送のために順方向にシミュレート可能でなければならない。

欠落したリンク/数学的ギャップ: 既存の手法の限界を克服しつつ、生成モデリングとドメイン転送を統合する、因果的かつ計算効率の高い輸送写像を、ペアのないデータからいかに構築するかが、まさに欠落したリンクである。

このギャップを埋めるための従来のアプローチは、いくつかの問題に直面していた。
1. 単純な線形補間: 単純な線形補間 $X_t = tX_1 + (1-t)X_0$ は直線的な経路を提供するが、「非因果的（先読み的）」である。$X_t$ を決定するために最終点 $X_1$ を知る必要があり、新規データ生成のための順方向シミュレーションが不可能である。
2. 最適輸送 (OT): OTは輸送コストを最小化する写像を見つけるための理論的に健全な枠組みを提供するが、高次元の連続測度に対しては「計算が極めて困難」であり、多くの機械学習タスクの特定の目的においては「直接的な関心の対象ではない」ことが多い。
3. 連続時間生成モデル (ODEs/SDEs): スコアベース生成モデルやDenoising Diffusion Probabilistic Models (DDPM) などの最近の進歩は印象的な結果を示している。しかし、これらのモデルは「実質的に『無限ステップ』」であり、ODE/SDEをシミュレートするために高コストなニューラルフォースフィールドを多数回呼び出す必要があるため、「推論時に高い計算コスト」を要する。

本論文は、この問題を単純な非線形最小二乗最適化として定式化することで、このギャップを埋めようと試みている。経験的にサンプリングされた点間の線形補間 $X_t = tX_1 + (1-t)X_0$ において、ODE $dZ_t = v(Z_t, t)dt$ が線形経路 $(X_1 - X_0)$ の方向を可能な限り忠実に辿るような速度場 $v(Z_t, t)$ を学習することを目指す。これは次のように表現される。
$$ \min_v \mathbb{E} \left[ \int_0^1 \|(X_1 - X_0) - v(X_t, t)\|^2 dt \right] $$
この定式化は、線形補間の直線経路を「因果化」し、シミュレーション可能にすることを目的としている。

制約と失敗モード

分布間の輸送写像を学習する問題は、いくつかの過酷で現実的な壁によって制約されている。

物理的、計算的、またはデータ駆動型の制約:
* ペアのないデータ: 教師なし学習環境における「ペアとなる入出力データの欠如」は、最も重大なデータ駆動型の制約である。これは、モデルが $X_0$ から $X_1$ への直接的な回帰を単純に学習できないことを意味する。
* データの高次元性: 現実世界のデータ、特に画像は非常に高次元の空間（$d$ が数百万に達することもある $\mathbb{R}^d$）に存在する。これにより、直接的な最適輸送計算が実行不可能となり、数値ODE/SDEソルバーの計算コストが増大する。
* ODE/SDEソルバーの計算コスト: 既存の連続時間モデルは、推論中に「高コストなニューラルフォースフィールドを多数回呼び出す」必要がある。これは多くのアプリケーションにおいて厳格なリアルタイムレイテンシ要件を課すことになり、数百から数千ステップで画像を生成することはあまりに低速である。
* ODEの非交差性: 適切に定義されたODEにおいて、その解は一意でなければならず、異なる経路が交差してはならない。これは、交差する可能性のある単純な線形補間とは異なり、学習されたフローが満たさなければならない基本的な数学的制約である。

なぜこのアプローチなのか

選択の必然性

著者らは、従来の生成モデル、特にGANと拡散モデルが、推論速度に関して根本的な「計算の壁」に突き当たっていることを特定した。GANは高速ではあるが、悪名高い学習の不安定性とモード崩壊に苦しむ。対照的に、拡散モデル（およびそのODEベースの変種であるPF-ODEs）は数学的に堅牢だが、精度を維持するために多くの離散化ステップを必要とする複雑で曲がった軌道を解く必要があるため、計算コストが高い。著者らは、これらの軌道の「曲がった」性質こそが主要なボトルネックであると認識した。もし2つの分布間の輸送経路を「直線」にできれば、ODEは最小限の離散化、あるいは単一ステップでさえも解くことができる。この認識により、焦点は単なる分布のマッチングから、それらの間の最短かつ最も直線的な経路を見つけることへとシフトした。

比較優位性

Rectified Flowは、輸送問題を単純でスケーラブルな制約のない最小二乗最適化へと変換するため、質的に優れている。繊細なミニマックスバランスを必要とするGANや、複雑なSDE/ODEソルバーに依存する拡散モデルとは異なり、Rectified Flowは「Reflow」手順を使用する。この手順は、フローの軌道を反復的に直線化するものである。構造的に、これは離散化誤差を大幅に低減する。標準的な拡散モデルが高品質な画像を生成するために数百回の関数評価（NFE）を必要とする可能性があるのに対し、Rectified Flow（特にReflow後）は、単一のEulerステップで同等以上の結果を生成できる。これは、1ステップモデル（VAEなど）と連続時間モデルの間のギャップを効果的に埋め、後者の高品質さと前者の高速さを両立させるものである。

制約への適合

問題の制約には、GANの不安定性や拡散モデルの法外な推論コストなしに、高次元データ（画像など）を扱えるモデルが必要であった。Rectified Flowは、輸送経路の「因果化」を通じてこれらの制約に適合している。ドリフトフォース $v$ が線形補間 $X_t = tX_1 + (1-t)X_0$ を辿るように学習することで、モデルは近視眼的かつ非交差的で決定論的な方法で質量を輸送することを学習する。ODEの枠組みと直線目標のこの「融合」は、モデルが計算効率（直線経路による）と理論的健全性（周辺分布を保持し、輸送コストを低減するため）の両方を備えていることを保証する。

数学的・論理的メカニズム

マスター方程式

Rectified Flowの核心的なメカニズムは、ソース分布 $\pi_0$ をターゲット分布 $\pi_1$ へと直線経路を辿って変換する速度場 $v(z, t)$ を学習することである。この速度場の学習に使用される目的関数は以下の通りである。

$$\min_{v} \int_{0}^{1} \mathbb{E} \left[ \left\| (X_1 - X_0) - v(X_t, t) \right\|^2 \right] dt, \quad \text{with } X_t = tX_1 + (1 - t)X_0$$

フローのステップ

1. 初期化: データ分布からペア $(X_0, X_1)$ がサンプリングされる。
2. 補間: システムは、ランダムにサンプリングされた時刻 $t$ における中間点 $X_t$ を計算する。
3. 速度予測: ニューラルネットワーク $v$ は、現在の状態 $X_t$ と時刻 $t$ を入力として受け取り、予測速度ベクトルを出力する。
4. 回帰: モデルは、予測された速度をターゲットの方向 $(X_1 - X_0)$ と比較する。
5. 更新: ネットワークパラメータは、差を最小化するために勾配降下法を通じて更新される。
6. 推論: サンプリング中、モデルは $Z_0 \sim \pi_0$ から開始し、数値ソルバー（Euler法など）を用いてODE $dZ_t = v(Z_t, t)dt$ を解き、$Z_1 \sim \pi_1$ に到達する。

最適化ダイナミクス

このメカニズムは、線形補間を「因果化」することで学習する。単純な経路 $X_t$ は未来（$X_1$）の知識を必要とするが、学習された速度場 $v(Z_t, t)$ は現在の状態と時刻のみの関数であり、妥当な因果的ODEとなっている。

「Reflow」手順は、重要な最適化ダイナミクスである。初期モデルを学習した後、そのモデルを使用して学習済みフローをシミュレートし、新しいペア $(Z_0, Z_1)$ を生成する。これらの新しいペアは、モデルの再学習に使用される。最初のモデルによって生成されたフローは、生のデータCouplingよりも既に「直線的」であるため、2回目の反復ではさらに直線的な経路が生成される。この反復プロセスはフローを効果的に「直線化」し、数値ソルバーの離散化誤差を低減する。その結果、損失地形はますます滑らかになり、モデルは非常に少ない（あるいは1つの）Eulerステップで高品質なサンプルを生成できる状態へと収束する。

結果、限界、および結論

実験計画とベースライン

著者らは、CIFAR-10データセットを用いた無条件画像生成、および高解像度データセット（LSUN, CelebA-HQ, AFHQ）を用いてRectified Flowを評価している。厳密なベースラインを確立するため、DDPM++フレームワーク（Song et al., 2020b）のU-Netアーキテクチャを採用した。実験計画は、「Reflow」手順の有効性と、学習されたODE軌道の「直線度」を検証するように構成されている。

エビデンスが証明するもの

提供されたエビデンスは、特にReflow手順の「直線化」効果に関して説得力がある。著者らは、初期の（1-rectified）フローは有効であるものの、完全には直線的ではないことを示している。前回のフローによって生成されたデータでモデルを再学習するReflow手順を適用することで、軌道はますます線形になる。

このメカニズムの決定的なエビデンスは2点ある。
* 定量的: CIFAR-10において、蒸留された2-rectified flowはFID 4.85を達成し、既知の最高性能の1ステップ生成モデル（TDPM, FID 8.91）を大幅に上回る。さらに、Recall 0.51はStyleGAN2+ADA（0.49）を超えており、この手法が高い多様性を維持していることを証明している。
* 視覚的/幾何学的: 図4および図18は、2-rectified flowの軌道がほぼ直線であることを視覚的に証明している。外挿 $\hat{z}_1^t = z_t + (1-t)v(z_t, t)$ は $t$ にかかわらずほぼ一定であり、これは直線ODEの特徴である。これは、モデルが輸送プロセスを正常に「因果化」し、最小限の離散化ステップで正確なシミュレーションを可能にしたことを裏付けている。

限界と今後の展望

本研究の今後の展望には以下が含まれる。
* 理論的洗練: 軌道直線化の利点が数値誤差の蓄積を上回る前に、Reflowステップ数に理論的な限界が存在するかどうかの探求。
* より広範な応用: 「直線化」特性が、物理システムモデリングや時系列予測など、非生成タスクにおいて活用できるかどうかの調査。
* 最適輸送の統合: 著者らが言及するように、Rectified Flowは特定のコスト関数 $c$ に対する $c$-最適輸送を厳密に保証するものではない。今後の研究では、速度場 $v$ を勾配場（例：$v = \nabla f$）に制約し、最適性を明示的に強制することに焦点を当てる可能性がある。

これらの知見は、生成モデリングにおけるパラダイムシフトを示唆している。「ノイズからデータへ」という拡散パラダイムから、「直線」輸送パラダイムへの移行であり、これは計算効率が高く、理論的にもより透明性が高い。