【深層学習】VAEとは?VAEの理論を徹底的に分かりやすく解説する

「画像を学習して、見たことのない新しい画像を生成したい」——そう思ったとき、通常のオートエンコーダ(AE)では不十分です。AEの潜在空間は構造がなく、任意の点をデコーダに入力しても意味のある画像が出るとは限りません。

この問題を解決するのが VAE(Variational Auto Encoder, 変分オートエンコーダー) です。VAEは潜在空間全体を正規分布で覆うことで、「任意の点からサンプリングすれば意味のある出力が得られる」という生成モデルとしての性質を実現します。

VAEが使われている場面は幅広く、たとえば次のような応用があります。

  • 画像生成・データ拡張: 医療画像の少ないデータから多様なサンプルを生成し、診断モデルを強化する
  • 異常検知: 正常データだけでVAEを学習し、再構成誤差が大きいサンプルを異常として検出する
  • 表現学習(β-VAE): 潜在変数の各次元に独立した意味を持たせ、「色」「形状」「向き」などの要因を分離する

本記事では、AEとの違いから出発し、確率的エンコーダ・デコーダの設計、ELBO(Evidence Lower Bound)の導出、再パラメータ化トリック、KL項の閉形式、PyTorchでの実装まで、数式を1行ずつ丁寧に追います。

本記事で学ぶこと

  • 通常のオートエンコーダとVAEの本質的な違い
  • 確率的エンコーダ $q_\phi(\bm{z}|\bm{x})$ とデコーダ $p_\theta(\bm{x}|\bm{z})$ の設計
  • ELBOの導出(再構成項 + KL正則化項)
  • 再パラメータ化トリックの必要性と実装
  • ガウス分布のKL閉形式とその意味
  • PyTorchでのVAE実装と潜在空間の可視化

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。


通常のオートエンコーダとVAEの違い

VAEを理解するための最初の一歩は、通常のオートエンコーダ(AE)の限界を知ることです。

通常のAEでは、エンコーダが入力 $\bm{x}$ をそのまま潜在変数 $\bm{z}$(ひとつの固定した点)にマッピングします。この「固定した点」というところが問題の核心です。訓練データをうまく再構成できても、潜在空間の中には「意味のある点」と「意味のない点」が混在していて、その境界が明確ではありません。

たとえば、数字「3」の画像が潜在空間の点 $A$ に、数字「8」が点 $B$ にマッピングされているとき、$A$ と $B$ の中間点をデコーダに入力すると何が出てくるでしょうか。AEでは、その中間に「正常な画像」がある保証はなく、ノイズのような出力になりがちです。

AEとVAEの構造比較 — 通常オートエンコーダは潜在変数が固定値、VAEは正規分布からサンプリング

上の図を見ると、AEは潜在変数 $\bm{z}$ が「一点」であるのに対して、VAEは $\bm{z}$ が正規分布に従う確率変数 として設計されています。エンコーダが出力するのは「潜在変数そのもの」ではなく、「潜在変数が従う正規分布のパラメータ(平均 $\bm{\mu}$ と分散 $\bm{\sigma}^2$)」です。

VAEの核心的なアイデアは「潜在空間の全体を正規分布で覆う」ことにあります。訓練時に $q_\phi(\bm{z}|\bm{x})$ を標準正規分布 $\mathcal{N}(\bm{0}, \bm{I})$ に近づけることで、潜在空間のどの点からサンプリングしても、デコーダが意味のある出力を返せるようになります。これが「生成モデル」としての本質的な性質です。

次節では、このアイデアを数式として形式化します。


VAEのアーキテクチャ全体像

VAEは大きく エンコーダ(推論モデル)と デコーダ(生成モデル)の2つから構成されます。まず全体の流れを把握しましょう。

$$ \bm{x} \xrightarrow{\text{エンコーダ}} (\bm{\mu}, \bm{\sigma}^2) \xrightarrow{\text{サンプリング}} \bm{z} \xrightarrow{\text{デコーダ}} \hat{\bm{x}} $$

