Variational Auto Encoder(VAE)を理解したい。

はるか昔に流行った生成AIとしてVariational Auto Encoder(VAE)というものがありますが、どういう仕組みなのかを理解しようとするたびに挫折した経験があるので今度こそは理解したいと思い再トライしてみました。

VAE概観

• Auto Encoder同様にEncoder-Decoder構造

• 潜在変数に確率分布(正規分布)を設定、Encoderは分布のパラメータ(𝜇, 𝜎)を推定

• 対数周辺尤度を直接最大化するのは大変なのでELBO(変分下限)の最大化を考える

• Encoder出力から潜在変数のサンプリングはニューラルネットが学習できる(誤差逆伝播法を適用できる)ように変数変換トリックを使う

Decoder

そもそも話~Decoderの導入~

やりたかったこと：生成モデルの作成！

VAEでは…

• 潜在変数 𝑧 を固定の正規分布から生成する (正規分布は標準正規分布 𝒩(𝑧;0, 𝐼) を仮定する)
• ニューラルネットワークによって潜在変数 𝑧 から観測変数 𝑥 へと変換する

目標：観測変数 𝑥 の確率分布 $𝑝_𝜃 (𝑥)$ を得ること

• ニューラルネットワークは確率分布として $𝑝_𝜃 (𝑥|𝑧)$ をモデル化する必要がある
• ニューラルネットワークの出力はベクトルなので「ニューラルネットワークの出力を平均ベクトルとする正規分布を考える」
• 分散は定数(ここでは単位行列)を使ったり、学習パラメータにしたり諸説あるらしい

定式化：

• $\hat{𝑥} = 𝑁𝑒𝑢𝑟𝑎𝑙𝑁𝑒𝑡(𝑧;𝜃)$
• $𝑝_𝜃 (𝑥│𝑧)=\mathcal{N}(𝑥;\hat{𝑥}, 𝐼)$

問題点：
最尤推定をするなら $\log p_\theta(x) = \log \int p_\theta(x|z)p(z)$ が必要になるが、積分があり一般的には行えない…。
最尤推定以外の手段を考える必要がある！

EMアルゴリズム

最尤推定以外の手段が必要と分かったがいったいどうすればよいのか？ということでEMアルゴリズム(本当は変分ベイズ)の登場です。

任意の確率分布 𝑞(𝑥) を使って対数尤度を変形する (結果的には $𝑝_𝜃 (𝑧|𝑥)$ を近似するものになる)

手に負えない $\log⁡ 𝑝_𝜃 (𝑥)$ の代わりに
ELBO： $\int 𝑞(𝑧) \log⁡\frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞(𝑧)} 𝑑𝑧$
の最大化を考えることでお茶を濁す作戦(最適解になるとは言っていない)

EMアルゴリズムではどうやってELBO： $\int 𝑞(𝑧) \log⁡\frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞(𝑧)} 𝑑𝑧$を最大化する？
→2つのパラメータ 𝑞, 𝜃 を交互に更新する！
(EステップとMステップを繰り返す)

Eステップ：𝑞を更新する

• $\log ⁡𝑝_𝜃 (𝑥) = 𝐸𝐿𝐵𝑂(𝑥;𝑞, 𝜃)+𝐷_𝐾𝐿 (𝑞(𝑧)||𝑝_𝜃 (𝑧|𝑥))$ の左辺は 𝑞 によらないので、𝜃 を固定するとELBO項とKL項の和は一定である
• よって 𝜃 を固定すると、KL項を小さくするほどELBO項は大きくなる
• ELBO項が最大となるのは「KL項=0」の時であり、$𝑞(𝑧)=𝑝_𝜃 (𝑧|𝑥)$ の時である
• さらに、KL項は0なので、 $\log⁡𝑝_𝜃 (𝑥) = 𝐸𝐿𝐵𝑂(𝑥;𝑞, 𝜃)$ となる
• 以上より 𝜃 を固定すると 𝑞(𝑧) の更新式は、$𝑞(𝑧)=𝑝_𝜃 (𝑧|𝑥)$ と表せる

https://qiita.com/kenmatsu4/items/59ea3e5dfa3d4c161efb

Mステップ：𝜃 を更新する

• ELBO項の 𝜃 に関する微分を求めて、それを0にした方程式を解くことで解析的に(現状の)最適解を求められる

