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 入門 (4) PyTorchによる Flow Matching 学習の流れ

1
Last updated at Posted at 2026-06-27

はじめに

この記事は、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は、学習済みの地図をたどる生成側の部品である。

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

6. 2D toy実験で何を見るか

この章で扱う概念: 2D toy実験は、低次元の既知分布を使って、経路、速度場、生成軌道、mode coverageを直接観察する検証環境である。高品質な最終モデルを作るためではなく、数式と実装の対応や失敗原因を目で確認するために使う。

主な段階: 学習 → 生成 → 評価 — 速度場を学習し、Euler samplerで生成し、分布と軌道を可視化して確かめる。

6.1 直感の地図

最小実装を動かしたら、最初に見るべきものはlossだけではない。Flow Matchingは「速度場を学ぶ」方法なので、次の3つを同時に見る。

  • source分布から出た点が、生成後にtarget分布へ近づいているか。
  • 粒子の軌道が、途中で大きく曲がりすぎたり発散したりしていないか。
  • 学習した速度場が、直線経路の教師信号と矛盾しない向きを向いているか。

2D toy実験では、targetを8個の小さなGaussianからなるGaussian混合分布(Gaussian mixture)にする。これは画像生成よりずっと小さいが、Flow Matchingが「ノイズの雲をデータの形へ運ぶ」様子を見るには十分である。

source、target、ペアごとの直線pathの対応を図2に示す。
toy_linear_path_concept.png
図2 2D toyにおけるsource、target、ペアごとのLinear path

図2は学習済みモデルの出力ではなく、学習で使う教師信号の作り方を表している。青い点がsourceサンプル $x_0$、赤い点がtargetサンプル $x_1$、灰色の線がペアごとのlinear path、緑の点が途中時刻の $x_t$ である。実際の学習済みモデルが作る軌道は、この教師pathの平均的な速度場をODEとして積分したものになる。

6.2 学習: 速度場とcheckpointを作る

train_minimal_2d.py は、次の処理を行う。

source x0 ~ N(0, I)
target x1 ~ eight Gaussians
t ~ Uniform(0, 1)
x_t = (1 - t) x0 + t x1
u_t = x1 - x0
loss = ||v_theta(t, x_t) - u_t||^2

上の処理は既定値の linear path、independent coupling、uniform time samplerの場合である。このスクリプトは、学習後に以下を保存する。

  • model: 学習済み MLPVelocity の重み。
  • model_config: MLPの幅、深さ、時刻埋め込み、活性化関数。
  • samples: Euler samplerで生成した点。
  • path, coupling, time_sampler, seed: 実験条件。

実行例:

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

6.3 生成・評価: サンプルと軌道を可視化する

plot_minimal_2d.py は、保存済みモデルを読み込み、source、target、generated samples、trajectoryを1枚の図にする。

PyTorch環境で学習と可視化を実行すると、図3の完全な確認図が得られる。
minimal_2d_result.png
図3 2D toyのsource分布、target分布、生成結果、粒子軌道

図3の各パネルは次のように読む。

  • 左上: source分布。最初は標準正規の丸い雲である。
  • 右上: target分布。8個のモードが円周上に並ぶ。
  • 左下: 生成結果。Euler samplerで運ばれた点がtargetの8モードへ近づくかを見る。
  • 右下: 軌道図。いくつかの粒子が、$t=0$ から $t=1$ までどのように動くかを見る。

図3は、lossの数字だけでは見えない失敗を見つけるために重要である。たとえば、lossが下がっていても、軌道が大きく蛇行していれば少ないNFEでの生成は難しくなる。また、軌道は学習済みの周辺速度場をODEとして積分した結果なので、ペアごとの教師pathそのものとは一致しない。

図4は、保存済みcheckpoint minimal_2d.pt 内の samples を描画した生成サンプルである。学習過程や軌道を表す図ではなく、checkpointに保存された生成終点だけを確認するための図である。

minimal_2d_checkpoint_samples.png
図4 minimal_2d.pt に保存された生成終点とtargetの8 mode中心

