Preface

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

Preface
Understanding Diffusion Models: A Unified Perspective

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

Introduction: Generative Models

　注目する分布から観測されたサンプル $x$ が与えられたとき、生成モデルの目標は、その真のデータ分布 $p(x)$ をモデルを学習することである。一度学習すれば、近似モデルから新しいサンプルを自由に生成することができる。さらに、ある定式化の下では、学習したモデルを用いて、観測データやサンプリングされたデータの尤度を評価することも可能である。

　ここでは、現在の文献でよく知られているいくつかの方向性を、高いレベルで簡単に紹介する。敵対的生成ネットワーク（GAN）は、複雑な分布のサンプリング手順をモデル化し、敵対的な方法で学習させる。"尤度ベース"と呼ばれる生成モデルの別のクラスは、観測されたデータサンプルに高い尤度を割り当てるモデルを学習しようとするものである。これには自己回帰モデル、正規化フロー、変分オートエンコーダ（VAE）などがある。また、類似の手法としてエネルギーベースモデリングがあり、これは分布を任意に柔軟なエネルギー関数として学習し、それを正規化するものである。

　スコアベース生成モデルは、エネルギー関数そのものをモデル化するのではなく、エネルギーベースモデルの"スコア"をニューラルネットワークとして学習するもので、非常に関連性の高いモデルである。この研究では、尤度ベースとスコアベースの両方の解釈を持つ拡散モデル（Diffusion Model）について検討し、レビューする。このようなモデルの背後にある数学は、誰もが拡散モデルとは何か、どのように機能するかを理解できることを目的として、非常に詳細に紹介する。

Background: ELBO, VAE, and Hierarchical VAE

　多くのモダリティでは、観測されたデータは、関連する見えない潜在変数（確率変数 $z$ で表すことができる）によって表される、あるいは生成されると考えることができる。この考えを最も良く直観的に表現するのは、プラトンの洞窟の寓話である。この寓話では、ある人々が生涯洞窟の中に鎖でつながれ、目の前の壁に映し出される二次元の影しか見ることができず、それは火の前を通る目に見えない三次元の物体によって生成されるとされている。そのような人たちにとって、自分が見ているものはすべて、実は自分には決して見ることのできない高次元の抽象的な概念によって決定されているのである。

　同様に、私たちが現実の世界で出会う物体も、例えば、色や大きさ、形といった抽象的な性質が内包された高次元の表現（representation）の関数として生成されている可能性がある。洞窟につながれた人が観測しているものが実は三次元の物体の二次元の投影であったように、私たちが観察しているものは、その抽象概念の三次元的な投影、あるいは実体、として解釈することができる。洞窟につながれた人たちは、隠されたものを決して見ることはできず、あるいは完全に理解することもできないが、それでも私たちが観測したデータを記述する潜在的表現を近似するのと同じ方法で、それについて推論を巡らせることはできる。

　プラトンの寓話は、潜在変数が観測を決定する潜在的に観測不可能な表現であるという考え方を示しているが、この類推の注意点は、生成モデルにおいては、一般的に高次元の潜在的表現よりも低次元の潜在的表現を学ぼうとすることである。これは、観測よりも高い次元の表現を学習しようとしても、強い事前分布を持たなければ実りのない努力になるからである。一方、低次元の潜在的表現の学習は圧縮の一形態とみなすこともでき、観測を記述する意味的（Semantically）に有意義な構造を発見することができる可能性がある。

Evidence Lower Bound

　数学的には、潜在変数と我々が観測するデータは、同時分布 $p(x, z)$ によってモデル化されると想像することができる。”尤度ベース”と呼ばれる生成モデリングの1つのアプローチは、すべての観測された $x$ の尤度 $p(x)$ を最大化するモデルを学習することであることを思い出してほしい。純粋な観測データの尤度 $p(x)$ を復元するために、この同時分布を操作する方法が2つある。まず、以下のように潜在変数 $z$ を明示的に周辺化する方法がある。

$\displaystyle p(x) = \int p(x, z) dz \tag{1}$

もう一つの方法は、確率の連鎖法則を用いて以下のように求めることである。

$\displaystyle p(x) = \frac {p(x,z)}{p(z|x)} \tag{2}$

