変分自動エンコーダー(VAE)の原理解説とMNISTデータセットを用いた実践的構築

1. 数理的基礎

1.1 KL散度

KL散度(クールバック・ライブラー情報量)は、確率分布 \(P\) と \(Q\) の間の統計的な距離を測る指標です。あるモデル分布 \(Q\) が真の分布 \(P\) からどれだけ乖離しているかを定量化します。

離散確率変数における定義は以下の通りです:

\[D_{\text{KL}}(P \parallel Q) := \sum_{x} P(x) \log \frac{P(x)}{Q(x)} = -\sum_{x} P(x) \log Q(x) + \sum_{x} P(x) \log P(x)\]

情報の観点では、分布 \(Q\) に基づく最適符号長で \(P\) を圧縮しようとした場合の過剰ビット数(\(P\) 自体のエントロピーとの差)に対応します。両者が等しい分布であれば差は \(0\) となり、一般に \(D_{\text{KL}}(P \parallel Q) \ge 0\) が成り立ちます。

1.2 証拠量の下界(ELBO)

潜在変数モデルにおいて、事後分布 \(P(z \mid x)\) は解析的に計算困難なため、近似分布 \(Q(z)\) を用いて KL散度を最小化する考え方が採用されます。以下のように展開します:

\[\begin{aligned} D_{\text{KL}}(Q \parallel P) &= \sum_{z} Q(z) \log \frac{Q(z)P(x)}{P(x,z)} \\ &= \sum_{z} Q(z) \left( \log \frac{Q(z)}{P(x,z)} + \log P(x) \right) \\ &= \sum_{z} Q(z)(\log Q(z) - \log P(x,z)) + \log P(x) \end{aligned}\]

期待値表記 \(\mathbb{E}_Q[f(z)]\) を用いると:

\[D_{\text{KL}}(Q(z) \parallel P(z \mid x)) = \mathbb{E}_Q[\log Q(z) - \log P(x,z)] + \log P(x)\]

移項すると対数周辺尤度 \(\log P(x)\) が現れます:

\[\log P(x) = D_{\text{KL}}(Q(z) \parallel P(z \mid x)) + \mathcal{L}(Q)\]

ここで \(\mathcal{L}(Q)\) を証拠量の下界(Evidence Lower Bound: ELBO)と呼びます。KL散度が非負であることを利用すれば、次の不等式が得られます:

\[\log P(x) \ge \mathcal{L}(Q)\]

この \(\mathcal{L}(Q)\) を最大化することで、間接的に周辺尤度を最適化できます。

2. モデルアーキテクチャ

2.1 基本的な前提

観測データセット \(\mathbf{X} = \{x^{(i)}\}_{i=1}^N\) は、隠れ変数 \(z\) を経由して生成されると仮定します:

  • 事前分布 \(p_{\theta^*}(z)\) から潜在ベクトル \(z^{(i)}\) をサンプリング
  • 条件付き生成分布 \(p_{\theta^*}(x \mid z = z^{(i)})\) から観測値 \(x^{(i)}\) を生成

これらはパラメータ族 \(p_{\theta}(z)\) および \(p_{\theta}(x \mid z)\) に含まれ、滑らかさと微分可能性を満たすものとします。未知の事後分布 \(p_{\theta}(z \mid x)\) を推定するために、認識モデル(エンコーダ) \(q_{\phi}(z \mid x)\) を導入し、同時学習を行います。\n\nencode过程对应于从 \(q_{\phi}(z \mid x)\) 采样 \(z\),decode过程对应于从 \(p_{\theta}(x \mid z)\) 采样 \(x\)。

2.2 周辺尤度の導出

パラメータ \(\theta\) の良否を評価するため、データの生成確率(周辺尤度)に着目します:

\[p_{\theta}(x) = \int p_{\theta}(x \mid z)p_{\theta}(z) \mathrm{d}z\]

対数尤度は各サンプルの和として表せ、先述のELBOの関係式を適用すると:

\[\log p_{\theta}(x^{(i)}) = D_{\text{KL}}(q_{\phi}(z \mid x^{(i)}) \parallel p_{\theta}(z \mid x^{(i)})) + \mathcal{L}(\theta,\phi; x^{(i)})\]

さらに展開すると:

\[\mathcal{L} = \mathbb{E}_{q_{\phi}}[-\log q_{\phi}(z \mid x) + \log p_{\theta}(z) + \log p_{\theta}(x \mid z)]\]

これを分解すると:

\[\mathcal{L} = -D_{\text{KL}}(q_{\phi}(z \mid x) \parallel p_{\theta}(z)) + \mathbb{E}_{q_{\phi}}[\log p_{\theta}(x \mid z)]\]