ここで重要なのは、エンコーダの出力が「潜在変数そのもの」ではなく、潜在変数が従う分布のパラメータ である点です。

確率的エンコーダ $q_\phi(\bm{z}|\bm{x})$

エンコーダは条件付き確率分布 $q_\phi(\bm{z}|\bm{x})$ として定義されます。「入力 $\bm{x}$ が与えられたとき、潜在変数 $\bm{z}$ がどんな分布に従うか」を表します。

VAEでは、この分布を対角ガウス分布と仮定します。

$$ q_\phi(\bm{z}|\bm{x}) = \mathcal{N}\!\left(\bm{z};\, \bm{\mu}_\phi(\bm{x}),\, \text{diag}(\bm{\sigma}^2_\phi(\bm{x}))\right) $$

$\bm{\mu}_\phi(\bm{x})$ と $\bm{\sigma}^2_\phi(\bm{x})$ はニューラルネットワーク(パラメータ $\phi$)が出力します。出力層を2つに分け、一方が平均 $\bm{\mu}$、もう一方が対数分散 $\ln \bm{\sigma}^2$ を出力するのが一般的な実装です。

確率的エンコーダ q(z|x) — 入力を平均μと分散σ²に変換し正規分布としてマッピング

この図が示すように、ひとつの入力画像 $\bm{x}$ に対してひとつの点ではなく、潜在空間上の「分布の山」が対応します。「この画像はだいたいこのあたりの潜在変数に対応する」という不確かさを確率的に表現するわけです。

分散がゼロに近ければ通常のAEに近づき(一点に収束)、分散が大きければ潜在変数は広い範囲から来得ることを意味します。VAEは学習を通してこの分散を適切な大きさに調整します。

確率的デコーダ $p_\theta(\bm{x}|\bm{z})$

デコーダは条件付き確率分布 $p_\theta(\bm{x}|\bm{z})$ として定義されます。「潜在変数 $\bm{z}$ が与えられたとき、どんな $\bm{x}$ を生成するか」を表します。

$$ p_\theta(\bm{x}|\bm{z}) = \mathcal{N}\!\left(\bm{x};\, \bm{\mu}_\theta(\bm{z}),\, \bm{\sigma}^2_\theta(\bm{z})\right) $$

実装では出力をそのままガウス分布の平均とみなし、分散を固定(例: $\sigma^2=1$)することが多いです。この場合、再構成損失は二乗誤差に対応します。

事前分布 $p(\bm{z})$

潜在変数の事前分布には標準正規分布を使います。

$$ p(\bm{z}) = \mathcal{N}(\bm{z};\, \bm{0},\, \bm{I}) $$

この選択は計算上の都合(KL項に閉形式が使える)だけでなく、「潜在空間が原点付近にまとまる」という正則化効果もあります。

これでVAEの3つの構成要素が揃いました。次は、これらのパラメータ $\phi$(エンコーダ)と $\theta$(デコーダ)をどうやって学習するかを見ていきます。


VAEの最適化 — ELBO(Evidence Lower Bound)の導出

VAEの目標は データの周辺対数尤度 $\ln p(\bm{X})$ を最大化すること です。しかし $\ln p(\bm{X})$ を直接計算するのは、潜在変数 $\bm{z}$ の周辺化が必要で計算困難です。

$$ \ln p(\bm{X}) = \ln \int p(\bm{X}, \bm{z})\, d\bm{z} $$

この積分は高次元では解析的に扱えないため、変分推論を使って下から抑えることを考えます。

イェンゼンの不等式で下界を作る

任意の確率分布 $q(\bm{z}|\bm{X})$ を導入し、式変形を進めます。

積分の中に $q(\bm{z}|\bm{X})$ を掛けて割ることで(同じ値を掛けているので等号は保たれます)、

$$ \ln p(\bm{X}) = \ln \int q(\bm{z}|\bm{X})\, \frac{p(\bm{X}, \bm{z})}{q(\bm{z}|\bm{X})}\, d\bm{z} $$

