はじめに

画像生成AIの世界でDiffusion Model（拡散モデル）は近年注目を集めています。特に、Jonathan Hoらによる論文 "Denoising Diffusion Probabilistic Models" (NeurIPS 2020) は、その後の研究に大きな影響を与え、高品質な画像生成技術のブレイクスルーとなりました。

本記事では、拡散モデル(DDPM) の論文について、冗長になりすぎないように理論的な部分を中心にかいつまんで解説します。

拡散モデルの基本的な考え方

拡散モデルのポイントは、以下の2つのプロセスです。

拡散過程 (Forward Process / Diffusion Process):
きれいな元画像に少しずつノイズを加えていき、最終的に完全なノイズ画像にする過程
逆過程 (Reverse Process / Denoising Process):
完全なノイズ画像から出発し、学習したモデルを使って少しずつノイズを除去していき、元の画像のような新しい画像を生成する過程

拡散モデルのイメージ図(論文中のFigure 2)。左から右へノイズを加えていく拡散過程、右から左へノイズを除去していく逆過程を示す。

1. 拡散過程：データを徐々に壊していく

まず、データをノイズに変換していく「拡散過程」を見ていきます。これはモデルが学習する部分ではなく、あらかじめ設計・固定されたプロセスです。

元のデータ（画像）を $x_0$ とします。これに $T$ ステップかけて徐々にノイズを加えていきます。ステップ $t$ での状態を $x_t$ とすると、$x_t$ は一つ前の状態 $x_{t-1}$ から次のように生成されます。このとき、$x_t$ は直前の $x_{t-1}$ だけに依存して決まります。これをマルコフ性と言い、このような確率過程の連鎖をマルコフ連鎖と呼びます。

$$
q(x_t | x_{t-1}) := \mathcal{N}(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_t I) \tag{2}
$$

*以下、式番号は論文のものを流用します

ここで、

$q(x_t | x_{t-1})$: $x_{t-1}$ が与えられたときの $x_t$ の条件付き確率分布。
$\mathcal{N}(x; \mu, \Sigma)$: 平均 $\mu$、共分散行列 $\Sigma$ の正規分布（ガウス分布）。
$\beta_t$: ステップ $t$ で加えるノイズの分散（ノイズの強さ）。$0 < \beta_t < 1$ で、通常はステップが進む（$t$ が大きくなる）につれて大きくなるように設定される（例: $\beta_1=0.0001$ から $\beta_T=0.02$ へ線形に増加）。この $\beta_1, ..., \beta_T$ の系列を分散スケジュールと呼びます。これは重要なハイパーパラメータです。
$I$: 単位行列。
$\sqrt{1-\beta_t}x_{t-1}$: 平均値。前の状態 $x_{t-1}$ を少しだけ縮小しています。
$\beta_t I$: 共分散行列。各次元（ピクセル）に独立に、分散 $\beta_t$ のガウスノイズが加えられます。

このプロセスを $T$ 回繰り返すと、$x_T$ はほぼ完全なノイズ（具体的には $\mathcal{N}(0, I)$ に近い分布）になります。

また、途中のステップ $t$ の状態 $x_t$ は、元の $x_0$ から直接計算することもできます。$\alpha_t := 1 - \beta_t$、$\bar{\alpha_t} := \prod_{s=1}^{t} \alpha_s$ と定義すると、

$$
q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)I) \tag{4}
$$

これは、$x_t$ が $x_0$ を $\sqrt{\bar{\alpha}_t}$ 倍に縮小したものに、分散 $(1-\bar{\alpha}_t)I$ のガウスノイズを加えたものであることを意味します。サンプリングする際は、標準正規分布からのノイズ $\epsilon \sim \mathcal{N}(0, I)$ を使って次のように計算できます。

$$
x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon
$$

(これは式(4)をサンプリング形式で書いたものです)

2. 逆拡散過程：ノイズからデータを復元（生成）する

