Tutorial on Variational Autoencodersを読む

Tutorial on Variational Autoencoders(VAEs)を読み解いていこうと思う。
先人たちによる日本語の詳細な解説はネット上にゴロゴロあるので、
本記事は自分自身の理解のためのメモという位置づけ。

潜在変数

生成モデルは、高次元空間におけるデータ{X}の分布{P(X)}自体をモデル化するもののこと。

例えば、MNISTのような0〜9からなる手書き文字を生成する目的の生成モデルがある. このとき、モデルの生成対象は0〜9の文字なので、 モデルはまずどの文字を生成するかを決めて、その決定のもとでデータを生成するものだ,と仮定してみる。

このときの予め決定しておくパラメータのことを潜在変数(latent variable) {z}と呼ぶ。 モデルが文字を描く前に、{z}を[0,…,9]からランダムにサンプリングするというイメージ. なぜ「潜在」と呼んでいるかというと、どの潜在変数が文字を生成したかについては必ずしも知っておく必要がないため。

潜在変数{z}を決める確率分布を{P(z)}とする。 固定のパラメータ{\theta}のもとで、ランダムにサンプリングした{z}から、{X}を生成する。 この{X}が訓練データのものとよく一致すればよい生成モデルと言えそう.

数学的に記述すると、次のようになる。 {} $$ \begin{equation} P(X) = \int P(X \, | \, z;\theta) P(z) dz \tag{1} \end{equation} $$

訓練データ{X}の周辺尤度を最大化することで,潜在変数{z}を推定する.

VAEでは、{P(X \, | \, z;\theta)}ガウス分布をよく使う。
この場合のガウス分布{P(X \, | \, z;\theta) = \mathcal{N}(X \, | \, f(z;\theta), σ^{2}*I)}となる。

VAEs

VAEはいわゆるAutoencoderとはあまり関係がない。 ではなぜAutoencoderと呼ばれているかというと,最終的な目的関数を構成するのがencoderとdecoderからなるから.

上記の(1)を解くにあたって、VAEは以下の2つの問題を解決しなくてはならない。

  • 潜在変数{z}をどう定義するか(どのような情報を表現するものと定義するか)
  • 潜在変数{z}積分をどう扱うか。

1つめの問題は「文字の角度」「文字のストローク」「ストロークの太さ」「文字のスタイル」…など様々な性質が関わって複雑なため、できれば直接{z}の表す情報を決定するのは避けたい。 VAEではこの問題に対して、変わったアプローチをとっている。 それは、{z}を単純な分布、例えば正規分布{\mathcal{N}(0,I)}からサンプリングするというもの。 本当にこのアプローチは正しいのだろうか。

実は{d}次元のどのような分布であっても、正規分布からサンプルされた{d}個の変数を複雑な関数でマッピングすれば表現可能なのである。 例として以下の図のようなリング状の二次元空間状の分布を得たい場合を考える。 正規分布からサンプルされた{z}を、{g(z) = z/10 + z/ ||z||}写像すると、うまいことリング状になる。

f:id:yusuke_ujitoko:20170525215054p:plain (arxiv論文から引用)

したがって、表現力豊かな関数があるのであれば、normally-distributedな{z}を用いて、目的の分布を得ることができ、訓練データ{X}を表現できる。

{P(X)}を計算するためのコンセプトは明快。

  1. {z}をサンプリングする( {\left[z_{1}, \cdots, z_{n} \right] })
  2. {P(X) = \frac{1}{n} \sum P(X \mid z_{i}) }を計算する

ここでの問題は、 高次元空間の場合、正確に{P(X)}を推定するには、{n}が非常に大きくなってしまうということ。

目的関数の設定

(1)式を簡単に計算する方法はなにかないだろうか? よく考えると、大抵の{z}に対して、{P(X \mid z)}はほぼ0となることに気づく。 大抵の{z}はなにも寄与しない。

VAEはこの考え方を利用していて、{X}を生成するような{z}のみをサンプリングしようとする。
そこで{Q(z \mid X)}という新たな関数を考える。 この関数は、{X}のもとで、{z}の分布を与えるもので、逆から考えるとその{z}は結局{X}を生成するようなものである。 このような{z}の空間は事前確率{P(z)}において極めて小さい。 これにより、例えば{\mathbb{E}_{z \sim Q}P(X \mid z)}の計算が簡単になる。