$\ln$ は凹関数なのでイェンゼンの不等式 $\ln \mathbb{E}[f] \geq \mathbb{E}[\ln f]$ が使えます。 $\mathbb{E}$ の添え字は $q(\bm{z}|\bm{X})$ についての期待値です。これを適用すると、

$$ \ln p(\bm{X}) \geq \int q(\bm{z}|\bm{X})\, \ln \frac{p(\bm{X}, \bm{z})}{q(\bm{z}|\bm{X})}\, d\bm{z} =: \mathcal{L}(\bm{X}, \bm{z}) $$

右辺を ELBO(Evidence Lower Bound, 変分下界) と呼びます。これが最大化すべき目標関数です。

ELBOを再構成項とKL項に分解する

ELBOをさらに読み解くために、$p(\bm{X}, \bm{z}) = p_\theta(\bm{X}|\bm{z})\, p(\bm{z})$ を代入して分解します。

まず ELBOの定義から始めます。

$$ \mathcal{L} = \int q_\phi(\bm{z}|\bm{X})\, \ln \frac{p_\theta(\bm{X}|\bm{z})\, p(\bm{z})}{q_\phi(\bm{z}|\bm{X})}\, d\bm{z} $$

対数の分数を分割します。

$$ = \int q_\phi(\bm{z}|\bm{X})\, \ln p_\theta(\bm{X}|\bm{z})\, d\bm{z} + \int q_\phi(\bm{z}|\bm{X})\, \ln \frac{p(\bm{z})}{q_\phi(\bm{z}|\bm{X})}\, d\bm{z} $$

第1項は $q_\phi(\bm{z}|\bm{X})$ についての期待値なので、第2項は KL ダイバージェンスの定義から(符号に注意)、

$$ \mathcal{L} = \underbrace{\mathbb{E}_{q_\phi(\bm{z}|\bm{X})}\!\left[\ln p_\theta(\bm{X}|\bm{z})\right]}_{\text{再構成項}} – \underbrace{D_{\mathrm{KL}}\!\left[q_\phi(\bm{z}|\bm{X})\,\|\, p(\bm{z})\right]}_{\text{KL正則化項}} $$

ELBO分解 — 再構成項とKL正則化項の関係と最大化の仕組み

この図が示すように、$\ln p(\bm{X})$(定数)はELBO $\mathcal{L}$ と KLギャップの和に等しく、ELBOを最大化すると同時に $q$ と $p$ のズレが縮まります。

ELBOの2つの項が何をしているかをそれぞれ直感的に説明します。

再構成項 $\mathbb{E}_{q_\phi}[\ln p_\theta(\bm{x}|\bm{z})]$ は「エンコーダが $\bm{z}$ をサンプリングし、デコーダがそこから $\bm{x}$ を復元できるか」を測ります。ガウスデコーダ(固定分散)の場合、これは平均二乗誤差(MSE)に比例します。値が大きいほど再構成がうまくいっています。

KL正則化項 $D_{\mathrm{KL}}[q_\phi(\bm{z}|\bm{x})\|p(\bm{z})]$ は「エンコーダが出力する分布が事前分布 $p(\bm{z})=\mathcal{N}(\bm{0}, \bm{I})$ からどれだけ離れているか」を測ります。ELBOを最大化するにはこれを 最小化(符号が負なので)する必要があり、潜在分布を標準正規分布に引き戻す正則化の役割を果たします。

次は、このELBOをニューラルネットワークで実際に最適化するための「再パラメータ化トリック」を説明します。


再パラメータ化トリック(Reparametrization Trick)

ELBOの再構成項には $q_\phi(\bm{z}|\bm{x})$ からのサンプル $\bm{z}$ が必要です。しかし「サンプリング操作」は微分できません。勾配が $\phi$ に流れないため、誤差逆伝播法(バックプロパゲーション)でエンコーダを学習できないのです。

再パラメータ化トリック — サンプリングを決定論的変換に置き換えて勾配を流す

