拡散モデル：これまでにない生成品質と多様性を達成する生成モデル

　深層生成モデルの1つである拡散モデル（Diffusion Model）の利用が急速に広がっている。2015年に米Stanford UniversityのSohl-Dickstein氏らによって最初のモデルが提案されてから^1）しばらく注目されていなかったが、2020年に米University of California BerkeleyのJonathan Ho氏ら^2）によってネットワークアーキテクチャや学習手法を工夫すれば他の深層生成モデルと同じ、もしくは上回る品質を達成できること、またデノイジングスコアマッチングなどとつながりがあることが示され（本コラム 2020年9月号も参照）、再注目された。

　拡散モデルは生成品質が高く、最尤推定によって学習するので多様なデータを生成でき、かつ安定して学習させることができる。また部分情報から残りを復元したり、条件付けで生成を制御したりすることもできる。超解像^3）、着色、修復、JPEG圧縮後修正^4）などのタスクで最高精度を達成している。特に2022年4月に米OpenAIから発表されたDALL-E 2^5）は拡散モデルを使ってテキストからの画像生成をこれまでにない品質と多様性で実現した。DALL-E 2は内部で2つの拡散モデル（テキストからCLIP画像埋め込みベクトル、CLIP画像埋め込みベクトルから画像）を組み合わせて使っている。これについてはまた詳しく紹介したい。また、これまで難しかった動画生成も拡散モデルを使うことで実現できたと報告されている^6）。本稿では拡散モデルとは何か、なぜ優れているのか、どのような改良がなされてきたのかについて解説する。

拡散モデルとは

　拡散モデルは、観測データ$\mathbf{x}$にノイズを徐々に加えていき、観測変数と同じ次元数を持つ潜在変数${\mathbf{z}_t | t \in [0, 1]}$を得る。ここでの添字$t$はデータセット中の番号ではなく、ノイズを加えていった時の時刻である。このデータにノイズを入れて破壊していく過程を拡散過程とよび、次のように定義する。

\[ q(\mathbf{z}_t | \mathbf{x}) = \mathcal{N}(\mathbf{z}_t; \alpha_t\mathbf{x}, \sigma_t^2 \mathbf{I}) \] \[ q(\mathbf{z}_t | \mathbf{z}_s) = \mathcal{N}(\mathbf{z}_t; (\alpha_t/ \alpha_s) \mathbf{z}_s, \sigma^2_{t|s} \mathbf{I}) \]

ただし$0 \leq s < t \leq 1, \sigma^2_{t|s} = (1 - e^{\lambda_t -\lambda_s}) \sigma^2_t$、$\alpha_t,\sigma_t$は各時刻のシグナルとノイズの大きさを指定し、$\lambda_t = \log (\alpha_t^2 / \sigma_t^2) $はシグナルノイズ比（S/N比）の対数を表す。このS/N比は時刻が$1$に近づくとともに$0$に近づき、$t=1$では$\mathbf{z}_1$が完全なノイズになるように、つまり$q(\mathbf{z}_1) \approx \mathcal{N}(\mathbf{0},\mathbf{I})$となるようにする。この拡散過程によりデータ$\mathbf{x}$は徐々にノイズが増えていき、最終的に完全なノイズ$\mathbf{z}_1 \sim \mathcal{N}(\mathbf{0},\mathbf{I})$に変換される。正規分布に従うノイズを加えているため、任意の時刻の潜在変数が上記のように解析的に閉じた形で得られる特徴がある。

　拡散モデルはこの拡散過程の逆向きをたどることで生成モデルを学習する。この完全なノイズ$\mathbf{z}_1$がデータ$\mathbf{x}$に徐々に変わっていく過程を逆拡散過程とよび、生成モデルを定義する。逆拡散過程は時刻$t$においてノイズ$\mathbf{\epsilon}$が加わったデータ$\mathbf{z}_t \sim q(\mathbf{z}_t|\mathbf{x})$からノイズを除去（デノイズ）したデータ$\hat{\mathbf{x}}(\mathbf{z}_t,\lambda_t; \theta) \approx \mathbf{x}$を推定することで実現される（$\lambda_t$は以降省略する）。このデノイジングモデルは一般に$\hat{\mathbf{x}}(\mathbf{z}_t;\theta) = (\mathbf{z}_t - \sigma_t \mathbf{\epsilon}(\mathbf{z}_t ;\theta)) / \alpha_t$のようにモデル化し、ノイズ$\mathbf{\epsilon}$を$ \mathbf{z}_t $から$\mathbf{\epsilon}(\mathbf{z}_t ; \theta)$で予測するようにする。このデノイジングモデルは次の重み付き二乗誤差からなる目的関数を最小化するようにして学習する。

\[ \mathbb{E}_{\mathbf{\epsilon}, t} [w(\lambda_t) ||\hat{\mathbf{x}}(\mathbf{z}_t; \theta) - \mathbf{x}||_2^2 ] \]

この重み設定には様々な方法が提案されており、それらによって学習結果が異なってくる。このモデルを使ってデータをサンプリングする場合は、$s < t$の時、次が成り立つことを利用する^7）。

