0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@suujyou3

Rectified Flow 入門

0
Posted at

はじめに

この記事は、Rectified Flowについて学ぶために、作成した個人的なノートの一部を整理したものです。

内容は論文や公開資料をもとに、Codexを活用しながらまとめていますが、内容や構成の確認、整理は著者自身が行っています。

現在、ネット上にはあまりRectifiedFlowの解説がありません。本記事は、入門のために理論に重きを置かず、直観的な理解に重きを置いています。

数学が全くいらないということではないですが、偏微分方程式とか、学習の理論について詳しくなることではなく、ざっくり理解して、コードを読める/書けるようにするという目的で作りました。

また、本資料の全文、コードはGitHubにあげています。

想定読者

本資料は、Rectified Flow を学びたい方を対象としています。

特に、

深層学習の基礎を一通り学んだ方
論文を読む、PyTorchで実装を行うための足掛かりとなる解説を探している方
数式だけでなく、直感的な理解も得たい方
を主な対象としています。

厳密な理論よりも、「なぜその式になるのか」「どのような意味を持つのか」を重視して説明しています。

ただし、数学については以下を前提としています。

  • 偏微分について、理解している
  • 微分方程式が読める
  • 統計記号が読める

また、本記事は現代的な見方でFlow MacthingとRectified Flow の関連性についても説明しています。このため、過去に書いたFlow Matchingの解説記事を読むことをおすすめします。

特に重要なのは、(1)、(2)、(3)です。pathやCFMを理解していることが前提です。

Rectified FlowはFlow Matchingの何を変えるのか

この章で扱う概念: 現在の統一的な見方では、Flow Matchingは「pathと教師速度を決め、速度場を学ぶ」一般的な枠組みとして読める。Rectified Flowは、直線pathを使った速度回帰から始め、生成時の実軌道をより直線的にして少ないNFEでも追いやすくすることを重視する関連手法である。1-Rectified Flowのlossだけを見ると、本教材のLinear Path CFMと同じ形である。違いは、学習後の軌道直線性を評価し、そのflowが作るcouplingでペアを作り直して再学習するReflowまで含めて設計する点にある。

まず結論: lossを別物にするのではない

Flow MatchingとRectified Flowの違いは、最初に次の三層へ分けると読みやすい。ただし、これは研究史の順番を表す説明ではなく、本教材で概念を整理するための見方である。研究史については改めて述べる。

  1. 広い枠組みとしてのFlow Matching
    Flow Matchingは、確率経路 $p_t$ と教師速度を設計し、その速度をニューラルネットワーク $v_\theta(t,x)$ で回帰する一般的な学習方法である。pathは直線でもよいし、Gaussian pathでもよい。couplingも独立coupling、OT coupling、ほかの設計を選べる。

  2. 1回目の学習としての1-Rectified Flow
    1-Rectified Flowは、端点ペア $(x_0,x_1)$ を直線で結び、教師速度 $x_1-x_0$ を学ぶ。この部分だけを見ると、本教材のLinear Path CFMとほぼ同じである。したがって、ここで新しいlossが突然出てくるわけではない。

  3. 研究上の狙いとしてのRectified Flow/Reflow
    Rectified Flowで重要なのは、学習済みODEの実軌道を直線に近づけ、少ないNFEでも生成しやすくすることである。そのために、学習済みflowでsource点を実際に運び、到着点を使って新しいペアを作り直す。この再ペアリング手続きがReflowである。

つまり、Flow Matchingとの差分は「速度回帰をやめること」ではない。同じ速度回帰の形を使いながら、どのcouplingで学ぶか、学習後の軌道をどう診断するか、次の学習ペアをどう作るかを設計対象にすることである。

ReflowなしのRectified FlowとLinear Path CFMを、数式だけで明確に区別することは難しい。どちらも、端点を直線で結び、その直線速度をニューラルネットワークで回帰する形として書けるからである。違いは、主に研究上の文脈と、その後に何を問題として扱うかに現れる。

見方 主な役割
Linear Path CFM 線形の条件付き経路を使うFlow Matchingの学習目的
Rectified Flow 同じ形の速度回帰を、輸送のrectification、軌道の直線化、coupling更新、Reflowへ展開する方法論

したがって、本文で「Rectified FlowはLinear Path CFMと同じ形で読める」と言うとき、それは数式上の初回学習に関する話である。一方、「Rectified Flowらしさ」と言うときは、輸送をまっすぐにする、少ないNFEで生成しやすくする、学習済みflowでペアを作り直す、という研究上の狙いまで含めて読んでいる。

本教材の最小Linear Path CFMと1-Rectified Flowを比べると、最初の学習では多くの部品が同じである。

