記事の要約
Flow Matching for Generative Modelingの要約である．
実験の部分は省略し，数学的な部分の要約をする．
特に定理1について，論文では省略された部分を丁寧に示す．
次の記事でトイモデルでの具体的な実装をする予定である．
導入
Denoising Diffusion Probabilistic Models
などで考察された拡散モデルは単純な分布から元のデータの分布を得るものである．
この拡散モデルとは異なる方法で単純な分布から元のデータの分布を得る手法として，この論文ではFlow matching（FM）を提案してる．
この論文では次のことをしている．
- フローの定義と分布の押し出しの定義
- FMの損失・目的関数の提案
- フローの変種の定義
- 実験
フローの前提知識
この章ではFlow matchingを定めるために必要なフローや押し出しを定義する．
時間依存のベクトル場
$\mathbb{R}^d$をデータの空間とする．データ点$(x^1,\dots,x^d) \in \mathbb{R}^d$と表す．
滑らか写像$v:[0,1]\times\mathbb{R}^d \rightarrow \mathbb{R}^d$を時間依存のベクトル場と呼ぶ．ただし，ここでは接ベクトル空間$T_x \mathbb{R}^d$と$\mathbb{R}^d$を同一視している．
時間依存ベクトル場$v_t$のフロー$\phi:[0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$とは，以下の常微分方程式から定義される積分曲線のことである：
\begin{align}
\frac{d}{dt} \phi_t(x) &= v_t(\phi_t(x)) \\
\phi_0(x) &= x
\end{align}
フローは任意の$t\in[0,1]$に対して微分同相写像であると仮定する．
微分方程式の解の一意性から$\phi_{s+t} = \phi_s(\phi_t)$である．
注意
数学的には$v_t$に対してフロー$\phi$が閉区間$[0,1]$全体で定義できるとは限らない．
ここではフロー$\phi$が閉区間全体$[0,1]$に伸びるもののみを考える．
例えば，$v:[0,1] \times \mathbb{R} \rightarrow \mathbb{R}$として，$v_t(x) =x^2$となるものを選ぶ．フローは
$$\phi_t(x_0) = \frac{x_0}{1-x_0t}$$
になり，これは$t = \frac{1}{x_0}$で定義されない．
確率密度と押し出し
微分同相写像$f:\mathbb{R}^d \rightarrow \mathbb{R}^d$に対して，$\mathbb{R}^d$の確率密度関数$p$の$f$による 押し出し $f_* p$は以下で定義される．
$$ f_* p(x) = p(f^{-1}(x)) det \left[ \frac{\partial f^{-1}}{\partial x} (x)\right] $$
$p_t(x): [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}_{>0}$が 確率密度のパス(probability density path) であるとは，$p_t(x)$が$t$に対して連続であり，各$t\in [0,1]$に対して$p_t$が$\mathbb{R}^d$上の確率密度関数になってることをいう．
フロー$\phi:[0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d$と$\mathbb{R}^d$の確率密度関数$p_0$に対して，
$$p_t = {\phi_t}_* p_0$$
とすると，$p_t$は確率密度のパスになる．
実際に$\phi_t$が微分同相なら，$p_0$に対応する測度$\mu$の像測度${\phi_t}_* \mu$が$p_t$の測度になっている．加えて，$\phi_t$が微分同相でないと像測度${\phi_t}_*\mu$のラドン・ニコディム微分が定義できない．
データの空間として$\mathbb{R}^d$を考えているが，体積形式をもつ多様体などボレル集合上の測度を持つ多様体$M$であれば，上と同じように可測関数$f_t:M \rightarrow [0, \infty)$のパスを定義することができる．
Flow Matching
FM目的関数を定義し，条件付き確率密度パス(conditional probability path)とベクトル場を構成する．
FM目的関数
$q$をデータ空間$\mathbb{R}^d$上の確率密度関数とする．
データ空間上の時間依存のベクトル場$u_t$とそれによって誘導される確率密度のパス$p_t:[0,1]\times\mathbb{R}^d \rightarrow \mathbb{R}^d$で
$$p_1 =q$$
を満たすものを考える．
このとき，$v_t(x,\theta)$に対して，FM目的関数を
$$
\mathcal{L}_{\mathrm{FM}}(\theta) = {\mathbb{E}}[v_t(x,\theta) - u_t(x)]
$$
と定める．ここで$t\sim \mathcal{U}[0,1]$，$x \sim P_t$に従う．ただし，$P_t$は確率密度$p_t$に対応する分布である．
条件付き確率分布と周辺分布
有限のデータサンプル$D \subset \mathbb{R}^d$をとる．
$p(x)$を$\mathbb{R}^d$上の確率密度関数とする．
条件付き確率密度のパス$ p_t(- | -) : \mathbb{R}^d \times D \rightarrow \mathbb{R}$で以下を満たすのをとる．
\begin{align}
p_0(x|x_1) = f_1(x,x_1)\\
p_1(x|x_1) = f_\sigma (x,x_1)\\
\end{align}
ただし，$\sigma \in {\mathbb{R}}$であり，$f_{\sigma}(x,x_1)$は正規分布$\mathcal{N}(x | x_1,\sigma^2I)$の確率密度関数である．
以降，このような$p_t(-|-)$と確率密度関数を誘導するベクトル場$u_t(-|-)$を考える．
$q_D(x)$を次で定まる密度関数とする．
\begin{equation}
q_D(x)=
\begin{cases}
\frac{1}{n} & x \in D \\
0 & \text{otherwise}
\end{cases}
\end{equation}
ここで$n$は$D$の元の個数．
この時，以下の周辺分布パスと周辺ベクトル場を以下のように定義する．
\begin{align}
p_t(x) = \int_D p_t(x | x_1) q_D(x_1)dx_1 \\
u_t(x) = \int_D {u_t(x | x_1) \frac{p_t(x | x_1) q_D(x_1)}{p_t(x)}dx_1}
\end{align}
この時，次が成り立つ．
定理1
周辺ベクトル場$u_t(x)$は$p_t(x)$を生成する．
定理1
上の状況において，
周辺ベクトル場$u_t(x)$は$p_t(x)$を生成する
証明には以下の補題が必要である．
補題2
以下の補題を示す[3,Theorem 5.34]．
$M = \mathbb{R}^d$とおき．
$v_t$を$M$上のベクトル場とする．
この時，滑らかベクトル場$v_t$により定まるフローが$T_t \colon [0,T)\times M \rightarrow M$を各$t$に対して$T_t$が微分同相と仮定する．
$p_0$を$M$上の$C^{\infty}$写像として，$p_t = {T_0}_* p_0$と置くと次が成り立つ．
$$\frac{d}{dt} p_t(x) + \mathrm{div}(p_t(x) v_t(x) ) = 0$$
証明
コンパクトな台を持つ滑らかな写像$\varphi:M \rightarrow \mathbb{R}$をとる．
この$\varphi$に対して以下を示せば十分．
\begin{align}
\frac{d}{dt}\int_M \varphi(x) p_t(x)dx &= \int_M \varphi(x) \mathrm{div}(p_t(x)v_t(x))dx
\end{align}
これは具体的な計算によって得る．
\begin{align}
\frac{d}{dt}\int_M \varphi(x) p_t(x)dx &= \frac{d}{dt} \int_M \varphi(T_t(x)) p_0(x)dx \\
&= \int_M \frac{d}{dt}\varphi(T_t(x)) p_0(x)dx \\
&= \int_M \left\{ (\mathrm{grad}\varphi)(T_t(x)) \cdot \frac{dT_t}{dt} \right\} p_0(x)dx\\
&= \int_M \left\{ (\mathrm{grad}\varphi)(T_t(x)) \cdot p_0(x)v_t(T_t(x)) \right\}dx\\
&= \int_M \left\{ (\mathrm{grad}\varphi)(x) \cdot p_t(x)v_t(x) \right\}dx \\
&= \int_M \mathrm{div}(\varphi(x)p_t(x) v_t(x)) dx - \int_M \varphi(x) \mathrm{div}(p_t(x) v_t(x))dx
\end{align}
初めの式変形は$p_t$の定義と変数変換，
次の変形は微分と積分の交換は連続関数が可測であることから従う．
3行目は連鎖律，4行目は$v_t$の定義である．
5行目は変数変換から6行目は勾配と発散の性質から従う．
ただし，$\varphi$の台を含む閉集合$V$と任意のベクトル場$F$に対してガウスの発散定理により
\begin{align}
\int_{M} \mathrm{div}{(\varphi(x)F(x))} dx
&= \int_V \mathrm{div}{(\varphi(x)F(x))} dx\\
&= \int_{\partial V} \varphi(x)(F(x) \cdot \boldsymbol{n} )dV\\
&=0
\end{align}
である．
よって補題を得る．
補題3
補題1と同じ状況で，
偏微分方程式
$$\frac{d}{dt} p_t(x) + \mathrm{div}(p_t(x) v_t(x) ) = 0$$
を満たすものは一意である．
証明
上の偏微分方程式の解$q_t\colon M \rightarrow \mathrm{R}$が$q_0 = 0$を満たすとする．
この時，各$t$で$q_t=0$を示せばいい．
連鎖律と定義から
$$ \frac{d}{dt}q_t(T_t(x)) =0$$
である．
よって，$ q_t(T_t(x)) = q_0(x) =0$になる．
$T_t$は微分同相だったから，$q_t =0$.
定理1の証明
微分を計算する．
\begin{align}
\frac{d}{dt} p_t(x) &= \int_M \left( \frac{d}{dt}p_t(x | x_1) \right) q_D(x_1)dx_1\\
&=-\int_M \mathrm{div} \left( u_t(x | x_1) p_t(x | x_1) \right) q_D(x_1)dx_1 \\
&= -\mathrm{div} \left(\int_M u_t(x | x_1) p_t(x | x_1) q_D(x_1)dx_1 \right) \\
&= -\mathrm{div} \left( u_t(x) p_t(x)\right)
\end{align}
補題3より，主張は得られた．
主定理を用いることで，$u_t$ をよりシンプルな$u_t(|x_1)$によって計算できる．
CFM目的関数
FM目的関数は$u_t$に着目したものであった．
条件付き確率に付随するベクトル場$u_t(x|x_1)$に着目することで，CFM目的関数（Conditional Flow Matching）が定義できる．
$$ \mathcal{L}_{\mathrm{CFM}}(\theta) = {\mathbb{E}}[v_t(x,\theta) - u_t(x|x_1)] $$
定理4
$p_t(x)>0$を仮定する．
このとき，CFM目的関数とFM目的関数はスカラー倍を除いて一致する．
従って
$$\nabla_{\theta}{\mathcal{L}}_{\mathrm{FM}} = \nabla_{\theta}{\mathcal{L}}_{\mathrm{CFM}}$$
を得る．
証明
具体的に計算する（省略）．
参考文献
- Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).
- Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851.
- Villani, Cédric. Topics in optimal transportation. Vol. 58. American Mathematical Soc., 2021.