Denoising Diffusion Probabilistic Models

はじめに

　2020年の以下の論文の式変形を記載する。 arxiv.org

拡散モデル

　この論文で考える「拡散モデル」はマルコフ連鎖モデル（未来の状態は現在の状態だけで決まるモデル）である（下図参照）。

1番左の状態 $x_T$ はガウスノイズ画像、1番右の状態 $x_0$ は物体が明確に写っている画像である。いま、 $x_0$ の状態にわずかなガウスノイズを載せた画像を作り、それを状態 $x_1$ とする。これを $T$ 回繰り返すと $x_T$ になるとする。 $x_0$ から $x_T$ に至る過程をForward過程（拡散過程）と呼ぶ。一方、 $x_T$ から $x_0$ に至る逆向きの過程も考えることができる。この過程は、ガウスノイズ画像が徐々にきれいな画像に変化していく過程である。この過程をReverse過程（生成過程）と呼ぶ。

　拡散過程の各ステップでは画像にガウスノイズを載せるので、 $q(x_t|x_{t-1})$ を次式で与えることができる。

${\displaystyle \begin{align} q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_tI) \label{eq0} \end{align} }$

ここで、 $x_t$ はベクトル、 $I$ は単位行列、 $\beta_t\in(0,1)$ である。本論文では $\beta_t$ （ハイパーパラメータ）として $10^{-4}\sim10^{-2}$ 程度の値を与えている。 $\beta_t$ が小さな値のとき生成過程 $p_{\theta}(x_{t-1}|x_t)$ もガウス分布関数で近似することができる。

${\displaystyle \begin{equation} p_{\theta}(x_{t-1}|x_{t})=\mathcal{N}\left(x_{t-1};\mu_{\theta}(x_t,t),\sigma_t^2I\right) \end{equation} }$

ここで、 $\mu_{\theta}(x_t,t)$ は $t$ と $x_t$ から計算される平均ベクトル、 $\sigma_{t}^2$ （ハイパーパラメータ）は分散である。拡散モデルでは、 $\mu_{\theta}(x_t,t)$ をニューラルネットワークを用いて学習する。

変分推論

　学習を行うには最小化する関数が必要である。拡散モデルでは次式を最小化する。

${\displaystyle \begin{align} -\ln{p_{\theta}(x_0)}=-\ln{\int dx_{1:T}p_{\theta}(x_{0:T})} \label{eq1} \end{align} }$

ここで、 $dx_{1:T}$ は $dx_1dx_2\cdots dx_T$ を表し、 $p_{\theta}(x_{0:T})$ は $p_{\theta}(x_{0},x_1,\cdots,x_T)$ を表す。マルコフ連鎖モデルなので

${\displaystyle \begin{equation} p_{\theta}(x_{0:T})=p(x_T)\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t) \end{equation} }$

と書くことができる。これを\eqref{eq1}に代入すると

${\displaystyle \begin{align*} -\ln{p_{\theta}(x_0)}=-\ln{\int dx_{1:T}p(x_T)\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t)} \end{align*} }$

を得る。ここで、 $q(x_{1:T}|x_0)$ を用いて上式を書き換える。

${\displaystyle \begin{align*} -\ln{p_{\theta}(x_0)}&=-\ln{\int dx_{1:T}p(x_T)\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t)}\\ &=-\ln{\int dx_{1:T}q(x_{1:T}|x_0)\dfrac{p(x_T)}{q(x_{1:T}|x_0)}\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t)}\\ &\leq-\int dx_{1:T}q(x_{1:T}|x_0)\ln{\dfrac{p(x_T)}{q(x_{1:T}|x_0)}\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t)}\equiv L_1 \end{align*} }$

3行目でイェンセンの不等式を用いた。左辺を最小化するには右辺（ $L_1$ ）を最小化すれば良いので、これ以降の議論では $L_1$ の最小化を考える。 $L_1$ を変形すると

${\displaystyle \begin{align} L_1&=-\int dx_{1:T}q(x_{1:T}|x_0)\ln{\dfrac{p(x_T)}{q(x_{1:T}|x_0)}\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t)}\notag \\ &=-\int dx_{1:T}q(x_{1:T}|x_0)\ln{p(x_T)}-\int dx_{1:T}q(x_{1:T}|x_0)\ln{\dfrac{\prod_{t=1}^Tp_{\theta}(x_{t-1}|x_t)}{q(x_{1:T}|x_0)}} \label{eq2} \end{align} }$

