VAE（Variational Autoencoder）

VAEは、オートエンコーダの「再構成能力」に、生成モデルとしての「サンプリング可能性」を加えたモデルです。潜在変数を確率分布で扱うことで、学習後に新しいデータを生成できるのが最大の特徴です。

VAE は「圧縮して戻す」だけでなく、「そこからサンプルできる空間」を作る

普通のオートエンコーダは再構成が上手ければ十分ですが、VAE は生成モデルなので、潜在空間に乱数を入れてもそれらしいサンプルが出てほしいという追加条件があります。そのために ELBO と KL 項が入ります。

このノートでは、まず VAE の目的関数を再構成と正則化の綱引きとして捉え、次に再パラメータ化が勾配計算をどう助けるかを見ます。そのあとで beta-VAE と posterior collapse まで進み、潜在空間を整えることの代償も確認します。

読み筋は「再構成したい」と「生成しやすくしたい」の両立です

VAE の面白さは、うまく戻したいだけなら不要な KL 項を、あえて入れているところにあります。後半では、この余計に見える制約が潜在空間をどう変え、何を失わせるかを見ます。

最初の山場は ELBO の 2 項の役割分担です

再構成項は入力に忠実であることを押し、KL 項は潜在分布を標準正規へ寄せます。どちらか一方だけを見ても VAE の振る舞いは理解しにくいので、常に 2 項を並べて読みます。

再パラメータ化は、乱数を消すのではなく位置をずらす書き方です

z を直接サンプルすると勾配が通りにくいので、mu + sigma * eps に書き換えます。ここで分けたいのは、学習したい量と、ただのノイズ源です。

後半では、VAE がうまくいきすぎて潜在を使わなくなる失敗も見る

posterior collapse は、潜在変数モデルとしてはかなり本質的な失敗です。潜在を入れたのに潜在を無視する、という事態がどう起こるかを最後に確認します。

ここでは VAE の核だけをむき出しで見る

高品質画像生成のための巨大な実装ではなく、ELBO、再パラメータ化、beta のトレードオフ、collapse という VAE の骨格を読み取るための notebook です。

まずは ELBO の中身を並べる

最初の節では、再構成項と KL 項がどんな引っ張り合いをしているかを、数字として見ます。

import math
import random
from statistics import mean

random.seed(33)

まず、なぜ普通のオートエンコーダだけでは不十分かを整理します。

通常のAEは x -> z -> x_hat の再構成は上手でも、潜在空間 z がバラバラだと「どこからサンプルを引けば自然なデータが出るか」が分かりません。VAEはここに事前分布 p(z) を入れて、潜在空間を生成に使える形へ整えます。

1. VAEの最小数式

VAEでは次の3つを同時に扱います。

• エンコーダ（近似事後分布）: $q_phi(z|x)$
• デコーダ（生成分布）: $p_theta(x|z)$
• 潜在の事前分布: p(z)=N(0,I)

学習目標はELBO最大化で、次の2項に分かれます。

• 再構成項: E_{q_phi(z|x)}[log p_theta(x|z)]
• 正則化項: KL(q_phi(z|x) || p(z))

要するに、再構成を良くしつつ、潜在分布を標準正規に寄せる、という設計です。

def gaussian_log_prob(x, mu, var):
    var = max(var, 1e-8)
    return -0.5 * ((x - mu) ** 2 / var + math.log(2 * math.pi * var))


def kl_normal_to_standard(mu, logvar):
    # KL( N(mu, sigma^2) || N(0,1) )
    return 0.5 * (mu * mu + math.exp(logvar) - 1.0 - logvar)


x = 1.2
mu_dec = 1.0
var_dec = 0.3
mu_enc = 0.4
logvar_enc = -0.2

recon = gaussian_log_prob(x, mu_dec, var_dec)
kl = kl_normal_to_standard(mu_enc, logvar_enc)
elbo = recon - kl

print('reconstruction term =', round(recon, 4))
print('KL term             =', round(kl, 4))
print('ELBO                =', round(elbo, 4))

普通のオートエンコーダは「再構成が良いこと」だけを押しますが、VAE はそこに KL を足して「生成しやすい潜在空間であること」も同時に求めます。だから reconstruction term が良くても KL を無視しすぎると、サンプルしやすい潜在空間になりません。

2. 再パラメータ化トリック

z ~ q_phi(z|x) をそのままサンプリングすると、勾配が流しにくくなります。VAEでは

z = mu + sigma * eps, eps ~ N(0,1)

と書き換え、乱数を eps 側へ分離します。これにより mu, sigma（=ネットワーク出力）へ勾配を通せます。

def reparameterize(mu, logvar, eps):
    sigma = math.exp(0.5 * logvar)
    return mu + sigma * eps


mu = 0.7
logvar = -0.8
eps_samples = [random.gauss(0.0, 1.0) for _ in range(5)]
zs = [reparameterize(mu, logvar, e) for e in eps_samples]

print('eps samples =', [round(e, 3) for e in eps_samples])
print('z samples   =', [round(z, 3) for z in zs])

3. 1次元トイデータでVAEを学習する

