VAEと拡散モデル

世界モデルでは、状態を圧縮して予測しつつ、必要なら高品質に観測を復元する必要があります。VAEと拡散はこの目的に対して異なる強みを持ちます。

圧縮してから戻すか、ノイズを削りながら戻すか

世界モデルの観測復元には、速く扱いやすい潜在空間が欲しい場面と、多少重くても見た目の品質を上げたい場面があります。VAE 系と拡散系は、その二つの要求に対して違う答えを出します。

ここでは 2 次元データを 1 次元に圧縮し、VAE 風の再構成と拡散風のデノイズを並べます。厳密な本実装ではありませんが、そのぶん『どこで速度を取り、どこで復元力を取るか』という設計の勘所が見えやすくなります。

世界モデルの中で両者はどう住み分けるか

計画用の潜在状態には軽さと一貫性が欲しく、観測を見せる出口には細部を戻す力が欲しくなります。VAE 系は前者に、拡散系は後者に寄りやすいことが多く、その差を小さな例で確認します。

見るべきポイント

z がどう圧縮されるか、PCA ベースの再構成で何が落ちるか、ノイズ除去反復でどこまで戻るかを比べながら読みます。

この notebook の数値は『どちらが絶対に優れているか』を決めるためではなく、同じデータに対して復元の作法がどう違うかを掴むためのものです。

近似実装で何を見ているか

VAE 側は PCA を encoder/decoder の近似に使い、拡散側は簡易スコア反復でデノイズを模します。細部は省いていますが、圧縮主体か高品質復元主体かという設計差は十分見えます。

読み方の軸

誤差が小さいかだけでなく、潜在空間を扱う軽さと、観測を戻す粘り強さのどちらを優先したい場面かを意識して見てください。

潜在を 1 次元へ落としてみる

まずは入力データの主な変動方向を 1 本に集めます。ここで失われる情報量が、あとで復元品質の差として現れます。

import numpy as np
np.random.seed(3)

n = 300
theta = np.random.rand(n) * 2 * np.pi
x = np.stack([np.cos(theta), np.sin(theta)], axis=1) + 0.08 * np.random.randn(n, 2)

再構成とデノイズを並べてみる

次に、圧縮してすぐ戻す流れと、ノイズを少しずつ削って戻す流れを比較します。世界モデルの出口設計を考えるときの分岐点です。

# VAE風: PCAを encoder/decoder の近似として利用
mu = x.mean(axis=0, keepdims=True)
xc = x - mu
u, s, vt = np.linalg.svd(xc, full_matrices=False)
pc1 = vt[0:1].T  # 2x1
z = xc @ pc1      # encoder
x_rec = z @ pc1.T + mu  # decoder
vae_mse = np.mean((x_rec - x) ** 2)

# Diffusion風: 潜在にノイズを加え、簡易スコア(ガウス事前)で反復デノイズ
z_noisy = z + 0.5 * np.random.randn(*z.shape)
z_den = z_noisy.copy()
prior_var = np.var(z_noisy) + 1e-8
for _ in range(30):
    score = -z_den / prior_var  # log N(0,var) の勾配
    z_den = z_den + 0.04 * score
x_den = z_den @ pc1.T + mu
diff_mse = np.mean((x_den - x) ** 2)

print('VAE-like reconstruction MSE       =', round(vae_mse, 6))
print('Diffusion-like latent denoise MSE =', round(diff_mse, 6))

この notebook の要点は、VAE と拡散を対立で覚えることではありません。計画に使う軽い潜在と、見せるための高品質復元を同じモデルにどう分担させるか、その設計感覚を掴むことにあります。