最近話題の"Diffusion Model（拡散モデル）"について、簡潔にまとめてみた

　OpenAIが先日発表したGLIDEにDiffusion Moldelが使用されているとのことで、最近話題のDiffusion Model（拡散モデル）について、まとめました。

サマリー

・Diffusion Model（拡散モデル）は、元データにノイズが徐々に付加されていき、最終的にガウシアンノイズとなるという前提を置き、その逆のプロセスをモデル化することでデータを生成する。

・GANやVAEよりも高品質の画像を生成することに成功しており、様々な分野への応用が期待される。

生成モデルとは

　生成モデルを使用することで、データ（ex. 画像）の生成プロセスをモデル化し、ノイズ等から新たなインスタンスを生成することが可能です。下記の図は、生成モデルによって生成された画像です。もはや本物としか思えないほど、精巧な画像を生成できています。

出典：https://arxiv.org/abs/1710.10196

　代表的な手法として、GAN・VAE・Flowがあります。Diffusion Modelはこれらの生成モデルと同等、或いはそれ以上に精巧にデータ生成できることが、昨今の研究で判明しています。

Diffusion Model（拡散モデル）とは

　Diffusion Modelはノイズからスタートし、徐々にノイズを除去していくことで、データを生成するモデルです。下記の図では、xT$x_T$からスタートし、T回のノイズ除去ステップを踏むことで、顔画像x0$x_0$を生成しています。もちろん、どのようにノイズを除去すればよいのか、言い換えるならば、どのような過程でノイズが付与されたのか、ということが分からなければ、データを生成することはできません。ノイズを付与していく過程をForward diffusion process、除去していく過程をReverse diffusion processと呼びます。以下、この2つのプロセスについて、考えていきましょう。

出典：https://arxiv.org/pdf/2006.11239.pdf

Forward diffusion process

　このプロセスは、状態xt−1$x_{t-1}$にガウシアンノイズを付与することで、次の状態xt$x_t$に遷移させるプロセスです。つまり、このプロセスはノイズの強さをβt${\beta }_t$とすると、以下のように表現できます。なお、初期状態x0$x_0$はq$q$からサンプルされているとします。

q(xt|xt−1)=N(xt;1−βt−−−−−√xt−1,βtI)q(x1:T|x0)=∏Tt=1q(xt|xt−1)$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$

これは、任意のステップt$t$の状態xt$x_t$が前の状態xt−1$x_{t-1}$の関数で表現できるということです。つまり、これをt$t$回繰り返すことで、xt$x_t$を初期状態x0$x_0$で表現できます。なお、αt=1−βt,α¯¯¯t=∏Ti=1αi${\alpha }_t=1-{\beta }_t,{\overline{\alpha }}_t=\prod _{i=1}^T{\alpha }_i$と定義しています。

xtq(xt|x0)=αt−−√xt−1+1−αt−−−−−√zt−1=αtαt−1−−−−−√xt−2+1−αtαt−1−−−−−−−−−√z¯t−2=…=α¯t−−√x0+1−α¯t−−−−−√z=N(xt;α¯t−−√x0,(1−α¯t)I) ;where zt−1,⋯∼N(0,I) ;where z¯t−2 merges two Gaussians $$\begin{array}{rlr}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{1-{\alpha }_t}{\mathbf{z}}_{t-1}& \text{ ;where }{\mathbf{z}}_{t-1},\cdots \sim \mathcal{N}(\mathbf{0},\mathbf{I})\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{\mathbf{z}}}_{t-2}& \text{ ;where }{\overline{\mathbf{z}}}_{t-2}\text{ merges two Gaussians }\\ & =\dots \\ & =\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}\mathbf{z}\\ q({\mathbf{x}}_t|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0,(1-{\overline{\alpha }}_t)\mathbf{I})\end{array}$$

Reverse diffusion process

