Diffusion Probablistic Modelについて調べてみました。

こんにちは、技術開発の三浦です。

最近色々なところでテキストから画像を生成するAIが話題になっています。色々なAIが発表されていますが、その中で私にとって、特に印象に残っているのが"Stable Diffusion"です。この技術の中核にあるのが"Latent Diffusion Model"で、このベースになっているのが"Diffusion Probablistic Model(Diffusion Model)"です。"Diffusion Model"は2015年の論文に登場しているのですが、2020年の論文でシンプルな損失関数による学習手法が提案され、CIFAR10データセットにおける画像生成の精度指標(FID)でStyleGAN2を上回り、LSUN 256x256データセットではProgressiveGANと同じくらい高品質な画像を生成することに成功しています。

今回"Diffusion Model"について調べるため、2020年に発表された「Denoising Diffusion Models」という論文を調べてみました。Tensorflowの実装がGithubに上がっており、それと比較しながらどんな技術なのか調べて、分かったことをまとめてみたいと思います。

今回読んだ論文

Denoising Diffusion Probabilistic Models
Jonathan Ho, Ajay Jain, Pieter Abbeel
Submitted on 19 Jun 2020 (v1), last revised 16 Dec 2020
https://arxiv.org/abs/2006.11239

Github

github.com

Diffusion Modelとは

"Diffusion Model"はデータの分布に潜在変数の存在を仮定した生成モデルです。潜在変数の間にはマルコフ連鎖の関係が仮定されていて、 $x_{0}$ を観測データ、 $x_{1}...x_{T}$ を潜在変数とすると、以下のような関係が仮定されています。

$\displaystyle p_{\theta}(x_{0:T}) := p(x_{T})\prod^{T}_{t=1}p_{\theta}(x_{t-1}|x_{t})\\ p_{\theta} (x_{t-1}|x_{t}) := \mathcal{N}(x_{t-1}; \mu_{\theta}(x_{t}, t),\textstyle{\sum}_{\theta}(x_{t},t)) \\ p_{T} := \mathcal{N}(x_{T}; \mathbf{0}, \mathbf{I})$

上の式の $p_{\theta}(x_{0:T})$ は"Diffusion Model"において、"reverse process"と呼ばれています。標準正規分布に従うノイズ $x_T$ からスタートし、現在の潜在変数の値 $x_{t}$ と時点 $t$ をパラメータ $\theta$ で決まる関数(モデル)に入力することによって得られる平均 $\mu_{\theta}(x_{t}, t)$ と分散 ${\sum}_{\theta}(x_{t},t)$ で定まる正規分布に従い、次の潜在変数 $x_{t-1}$ が生成される。これを $x_{t}, x_{t-1},...,x_{1}$ と繰り返すことで観測データが生成されることを表しています。分散 ${\sum}_{\theta}(x_{t},t)$ はモデルに推計させることも出来ますが、定数にして固定してしまっても問題がないようです。論文においても ${\sum}_{\theta}(x_{t},t)$ は定数にして固定してしまい、平均 $\mu_{\theta}(x_{t}, t)$ を推計対象として取り扱っています。

潜在変数の数 $T$ はハイパーパラメータです。ある程度大きくする必要があるようで論文の中では $T=1000$ が設定されていました。

"Diffusion Model"ではさらに逆向きのマルコフ連鎖の関係が仮定されています。こちらは各潜在変数は観測データ $x_{0}$ から徐々に正規分布に従うノイズを加えて生成されることを表しています。

$\displaystyle q(x_{1:T}|x_{0}) := \prod^{T}_{t=1}q(x_{t}|x_{t-1})\\ q(x_{t}|x_{t-1}) := \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}\mathbf{I})$

$q(x_{1:T}|x_{0})$ は"diffusion process"または"forward process"と呼ばれています。 $\beta_{t}$ は $t$ ごとにあらかじめ設定しておく定数です。 $t$ が小さい時は小さく取り、 $T$ に近づくほど大きくなるようにします。論文では $\beta_{1} = 10^{-4}$ , $\beta_{T} = 0.02$ が設定されていました。

これらをまとめると、論文に掲載されている以下の図になります。

Denoising Diffusion Probabilistic Models, Jonathan Ho, Ajay Jain, Pieter Abbeelhttps://arxiv.org/abs/2006.11239, Figure2より掲載

"diffusion"は"拡散"や"発散"を意味しますが、元の画像に徐々にノイズを加え、最終的に標準正規分布に従うノイズになる様子を表現しているようです。

損失関数