$q(x_{1:T}|x_0)$ は以下のように書くことができる。

${\displaystyle \begin{align*} q(x_{1:T}|x_0)&=q(x_1,\cdots,x_T|x_0)\\ &=q(x_{T}|x_{T-1})\cdots q(x_{1}|x_{0})\\ &=\prod_{t=1}^{T}q(x_{t}|x_{t-1}) \end{align*} }$

これを式\eqref{eq2}に代入すると

${\displaystyle \begin{align*} L_1&=-\int dx_{1:T}q(x_{1:T}|x_0)\ln{p(x_T)}-\int dx_{1:T}q(x_{1:T}|x_0)\ln{\prod_{t=1}^T\dfrac{p_{\theta}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1})}}\\ &=-\int dx_{1:T}q(x_{1:T}|x_0)\ln{p(x_T)}-\int dx_{1:T}q(x_{1:T}|x_0)\sum_{t=1}^T\ln{\dfrac{p_{\theta}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1})}} \end{align*} }$

を得る。上式の両辺に $q(x_0)$ をかけて $x_0$ で積分すると

${\displaystyle \begin{align*} \int dx_0q(x_0)L_1&=-\int dx_{0:T}q(x_{0:T})\ln{p(x_T)}-\int dx_{0:T}q(x_{0:T})\sum_{t=1}^T\ln{\dfrac{p_{\theta}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1})}}\equiv L_2 \end{align*} }$

となる。いま、期待値

${\displaystyle \begin{align*} E_{q(0:T)}[f]=\int dx_{0:T}q(x_{0:T})\;f \end{align*} }$

を定義すると、上式の右辺は

${\displaystyle \begin{align*} L_2=-E_{q(0:T)}\bigl[\ln{p(x_T)}\bigr]-E_{q(0:T)}\Bigl[\sum_{t=1}^T\ln{\dfrac{p_{\theta}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1})}}\Bigr] \end{align*} }$

と書くことができる。 $L_2$ を少し変形して

${\displaystyle \begin{align} L_2&=-E_{q(0:T)}\bigl[\ln{p(x_T)}\bigr]-E_{q(0:T)}\Bigl[\ln{\dfrac{p_{\theta}(x_{0}|x_1)}{q(x_{1}|x_{0})}}\Bigr]-E_{q(0:T)}\Bigl[\sum_{t=2}^T\ln{\dfrac{p_{\theta}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1})}}\Bigr] \label{eq3} \end{align} }$

を得る。ここで

${\displaystyle \begin{align*} q(x_t|x_{t-1})=q(x_t|x_{t-1},x_0)=\dfrac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)} \end{align*} }$

が成り立つので（後半の等式はベイズの定理）これを式\eqref{eq3}の右辺第3項に代入すると

${\displaystyle \begin{align} E_{q(0:T)}\Bigl[\sum_{t=2}^T\ln{\dfrac{p_{\theta}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1})}}\Bigr] &=E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } \dfrac{ q(x_{t-1}|x_0) }{ q(x_{t}|x_0) } } \Bigr]\notag\\ &= E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } } \Bigr] + E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ q(x_{t-1}|x_0) }{ q(x_{t}|x_0) } } \Bigr] \label{eq4} \end{align} }$

となる。式\eqref{eq4}の第2項は

${\displaystyle \begin{align*} E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ q(x_{t-1}|x_0) }{ q(x_{t}|x_0) } } \Bigr] &= E_{q(0:T)} \Bigl[ \ln{ \dfrac{ q(x_{1}|x_0) }{ q(x_{2}|x_0) }\dfrac{ q(x_{2}|x_0) }{ q(x_{3}|x_0) }\cdots\dfrac{ q(x_{T-1}|x_0) }{ q(x_{T}|x_0) } } \Bigr]\\ &= E_{q(0:T)} \Bigl[ \ln{ \dfrac{ q(x_{1}|x_0) }{ q(x_{T}|x_0) } } \Bigr] \end{align*} }$

と変形される。従って

${\displaystyle \begin{align*} L_2=-E_{q(0:T)}\bigl[\ln{p(x_T)}\bigr]-E_{q(0:T)}\Bigl[\ln{\dfrac{p_{\theta}(x_{0}|x_1)}{q(x_{1}|x_{0})}}\Bigr]-E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } } \Bigr] -E_{q(0:T)} \Bigl[ \ln{ \dfrac{ q(x_{1}|x_0) }{ q(x_{T}|x_0) } } \Bigr] \end{align*} }$

