変分オートエンコーダの理論と実装

論文を読みながらVAEを徹底的に理解する。

論文情報

Kingma, Diederik P., and M. Welling. "Auto-Encoding Variational Bayes." International Conference on Learning Representations, 2013.

VAE（変分オートエンコーダ）は2014年4月に行われたICLR2014にて発表されたモデルである。論文は2013年12月にarxivにて公開された。この研究は10年後のICLR2024にてTest of Time Awardに輝いている。10年経ってもなお、その貢献が学術分野に大きな影響を残していると言える。筆頭著者のKingma氏はこのVAE以外にもAdam最適化（ICLR2015）やGlow（NeurIPS2018）など深層学習という分野に多くの貢献を残している。

確率モデルのパラメータ推定とその課題

VAEが解決する課題を理解するために、その背景をまず理解する。

得られたデータ$X=\{x^{(1)},x^{(2)},\ldots,x^{(N)}\}$の分布を推定する代表的な手法として、最尤推定がある。最尤推定は生成モデルの学習にも用いられ、GANのような異端児を除いた多くのモデルが最尤推定に基づいた学習を行っている。

最尤推定では尤度が最大になるパラメータを求める。

最尤推定以外にも分布を推定する方法は存在し、例えばMAP推定では尤度ではなく事後確率が最大となるパラメータを求める。

最尤推定やMAP推定のようなパラメータのある一点を求める推定は点推定と括られる。一方で、パラメータの一点ではなくパラメータの分布を求める手法も存在し、それはベイズ推定と呼ばれる。

いずれの推定手法も初めにモデル化を行う必要がある。ここでのモデル化とは、確率変数$x$の発生過程をパラメータ$\theta$で記述することを指す。最尤推定ではパラメータ$\theta$に従う確率分布$p_\theta(x)$、MAP推定・ベイズ推定ではそれに加えて事前分布$p(\theta)$を定義する。

モデル化は$x$の性質や背景をよく観察した上で行われる。$x$がシンプルである場合、そのモデル$p(x)$として我々がよく知る扱いやすい分布を仮定できる。その場合は各種推定を解析的に行うことができるので、とても嬉しい。しかしそのようなケースは稀で、特に生成モデルが欲しいと感じるような高次元で複雑なデータに対して、扱いやすいシンプルな分布を仮定することは不可能に等しい。

Loading image...

そのような場面によく用いられる手法が、潜在変数の仮定である。$x$の裏に、$x$をよく表したより本質的な情報$z$が存在すると仮定し、その周辺化によって$x$の分布を表す。この$z$を潜在変数と呼ぶ。

この形は非常に表現力が高く、各分布$p(z)$、$p(x|z)$に対してそこまで複雑な分布を仮定せずとも、複雑な$x$をよく表せる。実際、上の図の右側は混合ガウスモデルという潜在変数モデルの可視化である。このモデル化ではデータ$x$が次のような流れで得られることを仮定している。

1. $p(z)$から$z$がサンプリングされる
2. $p(x|z)$から$x$がサンプリングされる

Loading image...

しかしこの分布の和の形は扱うのが非常に困難であり、各種推定において解析解を得ることは難しくなる。

EMアルゴリズム

潜在変数モデルの点推定を行う代表的な手法として、EMアルゴリズムが存在する。この手法は混合ガウスモデルのように潜在変数が離散的な場合に適用できる。

最尤推定を例に説明する。EMアルゴリズムでは対数尤度を次のように分解する。

$q$は何らかの確率分布。$\mathcal L$は変分下界と呼ばれ、次のように表される。

下界という言葉が示すように、以下が常に成り立つ。KL情報量の非負性から容易に納得できる。

これらの関係は次のように表せる。

Loading image...

EMアルゴリズムでは以下を繰り返すことで対数尤度を大きくする:

1. KL情報量を$q$について最小化
2. 変分下界を$\theta$について最大化

KL情報量は2つの確率分布が完全に一致するときにのみ最小値0を取る。つまりステップ1では次を行う。

対数尤度自体は$q$に依存しないので、$q$を更新することで変化するのは下界とKL情報量の比率になる。KL情報量を小さくすることで、下界はその分大きくなる。

ステップ2ではこの$q$の下で下界を最大化する。下界は数学的に扱いやすい性質を持ち、$z$が離散的な場合は$\theta$に関する最大化を解析的に行える。

