1
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

はじめに

この記事は、Flow Matching および関連する生成モデルについて学ぶために、作成した個人的なノートの一部を整理したものです。

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

現在、ネット上にある多くのFlow Matching解説が高度な数学が必要です。本記事は、入門のために理論に重きを置きません。

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

また、本資料の全文、コードはGitHubにあげています。Rectified Flowについては、本記事では扱っていませんが、全文には含まれています。

Githubに載せた全体の第0部と第1部を数回に分割して、記事にしていきます。前回の記事は以下になります。

誤りや改善点などがありましたら、コメントでご指摘いただけると幸いです。

想定読者

本資料は、Flow Matching および関連する生成モデルを学びたい方を対象としています。

特に、

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

を主な対象としています。

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

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

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

第1部 Flow Matchingの基礎・理論・最小実装

1.2 学習と生成の全体像(再掲)

Flow Matchingの全体像を理解するために再掲した。これだけで、全体の役割分担がわかるので、まだ読んでいない人はぜひ読んでほしい。

Flow Matchingには、「矢印の地図を作る学習」と「その地図を使って新しいデータを作る生成」という二つの流れがある。この二つを分けて考えると、全体像をつかみやすい。

学習は、簡単な乱数の集まりである source分布 $p_0$ と、学びたいデータの集まりである target分布 $p_1$ から始まる。sourceからノイズ点 $x_0$、targetからデータ点 $x_1$ を取り出し、どの2点を一組にするかを coupling が決める。couplingは、出発点と目的地の組み合わせを作る「ペア係」だと考えればよい。

ペアが決まると、time sampler が練習する時刻 $t$ を選び、path(経路) がその時刻の途中地点 $x_t$ を置く。同じpathから、「その途中地点ではどちらへ進むべきか」という 教師速度 $u_t$ も得られる。pathは練習用の道を描く部品、教師速度はその道に沿った正解の矢印である。

時刻 $t$ と途中地点 $x_t$ を受け取り、矢印を予測するニューラルネットワークが velocity model(速度場モデル) $v_\theta(t,x_t)$ である。loss(損失関数) は予測矢印と教師速度のずれを採点し、optimizer はその点数が小さくなるようにモデルを調整する。この練習を繰り返して得た「矢印の地図」を、重みや設定とともに保存したものが checkpoint である。

学習:
source x0 + target x1
  -> coupling: 出発点と目的地をペアにする
  -> time sampler: 練習する時刻tを選ぶ
  -> path: 途中地点x_tと教師速度u_tを作る
  -> velocity model: v_theta(t, x_t)を予測する
  -> loss: 予測と教師速度のずれを測る
  -> optimizer: モデルを直す
  -> checkpoint: 学習済みの矢印の地図を保存する

生成時には、学習に使ったtargetデータ、coupling、教師速度を使わない。source分布から新しいノイズ点を取り、checkpointから復元したvelocity modelへ「今の時刻と場所では、どちらへ進めばよいか」と問い合わせる。

連続した動きをコンピュータ上でたどるため、time grid が問い合わせる時刻を並べ、ODE solver がモデルの矢印を使って点を少しずつ更新する。Euler、Heun、RK4は、この「一歩の進み方」が異なるsolverである。初期ノイズの用意からsolverによる移動、生成結果の返却までをまとめた手続き全体を sampler と呼ぶ。

生成:
新しいsource noise x0
  -> checkpointからvelocity modelを読み込む
  -> time grid: 問い合わせる時刻を並べる
  -> ODE solver: モデルに矢印を聞いて少し進む
  -> sampler: この更新を繰り返して生成サンプルを返す
  -> evaluation: target分布に近いかを評価する

モデルへ矢印を問い合わせた回数を NFE(Number of Function Evaluations) と呼ぶ。NFEが少なければ生成は速くなりやすいが、一歩が粗すぎると正しい道から外れる。evaluation(評価) では、生成結果がtarget分布へ近づいたかだけでなく、同じNFEでどのsolverがよかったかも調べる。