を得る。右辺第2項と第4項をまとめて

${\displaystyle \begin{align*} L_2= -E_{q(0:T)}\bigl[\ln{p(x_T)}\bigr] -E_{q(0:T)}\Bigl[\ln{\dfrac{p_{\theta}(x_{0}|x_1)}{q(x_{T}|x_{0})}}\Bigr] -E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } } \Bigr] \end{align*} }$

右辺第1項と第2項を変形して

${\displaystyle \begin{align} L_2= -E_{q(0:T)}\bigl[\ln{p_{\theta}(x_0|x_1)}\bigr] -E_{q(0:T)}\Bigl[\ln{\dfrac{p_{\theta}(x_{T})}{q(x_{T}|x_{0})}}\Bigr] -E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } } \Bigr] \label{eq5} \end{align} }$

を得る。式\eqref{eq5}の右辺第2項を変形すると

${\displaystyle \begin{align*} E_{q(0:T)}\Bigl[\ln{\dfrac{p_{\theta}(x_{T})}{q(x_{T}|x_{0})}}\Bigr] &=\int dx_0dx_1\cdots dx_Tq(x_0,x_1,\cdots,x_T)\ln{\dfrac{p_{\theta}(x_{T})}{q(x_{T}|x_{0})}}\\ &=\int dx_0dx_T\Bigl[\int dx_1\cdots dx_{T-1}q(x_0,x_1,\cdots,x_T)\Bigr]\ln{\dfrac{p_{\theta}(x_{T})}{q(x_{T}|x_{0})}}\\ &=\int dx_0dx_Tq(x_0,x_T)\ln{\dfrac{p_{\theta}(x_{T})}{q(x_{T}|x_{0})}}\\ &=\int dx_0q(x_0)\Bigl[\int dx_Tq(x_T|x_0)\ln{\dfrac{p_{\theta}(x_{T})}{q(x_{T}|x_{0})}}\Bigr]\\ &=-\int dx_0q(x_0)D_{KL}\bigl(q(x_T|x_0)\;\|\;p(x_T)\bigr) \end{align*} }$

式\eqref{eq5}の右辺第3項を変形すると

${\displaystyle \begin{align*} E_{q(0:T)} \Bigl[ \sum_{t=2}^T \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } } \Bigr] &= \sum_{t=2}^T \int dx_0\cdots dx_T q(x_0,\cdots,x_T) \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } }\\ &= \sum_{t=2}^T \int dx_0dx_{t-1}dx_t q(x_0,x_{t-1},x_t) \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } }\\ &= \sum_{t=2}^T \int dx_0dx_{t}q(x_0,x_t)\int dx_{t-1} q(x_{t-1}|x_0,x_t) \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } }\\ &= -\sum_{t=2}^T \int dx_0dx_{t}q(x_0,x_t) D_{KL}\bigl(q(x_{t-1}|x_t,x_0)\;\|\;p_{\theta}(x_{t-1}|x_t) \bigr) \end{align*} }$

を得る。いま

${\displaystyle \begin{align*} E_{q(0)}[f]&\equiv \int dx_0 q(x_0)\;f\\ E_{q(0,t)}[f]&\equiv \int dx_0dx_t q(x_0,x_t)\;f \end{align*} }$

を定義すると式\eqref{eq5}は次式となる。

${\displaystyle \begin{align} L_2= -E_{q(0:T)}\bigl[\ln{p_{\theta}(x_0|x_1)}\bigr] &+E_{q(0)}\Bigl[D_{KL}\bigl(q(x_T|x_0)\;\|\;p(x_T)\bigr)\Bigr]\notag\\ &+\sum_{t=2}^TE_{q(0,t)}\Bigl[D_{KL}\bigl(q(x_{t-1}|x_t,x_0)\;\|\;p_{\theta}(x_{t-1}|x_t)\bigr)\Bigr] \label{eq7} \end{align} }$

この式が論文中の式(5)である。

$q(x_t|x_0)$ の計算

　式\eqref{eq0}で $\alpha_t=1-\beta_t$ とおく。

${\displaystyle \begin{align*} q(x_t|x_{t-1})=\mathcal{N}\bigl(x_t;\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I\bigr) \end{align*} }$

いま