拡散モデルの核心は、この拡散過程を逆向きにたどるプロセス、つまりノイズ除去プロセスを学習することです。

完全なノイズ $x_T \sim \mathcal{N}(0, I)$ から出発し、$x_{T-1}, x_{T-2}, ..., x_1, x_0$ とステップを逆に進めていきます。ステップ $t$ で $x_t$ から $x_{t-1}$ を生成する確率 $p_\theta(x_{t-1}|x_t)$ を学習します。$\theta$ は学習対象のパラメータ（ニューラルネットワークの重みなど）を表します。

逆過程もマルコフ連鎖としてモデル化されます。

$$
p_\theta(x_{0:T}) := p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} | x_t)
$$

そして、各ステップの遷移確率 $p_\theta(x_{t-1}|x_t)$ は、拡散過程がガウスノイズに基づいていたことを考慮し、これもガウス分布でモデル化します。

$$
p_\theta(x_{t-1} | x_t) := \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) \tag{1}
$$

ここで、モデル（ニューラルネットワーク）が学習するのは、現在の状態 $x_t$ と時刻 $t$ から、一つ前の状態 $x_{t-1}$ の分布の平均 $\mu_\theta(x_t, t)$ と共分散 $\Sigma_\theta(x_t, t)$ を予測する方法です。

ここでのポイント

拡散過程、逆拡散過程の各ステップはガウス分布を仮定します。これは、拡散過程で加えたノイズがガウス分布であり、かつノイズの量が小さい（$\beta_t$ が小さい）場合、理論的に最適な逆ステップもガウス分布に近くなるため妥当な仮定とされています。
拡散過程は固定し、学習を行いません。パラメータ（$\beta_t$ スケジュール）は事前に決められています。
逆過程 $p_\theta$ は学習によって獲得されます。モデルはどうすればノイズ $x_t$ から元の画像らしい $x_{t-1}$ を復元できるかを平均 $\mu_\theta(x_t, t)$と共分散 $\Sigma_\theta(x_t, t)$を予測することによって学習します。

どうやって学習するのか？

では、具体的に逆過程 $p_\theta$ をどう学習させるのでしょうか？　理想的には、モデルが生成するデータの尤度 $p_\theta(x_0)$ を最大化したいのですが、これは直接計算することが非常に困難です。

そこで、変分推論 (Variational Inference) という手法を用います。これは、直接計算できない尤度の代わりに、計算可能な変分下限 (Variational Lower Bound) または ELBO (Evidence Lower Bound) と呼ばれる量を最大化（あるいはその負 $L$ を最小化）する手法です。

論文では、この損失関数 $L$ を次のように定義しています。

L = \mathbb{E}_{q(x_{1:T}|x_0)}\Bigl[-log~p(x_{T})-\sum_{t\ge1}log\frac{p_{\theta}(x_{t-1}|x_{t})}{q(x_{t}|x_{t-1})}\Bigr] \tag{3}

さらに計算しやすくするため、より理解しやすい形に変形します。

L = \mathbb{E}_{q}\Bigl[\underbrace{D_{KL}(q(x_{T}|x_{0})||p(x_{T}))}_{L_T} +  \sum_{t=2}^{T}\underbrace{D_{KL}(q(x_{t-1}|x_{t},x_{0})||p_{\theta}(x_{t-1}|x_{t}))}_{L_{t-1}} - \underbrace{log~p_{\theta}(x_{0}|x_{1})}_{L_0}\Bigr] \tag{5}

このとき、

$\mathbb{E}_q[\dots]$: 拡散過程 $q$ に従って $x_0$ からノイズ系列 $(x_1, ..., x_T)$ を生成し、その系列について [...] の中の値の期待値、つまり平均を取る
$D_{KL}(A||B)$: KLダイバージェンス。分布 A が分布 B にどれだけ近いかを表す（近いほど小さい）

このように、損失関数を$L_T$、$L_{t-1}$、$L_0$の３つに分解します。