第1項は正規化誤差(KL正則化)、第2項は再構成誤差に対応します。

2.3 再パラメータ化とAEVBアルゴリズム

単なる Monte Carlo 近似では勾配の分散が大きくなるため、再パラメータ化技巧を用います。 \(z \sim q_{\phi}(z \mid x)\) を \(z = g_{\phi}(\epsilon, x)\) (ただし \(\epsilon \sim p(\epsilon)\))と表し、期待値を以下のように書き換えます:

\[\mathbb{E}_{q_{\phi}}[f(z)] \approx \frac{1}{L} \sum_{\ell=1}^L f(g_{\phi}(\epsilon^{(\ell)}, x))\]

この近似を利用したミニバッチ版 Auto-Encoding VB (AEVB) アルゴリズムのフローは以下の通りです:

\[\begin{array}{rl} 0 & M, L \gets 100, 1 \\ 1 & \text{事前分布、認識モデル、生成モデルを設定} \\ 2 & \theta, \phi \gets \text{初期化} \\ 3 & \textbf{while not converged do} \\ 4 & \quad \text{ミニバッチ } \mathbf{X}_M \text{ をサンプリング} \\ 5 & \quad \epsilon \sim p(\epsilon) \text{ をサンプリング} \\ 6 & \quad \nabla_{\theta,\phi} \hat{\mathcal{L}}^M \gets \text{再パラメータ化による勾配計算} \\ 7 & \quad \theta, \phi \gets \text{勾配降下法で更新} \\ 8 & \textbf{end while}\end{array}\]

2.4 VAEにおける具体的な分布設定

標準的なVAEでは以下の分布を選択します:

\[\begin{aligned} p(\epsilon) &= \mathcal{N}(\epsilon; \mathbf{0}, \mathbf{I}) \\ q_{\phi}(z \mid x^{(i)}) &= \mathcal{N}(z; \mu^{(i)}, (\sigma^2)^{(i)}\mathbf{I}) \\ p_{\theta}(z) &= \mathcal{N}(z; \mathbf{0}, \mathbf{I}) \\ g_{\phi}(\epsilon^{(\ell)}, x^{(i)}) &= \mu^{(i)} + \sigma^{(i)} \odot \epsilon^{(\ell)} \end{aligned}\]

生成部 \(p_{\theta}(x \mid z)\) はデータ特性に応じて選定します。画像など連続値にはガウス分布、バイナリ値にはベルヌーイ分布が適します。ここではベルヌーイ分布を用い、その対数確率を見積もります。

ガウス分布間のKL散度は閉形式で求解可能です。次元 \(J\) の場合:

\[-D_{\text{KL}}(q_{\phi}(z) \parallel p_{\theta}(z)) = \frac{1}{2} \sum_{j=1}^J \left(1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2\right)\]

最終的な損失関数は:

\[\hat{\mathcal{L}} = \underbrace{\frac{1}{2} \sum_{j=1}^J \left(1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2\right)}_{\text{KL Regularizer}} + \underbrace{\frac{1}{L}\sum_{\ell=1}^L \log p_{\theta}(x \mid z^{(\ell)})}_{\text{Reconstruction Loss}}\]

3. MNISTデータセットを用いた実践的実装

3.1 データパイプラインの構築

MNIST画像は \(28 \times 28\) ピクセルのグレイスケールです。テンソル化により値域を \([0, 1]\) に正規化し、ミニバッチローダーで供給します。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torch.utils.data import DataLoader

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

dataset_cfg = {
    'root': './mnist_data',
    'train': True,
    'download': True,
    'transform': torchvision.transforms.ToTensor()
}
raw_dataset = torchvision.datasets.MNIST(**dataset_cfg)
DATA_LOADER = DataLoader(raw_dataset, batch_size=64, shuffle=True, num_workers=2)

3.2 ネットワーク構成と損失関数の定義

エンコーダは平均 \(\mu\) と対数分散 \(\log \sigma^2\) を出力します。デコーダはシグモイド活性化を用いてベルヌーイパラメータを生成します。勾配安定化のため、ELBO計算ロジックを独立モジュールとして定義します。

class VariationalEncoder(nn.Module):
    def __init__(self, latent_dim):
        super().__init__()
        self.latent_dim = latent_dim
        self.shared_proj = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU()
        )
        self.mean_head = nn.Linear(256, latent_dim)
        self.logvar_head = nn.Linear(256, latent_dim)

    def forward(self, x):
        flat_x = x.view(x.size(0), -1)
        h = self.shared_proj(flat_x)
        mu = self.mean_head(h)
        logvar = self.logvar_head(h)
        return mu, logvar

