メインコンテンツまでスキップ

DDPMにおける変分下界(ELBO)の導出とノイズ予測 MSE への帰着

概要 (画像は、Geminiで作成されたものです)

概要と全体像

DDPM(Denoising Diffusion Probabilistic Models)において、 負の対数尤度の最小化は変分下界(ELBO: Evidence Lower Bound) の最大化(=ELBO損失 LVLBL_{\text{VLB}} の最小化)と同値になります。

一見複雑な確率分布のカルバック・ライブラー(KL)ダイバージェンスの総和である LVLBL_{\text{VLB}} は、以下の 3 つのステップを経て、 最終的に「実際に加えたノイズ ϵ\boldsymbol{\epsilon}」と「モデルが予測したノイズ ϵθ\boldsymbol{\epsilon}_\theta」の二乗誤差(MSE)へと奇跡的に帰着します。

  1. マルコフ性の利用によるELBOの項別分解
  2. 条件付き逆過程 q(xt1xt,x0)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t, \mathbf{x}_0) のガウス分布による解析的表示
  3. 平均 μθ\boldsymbol{\mu}_\theta からノイズ ϵθ\boldsymbol{\epsilon}_\theta へのパラメータ再表記

ELBO の展開と項の分解

なぜ負の対数尤度を最小化するのか?

生成モデルの根本的な目的は、モデルが学習データ(実画像など)x0\mathbf{x}_0 を生成する確率、すなわち尤度 pθ(x0)p_\theta(\mathbf{x}_0) を最大化することです。 これは統計学における最尤推定(Maximum Likelihood Estimation: MLE)のアプローチであり、 数学的にはその対数をとってマイナスを掛けた負の対数尤度 logpθ(x0)-\log p_\theta(\mathbf{x}_0) を最小化することと同値になります。 モデルが真のデータ分布に近づくほど、この値は小さくなります。

潜在変数と変分下界(ELBO)の導入

しかし、データ x0x_0 に対する負の対数尤度 logpθ(x0)-\log p_\theta(x_0) を直接計算・最小化することは困難です。 モデルの生成過程は最終的な x0x_0 だけでなく、そこに至るまでの無数のノイズ状態(潜在変数)x1:Tx_{1:T} を経由するため、以下の式のようにすべての可能な経路を考慮した積分(周辺化)が必要になるからです。

logpθ(x0)=logpθ(x0:T)dx1:T\log p_\theta(x_0) = \log \int p_\theta(x_{0:T}) dx_{1:T}

数学や機械学習の解析において、このように対数関数の内側に積分や総和が含まれる形(logf(x)dx\log \int f(x) dx)は非常に扱いづらいとされています。 対数の内側で足し合わせが行われていると、各変数に関する項を独立して分解することができず、解析的に数式を解くことができないためです。 また、高次元の潜在空間ですべての経路を計算することも現実的ではありません。

そこで、この複雑な積分を回避するため、拡散過程(順過程)の確率分布 q(x1:Tx0)q(x_{1:T} \mid x_0) を分母と分子に掛け合わせ、積分の式を期待値の形に書き換えます。

logpθ(x0)=logq(x1:Tx0)pθ(x0:T)q(x1:Tx0)dx1:T=logEq(x1:Tx0)[pθ(x0:T)q(x1:Tx0)]\log p_\theta(x_0) = \log \int q(x_{1:T} \mid x_0) \frac{p_\theta(x_{0:T})}{q(x_{1:T} \mid x_0)} dx_{1:T} = \log \mathbb{E}_{q(x_{1:T} \mid x_0)} \left[ \frac{p_\theta(x_{0:T})}{q(x_{1:T} \mid x_0)} \right]

