　お久しぶりです．ソリングです．ソリングは大学院では自己教師あり学習に関するマルチモーダル手法（特に，Lip SyncingやLip Reading）について研究していますが，それとは関係なく生成モデル，特に拡散モデルについて興味を抱いています．そこで，本記事は拡散モデルを理解するというテーマをもとに執筆します．

　本記事はあくまで自己流の理解です．もしかしたら内容に不明瞭な記述や誤りを含んでいるかもしれませんので，発見された場合はお手数ですがコメントお願いします．

拡散モデルとは？

　近年話題になっている生成モデル(Generative Model)においては，拡散モデル(Diffusion Model)が主に使用されています．このモデルの良い点として，従来用いられていた生成的敵対ニューラルネットワーク（GAN; Generative Adversarial Network）の欠点である「学習が不安定である」点や，変分オートエンコーダ（VAE; Variational Auto Encoder）の「出力された画像がぼやけやすい」欠点を解消した点が挙げられます．

　本記事では，拡散モデルを数式的・お気持ち的療法の両方から解説します．

拡散モデルのお気持ち

　お気持ちを一言で伝えると「加えたノイズを取り除く」です．~~あるある探検隊のリズムですね．~~

　実際に画像を使って説明します．画像に少しずつノイズを加えると，左の台湾まぜそばの画像から「ちょっとノイズが乗っかった画像」「だいぶノイズが加わった画像」「完全なノイズ」が生成されます．拡散モデルの目的は，一番右のノイズから一番左の台湾まぜそばを復元するために取り除くべきノイズを学習することです．

　ポイントは，ノイズを少しずつ加えていくということです．後述しますが，少しずつ加えていくことにより，数式的に操作が扱いやすくなります．

数式を用いた理解

ノイズを少しずつ加えていく過程のことを「拡散過程」と呼び，ノイズを取り除いていく過程のことを「逆拡散過程」と呼びます．ニューラルネットワークでは逆拡散過程のパラメーターを学習することになります．これを順番に見ていきましょう．

拡散過程

拡散過程ではノイズを加えていきます．Diffusion Modelはノイズの「取り除き方」を学ぶモデルですので，Modelはこの課程においてはパラメーターを持っていません．

入力のベクトル$x_0 \in \mathbb{R}^d$に対し，合計で$T$回ノイズを少しずつ加えていくことを考えます．

これに平均0分散1の独立な多変数正規分布(実際は正規分布でなくても良いかもしれませんが，数学的な扱いやすさの観点から正規分布としています)に従う$\epsilon_{t-1}\in \mathbb{R}^d$を足していきます．数式で書くと以下のとおりです．

\begin{equation}
\label{kakusanhouteisiki}
x_{t}=\sqrt{1-\beta_t} x_{t-1}+\sqrt{\beta_t}\epsilon_{t-1} \quad (t=1,2,\cdots, T) \tag{1}
\end{equation}

$\beta_t$とはノイズの加え方を指定する値のことです．この値が大きいほど1度に加えられるノイズが大きくなりますが，一般的には0.01のような小さな値が設定されます．$x_{t}$は$x_{t-1}$のみに依存し，$x_{t-2}$などそれ以外のステップの情報は必要としません（$x_{t-1}$は$x_{t-2}$に従うので，結果として$x_t$も$x_{t-2}$に依存しますが，$x_t$を得るためには$x_{t-1}$だけでOKです．）このような過程のことをマルコフ過程といいます．

さて，式(1)は，そのままでは$T$回の代入計算を行わなければならず，少しばかり大変です．そのため，以下の方法を用いて計算を簡略化することができます．ただし，$\alpha_t+\beta_t = 1$です．

x_{t}=x_{0}\prod^{t}_{t'=1} \sqrt{\alpha_{t'}} + \epsilon\sqrt{1-\prod^{t}_{t'=1} \alpha_{t'}}, =: x_0 \sqrt{\bar{\alpha}} + \epsilon\sqrt{1-\bar{\alpha}}  \tag{2} A

証明は記事が冗長にならないように折りたたみます．

証明

正規分布の再生性，つまり独立な正規分布の和はまた正規分布に従い，その平均，分散はそれぞれ元の平均，分散の和に一致するという性質を用います．