${\displaystyle \begin{align*} q(x_1|x_0)&=\mathcal{N}\bigl(x_1;\sqrt{\alpha_1}x_{0},(1-\alpha_1)I\bigr)\\ q(x_2|x_1)&=\mathcal{N}\bigl(x_2;\sqrt{\alpha_2}x_1,(1-\alpha_2)I\bigr) \end{align*} }$

を考え、次式を計算する。

${\displaystyle \begin{align} q(x_2|x_0)&=\int dx_1q(x_2|x_1)q(x_1|x_0)\notag\\ &=\int dx_1\mathcal{N}\bigl(x_2;\sqrt{\alpha_2}x_1,(1-\alpha_2)I\bigr)\mathcal{N}\bigl(x_1;\sqrt{\alpha_1}x_{0},(1-\alpha_1)I\bigr)\notag\\ &\propto\int dx_1 \exp{ \Bigl\{ -\dfrac{1}{2} \Bigl[ \dfrac{\left(x_1-\sqrt{\alpha_1}x_0 \right)^T\left(x_1-\sqrt{\alpha_1}x_0\right)}{1-\alpha_1} \Bigr. \Bigr.\\ \Bigl. \Bigl. +\dfrac{\left(x_2-\sqrt{\alpha_2}x_1\right)^T\left(x_2-\sqrt{\alpha_2}x_1\right)}{1-\alpha_2} \Bigr] \Bigr\} } \label{eq6} \end{align} }$

3行目は $x_1$ に依存する項だけを残した。指数関数の肩の $[\cdots$ ]に注目して

${\displaystyle \begin{align*} [\cdots]=a\;x_1^Tx_1-b^Tx_1-x_1^Tb+\dfrac{\alpha_1}{1-\alpha_1}x_0^Tx_0+\dfrac{1}{1-\alpha_2}x_2^Tx_2 \end{align*} }$

ここで

${\displaystyle \begin{align*} a&\equiv\dfrac{1}{1-\alpha_1}+\dfrac{\alpha_2}{1-\alpha_2}\\ b&\equiv\dfrac{\sqrt{\alpha_1}}{1-\alpha_1}x_0+\dfrac{\sqrt{\alpha_2}}{1-\alpha_2}x_2 \end{align*} }$

と置いた。 $a$ はスカラー、 $b$ はベクトルである。さらに計算すると

${\displaystyle \begin{align*} [\cdots]=a\left(x_1-\dfrac{b}{a}\right)^T\left(x_1-\dfrac{b}{a}\right)-\dfrac{b^Tb}{a} +\dfrac{\alpha_1}{1-\alpha_1}x_0^Tx_0 +\dfrac{1}{1-\alpha_2}x_2^Tx_2 \end{align*} }$

を得る。上式を式\eqref{eq6}に代入し、 $x_1$ について積分すると

${\displaystyle \begin{align*} q(x_2|x_0)&\propto \exp{\Bigl\{-\dfrac{1}{2}\Bigl[ -\dfrac{b^Tb}{a} +\dfrac{\alpha_1}{1-\alpha_1}x_0^Tx_0 +\dfrac{1}{1-\alpha_2}x_2^Tx_2 \Bigr]}\Bigr\}\\ &\propto \exp{\Bigl\{-\dfrac{1}{2}\Bigl[ \dfrac{\left(x_2-\sqrt{\alpha_1\alpha_2}x_0\right)^T\left(x_2-\sqrt{\alpha_1\alpha_2}x_0\right)}{1-\alpha_1\alpha_2} \Bigr]}\Bigr\}\\ \end{align*} }$

となる。2行目への変形では $x_2$ に依存する項だけを残した。以上より

${\displaystyle \begin{align*} q(x_2|x_0)=\mathcal{N}\bigl(x_2;\sqrt{\bar{\alpha}_2}\;x_{0},(1-\bar{\alpha}_2)I\bigr) \end{align*} }$

が成り立つことがわかる。ここで

${\displaystyle \begin{align*} \bar{\alpha}_2=\alpha_1\alpha_2 \end{align*} }$

とした。同じ計算を繰り返していけば

${\displaystyle \begin{align} q(x_t|x_0)&=\mathcal{N}\bigl(x_t;\sqrt{\bar{\alpha}_t}\;x_{0},(1-\bar{\alpha}_t)I\bigr)\label{eq9}\\ \bar{\alpha}_t&=\prod_{s=1}^t\alpha_s\notag \end{align} }$