そしてその新たな$\theta$の下で再度$q$を更新し、と繰り返すことで目的の対数尤度を大きくすることが可能である。この更新によって対数尤度が単調増加（厳密には単調非減少）することは数学的に示せる。

EMアルゴリズムは、パラメータ$\theta$の下での潜在変数$Z$に関する事後分布$p_\theta(Z|X)$が計算可能であるという好ましい状況によって成り立つ手法である。$Z$が離散的であることで事後分布の計算に必要な周辺尤度（分母）を総当たりによって求められるようになる。

変分ベイズ

EMアルゴリズムは潜在変数モデルの点推定を行うが、潜在変数モデルのベイズ推定を行いたい場面もある。それを行う手法が変分ベイズ（Variational Bayse）である。変分推論と呼ばれたりもする。

潜在変数モデルにおけるベイズ推定では、パラメータと潜在変数を同一視する。

潜在変数$Z$はデータ$X$の発生に関わる確率変数であるため、パラメータを確率変数とするベイズ的な立場の下ではパラメータ$\theta$も潜在変数に含まれる。

変分ベイズでは、データ$X$が得られた下での潜在変数$Z$の事後分布を求める。

しかし複雑なモデル化の下でこの事後分布を求めることは困難である。そこで、事後分布を別の分布$q(Z)$で近似することを考える。

近似精度を表す指標としてKL情報量を採用し、その最小化を目指す。

そしてこのKL情報量はこのように表せる。

この式はEMアルゴリズムで出てきた分解をパラメータ$\theta$で周辺化することで得られる。ここで、周辺対数尤度$\ln p(X)$は$q$に依らない定数であることがわかる。つまりKL情報量の最小化は下界の最大化と同じ意味になる。

変分ベイズでは、変分下界を$q$について最大化することで事後分布を近似する。しかし多くの場合でその最大化が素直にできない。そこで平均場近似という簡略化を行う。これはいくつかの潜在変数が独立であるという仮定で、この仮定の下、座標上昇法による下界の最大化を行う。なおこの最大化の過程で変分法という汎関数の停留点を求める手法を用いるため、変分ベイズという名前がついている。

VAEの目的

潜在変数モデルにおける上記パラメータ推定手法の課題を解決するモデルとして変分オートエンコーダ（VAE; Variational Autoencoder）というモデルが提案された。

EMアルゴリズムでは潜在変数モデルにおける点推定を実現したが、適用場面は潜在変数が離散的であるという特殊なケースに限られる。また変分ベイズでは潜在変数モデルにおけるベイズ推定を実現したが、多くの場合で平均場近似という強い簡略化が必要になり、複雑なモデルに適用することは難しい。

VAEはそれらの問題をまとめて解決する手法として提案された。連続的な潜在変数を持つモデルにおけるパラメータに関する点推定と潜在変数に関するベイズ推定を同時に行う。この目的は論文にも分かりやすく書いてある。

The strategy in this section can be used to derive a lower bound estimator (a stochastic objective function) for a variety of directed graphical models with continuous latent variables. We will restrict ourselves here to the common case where we have an i.i.d. dataset with latent variables per datapoint, and where we like to perform maximum likelihood (ML) or maximum a posteriori (MAP) inference on the (global) parameters, and variational inference on the latent variables.

本稿では点推定の中でも最尤推定に絞ってVAEを説明する。つまり目的を次の2つとする。

1. 尤度$p_\theta(X)$を大きくする$\theta$を求める
2. その下での潜在変数$Z$に関する事後分布$p_\theta(Z|X)$を求める

潜在変数が離散的である場合、目的1はEMアルゴリズムによって実現でき、目的2についてはいかなる$\theta$について解析的に行える。ここから考えると、VAEはEMアルゴリズムを進化させて連続的な潜在変数に対応させたものとも見られる。EMアルゴリズム・変分ベイズ・VAEの位置関係をまとめると次のようになる。

VAEのモデル化

これらの目的を達成するための具体的なモデル化を見ていこう。

得られた$x$が複雑でシンプルなモデルでは記述できない場面、潜在変数$z$を導入することで高い表現力をモデルに与えようとするのであった。離散的な$z$ではこうなる。

しかしこれでも表現力が足りないことがある。得られたデータが画像のような高次元データであると、離散的な潜在変数では十分な表現力が得られない。そこで連続的な潜在変数を仮定する。VAEではこれを考える。

