Diffusion model を理解するためのノート

#diffusion model

2022-06-14

ぼんやりと過ごしているうちに、世間では DALL·E 2 や Imagen のような高性能の生成モデルが話題になっていました。いつの間にか生成部分にGANではなく、diffusion model というものが使われるようになっていたのも知りました。

このままでは時代に取り残されてしまうと思い diffusion model の草分けと思われる論文の一つarXiv:2006.11239 を理解しようと読み始めたのですが、式の導出が思ったより大変だったので過程をまとめておきたい、というのがこの記事の趣旨です。

目標は arXiv:2006.11239 で提案された diffusion model の損失関数

\begin{equation} \label{eq:simple-loss} L_{\mathrm{simple}}(\theta) = \mathbb{E}_{t,\boldsymbol{x}_0,\boldsymbol{\epsilon}} \left[\left| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar\alpha_t} \boldsymbol{x}_0 + \sqrt{1-\bar\alpha_t}\boldsymbol{\epsilon},t) \right|^2\right] \end{equation}

を導出することです。

Diffusion model とは？というところや、具体的な式の導出については以下のブログにも詳しく書かれています。

What are Diffusion Models?

日本語訳もされているので、そもそも diffusion model って何？という人はそちらを見ると良いと思います。

また上記論文でも appendix や引用文献を詳しく読めばわかるようにはなっています。以下では、それらを読む中で多少ギャップを感じた部分を埋めたり、期待値の定義が曖昧に書かれている部分を explicit にしたりということをしていきます。

問題のセットアップ

$\boldsymbol{x}_0$ をデータ、$\boldsymbol{x}_T$ は $p(\boldsymbol{x}_T) = \mathcal{N}(\boldsymbol{x}_T; \boldsymbol{0},\mathbf{I})$ から生成される正規乱数とします。学習時にはデータ $\boldsymbol{x}_0$ に少しずつノイズを加えて $\boldsymbol{x}_T$ が得られるようにモデルを学習し (diffusion process) 、推論時には逆に $\boldsymbol{x}_T$ を少しずつ denoize することで (reverse process) 学習データ分布からの生成を実現します。

Diffusion process

Diffusion process は以下のようなマルコフ過程とします \begin{align} \label{eq:diffusion-process} \boldsymbol x_0 &\sim q(\boldsymbol x_0),\notag \\ \boldsymbol x_t &\sim q(\boldsymbol x_t | \boldsymbol x_{t-1}) = \mathcal{N}(\boldsymbol x_t;\sqrt{1-\beta_t}\boldsymbol x_{t-1},\beta_t)\quad (t=1,2,\dots,T). \end{align} 元々のデータ $\boldsymbol x_0$ を生成する分布を $q(\boldsymbol x_0)$ としています。時刻 $t=0$ から初めて $t=1,2 \dots,T$ で少しずつノイズを加えていく様子を記述しています。$\beta_t$ は $t$ に対して単調増加な関数に選んで、徐々に加えるノイズが大きくなり、より分布の中心がゼロに近づくようにしています。¹

Reverse process

Reverse process は以下のようなマルコフ過程とします \begin{align} \boldsymbol x_T &\sim p(\boldsymbol x_T) = \mathcal{N}(\boldsymbol{x}_T; \boldsymbol{0},\mathbf{I}),\notag\\ \boldsymbol x_{t-1} &\sim p_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t) = \mathcal{N}(\boldsymbol x_{t-1};\boldsymbol\mu_\theta(\boldsymbol x_t, t), \boldsymbol\Sigma_\theta(\boldsymbol x_t, t))\quad (t=T-1,\dots,0) .\label{eq:reverse-process} \end{align} ここで $\theta$ は学習パラメータです。Reverse process は正規乱数 $\boldsymbol x_T$ から始まり、これからノイズを除去して、学習データ分布からのサンプルに近づくように学習します。

