Diffusion Modelについてこの論文が分かりやすく解説していたので、備忘として簡潔に要点をまとめておく。
ざっくりまとめ
- Diffusion Modelは、VAEと同様の潜在変数を持つ生成モデルとして理解できる
- 学習は「ノイズの乗った画像から、その画像に乗ったノイズを予測する」ことに帰着する
モデルの特徴
Diffusion Modelは本質的にはVAEと同じく、潜在変数を持つ生成モデルとして捉えられる。ただし、以下の点で特徴的。
- 潜在変数が階層的になっている
- VAEの潜在変数は階層的な構造を成していない
- 各潜在変数の次元は全て元の画像と同じ
- VAEでは通常潜在変数の次元は元画像より低い
- エンコーダーでは画像に逐次的にガウシアンノイズを加えるだけで、学習するパラメータはない
- VAEではデコーダーとエンコーダーの双方をNNでモデリングする
- 最終的な潜在変数が（多変量）標準正規分布に近づくように元画像をエンコーディングする
モデルの定式化
元の画像に対応する確率変数を$x_0$、これと同じ次元の階層的な潜在変数を$x_1, \dots, x_T$とする。
デコーダー
VAEと同様に、潜在変数を持つ生成モデルとしてデコーダーを定義する。ただし、潜在変数は次のように階層的な構造を持ち、潜在変数の系列はマルコフ性を満たすものとする。
x_T \sim N(\mathbf{0}, \mathbf{I}), \quad x_{t-1} | x_t \sim p_\theta(x_{t-1} | x_t) \quad (t = 1, \dots, T)
各条件付き分布$p_\theta(x_{t-1} | x_t)$はNNを用いてモデリングするが、詳細については後述する。
エンコーダー
デコーダーで定義したモデルの尤度を最大化するようにパラメータを学習するが、VAEと同様に直接的な尤度の最大化は困難なため、変分下界の最大化を通して間接的に最大化を行う（いわゆる変分推論）。それにあたり、潜在変数の事後分布$p_\theta(x_{1:T}|x_0)$近似するようなモデル$q(x_{1:T}|x_0)$をエンコーダーとして用意する。
VAEではここをNNでモデリングする一方、Diffusion Modelでは（やや天下り的なようにも感じるが）潜在変数の系列にマルコフ性
q(x_{1:T}|x_0) = \Pi_{t=1}^{T}q(x_t|x_{t-1})
を課した上で、各条件付き分布を次のような形でモデル化する。
x_t = \sqrt{\alpha_t}x_{t-1} + \sqrt{1-\alpha_t} \epsilon_t, \quad \epsilon_t \sim N(\mathbf{0}, \mathbf{I}) \quad (t = 1, \dots, T)
ここで、$\alpha_t$はある定数とする。つまり、デコーダに学習するパラメータは存在しない。このようなモデル化の下、エンコーダーは「元の画像に逐次的に独立なガウスノイズを乗せる」という手続きとして見ることができる。
なお、このエンコーダーの逐次的な手続きをひとまとめにした$q(x_t|x_0)$は以下のように表される。
x_t = \sqrt{\prod_{i=1}^{t} \alpha_t}x_0 + \sqrt{1 - \prod_{i=1}^{t} \alpha_t} \epsilon_0 =: \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha_t}}\epsilon_0, \quad \epsilon_0 \sim N(\mathbf{0}, \mathbf{I}) \quad (\ast)
$t$が十分大きいと$\bar{\alpha}_t \approx 0$となるため、任意の入力画像$x_0$に対して最終的な潜在変数の分布は（多変量）標準正規分布に近くなる。
以上を踏まえると、Diffusion Modelは元の画像に逐次的にノイズを加えることで標準正規分布に近づけるエンコーダーと、標準正規分布に従うノイズそのものから画像を再構成するデコーダーからなるモデルとみなせる。
目的関数
エンコーダー、デコーダーとして定義したモデルがいずれもマルコフ性を満たすことを活かして頑張って計算すると、サンプル$x_0$に対し、変分下界は以下のように導出できる。
\begin{align*}
E_{q(x_{1:T}|x_0)}\left[\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}\right] &=
E_{q(x_1|x_0)}\left[\log p_{\theta}(x_0|x_1)\right] \\
&- D_{KL}(q(x_T|x_0)\ || \ p(x_T))\\ &- \sum_{t=2}^{T}E_{q(x_t|x_0)}\left[D_{KL}(q(x_{t-1}|x_t, x_0) \ || \ p_{\theta}(x_{t-1}|x_{t})) \right]
\end{align*}
各項の意味は以下の通り。
1項目（reconstruction term）
VAEで言う再構成誤差に対応する項。VAE同様モンテカルロ近似により計算可能。
2項目（prior matching term）
最終的な潜在変数の分布がどれほど標準正規分布に近いかを表す項。学習可能なパラメータを含まないので、学習の際には無視できる。また、エンコーディングの手続きを経た最終的な潜在変数の分布は標準正規分布に近いため、この項はほぼ0と等しくなる。
3項目（denoising matching term）
ここがDiffusion Model特有の項。デコーダーのモデル$p_{\theta}(x_{t-1}|x_t)$を、エンコーダーの手続きを逆にした（＝デノイジングに対応）$q(x_{t-1}|x_t, x_0)$に近づけることでこの項を小さくできる。
以下では、denoizing matching termのKLダイバージェンスの部分をより詳細に見る。エンコーダーを正規分布でモデル化していることから、$q(x_{t-1}|x_t, x_0)$も同様に正規分布となり、次のような解析的な形で書ける。
q(x_{t-1}|x_t, x_0) = N(\mu_q(x_t, x_0), \Sigma_q(t))
\mu_q(x_t, x_0) := \frac{1}{\sqrt{\alpha_t}} x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha}_t} \epsilon_0, \quad \Sigma_q(t) := \frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} \mathbf{I} =: \sigma_q^2(t) \mathbf{I}
ここで、$\bar{\alpha}_t$や$\epsilon_0$は式$(*)$で定義された量である。
ここにきて、正体をぼかしていたエンコーダモデル$p_{\theta}(x_{t-1}|x_t)$の話に立ち返る。denoising matching termを小さくするためには、$p_{\theta}(x_{t-1}|x_t)$が$q(x_{t-1}|x_t, x_0)$の近似となるようにモデリングすることが肝要だが、この目的に沿うよう、エンコーダーを次のようにモデリングする。
p_{\theta}(x_{t-1}|x_t) = N(\mu_\theta(x_t), \Sigma_q(t)), \quad \mu_\theta(x_t) := \frac{1}{\sqrt{\alpha_t}} x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha}_t} \hat{\epsilon}_\theta(x_t, t)
つまり、元の画像に付加されたノイズ$\epsilon_0$を、ノイズが乗った画像$x_t$の情報のみから予測するNN$\hat{\epsilon}_\theta(x_t, t)$により近似する。
このモデルの設定の下で、KLダイバージェンスは
D_{KL}(q(x_{t-1}|x_t, x_0) \ || \ p_{\theta}(x_{t-1}|x_t)) = \frac{1}{2 \sigma_q^2(t)} \frac{(1-\alpha_t)^2}{(1-\bar{\alpha}_t) \alpha_t} ||\epsilon_0 - \hat{\epsilon}_\theta(x_t, t)||^2
という形で書き下せる。
以上をまとめると、Diffusion Modelの学習は「様々なレベルでノイズの乗った画像$x_t \ (t = 1, \dots, T)$それぞれから、その画像に乗っているノイズを予測する」ようなNNを学習することに帰着する、と理解できる。