以後観測データ$x$と潜在変数$z$はベクトルとする。

まずは事前分布$p(\bm z)$を決める。といってもこれは割と適当で良い。というか考えようがない。VAEにおける潜在変数の仮定とは、$\bm x$のより本質的な情報を表した何かが存在するというものであり、それ以上の仮定は行わない。ただかといって、具体的にモデル化しないと先へ進めないので、実装時は我々にとって都合のいいようなモデル化を適当に行なっておく。基本的には標準正規分布$\mathcal N(\bm 0, I)$かな。サンプリングもしやすいし、扱いやすいし、これでいいでしょ。

実際、$\bm z$が連続である程度の次元数を持っていればその時点でそれなりの表現力を持つため、どのような分布であっても、現実世界で観測されるデータを表すにはほとんど場合で十分である（多様体仮説）。そのためVAEでは、$\bm z$には$\bm x$よりもずっと低い次元の空間を仮定することが多い（$D_z < D_x$）。ちなみに、この「生成したいデータの裏に、サンプリングが容易な何らかの潜在変数が存在する」という仮定は多くの生成モデルの出発点である（VAE、GAN、Flow、DDPM）。

次にその$\bm z$から$\bm x$への対応として条件付き分布$p(\bm x|\bm z)$をモデル化する。しかし$\bm z$は多次元連続ベクトルなので自力で細かく記述することは難しい。そこで、DNN（Deep Neural Network）を用いる。DNNは表現力が高いので、多様な表現を持つ$\bm z$をうまく扱えると仮定する。具体的には、$\bm z$を受け取って何らかの確率分布のパラメータ$\hat\psi$を出力するDNNを定義し、条件付き分布をモデル化する。

そしてこのときの$p_\theta(\bm x|\bm z)$にもシンプルな分布を仮定してよい。どれくらいシンプルにしていいかというと、まず$\bm x$の各要素$x_i$がそれぞれ独立であると仮定してよい。

さらにその上で$p(x_i;\hat\psi_{\theta,i})$にはベルヌーイ分布や正規分布といった我々がよく知る扱いやすい分布を仮定する。$x_i$が二値ないしそう見てしまっても問題ないような場面ではベルヌーイ分布を仮定し、

それ以外の場面では分散固定の正規分布を仮定する。

大抵のモデル化はこのいずれかになる。

ちなみに、$x_i$が独立であると仮定しているが、実際は当然そんな訳がない。例えば$\bm x$が画像を表したベクトルで$x_i$がある一つのピクセルの色を表しているとすると、隣り合うピクセルには強い相関があるはずである。それらが独立であると仮定した場合、$p_\theta(\bm x|\bm z)$は観測データを表すのにやや不正確であると言える。

そのような問題もあってか、VAEを用いてデータを生成する際は、$p_\theta(\bm x|\bm z)$からのサンプリングに通常の確率的な手法ではなく、最も確率密度の高い点を出力するという決定的な手法を取ることが多い。そのような点はシンプルな$p_\theta(\bm x|\bm z)$の下では解析的に得られる、つまりDNNが出力したパラメータ$\hat{\bm\psi}_\theta$によって一意に求められる。そしてその$\hat{\bm\psi}_\theta$の各要素$\hat\psi_{\theta,i}$は独立ではないため、これによって比較的正確なサンプルが得られるようになる。これは$\bm z$と$\bm x$の対応関係の記述を全てDNNに任せるということである。

例えば$p_\theta(\bm x|\bm z)$に正規分布を仮定した場合、最も確率密度の高い点はその平均$\hat{\bm\mu}_\theta=\text{DNN}_\theta(\bm z)$と一致する。つまり、サンプリングした$\bm z$をDNNに入れて出てきたものがそのまま生成結果となる。分散を固定する理由はこれで、そもそも分散を$\bm z$によって変化させたい場面が存在しないからである。

連続的な潜在空間が持つ多様性により、このような一部が決定的なサンプリング手法でもまあまあ多様なデータが生成できる。なおこの決定的なサンプリングは、独立性に加えて分散が非常に小さい分布であるという仮定を行い、その分布からサンプリングしている、と見ることもできる。

一旦ここまでのモデル化を整理する。

$p(\bm z)$はパラメータ$\theta$に依存しない固定分布。$\bm x$の発生過程を図で表すとこう。

Loading image...

