What are Diffusion Models?
本記事はWhat are diffusion Models?を許可を得て翻訳したものです。Diffusion Modelの日本語の記事がまったくなかったので勉強がてら、翻訳してみました。誤訳や誤植などありましたらお知らせください。
Diffusion modelはとても新しいタイプの生成モデルで、どのような複雑なデータ分布でも学習することができ、なおかつその分布を解析的に評価することができます。
近年ではDiffusion modelは高精度な画像を生成でき、GANでSOTAなモデルより高い精度を達成しています。
これまでGAN、VAE、Flowといった生成モデルについての記事を書いてきました。どれも高いクオリティの画像を生成できますが、一方でそれぞれ個々の問題を抱えていました。GANは敵対的学習の構造そのものに、学習の不安定さと多様性の欠如があります。VAEは代理的に用いているロス関数に依存しています。Flowでは、逆変換を行えるモデル構造をしていなければいけません。
Diffusion Modelは非平衡熱力学から発想を得ています。マルコフ連鎖を用いて各diffusion stepで少しずつ画像などの実データにランダムノイズを足していきます。モデルは反対にノイズから少しずつ実データに近づけていくように学習します。VAEやFlowとは異なり、学習はすべて固定されて行われ、偏在変数は元データと同じ分だけの大きさを持ちます。
Fig. 1. 異なる生成モデルの概念図
What are Diffusion Models?
Diffusiom modelの基本的なアイディアに基づいて、数種類のモデルが提案されてきました。
Forward diffusion process
まずforward diffusion stepについて説明しましょう。ここでは実データxを実データ分布qからサンプリングします。つまり\mathbf{x}_0 \sim q(\mathbf{x})です。diffusion stepはT回行うとします。すべてのステップでx_0に小さなガウシアンノイズを足していってノイズに近づいていくx_1,...,x_Tを作成します。ステップの幅は\{\beta_t \in (0, 1)\}_{t=1}^tでコントロールします。
q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I}) \quad
q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod^T_{t=1} q(\mathbf{x}_t \vert \mathbf{x}_{t-1})
実データx_0はステップが進むごとに徐々に自身の特徴を失っていきます。T \to \inftyであるときx_Tはガウシアンノイズそのものになります。
図２ マルコフ連鎖のforward (reverce) difusion の進行によりノイズが加算（除去）されている。画像はHo et al. 2020に文字をつけた。
上記のプロセスが優れている点はどのようなstep数Tでもreparameterization trickを用いて 任意の時間tにおけるx_tを閉形式（訳注：ここでは加減乗除や初等関数の合成関数による解の表し方を意味する）で表せることです。
\begin{aligned}
\mathbf{x}_t
&= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\mathbf{z}_{t-1} & \text{ ;where } \mathbf{z}_{t-1}, \mathbf{z}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\
&= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\mathbf{z}}_{t-2} & \text{ ;where } \bar{\mathbf{z}}_{t-2} \text{ merges two Gaussians (*).} \\
&= \dots \\
&= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z} \\
q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I})
\end{aligned}
(*)
ステップが進むごとにノイズが大きくなるようにしたいので、\beta_1 < \beta_2 < \dots < \beta_Tであり\bar{\alpha}_1 > \dots > \bar{\alpha}_Tです。
確率的勾配Langevin動力学とのつながり
Langevin動力学は物理学において静的な分子系のモデリングに寄与しています。これが確率的勾配法と組み合わされた、確率的勾配Langevin動力学(SGLD)(Welling & Teh 2011)は確率分布p(x)からマルコフ連鎖の更新を重ねることで、サンプリングを行うことができます。そしてその時必要なものは勾配\nabla_\mathbf{x} \log p(\mathbf{x})だけなのです。具体的には以下の反復式に従ってx_tを更新していきます。
\mathbf{x}_t = \mathbf{x}_{t-1} + \frac{\epsilon}{2} \nabla_\mathbf{x} p(\mathbf{x}_{t-1}) + \sqrt{\epsilon} \mathbf{z}_t
,\quad\text{where }
\mathbf{z}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})
ここで\epsilonはステップサイズです。T \to \infty, \epsilon \to 0のときx_tは確率密度関数p(x)と等しくなります。
通常のSGDと比べると、SGLDはパラメータの更新の際にガウシアンノイズを加算しているのがわかります。これにより局所解への陥りを防ぐことができます。
Reverse diffusion process
もし上記のプロセスを反対にしてq(x_{t-1}|x_t)からサンプルできるなら、ガウシアンノイズ\mathbf{x_T} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})から実データを復元できます。（\beta_tが十分小さければq(x_{t-1}|x_t)もまたガウシアンノイズです）しかし残念がなら、qを推定するのは簡単ではありません。なぜならそのためにはデータセット全体を用いて条件付き確率を推定するモデルp_{\theta}を学習する必要があるからです。
p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod^T_{t=1} p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) \quad
p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))
図３ ２次元のスイスロールモデルを学習するdiffusion model
逆条件付き確率はx_0を用いると扱いやすくなります。
q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I})
ベイズの定理を用いて
\begin{aligned}
q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)
&= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\
&\propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\
&= \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_t}}{1 - \bar{\alpha}_t} \mathbf{x}_0)} \mathbf{x}_{t-1} + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)
\end{aligned}
ここでC(x_t,x_0)はx_{t-1}を含まない関数で、詳細は省略します。正規分布の確率密度関数より平均と分散は以下のように求められます。
\begin{aligned}
\tilde{\beta}_t &= 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t \\
\tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0)
&= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_t}}{1 - \bar{\alpha}_t} \mathbf{x}_0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})
= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0\\
\end{aligned}
前に行った計算で、\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t)という素晴らしい式がありました。これを上の式に代入します。
\begin{aligned}
\tilde{\boldsymbol{\mu}}_t
&= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t) \\
&= \color{cyan}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_t \Big)}
\end{aligned}
図２でお見せしたように、このモデルはVAEにとても似ていて、同じように尤度最適化のために変分下限を用いることができます。
\begin{aligned}
- \log p_\theta(\mathbf{x}_0)
&\leq - \log p_\theta(\mathbf{x}_0) + D_\text{KL}(q(\mathbf{x}_{1:T}\vert\mathbf{x}_0) \| p_\theta(\mathbf{x}_{1:T}\vert\mathbf{x}_0) ) \\
&= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_{\mathbf{x}_{1:T}\sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T}) / p_\theta(\mathbf{x}_0)} \Big] \\
&= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_q \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} + \log p_\theta(\mathbf{x}_0) \Big] \\
&= \mathbb{E}_q \Big[ \log \frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\
\text{Let }L_\text{VLB}
&= \mathbb{E}_{q(\mathbf{x}_{0:T})} \Big[ \log \frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \geq - \mathbb{E}_{q(\mathbf{x}_0)} \log p_\theta(\mathbf{x}_0)
\end{aligned}
同じようにJansenの不等式を使った結果を得ることもできます。そこでcross entropyを最小化したいと思います。
\begin{aligned}
L_\text{CE}
&= - \mathbb{E}_{q(\mathbf{x}_0)} \log p_\theta(\mathbf{x}_0) \\
&= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T} \Big) \\
&= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \int q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} d\mathbf{x}_{1:T} \Big) \\
&= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \mathbb{E}_{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \Big) \\
&\leq - \mathbb{E}_{q(\mathbf{x}_{0:T})} \log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \\
&= \mathbb{E}_{q(\mathbf{x}_{0:T})}\Big[\log \frac{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})} \Big] = L_\text{VLB}
\end{aligned}
各項を解析的に計算可能にするために、目的関数をさらに幾つかのKL divergenceとエントロピー項で置き換えましょう。(付録Bの式変形のより詳しい内容はここを見てね)
\begin{aligned}
L_\text{VLB}
&= \mathbb{E}_{q(\mathbf{x}_{0:T})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\
&= \mathbb{E}_q \Big[ \log\frac{\prod_{t=1}^T q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{ p_\theta(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) } \Big] \\
&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=1}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} \Big] \\
&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \Big( \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\
&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big]\\
&= \mathbb{E}_q \Big[ \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \Big] \\
&= \mathbb{E}_q [\underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} ]
\end{aligned}
個々の項にラベルをつけましょう。
\begin{aligned}
L_\text{VLB} &= L_T + L_{T-1} + \dots + L_0 \\
\text{where } L_T &= D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T)) \\
L_t &= D_\text{KL}(q(\mathbf{x}_t \vert \mathbf{x}_{t+1}, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_t \vert\mathbf{x}_{t+1})) \text{ for }1 \leq t \leq T-1 \\
L_0 &= - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)
\end{aligned}
上の式におけるL_0を除くすべてのKL項は２つの正規分布を比較しているので、閉形式で求めることができます。またL_Tは定数であり、学習可能なパラメタを含んでいないので学習中は無視することができます。なぜならqは学習可能なパラメタを含んでおらず、x_Tはただのガウシアンノイズだからです。Ho et al. 2020ではL_0を\mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1, 1))から派生した別の離散的なデコーダを使ってモデル化しました。
Training Lossのためのパラメータ化
反対向きのdiffuison processにおける条件付き確率分布p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))の近似のためにニューラルネットワークを学習しましょう。
ここで\tilde{\boldsymbol{\mu}}_t = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_t \Big)を予測するように\mu_{\theta}を学習させましょう。x_tは学習中に利用することができるので、x_tとtからz_tを推測する代わりにガウシアンノイズ項を求めることができます。
\begin{aligned}
\boldsymbol{\mu}_\theta(\mathbf{x}_t, t) &= \color{cyan}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_\theta(\mathbf{x}_t, t) \Big)} \\
\text{Thus }\mathbf{x}_{t-1} &= \mathcal{N}(\mathbf{x}_{t-1}; \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_\theta(\mathbf{x}_t, t) \Big), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))
\end{aligned}
ロス項L_tは\tilde{\mu}との差を最小化します。
\begin{aligned}
L_t
&= \mathbb{E}_{\mathbf{x}_0, \mathbf{z}} \Big[\frac{1}{2 \| \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) \|^2_2} \| \color{blue}{\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0)} - \color{green}{\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)} \|^2 \Big] \\
&= \mathbb{E}_{\mathbf{x}_0, \mathbf{z}} \Big[\frac{1}{2 \|\boldsymbol{\Sigma}_\theta \|^2_2} \| \color{blue}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \mathbf{z}_t \Big)} - \color{green}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\mathbf{z}}_\theta(\mathbf{x}_t, t) \Big)} \|^2 \Big] \\
&= \mathbb{E}_{\mathbf{x}_0, \mathbf{z}} \Big[\frac{ \beta_t^2 }{2 \alpha_t (1 - \bar{\alpha}_t) \| \boldsymbol{\Sigma}_\theta \|^2_2} \|\mathbf{z}_t - \mathbf{z}_\theta(\mathbf{x}_t, t)\|^2 \Big] \\
&= \mathbb{E}_{\mathbf{x}_0, \mathbf{z}} \Big[\frac{ \beta_t^2 }{2 \alpha_t (1 - \bar{\alpha}_t) \| \boldsymbol{\Sigma}_\theta \|^2_2} \|\mathbf{z}_t - \mathbf{z}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t, t)\|^2 \Big]
\end{aligned}
Simplification
経験的にHo et al. (2020)は以下の係数を無視したロス関数を使うとdiffusion modelの学習がよりうまく行くことを発見しました。
L_t^\text{simple} = \mathbb{E}_{\mathbf{x}_0, \mathbf{z}_t} \Big[\|\mathbf{z}_t - \mathbf{z}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z}_t, t)\|^2 \Big]
最終的に簡単化されたロス関数は以下のようになりました。
L_{simple}=L^{simple}_{t}+C
ここでCは\thetaに依存しない定数項です。
図４ DDPMにおけるtrainとsamplingのアルゴリズム(Ho et al. 2020より)
noise-conditionsed score networks(NCSN)との関係
Song & Ermon (2019)はスコア推定に基づく生成モデルを提案しました。これはスコアマッチングにより推定されたデータ分布の勾配を用いて、Langevin 動力学を通して生成されます。個々のデータxの確率密度のスコアは\nabla_{\mathbf{x}} \log p(\mathbf{x})で定義され、モデルs_{\theta}:\mathbb{R}^D \to \mathbb{R}^Dはこれを推定するように学習されます。深層学習で扱うために高次元へ拡張するために、denoising score matching(事前に決めておいた小さなノイズをデータに加算する;(Vincent, 2011))やslided score matching(Random Projectionを使う;(Yang et al., 2019))が提案されました。
Langevin 動力学はその反復的なプロセスにおいて、スコアの勾配\nabla_{\mathbf{x}} \log p(\mathbf{x})のみを用いて確率密度分布からデータをサンプリングできましたね。
しかしながら多様体仮説によれば、たとえ観測されたデータがどんな高次元にいるように見えたとしても、高次元空間では多くのデータは低次元多様体、かなり小さい空間に集中しています。データが空間全体をカバーしないので、スコア推定に問題が生じます。データの密度が小さい空間では、スコア推定は精度が悪くなるのです。そのため空間全体\mathbf{R^D}上のデータに僅かなガウシアンノイズを足して,「ぼかす」ことでスコア推定の学習はより安定するようになりました。Song & Ermon (2019)はノイズの強さを学習中に変えて、ノイズの強さが異なるデータをまとめて学習することで改善させました。
ノイズの強さが増大する様子は、今回のdiffusion processに似ています。
スケジューリングのパラメタ化
Ho et al.(2020)では実データからガウシアンノイズへの近づけかたを決定する分散\betaのスケジューリングは\beta_1=10^{-4}から\beta_T=0.02へ線形に増加するようにされていました。これは[-1,1]へ正規化された画像データと比べるととても小さいですよね。Diffusion modelは高い精度で生成できていましたが、依然として他の生成モデルに尤度で劣っていました。
Nichol & Dhariwal(2021)はより優れた尤度を達成するために幾つかのテクニックを提案しました。その一つに分散のスケジューリングを初めと終わり(t=0とt=T付近)は分散の変化が大きく、途中では線形にすると言ったものがあります。式で示すと以下のようになります。
\beta_t = \text{clip}(1-\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}, 0.999) \quad\bar{\alpha}_t = \frac{f(t)}{f(0)}\quad\text{where }f(t)=\cos\Big(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2}\Big)
ここでオフセットsはt=0付近で\beta_tが小さくなりすぎることを防いでいます。
図５ 学習中の\beta_tの線形とcosine-basedのスケジューリングの比較
Ho et al.(2020)は\beta_tを学習可能にはせず指定した値を用いました。また\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) = \sigma^2_t \mathbf{I}において\sigma_tを学習可能にせず\beta_tか\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_tを用いました。なぜなら対角分散\boldsymbol{\Sigma}_\thetaを学習可能にすると学習が不安定になり生成の質が落ちたからだそうです。
Nichol & Dhariwal(2021)は\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t)をモデルが推定する混合ベクトルvを用いて、\beta_tと\tilde{\beta_t}で加重平均を取ることを提案しました。
\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) = \exp(\mathbf{v} \log \beta_t + (1-\mathbf{v}) \log \tilde{\beta}_t)
しかしながら、L_{simple}は\boldsymbol{\Sigma}_\thetaを式に含んでいません。このため新たなロス関数L_{hybrid}=L_{simple}+\lambda L_{VLB}を用いています。このとき{\lambda=0.001}と小さく、{L_{VLB}}の項にて\boldsymbol{\Sigma}_\thetaのみについて学習するように\mu_{\theta}は勾配が止められています。経験的にL_{VLB}の学習はノイジーな勾配などの原因により難しいことが知られていたので、L_{VLB}の時間平均をしたものを学習することが提案されています。
図６ 改善されたDDPMとほかの尤度最適化に基づく生成モデルのNegative log-likelihood(NLL)の比較。NLLの単位はbits/dim(Nichol & Dhariwal, 2021より)
Duffusion modelのサンプリングの高速化
DDPMではサンプリングはマルコフ連鎖に基づき、Tが１や数千ステップになるまで繰り返されますが、とても時間がかかります。Song et al.(2020)によれば以下のように述べられています。
DDPMではNvidia 2080 Ti GPUでも32x32を50k枚サンプリングするのに20時間もかかる。でもGANなら1分とかからない。
単純な解決方法はサンプリングを飛ばし飛ばしで行うことです(Nichol & Dhariwal, 2021)。ステップ数をTからSに減らして、[T/S]ステップ毎にサンプリングを行い、新たなスケジューリング\{\tau_1, \dots, \tau_S\} where \tau_1 < \tau_2 < \dots <\tau_S \in [1, T] and S<Tでサンプリングされます。
他のアプローチとして、q_\sigma(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)を前に計算した素晴らしい式を用いて望ましいものに書き換えて見ましょう。
\begin{aligned}
\mathbf{x}_{t-1}
&= \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1}}\mathbf{z}_{t-1} \\
&= \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \mathbf{z}_t + \sigma_t\mathbf{z} \\
&= \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \frac{\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}_0}{\sqrt{1 - \bar{\alpha}_t}} + \sigma_t\mathbf{z} \\
q_\sigma(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)
&= \mathcal{N}(\mathbf{x}_{t-1}; \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \frac{\mathbf{x}_t - \sqrt{\bar{\alpha}_t}\mathbf{x}_0}{\sqrt{1 - \bar{\alpha}_t}}, \sigma_t^2 \mathbf{I})
\end{aligned}
ここでq(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I})なので、以下を得ます。
\tilde{\beta}_t = \sigma_t^2 = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t
サンプリングの乱雑さをコントロールするようなハイパラメタ\eta \in \mathbb{R}^+を定めて、\sigma_t^2 = \eta \cdot \tilde{\beta}_tのようにしましょう。例えば\eta = 0のような特別なケースでは、サンプリングはランダムさが排除され決定的になります。このようなモデルはdenoising diffusion implicit model(DDIM;Song et al., 2020)と名付けられました。DDIMは同じ結合ノイズ分布を持ちますが、決定的にノイズと実データを対応付けます。
生成する際、diffusion step S=\{\tau_1, \dots, \tau_S\}の部分集合のみでサンプリングできるので、生成の手順は
以下のようになります。
q_{\sigma, \tau}(\mathbf{x}_{\tau_{i-1}} \vert \mathbf{x}_{\tau_t}, \mathbf{x}_0)
= \mathcal{N}(\mathbf{x}_{\tau_{i-1}}; \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \frac{\mathbf{x}_{\tau_i} - \sqrt{\bar{\alpha}_t}\mathbf{x}_0}{\sqrt{1 - \bar{\alpha}_t}}, \sigma_t^2 \mathbf{I})
実験ではすべてのモデルがT=1000で学習され、Sが小さい時DDPM(\eta=1)は精度が低く、反対にDDIM(\eta=0)は最も良い精度でした。
図７ CIFAR10とCelebAにおけるDDIM(\eta=0)とDDPM(\hat{\sigma})FIDの比較実験結果(Song et al., 2020)
DDPMと比べてDDIMは
- より少ないステップ数でよりクオリティが高い生成が可能
- 生成のプロセスが決定的なので一貫性があり、近い潜在変数を用いると同じような画像を生成できる
- この一貫性によりDDIMは潜在変数の意味的な補完が可能
クラス付き生成
ImageNetを使う時は、クラスを指定して生成することはもはや普通ですよね。diffusion modelでクラス情報を扱うために、Dhariwal & Nichol(2021)はノイズ入画像x_tを用いて分類器f_\phi(y \vert \mathbf{x}_t, t)を学習させて、サンプリングではクラスyへ誘導するために勾配\nabla_\mathbf{x} \log f_\phi(y \vert \mathbf{x}_t, t)を用いました。
彼らのablated diffusion model(ADM)とクラスガイド付きモデル(ADM-G)はBig GANなどのSOTAな生成モデルより優れた精度を達成しました。
図８ クラスガイドをつけたDDPMとDDIMのアルゴリズム(Dhariwal & Nichol, 2021)
Dhariwal & Nichol(2021)はU-Net型のモデルを改造してdiffusion modelがGANより優れていることを示しました。主な変更点はより大きなdepthとwidth、より多いattention head、複数解像度におけるattention、BigGANのresidual blockを用いたup/down sampling、1/\sqrt{2}ごとのresidual connection、adaptive groupe normalization(AdaGN)など多岐に渡ります。
Quick Summary
-
良い点: 生成モデルでは扱いやすさと柔軟さの両立は難しいです。扱いやすいモデルは解析的に評価ができ、例えばGaussian やLaplaceを通して安価に学習が行えますが、リッチな画像のデータセットでの表現は難しいものがあります。柔軟なモデルは任意のデータ構造に対して学習することができますが、モデルの学習、評価、生成はいつも大変です。diffusion modelは扱いやすさと柔軟さが両立していると言えるでしょう。
-
悪い点: doffision modelの生成ステップは長いマルコフ連鎖に依存していて、処理に多大な時間がかかります。高速化手法は新たに提案されていますが、GANと比較すると依然として遅いです。
References
[1] Jascha Sohl-Dickstein et al. “Deep Unsupervised Learning using Nonequilibrium Thermodynamics.” ICML 2015.
[2] Max Welling & Yee Whye Teh. “Bayesian learning via stochastic gradient langevin dynamics.” ICML 2011.
[3] Yang Song & Stefano Ermon. “Generative modeling by estimating gradients of the data distribution.” NeurIPS 2019.
[4] Yang Song & Stefano Ermon. “Improved techniques for training score-based generative models.” NeuriPS 2020.
[5] Jonathan Ho et al. “Denoising diffusion probabilistic models.” arxiv Preprint arxiv:2006.11239 (2020). code
[6] Jiaming Song et al. “Denoising diffusion implicit models.” arxiv Preprint arxiv:2010.02502 (2020). code
[7] Alex Nichol & Prafulla Dhariwal. “ Improved denoising diffusion probabilistic models” arxiv Preprint arxiv:2102.09672 (2021). code
[8] Praffula Dhariwal & Alex Nichol. “Diffusion Models Beat GANs on Image Synthesis.” arxiv Preprint arxiv:2105.05233 (2021). code
Discussion
元記事の通りなので誤訳ではありませんが、ベータのスケジューリングの元論文Nichol & Dhariwal(2021)ではf(t)の定義でcosに二乗が付きます。