変分オートエンコーダVAEに入門してみる

# はじめに

世間では生成AIのニュースで賑やかです。生成AIには、①GAN、②VAE、③Diffusion modelなどがあります。変分オートエンコーダ(Variational Autoencoder:VAE)を知らなかったので調査しました。本記事では、②VAEについて簡単に紹介します。


![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/139678/389d6bd7-ccb5-6926-9dcc-3f27a3a217d4.png)
[What are Diffusion Models?](https://lilianweng.github.io/posts/2021-07-11-diffusion-models/)

# VAEの実装

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/139678/a0677211-04a2-106d-7c7e-6b0907b84531.png)
[VAEのPytorch実装](https://github.com/beginaid/VAE)

# 変分下限(Evidence Lower Bound:ELBO)
対数尤度$\log p(x)$を直接最大化することは困難です。そこで、対数尤度$\log p(x)$のELBO$\mathcal{L}(x,z)$を最大化することで、間接的に対数尤度$\log p(x)$を最大化することを考えます。

```math
\begin{eqnarray}
\log p(x)
&=&\log \int p(x,z)dz\\
&=&\log \int \dfrac{p(x,z)q_{\phi}(z|x)}{q_\phi(z|x)}dz\\
&=&\log \mathbb{E}_{q_\phi}(z|x)\Biggl\lbrack\dfrac{p(x,z)}{q_\phi(z|x)}\Biggr\rbrack\\
&\geq&\mathbb{E}_{q_\phi}(z|x)\Biggl\lbrack\log \dfrac{p(x,z)}{q_\phi(z|x)}\Biggr\rbrack \ \ (\because \text{Jensen's Inequality})\\
&=&\mathcal{L}(x,z)
\end{eqnarray}
```

# VAEの学習

対数尤度$\log p(x)$を直接最大化することは困難です。そこで、対数尤度$\log p(x)$のELBO$\mathcal{L}(x,z)$を最大化することで、VAEのEncoder$q_\phi(z|x)$を構成するパラメータ$\phi$と、Decoder$p_\theta(x|z)$を構成するパラメータ$\theta$を最適化することを考えます。間接的に、対数尤度$\log p(x)$を最大化できるはずです。

```math
\begin{eqnarray}
\mathcal{L}(x,z)
&=&\mathbb{E}_{q_\phi(z|x)}\Bigl\lbrack \log \dfrac{p(x,z)}{q_\phi(z|x)}\Bigr\rbrack\\
&=&\mathbb{E}_{q_\phi(z|x)}\Bigl\lbrack \log \dfrac{p_{\theta}(x|z)p(z)}{q_\phi(z|x)}\Bigr\rbrack\\
&=&\mathbb{E}_{q_\phi(z|x)}\Bigl\lbrack \log p_{\theta}(x|z)\Bigr\rbrack
 + \mathbb{E}_{q_\phi(z|x)}\Bigl\lbrack \log \dfrac{p(z)}{q_\phi(z|x)}\Bigr\rbrack\\
&=&\mathbb{E}_{q_\phi(z|x)}[\log p_{\theta}(x|z)]-\text{KL}(q_\phi(z|x)||p(z))
\end{eqnarray}
```

ELBOを最大化することを考えます。ELBOと目的関数の$\pm$の符号が違うことに注意すれば、式(1)の最小化問題を解けば、EncoderとDecoderのパラメータ$\theta,\phi$が求まります。

```math
\min_{\theta,\phi} \dfrac{1}{N}\sum_{n=1}^{N}L_D(x_n)+L_E(x_n) \tag{1}
```
ただし、

```math
\begin{eqnarray}
L_D(x)&=&-\mathbb{E}_{z\sim q_\phi}[\log p_\theta(x|z)]\\
L_E(x)&=&\text{KL}(q_\phi(z|x)||p(z))
\end{eqnarray}
```

ここで、$L_D(x)$は、Decoderが$x$をどれだけ再構成ができているかを評価しています。$x$が2値をとる場合は、モンテカルロ法から$L_D(x)$は交差エントロピーとなります。$L_E(x)$は、Encoderの$q_\phi(z|x)$が事前分布$p(z)$とどれだけ同じかを評価しています。事前分布$p(z)$は、人が勝手に決める設計情報で、よく正規分布などが用いられます。そのときは、$q_\phi(z|x)$も正規分布なので、$L_E(x)$は解析解をえることができます。


```math
\begin{eqnarray}
q_\phi(z|x)&=&\mathcal{N}(z;\mu_{\phi}(x),\sigma^{2}_{\phi}(x)I)\\
p(z) &=& \mathcal{N}(z;0,I)
\end{eqnarray}
```

デコーダーの項$L_D(x)$は、モンテカルロ法から近似計算をします。

```math
\begin{eqnarray}
&&\min_{\theta,\phi} \dfrac{1}{N}\sum_{n=1}^{N}L_D(x_n)+L_E(x_n)\\
&&\simeq \min_{\theta,\phi} -\sum_{l=1}^{L}\log p_{\theta}(x|z^{(l)}) 
        + \text{KL}(q_\phi(z|x)||p(z))
\end{eqnarray}
```
$\lbrace z^{(l)}\rbrace_{l=1}^{L}$は、$q_\phi(z|x)$からサンプリングしたものです。
しかし、$z$をサンプリングする操作は微分不可能であり、Encoder側に勾配を伝えることが出来ません。そこで、reparameterization trickを用いることで微分可能にします。これは、標準正規分布から$\varepsilon$をサンプリングし、$z$を再構成することを考えます。

```math
z = \mu_{\phi}(x)+\varepsilon\circ\sigma_{\phi}(x) \ \ (\text{with} \ \ \varepsilon\sim \mathcal{N}(\varepsilon;0,I))
```


```python
# Encoder側からサンプリングする
mean, log_var = self.encoder(x)

# reparameterization trick
z = self.sample_z(mean, log_var, device)
y = self.decoder(z)

# Lower Boundを計算する
L_D = -torch.sum(x * torch.log(y + self.eps) + (1 - x) * torch.log(1 - y + self.eps))
L_E = -0.5 * torch.sum(1 + log_var - mean**2 - torch.exp(log_var))

loss = L_E + L_D
```

# 参考文献
- [深層生成モデルを巡る旅(2): VAE](https://qiita.com/shionhonda/items/e2cf9fe93ae1034dd771)
- [【徹底解説】EMアルゴリズムをはじめからていねいに](https://academ-aid.com/ml/em)
- [【徹底解説】変分ベイズをはじめからていねいに](https://academ-aid.com/ml/vb)
- [【徹底解説】VAEをはじめからていねいに](https://academ-aid.com/ml/vae)
- [VAEのPytorch実装](https://github.com/beginaid/VAE)
- 深層学習、岡谷貴之、MLP、講談社　
- [What are Diffusion Models?](https://lilianweng.github.io/posts/2021-07-11-diffusion-models/)
- [Understanding Diffusion Models: A Unified Perspective](https://arxiv.org/abs/2208.11970)
- [PyTorchでCVAEを実装してみた](https://qiita.com/isuya/items/a856971647ba2390f5fe)
- [多変量正規分布の場合のKullback Leibler Divergenceの導出](https://qiita.com/kenmatsu4/items/c107bd51503462fb677f)