変分オートエンコーダ (VAE)

変分オートエンコーダ(VAE)のELBOをJensenの不等式から導出し、ガウス分布デコーダとの対応、KLダイバージェンスの解析解を導出。手動で誤差逆伝播を導出したNumPy実装で実際にELBOを最大化し、非識別性という落とし穴も検証します。

変分オートエンコーダ (Variational AutoEncoder, VAE) は、生成モデルの一種であり、データの潜在的な構造を学習し、新しいデータを生成することを目的としています。EMアルゴリズムが対数尤度下界を最大化するのと同様に( EMアルゴリズム 参照)、VAEも変分下界(ELBO)を最大化することで学習を行います。

VAEの大きな特徴は、**確率的なエンコーダ(認識モデル)確率的なデコーダ(生成モデル)**を持つ点です。これにより、潜在空間が滑らかになり、意味のあるデータ生成が可能になります。

  • 認識モデル(エンコーダ): 入力データ \(x\) から、潜在変数 \(z\) の確率分布 \(q_\phi(z|x)\) を推定します。パラメータ \(\phi\) を持ちます。
  • 生成モデル(デコーダ): 潜在変数 \(z\) から、データ \(x\) の確率分布 \(p_\theta(x|z)\) を生成します。パラメータ \(\theta\) を持ちます。

潜在変数 \(z\) は、入力データ \(x\) の持つ情報をより低次元で抽象的な「潜在表現」または「潜在コード」として表現していると解釈できます。

オートエンコーダ (AutoEncoder)

オートエンコーダは、入力データを圧縮・符号化し、それを再構築(復号化)することで、入力データと同じ内容を出力するように学習するニューラルネットワークです。中間層(潜在空間)は、入力データの重要な特徴を捉えた圧縮表現(符号)となります。

VAEは、このオートエンコーダの枠組みに変分推論の概念を導入したものです。

VAEにおけるELBOの導出

VAEの目標は、周辺尤度 \(p_\theta(x) = \int p_\theta(x|z)p(z)dz\) を最大化することですが、潜在変数 \(z\) に関する積分は一般に解析的に計算できません。EMアルゴリズムの記事( EMアルゴリズム )で用いたのと同じ Jensen の不等式による議論を、認識モデル \(q_\phi(z|x)\) を使って適用します。

\[ \log p_\theta(x) = \log \int p_\theta(x,z)dz = \log \mathbb{E}_{q_\phi(z|x)}\left[\frac{p_\theta(x,z)}{q_\phi(z|x)}\right] \ge \mathbb{E}_{q_\phi(z|x)}\left[\log\frac{p_\theta(x,z)}{q_\phi(z|x)}\right] =: \mathcal{L}(\theta,\phi;x) \]

(対数関数が凹関数であることによる Jensen の不等式です。)EMのときと同様、対数尤度とELBOの差はKLダイバージェンスとして厳密に評価できます。

\[ \log p_\theta(x) - \mathcal{L}(\theta,\phi;x) = KL(q_\phi(z|x) \| p_\theta(z|x)) \]

\(KL \ge 0\) なので \(\mathcal{L}(\theta,\phi;x)\) は対数尤度の下界(ELBO)であり、等号成立は \(q_\phi(z|x)=p_\theta(z|x)\) (真の事後分布に一致する)ときのみです。EMアルゴリズムとの違いは、EMでは各データ点ごとに厳密な事後分布 \(p(z|x,\hat\theta)\) を計算していたのに対し、VAEでは \(q_\phi(z|x)\) をニューラルネットワークで近似し、全データ点に共通のパラメータ \(\phi\) で「まとめて」推論する(償却推論, amortized inference)点です。これにより厳密な事後分布計算は不要になりますが、その代償として \(q_\phi\) と真の事後分布の間に恒常的なギャップ(\(KL>0\) )が残ります。

ELBOを展開すると、おなじみの2項に分解できます。

\[ \mathcal{L}(\theta,\phi;x) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - KL(q_\phi(z|x) \| p(z)) \]

