もし生物情報科学専攻の大学院生が "StableDiffusion" を理解しようとしたら 6 ~VAE~

前回

cake-by-the-river.hatenablog.jp

~~今回は、拡散モデルの最も重要な論文である Denoising Diffusion Probablistic Models （DDPM）を解説します。~~

DDPMを解説する予定だったのですが、基礎となるVariational AutoEncoder （VAE）など潜在変数モデルとその学習に関する話を書いたところ記事が膨らんでしまったので、今回はVAEをメインに説明することとします。なお、今回の内容は、ベイズ深層学習の本を主に参考にしています。

ベイズ深層学習 (機械学習プロフェッショナルシリーズ)

作者:須山敦志
講談社

Amazon

潜在変数モデル

この記事を読んでいる人は、潜在変数モデルについても良く知っているかもしれないですが、一応ざっくりと解説していこうと思います。

潜在変数モデルとは、データが何か目に見えない（観測できない）変数に依存して決まるものとし、それらの関係をモデル化したものです。特に生成モデルの文脈では、データが何かより上位の（確率）変数に基づいて生成されるということを明示的に考えることになります。特にそれらの間の依存関係をグラフにして考えることが多いです。

グラフィカルモデルの例。確率変数は丸で囲まれ、パラメータは裸で書かれている。実線の矢印は依存関係（生成など）を表し、破線の矢印は変分推論を表す。灰色の四角で囲われた領域では、右下の数字（N）の数だけ生成を行ったということを表す（プレート表現）。VAEの元論文より引用。

上図のモデルでは、パラメータ $\theta$ に基づき、潜在変数 $z$ がまず生成されます。その後、観測できる変数 $x$ がパラメータ $\theta$ と潜在変数 $z$ によって生成されています。また、灰色の四角（プレート）で囲われた変数はN個分生成されるものとします。したがって、上図の生成過程を数式で表すと以下のようになります。

$z_i \sim p(z | \theta), x_i \sim p(x | z_i, \theta), i = 1, ..., N$

画像生成の場合、画像が観測できる変数 $x$ に相当し、画像の中で記述されている物体の情報（キリンの首の長さ・青空の色・木の位置、などなど）が潜在変数 $z$ に相当すると言えます。今、画像生成や画像内の物体の認識などの多くのタスクで、画像から元の潜在変数に関する情報を推定することは最も重要な問題だと言えますが、これは事後確率 $p(z|x, \theta)$ を求めることに相当します（上図の破線矢印）。事後確率はベイズの定理から、以下の計算で求まります。

$\displaystyle p(z | x, \theta) = \frac{p(x | z, \theta)p(z | \theta)}{p(x | \theta)},\ \ p(x | \theta) = \int p(x | z, \theta) dz$

問題は、これらの計算に含まれる積分（周辺尤度, モデルエビデンス）が一般に求まらないことです。これは前回の最尤推定における分配関数の困難性と同じ話です。特に、 $z$ から $x$ を生成する過程が（実世界の画像の作成のように）非常に複雑なものである場合、逆方向の確率計算は全く追うことが出来ないと言えます。

ELBOと変分推論

上記の問題において、解けない積分である周辺尤度自体を求めるのではなく、可能な限り精度よく近似するというアプローチを考えることにしましょう。つまり、何かしら潜在変数 $z$ を $x$ から推測する近似モデル $q(z | x, \phi)$ を用意し、その対数尤度を可能な限り近づけるように考えます（以降、 $q$ を"近似モデル", $p$ を "真のモデル" と呼ぶことにします。 $p$ は自然界が行っている本当のメカニズム自体ではないことに注意）。

実は、対数尤度は ELBO （Evidence Lower Bound, 変分下限） $L[q(z | x, \phi), \theta; x$ ] と呼ばれる下限によって下から抑えられることが示せます。

$\displaystyle \mathrm{log}\ p(x | \theta) \ge \int q(z | x, \phi) \mathrm{log}\ \frac{p(x | z, \theta)p(z | \theta)}{q(z | x, \phi)} dz \equiv L[q(z | x, \phi), \theta; x$ ]

実際、対数尤度とELBOの差を計算すると、