\begin{align*}
x_t &= \sqrt{\alpha_t} x_{t-1} + \sqrt{\beta_t} \epsilon_{t-1}\\
&= \sqrt{\alpha_t} (\sqrt\alpha_{t-1} x_{t-2}+\sqrt{\alpha_{t-1}} \epsilon_{t-2} )+ \sqrt{\beta_t} \epsilon_{t-1}\\
&= \sqrt{\alpha_t} (\sqrt{\alpha_{t-1}} x_{t-2} + \sqrt{1-\alpha_{t-1}} \epsilon_{t-2})+ \sqrt{1-\alpha_t} \epsilon_{t-1}\\
&= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{\alpha_t - \alpha_t \alpha_{t-1}} \epsilon_{t-2} + \sqrt{1-\alpha_t} \epsilon_{t-1} \\
&= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{\alpha_t - \alpha_t\alpha_{t-1}+1-\alpha_t} \epsilon\\
&= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1- \alpha_t\alpha_{t-1}} \epsilon & (正規分布の再生性)\\
&= \cdots = x_{0}\prod^{t}_{t'=1} \sqrt{\alpha_{t'}} + \epsilon\sqrt{1-\prod^{t}_{t'=1} \alpha_{t'}}
\end{align*}

とできます．厳密には数学的帰納法を用いると良いでしょう．

式(2)の形を見るとわかるように，$t$が大きくなるに連れて徐々に$x_t$がノイズになっていきます．たとえば，$\beta_t$が等差数列だとして，$T=1000$で，$\beta_1=1\times10^{-4}, \beta_{1000}=0.02$とすると，

\begin{align*}
\prod^{T}_{t=1} \sqrt{\alpha_{t}}\approx 0.006,\,\,  \sqrt{1-
\prod^{T}_{t=1}\alpha_t}\approx 0.99997
\end{align*}

となり，元の画像の項はほぼ消え，殆どがノイズに由来する項から来ていることがわかります．次の項では，このノイズからもとの画像を復元する方法について解説します．

逆拡散過程

節では$x_0$から完全なノイズ$x_T$を得ました．今度は逆に，完全なノイズ$x_T$から$x_0$を得ることを考えます．

まず，上述した拡散分布の式は以下の式で書けることがわかります．

\begin{equation}
q(x_t \mid x_0) = \mathcal{N} \left(x_t;\, x_0\prod^{t}_{t'=1} \sqrt{\alpha_{t'}},\,\, I\sqrt{1-\prod^{t}_{t'=1} \alpha_{t'}}\right)
\end{equation}

この式は複雑に見えるかもしれませんが，直感的には「$x_0$が与えられたときの$x_t$の分布は，平均が$ \prod_{t'=1}^{t} \sqrt{\alpha_{t'}}x_0 $ で，分散が$ I\sqrt{1-\prod_{t'=1}^{t} \alpha_{t'}}$の正規分布に従うよ」ということを言っています．ここで，$I$は$d$次元の単位行列であり，正規分布が独立であることを示します．

これが拡散課程ですが，逆拡散過程は逆に$x_t$を与えたときに$x_{t-1}$を求めることをいいます．すなわち，$x_t$を与えたときに，少しだけノイズを取り除いた分布$q(x_{t-1} \mid x_t)$が知りたい情報です．

しかしながら，これを求めることは容易ではありません（ニューラルネットワークで近似することになります）．そのため，ノイズを一切加えていない情報$x_0$を条件に加えた分布$q(x_{t-1} \mid x_t，x_0)$を考えます．これにより，逆拡散過程の分布を考えることができます．これは，損失関数を導出する際に用いることができます．

複雑な式変形になりますので，スルーしてもOKですが，可能な限り頑張って追ってみてください．なお，証明では条件付き確率の定義$P(A\mid B)P(B)=P(A, B)$を使用し，拡散過程のみが式に現れるよう，$q(x_{p}\mid x_{q})$（ただし，$p>q$）の形だけであらわされるように変形していきます．

\begin{align*}
q(x_{t-1} \mid x_t，x_0) &= \dfrac{q(x_{t-1} \mid x_t，x_0)q(x_{t}, x_0)}{q(x_{t}, x_0)}&(上と下にq(x_t, x_0)をかけた)\\
&= \dfrac{q(x_{t-1}, x_t, x_0)}{q(x_t, x_0)}&(条件付き確率の定義)\\
&= \dfrac{q(x_{t-1}, x_t, x_0) q(x_0)}{q(x_t, x_0) q(x_0)}   &(上と下にq(x_0)をかけた)\\
&= \dfrac{q(x_{t-1}, x_t, x_0) }{q(x_t\mid x_0) q(x_0)}   &(条件付き確率の定義)\\
&= \dfrac{q(x_{t-1}, x_t, x_0) q(x_0, x_{t-1}) }{q(x_t\mid x_0) q(x_0, x_{t-1})q(x_0)}   &(上と下にq(x_0, x_{t-1})をかけた)\\
&= \dfrac{q(x_{t}\mid x_{t-1}, x_0) q(x_{t-1}\mid x_0) }{q(x_t\mid x_0)}   &(条件付き確率の定義)\\
&= \dfrac{q(x_{t}\mid x_{t-1}) q(x_{t-1}\mid x_0) }{q(x_t\mid x_0)}   &(拡散過程のマルコフ性)
\end{align*}

