はじめに

「Introduction to Flow Matching and Diffusion Models」(MIT IAP 2026) のコース資料を基にフローマッチングと拡散モデルについて勉強した際の日本語メモです。

生成モデルとSDE

データ分布からのサンプリング = “生成”

$$ z_\text{new} \sim p_\text{data} $$

データ分布 $p_\text{data}$ の確率密度は未知

$$ p_\text{data}: \mathbb{R}^d \rightarrow \mathbb{R}_{\geq 0}, \\ z \mapsto p_\text{data}(z) $$

条件付き生成 = 条件付き確率分布からのサンプリング

条件ベクトル $y$ を用いて

$$ z_\text{new} \sim p_\text{data}(\cdot|y) $$

ピカール・リンデレフの定理

ベクトル場 $u_t = u_t(x)$ について $u$ が $x$ について連続微分可能であるなら初期値条件 $X_0=x_0$ を満たす常微分方程式

$$ \frac{\text{d}}{\text{dt}}X_t = u_t(X_t) $$

の解は一意に存在する。

  • ベクトル場

    空間中の1点 $\bold{r}$ を指定すると $\bold{r}$ の関数としてベクトルが1つ決まるような関数。

例

$u_t(x) = -\theta x$ のときを考えると常微分方程式 $\frac{\text{d}}{\text{dt}}X_t = u_t(X_t)$ の解は、$X_t = Ce^{-\theta t}$ となる。初期条件 $X_0=x_0$ より $C=x_0$ だから $X_t = x_0e^{-\theta t}$

オイラー法

十分に小さい刻み幅 $h$ を設定し

$$ X_i = X_{i-1}+hu_t(X_{i-1}) $$

と計算して $X_0,...,X_T$ の近似値を求める。

フローモデルのサンプリング

ニューラルネットワークベクトル場を $u = u^\theta (x,t)$, 反復回数を $n$ としたとき

t = 0
h = 1/n
X_t = sample(p_init) # 初期ノイズのサンプリング

for i in range(n):
    X_t = X_t + h * u(X_t, t) 
    t += h 

return X_t # 最終的な生成サンプル

確率微分方程式とブラウン運動

ブラウン運動 $B(t)_{t\in[0,T]}$ で駆動する確率微分方程式は、$t,x$ を変数とする関数 $u,b$ を用いて

$$ X_t = X_0 + \int_0^t u(X(s),t) ds +\int_0^t b(X(s),t) dB(s) $$

のかたちをとり、微分形で表すと

$$ dX_t = u(X(t),t)dt + b(X(t),t)dB(t) $$

オイラー・丸山法

$$ X_{t+h} - X_t =\int_t^{t+h} u(X(s),t) ds +\int_t^{t+h} b(X(s),t) dB(s) $$

$h$ が十分に小さければ上式の右辺は

$$ u(t,X_t)h + b(t,X_t)(B(t+h)-B(t)) $$

と近似できる。

ここで確率変数 $B(t+h) - B(t)$ は正規分布 $\mathcal N(0, h)$ に従うので、$b(t, X_t)$ が状態 $X_t$ に依存しない関数 $\sigma_t$ で表される場合

$$ X_{t+h} - X_t = u(t, X_t)h + \sigma_t\sqrt h \epsilon, \quad \epsilon \sim \mathcal N(0, I_d) $$

と計算できる。

オルンシュタイン・ウーレンベック過程

確率微分方程式

$$ dX_t = -\theta(X_t-\mu)dt+ \sigma d B_t $$

の解として与えられる確率過程のこと。

ブラウン運動で駆動する確率微分方程式について $u_t(x) = -\theta x$, 平均 $\mu = 0$ と設定した場合のかたちに等しい。

フローマッチング

確率パス

時刻 $t$ におけるデータ $x$ の確率密度を表す $p_t(x)$

条件付き確率パス $p_t(\cdot | z)$

任意の点 $z \in \mathbb R^d$ について

$$ p_0(\cdot | z) = p_\text{init}, \quad p_1(\cdot | z) = \delta_z $$

が成立し、時間が $t=1$ のとき、$z$ の位置が確定する。

例えば、ガウス分布 $p_t(\cdot | z) = \mathcal N(\alpha_tz, \beta_t^2I_d)$ に従う場合

  • $\alpha_0=0, \space \beta_0 = 1$ の場合 $\mathcal N(0, I_d) = p_\text{init}$
  • $\alpha_1=1, \space \beta_1 = 0$ の場合 $\mathcal N(z, 0) = \delta_z$