class GenerativeDecoder(nn.Module):
    def __init__(self, latent_dim):
        super().__init__()
        self.decoder_net = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 784),
            nn.Sigmoid()
        )

    def forward(self, z):
        return self.decoder_net(z)

class VAETrainer(nn.Module):
    def __init__(self, z_dim):
        super().__init__()
        self.encoder = VariationalEncoder(z_dim)
        self.decoder = GenerativeDecoder(z_dim)
        self.eps_tol = 1e-7

    def forward(self, x_batch):
        mu, logvar = self.encoder(x_batch)
        std = torch.exp(0.5 * logvar)
        
        # 再パラメータ化技巧
        epsilon = torch.randn_like(std)
        z_sample = mu + std * epsilon
        
        recon_pred = self.decoder(z_sample)
        
        # KLダイバージェンス(ベクトル演算化)
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        
        # ベルヌーイ対数尤度(再構成誤差)
        bce_numerator = x_batch * torch.log(recon_pred + self.eps_tol)
        bce_denominator = (1 - x_batch) * torch.log(1 - recon_pred + self.eps_tol)
        rec_loss = -torch.sum(bce_numerator + bce_denominator)
        
        # ELBO(最大化ではなく最小化のため符号を反転)
        elbo = rec_loss - kl_loss
        return elbo, rec_loss, kl_loss

3.3 トレーニングループの実装

Adamオプティマイザーでパラメータを更新します。各ステップでELBO、再構成誤差、KL正則化項をモニタリングします。

def run_training(model_obj, dataloader, lr_rate=1e-3, n_epochs=15):
    opt = torch.optim.Adam(model_obj.parameters(), lr=lr_rate)
    model_obj.train()
    
    for ep_idx in range(n_epochs):
        epoch_accum_loss = 0.0
        batch_cnt = 0
        
        for inputs, _ in dataloader:
            inputs = inputs.to(DEVICE).float()
            opt.zero_grad()
            
            elbo_val, rec_val, kl_val = model_obj(inputs)
            total_cost = elbo_val
            
            total_cost.backward()
            opt.step()
            
            epoch_accum_loss += elbo_val.item()
            batch_cnt += 1
            
            if batch_cnt % 50 == 0:
                print(f"[Batch {batch_cnt}] Cost:{elbo_val.item():.4f} | Rec:{rec_val.item():.4f} | KL:{kl_val.item():.4f}")
                
        avg_cost = epoch_accum_loss / batch_cnt
        print(f"--- Epoch {ep_idx+1}/{n_epochs} Completed. Avg Loss: {avg_cost:.4f} ---")
        
    return model_obj

3.4 生成結果と潜在空間の可視化

訓練後、潜在空間から直接サンプリングして画像を生成し、エンコード後の分布状態を確認します。

def visualize_generated_samples(vae_model, rows=4, cols=4):
    vae_model.eval()
    z_rand = torch.randn(rows * cols, vae_model.encoder.latent_dim).to(DEVICE)
    img_preds = vae_model.decoder(z_rand)
    grid_tensors = img_preds.view(-1, 1, 28, 28).cpu()
    torchvision.utils.save_image(grid_tensors, 'generated_samples.png', nrow=rows, normalize=True)

def inspect_latent_distribution(vae_model, val_loader):
    vae_model.eval()
    latent_coords = []
    true_labels = []
    
    with torch.no_grad():
        for data, labels in val_loader:
            data = data.view(data.size(0), -1).to(DEVICE)
            mu, _ = vae_model.encoder(data)
            latent_coords.append(mu.cpu().numpy())
            true_labels.extend(labels.numpy())
            
    latent_matrix = np.concatenate(latent_coords, axis=0)
    plt.figure(figsize=(8, 8))
    scatter_plot = plt.scatter(latent_matrix[:, 0], latent_matrix[:, 1], 
                               c=true_labels, cmap='tab10', alpha=0.6, s=15)
    plt.colorbar(scatter_plot)
    plt.title('Projected Latent Space (Mean Coordinates)')
    plt.xlabel('Dimension 0')
    plt.ylabel('Dimension 1')
    plt.tight_layout()
    plt.savefig('latent_projection.png')
    plt.close()

生成されたグリッド画像は典型的手書き数字の特徴を保持しており、潜像空間の散布図からはクラスターごとの分離傾向が確認できます。特に類似する字形(例えば4と9)については重なりが生じますが、全体的なクラスター構造は解釈可能です。

4. 参考文献

[1] Kingma, D. P., & Welling, M. (2013). Auto-encoding variational bayes.

[2] Doersch, C. (2016). Tutorial on variational autoencoders.

[3] Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning. MIT Press.

タグ: variational-autoencoder vae deep-learning-mathematics reparameterization-trick pixel-generation

8月28日 05:40 投稿