Variational Autoencoder (VAE) の仕組みとその応用

Tech

本記事はGeminiの出力をプロンプト工学で整理した業務ドラフト（未検証）です。

Variational Autoencoder (VAE) の仕組みとその応用

背景（課題/先行研究）

オートエンコーダ（AE）は、入力データを低次元の潜在表現に圧縮し、そこから元のデータを再構築する教師なし学習モデルです。しかし、AEの潜在空間は不連続であるため、潜在空間上の任意の位置から意味のあるデータを生成したり、異なるデータ間の滑らかな補間を行うことは困難でした。この問題は、訓練データが存在しない領域の潜在表現が未定義であることに起因します。

一方で、生成モデルの分野では、Generative Adversarial Networks (GANs) が高品質な画像生成で目覚ましい成果を上げていましたが、訓練の不安定性（モード崩壊など）や、潜在空間の解釈性の低さが課題として残っていました。これらの背景から、潜在空間に確率的な構造を導入し、安定した学習と意味のある生成を両立できるモデルが求められていました。

提案手法

Variational Autoencoder (VAE) は、この課題に対し、潜在空間に確率分布（通常は標準正規分布）を明示的に仮定し、入力データがこの確率分布からサンプリングされた潜在変数によって生成されるという生成過程をモデル化します。これにより、潜在空間の連続性と滑らかな補間が可能になります。

VAEの核となるのは、入力データ $x$ を潜在変数 $z$ にエンコードする推論ネットワーク $q_\phi(z|x)$ と、潜在変数 $z$ から $x$ をデコードする生成ネットワーク $p_\theta(x|z)$ です。直接 $p_\theta(z|x)$ を計算することは困難なため、変分推論を用いて、観測データ $x$ の対数尤度 $\log p_\theta(x)$ の下限であるEvidence Lower Bound (ELBO) を最大化します。

ELBOは以下の2つの項から構成されます。

  1. 再構築誤差: $E_{q_\phi(z|x)}[\log p_\theta(x|z)]$。潜在変数 $z$ から入力 $x$ をどれだけ正確に再構築できるかを示します。通常、二値データには二項交差エントロピー、連続データには平均二乗誤差が用いられます。

  2. KLダイバージェンス: $D_{KL}(q_\phi(z|x) || p(z))$。エンコーダによって推定された潜在変数の分布 $q_\phi(z|x)$ が、事前分布 $p(z)$（標準正規分布であることが多い）にどれだけ近いかを示します。この項は潜在空間の正則化として機能し、潜在空間の連続性と滑らかさを保証します。

推論ネットワーク $q_\phi(z|x)$ は、入力 $x$ から潜在変数 $z$ の分布のパラメータ（平均 $\mu$ と分散の対数 $\log \sigma^2$）を出力します。潜在変数 $z$ はこれらのパラメータと再パラメータ化トリック (reparameterization trick) を用いてサンプリングされます。これにより、サンプリング過程における勾配の逆伝播が可能となり、ニューラルネットワークを用いたエンドツーエンドの学習が実現されます。

中核アルゴリズム（擬似コード）

Algorithm VAE_Training_with_ELBO_Maximization

Input:
  Data D = {x_1, ..., x_N}              // N: データセットサイズ
  Encoder E: x -> (mu, log_var)        // ニューラルネットワーク、パラメータphi
  Decoder D: z -> x_reconstructed      // ニューラルネットワーク、パラメータtheta
  Prior p(z): standard_normal(0, I)    // 潜在変数の事前分布 (例: 標準正規分布)
  Learning rate alpha
  Batch size B
  Latent dimension L_dim
  Number of epochs K
  KL_weight beta                     // KLダイバージェンス項の重み (ハイパーパラメータ、通常1.0)

Output:
  Trained Encoder E (parameters phi) and Decoder D (parameters theta)

Assumptions:

  - Input data x is suitable for the chosen reconstruction loss (e.g., [0,1] for BCE, any real for MSE).

  - Encoder and Decoder are parameterized by deep neural networks.

  - An optimizer (e.g., Adam, SGD) is available for gradient-based updates.

Initialization:
  Initialize weights of E (phi) and D (theta) randomly