L_Tの扱い

$L_T = D_{KL}(q(x_T|x_0)||p(x_T))$ は、拡散過程の最終ノイズ状態 $x_T$ の分布 $q(x_T|x_0)$ が、事前分布 $p(x_T)$（通常 $\mathcal{N}(0, I)$）からどれだけ離れているかを表します。
この論文では、拡散過程 $q$ を定義する分散スケジュール $\beta_t$ は学習せず、固定値として扱います。
したがって、パラメータ $\theta$ を最適化する際には$L_T$ は無視することができます。

L_t-1の扱い

$L_{t-1} = D_{KL}(q(x_{t-1}|x_t, x_0)||p_\theta(x_{t-1}|x_t))$ は、学習する逆過程 $p_\theta(x_{t-1}|x_t)$ が、理想的な逆ステップ $q(x_{t-1}|x_t, x_0)$ にどれだけ近いかを測ります。ここが逆過程の学習の核となる部分になりますので、次のセクションで詳細を記載します。

L_0 の扱い

$L_0 = -\log p_\theta(x_0|x_1)$ は、逆過程の最後のステップで、元のデータ $x_0$ をどれだけうまく復元できるか（尤度）を表す項です。
元のデータ $x_0$ は通常、離散的なピクセル値（例：0-255）を持っていますが、モデルが予測する $\mu_\theta(x_1, 1)$ は連続値です。
そのため、論文では、離散データに対する正しい対数尤度を得るために、特殊なデコーダーを定義します。これは、モデルが予測した連続的なガウス分布 $\mathcal{N}(x; \mu_\theta^i(x_1, 1), \sigma_1^2)$ を、各ピクセル $i$ の離散値 $x_0^i$ が対応する微小区間（ビン）で積分することで、その離散値が得られる確率 $p_\theta(x_0^i|x_1)$ を計算します。
理論的な変分下限 $L$ を直接最適化する場合は、この積分を含む $L_0$ 項を計算します。しかし、後述する単純化された目的関数 $L_{simple}$ を使う場合は、この $L_0$ 項も、$t=1$ におけるノイズ予測誤差 $||\epsilon - \epsilon_\theta||^2$ として近似的に扱われます。

ここでのポイント

なぜ変分推論なのか？
尤度 $p_\theta(x_0)$ を直接最大化できないため、代わりの計算可能な目的関数（ELBOの負 $L$）を最小化する。
変分下限を３つに分解する
- $L_T$: 無視できる定数項。
- $L_{t-1}$ ($t=2..T$): 学習する逆過程が、理想的な逆過程にどれだけ近いかを測る。
- $L_0$: 最終的なデータ復元項。離散的な尤度を得るための特別なデコーダー。

DDPM論文の何がすごいのか？

DDPM論文の特に重要な貢献は、式(5)の損失関数の学習を高性能化・安定化させるため以下の２つの工夫を行った点です。

1. 逆過程 L_t-1 のパラメータ化

DDPMでは逆過程のモデル $p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$ を学習させやすくするために、パラメータ化を行っています。

論文では$L_{t-1}$を最小化することは、$\Sigma_\theta$ を固定 ($\sigma_t^2 I$) すれば、モデルの予測平均 $\mu_\theta$ を拡散過程の平均 $\tilde{\mu}_t = \tilde{\mu}_t(x_t, x_0)$ に近づける（二乗誤差を最小化する）こととほぼ等価であることが示されています。

しかし、DDPMではここからさらに分析を進め、ターゲット $\tilde{\mu}_t$ が、現在の $x_t$ とそれを生成するのに使われたノイズ $\epsilon$ を使って書き直せることを示しています。

