Variational AutoEncoderのお気持ち

VAEの元論文を読んだ際のほぼ自分用まとめ記事。 VAE自体の分かりやすい解説記事は下記リンクを参照。

猫でも分かるVariational AutoEncoder from Sho Tatsuno

www.slideshare.net

1. Auto-Encoding Variational Bayes (Diederik P Kingma, Max Welling)

https://arxiv.org/abs/1312.6114
連続潜在変数を含んだ確率モデルに対する、スケーラブルな変分推論の方法の提案。
reparameterization trickの導入により、確率的勾配降下法で最適化可能な変分下限を導出した。これにより、潜在変数モデルを効率的に近似推論できるようになった。

MCMCは計算量が多く、平均場近似による変分推論では解析的に事後分布を導出することができないことが多い。一方、Stochastic Gradient Variational Bayesは連続変数を含むほとんどどのようなモデルにも適用可能。

潜在変数zの事前分布を設定。xのzによる条件付き分布からxが生成。いずれもパラメータθで規定された確率分布であり、pdfが微分可能とする。

f:id:kibya:20200518215705p:plain — VAEのグラフィカルモデル

推論の目的は下記の3つ。 1. パラメータθの推定パラメータ自体に興味がある場合。また、人工データを生成したい場合。 2. 潜在変数zの事後分布の推定暗号化やデータ表現のタスクに有用。 3. xの周辺尤度の推論画像のノイズ除去など。

潜在変数zの事後分布の近似分布q(z|x)を考える。確率分布のパラメータφとする。 q(z|x)をエンコーダ、p(x|z)をデコーダと呼ぶ。

変分下限を変形すると、第１項のKLダイバージェンスはq(z|x)の分布に関する制約である正規化項。第２項はEncoder, Decoderを表す項であり、Encoderで推論したzの確率分布を用いて、対数尤度logp(X|z)の期待値を取っている。

f:id:kibya:20200518220306p:plain

2. Semi-Supervised Learning with Deep Generative Models (Diederik P. Kingma, Danilo J. Rezende, Shakir Mohamed, Max Welling)

https://arxiv.org/abs/1406.5298
生成モデルを用いた半教師あり学習。ラベル割り当てのコストが高い、または不可能であり、一部のデータにしかラベルがないケースがある。ラベルがないデータも活用して、どのように分類精度を上げるかが問題。

3. Tutorial on Variational Autoencoders (Carl Doersch)

[1606.05908] Tutorial on Variational Autoencoders

1. Introduction

生成モデルは確率分布p(X)を取り扱うモデル。Xは高次元空間におけるdatapoint。
例えば画像生成の場合、datapoint Xは画像であり、Xは多数のpixelからなる高次元データ。
生成モデルの用途としては、既存のデータに似たデータを生成したいケースがある。
ある確率分布にしたがって生成されたサンプルから、なるべく元の確率分布に近いモデルを学習することで、入力と類似したサンプルを生成できる。

潜在変数モデル

0-9の数字の画像を生成する問題を考える。モデルから生成する際は、まず0-9の数字から生成対象とする数字zをランダムに決定し、その数字の画像になるように各ピクセルの値を生成する。zを潜在変数と呼ぶ。簡単にサンプリング可能な確率分布p(z)に従う潜在変数zと決定論的な関数f(z;theta)を用いて観測変数xのサンプル（データ点）が得られる。 f(z;theta)が高確率で実際のデータに近くなるように、パラメータthetaを最適化したい。

f:id:kibya:20200531153722p:plain

f(z;theta)をP(x|z;theta)に修正すると、(1)式のP(X)を最大化する問題となる。これは最尤推定である。 VAEではP(x|z;theta)は下記のように平均f(z;theta)の正規分布とすることが多いが、例えばXが二値の場合はベルヌーイ分布など、他の分布でも良い。

f:id:kibya:20200531211011p:plain

Variational Autoencoders

(1)式を解く上では、潜在変数zをどう決めるか、zに関する積分をどう取り扱うか、の２点を決める必要がある。