学習では以下の negative log likelihood を最小にします² \begin{align} \label{eq:neg-log-likelihood} \mathbb{E}_{\boldsymbol x_0}[-\log p_\theta(\boldsymbol x_0)] &= -\int d\boldsymbol x_0 q(\boldsymbol x_0)\log\left[\int d\boldsymbol x_{1:T} p_\theta(\boldsymbol x_{0:T})\right],\\ \mathrm{where}\quad p_\theta(\boldsymbol x_{0:T}) &= p(\boldsymbol x_T)\prod_{t=1}^Tp_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t).\notag \end{align} データ $\boldsymbol x_0$ の確率ができるだけ高くなるように、$\theta$ を調整することになります。$\boldsymbol x_0 \sim q(\boldsymbol x_0)$ としたので、期待値は $q(\boldsymbol x_0)$ に関する積分になります。

Variational bound を導出する

式 \eqref{eq:neg-log-likelihood} を学習に便利なように変形していきます。

Jensen の不等式から

\begin{align} -\log p_\theta(\boldsymbol x_0) &= -\log\left[\int d\boldsymbol x_{1:T} p_\theta(\boldsymbol x_{0:T})\right] \notag\\ &= -\log\left[\int d\boldsymbol x_{1:T} q(\boldsymbol x_{1:T}|\boldsymbol x_0) \frac{p_\theta(\boldsymbol x_{0:T})}{q(\boldsymbol x_{1:T}|\boldsymbol x_0)}\right]\notag\\ &\leq -\int d\boldsymbol x_{1:T} q(\boldsymbol x_{1:T}|\boldsymbol x_0)\log\left[\frac{p_\theta(\boldsymbol x_{0:T})}{q(\boldsymbol x_{1:T}|\boldsymbol x_0)}\right]. \end{align}

なので

\begin{align} \mathbb{E}_{\boldsymbol x_0}[-\log p_\theta(\boldsymbol x_0)] &\leq -\int d\boldsymbol x_{0:T} q(\boldsymbol x_{0:T})\log\left[\frac{p_\theta(\boldsymbol x_{0:T})}{q(\boldsymbol x_{1:T}|\boldsymbol x_0)}\right]\notag\\ &=\int d\boldsymbol x_{0:T} q(\boldsymbol x_{0:T}) \left[-\log p(\boldsymbol x_T) - \sum_{t=1}^T\log\frac{p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t)}{q(\boldsymbol x_t|\boldsymbol x_{t-1})}\right]\notag\\ &= \mathbb{E}_q\left[-\log p(\boldsymbol x_T) - \sum_{t=1}^T\log\frac{p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t)}{q(\boldsymbol x_t|\boldsymbol x_{t-1})}\right]\notag\\ &=:L. \end{align}

$\mathbb{E}_q$ は $q(\boldsymbol x_{0:T})$ の下での平均としました。変分法でよくある式変形だと思います。以降は式 \eqref{eq:neg-log-likelihood} の上限である $L$ を最小化することにします。

天下りですが、$q(\boldsymbol x_t|\boldsymbol x_{t-1})$ を以下のように分解します \begin{align} q(\boldsymbol x_t|\boldsymbol x_{t-1}) =q(\boldsymbol x_t|\boldsymbol x_{t-1}, \boldsymbol x_0) =q(\boldsymbol x_{t-1}|\boldsymbol x_t, \boldsymbol x_0)\frac{q(\boldsymbol x_t|\boldsymbol x_0)}{q(\boldsymbol x_{t-1}|\boldsymbol x_0)}. \end{align}

これを $L$ の中の和の $t>1$ の部分に適用すると