この図の左側(再パラメータ化なし)では、サンプリング操作 $\bm{z} \sim q_\phi(\bm{z}|\bm{x})$ がネットワークの中に埋め込まれており、勾配がそこで止まってしまいます。

解決策は「確率性を外に出す」ことです。

$$ \bm{z} = \bm{\mu}_\phi(\bm{x}) + \bm{\sigma}_\phi(\bm{x}) \odot \bm{\varepsilon}, \quad \bm{\varepsilon} \sim \mathcal{N}(\bm{0}, \bm{I}) $$

ノイズ $\bm{\varepsilon}$ をネットワークの外から「定数として」サンプリングし、$\bm{z}$ を $\bm{\mu}$ と $\bm{\sigma}$ の決定論的な関数として表現します。これで $\bm{z}$ は $\phi$ について微分可能となり、勾配が逆伝播できます(図の右側)。

数学的には次のように確認できます。$\bm{z}$ の微分は、

$$ \frac{\partial \bm{z}}{\partial \bm{\mu}} = \bm{1}, \quad \frac{\partial \bm{z}}{\partial \bm{\sigma}} = \bm{\varepsilon} $$

どちらも確定した値を持つため、バックプロパゲーションが問題なく適用できます。$\bm{\varepsilon}$ 自体は $\phi$ に依存しないので、勾配の計算に影響しません。

実装では log_var(対数分散 $\ln \bm{\sigma}^2$)を出力させることが多いです。これは $\bm{\sigma}^2 > 0$ の制約を自然に満たすためです。

import torch
import torch.nn as nn

class VAEEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super().__init__()
        self.fc = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_log_var = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        h = torch.relu(self.fc(x))
        mu = self.fc_mu(h)
        log_var = self.fc_log_var(h)
        return mu, log_var

    def reparametrize(self, mu, log_var):
        # 再パラメータ化トリック: z = mu + sigma * eps
        sigma = torch.exp(0.5 * log_var)   # sigma = sqrt(exp(log_var))
        eps = torch.randn_like(sigma)        # eps ~ N(0, I)
        return mu + sigma * eps

torch.exp(0.5 * log_var) で $\bm{\sigma} = \sqrt{e^{\ln \bm{\sigma}^2}}$ を得ます。torch.randn_like(sigma) は $\sigma$ と同じ形状・デバイスでノイズをサンプリングします。

推論時は $\bm{\varepsilon} = \bm{0}$(つまり $\bm{z} = \bm{\mu}$)として確定的に使うことが多いです。

コードの reparametrize が行っていることを確認します。$\mu = 1.5, \ln\sigma^2 = 0.2$ のとき、$\sigma = e^{0.1} \approx 1.105$、$\bm{\varepsilon} \sim \mathcal{N}(0,1)$ を引いてくれば $\bm{z} = 1.5 + 1.105 \times \varepsilon$ となり、これは $\mathcal{N}(1.5, 1.105^2)$ からのサンプルに相当します。サンプリング前と後で分布が一致しており、トリックが正しく動作していることがわかります。

再パラメータ化トリックを使えばELBOの勾配を($\phi$, $\theta$)の両方について計算できます。次は、KL項の具体的な計算式を見ていきましょう。


KL項の閉形式 — ガウス分布の場合

ELBOのKL項 $D_{\mathrm{KL}}[q_\phi(\bm{z}|\bm{x})\|p(\bm{z})]$ は、一般には積分を数値的に近似する必要がありますが、$q$ と $p$ が両方ガウス分布ならば 解析的に閉形式 で計算できます。これがVAEにガウス分布を仮定するもうひとつの理由です。

$q_\phi(\bm{z}|\bm{x}) = \mathcal{N}(\bm{\mu}, \text{diag}(\bm{\sigma}^2))$、$p(\bm{z}) = \mathcal{N}(\bm{0}, \bm{I})$ のとき、KLダイバージェンスは次のように計算できます。

まず KL の定義から始めます。

$$ D_{\mathrm{KL}}[q \| p] = \int q(\bm{z})\, \ln \frac{q(\bm{z})}{p(\bm{z})}\, d\bm{z} $$