\[ q(\mathbf{z}_s | \mathbf{z}_t, \mathbf{x}) = \mathcal{N}(\mathbf{z}_s;\tilde{\mathbf{\mu}}_{s|t}(\mathbf{z}_t, \mathbf{x}),\tilde{\sigma}^2_{s|t} \mathbf{I}) \] \[ \tilde{\mathbf{\mu}}_{s|t}(\mathbf{z}_t, \mathbf{x}) = e^{\lambda_t -\lambda_s}(\alpha_s / \alpha_t) \mathbf{z}_t + (1 - e^{\lambda_t -\lambda_s}) \alpha_s \mathbf{x} \] \[ \tilde{\sigma}^2_{s|t} = (1 - e^{\lambda_t - \lambda_s}) \sigma^2_s \]

これを使って、はじめに、$\mathbf{z}_1 \sim \mathcal{N}(\mathbf{0},\mathbf{I})$とサンプリングし、次に$1$から$0$までを適当に離散化（均等に分割することが多い）した上で、伝承サンプリングを使って$\mathbf{z}_t$から$\mathbf{z}_s$を次のように求める。

\[ \mathbf{z}_s = \tilde{\mathbf{\mu}}_{s|t}(\mathbf{z}_t,\hat{\mathbf{x}}(\mathbf{z}_t; \theta)) +\sqrt{(\tilde{\sigma^2}_{s|t})^{1-\gamma}(\sigma^2_{t|s})^{\gamma}}\mathbf{\epsilon} \]

ただし$\mathbf{\epsilon}$は標準正規分布ノイズであり、$\gamma$はサンプリングの乱雑度合いを制御する変数である。

拡散モデルはVAEの特殊版とみなせる

　この拡散モデルによる生成モデルはノイズが加えられた変数を潜在変数とした、潜在変数モデルとみなすことができる。先程の連続時刻を$T$箇所に離散化し、各時刻のノイズが加わったデータを$\mathbf{x}_1,\mathbf{x}_2, \ldots, \mathbf{x}_T$（$\mathbf{x}_T =\mathbf{z}_1$）とする。$\mathbf{x}_{i:j} := \mathbf{x}_i, \mathbf{x}_{i+1}\ldots, \mathbf{x}_j$とする。また観測データを$\mathbf{x}_0$とする。このとき、観測データ$\mathbf{x}_0$の尤度は潜在変数を周辺化することで求められる。

\[ p(\mathbf{x}_0) = \int_{\mathbf{x}_{1:T}} p(\mathbf{x}_{0:T}) d \mathbf{x}_{1:T} \] このままでは尤度は積分が含まれているので直接最大化することはできないが、VAE（変分自己符号化器）と同様に拡散過程$q$を認識モデル、逆拡散過程の$p$を生成モデルとみなし、対数尤度の下限であるELBO最大化で対数尤度最大化を実現する。 \[ \log p( \mathbf{x}_0) \geq \mathbb{E}_{q} [ \log \frac{ p(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}] \] \[ \hspace{110pt}= \mathbb{E}_q [ \log p(\mathbf{x}_T) + \sum_{t \geq 1}^T \log\frac{p(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_t |\mathbf{x}_{t-1})} ] \]

ここで期待値をとっている$q$は$\mathbf{x}_{1:T} \sim q(\mathbf{x}_{1:T}|\mathbf{x}_0)$の略である。この右辺のうち$p(\mathbf{x}_T)$は完全なノイズなので無視できる。第二項が重要で、この項は$x_{t-1}$にノイズを加えて$x_{t}$となった時に、その逆向きを予測していることに対応する。この部分を変形すると先程の重み付き二乗誤差が導出される。

　上記のように拡散モデルはVAEと同様に潜在変数モデルを使った生成モデルであり、ELBO最大化によって学習しているともみなせる。だが、重要な部分で異なる。1つ目は拡散モデルは生成モデルのみ学習し、認識モデルは固定の拡散モデル（ノイズを加えるだけ）を使っていることである。生成モデルの学習において、その事後確率分布を予測する認識モデルの学習が一般に難しいが、拡散モデルはその必要がない。

　2つ目に拡散モデルは非常に多くのステップを使ってデータを徐々に変換する。これにより各ステップの変換が単純になり学習しやすくなる。一方ステップ数が増えるとパラメータ数が多くなってしまうが拡散モデルは全ステップで同じモデルを使い入力時刻$t$だけ変えて挙動を変えている。

　3つ目に、拡散モデルは学習時に全ステップを展開する必要がなく一部の時刻をサンプリングし学習できる。前述のように任意の時間差のノイズ、デノイズを閉じた式で導出できるためであり、任意時刻$t$をサンプリングして、それに対するデノイズを学習するだけで良い。

拡散モデルの改良

　拡散モデルは逐次的に生成する自己回帰モデルとは違ってデータ全体を並列に生成できるが、生成時に多くのステップを必要とするため生成が遅いのが問題である。この問題を解決するための手法もいくつか登場している。例えば、入力を自己符号化器などで低次元に変換しておき、この低次元の変数に対する拡散モデルを学習することで高速化できる^8）。もう1つが少ない生成ステップ数で生成できるよう蒸留していく方法である^9）。拡散モデルは連続的な生成過程を離散化しているため、ステップ数を少なくすると生成品質が急速に悪くなることが知られている。これに対し文献^9）では$T$ステップで生成できる拡散モデルがあった時にその途中潜在変数を目標に$T/2$ステップで生成できる拡散モデルを学習する。これを再帰的に繰り返していき少ないステップ数（4など）でも高品質の生成ができるようになる。今後も拡散モデルの改良が進み、より多くの問題で利用されるようになると考えられる。