\begin{align} L &= \mathbb{E}_q\left[ -\log p(\boldsymbol x_T) - \sum_{t=2}^T\log\frac{p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t)}{q(\boldsymbol x_{t-1}|\boldsymbol x_t, \boldsymbol x_0)} + \sum_{t=2}^T\log\frac{q(\boldsymbol x_t|\boldsymbol x_0)}{q(\boldsymbol x_{t-1}|\boldsymbol x_0)} - \log\frac{p_\theta(\boldsymbol x_0|\boldsymbol x_1)}{q(\boldsymbol x_1|\boldsymbol x_0)} \right]\notag\\ &= \mathbb{E}_q\left[ -\log\frac{p(\boldsymbol x_T)}{q(\boldsymbol x_T|\boldsymbol x_0)} - \sum_{t=2}^T\log\frac{p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t)}{q(\boldsymbol x_{t-1}|\boldsymbol x_t, \boldsymbol x_0)} - \log p_\theta(\boldsymbol x_0|\boldsymbol x_1) \right]\notag\\ &\equiv L_T + \sum_{t=2}^T L_{t-1} + L_0. \end{align} 以下で見るように、この形にすることで $L$ を tractable な形で書くことができるようになります。

なお、この式の第一項 $L_T$ は学習パラメータを含まないため、考える必要はありません。以下では第二項、第三項を見ていきます。

$L_{t-1}$ の計算

$L_{t-1}$ を計算するには、$q(\boldsymbol x_t|\boldsymbol x_0)$ と $q(\boldsymbol x_{t-1}|\boldsymbol x_t, \boldsymbol x_0)$ の関数形を求める必要があります。

$q(\boldsymbol x_t|\boldsymbol x_0)$ の計算

任意の時刻 $t$ での diffusion の分布 $q(\boldsymbol x_t|\boldsymbol x_0)$ は、単なる Gaussian noise の合成なので、解析的に計算できます

\begin{align} q(\boldsymbol x_t|\boldsymbol x_0) &= \int d\boldsymbol x_{1:t-1}\prod_{s=1}^{t}q(\boldsymbol x_s|\boldsymbol x_{s-1})\notag\\ &= \int d\boldsymbol x_{1:t-1}\prod_{s=1}^{t} \frac{1}{\sqrt{2\pi}\beta_s} \exp\left[-\frac{1}{2\beta_s}\left|\boldsymbol x_s - \sqrt{1-\beta_s}\boldsymbol x_{s-1}\right|^2\right]. \end{align} まず $\boldsymbol x_{t-1}$ による積分を考えます。関連する項だけ抜き出すと

\begin{align} \int d\boldsymbol x_{t-1} \exp\left[ -\frac{1}{2(1-\alpha_t)}\left|\boldsymbol x_t - \sqrt{\alpha_t}\boldsymbol x_{t-1}\right|^2 -\frac{1}{2(1-\alpha_{t-1})}\left|\boldsymbol x_{t-1} - \sqrt{\alpha_{t-1}}\boldsymbol x_{t-2}\right|^2 \right] \end{align} となります。論文にならって $\alpha_t = 1-\beta_t$ としました。$\boldsymbol x_{t-1} = \sqrt{\alpha_{t-1}}\boldsymbol x_{t-2} + \sqrt{1-\alpha_{t-1}}\boldsymbol z_{t-1}$ と変数変換し $\boldsymbol z_{t-1}$ について平方完成して積分すると \begin{align} &\int d\boldsymbol z_{t-1} \exp\left[ -\frac{1}{2(1-\alpha_t)}\left|\sqrt{\alpha_t(1-\alpha_{t-1})}\boldsymbol z_{t-1} - (\boldsymbol x_t - \sqrt{\alpha_t\alpha_{t-1}}\boldsymbol x_{t-2})\right|^2 -\frac{1}{2}\boldsymbol z_{t-1}^2 \right]\notag\\ &\propto \exp\left(-\frac{1}{2(1-\alpha_t\alpha_{t-1})}|\boldsymbol x_t - \sqrt{\alpha_t\alpha_{t-1}}\boldsymbol x_{t-2}|^2\right) \end{align} となるので、これを繰り返すと \begin{align} q(\boldsymbol x_t|\boldsymbol x_0) &\propto \exp\left(-\frac{1}{2(1-\bar\alpha_t)}|\boldsymbol x_t - \sqrt{\bar\alpha_t}\boldsymbol x_0|^2\right), \mathrm{where}\quad \bar\alpha_t = \prod_{s=1}^t\alpha_t \end{align} つまり \begin{align} \label{eq:q-at-t} q(\boldsymbol x_t|\boldsymbol x_0) =\mathcal{N}(\boldsymbol x_t; \sqrt{\bar\alpha_t}\boldsymbol x_0, (1-\bar\alpha_t)\mathbf{I}) \end{align} となることがわかります。