観点 Linear Path CFM 1-Rectified Flowで同じ点 Rectified Flow/Reflowで加わる視点
path 端点を直線で結ぶ 同じ直線補間を使う 直線pathが作る学習済み軌道が本当に直線的かを見る
教師速度 $x_1-x_0$ 同じ教師速度を使う 矛盾する矢印が減るように、次のペアを作り直す
loss 速度のMSE 1回目は同じMSE loss式よりも、coupling更新と少NFE生成を重視する
velocity model $v_\theta(t,x)$ 入出力は同じ モデル名ではなく、軌道直線性の実験として読む
生成 ODE solverで速度場を積分する 同じODE生成を使う 少ないNFEで品質が保てるかを評価する
次の学習ペア 通常は元のcouplingを使う 1回目は元のペアでよい 学習済みflowが作るcouplingへ更新する

したがって、「Rectified FlowはCFM lossを新しい式へ交換した手法」と理解するとずれる。1回目の学習はLinear Path CFMと同じ形である。そのうえで、生成時の軌道が曲がる原因をcouplingの作り方に戻って考え、Reflowでペアを更新するところがRectified Flowらしい読みどころである。

研究史: 包含関係は後から整理された見方である

「Flow Matchingという一般枠組みが先にあり、そこからRectified Flowが特殊例として派生した」と読むと、研究史としては単純化しすぎである。

実際には、Rectified Flowを提案した "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow" と、Flow Matchingを一般的な学習原理として提示した "Flow Matching for Generative Modeling" は、どちらも2022年にarXivで公開された非常に近い時期の関連研究である。Rectified Flow論文は、輸送と軌道の直線化、Reflowによるrectificationを中心にしている。一方、Flow Matching論文は、確率経路と速度場回帰をより一般的な生成モデリング原理として整理している。

そのため、本教材で「Rectified FlowをFlow Matchingの直線path版として読める」と説明するのは、後から見た統一的な整理である。歴史的に「Flow Matchingが先にあり、その特殊例としてRectified Flowが生まれた」と主張しているわけではない。

研究用に読むときは、次の二つを分ける。

統一的な理解:
  Rectified Flowの初回学習は、Linear Path CFMと同じ形で書ける

研究史・方法論:
  Rectified Flowは、輸送のrectificationと軌道直線化を中心に発展した
  Flow Matchingは、より一般的な速度場学習原理として整理された

この区別を入れると、同じ数式が出てくるのに論文名や実装名が違う理由を誤解しにくくなる。

直感: 練習用の直線と実際に走る道は違う

source点とtarget点を直線で結べば、ペアごとの教師矢印はまっすぐになる。しかし、多数の直線が途中で交差すると、同じ場所に異なる向きの教師矢印が集まる。velocity modelはそれらを平均して学ぶため、生成時にsolverがたどる実軌道は曲がることがある。

道路にたとえると、最初のFlow Matchingは、ばらばらに選んだ出発地と目的地を定規で結び、その線上の矢印を大量に見せる学習である。線が交差する交差点では、モデルは複数の進行方向を一つにまとめなければならない。その結果、完成した道路は、個々の定規の線より曲がることがある。

Rectified Flowは、いったん完成した道路を実際に走ってみて、「この出発点は、このflowではどこへ到着したか」を新しい対応関係として使う。その対応で再び直線の教師を作ると、互いに矛盾する矢印が減り、実際の道路も直線に近づきやすくなる。このペアの作り直しが次章の Reflow である。

1回目:
独立に作ったペア
  -> 直線の教師矢印で学習
  -> 学習済みflowの実軌道は曲がることがある

次の段階:
学習済みflowでsourceを実際に運ぶ
  -> flowが作った到着点を新しいペアにする
  -> そのペアで再学習し、実軌道の直線化を狙う

ここでいう「まっすぐ」は、訓練用pathが直線であることだけではない。学習済み速度場で粒子を生成したときの軌道が、遠回りせず端点へ向かうことを意味する。

コードで見る: 最初のlossは同じである

第1部のLinear Path CFMと、1回目のRectified Flow（1-Rectified Flow）は、どちらも

x_t=(1-t)x_0+tx_1,\qquad u_t=x_1-x_0

を使った速度回帰である。コード上でも、二つのlossは同じ処理として読める。

# Linear Path CFMとして読む場合
loss_fm = conditional_flow_matching_loss(
    model, LinearPath(), x0, x1
)

# 1-Rectified Flowとして読む場合
loss_rf = rectified_flow_loss(model, x0, x1)

rectified_flow_loss は内部でLinear PathのCFM lossを呼ぶ。RectifiedFlowMLP も構造としては MLPVelocity と同じ小さなMLPである。名前を分けているのは、新しいモデル構造を導入したからではなく、軌道直線性とReflowを重視する実験であることを明示するためである。

