google Colab で機械学習を始める人向けの簡単なメモ (VAE編)

# はじめに

google Colab で、変分オートエンコーダ（VAE）を理解するための記事です。

https://qiita.com/yamadasuzaku/items/302f5cd7d328c957c0b4

を読んでいることを前提としています。

MNISTデータセットを用いて、VAE（Variational Autoencoder）をKerasで実装したチュートリアルは多くありますが、本記事では、google Colab でVAEのコードを実行しながら、以下の数式的な観点からVAEの構造を丁寧に解説します。

https://colab.research.google.com/drive/1EgWfJNNErQ6Debfk_wjvYG73H5Hk8oKI?usp=sharing

にあるコードは、2025.5.29に動作確認したので、その前後で時期であれば動くはずです。

コードの前に、まずは基本事項の確認です。

* 変分推論とは？
* VAEの損失関数はなぜあの形になるのか？
* 再パラメータ化トリックとは？

---

## なぜAutoencoderではなくVAEなのか？

通常のAutoencoderは入力データ \$\mathbf{x}\$ を低次元ベクトル \$\mathbf{z}\$ に圧縮（encode）し、再び復元（decode）します。

```math
\begin{aligned}
\mathbf{z} &= f_{\text{enc}}(\mathbf{x}) \\
\hat{\mathbf{x}} &= f_{\text{dec}}(\mathbf{z})
\end{aligned}
```

このとき、エンコーダ \$f\_{\text{enc}}\$ とデコーダ \$f\_{\text{dec}}\$ はニューラルネットで表現され、損失関数としては単純な再構成誤差（例：2乗誤差やバイナリ交差エントロピー）を使います。

### 問題点：

このモデルでは \$\mathbf{z}\$ が固定的なベクトルにしかならず、「生成モデル」として使うにはランダム性が不足します。

---

## 変分オートエンコーダの基本アイデア

VAEでは、エンコーダは潜在変数 \$\mathbf{z}\$ を「確率分布」として出力します：

```math
q_{\phi}(\mathbf{z}|\mathbf{x}) = \mathcal{N}(\mathbf{z}; \boldsymbol{\mu}(\mathbf{x}), \mathrm{diag}(\boldsymbol{\sigma}^2(\mathbf{x})))
```

ここで \$\phi\$ はエンコーダのパラメータ（ニューラルネットの重み）です。

---

## 目的関数（ELBO）の導出

生成モデルの最大目的は **データの対数尤度** を最大化すること：

```math
\log p_{\theta}(\mathbf{x}) = \log \int p_{\theta}(\mathbf{x}, \mathbf{z}) d\mathbf{z}
```

これは積分の計算が困難なので、「変分下限 (ELBO: Evidence Lower Bound)」を最大化します。

```math
\log p_{\theta}(\mathbf{x}) \ge \mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})}[\log p_{\theta}(\mathbf{x}|\mathbf{z})] - D_{\mathrm{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x}) \| p(\mathbf{z}))
```

### 各項の意味：

* 再構成誤差：
``` math
\mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})}[\log p_{\theta}(\mathbf{x}|\mathbf{z})]
```
* 正則化項：\$D\_{\mathrm{KL}}(q\_{\phi}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}))\$
  → 潜在空間を正規分布 \$p(\mathbf{z}) = \mathcal{N}(0, I)\$ に近づける

---

## 再パラメータ化トリックとは？

エンコーダの出力は平均 \$\boldsymbol{\mu}\$ と分散 \$\boldsymbol{\sigma}^2\$ ですが、ここからサンプリングを行うとニューラルネットの勾配が計算できなくなります。

:::note info

**再パラメータ化トリック（reparameterization trick）** の数学的な背景から丁寧に説明します。

なぜ \$\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}\$ とすると微分可能になるのか？

### 問題の背景

VAEでは、潜在変数 \$\mathbf{z}\$ は確率的に生成される：

```math
\mathbf{z} \sim \mathcal{N}(\boldsymbol{\mu}, \mathrm{diag}(\boldsymbol{\sigma}^2))
```

しかし、**確率変数のサンプリング操作（\$\sim\$）は非微分的**です。つまり、サンプリング操作そのものには勾配が定義されません。これは、ニューラルネットの学習に必要な **誤差逆伝播(backpropagation)** の妨げになります。

:::


# コードの解説

Autoencoderから始めて、VAEの仕組み・損失関数・潜在空間の構造を解説します。


## Part 1: Autoencoderの基礎

Autoencoder（オートエンコーダ）は、入力データを一度圧縮（encode）し、そこから復元（decode）することでデータの特徴を抽出するニューラルネットワークです。

### 🔧 1-1. MNISTデータの読み込みと前処理

```python
from keras.layers import Input, Dense
from keras.models import Model
from keras.datasets import mnist
import numpy as np
import matplotlib.pyplot as plt

(x_train, _), (x_test, _) = mnist.load_data()
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))
```