ここで最も大切な関係は、次の三つである。

  • coupling、time sampler、path、教師速度、loss、optimizerは、矢印の地図を作る学習側の部品である。
  • velocity modelは、学習時に矢印を練習し、生成時に矢印を答える共有部品である。
  • time grid、ODE solver、sampler、NFEは、学習済みの地図をたどる生成側の部品である。

この先の章では、これらを一つずつ数式とコードへ落としていく。今は用語を暗記する必要はない。「ペアを作る → 練習経路と正解矢印を作る → 矢印の地図を学ぶ → 新しいノイズを地図に沿って歩かせる」という関係だけをつかめばよい。

2. Linear Pathの数式とPyTorch実装

この章で扱う概念: Linear Pathは、source点 $x_0$ とtarget点 $x_1$ を線分で結び、時刻 $t$ に応じた途中点を線形補間で決める条件付き経路である。分布全体を一本の線にするのではなく、サンプルの各ペアに対して一本ずつ直線経路を定める。

主な段階: 学習 — path、time sampler、途中地点、教師速度は、velocity modelへ練習問題を作るために使う。

2.1 直感の地図

Linear pathは、2点をまっすぐ結ぶ最も簡単な道である。

時刻 $t=0$ ではsource点 $x_0$ にいる。時刻 $t=1$ ではtarget点 $x_1$ にいる。途中の $t=0.25$ なら4分の1だけ進んだ場所、$t=0.5$ なら中点、$t=0.75$ なら4分の3だけ進んだ場所にいる。

この直感をそのまま式にすると、

x_t = (1 - t)x_0 + t x_1

である。

2.2 小さな数値例

1次元で考える。$x_0 = -2$, $x_1 = 6$ とする。

t = 0.00 -> x_t = -2
t = 0.25 -> x_t = 0
t = 0.50 -> x_t = 2
t = 0.75 -> x_t = 4
t = 1.00 -> x_t = 6

速度は常に

u_t = x_1 - x_0 = 8

である。つまり、この点は単位時間で8だけ進む等速運動をしている。

2次元でも同じで、$x_0$, $x_1$, $x_t$, $u_t$ がスカラーではなくベクトルになるだけである。

2.3 テンソルshape

ミニバッチで実装すると、典型的なshapeは次のようになる。

数式 コード変数 shape 意味
$x_0$ x0 [B, D] source分布からの点
$x_1$ x1 [B, D] データ分布からの点
$t$ $t$ [B, 1] 各サンプルの時刻
$x_t$ $x_t$ [B, D] 途中地点
$u_t$ $u_t$ [B, D] 正解速度

$t$ を [B] ではなく [B, 1] にするのは、各サンプルに時刻を一つずつ対応させるためである。2Dの [B,D] にはこのままbroadcastできる。画像の [B,C,H,W] では、path内部で [B,1,1,1] へ拡張する必要があり、教材コードはデータの次元数に合わせてこの拡張を自動で行う。

2.4 経路と教師速度のコード対応

losses.py では、Linear pathに関する部分を次の2行として読める。

x_t = path.sample(t, x0, x1)
u_t = path.velocity(t, x0, x1)

ここでは、ペアと時刻から「モデルに見せる途中地点」と「その地点での正解の矢印」を作っている。損失関数そのものは、条件付きという言葉を定義した後の4.3節で、数式・直感・コードをまとめて説明する。

コード 数式 説明
path.sample(t, x0, x1) $x_t = (1-t)x_0 + tx_1$ 途中地点を作る
path.velocity(t, x0, x1) $u_t = x_1 - x_0$ 正解速度を作る

この2つを分けておくと、pathを交換したときに「途中地点」と「教師速度」が同時に変わることが見やすい。

2.5 time sampler: どの時刻を学ぶか

time sampler は、学習用の時刻 $t$ をどの確率分布から選ぶかを担当する部品である。生成時にODEを進めるsamplerとは別物であり、学習時にだけ使う。

CFM lossには時刻に関する期待値が含まれる。

\mathcal{L}_{\mathrm{CFM}}
=
\mathbb{E}_{t,x_0,x_1}
\left[
\lVert v_\theta(t,x_t)-u_t\rVert^2
\right].