尤度 $p(x)$ を直接計算し最大化することは、式(1)においては、すべての潜在変数 $z$ を積分する必要があり、複雑なモデルでは困難であり、また式(2)においても、ground-truthであるエンコーダ $p(z|x)$ を求める必要があるため、困難である。しかし、この2つの式を用いると、その名の通りエビデンスの下界であるEvidence Lower Bound （ELBO）という項を導き出すことができる。エビデンスはこの場合、観測データの対数尤度として定量化される。そして、ELBOを最大化することは、潜在変数のモデルを最適化するための代理目的となる。ELBOが強力にチューニングされ、完全に最適化された最良の場合、それはエビデンスと厳密に等価となる。形式的には、ELBOは以下の式で表される。

$\displaystyle \mathbb E_{q_{φ}(z|x)} \left[ \frac {p(x,z)}{q_{φ}(z|x)} \right] \tag{3}$

エビデンスとの関係を明示的に表したい場合は、以下のように書ける。

$\displaystyle \log p(x) \geq \mathbb E_{q_{φ}(z|x)} \left[ \frac {p(x,z)}{q_{φ}(z|x)} \right] \tag{4}$

ここで、 $q_{φ}(z|x)$ は，パラメータ $φ$ を持つ柔軟な近似変分であり、我々が最適化しようとする分布である。直感的には，与えられた観測 $x$ の潜在変数上の真の分布を推定するために学習されるパラメータを持つモデルとして考えることができる。次節で変分オートエンコーダについて探る際に見るように、ELBO を最大化するためにパラメータ $φ$ をチューニングして下限を大きくすると、真のデータ分布をモデル化し、そこからサンプリングするために使用できる要素にアクセスできるようになり、生成モデルを学習することができる。ここで、なぜELBOが最大化したい目的関数であるのかについて、もう少し掘り下げて考える。まず、式 (1)を用いて、ELBO を導出する。

$\displaystyle \begin{align} \log p(x) &= \log \int p(x, z) dz &(式(1)を用いた) \tag{5} \\ &= \log \int \frac{p(x, z) q_{φ}(z|x)}{ q_{φ}(z|x)}dz &(1=\frac {q_{φ}(z|x)}{ q_{φ}(z|x)}を掛けた) \tag{6} \\ &= \log \mathbb E_{q_{φ}(z|x)} \left[ \frac {p(x,z)}{q_{φ}(z|x)} \right] &(期待値の定義) \tag{7} \\ & \geq \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)}{q_{φ}(z|x)} \right] &(Jensenの不等式を用いた) \tag{8} \end{align}$

この導出では、Jensenの不等式を適用することにより、直接的に下界に到達している。しかし、これは、実際に何が起こっているかについての有益な情報をあまり与えていない。決定的なのは、この証明が、ELBOが実際にエビデンスの下界であるのは何故か、という問いに対し、正確な直感を与えないことであり、Jensenの不等式がそれを遠ざけてしまうことである。さらに、ELBOが本当にデータの下界であることを単に知るだけでは、何故それを目的関数として最大化したいのかがよく分からない。エビデンスと ELBO の関係をよりよく理解するために、今度は式(2)を使って別の導出を行う。

$\displaystyle \begin{align} \log p(x) &= \log p(x) \int q_{φ}(z|x) dz &( 1= \int q_{φ}(z|x) dz を掛けた) \tag{9} \\ &= \int q_{φ}(z|x) (\log p(x)) dz &(エビデンスを積分の中に入れた) \tag{10} \\ &= \mathbb E_{q_{φ}(z|x)} \left[ \log p(x) \right] &(期待値の定義) \tag{11} \\ &= \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)}{p(z|x)} \right] &(式(2)を用いた) \tag{12} \\ &= \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)q_{φ}(z|x)}{p(z|x)q_{φ}(z|x)} \right] &(1=\frac {q_{φ}(z|x)}{ q_{φ}(z|x)}を掛けた) \tag{13} \\ &= \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)}{q_{φ}(z|x)} \right] + \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {q_{φ}(z|x)}{p(z|x)} \right] &(期待値を分割した) \tag{14} \\ &= \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)}{q_{φ}(z|x)} \right] + D_{\rm{KL}} ( q_{φ}(z|x) || p(z|x) ) &(\rm{KL}ダイバージェンスの定義) \tag{15} \\ & \geq \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)}{q_{φ}(z|x)} \right] &(\rm{KL}ダイバージェンスは常に\geq 0) \tag{16} \end{align}$