src/train_rectified_flow_2d.py は、1-Rectified Flowを2D toyデータで学習するスクリプトである。

python .\src\train_rectified_flow_2d.py --steps 2000 --batch 512 --device cpu

モデルは RectifiedFlowMLP である。

model = RectifiedFlowMLP()
loss = rectified_flow_loss(model, x0, x1)

学習後は、生成サンプルだけでなく軌道を取り出し、straightness_ratio を計算する。

trajectory = euler_solve(
    lambda t, x: model(t, x),
    x0,
    steps=64,
    return_trajectory=True,
)
ratio = straightness_ratio(trajectory)

つまり、最初のコード変更の中心はlossではなく、軌道を保存し、直線性を評価し、次のReflowへ渡せるようにすることである。

直感: Reflowは「学習後にペアを作り直す」

通常のFlow Matchingでは、最初に用意したsource点とtarget点のペアを使って学習する。たとえば独立couplingなら、sourceから取った点とtargetから取った点を、かなり機械的に結ぶ。そのペアに対して直線pathを作り、各時刻で「こちらへ進め」という教師矢印を与える。

この方法は単純で扱いやすい。一方で、ペアの結び方が悪いと、途中で多くの直線が交差する。同じ場所の近くに、右へ行く矢印、左へ行く矢印、上へ行く矢印が集まる。ニューラルネットワークはそれらを一つの速度場として学ぶため、生成時の粒子は平均された矢印に従って曲がった道を進むことがある。

Reflowの発想は、ここで一度立ち止まることである。最初の学習でできたflowを実際に使い、source点を生成側へ流してみる。

ここで「流す」とは、source点 $x_0$ を学習済み速度場 $v_\theta(t,x)$ の上に置き、生成時と同じようにODE solverで少しずつ動かす、という意味である。たとえばEuler samplerなら、現在位置で速度を予測し、その向きへ少し進む。この更新を $t=0$ から $t=1$ まで繰り返すと、source点はある終点へ到着する。Reflowでは、この終点を「このsource点が、今の学習済みflowで到着した場所」として扱う。

すると、「このsource点は、この学習済みflowではこの場所へ到着した」という対応が得られる。Reflowでは、その対応を次の訓練ペアとして使う。

通常のFlow Matchingとの違いは、次のように整理できる。

観点 通常のFlow Matching Reflowを使うRectified Flow
ペアの作り方 学習前に決めたcouplingを使う 学習済みflowを生成時と同じようにODE solverでたどり、到着点を使ってペアを作り直す
学習で変えるもの 主に速度場のパラメータ 速度場に加えて、次の学習に使うペアの対応関係も変える
直線pathの意味 与えられたペアを直線で結ぶ 学習済みflowが作った対応を直線で結び直す
期待する効果 target分布へ運ぶ速度場を学ぶ 矛盾する矢印を減らし、生成時の実軌道をより直線的にする

この手順の違いを図6に示す。

通常のFlow MatchingとReflowを使うRectified Flowの学習方法の違い

図　通常のFlow MatchingとReflowを使うRectified Flowの学習方法の違い

図の読み方: 左は通常のFlow Matchingである。学習前に決めたsourceとtargetのペアを使い、そのペアから直線pathと教師速度を作って速度場を学ぶ。生成軌道が曲がっても、通常はその結果を使って次の学習ペアを作り直さない。右はReflowを使うRectified Flowである。まず1回目の速度場を学習し、その学習済みflowを生成時と同じようにODE solverでたどる。得られた到着点を使ってペアを作り直し、その新しいペアで次の速度場を学ぶ。したがって、違いはloss式ではなく、次の学習に使うcouplingを更新するかどうかにある。

重要なのは、Reflow後も「直線pathで速度を学ぶ」という学習の形そのものは大きく変わらない点である。変わるのは、どのsource点とどの到着点をペアとして結ぶかである。つまりReflowは、新しいlossを導入するというより、学習済みflowを使って次の教材データを作り直す操作である。

この節では直感だけを押さえればよい。数式としてのmarginal preserving propertyは次節で扱う。

理論: marginal preserving property

Rectified Flowの理論論文で重要なのは、marginal preserving property、つまり「各時刻の周辺分布を保つ」という見方である。

まず、任意のcouplingからペア $(X_0,X_1)$ を取り、直線補間

X_t = (1-t)X_0 + tX_1

を考える。このとき $X_t$ の分布を $p_t$ と書く。

Rectified Flowは、条件付き速度 $X_1-X_0$ をそのまま全点に貼るのではなく、同じ中間点 $x$ に来る可能性のあるペアの速度を平均して、周辺速度場を作る。

v(t,x) = \mathbb{E}[X_1 - X_0 \mid X_t=x].