"Diffusion Model"はこのようなマルコフ連鎖で定義されたモデルですが、このモデルを決定するために先ほどの"reverse process"の式に出てきた $\mu_{\theta}(x_{t}, t)$ を推計する必要があります。

まず"diffusion process"に沿って考えると、条件付き分布 $q(x_{t}|x_{0})$ は $x_{0}$ と、 $\beta_{t}$ を使って計算できる $\bar{\alpha}_{t}$ によって以下の様に表すことが出来ます。

$\displaystyle q(x_{t}|x_{0}) = \mathcal{N}(x_{t}; \sqrt{\bar{\alpha}_{t}}x_{0}, (1-\bar{\alpha}_{t})\mathbf{I})\\\ \alpha_{t} := 1 - \beta_{t}\\\ \bar{\alpha}_{t} := \prod_{s=1}^{t}\alpha_{s}\\\$

このことから、 $\epsilon ～ \mathcal{N}(x_{T}; \mathbf{0}, \mathbf{I})$ を使って $x_{t} = \sqrt{\bar{\alpha}_{t}}x_{0} + \sqrt{1-\bar{\alpha}_{t}}\epsilon$ と表すことが出来ます。一方"reverse process"に沿って考えた場合、正規分布の平均 $\mu_{\theta}(x_{t}, t)$ は以下のような式を満たす必要があります。

$\displaystyle \mu_{\theta}(x_{t}, t) = \frac{1}{\sqrt{\alpha_{t}}}(x_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(x_{t},t)) \\$

$\epsilon_{\theta}$ は $x_{t}$ から $\epsilon$ を推計する関数(モデル)です。言い換えると $\epsilon_{\theta}$ をこの条件を満たすように学習できれば、欲しかった $\mu_{\theta}(x_{t}, t)$ を得ることが出来ます。

このことを用いれば損失関数は以下のようなシンプルな形になることが述べられています。

$\displaystyle L_{simple}(\theta) := \mathbb{E}_{t, x_{0},\epsilon}[ \|\epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_{t}}x_{0} + \sqrt{1 - \bar{\alpha}_{t}}\epsilon, t)\|^2$ ]

$t$ は $1～T$ の一様分布に従う変数として捉えています。

学習アルゴリズムと画像生成アルゴリズム

先ほどの損失関数 $L_{simple}(\theta)$ を最小化する $\epsilon_{\theta}$ をニューラルネットワークで推計出来るようにすることが学習の目標になります。 $\epsilon_{\theta}$ を推計するネットワークとして"U-Net"というネットワーク構造が使われています。学習アルゴリズム、画像生成アルゴリズムは以下の様になります。

Denoising Diffusion Probabilistic Models, Jonathan Ho, Ajay Jain, Pieter Abbeelhttps://arxiv.org/abs/2006.11239より掲載

画像生成アルゴリズムの方は、標準正規分布のノイズに対して徐々にノイズを取り除いていき、画像を生成する流れになっています。

2つの画像の特徴を持った画像の生成

"Diffusion Model"では、画像の生成過程が細かく分かれています。そのため画像がどのように生成されるのかを追跡することが出来ますし、どこかのタイミングで人為的に介入することも出来ます。たとえば $t=1～T$ のどこかのタイミングで意図的にデータを挿入したら、最終的にどのような画像が生成されるのでしょうか。

論文では"diffusion process"によって画像を潜在空間に写すことが出来ると捉え、潜在空間上で重みを付けて合成したデータを"reverse process"によって元の画像空間に戻した時の画像の評価を行っています。単純に2つの画像を画像空間上で合成するとちぐはぐな画像になってしまいますが、潜在空間上で合成し、"reverse process"を通じて画像空間に写すことでノイズが除去され、自然な合成画像が生成出来ることを示しています。

Denoising Diffusion Probabilistic Models, Jonathan Ho, Ajay Jain, Pieter Abbeelhttps://arxiv.org/abs/2006.11239, Figure 8より掲載

こういう結果を見ると、「じゃあこんな操作を加えたらどうなるんだろう」とワクワクします。おそらくこの辺りの"人がコントロールしやすい"特徴が、"Diffusion Model"の面白いところなのだと思います。

まとめ

今回は"Denoising Diffusion Probabilistic Models"という論文を読んで、最近画像生成の領域で高い精度を上げ、テキストから画像を生成するAIにも活用されている"Diffusion Model"について調べてみたことをまとめてみました。とても面白い技術だな、と思いました。今度は自分でも動かしてみて、どんな画像がどのような過程で生成されるのかを実験して確かめてみたいと思います！