一方、{P(X)}の計算は簡単になるだろうか。 とりあえず最初にしてみるのは{\mathbb{E}_{z \sim Q}P(X \mid z)}{P(X)}の関係付けである。

{\mathbb{E}_{z \sim Q} P(X \mid z)}{P(X)}の関係は、変分ベイズで片付けられる。 まず{P(X \mid z)}{Q(z)}間のKL divergenceを任意の{Q}に対して定義する。 {} $$ D \left[ Q(z) \mid\mid P(z \mid X) \right] = \mathbb{E}_{z \sim Q} \left[ \log Q(z) - log P(z \mid X) \right] $$

この式に対して、ベイズの定理を適用して変換して、さらに{z}に無関係な項を期待値からはずす操作を施す。 {} $$ D \left[ Q(z) \mid\mid P(z \mid X) \right] = \mathbb{E}_{z \sim Q} \left[ \log Q(z) - \log P(X \mid z) - \log P(z) \right] + \log P(X) $$

両辺にマイナスをかけて、期待値の項の中身の一部をKL divergenceに書き換える。 {} $$ \log P(X) - D \left[ Q(z) \mid\mid P(z \mid X) \right] = \mathbb{E}_{z \sim Q} \left[ \log P(X \mid z) \right] - D \left[ Q(z) \mid\mid P(z) \right] $$

{X}は固定値であり、{Q}はあらゆる分布がありうることに注意しておく。 {P(X)}を推論することに興味があるため、{X}に依存するような{Q}を構築したい。 それはつまり{D \left[ Q(z) \mid\mid P(z \mid X) \right]}が小さくなるような場合。 {} $$ \begin{equation} \log P(X) - D \left[ Q(z \mid X) \mid\mid P(z \mid X) \right] = \mathbb{E}_{z \sim Q} \left[ \log P(X \mid z) \right] - D \left[ Q(z \mid X) \mid\mid P(z) \right] \tag{2} \end{equation} $$

この式はVAEのコアになる部分なので詳しくみていく。

  • 左辺の{P(X)}は最大化したい。 (一方で、左辺の{D \left[ Q(z \mid X) \mid\mid P(z \mid X) \right]}を最小化しつつ)
  • 右辺は、勾配降下法でうまく{Q}を選んでいけば最適化できる。 この右辺の項はAutoencoderとみなせる。 {Q}{X}{z}にencodeし、一方{P}{z}から{X}を復元しようとしている。

左辺では{ D \left[ Q(z \mid X) \mid\mid P(z \mid X) \right]}を最小化しつつ、{P(X)}を最大化する。

{P(z | X)}は解析的に計算できない。 しかし、左辺の二項目の{Q(z | x)}{P(z | x)}に近づいていくことを考えてみる。 もし、{Q(z | x)}の表現力が高く、{P(z | x)}と近い場合には、これらの間のKL divergenceは0になる。 この仮定のもとでは {\log P(X)}のみを最適化すればよい。

目的関数を最適化する

(2)の右辺をどのように最適化すればよいだろうか? {Q(z | X)}をどうするかという問いになるが、 普通は {Q(z | X) = \mathcal{N} (z | \mu(X;\theta), Σ(X;\theta))}とする。 ここで{\mu}{Σ}はパラメータ{\theta}を持つ任意の関数でデータから学習する。 実際、{\mu}{Σ}ニューラルネットワークで決める。 この手法は計算量は小さいのが利点。