この式は、本教材で何度も出てきた「二乗誤差は条件付き平均を学ぶ」という話と同じである。ニューラルネットワーク $v_\theta(t,x)$ は、この $v(t,x)$ を回帰で近似している。

この速度場でODE

\frac{dZ_t}{dt}=v(t,Z_t),\quad Z_0 \sim p_0

を解くと、理想的には $Z_t$ の周辺分布が $p_t$ と一致する。特に $t=1$ では、$Z_1$ の分布はtarget分布 $p_1$ になる。

ここが重要である。Rectified Flowは、個々のサンプル軌道を元の直線補間と完全に同じにする方法ではない。多くの直線補間が作る時刻ごとの分布を保つように、平均速度場でODEを作る方法である。したがって、「訓練ペアごとの直線」と「学習済みODEの粒子軌道」は区別して読む必要がある。

理論と実用の接続: なぜ少ないステップと関係するのか

Euler samplerは、短い直線の矢印をつないでODEを近似する。

x_{k+1} = x_k + \Delta t\, v_\theta(t_k, x_k).

本当の軌道が曲がっていると、少ないステップでは曲がり角を取り逃がす。逆に、軌道がほぼ直線なら、粗いステップでも端点へ近づきやすい。

Rectified Flowの「straight is fast」という考え方はここにある。速度場の軌道がまっすぐなら、NFEを減らしても生成品質が落ちにくい。研究上は、これはsampling高速化のための設計原理になる。

ただし、注意点がある。訓練時に直線pathを使っていても、学習された周辺速度場の軌道が完全な直線になるとは限らない。多くのペアから来る教師速度を平均するため、同じ場所の近くで複数方向の速度が混ざることがある。

この注意は、marginal preserving propertyと矛盾しない。周辺分布を正しく運ぶことと、各粒子の軌道が完全に直線であることは別の性質である。Rectified Flowの研究では、この2つを混同しないことが大切である。

Reflow: 学習済みflowでペアを作り直す

この章で扱う概念: Reflowは、学習済みRectified Flowでsource標本を実際に輸送し、その始点と終点を新しい学習ペアとして次のRectified Flowを学ぶ反復操作である。モデル構造の名前ではなく、学習データとなるcouplingを更新する手続きである。

直感の地図

Reflowは、学習済みのflowを使って、新しい訓練ペアを作り直す操作である。

最初の学習では、sourceから点を取り、targetから点を取り、両者を結んだ。この段階のペアは、データの本当の対応関係を知っているわけではない。多くの場合、「sourceのこの点とtargetのこの点を結ぶことにする」という学習上の仮の対応である。

Reflowでは、まず学習済みモデルでsource点を実際に流す。すると、モデル自身が「このsource点なら、生成過程ではこのあたりへ到着する」と対応を作る。その到着点を、次の学習で使う新しい相手として扱う。

通常のFlow Matchingとの違いは、ペアを決めるタイミングにある。

通常のFlow Matching:
  学習前にペアを決める
  -> そのペアを直線で結ぶ
  -> 速度場を学ぶ

Reflow:
  まず一度flowを学ぶ
  -> 学習済みflowでsource点を実際に運ぶ
  -> 到着点を使ってペアを作り直す
  -> もう一度、直線Flow Matchingを学ぶ

直感的には、最初のモデルが「このsource点はだいたいこのあたりへ行く」と教えてくれる。その対応を使ってペアを結び直すことで、次の学習ではより素直な直線を教師にできる。ここで変わるのはlossの式ではなく、訓練データとして見せる「出発点と到着点の組」である。

ReflowとOTの関係

Reflowは、学習済みODEが作る写像を使って、新しいcouplingを誘導する。

最初の学習済みモデルで、source側の点を $t=0$ から $t=1$ まで流す。

\frac{dX_t}{dt} = v_\theta(t, X_t),\quad X_0 = x_0.

その終点を、新しいtarget側の点として扱う。

Z_0 = X_0,\quad Z_1 = \mathrm{ODESolve}(v_\theta, X_0).

このペア $(Z_0,Z_1)$ は、source側の周辺分布を保つ。理想的なRectified Flowなら、target側の周辺分布も保つ。つまり、Reflowは「target分布を捨ててモデルの出力だけを信じる」操作ではなく、分布を運ぶflowが誘導したcouplingを次の訓練に使う操作である。

Qiang Liuの "Rectified Flow: A Marginal Preserving Approach to Optimal Transport" で強調されるのは、この操作がOTと関係する点である。大まかには、Rectified Flowは任意の初期couplingから出発し、marginalを保ったまま、より輸送コストの小さい、より直線的なcouplingへ近づけることを狙う。