$\displaystyle \begin{eqnarray} \mathrm{log}\ p(x | \theta) - \int q(z | x, \phi) \mathrm{log}\ \frac{p(x | z, \theta)p(z | \theta)}{q(z | x, \phi)} dz &=& \mathrm{log}\ p(x | \theta) - \int q(z | x, \phi) \mathrm{log}\ \frac{p(z | x, \theta)p(x | \theta)}{q(z | x, \phi)} dz \\ &=& \mathrm{log}\ p(x | \theta) - \int q(z | x, \phi) \mathrm{log}\ \frac{p(z | x, \theta)}{q(z | x, \phi)} dz - \int q(z | x, \phi) \mathrm{log}\ p(x | \theta) dz \\ &=& \int q(z | x, \phi) \mathrm{log}\ \frac{q(z | x, \phi)}{p(z | x, \theta)} dz \\ &=& D_{KL}[q(z | x, \phi) | p(z | x, \theta)] \end{eqnarray}$

このように、近似モデル分布 $q(z | x, \phi)$ と真モデル分布 $p(z | x, \theta)$ のKLダイバージェンスとなりました。KLダイバージェンスは分布間の"距離のようなもの"であり、完全に一致した場合を除いて常に正の値を取ります。したがって、ELBOは確かに下限を表していると言えます。また、上手く真のモデル分布を近似モデル分布で推定することで、ELBOと対数尤度の間の差を減らすことが出来ることも分かりました。

さて、（最尤推定などで）対数尤度を最大化するためには、ELBOと対数尤度の差を減らすのみならず、ELBO自身の値も最大化することが考えられます。ELBOは以下の形にさらに分解することが可能です。

$\displaystyle \begin{eqnarray} L[q(z | x, \phi), \theta; x] &=& \mathbb{E}_{q(z | x, \phi)}[\mathrm{log}\ \frac{p(x | z, \theta)p(z | \theta)}{q(z | x, \phi)}] \\ &=& \mathbb{E}_{q(z | x, \phi)}[\mathrm{log}\ p(x | z, \theta)] - D_{KL}[q(z | x, \phi) | p(z | \theta)] \end{eqnarray}$

第1項は、近似モデル $q(z | x, \phi)$ により推定した潜在変数 $z$ を用いた際に、真のモデルによりデータ $x$ が得られる確率（尤度）の期待値であり、近似モデルを用いてデータを再構成するときの誤差、あるいは不確かさの指標と言えます。一方、第2項は近似事後確率 $q(z | x, \phi)$ を真の事前分布 $p(z | \theta)$ に近づけようというもので、 $\phi$ に対する制約ともいえます。

対数尤度を最大化するには、ELBOと対数尤度の差を減らす（近似モデルと真のモデルを一致させる）こととELBOを最大化することの二つが絡みそうですが、変分法を利用することで、実はこの二つが同じことを表すことが示せます（そのため変分下限と呼ばれている？）。変分法は解析力学などで利用される最適化の理論で、関数を入力とする関数（汎関数）について、その値を最大化できるような入力関数の条件が求められます。ここではその詳細を割愛しますが、汎関数をまるで"関数で微分”するかのように扱うことで、極値の条件式を導出できます。今回の場合は、ELBO $L[q(z | x, \phi), \theta; x]$ を最大化する $q(z)$ という関数の極値の条件が以下と等しくなることが示せます（オイラー=ラグランジュ方程式）。

$q(z) = \mathrm{argmax}_{q} L[q(z | x, \phi)] = p(z)$

ここでは、あえてパラメータ $\phi, \theta$ を用いませんでしたが（パラメータ $\phi$ の最適化ではなく関数 $q$ の最適化の理論なので）、実際に最適化するのはパラメータの方であり、上記の結果からELBOが最大化されるとき、ELBOと対数尤度の差を表すKLダイバージェンスも0となり、対数尤度も最大化できることが分かります。したがって、ELBO最大化が達成できれば理論上対数尤度も最大化できます。このように、変分下限の最適化を行うことで近似的に推論するアルゴリズムは変分推論法と呼ばれています。

一般に、近似モデルのパラメータ $\phi$ はもちろん、真のモデルのパラメータ $\theta$ についても推定する場合（VAEなど）がありますが、その際は両方の変数を最適化するアプローチが取られます。