ここで、対数関数(log\log)は上に凸な関数です。そのため、イェンゼンの不等式(Jensen's Inequality)logE[f(x)]E[logf(x)]\log \mathbb{E}[f(x)] \ge \mathbb{E}[\log f(x)] を適用して、期待値と対数の順序を入れ替えます。
これにより対数が式の内側に入るため、数式の展開(分解)が可能になり、計算可能な対数尤度の下界(変分下界:ELBO)が得られます。

logpθ(x0)Eq(x1:Tx0)[logpθ(x0:T)q(x1:Tx0)]\log p_\theta(\mathbf{x}_0) \ge \mathbb{E}_{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \left[ \log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \right]

この不等式は、特定の1つのデータ点 x0\mathbf{x}_0 に対する対数尤度の下界を示しています。 しかし、モデルの学習において最小化したいのは、特定の1点に対する誤差ではなく、データセット全体(すなわちデータ分布 q(x0)q(\mathbf{x}_0))にわたる平均的な負の対数尤度 です。

そのため、まず両辺にマイナスを掛けて不等号の向きを反転させた上で、データ分布 q(x0)q(\mathbf{x}_0) の下での期待値 Eq(x0)\mathbb{E}_{q(\mathbf{x}_0)} を取ります。

Eq(x0)[logpθ(x0)]Eq(x0)[Eq(x1:Tx0)[logpθ(x0:T)q(x1:Tx0)]]\mathbb{E}_{q(\mathbf{x}_0)}\left[ -\log p_\theta(\mathbf{x}_0) \right] \le \mathbb{E}_{q(\mathbf{x}_0)} \left[ \mathbb{E}_{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \left[ -\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \right] \right]

確率の乗法定理 q(x0:T)=q(x0)q(x1:Tx0)q(\mathbf{x}_{0:T}) = q(\mathbf{x}_0) q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) から、右辺のデータ分布 q(x0)q(\mathbf{x}_0) と条件付き拡散過程 q(x1:Tx0)q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) の二重の期待値は、データとすべてのノイズ状態の同時分布 q(x0:T)q(\mathbf{x}_{0:T}) による期待値として、以下のように1つにまとめることができます。

Eq(x0)[logpθ(x0)]Eq(x0:T)[logpθ(x0:T)q(x1:Tx0)]\mathbb{E}_{q(\mathbf{x}_0)}\left[ -\log p_\theta(\mathbf{x}_0) \right] \le \mathbb{E}_{q(\mathbf{x}_{0:T})}\left[ -\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \right]

これによって、左辺(本来最小化したい、データセット全体の平均的な負の対数尤度)を間接的に最小化するための、計算可能な上界(変分下界:Variational Upper Bound)が得られます。 これが「負の対数尤度の変分下界(LVLBL_{\text{VLB}})」の正体です。

LVLB=Eq[logpθ(x0:T)q(x1:Tx0)]L_{\text{VLB}} = \mathbb{E}_{q}\left[ -\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \right]

(※以降、数式を簡潔にするため、データとすべてのノイズ状態の同時分布 q(x0:T)q(\mathbf{x}_{0:T}) による期待値 Eq(x0:T)\mathbb{E}_{q(\mathbf{x}_{0:T})} を、単に Eq\mathbb{E}_{q} と表記します)

直接計算できない負の対数尤度の代わりに、この LVLBL_{\text{VLB}} を最小化することで、間接的に尤度を最大化します。

マルコフ連鎖の性質やベイズの定理を用いてこの式を展開していくと、各タイムステップ tt ごとの KL ダイバージェンスの和に分解することができます。

この変形プロセスを、ステップ・バイ・ステップで追ってみましょう。

LVLBL_{\text{VLB}} の各タイムステップへの分解プロセス

LVLBL_{\text{VLB}} の定義式の中身にある分数(対数の引数) pθ(x0:T)q(x1:Tx0)\frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} を分解していきます。

Step 1: マルコフ連鎖の定義の適用

生成過程(逆過程)および拡散過程(順過程)はどちらもマルコフ連鎖であるため、同時確率および条件付き確率は以下のように展開できます。

  • 生成過程: pθ(x0:T)=p(xT)t=1Tpθ(xt1xt)p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)
  • 拡散過程: q(x1:Tx0)=t=1Tq(xtxt1)q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod_{t=1}^T q(\mathbf{x}_t \vert \mathbf{x}_{t-1})

これらを代入すると、以下のようになります。

pθ(x0:T)q(x1:Tx0)=p(xT)t=1Tpθ(xt1xt)t=1Tq(xtxt1)\frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} = \frac{p(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)}{\prod_{t=1}^T q(\mathbf{x}_t \vert \mathbf{x}_{t-1})}

ここで、積の最初(t=1t=1)とそれ以外(t2t \ge 2)を分けて書き出します。

pθ(x0:T)q(x1:Tx0)=p(xT)pθ(x0x1)t=2Tpθ(xt1xt)q(x1x0)t=2Tq(xtxt1)\frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} = \frac{p(\mathbf{x}_T) p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \prod_{t=2}^T p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)}{q(\mathbf{x}_1 \vert \mathbf{x}_0) \prod_{t=2}^T q(\mathbf{x}_t \vert \mathbf{x}_{t-1})}

Step 2: ベイズの定理による拡散過程の逆向き変形

拡散過程の遷移確率 q(xtxt1)q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) は、元のデータ x0\mathbf{x}_0 が与えられた状態でもマルコフ連鎖の性質を満たすため、q(xtxt1)=q(xtxt1,x0)q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) です。これをベイズの定理を用いて、条件が逆向きの確率に書き換えます。