さて、最初に掲げた目的をもう一度見てみよう。

1. 尤度$p_\theta(X)$を大きくする$\theta$を求める
2. その下での事後分布$p_\theta(Z|X)$を求める

ここで、先のモデル化では2を達成することができない。以下の事後分布を求めることは事実上不可能であるため。

$p_\theta(X|Z)$は多層のニューラルネットで記述されるので、その積分なんて求められるわけがない。$Z$は多次元連続ベクトルなので、GMMのように取りうるすべての$Z$で総当たりをすることもできない。

ここで変分ベイズの手法を踏襲する。変分ベイズでは、直接求められない事後分布$p(Z|X)$の代わりに別の分布を持ち出して事後分布を近似した。これを真似して、$p_\theta(Z|X)$を近似する分布として$q_\phi(Z|X)$を定義する。そしてこの$q$を先ほどと同じようにDNNでモデル化する。

$q_\phi(\bm z|\bm x)$にはほとんどの場合で独立な正規分布を仮定する。$p(\bm z)$に標準正規分布を仮定する場合、同じ分布を仮定しておくと後の計算が楽になる。この場合、DNNは正規分布の平均と分散を出力する。

最後にモデル化をまとめる。

まず確率モデルを次のように定義する。

そして事後分布は素直に求められないので、次のように近似する。

VAEはこれら二つのモデルによって構成される。

VAEの学習

学習目標を再再掲する。

1. 尤度$p_\theta(X)$を大きくする$\theta$を求める
2. その下での事後分布$p_\theta(Z|X)$を求める

2については別の分布で近似するものとしたので、次のように書き換える。

1. 尤度$p_\theta(X)$を大きくする$\theta$を求める
2. その下での事後分布$p_\theta(Z|X)$をよく表す近似分布$q_\phi(Z|X)$を求める

この目標に向かって学習を進めよう。

まず、これらの目標をより厳密に記述する。この二つの達成度合いを表す具体的な値を考える。

1の尤度の大きさは対数尤度で表すと良さそう:

2の近似の精度については、変分ベイズに倣ってKL情報量を使おう:

さてこうすると、我々の目標は、1を最大化しながら2を最小化することとなる。ここで、片方にマイナスをかけて目標の方向（最大化 or 最小化）を揃え、この2つを足すことを考える。そうすると、その足した値を最大化ないし最小化すればよいことになる。

KL情報量にマイナスをかけて足してみよう。すると目標はその足した値の最大化となる。

さてこの式、どこかで似たものを見た気がする。EMアルゴリズムで出てきた変形かな。

パラメータの書き方を合わせて、移項すると？

先の目標が変分下界$\mathcal L(\theta, \phi)$という一つの項で記述できた。つまり、1. 尤度の最大化と、2. KL情報量の最小化、という2つの目標を、変分下界の最大化という一つの目標で表すことができた。そしてこれがVAEの学習目標となる。

下界の最大化は勾配法で行う。勾配を求めるために、下界を分かりやすい形に変形する。

一項目は近似事後分布$q_\phi(Z|X)$の下での条件付き対数尤度$\ln p_\theta(X|Z)$の期待値。二項目は潜在変数$Z$に関する近似事後分布と事前分布のKL情報量。この勾配を計算し、その勾配に基づいてパラメータの更新を繰り返すことで、下界を最大化する。

Reparameterization Trick

変分下界を期待値とKL情報量の二つの項に分解した。求めたいものはこれらの勾配である。KL情報量の項については解析的に求められることが多く、勾配もそこからの逆伝播で求められるが、期待値の項についてはそうもいかない。まだ$\theta$については比較的簡単な近似ができるが、$\phi$に関する微分は非常に厄介で、期待値の確率分布$q_\phi(Z|X)$を記述しているため$\theta$で使える近似手法も使えない。

ここで、ある期待値のその分布のパラメータに関する勾配を求めることを考える。

これは変分下界における期待値の項の勾配を求める問題の一般化である。そして値は次のように変形を進めることで近似できる。

2-3行目の変換には対数微分のトリックを使った。

これで期待値の勾配を推定できる。この推定手法はREINFORCEアルゴリズムにおける方策勾配の近似にも使われており、そのままREINFORCEと呼ばれたりする。スコア関数$\nabla_\theta \ln p_\theta(x)$を用いるのでスコア関数推定量と呼ばれたりもする。本稿でもそう呼ぼうと思う。