これを1行にまとめると以下のようになります．

\begin{equation}
q(x_{t-1} \mid x_t，x_0) =  \dfrac{q(x_{t}\mid x_{t-1}) q(x_{t-1}\mid x_0) }{q(x_t\mid x_0)}   
\end{equation}

ここで，

\begin{align*}
q(x_{t}\mid x_{t-1}) &= \mathcal{N}(x_{t}; \sqrt{1-\beta_t} x_{t-1} \sqrt{\alpha_t}, I \sqrt{\beta_t})\\
q(x_{t-1}\mid x_0) &=\mathcal{N} \left(x_t;\, x_0\prod^{t-1}_{t'=1} \sqrt{\alpha_{t'}},\,\, I\sqrt{1-\prod^{t-1}_{t'=1} \alpha_{t'}}\right)\\
q(x_{t}\mid x_0) &=\mathcal{N} \left(x_t;\, x_0\prod^{t}_{t'=1} \sqrt{\alpha_{t'}},\,\, I\sqrt{1-\prod^{t}_{t'=1} \alpha_{t'}}\right)
\end{align*}

であることを思い出すと，

\begin{equation}
\label{gyakubunpu}
q(x_{t-1} \mid x_t，x_0) = \mathcal{N}\left(x_{t-1}; \dfrac{\sqrt{\alpha_t}(1-\bar{\alpha_{t-1}})}{1-\bar{\alpha_t}}x_t+\dfrac{\sqrt{\bar{\alpha}}\beta_t }{1-\bar{\alpha}}x_0, I \dfrac{1-\bar{\alpha_{t-1}}}{1-\bar{\alpha_t}}\beta_t\right)
\end{equation}

と書けますが，この導出は~~単純計算でめんどくさいだけなので~~省略します．

何を学習するのか

さて，そろそろ本題の逆拡散過程$q(x_{t-1} \mid x_t)$について考えていきたいところです．これがどのような分布かがわかれば一挙解決なのですが，そうは問屋が卸してくれません．そこで，これをニューラルネットワークを用いて近似します．詳細は省略しますが，1回に加えるノイズの量が少ない場合，逆拡散過程と拡散過程が同じ正規分布に従うことが知られています．

ニューラルネットワーク$\theta$を用いて，その正規分布の平均を学習します．具体的には，以下の式で書けます．なお，$\mu_\theta(x_t, t)$は時刻$t$と時刻$t$におけるノイズが加わった画像$x_t$を入力として正規分布の平均を出力とします．

\begin{equation}
q(x_{t-1} \mid x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \beta_t I)
\end{equation}

なお，学習するのは平均だけで，分散は学習しません．これは，分散は学習しなくても精度が良いことが実証的に示されているからです．この式から損失関数を定め，その損失関数を最大化するように学習を行います．

損失関数・変分下界

さて，ニューラルネットワーク$\theta$で$\mathbb{E}{q(x_0)} [p_{\theta}(x_0)]$（$x_0$が分布$q(x_0)$に従って生成されるときの$p_{\theta}(x_0)$の期待値）を最大化したいのですが，これを直接計算するのは極めて困難です．そのため，これを下から評価した下界を与え，その下界を最大化することで$p_{\theta}(x)$を最大化することを考えます．これを変分下界といいます．

計算の簡略化のため，対数尤度$-\mathbb{E}{q(x_0)} [\log p_{\theta}(x_0)]$を最小化することを考えます．以下，下界を計算できる関数で表示します．複雑な式変形が続きますが，頑張って追ってみてください．ただし，

x_{a:b}=
\begin{cases}
(x_a, x_{a+1},\cdots, x_{b})\quad\quad (a\leq bのとき)\\
(x_a, x_{a-1},\cdots, x_{b})\quad\quad (a\geq bのとき)\\
\end{cases}

とします．要するに，略記用の記号だと思ってください．不等号のところではイェンゼンの不等式$\displaystyle \int \log f(x) dx \geq \log \left( \int f(x) dx\right) $を用いています．

