Preface

　このシリーズでは、Diffusionモデルについて学習する時にノート代わりに記事を書いていく。これはその第二弾で以下の第一弾の続き。手始めにめちゃ分かりやすいと巷で話題の（そして実際分かりやすかった）以下のDiffusionモデルの解説論文を少しずつ翻訳していき、脳に焼き付けていく。後々より詳しい解説とか、自分でJAXで実装とかができたらいいなと思っている。

izmyon.hatenablog.com

Preface
Understanding Diffusion Models: A Unified Perspective
- Variational Diffusion Model
- Reference

arxiv.org

Calvin Luo: Understanding Diffusion Models: A Unified Perspective, arXiv: 2208:11970, doi: 10.48550/ARXIV.2208.11970

©Calvin Luo, Originally posted in arXiv(https://arxiv.org/abs/2208.11970), 25 Aug 2022

License: Creative Commons Attribution 4.0 International (CC-BY)

以下は、原文の一部を翻訳したもので、以下の図はそこから引用したものです。

The following is the translation of part of the original content and the figures below are retrieved from it.

Understanding Diffusion Models: A Unified Perspective

Variational Diffusion Model

　変分拡散モデル (VDM: Variational Diffusion Model) [4, 5, 6] を考える最も簡単な方法は、単純に以下の三つの重要な制約を持つマルコフ型階層変分オートエンコーダとして考えることである。

潜在次元はデータ次元と正確に等しい。
各タイムステップにおける隠れエンコーダの構造は学習されず、線形ガウスモデルとして予め定義されている。言い換えれば、それは前のタイムステップの出力を中心とするガウス分布である。
隠れエンコーダのガウスパラメータは、最終タイムステップTにおける潜在変数 $x_T$ の分布が標準ガウス分布になるように時間的に変化する。

さらに、標準的なマルコフ型階層変分オートエンコーダにおいける階層的遷移間のマルコフ特性は、明示的に維持するものとする。これらの仮定が意味するところを拡大してみよう。最初の仮定から、多少の表記法の乱用はあるが、真のデータサンプルと潜在変数の両方を $x_t$ で表すものとする。ここで、 $t = 0$ は真のデータサンプルを表し、 $t∈[1,T$ ]はインデックス $t$ の階層を持つ対応する潜在変数を表す。VDMの事後分布は、MHVAEの事後分布（式(24)）と同じだが、今は次のように書き換えることができる。

$\displaystyle q(x_{1:T} | x_0 ) = \prod_{t=1}^T q( x_t | x_{t-1} ) \tag{30}$

第二の仮定から、エンコーダの各潜在変数の分布は、その前のHVAEの出力を中心とするガウス分布であることがわかる。マルコフ型HVAEとは異なり、各タイムステップ $t$ におけるエンコーダの構造は学習されない。それは線形ガウスモデルとして固定され、平均と標準偏差はハイパーパラメータとしてあらかじめ設定されるか[5]、またはパラメータとして学習される[6]。ここでは、平均 $\boldsymbol{μ_ t} (x_ t) = \sqrt{ α_ t } x_ {t-1}$ 、分散 $\sum_t (x_t) = (1 - α_t) \boldsymbol{\rm{I}}$ のガウシアンエンコーダとしてパラメータ化する。ここで係数の形式は，潜在変数の分散が同じようなスケールにとどまるように選択される。言い換えれば、このエンコーディングの手順は”分散保存的（variance-preserving）”である。ここで、ガウス分布のパラメータ化は他の方法も可能であり、同様の導出ができることに注意。このやり方では、 $α_t$ は（潜在的に学習可能な）係数であり、階層的深さ $t$ に応じて柔軟に変化することである。数学的には、エンコーダの遷移は次のように表される。

$\displaystyle q(x_t | x_{t-1} ) = \mathcal{N}(x_t; \sqrt{α_t} x_{t-1}, (1 - α_t) \boldsymbol{\rm{I}} ) \tag{31}$

第三の仮定から、 $α_t$ は固定された、あるいは学習可能なスケジュールに従って時間と共に進化し、最終的な潜在変数 $p(x_T)$ の分布が標準ガウスとなることがわかる。そして、マルコフ型HVAEの同時分布（式(23)）を更新して、VDMの同時分布を次のように書くことができる。

$\displaystyle \begin{align} p(x_{0:T}) &= p(x_T) \prod_{t=1}^T p_{\theta}( x_{t-1} | x_t ) \tag{32} \\ where, \\ p(x_T) &= \mathcal{N}(x_T ; \boldsymbol{0, \rm{I} } ) \tag{33} \end{align}$

つまり、これらの制約が意図しているのは、時間経過とともに入力画像を徐々にノイズに変えていくということであり、より詳しく言うと、ガウスノイズを加えていくことで入力画像を徐々に崩壊させ、最終的には純粋なガウスノイズと完全に同じにさせるということである。この過程を視覚的に表現したのが図3である。

図3：変分拡散モデルの視覚的表現。 $x_0$ は自然画像などの真のデータ観測、 $x_T$ は純粋なガウスノイズ、 $x_t$ は $x_0$ にノイズを加えた $x_T$ との中間的な状態の画像である。各 $q(x_t | x_{t-1})$ は直前の状態の出力を平均とするガウス分布としてモデル化される。

図3：変分拡散モデルの視覚的表現。 $x_0$ は自然画像などの真のデータ観測、 $x_T$ は純粋なガウスノイズ、 $x_t$ は $x_0$ にノイズを加えた $x_T$ との中間的な状態の画像である。各 $q(x_t | x_{t-1})$ は直前の状態の出力を平均とするガウス分布としてモデル化される。

なお、エンコーダの分布 $q(x_t | x_{t-1})$ は、各タイムステップにおいて、定義済みの平均と分散のパラメータを持つガウシアンとして完全にモデル化されるため、もはや $φ$ によってパラメータ化されないことに注意する。したがって、VDMでは、新しいデータをシミュレートできるように、条件 $p_{\theta} (x_{t-1} | x_t )$ を学習することにのみ興味がある。VDM を最適化した後のサンプリング手順は、ガウシアンノイズを $p(x_T)$ からサンプリングし、繰り返しノイズ除去遷移 $p_{\theta} (x_{t-1} | x_t)$ をTステップ通し、新しい $x_0$ を生成する。

他のHVAEと同様に、VDMはELBOを最大化することで最適化でき、次のように導出される。

ELBOの導出形式は、その個々の構成要素で解釈することができる。

$\mathbb{E}_{q(x_0 | x_1)} \left[ \log p_{\theta} ( x_0 | x_1 ) \right]$ は再構成項と解釈され、一段階目の潜在変数が与えられたときの元のデータサンプルの対数尤度を予測する。この項はバニラVAEにも現れ、同様に学習させることができる。
$\mathbb{E}_{q(x_{T-1} | x_0 )} \left[ D_{KL} (q(x_T | x_{T-1} ) || p(x_T) ) \right]$ は事前マッチング項であり、最終的な潜在分布がガウス事前分布と一致するときに最小化される。この項は、学習可能なパラメータを持たないため、最適化の必要がない。さらに、最終的な分布がガウス分布となるように、十分に大きな $T$ を仮定しているため、この項は実質的にゼロになる。
$\mathbb{E}_{q \left( x_{t-1}, x_{t+1} | x_0 \right) } \left[ D_{KL} \left( q(x_t | x_{t-1} ) || p_{\theta}(x_t | x_{t+1} ) \right) \right]$ は整合項であり、 $x_t$ の分布が前方、後方の両プロセスで整合するように努める。つまり、各中間タイムステップごとに、ノイズの多い画像からのノイズ除去ステップは、対応するよりクリーンな画像からのノイズ加算ステップと一致すべきであり、これはKLダイバージェンスによって数学的に反映される。この項は、式(31)で定義されているように、 $p_{\theta} (x_t | x_{t+1})$ がガウス分布 $q(x_t| x_{t-1})$ と一致するように学習されるとき最小となる。

ELBOのこの解釈は、図4に視覚的に描かれている。すべてのタイムステップ $t$ に対して最適化しなければならないため、VDMを最適化するコストは、主に第3項によって支配される。

図4：最初の導出のもとで、VDMは、各潜在変数 $x_t$ について、その後の潜在変数からの事後分布 $p_{\theta} (x_t | x_{t+1})$ が、その前の潜在変数からのガウシアンによる劣化 $q( x_t |x_{t-1} )$ と一致するように最適化することができる。この図では、各潜在変数 $x_t$ について、ピンクと緑の矢印で示される分布の差を最小化している。

図4：最初の導出のもとで、VDMは、各潜在変数 $x_t$ について、その後の潜在変数からの事後分布 $p_{\theta} (x_t | x_{t+1})$ が、その前の潜在変数からのガウシアンによる劣化 $q( x_t |x_{t-1} )$ と一致するように最適化することができる。この図では、各潜在変数 $x_t$ について、ピンクと緑の矢印で示される分布の差を最小化している。

この導出の下では、ELBOのすべての項は期待値として計算され、したがって、モンテカルロ推定を使用して近似することができる。しかし、今導出した項を用いてELBOを実際に最適化すると、最適とは言えないかもしれない。整合項は、すべてのタイムステップごとに2つの確率変数 $\left\{ x_{t-1}, x_{t+1} \right\}$ に対する期待値として計算されるので、モンテカルロ推定の分散は、タイムステップごとに1つの確率変数のみを用いて推定される項よりも大きくなる可能性がある。 $T-1$ 個の整合項の合計で計算されるため、最終的なELBOの推定値は $T$ の値が大きいと分散が大きくなる可能性がある。

その代わりに、各項が一度に1つの確率変数に対する期待値として計算されるELBOの形式を導出することを試みましょう。重要なのは、エンコーダの遷移を $q(x_t | x_{t-1}) = q(x_t | x_{t-1}, x_0)$ と書き換えることで、マルコフ特性により、余分な条件項は不要になることである。そして、ベイズ則に従って、各遷移を次のように書き換えることができる。

$\displaystyle q\left( x_t \mid x_{t-1}, x_{0} \right) = \frac{ q \left( x_{t-1} \mid x_{t}, x_{0}\right) q\left( x_{t} \mid x_{0} \right) }{q \left( x_{t-1} \mid x_{0} \right)} \tag{46}$

この新しい式を使って、式(37)のELBOから再開して導出を試みることができる。

以上より、低い分散で推定できるELBOの解釈を導き出すことに成功し、各項が一度に最大でも一つの確率変数の期待値として計算されることがわかる。この定式化は、個々の項を調べることでエレガントに解釈することができる。

$\mathbb{E}_{q(x_1 | x_0)} \left[ \log p_{\theta} \left( x_0 | x_1 \right) \right]$ は再構成項と解釈できる。バニラ VAEのELBO における類似の項と同様、この項はモンテカルロ推定を用いて近似および最適化できる。
$D_{KL} \left( q \left( x_T | x_0 \right) ‖ p ( x_T) \right)$ はノイズ化された入力の最後の分布が標準ガウス事前分布にどれくらい近いかを示す。学習可能なパラメータは無く、ここでの仮定の下では０に等しい。
$\mathbb{E}_{q \left( x_t | x_0 \right) } \left[ D_{KL} ( q \left( x_{t-1} | x_t, x_0 \right) \right]$ は"ノイズ除去マッチング項"である。望みのノイズ除去遷移ステップ $p_{\theta} ( x_{t-1} | x_t )$ を、ground-truthのノイズ除去遷移ステップ $q \left( x_{t-1} | x_t, x_0 \right)$ の扱いやすい近似として学習する。 $q \left( x_{t-1} | x_t, x_0 \right)$ はノイズの多い画像 $x_t$ をどのようにノイズ除去するかを定義し、最終的に完全にノイズ除去された画像 $x_0$ がどうあるべきかを知っているため、ground-truth信号として機能することができる。したがって、この項は、KLダイバージェンスによって測定されているように、2つのノイズ除去ステップができるだけ一致するときに最小化される。

余談だが、二つのELBOの導出過程（式(45)および式(58)）において、マルコフ仮定のみが用いられており、その結果、これらの式は任意のマルコフ型HVAEに対して成り立つことがわかる。さらに、 $T=1$ とすると、二つのVDMのELBOの解釈は、いずれも式(19)で書かれるように、バニラVAEのELBO方程式を正確に再現する。

このELBOの導出では、最適化コストの大部分が再び総和項にあり、再構成項に対して支配的であった。各KLダイバージェンスの項 $D_{KL} \left( q \left( x_{t-1} \mid x_t, x_0 \right) ‖ p_{\theta} \left( x_{t-1}｜x_t \right) \right)$ は、エンコーダを同時に学習するという複雑さが加わるため、任意に複雑なマルコフ型HVAEでは任意の事後分布に対して最小化することが難しいが、VDMではガウス推移仮定を利用して最適化を扱いやすくすることが可能である。ベイズの定理により、以下のようになる。

$\displaystyle q \left( x_{t-1} \mid x_{t}, x_{0} \right) = \frac{ q \left( x_{t} \mid x_{t-1}, x_{0} \right) q \left( x_{t-1} \mid x_{0} \right) } { q \left( x_{t} \mid x_{0} \right)}$

エンコーダの遷移に関する仮定（式(31)）から、 $q( x_t | x_{t-1}, x_0 ) = q( x_t | x_{t-1} ) = N \left( x_t ; \sqrt{α_t} x_{t-1}, (1-α_t) \boldsymbol{\rm{I}} \right)$ が既に分かっているので、後は $q( x_t | x_0 )$ と $q(x_{t-1} | x_0 )$ の形を導けば良い。幸運なことに、VDMのエンコーダ遷移は直線ガウスモデルであるという事実を利用して、これも扱いやすくすることができる。再パラメータ化トリックの下で、サンプル $x_t \sim q ( x_t | x_{t-1} )$ は以下のように書き換えることができる。

$\displaystyle x_t = \sqrt{α_t} x_{t-1} + \sqrt{1 - α_t} \epsilon \quad with \quad \epsilon \sim \mathcal{N}(\boldsymbol{\epsilon ; 0, \rm{I}} ) \tag{59}$

同様に、サンプル $x_{t-1} \sim q \left( x_{t-1} | x_{t-2} \right)$ は次のように書き換えることができることがわかる。

$\displaystyle x_t = \sqrt{α_{t-1}} x_{t-2} + \sqrt{1 - α_{t-1}} \epsilon \quad with \quad \epsilon \sim \mathcal{N}(\boldsymbol{\epsilon ; 0, \rm{I}} ) \tag{60}$

図5：VDMを小さい分散で最適化するための代替方法を図示した。ベイズの定理を用いてground-truthのノイズ除去ステップ $q(x_{t-1} | x_t, x_0 )$ を計算し、そのKLダイバージェンスを近似ノイズ除去ステップ $p_{\theta} (x_{t-1} | x_t )$ で最小化する。ここでは再び、マッチングさせる分布を緑色の矢印とピンク色の矢印で視覚的に表している。ここで、全体像を正確に書くと、各ピンクの矢印は、条件付けの項でもあるため、 $x_0$ からも派生しているはずであるが、ここでは割愛した。

図5：VDMを小さい分散で最適化するための代替方法を図示した。ベイズの定理を用いてground-truthのノイズ除去ステップ $q(x_{t-1} | x_t, x_0 )$ を計算し、そのKLダイバージェンスを近似ノイズ除去ステップ $p_{\theta} (x_{t-1} | x_t )$ で最小化する。ここでは再び、マッチングさせる分布を緑色の矢印とピンク色の矢印で視覚的に表している。ここで、全体像を正確に書くと、各ピンクの矢印は、条件付けの項でもあるため、 $x_0$ からも派生しているはずであるが、ここでは割愛した。

そして、 $q( x_t | x_0 )$ の形式は再パラメータ化のトリックを繰り返し適用することで再帰的に導出することができる。ここで、 $2T$ のランダムノイズ変数 $\left\{ \epsilon_t^{*}, \epsilon_{t=0} \right\}_{t=0}^T \sim i.i.d \quad \mathcal{N} \left( \boldsymbol{\epsilon ; 0, \rm{I}} \right)$ を取得できたとする。そして、任意のサンプル $x_t \sim q( x_t | x_0)$ に対して、以下のように書き換えることができる。

ここで、式(64)では2つの独立なガウス確率変数の和がいまだガウシアンであり、平均は2つの平均の和、分散は2つの分散の和であることを利用している。 $\sqrt{1- α_t} \epsilon_{t-1}^{*}$ をガウシアン $\mathcal{N} \left( \boldsymbol{0}, (1-α_t) \boldsymbol{\rm{I}} \right)$ からのサンプル、そして $\sqrt{ α_t - α_t α_{t-1}} \epsilon_{t-2}^{*}$ をガウシアン $\mathcal{N} \left( \boldsymbol{0}, (α_t - α_t α_{t-1} ) \boldsymbol{\rm{I}} \right)$ の標本とすると、それらの和はガウシアン $\mathcal{N}\left( \boldsymbol{0}, (1-α_t + α_t - α_t α_{t-1}) \boldsymbol{\rm{I}} \right) =\mathcal{N} \left( \boldsymbol{0}, (1-α_t α_{t-1}) \boldsymbol{\rm{I}} \right)$ からサンプリングした確率変数として扱うことができる。この分布からのサンプルは、再パラメータ化のトリックを使って、 $\sqrt{1-α_t α_{t-1}} \epsilon_{t-2}$ として、式(66)のように表現される。

したがって、ガウシアン形式 $q(x_t | x_0 )$ を導出しました。この導出は $q(x_{t-1} | x_0)$ を記述するガウス分布のパラメータを得るために修正することができる。ここで、 $q(x_t | x_0 )$ と $q(x_{t-1} | x_0 )$ の両方の形式を知っているので、ベイズ則展開に代入して $q \left( x_{t-1} | x_t, x_0 \right)$ の形式の計算に進むことができる。

ここで式(75)の $C(x_t, x_0)$ は $x_t$ 、 $x_0$ と $α$ の値のみの組み合わせとして計算された $x_{t-1}$ のそれぞれに関する定数項であり、この項は式(84)に暗黙的に返されて平方完成される。

したがって、各ステップにおいて $x_{t-1} \sim q( x_{t-1} | x_t, x_0)$ は $x_t$ と $x_0$ の関数である平均 $μ_q (x_t, x_0)$ と係数 $α$ の関数である分散 $\sum_q (t)$ で正規分布することが示された。これらの係数 $α$ は既知で各タイムステップで固定されており、ハイパーパラメータとしてモデル化された場合は恒常的に固定化されるか、またはモデル化しようとするネットワークの現在の推論出力として扱われる。式(84)に従うと、分散方程式を ${\sum}_q (t) = {\sigma}_q^{2} (t)$ と書き直すことができ、以下が成り立つ。

$\displaystyle \sigma_q^2 (t) = \frac {(1-α_t)(1-α_{t-1})}{1-α_t} \tag{85}$

近似ノイズ除去遷移ステップ $p_{\theta} ( x_{t-1} | x_t )$ をground-truthのノイズ除去遷移ステップ $q_(x_{t-1} | x_t, x_0 )$ にできるだけ近づけるために、ガウシアンとしてモデル化することもできる。さらに、すべてのα項は各タイムステップで凍結されることが知られているため、近似されたノイズ除去遷移ステップの分散も $\sum_ q (t) = \sigma_ q ^2 (t) \boldsymbol{\rm{I}}$ となるように直ちに構築することができる。しかし、 $p_{\theta} (x_{t-1} |x_t )$ は $x_0$ を条件としないので、その平均 $μ_{\theta} ( x_t, t )$ を $x_t$ の関数としてパラメータ化しなければならない。

ここで、2つのガウス分布の間のKL収束は次の通りであることを思い出してほしい。

$\displaystyle D_{KL} \left( \mathcal{N} (x; μ_x, {\sum}_x ) || \mathcal{N} ( y; μ_y, {\sum}_y ) \right) = \frac{1}{2} \left[ \log \frac{|{\sum}_y|}{|{\sum}_x|} - d + tr({\sum}_y^{-1} {\sum}_x ) + (μ_y - μ_x ) ^T {\sum}_y^{-1} ( μ_y - μ_x ) \right] \tag{86}$

この場合、2つのガウス分布の分散を正確に一致させることができるので、KLダイバージェンス項の最適化は、2つの分布の平均の差を最小化させるために減少させることになる。

ここで、 $μ_q$ は $μ_q(x_t, x_0 )$ の略記、 $μ_{\theta}$ は $μ_{\theta} (x_t, t)$ の略記として簡略化して書いている。言い換えれば、我々は、 $μ_q(x_t，x_0)$ に一致する $μ_{\theta}(x_t，t)$ を最適化したいのであって、我々の導き出した式(84)から、次のような形をとる。

$\displaystyle \boldsymbol{μ}_{q} (x_t, x_0) = \frac{ \sqrt{ α_t }(1- \hat{α_{t-1}} ) x_t + \sqrt{ \hat{α_{t-1}} } (1-α_t)x_0}{1- \hat{α_{t}} } \tag{93}$

$μ_{\theta} (x_t, t)$ も $x_t$ を条件としているので、以下の形に設定することで $μ_q(x_t，x_0)$ に近い形で一致させることができる。

$\displaystyle \boldsymbol{μ}_{\theta} (x_t, t) = \frac{\sqrt{α_t}(1- \hat{α_{t-1}} ) x_t + \sqrt{\hat{α_{t-1}}} (1-α_t) \hat{x_{\theta}}(x_t, t) }{1- \hat{α_{t}} } \tag{94}$

ここで、 $\hat{x_{\theta}} (x_t, t)$ は、ノイズの多い画像 $x_t$ と時間インデックス $t$ から $x_0$ を予測しようとするニューラルネットワークによってパラメータ化される。そして、最適化問題は次のように単純化される。

したがって、VDMを最適化することは、任意にノイズ化された画像から元のground-truthとなる画像を予測するニューラルネットワークの学習に帰結する[5]。さらに、すべてのノイズレベルにわたって、我々の導き出したELBO目的関数（式(58)）の総和項を最小化することは、すべてのタイムステップにわたって、この式を最小化することによって近似できる。

$\displaystyle \arg \min_{\theta} \mathbb{E}_{t \sim U \left[ 2, T \right] } \left[ \mathbb{E}_{q(x_t | x_0 )} \left[ D_{KL} \left( q ( x_{t-1} | x_t, x_0 ) || p_{\theta} ( x_{t-1} | x_t ) \right) \right] \right] \tag{100}$

これは次に、時間ステップにわたる確率的なサンプルを使うことで最適化される。

Reference

[4] ] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256–2265.PMLR, 2015.

[5] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.

[6] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696–21707, 2021.