変分下界の期待値の項の勾配もスコア関数推定量で近似できる。$p_\theta(x)=q_\phi(Z|X)$、$f(x)=\ln p_\theta(X|Z)$を代入すればいい。

しかしこの推定量には問題点がある。

この値は期待値の不偏推定量なので、$L\to\infty$の極限で真の値と一致するが、現実的には$L$を何らかの有限な値に設定する必要がある。$L$を有限な値に設定する場合、気になるのはそのばらつきである。もし推定量が大きな分散を持つ場合、真の値から外れた値を取ることが多くなる。

ここから、勾配法における勾配を不偏推定で近似する場面において、その分散が大きいことは、「各ステップで行うパラメータ更新が正しい方向への更新にならない可能性が高い」と言える。そして実際に、スコア関数推定量は大きな分散を持つことが知られている。これではVAEの学習が効率的に行えない。

この問題を解決するため、Kingmaらは新たな勾配近似手法としてReparameterization Trickを提案した。日本語では再パラメータ化トリックと呼ばれる。

改めて目的を整理すると、我々が欲しいものは「期待値の勾配」であった。この時、この期待値をモンテカルロ法で近似してそれを微分する、というアプローチが最も愚直に感じる。

しかし、この近似を成立させるためには下段のサンプリング過程$x^{(l)} \sim p_\theta(x)$がパラメータ$\theta$について微分できなければならず、その要件は一般的な場面で満たせない。これに対しKingmaらは、いくつかの緩い条件の下ではその要件をクリアできると主張した。具体的には、サンプリング過程を次のように書き直せる場合、サンプリング過程をパラメータで微分できるようになることを示した。

適当な分布からノイズ$\epsilon$を取得し、それを用いて確率変数$\tilde x$をパラメータ$\theta$でdeterministic=決定論的に記述した。これによって確率変数$\tilde x$がパラメータ$\theta$で微分できるようになり、$\tilde x \sim p_\theta(x)$となるようなノイズ分布$p(\epsilon)$と$\theta$について微分可能な関数$g_\theta(\epsilon)$をうまく見つけることができれば、$p_\theta(x)$からのサンプリング過程が$\theta$で微分できるようになる。例えば正規分布の場合、$p(\epsilon)$を標準正規分布にすることで再パラメータ化が可能である。

ある分布や変数をパラメータを用いて記述することをパラメータ化（Parameterization）と呼ぶ。先の書き直しは、本来微分できなかったサンプリング過程を微分できる形で改めて（再び）パラメータ化したことになる。Reparameterization（再パラメータ化）という名前はそこからきているのでしょう。

直感的には、この書き直しは、サンプリングという微分ができない操作をパラメータ$\theta$とは関係のない部分に逃がしたものと言える。

Loading image...

$\theta$と$x$が確率的な過程を経ずに繋がることで、微分ができるようになる。

また、具体例をもとに別の捉え方もしてみよう。$y=f(x)$の$x$についての微分は、「$x$が少し違う値だった時に$y$がその何倍変わっていたか」と捉えられる。これをサンプリングについても考えてみる。例えば、平均$\mu$の正規分布から一つのサンプル$x$を取得したとして、仮に$\mu$が少しずれていた場合に$x$がどれくらいずれていたかを考える。もしこれが分かれば、サンプリングという確率的な過程を微分できそう。すると、$\mu$がずれることで正規分布はそのまま平行移動するだけなので、$x$も同じ方向に同じ量だけずれていただろうと想像できる。$\mu$が少しずれたときに$x$も同じだけずれるとすると、この場合の$x$の$\mu$に関する微分を1と定義できそう。

実際に、先の$\tilde x$を$\mu$について微分すると1になるね。このように、分布のパラメータが確率変数に与える影響が比較的シンプルに表せるとき、このような再パラメータ化が可能になる。

では、このReparameterization Trickを用いて変分下界の期待値の項の微分を近似することを考えよう。期待値を次のように近似し、各パラメータで微分する。

$\theta$に関する微分は素直にできるのでヨシ。$\phi$に関する微分はReparameterization Trickを使用する。$Z^{(l)}$を再パラメータ化することで$\phi$に関する微分を可能にする。

例えば、正規分布を仮定する場合はこう。

これで、勾配をより正確に推定することが可能になった。実際に、1サンプルあたりの$Z$のサンプル数$L$について、論文では次のように述べられている。