ただし、ここで「Rectified Flowは常に真のOT mapを学ぶ」と読んではいけない。有限データ、有限モデル、有限ステップのsolver、最適化誤差がある実装では、得られるのはOTに関係した近似的なcouplingである。本教材では、Reflowを次のように読む。

理論上の役割: marginalを保つflowから、新しいcouplingを誘導する
実装上の役割: 学習済みモデルでsourceを流し、次の訓練ペアを作る
実験上の確認: NFE、生成品質、軌道直線性、計算コストを見る

この整理を入れると、OT-CFMとの違いもはっきりする。OT-CFMはサンプル間の距離コストからcouplingを選ぶ。Reflowは、学習済みflowが作った到達点からcouplingを作る。どちらもcoupling設計だが、使う情報が異なる。

なぜペアを作り直すと直線的な軌道を学びやすいのか

Reflowは、「軌道を曲げたら罰する」という項をlossに足しているわけではない。直線性を直接の罰則として最適化するのではなく、次の学習で見せるペアを変えることで、直線的な速度場を学びやすくする。

通常のFlow Matchingで独立couplingを使うと、source点とtarget点の対応はかなり任意になる。たとえば、左下のsource点が右上のtarget点へ結ばれ、近くの別のsource点が左上のtarget点へ結ばれることがある。このようなペアをすべて直線で結ぶと、途中で線が交差しやすい。

線が交差すると、その交差付近では次のようなことが起きる。

同じ場所の近くに来た点A: 右上へ進みたい
同じ場所の近くに来た点B: 左上へ進みたい
同じ場所の近くに来た点C: 右下へ進みたい

ニューラルネットワーク:
  その場所と時刻では、複数の矢印を一つの矢印として学ぶ
  -> 平均された向きになる
  -> 個々の直線とは違う向きになり、生成軌道が曲がる

Reflowでは、まず学習済みflowでsource点を実際に運ぶ。すると、各source点に対して「このflowが実際に到着させた点」が得られる。この到着点は、完全な正解対応ではないが、少なくとも現在のflowが作った移動の結果である。

次の学習では、このsource点と到着点を直線で結び直す。すると、前回のように任意に結ばれたペアよりも、学習済みflowの動きに沿った対応になりやすい。結果として、直線path同士の無理な交差や、同じ場所で矛盾する教師矢印が減りやすい。

この違いを図に示す。

Reflowでペアを作り直すと軌道が直線的になりやすい理由

図　Reflowによるペア更新と軌道直線化の直感

図の読み方: 左は通常のFlow Matchingである。学習前に決めたsourceとtargetのペアを直線で結ぶため、途中で複数の教師矢印が同じ場所に集まり、ニューラルネットワークはそれらを平均して学ぶ。その結果、生成時の実軌道は曲がり得る。右はReflow後の学習である。学習済みflowでsource点を実際に流し、その到着点を新しい相手としてペアを結び直す。これにより、次の学習で同じ場所に集まる矛盾した教師矢印が減り、実軌道が直線的になりやすい。

この仕組みを一文で言えば、Reflowは次の操作である。

曲がった生成軌道を観察する
  -> その軌道の始点と終点を新しいペアにする
  -> そのペアを直線で結ぶ教師に置き換える
  -> 次のモデルは、前より矛盾の少ない直線教師から学ぶ

したがって、Reflowで「直線になるように学習させる」とは、lossにstraightness ratioを入れることではない。ペアの対応関係を更新し、速度回帰で平均される矢印の衝突を減らすことで、生成時の実軌道が直線に近づくように誘導する、という意味である。

ただし、これは保証ではない。base modelの精度が低い、pair-steps が粗すぎる、モデル容量が足りない、あるいはReflow後のペアでもpath交差が多い場合、軌道は十分に直線化しない。そのため、
straightness ratioで、Reflow前後の軌道を評価する必要がある。

最小コード

src/fm_minimal/reflow.py は、Reflowで使用する新しい始点・終点ペアを生成する。

@torch.no_grad()
def make_reflow_pairs(model, x0, steps=64, solver=euler_solve):
    x1_hat = solver(lambda t, x: model(t, x), x0, steps=steps)
    return x0.detach(), x1_hat.detach()

この関数は学習済みモデルを使って x0 を流し、到達点 x1_hat を返す。ここで detach() しているのは、Reflow用のペア作成を、次のモデル学習の計算グラフから切り離すためである。

使い方は次のようになる。

from fm_minimal import LinearPath, make_reflow_pairs

path = LinearPath()

# model_1は、すでに1-Rectified Flowとして学習済みだとする。
x0 = sample_standard_normal(batch, dim=2, device=device)
z0, z1 = make_reflow_pairs(model_1, x0, steps=64)

# model_2を、新しいペアに対して学習する。
loss = conditional_flow_matching_loss(model_2, path, z0, z1)