VAE

VAE（Variational AutoEncoder, 変分自己符号化器）は変分推論法における近似モデルや真のモデルの生成過程に深層学習による回帰モデルを組み込んだものになります。元論文では、AEVBアルゴリズムというより抽象的なアルゴリズムの一部としてVAEを挙げています。

arxiv.org

Stable Diffusionでは潜在表現から画像を実際に生成する際の品質向上目的で利用されています（また、そもそも全体を巨大な潜在変数モデルともみなせる）。なお、個人的に以下の記事が（恐ろしく）丁寧に元論文に従って説明していると思うので、より詳細な説明は下記をご覧ください。

academ-aid.com

VAEでは、潜在変数 $z$ がガウス分布に従うものとし、この潜在変数が何かしらの関数 $f(z | W)$ を平均としたノイズを乗せた形で生成されると仮定します（ここではパラメータ $W$ も確率変数として扱う）。実際にはN個のデータがＮ個の潜在変数のもと生成されるとするので、以下の式で表せます。

$p(x_n | z_n, W) = \mathcal{N}(x_n | f(z_n ; W), \lambda_x^{-1} \mathbb{I})$

もし観測データが二値の場合は、ガウス分布ではなくベルヌーイ分布を仮定します。近似分布は色々な形式が考えられますが、 $z_n$ と $W$ の分布がそれぞれ独立したガウス分布の積で書ける（平均場近似）として

$\displaystyle q(Z, W | X, \theta, \phi) = q(W | \theta) q(Z | X, \phi) = \prod_{i,j,l} \mathcal{N}(w_{i,j}^{(l)} | m_{i,j}^{(l)}, v_{i,j}^{(l)}) \prod_{n} \mathcal{N}(z_n | m(x_n ; \phi), v(x_n ; \phi)),\ \ \theta = \{ m_{i,j}^{(l)}, v_{i,j}^{(l)} \}$

とし、さらに、近似モデルにおける $m, v$ がいずれも関数 $f(x_n; \phi)$ の出力値（ペア）であると仮定します。このモデルのポイントは、真のモデルによる各データの生成と近似モデルにおける潜在変数の事後分布が、いずれも関数により決定されるものと考えることです。この関数をニューラルネットワークによって学習させることで、複雑な生成過程であっても学習ができると考えるわけです。

また、サンプルされる各データ・潜在変数 $x_n, z_n$ それぞれの出力について計算するのではなく、それぞれがより上位の関数によって式変形されるとし、その関数のパラメータを求めるように問題を設定したことで、計算コストを抑えることが出来ます（このような手法は償却推論と呼ばれています）。

図からもわかるように、近似モデル側はデータから潜在表現へと情報を圧縮するEncoder, 真のモデル側は潜在表現からデータを生成するDecoderとして捉えることができ、こうした構造のモデルはAutoEncoder（自己符号化器）と呼ばれます。

このモデルを変分推論法で学習するには、ELBOを最大化すればよいです。具体的にはパラメータ（の集合） $\theta, \phi$ それぞれについてELBOを分解した各項について勾配を求め、勾配降下法を適用します。今回はニューラルネットワークの学習なので、データからミニバッチ $\mathcal{S}$ を構成し、確率的勾配降下法を適用して各パラメータごとに更新することとなります。

$\displaystyle L_\mathcal{S}(\theta, \phi) = \frac{N}{M}\sum_\mathcal{S} [\mathbb{E}[\mathrm{log}\ p(x_n|z_n, W)] +\mathbb{E}[\mathrm{log}\ p(z_n)] - \mathbb{E}[\mathrm{log}\ q(z_n | x_n, \phi)] ] + \mathbb{E}[\mathrm{log}\ p(W)] - \mathrm{E}[\mathrm{log}\ q(W | \theta)]$

$\displaystyle \nabla_\theta L_\mathcal{S}(\theta, \phi) = \frac{N}{M}\sum_\mathcal{S} \nabla_\theta \mathbb{E}[\mathrm{log}\ p(x_n|z_n, W)] +\nabla_\theta \mathbb{E}[\mathrm{log}\ p(W)] - \nabla_\theta \mathbb{E}[\mathrm{log}\ q(W | \theta)]$