を得る。これが論文の式(4)である。また、 $t$ を大きくしていくと $\bar{\alpha}_t$ は小さくなるので $q(x_t|x_0)$ は標準正規分布に近づいていくことが分かる。

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

　ベイズの定理より

${\displaystyle \begin{align*} q(x_{t-1}|x_t,x_0)=\dfrac{q(x_{t-1}|x_0)q(x_t|x_{t-1},x_0)}{q(x_t|x_0)} \end{align*} }$

ここに\eqref{eq0}と\eqref{eq9}を代入する。

${\displaystyle \begin{align*} q(x_{t-1}|x_t,x_0)&=\dfrac{\mathcal{N}\bigl(x_{t-1};\sqrt{\bar{\alpha}_{t-1}}\;x_{0},(1-\bar{\alpha}_{t-1})I\bigr)\mathcal{N}(x_t;\sqrt{\alpha_t}x_{t-1},\beta_tI) }{\mathcal{N}\bigl(x_t;\sqrt{\bar{\alpha}_t}\;x_{0},(1-\bar{\alpha}_t)I\bigr)}\\ &\propto \exp{\Bigl\{-\dfrac{1}{2}\Bigl[ \dfrac{\left(x_{t-1}-\sqrt{\bar{\alpha}_{t-1}}\;x_0\right)^T \left(x_{t-1}-\sqrt{\bar{\alpha}_{t-1}}\;x_0 \right)}{1-\bar{\alpha}_{t-1}} \\ \;\;\;\;\;\;+\dfrac{\left(x_{t}-\sqrt{\alpha_{t}}\;x_{t-1}\right)^T \left(x_{t}-\sqrt{\alpha_{t}}\;x_{t-1} \right)}{\beta_t} \Bigr]\Bigr\}} \end{align*} }$

2行目では $x_{t-1}$ に依存する項だけを残した。指数関数の肩の $[\cdots$ ]に注目して

${\displaystyle \begin{align*} [\cdots]=a\;x_{t-1}^Tx_{t-1}-b^Tx_{t-1}-x_{t-1}^Tb+\cdots \end{align*} }$

ここで

${\displaystyle \begin{align*} a&\equiv\dfrac{1}{1-\bar{\alpha}_{t-1}}+\dfrac{\alpha_t}{\beta_t}\\ b&\equiv\dfrac{\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}}x_0+\dfrac{\sqrt{\alpha_t}}{\beta_t}x_{t} \end{align*} }$

と置いた。 $a$ はスカラー、 $b$ はベクトルである。さらに変形すると

${\displaystyle \begin{align*} [\cdots]=a\left(x_{t-1}-\dfrac{b}{a}\right)^T\left(x_{t-1}-\dfrac{b}{a}\right)+\cdots \end{align*} }$

ここに上の $a$ と $b$ を代入すると

${\displaystyle \begin{align*} [\cdots]&=\dfrac{1-\bar{\alpha}_t}{(1-\bar{\alpha}_{t-1})\beta_t} \left(x_{t-1}-\left(\dfrac{\beta_t\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0+\dfrac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}x_t\right)\right)^T\\ &\;\;\;\;\;\;\;\;\;\;\;\;\times \left(x_{t-1}-\left(\dfrac{\beta_t\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0+\dfrac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}x_t\right)\right)\\ &=\dfrac{\left(x_{t-1}-\tilde{\mu}_t(x_t,x_0)\right)^T\left(x_{t-1}-\tilde{\mu}_t(x_t,x_0)\right)}{\tilde{\beta}_t} \end{align*} }$

となる。ここで