\begin{align*}
-\mathbb{E}_{q(x_0)} [\log p_{\theta}(x_0)] &= -\int dx_{0} q(x_0)\log p_{\theta}(x_0) \\
&= -\int dx_{0} q(x_0) \log\left(\int  p_{\theta}(x_0)p_{\theta}(x_{1:T})dx_{1:T}\right) &(3)\\ 
&= -\int dx_{0} q(x_0)  \log \left(\int p_{\theta}(x_{0:T})dx_{1:T}\right)\\
&= -\int dx_{0} q(x_0)  \log\left(\int q(x_{1:T}|x_0)\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{1:T}\right)\\
&\leq -\int dx_{0} q(x_0)  \left(\int q(x_{1:T}|x_0)\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{1:T}\right)\\
&= -\int dx_{0}   \left(\int q(x_0)q(x_{1:T}|x_0)\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{1:T}\right)\\
&= -\int dx_{0}   \left(\int q(x_0, x_{1:T})\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{1:T}\right)\\
&= -\int dx_{0}   \left(\int q(x_{0:T})\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{1:T}\right)\\
&= -\int   \left(\int q(x_{0:T})\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{0} dx_{1:T}\right) \\
&= \int q(x_{0:T})\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{0} dx_{1:T} \\
&= \int q(x_{0:T})\log\dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} dx_{0:T} &(脚注参照) \\
&=-\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)} \right]\\
&=\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})} \right] =  \mathcal{L}_{VLB}
\end{align*}

(注)
(1) $\log$の引数の中にインテグラルは1個しかありませんが，本質的には$x_1$から$x_T$まで$T$個の変数で積分しているため，インテグラルが$T$個あるものだと思うと良いです．本書の作者は，そこの理解が浅く，この数式に苦しみました．

(2) 上であった$T$個のインテグラルと1個のインテグラルをまとめ，$T+1$個のインテグラルを1個に統合しました．本質的にはインテグラルが$T+1$個あると思ったほうがいいです．ここで，VLBはVariational Lower Bound（変分下界）の略です．

さらに変形

上で述べたVLBを最小化する事で，$-\mathbb{E}{q(x_0)} [\log p_{\theta}(x_0)]$を小さくすることができます．以下，これをさらに変形します．

\begin{align*}
\mathcal{L}_{VLB}&=\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})} \right]\\
&= \mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})} \right]
\end{align*}

拡散仮定・逆拡散過程にはマルコフ性を仮定しましたので，$ q(x_{0}|x_T)= \prod^T_{t=1} q(x_{t-1}|x_t)$，$\displaystyle q(x_{T}|x_0)= \prod^T_{t=1} q(x_{t}|x_{t-1})$が成り立ちます．これを用いて式変形を続けていきましょう．

\begin{align*}
\mathcal{L}_{VLB}&=\mathbb{E}{q(x_{0:T})} \left[\log \dfrac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})} \right]\\
&=\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_1, x_2, \cdots, x_T|x_0)}{p_{\theta}(x_{0:T})} \right]\\
&= \mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_1|x_0)q(x_2, \cdots, x_t|x_0, x_1)}{p_{\theta}(x_{0:T})} \right]\\
&= \mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_1|x_0)q(x_2, \cdots, x_t|x_1)}{p_{\theta}(x_{0:T})} \right]\\
&= \cdots = \mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{\prod^T_{t=1} q(x_t|x_{t-1})}{p_{\theta}(x_{0}, x_1, \cdots, x_T)} \right]\\
&=\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{\prod^T_{t=1} q(x_t|x_{t-1})}{p_{\theta}(x_T) p_{\theta} (x_0, x_1,\cdots, x_{T-1}|x_T) } \right]\\
&=\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{\prod^T_{t=1} q(x_t|x_{t-1})}{p_{\theta}(x_T)p_{\theta}(x_{T-1} | x_T) p_{\theta} (x_0, x_1,\cdots, x_{T-2}|x_T, x_{T-1}) } \right]\\
&=\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{\prod^T_{t=1} q(x_t|x_{t-1})}{p_{\theta}(x_T)p_{\theta}(x_{T-1} | x_T) p_{\theta} (x_0, x_1,\cdots, x_{T-2}|x_{T-1}) } \right]\\
&=\cdots =\mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{\prod^T_{t=1} q(x_t|x_{t-1})}{p_{\theta}(x_T)\prod^T_{t=1} p_{\theta}(x_{t-1}|x_{t})} \right]\\
\end{align*}

となります．ここで，$\log$の中に$\prod$が入っているものは$\log$外の$\sum$に置き換えられるので，