* 画像を `[0, 1]` の範囲に正規化
* 28x28 = 784 の一次元ベクトルに変換

### 🔧 1-2. Autoencoderモデルの構築と学習

```python
encoding_dim = 32
input_img = Input(shape=(784,))
x1 = Dense(256, activation='relu')(input_img)
x2 = Dense(64, activation='relu')(x1)
encoded = Dense(encoding_dim, activation='relu')(x2)
x3 = Dense(64, activation='relu')(encoded)
x4 = Dense(256, activation='relu')(x3)
decoded = Dense(784, activation='sigmoid')(x4)

autoencoder = Model(inputs=input_img, outputs=decoded)
autoencoder.compile(optimizer='adam', loss='binary_crossentropy')
```

### 🏋️‍♂️ 1-3. モデルの学習と可視化

```python
autoencoder.fit(x_train, x_train,
                epochs=50,
                batch_size=256,
                shuffle=True,
                validation_data=(x_test, x_test))

decoded_imgs = autoencoder.predict(x_test)
```

* 入力＝出力の構造で学習することで、特徴抽出を行う

---

## Part 2: Variational Autoencoder (VAE)の構築

Autoencoderを発展させ、潜在変数を**確率的に**扱うようにしたのがVAEです。

### 🔧 2-1. 再パラメータ化トリック

```python
def sampling(z_mean, z_logvar):
    epsilon = tf.random.normal(shape=tf.shape(z_mean), seed=42)
    return z_mean + tf.exp(0.5 * z_logvar) * epsilon
```

* `z_mean` と `z_logvar` を使ってガウス分布からサンプリング

### 🔧 2-2. エンコーダの定義

```python
def build_encoder():
    encoder_inputs = layers.Input(shape=(784,))
    x = layers.Dense(256, activation="relu")(encoder_inputs)
    x = layers.Dense(64, activation="relu")(x)
    z_mean = layers.Dense(2)(x)
    z_logvar = layers.Dense(2)(x)
    z = layers.Lambda(lambda x: sampling(*x))([z_mean, z_logvar])
    return models.Model(encoder_inputs, [z_mean, z_logvar, z], name="encoder")
```

### 🔧 2-3. デコーダの定義

```python
def build_decoder():
    latent_inputs = layers.Input(shape=(2,))
    x = layers.Dense(64, activation="relu")(latent_inputs)
    x = layers.Dense(256, activation="relu")(x)
    outputs = layers.Dense(784, activation="sigmoid")(x)
    return models.Model(latent_inputs, outputs, name="decoder")
```

---

## Part 3: VAEクラスの中身を覗く

Kerasのカスタムモデルとして定義。

```python
class VAE(tf.keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = tf.keras.metrics.Mean(name="loss")
        self.reconstruction_loss_tracker = tf.keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = tf.keras.metrics.Mean(name="kl_loss")

    def call(self, inputs):
        z_mean, z_logvar, z = self.encoder(inputs)
        return self.decoder(z)
```

### 💡 損失関数の定義（train\_step）

```python
    def train_step(self, data):
        if isinstance(data, tuple):
            data = data[0]
        with tf.GradientTape() as tape:
            z_mean, z_logvar, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.keras.losses.binary_crossentropy(data, reconstruction)
            ) * 784
            kl_loss = -0.5 * tf.reduce_mean(
                tf.reduce_sum(1 + z_logvar - tf.square(z_mean) - tf.exp(z_logvar), axis=1)
            )
            total_loss = reconstruction_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }
```

* **再構成誤差（reconstruction\_loss）**：元画像と出力の誤差
* **KLダイバージェンス（kl\_loss）**：正規分布とのズレ
* 合わせて `total_loss` を最小化する

---

## Part 4: 学習と結果の可視化

```python
encoder = build_encoder()
decoder = build_decoder()
vae = VAE(encoder, decoder)
vae.compile(optimizer="adam", loss=lambda y_true, y_pred: 0.0)
vae.fit(x_train, epochs=50, batch_size=256, validation_data=(x_test, x_test))
```

### 🖼️ 潜在空間の可視化

```python
def plot_results(encoder, decoder, x_test, y_test):
    z_mean, _, _ = encoder.predict(x_test)
    plt.scatter(z_mean[:, 0], z_mean[:, 1], c=y_test, cmap='tab10')
    plt.colorbar()
    plt.xlabel("z[0]")
    plt.ylabel("z[1]")
    plt.show()
```

---

## 🎓 まとめ：VAEの魅力と応用

VAEは、次のような特徴を持ちます：

* 潜在空間が滑らかで、意味のある連続性を持つ
* 潜在空間上を動かすことで、新しい画像を生成できる
* 応用例：異常検知、画像生成、教師なしクラスタリングなど

