Flow Matching

Flow Matching(FM)とは？

FLUX.1、Qwen-Image*1やWAN*2でDiTとともに採用されている画像生成のアルゴリズム。
SD3はその派生であるRectified Flowを採用している。
従来の潜在拡散モデルと比較してのメリットは、色彩の豊かさ等が大幅に改善されており、
仕組みがシンプルなため訓練効率が上がりやすく*3低ステップ化に向いている。

数学的な背景はともかくとして実装上は、生成時にサンプラーが行うことはV-Predの潜在拡散モデルに近い。
ただし、学習時の処理は大きく異なる。

潜在拡散モデルは各ステップのノイズ画像（仕掛品）をその場その場でノイズ「除去」していたが、
Flow Matchingではベクトル場（単位時間あたりの変化率。現在の画像生成AIでは一般的に入力画像と出力画像との差分に相当）の推定を学習する。
Flow Matchingは出発分布に依存せず、ノイズ画像以外の初期分布（例えば別の画像）を起点とすることも可能である。
ただし、AIにとっての「何も描いてないまっさらなキャンバス」としてはガウシアンノイズが扱いやすい。

WAN、Qwen-Image-EditやFLUX.1 Kontextなどは出発点をノイズ画像ではなく別の画像にしたFlow Matchingといえる。

論文

Yaron Lipman他「Flow Matching for Generative Modeling」arXiv

理論

U-Net拡散モデルでは「1枚1枚の教師画像に1000ステップかけてガウスノイズを撒いていき、各ステップの画像について除去すべきノイズ全体を覚えさせる」という訓練を行ってAIにノイズの中から画像を削りだすやり方を学習させていた。
初期の拡散モデルは技術的な制約でそのような手間をかけなければ実現できなかったからだが、訓練コスト以外にも明確な弱点がある。

各ステップで正規分布に基づいてノイズを撒くため、その過程を学習したモデルは「一度描いたものを消してまた描き直す」といったようなふらふらした描画過程を再現することになる。
つまり例えばピクセル単位で見ると（潜在拡散モデルはピクセル単位で処理を行っているわけではないが、比喩として）、
一度赤で塗られたマスが少し経つと青になり、今度は白になったりという挙動を示す。
このような「曲がりくねった」描画過程は大変扱いづらく、以下のような問題が発生していた。

• 画像の表現範囲に強い制約を課す必要がある
  AIにとってもノイズ予測（ε予測）は難しいため、画像全体のデータ分布を平均0・分散1に正規化してモデルの負担を抑える必要があった。
  潜在拡散モデルはVAEによる変換を挟んでいるため平均化すると灰色になるような画像しか描けないわけではないが、構造的に中間色に寄る。
  このためSDXLモデルでは高めのCFG Scaleで強引にブーストしないとメリハリのある画像にならない。

• 高速化が困難。画像生成モデルの挙動を近似するのが難しい
• ステップ数を減らすと品質が極端に劣化する
  画像生成モデルを1000ステップ呼んで生成を行うのは時間がかかりすぎるため、
  現在ではそのうちの数十ステップ程度のみをモデルに推論させ、
  間のステップについてはサンプラーが軽い計算を行って補間するのが当たり前となっている。
  しかし拡散モデルの描画過程は気まぐれで曲がりくねっているため、この補間を行うのが難しい。

そもそもサンプラーがモデルの代わりに補間を行えるのは、最初の拡散モデル（DDPM）の発表後、DDIMによって
「拡散モデルの描画過程は書こうと思えば一つの数式にまとめられる」ことが示されたためである。
1000ステップの拡散過程なんてやめてノイズを画像にする数式を直接学習させるだけで性能を出すにはどうすればいいか、という発想の到達点がFlow Matching。

Flow Matchingではノイズ拡散過程ではなく、教師画像の上に起点画像（ノイズ画像等）を透過率100%で重ね少しずつ透過率を下げていくような過程を学習している。
予測対象もノイズではなく、ベクトル場（起点画像と出力画像との差分）に変わった。
これはv-predのベロシティと似ているが、損失関数（訓練時の採点基準）が異なりモデルの負担がより少ないため、学習が安定しやすい。
これによりFMモデルは潜在拡散モデルのような気まぐれな挙動ではなく、海の底から少しずつ画像を浮かび上がらせていくような直線的で安定した描画過程を示す。

（これには、初期拡散モデルではプロンプト等による描画指示が実用化されていなかったため正真正銘ガチャに頼るしかなかったが、
　現在では乱数要素を減らしても入力値を変えれば多様な出力ができるようになっているという点も大きい。）

これにより次のようなメリットがある。

• ダイナミックレンジが大幅に改善された
  AIにとってはノイズεよりベクトル場（起点画像と出力画像との差分）の予測のほうが負担が少ない。
  またFMモデルは描画過程も直線的で安定しているため、拡散モデルのようにいきなり原色で塗りたくるような挙動を示すことも少ない。
  これにより、少なくともサンプラーの処理中は画像全体のデータ分布についての制約が緩和された。
  これが生成画像の色彩の豊かさにつながっている。

• Euler法のような単純なサンプラーでも補間精度を保ちやすい
• 蒸留による低ステップ化と相性が良い
  FMモデルは直線的な描画過程を学んでいるため、CFG値が穏当であれば拡散モデルのようにいきなり原色で塗りたくるような挙動は少なく、簡素なサンプラーでも品質を維持しやすい。
  （計算誤差によって生成案を「切り替える」ような挙動を見せるときもあるが、FMのベクトル場は何らかの完成画像への一貫した流れを持つため、最終的な収束は安定しており破綻しにくい。）
  理想的には一直線になるものの実際の画像の分布は複雑なため、通常のモデルでは収束までのステップ数は拡散モデルと変わらないが、
  蒸留によって低ステップ化しても品質を維持しやすい。

Flow Matchingモデルの弱点

• 焼きつきやすい
  Flow Matchingモデルというよりは完成画像への方向性を予測する方式の弱点であり、v-predのSDXLモデルと共通する。
  Flow Matchingモデルは直線的な描画過程を前提としている為、仕掛かり画像の状態を肯定的にとらえる傾向があり一度勢いがつくと修正が効きづらい。
  このためU-Net拡散モデルよりCFG耐性がやや低い。

• 一度大きく崩れると画像全体が発散する
  FMというよりはDiT系FMの弱点だが、DiTはエラーが拡散しやすいため大きく崩れるとそれが全体に波及し、FMの直線的な描画傾向と相まって一気に崩壊する。