両方ガウス分布なので対数比は二次形式になります。次元 $j$ ごとに独立なので、各次元の KL を足し合わせます。

$$ = \sum_{j=1}^{d} \int \mathcal{N}(z_j;\mu_j, \sigma_j^2)\, \ln \frac{\mathcal{N}(z_j;\mu_j, \sigma_j^2)}{\mathcal{N}(z_j; 0, 1)}\, dz_j $$

各次元の積分を計算します。$\ln \mathcal{N}(z; \mu, \sigma^2) = -\frac{1}{2}\ln(2\pi\sigma^2) – \frac{(z-\mu)^2}{2\sigma^2}$ を代入し、$q$ の期待値をとると、

$$ = \sum_{j=1}^{d} \left( -\frac{1}{2} + \frac{\sigma_j^2}{2} + \frac{\mu_j^2}{2} + \frac{1}{2}\ln\sigma_j^2 \cdot (-1) \right) $$

整理すると($\ln \sigma_j^2 = \log\_var_j$ と書けば)、

$$ D_{\mathrm{KL}} = \frac{1}{2} \sum_{j=1}^{d} \left( \sigma_j^2 + \mu_j^2 – 1 – \ln \sigma_j^2 \right) $$

これがKL項の 閉形式 です。積分を数値的に解く必要はなく、$\bm{\mu}$ と $\bm{\sigma}^2$ さえわかれば即座に計算できます。

def kl_loss(mu, log_var):
    # KL(N(mu, sigma^2) || N(0, I)) の閉形式
    # = 0.5 * sum(sigma^2 + mu^2 - 1 - log_var)
    return 0.5 * torch.sum(torch.exp(log_var) + mu**2 - 1 - log_var, dim=1).mean()

log_var を使って実装すると、torch.exp(log_var) が $\sigma^2$ に、log_var が $\ln \sigma^2$ に対応します。dim=1 で潜在次元方向に和を取り、バッチ平均を返します。

KL項の性質を確認しましょう。

  • $\bm{\mu} = \bm{0}$, $\bm{\sigma}^2 = \bm{1}$($\ln\bm{\sigma}^2 = \bm{0}$)のとき: $D_{\mathrm{KL}} = \frac{1}{2}(1 + 0 – 1 – 0) = 0$(完全一致)
  • $\mu_j$ が大きいほど KL が大きくなる(潜在分布が原点から遠ざかる)
  • $\sigma_j^2$ が1から離れるほど KL が大きくなる(分散が崩れる)

KL項の閉形式 — ガウス分布間のKLダイバージェンスとσ・μへの依存

左のパネルを見ると、$q_\phi(\bm{z}|\bm{x})$(色付きの山)が事前分布 $p(\bm{z})$(黒破線)から離れるほど KL が大きくなることが視覚的にわかります。右のパネルでは、KL閉形式のグラフを見ると $\sigma=1, \mu=0$ でちょうど KL=0 になり、そこから $\sigma$ が大きくなるにつれ KL も増加することが確認できます。学習中はこのKL値がほどよく小さくなるよう $\phi$ が更新されていきます。


VAEの損失関数まとめ

ELBOを最大化することは、その負の値(損失)を最小化することと等価です。

$$ \mathcal{L}_{\mathrm{VAE}} = \underbrace{\mathbb{E}_{q_\phi(\bm{z}|\bm{x})}\!\left[-\ln p_\theta(\bm{x}|\bm{z})\right]}_{\text{再構成損失 (MSEなど)}} + \underbrace{D_{\mathrm{KL}}\!\left[q_\phi(\bm{z}|\bm{x})\,\|\, p(\bm{z})\right]}_{\text{KL正則化損失}} $$

実装では次のように書きます。

def vae_loss(x, x_recon, mu, log_var):
    # 再構成損失: 平均二乗誤差(ガウスデコーダに対応)
    recon = nn.functional.mse_loss(x_recon, x, reduction="sum") / x.size(0)

    # KL正則化損失: 閉形式
    kl = 0.5 * torch.sum(torch.exp(log_var) + mu**2 - 1 - log_var, dim=1).mean()

    return recon + kl

