DDPM(Denoising Diffusion Probabilistic Models)の学習方法を3つの特徴で理解する
この記事について
書籍「拡散モデル データ生成技術の数理」と書籍「ゼロから作るDeep Learning ❺ ―生成モデル編」を読みながら、拡散モデルを理解しようとしています。この記事では、特にDDPMの理解を深めるために、DDPMの学習の特徴を3つの要素に自分なりに整理してみました。
1. DDPMの概要
1.1 DDPMとは
DDPM(Denoising Diffusion Probabilistic Models)は、2020年にHoらが提案した拡散モデルの一つです。
拡散モデルは、2015年にSohl-Dicksteinらが提案した生成モデルの一つであり、その定義は対になる二つの過程に基づいています。一つは、元のデータに段階的にノイズを加え、最終的にその構造を失ったランダムノイズへと変換する、あらかじめ定められた「順過程」です。もう一つは、その逆をたどり、純粋なノイズから出発して少しずつノイズを除去し、元のデータを復元する「逆過程」であり、この過程をAIが学習します。この学習された復元プロセスが生成モデルとして機能し、新たなデータを創り出すことを可能にします。
DDPM(Denoising Diffusion Probabilistic Models)は、拡散モデルでガウスノイズを扱う場合に相当します。ガウス過程の性質の良さから、理論的に複雑な尤度関数を最適化する代わりに「ノイズが加えられた画像を見て、そこに含まれるノイズ成分そのものを直接予測する」という、シンプルな課題を解くことでDDPMは学習できます。このシンプルな学習方法でも、高品質な画像生成が達成され、その後の拡散モデル研究の基礎となりました。
1.2 DDPMの学習における3つの特徴
DDPMの学習には、3つの特徴があると考えています。
第一は、逐次的学習と理想的な「お手本」の活用です。 複雑な生成問題を小さなステップに分解し、各ステップで解析的に求められる理想的な目標$q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$を設定します。これにより、「何を学習すべきか」が数学的に明確になります。
第二は、解析的表現によりエンコーダ学習を不要にすることです。 順方向過程が学習パラメータを持たない固定の数学関数であるため、VAEのようなモデルがエンコーダ$q_\phi(\mathbf{z}|\mathbf{x})$とデコーダ$p_\theta(\mathbf{x}|\mathbf{z})$の両方を学習する複雑な最適化問題に直面するのに対し、DDPMは学習問題の半分を排除しています。
第三は、ノイズ除去学習とノイズ推定学習の等価性を利用することです。 画像のノイズの除去方法を学習すること、付加されたノイズを予測できるように学習することが数学的に等価です。
1.3 この記事の構成
この記事では、まずDDPMを数学的に定式化し、次に3つの特徴を詳しく解説します。その後、実際の学習方法と推論方法(生成方法)について説明します。
2. DDPMの数学的定式化
2.1 拡散過程の定義
都合の良い逆拡散過程を得るために、拡散過程を以下のように定義します。データ$\mathbf{x}_0$に対して、少しずつノイズを加えていくマルコフ過程です。
q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t \mathbf{I})
q(\mathbf{x}_{1:T} \mid \mathbf{x}_{0}) = \prod_{t=1}^{T} q(\mathbf{x}_t \mid \mathbf{x}_{t-1})
ここで$0< {\beta_t} < 1$は事前に定義された分散スケジュールです。この過程により、複雑な生成問題を扱いやすい小さなステップに分解できます。
任意の時刻での分布を直接表現することも可能です。
\alpha_t = 1 - \beta_t, \ \bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s
とすると、
q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t) \mathbf{I})
これは以下のようにも書けます(reparameterization trickと呼ばれます)。
\mathbf{x}_t = \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}
ここで$\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$は標準正規ノイズです。
任意の時刻での分布が解析的に得られることは数学的帰納法より証明できます。
2.2 お手本となる逆拡散過程の定式化
データ$\mathbf{x}_0$で条件づけた逆拡散過程はガウス分布で表現できます。この事後分布は解析的に計算可能であり、これがモデルの学習すべき理想的な目標(お手本)となります。
q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I})
ここで:
- 平均:
\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t
- 分散:
\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t
です。
解析的に素直なのは$q(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$という事後分布です。
しかし、この事後分布を計算するにあたり周辺尤度$q(\mathbf{x}_{t-1})$も求める必要があります。
これはデータ分布 $p(\mathbf{x}_0)$ 全体での積分が必要なため、解析的に求めることができません。そのため、特定のデータで条件づけて解析的に表現できるようにしています。
平均と分散の導出の概要を示します。ベイズの定理により、
q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \propto q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0) q(\mathbf{x}_{t-1} \mid \mathbf{x}_0)
マルコフ性により
q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \mathbf{x}_0) = q(\mathbf{x}_t \mid \mathbf{x}_{t-1})
です。
q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{\alpha_t} \mathbf{x}_{t-1}, \beta_t \mathbf{I})
q(\mathbf{x}_{t-1} \mid \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0, (1-\bar{\alpha}_{t-1}) \mathbf{I})
であり、2つのガウス分布の積は、再びガウス分布になります。
一般に、$\mathcal{N}(\mathbf{x}; \boldsymbol{\mu}_1, \Sigma_1) \cdot \mathcal{N}(\mathbf{x}; \boldsymbol{\mu}_2, \Sigma_2) \propto \mathcal{N}(\mathbf{x}; \boldsymbol{\mu}, \Sigma)$において、
\Sigma^{-1} = \Sigma_1^{-1} + \Sigma_2^{-1}
\boldsymbol{\mu} = \Sigma(\Sigma_1^{-1}\boldsymbol{\mu}_1 + \Sigma_2^{-1}\boldsymbol{\mu}_2)
この公式を適用すると、上記の平均と分散の式が得られます。
2.3 逆拡散過程のモデルの定義
推論時には、$\mathbf{x}_0$は与えられません。
そこで$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$をモデル化する必要があります。
ノイズのかかった画像$\mathbf{x}_t$から元の画像に近づけていく復元過程です。
データで条件付けない逆拡散過程は解析的に求められないものの、$\beta_{t}$の幅が小さい場合、正規分布で表現できることが知られています。そこで、ある時刻から一つ前の時刻のデータの正規分布について、その形をニューラルネットワークで学習させます。
また、推論のスタートは完全なノイズとします。
p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \Sigma_\theta(\mathbf{x}_t, t))
p_\theta(\mathbf{x}_{T}) = \mathcal{N}(\mathbf{x}_{T}; \mathbf{0}, \mathbf{I})
p_\theta(\mathbf{x}_{0:T}) = p_\theta(\mathbf{x}_{T}) \prod_{t=1}^{T} p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)
通常、分散$\Sigma_\theta$は固定し、平均$\boldsymbol{\mu}_\theta$をニューラルネットで予測します。
Hoら (2020)の原論文では、分散を固定値$\sigma_t^2 \mathbf{I}$(具体的には$\beta_t \mathbf{I}$または$\tilde{\beta}_t \mathbf{I}$)とし、
平均$\boldsymbol{\mu}_\theta$の学習に集中することで良い結果が得られたと報告されています。
この設計選択により、学習すべきパラメータを削減でき、実践的にも優れた性能を示しました。
DDPMでは、このニューラルネットワークについて、各タイムステップでお手本に近づくように学習します。
p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) \rightarrow q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)
3. DDPMの3つの学習特徴
3.1 特徴1:逐次的学習と理想的な「お手本」の活用
DDPMでは、複雑な生成過程を小さなステップに分解し、各ステップで理想的な目標を設定します。
3.1.1 逐次的分解による学習の単純化
従来の生成モデルは「一度にデータ全体を生成」しようとしていました。DDPMは「少しずつノイズを除去する」というタスクに分解することで、学習を単純化しています。
各ステップでの学習目標は2.2、2.3で先述したように、以下のようになります。
p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) \rightarrow q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)
3.1.2 変分下界(ELBO)最大化から導かれる理想的な「お手本」
上記のDDPMの学習目標は、対数尤度の変分下界(VLB)の最大化という観点から導出できます。
学習データと拡散過程を前提に置くことで、尤度最大化が各時刻での$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$を
$q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$にKLの意味で近づけることに相当することが示されます。
3.1.2.1 ELBO導出の詳細計算
あるデータ$\mathbf{x}_{0}$の尤度は、
逆拡散過程の同時確率 $p_{\theta}(\mathbf{x}_{0:T})$について、潜在変数を周辺化することで得られます。
p_{\theta}(\mathbf{x}_0) = \int p_{\theta}(\mathbf{x}_{0:T})\,d\mathbf{x}_{1:T}
対数尤度は:
\log p_{\theta}(\mathbf{x}_0) = \log \int p_{\theta}(\mathbf{x}_{0:T})\,d\mathbf{x}_{1:T}
Jensen不等式により変分下界(ELBO)を得ます:
\log p_{\theta}(\mathbf{x}_0) = \\
\log \int q(\mathbf{x}_{1:T} | \mathbf{x}_0) \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)} d\mathbf{x}_{1:T}\\
\geq \int q(\mathbf{x}_{1:T} | \mathbf{x}_0) \log \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)} d\mathbf{x}_{1:T}\\
= \mathbb{E}_{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\left[ \log \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)} \right]
3.2.1.2 マルコフ分解による展開
生成過程$p_{\theta}$と拡散過程$q$がともにマルコフ過程であることから、同時確率を分解できます:
\begin{aligned}
&\mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)} \left[ \log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} | \mathbf{x}_0)} \right] \\
&= \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)} \left[ \log \frac{p_\theta(\mathbf{x}_{0} | \mathbf{x}_{1}) p_\theta(\mathbf{x}_{1} | \mathbf{x}_{2}) ... p_\theta(\mathbf{x}_{T-1} | \mathbf{x}_{T}) p_\theta(\mathbf{x}_{T})}{q(\mathbf{x}_{T} | \mathbf{x}_{T-1})q(\mathbf{x}_{T-1} | \mathbf{x}_{T-2})...q(\mathbf{x}_{1} | \mathbf{x}_{0})} \right] \\
&= \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[ \log p_{\theta}(\mathbf{x}_T) + \sum_{t=1}^{T} \log \frac{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_t|\mathbf{x}_{t-1})} \right] \\
&= \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[ \log p_{\theta}(\mathbf{x}_T) + \sum_{t=2}^{T} \log \frac{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_t|\mathbf{x}_{t-1})} + \log \frac{p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1})}{q(\mathbf{x}_{1}|\mathbf{x}_{0})} \right] \\
&= \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[ \log p_{\theta}(\mathbf{x}_T) + \sum_{t=2}^{T} \log \frac{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})} \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{0})}{q(\mathbf{x}_{t}|\mathbf{x}_{0})} + \log \frac{p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1})}{q(\mathbf{x}_{1}|\mathbf{x}_{0})} \right] \\
&= \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[ \log \frac{p_{\theta}(\mathbf{x}_T)}{q(\mathbf{x}_{T}|\mathbf{x}_{0})} + \sum_{t=2}^{T} \log \frac{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})} + \log p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1}) \right]
\end{aligned}
ここで、4つめの等号では
q(\mathbf{x}_t|\mathbf{x}_{t-1}) = q(\mathbf{x}_t|\mathbf{x}_{t-1},\mathbf{x}_0)
のマルコフ性とベイズの定理を利用しています。
3.2.1.3 3項への整理
最大化問題を最小化問題に置き換えます。つまり、負の対数尤度を最小化するために、変分上界の上限(以下$L(\theta)$)を最小化する問題に読み替えます。
\begin{aligned}
L(\theta) :&= \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[ -\log \frac{p_{\theta}(\mathbf{x}_T)}{q(\mathbf{x}_{T}|\mathbf{x}_{0})} - \sum_{t=2}^{T} \log \frac{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})} - \log p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1}) \right] \\
&= \underbrace{\mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[-\log \frac{p_{\theta}(\mathbf{x}_T)}{q(\mathbf{x}_{T}|\mathbf{x}_{0})} \right]}_{L_T} \\
&+ \sum_{t=2}^{T} \underbrace{\mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[-\log \frac{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})} \right]}_{L_{t-1}} \\
&+ \underbrace{\mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[- \log p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1}) \right]}_{L_{0}} \\
\end{aligned}
計算していくと最終的に、ELBOは以下の3項に分解されます:
\begin{aligned}
L(\theta) &= \underbrace{D_{\mathrm{KL}}(q(\mathbf{x}_T|\mathbf{x}_0) \| p_{\theta}(\mathbf{x}_T))}_{L_T} \\
&+ \sum_{t=2}^{T}\underbrace{\mathbb{E}_{q(\mathbf{x}_{t} | \mathbf{x}_0)}[D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) \| p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t))]}_{L_{t-1}} \\
&- \underbrace{\mathbb{E}_{q(\mathbf{x}_{1} | \mathbf{x}_0)}[\log p_{\theta}(\mathbf{x}_0|\mathbf{x}_1)]}_{L_0}
\end{aligned}
3.2.1.4 各項の意味と計算方法
中間KL項 $L_t$:学習の中核をなす最も重要な項
\begin{aligned}
\sum_{t=2}^{T} L_t &= \sum_{t=2}^{T} \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_0)}\left[\log \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)} \right] \\
&= \sum_{t=2}^{T} \int q(\mathbf{x}_{1:T} | \mathbf{x}_0) \log \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)} d\mathbf{x}_{1:T} \\
&= \sum_{t=2}^{T} \int q(\mathbf{x}_{t-1}, \mathbf{x}_{t} | \mathbf{x}_0) \log \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)} d\mathbf{x}_{t-1} d\mathbf{x}_{t} \\
&= \sum_{t=2}^{T} \int q(\mathbf{x}_{t-1} | \mathbf{x}_{t}, \mathbf{x}_0) q(\mathbf{x}_{t} | \mathbf{x}_0) \log \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)} d\mathbf{x}_{t-1} d\mathbf{x}_{t} \\
&= \sum_{t=2}^{T} \int d\mathbf{x}_{t} \ q(\mathbf{x}_{t} | \mathbf{x}_0) \int q(\mathbf{x}_{t-1} | \mathbf{x}_{t}, \mathbf{x}_0) \log \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_{t}, \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)} d\mathbf{x}_{t-1} \\
&+ \sum_{t=2}^{T} \mathbb{E}_{q(\mathbf{x}_{t} | \mathbf{x}_0)}[D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) \| p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t))] \\
\end{aligned}
この$L_t$項が拡散モデルの学習の中核をなします。
モデルの逆拡散過程 $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$が 真の逆拡散過程$q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)$に近づくように学習することが目的です。
これが「3.1 特徴1:逐次的学習と理想的な「お手本」の活用」最も伝えたかったことです。
最終復元項 $L_0$:データ再構築品質の評価
L_0 = \mathbb{E}_{q(\mathbf{x}_{1} | \mathbf{x}_0)}[-\log p_{\theta}(\mathbf{x}_0|\mathbf{x}_1)]
この項は、モデルがノイズの多い$\mathbf{x}_1$から
元のデータ$\mathbf{x}_0$をどれだけうまく再構築できるかを示します。
$p_\theta(\mathbf{x}_0|\mathbf{x}_1)$は、モデルが$\mathbf{x}_1$から$\mathbf{x}_0$を予測する確率分布で、通常はガウス分布としてモデル化されます:
p_\theta(\mathbf{x}_0|\mathbf{x}_1) = \mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1,1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1,1))
この負の対数尤度を最小化することは、ニューラルネットワークで予測された平均$\boldsymbol{\mu}_\theta(\mathbf{x}_1,1)$が実際の$\mathbf{x}_0$に近づくように学習することを意味します。
こちらも、$\mathbf{x}_0$をお手本としており、理想的な「お手本」の活用して学習していると言えます。
終端KL項 $L_T$:拡散過程の収束性を保証
\begin{aligned}
L_T &= \int q(\mathbf{x}_{1:T} | \mathbf{x}_0) \log \frac{q(\mathbf{x}_{T} | \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{T})} d\mathbf{x}_{1:T} \\
&= \int q(\mathbf{x}_{T} | \mathbf{x}_0) \log \frac{q(\mathbf{x}_{T} | \mathbf{x}_{0})}{p_{\theta}(\mathbf{x}_{T})} d\mathbf{x}_{T} \\
&= D_{\mathrm{KL}}(q(\mathbf{x}_T|\mathbf{x}_0) \| p_{\theta}(\mathbf{x}_T))
\end{aligned}
ここで$p_{\theta}(\mathbf{x}_T) = \mathcal{N}(\mathbf{0}, \mathbf{I})$です。
$p_{\theta}(\mathbf{x}_T) = \mathcal{N}(\mathbf{0}, \mathbf{I})$のため、パラメータに依存しないこうのため、$L(\theta)$の最小化をする上では無視します。
3.2.1.5 解析的に求められる理想的な「お手本」
ここまでで示したように、負のVLBは各時間ステップ$t$に対応する項の和に分解できます。
その項の多く($L_t$, $t\neq0$, $t\neq T$)はKLダイバージェンス
D_{KL}(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t))
の形をとります。
この上界を最小化することが、
各時刻での$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$と
$q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$のKLダイバージェンスを最小化することに対応します。
$q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$が解析的に求められるため、理論的に完璧なノイズ除去の手順を導出できます。
ニューラルネットワークは、この理想的な過程を模倣します。
\text{学習目標:} \quad p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) \rightarrow q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)
これにより、何を学習すべきかが数学的に明確になります。理想的な「答え」が既に分かっているため、効率的な学習が可能です。
3.2 特徴2:エンコーダ学習の不要化
VAEでは、エンコーダ$q_\phi(\mathbf{z} \mid \mathbf{x})$とデコーダ$p_\theta(\mathbf{x} \mid \mathbf{z})$の両方を学習する複雑な最適化問題に直面します。
DDPMでは順方向過程が数学的に明確に定義されているため、
順方向$q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1-\bar{\alpha}_t) \mathbf{I})$は学習不要で、
逆方向$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$のみが学習対象となります。
これは、順方向過程$q(\mathbf{x}_t|\mathbf{x}_0)$が、
VAEのエンコーダ$q_\phi(\mathbf{z}|\mathbf{x})$のように学習可能な重み$\phi$によってパラメータ化されておらず、
固定の分散スケジュール${\beta_t}$のみによって決定されるためです。
これにより、学習対象が明確になり、エンコーダ・デコーダ間の不整合が発生しません。単一の目的関数による学習(ノイズ除去)だけに集中すれば良いことになります。
3.3 特徴3:ノイズ除去学習とノイズ推定学習の等価性
ノイズ除去学習とノイズの推定が数学的に等価であることを示します。
これまで述べた通り、DDPMは各ステップでノイズ除去の方法を解析的に求めたお手本をもとに学ぶのでした。
L_{t-1}
= \mathbb{E}_{q(\mathbf{x}_{t}\mid\mathbf{x}_0)}
\bigl[
D_{\mathrm{KL}}\!\bigl(
q(\mathbf{x}_{t-1}\mid\mathbf{x}_t,\mathbf{x}_0)
\,\big\|\,
p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_t)
\bigr)
\bigr]
p_\theta\!\bigl(\mathbf{x}_{t-1}\mid\mathbf{x}_t\bigr)
\;\xrightarrow[\text{minimize }L_{t-1}]{}
q\!\bigl(\mathbf{x}_{t-1}\mid\mathbf{x}_t,\mathbf{x}_0\bigr)
L_{0}
= \mathbb{E}_{q(\mathbf{x}_{1}\mid\mathbf{x}_0)}
\Bigl[-\log p_{\theta}(\mathbf{x}_0\mid\mathbf{x}_1)\Bigr] \\
\boldsymbol{\mu}_{\theta}(\mathbf{x}_1,1)
\;\xrightarrow[\text{minimize }L_{0}]{}
\mathbf{x}_0
ノイズ除去ができるということは、ノイズ有のデータからノイズ無のデータを作り、その差分としてノイズを得ることもできます。
逆にノイズを推定できば、ノイズ除去もできます。
そのため、直観的には、ノイズ除去を学ぶこととノイズ推定を学ぶことは等価に思えます。
これを数式を用いて丁寧に証明します。
$L_{t-1}$については、まずガウス分布の性質の良さから、KLダイバージェンスの部分をさらに簡略化できることを示します。
そして、ノイズ除去の学習が、お手本となる逆拡散過程の平均を推定する問題に置き換わることをみます。
さらに、お手本となる逆拡散過程の平均をノイズ部分とそれ以外に分けて定式化することで、モデルの学習をノイズ部分の推定にフォーカスできることを示します。
$L_0$については、
p_\theta(\mathbf{x}_0|\mathbf{x}_1) = \mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1,1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1,1))
がガウス分布なので対数尤度の計算が可能です。
3.3.1 $L_{t-1}$について
KLダイバージェンスの簡略化:
KLダイバージェンスは、
モデルの逆拡散過程$p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$と
真の逆拡散過程$q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)$がともにガウス過程であることから、さらに簡略化できます。
各時間ステップでの事後分布$q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)$はガウス分布で表現できます:
q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0), \tilde{\beta}_t\mathbf{I})
モデルの逆拡散過程も同様に固定された分散を持つガウス分布として定義しています:
p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t,t), \sigma_t^2\mathbf{I})
分散が固定されている場合、ガウス分布間のKLダイバージェンスは平均の差の二乗誤差の形になります:
D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) \| p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)) = \frac{1}{2\sigma_t^2} \|\boldsymbol{\mu}_\theta(\mathbf{x}_t,t) - \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0)\|^2 + Const.
ノイズ予測への変換:
この平均の予測をノイズの予測に置き換えることができます。
真の平均$\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0)$は以下のように変形できます(本節最後に計算方法を記載):
\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0) = \frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\boldsymbol{\epsilon}\right)
この形に合わせるように、ニューラルネットワークの学習対象を再定義できます。
\boldsymbol{\mu}_\theta(\mathbf{x}_t,t) := \frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)\right)
これをKLダイバージェンスの式に代入すると、ニューラルネットワークにより予測されたノイズ$\boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)$と真のノイズ$\boldsymbol{\epsilon}$との間の平均二乗誤差を最小化することに帰着します:
\begin{aligned}
L_{t-1} - Const. &= \mathbb{E}_{q(\mathbf{x}_{t} | \mathbf{x}_0)}[\frac{1}{2\sigma_t^2} \|\boldsymbol{\mu}_\theta(\mathbf{x}_t,t) - \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0)\|^2] \\
&= \mathbb{E}_{q(\mathbf{x}_{t} | \mathbf{x}_0)}[\frac{{\beta_t}^2}{2\sigma_t^2 \alpha_t (1-\bar{\alpha}_t)} \|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t,t)\|^2] \\
&= \mathbb{E}_{\mathbf{x}_0 \sim q(\mathbf{x}_0), \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0},\mathbf{I})}[\frac{{\beta_t}^2}{2\sigma_t^2 \alpha_t (1-\bar{\alpha}_t)} \|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon},t)\|^2] \\
\end{aligned}
3.3.2 $L_{0}$について
データの次元をDとして、対数尤度を計算すると、
\begin{aligned}
-\log p_{\theta}(\mathbf{x}_0|\mathbf{x}_1) &= -\log \left( \frac{1}{\sqrt{(2\pi\sigma_1^2)^D}} \exp\left(-\frac{||\mathbf{x}_0 - \mathbf{\mu}_\theta(\mathbf{x}_1, 1)||^2}{2\sigma_1^2}\right) \right) \\
&= \frac{D}{2}\log(2\pi\sigma_1^2) + \frac{1}{2\sigma_1^2}||\mathbf{x}_0 - \mathbf{\mu}_\theta(\mathbf{x}_1, 1)||^2
\end{aligned}
であるので、
L_0 = \mathbb{E}_{q(\mathbf{x}_{1} | \mathbf{x}_0)}[-\log p_{\theta}(\mathbf{x}_0|\mathbf{x}_1)]
の式に代入します。
L_0 = \mathbb{E}_{q(\mathbf{x}_{1} | \mathbf{x}_0)}\left[ \frac{D}{2}\log(2\pi\sigma_1^2) + \frac{1}{2\sigma_1^2}||\mathbf{x}_0 - \mathbf{\mu}_\theta(\mathbf{x}_1, 1)||^2 \right]
第1項 $\frac{D}{2}\log(2\pi\sigma_1^2)$ はモデルのパラメータ $\theta$ に依存しない定数なので、最適化の観点からは無視できます。したがって、$L_0$ の最小化は、実質的に以下の項の最小化と等価です。
L_0^{'} = \mathbb{E}_{q(\mathbf{x}_{1} | \mathbf{x}_0)}\left[ ||\mathbf{x}_0 - \mathbf{\mu}_\theta(\mathbf{x}_1, 1)||^2 \right]
さらに計算を推し進めます。$L_{t-1}$の場合と同じようにニューラルネットワークの学習対象を再定義しますと、
\boldsymbol{\mu}_\theta(\mathbf{x}_1,1) := \frac{1}{\sqrt{\alpha_1}}\left(\mathbf{x}_1 - \frac{\beta_1}{\sqrt{1-\bar{\alpha}_1}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_1,1)\right)
これと
\mathbf{x}_1 = \sqrt{\bar{\alpha}_1} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_1} \boldsymbol{\epsilon}
を$L_0^{'}$に代入すると、$L_{t-1}$の場合と同じ形が出に入ります:
L_0^{'} = \mathbb{E}_{\mathbf{x}_0 \sim q(\mathbf{x}_0), \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0},\mathbf{I})} \left[ \frac{1-\bar{\alpha}_1}{\bar{\alpha}_1} ||\mathbf{\epsilon} - \mathbf{\epsilon}_\theta(\sqrt{\bar{\alpha}_1}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_1}\mathbf{\epsilon}, 1)||^2 \right]
3.3.3 $L = L_0 + \sum_{t=2}^{T} L_{t-1} + L_{T}$ について
以上の議論から$\theta$に依存しない部分($L_{T}$等)無視すれば、
L = \sum_{t=1}^{T} \mathbb{E}_{\mathbf{x}_0 \sim q(\mathbf{x}_0), \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0},\mathbf{I})}[\omega_t \|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon},t)\|^2] \\
$w_t$は各時刻での2乗項の重みです。具体的案値はこれまでの式変形を参照してください。
ここで終わってもよいのですが、Hoら (2020)はさらに単純化を進め、$w_t = 1$としました。
つまり、
L = \sum_{t=1}^{T} \mathbb{E}_{\mathbf{x}_0 \sim q(\mathbf{x}_0), \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0},\mathbf{I})}[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon},t)\|^2]
であり、T個の和は一様分布$U(t)$の期待値とみることができるので、
L = \mathbb{E}_{t \sim U(t), \mathbf{x}_0 \sim q(\mathbf{x}_0), \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0},\mathbf{I})}[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon},t)\|^2] =: L_{\text{simple}}
となる。
Hoら (2020)はこの単純化された目的関数を導入し、「この目的関数で学習した結果、より良いサンプル品質が得られた」と報告しています。
以下は参考まで。
真の平均$\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0)$は以下のように変形できることについて、導出過程を示します。
$\mathbf{x}_t = \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}$から、
\mathbf{x}_0 = \frac{\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}}{\sqrt{\bar{\alpha}_t}}
です。
\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t
に$\mathbf{x}_0 = \frac{\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}}{\sqrt{\bar{\alpha}_t}}$を代入すると、
\tilde{\boldsymbol{\mu}}_t = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} \cdot \frac{\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}}{\sqrt{\bar{\alpha}_t}} + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t
この式を$\mathbf{x}_t$と$\boldsymbol{\epsilon}$について整理すると、
\tilde{\boldsymbol{\mu}}_t = \frac{1}{\sqrt{\alpha_t}} \mathbf{x}_t - \frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}
$\sqrt{\alpha_t} = \sqrt{1-\beta_t}$であることを考慮すると、最終的に
\tilde{\boldsymbol{\mu}}_t = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon} \right)
が得られます。
4. 学習方法
特徴3により、学習アルゴリズムは以下のようになります。
- データバッチ${\mathbf{x}_0}$をサンプル
- 時刻$t \sim \text{Uniform}(1, T)$をサンプル
- ノイズ$\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$をサンプル
- $\mathbf{x}_t = \sqrt{\bar{\alpha}_t} \mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon}$を計算
- 損失$L_{\text{simple}} = |\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)|^2$を計算
- 勾配降下法で$\theta$を更新
単純化された目的関数$L_{\text{simple}}$を用いることで、学習ループが非常に簡潔になっていることがわかります。
5. 推論過程
5.1 サンプリング手順
推定したノイズを除去する作業を続けていくことで、標準正規分布に従うノイズから画像を生成します。
1. \mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\
2. for t = T, ..., 1 do \\
3.
\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \text{ if t > 1, else } \mathbf{z} = \mathbf{0}
4.
\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \mathbf{\epsilon}_\theta(\mathbf{x}_t, t) \right) + \sigma_t \mathbf{z}
5. end for
6. return \mathbf{x}_0
※$\mathbf{x}_1$から$\mathbf{x}_0$をサンプリングする際は、ノイズを追加せずにサンプリングする方が性能が良いことが知られています。
※数式が汚くてごめんなさい。。。ドルマークがうまく機能しないときがあり、なくなくこの形に。
DDPMのサンプリングは数百から数千回の逐次的なモデル評価を必要とするため、非常に遅いことが広く知られた制約です。
5.2 高速化手法 (generated by AI)
DDPMの基本的なサンプリングは多くのステップを必要としますが、様々な高速化手法が提案されています。
DDIMは、通常1,000ステップ以上かけて画像を生成する拡散モデル(DDPM)の「1ステップごとに少しずつノイズを足しては消す」という流れを、η(イータ)というパラメータで“途中のノイズ付加をカット”しながら進む仕組みに置き換えたものです。η=1なら従来どおり、ステップごとにランダムなノイズを混ぜつつ後戻りし、η=0なら「学習済みの平均(μ)のみ」をたどる完全決定論的な更新になります。これにより、本来マルコフ的に1ステップずつしか動けない逆拡散過程を、いくつか飛ばしながら一気に進める非マルコフ的なやり方で実行できます。
なぜ途中のノイズを省いても問題ないかというと、モデルは訓練時に「ノイズありデータから元のデータへ戻る平均的な最短経路」をしっかり学んでおり、その平均値(μ)だけを追えば、理論上はデータの分布を再現できると保証されているからです。また生成の多様性は、最初に与える“最終ステップの乱数(初期ノイズ)”だけで生まれるため、途中のノイズ省略が画像のバリエーションに影響を与えません。この結果、ステップ数を数十程度に減らしても高品質を維持しつつ、確実に高速で安定したサンプリングが可能になります。
6. まとめ
この記事では、DDPMの学習を3つの特徴に焦点を当てて説明しました。
逐次的学習と理想的なお手本の活用により、複雑な生成問題を扱いやすくし、効率的な学習を可能にしています。各時間ステップで解析的に求められる理想的な目標$q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$が、明確な学習指針になります。
エンコーダ学習の不要化により、学習対象を明確にしています。VAEのような複雑な同時最適化問題を回避し、モデルの全能力をノイズ除去タスクに集中投入できます。
ノイズ除去学習とノイズ推定学習の等価性により、単純化された目的関数$L_{\text{simple}}$による効率的な学習を可能にしました。
参考文献
- Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., & Ganguli, S. (2015) "Deep Unsupervised Learning using Nonequilibrium Thermodynamics", ICML 2015. [arXiv:1503.03585] - 拡散モデルの理論的基礎を築いた原点となる論文
- Ho, J., Jain, A., & Abbeel, P. (2020) "Denoising Diffusion Probabilistic Models", NeurIPS 2020. [arXiv:2006.11239] - 本記事の紹介するDDPMの原論文
- Song, J., Meng, C., & Ermon, S. (2021) "Denoising Diffusion Implicit Models (DDIM)", ICML 2021. [arXiv:2010.02502] - サンプリングの高速化
- 岡野原 大輔 (2023) 拡散モデル データ生成技術の数理
- 斎藤 康毅 (2024) ゼロから作るDeep Learning ❺ ―生成モデル編