In our experiments we found that the number of samples $L$ per datapoint can be set to $1$ as long as the minibatch size $M$ was large enough, e.g. $M = 100$.

バッチサイズがある程度大きければ$L=1$で十分な近似ができるらしい。すごいね。

最後に変分下界の近似を整理する。バッチサイズをある程度大きくし、$L=1$とすることを想定する。VAE学習の各ステップでは、得られたミニバッチの下で次の下界を求め、各パラメータ$\theta$、$\phi$の勾配を誤差逆伝播で計算する。

ここで、この算出過程をよく見ると、このモデルがオートエンコーダと見なせることがわかる。分かりやすく、ある一つのデータ$\bm x^{(n)}$が得られた下での近似変分下界を定義し、算出の流れを見てみる。

特に一項目の条件付き対数尤度に着目する。この項は次の流れで算出される。

1. $\text{DNN}_\phi$に$\bm x^{(n)}$を与えて事後分布のパラメータ$\hat{\bm\psi}_\phi$を取得
2. Reparameterization Trickを用いて$\hat{\bm\psi}_\phi$に従う確率分布$p(\bm z;\hat{\bm\psi}_\phi)$からのサンプル$\bm z’$を取得
3. $\text{DNN}_\theta$に$\bm z'$を与えて条件付き分布のパラメータ$\hat{\bm\psi}_\theta$を取得
4. $\hat{\bm\psi}_\phi$の下での$\bm x^{(n)}$の確率密度$p(\bm x=\bm x^{(n)}; \hat{\bm\psi}_\phi)$を求め、対数に変換する

ここで、ステップ4で得られる条件付き対数尤度について考える。条件付き分布$p(\bm x|\bm z)$に正規分布を仮定した場合、対数尤度は初めに与えたデータ$\bm x^{(n)}$とその正規分布の平均$\hat{\bm\mu}_\theta$の負の二乗和誤差になる。なおこの平均$\hat{\bm\mu}_\theta$は$\text{DNN}_\theta$が出力するベクトルそのものである。

次に、条件付き分布にベルヌーイ分布を仮定した場合でも考えてみる。この場合の対数尤度は、初めに与えたデータ$\bm x^{(n)}$とそのベルヌーイ分布のパラメータ$\hat{\bm\lambda}_\theta$の負のバイナリ交差エントロピー誤差になる。このパラメータ$\bm{\hat\lambda}_\theta$も$\text{DNN}_\theta$が出力するベクトルそのものである。

$\text{DNN}_\theta$から得られるデータを$\hat{\bm x}$と表すことにしよう。今まで$\hat{\bm\psi}_\theta$と表していたやつだね。するとステップ4では、（基本的なモデル化の下では、）適当な距離関数を用いて$\bm x^{(n)}$と$\hat{\bm x}$の誤差を求めていることになる。これは再構成誤差と呼べる。このステップ4を踏まえると、ステップ1から4の流れは、与えられたデータを一度低次元の潜在表現に変換したのち、そこから元のデータを復元するように見える。

Loading image...

そしてこの流れは正にオートエンコーダである。変分オートエンコーダという名前はここから来ている。またここから、分布の記述に用いるDNNはエンコーダ・デコーダと呼ばれる。$\text{DNN}_\phi$がエンコーダ、$\text{DNN}_\theta$がデコーダということ。

また変分下界の二項目のKL情報量は、エンコーダが出力する事後分布のパラメータ$\hat{\bm\psi}_\phi$ を調整している。これは正則化項とも呼ばれ、エンコーダを介して得られる潜在変数が指定した分布に従うような調整を行う。

深層学習モデルの学習は、通常、目的関数を最小化の対象として記述し、勾配降下法による最小化を行う。VAEの目的関数を最小化の目標として記述し直すと負の変分下界となり、それを構成する二つの項はオートエンコーダに倣うと次のように呼べる。

このことから、VAEは潜在変数を指定した分布に制限したオートエンコーダと見ることができる。

生成モデルとしてのVAE

VAEの学習では変分下界の最大化を目指すが、そのモチベーションは以下であった。

1. 尤度$p_\theta(X)$を大きくする$\theta$を求める
2. その下での事後分布$p_\theta(Z|X)$を求める

この目標を一つにまとめたものが変分下界の最大化だったね。