$\displaystyle \nabla_\phi L_\mathcal{S}(\theta, \phi) = \frac{N}{M}\sum_\mathcal{S} \nabla_\phi \mathbb{E}[\mathrm{log}\ p(x_n|z_n, W)] +\nabla_\phi \mathbb{E}[\mathrm{log}\ p(z_n)] - \nabla_\phi \mathbb{E}[\mathrm{log}\ q(z_n | x_n, \phi)]$

ところが、このままでは上手く学習させられません。なぜならば、ELBOの各項は期待値=積分であり、MCMCのようなサンプリング法を導入する必要がありますが、誤差逆伝播法などで微分を計算する際にサンプリング自体は一般に微分が出来ないからです。例えば、上記で $W$ や $z_n$ はそれぞれ $q(W | \theta)), q(z_n | x_n, \phi)$ という分布から従う変数であり、これらを含む期待値の計算はサンプリングを用いる必要がありそうです。

Reparameterization Trick

実際には 再パラメータ (reparameterization trick) というテクニックを用いることでこの問題を回避できます。これを今回の例で説明します。

今、ある関数 $f(z | \phi)$ を、確率 $q(z | x, \phi)$ のもとに期待値計算したいとします。ここで、 $\epsilon \sim p(\epsilon)$ という確率変数を用いて以下のように $q(z | x, \phi)$ が書けるとします。

$z = g(\epsilon, x ; \phi),\ \ q(z | x, \phi)dz = p(\epsilon)d\epsilon$

このとき、 $q(z | x)$ による期待値計算は、以下のように $\epsilon$ の期待値計算に置き換えることが出来ます。

$\int q(z|x, \phi)f(z, \phi)dz = \int p(\epsilon)f(g(\epsilon, x ; \phi)d\epsilon$

すなわち、パラメータや入力によって制約が課せられた（そのため複雑なサンプリングが発生しうる） $q(z | x, \phi)$ による期待値から、外部の制約がない確率変数である $\epsilon$ による期待値にすり替えることが出来ます。このとき、後半の期待値計算は単に $p(\epsilon)$ に従って用意した数値 $\tilde{\epsilon}$ を関数 $g(\tilde(\epsilon), x ; \phi)$ で変換したもので近似できるようになります。もし関数 $g$ が微分可能であれば、このプロセス全体もパラメータ $\phi$ で微分可能となり、サンプリングなしで期待値計算に相当する計算を行ったことになります。

抽象的だったので、実際の具体例で考えましょう。VAEの例では、パラメータ $W$ の分布は $q(W | \theta) = \prod_{i,j,l} \mathcal{N}(w_{i,j}^{(l)} | m_{i,j}^{(l)}, v_{i,j}^{(l)})$ でしたが、これは（正規分布の性質から）

$g(\epsilon, | m, v) = m + \epsilon v,\ \ \epsilon \sim \mathcal{N}(\mathbb{0}, \mathbb{I})$

という式変形で表せます（標準ガウス分布を平均 $m$ に平行移動し、標準偏差 $v$ に拡大したもの）。この式は各パラメータ $m, v$ に対し線形で微分可能なので、通常の誤差逆伝播法を適用することができ、パラメータの更新が可能です。同様に、 $z_n$ の分布もガウス分布だったため、疑似変数 $\epsilon$ を導入して微分可能な形に変えられます。

この再パラメータを用いて各種パラメータを変分下限について最適化するのが、VAEとなります。変分推論の際説明したように、ELBOの最大化は、近似モデル $q$ のEncoderで推論した潜在変数 $z$ を真のモデルに対応するDecoderで再び入力データに戻したときの誤差（再構成誤差）と、潜在変数の分布の差を抑えるKLダイバージェンスの項（正則化）からなるため、VAEの学習もこれらモデルによる生成プロセスを合体させたAutoEncoder型のニューラルネットワークとして行うことで可能となります。

今回は潜在変数モデルの基本とVAEについて解説しました。次回こそついに拡散モデル・DDPMの解説を行いたいと思います。

cake-by-the-river.hatenablog.jp

余談：変分推論と脳の自由エネルギー原理

そのうち書きます。