ここではPyTorchを使わず、あえて超小型モデルを手書きで最適化します。目的は、VAEの損失構造（再構成 vs KL）を体感することです。

• エンコーダ: mu(x)=a*x+b, logvar(x)=c*x+d
• デコーダ: $x_hat(z)=m*z+n$

学習は有限差分で行い、目的関数（再構成項 - beta * KL）を直接最大化します（教育用）。
注意: このトイモデルは線形デコーダなので表現力を意図的に制限しています。2峰性データを完全に表現するのが目的ではなく、VAEの学習原理を確認するための設定です。

def make_toy_data(n1=60, n2=60):
    left = [random.gauss(-2.0, 0.45) for _ in range(n1)]
    right = [random.gauss(2.2, 0.55) for _ in range(n2)]
    return left + right


data = make_toy_data()
print('dataset size =', len(data))
print('mean(data)   =', round(mean(data), 3))
print('min/max      =', round(min(data), 3), round(max(data), 3))

再パラメータ化で勾配を通す

ここでは eps と z の関係を観察し、ノイズを含んだ潜在サンプルでもパラメータへ勾配が流れる書き方を確認します。

def encode(x, params):
    a, b, c, d, _, _ = params
    mu = a * x + b
    logvar = c * x + d
    logvar = max(min(logvar, 4.0), -6.0)  # 数値安定
    return mu, logvar


def decode(z, params):
    _, _, _, _, m, n = params
    return m * z + n


def elbo_dataset(data, params, beta=1.0, recon_var=0.25, eps_list=None):
    # recon_var: p_theta(x|z)=N(x_hat, recon_var) の固定分散
    # 小さいほど再構成誤差に厳しくなる
    if eps_list is None:
        # 通常のモンテカルロ近似では毎回epsをサンプルする
        eps_list = [random.gauss(0.0, 1.0) for _ in range(len(data))]

    recon_terms = []
    kl_terms = []
    objectives = []

    for x, eps in zip(data, eps_list):
        mu, logvar = encode(x, params)
        z = reparameterize(mu, logvar, eps)
        x_hat = decode(z, params)

        recon = gaussian_log_prob(x, x_hat, recon_var)
        kl = kl_normal_to_standard(mu, logvar)
        obj = recon - beta * kl

        recon_terms.append(recon)
        kl_terms.append(kl)
        objectives.append(obj)

    return mean(objectives), mean(recon_terms), mean(kl_terms)


params0 = [0.2, 0.0, -0.1, -0.5, 0.7, 0.0]
fixed_eps = [random.gauss(0.0, 1.0) for _ in range(len(data))]

obj0, recon0, kl0 = elbo_dataset(data, params0, beta=1.0, eps_list=fixed_eps)
print('initial objective =', round(obj0, 4))
print('initial recon=', round(recon0, 4), 'initial KL=', round(kl0, 4))

1 次元の toy VAE を実際に回す

小さなモデルで十分なので、目的関数がどう下がり、2 項のバランスがどう動くかを見ます。

def finite_diff_grad(data, params, beta, recon_var, eps_list, h=1e-3):
    grads = [0.0] * len(params)
    for i in range(len(params)):
        p_plus = params[:]
        p_minus = params[:]
        p_plus[i] += h
        p_minus[i] -= h

        f_plus, _, _ = elbo_dataset(data, p_plus, beta=beta, recon_var=recon_var, eps_list=eps_list)
        f_minus, _, _ = elbo_dataset(data, p_minus, beta=beta, recon_var=recon_var, eps_list=eps_list)
        grads[i] = (f_plus - f_minus) / (2 * h)
    return grads


def train_toy_vae(data, beta=1.0, steps=120, lr=0.03, recon_var=0.25):
    params = [0.2, 0.0, -0.1, -0.5, 0.7, 0.0]  # a,b,c,d,m,n

    # 有限差分のノイズを減らすため、共通乱数(common random numbers)でepsを固定
    eps_list = [random.gauss(0.0, 1.0) for _ in range(len(data))]

    trace = []
    for t in range(steps):
        grads = finite_diff_grad(data, params, beta, recon_var, eps_list)
        for i in range(len(params)):
            params[i] += lr * grads[i]  # objective最大化

        obj, recon, kl = elbo_dataset(data, params, beta=beta, recon_var=recon_var, eps_list=eps_list)
        trace.append((obj, recon, kl))

        if t % 20 == 0 or t == steps - 1:
            print(f'step={t:03d} objective={obj:.4f} recon={recon:.4f} KL={kl:.4f}')

    return params, trace


params_beta1, trace_beta1 = train_toy_vae(data, beta=1.0, steps=140, lr=0.025, recon_var=0.30)
print('trained params (beta=1) =', [round(v, 4) for v in params_beta1])

上のログで、目的関数（recon - beta * KL）が改善しつつ、再構成項とKL項のバランスが動くことが確認できます。これがVAE学習のダイナミクスです。

4. beta-VAEでトレードオフを見る

beta を大きくするとKLを強く罰するため、潜在分布は事前分布に近づきます。その代わり再構成が悪化しやすくなります。