\begin{align*}
\mathcal{L}_{VLB} &= \mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{\prod^T_{t=1} q(x_t|x_{t-1})}{p_{\theta}(x_T)\prod^T_{t=1}  p_{\theta}(x_{t-1}|x_{t})} \right]\\
&= \mathbb{E}_{q(x_{0:T})} \left[\sum^{T}_{t=1} \log q(x_t|x_{t-1}) - \log p_{\theta} (x_T) - \sum^{T}_{t=1} \log  p_{\theta}(x_{t-1}|x_{t})   \right]\\
&= \mathbb{E}_{q(x_{0:T})} \left[ - \log p_{\theta} (x_T)+\sum^{T}_{t=1} \log q(x_t|x_{t-1}) - \sum^{T}_{t=1} \log  p_{\theta}(x_{t-1}|x_{t})   \right]\\
&= \mathbb{E}_{q(x_{0:T})}  \left[- \log p_{\theta} (x_T)+\log \dfrac{q(x_1|x_0)}{ p_{\theta}(x_0|x_1)}+\sum^{T}_{t=2} \log \dfrac{q(x_t|x_{t-1})}{p_{\theta}(x_{t-1}|x_t)}   \right]\\
&= \int_{q(x_{0:T})} d q(x_{0:T}) \left(-\log p_{\theta}(x_0|x_1)+\log \dfrac{q(x_1|x_0)}{ p_{\theta} (x_T)}+\sum^{T}_{t=2} \log \dfrac{q(x_t|x_{t-1})}{p_{\theta}(x_{t-1}|x_t)}   \right)\\
&= - \int dx_{0:T} q(x_{0:T}) \log p_{\theta} (x_0 | x_1) +\underbrace{ \int}_{T+1個}d x_0 \cdots dx_T q(x_0,\cdots, x_T)\log \dfrac{q(x_1|x_0)}{ p_{\theta}(x_T)}\\
&+ \underbrace{ \int}_{T+1個}d x_0 \cdots dx_T q(x_0,\cdots, x_T)\left(\sum^{T}_{t=2} \log \dfrac{q(x_t|x_{t-1})}{p_{\theta}(x_{t-1}|x_t)}\right)\\
&= -  \mathbb{E}_{q(x_{0:T})} [\log_{p_{\theta}} (x_0|x_1)] + \mathbb{E}_{q(x_{0:T})} \left[\log \dfrac{q(x_1|x_0)}{p_{\theta}(x_T)}\right]+ \mathbb{E}_{q(x_{0:T})} \left[\sum^{T}_{t=2} \log \dfrac{q(x_t|x_{t-1})}{p_{\theta}(x_{t-1}|x_t)}\right]\\
&= \mathbb{E}_{x_{0:T}} \left[\underbrace{D_{\mathrm{KL}} (q(x_T|x_0)||p_{\theta}(x_T))}_{L_Tとおく}+\sum^{T}_{t=2}\underbrace{ D_{\mathrm{KL}} (q(x_{t-1}|x_t, x_0) || p_{\theta}(x_{t-1}|x_t))}_{L_tとおく}\underbrace{-\log p_{\theta} (x_0 |x_1)}_{L_0とおく}\right]\\
&= L_T + L_{T-1}+\cdots+L_0
\end{align*}

と書けます．ここで，$D_{\mathrm{KL}}$はKLダイバージェンスを表します．これは2つの確率密度関数$p(x), q(x)$に対して定義される値のことで，その式は以下で書けます．

$$ D_{\mathrm{KL}} (p||q) = \int p(x) \log \dfrac{p(x)}{q(x)}dx $$

これを踏まえて，$L$の性質を見ていきましょう．

$x_T$は完全なガウスノイズですから，これはパラメータ$\theta$に依存しません．また，$L_0$は頑張れば計算できます．

残った$L_{t}$について考えます．学習するのは平均で，分散は決め打ちですから，$p_{\theta} (x_{t-1}|x_t)=\mathcal{N}(\mu_{\theta}(x_t, t), \beta_t I)$であり，$L_t$は2つの正規分布のKLダイバージェンスですから，解析的に計算することが可能です．

すると，$L_t$は，$x_0$から$x_t$までが与えられたときの「時刻$t$のノイズから得られた時刻$t-1$のノイズ（$\theta$に依存する変数）」と「時刻0と時刻$t$の情報から得られた時刻$t-1$のノイズ」の距離の期待値を，$L_0$は$x_1$から$x_0$が得られる期待値のマイナス1倍を表しています．これらを共に小さくすることで，2つのノイズの距離を近くし，期待値を大きくすることができ，結論として理想的な分布を得ることができます．

次回予告

次回は，今回の理論に基づいたpytorchでの実装にとりかかります．

参考文献

"Denoising Diffusion Probabilistic Models", Jonathan Ho, Ajay Jain, Pieter Abbeel, 2020.