画像を生成する際に、潜在変数zの各次元が表す内容や次元間の相関を明示的に決めるのは難しい。
そこで、VAEでは潜在変数zは単純な正規分布（共分散行列が単位行列）からサンプル可能とする（P(z)=N(0, I))。単純な分布からのサンプルを、複雑な関数で変換することで、モデルが必要とする潜在変数に変換できる、と仮定している。

f(z;theta)が多層ニューラルネットワークとすれば、最初の何層かで正規分布に従うzを潜在変数の値（文字の太さ、角度など）にマッピングし、残りの層で潜在変数を観測される数字の画像に変換することを想定できる。（この点については、理論的な保証はなさそう？）
実際に上記のような潜在的な構造があるかは心配する必要がなく、もしそのような構造が数字の画像の生成に役に立つならばいずれかの層で学習されるだろう。

残った問題は、zが標準正規分布に従う時、(1)式のP(x)をどうやって最大化するかである。機械学習の分野で一般的なように、P(X)の勾配をとって、SGDで最適化できる。そのためには、P(x)の計算式が必要になる。概念的には、zのサンプルを多数取得し、P(x|z)の期待値を取れば良いが（単純モンテカルロ法）、高次元空間においてP(X)を正確に計算するには無数のサンプルが必要になるため難しい。

実際にところ、ほとんどのzについて、P(X|z)は0に近い値になる。 VAEのキーとなる考え方は、Xを生成したであろうzのみを用いて、P(x)を計算することである。そこで、新しい関数Q(z|X)を導入する。QはXを受け取り、Xを生成しやすいようなｚの分布を与える。 Qの元でもっともらしいzの空間は、事前分布p(z)のもとでのzの空間よりも小さいことが期待される。それにより、例えば $E _ {z \sim Q} P(X \mid z)$ の計算が容易になる。

次に、 $E _ {z \sim Q} P(X \mid z)$ とP(X)と関係が、変分ベイズ法のポイントとなる。

簡単な式変形により、logP(X)の下限（変分下限、ELBO）は下式で表されることが分かる。

f:id:kibya:20200531215113p:plain

左辺の第2項のKLダイバージェンスを最小化することが、ELBOの最大化と同値となる。P(z|X)は解析的に求めることはできないが、Q(z|X)を任意の精度を持つモデルで表現すると、P(z|X)に一致することが期待され、KLダイバージェンスは０となる。そのため、右辺の最大化はP(X)の最大化と同値となる。

ELBOの最大化の際は、上式の右辺をSGDにより最大化する。右辺はencode, decoderを表す第１項と、zの近似事後分布と事前分布のKLダイバージェンスの第２項に分解できる。
zの分布系を正規分布とする場合、第２項は正規分布どうしのKLダイバージェンスは解析的に計算可能。

第１項の期待値を正確に計算するためにはzのサンプルを多数得る必要があるが、計算コストが高い。
そこで、SGDで標準的な通り、zのサンプルを1つ得て、そのzに対するP(X|z)を $E _ {z \sim Q} P(X \mid z)$ の期待値として取り扱う。

実際は、データセットDのサンプルであるXを用いて、下式を最大化する事になる。 f:id:kibya:20200531220839p:plain

この式の勾配をとると、勾配の計算は期待値の中に移動することができる。そのため、まずXを1つサンプルし、Q(z|X)からzを1つサンプルし、下式の勾配を求めれば良い。

f:id:kibya:20200531221636p:plain

ただし、この式で表現されるネットワークは下式に左図の場合にあたる。この場合、順伝播は問題ないが、バックプロパゲーションの際、Q(Z|X)からzをサンプルする層が非連続的となり、勾配が計算できない。
そこでreparametrization trickが必要になる。

reparametrization trickは、 $N( \mu(X), \Sigma(X))$ を計算する際に、最初に $\varepsilon \sim N(0, I)$ を得て、次に $z = \mu(X) + \Sigma ^ {1/2}(X) * \varepsilon$ を計算する。この結果、勾配を求める式は下記のようになる。上図の右側の場合に相当する。