実装では連続するすべての時刻でlossを計算できないため、各サンプルに一つずつ $t$ を割り当てて期待値を近似する。$t$ のshapeを [B,1] にする理由は2.3節で説明したbroadcastである。

time_samplers.py には、比較用に三つの分布がある。

名前 分布 多く選ばれる場所 用途
uniform 一様分布 全区間を同じ頻度 最初に使うbaseline
center 対称なBeta(2,2)分布 $t=0.5$付近 中間状態を重点的に学ぶ比較
endpoint 対称なBeta(0.5,0.5)分布 $t=0$ と $t=1$付近 両端を重点的に学ぶ比較

Beta分布の形を読む

Beta分布は $[0,1]$ 上だけで定義されるため、Flow Matchingの時刻を選ぶ分布として使いやすい。二つの正の値 $\alpha,\beta$ によって、確率密度の形が決まる。

p(t)
=
\frac{1}{B(\alpha,\beta)}
t^{\alpha-1}(1-t)^{\beta-1},
\qquad 0<t<1.

$B(\alpha,\beta)$ は、密度全体の面積を1にするための正規化定数(ベータ関数)である。time samplerを理解する段階では、この定数を自分で計算する必要はない。重要なのは、$t$ と 1-t の指数が分布の形を変えることである。

ここでの $\alpha,\beta$ はBeta分布の形を決める値であり、第3章で使うpath係数 $\alpha(t),\sigma(t)$ とは別物である。

この教材コードでは $\alpha=\beta$ とした左右対称のBeta分布を使う。

  • $\alpha=\beta=1$: 式の $t$ に依存する部分が1になり、一様分布になる。
  • $\alpha=\beta>1$: $t=0$ と $t=1$ 付近の密度が小さく、中央が山になる。値を大きくするほど $t=0.5$ 付近へ集中する。
  • $0<\alpha=\beta<1$: 中央の密度が小さく、両端が高いU字型になる。値を0へ近づけるほど端点付近へ強く集中する。

したがって、$\mathrm{Beta}(2,2)$ は中央重視、$\mathrm{Beta}(0.5,0.5)$ は両端重視になる。もし $\alpha$ と $\beta$ を異なる値にすれば、source側とtarget側のどちらか一方へ偏らせることもできるが、最初の教材コードでは比較を単純にするため左右対称だけを扱う。

time samplerを変えると何が変わるか

同じbatch sizeと学習step数なら、計算できるlossの個数は変わらない。time samplerを変える操作は、その限られた学習回数を時間区間のどこへ多く配るかを変える操作である。

たとえば center を使うと、モデルはsourceとtargetが強く混ざった中央付近の状態を何度も見る。その区間の速度予測は改善しやすい一方、端点付近を見る回数は減る。endpoint では逆に、sourceから動き始める区間とtargetへ到着する区間を多く学ぶが、中央付近の学習回数は減る。

数式上も、time samplerを $q(t)$ に変えると、近似する目的関数は

\mathcal{L}_{q}
=
\mathbb{E}_{t\sim q(t),x_0,x_1}
\left[
\lVert v_\theta(t,x_t)-u_t\rVert^2
\right]

となる。つまり、単に乱数の生成方法を変えるだけでなく、時刻ごとの誤差の重み付けを変えている。多く選ばれる時刻の誤差はoptimizerへ何度も渡されるため、モデルはその区間へ容量を多く使うようになる。

この変更には次の効果と副作用がある。

sampler 期待する効果 起こり得る副作用
uniform 全区間を偏りなく学ぶ 特に難しい区間へ十分な更新を割けないことがある
center 中央付近の複雑な速度場を重点的に学ぶ source直後・target直前の誤差が残ることがある
endpoint 軌道の出発・到着付近を重点的に学ぶ 中央付近の速度場が粗くなることがある

どの区間が難しいかは、path、coupling、data分布、モデルによって変わる。したがって、centerendpoint が常に uniform より優れているわけではない。時刻をいくつかの区間に分けてlossを記録し、誤差が集中している場所と、生成軌道が崩れる場所を確認して選ぶ。

また、time samplerは学習時にどの時刻を見るかを変える部品である。生成時のODE step数、solver、評価時刻の並べ方は変えない。学習用time samplerと生成用scheduleを混同しないことが重要である。