この導出より、式(15)から、エビデンスがELBOに近似事後分布 $q_{φ}(z|x)$ と真の事後分布 $p(z|x)$ 間のKLダイバージェンスを加えたものであることは明らかである。実は、最初の導出の式(8)でJensenの不等式によって魔法のように取り除かれていたのは、このKLダイバージェンスの項なのである。この項を理解することは、ELBOとエビデンスの関係だけでなく、ELBOを最適化することが全く適切な目的である理由を理解するための鍵となる。

まず、我々はここでELBOが本当に下界である理由がわかった。エビデンスとELBOの差は厳密に非負のKL項であり、したがってELBOの値は決してエビデンスを超えることができないのである。

図1：変分オートエンコーダを視覚的に表した。ここで、エンコーダ $q(z|x)$ は観測 $x$ に対する潜在変数 $z$ の分布を表しており、 $p(x|z)$ は潜在変数を観測にデコードする。

図1：変分オートエンコーダを視覚的に表した。ここで、エンコーダ $q(z|x)$ は観測 $x$ に対する潜在変数 $z$ の分布を表しており、 $p(x|z)$ は潜在変数を観測にデコードする。

次に、何故ELBOを最大化しようとするのかを探る。モデル化したい潜在変数 $z$ を導入した我々の目標は、観測データを記述するこの潜在的な構造を学習することである。言い換えれば、我々は変分事後分布 $q_{φ}(z|x)$ のパラメータを最適化し、真の事後分布 $p(z|x)$ に正確に一致させたいのであり、これはKLダイバージェンスを最小化（理想的には０に）することによって実現されるのである。残念ながら、我々はground-truthである $p(z|x)$ の分布を求めることはできないため、このKLダイバージェンス項を直接最小化することは困難である。しかしながら、式(15)の左側で、我々のデータの尤度（したがって我々のエビデンス項 $\log p(x)$ ）は、同時分布 $p(x,z)$ から全ての潜在変数 $z$ を周辺化して計算されるので、 $φ$ に対して常に定数であり、全く $φ$ に依存していないことに注目されたい。ELBO項とKLダイバージェンス項の和は定数になるので、φに関してELBO項を最大化すると、必然的にKLダイバージェンス項が等しく最小化されることになる。したがって、ELBOは真の潜在的な事後分布を完全にモデル化する方法を学習するために、代理として最大化する目的関数とすることができる。ELBOを最適化すればするほど、我々の近似的な事後分布は真の事後分布により近くなる。さらに、一度学習すると、モデルのエビデンス $\log p(x)$ を近似するように学習するので、ELBOは観測データまたは生成データの尤度を推定するためにも使用することができる。

Variational Autoencoders

Variational Autoencoder (VAE) [1]のデフォルトの定式化では、ELBOを直接的に最大化します。このアプローチは変分法であり、 $φ$ によってパラメータ化された潜在的な事後分布族の中から最も良い $q_{φ}(z|x)$ に最適化する。オートエンコーダと呼ばれるのは、入力データが中間的なボトルネック表現を経た後、入力データ自身を予測するように学習されるという、従来のオーエンコーダモデルを彷彿とさせるからである。この関係を明示するために、ELBO項をさらに分解してみる。

$\displaystyle \begin{align} \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(x,z)}{q_{φ}(z|x)} \right] &=\mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p_θ(x|z)p(z)}{q_{φ}(z|x)} \right] \quad &(確率の連鎖率) \tag{17} \\ &=\mathbb E_{q_{φ}(z|x)} \left[ \log p_θ(x|z) \right] + \mathbb E_{q_{φ}(z|x)} \left[ \log \frac {p(z)}{q_{φ}(z|x)} \right] \quad &(期待値を分割) \tag{18} \\ &=\mathbb E_{q_{φ}(z|x)} \left[ \log p_θ(x|z) \right] - D_{\rm{KL}} ( q_{φ}(z|x) || p(z) ) \quad &(\rm{KL}ダイバージェンスの定義) \tag{19} \end{align}$

この場合、我々は中間的なボトルネック分布 $q_{φ}(z|x)$ はエンコーダとして扱うことができ、入力を潜在変数が取りうる値の分布に変換する。同時に、決定論的関数 $p_{θ}(x|z)$ を学習し、与えられた潜在ベクトル $z$ を観測値 $x$ に変換するが、これはデコーダと解釈される。式(19)の2項はそれぞれ直感的に説明できる。第一項は我々の変分分布からデコーダの再構成された尤度を測定する。これは学習した分布が元のデータを再生成できる有効な潜在空間をモデル化していることを保証するものである。第二項は、学習された変分分布が潜在変数に関する事前信念にどれだけ似ているかを測定する。この項を最小化することで、エンコーダはディラック・デルタ関数に陥ることなく、実際に分布を学習するようになる。したがって、ELBOを最大化することは、その第一項を最大化し、第二項を最小化することと等価である。VAEを定義づける特徴は、パラメータ $φ$ と $θ$ に対して同時にELBOを最適化する方法である。一般的には、VAEのエンコーダには対角共分散を持つ多変量ガウシアンをモデル化するもの、事前分布は標準的な多変量ガウシアンが用いられることが多い。