このコードで重要なのは、Reflow後の学習でもlossの形が変わらないことである。変わるのは訓練ペアである。

1回目: source sample と target sample を結ぶ
Reflow用ペア作成: source sample を model_1 で流す
2回目: source sample と model_1 の到達点を結ぶ

この意味で、Reflowは「モデル構造を変える方法」ではなく、「次の学習に使うcouplingを、学習済みflowから作る方法」と読むと見通しがよい。

一連の実験は、まず1-Rectified Flowを学習し、そのcheckpointを train_reflow_2d.py へ渡す。

python .\src\train_rectified_flow_2d.py `
  --out .\_outputs\rectified_flow

python .\src\train_reflow_2d.py `
  --base-checkpoint .\_outputs\rectified_flow\rectified_flow_2d.pt `
  --pair-steps 64 `
  --out .\_outputs\reflow

pair-steps はbase modelを積分して新しい終点を作る精度を決める。新モデルの学習step数や評価時のNFEとは別である。新モデルはbase checkpointの model_config を引き継ぐため、Reflowとモデル構造変更を同時に行わず、まずcoupling更新だけの効果を比較できる。

Rectified FlowとReflowの違い

Rectified FlowとReflowは同じものではない。Rectified Flowは、分布を運ぶ速度場を直線的にしようとする学習枠組みである。Reflowは、その枠組みの中で、学習済みflowを使って新しい訓練ペアを作る操作である。

観点 Rectified Flow Reflow
何を指すか 速度場を学ぶ枠組み 学習済みflowでペアを作り直す操作
入力ペア 初期couplingの $(x_0, x_1)$ $x_0$ と modelの到達点 $\tilde{x}_1$
loss 直線補間の速度回帰 loss自体は同じ直線速度回帰
変わるもの 学習される速度場 次の訓練に使うcoupling
目的 周辺分布を運び、軌道を扱いやすくする 次のflowをよりまっすぐにしやすいペアへ更新する

短く言えば、Rectified Flowは「何を学ぶか」、Reflowは「次の学習データをどう作るか」である。1回目のRectified Flowを学習しただけでは、まだReflowをしたことにはならない。学習済みモデルでsourceを流し、その到達点を使って次のペアを作った段階でReflowになる。

直線性を測る

Rectified Flowを扱うとき、lossだけでは「軌道がまっすぐになったか」はわからない。そこで、生成軌道の長さと、始点から終点までの直線距離を比べる。

ここでのstraightness ratioは、損失関数の一部ではなく、Reflowの効果を確認するための評価・診断指標である。Reflow後の学習でも、基本のlossは直線path上の速度回帰である。straightness ratioをlossに足してモデルを直接最適化しているわけではない。

Reflowでの使い方は次のように読む。

1回目のRectified Flowを学習する
  -> 生成軌道を取り出す
  -> straightness ratioで、軌道がどの程度曲がっているかを測る
  -> 学習済みflowでReflowペアを作る
  -> 2回目のモデルを学習する
  -> もう一度straightness ratioを測り、直線化したかを比較する

つまり、straightness ratioは「Reflowによって軌道がまっすぐになったか」を見るための物差しである。値が下がれば、同じsource点からtarget側へ向かう道が短く、遠回りが少なくなったと解釈できる。値が下がらない場合は、Reflowペアの作り方、pair-steps、学習step数、モデル容量、あるいは元のcouplingの交差が問題になっている可能性がある。

\text{straightness ratio}
=
\frac{\text{trajectory length}}{\|X_1 - X_0\|}.

完全に直線なら、この値は1に近い。曲がったり遠回りしたりすると、軌道長が増えるので1より大きくなる。

コードでは次のように計算する。

trajectory = euler_solve(
    lambda t, x: model(t, x),
    x0,
    steps=64,
    return_trajectory=True,
)
ratio = straightness_ratio(trajectory)

ここで trajectory のshapeは (time, batch, dim) である。straightness_ratio は各サンプルごとの比を返すので、平均、分位点、外れ値を見るとよい。

ratio.mean(), ratio.quantile(0.9), ratio.max()

平均だけを見ると、少数の大きく曲がったサンプルを見落とすことがある。研究用の診断では、生成サンプル、軌道図、straightness ratioの分布を合わせて見る。

Reflowは何回必要か

Reflowは理論上、繰り返すほどcouplingを更新できる。しかし、実装では回数を増やせば必ずよくなるわけではない。重要なのは、Reflow後に作られた直線path同士が、同じ時刻・同じ場所で異なる向きの教師速度を要求するかどうかである。

