拡散モデル（Diffusion Model）は、画像にノイズを段階的に加えていく過程（前方過程）と、そのノイズを逆に除去していく過程（逆過程）を学習する生成モデルです。2020年に Ho らが提案した DDPM（Denoising Diffusion Probabilistic Models） は、GAN に匹敵する高品質な画像生成を実現し、その後の Stable Diffusion や DALL-E 2 などの基盤技術となりました。
拡散モデルの理論的な魅力は、前方過程がガウスノイズの逐次的な付加という単純な操作でありながら、逆過程の学習が変分推論とデノイジングスコアマッチングの両面から美しく定式化される点にあります。本記事では、前方過程の定式化から変分下界（ELBO）の導出、そして簡略化された損失関数 $L_{\text{simple}}$ の導出までを省略なく行い、Python で2D分布上の拡散モデルを実装します。
本記事の内容
- 拡散モデルの直感的な理解（ノイズの付加と除去）
- 前方過程 $q(\bm{x}_t | \bm{x}_{t-1})$ の定義
- 任意時刻の分布 $q(\bm{x}_t | \bm{x}_0)$ の閉形式導出
- 逆過程 $p_\theta(\bm{x}_{t-1} | \bm{x}_t)$ の定式化
- 変分下界（ELBO）の完全な導出
- 簡略化された損失関数 $L_{\text{simple}} = \mathbb{E}[\|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\|^2]$ の導出
- ノイズスケジュールとサンプリングアルゴリズム
- Python で2D分布での拡散モデルの実装と可視化
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
拡散モデルの直感
拡散モデルの基本アイデアは驚くほど単純です。次の2つのステップからなります。
前方過程（Forward Process / Diffusion Process）: データ $\bm{x}_0$ に少しずつガウスノイズを加えていき、最終的に純粋なガウスノイズ $\bm{x}_T \sim \mathcal{N}(\bm{0}, \bm{I})$ にします。これは $T$ ステップのマルコフ連鎖として定式化されます。
逆過程（Reverse Process / Denoising Process）: 純粋なガウスノイズ $\bm{x}_T$ から出発し、段階的にノイズを除去していくことで、元のデータ分布に従うサンプル $\bm{x}_0$ を生成します。この逆過程をニューラルネットワークで学習します。
イメージとしては、インクを水に一滴垂らすと徐々に拡散して均一になる過程が前方過程、そしてその拡散を時間反転させてインクの一滴を復元する過程が逆過程です。
前方過程の定式化
1ステップの前方過程
前方過程は、各タイムステップ $t = 1, 2, \dots, T$ で以下のガウス遷移核を適用するマルコフ連鎖です。
$$ q(\bm{x}_t | \bm{x}_{t-1}) = \mathcal{N}(\bm{x}_t; \sqrt{1 – \beta_t}\,\bm{x}_{t-1},\; \beta_t \bm{I}) $$
ここで $\beta_t \in (0, 1)$ はノイズスケジュールと呼ばれるハイパーパラメータで、各ステップで加えるノイズの量を制御します。典型的には $\beta_1 = 10^{-4}$ から $\beta_T = 0.02$ まで線形に増加させます。
再パラメータ化トリック（reparameterization trick）を用いると、この分布からのサンプリングは次のように書けます。
$$ \bm{x}_t = \sqrt{1 – \beta_t}\,\bm{x}_{t-1} + \sqrt{\beta_t}\,\bm{\epsilon}_t, \quad \bm{\epsilon}_t \sim \mathcal{N}(\bm{0}, \bm{I}) $$
この式の直感は明快です。$\bm{x}_{t-1}$ を $\sqrt{1-\beta_t}$ 倍に縮小し（信号を弱め）、$\sqrt{\beta_t}$ の標準偏差を持つガウスノイズを加えています。
$T$ ステップの同時分布
前方過程はマルコフ連鎖ですので、$T$ ステップの同時分布は
$$ q(\bm{x}_{1:T} | \bm{x}_0) = \prod_{t=1}^{T} q(\bm{x}_t | \bm{x}_{t-1}) $$
と分解されます。
任意時刻の分布 $q(\bm{x}_t | \bm{x}_0)$ の導出
記号の準備
導出を簡潔にするため、以下の記号を導入します。
$$ \alpha_t = 1 – \beta_t, \quad \bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s $$
$\alpha_t$ は各ステップでの「信号保持率」、$\bar{\alpha}_t$ はステップ $0$ から $t$ までの累積的な信号保持率を表します。
再パラメータ化の繰り返し
ここで重要な性質を導出します。$\bm{x}_t$ を $\bm{x}_0$ と独立なノイズ $\bm{\epsilon}$ のみで直接表すことができるのです。
ステップ $t=1$:
$$ \bm{x}_1 = \sqrt{\alpha_1}\,\bm{x}_0 + \sqrt{1 – \alpha_1}\,\bm{\epsilon}_1, \quad \bm{\epsilon}_1 \sim \mathcal{N}(\bm{0}, \bm{I}) $$
ステップ $t=2$:
$$ \begin{align} \bm{x}_2 &= \sqrt{\alpha_2}\,\bm{x}_1 + \sqrt{1 – \alpha_2}\,\bm{\epsilon}_2 \\ &= \sqrt{\alpha_2}\left(\sqrt{\alpha_1}\,\bm{x}_0 + \sqrt{1 – \alpha_1}\,\bm{\epsilon}_1\right) + \sqrt{1 – \alpha_2}\,\bm{\epsilon}_2 \\ &= \sqrt{\alpha_2 \alpha_1}\,\bm{x}_0 + \sqrt{\alpha_2(1 – \alpha_1)}\,\bm{\epsilon}_1 + \sqrt{1 – \alpha_2}\,\bm{\epsilon}_2 \end{align} $$
ここで、$\bm{\epsilon}_1$ と $\bm{\epsilon}_2$ は独立な標準正規分布に従うので、2つのガウスノイズ項の和は次のガウス分布に従います。
$$ \sqrt{\alpha_2(1 – \alpha_1)}\,\bm{\epsilon}_1 + \sqrt{1 – \alpha_2}\,\bm{\epsilon}_2 \sim \mathcal{N}\left(\bm{0},\; \left[\alpha_2(1-\alpha_1) + (1-\alpha_2)\right]\bm{I}\right) $$
分散を計算します。
$$ \alpha_2(1-\alpha_1) + (1-\alpha_2) = \alpha_2 – \alpha_2\alpha_1 + 1 – \alpha_2 = 1 – \alpha_1\alpha_2 = 1 – \bar{\alpha}_2 $$
したがって
$$ \bm{x}_2 = \sqrt{\bar{\alpha}_2}\,\bm{x}_0 + \sqrt{1 – \bar{\alpha}_2}\,\bm{\epsilon}, \quad \bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I}) $$
一般のステップ $t$（帰納法）:
帰納法の仮定として、$\bm{x}_{t-1} = \sqrt{\bar{\alpha}_{t-1}}\,\bm{x}_0 + \sqrt{1 – \bar{\alpha}_{t-1}}\,\bar{\bm{\epsilon}}_{t-1}$ が成り立つとします。$\bar{\bm{\epsilon}}_{t-1} \sim \mathcal{N}(\bm{0}, \bm{I})$ です。
$$ \begin{align} \bm{x}_t &= \sqrt{\alpha_t}\,\bm{x}_{t-1} + \sqrt{1 – \alpha_t}\,\bm{\epsilon}_t \\ &= \sqrt{\alpha_t}\left(\sqrt{\bar{\alpha}_{t-1}}\,\bm{x}_0 + \sqrt{1 – \bar{\alpha}_{t-1}}\,\bar{\bm{\epsilon}}_{t-1}\right) + \sqrt{1 – \alpha_t}\,\bm{\epsilon}_t \\ &= \sqrt{\alpha_t \bar{\alpha}_{t-1}}\,\bm{x}_0 + \sqrt{\alpha_t(1 – \bar{\alpha}_{t-1})}\,\bar{\bm{\epsilon}}_{t-1} + \sqrt{1 – \alpha_t}\,\bm{\epsilon}_t \end{align} $$
後ろの2つのノイズ項は独立なので、分散を合算します。
$$ \alpha_t(1 – \bar{\alpha}_{t-1}) + (1 – \alpha_t) = \alpha_t – \alpha_t\bar{\alpha}_{t-1} + 1 – \alpha_t = 1 – \alpha_t\bar{\alpha}_{t-1} = 1 – \bar{\alpha}_t $$
よって、次の閉形式を得ます。
$$ \boxed{q(\bm{x}_t | \bm{x}_0) = \mathcal{N}(\bm{x}_t;\; \sqrt{\bar{\alpha}_t}\,\bm{x}_0,\; (1 – \bar{\alpha}_t)\bm{I})} $$
すなわち
$$ \bm{x}_t = \sqrt{\bar{\alpha}_t}\,\bm{x}_0 + \sqrt{1 – \bar{\alpha}_t}\,\bm{\epsilon}, \quad \bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I}) $$
この結果は非常に重要です。$T$ ステップの逐次的なノイズ付加を経ることなく、任意の時刻 $t$ のノイズ付きデータ $\bm{x}_t$ を $\bm{x}_0$ から1ステップで直接サンプリングできるのです。これにより学習時の計算が大幅に効率化されます。
逆過程の定式化
逆過程の必要性
生成を行うには、前方過程の逆、すなわち $q(\bm{x}_{t-1} | \bm{x}_t)$ を計算する必要があります。しかし、$q(\bm{x}_{t-1} | \bm{x}_t)$ は $q(\bm{x}_{t-1} | \bm{x}_t) = q(\bm{x}_t | \bm{x}_{t-1}) q(\bm{x}_{t-1}) / q(\bm{x}_t)$ とベイズの定理で書けますが、周辺分布 $q(\bm{x}_t)$ の計算は困難（全データの積分が必要）です。
そこで、逆過程をパラメータ $\theta$ を持つニューラルネットワークで近似します。
$$ p_\theta(\bm{x}_{t-1} | \bm{x}_t) = \mathcal{N}(\bm{x}_{t-1};\; \bm{\mu}_\theta(\bm{x}_t, t),\; \sigma_t^2 \bm{I}) $$
ここで $\bm{\mu}_\theta(\bm{x}_t, t)$ はニューラルネットワークが予測する平均、$\sigma_t^2$ は後述するように既知の値に固定します。
逆過程の同時分布は
$$ p_\theta(\bm{x}_{0:T}) = p(\bm{x}_T) \prod_{t=1}^{T} p_\theta(\bm{x}_{t-1} | \bm{x}_t) $$
で、$p(\bm{x}_T) = \mathcal{N}(\bm{x}_T; \bm{0}, \bm{I})$ です。
真の逆過程（$\bm{x}_0$ が既知の場合）
$\bm{x}_0$ が条件として与えられた場合の逆遷移確率 $q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0)$ は解析的に計算できます。ベイズの定理を適用すると
$$ q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0) = \frac{q(\bm{x}_t | \bm{x}_{t-1}, \bm{x}_0)\, q(\bm{x}_{t-1} | \bm{x}_0)}{q(\bm{x}_t | \bm{x}_0)} $$
マルコフ性より $q(\bm{x}_t | \bm{x}_{t-1}, \bm{x}_0) = q(\bm{x}_t | \bm{x}_{t-1})$ です。各項はガウス分布として既知です。
$$ q(\bm{x}_t | \bm{x}_{t-1}) = \mathcal{N}(\bm{x}_t;\; \sqrt{\alpha_t}\,\bm{x}_{t-1},\; (1-\alpha_t)\bm{I}) $$
$$ q(\bm{x}_{t-1} | \bm{x}_0) = \mathcal{N}(\bm{x}_{t-1};\; \sqrt{\bar{\alpha}_{t-1}}\,\bm{x}_0,\; (1-\bar{\alpha}_{t-1})\bm{I}) $$
$$ q(\bm{x}_t | \bm{x}_0) = \mathcal{N}(\bm{x}_t;\; \sqrt{\bar{\alpha}_t}\,\bm{x}_0,\; (1-\bar{\alpha}_t)\bm{I}) $$
ガウス分布の指数部を展開して $\bm{x}_{t-1}$ について平方完成を行います。対数を取ると
$$ \begin{align} &\log q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0) \\ &\propto -\frac{1}{2}\left[\frac{(\bm{x}_t – \sqrt{\alpha_t}\,\bm{x}_{t-1})^2}{1-\alpha_t} + \frac{(\bm{x}_{t-1} – \sqrt{\bar{\alpha}_{t-1}}\,\bm{x}_0)^2}{1-\bar{\alpha}_{t-1}}\right] \\ &= -\frac{1}{2}\left[\frac{\alpha_t}{1-\alpha_t}\bm{x}_{t-1}^2 – \frac{2\sqrt{\alpha_t}\,\bm{x}_t}{1-\alpha_t}\bm{x}_{t-1} + \frac{1}{1-\bar{\alpha}_{t-1}}\bm{x}_{t-1}^2 – \frac{2\sqrt{\bar{\alpha}_{t-1}}\,\bm{x}_0}{1-\bar{\alpha}_{t-1}}\bm{x}_{t-1}\right] + C \end{align} $$
ここで $C$ は $\bm{x}_{t-1}$ に依存しない定数です。$\bm{x}_{t-1}^2$ の係数と $\bm{x}_{t-1}$ の係数を集めます。
$\bm{x}_{t-1}^2$ の係数:
$$ \frac{\alpha_t}{1-\alpha_t} + \frac{1}{1-\bar{\alpha}_{t-1}} = \frac{\alpha_t(1-\bar{\alpha}_{t-1}) + (1-\alpha_t)}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})} $$
分子を計算します。
$$ \alpha_t – \alpha_t\bar{\alpha}_{t-1} + 1 – \alpha_t = 1 – \alpha_t\bar{\alpha}_{t-1} = 1 – \bar{\alpha}_t $$
したがって $\bm{x}_{t-1}^2$ の係数は
$$ \frac{1 – \bar{\alpha}_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})} $$
これはガウス分布 $\mathcal{N}(\tilde{\bm{\mu}}_t, \tilde{\beta}_t \bm{I})$ の $1/(2\tilde{\beta}_t)$ に対応するので、後方分散は
$$ \boxed{\tilde{\beta}_t = \frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1 – \bar{\alpha}_t} = \frac{\beta_t(1-\bar{\alpha}_{t-1})}{1 – \bar{\alpha}_t}} $$
次に $\bm{x}_{t-1}$ の係数から後方平均を求めます。
$$ \tilde{\bm{\mu}}_t(\bm{x}_t, \bm{x}_0) = \tilde{\beta}_t \left(\frac{\sqrt{\alpha_t}}{1-\alpha_t}\bm{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}}\bm{x}_0\right) $$
$\tilde{\beta}_t$ を代入して整理すると
$$ \boxed{\tilde{\bm{\mu}}_t(\bm{x}_t, \bm{x}_0) = \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\bm{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1-\bar{\alpha}_t}\bm{x}_0} $$
つまり、$q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0) = \mathcal{N}(\bm{x}_{t-1}; \tilde{\bm{\mu}}_t, \tilde{\beta}_t \bm{I})$ はガウス分布です。
変分下界（ELBO）の導出
目的: 対数尤度の最大化
生成モデルの学習目標は、データの対数尤度 $\log p_\theta(\bm{x}_0)$ を最大化することです。直接計算は困難なので、変分下界（Evidence Lower Bound, ELBO）を導出します。
$$ \begin{align} \log p_\theta(\bm{x}_0) &= \log \int p_\theta(\bm{x}_{0:T})\,d\bm{x}_{1:T} \\ &= \log \int \frac{p_\theta(\bm{x}_{0:T})}{q(\bm{x}_{1:T}|\bm{x}_0)} q(\bm{x}_{1:T}|\bm{x}_0)\,d\bm{x}_{1:T} \\ &= \log \mathbb{E}_{q(\bm{x}_{1:T}|\bm{x}_0)}\left[\frac{p_\theta(\bm{x}_{0:T})}{q(\bm{x}_{1:T}|\bm{x}_0)}\right] \\ &\geq \mathbb{E}_{q(\bm{x}_{1:T}|\bm{x}_0)}\left[\log \frac{p_\theta(\bm{x}_{0:T})}{q(\bm{x}_{1:T}|\bm{x}_0)}\right] \quad (\text{イェンセンの不等式}) \end{align} $$
この下界を $L_{\text{VLB}}$ とおきます。
$$ L_{\text{VLB}} = \mathbb{E}_q\left[\log \frac{p_\theta(\bm{x}_{0:T})}{q(\bm{x}_{1:T}|\bm{x}_0)}\right] $$
ELBOの分解
同時分布を展開します。
$$ p_\theta(\bm{x}_{0:T}) = p(\bm{x}_T) \prod_{t=1}^{T} p_\theta(\bm{x}_{t-1}|\bm{x}_t) $$
$$ q(\bm{x}_{1:T}|\bm{x}_0) = \prod_{t=1}^{T} q(\bm{x}_t|\bm{x}_{t-1}) $$
比を対数に入れます。
$$ \begin{align} L_{\text{VLB}} &= \mathbb{E}_q\left[\log \frac{p(\bm{x}_T) \prod_{t=1}^{T} p_\theta(\bm{x}_{t-1}|\bm{x}_t)}{\prod_{t=1}^{T} q(\bm{x}_t|\bm{x}_{t-1})}\right] \\ &= \mathbb{E}_q\left[\log p(\bm{x}_T) + \sum_{t=1}^{T} \log \frac{p_\theta(\bm{x}_{t-1}|\bm{x}_t)}{q(\bm{x}_t|\bm{x}_{t-1})}\right] \end{align} $$
ここで、$t \geq 2$ の項についてベイズの定理を用いて $q(\bm{x}_t|\bm{x}_{t-1})$ を書き換えます。
$$ q(\bm{x}_t|\bm{x}_{t-1}) = \frac{q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0)\, q(\bm{x}_t|\bm{x}_0)}{q(\bm{x}_{t-1}|\bm{x}_0)} $$
これを代入すると
$$ \begin{align} L_{\text{VLB}} &= \mathbb{E}_q\left[\log p(\bm{x}_T) + \sum_{t=2}^{T} \log \frac{p_\theta(\bm{x}_{t-1}|\bm{x}_t)}{q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0)} \cdot \frac{q(\bm{x}_{t-1}|\bm{x}_0)}{q(\bm{x}_t|\bm{x}_0)} + \log \frac{p_\theta(\bm{x}_0|\bm{x}_1)}{q(\bm{x}_1|\bm{x}_0)}\right] \end{align} $$
テレスコーピング和（$q(\bm{x}_{t-1}|\bm{x}_0)/q(\bm{x}_t|\bm{x}_0)$ の積が $q(\bm{x}_1|\bm{x}_0)/q(\bm{x}_T|\bm{x}_0)$ に縮約する）を利用すると
$$ \begin{align} L_{\text{VLB}} = \mathbb{E}_q\bigg[&\underbrace{\log \frac{p(\bm{x}_T)}{q(\bm{x}_T|\bm{x}_0)}}_{L_T} + \sum_{t=2}^{T} \underbrace{\log \frac{p_\theta(\bm{x}_{t-1}|\bm{x}_t)}{q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0)}}_{L_{t-1}} + \underbrace{\log p_\theta(\bm{x}_0|\bm{x}_1)}_{L_0}\bigg] \end{align} $$
各項の意味を整理します。
- $L_T = -D_{\text{KL}}(q(\bm{x}_T|\bm{x}_0) \| p(\bm{x}_T))$: 前方過程の終端と事前分布のKLダイバージェンス。$\beta_t$ が固定されているため学習パラメータを含まない定数項です。
- $L_{t-1} = -D_{\text{KL}}(q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) \| p_\theta(\bm{x}_{t-1}|\bm{x}_t))$ （$t = 2, \dots, T$）: 真の逆過程とモデルの逆過程のKLダイバージェンス。
- $L_0 = \log p_\theta(\bm{x}_0|\bm{x}_1)$: 再構成項。
$L_{t-1}$ の計算
2つのガウス分布間のKLダイバージェンスを計算します。$q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) = \mathcal{N}(\tilde{\bm{\mu}}_t, \tilde{\beta}_t \bm{I})$ と $p_\theta(\bm{x}_{t-1}|\bm{x}_t) = \mathcal{N}(\bm{\mu}_\theta, \sigma_t^2 \bm{I})$ について、$\sigma_t^2 = \tilde{\beta}_t$ と設定すると
$$ D_{\text{KL}}(q \| p_\theta) = \frac{1}{2\tilde{\beta}_t} \|\tilde{\bm{\mu}}_t(\bm{x}_t, \bm{x}_0) – \bm{\mu}_\theta(\bm{x}_t, t)\|^2 $$
簡略化損失関数 $L_{\text{simple}}$ の導出
$\bm{x}_0$ による平均のパラメータ化
$\bm{x}_t = \sqrt{\bar{\alpha}_t}\,\bm{x}_0 + \sqrt{1-\bar{\alpha}_t}\,\bm{\epsilon}$ より、$\bm{x}_0$ を $\bm{x}_t$ と $\bm{\epsilon}$ で表すと
$$ \bm{x}_0 = \frac{\bm{x}_t – \sqrt{1-\bar{\alpha}_t}\,\bm{\epsilon}}{\sqrt{\bar{\alpha}_t}} $$
これを後方平均 $\tilde{\bm{\mu}}_t$ に代入します。
$$ \begin{align} \tilde{\bm{\mu}}_t &= \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\bm{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1-\bar{\alpha}_t}\bm{x}_0 \\ &= \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\bm{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1-\bar{\alpha}_t} \cdot \frac{\bm{x}_t – \sqrt{1-\bar{\alpha}_t}\,\bm{\epsilon}}{\sqrt{\bar{\alpha}_t}} \end{align} $$
$\bm{x}_t$ の係数を整理します。$\bar{\alpha}_t = \alpha_t \bar{\alpha}_{t-1}$ より $\sqrt{\bar{\alpha}_{t-1}}/\sqrt{\bar{\alpha}_t} = 1/\sqrt{\alpha_t}$ なので
$$ \begin{align} \tilde{\bm{\mu}}_t &= \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\bm{x}_t + \frac{\beta_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}\bm{x}_t – \frac{\beta_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}} \cdot \sqrt{1-\bar{\alpha}_t}\,\bm{\epsilon} \\ &= \frac{\alpha_t(1-\bar{\alpha}_{t-1}) + \beta_t}{(1-\bar{\alpha}_t)\sqrt{\alpha_t}}\bm{x}_t – \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}\,\sqrt{\alpha_t}}\bm{\epsilon} \end{align} $$
$\bm{x}_t$ の係数の分子を計算します。$\beta_t = 1 – \alpha_t$ なので
$$ \alpha_t(1-\bar{\alpha}_{t-1}) + (1-\alpha_t) = \alpha_t – \bar{\alpha}_t + 1 – \alpha_t = 1 – \bar{\alpha}_t $$
よって
$$ \boxed{\tilde{\bm{\mu}}_t(\bm{x}_t, \bm{\epsilon}) = \frac{1}{\sqrt{\alpha_t}}\left(\bm{x}_t – \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\bm{\epsilon}\right)} $$
ノイズ予測によるパラメータ化
Ho et al. (2020) は、ニューラルネットワーク $\bm{\epsilon}_\theta(\bm{x}_t, t)$ でノイズ $\bm{\epsilon}$ を予測する形でモデルの平均をパラメータ化しました。
$$ \bm{\mu}_\theta(\bm{x}_t, t) = \frac{1}{\sqrt{\alpha_t}}\left(\bm{x}_t – \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\bm{\epsilon}_\theta(\bm{x}_t, t)\right) $$
この形を $L_{t-1}$ に代入すると
$$ \begin{align} L_{t-1} &= \frac{1}{2\tilde{\beta}_t}\|\tilde{\bm{\mu}}_t – \bm{\mu}_\theta\|^2 \\ &= \frac{1}{2\tilde{\beta}_t} \left\|\frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1-\bar{\alpha}_t}}\right\|^2 \|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\|^2 \\ &= \frac{\beta_t^2}{2\tilde{\beta}_t \alpha_t (1-\bar{\alpha}_t)} \|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\|^2 \end{align} $$
Ho et al. は実験的に、時刻依存の重み係数を取り除いた簡略化損失関数がよりよく機能することを報告しました。
$$ \boxed{L_{\text{simple}} = \mathbb{E}_{t, \bm{x}_0, \bm{\epsilon}}\left[\|\bm{\epsilon} – \bm{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\,\bm{x}_0 + \sqrt{1-\bar{\alpha}_t}\,\bm{\epsilon},\; t)\|^2\right]} $$
ここで $t \sim \text{Uniform}\{1, \dots, T\}$、$\bm{x}_0 \sim q(\bm{x}_0)$、$\bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I})$ です。
この損失関数は驚くほど単純です。「ノイズ付きデータからノイズを予測し、真のノイズとのMSEを最小化する」という明快な目標になっています。
ノイズスケジュール
ノイズスケジュール $\{\beta_t\}_{t=1}^T$ は拡散モデルの品質に大きく影響します。
線形スケジュール: Ho et al. (2020) が使用した最も基本的なスケジュールです。
$$ \beta_t = \beta_{\min} + \frac{t-1}{T-1}(\beta_{\max} – \beta_{\min}) $$
典型的に $\beta_{\min} = 10^{-4}$, $\beta_{\max} = 0.02$, $T = 1000$ です。
コサインスケジュール: Nichol & Dhariwal (2021) が提案した改良版で、$\bar{\alpha}_t$ がよりなだらかに減少します。
$$ \bar{\alpha}_t = \frac{f(t)}{f(0)}, \quad f(t) = \cos\left(\frac{t/T + s}{1+s} \cdot \frac{\pi}{2}\right)^2 $$
ここで $s = 0.008$ はオフセットパラメータです。
サンプリングアルゴリズム
学習済みモデルからのサンプリングは以下のアルゴリズムで行います。
アルゴリズム: DDPMサンプリング
- $\bm{x}_T \sim \mathcal{N}(\bm{0}, \bm{I})$ をサンプリング
- $t = T, T-1, \dots, 1$ について: – $\bm{z} \sim \mathcal{N}(\bm{0}, \bm{I})$ （$t > 1$ のとき）、$\bm{z} = \bm{0}$ （$t = 1$ のとき） – $\bm{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left(\bm{x}_t – \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\bm{\epsilon}_\theta(\bm{x}_t, t)\right) + \sqrt{\tilde{\beta}_t}\,\bm{z}$
- $\bm{x}_0$ を返す
Pythonでの実装
ここでは、2D分布（2つのガウス混合分布）を対象に、拡散モデルのフルパイプラインを実装します。
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
# --- データ生成 ---
np.random.seed(42)
def generate_data(n_samples=2000):
"""2つの月形分布（make_moons）からデータを生成"""
X, _ = make_moons(n_samples=n_samples, noise=0.05)
return X.astype(np.float64)
data = generate_data()
# --- ノイズスケジュール ---
T = 100 # 2Dデータなのでステップ数は少なめ
beta_min, beta_max = 1e-4, 0.02
betas = np.linspace(beta_min, beta_max, T)
alphas = 1.0 - betas
alpha_bars = np.cumprod(alphas)
# --- 前方過程の可視化 ---
fig, axes = plt.subplots(1, 5, figsize=(20, 4))
timesteps_to_show = [0, 10, 30, 60, 99]
for i, t in enumerate(timesteps_to_show):
if t == 0:
x_t = data
else:
# q(x_t | x_0) から直接サンプリング
eps = np.random.randn(*data.shape)
x_t = np.sqrt(alpha_bars[t]) * data + np.sqrt(1 - alpha_bars[t]) * eps
axes[i].scatter(x_t[:, 0], x_t[:, 1], s=1, alpha=0.5)
axes[i].set_title(f't = {t}, alpha_bar = {alpha_bars[t]:.3f}')
axes[i].set_xlim(-3, 3)
axes[i].set_ylim(-3, 3)
axes[i].set_aspect('equal')
plt.suptitle('Forward Process: Gradual Noise Addition', fontsize=14)
plt.tight_layout()
plt.show()
次に、簡易的なニューラルネットワーク（MLP）でノイズ予測モデルを学習します。
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
np.random.seed(42)
# --- データ・スケジュール再定義 ---
def generate_data(n=3000):
X, _ = make_moons(n_samples=n, noise=0.05)
return X.astype(np.float64)
T = 100
beta_min, beta_max = 1e-4, 0.02
betas = np.linspace(beta_min, beta_max, T)
alphas = 1.0 - betas
alpha_bars = np.cumprod(alphas)
# --- 簡易MLPの定義（NumPyのみ） ---
class SimpleMLP:
"""2D拡散モデル用の3層MLP（ノイズ予測ネットワーク）"""
def __init__(self, hidden=128, lr=1e-3):
# 入力: x(2次元) + t_emb(1次元) = 3次元
self.W1 = np.random.randn(3, hidden) * 0.1
self.b1 = np.zeros(hidden)
self.W2 = np.random.randn(hidden, hidden) * 0.1
self.b2 = np.zeros(hidden)
self.W3 = np.random.randn(hidden, 2) * 0.1
self.b3 = np.zeros(2)
self.lr = lr
# Adam用
self.params = [self.W1, self.b1, self.W2, self.b2, self.W3, self.b3]
self.m = [np.zeros_like(p) for p in self.params]
self.v = [np.zeros_like(p) for p in self.params]
self.step = 0
def forward(self, x, t_emb):
"""順伝播"""
inp = np.concatenate([x, t_emb.reshape(-1, 1)], axis=1) # (N, 3)
self.inp = inp
self.h1_pre = inp @ self.W1 + self.b1
self.h1 = np.maximum(0, self.h1_pre) # ReLU
self.h2_pre = self.h1 @ self.W2 + self.b2
self.h2 = np.maximum(0, self.h2_pre) # ReLU
out = self.h2 @ self.W3 + self.b3 # (N, 2)
return out
def backward(self, eps_true, eps_pred):
"""逆伝播（MSE損失）"""
N = eps_true.shape[0]
# dL/d(eps_pred) = 2/N * (eps_pred - eps_true)
d_out = 2.0 / N * (eps_pred - eps_true) # (N, 2)
dW3 = self.h2.T @ d_out
db3 = d_out.sum(axis=0)
d_h2 = d_out @ self.W3.T
d_h2 *= (self.h2_pre > 0) # ReLU逆伝播
dW2 = self.h1.T @ d_h2
db2 = d_h2.sum(axis=0)
d_h1 = d_h2 @ self.W2.T
d_h1 *= (self.h1_pre > 0)
dW1 = self.inp.T @ d_h1
db1 = d_h1.sum(axis=0)
grads = [dW1, db1, dW2, db2, dW3, db3]
# Adam更新
self.step += 1
beta1, beta2, eps_adam = 0.9, 0.999, 1e-8
for i, g in enumerate(grads):
self.m[i] = beta1 * self.m[i] + (1 - beta1) * g
self.v[i] = beta2 * self.v[i] + (1 - beta2) * g**2
m_hat = self.m[i] / (1 - beta1**self.step)
v_hat = self.v[i] / (1 - beta2**self.step)
self.params[i] -= self.lr * m_hat / (np.sqrt(v_hat) + eps_adam)
self.W1, self.b1, self.W2, self.b2, self.W3, self.b3 = self.params
# --- 学習ループ ---
data = generate_data(3000)
model = SimpleMLP(hidden=128, lr=1e-3)
n_epochs = 300
batch_size = 256
losses = []
for epoch in range(n_epochs):
# ミニバッチサンプリング
idx = np.random.choice(len(data), batch_size)
x0 = data[idx]
# ランダムなタイムステップ
t = np.random.randint(0, T, size=batch_size)
# ノイズ付加: x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * eps
eps = np.random.randn(*x0.shape)
ab = alpha_bars[t]
x_t = np.sqrt(ab).reshape(-1, 1) * x0 + np.sqrt(1 - ab).reshape(-1, 1) * eps
# 時刻埋め込み（正規化）
t_emb = t.astype(np.float64) / T
# ノイズ予測
eps_pred = model.forward(x_t, t_emb)
loss = np.mean((eps - eps_pred)**2)
losses.append(loss)
model.backward(eps, eps_pred)
# --- サンプリング（逆過程） ---
def sample_ddpm(model, n_samples=1000):
"""DDPMサンプリングアルゴリズム"""
x = np.random.randn(n_samples, 2) # x_T ~ N(0, I)
for t in reversed(range(T)):
t_emb = np.full(n_samples, t / T)
eps_pred = model.forward(x, t_emb)
# x_{t-1} の計算
coef1 = 1.0 / np.sqrt(alphas[t])
coef2 = betas[t] / np.sqrt(1.0 - alpha_bars[t])
mean = coef1 * (x - coef2 * eps_pred)
if t > 0:
# 後方分散
beta_tilde = betas[t] * (1 - alpha_bars[t-1]) / (1 - alpha_bars[t])
z = np.random.randn(n_samples, 2)
x = mean + np.sqrt(beta_tilde) * z
else:
x = mean
return x
samples = sample_ddpm(model, n_samples=2000)
# --- 可視化 ---
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
# 学習損失
axes[0].plot(losses)
axes[0].set_xlabel('Iteration')
axes[0].set_ylabel('MSE Loss')
axes[0].set_title('Training Loss')
axes[0].set_yscale('log')
# 元データ
axes[1].scatter(data[:, 0], data[:, 1], s=1, alpha=0.5, c='blue')
axes[1].set_title('Original Data')
axes[1].set_xlim(-3, 3)
axes[1].set_ylim(-3, 3)
axes[1].set_aspect('equal')
# 生成サンプル
axes[2].scatter(samples[:, 0], samples[:, 1], s=1, alpha=0.5, c='red')
axes[2].set_title('Generated Samples (DDPM)')
axes[2].set_xlim(-3, 3)
axes[2].set_ylim(-3, 3)
axes[2].set_aspect('equal')
plt.tight_layout()
plt.show()
上のコードでは、前方過程で2Dデータに段階的にガウスノイズを加え、3層MLPでノイズを予測するモデルを学習し、逆過程で新しいサンプルを生成しています。make_moons の月形分布を復元できていることが確認できます。
学習アルゴリズムのまとめ
アルゴリズム: DDPM学習
- データ $\bm{x}_0 \sim q(\bm{x}_0)$ をサンプリング
- タイムステップ $t \sim \text{Uniform}\{1, \dots, T\}$ をサンプリング
- ノイズ $\bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I})$ をサンプリング
- ノイズ付きデータ $\bm{x}_t = \sqrt{\bar{\alpha}_t}\,\bm{x}_0 + \sqrt{1-\bar{\alpha}_t}\,\bm{\epsilon}$ を計算
- 損失 $\|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\|^2$ を最小化
これを繰り返すことでモデルが学習されます。
まとめ
本記事では、拡散モデル（DDPM）の理論を導出しました。
- 前方過程 $q(\bm{x}_t|\bm{x}_{t-1}) = \mathcal{N}(\sqrt{1-\beta_t}\,\bm{x}_{t-1}, \beta_t\bm{I})$ は段階的なノイズ付加
- 任意時刻の分布 $q(\bm{x}_t|\bm{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t}\,\bm{x}_0, (1-\bar{\alpha}_t)\bm{I})$ を閉形式で導出
- 変分下界（ELBO）をKLダイバージェンスの和に分解
- 後方平均のノイズ予測パラメータ化により、簡略化損失関数 $L_{\text{simple}} = \mathbb{E}[\|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\|^2]$ を導出
- DDPMのサンプリングアルゴリズムは逆過程を $T$ ステップ逐次実行
次のステップとして、以下の記事も参考にしてください。