q(xtxt1,x0)=q(xt1xt,x0)q(xtx0)q(xt1x0)q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) = \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)}

この変形を、分母の t2t \ge 2 の積 t=2Tq(xtxt1)\prod_{t=2}^T q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) に代入します。

t=2Tq(xtxt1)=t=2Tq(xt1xt,x0)q(xtx0)q(xt1x0)\prod_{t=2}^T q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \prod_{t=2}^T \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)}

この式の右側の分数の積を展開すると、隣り合う項の分子と分母の間で q(xsx0)q(\mathbf{x}_s \vert \mathbf{x}_0) が順番に相殺(テレスコーピング:Telescoping)されます。

t=2Tq(xtx0)q(xt1x0)=q(x2x0)q(x1x0)q(x3x0)q(x2x0)q(xTx0)q(xT1x0)=q(xTx0)q(x1x0)\prod_{t=2}^T \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} = \frac{q(\mathbf{x}_2 \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} \cdot \frac{q(\mathbf{x}_3 \vert \mathbf{x}_0)}{q(\mathbf{x}_2 \vert \mathbf{x}_0)} \dots \frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_{T-1} \vert \mathbf{x}_0)} = \frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)}

これによって、分母の積は以下のように整理されます。

t=2Tq(xtxt1)=(t=2Tq(xt1xt,x0))q(xTx0)q(x1x0)\prod_{t=2}^T q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \left( \prod_{t=2}^T q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \right) \cdot \frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)}

Step 3: 分数全体の再構築と約分

これを Step 1 の式の分母に代入します。すると、分母にある q(x1x0)q(\mathbf{x}_1 \vert \mathbf{x}_0) がちょうど相殺されて消えます。

pθ(x0:T)q(x1:Tx0)=p(xT)pθ(x0x1)t=2Tpθ(xt1xt)q(x1x0)(t=2Tq(xt1xt,x0))q(xTx0)q(x1x0)=p(xT)pθ(x0x1)t=2Tpθ(xt1xt)q(xTx0)t=2Tq(xt1xt,x0)\begin{aligned} \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} &= \frac{p(\mathbf{x}_T) p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \prod_{t=2}^T p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)}{q(\mathbf{x}_1 \vert \mathbf{x}_0) \cdot \left( \prod_{t=2}^T q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \right) \cdot \frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)}} \\ &= \frac{p(\mathbf{x}_T) p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \prod_{t=2}^T p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)}{q(\mathbf{x}_T \vert \mathbf{x}_0) \prod_{t=2}^T q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)} \end{aligned}

この式を、同じタイムステップの変数同士になるようにグループ分け(並び替え)します。

pθ(x0:T)q(x1:Tx0)=(p(xT)q(xTx0))(t=2Tpθ(xt1xt)q(xt1xt,x0))pθ(x0x1)\frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} = \left( \frac{p(\mathbf{x}_T)}{q(\mathbf{x}_T \vert \mathbf{x}_0)} \right) \cdot \left( \prod_{t=2}^T \frac{p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)} \right) \cdot p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)

Step 4: 負の対数(log-\log)と期待値 Eq\mathbb{E}_q の適用

この結果を負の対数尤度の上界(LVLBL_{\text{VLB}})の定義式に適用し、対数の性質を用いて「積を和(\sum)に展開」し、さらにマイナスの符号を「対数の分子と分母の反転」に利用します。

LVLB=Eq[logpθ(x0:T)q(x1:Tx0)]=Eq[log(p(xT)q(xTx0))t=2Tlog(pθ(xt1xt)q(xt1xt,x0))logpθ(x0x1)]=Eq[logq(xTx0)p(xT)]+t=2TEq[logq(xt1xt,x0)pθ(xt1xt)]Eq[logpθ(x0x1)]\begin{aligned} L_{\text{VLB}} &= \mathbb{E}_q \left[ -\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \right] \\ &= \mathbb{E}_q \left[ -\log \left( \frac{p(\mathbf{x}_T)}{q(\mathbf{x}_T \vert \mathbf{x}_0)} \right) - \sum_{t=2}^T \log \left( \frac{p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)} \right) - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \right] \\ &= \mathbb{E}_q \left[ \log \frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p(\mathbf{x}_T)} \right] + \sum_{t=2}^T \mathbb{E}_q \left[ \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)} \right] - \mathbb{E}_q \left[ \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \right] \end{aligned}

Step 5: KLダイバージェンスの定義への書き換え