reduction="sum" でバッチ内全ピクセルの二乗誤差を合計してからバッチサイズで割ります。KLも同じスケールで計算します。


潜在空間の構造と生成の仕組み

VAEを学習した後の潜在空間は「きれいな構造」を持ちます。

AEとVAEの潜在空間の違い — VAEは正規分布に整列しクラスが連続的に分布

左のAEでは各クラスが空間のあちこちに散在しており、クラス間の補間点が「何もない領域」を通ることがあります。右のVAEでは標準正規分布の輪郭(黒破線)の内側に各クラスがまとまり、クラス間の補間が意味のある変化をたどります。

トイデータの元データ空間とVAEの潜在空間 — 正規分布に整列した2次元マップ

左が元の2次元データ空間(月型の分布)、右がVAEの潜在空間です。潜在空間では2つのクラスが原点付近に整列し、標準正規分布の$1\sigma$/2$\sigma$輪郭(黒破線)の中に収まっています。このように整列した潜在空間からサンプリングすれば、どの点でも意味のある出力が得られます。

潜在空間からの新しいサンプル生成

生成時は次のように行います。

  1. 事前分布からサンプリング: $\bm{z} \sim \mathcal{N}(\bm{0}, \bm{I})$
  2. デコーダに入力して出力を得る: $\hat{\bm{x}} = \text{Decoder}(\bm{z})$
@torch.no_grad()
def generate(decoder, n_samples=16, latent_dim=2):
    # 標準正規分布から潜在変数をサンプリング
    z = torch.randn(n_samples, latent_dim)
    # デコーダで生成
    x_gen = decoder(z)
    return x_gen

また潜在空間をグリッド状にサンプリングすることで、「潜在変数を連続的に動かしたときに出力がどう変化するか」を可視化できます。

@torch.no_grad()
def grid_decode(decoder, n_grid=10, z_range=2.5):
    # 2次元潜在空間をグリッドサンプリング
    z1 = torch.linspace(-z_range, z_range, n_grid)
    z2 = torch.linspace(-z_range, z_range, n_grid)
    Z1, Z2 = torch.meshgrid(z1, z2, indexing="ij")
    z = torch.stack([Z1.ravel(), Z2.ravel()], dim=1)  # (n_grid^2, 2)
    return decoder(z)

VAEの生成サンプル — 潜在空間グリッドサンプリングとデコーダ出力の概念図

左パネルの各点 $\bm{z}$ をデコーダに通した結果が右パネルです。右パネルでは色が滑らかに変化しており、潜在空間上で隣り合う点ほど類似した出力が得られることを示しています。この「補間の滑らかさ」がVAEの生成モデルとしての強みです。


PyTorch実装 — 完全なVAEコード

これまでの理論をすべて組み込んだ完全なVAE実装です。2次元の合成データ(2クラスのガウス分布)で動作を検証します。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

torch.manual_seed(42)
np.random.seed(42)

# --- データ生成(2次元トイデータ)---
n = 500
X_np = np.vstack([
    np.random.randn(n, 2) * 0.5 + [2.0,  2.0],   # クラス0
    np.random.randn(n, 2) * 0.5 + [-2.0, -2.0],   # クラス1
]).astype(np.float32)
X = torch.from_numpy(X_np)
y_label = torch.cat([torch.zeros(n), torch.ones(n)]).long()

INPUT_DIM  = 2
HIDDEN_DIM = 32
LATENT_DIM = 2
BATCH_SIZE = 64
EPOCHS     = 200
LR         = 1e-3

# --- エンコーダ ---
class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.shared = nn.Sequential(nn.Linear(INPUT_DIM, HIDDEN_DIM), nn.ReLU())
        self.fc_mu  = nn.Linear(HIDDEN_DIM, LATENT_DIM)
        self.fc_lv  = nn.Linear(HIDDEN_DIM, LATENT_DIM)  # log_var

    def forward(self, x):
        h = self.shared(x)
        mu = self.fc_mu(h)
        log_var = torch.clamp(self.fc_lv(h), -6, 6)  # 数値安定化
        return mu, log_var