もうひとつの項は、{D \left[ Q(z|X || P(z) \right]}はKL divergenceである。 多変量ガウス分布のKL divergenceは、次のようなclosed formで書ける。 {} $$ D [ \mathcal{N}(\mu_{0}, Σ_{0}) || \mathcal{N}(\mu_{1}, Σ_{1}) ] = \frac{1}{2} \bigl( tr(Σ_{1}^{-1} Σ_{0}) + (\mu_{1} - \mu_{0})^{T} Σ_{1}^{-1} (\mu_{1} - \mu_{0}) -k + \log(\frac{det Σ_{1}}{det Σ_{0}}) \bigr) $$ 今回に適用するとシンプルになる。 {} $$ D [ \mathcal{N}(\mu(X), Σ(X)) || \mathcal{N}(0,I) ] = \frac{1}{2} \bigl( tr(Σ(X)) + (\mu(X))^{T} (\mu(X) -k - \log det (Σ(X)) \bigr) $$

さて、(2)の式に戻る。 (2)の右辺第一項に対する対応は、{z}を1つサンプルして、それに基づく{P(X \mid z)}{\mathbb{E}_{z \sim Q}[log P(X \mid z)}とみなすこととする。

最後に、データセット全体からサンプル{X}をとって期待値を計算すると、 最適化したい式は次のようになる。 {} $$ \begin{equation} E_{X \sim D} [ log P(X) - D[Q(z|X) || P(z|X)] ] = E_{X \sim D} [ E_{z \sim Q} [ log P(X | z)] - D[Q(z|X) || P(z) ]] \tag{3} \end{equation} $$ この式に基づいて勾配計算をするときには、期待値の中身を計算すればよい。
つまり{X}{Q(z|X)}をサンプルして、以下の勾配を計算する。 {} $$ log P(X|z) - D[ Q(z|X) || P(z) ] $$

ところが実は{\mathbb{E}_{z \sim Q}\left[log P(X \mid z)\right]}の計算に際して 深刻な問題 がある。 {log P(X \mid z)}{P}だけでなく{Q}にも依存している。 しかし、(3)ではこの依存が消滅してしまっている。 VAEを上手く機能させるためには、{Q}を経た上で、{P}{X}を復元させるようにする必要がある。

下図の(左)が問題を示している。

f:id:yusuke_ujitoko:20170525224357p:plain (arxiv論文から引用)

forward passは問題ない。 しかしbackpropagate時に、計算グラフが途切れているため、誤差を{z}から{Q(z \mid X}へ伝播することができない。

ここで reparameterization trick を使う。
この手法ではサンプリング操作の代わりに{z}{Q}の出力結果を元に関数で生成する。 具体的には、{\mu(X)}{Σ(X)}のもとで、{z = \mu(X) + Σ^{\frac{1}{2}} \odot \epsilon}として{z}を計算する。(ここで{\epsilon}は最適化とは無関係のハイパーパラメータ) よって実際に勾配を取る式は次のようになる。 {} $$ E_{X \sim D} \left[ E_{\epsilon \sim N(0, I)}[log P(X|z = \mu(X) + Σ^{½} \odot \epsilon)] - D[Q(z|X) || P(z) ] \right] $$

学習済みモデルを試すためには

新しいデータを生成したい時には、decoderに{z \sim N(0,I)}を投入する。 単純にVAEからencoderを取り除く操作となる。

いらすとや画像を生成してみる

VAEをtensorflowで実装し, いらすとや画像を生成してみた.
実装に誤りがあるのか, 現状うまく生成できなかった..

f:id:yusuke_ujitoko:20170526000145p:plain

mode collapseが起きている。

(追記)

上の記述だと, 式変形を追いにくいので,全体の変形をまとめて書いてみる. 上よりも式変形を補った. {} $$ \begin{align} D \left[ Q(z) \mid\mid P(z \mid X) \right] &= \int Q(z) [\log Q(z) - \log P(z \mid X)] dz \\ &= \mathbb{E}_{z \sim Q} \left[ \log Q(z) - \log P(z \mid X) \right] \\ &= \mathbb{E}_{z \sim Q} \left[ \log Q(z) - \log P(X \mid z) - \log P(z) + \log P(X) \right] \\ &= \mathbb{E}_{z \sim Q} \left[ \log Q(z) - \log P(X \mid z) - \log P(z) \right] + \log P(X) \\ &= \mathbb{E}_{z \sim Q} \left[ \log Q(z) - \log P(z)\right] - \mathbb{E}_{z \sim Q} \left[ \log P(X \mid z) \right] + \log P(X) \\ &= D \left[ Q(z) \mid\mid P(z) \right] - \mathbb{E}_{z \sim Q} \left[ \log P(X \mid z) \right] + \log P(X) \end{align} $$ 左辺と右辺を入れ替えて, {} $$ \begin{align} \log P(X) - D \left[ Q(z \mid X) \mid\mid P(z \mid X) \right] &= \mathbb{E}_{z \sim Q} \left[ \log P(X \mid z) \right] - D \left[ Q(z \mid X) \mid\mid P(z) \right] \\ \log P(X) & \geq \mathbb{E}_{z \sim Q} \left[ \log P(X \mid z) \right] - D \left[ Q(z \mid X) \mid\mid P(z) \right] = L(\theta, \phi; x) \end{align} $$ 左辺を最大化する代わりに右辺を最大化する.