拡散モデル(Diffusion Model)の概要と式定義まとめ

拡散とDenoisingに基づく拡散モデル(Diffision Model)は多くの生成モデル(generative model)に導入される概念です。当記事では拡散モデルの概要と式定義、イェンセンの不等式などを用いるloss関数の導出などについて取りまとめを行いました。
DDPM論文や「拡散モデル ーデータ生成技術の数理(岩波書店)」の$2$章の「拡散モデル」などを参考に作成を行いました。

・用語/公式解説
https://www.hello-statisticians.com/explain-terms

拡散モデルの概要

拡散モデルの概要は上図を元に理解すると良い。$\mathbf{x}_{0}$が画像、$\mathbf{x}_{T}$が$\mathbf{x}_{0}$と次元が同じ潜在変数を表す。また、図の$q$が拡散過程(diffusion process)、$p$が逆向き過程(reverse process)にそれぞれ対応する。

拡散モデルでは元データに対し、拡散過程を適用した結果を逆向き過程で復元できるように逆向き過程$p$の学習を行うことで、新たに潜在変数の$\mathbf{x}_{T}$が与えられた際に$\mathbf{x}_{0}$を計算し、生成を行うことができる。

拡散モデルの式定義

拡散過程と逆拡散過程の式定義

拡散モデル(Diffusion model)は潜在変数モデルであり、下記のような式で定義される。
$$
\large
\begin{align}
p_{\theta}(\mathbf{x}_{0}) = \int p_{\theta}(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T} \quad (1)
\end{align}
$$

$(1)$式は生成される画像の分布は同時分布$p_{\theta}(\mathbf{x}_{0:T})$を$\mathbf{x}_{1}, \cdots , \mathbf{x}_{T}$について積分し、周辺分布を得たと解釈すると良い。ここで式に出てくる同時分布$p_{\theta}(\mathbf{x}_{0:T})$は下記のように定義する。
$$
\large
\begin{align}
p_{\theta}(\mathbf{x}_{0:T}) &= p(\mathbf{x}_{T}) \prod_{t=1}^{T} p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) \quad (2) \\
p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) &= \mathcal{N}(\boldsymbol{\mu}_{\theta}(\mathbf{x}_{t},t), \boldsymbol{\Sigma}_{\theta}(\mathbf{x}_{t}, t)) \quad (3) \\
p(\mathbf{x}_{T}) &= \mathcal{N}(\mathbf{0}, \mathbf{I})
\end{align}
$$

$(2)$式の同時確率分布の分解(factorization)ではマルコフ性(Markov Property)を仮定している。また、$(3)$式は前節における逆向き過程を表す。多次元正規分布の平均ベクトルと共分散行列はDeep Learningを用いて学習を行ったパラメータ$\theta$の関数から出力される値である。

拡散モデルの特徴は$p$で表される逆向き過程(reverse process)に対応する拡散過程(diffusion process)の$q$を定義することである。拡散過程$q$はvariance scheduleの$\beta_1, \cdots , \beta_{T}$を用いて下記のように定義される。
$$
\large
\begin{align}
q(\mathbf{x}_{1:T}|\mathbf{x}_{0}) &= \prod_{t=1}^{T} q(\mathbf{x}_{t}|\mathbf{x}_{t-1}) \quad (4) \\
q(\mathbf{x}_{t}|\mathbf{x}_{t-1}) &= \mathcal{N} \left( \sqrt{1-\beta_{t}} \mathbf{x}_{t}, \beta_{t} \mathbf{I} \right) \quad (5)
\end{align}
$$

任意時刻の拡散条件付き確率の証明

$$
\large
\begin{align}
q(\mathbf{x}_{1:T}|\mathbf{x}_{0}) &= \prod_{t=1}^{T} q(\mathbf{x}_{t}|\mathbf{x}_{t-1}) \quad (4) \\
q(\mathbf{x}_{t}|\mathbf{x}_{t-1}) &= \mathcal{N} \left( \sqrt{1-\beta_{t}} \mathbf{x}_{t-1}, \beta_{t} \mathbf{I} \right) \quad (5)
\end{align}
$$

拡散過程$q$が上記のように表される時、任意時刻$t$における拡散条件付き確率$q(\mathbf{x}_{t}|\mathbf{x}_{0})$は下記のように表される。
$$
\large
\begin{align}
q(\mathbf{x}_{t}|\mathbf{x}_{0}) &= \mathcal{N}(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}, (1 \, – \, \bar{\alpha}_{t}) \mathbf{I}) \quad (6) \\
\alpha_{t} &= 1-\beta_{t} \quad (7) \\
\bar{\alpha}_{t} &= \prod_{s=1}^{t} \alpha_{s} \quad (8)
\end{align}
$$

上記の$(6)$式はDDPM論文の$\mathrm{Eq}. \: (4)$に対応する。以下、数学的帰納法を用いて上記の証明を行う。

・$[1] \,$ $t=1$の時
$(7)$式と$(8)$式より下記が成立する。
$$
\large
\begin{align}
\sqrt{\bar{\alpha}_{1}} &= \sqrt{\alpha_{1}} = \sqrt{1-\beta_{1}} \\
1 \, – \, \bar{\alpha_{1}} &= 1 \, – \, \alpha_{1} = \beta_{1}
\end{align}
$$