対数尤度が一つ目の目標、KL情報量が二つ目の目標に対応する。ではここで、「生成モデルだけが欲しい」という状況を考えてみる。生成モデルの学習時に気にするのは尤度の大きさのみなので、二つ目の目標は無くなる。

しかしこの場合でも目的関数は変分下界のままである。下界を二種類のパラメータ$\theta, \phi$について最大化することで間接的な対数尤度の最大化を図る。理由はごく単純で、対数尤度を直接大きくする方法がないから。

さて、対数尤度は手に負えないけど下界なら扱えるという点は一旦納得するとして、果たしてその下界の最大化は本当に対数尤度の最大化と繋がるのだろうか？下界はあくまでも二つの目標をまとめたもので、片方が実現できなくてももう片方が強く実現されていれば全体としては達成される。要は、仮に対数尤度が減少していても、その減少幅を負のKL情報量の増加幅が上回っていればそれらの和である下界は大きくなるということ。つまり下界を大きくしても対数尤度が大きくなっているとは限らない。

Loading image...

また、$\ln p_\theta(X) \geq \mathcal L(\theta, \phi)$という関係性を用いた、「下からの評価を最大化することで間接的に対数尤度を最大化する」という説明もよく見られるが、やはりこれも下界の最大化と対数尤度の最大化が繋がることは示しておらず、依然として腑に落ちない。「多くの場合で成り立っているので経験則的に成り立つ」という説明はまだ納得感があるものの、できることならそれを数学的に示したいものである。ということで、本章ではそれを示そうと思う。

対数尤度の微分を見てみよう。

微分を整理する

さてここで、

より、こうなる。

これを代入すると

同時分布の勾配の期待値になった。対数尤度の微分なので、これが対数尤度を大きくするパラメータ$\theta$の方向ということになる。

さて今度は、変分下界を大きくするパラメータ$\theta$の方向を見てみよう。まず下界を少し変形する。

これを$\theta$で偏微分すると

こちらも同時分布の勾配の期待値になった。これらを並べてみよう。

どちらも同時分布の勾配の期待値であるが、確率分布が異なることがわかる。対数尤度の勾配は$\theta$の下での真の事後分布$p_\theta(Z|X)$、下界の勾配は近似事後分布$q_\phi(Z|X)$。もし近似事後分布が真の事後分布と完全に一致していれば、当然これらの勾配も一致する。つまりその地点において下界を大きくする方向が対数尤度を大きくする方向と完全に一致する。実際その場合のKL情報量は0になるので、対数尤度=下界が成り立つ。

近似事後分布が真の事後分布と完全に一致していなくとも、それらがある程度近ければ、これらの勾配もある程度近い値になるだろうと考えられる。つまり下界を大きくする方向に$\theta$を動かせば、対数尤度もそれなりに大きくなることが期待できる。反対に、近似事後分布が真の事後分布から離れている場合はそれらの勾配も離れるだろうと予想できる。つまり、下界を大きくする方向に$\theta$を動かしても対数尤度が大きくならない可能性が高くなる。

以上を踏まえ、下界 $\mathcal L(\theta,\phi)$を$\theta$と$\phi$について更新するとどのようなことが起こるかを考えてみる。$\theta, \phi$をランダムに初期化した場合、初めから$\phi$による事後分布の近似精度が良いことはなかなかないだろう。事後分布の近似精度が悪い場合、下界の勾配と対数尤度の勾配が揃いづらい。つまり学習が開始された初期の段階では、下界を更新しても対数尤度$\ln p_\theta(X)$はあまり上昇しないと予想できる。

ではKL情報量$D_{KL}[q_\phi(Z|X)\|p_\theta(Z|X)]$はどうだろう。下界の中で$\phi$が関わっているのはこの項だけなので、少なくとも$\phi$の勾配はこのKL情報量を小さくする方向に向いていると言える。問題は$\theta$も同時に更新されることであるが、$\theta$に関する勾配もある程度はこのKL情報量を小さくする方向を向いているはずである。対数尤度の勾配がどれだけKL情報量の最小化を邪魔するかは分からないが、$\phi$による更新を完全に打ち消すような影響を毎ステップで与えられるとは考えづらく、更新を繰り返すことでKL情報量の最小化は徐々に達成されていくと考えられる。