${\displaystyle \begin{align} \tilde{\mu}_t(x_t,x_0)&\equiv \dfrac{\beta_t\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0+\dfrac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}x_t\label{eq11}\\ \tilde{\beta}_t&\equiv \dfrac{(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\beta_t\notag \end{align} }$

と置いた。上式を指数関数の肩に戻すと

${\displaystyle \begin{align*} q(x_{t-1}|x_t,x_0) &\propto \exp{\Bigl\{-\dfrac{1}{2}\Bigl[ \dfrac{\left(x_{t-1}-\tilde{\mu}_t(x_t,x_0)\right)^T\left(x_{t-1}-\tilde{\mu}_t(x_t,x_0)\right)}{\tilde{\beta}_t} \Bigr]\Bigr\}} \end{align*} }$

となる。従って

${\displaystyle \begin{align*} q(x_{t-1}|x_t,x_0) =\mathcal{N}\left(x_{t-1};\tilde{\mu}_t(x_t,x_0),\tilde{\beta}_tI\right) \end{align*} }$

を得る。これらが論文中の式(6)と式(7)に相当する。

最小化すべき関数の計算

　最小化すべき関数\eqref{eq7}は以下であった。

${\displaystyle \begin{align*} L_2= -E_{q(0:T)}\bigl[\ln{p_{\theta}(x_0|x_1)}\bigr] &+E_0\Bigl[D_{KL}\bigl(q(x_T|x_0)\;\|\;p(x_T)\bigr)\Bigr]\notag\\ &+\sum_{t=2}^TE_{0,t}\Bigl[D_{KL}\bigl(q(x_{t-1}|x_t,x_0)\;\|\;p_{\theta}(x_{t-1}|x_t)\bigr)\Bigr] \end{align*} }$

右辺第2項は学習から決めるパラメータ $\theta$ に依存しないので無視できる。

${\displaystyle \begin{align} L_3\equiv -E_{q(0:T)}\bigl[\ln{p_{\theta}(x_0|x_1)}\bigr] +\sum_{t=2}^TE_{0,t}\Bigl[D_{KL}\bigl(q(x_{t-1}|x_t,x_0)\;\|\;p_{\theta}(x_{t-1}|x_t)\bigr)\Bigr] \label{eq16} \end{align} }$

最初に $L_3$ の第2項を考える。期待値の中のKullback Leibler divergenceを取り出して

${\displaystyle \begin{align*} D_{KL}\bigl(q(x_{t-1}|x_t,x_0)\;\|\;p_{\theta}(x_{t-1}|x_t)\bigr)= -\int dx_{t-1} q(x_{t-1}|x_0,x_t) \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } } \end{align*} }$

ここに

${\displaystyle \begin{align*} q(x_{t-1}|x_t,x_0)&=\mathcal{N}(x_{t-1};\tilde{\mu}_t(x_t,x_0),\tilde{\beta}_tI)\\ p_{\theta}(x_{t-1}|x_t)&=\mathcal{N}(x_{t-1};\mu_{\theta}(x_t,x_0),\sigma_t^2I) \end{align*} }$

を代入すると

${\displaystyle \begin{align} &-\int dx_{t-1} q(x_{t-1}|x_0,x_t) \ln{ \dfrac{ p_{\theta}(x_{t-1}|x_t) }{ q(x_{t-1}|x_{t},x_0) } }\notag\\ &=\int dx_{t-1} q(x_{t-1}|x_0,x_t)\Bigl\{\dfrac{1}{2}\Bigl[ \dfrac{\left(x_{t-1}-\mu_{\theta}\right)^T\left(x_{t-1}-\mu_{\theta}\right)}{\sigma_t^2} -\dfrac{\left(x_{t-1}-\tilde{\mu}_{t}\right)^T\left(x_{t-1}-\tilde{\mu}_{t}\right)}{\tilde{\beta}_t} \Bigr] \Bigr\}+C \label{eq10} \end{align} }$

となる。 $C$ は $x_{t-1}$ に依存しない項である。 $\tilde{\beta}_t=\sigma_t^2$ の場合を考えて $\{\cdots\}$ の中を計算すると

${\displaystyle \begin{align*} \{\cdots\}&= \dfrac{1}{2\sigma_t^2}\Bigl[ \left(x_{t-1}-\mu_{\theta}\right)^T\left(x_{t-1}-\mu_{\theta}\right) -\left(x_{t-1}-\tilde{\mu}_{t}\right)^T\left(x_{t-1}-\tilde{\mu}_{t}\right) \Bigr]\\ &= \dfrac{1}{2\sigma_t^2}\Bigl[ -\left(x_{t-1}-\tilde{\mu}_{t}\right)^T\left(\mu_{\theta}-\tilde{\mu}_{t}\right) -\left(\mu_{\theta}-\tilde{\mu}_{t}\right)^T\left(x_{t-1}-\tilde{\mu}_{t}\right) +\left(\tilde{\mu}_t-\mu_{\theta}\right)^T\left(\tilde{\mu}_t-\mu_{\theta}\right) \Bigr] \end{align*} }$

上式を\eqref{eq10}に戻して $x_{t-1}$ について積分すると、上式の第1項と第2項はガウス関数と掛け算されて奇関数となるので積分の寄与はゼロになる。以上より