$\displaystyle \begin{align} \displaystyle q_{φ}(z|x) &= \mathcal{N}(z; \boldsymbol{μ}_φ(x), \boldsymbol{\sigma}_φ^2 (x) \boldsymbol {x} )\tag{20} \\ p(z) &= \mathcal{N}(z; \boldsymbol{0}, \boldsymbol{I} \tag{21}) \end{align}$

そして、ELBOのKLダイバージェンス項は解析的に計算され、再構成項はモンテカルロ推定で近似できる。そして、我々の目的関数は次のように書き換えられる。

$\displaystyle \arg \max_{φ, θ} \mathbb E_{q_{φ}(z|x)} \left[ \log p_θ(x|z) \right] - D_{\rm{KL}} ( q_{φ}(z|x) || p(z) )$ $\simeq \arg \max_{φ, θ} \sum_{l=1}^L log p_θ(x|z^{(l)}) - D_{\rm{KL}} ( q_{φ}(z|x) || p(z) ) \tag{22}$

ここで潜在変数 $\left\{ z^{(l)} \right\} _{l=1}^{L}$ は、データセット中のすべての観測 $x$ に対して $q_{φ}(z|x)$ からサンプリングされる。しかし、このデフォルトの設定では問題が発生する。というのも、我々の損失が計算される各 $z^{(l)}$ は確率的サンプリング手順で生成され、それは一般的に微分可能ではないのである。幸い、 $q_{φ}(z|x)$ が多変量ガウス分布など、ある種の分布をモデル化するように設計されている場合、”再パラメータ化のトリック”によって対処することができる。

再パラメータ化のトリックとは、ランダムな変数をノイズ変数の決定論的関数に書き換えることである。これにより、勾配降下法により非確率項を最適化することを可能にする。例えば、任意の平均 $μ$ と分散 $σ^2$ を持つ正規分布 $x ∼N (x ; μ, σ^2)$ からのサンプルは以下のように書き換えることができる。

$\displaystyle x = μ + \sigma \epsilon \quad with \quad \epsilon \sim \mathcal{N}(\epsilon; 0, \boldsymbol{I})$

つまり，任意のガウス分布は，その平均を0から目標の平均 $μ$ に加算シフトし，分散を目標の分散 $σ^2$ だけ伸ばした標準ガウス分布（のサンプル）と解釈することができる。したがって、再パラメータ化のトリックにより、任意のガウス分布からのサンプリングは、標準ガウスからサンプリングし、その結果を目標の標準偏差でスケーリングし、目標の平均でシフトすることで行うことができる。

VAEでは、このように各 $z$ は入力 $x$ と補助雑音変数 $\epsilon$ の決定論的関数として計算される。

$\displaystyle \boldsymbol{z} =\boldsymbol{μ_φ} (x)+ \boldsymbol{\sigma_φ} (x) \odot \epsilon \quad with \quad \boldsymbol{\epsilon} \sim \boldsymbol{\mathcal{N}} (\boldsymbol{ \epsilon }; \boldsymbol{0, I})$

ここで、 $\odot$ は要素ごとの積を表す。このように再パラメータ化された $z$ のもとで、 $φ$ に対して勾配を計算し、 $μ_φ$ と $\sigma_{φ}$ を最適化する。それにより、VAEは再パラメータ化のトリックとモンテカルロ推定を利用し、 $φ$ と $θ$ に対してELBOを同時に最適化することができる。 VAEの学習後、潜在空間 $p(z)$ から直接サンプリングし、デコーダを通すことで、新しいデータを生成することができる。変分オートエンコーダは、 $z$ の次元が入力 $x$ の次元よりも小さい場合に特に興味深いもので、コンパクトで有用な表現を学習することができる。さらに、意味的に（semantically）有用な潜在空間を学習した場合、潜在ベクトルをデコーダに渡す前に編集し、生成されるデータをより正確に制御することができる。