　ノイズを除去していくプロセスをモデル化できれば、状態xT$x_T$からデータを生成できます。つまり、Forward processの逆プロセスであるq(xt−1|xt)$q(x_{t-1}|x_t)$を表現できれば良いということです。βt${\beta }_t$が十分小さいとすると、q(xt−1|xt)$q(x_{t-1}|x_t)$はガウス分布に従います。これをパラメータθ$\theta$のニューラルネットワークpθ$p_{\theta }$で近似します。

pθ(x0:T)=p(xT)∏Tt=1pθ(xt−1|xt)pθ(xt−1|xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$

q(xt−1|xt)$q(x_{t-1}|x_t)$は、このままでは非常に扱いにくいので、x0$x_0$で条件付けします。このq(xt−1|xt,x0)$q(x_{t-1}|x_t,x_0)$もガウス分布に従うと仮定し、平均及び分散を何かしらの形で表現することを目指します。詳細は割愛しますが、ベイズの定理を用いて、平均及び分散は以下となります。

μ~tβ~t=1αt−−√(xt−βt1−α¯t−−−−−√zt)=1/(αtβt+11−α¯t−1)=1−α¯t−11−α¯t⋅βt$$\begin{array}{rl}{\tilde{\mathit{\mu }}}_t& =\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\mathbf{z}}_t)\\ {\tilde{\beta }}_t& =1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\cdot {\beta }_t\end{array}$$

学習

　上記を踏まえて、Diffusion Modelを学習させるために必要な損失関数を考えます。詳細は割愛しますが、クロスエントロピーの最小化、或いはVAEと同様に負の対数尤度の最小化を目指すことで、損失関数は以下となります。なお、DKL$D_{KL}$はカルバックライブラーダイバージェンスを表しています。

LVLBwhere LTLtL0=LT+LT−1+⋯+L0=DKL(q(xT|x0)∥pθ(xT))=DKL(q(xt|xt+1,x0)∥pθ(xt|xt+1)) for 1≤t≤T−1=−logpθ(x0|x1)$$\begin{array}{rl}{L}_{\text{VLB}}& =L_T+L_{T-1}+\cdots +L_0\\ \text{where }{L}_T& =D_{\text{KL}}(q({\mathbf{x}}_T|{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_T))\\ L_t& =D_{\text{KL}}(q({\mathbf{x}}_t|{\mathbf{x}}_{t+1},{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_t|{\mathbf{x}}_{t+1}))\text{ for }1\le t\le T-1\\ L_0& =-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)\end{array}$$

DKL$D_{KL}$の最小化は、2つのガウス分布の平均及び分散の差分の最小化と等しくなります。分散を固定のパラメータとすると、ガウス分布の平均の差を小さくするように学習すれば良いことになります。

※分散を固定のパラメータとせず、学習可能なパラメータとしたほうが良い結果が得られるとの実験結果が出ています。詳細はこちらをご覧ください。

応用例

画像生成

　DhariwalとNicholはDiffusion Modelを使用して、BigGANと呼ばれる非常に高品質の画像と同等の品質の画像を生成することに成功しています。

出典：https://arxiv.org/pdf/2105.05233.pdf

条件付き画像生成

　DhariwalとNicholらはテキストに沿った高品質の画像を生成することに成功しています。学習データセットに含まれないような画像を生成できており、モデルが人間のように言葉を理解し、それを画像に反映させているようにも見えます。

出典：https://arxiv.org/pdf/2112.10741.pdf

画像変換

　DhariwalとNicholらは上記と同様のDiffusion Modelを使用して、指定した個所をテキストに沿って変換するタスクを実施しています。結果を見ると、非常に高精度に画像を変換できていることが分かります。

出典：https://arxiv.org/pdf/2112.10741.pdf

参考文献

・Inject Noise to Remove Noise: A Deep Dive into Score-Based Generative Modeling Techniques

・What are Diffusion Models?

・A new SotA for generative modelling — Denoising Diffusion Probabilistic Models | by Sebastian Orbell | Graphcore | Medium

・Denoising Diffusion Probabilistic Models

・Diffusion Models Beat GANs on Image Synthesis 

・GLIDE: Towards Photorealistic Image Generation and Editing with
Text-Guided Diffusion Models