複数のpathが同じ途中点で交差すると、その場所では異なる教師速度が混ざる。MSEが学ぶのはそれらの条件付き平均なので、学習される速度は各ペアの直線方向からずれ、生成軌道が曲がり得る。これは「二乗誤差は平均を学ぶ」という性質の、Reflowにおける現れである。

一方、学習済みflowが作ったペアの対応がほぼ一対一で、直線path同士の交差が少なければ、途中点 $x_t$ から元のペアをほぼ特定できる。このとき条件付き平均に混ざる速度が少なくなるため、1回のReflow後に学ぶflowでも軌道は直線に近づきやすい。

pathが多く交差する
  -> 同じ途中点に複数の教師速度が来る
  -> MSEが方向を平均する
  -> 軌道が曲がりやすい

pathの交差が少ない
  -> 途中点と教師速度がほぼ一対一に対応する
  -> 平均による方向のずれが小さい
  -> 1回のReflowでも直線化しやすい

したがって、Reflow回数は固定の正解ではない。まず1回のReflow後にstraightness ratioとNFE別生成品質を測り、まだ軌道が曲がる場合に、追加Reflowと学習方法の改善を比較する。この考え方は “Improving the Training of Rectified Flows” で詳しく検討されているが、本教材では「Reflow回数より先に、pathの交差と学習誤差を診断する」という一般的な設計原理として扱う。

まとめ

この記事では、Rectified FlowをFlow Matchingとの関係から整理し、Reflowが果たす役割を説明した。

要点は次のとおりである。

  • 1-Rectified Flowの初回学習は、端点を直線で結び、その速度を回帰するという点でLinear Path CFMと同じ形で書ける。
  • 直線のpathを教師に使っても、複数の教師速度が混ざると、学習済みODEの実軌道まで直線になるとは限らない。
  • Rectified Flowでは、target分布へ輸送できることに加えて、生成軌道を直線に近づけ、少ないNFEでも追いやすくすることを重視する。
  • Reflowは新しいlossやモデル構造ではなく、学習済みflowでsource点を運び、その到達点との組を次の訓練ペアとして使う操作である。
  • Reflowの効果はlossだけで判断せず、生成品質、NFE、straightness ratio、軌道図などを組み合わせて確認する必要がある。

つまり、Rectified Flowを理解するうえでは、訓練に使う直線path、学習された速度場、生成時にODE solverがたどる実軌道を分けて考えることが重要である。さらにReflowでは、速度場だけでなく、次の学習に使うcouplingそのものを更新する。この視点を押さえると、Rectified Flowが少ステップ生成と結び付く理由や、Linear Path CFMとの共通点・相違点を整理しやすくなる。

参考文献