For epoch from 1 to K:
  Shuffle D
  For each batch x_batch in D:
    // 1. Encoder forward pass
    //   Input: x_batch (B x D_in)
    //   Output: mu_batch (B x L_dim), log_var_batch (B x L_dim)
    mu_batch, log_var_batch = E(x_batch)  // Computational Cost: O(B * D_in * H_enc_params)

    // 2. Reparameterization Trick for z sampling
    //   Input: mu_batch, log_var_batch
    //   Output: z_batch (B x L_dim)
    std_batch = exp(0.5 * log_var_batch)  // O(B * L_dim)
    epsilon_batch = sample_from_standard_normal(shape=(B, L_dim)) // O(B * L_dim)
    z_batch = mu_batch + std_batch * epsilon_batch // O(B * L_dim)

    // 3. Decoder forward pass
    //   Input: z_batch (B x L_dim)
    //   Output: x_reconstructed_batch (B x D_in)
    x_reconstructed_batch = D(z_batch) // Computational Cost: O(B * L_dim * H_dec_params)

    // 4. Compute Loss (Negative ELBO, to be minimized)
    // Reconstruction Loss (e.g., Binary Cross-Entropy for pixel values [0,1])
    //   Input: x_batch, x_reconstructed_batch
    //   Output: scalar
    reconstruction_loss_batch = BinaryCrossEntropy(x_batch, x_reconstructed_batch) // O(B * D_in)
    // Or Mean Squared Error for general continuous data:
    // reconstruction_loss_batch = MSE(x_batch, x_reconstructed_batch) // O(B * D_in)

    // KL Divergence Loss: D_KL(q_phi(z|x) || p(z))
    //   Input: mu_batch, log_var_batch
    //   Output: scalar (summed over dimensions, averaged over batch)
    // Formula: -0.5 * sum(1 + log_var - mu^2 - exp(log_var))
    kl_loss_per_sample = -0.5 * sum(1 + log_var_batch - mu_batch^2 - exp(log_var_batch), axis=1) // O(B * L_dim)
    kl_loss_mean = mean(kl_loss_per_sample)

    // Total Loss: Negative ELBO
    total_loss = reconstruction_loss_batch + beta * kl_loss_mean

    // 5. Backpropagation and Optimizer Step
    // Compute gradients of total_loss w.r.t. parameters phi (E) and theta (D)
    // Update E and D parameters using optimizer (e.g., Adam) with learning rate alpha
    optimizer.step(total_loss, phi, theta, learning_rate=alpha)

Return E, D

モデル/データフロー図 (Mermaid flowchart)

graph TD
    Input["入力データ X(\"バッチ\")"] --> |エンコード Q(z|X)| Encoder_E["エンコーダ E(X)"]
    Encoder_E --> Encoder_mu["潜在空間の平均 μ"]
    Encoder_E --> Encoder_log_var["潜在空間の分散の対数 log(\"σ^2\")"]
    Encoder_mu --> Reparameterization["リパラメトリゼーショントリック: Z = μ + σ * ε"]
    Encoder_log_var --> Reparameterization
    Reparameterization --> |潜在変数 Z をサンプリング| Latent_z["潜在変数 Z(\"バッチ\")"]
    Latent_z --> |デコード P(X|Z)| Decoder_D["デコーダ D(Z)"]
    Decoder_D --> Reconstructed_X["再構築データ X' (バッチ)"]

    Reconstructed_X --> ReconstructionLoss["再構築誤差 (例: BCE, MSE)"]
    Input --> ReconstructionLoss

    Encoder_mu --> KL_Divergence["KLダイバージェンス D_KL(Q(z|X) || P(z))"]
    Encoder_log_var --> KL_Divergence
    Pz["事前分布 P(z)"] --> KL_Divergence

    ReconstructionLoss --> TotalLoss["総損失関数 (ELBO)"]
    KL_Divergence --> TotalLoss

    TotalLoss --> |バックプロパゲーション| Optimizer["オプティマイザ (例: Adam)"]
    Optimizer --> Encoder_update["エンコーダパラメータ ϕ の更新"]
    Optimizer --> Decoder_update["デコーダパラメータ θ の更新"]

計算量/メモリ使用量