図4の青い点が保存済みモデルから生成されたサンプル、赤い点がtargetの8 mode中心である。これは簡易確認用であり、図3の左下パネルを単独で見ているものに近い。図4だけでは軌道の曲がり方やNFE依存性はわからない。

実行例:

python .\src\plot_minimal_2d.py `
  --checkpoint .\_outputs\minimal_2d.pt `
  --out .\figures\minimal_2d_result.png

6.4 学習・生成・評価の対応

観察するもの Flow Matchingでの意味 次に疑う点
generated samplesがtargetに乗らない 速度場が分布を正しく運べていない 学習不足、モデル容量不足、時刻埋め込み
軌道が大きく曲がる Euler少ステップで誤差が出やすい path設計、coupling、solver
一部のmodeにしか乗らない target分布のcoverageが悪い データサンプリング、モデル容量、学習時間
sourceの形が残る ODEで十分に運べていない loss、学習率、ステップ数、NFE

この時点では、完璧な生成品質は目標ではない。目標は、「Flow Matchingの式が、実際の点群の動きとして何を起こすのか」を目で確認することである。

6.5 2D toyを公式実装と対応づける

Meta flow_matching へ移るときは、まず2D continuous exampleを見る。教材との対応は次である。

具体的なURL:

教材で見たもの 公式実装で見る場所 読み方
LinearPath flow_matching/path $x_t$ と教師速度を作る部品
conditional_flow_matching_loss flow_matching/loss またはexample内のloss計算 モデル出力と教師信号の比較
euler_solve flow_matching/solver 学習済み速度場で点を動かす
2D toy可視化 examples/2d_flow_matching.ipynb path、loss、solverを小さく確認する入口

公式コードを読むときも、最初はモデル構造ではなく path -> loss -> solver の順に探す。

ここまでで最小実装を一周した。第7章では、第1章から第6章で使ったCFMを実装部品として整理する。第8章では、点をODEで動かすことが分布の変形へどうつながるかを、連続の式とpushforwardで補う。その後、第9章と第10章で「なぜ周辺速度場を直接教師にできないのか」「CFMは何を条件付きにしているのか」を理論的に掘り下げる。先に動くコードの役割を固定し、分布が動く直感を得てから、CFMの理由を数式で確かめる順である。

7. CFMを実装単位で読む

この章で扱う概念: CFM実装は、時刻を選ぶ、経路上の点を作る、教師速度を作る、モデルで予測する、二乗誤差を取る、という交換可能な部品から成る。実装単位で読むとは、数式上の選択がどのクラスや関数に閉じ込められているかを確認することである。

主な段階: 学習の実装 — time sampler、path、velocity model、lossを一つの学習stepとして接続する。

7.1 直感の地図

CFMの実装は、モデル本体よりも「訓練データをどう作るか」が重要である。モデルがMLPでもU-NetでもDiTでも、最小の訓練ステップは次の形になる。

sourceを取る
targetを取る
tを取る
x_tを作る
u_tを作る
v_theta(t, x_t)を出す
MSEを取る

この順番で読めば、Flow Matching系の多くのコードはかなり見通しがよくなる。

7.2 最小実装の役割分担

この教材の最小実装では、役割を次のように分けている。

ファイル 役割 対応する概念
paths.py $x_t$ と $u_t$ を作る 条件付き経路
losses.py CFM lossを計算する 速度場の教師あり回帰
models.py $v_\theta(t, x)$ を定義する ニューラル速度場
solvers.py ODEを数値的に解く 生成サンプラー
data.py source/targetサンプルを作る $p_0$, $p_1$

研究コードを読むときも、この分解を探すとよい。ファイル名は違っても、だいたい同じ役割が存在する。

7.3 LinearPathを差し替えると何が変わるか

今の LinearPath

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

だけを実装している。

もしGaussian pathやvariance preserving pathに変えるなら、変わるのは主に paths.pysamplevelocity である。lossの外形は大きく変わらない。

x_t = path.sample(t, x0, x1)
u_t = path.velocity(t, x0, x1)
pred = model(t, x_t)
loss = ((pred - u_t) ** 2).flatten(1).sum(dim=1).mean()

この構造が見えていると、「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?