以下には、本教材の内容、説明順序、実装対応、用語選択の検討に使用した論文、公式実装、技術文書、解説資料、動画、既存教材を掲載する。

  • Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le, "Flow Matching for Generative Modeling", ICLR 2023. Flow Matchingの基本定式化、Gaussian path、CFMの基礎になる論文。
  • Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, Yoshua Bengio, "Conditional Flow Matching: Simulation-Free Dynamic Optimal Transport", 2023. CFM、OT-CFM、simulation-freeなdynamic OT近似を扱う論文。
  • Xingchao Liu, Chengyue Gong, Qiang Liu, "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow", 2022. Rectified Flow、Reflow、少ステップ生成の実用的な入口になる論文。
  • Qiang Liu, "Rectified Flow: A Marginal Preserving Approach to Optimal Transport", 2022. Rectified Flowのmarginal preserving propertyとOTとの関係を説明する理論寄りの論文。
  • Sangyun Lee, Zinan Lin, Giulia Fanti, "Improving the Training of Rectified Flows", NeurIPS 2024. Reflowを繰り返す前に、pathの交差、時刻分布、loss設計を改善する考え方を扱う。
  • Michael S. Albergo, Eric Vanden-Eijnden, "Building Normalizing Flows with Stochastic Interpolants", 2023. 補間過程を設計し、その確率流から速度場を学ぶStochastic Interpolantsの基礎論文。
  • Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, "Stochastic Interpolants: A Unifying Framework for Flows and Diffusions", 2023. FlowとDiffusionを補間・速度・scoreの枠組みで統一的に見る論文。
  • Patrick Esser et al., "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis", 2024. 高解像度画像生成でRectified Flow Transformerをスケールさせる実モデル系の論文。
  • Yaron Lipman et al., "Flow Matching Guide and Code", 2024. Flow Matchingの用語、実装、公式コード対応をまとめた包括的ガイド。
  • Peter Holderrieth, Ezra Erives, "An Introduction to Flow Matching and Diffusion Models", MIT 6.S184 lecture notes. DiffusionとFlow MatchingをODE/SDE、score、velocityの観点でつなぐ講義ノート。
  • Will Grathwohl et al., "FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models", ICLR 2019. continuous normalizing flowのinstantaneous change-of-variables、発散推定、NLL計算の基礎資料。
  • Cheng Lu et al., "DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps", NeurIPS 2022. NFEをそろえた高速sampling評価と、少stepでの品質比較を考えるための代表的なsolver研究。
  • Tero Karras et al., "Elucidating the Design Space of Diffusion-Based Generative Models", NeurIPS 2022. sampler設定、NFE、FID、計算量を分けて評価する実験設計の参考資料。
  • Jonathan Ho, Ajay Jain, Pieter Abbeel, "Denoising Diffusion Probabilistic Models", 2020. DDPMの基本となる前向きノイズ化、epsilon prediction、逆過程の論文。
  • Martin Heusel et al., "GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium", NeurIPS 2017. Fréchet Inception Distance（FID）を導入した論文。
  • Mikołaj Bińkowski et al., "Demystifying MMD GANs", ICLR 2018. Kernel Inception Distance（KID）とその不偏推定を導入した論文。
  • Catherine F. Higham, Desmond J. Higham, Peter Grindrod, "Diffusion Models for Generative Artificial Intelligence: An Introduction for Applied Mathematicians", 2023. 応用数学寄りにDiffusionモデルを導入する入門論文。
  • TorchCFM documentation, atong01/conditional-flow-matching. CFM variantsをFlowMatcherクラスとして読むための実装資料。
  • Diff2Flow documentation, "Training Flow Matching Models via Diffusion Model Alignment". Diffusion実装をFlow Matchingへ変換・整合させる観点の実装資料。
  • Black Forest Labs, FLUX official repository. FLUX.1の公式推論コード、モデル一覧、利用条件への入口。
  • Black Forest Labs, FLUX.1-schnell model card. 公式checkpoint、pipeline構成、Apache-2.0表示の確認元。
  • Black Forest Labs, FLUX.1-dev model card. gated checkpointと独自ライセンスへの入口。
  • Black Forest Labs, FLUX.1 [dev] Non-Commercial License. weight、派生物、outputを区別して利用条件を確認する一次資料。
  • Stability AI, Stable Diffusion 3.5 official repository. SD3.5モデル、推論例、モデルカードへの公式入口。
  • Stability AI, Stable Diffusion 3.5 Medium and Large model cards. MMDiT構成、用途、Community Licenseの確認元。
  • Stability AI, Community License. SD3.5を含む対象モデルの利用・再配布条件を確認する一次資料。
  • Tencent, HunyuanVideo official repository. 公開video generation model、checkpoint、推論条件、ライセンスの一次資料。
  • William Peebles et al., SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers. Stochastic Interpolant Transformerの公式実装とImageNet checkpoint。
  • Hugging Face Diffusers, DreamBooth training for FLUX. full/LoRA学習、対象module、text encoder学習、VRAM上の注意。
  • Hugging Face Diffusers, DreamBooth training for Stable Diffusion 3. SD3系MMDiTのLoRA学習例。
  • Hugging Face Diffusers, LoRA training and loading adapters. adapterの学習、読込、複数adapter、fuse/unfuseの公式文書。
  • Hugging Face PEFT, LoRA developer guide. rank/alpha、対象module、merge/unmerge、発展的LoRA設定の公式文書。
  • Hugging Face, safetensors documentation. tensorとmetadataを保存する形式の仕様・利用法。
  • Edward J. Hu et al., LoRA: Low-Rank Adaptation of Large Language Models. 低ランク更新の原論文。

補足的に使用した研究論文・講義資料

  • Michael S. Albergo et al., "Stochastic Interpolants with Data-Dependent Couplings", 2023. couplingを固定的な独立ペアリングに限定せず、データ依存に設計する観点の参考資料。
  • Stanley H. Chan, "Tutorial on Diffusion Models for Imaging and Vision", 2024. 画像・視覚問題におけるDiffusionの前向き過程、逆過程、score、samplingの説明を確認するために使用した。
  • Itai Gat et al., "Discrete Flow Matching", 2024. 連続空間のFlow Matchingを離散データへ拡張する研究資料。
  • Aram-Alexandre Pooladian et al., "Multisample Flow Matching", ICML 2023. ミニバッチ内の複数標本とcouplingを使って輸送を改善する観点の参考資料。
  • Nikita Kornilov et al., "Optimal Flow Matching: Learning Straight Trajectories in Just One Step", NeurIPS 2024. 直線的な軌道と少ステップ生成を比較するための発展資料。
  • Xixi Hu et al., "Improving Rectified Flow with Boundary Conditions", ICCV 2025. Rectified Flowの境界付近の挙動と改善方法を扱う資料。

公式実装・実装補助資料

直感説明・可視化・動画の参考資料

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?