EMアルゴリズムの重要な性質：$\log ⁡𝑝_{𝜃_{𝑛𝑒𝑤}} (𝑥) ≥ \log ⁡𝑝_{𝜃_{𝑜𝑙𝑑}} (𝑥)$

→ 対数周辺尤度は単調に増加する(減少しない)[証明略]

複数データに対するEMアルゴリズム

N個の観測データ $(𝑥^{(1)}, 𝑥^{(2)}, …, 𝑥^{(𝑁)} )$ に対しての拡張
各データに対応した任意の確率分布を $(𝑞^{(1)}, 𝑞^{(2)}, …, 𝑞^{(𝑁)} )$ としてN個用意する
このときの対数尤度とELBOの関係は以下：

\sum_{n=1}^N \log ⁡𝑝_𝜃 (𝑥) = \sum_{n=1}^N ( 𝐸𝐿𝐵𝑂(𝑥;𝑞, 𝜃)+𝐷_𝐾𝐿 (𝑞(𝑧)||𝑝_𝜃 (𝑧|𝑥)) )

Eステップにおいて、$(𝑥^{(1)}, 𝑥^{(2)}, …, 𝑥^{(𝑁)} )$の各データに対する最適な分布は $𝑞^{(𝑛)} (𝑧) = 𝑝_𝜃 (𝑧|𝑥^{(𝑛)})$ であり、その時に総和 ($∑_{𝑛=1}^𝑁$・) も最適になる

以上より、EMアルゴリズムは以下のようにまとめられる：

• Eステップ： $(𝑞^{(1)}, 𝑞^{(2)}, …, 𝑞^{(𝑁)} )$ の更新 (𝜃 は固定)
  各 𝑛 に対し $𝑞^{(𝑛)} (𝑧) = 𝑝_𝜃 (𝑧|𝑥^{(𝑛)})$ とする
• Mステップ：𝜃 の更新 ($(𝑞^{(1)}, 𝑞^{(2)}, …, 𝑞^{(𝑁)} )$ は固定)
  $\sum_{n=1}^N ELBO(x^{(n)}; q^{(n)}, \theta)$ が最大になる 𝜃 を解析的に求める
• 終了判定：対数尤度を計算し、前回の対数尤度と比較する

Encoder

EMアルゴリズムのVAEへの適用

EMアルゴリズムのVAEへの適用(1データ)は以下の流れ：

これまでのものをそのまま適用しようとすると、

\log⁡𝑝_𝜃 (𝑥) = \int 𝑞(𝑧) \log⁡ \frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞(𝑧) } 𝑑𝑧 + 𝐷_{𝐾𝐿} (𝑞(𝑧)||𝑝_𝜃 (𝑧|𝑥)) ≥ \int 𝑞(𝑧) \log⁡\frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞(𝑧)} 𝑑𝑧  (=𝐸𝐿𝐵𝑂)

の最適化時に問題点が発生する。

• EMアルゴリズムのEステップでは、𝑞(𝑧) を $𝑞(𝑧)=𝑝_𝜃 (𝑧|𝑥)$ と事後分布で更新する
• しかし、VAEでは事後分布を直接求めることは困難…。
• そのため以下のアプローチをとる：
  1. 𝑞(𝑧) を簡単な分布 (正規分布：パラメータは 𝜓={𝜇, Σ}) に限定する ← この分布を与えて近似する手法が変分ベイズ、変分近似と呼ばれる
  2. その限定された分布の中でELBOの最大化を行う→ $𝑞_𝜓 (𝑧)=\mathcal{N}(𝑧;𝜇, Σ)$

上記より、

\log⁡𝑝_𝜃 (𝑥) = \int 𝑞_\psi(𝑧) \log⁡ \frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞_\psi(𝑧) } 𝑑𝑧 + 𝐷_{𝐾𝐿} (𝑞_\psi(𝑧)||𝑝_𝜃 (𝑧|𝑥)) ≥ \int 𝑞_\psi(𝑧) \log⁡\frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞_\psi(𝑧)} 𝑑𝑧  (=𝐸𝐿𝐵𝑂)

に書き換わる。

以上より、目標はELBOを 𝜃, 𝜓 に関して最大化することになる。

また、ELBOの最大化の過程で、$𝑝_𝜃 (𝑧, 𝑥)$ にフィットする $𝑞_𝜓 (𝑧)$ が得られる。