(\(p_\theta(x,z) = p_\theta(x|z)p(z)\) を代入し、\(\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x,z)] - \mathbb{E}_{q_\phi(z|x)}[\log q_\phi(z|x)]\) を整理すると得られます。)

  1. 再構成誤差(第一項): \(\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]\) エンコーダが生成した潜在変数 \(z\) を使って、デコーダが元の入力 \(x\) をどれだけ正確に再構成できるかを表します。この項を最大化することは、再構成誤差を最小化することに相当します。

  2. 正則化項(第二項): \(KL(q_\phi(z|x) \| p(z))\) エンコーダが推定した潜在変数 \(z\) の分布 \(q_\phi(z|x)\) が、事前に定めた潜在変数の事前分布 \(p(z)\) にどれだけ近いかを表します。この項を最小化することは、潜在空間を滑らかにし、意味のあるデータ生成を可能にする役割があります。

認識モデル \(q_\phi(z|x)\)

認識モデルは、入力 \(x\) から潜在変数 \(z\) の分布を推定します。一般的に、ニューラルネットワークを用いて、潜在変数 \(z\) が従う多変量正規分布の平均 \(\mu(x)\) と分散 \(\sigma^2(x)\) (または対数分散 \(\log \sigma^2(x)\) )を出力します。

\[ q_\phi(z|x) = \mathcal{N}(z \mid \mu_\phi(x), \text{diag}(\sigma^2_\phi(x))) \]

ここで、\(\mu_\phi(x)\) と \(\sigma^2_\phi(x)\) は、入力 \(x\) を受け取るニューラルネットワークの出力です。

生成モデル \(p_\theta(x|z)\) とガウス分布デコーダの再構成誤差

生成モデルは、潜在変数 \(z\) からデータ \(x\) の分布を生成します。どのような確率分布を用いるかは、データ \(x\) の種類に依存します。

  • 二値データ(例: 白黒画像): ベルヌーイ分布(またはカテゴリカル分布)
  • 連続値データ(例: グレースケール画像): ガウス分布

連続値データの場合、分散を固定した多変量正規分布を用いることがよくあります。

\[ p_\theta(x|z) = \mathcal{N}(x \mid \nu_\theta(z), \sigma_x^2 I) \]

ここで、\(\nu_\theta(z)\) は潜在変数 \(z\) を受け取るニューラルネットワークの出力です。このとき、再構成誤差の項を具体的に計算すると、二乗誤差損失(MSE)と等価であることが導出できます。\(D\) 次元のデータ \(x\) に対して、

\[ \log p_\theta(x|z) = \log \mathcal{N}(x \mid \nu_\theta(z), \sigma_x^2 I) = -\frac{D}{2}\log(2\pi\sigma_x^2) - \frac{1}{2\sigma_x^2}\lVert x - \nu_\theta(z) \rVert^2 \]

右辺の第1項は \(z\) に依存しない定数なので、ELBOを \(\theta,\phi\) について最大化する際には無視でき、

\[ \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] = -\frac{1}{2\sigma_x^2}\mathbb{E}_{q_\phi(z|x)}\left[\lVert x - \nu_\theta(z) \rVert^2\right] + \text{const} \]

となります。つまり再構成誤差項を最大化することは、デコーダの出力 \(\nu_\theta(z)\) と入力 \(x\) の平均二乗誤差を最小化することと定数倍・符号を除いて同じです。実装上「再構成損失にMSEを使う」というのは、暗黙にガウス分布デコーダ(分散固定)を仮定していることを意味します。

潜在変数の事前分布 \(p(z)\) とKLダイバージェンスの解析解

潜在変数の事前分布は、通常、標準正規分布(平均0、分散1)の積として定義されます。

\[ p(z) = \prod_{j=1}^k \mathcal{N}(z_j | 0, 1) \]

この事前分布は、学習中に固定され、パラメータ \(\theta\) や \(\phi\) には依存しません。\(q_\phi(z|x)\) と \(p(z)\) がともにガウス分布であるため、正則化項のKLダイバージェンスは積分計算をせずに解析的な閉形式で求まります。1次元の場合(\(q=\mathcal{N}(\mu,\sigma^2)\) 、\(p=\mathcal{N}(0,1)\) )を導出します。

\[ KL(q\|p) = \mathbb{E}_q[\log q(z)] - \mathbb{E}_q[\log p(z)] \]

\(\log q(z) = -\frac{1}{2}\log(2\pi\sigma^2) - \frac{(z-\mu)^2}{2\sigma^2}\) の期待値は、\(\mathbb{E}_q[(z-\mu)^2]=\sigma^2\) より

\[ \mathbb{E}_q[\log q(z)] = -\frac{1}{2}\log(2\pi\sigma^2) - \frac{1}{2} \]