$q(\boldsymbol x_{t-1}|\boldsymbol x_t,\boldsymbol x_0)$ の計算

式 \eqref{eq:q-at-t} を使うことで \begin{align} q(\boldsymbol x_{t-1}|\boldsymbol x_t,\boldsymbol x_0) &=\frac{q(\boldsymbol x_t|\boldsymbol x_{t-1})q(\boldsymbol x_{t-1}|\boldsymbol x_0)}{q(\boldsymbol x_t|\boldsymbol x_0)} \notag\\ &\propto \exp\left[ -\frac{1}{2}\left( \frac{1}{1-\alpha_t}|\boldsymbol x_t - \sqrt{\alpha_t}\boldsymbol x_{t-1}|^2 +\frac{1}{1-\bar\alpha_{t-1}}|\boldsymbol x_{t-1} - \sqrt{\bar\alpha_{t-1}}\boldsymbol x_0|^2 \right) \right]\notag\\ &\propto \exp\left[ -\frac{1}{2}\left( \frac{1-\bar\alpha_t}{(1-\alpha_t)(1-\bar\alpha_{t-1})}\boldsymbol x_{t-1}^2 -2\left(\frac{\sqrt{\alpha_t}}{1-\alpha_t}\boldsymbol x_t + \frac{\sqrt{\bar\alpha_{t-1}}}{1-\bar\alpha_{t-1}}\boldsymbol x_0\right)\boldsymbol x_{t-1} \right) \right], \end{align} となります。$\boldsymbol x_{t-1}$ について二次なので Gaussian です。 $q(\boldsymbol x_{t-1}|\boldsymbol x_t,\boldsymbol x_0) = \mathcal{N}(\boldsymbol x_{t-1}|\tilde{\boldsymbol\mu}_t(\boldsymbol x_t, \boldsymbol x_0), \tilde\beta_t\mathbf{I})$ とおくと

\begin{align} \tilde{\boldsymbol\mu}_t(\boldsymbol x_t, \boldsymbol x_0) &=\frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar\alpha_t} \left( \frac{\sqrt{\alpha_t}}{1-\alpha_t}\boldsymbol x_t + \frac{\sqrt{\bar\alpha_{t-1}}}{1-\bar\alpha_{t-1}}\boldsymbol x_0 \right)\notag\\ &= \frac{(1-\alpha_t)\sqrt{\bar\alpha_{t-1}}}{1-\bar\alpha_t}\boldsymbol x_0 +\frac{(1-\bar\alpha_{t-1})\sqrt\alpha_t}{1-\bar\alpha_t}\boldsymbol x_t\notag\\ \tilde\beta_t &=\frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar\alpha_t} \end{align} となります。

Variational bound の計算

論文にならって $\Sigma_\theta(\boldsymbol x_t, t) = \sigma^2_t\mathbf{I}$ という ansatz をおくと

\begin{align} L_{t-1} &=-\mathbb{E}_q\left[\log\frac{p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t)}{q(\boldsymbol x_{t-1}|\boldsymbol x_t, \boldsymbol x_0)}\right]\notag\\ &=-\int d\boldsymbol x_0 d\boldsymbol x_{t-1}d\boldsymbol x_t q(\boldsymbol x_0)q(\boldsymbol x_t|\boldsymbol x_0)q(\boldsymbol x_{t-1}| \boldsymbol x_0, \boldsymbol x_t) \log p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t) + \mathrm{const.}\notag\\ &=\int d\boldsymbol x_0 d\boldsymbol x_{t-1}d\boldsymbol x_t q(\boldsymbol x_0)q(\boldsymbol x_t|\boldsymbol x_0)q(\boldsymbol x_{t-1}| \boldsymbol x_0, \boldsymbol x_t) \frac{1}{2\sigma^2_t}|\boldsymbol x_{t-1} - \boldsymbol\mu_\theta(\boldsymbol x_t, t)|^2 + \mathrm{const.}\notag\\ &=\int d\boldsymbol x_0 d\boldsymbol x_t q(\boldsymbol x_0)q(\boldsymbol x_t|\boldsymbol x_0) \frac{1}{2\sigma^2_t}|\tilde{\boldsymbol\mu}_t(\boldsymbol x_t, \boldsymbol x_0) - \boldsymbol\mu_\theta(\boldsymbol x_t, t)|^2 + \mathrm{const.} \end{align} $\mathrm{const.}$ は学習パラメータ $\theta$ を含まないという意味です。$\sigma^2_t$ の関数形は任意性があります。詳しくは元論文を見てください。