Hierarchical Variational Autoencoders

階層的オートエンコーダ（HVAE）[2, 3] は、潜在変数に対する複数個の階層に拡張されたVAE の一般化である。この定式化のもとでは、潜在変数それ自体は、他のより高いレベルの、より抽象的な潜在変数から生成されたものとして解釈される。直感的には、我々が三次元の観測対象をより抽象的な潜在表現から生成されたものとして扱うように、プラトンの洞窟の人々は三次元の対象を二次元の観測を生成する潜在表現として扱っている。したがって、プラトンの洞窟の住人から見れば、彼らの観測は深さ２（またはそれ以上）の潜在的階層によってモデル化されたものとして扱えられる。

図2： $T$ 個の隠れ階層を持つマルコフ階層変分オートエンコーダ。生成プロセスはマルコフ連鎖としてモデル化され、それぞれの潜在変数 $z_t$ は前の潜在変数 $z_{t+1}$ によってのみ条件付けられる。

$T$ 個の階層を持つ一般的なHVAEでは、各潜在変数は以前の全ての潜在変数に条件付けすることができるが、本研究では、マルコフHVAE（MHVAE）と呼ぶ特別なケースに注目する。MHVAEでは、生成過程がマルコフ連鎖である。つまり、階層を下る各遷移はマルコフ連鎖であり、各潜在変数 $z_{t}$ は前の潜在変数 $z_{t+1}$ にのみ条件付けられる。直感的、視覚的には、これは図2に描かれているように、単にVAEを積み重ねたものと見ることができる。このモデルを表現する別の適切な用語は再帰的VAEである。数学的には、マルコフHVAEの同時分布と事後分布を次のように表現する。

$\displaystyle \begin{align} p(x, z_{1:T}) &= p(x_T)p_{\theta}(x|z_1) \prod_{t=2}^T p_{\theta}(z_{t-1} | z_t) \tag{23} \\ q_{φ}(z_{1:T}|x) &= q_{φ}(z_1|x) \prod_{t=2}^T q_{φ}(z_t | z_{t-1}) \tag{24} \end{align}$

そして以下のようにして、ELBOを簡単に拡張することができる。

$\displaystyle \begin{align} \log p(x) &= \log \int p( \boldsymbol{x}, z_{1:T}) dz_{1:T} &(式(1)を用いた。) \tag{25} \\ &= \log \int \frac{p( \boldsymbol{x}, z_{1:T}) q_{φ}(z_{1:T}|\boldsymbol{x}) }{ q_{φ}(z_{1:T}|\boldsymbol{x}) }dz &(1=\frac {q_{φ}(z_{1:T}|\boldsymbol{x})}{ q_{φ}(z_{1:T}|\boldsymbol{x})}を掛けた) \tag{26} \\ &= \log \mathbb E_{q_{φ}(z_{1:T}|\boldsymbol{x})} \left[ \frac{p( \boldsymbol{x}, z_{1:T})}{q_{φ}(z_{1:T}|\boldsymbol{x})} \right] &(期待値の定義) \tag{27} \\ & \geq \mathbb E_{q_{φ}(z_{1:T}|\boldsymbol{x})} \left[ \log \frac{p( \boldsymbol{x}, z_{1:T})}{q_{φ}(z_{1:T}|\boldsymbol{x})} \right] &(Jensenの不等式を用いた) \tag{28} \end{align}$

次に、同時分布(式(23))と事後分布(式(24))を式(28)に代入すると、別の形になる。

$\displaystyle \mathbb E_{q_{φ}(z_{1:T}|\boldsymbol{x})} \left[ \log \frac{p( \boldsymbol{x}, z_{1:T})}{q_{φ}(z_{1:T}|\boldsymbol{x})} \right] = \mathbb E_{q_{φ}(z_{1:T}|\boldsymbol{x})} \left[ \log \frac{p(x_T)p_{\theta}(x|z_1) \prod_{t=2}^T p_{\theta}(z_{t-1} | z_t) }{q_{φ}(z_1|x) \prod_{t=2}^T q_{φ}(z_t | z_{t-1})} \right] \tag{29}$

後述するように、変分拡散モデル（Variational Diffusion Model）を検討する場合、この目的関数はさらに解釈可能な要素に分解することが可能である。

Reference

[1] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint, arXiv:1312.6114, 2013.

[2] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 29, 2016.

[3] Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Advances in neural information processing systems, 29, 2016.