最後に、KLダイバージェンスの定義 DKL(PQ)=EP[logP(x)Q(x)]D_{\text{KL}}(P \parallel Q) = \mathbb{E}_P \left[ \log \frac{P(x)}{Q(x)} \right] に基づいて各項を整理します。

期待値 Eq\mathbb{E}_q (同時分布 q(x0:T)q(\mathbf{x}_{0:T}))の下では、それぞれの項において該当しない確率変数は周辺化によって消去されるため、以下のように最終的に美しいKLダイバージェンスの期待値として分解された式が得られます。

LVLB=EqDKL(q(xTx0)p(xT))LT (学習パラメータ無)+t=2TEqDKL(q(xt1xt,x0)pθ(xt1xt))Lt1 (主となる拡散損失)Eq[logpθ(x0x1)]L0 (再構成誤差)L_{\text{VLB}} = \mathbb{E}_q \underbrace{D_{\text{KL}}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p(\mathbf{x}_T))}_{L_T \text{ (学習パラメータ無)}} + \sum_{t=2}^T \mathbb{E}_q \underbrace{D_{\text{KL}}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t))}_{L_{t-1} \text{ (主となる拡散損失)}} - \underbrace{\mathbb{E}_q \left[ \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \right]}_{L_0 \text{ (再構成誤差)}}

ここで中心となるのは、任意のステップ tt における Lt1L_{t-1} の項です。

条件付き真の逆過程 q(xt1xt,x0)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t, \mathbf{x}_0) の導出

本来、単なる q(xt1xt)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t) はデータ全体分布を知る必要があるため計算不能です。 しかし、元データ x0\mathbf{x}_0 で条件付けすることで、ベイズの定理により解析的に解くことが可能になります。

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t, \mathbf{x}_0) = q(\mathbf{x}_t \vert{} \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_0)}{q(\mathbf{x}_t \vert{} \mathbf{x}_0)}

ここで、拡散過程の一括サンプリング公式(αt=1βt,αˉt=s=1tαs\alpha_t = 1 - \beta_t, \bar{\alpha}_t = \prod_{s=1}^t \alpha_s)より:

  • q(xtxt1)=N(xt;αtxt1,(1αt)I)q(\mathbf{x}_t \vert{} \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{\alpha_t}\mathbf{x}_{t-1}, (1-\alpha_t)\mathbf{I})
  • q(xt1x0)=N(xt1;αˉt1x0,(1αˉt1)I)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \sqrt{\bar{\alpha}_{t-1}}\mathbf{x}_0, (1-\bar{\alpha}_{t-1})\mathbf{I})
  • q(xtx0)=N(xt;αˉtx0,(1αˉt)I)q(\mathbf{x}_t \vert{} \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_0, (1-\bar{\alpha}_t)\mathbf{I})

これらのガウス分布の確率密度関数を代入して指数部を整理すると、q(xt1xt,x0)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t, \mathbf{x}_0) もまたガウス分布 N(xt1;μ~t(xt,x0),β~tI)\mathcal{N}(\mathbf{x}_{t-1}; \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0), \tilde{\beta}_t \mathbf{I}) になることが分かります。 その真の平均 μ~t\tilde{\boldsymbol{\mu}}_t は以下のように定まります。

μ~t(xt,x0)=αˉt1βt1αˉtx0+αt(1αˉt1)1αˉtxt\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

ガウス分布同士の KL ダイバージェンスの計算

生成モデルの逆過程 pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t))p_\theta(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \Sigma_\theta(\mathbf{x}_t, t)) と置きます。(※DDPMでは簡単のため分散 Σθ\Sigma_\theta は固定値 σt2I\sigma_t^2 \mathbf{I} とします)
2つの多変量ガウス分布間の KL ダイバージェンスの公式を適用すると、定数項を除き、「平均の二乗誤差」に帰着します。

Lt1=DKL(q(xt1xt,x0)pθ(xt1xt))=Eq[12σt2μ~t(xt,x0)μθ(xt,t)2]L_{t-1} = D_{\text{KL}}(q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t)) = \mathbb{E}_q \left[ \frac{1}{2\sigma_t^2} \left\Vert{} \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) \right\Vert{}^2 \right]

つまり、モデル pθp_\theta の学習目標は「目標の平均 μ~t\tilde{\boldsymbol{\mu}}_t を予測すること」となります。

ノイズ予測 ϵθ\boldsymbol{\epsilon}_\theta への変換と最終形の導出

ここで、順過程のサンプリング式 xt(x0,ϵ)=αˉtx0+1αˉtϵ\mathbf{x}_t(\mathbf{x}_0, \boldsymbol{\epsilon}) = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon} (ただし ϵN(0,I)\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}))を利用して、x0\mathbf{x}_0xt\mathbf{x}_tϵ\boldsymbol{\epsilon} で書き換えます。