本記事では、Autoencoderから始めて、VAEの数理的な仕組みや損失関数の直感的理解まで網羅しました。ぜひGoogle Colabなどでコードを実行しながら、モデルの挙動を観察してみてください！


## 🧠 解決法：再パラメータ化トリック

### アイデア：

サンプリングを、「μとσに依存する**関数**」として書き直す！

次のように書き換えます：

```math
\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}, \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, I)
```

この式では、

* \$\boldsymbol{\epsilon}\$ は標準正規分布に従う「固定された乱数」
* \$\boldsymbol{\mu}\$ と \$\boldsymbol{\sigma}\$ はニューラルネットの出力（パラメータに依存）

---

## ✨ 微分可能になる理由

この式は、\$\mathbf{z}\$ を \$\boldsymbol{\mu}\$ と \$\boldsymbol{\sigma}\$ の **明示的な関数**として書いた形です：

```math
\mathbf{z} = f(\boldsymbol{\mu}, \boldsymbol{\sigma}; \boldsymbol{\epsilon}) = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}
```

* ここで \$\boldsymbol{\epsilon}\$ は固定された乱数としてみなす（外から与える）
* \$\boldsymbol{\mu}, \boldsymbol{\sigma}\$ に対して微分可能な演算（加算・乗算）だけで構成されている

よって、**チェインルール**を使って \$\mathbf{z}\$ の勾配を \$\boldsymbol{\mu}\$ や \$\boldsymbol{\sigma}\$ に流すことが可能になります。

---

## 📌 まとめ図（概念）

```
x ──→ μ, σ ──→ z = μ + σ * ε ──→ decoder ──→ x_hat
                 ↑                          ↑
                 |                          |
            微分可能!                    損失関数
```

---

## ✅ 重要なポイント

* `z ~ N(μ, σ²)` のままだと、勾配が流れず学習できない
* `z = μ + σ * ε` と書き換えることで、μとσを通じて勾配が流れる
* この操作が **再パラメータ化トリック** と呼ばれ、VAEの核心技術です

---

必要であれば、この部分だけをQiita記事の「再パラメータ化トリックとは？」の章として組み込んだ全文も提供できます。ご希望があればお知らせください。

:::


### 対処法：

```math
\mathbf{z} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon},\quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, I)
```

こうすれば、\$\mathbf{z}\$ を微分可能な形にできます。

## 💻 コードで見る構造（簡略）

### エンコーダ

```python
z_mean = Dense(latent_dim)(x)
z_logvar = Dense(latent_dim)(x)
z = Lambda(lambda x: sampling(*x))([z_mean, z_logvar])
```

### 再構成損失

```python
reconstruction_loss = binary_crossentropy(x, x_hat) * original_dim
```

### KLダイバージェンス

```math
\mathrm{KL} = -\frac{1}{2} \sum_{i=1}^{d} \left( 1 + \log \sigma_i^2 - \mu_i^2 - \sigma_i^2 \right)
```

Pythonでの計算例：

```python
kl_loss = -0.5 * tf.reduce_mean(
    tf.reduce_sum(1 + z_logvar - tf.square(z_mean) - tf.exp(z_logvar), axis=1)
)
```

---

## 🖼️ 結果の可視化

### 再構成画像（上：元画像、下：VAEによる再構成）

```python
decoded_imgs = vae.predict(x_test)
```

### 潜在空間の2次元マップ

```python
z_mean, _, _ = encoder.predict(x_test)
plt.scatter(z_mean[:, 0], z_mean[:, 1], c=y_test)
```

### 潜在空間のメッシュグリッドから生成

```python
z_sample = np.array([[xi, yi]])
digit = decoder.predict(z_sample).reshape(28, 28)
```

### 生成される図

内在変数と出力の関係。

![vae.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/179652/acab10f7-8683-4e62-aa4d-f6a8d3c18f2e.png)


## 📌 まとめ

| 用語                  | 意味                                             |
| ------------------- | ---------------------------------------------- |
| 潜在変数 \$\mathbf{z}\$ | データの背後にある「潜在的な表現」                              |
| エンコーダ               | \$\mathbf{x} \to (\mu, \log \sigma^2)\$ のマッピング |
| デコーダ                | \$\mathbf{z} \to \hat{\mathbf{x}}\$ の再構成       |
| ELBO                | VAEの損失関数（再構成 + KL項）                            |
| 再パラメータ化             | 勾配を流せるように工夫した手法                                |

# まとめ

理解できた人もそうでない人も、まずは下記を実行して遊んでみてください。

https://colab.research.google.com/drive/1EgWfJNNErQ6Debfk_wjvYG73H5Hk8oKI?usp=sharing#scrollTo=5644KlCYO8Gv


latent 空間の大きさや、周辺情報などを使って、正答率が上がらないかなど、色々と試して理解を深めてもらえればと思います。