${\displaystyle \begin{align*} D_{KL}\bigl(q(x_{t-1}|x_t,x_0)\;\|\;p_{\theta}(x_{t-1}|x_t)\bigr)= \int dx_{t-1} q(x_{t-1}|x_0,x_t) \Bigl[\dfrac{1}{2\sigma_t^2}\left(\tilde{\mu}_t-\mu_{\theta}\right)^T\left(\tilde{\mu}_t-\mu_{\theta}\right)\Bigr] \end{align*} }$

を得る。ただし定数項は落とした。 $L_3$ の第2項に戻すと

${\displaystyle \begin{align*} \sum_{t=2}^T&E_{q(0,t)}\Bigl[ \int dx_{t-1} q(x_{t-1}|x_0,x_t) \Bigl[\dfrac{1}{2\sigma_t^2}\left(\tilde{\mu}_t-\mu_{\theta}\right)^T\left(\tilde{\mu}_t-\mu_{\theta}\right)\Bigr]\Bigr]\\ &= \sum_{t=2}^T \int dx_0dx_tq(x_0,x_t)\int dx_{t-1} q(x_{t-1}|x_0,x_t) \Bigl[\dfrac{1}{2\sigma_t^2}\left(\tilde{\mu}_t-\mu_{\theta}\right)^T\left(\tilde{\mu}_t-\mu_{\theta}\right)\Bigr]\\ &= \sum_{t=2}^T \int dx_{0:t}q(x_{0:t}) \Bigl[\dfrac{1}{2\sigma_t^2}\left(\tilde{\mu}_t-\mu_{\theta}\right)^T\left(\tilde{\mu}_t-\mu_{\theta}\right)\Bigr]\\ &= \sum_{t=2}^T E_{q(0:t)}\Bigl[ \dfrac{1}{2\sigma_t^2}\left(\tilde{\mu}_t-\mu_{\theta}\right)^T\left(\tilde{\mu}_t-\mu_{\theta}\right) \Bigr]\equiv L_4 \end{align*} }$

を得る。これが論文中の式(8)である。先に求めた式

${\displaystyle \begin{align*} q(x_t|x_0)=\mathcal{N}\bigl(x_t;\sqrt{\bar{\alpha}_t}\;x_{0},(1-\bar{\alpha}_t)I\bigr) \end{align*} }$

に再パラメータ化トリックを適用して

${\displaystyle \begin{align} x_t&=\sqrt{\bar{\alpha}_t}\;x_0+\sqrt{1-\bar{\alpha}_t}\;\epsilon\label{eq14}\\ \epsilon&\sim\mathcal{N}(\epsilon;0,I)\notag \end{align} }$

これを $x_0$ について解くと

${\displaystyle \begin{align} x_0=\frac{1}{\sqrt{\bar{\alpha}}_t}\left(x_t-\sqrt{1-\bar{\alpha}_t}\;\epsilon\right) \label{eq12} \end{align} }$

となる。ここで、上で求めた式\eqref{eq11}を考える。

${\displaystyle \begin{align*} \tilde{\mu}_t= \dfrac{\beta_t\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0+\dfrac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}x_t \end{align*} }$

右辺の $x_0$ に式\eqref{eq12}を代入すると

${\displaystyle \begin{align*} \tilde{\mu}_t&= \dfrac{\beta_t\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_t}x_0+\dfrac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}x_t\\ &= \frac{1}{\sqrt{\alpha_t}}\left( x_t(x_0,\epsilon)-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon \right) \end{align*} }$

となる。従って $L_4$ は

${\displaystyle \begin{align} L_4= \sum_{t=2}^T E_{q(0:t)} \left[ \dfrac{1}{2\sigma_t^2} \left\| \frac{1}{\sqrt{\alpha_t}}\left( x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon \right) -\mu_{\theta} \right\|^2 \right]\label{eq13} \end{align} }$

となる。これが論文中の式(10)である。上式が最小になるのは

${\displaystyle \begin{align*} \mu_{\theta} =\frac{1}{\sqrt{\alpha_t}}\left( x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon \right) \end{align*} }$

のときである。 $\mu_{\theta}$ は学習で決めるパラメータであるから上式右辺には学習で決める項がなければならない。そこで、新たな量 $\epsilon_{\theta}$ を導入する。