EMアルゴリズムのVAEへの適用(データセット全体)は以下の流れ：

N個の観測データ $(𝑥^{(1)}, 𝑥^{(2)}, …, 𝑥^{(𝑁)} )$ に対しての拡張：

• $
  \sum_{n=1}^N ELBO(x^{(n)};\psi^{(n)}, \theta) = \int 𝑞_\psi(𝑧) \log⁡\frac{𝑝_𝜃 (𝑥,𝑧)}{𝑞_\psi(𝑧)} 𝑑𝑧
  $ を目的関数とする

• 各データ $𝑥^{(𝑛)}$ に対応した確率分布 $𝑞_(𝜓^{(𝑛)} ) (𝑧)$ を用意する。
  ※$𝑞_(𝜓^{(𝑛)} ) (𝑧)$ はパラメータが $𝜓^{(𝑛)}=(𝜇^{(𝑛)}, Σ^{(n)})$ の正規分布

• ELBOを最大化することで $p_\theta (z^{(n)}, x^{(n)})$ にフィットする $𝑞_{𝜙^{(𝑛)} } (𝑧)$ が得られる

• しかし、めでたしめでたし…とはいかない…。

データセット全体時の問題点

• 各データ 𝑥^((𝑛)) に対しパラメータ 𝜓^((𝑛)) が必要になる
  ※データセットが1億データなら1億個のパラメータが必要

• それはちょっとしんどいんですけど…。

• せや、 $𝑥^{(𝑛)} から 𝜓^{(𝑛)}$ を出力するニューラルネットワークがあればいいじゃん！ → Encoderの必要性！

Encoderについて

• 入力を 𝑥 、出力を 𝜇, Σ とする

• 計算コスト削減のため、分散共分散行列は以下のように対角行列とする

\Sigma=\begin{pmatrix}
\sigma_1^2 & 0 & \cdots & 0\\
0 & \sigma_2^2 & \cdots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0 & \cdots & \sigma_H^2
\end{pmatrix}

• 以下のように$\sigma$を定義すると、$𝜎^2 𝐼 = \Sigma$ であり学習パラメータは 𝜎 と考えられる

\sigma=\begin{pmatrix}
\sigma_1^2\\
\sigma_H^2
\end{pmatrix}

定式化：

• $𝜇, 𝜎=𝑁𝑒𝑢𝑟𝑎𝑙𝑁𝑒𝑡(𝑥;𝜙)$
• $𝑞_𝜙 (𝑧│𝑥) = \mathcal{N}(𝑧;𝜇, 𝜎^2 𝐼)$

ELBO評価

これまでのまとめ

登場人物は以上なのでいったん整理する。

• Decoderの定式化：
  • $\hat{𝑥} = 𝑁𝑒𝑢𝑟𝑎𝑙𝑁𝑒𝑡(𝑧;𝜃)$
  • $𝑝_𝜃 (𝑥│𝑧)=\mathcal{N}(𝑥;\hat{𝑥},𝐼)$

• Encoderの定式化：
  • $𝜇, 𝜎=𝑁𝑒𝑢𝑟𝑎𝑙𝑁𝑒𝑡(𝑥;𝜙)$
  • $𝑞_𝜙 (𝑧│𝑥)=\mathcal{N}(𝑧;𝜇, 𝜎^2 𝐼)$

• 目的関数

  • ELBO

• 最適化手法

  • EMアルゴリズム

んっ？学習パラメータってニューラルネットワークだけ？
EMアルゴリズム2ステップ踏まないといけないしめんどくない？
それ、勾配法で誤差逆伝播できませんか？

ということで勾配法での誤差逆伝播の導入を考えます。

ELBO整理

勾配法で頑張りたいので、ELBOについてちゃんと考えます。

変数変換トリック(Reparameterization Trick)

まとめ

• (結果的に)Auto Encoder同様にEncoder-Decoder構造

• 潜在変数に確率分布(正規分布)を設定、Encoderは分布のパラメータ(𝜇, 𝜎)を推定

• 対数周辺尤度を直接最大化するのは大変なのでELBO(変分下限)の最大化を考える、最適解が得られる保証はない

• Encoder出力から潜在変数のサンプリングはニューラルネットが学習できる(誤差逆伝播法を適用できる)ように変数変換トリックを使う

参考