同様に \(\log p(z) = -\frac{1}{2}\log(2\pi) - \frac{z^2}{2}\) の期待値は、\(\mathbb{E}_q[z^2] = \sigma^2+\mu^2\) (分散の定義 \(\text{Var}(z)=\mathbb{E}[z^2]-\mu^2\) より)を用いて

\[ \mathbb{E}_q[\log p(z)] = -\frac{1}{2}\log(2\pi) - \frac{\sigma^2+\mu^2}{2} \]

両者の差を取ると、

\[ KL(q\|p) = \left[-\frac{1}{2}\log(2\pi\sigma^2) - \frac{1}{2}\right] - \left[-\frac{1}{2}\log(2\pi) - \frac{\sigma^2+\mu^2}{2}\right] = \frac{1}{2}\left(\sigma^2 + \mu^2 - 1 - \log\sigma^2\right) \]

\(q_\phi(z|x)\) の各次元が独立(対角共分散)なので、多次元の場合はこれを次元ごとに足し合わせるだけです。

\[ KL(q_\phi(z|x) \| p(z)) = \frac{1}{2}\sum_{j=1}^k\left(\sigma_j^2 + \mu_j^2 - 1 - \log\sigma_j^2\right) \]

この式が積分を経ずに閉形式で書けることが、VAEの学習を実用的にしている大きな理由の一つです(再構成誤差項はモンテカルロサンプリングで近似する必要がありますが、KL項は正確に計算できます)。

勾配降下法と再パラメータ化トリック

VAEの学習は、変分下界 \(\mathcal{L}(\theta, \phi)\) を最大化するために、勾配降下法(Adamなどの最適化アルゴリズム)を用いてパラメータ \(\theta\) と \(\phi\) を更新します。

再構成誤差項の勾配は比較的容易に計算できますが、期待値 \(\mathbb{E}_{q_\phi(z|x)}[\cdot]\) の計算自体が \(\phi\) に依存する分布の上で行われるため、素朴にモンテカルロサンプリングした \(z\sim q_\phi(z|x)\) を使って勾配を取ろうとすると、サンプリング操作を通り抜けて \(\phi\) に逆伝播できないという問題が生じます。

この問題を解決するために、再パラメータ化トリック (Reparameterization Trick) が用いられます。これは、潜在変数 \(z\) を、パラメータ \(\phi\) に依存しない確率変数 \(\epsilon\) と、パラメータ \(\phi\) に依存する決定的な関数 \(g(\epsilon, x, \phi)\) を用いて表現する手法です。ガウス分布の場合、\(z = \mu_\phi(x) + \sigma_\phi(x) \cdot \epsilon\) (\(\epsilon \sim \mathcal{N}(0, I)\) )と表現できます。トリックの数式的な詳細や多次元・非ガウス分布への一般化は リパラメータ化トリック で扱っているので、そちらを参照してください。

ここでは、このトリックが単に「微分可能にする」だけでなく、勾配推定の分散を大きく下げるという実用上決定的な利点を持つことを補足します。比較対象となるのが、サンプリング操作を含む期待値の勾配を得るもう一つの一般的な手法、**スコア関数推定量(REINFORCE)**です。

\[ \nabla_\phi \mathbb{E}_{q_\phi(z)}[f(z)] = \mathbb{E}_{q_\phi(z)}\left[f(z)\nabla_\phi \log q_\phi(z)\right] \]

この推定量は \(f\) が微分不可能でも使える汎用性がありますが、\(f(z)\) の値そのものが勾配の重みにかかるため、\(f(z)\) の分散が大きいと推定量の分散も大きくなりやすいという欠点があります。一方、再パラメータ化トリックによる勾配

\[ \nabla_\phi \mathbb{E}_{q_\phi(z)}[f(z)] = \mathbb{E}_{\epsilon}\left[\nabla_z f(z)\cdot\nabla_\phi g(\epsilon,\phi)\right] \]

は \(f\) の勾配情報(パスワイズ微分)を直接利用するため、\(f\) が滑らかな場合には経験的にも理論的にも分散が小さくなることが知られています(Kingma & Welling, 2014; Rezende et al., 2014)。VAEのデコーダは通常ニューラルネットワーク(微分可能)なので、この低分散な再パラメータ化勾配が使える状況にちょうど当てはまります。