$$
\tilde{\mu}_t(x_t(x_0, \epsilon), x_0) = \frac{1}{\sqrt{\alpha_t}}\Bigl(x_t(x_0, \epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\Bigr)
$$

この関係性に着目し、モデル $\mu_\theta$ で直接 $\tilde{\mu_t}$ を予測するのではなく、代わりにノイズ $\epsilon$ を予測するモデル $\epsilon_\theta(x_t, t)$ を導入し、$\mu_\theta$ を次のように計算する設計（パラメータ化）を提案しました。

\mu_{\theta}(x_{t},t) = \frac{1}{\sqrt{\alpha_{t}}}\Bigl(x_{t}-\frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}\epsilon_{\theta}(x_{t},t)\Bigr)

2. 目的関数の単純化

理論的な損失 $L$ (式(5))は、各項に複雑な重みが付いています。DDPM論文では学習を行いやすくするため、この重みを取り除いた、よりシンプルな目的関数 $L_{simple}$ を提案しています。

L_{simple}(\theta):=\mathbb{E}_{t,x_{0},\epsilon}\Bigl[||\epsilon-\epsilon_{\theta}(x_t(x_0, \epsilon),t)||^{2}\Bigr]

これは、ランダムな時刻 $t$、画像データ $x_0$、ノイズ $\epsilon$ について、真のノイズ $\epsilon$ とモデルが予測したノイズ $\epsilon_\theta$ の間の単純な平均二乗誤差（MSE） を最小化することを意味します。

ここでのポイント

モデルのタスクが「平均 $\tilde{\mu}_t$ を予測する」ことから、「ノイズ $\epsilon$ を予測する」ことに変わっています。これは逆拡散過程の本質的な意味であるデノイジング（ノイズ除去）タスクに近いと言えます。
理論的な損失関数から重みの排除をして単純化を行っています。理論的に導かれた重みを無視しても大丈夫なのか？と思われますが、実際は実装が単純になるだけでなく、経験的にサンプル品質が向上したことが報告されています。

画像をどうやって生成するのか？

学習済みのモデル $\epsilon_\theta$ を使って実際に画像を生成する手順は論文中の Algorithm 2 にまとめられています。

論文中のAlgorithm 2

具体的には以下のようになります。

$x_T$ を $\mathcal{N}(0, I)$ からサンプリング（純粋なノイズから開始）。
$t = T$ から $1$ まで、以下のステップを繰り返す。
i. $z$ を $\mathcal{N}(0, I)$ からサンプリング（ただし $t=1$ なら、$z=0$）。
ii. ニューラルネットワーク $\epsilon_\theta(x_t, t)$ を使ってノイズを予測する。
iii. 次の状態 $x_{t-1}$ を計算する：

        x_{t-1} = \frac{1}{\sqrt{\alpha_{t}}}\Bigl(x_{t}-\frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}\epsilon_{\theta}(x_{t},t)\Bigr)+\sigma_{t}z

3 出力: 最終的に得られた $x_0$ が生成された画像になる。

実際に生成された画像の例。論文中のFigure 3, Figure 4

ベンチマーク結果について

DDPMは、提案手法（特に $\epsilon$ 予測と $L_{simple}$ の組み合わせ）を用いることで、CIFAR10データセットにおいて当時の最先端となる FID=3.17 という驚異的なスコアを達成しました。

また、ここでは紹介しませんが、論文では拡散モデルの対数尤度（NLL、ロスレス圧縮性能）は最高レベルではないものの、サンプル品質が高いことから、「損失のある圧縮器としては非常に優れている」という考察や、拡散モデルが自己回帰モデルの一般化と見なせる可能性についても議論がなされています。

まとめ

拡散モデルは、データにノイズを加えていく拡散過程と、それを逆向きにたどってノイズを除去していく逆過程を学習させる手法であり、

ノイズ $\epsilon$ を予測するパラメータ化
単純化された目的関数 $L_{simple}$

といった改良を導入することで、画像生成の品質を飛躍的に向上させた、ということが分かりました。生成モデルについてはまだまだ勉強不足なので今後も継続的にアウトプットを行っていきたいと思います。

参考文献