# --- デコーダ ---
class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(LATENT_DIM, HIDDEN_DIM), nn.ReLU(),
            nn.Linear(HIDDEN_DIM, INPUT_DIM)
        )

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

# --- VAEモデル ---
class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def reparametrize(self, mu, log_var):
        sigma = torch.exp(0.5 * log_var)
        eps   = torch.randn_like(sigma)
        return mu + sigma * eps

    def forward(self, x):
        mu, log_var = self.encoder(x)
        z = self.reparametrize(mu, log_var)
        x_recon = self.decoder(z)
        return x_recon, mu, log_var

def vae_loss(x, x_recon, mu, log_var):
    recon = F.mse_loss(x_recon, x, reduction="sum") / x.size(0)
    kl    = 0.5 * torch.sum(torch.exp(log_var) + mu**2 - 1 - log_var, dim=1).mean()
    return recon + kl, recon.item(), kl.item()

# --- 学習ループ ---
model = VAE()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
dataset = torch.utils.data.TensorDataset(X)
loader  = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

history_recon, history_kl = [], []
for epoch in range(EPOCHS):
    for (xb,) in loader:
        optimizer.zero_grad()
        x_recon, mu, log_var = model(xb)
        loss, r, k = vae_loss(xb, x_recon, mu, log_var)
        loss.backward()
        optimizer.step()
    if epoch % 20 == 0:
        history_recon.append(r)
        history_kl.append(k)
        print(f"Epoch {epoch:3d}: recon={r:.4f}, KL={k:.4f}")

このコードを実行すると次のような出力が得られます(値は実行ごとに若干変動します)。

Epoch   0: recon=8.4967, KL=0.1730
Epoch  20: recon=1.0815, KL=1.4111
Epoch  40: recon=0.7551, KL=1.1972
...
Epoch 180: recon=0.5155, KL=1.0519

学習が進むにつれて再構成損失は単調に減少し、KL損失は最初ほぼゼロから徐々に適切な値へ増加します。初期段階では再構成を優先し、学習後半になるほど正規分布への整列が進む様子が読み取れます。

再構成損失が下がるのは「エンコーダとデコーダが連携してデータを圧縮・復元できるようになった」ことを示します。KL損失が増えることは「潜在分布が事前分布から離れつつある」ように見えますが、これは最初に KL=0($\sigma \to 0$, 一点に収束)から始まるため一時的に上昇するのであり、最終的にはバランスのとれた値に落ち着きます。

VAE学習曲線 — 再構成損失とKL損失の推移とKLアニーリング

左パネルが標準的なVAEの学習曲線、右パネルが「KLアニーリング」の効果を示しています。KLアニーリングとは、最初は $\beta=0$(KL項を無視)として再構成だけを学習し、徐々に $\beta$ を1まで増やす手法です。再構成の基盤ができてからKL正則化を加えるため、潜在変数が崩壊(KL=0 に張り付いて何も学習しない「後部崩壊」)するリスクが下がります。


β-VAE — KL正則化の強さを調整する

通常のVAEは再構成とKL項をそのまま足し合わせますが、KL項の重み $\beta$ を調整したものが $\beta$-VAE です。

$$ \mathcal{L}_{\beta\text{-VAE}} = \mathbb{E}\!\left[-\ln p_\theta(\bm{x}|\bm{z})\right] + \beta \cdot D_{\mathrm{KL}}\!\left[q_\phi(\bm{z}|\bm{x})\,\|\, p(\bm{z})\right] $$

$\beta = 1$ が通常のVAE、$\beta > 1$ にするほど正則化が強くなり、各潜在次元が統計的に独立した意味をもちやすくなる(表現の絡まりほぐし, disentanglement) という性質があります。

たとえば顔画像のVAEで $\beta$ を大きくすると、$z_1$ が「笑顔かどうか」、$z_2$ が「顔の向き」のように、次元ごとに独立した要因を表すことが報告されています。