KL情報量がある程度小さくなる、すなわち事後分布の近似精度がある程度向上すれば、下界の勾配と対数尤度の勾配が同じ方向を向くことになり、下界の最大化を通じた対数尤度の最大化が実現されるようになる。これが、下界の最大化が対数尤度を大きする説明である。下界の最大化を開始すると、学習初期の段階でこそ対数尤度は上昇しないが、次第に事後分布の近似精度が上がり、対数尤度が上昇するようになる。

VAEにおいて、対数尤度の最大化+事後分布の近似推論を目指す場合でも、対数尤度の最大化のみを目指す場合でも、やることは変わらず下界の最大化であることが説明できた。

VAEの実装

実際に適当なデータを用いてVAEを学習させ、生成モデルを獲得しよう。

MNIST手書き数字画像

MNIST手書き数字画像を学習させ、手書き数字画像の生成モデルを作る。MNIST手書き数字は以下のような画像データ。

Loading image...

実装にあたり、具体的なモデル化を行う必要がある。次のモデルを具体的に定義する。

• 事前分布$p(\bm z)$
• 事後分布$q_\phi(\bm z|\bm x)$
• 条件付き分布$p_\theta(\bm x|\bm z)$

モデル化の章でも述べた通り、事前分布には特に理由もなく標準正規分布を仮定する。また変分下界のKL情報量の計算を楽にしたいので、事後分布にも正規分布を仮定する。

条件付き分布について、今回は、手書き数字画像の各ピクセルが黒か白の二値であると仮定し、ベルヌーイ分布を仮定する。実際はその間の色も多く含まれているが、手書き数字画像における各ピクセルの本質的な情報は「そのピクセルに線が引かれているかどうか」というbooleanなものであるため、二値とみなしてしまう。なお実際に二値に変換するわけではない。あくまでもそう見なすだけ。0と1の間の値でもベルヌーイ分布の確率密度はそれっぽく求められるので、別にそれでいい。

ということでモデル化まとめ。

ではこれを実装していこう。まずはエンコーダから。なお本章で実装するニューラルネットのネットワーク構造は論文の実験で使用されたものと同じ。

class Encoder(nn.Module):
    def __init__(self, w, h, h_dim, z_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(w * h, h_dim),
            nn.Tanh(),
            nn.Linear(h_dim, z_dim * 2)
        )
 
    def forward(self, x):
        x = x.view(len(x), -1)
        params = self.net(x)
        mu, log_sigma = params.chunk(2, dim=-1)
        return mu, log_sigma

画像を受け取って正規分布のパラメータ$\hat{\bm\mu}_\phi, \hat{\bm\sigma}_\phi$を出力するモデルを適当なニューラルネットで定義した。なお$\hat{\bm\sigma}$の非負性を満たすことが難しいため、対数領域で出力している。後で元に戻せばいい。

次にデコーダ。

class Decoder(nn.Module):
    def __init__(self, w, h, h_dim, z_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(z_dim, h_dim),
            nn.Tanh(),
            nn.Linear(h_dim, w * h),
            nn.Sigmoid(),
        )
        self.w = w
        self.h = h
 
    def forward(self, z):
        x_hat = self.net(z)
        x_hat = x_hat.view(-1, 1, self.w, self.h)
        return x_hat

出力されるx_hat が$\hat{\bm\lambda}_\theta$ってこと。

最後にこれをまとめたモデルVAEを実装する。

class VAE(nn.Module):
    def __init__(self, w, h, h_dim, z_dim):
        super().__init__()
        self.encoder = Encoder(w, h, h_dim, z_dim)
        self.decoder = Decoder(w, h, h_dim, z_dim)
 
    @staticmethod
    def reparameterization(mu, log_sigma):
        eps = torch.randn_like(mu)
        z = mu + torch.exp(log_sigma / 2) * eps
        return z
 
    def forward(self, x):
        mu, log_sigma = self.encoder(x)
        z = self.reparameterization(mu, log_sigma)
        x_hat = self.decoder(z)
        return x_hat, mu, log_sigma

エンコーダで事前分布のパラメータを求め、Reparameterization Trickで潜在変数$z'$をサンプリングし、デコーダで$\hat{\bm x}$を得る。最終的に得られた$\hat{\bm x}$に加え、lossの計算に必要な$\hat{\bm\mu}_\phi$と$\hat{\bm\sigma}_\phi$も返す。

ではこれを実際に学習していく。まず損失関数を定義する。負の変分下界がこちら。

再構成誤差はこうなる。バイナリ交差エントロピー。