この式のままでも良いのですが、実用上は以下のようなリパラメトライゼーションをするとより性能がよくなるそうです。

式 \eqref{eq:q-at-t} を思い出すと、$\boldsymbol x_t$ は $\boldsymbol\epsilon \sim \mathcal{N}(\boldsymbol\epsilon;\boldsymbol 0,\mathbf{I})$ によって $\boldsymbol x_t=\sqrt{\bar\alpha_t}\boldsymbol x_0 + \sqrt{1-\bar\alpha_t}\boldsymbol\epsilon$ とパラメトライズでき、これにより $\tilde{\boldsymbol\mu}_t = (\boldsymbol x_t - \beta_t/\sqrt{1-\bar\alpha_t}\boldsymbol\epsilon) / \sqrt{\alpha_t}$ と書けます。よって $\boldsymbol\mu_\theta(\boldsymbol x_t, t)$ も同様に \begin{align} \label{eq:reverse-mu} \boldsymbol\mu_\theta(\boldsymbol x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( \boldsymbol x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\boldsymbol\epsilon_\theta(\boldsymbol x_t, t) \right) \end{align} とパラメトライズすると

\begin{align} L_{t-1} &=\int d\boldsymbol x_0 d\boldsymbol\epsilon q(\boldsymbol x_0)\mathcal{N}(\boldsymbol\epsilon;\boldsymbol 0,\mathbf{I}) \frac{\beta^2_t}{2\sigma^2_t\alpha_t(1-\bar\alpha_t)}|\boldsymbol\epsilon - \boldsymbol\epsilon_\theta(\boldsymbol x_t, t)|^2 + \mathrm{const.}\notag\\ &=\mathbb{E}_{\boldsymbol x_0, \boldsymbol\epsilon} \left[ \frac{\beta^2_t}{2\sigma^2_t\alpha_t(1-\bar\alpha_t)} \left|\boldsymbol\epsilon - \boldsymbol\epsilon_\theta(\sqrt{\bar\alpha_t}\boldsymbol x_0 + \sqrt{1-\bar\alpha_t}\boldsymbol\epsilon, t)\right|^2 \right] + \mathrm{const.} \label{eq:likelihood-t-1} \end{align}

となります。各学習データ $\boldsymbol x_0$ に対して、標準正規分布から乱数 $\boldsymbol\epsilon$ を生成し、モデル $\boldsymbol\epsilon_\theta$ に入力した結果が元の乱数 $\boldsymbol\epsilon$ を再現するように学習することになります。

$L_0$ の計算

$L_0=\mathbb{E}_q[-\log p_\theta(\boldsymbol x_0|\boldsymbol x_1)]$ は denoizing の最終レイヤーであり、デコーダの役割を果たします。ここだけさらに複雑なモデルを設定することもできるらしいですが、論文では単に independent discrete decoder、つまり $\boldsymbol\mu_\theta(\boldsymbol x_1,1)$ をそのまま $\boldsymbol x_0$ として使うようなデコーダを考えています。

論文では画像データの生成を考えているので各ピクセルは $0, 1, \dots, 255$ の離散値を取り、$x_0$ の各成分はそれを $[-1, 1]$ へ線形にスケールしたものとしています。なので、取りうる値は $-1, -1+2/255, -1 + 4/255, \dots, 1$ であり、$2/255$ ずつになっています。

論文では以下のような likelihood を設定しています

\begin{equation} \label{eq:decoder-likelihood} p_\theta(\boldsymbol x_0|\boldsymbol x_1) = \prod_{i=1}^D\int_{\delta_-(x_0^i)}^{\delta_+(x_0^i)} dx\mathcal{N}(x;\mu_\theta^i(\boldsymbol x_1, 1), \sigma_1^2), \end{equation}

\begin{align} \delta_+(x) = \begin{cases} \infty\quad(x = 1)\\ x + 1/255\quad(x<1) \end{cases}\quad \delta_-(x) = \begin{cases} -\infty\quad(x = -1)\\ x - 1/255\quad(x>-1) \end{cases} \end{align}

直観的には、データの値 $x_0^i$ の $\pm 1/255$ の区間に $\mu_\theta^i$ が入ると likelihood が高くなるということです。また $x_0^i=+1$ $(-1)$ の時は、上 (下) に外す分には構わないという式になっています。

式 \eqref{eq:decoder-likelihood} のままでは扱いづらいので、$x_0^i=\pm1$ のコーナーケースを無視し、積分を区間の長さ $2/255$ が短いと思って近似してみます。すると

\begin{equation} -\log p_\theta(\boldsymbol x_0|\boldsymbol x_1) \simeq \frac{1}{2\sigma_1^2}|\boldsymbol x_0 - \boldsymbol\mu_\theta(\boldsymbol x_1, 1)|^2 +\mathrm{const.} \end{equation}

と書けるので

\begin{equation} \label{eq:likelihood-0} L_0 = \mathbb{E}_{\boldsymbol x_0, \boldsymbol\epsilon}\left[ \frac{1-\alpha_1}{2\sigma_1^2\alpha_1} \left|\boldsymbol\epsilon - \boldsymbol\epsilon_\theta(\sqrt{\bar\alpha_1}\boldsymbol x_0 + \sqrt{1-\bar\alpha_1}\boldsymbol\epsilon, 1)\right|^2 \right] +\mathrm{const.} \end{equation}

となり、$L_{t-1}$ に $t=1$ を代入した結果と一致することがわかります。

目的関数の簡略化

実用上は式 \eqref{eq:likelihood-t-1} と \eqref{eq:likelihood-0} の二乗誤差の前の係数は無視しても良いらしいです。$t$ についての和も、$[1,T]$ からの一様サンプリングに置き換えることにすると、損失関数は結局

\begin{equation} \label{eq:likelihood-simple} L_{\mathrm{simple}}(\theta) = \mathbb{E}_{t,\boldsymbol{x}_0,\boldsymbol{\epsilon}} \left[\left| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar\alpha_t} \boldsymbol{x}_0 + \sqrt{1-\bar\alpha_t}\boldsymbol{\epsilon},t) \right|^2\right] \end{equation}

で良くなります。これは各時刻における重みを変えていることになります。

式 \eqref{eq:likelihood-t-1} のままではノイズの分散が小さい small $t$ における損失の weight が大きくなります。式 \eqref{eq:likelihood-simple} のように取り直すことで、より難しい、large $t$ における学習にフォーカスできるようになるらしいです。

実際に学習する際には、まず $\boldsymbol{\epsilon}_\theta$ に使うモデルを決め、データ $\boldsymbol x_0$ と乱数 $t$, $\boldsymbol \epsilon$ を使って、式 \eqref{eq:likelihood-simple} を損失関数とした optimize を行います。推論ではまず正規乱数 $\boldsymbol x_T$ をサンプルします。そして学習済みのモデル $\boldsymbol{\epsilon}_\theta$ を使って式 \eqref{eq:reverse-mu} から $\boldsymbol\mu_\theta$ を計算、式 \eqref{eq:reverse-process} を使って $\boldsymbol x_{t-1}$ をサンプリング、というステップを $t=0$ まで繰り返すことでデータ生成を行います。