拡散モデルに入門してみる

はじめに

世間では生成AIのニュースで賑やかです。生成AIには、①GAN、②VAE、③Diffusion modelなどがあります。拡散モデル(Diffusion model)を知らなかったので調査しました。本記事では、③Diffusion modelについて簡単に紹介します。

What are Diffusion Models?

GANやVAEとなにがちがうの？

拡散モデルは、データからノイズ(潜在変数)への変換の逆変換と考えるところがGANやVAEと違います。拡散モデルでは、データに徐々にノイズを付加する過程(拡散過程)と、徐々にノイズを除去する過程(逆拡散過程)を考えます。拡散過程を仮定したことで、Encoderのパラメータの学習が不要になりました。

拡散モデルの長所と短所

長所

• 2乗誤差最小化であるため最適化がシンプルである
• 多様なデータの生成に強い
• 理論的な背景が明確

短所

• 生成が遅い
• 高次元データの生成する際は、潜在変数は同じ次元であるため、次元が高くなる

変分下限(Evidence Lower Bound:ELBO)

まず、順方向の拡散過程を考えます。
$x_0$を観測値とし、$x_T$はガウシアンノイズとします。$x_t$は、$x_0$から$x_T$の途中経過です。$q(x_t|x_{t-1})$は、前時刻$x_{t-1}$を用いてガウス分布でモデル化します。

q(x_t|x_{t-1}) = \mathcal{N}(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I)

q(x_{1:T}|x_{0})=\Pi_{t=1}^{T}q(x_t|x_{t-1})

ここで、データに徐々にノイズを付加する過程(拡散過程)は、もはや確率分布$q$をパラメータ$\phi$でパラメータ化する必要がないので、Encoder側はモデル化も学習する必要もないです！

つぎに、逆方向の拡散過程を考えます。
ノイズを除去する過程(逆拡散過程)での条件付き確率$p_{\theta}(x_{t-1}|x_{t})$をどのように学習すればよいかがわかっていません。つまり、Decoder側はどのようにモデル化するか？どのように学習すればよいか？が問題として残されています。

表記を簡単にするために、$p(x_{0:T})$と$p(x_{T})$を導入します。

p(x_{0:T})=p(x_T)\Pi_{t=1}^{T}p_{\theta}(x_{t-1}|x_{t})

p(x_{T})=\mathcal{N}(x_T;0,I)

対数尤度$\log p(x)$が最大となるように、条件付き確率$p_\theta(x_{t-1}|x_t)$のパラメータ$\theta$を推定すればよさそうです。しかし、残念ながら対数尤度$\log p(x)$を直接最大化することは困難です。そこで、対数尤度$\log p(x)$のELBO$\mathcal{L}(x)$を最大化することで、間接的に対数尤度$\log p(x)$を最大化することを考えます。

\begin{eqnarray}
\log p(x)
&=&\log \int p(x_{0:T})dx_{1:T}\\
&=&\log \int \dfrac{p(x_{0:T})q(x_{1:T}|x_0)}{q(x_{1:T}|x_0)}dx_{1:T}\\
&=&\log \mathbb{E}_{q(x_{1:T}|x_0)}\Biggl\lbrack\dfrac{p(x_{0:T})}{q(x_{1:T}|x_0)}\Biggr\rbrack\\
&\geq&\mathbb{E}_{q(x_{1:T}|x_0)}\Biggl\lbrack\log\dfrac{p(x_{0:T})}{q(x_{1:T}|x_0)}\Biggr\rbrack\ \ 
(\because \text{Jensen's Inequality})\\
&=&\mathcal{L}(x)
\end{eqnarray}

拡散モデルの学習

式変形の詳細は下記を参考にしてください。

• Understanding Diffusion Models: A Unified Perspective

VAEと同様にELBO$\mathcal{L}(x)$を最大化することを考えます。ELBO$\mathcal{L}(x)$を式変形すると、

\begin{eqnarray}
\mathcal{L}(x)
&=&\mathbb{E}_{q(x_{1:T}|x_0)}\Biggl\lbrack\log\dfrac{p(x_{0:T})}{q(x_{1:T}|x_0)}\Biggr\rbrack\\
&=&\mathbb{E}_{q(x_{1}|x_0)}\Biggl\lbrack\log p_{\theta}(x_0|x_1)\Biggr\rbrack
 - \text{D}_{\text{KL}}(q(x_T|x_0)||p(x_T))
 - \sum_{t=2}^{T}\mathbb{E}_{q(x_{t}|x_{0})}
   \Biggl\lbrack\text{D}_{\text{KL}}(q(x_{t-1}|x_t,x_0)||p_{\theta}(x_{t-1}|x_{t}))\Biggr\rbrack
\end{eqnarray}

ELBO$\mathcal{L}(x)$の第1項は計算でき、第2項は定数であり、第3項はdenoising matchingになっています。拡散モデルでは、第3項に着目し、パラメータ$\theta$を推定します。

第3項をさらに式変形すると、

\begin{eqnarray}
& &\max_{\theta}-\text{D}_{\text{KL}}(q(x_{t-1}|x_t,x_0)||p_{\theta}(x_{t-1}|x_{t}))\\
&=&\min_{\theta}\text{D}_{\text{KL}}(q(x_{t-1}|x_t,x_0)||p_{\theta}(x_{t-1}|x_{t}))\\
&=&\min_{\theta}\dfrac{1}{2\sigma_{q}^2}\dfrac{(1-\alpha_t)^2}{1-\bar{\alpha}_t)\alpha_t}
\Biggl\lbrack\|\epsilon_0-\hat{\epsilon}_{\theta}(x_t,t)\|_2^2\Biggr\rbrack
\end{eqnarray}

数値実験で、係数を無視しても問題がないことがわかっているので、えいや！と無視してしまいます。

\begin{eqnarray}\min_{\theta}
\|\epsilon_0-\hat{\epsilon}_{\theta}(x_t,t)\|_2^2
\end{eqnarray}

推定したノイズ$\hat{\epsilon}_{\theta}(x_t,t)$とノイズ$\epsilon_0$が一致するように、ニューラルネットワークを学習することで、ELBO$\mathcal{L}(x)$を最大化できます。間接的に、対数尤度$\log p(x)$を最大化していることになります。

データの生成

参考文献

• データ生成・変換のための機械学習　第７回前編「Diffusion models」
• Understanding Diffusion Models: A Unified Perspective
• 拡散モデル　データ生成技術の数理、岡野原大輔、岩波書店
• 拡散モデルのPytorch実装
• 多変量正規分布の場合のKullback Leibler Divergenceの導出

おわりに

ELBO$\mathcal{L}(x)$の第1項を考えなくてよい理由がよくわからなかったです。