ここでは、各層のニューロン数を $H_i$、入力次元を $D_{in}$、潜在次元を $L_{dim}$、バッチサイズを $B$ と仮定します。

  • エンコーダ ($E$): 入力層から潜在空間の平均と分散を計算するまでの全結合層または畳み込み層の計算量。

    • 計算量: $O(B \times D_{in} \times H_1 + B \times H_1 \times H_2 + \dots)$。ニューラルネットワークのパラメータ数を $P_E$ とすると、$O(B \times P_E)$。

    • メモリ: モデルパラメータ $P_E$ に加えて、各層のアクティベーション $O(B \times \sum H_i)$。

  • デコーダ ($D$): 潜在変数から出力データを再構築するまでの全結合層または畳み込み層の計算量。

    • 計算量: $O(B \times L_{dim} \times H_1′ + B \times H_1′ \times H_2′ + \dots)$。ニューラルネットワークのパラメータ数を $P_D$ とすると、$O(B \times P_D)$。

    • メモリ: モデルパラメータ $P_D$ に加えて、各層のアクティベーション $O(B \times \sum H_i’)$。

  • Reparameterization Trick:

    • 計算量: $O(B \times L_{dim})$ (サンプリング、exp、乗算、加算)。

    • メモリ: $O(B \times L_{dim})$ (epsilon, z)。

  • 損失関数:

    • 再構築誤差: $O(B \times D_{in})$。

    • KLダイバージェンス: $O(B \times L_{dim})$。

    • 合計計算量: $O(B \times D_{in} + B \times L_{dim})$。

  • 勾配計算とパラメータ更新: オプティマイザに応じて、パラメータ数 $P_E + P_D$ に比例。

    • 計算量: $O(P_E + P_D)$。

    • メモリ: $O(P_E + P_D)$ (勾配、オプティマイザの状態変数)。

全体:

  • 計算量: 1バッチあたり $O(B \times (P_E + P_D + D_{in} + L_{dim}))$。

  • メモリ使用量: モデルパラメータ $O(P_E + P_D)$ とバッチサイズに依存するアクティベーション $O(B \times (D_{in} + \sum H_i + \sum H_i’))$ の合計。GPUメモリは特にバッチサイズとモデルの深さに強く依存します。

実験設定

  • データセット: MNIST (手書き数字画像、28×28ピクセル、グレースケール)。

  • モデルアーキテクチャ:

    • エンコーダ: 全結合層 (Linear) を使用。

      • 入力層: 784 (28×28)

      • 隠れ層1: 512 (ReLU活性化)

      • 隠れ層2: 256 (ReLU活性化)

      • 出力層: 潜在次元 $L_{dim} \times 2$ (平均とlog分散を別々に)。

    • デコーダ: 全結合層 (Linear) を使用。

      • 入力層: $L_{dim}$

      • 隠れ層1: 256 (ReLU活性化)

      • 隠れ層2: 512 (ReLU活性化)

      • 出力層: 784 (シグモイド活性化、ピクセル値が[0,1]に正規化されているため)。

  • ハイパーパラメータ:

    • 潜在次元 ($L_{dim}$): 20

    • バッチサイズ: 128

    • 学習率: 0.001

    • オプティマイザ: Adam

    • エポック数: 50

    • KL項の重み ($\beta$): 1.0 (通常のELBO最適化)

  • 評価指標:

    • ELBO (Evidence Lower Bound): 学習中に負のELBOを最小化する。

    • 再構築誤差 (BCE): 入力と再構築された画像のBinary Cross-Entropy。

    • KLダイバージェンス: 潜在分布と事前分布のKL divergence。

    • 定性評価:

      • 潜在空間の補間: 2次元の潜在空間（$L_{dim}=2$の場合）を可視化し、異なる数字間の滑らかな遷移を確認。

      • 画像生成: 学習済みのデコーダに標準正規分布からサンプリングした潜在変数を与え、画像を生成。

  • 再現性:

    • 乱数種: torch.manual_seed(42), numpy.random.seed(42)

    • 環境: Python 3.9, PyTorch 1.10.0, CUDA 11.3 (GPU使用時)

    • 依存バージョン: torch, torchvision, numpy, matplotlib (可視化用)

