MIT 6.7960 L15：變分自編碼器（VAE）—— ELBO、再參數化技巧與潛變數表示

🌏 English version

教材版本：基於 MIT 6.7960 Fall 2024 OCW（對應 OCW Lec 15）。影片、投影片、作業全公開於 MIT OCW。本講由 Phillip Isola 授課。

L14 留下的問題：p(x) 算不出來

L14 講了生成模型三大家族（likelihood、自回歸、隱變數）。其中隱變數模型 p(x) = ∫ p(x|z) p(z) dz 看起來漂亮，但對 z 的高維積分是 intractable 的，沒辦法直接最大化 log p(x)。這講就講怎麼繞過去。

ELBO：把 log p(x) 換成可優化的下界

引進一個近似後驗 q(z|x)（用神經網路學），做 Jensen 不等式：

log p(x) = log ∫ p(x|z) p(z) dz
         = log ∫ q(z|x) [p(x|z) p(z) / q(z|x)] dz
         ≥ E_{q(z|x)}[log p(x|z)] − KL(q(z|x) ‖ p(z))
        ≡ ELBO

右邊就是 Evidence Lower Bound (ELBO)：第一項是「用 z 重構 x 的期望對數似然」，第二項是「近似後驗跟先驗的 KL」。最大化 ELBO = 同時（1）讓 q 估得準、（2）讓 p(x|z) 能從 z 重建 x。

把這兩項寫成 loss：

loss = E_q[−log p(x|z)] + KL(q(z|x) ‖ p(z))
     = recon_loss + kl_loss

p(z) 常用標準常態 N(0, I)，q(z|x) 設成 diagonal Gaussian N(μ(x), diag σ²(x))，KL 就解析有閉式。

再參數化技巧：讓採樣可反向傳播

ELBO 第一項 E_q[log p(x|z)] 期望的 z ~ q(z|x) = N(μ, σ²)，而採樣是隨機節點，梯度穿不過去。再參數化技巧把它改寫成

z = μ(x) + σ(x) ⊙ ε,    ε ~ N(0, I)

隨機性從 z 移到了 ε，z 變成 μ, σ 的確定性函數 → 梯度可以從 decoder 透過 z 流回 encoder。這個「線從 N(0,1) 流入」的小技巧是 VAE 能用 SGD 訓練的關鍵。

一個最小的 VAE 訓練骨架

import torch, torch.nn as nn

class VAE(nn.Module):
    def __init__(self, z_dim=16):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(784, 256), nn.ReLU())
        self.mu  = nn.Linear(256, z_dim)
        self.lv  = nn.Linear(256, z_dim)  # log-variance
        self.dec = nn.Sequential(nn.Linear(z_dim, 256), nn.ReLU(),
                                 nn.Linear(256, 784), nn.Sigmoid())
    def forward(self, x):
        h = self.enc(x.view(-1, 784))
        mu, lv = self.mu(h), self.lv(h)
        z = mu + (0.5*lv).exp() * torch.randn_like(mu)  # reparameterize
        return self.dec(z), mu, lv

def loss(x, xh, mu, lv):
    recon = ((x.view(-1,784) - xh)**2).sum(-1).mean()
    kl = (-0.5 * (1 + lv - mu**2 - lv.exp())).sum(-1).mean()  # 對 N(0,I) 解析
    return recon + kl, recon, kl

訓練時 loss.backward() 會一次更新 encoder 跟 decoder；採 x̂ 就是 decode 一個隨機 z ~ N(0,I) 的結果。

實務上會踩到的坑

• Posterior collapse：decoder 太強，KL 退化到 0、忽略 z。常見解法是「warm-up KL 權重」或「free bits」。
• 圖片偏糊：VAE 的 likelihood 在像素層級用 MSE/Gaussian 假設，常態分數不如 GAN。但 latent 結構清楚、利於插值/編輯，這是它的取捨。
• β-VAE：把 KL 加權 β > 1，壓迫 latent → 更解耦但重建變差。
• VQ-VAE（van den Oord et al.）：潛變數量化為離散 codebook，避免後驗 collapse，後來 DALL·E / Stable Diffusion 的 latent diffusion 都靠它。
• Hierarchical / NVAE：多層 latent，每層都做 ELBO，收斂更穩、解析度更高。

VAE 的位置：和 L14 三大家的對照

VAE 是「likelihood 路線最便宜、latent 結構最好」的折衷——後面的 latent diffusion / Stable Diffusion 就是站在這個肩膀上。

參考資料

• MIT 6.7960 OCW（Fall 2024）：課程首頁
• Kingma & Welling, Auto-Encoding Variational Bayes：arXiv:1312.6114
• Doersch, Tutorial on Variational Autoencoders：arXiv:1606.05908
• Rezende et al., Stochastic Backpropagation and Approximate Inference in Deep Generative Models：arXiv:1401.4082