拡散モデル(Diffusion Model)は、近年の画像生成 AI の中核をなす生成モデルです。Stable Diffusion や DALL·E 2 など、社会的にも大きなインパクトを与えたモデルの基盤技術であり、GAN や VAE に続く第三の生成モデルとして注目されています。
拡散モデルのアイデアは非常に直感的です。データにノイズを少しずつ加えて完全なガウスノイズに変換する前向き過程と、そのノイズから元のデータを少しずつ復元する逆過程の 2 つのプロセスから構成されます。本記事では、Ho et al. (2020) が提案した DDPM(Denoising Diffusion Probabilistic Models) の理論を数式レベルで丁寧に導出し、Python での実装まで行います。
本記事の内容
- 前向き過程と逆過程の数学的定義
- 任意の時刻 $t$ で直接サンプリングできる公式の導出
- 変分下界(ELBO)から簡略化された損失関数 $L_{\text{simple}}$ の導出
- ノイズ予測ネットワーク $\epsilon_\theta$ とサンプリングアルゴリズム
- Python での 1 次元ガウス混合分布への適用と可視化
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
拡散モデルとは
拡散モデルは、データ分布からのサンプリングを目的とする生成モデルです。GAN のように敵対的に訓練する必要がなく、VAE のようにエンコーダ・デコーダの構造を明示的に設計する必要もありません。その代わりに、「ノイズを加える過程」と「ノイズを除去する過程」という物理的に直感的な枠組みを利用します。
大まかな流れは以下のとおりです。
- 前向き過程(Forward Process): データ $\bm{x}_0$ に対して、$T$ ステップかけて少しずつガウスノイズを加え、最終的に純粋なガウスノイズ $\bm{x}_T \sim \mathcal{N}(\bm{0}, \bm{I})$ に変換します。
- 逆過程(Reverse Process): ガウスノイズ $\bm{x}_T$ から出発し、学習済みのニューラルネットワークを使って 1 ステップずつノイズを除去していき、最終的にデータ分布に従うサンプル $\bm{x}_0$ を生成します。
イメージとしては、インクを水に垂らすと徐々に拡散していく過程(前向き過程)を記録しておき、その逆再生(逆過程)をニューラルネットワークに学習させるようなものです。
前向き過程(Forward Process)
定義
前向き過程は、データ $\bm{x}_0 \sim q(\bm{x}_0)$ に対して、以下のマルコフ連鎖として定義されます。
$$ 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}_{t-1}) = \mathcal{N}(\bm{x}_t ; \sqrt{1 – \beta_t}\, \bm{x}_{t-1},\, \beta_t \bm{I}) $$
ここで、$\beta_t \in (0, 1)$ は各時刻で事前に決められたノイズスケジュール(variance schedule)です。$\beta_t$ が小さいほど、1 ステップあたりに加わるノイズの量は少なくなります。
この定義から、$\bm{x}_t$ は再パラメータ化トリックを用いて以下のように書けます。
$$ \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}) $$
つまり、前の時刻の値を少し縮小し、そこにノイズを加える操作を繰り返します。$T$ が十分大きく、$\beta_t$ が適切にスケジュールされていれば、$\bm{x}_T$ はほぼ標準正規分布に収束します。
任意の時刻 $t$ で直接サンプリングする公式
動機
前向き過程の定義では、$\bm{x}_t$ を得るために $\bm{x}_0$ から 1 ステップずつ逐次的にノイズを加える必要があります。しかし、学習時には任意の時刻 $t$ における $\bm{x}_t$ を効率的に生成したいため、$\bm{x}_0$ から直接 $\bm{x}_t$ をサンプリングする公式が必要です。
導出
$\alpha_t = 1 – \beta_t$ と定義し、さらに累積積を
$$ \bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s = \prod_{s=1}^{t} (1 – \beta_s) $$
と定義します。この記法を用いて、$q(\bm{x}_t | \bm{x}_0)$ を導出しましょう。
まず $t = 1$ のとき、
$$ \bm{x}_1 = \sqrt{\alpha_1}\, \bm{x}_0 + \sqrt{1 – \alpha_1}\, \bm{\epsilon}_1 $$
です。$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_1 \alpha_2}\, \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$ は独立な標準正規分布に従うので、ノイズ部分の分散を合算できます。
$$ \begin{align} \text{Var}[\text{ノイズ部分}] &= \alpha_2(1 – \alpha_1) + (1 – \alpha_2) \\ &= \alpha_2 – \alpha_1 \alpha_2 + 1 – \alpha_2 \\ &= 1 – \alpha_1 \alpha_2 \\ &= 1 – \bar{\alpha}_2 \end{align} $$
したがって、独立な正規分布の和の性質(分散の加法性)を利用すると、
$$ \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$ について以下が成り立ちます。
$$ \boxed{q(\bm{x}_t | \bm{x}_0) = \mathcal{N}\!\left(\bm{x}_t;\, \sqrt{\bar{\alpha}_t}\, \bm{x}_0,\, (1 – \bar{\alpha}_t)\, \bm{I}\right)} $$
再パラメータ化すると、
$$ \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}) $$
この公式のおかげで、学習時には $\bm{x}_0$ と時刻 $t$ を与えるだけで、逐次計算なしに $\bm{x}_t$ を直接生成できます。$\bar{\alpha}_t$ は $t$ が大きくなるにつれて 0 に近づくため、$\bm{x}_T$ は $\bm{x}_0$ の情報をほとんど失い、標準正規分布に近づくことが確認できます。
逆過程(Reverse Process)
定義
逆過程では、前向き過程を「巻き戻す」ことで、ノイズから元のデータを復元します。逆過程もまたマルコフ連鎖として定義されます。
$$ 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{0}, \bm{I})$ であり、各ステップの遷移はパラメータ $\theta$ を持つニューラルネットワークで以下のようにパラメータ化されます。
$$ p_\theta(\bm{x}_{t-1} | \bm{x}_t) = \mathcal{N}\!\left(\bm{x}_{t-1};\, \bm{\mu}_\theta(\bm{x}_t, t),\, \sigma_t^2 \bm{I}\right) $$
つまり、ニューラルネットワークは各時刻 $t$ における $\bm{x}_t$ を受け取り、1 ステップ前の $\bm{x}_{t-1}$ の分布の平均 $\bm{\mu}_\theta(\bm{x}_t, t)$ を予測します。分散 $\sigma_t^2$ は DDPM の原論文では学習せず、固定値を使います。
前向き過程の事後分布
逆過程を学習するためのターゲットとして、$\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})$ です。右辺の 3 つの分布はすべてガウス分布であるため、この事後分布もガウス分布になります。
$$ q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0) = \mathcal{N}\!\left(\bm{x}_{t-1};\, \tilde{\bm{\mu}}_t(\bm{x}_t, \bm{x}_0),\, \tilde{\beta}_t \bm{I}\right) $$
具体的に計算していきましょう。ガウス分布の指数部分を展開します。
$$ \begin{align} &\quad -\frac{1}{2}\left[\frac{(\bm{x}_t – \sqrt{\alpha_t}\,\bm{x}_{t-1})^2}{\beta_t} + \frac{(\bm{x}_{t-1} – \sqrt{\bar{\alpha}_{t-1}}\,\bm{x}_0)^2}{1 – \bar{\alpha}_{t-1}} – \frac{(\bm{x}_t – \sqrt{\bar{\alpha}_t}\,\bm{x}_0)^2}{1 – \bar{\alpha}_t}\right] \end{align} $$
$\bm{x}_{t-1}$ に関する項を整理すると($\bm{x}_{t-1}$ についての二次式として完全平方完成)、
$$ \tilde{\beta}_t = \frac{1}{\frac{\alpha_t}{\beta_t} + \frac{1}{1 – \bar{\alpha}_{t-1}}} = \frac{(1 – \bar{\alpha}_{t-1})\,\beta_t}{1 – \bar{\alpha}_t} $$
$$ \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 $$
ここで、$\bm{x}_t = \sqrt{\bar{\alpha}_t}\,\bm{x}_0 + \sqrt{1 – \bar{\alpha}_t}\,\bm{\epsilon}$ の関係式から $\bm{x}_0$ を逆算すると、
$$ \bm{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}\left(\bm{x}_t – \sqrt{1 – \bar{\alpha}_t}\,\bm{\epsilon}\right) $$
これを $\tilde{\bm{\mu}}_t$ に代入すると、
$$ \begin{align} \tilde{\bm{\mu}}_t(\bm{x}_t, \bm{\epsilon}) &= \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{1}{\sqrt{\bar{\alpha}_t}}\left(\bm{x}_t – \sqrt{1 – \bar{\alpha}_t}\,\bm{\epsilon}\right) \\ &= \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}{\sqrt{\alpha_t}\sqrt{1 – \bar{\alpha}_t}}\,\bm{\epsilon} \end{align} $$
$\bm{x}_t$ の係数をまとめると、
$$ \begin{align} \frac{\sqrt{\alpha_t}(1 – \bar{\alpha}_{t-1})}{1 – \bar{\alpha}_t} + \frac{\beta_t}{(1 – \bar{\alpha}_t)\sqrt{\alpha_t}} &= \frac{\alpha_t(1 – \bar{\alpha}_{t-1}) + \beta_t}{(1 – \bar{\alpha}_t)\sqrt{\alpha_t}} \\ &= \frac{\alpha_t – \bar{\alpha}_t + 1 – \alpha_t}{(1 – \bar{\alpha}_t)\sqrt{\alpha_t}} \quad (\because \beta_t = 1 – \alpha_t) \\ &= \frac{1 – \bar{\alpha}_t}{(1 – \bar{\alpha}_t)\sqrt{\alpha_t}} \\ &= \frac{1}{\sqrt{\alpha_t}} \end{align} $$
したがって、
$$ \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)} $$
この式は極めて重要です。逆過程の事後分布の平均が、現在のノイズ付きデータ $\bm{x}_t$ と、そこに混入したノイズ $\bm{\epsilon}$ で表現されることを示しています。
学習目的関数の導出
変分下界(ELBO)
生成モデルの学習では、データの対数尤度 $\log p_\theta(\bm{x}_0)$ を最大化したいのですが、直接計算は困難です。そこで、VAE と同様に変分下界(ELBO: Evidence Lower Bound) を導出します。
$$ \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} \\ &\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{Jensen の不等式}) \end{align} $$
右辺を $-L_{\text{VLB}}$(Variational Lower Bound の符号反転)として展開します。
$$ \begin{align} -L_{\text{VLB}} &= \mathbb{E}_q\!\left[\log \frac{p_\theta(\bm{x}_{0:T})}{q(\bm{x}_{1:T} | \bm{x}_0)}\right] \\ &= \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} $$
各時刻の KL ダイバージェンスへの分解
前向き過程の条件付き分布を $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)}$ と書き換えると、$L_{\text{VLB}}$ は以下の 3 つの項に分解できます。
$$ L_{\text{VLB}} = \underbrace{D_{\text{KL}}(q(\bm{x}_T | \bm{x}_0)\, \|\, p(\bm{x}_T))}_{L_T} + \sum_{t=2}^{T} \underbrace{D_{\text{KL}}(q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0)\, \|\, p_\theta(\bm{x}_{t-1} | \bm{x}_t))}_{L_{t-1}} + \underbrace{(-\log p_\theta(\bm{x}_0 | \bm{x}_1))}_{L_0} $$
各項の意味は以下のとおりです。
- $L_T$: 前向き過程の最終分布 $q(\bm{x}_T | \bm{x}_0)$ と事前分布 $p(\bm{x}_T) = \mathcal{N}(\bm{0}, \bm{I})$ の KL ダイバージェンス。$\beta_t$ は固定なので $L_T$ は定数であり、学習には関与しません。
- $L_{t-1}$ ($t = 2, \dots, T$): 前向き過程の事後分布 $q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0)$ と逆過程 $p_\theta(\bm{x}_{t-1} | \bm{x}_t)$ の KL ダイバージェンス。これが学習の主要な項です。
- $L_0$: 再構成誤差項。$\bm{x}_1$ からデータ $\bm{x}_0$ を復元する精度を表します。
$L_{t-1}$ の具体的な計算
$q(\bm{x}_{t-1} | \bm{x}_t, \bm{x}_0)$ と $p_\theta(\bm{x}_{t-1} | \bm{x}_t)$ はどちらもガウス分布なので、KL ダイバージェンスは閉じた形で計算できます。2 つのガウス分布 $\mathcal{N}(\bm{\mu}_1, \sigma_1^2 \bm{I})$ と $\mathcal{N}(\bm{\mu}_2, \sigma_2^2 \bm{I})$ の KL ダイバージェンスは、
$$ D_{\text{KL}} = \frac{1}{2\sigma_2^2} \|\bm{\mu}_1 – \bm{\mu}_2\|^2 + (\text{定数項}) $$
分散 $\sigma_t^2$ を固定する($\sigma_t^2 = \tilde{\beta}_t$ とする)と、学習すべきは平均 $\bm{\mu}_\theta(\bm{x}_t, t)$ のみです。したがって、
$$ L_{t-1} = \frac{1}{2\tilde{\beta}_t} \left\|\tilde{\bm{\mu}}_t(\bm{x}_t, \bm{x}_0) – \bm{\mu}_\theta(\bm{x}_t, t)\right\|^2 + C $$
簡略化された損失関数 $L_{\text{simple}}$
先ほど導出した事後分布の平均 $\tilde{\bm{\mu}}_t$ の表式を使い、ニューラルネットワークの出力としてノイズを予測するパラメータ化を採用します。
$$ \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) $$
ここで $\bm{\epsilon}_\theta(\bm{x}_t, t)$ は、$\bm{x}_t$ に混入したノイズ $\bm{\epsilon}$ を予測するニューラルネットワークです。
これを $L_{t-1}$ に代入すると、$\bm{x}_t$ の項がキャンセルされ、
$$ L_{t-1} = \frac{\beta_t^2}{2\tilde{\beta}_t \alpha_t (1 – \bar{\alpha}_t)} \left\|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\right\|^2 + C $$
Ho et al. (2020) は、時刻 $t$ に依存する係数を取り除いた簡略化された損失関数を提案し、これが実験的に良い結果を与えることを示しました。
$$ \boxed{L_{\text{simple}} = \mathbb{E}_{t, \bm{x}_0, \bm{\epsilon}}\left[\left\|\bm{\epsilon} – \bm{\epsilon}_\theta\!\left(\sqrt{\bar{\alpha}_t}\,\bm{x}_0 + \sqrt{1 – \bar{\alpha}_t}\,\bm{\epsilon},\, t\right)\right\|^2\right]} $$
この損失関数は非常にシンプルで、「ノイズ付きデータから、加えたノイズを当てる」という回帰問題に帰着しています。
ノイズ予測ネットワーク $\epsilon_\theta$ の役割
学習の全体像をまとめましょう。ニューラルネットワーク $\epsilon_\theta$ は、以下の入力を受け取ります。
- ノイズ付きデータ $\bm{x}_t$: 元データ $\bm{x}_0$ に時刻 $t$ 分のノイズを加えたもの
- 時刻情報 $t$: 現在の拡散ステップ
そして、$\bm{x}_t$ に含まれるノイズ成分 $\bm{\epsilon}$ を予測します。学習は以下の手順で行われます。
学習アルゴリズム(Training):
- データセットから $\bm{x}_0$ をサンプリング
- 時刻 $t$ を $\{1, 2, \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$ に対して勾配降下法を実行
画像生成タスクでは、$\epsilon_\theta$ として U-Net アーキテクチャが広く使われます。時刻 $t$ の情報は Transformer で使われる正弦波位置エンコーディング(sinusoidal positional encoding)によって埋め込まれます。
サンプリングアルゴリズム
学習が完了したら、以下のアルゴリズムで新しいサンプルを生成します。
サンプリングアルゴリズム(Sampling):
- $\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) + \sigma_t \bm{z}$
- $\bm{x}_0$ を返す
ここで $\sigma_t = \sqrt{\tilde{\beta}_t} = \sqrt{\frac{(1 – \bar{\alpha}_{t-1})\beta_t}{1 – \bar{\alpha}_t}}$ です($t = 1$ のときは $\sigma_1 = \sqrt{\beta_1}$ とします)。
サンプリングでは $T$ 回のニューラルネットワーク評価が必要なため、GAN と比べて生成速度が遅いことが DDPM の弱点です。この問題は後続の研究(DDIM、Distillation など)で改善されています。
ノイズスケジュール $\beta_t$
$\beta_t$ のスケジュールは拡散モデルの性能に大きな影響を与えます。代表的なスケジュールを 2 つ紹介します。
線形スケジュール(Linear Schedule)
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$ が使われました。
コサインスケジュール(Cosine Schedule)
Nichol & Dhariwal (2021) が提案した改良版です。線形スケジュールでは $\bar{\alpha}_t$ が終盤で急激に 0 に近づくため、情報の損失が激しくなります。コサインスケジュールは $\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$ はオフセットパラメータです。$\beta_t$ は $\bar{\alpha}_t$ の変化率から逆算します。
$$ \beta_t = 1 – \frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}} $$
コサインスケジュールは特に低解像度の画像生成で効果的であることが報告されています。
Python 実装:1 次元ガウス混合分布への適用
理論を確認するために、1 次元のガウス混合分布に対して DDPM を訓練し、サンプリング過程を可視化してみましょう。
import numpy as np
import matplotlib.pyplot as plt
# --- 1. データ生成:ガウス混合分布 ---
np.random.seed(42)
n_samples = 10000
# 3つのガウス分布の混合
means = [-3.0, 0.0, 3.0]
stds = [0.5, 0.8, 0.5]
weights = [0.3, 0.4, 0.3]
data = []
for _ in range(n_samples):
k = np.random.choice(len(means), p=weights)
data.append(np.random.normal(means[k], stds[k]))
data = np.array(data, dtype=np.float32)
# --- 2. ノイズスケジュール ---
T = 200 # 拡散ステップ数(1次元なので少なめ)
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)
# --- 3. ノイズ予測ネットワーク(シンプルなMLP) ---
class NoisePredictor:
"""
ε_θ(x_t, t) を予測するシンプルなニューラルネットワーク。
1次元データなので、NumPyで手書きのMLPで十分。
"""
def __init__(self, hidden_dim=128, time_emb_dim=32):
self.hidden_dim = hidden_dim
self.time_emb_dim = time_emb_dim
input_dim = 1 + time_emb_dim
# Xavier初期化
scale1 = np.sqrt(2.0 / input_dim)
scale2 = np.sqrt(2.0 / hidden_dim)
scale3 = np.sqrt(2.0 / hidden_dim)
self.W1 = np.random.randn(input_dim, hidden_dim).astype(np.float32) * scale1
self.b1 = np.zeros(hidden_dim, dtype=np.float32)
self.W2 = np.random.randn(hidden_dim, hidden_dim).astype(np.float32) * scale2
self.b2 = np.zeros(hidden_dim, dtype=np.float32)
self.W3 = np.random.randn(hidden_dim, 1).astype(np.float32) * scale3
self.b3 = np.zeros(1, dtype=np.float32)
# 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_count = 0
def time_embedding(self, t, dim):
"""正弦波による時刻の埋め込み"""
half = dim // 2
freqs = np.exp(-np.log(10000.0) * np.arange(half) / half)
args = t[:, None] * freqs[None, :]
return np.concatenate([np.sin(args), np.cos(args)], axis=-1).astype(np.float32)
def forward(self, x_t, t):
"""順伝播"""
t_emb = self.time_embedding(t, self.time_emb_dim)
inp = np.concatenate([x_t[:, None], t_emb], axis=-1)
# 隠れ層1(ReLU)
self.z1 = inp @ self.W1 + self.b1
self.a1 = np.maximum(self.z1, 0)
# 隠れ層2(ReLU)
self.z2 = self.a1 @ self.W2 + self.b2
self.a2 = np.maximum(self.z2, 0)
# 出力層
out = self.a2 @ self.W3 + self.b3
# キャッシュ
self.inp = inp
return out.squeeze(-1)
def backward(self, x_t, t, eps_true):
"""逆伝播と勾配計算"""
eps_pred = self.forward(x_t, t)
batch_size = len(x_t)
# dL/d(eps_pred)
diff = eps_pred - eps_true
loss = np.mean(diff ** 2)
d_out = (2.0 / batch_size) * diff[:, None]
# 出力層
dW3 = self.a2.T @ d_out
db3 = d_out.sum(axis=0)
d_a2 = d_out @ self.W3.T
# 隠れ層2
d_z2 = d_a2 * (self.z2 > 0).astype(np.float32)
dW2 = self.a1.T @ d_z2
db2 = d_z2.sum(axis=0)
d_a1 = d_z2 @ self.W2.T
# 隠れ層1
d_z1 = d_a1 * (self.z1 > 0).astype(np.float32)
dW1 = self.inp.T @ d_z1
db1 = d_z1.sum(axis=0)
self.grads = [dW1, db1, dW2, db2, dW3, db3]
return loss
def update(self, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8):
"""Adamオプティマイザ"""
self.step_count += 1
for i in range(len(self.params)):
g = self.grads[i]
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_count)
v_hat = self.v[i] / (1 - beta2 ** self.step_count)
self.params[i] -= lr * m_hat / (np.sqrt(v_hat) + eps)
self.W1, self.b1, self.W2, self.b2, self.W3, self.b3 = self.params
# --- 4. 学習 ---
model = NoisePredictor(hidden_dim=128, time_emb_dim=32)
n_epochs = 300
batch_size = 256
losses = []
for epoch in range(n_epochs):
perm = np.random.permutation(n_samples)
epoch_loss = 0.0
n_batches = 0
for i in range(0, n_samples, batch_size):
idx = perm[i:i+batch_size]
x0 = data[idx]
# ランダムな時刻をサンプリング
t = np.random.randint(0, T, size=len(x0))
# ノイズをサンプリング
eps = np.random.randn(len(x0)).astype(np.float32)
# x_t を計算
sqrt_ab = np.sqrt(alpha_bars[t]).astype(np.float32)
sqrt_one_minus_ab = np.sqrt(1.0 - alpha_bars[t]).astype(np.float32)
x_t = sqrt_ab * x0 + sqrt_one_minus_ab * eps
# 勾配計算と更新
loss = model.backward(x_t, t.astype(np.float32), eps)
model.update(lr=1e-3)
epoch_loss += loss
n_batches += 1
losses.append(epoch_loss / n_batches)
if (epoch + 1) % 50 == 0:
print(f"Epoch {epoch+1}/{n_epochs}, Loss: {losses[-1]:.4f}")
# --- 5. サンプリング ---
def sample_ddpm(model, n_samples, T, betas, alphas, alpha_bars):
"""DDPMサンプリングアルゴリズム"""
x = np.random.randn(n_samples).astype(np.float32)
trajectory = [x.copy()]
for t_idx in reversed(range(T)):
t_arr = np.full(n_samples, t_idx, dtype=np.float32)
eps_pred = model.forward(x, t_arr)
# 平均を計算
coeff = betas[t_idx] / np.sqrt(1.0 - alpha_bars[t_idx])
mean = (1.0 / np.sqrt(alphas[t_idx])) * (x - coeff * eps_pred)
if t_idx > 0:
# 分散
beta_tilde = (1.0 - alpha_bars[t_idx - 1]) * betas[t_idx] / (1.0 - alpha_bars[t_idx])
sigma = np.sqrt(beta_tilde)
z = np.random.randn(n_samples).astype(np.float32)
x = mean + sigma * z
else:
x = mean
trajectory.append(x.copy())
return x, trajectory
generated, trajectory = sample_ddpm(model, 5000, T, betas, alphas, alpha_bars)
# --- 6. 可視化 ---
fig, axes = plt.subplots(2, 3, figsize=(15, 9))
# (a) 学習曲線
axes[0, 0].plot(losses, color='steelblue')
axes[0, 0].set_xlabel('Epoch')
axes[0, 0].set_ylabel('MSE Loss')
axes[0, 0].set_title('Training Loss')
axes[0, 0].grid(True, alpha=0.3)
# (b) alpha_bar の推移
axes[0, 1].plot(range(T), alpha_bars, color='darkorange')
axes[0, 1].set_xlabel('t')
axes[0, 1].set_ylabel(r'$\bar{\alpha}_t$')
axes[0, 1].set_title('Noise Schedule ($\\bar{\\alpha}_t$)')
axes[0, 1].grid(True, alpha=0.3)
# (c) 元データ vs 生成データ
x_range = np.linspace(-6, 6, 300)
axes[0, 2].hist(data, bins=80, density=True, alpha=0.5, label='Original', color='steelblue')
axes[0, 2].hist(generated, bins=80, density=True, alpha=0.5, label='Generated', color='coral')
axes[0, 2].set_xlabel('x')
axes[0, 2].set_ylabel('Density')
axes[0, 2].set_title('Original vs Generated')
axes[0, 2].legend()
axes[0, 2].grid(True, alpha=0.3)
# (d)-(f) サンプリング過程の可視化(複数の時刻でのヒストグラム)
time_snapshots = [T, int(T * 0.75), int(T * 0.5), int(T * 0.25), int(T * 0.1), 0]
snapshot_labels = [f't={T-s}' if s < T else 't=0 (final)' for s in time_snapshots]
for ax_idx, (snap_idx, label) in enumerate(zip(time_snapshots[:3], snapshot_labels[:3])):
traj_idx = T - snap_idx # trajectory配列のインデックス
axes[1, ax_idx].hist(trajectory[traj_idx], bins=80, density=True, alpha=0.7, color='mediumpurple')
axes[1, ax_idx].set_title(f'Sampling at {label}')
axes[1, ax_idx].set_xlabel('x')
axes[1, ax_idx].set_ylabel('Density')
axes[1, ax_idx].set_xlim(-6, 6)
axes[1, ax_idx].grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('ddpm_1d_result.png', dpi=150, bbox_inches='tight')
plt.show()
このコードを実行すると、以下の結果が得られます。
- 学習曲線: MSE 損失が安定して減少していくことが確認できます。
- ノイズスケジュール: $\bar{\alpha}_t$ が 1 から 0 に向けて単調減少する様子がわかります。
- 元データ vs 生成データ: 3 つの山を持つガウス混合分布が、拡散モデルによって正確に再現されることが確認できます。
- サンプリング過程: 純粋なガウスノイズから出発し、逆過程を経て徐々にデータ分布の構造が浮かび上がる様子を可視化しています。
まとめ
本記事では、拡散モデル(DDPM)の理論を数式レベルで丁寧に導出しました。
- 前向き過程 $q(\bm{x}_t | \bm{x}_{t-1})$ は、データにガウスノイズを逐次的に加えるマルコフ連鎖であり、$q(\bm{x}_t | \bm{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t}\,\bm{x}_0, (1-\bar{\alpha}_t)\bm{I})$ により任意の $t$ で直接サンプリングできます。
- 逆過程 $p_\theta(\bm{x}_{t-1} | \bm{x}_t)$ は、ニューラルネットワークでパラメータ化されたガウス分布であり、前向き過程の事後分布に一致するように学習されます。
- 学習目的関数は、変分下界(ELBO)から導出され、最終的にはノイズ予測の MSE $L_{\text{simple}} = \mathbb{E}\!\left[\|\bm{\epsilon} – \bm{\epsilon}_\theta(\bm{x}_t, t)\|^2\right]$ という極めてシンプルな形になります。
- サンプリングは $T$ ステップの逐次ノイズ除去で行われます。
DDPM は理論的にエレガントであり、学習も安定しています。一方で、サンプリングに $T$ ステップ(通常 1000 回)のネットワーク評価が必要な点が弱点です。
次のステップとして、以下のトピックも参考にしてください。
- スコアベースモデル(Score-based Models): 拡散モデルをスコアマッチングの観点から理解するアプローチ
- DDIM: 決定論的サンプリングにより少ないステップで高品質な生成を実現する手法
- 条件付き生成(Classifier-Free Guidance): テキストやクラスラベルに基づいた生成を可能にする技術