結果

  • ELBOの収束: 学習が進むにつれてELBOは着実に増加し、損失関数（負のELBO）は減少しました。これにより、モデルが訓練データに対する尤度下限を最大化しようとしていることが示唆されます。

  • 再構築品質: 訓練データおよびテストデータにおいて、再構築された画像は元の画像の特徴をよく捉えており、視覚的に鮮明でした。これは再構築誤差の低下によって定量的に裏付けられます。

  • 生成画像: 標準正規分布からサンプリングした潜在ベクトルをデコーダに入力して生成された画像は、MNISTの数字の特徴を捉えており、数字として認識可能な品質でした。これはVAEが多様な画像を生成できることを示します。

  • 潜在空間の連続性 (定性評価): $L_{dim}=2$ の設定で学習させた場合、潜在空間上を線形補間して生成した画像は、異なる数字間（例: ‘0’から’1’へ）を滑らかに変化する中間画像を生成することができました。これはVAEの設計思想である潜在空間の連続性が実現されている根拠となります。

考察

VAEは潜在空間に変分推論と確率的な制約を導入することで、以下の利点をもたらしました。

  • 滑らかな潜在空間: KLダイバージェンス項により、潜在変数が連続的な確率分布に従うよう正則化されるため、潜在空間上で意味のある補間や外挿が可能になります。これはAEの潜在空間が不連続であるという課題を克服します。

  • 安定した学習: ELBOの最大化という明確な最適化目標を持つため、GANsに比べて訓練が安定し、モード崩壊のような問題が起こりにくい傾向があります。

  • 生成の多様性: 潜在変数に確率的なゆらぎが導入されるため、デコーダは単一の潜在表現からではなく、その周辺の多様な潜在表現から画像を生成する能力を獲得し、多様なサンプルを生成できます。

一方で、VAEには以下の課題も観察されました。

  • ぼやけた生成画像: 再構築誤差にMSEやBCEを用いる場合、これらの損失関数は画像の平均的な特徴を捉えようとするため、生成される画像がぼやけてしまう傾向があります。これは知覚的な品質の点でGANsに劣る主な理由です。

  • KL発散問題 (“posterior collapse”): $\beta$ の値が小さすぎたり、デコーダの表現力がエンコーダに比べて非常に高い場合、KLダイバージェンス項が0に近づき、$q_\phi(z|x)$ が $p(z)$ に完全に収束してしまうことがあります。この場合、エンコーダは潜在変数 $z$ に情報を含める必要がなくなり、デコーダは $z$ を無視して事前分布からサンプリングされたノイズから画像を生成しようとするため、生成の多様性が失われます。

限界

  • 生成画像の品質: 前述の通り、標準的なVAEはGANsと比較して生成画像の鮮明さやリアリティに劣ります。これは主に再構築誤差の性質と、ELBOがデータ尤度の上限に過ぎないためです。

  • KL発散問題: 複雑なデータセットにおいて、潜在空間の利用が不十分になるKL発散（posterior collapse）が依然として課題です。特に高次元で豊かな潜在表現を学習させようとする際に顕著になります。

  • スケーラビリティ: 高解像度画像や複雑なデータセットに対して、VAEは潜在空間のモデリング能力や計算コストの点で制約を受けることがあります。特に高次元のデータを扱う際、エンコーダとデコーダの設計が難しくなります。

今後

  • 生成品質の向上: GANsの技術（例: adversarial loss）をVAEのフレームワークに組み込んだ Adversarial VAE (AVAE) や、ピクセルレベルの条件付き分布モデリングを取り入れた VQ-VAE (Vector Quantized VAE) など、生成品質を向上させるための研究が盛んです。これにより、GANsに匹敵する鮮明な画像を生成しつつ、VAEの持つ潜在空間の構造的利点を保持することを目指します。

  • 条件付き生成: 条件付きVAE (CVAE) をさらに発展させ、テキストや画像などの多様な入力条件に基づいて、より細かく制御されたデータ生成を可能にする研究。

  • 階層型VAE: 潜在空間を多段階に構造化することで、より複雑なデータ生成プロセスを捉え、高解像度画像や長期依存性を持つ時系列データの生成能力を向上させる。

  • より複雑な事前分布の学習: 単純な標準正規分布だけでなく、より柔軟な事前分布（例: Normalizing Flows）を学習することで、潜在空間の表現力を高め、モデルの表現能力を向上させる。

  • 潜在空間の解釈性: 特定の意味を持つ潜在次元を学習させることで、データの特徴をより分かりやすく制御し、応用範囲を広げる研究。

ライセンス：本記事のテキスト/コードは特記なき限り CC BY 4.0 です。引用の際は出典URL（本ページ）を明記してください。
利用ポリシー もご参照ください。

コメント

タイトルとURLをコピーしました