注意: beta を変えると目的関数そのものが変わるので、objective 値を横比較して優劣を決めるのは不適切です。比較は主に再構成項とKL項で行います。

params_beta4, trace_beta4 = train_toy_vae(data, beta=4.0, steps=140, lr=0.02, recon_var=0.30)

final1 = trace_beta1[-1]
final4 = trace_beta4[-1]

print('beta=1 final: objective={:.4f}, recon={:.4f}, KL={:.4f}'.format(*final1))
print('beta=4 final: objective={:.4f}, recon={:.4f}, KL={:.4f}'.format(*final4))

beta を強めたときの代償を見る

潜在空間を整える圧力を強くすると、再構成はどう崩れるのか。beta-VAE の典型的なトレードオフをここで確認します。

def summarize_latent_stats(data, params):
    mus = []
    vars_ = []
    for x in data:
        mu, logvar = encode(x, params)
        mus.append(mu)
        vars_.append(math.exp(logvar))

    emu = mean(mus)
    emu2 = mean(m * m for m in mus)
    evar = mean(vars_)
    return emu, emu2, evar


mu1, mu21, var1 = summarize_latent_stats(data, params_beta1)
mu4, mu24, var4 = summarize_latent_stats(data, params_beta4)

print('latent summary beta=1: E[mu]=', round(mu1, 4), 'E[mu^2]=', round(mu21, 4), 'E[var]=', round(var1, 4))
print('latent summary beta=4: E[mu]=', round(mu4, 4), 'E[mu^2]=', round(mu24, 4), 'E[var]=', round(var4, 4))

beta を上げると、q(z|x) が標準正規に近づきやすくなり、潜在空間は整いますが、入力ごとの情報を削りすぎると再構成が崩れます。この綱引きを調整するのがbeta-VAEの本質です。

5. 生成と補間

学習後は z ~ N(0,1) を引いてデコーダに通すと新サンプルを生成できます。また、2つの入力の潜在平均を線形補間してデコードすると、連続的な変化を観察できます。

def generate_from_prior(params, n=8):
    out = []
    for _ in range(n):
        z = random.gauss(0.0, 1.0)
        x_hat = decode(z, params)
        out.append((z, x_hat))
    return out


gen = generate_from_prior(params_beta1, n=10)
print('samples from prior z~N(0,1):')
for z, xh in gen:
    print('z=', round(z, 3), '-> x_hat=', round(xh, 3))

潜在からの生成と補間を試す

学習後は、再構成だけでなく、潜在空間を移動したときに出力がどう変わるかを見ることで、空間が滑らかに整理されているかを確かめます。

def latent_mean(x, params):
    mu, _ = encode(x, params)
    return mu

x_a = data[5]
x_b = data[-5]
mu_a = latent_mean(x_a, params_beta1)
mu_b = latent_mean(x_b, params_beta1)

print('x_a=', round(x_a, 3), 'mu_a=', round(mu_a, 3))
print('x_b=', round(x_b, 3), 'mu_b=', round(mu_b, 3))
print('interpolation in latent mean:')

for t in [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]:
    z = (1 - t) * mu_a + t * mu_b
    x_hat = decode(z, params_beta1)
    print('t=', round(t, 1), 'z=', round(z, 3), 'x_hat=', round(x_hat, 3))

6. よくある失敗: Posterior Collapse

KLを強くしすぎたり、デコーダが強すぎたりすると、q(z|x) がほぼ N(0,1) になって入力情報を使わなくなることがあります。これをposterior collapseと呼びます。

対策としては、KL warm-up、free bits、デコーダ容量調整、学習率調整などが使われます。

注意: 実務での厳密判定は、KL値だけでは不十分です。I(x;z)（相互情報量）や、潜在をシャッフルしたときの再構成劣化など、情報保持の指標を併用します。ここで示す判定は教育用ヒューリスティックで、collapseの必要条件でも十分条件でもありません。

# collapse の簡易判定（toy）: KLが極端に小さく、再構成が悪化していないか
# 厳密判定ではなく、兆候を見るための簡易チェック（必要条件・十分条件ではない）
kl_beta1 = trace_beta1[-1][2]
kl_beta4 = trace_beta4[-1][2]
recon_beta1 = trace_beta1[-1][1]
recon_beta4 = trace_beta4[-1][1]

print('final KL beta=1 =', round(kl_beta1, 6), '| recon=', round(recon_beta1, 6))
print('final KL beta=4 =', round(kl_beta4, 6), '| recon=', round(recon_beta4, 6))

if kl_beta4 < 0.01 and recon_beta4 < recon_beta1 - 0.3:
    print('beta=4 run: collapse risk is high (toy criterion).')
else:
    print('beta=4 run: collapse risk is not extreme in this toy run.')

print('For strict diagnosis, add mutual-information or ablation-based checks.')

VAEを一言で言うと、「再構成したい」という要求と「生成しやすい潜在空間にしたい」という要求を、ELBOで同時に満たすモデルです。

この後のGANや拡散モデルに進むときも、何を近似し、どの損失でその近似を押すのか、という観点で比較すると整理しやすくなります。