はじめに
拡散モデルを解説するッ!!!①の続きです.
今回は拡散確率モデル(Diccusion Probablistic Models, DPMs)(以後単に拡散モデルと呼ぶ)の解説をやっていきたいと思います.
前回はSBMの解説をやったので,今回はDDPMの概要を説明し,,,と思ったのですが,その前段階の拡散モデルの解説を論文[3]を元にしつつやっていきます.
(なんか,バグなのかわかんないんですけどちょいちょい確率分布p, qが太字になっているかもです.うまいこと読み替えてください。。。)
拡散モデルの概要(+この辺りを調べての所感)
拡散モデルは非平衡統計物理の文脈から着想を得て作られたモデルです(物理系と情報系が交わるっていいですよね,Hinton先生ノーベル賞おめでとうございます).
時系列的に言えば,このモデルが2015年くらいなので前に説明したSBMよりも前に拡散モデルは発表されています.記事は前後していますが.それなのでこの辺りを整理すると,拡散モデル→SBM→DDPMという流れです.しかし,この辺りは著者に直接確認しないとなんとも言えないのですが,SBMはこの時点では拡散モデルは意識していなかったと思います.(論文の本文中にdiffusionという単語は3回のみ)一方で,SongらがDDPMが発表された後に出した[6]では38回もdiffusionという単語が使われているので,DDPMの後に意識しだした感じでしょうかね.正直拡散モデルの発想を知ってから読むとSBMも拡散モデルも同じように見えてしまいますが。。。(どちらも元データをノイズで攪乱するモデルという意味で)となると,これらを統一したくなるというモチベーションがDDPMにはあったのかもしれないですね(本当のモチベーションは後ほど).
とまぁ,拡散モデルは発表されてから色々やられているのですが,それはさておき拡散モデルの発想としてはシンプルです.データにノイズを乗っけて,各拡散ステップで加えられたノイズを推定し,ノイズの乗ったデータから推定されたノイズをさっ引いてあげる(デノイジングする)ことで元データに近いデータを生成する,と.この流れを数式込みで追いかけてみようと思います.
拡散モデルの解説
拡散過程
データ分布を$q(\bf{x}^{(0)})$とします.この時,$\bf{x}$をマルコフ過程(ひとつ前のものにのみ依存する確率過程)とすると$T$ステップの拡散プロセスは
q(\bf{x}^{(0...T)})=q(\bf{x}^{(0)})\prod_{t=1}^T q(\bf{x}^{(t)}|\bf{x}^{(t-1)})
とかけます.この式を分解してみると直感的な理解は難しくないと思います.
\begin{align}
q(\bf{x}^{(0...T)})&=q(\bf{x}^{(0)})\prod_{t=1}^T q(\bf{x}^{(t)}|\bf{x}^{(t-1)}) \\
&=q(\bf{x}^{(T)}|\bf{x}^{(T-1)})q(\bf{x}^{(T-1)}|\bf{x}^{(T-2)})\cdots q(\bf{x}^{(2)}|\bf{x}^{(1)})q(\bf{x}^{(1)}|\bf{x}^{(0)})q(\bf{x}^{(0)})
\end{align}
なんか,$0$から$T$へ遷移してそうな式になっていそうですね.笑
一応式の意味を説明しておくと,$q(\bf{x}^{(0)})$はデータ分布ですね.つまり,この分布からサンプリングすれば元データが得られます.例えば,コイン投げを表すデータ分布であればサンプリングすると表と裏のどっちかが確率$\frac{1}{2}$で得られます.ネコチャァンの画像を表すデータ分布であれば,サンプリングによってなんらかのネコチャァンの画像が得られますニャア.では,
q(\bf{x}^{(t)}|\bf{x}^{(t-1)})
は何かっていうのを説明する、、、前に.$\bf{x}^{(t)}$を定義しておきましょうか.
$\bf{x}^{(t)}$は,$\beta_t$を$0<\beta_1<\beta_2\cdots<\beta_{T-1}<\beta_T<1$を満たす数として,
\bf{x}^{(t)} = \sqrt{1-\beta_t}\bf{x}^{(t-1)}+\beta_t \sigma \,\,\, (\sigma \sim\mathcal{N}(\bf{0}, \bf{I}))
です(もしかしたら$\beta_T=1$かも?,まぁ一旦細かいことはさておき←).つまり,$\bf{x}^{(t)}$は前のデータに標準正規分布からのノイズを加えたものです.さらに$\beta_t$が0〜1まで大きくなるので,徐々に元データの成分が小さくなっていきます.わかりますか(式から読み取れますか)?つまり,最終的には,$\beta_t$が1に近くなるので,ノイズが支配的になります.極端なケースだと,
\begin{align}
\bf{x}^{(T)} &= \sqrt{1-1}\bf{x}^{(T-1)}+1\times \sigma \,\,\, (\sigma \sim\mathcal{N}(\bf{0}, \bf{I}))\\
&=0+\sigma \,\,\, (\sigma \sim\mathcal{N}(\bf{0}, \bf{I}))
\end{align}
ですね.なので$\bf{x}^{(t)}$は,徐々にノイズを加えて完全なノイズにするまでの拡散過程を表しています.さて話を戻して,
q(\bf{x}^{(t)}|\bf{x}^{(t-1)})
が何か,ですがここまで来たらわかると思います.
これは,$\bf{x}^{(t-1)}$にノイズを乗っけた後のデータの分布を表しています.すなわち,$\bf{x}^{(t-1)}$にノイズを加えた$\bf{x}^{(t)}$がサンプリングによって得られます.で,最後に左辺の$q(\bf{x}^{(0...T)})$は$\bf{x}^{(0...T)}=\bf{x}^{(0)}, \bf{x}^{(1)},...,\bf{x}^{(T)}$の同時確率でこの分布からのサンプリングにより,一本のノイズを加えていく経路がサンプリングされます.ひとます,式の意味はいいですね.これが拡散過程を表している(表してそうだ)というのを理解していただけましたかね?(質問があればぜひ。。。)
さて,ここで欲しいのは$q(\bf{x}^{(0)})$です.なぜなら,この分布がわかれば現在持っているデータ以外に無限にデータ生成できますからね.ネコチャァンの画像を生成できる分布を知っていればマンチカンの画像だろうが,スコティッシュフォールドの画像だろうが生成できます.ですが,僕らは知りません.データは手に持っているでしょって?マンチカンとスコティッシュフォールドの画像しかなかったらそれしか機械は生成できません.だからもっと多様性のある画像生成をするために現在持っているデータを使って$q(\bf{x}^{(0)})$をうまく近似していくことにします.どうやって近似するか?逆拡散過程を用いてデータ分布$q(\bf{x}^{(0)})$をモデル分布$p(\bf{x}^{(0)})$に近似します.
逆拡散過程
データを生成するモデル分布を$p(\bf{x}^{(0)})$とします.
ここで,
\begin{align}
p(\bf{x}^T) &\sim \mathcal{N}(\bf{0}, \bf{I}) \\
p(\bf{x}^{(0...T)}) &= p(\bf{x}^T) \prod_{t=1}^T p(\bf{x}^{(t-1)}|\bf{x}^{(t)})
\end{align}
とします.ここで式をみてみると,$p(\bf{x}^{(0...T)})$は$\bf{x}^{(0...T)}=\bf{x}^{(0)}, \bf{x}^{(1)},...,\bf{x}^{(T)}$の同時分布であり,サンプリングによって拡散過程と逆向きの経路(ノイズからデータへのデータ生成経路)が得られる確率分布を表します.これを踏まえて逆拡散過程を用いて$p(\bf{x}^{(0)})$は
p(\bf{x}^{(0)}) = \int d \bf{x}^{(1...T)} p(\bf{x}^{(0...T)})
と書けます.天下り的な説明で恐縮ですが,一応この式でやっていることを説明します.これは周辺化をしています.逆拡散過程なので,$1...T$までの$\bf{x}$で積分をとってあげると$p(\bf{x}^{(0)})$が得られます.当たり前と言えば当たり前ですが.で,この積分が計算できれば求めたい$p(\bf{x}^{(0)})$が得られるんですが,この積分は``intractable(計算困難)''です.なので,ここから拡散モデルではジャルジンスキー等式という非平衡統計物理の文脈からきた等式から着想を得て計算可能な形に持っていくそうです(よく知りませんごめんなさい).式変形自体は難しいことはない(?)のでご安心を.さてでは引き続き式変形していくと,
\begin{align}
p(\bf{x}^{(0)}) &= \int d\bf{x}^{(1...T)} p(\bf{x}^{(0...T)}) \\
&= \int d\bf{x}^{(1...T)} p(\bf{x}^{(0...T)}) \frac{q(\bf{x}^{(1...T)}|\bf{x}^{(0)})}{q(\bf{x}^{(1...T)}|\bf{x}^{(0)})} \\
&= \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) \frac{p(\bf{x}^{(0...T)})}{q(\bf{x}^{(1...T)}|\bf{x}^{(0)})} \\
&= \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) p(\bf{x}^T) \frac{\prod_{t=1}^T p(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{\prod_{t=1}^T q(\bf{x}^{(t)}|\bf{x}^{(t-1)})} \\
&= \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) p(\bf{x}^T) \prod_{t=1}^T \frac{p(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}
\end{align}
のようになります.最後の式変形は条件付き確率の定義を用いて以下のように変形してます
\begin{align}
q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) &= \frac{q(\bf{x}^{(1...T)},\bf{x}^{(0)})}{q(\bf{x}^{(0)})} \\
&= \frac{q(\bf{x}^{(0...T)})}{q(\bf{x}^{(0)})} \\
&= \frac{q(\bf{x}^{(0)})\prod_{t=1}^T p(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(0)})} \\
&=\prod_{t=1}^T p(\bf{x}^{(t-1)}|\bf{x}^{(t)})
\end{align}
.さて,ここでDDPMとの話の接続を考えて,モデル分布$p(\bf{x}^{(0)})$をモデルの尤度$p_\theta(\bf{x}^{(0)})$としておきます.ここから,対数尤度を考えていきます.確率モデルを考えるのに対数尤度を考える理由としては,一般的に尤度は積の形で与えられるので対数をとって和の形にしてあげた方が扱いやすいからですね.まぁこの話は確率モデル(生成モデル)の話で死ぬほど出てくるので,あえて補足する必要もない気がしますが.で,ここからパラメータ$\theta$を推定するという話に進んでいきたいんですが,話が分岐してしまいます.パラメータ推定のために,岡野原さん書かれた本[1]では,負の対数尤度の最小化を考えているのに対し,元論文[3]では,コスト関数$L$の次のように
\begin{align}
L &= \mathbb{E}_{q(\bf{x}^{(0)})}\left[ \log p_{\theta}(\bf{x}^{(0)})\right] \\
&= \int d \bf{x}^{(0)} q ( \bf{x}^{(0)} ) \log p_{\theta} ( \bf{x}^{(0)})
\end{align}
定義し,最大化することを考えています.どちらを使っても良いと思うんですが,一応論文を元にやります(つまり後者).
さて,それを踏まえて式変形を続けます.今,コスト関数を先ほどの$p(\bf{x}^{(0)})$の式を用いて変形すると,
\begin{align}
L &= \mathbb{E}_{q(\bf{x}^{(0)})}\left[ \log p_{\theta}(\bf{x}^{(0)})\right] \\
&= \int d \bf{x}^{(0)} q ( \bf{x}^{(0)} ) \log p_{\theta} ( \bf{x}^{(0)}) \\
&= \int d \bf{x}^{(0)} q ( \bf{x}^{(0)} ) \log \left[ \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right]
\end{align}
となります.ここで,イェンセン不等式を用います.イェンセン不等式は上に凸な関数$f$に対して($f$が上凸か下凸は重要です,不等号の向きが変わるので),
f\left( \int g(x)p(x)dx\right) \geq \int f(g(x))p(x)dx
です.今回はこの$f$を対数関数とします.そして,対数関数は上に凸です.そのためJensen不等式を$\log p_{\theta} ( \bf{x}^{(0)})$にそのまま用いて,
\begin{align}
\log p_{\theta} ( \bf{x}^{(0)}) &= \log \left[ \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
&\geq \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) \log \left[ p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right]
\end{align}
ですね.よって,
\begin{align}
L &= \int d \bf{x}^{(0)} q ( \bf{x}^{(0)} ) \log \left[ \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
&\geq \int d \bf{x}^{(0)} q ( \bf{x}^{(0)} ) \int d\bf{x}^{(1...T)} q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) \log \left[ p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
&= \int d \bf{x}^{(0)} d \bf{x}^{(1...T)} q ( \bf{x}^{(0)} )q(\bf{x}^{(1...T)}|\bf{x}^{(0)}) \log \left[ p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
&= \int d \bf{x}^{(0...T)} q ( \bf{x}^{(0)} )\frac{q(\bf{x}^{(0...T)})}{q(\bf{x}^{(0)})} \log \left[ p_{\theta}(\bf{x}^T) \prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
&= \int d \bf{x}^{(0...T)} q(\bf{x}^{(0...T)}) \log \left[ p_{\theta}(\bf{x}^T)\prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] := K
\end{align}
となります.記号が多くてイヤになりますが,難しいことはしていないと思います.さらに,式変形を続けます.ここからの話は[3]のAppendix B.の部分に相当します.さて,それでは$K$について考えます.今,$K$は
\begin{align}
K &= \int d \bf{x}^{(0...T)} q(\bf{x}^{(0...T)}) \log \left[ p_{\theta}(\bf{x}^T)\prod_{t=1}^T \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right]\\
&= \int d \bf{x}^{(0...T)} q(\bf{x}^{(0...T)}) \left[ \log p_{\theta}(\bf{x}^T) + \sum_{t=1}^T \log \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
&= \int d \bf{x}^{(0...T)} q(\bf{x}^{(0...T)}) \log p_{\theta}(\bf{x}^T) \\
&+ \int d \bf{x}^{(0...T)} q(\bf{x}^{(0...T)}) \left[ \sum_{t=1}^T \log \frac{p_{\theta}(\bf{x}^{(t-1)}|\bf{x}^{(t)})}{q(\bf{x}^{(t)}|\bf{x}^{(t-1)})}\right] \\
\end{align}
です.このコスト関数を最大化するのが,[3] の論文でやられたことです.どういうことかというと,$L$の形のままだと扱いづらいというか計算できないので,$L\geq K$という関係を用いて,$K$の最大化によって間接的に$L$を最大化していくという話に持っていったということです.具体的な話($K$をどうやってコードに落とし込むかなど)は次回のDDPMの時に補足しようかと思います.
まとめ
ひとまず,拡散モデルのコスト関数まで説明しました.このモデルでは,$K$直接最大化していきます.次回説明しますが,今回説明したモデルはカルバックライブラー情報量の形をコスト関数に持つため,複雑な形をしています.ですが,次回説明するDDPMでは,Denoisign Score Matchingと関係があることがわかり,それによってコスト関数がMSEの形(簡単な形)で書かれることが提案されました(前回を少し参照してみてください,あの形です).これによって,シンプルなコスト関数を用いて高品質な画像生成が行われるようになったと.なんかそんな感じの話の流れらしいです.
あとがき
なんかミスあったら教えてください、、、
アドバイスもぜひ.
参考文献
[1]岡野原 大輔,"拡散モデル データ生成技術の数理", 岩波書店
[2] Yang Song and Stefano Ermon, "Generative Modeling by Estimating Gradients of the Data Distribution", NeurIPS 2019
[3] Jascha Sohl-Dickstein et al.,"Deep Unsupervised Learning using Nonequilibrium Thermodynamics", ICML 2015
[4]Jonathan Ho et al.,"Denoising Diffusion Probabilistic Models", NeurIPS 2020
[5]Diederik Kingma et al.,"Variational Diffusion Models", NeurIPS 2021
[6]Yang Song et al., "SCORE-BASED GENERATIVE MODELING THROUGH
STOCHASTIC DIFFERENTIAL EQUATIONS",ICLR 2021