beta = 4.0  # デフォルト1より強い正則化

def beta_vae_loss(x, x_recon, mu, log_var, beta=4.0):
    recon = F.mse_loss(x_recon, x, reduction="sum") / x.size(0)
    kl    = 0.5 * torch.sum(torch.exp(log_var) + mu**2 - 1 - log_var, dim=1).mean()
    return recon + beta * kl

ただしトレードオフがあります。$\beta$ が大きいほど再構成品質は下がります。KL項が強すぎると潜在変数が $\mathcal{N}(\bm{0}, \bm{I})$ に押しつぶされ、$\bm{x}$ の情報が潜在空間に残らなくなるからです。

β-VAEのβ値とトレードオフ — 再構成品質・KL損失・絡まりほぐし度の変化

左パネルが潜在変数トラバーサル($z_1$ を動かしたときの出力への影響)の模式図で、$\beta=4$ の方がより明確な変化を示しています。右パネルが $\beta$ を変えたときの各指標の変化を示し、再構成品質・KL損失・絡まりほぐし度は互いにトレードオフの関係にあります。タスクに応じて適切な $\beta$ を選ぶことが重要です。


VAEのグラフィカルモデル

VAEの確率変数の関係をグラフィカルモデルで整理します。

$$ p_\theta(\bm{x}, \bm{z}) = p_\theta(\bm{x}|\bm{z})\, p(\bm{z}) $$

生成方向($\bm{z} \to \bm{x}$): 潜在変数 $\bm{z}$ が先に決まり、そこからデータ $\bm{x}$ が生成されます。

推論方向($\bm{x} \to \bm{z}$): データ $\bm{x}$ が観測され、潜在変数 $\bm{z}$ を事後分布から推定します。

VAEでは、この推論を真の事後分布 $p_\theta(\bm{z}|\bm{x})$ で行いたいのですが、直接計算できません(ベイズの定理の分母の周辺尤度が計算困難)。そのため、近似分布 $q_\phi(\bm{z}|\bm{x})$ でこれを近似します。これが変分推論の核心です。

パラメータ $\phi$(エンコーダ)と $\theta$(デコーダ)は両方ともニューラルネットワークのパラメータであり、ELBOを最大化することで同時に最適化されます。


まとめ

本記事では、VAE(変分オートエンコーダー)について以下の内容を解説しました。

  • AEとの違い: エンコーダが一点でなく「分布」を出力することで、潜在空間に構造が生まれ、生成モデルとして機能する
  • 確率的エンコーダ $q_\phi(\bm{z}|\bm{x})$: 入力を正規分布のパラメータ($\bm{\mu}$, $\bm{\sigma}^2$)にマッピングする
  • ELBO: 周辺対数尤度をイェンゼンの不等式で下から抑えた下界。再構成項とKL正則化項に分解できる
  • 再パラメータ化トリック: $\bm{z} = \bm{\mu} + \bm{\sigma} \odot \bm{\varepsilon}$ として確率性を外に出し、勾配を流せるようにする
  • KL閉形式: ガウス分布間のKLは $\frac{1}{2}\sum(\sigma_j^2 + \mu_j^2 – 1 – \ln\sigma_j^2)$ で解析的に計算できる
  • β-VAE: $\beta > 1$ にすることで絡まりほぐし(disentanglement)を促進できる

VAEは「生成」だけでなく、意味のある潜在表現の学習・データ補完・異常検知など幅広い場面で活用されています。次のステップとして、VAEを異常検知に応用した手法や、さらに強力な生成モデルである正規化フロー、拡散モデルも学んでみてください。

画像なし
オートエンコーダの種類と理論を体系的に解説
基本構造からスパースAE・デノイジングAE・収縮AEまで、各種オートエンコーダの理論とVAEへの接続を解説します。
KLダイバージェンス完全ガイド
KLダイバージェンスの定義・非対称性・ガウス分布での閉形式・VAEでの役割を直感と数式で解説します。