のようにして $\alpha_{0}, \alpha_{1}, \beta_0, \beta_1$ を適切に設定するようなノイズスケジュールを考えればよい。

周辺確率パス $p_t$

周辺確率パスについて、

$$ p_0 = p_\text{init}, \quad p_1=p_\text{data} $$

が成立する。

条件付き確率パス $p_t(\cdot | z)$ について確率の周辺化

$$ p_t(x) = \int p_t(x|z)p_\text{data}(z)dz $$

を考えて求める。

条件付きベクトル場 $u_t(x|z)$

任意の条件付き確率パス $p_t(\cdot | z)$ には等価なベクトル場・常微分方程式が存在する。

ガウス分布に従う条件付き確率パスの場合

まず条件付きフロー $\psi^\text{target}_t(x|z) = \alpha_t z+\beta_tx$ を定義する。

ここでフローの定義より任意の初期点 $X_0 \sim \mathcal p_\text{init}$ について

$$ X_t = \psi^\text{target}_t(X_0|z) = \alpha_tz + \beta_t X_0 \sim \mathcal N(\alpha_tz, \beta_t^2I_d) = p_t(\cdot | z) $$

であるから、常微分方程式 $\frac{\text{d}}{\text{dt}}X_t = u_t(X_t|z)$ に代入して

$$ \frac{d}{dt}\psi^\text{target}_t(x|z) = u^\text{target}_t(\psi^\text{target}_t(x|z)\mid z) $$

$\psi^\text{target}_t(x|z)$ の定義から

$$ \dot{\alpha}_t z + \dot{\beta}_t x = u^\text{target}_t(\alpha_t z + \beta_t x \mid z) $$

$x = (X_t - \alpha_tz)/\beta_t$ を代入して

$$ \dot{\alpha}_t z + \dot{\beta}_t \left( \frac{X_t - \alpha_t z}{\beta_t} \right) = u^\text{target}_t(X_t | z) $$

$z,X_t$ についてまとめて

$$ \left( \dot{\alpha}_t - \frac{\dot{\beta}_t}{\beta_t} \alpha_t \right) z + \frac{\dot{\beta}_t}{\beta_t} X_t = u^\text{target}_t(X_t | z) $$

周辺ベクトル場 $u_t(x)$

確率の周辺化を用いて

$$ u^\text{target}_t(x) = \int u^\text{target}_t(x|z)p_t(z|x)dz $$

ベイズの定理より $p_t(z|x) = \frac{p_t(x|z)p_t(z)}{p_t(x)}$ であり、かつ $p_t(z)$ は実際には時間依存せず $p_t(z) = p_\text{data}(z)$ だから

$$ u^\text{target}_t(x) = \int u^\text{target}_t(x|z)\frac{p_t(x|z)p_\text{data}(z)}{p_t(x)}dz $$

連続の式

ベクトル場 $u_t$ が確率パス $p_t$ に対応することは、連続の式 (continuity equation)

$$ \frac{\partial}{\partial t}p_t(x) = -\text{div}(p_t u_t)(x) $$

を満たすことと同値関係にある。

損失関数

フローマッチングの損失関数は次のように定義される。

$$ \mathcal L(\theta) = ||u_t^\theta(x) - u^\text{target}_t(x|z)||^2 $$

ガウス分布に従う条件付きベクトル場を考えると

  • $x = \alpha_tz + \beta_t \epsilon, \quad \epsilon \sim \mathcal N (0, I_d)$
  • $\dot{\alpha}_t z + \dot{\beta}_t x = u^\text{target}_t(\alpha_t z + \beta_t x \mid z)$

を踏まえて

$$ ||u_t^\theta(x) - u^\text{target}_t(x|z)||^2 = ||u_t^\theta(\alpha_tz + \beta_t \epsilon) - (\dot{\alpha}_t z + \dot{\beta}_t \epsilon)||^2 $$

特に、$\alpha_t = t, \beta_t= 1-t$ を代入すれば

$$ \begin{align*}\mathcal L (\theta) &= \mathbb E_{t\sim U[0,1],z \sim p_\text{data}, \epsilon \sim \mathcal N(0, I_d)}[||u_t^\theta(\alpha_tz + \beta_t \epsilon) - (\dot{\alpha}_t z + \dot{\beta}t \epsilon)||^2] \\ &= \mathbb E_{t\sim U[0,1],z \sim p_\text{data}, \epsilon \sim \mathcal N(0, I_d)}[||u_t^\theta(tz + (1-t) \epsilon) - (z - \epsilon)||^2]\end{align*} $$

このようなノイズスケジュールを特に直線スケジュールと呼び、最適輸送と一致している。

参考