Python実行検証:線形ガウスVAEを手動導出した勾配で学習する

ELBOの最大化を実際に手を動かして確認するため、ニューラルネットワークのライブラリ(PyTorchなど)を使わず、すべての勾配を上記の導出に基づいて手計算し、NumPyだけで実装します。エンコーダ・デコーダを線形写像に単純化した「線形ガウスVAE」を、1次元の潜在変数から生成した2次元データで学習します。

データ生成過程: 真の潜在変数 \(z^* \sim \mathcal{N}(0,1)\) から、真の方向ベクトル \(w^*=(2,1)\) とオフセット \((1,-0.5)\) を使って2次元データ \(x = z^* w^* + (1,-0.5) + \text{noise}\) (\(\text{noise}\sim\mathcal{N}(0,0.3^2 I)\) )を生成します。

モデル: \(\mu_\phi(x)=w_\mu^\top x + b_\mu\) 、\(\log\sigma^2_\phi(x) = w_{lv}^\top x + b_{lv}\) (ともにスカラー出力)、\(\nu_\theta(z) = w_{dec}\,z + b_{dec}\) (2次元出力)、デコーダ分散 \(\sigma_x^2=0.3^2\) 固定。

勾配の導出: 損失 \(L=-\mathcal{L}\) (負のELBO、1サンプルあたり)は、再パラメータ化 \(z=\mu+\sigma\epsilon\) (\(\sigma=\exp(\frac12\log\sigma^2)\) )を用いて

\[ L = \underbrace{\frac{1}{2\sigma_x^2}\lVert x-\hat x \rVert^2}_{\text{再構成誤差}} + \underbrace{\frac{1}{2}\left(e^{\log\sigma^2}+\mu^2-1-\log\sigma^2\right)}_{\text{KL項}}, \qquad \hat x = w_{dec}z+b_{dec} \]

と書け、連鎖律を逆向きにたどると次の勾配が得られます(\(d(\cdot)\) は \(L\) による勾配を表します)。

\[ d\hat x = -\frac{1}{\sigma_x^2}(x-\hat x), \quad dw_{dec}=d\hat x\cdot z, \quad db_{dec}=d\hat x, \quad dz = d\hat x \cdot w_{dec} \] \[ d\mu = dz + \mu, \qquad d\log\sigma^2 = dz\cdot\epsilon\cdot\tfrac12\sigma + \tfrac12(e^{\log\sigma^2}-1) \] \[ dw_\mu = d\mu\cdot x,\quad db_\mu=d\mu,\quad dw_{lv}=d\log\sigma^2\cdot x,\quad db_{lv}=d\log\sigma^2 \]

(\(d\mu\) の第2項 \(\mu\) 、\(d\log\sigma^2\) の第2項 \(\tfrac12(e^{\log\sigma^2}-1)\) はKL項をそれぞれ \(\mu,\log\sigma^2\) で微分した寄与です。)

import numpy as np

rng = np.random.default_rng(0)

N = 500
true_direction = np.array([2.0, 1.0])
z_true = rng.normal(0, 1, size=N)
noise = rng.normal(0, 0.3, size=(N, 2))
X = z_true[:, None] * true_direction[None, :] + noise + np.array([1.0, -0.5])

sigma_x2 = 0.3**2

w_mu = rng.normal(0, 0.1, size=2); b_mu = 0.0
w_lv = rng.normal(0, 0.1, size=2); b_lv = 0.0
w_dec = rng.normal(0, 0.1, size=2); b_dec = rng.normal(0, 0.1, size=2)

lr = 0.005
n_epochs = 400
elbo_history = []