from fm_minimal import get_time_sampler

time_sampler = get_time_sampler("uniform")
t = time_sampler(batch_size, x0.device, x0.dtype)  # [B, 1]

loss関数はtime samplerを引数として受け取る。

loss = conditional_flow_matching_loss(
    model,
    path,
    x0,
    x1,
    time_sampler=time_sampler,
)

最初の学習では uniform を使う。time samplerを変える実験では、data、coupling、path、model、optimizerを同じに保ち、変更点を時刻分布だけにする。また、学習lossだけで優劣を決めず、同じNFEで生成品質も比較する。時刻分布が違うと「どの時刻の誤差を重く見たlossか」も変わるため、loss値の単純比較だけでは判断できない。

(おまけ) 3.Gaussian Pathの数式とPyTorch実装

Flow Matchingの基本理解には不要だが、比較対象の他のPathとして書いています。

この章で扱う概念: Gaussian pathは、データ成分とGaussian noise成分を時刻依存係数 $\alpha(t)$, $\sigma(t)$ で混合して中間分布を作る確率経路である。その速度は混合係数を時間微分して得られ、Linear Pathとは途中分布と教師速度の両方が異なる。

主な段階: 学習 — Linear Pathと同様に、途中地点 $x_t$ と教師速度 $u_t$ を作るpathの選択肢として使う。

3.1 直感の地図

Linear pathでは、途中点は

x_t = (1-t)x_0 + t x_1

であり、速度は一定だった。

Gaussian pathでは、ノイズとデータを混ぜる係数が $t$ によって非線形に変わる。

x_t = \alpha(t)x_1 + \sigma(t)x_0.

このとき速度は、係数を微分すればよい。

u_t
=
\frac{d x_t}{dt}
=
\dot{\alpha}(t)x_1 + \dot{\sigma}(t)x_0.

直感的には、$\alpha(t)$ はデータ成分の音量、$\sigma(t)$ はノイズ成分の音量である。時間が進むにつれてデータ成分の音量を上げ、ノイズ成分の音量を下げる。その変化率が速度になる。

3.2 三角関数スケジュール

この教材の最小コードでは、例として次のスケジュールを使う。

\alpha(t) = \sin\left(\frac{\pi t}{2}\right),
\quad
\sigma(t) = \cos\left(\frac{\pi t}{2}\right).

このとき

$$
\alpha(t)^2 + \sigma(t)^2 = 1
$$

なので、分散を保つ形のGaussian pathとして直感的に扱いやすい。

微分は

\dot{\alpha}(t)
=
\frac{\pi}{2}\cos\left(\frac{\pi t}{2}\right),
\quad
\dot{\sigma}(t)
=
-\frac{\pi}{2}\sin\left(\frac{\pi t}{2}\right).

したがって速度は

u_t
=
\frac{\pi}{2}
\cos\left(\frac{\pi t}{2}\right)x_1
-
\frac{\pi}{2}
\sin\left(\frac{\pi t}{2}\right)x_0.

対応するコードは paths.pyTrigGaussianPath である。

def sample(self, t, x0, x1):
    return self.alpha(t) * x1 + self.sigma(t) * x0

def velocity(self, t, x0, x1):
    return self.alpha_dot(t) * x1 + self.sigma_dot(t) * x0

このように、pathを変えてもCFM lossの外側は変わらない。変わるのは、$x_t$ と $u_t$ の作り方である。

3.3 Linear pathとの比較

項目 Linear path Gaussian path
中間点 $(1-t)x_0+tx_1$ $\alpha(t)x_1+\sigma(t)x_0$
速度 x1 - x0 $\dot{\alpha}(t)x_1 + \dot{\sigma}(t)x_0$
係数の変化 一定 スケジュール依存
Diffusionとの対応 直感的な直線補間 ノイズ混合式と対応しやすい

Linear pathは、最初にFlow Matchingを理解するには一番よい。

参考文献

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

  • 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の境界付近の挙動と改善方法を扱う資料。

公式実装・実装補助資料

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

1
1
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
1
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?