拡散モデル (Diffusion Models)
ノイズ除去による生成
拡散モデル(Diffusion Models)は、データに段階的にノイズを加える順過程と、 そのノイズを除去してデータを復元する逆過程の2つからなる生成モデルである。 順過程は固定された規則で、学習するのは逆過程だけである(図1)。 DALL-E、Stable Diffusion、Midjourneyなど、最先端の画像生成に使用されている。
前提知識:正規分布、期待値と分散、条件付き確率、ニューラルネットワークの基本。 マルコフ連鎖とベイズの定理は、必要な範囲を記事中で説明する。
本記事が数式で扱うのは、その基本形である DDPM(Denoising Diffusion Probabilistic Models)である。 以降 DDPM と書いたときは、この基本形とその標準的な設定を指す。
拡散モデルの基本概念
数学的定式化
順過程(Forward Process)
ノイズ付加過程
マルコフ連鎖で徐々にガウスノイズを加える。$q(x_t | x_{t-1})$ は「前の状態 $x_{t-1}$ が与えられたときの次の状態 $x_t$ の条件付き確率分布」である:
\[ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1-\beta_t} x_{t-1}, \beta_t I) \]右辺の \(\mathcal{N}(x;\, \mu,\, \Sigma)\) という書き方は「変数 \(x\) についての正規分布の密度で、平均が \(\mu\)、共分散行列が \(\Sigma\)」を表す。 セミコロンが変数とパラメータを区切っている。この式では次のように対応する:
- 変数 \(x \to x_t\):どの変数についての密度か。左辺 \(q(x_t | x_{t-1})\) が \(x_t\) の分布なので、それに対応する
- 平均 \(\mu \to \sqrt{1-\beta_t}\, x_{t-1}\):直前の状態を \(\sqrt{1-\beta_t}\) 倍に縮めた点が分布の中心になる
- 共分散行列 \(\Sigma \to \beta_t I\):\(I\) は単位行列なので各成分は独立で、分散はどれも \(\beta_t\) である(等方的)
第3引数が「分散」ではなく共分散行列(分散共分散行列ともいう)なのは、\(x_t\) が画像すなわちベクトルだからである。 1次元の正規分布なら第3引数はただの分散 \(\sigma^2\) で済むが、多次元では成分どうしの相関まで指定する必要があり、行列になる。 ただしこの順過程では \(\Sigma = \beta_t I\) と単位行列の定数倍なので、実質は「全成分が分散 \(\beta_t\) で互いに独立」という最も単純な場合である。 後の分散の計算で1成分だけ追えば足りるのは、このためである。
すなわち「\(x_{t-1}\) を少し縮めた点を中心に、分散 \(\beta_t\) のノイズをばらまく」ということで、後述の更新式 \(x_t = \sqrt{1-\beta_t}\, x_{t-1} + \sqrt{\beta_t}\, \epsilon_t\) と同じ内容である。
なお \(\epsilon_t \sim \mathcal{N}(0, I)\) のようにセミコロンが無く引数が2つの形は、密度ではなく分布そのものを指し、 引数は (平均, 共分散行列) である。変数が抜けるぶん、\(\mathcal{N}(x;\, \mu,\, \Sigma)\) と比べて引数の位置が1つずれる。
ここで \(\beta_t\) はノイズスケジュール(通常 \(0.0001 \sim 0.02\))である。
なぜ \(\sqrt{1-\beta_t}\) を掛けるのか — 分散の計算
この順過程は1ステップの更新式として書くと \[x_t = \sqrt{1-\beta_t}\, x_{t-1} + \sqrt{\beta_t}\, \epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, I)\] である。ここで \(\epsilon_t\) は \(x_{t-1}\) と独立に毎ステップ引き直す。\(\sqrt{1-\beta_t}\) は信号を縮小する係数で、この係数の役割は分散を追いかけると明確になる。 以下では \(x_t\) の第 \(i\) 成分だけを追い、その分散を \[v_t := \mathrm{Var}\big(x_t^{(i)}\big) = \mathbb{E}\Big[\big(x_t^{(i)} - \mathbb{E}[x_t^{(i)}]\big)^2\Big]\] と書く。加わるノイズが等方的(\(\Sigma = \beta_t I\))で係数もスカラーなので、漸化式の形はどの成分でも同じになる。 ただし値まで成分に依らないのは各成分の初期分散が揃っている場合で、そのとき共分散行列は \(v_t I\) と書ける。 揃っていなくても、成分間のばらつきは (iii) と同じ理屈で \(\bar{\alpha}_t\) 倍に縮んで消える。 以降は \(\mathrm{Var}(x_t)\) とも略記する。\(\mathrm{Var}(\epsilon_t) = 1\) である。
(i) 1ステップで分散がどう変わるか
独立な確率変数 \(X, Y\) と定数 \(a, b\) について \(\mathrm{Var}(aX + bY) = a^2\mathrm{Var}(X) + b^2\mathrm{Var}(Y)\) が成り立つ。 \(a = \sqrt{1-\beta_t}\)、\(b = \sqrt{\beta_t}\)、\(X = x_{t-1}\)、\(Y = \epsilon_t\) として \[ v_t = \left(\sqrt{1-\beta_t}\right)^2 v_{t-1} + \left(\sqrt{\beta_t}\right)^2 \cdot 1 = (1-\beta_t)\, v_{t-1} + \beta_t \tag{1} \] を得る。係数に平方根が付いているのは、分散が「係数の2乗」で効くためである。 \(\sqrt{1-\beta_t}\) と \(\sqrt{\beta_t}\) を掛けておけば、2乗した \((1-\beta_t)\) と \(\beta_t\) がちょうど足して 1 になる。
(ii) \(v = 1\) は不動点である
式 (1) に \(v_{t-1} = 1\) を代入すると \[ v_t = (1-\beta_t)\cdot 1 + \beta_t = 1 - \beta_t + \beta_t = 1 \] となる。よって \(\mathrm{Var}(x_0) = 1\) なら、帰納法により任意の \(t\) で \(\mathrm{Var}(x_t) = 1\) が保たれる(分散保存)。
ただし実データの分散が 1 である必要はない。 画像を \([-1, 1]\) に正規化しても分散が 1 になるわけではなく、 実際 MNIST を \([-1,1]\) に正規化して測ると分散は \(0.3797\)、平均は \(-0.7387\) である (背景の黒が大半を占めるため)。\([-1,1]\) へのスケーリングは、 主としてニューラルネットワークが扱いやすい数値範囲に収めるためのものである。 分散が 1 からずれていてよい理由は、次の (iii) が与える。
(iii) 初期分散がずれていても 1 に引き寄せられる
実データの分散はぴったり 1 ではない。式 (1) の両辺から 1 を引くと、\(\alpha_t = 1 - \beta_t\) として \[ v_t - 1 = (1-\beta_t)\, v_{t-1} + \beta_t - 1 = (1-\beta_t)(v_{t-1} - 1) = \alpha_t\,(v_{t-1} - 1) \] となり、これを \(t\) 回繰り返せば \[ v_t - 1 = \bar{\alpha}_t\,(v_0 - 1), \qquad \bar{\alpha}_t = \prod_{s=1}^{t}(1-\beta_s) \tag{2} \] である。つまり1 からのずれは \(\bar{\alpha}_t\) 倍に縮む。 DDPM の標準設定(\(\beta_t\) を \(10^{-4}\) から \(0.02\) まで線形、\(T = 1000\))では \(\bar{\alpha}_T \approx 4.0 \times 10^{-5}\) なので、 仮に \(\mathrm{Var}(x_0) = 4\) と大きくずれていても \(v_T = 1 + 4.0\times10^{-5}\times 3 \approx 1.00012\) となり、ずれは自動的に吸収される。 上で測った MNIST の \(v_0 = 0.3797\) を入れれば \(v_T = 0.999975\) である。 データを分散 1 に揃える手間は要らないというのが、この (iii) の実際的な意味である。
(iv) 係数がないと分散は発散する
もし \(x_t = x_{t-1} + \sqrt{\beta_t}\,\epsilon_t\) とすると、式 (1) の代わりに \(v_t = v_{t-1} + \beta_t\) となり、 \[ v_T = v_0 + \sum_{s=1}^{T}\beta_s \] と単調に増え続ける。上の標準設定では \(\sum_{s}\beta_s = 1000 \times \frac{10^{-4} + 0.02}{2} = 10.05\) なので、 \(v_0 = 1\) のとき \(v_T = 11.05\)(標準偏差 \(3.32\))まで膨らむ。 しかもこの値はスケジュールと \(v_0\) に依存するため、逆過程の出発点となる分布が事前に決まらない。 \(\sqrt{1-\beta_t}\) は、加えたノイズの分だけ信号を縮めてこの発散を止めるブレーキである。
(v) 平均は 0 に潰れる
分散が 1 に保たれるだけでは \(\mathcal{N}(0, I)\) にはならない。平均も見ると、\(\mathbb{E}[\epsilon_t] = 0\) より \[ \mathbb{E}[x_t] = \sqrt{1-\beta_t}\;\mathbb{E}[x_{t-1}] = \sqrt{\bar{\alpha}_t}\;\mathbb{E}[x_0] \] である。\(\bar{\alpha}_T \approx 4.0\times10^{-5}\) すなわち \(\sqrt{\bar{\alpha}_T} \approx 0.0064\) なので、平均は元の 0.64% まで縮む。
ここで「平均が 0、分散が 1 だから \(\mathcal{N}(0, I)\) である」とは言えないことに注意したい。 平均 0・分散 1 の分布は正規分布とは限らず、一様分布にもそのようなものがある。 平均と分散は正規分布へ近づくために必要な振る舞いを示しているだけで、 それだけで分布の形までは決まらない。
\(x_T\) が実際に標準正規分布に近づくことは、次の (vi) の閉じた形 \(x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) から直接分かる。 \(\bar{\alpha}_t \to 0\) のとき第1項は消え、第2項の係数は 1 に近づくので \[x_T \;\approx\; \epsilon \;\sim\; \mathcal{N}(0, I)\] となる。すなわち \(x_T\) はほぼ「加えたノイズそのもの」であり、 正規分布であることはノイズが正規分布だからである。 これによって、逆過程で「標準正規ノイズからスタート」という手順が正当化される。
(vi) 閉じた形との整合
次に示す \(q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t) I)\) でも同じ結論が確認できる。 \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) の分散を同じ規則で計算すると \[ \mathrm{Var}(x_t) = \bar{\alpha}_t \,\mathrm{Var}(x_0) + (1-\bar{\alpha}_t) \;\overset{\mathrm{Var}(x_0)=1}{=}\; \bar{\alpha}_t + 1 - \bar{\alpha}_t = 1 \] となり、式 (1) を \(t\) 回適用した結果と一致する。信号の重み \(\bar{\alpha}_t\) とノイズの重み \(1-\bar{\alpha}_t\) が常に足して 1 になるのが分散保存の正体である(図3)。 この性質から、この形式の拡散過程は variance preserving (VP) と呼ばれる。 分散を保存しない別系統(variance exploding, VE)との統一的な扱いは Stable Diffusion の数学的原理 の SDE による定式化を参照。
\(\alpha_t = 1 - \beta_t\)、\(\bar{\alpha}_t = \displaystyle\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) \]これにより任意のステップ \(t\) に直接ジャンプできる:
\[ x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \]逆過程(Reverse Process)
この節の地図:ここから式が少し長くなるが、たどり着く結論は単純である。
- \(x_t\) から一段戻るには、そこに含まれているノイズの量が分かればよい。
- そのノイズ \(\epsilon\) は未知なので、ニューラルネットワーク \(\epsilon_\theta(x_t, t)\) に推定させる。
- したがって学習は、自分で加えた既知のノイズを当てさせる教師あり学習として実行できる。
以下のベイズの定理と平方完成は、この3行が実際に成り立つことを確かめる作業である。 途中で見失ったら、ここへ戻ればよい。
逆過程はニューラルネットワーク \(\theta\) でパラメータ化する:
\[ p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) \]ここでセミコロンの前が \(x_t\) ではなく \(x_{t-1}\) になっていることに注意したい。 \(x_t\) を条件として、その1つ前の \(x_{t-1}\) の分布を書いているためで、 順過程と時間の向きが逆であることが変数の位置に現れている。
添字の \(\theta\) はニューラルネットワークの全重みをまとめた記号であって、引数ではない。 \(\mu_\theta\) は「\(\theta\) によって形が決まる関数」という意味で、\(x_t, t\) がサンプルごとに変わる入力であるのに対し、 \(\theta\) は推論中は固定で学習時にだけ変わる。\(p_\theta, \mu_\theta, \Sigma_\theta\) はすべて同じ \(\theta\)、 つまり1つのネットワークから来ている。順過程 \(q\) に添字が付かないのは、そちらが \(\beta_t\) だけで決まっていて学習しないからである。
ここで \(\epsilon_\theta\) はノイズ予測ネットワークである。ノイズを含む \(x_t\) と時刻 \(t\) を受け取り、 \(x_t\) と同じ形のテンソル(画像なら同じ幅・高さ・チャンネル数)を返す: \[\epsilon_\theta:\ (x_t,\ t)\ \longmapsto\ \hat{\epsilon}\ \in \mathbb{R}^{\dim(x_t)}\] その出力が表すのは、\(x_0\) から \(x_t\) を作ったときに加わったノイズ、すなわち閉じた形 \(x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) の \(\epsilon\) の推定値である。 実体は後述の U-Net(図5)である。
時刻 \(t\) は整数1つだが、そのまま数値として入れるのではなくsinusoidal 埋め込みでベクトルに変換してから各層に注入する。
これは \(t\) を、波長の異なる \(\sin\) と \(\cos\) を並べたベクトル
\[\big(\sin(\omega_1 t),\ \cos(\omega_1 t),\ \sin(\omega_2 t),\ \cos(\omega_2 t),\ \dots\big),
\qquad \omega_k = 1/10000^{2k/d}\]
に写す固定の変換で、学習するパラメータを持たない(後述の実装の SinusoidalPositionEmbeddings がこれにあたる)。
生の整数を入れると値域が広すぎて扱いにくいのに対し、この形なら各成分が \([-1,1]\) に収まり、
かつ \(t\) ごとに異なる波形の組み合わせになるので、ネットワークが時刻を区別しやすい。
Transformer で語の位置を表すのに使う位置エンコーディングと同じ仕組みで、
「位置 \(pos\)」の代わりに「時刻 \(t\)」を入れているだけである。
式の意味と、相対位置が回転行列で表せることの導出は
Attention機構 の位置エンコーディングの節で扱っている。
ただし \(x_t\) だけから \(\epsilon\) を言い当てることは原理的にできない。同じ \(x_t\) を作りうる \((x_0, \epsilon)\) の組が何通りもあるからである。 二乗誤差を最小にする最良の予測は条件付き期待値 \[\epsilon^{\star}(x_t, t) := \mathbb{E}\big[\,\epsilon \ \big|\ x_t,\, t\,\big]\] であり、学習が目指す先はこれである。個々の \(\epsilon\) を的中させることではない。
右肩の \(\star\) は理想の最適解という意味の印で、\(\epsilon_\theta\) との違いは次のとおりである。
- \(\epsilon_\theta(x_t, t)\):いま手元にある重み \(\theta\) のネットワークが実際に出す値
- \(\epsilon^{\star}(x_t, t)\):順過程とデータの分布だけで決まる理想の関数。 ネットワークが無くても定義でき、\(\theta\) に依存しない(だから添字 \(\theta\) を付けない)
\(\epsilon_\theta\) が \(\epsilon^{\star}\) にどれだけ近づけるかは、ネットワークの表現力と最適化がうまくいくかで決まる。 両者が一致する保証はないので、記号を分けている。
ネットワークが出力するのはこの \(\epsilon_\theta\) だけで、平均はそこから決まった式で計算する: \[\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon_\theta(x_t, t)\right)\] 共分散行列 \(\Sigma_\theta\) は、素の DDPM では学習させず \(\beta_t I\) に固定するのが普通である (添字 \(\theta\) は一般形の名残で、これを学習させる改良版もある)。 後述の実装でも、この式と固定分散をそのまま使っている。
\(\mu_\theta\) はどこから来るのか — 事後分布からの導出
上の \(\mu_\theta\) の式は天下りに置いたものではなく、順過程だけから決まる真の事後分布 \(q(x_{t-1} | x_t, x_0)\) に由来する。逆過程 \(q(x_{t-1} | x_t)\) 単独は扱えないが、 \(x_0\) も条件に加えると正規分布になり、閉じた形で書ける。
(i) ベイズの定理で既知の分布に分解する
順過程はマルコフ連鎖なので \(q(x_t | x_{t-1}, x_0) = q(x_t | x_{t-1})\) であり、 \[q(x_{t-1} | x_t, x_0) = \frac{q(x_t | x_{t-1})\; q(x_{t-1} | x_0)}{q(x_t | x_0)}\] と書ける。右辺の3つはいずれも既出の正規分布である:
- \(q(x_t | x_{t-1}) = \mathcal{N}(x_t;\ \sqrt{\alpha_t}\,x_{t-1},\ \beta_t I)\)(順過程の定義そのもの)
- \(q(x_{t-1} | x_0) = \mathcal{N}(x_{t-1};\ \sqrt{\bar{\alpha}_{t-1}}\,x_0,\ (1-\bar{\alpha}_{t-1}) I)\)
- \(q(x_t | x_0) = \mathcal{N}(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t) I)\)
(ii) 指数部を平方完成する
以下も1成分で書く(各成分で同じ計算になる)。まず分母は落とせる。 \(q(x_t | x_0)\) は \(x_{t-1}\) を含まないので、\(x_{t-1}\) の関数として見れば定数であり、規格化定数に吸収されるからである: \[q(x_{t-1} | x_t, x_0) \;\propto\; q(x_t | x_{t-1})\; q(x_{t-1} | x_0)\] 残る2つの指数部を書き下すと、\(x_{t-1}\) に依存しない項をまとめて \(\mathrm{const}\) として \[\log q(x_{t-1} | x_t, x_0) = -\frac{1}{2}\left[\frac{\big(x_t - \sqrt{\alpha_t}\,x_{t-1}\big)^2}{\beta_t} + \frac{\big(x_{t-1} - \sqrt{\bar{\alpha}_{t-1}}\,x_0\big)^2}{1-\bar{\alpha}_{t-1}}\right] + \mathrm{const}\] となる。
角括弧の中を展開し、\(x_{t-1}\) について整理する。2乗の係数を \(A\)、1次の係数を \(-2C\) と置くと \[A = \frac{\alpha_t}{\beta_t} + \frac{1}{1-\bar{\alpha}_{t-1}}, \qquad C = \frac{\sqrt{\alpha_t}}{\beta_t}\,x_t + \frac{\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}}\,x_0\] であり(\(x_t^2\) や \(x_0^2\) の項は \(x_{t-1}\) を含まないので \(\mathrm{const}\) 側へ回る)、 \[\log q = -\frac{1}{2}\Big[A\,x_{t-1}^2 - 2C\,x_{t-1}\Big] + \mathrm{const} = -\frac{A}{2}\left(x_{t-1} - \frac{C}{A}\right)^{2} + \mathrm{const}\] と平方完成できる。
これを正規分布 \(\mathcal{N}(x_{t-1};\ \tilde{\mu}_t,\ \tilde{\beta}_t)\) の指数部 \(-\dfrac{(x_{t-1}-\tilde{\mu}_t)^2}{2\tilde{\beta}_t}\) と見比べれば、係数の対応から直ちに \[\tilde{\beta}_t = \frac{1}{A}, \qquad \tilde{\mu}_t = \frac{C}{A} = \tilde{\beta}_t\, C\] が読み取れる。あとは \(A\) を計算するだけである。通分すると \[A = \frac{\alpha_t(1-\bar{\alpha}_{t-1}) + \beta_t}{\beta_t\,(1-\bar{\alpha}_{t-1})}\] で、分子は \[\alpha_t(1-\bar{\alpha}_{t-1}) + \beta_t = \underbrace{\alpha_t + \beta_t}_{=\,1} - \underbrace{\alpha_t\bar{\alpha}_{t-1}}_{=\,\bar{\alpha}_t} = 1 - \bar{\alpha}_t\] と畳まれる。したがって \[\tilde{\beta}_t = \frac{1}{A} = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\,\beta_t\] であり、\(\tilde{\mu}_t = \tilde{\beta}_t C\) に代入して \(x_t\) と \(x_0\) の係数を整理すれば \[q(x_{t-1} | x_t, x_0) = \mathcal{N}\big(x_{t-1};\ \tilde{\mu}_t(x_t, x_0),\ \tilde{\beta}_t I\big) \tag{3}\] \[\tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1-\bar{\alpha}_t}\,x_0 + \frac{\sqrt{\alpha_t}\,(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\,x_t\] を得る。平均が \(x_0\) と \(x_t\) の線形結合になっている点に注意したい。
ここで効いた \(\alpha_t + \beta_t = 1\) と \(\alpha_t\bar{\alpha}_{t-1} = \bar{\alpha}_t\) の2つは、 次の (iii) で \(x_t\) の係数が \(1/\sqrt{\alpha_t}\) に畳まれるときにも同じように効く。 \(\alpha_t = 1-\beta_t\) と \(\bar{\alpha}_t = \prod_s \alpha_s\) という定義の置き方が、 この2箇所で式を綺麗にしているわけである。
(iii) \(x_0\) を消す
式 (3) はこのままでは使えない。生成時に \(x_0\)(これから作ろうとしている画像そのもの)は分かっていないからである。 そこで閉じた形 \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon\) を \(x_0\) について解いた \[x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}\left(x_t - \sqrt{1-\bar{\alpha}_t}\,\epsilon\right)\] を代入する。\(\bar{\alpha}_t = \bar{\alpha}_{t-1}\alpha_t\) より \(\sqrt{\bar{\alpha}_{t-1}}/\sqrt{\bar{\alpha}_t} = 1/\sqrt{\alpha_t}\) なので、 \(x_t\) の係数は \[\frac{1}{\sqrt{\alpha_t}\,(1-\bar{\alpha}_t)}\Big[\beta_t + \alpha_t(1-\bar{\alpha}_{t-1})\Big] = \frac{1}{\sqrt{\alpha_t}\,(1-\bar{\alpha}_t)}\Big[\underbrace{\beta_t + \alpha_t}_{=\,1} - \underbrace{\alpha_t\bar{\alpha}_{t-1}}_{=\,\bar{\alpha}_t}\Big] = \frac{1}{\sqrt{\alpha_t}}\] とまとまる。結局 \(x_0\) が消えて \[\tilde{\mu}_t = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon\right) \tag{4}\] となる。\(\beta_t + \alpha_t = 1\) という定義がちょうど効いて綺麗に畳まれるところが、この導出の要である。
(iv) 未知の \(\epsilon\) をネットワークで置き換える
式 (4) に残った未知量は、\(x_0\) から \(x_t\) を作るときに加えたノイズ \(\epsilon\) ただ一つである。 これを予測値 \(\epsilon_\theta(x_t, t)\) で置き換えたものが、上の \(\mu_\theta(x_t, t)\) にほかならない。
ネットワークにノイズを予測させる理由がここにある。 逆過程の平均を組み立てるのに必要なものが \(\epsilon\) だけなので、学習目標は「加えたノイズを当てる」だけでよい。 次節の損失 \(L_{simple}\) が \(\|\epsilon - \epsilon_\theta(x_t, t)\|^2\) という単純な形になるのは、この事情による。
なお分散を \(\tilde{\beta}_t\) と \(\beta_t\) のどちらに固定してもほとんど違いは出ない。 両者の比 \((1-\bar{\alpha}_{t-1})/(1-\bar{\alpha}_t)\) は \(t\) が小さいうちだけ差があり (\(t=5\) で \(0.743\))、\(t\) が大きくなるとほぼ 1 になる(\(t=999\) で \(1.0000\))。
訓練目的関数
対数尤度を直接最大化するのは難しいため、その下から押さえる量(変分下界)を代わりに最大化する。 そこから導かれる損失を簡略化すると、次の形になる(導出は Stable Diffusion の数学的原理 にある):
\[ L_{simple} = \mathbb{E}_{t, x_0, \epsilon}\left[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \right] \]これは「加えたノイズ \(\epsilon\) を正確に予測せよ」という直感的な目標である。
U-Net アーキテクチャ
改良手法
DDPM vs DDIM
DDIM(Denoising Diffusion Implicit Models)は、DDPM と同じ訓練済みネットワークを使いながら、 サンプリングの規則だけを差し替えて生成を高速化する手法である。訓練をやり直す必要はない。
より正確に言えば、DDIM は DDPM のサンプリングを近似したものではない。 DDPM と同じ訓練目的を与える非マルコフな拡散過程の族を構成し、 その中に決定的なサンプリングを含める、という組み立てになっている。 訓練目的が同じなので、DDPM 用に学習した \(\epsilon_\theta\) をそのまま流用できるわけである。
Classifier-Free Guidance(CFG)
条件 \(c\)(テキストプロンプトなど)に基づく生成を強化するため、 条件付きと無条件の予測を組み合わせる:
\[ \tilde{\epsilon}_\theta(x_t, t, c) = \epsilon_\theta(x_t, t, \varnothing) + w \cdot (\epsilon_\theta(x_t, t, c) - \epsilon_\theta(x_t, t, \varnothing)) \]- \(w\):ガイダンススケール(従来の Stable Diffusion 系では 7 前後など比較的大きな値がよく使われたが、適切な値はモデルやサンプラーによって大きく異なる)
- \(w = 1\):通常の条件付き生成
- \(w > 1\):条件への適合を強化(品質向上、多様性低下)
\(w\) は何をしているのか — 条件方向への外挿
上の式は外挿の形をしている。差分 \(\Delta = \epsilon_\theta(x_t,t,c) - \epsilon_\theta(x_t,t,\varnothing)\) は 「条件 \(c\) を与えると予測がどれだけ変わるか」、つまりただの画像を \(c\) らしくする方向であり、 \(w\) はその方向へ進む倍率である。
- \(w = 0\):無条件そのもの。条件を完全に無視する
- \(w = 1\):\(\epsilon_\theta(\varnothing) + \big(\epsilon_\theta(c) - \epsilon_\theta(\varnothing)\big) = \epsilon_\theta(c)\) と \(\varnothing\) の項が打ち消え、条件付き予測そのものになる
- \(w > 1\):条件付き予測を通り越して外挿する。 すなわちネットワークが単独では出さない予測を人工的に作っている
誇張することで、各ステップの軌道が条件付き分布の密度の高い場所へ寄る。 より典型的な見本に落ちるぶん品質が上がり、同時に典型から外れた個性的な出力は出にくくなる。 「品質向上」と「多様性低下」は別々の現象ではなく、同じ操作の表と裏である。
代償は計算量である。\(w \neq 1\) では \(\epsilon_\theta(c)\) と \(\epsilon_\theta(\varnothing)\) の 2回の推論が要るため、1ステップあたりの計算が約2倍になる。
Latent Diffusion (Stable Diffusion)
PyTorchでの実装
ここまでの内容を、実際に動く最小限のコードにまとめる。 MNIST(28×28 の手書き数字)を対象とし、CPU だけで学習・生成が完結する規模にしてある。 クラス条件と Classifier-Free Guidance も組み込んであるので、 この1本で無条件生成から強い条件付き生成まで試せる。
import torch
import torch.nn as nn
import torch.nn.functional as F
class SinusoidalPositionEmbeddings(nn.Module):
"""時間ステップの位置エンコーディング"""
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, time):
device = time.device
half_dim = self.dim // 2
embeddings = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
embeddings = time[:, None] * embeddings[None, :]
embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
return embeddings
class ResBlock(nn.Module):
"""U-Net を構成する基本ブロック。
GroupNorm → SiLU → 畳み込み を2回通し、その間に時刻とクラスの埋め込みを足し込み、
最後に入力を残差として加える。正規化を畳み込みの「前」に置く形(pre-activation)に
してあるので、ブロックを積み重ねても勾配が通りやすい。
"""
def __init__(self, in_ch, out_ch, emb_dim, groups=8):
super().__init__()
self.norm1 = nn.GroupNorm(min(groups, in_ch), in_ch)
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
self.emb_proj = nn.Linear(emb_dim, out_ch) # 図5 の黄色い矢印にあたる
self.norm2 = nn.GroupNorm(min(groups, out_ch), out_ch)
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
# 入出力のチャンネル数が違うときだけ、残差側も 1×1 畳み込みで合わせる
self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
def forward(self, x, emb):
h = self.conv1(F.silu(self.norm1(x)))
# 埋め込みをチャンネルごとのバイアスとして加える(空間方向には一様)
h = h + self.emb_proj(F.silu(emb))[:, :, None, None]
h = self.conv2(F.silu(self.norm2(h)))
return h + self.skip(x)
class UNet(nn.Module):
"""クラス条件つきのノイズ予測器(約470万パラメータ)
ラベル y は 0..num_classes-1 のクラス、num_classes は「条件なし ∅」を表す。
∅ の埋め込みを1つ余分に持たせておくことで、同じ重みで条件付きと無条件の
両方を学習でき、生成時に Classifier-Free Guidance が使えるようになる。
"""
def __init__(self, base=64, emb_dim=256, num_classes=10, n_res=2):
super().__init__()
self.num_classes = num_classes
self.time_mlp = nn.Sequential(
SinusoidalPositionEmbeddings(emb_dim),
nn.Linear(emb_dim, emb_dim),
nn.SiLU(),
nn.Linear(emb_dim, emb_dim),
)
# +1 が「条件なし ∅」の枠
self.label_emb = nn.Embedding(num_classes + 1, emb_dim)
c1, c2, c3 = base, base * 2, base * 2 # 28×28 / 14×14 / 7×7 でのチャンネル数
self.stem = nn.Conv2d(1, c1, 3, padding=1)
# エンコーダ:各解像度に残差ブロックを n_res 個ずつ置く
self.down1 = nn.ModuleList([ResBlock(c1, c1, emb_dim) for _ in range(n_res)])
self.pool1 = nn.Conv2d(c1, c1, 3, stride=2, padding=1) # 28 -> 14
self.down2 = nn.ModuleList([ResBlock(c1 if i == 0 else c2, c2, emb_dim)
for i in range(n_res)])
self.pool2 = nn.Conv2d(c2, c2, 3, stride=2, padding=1) # 14 -> 7
self.down3 = nn.ModuleList([ResBlock(c2 if i == 0 else c3, c3, emb_dim)
for i in range(n_res)])
# ボトルネック
self.mid1 = ResBlock(c3, c3, emb_dim)
self.mid2 = ResBlock(c3, c3, emb_dim)
# デコーダ:skip connection は連結(cat)なので、最初のブロックだけ入力が2倍になる
self.up3 = nn.ModuleList([ResBlock(c3 * 2 if i == 0 else c3, c3, emb_dim)
for i in range(n_res)])
self.upsample2 = nn.ConvTranspose2d(c3, c2, 4, stride=2, padding=1) # 7 -> 14
self.up2 = nn.ModuleList([ResBlock(c2 * 2 if i == 0 else c2, c2, emb_dim)
for i in range(n_res)])
self.upsample1 = nn.ConvTranspose2d(c2, c1, 4, stride=2, padding=1) # 14 -> 28
self.up1 = nn.ModuleList([ResBlock(c1 * 2 if i == 0 else c1, c1, emb_dim)
for i in range(n_res)])
self.out_norm = nn.GroupNorm(8, c1)
self.out_conv = nn.Conv2d(c1, 1, 3, padding=1)
# 出力層をゼロで始める。学習の最初は「ノイズを予測しない」状態から入るので安定する
nn.init.zeros_(self.out_conv.weight)
nn.init.zeros_(self.out_conv.bias)
def forward(self, x, t, y):
# 時刻とクラスを足して1本の埋め込みにし、全ブロックで使い回す(図5 の黄色い矢印)
emb = self.time_mlp(t) + self.label_emb(y)
h = self.stem(x)
for block in self.down1:
h = block(h, emb)
skip1 = h
h = self.pool1(h)
for block in self.down2:
h = block(h, emb)
skip2 = h
h = self.pool2(h)
for block in self.down3:
h = block(h, emb)
skip3 = h
h = self.mid2(self.mid1(h, emb), emb)
# skip connection は「足す」のではなく「連結する」(図5 の赤い矢印)
h = torch.cat([h, skip3], dim=1)
for block in self.up3:
h = block(h, emb)
h = self.upsample2(h)
h = torch.cat([h, skip2], dim=1)
for block in self.up2:
h = block(h, emb)
h = self.upsample1(h)
h = torch.cat([h, skip1], dim=1)
for block in self.up1:
h = block(h, emb)
return self.out_conv(F.silu(self.out_norm(h)))
# 拡散スケジューラ
class DiffusionScheduler:
def __init__(self, num_timesteps=1000, beta_start=1e-4, beta_end=0.02):
self.num_timesteps = num_timesteps
# ノイズスケジュール
self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def add_noise(self, x_0, t, noise=None):
"""順過程: x_0 にノイズを加えて x_t を生成"""
if noise is None:
noise = torch.randn_like(x_0)
sqrt_alpha_cumprod = self.alphas_cumprod[t].sqrt()
sqrt_one_minus_alpha_cumprod = (1 - self.alphas_cumprod[t]).sqrt()
# ブロードキャスト用に形状を調整
sqrt_alpha_cumprod = sqrt_alpha_cumprod[:, None, None, None]
sqrt_one_minus_alpha_cumprod = sqrt_one_minus_alpha_cumprod[:, None, None, None]
return sqrt_alpha_cumprod * x_0 + sqrt_one_minus_alpha_cumprod * noise
@torch.no_grad()
def sample(self, model, shape, device, labels=None, guidance_scale=1.0):
"""逆過程: ノイズから画像を生成
labels : 生成したいクラス。None なら全て「条件なし ∅」で無条件生成
guidance_scale : Classifier-Free Guidance の w
w = 0 -> 条件を無視、w = 1 -> 条件付き予測そのもの、
w > 1 -> 条件方向へ外挿(品質向上・多様性低下)
"""
n = shape[0]
null = torch.full((n,), model.num_classes, device=device, dtype=torch.long) # ∅
if labels is None:
labels, guidance_scale = null, 1.0
x = torch.randn(shape, device=device)
for t in reversed(range(self.num_timesteps)):
t_batch = torch.full((n,), t, device=device, dtype=torch.long)
# ノイズ予測(w=1 のときは条件付きの1回で足りる)
if guidance_scale == 1.0:
predicted_noise = model(x, t_batch, labels)
else:
eps_cond = model(x, t_batch, labels)
eps_uncond = model(x, t_batch, null)
predicted_noise = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
# x_{t-1} を計算
alpha = self.alphas[t]
alpha_cumprod = self.alphas_cumprod[t]
beta = self.betas[t]
if t > 0:
noise = torch.randn_like(x)
else:
noise = 0
x = (1 / alpha.sqrt()) * (x - (beta / (1 - alpha_cumprod).sqrt()) * predicted_noise)
x = x + (beta.sqrt() * noise)
return x
# 訓練ループ
def train_diffusion(model, scheduler, dataloader, epochs, device, p_uncond=0.15):
"""p_uncond の確率でラベルを ∅ に落とし、条件付きと無条件を同じ重みで学習する。
これをやっておかないと、生成時に eps(∅) が作れず CFG が使えない。"""
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
for epoch in range(epochs):
for batch in dataloader:
x_0 = batch[0].to(device)
y = batch[1].to(device)
batch_size = x_0.size(0)
# 一定確率でラベルを「条件なし ∅」に置き換える
drop = torch.rand(batch_size, device=device) < p_uncond
y = torch.where(drop, torch.full_like(y, model.num_classes), y)
# ランダムな時間ステップをサンプル
t = torch.randint(0, scheduler.num_timesteps, (batch_size,), device=device)
# ノイズを加える
noise = torch.randn_like(x_0)
x_t = scheduler.add_noise(x_0, t, noise)
# ノイズを予測
predicted_noise = model(x_t, t, y)
# 損失計算
loss = F.mse_loss(predicted_noise, noise)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 使い方の例
if __name__ == "__main__":
model = UNet(base=64, num_classes=10)
scheduler = DiffusionScheduler()
# train_diffusion(model, scheduler, mnist_loader, epochs=20, device="cpu")
# 「3」を10枚、ガイダンス強めで生成する
labels = torch.full((10,), 3, dtype=torch.long)
imgs = scheduler.sample(model, (10, 1, 28, 28), "cpu",
labels=labels,
guidance_scale=5.0) # ← 式の w。1 より大きいほど条件に強く従う
# ラベルを渡さなければ無条件生成(w = 0 と同じ)
imgs_uncond = scheduler.sample(model, (10, 1, 28, 28), "cpu")
なぜこの作りなのか
正規化層を入れること、各解像度に畳み込みを複数層置くこと、そしてある程度のパラメータ数を確保することは、 どれも品質に直接効く。他をすべて同じにしたまま 正規化層を外し、各解像度の畳み込みを1層に減らし、パラメータを約200万に落とすと、 最終損失は $0.0219 \to 0.0320$ に悪化し、 ガイダンスなし($w = 1$)での正解率は後述の $0.920$ から $0.460$ まで落ちた。 上の版は2 エポック目で、削った版の 20 エポック分の損失に追いつく。
実際に動かした結果
上のコードを MNIST(手書き数字 60,000 枚、各画素を \([-1,1]\) に正規化)で 20 エポック学習させ、
\(T = 1000\) ステップの逆過程で生成した。パラメータ数は 4,702,593、
20 エポック後の訓練損失は $0.0219$ で、10 エポックを過ぎるとほぼ横ばいになる。
CPU のみで学習に約2時間45分、生成は 100 枚あたり3〜6分かかる
(\(w \neq 1\) では条件付きと無条件の推論が2回要るので約2倍になる)。
訓練時に確率 p_uncond でラベルを \(\varnothing\) に落としてあるので、
学習し直すことなく、\(w\) を変えるだけで無条件生成から強い条件付き生成まで切り替えられる。
指定した数字が実際に出たかを別途学習した分類器(テスト精度 0.979)で判定すると、次のようになった。 「多様性」は同じ数字を指定した10枚について、ピクセル空間での全ペアの \(L_2\) 距離を平均したもので、 大きいほど互いに似ていないことを表す。 本物の MNIST の行は、生成物と同じ手続き(各クラス10枚を訓練データから抜いて同じ式で計算)で測った。
| 設定 | 正解率 | 平均確信度 | 多様性 |
|---|---|---|---|
| 本物の MNIST | 0.983 | 0.983 | 18.13 |
| $w = 0$(条件を無視) | 0.150 | 0.911 | 21.14 |
| $w = 1$(ガイダンスなし) | 0.920 | 0.970 | 17.71 |
| $w = 2$ | 1.000 | 1.000 | 16.66 |
| $w = 3$ | 1.000 | 1.000 | 16.26 |
| $w = 5$ | 1.000 | 1.000 | 15.89 |
$w = 0$ は条件を完全に無視する設定である。正解率 0.150 は 10 クラスの偶然の水準 $0.10$ と 区別できない(100 枚の判定なので偶然のばらつきは $\pm 0.03$ ほどあり、 $0.150$ はその2倍の幅に収まる)。図8 の左上で行と数字が対応していないのがこれである。 $w = 0$ から $w = 5$ までが1つのモデルで連続的に切り替わる点が CFG の利点で、 無条件生成のために別のモデルを用意する必要がない。
$w$ を $0$ から $1$ に変えるだけで、正解率は $0.150 \to 0.920$ に上がる。 $w = 1$ は外挿を使わず条件付き予測 $\epsilon_\theta(x_t, t, c)$ をそのまま使う設定なので、 この上がり幅はガイダンスではなく、条件付けそのものの効果である。 4通りの初期ノイズで測り直した平均確信度は $0.976 \pm 0.010$ で、ばらつきよりはるかに大きい。 クラス条件を学習させておけば、ガイダンスに頼らずともモデルは指定どおりの数字を描く。
$w$ をさらに上げると、正解率と平均確信度は $w = 2$ で $1.000$ に達し、そこから先は動かない。 品質の側は飽和する。
いっぽう多様性は $w = 2$ を過ぎても下がり続ける。 4通りの初期ノイズで測り直すと $w = 2$ で $16.87 \pm 0.15$、$w = 5$ で $15.89 \pm 0.04$、 差 $-0.98$ はばらつきの目安($0.30$)の3倍以上あって、誤差では説明できない。 つまり $w$ を $2$ より大きくしても、品質はもう上がらないのに多様性だけが減っていく。 $w$ は大きいほど良い量ではなく、払うものだけが残る領域がある。
本物の MNIST の多様性 $18.13 \pm 0.27$ と並べると、$w$ の意味がさらにはっきりする。 ガイダンスなし($w = 1$)の $17.71$ が本物に最も近く、 $w$ を上げるほど本物より単調に「似すぎ」の側へ離れていく。 なお $w = 0$ の $21.14$ が本物を上回るのは、多様性が優れているからではない。 条件が効いていないため、同じ数字を指定した10枚に 別々の数字が混ざっているぶん距離が大きく出ているだけである。
$1.000$ は「本物より良い」ではない
$w \geq 2$ では正解率・平均確信度とも $1.000$ となり、本物の MNIST($0.983$)を上回る。 しかしこれは品質が本物を超えたという意味ではない。 分類器が迷わないほど「典型的な」字形ばかりになったということである。 本物の手書き数字には癖のある字が混ざるので確信度は 1.000 にならない。 多様性の数値が $21.14 \to 15.89$ と単調に下がっていることが、その裏づけである。
確信度のような指標は典型性を測っており、分布の再現度を測っていない。 生成モデルの評価に FID のような分布間の距離が使われるのは、この限界を避けるためである。
生成物の画素値は、学習データの範囲 $[-1,1]$ をわずかに超える。 逆過程の更新式には値域を $[-1,1]$ に抑える仕組みがないので、これ自体は当然で、 実装では保存や表示の直前に切り詰める(図8 もそうしている)。
ただし、その超過は $w$ にほとんど依存しない。 外挿をまったく使わない $w = 0$ の時点で $[-1.06,\ 1.08]$ であり、 $w = 5$ でも $[-1.09,\ 1.10]$ にしかならない。 すなわち外挿によって広がったぶんは $0.03$ 程度で、 「条件方向へ押し込みすぎて画素値が破綻する」という形の副作用は、この実装では見られなかった。 $w$ を上げた代償は、値域ではなく多様性の低下のほうに出ている。
この数値をどこまで一般化してよいか
上の表は初期ノイズ1通り・各設定100枚の測定である。 正解率は 100 枚の判定なので偶然のばらつきが $\pm 0.03$ ほど、 シード間のばらつきは平均確信度で $\pm 0.01$、多様性で $\pm 0.15$ ほどある。 この幅より小さい差を読み取ってはいけない。 本文で「差がある」と述べた3か所($w=0$ 対 $w=1$、$w=2$ 対 $w=5$ の多様性、 作りを削った版との比較)は、いずれもこの幅を超えていることを確かめてある。
また、飽和が $w = 2$ で起きたことを他の題材に持ち込むことはできない。 実際の画像生成モデルではガイダンススケールをもっと大きく取ることが多い (従来の Stable Diffusion 系では 7 前後など。適切な値はモデルによって大きく異なる)。 MNIST はクラスが10種類しかなく字形の変化も小さいため飽和が早いのだと考えられるが、 これは確かめていない推測である。
まとめ
- 拡散モデルはノイズ付加(順過程)とノイズ除去(逆過程)を学習
- 訓練目標は単純:加えたノイズを正確に予測する
- U-Netと時間埋め込みでノイズ予測器を構成
- DDIMで高速サンプリング、CFGで条件付き生成を強化
- Latent Diffusion:潜在空間での計算で効率化
- Stable Diffusion、DALL-E 3、Midjourneyなど実用化済み
Stable Diffusion の数学的原理 では、 SDE による統一的定式化、スコア関数とノイズ予測の等価性、DDIM の ODE 解釈、 Classifier-Free Guidance のベイズ的導出を厳密に展開している。
よくある質問(FAQ)
Q1. 拡散モデル(Diffusion Model)とは何か
データにノイズを徐々に加えて白色雑音にする(前向き過程)後、ノイズを予測・除去するネットワークを学習する(逆過程)生成モデルである。DDPM(Denoising Diffusion Probabilistic Models)が基礎形式で、Stable Diffusion等の基盤技術である。
Q2. 拡散モデルはGANに比べてなぜ高品質なサンプルを生成できるのか
GAN的な対抗訓練なく尤度ベースの最適化で安定した学習ができ、モード崩壊が起きにくい。スコアマッチングとの理論的等価性があり、多様なサンプルを生成できる。ただし推論時に多段階のデノイジングが必要で計算コストが高い欠点がある。
Q3. DDIMとは何か
Denoising Diffusion Implicit Models の略である。DDPM の確率的サンプリングを近似したものではなく、DDPM と同じ訓練目的をもつ非マルコフな拡散過程を構成し、その中に決定的なサンプリングを含める。DDPM の訓練済みネットワークをそのまま使えるうえ、ランダムなノイズ注入が無いため時刻を間引いても軌道が壊れにくく、推論ステップ数を大幅に削減できる(通常1000→10〜50)。これにより実用的な速度での生成が可能になった。
Q4. 順過程でなぜ \(\sqrt{1-\beta_t}\) を掛けるのか
分散を 1 に保つためである。分散は係数の2乗で効くので、更新式 \(x_t = \sqrt{1-\beta_t}\,x_{t-1} + \sqrt{\beta_t}\,\epsilon_t\) の分散は \(\mathrm{Var}(x_t) = (1-\beta_t)\mathrm{Var}(x_{t-1}) + \beta_t\) となり、\(\mathrm{Var}(x_{t-1}) = 1\) なら \((1-\beta_t) + \beta_t = 1\) がそのまま保たれる(分散保存, variance preserving)。この係数がないと分散は \(\mathrm{Var}(x_0) + \sum_s \beta_s\) と増え続け、DDPM の標準スケジュールでは 11.05 まで発散して、逆過程を標準正規ノイズから始められなくなる。