x0=1αˉt(xt1αˉtϵ)\mathbf{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}\left( \mathbf{x}_t - \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon} \right)

これを先ほど求めた真の平均 μ~t\tilde{\boldsymbol{\mu}}_t の式に代入して整理すると、驚くほど綺麗な形に代入・消去されます。

μ~t(xt,x0)=1αt(xtβt1αˉtϵ)\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 も同様の形式にパラメータ化(再定義)します。

μθ(xt,t)=1αt(xtβt1αˉtϵθ(xt,t))\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)

この μ~t\tilde{\boldsymbol{\mu}}_tμθ\boldsymbol{\mu}_\thetaLt1L_{t-1} の差分の式に代入します。

μ~tμθ=βtαt(1αˉt)(ϵϵθ(xt,t))\tilde{\boldsymbol{\mu}}_t - \boldsymbol{\mu}_\theta = \frac{\beta_t}{\sqrt{\alpha_t(1-\bar{\alpha}_t)}} \left( \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \right)

これを二乗誤差の式に書き戻すと、以下のようになります。

Lt1=Et,x0,ϵ[βt22σt2αt(1αˉt)ϵϵθ(αˉtx0+1αˉtϵ,t)2]L_{t-1} = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1-\bar{\alpha}_t)} \left\Vert{} \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}, t) \right\Vert{}^2 \right]

結論(簡略化損失関数 LsimpleL_{\text{simple}}

DDPMの論文「Denoising Diffusion Probabilistic Models」(Ho et al., 2020)では、上記損失関数の前についている複雑な係数(重み付け)を 1 に簡略化した方が、サンプル生成品質(FIDスコア)が大幅に向上することが実験的に示されました。
こうして最終的に得られたのが、提示されたLsimpleL_{\text{simple}} です。

Lsimple(θ):=Et,x0,ϵ[ϵϵθ(αˉtx0+1αˉtϵ,t)2]L_{\text{simple}}(\theta) := \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \left\Vert{} \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}, t) \right\Vert{}^2 \right]

これにより、確率論の複雑なKLダイバージェンスの計算は、 「実際に加えたガウスノイズ ϵ\boldsymbol{\epsilon} と U-Net 等のAIが予測したノイズ ϵθ\boldsymbol{\epsilon}_\theta の単純な二乗誤差(MSE)」 に完全に一致することが証明されます。

まとめ

本記事では、DDPM(Denoising Diffusion Probabilistic Models)の学習において、一見複雑に見える確率論的なアプローチが、なぜ最終的に極めてシンプルな「ノイズ予測の二乗誤差(MSE)」へと帰着するのか、その数学的な導出プロセスを解説しました。

具体的には、以下のステップに沿って数式を展開し、その本質を紐解きました。

  • 変分下界(ELBO)の導入と分解: 直接計算が困難な負の対数尤度を最小化するために、イェンゼンの不等式を用いて変分下界 LVLBL_{\text{VLB}} を導出し、各タイムステップにおけるKLダイバージェンスの和に分解しました。
  • 条件付き逆過程の解析的表示: 元のデータ x0\mathbf{x}_0 で条件付けることで、通常は計算できない逆過程 q(xt1xt,x0)q(\mathbf{x}_{t-1} \vert{} \mathbf{x}_t, \mathbf{x}_0) をガウス分布として解析的に導出し、その真の平均 μ~t\tilde{\boldsymbol{\mu}}_t を明らかにしました。
  • ノイズ予測へのパラメータ化: ガウス分布間のKLダイバージェンスが平均の二乗誤差に帰着することを示した上で、一括サンプリング公式を用いて「平均の予測」を「ノイズ ϵ\boldsymbol{\epsilon} の予測」へと変換しました。
  • 簡略化損失関数 LsimpleL_{\text{simple}} の導出: 複雑な係数を取り除いた LsimpleL_{\text{simple}} を採用することで、画像生成品質(FIDスコア)が向上する理由、精度よく単純なMSE損失に一致することを確認しました。

拡散モデルは、一見すると難解な確率微分方程式や数式に満ちていますが、その根底にはこのように美しく合理的な数式の展開が存在しています。この基盤となるメカニズムを理解することは、Stable Diffusionなどのより発展的な画像生成AIや、応用モデルを学ぶ上での非常に強力な武器になります。ぜひ本記事の数式展開を手元で動かしながら、そのエレガントな構造を深く体感してみてください。

本記事の文章・構成の一部に生成AIを使用しています。