上記より$t=1$の時$(6)$式は$(5)$式の定義式に一致するので$(6)$式は成立する。

・$[2] \,$ $t=k$で$(6)$式が成立する場合の$t=k+1$の時
$t=k$で$(6)$式が成立するとき、下記が成立する。
$$
\large
\begin{align}
q(\mathbf{x}_{k}|\mathbf{x}_{0}) = \mathcal{N}(\sqrt{\bar{\alpha}_{k}} \mathbf{x}_{0}, (1 \, – \, \bar{\alpha}_{k}) \mathbf{I}) \quad (6)’
\end{align}
$$

このとき$(6)’$式に基づいて$\mathbf{x}_{k}$は下記のように表せる。
$$
\large
\begin{align}
\mathbf{x}_{k} = \sqrt{\bar{\alpha}_{k}} \mathbf{x}_{0} + \sqrt{1 \, – \, \bar{\alpha}_{k}} \epsilon, \quad \epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \quad (9)
\end{align}
$$

ここで$(5)$式の$q$の定義式より、$q(\mathbf{x}_{k+1}|\mathbf{x}_{k})$は下記のように表される。
$$
\large
\begin{align}
q(\mathbf{x}_{k+1}|\mathbf{x}_{k}) = \mathcal{N} \left( \sqrt{1-\beta_{t+1}} \mathbf{x}_{t}, \beta_{t+1} \mathbf{I} \right) \quad (5)’
\end{align}
$$

ここで$(5)’$式と$(9)$式に基づいて$\mathbf{x}_{k+1}$は下記のように表せる。
$$
\large
\begin{align} \mathbf{x}_{k+1} &= \sqrt{1-\beta_{t+1}} \mathbf{x}_{t} + \sqrt{\beta_{t+1}} \epsilon, \quad \epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\
&= \sqrt{1-\beta_{t+1}} (\sqrt{\bar{\alpha}_{k}} \mathbf{x}_{0} + \sqrt{1 \, – \, \bar{\alpha}_{k}} \epsilon) + \sqrt{\beta_{t+1}} \epsilon \\
&= \sqrt{\alpha_{t+1}} (\sqrt{\bar{\alpha}_{k}} \mathbf{x}_{0} + \sqrt{1 \, – \, \bar{\alpha}_{k}} \epsilon) + \sqrt{\beta_{t+1}} \epsilon \\
&= \sqrt{\bar{\alpha}_{k+1}} \mathbf{x}_{0} + \sqrt{\alpha_{k+1}} \sqrt{1 \, – \, \bar{\alpha}_{k}} \epsilon + \sqrt{\beta_{t+1}} \epsilon \quad (10)
\end{align}
$$

$(10)$式より、$\mathbf{x}_{k+1}$の分散は$\mathbf{X}, \mathbf{Y} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$とする際の$\sqrt{\alpha_{k+1}} \sqrt{1 \, – \, \bar{\alpha}_{k}} \mathbf{X} + \sqrt{\beta_{t+1}} \mathbf{Y}$の分散に一致する。ここで正規分布の再生性より、$\sqrt{\alpha_{k+1}} \sqrt{1 \, – \, \bar{\alpha}_{k}} \mathbf{X} + \sqrt{\beta_{t+1}} \mathbf{Y}$の分散$\mathbf{\Sigma}$は下記のように計算できる。
$$
\large
\begin{align}
\mathbf{\Sigma} &= (\sqrt{\alpha_{k+1}} \sqrt{1 \, – \, \bar{\alpha}_{k}})^{2} \mathbf{I} + (\sqrt{\beta_{t+1}})^{2} \mathbf{I} \\
&= \alpha_{k+1} (1 \, – \, \bar{\alpha}_{k}) \mathbf{I} + \beta_{t+1} \mathbf{I} \\
&= \left( \cancel{\alpha_{k+1}} – \bar{\alpha}_{k+1} + 1 – \cancel{\alpha_{k+1}} \right) \mathbf{I} = (1-\bar{\alpha}_{k+1}) \mathbf{I}
\end{align}
$$

$(10)$式と$(11)$式より$q(\mathbf{x}_{k+1}|\mathbf{x}_{0})$は下記のように表せる。
$$
\large
\begin{align}
q(\mathbf{x}_{k+1}|\mathbf{x}_{0}) = \mathcal{N} \left( \sqrt{\bar{\alpha}_{k+1}} \mathbf{x}_{0}, (1-\bar{\alpha}_{k+1}) \mathbf{I} \right) \quad (12)
\end{align}
$$

$(12)$式より、$t=k$で$(6)$式が成立する場合、$t=k+1$でも$(6)$式が成立することが確認できる。

$[1], \, [2]$より、数学的帰納法に基づいて$(6)$式が成立することが示される。

モーメント母関数を用いた正規分布の再生性の導出は下記で詳しく取り扱った。

確率分布の再生性とその導出（二項分布・ポアソン分布・正規分布）