for epoch in range(n_epochs):
    eps = rng.normal(size=N)
    mu = X @ w_mu + b_mu
    logvar = np.clip(X @ w_lv + b_lv, -8, 8)
    sigma = np.exp(0.5 * logvar)
    z = mu + sigma * eps

    xhat = z[:, None] * w_dec[None, :] + b_dec
    recon = (1.0 / (2 * sigma_x2)) * np.sum((X - xhat) ** 2, axis=1)
    kl = 0.5 * (np.exp(logvar) + mu**2 - 1 - logvar)
    elbo_history.append(-(recon + kl).mean())

    d_xhat = -(1.0 / sigma_x2) * (X - xhat)
    d_w_dec = (d_xhat * z[:, None]).mean(axis=0)
    d_b_dec = d_xhat.mean(axis=0)
    d_z = d_xhat @ w_dec

    d_mu = d_z + mu
    d_logvar = d_z * eps * 0.5 * sigma + 0.5 * (np.exp(logvar) - 1)

    d_w_mu = (d_mu[:, None] * X).mean(axis=0)
    d_b_mu = d_mu.mean()
    d_w_lv = (d_logvar[:, None] * X).mean(axis=0)
    d_b_lv = d_logvar.mean()

    w_dec -= lr * d_w_dec; b_dec -= lr * d_b_dec
    w_mu  -= lr * d_w_mu;  b_mu  -= lr * d_b_mu
    w_lv  -= lr * d_w_lv;  b_lv  -= lr * d_b_lv

print(f"epoch   1: ELBO = {elbo_history[0]:.4f}")
print(f"epoch  10: ELBO = {elbo_history[9]:.4f}")
print(f"epoch  50: ELBO = {elbo_history[49]:.4f}")
print(f"epoch 100: ELBO = {elbo_history[99]:.4f}")
print(f"epoch 400: ELBO = {elbo_history[-1]:.4f}")

実行結果は以下の通りです。

epoch   1: ELBO = -36.9169
epoch  10: ELBO = -30.7112
epoch  50: ELBO = -5.4867
epoch 100: ELBO = -4.5815
epoch 400: ELBO = -3.3315

手計算した勾配だけでELBOが \(-36.92 \to -3.33\) と単調に近い形で改善していくことが確認できました。学習後のパラメータを真の生成過程と比較すると、

true decoder direction   w* = [2.0, 1.0]
learned decoder weight w_dec = [-1.4122, -0.7163]
cosine similarity(w_dec, w*) = -1.0000
learned b_dec = [0.804, -0.6096]  (true offset = [1.0, -0.5])
corr(encoder mean mu(x), true latent z_true) = -0.9920

エンコーダの出力 \(\mu(x)\) と真の潜在変数 \(z^*\) の相関は \(-0.992\) と極めて強く、潜在変数が正しく回復されていることが分かります。一方でコサイン類似度がちょうど \(-1.0\) であること、つまり学習された方向ベクトルが真の方向のちょうど逆符号になっていることに注目してください。これはバグではなく、次のエッジケースで説明する潜在変数モデルの符号の非識別性です。

学習曲線を可視化すると、以下のようになります。

線形ガウスVAEの学習曲線。手動導出した勾配のみ(NumPy実装)でELBOが単調に近い形で改善していく

エッジケース・注意点

  • 潜在変数の非識別性(符号・回転の任意性): 上記の実験で確認したように、\(z\) を \(-z\) に、\(w_{dec}\) を \(-w_{dec}\) に同時に反転しても、尤度 \(p_\theta(x|z)\) は変化しません(一般には、線形変換 \(z \to Az\) 、\(w_{dec} \to w_{dec}A^{-1}\) に対して不変です)。これは通常は無害な非識別性ですが、Wang, Blei & Cunningham (2023) は、この非識別性が極端な形で現れると事後分布崩壊 (posterior collapse) ——エンコーダが入力 \(x\) を無視して \(q_\phi(z|x) \approx p(z)\) に潰れてしまい、潜在表現が意味を失う現象——につながることを理論的に示しています。デコーダの表現力が高い(例えばRNNやTransformerのように、\(z\) を使わずとも系列の自己回帰だけで再構成できてしまう)場合に特に起こりやすいことが知られています。
  • KL項の消失とその緩和策: 事後分布崩壊が起きると、正則化項 \(KL(q_\phi(z|x)\|p(z))\) が0に近づき、\(z\) が入力の情報を運ばなくなります。緩和策として、KL項の重みを学習初期は小さくして徐々に増やす KLアニーリング や、KL項に下限を設ける free bits といった手法が使われます。
  • デコーダの分散 \(\sigma_x^2\) の選び方: 本記事の導出で見た通り、\(\sigma_x^2\) を固定するとMSE損失に帰着しますが、\(\sigma_x^2\) を大きくすると再構成誤差項の重みが相対的に小さくなり、KL項(正則化)が支配的になるため、\(\sigma_x^2\) はKL項とのバランスを左右するハイパーパラメータとして機能します(\(\beta\) -VAEの \(\beta\) と類似の役割)。

参考