${\displaystyle \begin{align*} \mu_{\theta} =\frac{1}{\sqrt{\alpha_t}}\left( x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(x_t,t) \right) \end{align*} }$

$\epsilon_{\theta}$ は $x_t$ と $t$ から計算される関数である。上式を式\eqref{eq13}に代入すると

${\displaystyle \begin{align} L_4 = \sum_{t=2}^T E_{q(0:t)}\left[ \dfrac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}_t)} \left\| \epsilon-\epsilon_{\theta}(x_t,t) \right\|^2 \right] \label{eq15} \end{align} }$

を得る。ここで、 $x_t$ は式\eqref{eq14}で与えられていることに注意する（次式）。

${\displaystyle \begin{align*} x_t&=\sqrt{\bar{\alpha}_t}\;x_0+\sqrt{1-\bar{\alpha}_t}\;\epsilon\\ \epsilon&\sim\mathcal{N}(\epsilon;0,I)\notag \end{align*} }$

式\eqref{eq15}は、標準正規分布で生成したノイズ $\epsilon$ を $x_t$ と $t$ を用いて予測する目的関数になっている。式\eqref{eq15}を実際に計算する際は、 $x_0$ と $\epsilon$ でサンプリングして期待値を取る。

${\displaystyle \begin{align*} L_4 = \sum_{t=2}^T E_{x_0,\epsilon}\left[ \dfrac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}_t)} \left\| \epsilon-\epsilon_{\theta}(x_t,t) \right\|^2 \right] \end{align*} }$

さらに、ノルムの係数も落とし、ステップ $t$ についてもサンプリングを取る。

${\displaystyle \begin{align*} L_4 &= E_{t,x_0,\epsilon}\left[ \left\| \epsilon-\epsilon_{\theta}(x_t,t) \right\|^2 \right]\\ &= E_{t,x_0,\epsilon}\left[ \left\| \epsilon-\epsilon_{\theta}(\sqrt{\bar{\alpha}_t}\;x_0+\sqrt{1-\bar{\alpha}_t}\;\epsilon,t) \right\|^2 \right] \end{align*} }$

これが論文中の式(14)である。最後に式\eqref{eq16}の第1項を考える。

${\displaystyle \begin{align*} -E_{q(0:T)}\bigl[\ln{p_{\theta}(x_0|x_1)}\bigr] &= -\int dx_{0:T}q(x_{0:T})\ln{p_{\theta}(x_0|x_1)}\\ &= -\int dx_1q(x_1)\int dx_{0}q(x_{0}|x_1)\ln{p_{\theta}(x_0|x_1)} \end{align*} }$

$p_{\theta}(x_0|x_1)$ は、生成過程の一番最後のステップである。 $x_0$ を各画素から構成される $D$ 次元のベクトルと考える。

${\displaystyle \begin{align*} &-\int dx_{0}q(x_{0}|x_1)\ln{p_{\theta}(x_0|x_1)}\\ &= -\int dx_0^1\cdots dx_0^D\;q(x_0^1,\cdots,x_0^D|x_1) \ln{p_{\theta}(x_0^1,\cdots,x_0^D|x_1)}\\ &= -\int dx_0^1\cdots dx_0^D\;q(x_0^1,\cdots,x_0^D|x_1) \sum_{i=1}^D\ln{p_{\theta}(x_0^i|x_1)}\\ &= -\sum_{i=1}^D\int dx_0^i\;q(x_0^i|x_1) \ln{p_{\theta}(x_0^i|x_1)}\\ &= -\sum_{i=1}^D\int dx_0^i\;q(x_0^i|x_1) \ln{\mathcal{N}(x_0^i;\mu_{\theta}^i(x_1,1),\sigma_1^2})\\ &= -\sum_{i=1}^D\int_{\mu_{\theta}^i-\delta}^{\mu_{\theta}^i+\delta} dx_0^i \ln{\mathcal{N}(x_0^i;\mu_{\theta}^i(x_1,1),\sigma_1^2}) \label{eq17} \end{align*} }$

最終行では積分範囲を平均値 $\mu_{\theta}^i$ の近傍に制限した（下図参照）。これが論文中の式(13)である（恐らく）。

まとめ

　拡散モデルの2020年の論文Denoising Diffusion Probabilistic Modelsの式変形を追ってみた。恐らく誤りや勘違いがあると思うので指摘いただければ幸いです。

Tech Blog

勉強したことをまとめます。