Bài 43: Variational Autoencoder (VAE) intuition — sinh dữ liệu mới

Sau bài học, bạn sẽ:

• Giải thích được vì sao autoencoder standard không phải generative model — latent space không "đẹp" để sample.
• Viết ra được probabilistic formulation của VAE: prior \( p(z) \), posterior \( q(z \mid x) \), likelihood \( p(x \mid z) \).
• Hiểu vì sao cần reparameterization trick và viết được \( z = \mu + \sigma \cdot \epsilon \).
• Trình bày được loss VAE = reconstruction \( + \) KL divergence với prior \( \mathcal{N}(0, I) \).
• Áp dụng được công thức closed-form \( D_{KL}(\mathcal{N}(\mu, \sigma^2) \,\|\, \mathcal{N}(0, 1)) \).
• Implement VAE bằng PyTorch trên MNIST, sample được ảnh mới từ \( z \sim \mathcal{N}(0, I) \).
• Phân biệt VAE vs GAN: ưu nhược, vị trí của diffusion model.
• Biết β-VAE (disentanglement) và CVAE (conditional generation).

Bài này nối Bài 42 — Denoising Autoencoder và mở đường cho Bài 44 — GAN intuition.

Autoencoder (B41) học hàm \( x \mapsto z \mapsto \hat{x} \) sao cho \( \hat{x} \approx x \). Latent \( z \in \mathbb{R}^{d_z} \) chứa biểu diễn nén của \( x \). Nhưng AE chưa phải generative model:

• Latent space không có cấu trúc rõ: encoder chỉ học sao cho mỗi training sample \( x_i \) map về một \( z_i \) cụ thể, không ràng buộc \( z \) phân bố ra sao trong không gian.
• Vùng không phủ: giữa các \( z_i \) là "khoảng trống" — sample một \( z \) ngẫu nhiên thường rơi vào khoảng trống đó, decoder map ra ảnh vô nghĩa.
• Không có phân phối để sample từ: muốn sinh ảnh mới, phải biết \( z \) "hợp lệ" trông như thế nào. AE không cho biết điều đó.

Ví dụ: train AE trên MNIST với \( d_z = 2 \), plot \( z_i \) của test set lên 2D — bạn sẽ thấy 10 cụm rời rạc cho 10 chữ số, các cụm cách xa nhau và xen lẫn vùng trống. Sample \( z = (0, 0) \) có thể không thuộc cụm nào → decoder ra ảnh nhiễu.

Để biến AE thành generative model thực sự, cần ép latent space tuân theo một phân phối biết trước — sample từ phân phối đó là đủ.

Kingma & Welling (2013) — "Auto-Encoding Variational Bayes" — đề xuất:

1. Đặt một prior trên latent: \( p(z) = \mathcal{N}(0, I) \). Đơn vị, đa biến, không tương quan giữa các chiều.
2. Encoder không trả về một vector \( z \) mà trả về phân phối \( q(z \mid x) = \mathcal{N}(\mu(x), \sigma^2(x)) \).
3. Loss có thêm KL divergence \( D_{KL}(q(z \mid x) \,\|\, p(z)) \) ép \( q(z \mid x) \) gần \( \mathcal{N}(0, I) \).
4. Khi train xong, sample \( z \sim \mathcal{N}(0, I) \) → đưa qua decoder → ảnh mới.

Autoencoder:
   x ─► encoder ─► z (1 vector) ─► decoder ─► x̂

VAE:
   x ─► encoder ─► (μ, σ) ─► sample z ~ N(μ, σ²) ─► decoder ─► x̂
                              │
                              └── ràng buộc q(z|x) ≈ N(0, I)

Generation (sau khi train):
   z ~ N(0, I) ─► decoder ─► sample mới

Trực giác: thay vì map mỗi \( x \) về một điểm, VAE map về một vùng Gaussian. Vùng của các \( x \) khác nhau "đè" lên nhau quanh gốc 0 → latent space dày, phủ kín, không còn khoảng trống.

VAE là latent variable model. Mô hình giả định data được sinh ra qua hai bước:

\[ z \sim p(z) = \mathcal{N}(0, I), \quad x \sim p_\theta(x \mid z) \]

3 phân phối cần để ý:

• Prior \( p(z) = \mathcal{N}(0, I) \) — chọn trước, không học. Đây là phân phối ta muốn ép latent tuân theo.
• Likelihood (decoder) \( p_\theta(x \mid z) \) — mạng neural có tham số \( \theta \), nhận \( z \) trả về phân phối trên \( x \). Cho ảnh, thường giả định Bernoulli (pixel nhị phân, dùng BCE) hoặc Gaussian (pixel liên tục, dùng MSE).
• Posterior (encoder) \( q_\phi(z \mid x) \) — xấp xỉ \( p(z \mid x) \) bằng mạng neural có tham số \( \phi \). Thường chọn Gaussian \( \mathcal{N}(\mu_\phi(x), \sigma_\phi^2(x)) \).

Muốn fit data, ta muốn maximize log-likelihood \( \log p_\theta(x) = \log \int p_\theta(x \mid z) p(z) \, dz \). Tích phân này intractable (không tính được kín). VAE giải bằng variational inference: dùng \( q_\phi(z \mid x) \) xấp xỉ \( p(z \mid x) \), tối ưu một lower bound (ELBO, mục 7) thay cho log-likelihood trực tiếp. Đó là lý do có chữ "Variational" trong tên.

Encoder trong AE: \( z = f_\phi(x) \) — output 1 vector \( z \in \mathbb{R}^{d_z} \).

Encoder trong VAE: output hai vector cùng kích thước \( d_z \):

\[ \mu(x) \in \mathbb{R}^{d_z}, \quad \log \sigma^2(x) \in \mathbb{R}^{d_z} \]

Hai vector này định nghĩa một Gaussian đa biến với covariance chéo (mỗi chiều độc lập):

\[ q_\phi(z \mid x) = \mathcal{N}\!\big(z;\ \mu(x),\ \mathrm{diag}(\sigma^2(x))\big) \]

Trong code, encoder thường có một MLP/CNN shared, rồi 2 linear head riêng:

self.encoder = nn.Sequential(
    nn.Linear(input_dim, 256), nn.ReLU(),
)
self.fc_mu     = nn.Linear(256, latent_dim)
self.fc_logvar = nn.Linear(256, latent_dim)

Lưu ý chính: cùng một \( x \) chạy qua encoder ra cùng \( (\mu, \log \sigma^2) \) (deterministic). Tính ngẫu nhiên chỉ xuất hiện khi sample \( z \) từ phân phối — bước tiếp theo.

Bài toán: cần sample \( z \sim \mathcal{N}(\mu, \sigma^2) \) rồi đưa qua decoder, sau đó backprop gradient về \( \mu, \sigma \). Nhưng bước "sample" thông thường không có gradient — node ngẫu nhiên cắt đứt đồ thị tính toán.

Reparameterization trick: viết lại bước sampling dưới dạng deterministic phụ thuộc \( (\mu, \sigma) \) cộng thêm một noise độc lập:

\[ z = \mu + \sigma \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \]

Ý tưởng:

• \( \epsilon \) lấy mẫu từ \( \mathcal{N}(0, I) \) — nguồn ngẫu nhiên duy nhất, không phụ thuộc \( \phi \), nên không cần gradient.
• \( z \) là phép biến đổi affine deterministic của \( \mu, \sigma, \epsilon \). Gradient \( \partial L / \partial z \) chảy được về \( \mu \) (qua phép cộng) và \( \sigma \) (qua phép nhân).
• Phân phối của \( z \) đúng là \( \mathcal{N}(\mu, \sigma^2) \) — phân phối ta muốn sample.

Naive sampling (KHÔNG backprop được):
  z ~ N(μ, σ²)     ← node ngẫu nhiên, không gradient
  decoder(z) ─► loss
  loss.backward()  ← gradient dừng tại z

Reparameterization (backprop được):
  ε ~ N(0, I)              ← noise độc lập, không cần gradient
  z = μ + σ * ε            ← deterministic theo μ, σ
  decoder(z) ─► loss
  loss.backward()  ← gradient chảy qua z về μ, σ → encoder

Code:

def reparameterize(self, mu, logvar):
    std = torch.exp(0.5 * logvar)       # σ = exp(0.5 * logvar)
    eps = torch.randn_like(std)         # ε ~ N(0, I)
    return mu + eps * std               # z = μ + σ * ε

Đây là đóng góp kỹ thuật quan trọng của paper Kingma & Welling. Cùng năm, Rezende, Mohamed, Wierstra (2014) đề xuất kỹ thuật tương tự ("stochastic backpropagation") cho deep latent Gaussian model.

Log-likelihood của data \( \log p_\theta(x) \) intractable. VAE tối ưu một lower bound — Evidence Lower BOund (ELBO):

\[ \log p_\theta(x) \ \ge\ \mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log p_\theta(x \mid z) \right] - D_{KL}\!\big(q_\phi(z \mid x) \,\|\, p(z)\big) \]

Có thể chứng minh ELBO chính xác bằng cách triển khai:

\[ \log p_\theta(x) = \mathcal{L}(\theta, \phi; x) + D_{KL}\!\big(q_\phi(z \mid x) \,\|\, p_\theta(z \mid x)\big) \]

Vì KL \( \ge 0 \), ELBO là lower bound chặt nhất khi \( q_\phi(z \mid x) = p_\theta(z \mid x) \). Maximize ELBO đồng thời:

• Đẩy \( \log p_\theta(x) \) lên (tăng likelihood).
• Kéo \( q_\phi(z \mid x) \) về gần posterior thật \( p_\theta(z \mid x) \).

Trong thực hành, ta minimize \( -\mathcal{L} \). Viết ngược dấu:

\[ L_{\text{VAE}}(x) = \underbrace{-\mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log p_\theta(x \mid z) \right]}_{\text{reconstruction loss}} \;+\; \underbrace{D_{KL}\!\big(q_\phi(z \mid x) \,\|\, p(z)\big)}_{\text{KL regularizer}} \]

2 thành phần, cả 2 cùng được minimize. Mục 8 và 9 chi tiết từng phần.

Term reconstruction là expectation log-likelihood của \( x \) cho \( z \):

\[ -\mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log p_\theta(x \mid z) \right] \]

Trong thực hành ước lượng bằng Monte Carlo với 1 mẫu \( z \) lấy qua reparameterization (chính là output của bước sampling):

\[ \approx -\log p_\theta(x \mid z), \quad z = \mu(x) + \sigma(x) \cdot \epsilon \]

Tùy giả định trên \( p_\theta(x \mid z) \):

• Bernoulli (pixel \( \in [0, 1] \), MNIST binary): \( -\log p_\theta(x \mid z) = \mathrm{BCE}(\hat{x}, x) \). Decoder cuối có sigmoid, output là xác suất pixel = 1.
• Gaussian với variance cố định: \( -\log p_\theta(x \mid z) \propto \| \hat{x} - x \|^2 = \mathrm{MSE} \). Decoder không có activation cuối (hoặc tanh), output là giá trị pixel.

Quy ước chung trong reference implementation: sum (không phải mean) trên các pixel của một sample, để cân với scale của KL (KL cũng là sum trên các chiều latent). Nếu dùng mean cho recon, KL sẽ áp đảo và model gần như không reconstruct.

recon_loss = F.binary_cross_entropy(x_recon, x, reduction="sum")
# hoặc với Gaussian assumption:
# recon_loss = F.mse_loss(x_recon, x, reduction="sum")

Với \( q_\phi(z \mid x) = \mathcal{N}(\mu, \mathrm{diag}(\sigma^2)) \) và \( p(z) = \mathcal{N}(0, I) \), KL có công thức kín — không cần Monte Carlo:

\[ D_{KL}\!\big(q_\phi(z \mid x) \,\|\, p(z)\big) = -\frac{1}{2} \sum_{j=1}^{d_z} \Big( 1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2 \Big) \]

Đặt \( \ell_j = \log \sigma_j^2 \) (logvar), viết theo logvar và tránh \( \sigma^2 \) âm:

\[ D_{KL} = -\frac{1}{2} \sum_{j=1}^{d_z} \Big( 1 + \ell_j - \mu_j^2 - e^{\ell_j} \Big) \]

Quan sát:

• KL đạt 0 khi \( \mu_j = 0 \) và \( \sigma_j^2 = 1 \) với mọi \( j \) — đúng prior \( \mathcal{N}(0, I) \).
• KL phạt encoder nếu \( \mu_j \ne 0 \) (\( -\mu_j^2 \) lớn theo trị tuyệt đối \( \mu_j \)) hoặc \( \sigma_j \) lệch xa 1.
• Tổng trên \( d_z \) chiều — KL scale theo latent dim.

kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

Một dòng. Đây là một trong những lý do người ta thường chọn prior \( \mathcal{N}(0, I) \) và posterior Gaussian — toàn bộ KL ra closed-form, không cần sampling thứ hai.

Tổng loss \( L = L_{\text{recon}} + L_{KL} \). Hai term kéo theo hai hướng đối lập:

Điểm cân bằng phụ thuộc dataset, latent dim, capacity của encoder/decoder. Vài quan sát thực hành:

• MNIST \( d_z = 20 \): cân bằng tự nhiên, không cần tuning nhiều.
• Ảnh phức tạp hơn (CIFAR-10, CelebA): VAE vanilla thường blur. Cần convolutional encoder/decoder, latent lớn hơn (\( d_z = 128 \)–\( 512 \)).
• Khi posterior collapse, một số dim của \( z \) "tắt" — \( \mu_j \approx 0, \sigma_j \approx 1 \) bất kể \( x \). Có thể dùng KL annealing (tăng dần KL weight từ 0 lên 1 qua epoch) hoặc free bits (chỉ phạt KL khi \( > \delta \) per dim).

β-VAE (mục 16) thêm hệ số \( \beta \) cho KL term để điều khiển trade-off này một cách tường minh.

Encoder dự đoán \( \log \sigma^2 \) (logvar) thay vì \( \sigma \) hay \( \sigma^2 \) trực tiếp. Lý do thực dụng:

• Range không ràng buộc: \( \log \sigma^2 \in \mathbb{R} \), output linear không cần activation. Còn \( \sigma^2 > 0 \), \( \sigma > 0 \) cần softplus/exp/abs để đảm bảo dương — gradient và stability kém hơn.
• \( \sigma^2 = e^{\text{logvar}} \) luôn dương: không bao giờ phải lo NaN do log của số âm.
• Numerical stable: với \( \sigma \) rất nhỏ (gần 0), \( \log \sigma^2 \) lớn theo phía âm — biểu diễn được; còn \( \sigma^2 \) trực tiếp có thể underflow xuống 0 và 1/\( \sigma^2 \) overflow.
• Khớp công thức: KL closed-form (mục 9) viết theo \( \log \sigma^2 \) trực tiếp. Reparameterization \( \sigma = \exp(0.5 \cdot \text{logvar}) \).

Quy ước này có ở hầu hết implementation chính thống (Kingma reference, PyTorch examples). Khi đọc code VAE, gặp logvar là chuyện bình thường.

VAE đơn giản cho MNIST (input 28×28 = 784, latent 20):

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=256, latent_dim=20):
        super().__init__()
        # Encoder shared trunk
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), nn.ReLU(),
        )
        # 2 heads: μ và log σ²
        self.fc_mu     = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, input_dim), nn.Sigmoid(),
        )

    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decoder(z)
        return x_recon, mu, logvar


def vae_loss(x_recon, x, mu, logvar):
    # Recon: BCE sum trên pixel
    recon_loss = F.binary_cross_entropy(x_recon, x, reduction="sum")
    # KL closed-form
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_loss

Training loop chuẩn:

model = VAE().to(device)
optim = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(50):
    for x, _ in train_loader:               # x: (B, 1, 28, 28), label bỏ qua
        x = x.view(-1, 784).to(device)
        x_recon, mu, logvar = model(x)
        loss = vae_loss(x_recon, x, mu, logvar)
        optim.zero_grad(); loss.backward(); optim.step()

Điểm cần để ý khi đọc/viết code:

• Input MNIST scale về \( [0, 1] \) (dùng ToTensor()), decoder cuối là Sigmoid() — phải khớp BCE.
• reduction="sum" trên cả batch — loss có scale theo batch_size; chia trung bình sau nếu muốn theo dõi per-sample.
• Latent 20 đủ cho MNIST. Quá nhỏ (1–2) thì reconstruction yếu; quá lớn (200+) thì KL term áp đảo, hoặc posterior collapse từng dim.

Sau khi train, chỉ cần decoder để sinh sample mới. Bỏ encoder, sample \( z \) trực tiếp từ prior:

model.eval()
with torch.no_grad():
    z = torch.randn(16, 20, device=device)      # 16 sample, latent_dim=20
    samples = model.decoder(z)                  # (16, 784)
    samples = samples.view(-1, 1, 28, 28)
    # plot samples bằng matplotlib hoặc save bằng torchvision.utils.save_image

Đây là khác biệt then chốt so với AE thường: ở VAE, sample \( z \sim \mathcal{N}(0, I) \) là hợp lệ vì training đã ép \( q(z \mid x) \) gần \( \mathcal{N}(0, I) \) — vùng có dữ liệu thật. Decoder map về một ảnh trong manifold đã học.

Chất lượng MNIST: VAE vanilla sinh ra số khá nhận diện được, hơi mờ ở vài chi tiết (do reconstruction term dùng BCE/MSE per-pixel, không phạt cấu trúc cao cấp). Để cải thiện sharpness, người ta thường dùng GAN hoặc diffusion model (mục 15, 19).

Một cách trực quan để kiểm tra latent space VAE có "smooth" hay không là nội suy giữa hai sample:

1. Chọn 2 ảnh \( x_1, x_2 \) (vd số "3" và số "8" từ MNIST).
2. Lấy \( z_1 = \mu(x_1), z_2 = \mu(x_2) \) qua encoder.
3. Nội suy tuyến tính: \( z_t = (1 - t) z_1 + t z_2 \) với \( t \in [0, 1] \).
4. Decode mỗi \( z_t \) → chuỗi ảnh chuyển dần từ \( x_1 \) sang \( x_2 \).

model.eval()
with torch.no_grad():
    mu1, _ = model.encode(x1.view(1, -1))
    mu2, _ = model.encode(x2.view(1, -1))
    ts = torch.linspace(0, 1, 10).unsqueeze(1)        # (10, 1)
    z_interp = (1 - ts) * mu1 + ts * mu2              # (10, latent_dim)
    imgs = model.decoder(z_interp).view(-1, 1, 28, 28)

Với VAE, chuỗi ảnh chuyển dần mượt — số "3" cong dần thành "8". Với AE thường, chuỗi gãy: giữa hai ảnh có thể có một số bước ra ảnh nhiễu (vì \( z_t \) rơi vào vùng không phủ).

Nội suy không phải luôn dùng linear. Trên Gaussian space, một số paper khuyến nghị slerp (spherical linear interpolation) — đi theo cung tròn trên hyper-sphere thay vì đường thẳng — giữ \( \|z_t\| \) ổn định gần \( \sqrt{d_z} \), khớp hơn với norm trung bình của sample từ \( \mathcal{N}(0, I) \).

VAE (B43) và GAN (B44) là hai paradigm generative khác nhau cho cùng task "sinh ảnh mới". Khác biệt chính:

Lựa chọn thực hành:

• Cần likelihood, latent space có encoder, anomaly detection → VAE.
• Cần ảnh sharp cao chất lượng → GAN (hoặc diffusion).
• Modern: Diffusion model (DDPM, 2020; sẽ chạm trong Series 4) đạt chất lượng GAN và ổn định như VAE — best of both worlds.
• Combo VAE + diffusion (Stable Diffusion, mục 19) đang là default cho image generation production.

Higgins và cộng sự (2017) — "β-VAE" — thêm một hệ số \( \beta \) cho KL term:

\[ L_{\beta\text{-VAE}} = -\mathbb{E}_{q_\phi(z \mid x)}\!\left[ \log p_\theta(x \mid z) \right] + \beta \cdot D_{KL}\!\big(q_\phi(z \mid x) \,\|\, p(z)\big) \]

• \( \beta = 1 \): VAE chuẩn.
• \( \beta > 1 \) (vd 4, 10): KL nặng hơn → ép latent chặt hơn về prior. Trên dataset có "factor of variation" rõ (vd dSprites: shape, scale, rotation, x, y), mỗi chiều latent có xu hướng ánh xạ vào một factor riêng — disentangled representation.
• \( \beta < 1 \): giảm KL, gần với AE — recon tốt hơn, latent kém structured.

Trade-off rõ: \( \beta \) lớn → disentangle tốt nhưng reconstruction blur. Paper gốc dùng \( \beta \) cỡ 4–250 tuỳ dataset.

Disentanglement hữu ích cho controllable generation: thay đổi một chiều \( z_j \) để đổi một factor (vd rotation) mà giữ các factor khác. Tuy nhiên, Locatello và cộng sự (2019) chứng minh disentanglement không identifiable nếu không có inductive bias hoặc giám sát — không có model nào "tự nhiên" disentangle hoàn hảo trên data unsupervised.

VAE vanilla sinh sample "ngẫu nhiên" từ \( \mathcal{N}(0, I) \) — không kiểm soát được class. Sohn và cộng sự (2015) — "Learning Structured Output Representation using Deep Conditional Generative Models" — đề xuất CVAE: thêm điều kiện \( y \) (vd nhãn class) vào cả encoder và decoder.

\[ q_\phi(z \mid x, y), \quad p_\theta(x \mid z, y), \quad p(z \mid y) \text{ hoặc } p(z) = \mathcal{N}(0, I) \]

ELBO tương tự, chỉ thêm \( y \):

\[ \mathcal{L}(\theta, \phi; x, y) = \mathbb{E}_{q_\phi(z \mid x, y)}\!\left[ \log p_\theta(x \mid z, y) \right] - D_{KL}\!\big(q_\phi(z \mid x, y) \,\|\, p(z \mid y)\big) \]

Implement bằng cách concat one-hot \( y \) vào input của encoder và decoder. Sau train, sinh ảnh của class cụ thể: chọn \( y \), sample \( z \sim \mathcal{N}(0, I) \), decode \( p(x \mid z, y) \).

# CVAE skeleton (sửa từ VAE skeleton mục 12)
def encode(self, x, y):
    h = self.encoder(torch.cat([x, y], dim=-1))  # y: one-hot, vd (B, 10)
    return self.fc_mu(h), self.fc_logvar(h)

def decode(self, z, y):
    return self.decoder(torch.cat([z, y], dim=-1))

CVAE là bước đệm conceptual cho text-to-image: thay \( y \) one-hot bằng text embedding (vd CLIP), ta được conditional generation theo prompt — chính là idea của Stable Diffusion (mục 19).

VAE được dùng trong các nhóm bài toán sau:

• Image generation: chất lượng giới hạn so với GAN/diffusion, hợp cho prototype hoặc bài toán có ràng buộc latent.
• Anomaly detection: train VAE trên data bình thường; lúc test, reconstruction error hoặc likelihood \( \log p(x) \) thấp → anomaly. Y tế (ảnh CT/MRI), industrial (lỗi sản phẩm), cybersecurity.
• Drug discovery: Gómez-Bombarelli và cộng sự (2018) train VAE trên SMILES string biểu diễn phân tử, dùng latent space để search molecule mới có thuộc tính mong muốn.
• Semi-supervised learning: Kingma và cộng sự (2014) M2 model — VAE với một phần label được giám sát, latent học representation tốt hơn cho classifier downstream.
• Representation learning: latent của VAE có thể dùng làm feature cho downstream task. Tuy nhiên, modern thường dùng contrastive learning (SimCLR, CLIP) thay vì VAE cho mục đích này.

Limitation:

• Output blur: per-pixel loss (BCE/MSE) trung bình hoá → mất chi tiết tần số cao.
• Gaussian assumption cho cả posterior và likelihood: không khớp với data đa modal — model bị "kéo về" giá trị trung bình.
• Posterior collapse: một số dim latent không học được gì.
• Likelihood không phải metric chất lượng: VAE có likelihood cao nhưng ảnh xấu (Theis và cộng sự, 2016 chỉ ra log-likelihood và visual quality có thể không tương quan).

Do limitation này, image generation hiện đại đã chuyển sang diffusion model. Nhưng VAE chưa "chết" — nó là component trong các pipeline lớn hơn (mục 19).

Rombach và cộng sự (2022) — "High-Resolution Image Synthesis with Latent Diffusion Models" (Stable Diffusion) — kết hợp VAE và diffusion:

1. Train một VAE (encoder + decoder) trên dataset ảnh lớn → latent space \( z \in \mathbb{R}^{c \times h/8 \times w/8} \) (giảm 8× mỗi chiều spatial so với pixel).
2. Diffusion model chạy trên latent \( z \) thay vì pixel \( x \) trực tiếp. Diffusion noise / denoise rẻ hơn nhiều lần.
3. Inference: text → noise latent → denoise qua diffusion → VAE decoder → ảnh pixel.

Pixel space (chậm, tốn bộ nhớ):
   x ──► add noise ──► denoise ──► x̂              [DDPM 2020 chạy trực tiếp ở đây]

Latent space (Stable Diffusion):
   x ──► VAE encoder ──► z ──► add noise ──► denoise ──► ẑ ──► VAE decoder ──► x̂
         (1 lần, train xong)    (diffusion model trên z)         (1 lần, train xong)

VAE ở đây không phải để generate trực tiếp mà để compress ảnh về latent rẻ hơn. Diffusion làm phần generative chính. Cách bố trí này:

• Giảm chi phí compute 4–16× so với diffusion pixel space.
• Cho phép training/inference ảnh 512×512 trên GPU consumer (12GB).
• Vẫn giữ được chất lượng vì VAE đủ tốt ở phạm vi reconstruction (không cần generate).

Hiểu VAE bài này là tiền đề để hiểu Stable Diffusion sau (Series 4 sẽ chạm). Diffusion model chi tiết sẽ vượt scope Series 3 — Series 4 module Generative AI sẽ deep dive.

Ghép skeleton mục 12 với data MNIST và 3 routine: reconstruction, generation, interpolation.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

transform = transforms.ToTensor()
train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform)
test_ds  = datasets.MNIST("./data", train=False, transform=transform)
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
test_loader  = DataLoader(test_ds, batch_size=128)

class VAE(nn.Module):
    def __init__(self, latent_dim=20):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(784, 256), nn.ReLU())
        self.fc_mu     = nn.Linear(256, latent_dim)
        self.fc_logvar = nn.Linear(256, latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid(),
        )

    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        return mu + torch.randn_like(std) * std

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar


def vae_loss(x_recon, x, mu, logvar):
    recon = F.binary_cross_entropy(x_recon, x, reduction="sum")
    kl    = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon + kl, recon, kl


# Train
model = VAE(latent_dim=20).to(device)
optim = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(30):
    model.train()
    total = 0.0
    for x, _ in train_loader:
        x = x.view(-1, 784).to(device)
        x_recon, mu, logvar = model(x)
        loss, _, _ = vae_loss(x_recon, x, mu, logvar)
        optim.zero_grad(); loss.backward(); optim.step()
        total += loss.item()
    print(f"Epoch {epoch+1}: loss = {total / len(train_ds):.2f}")


# 1. Reconstruction visual
model.eval()
with torch.no_grad():
    x, _ = next(iter(test_loader))
    x = x.view(-1, 784).to(device)
    x_recon, _, _ = model(x)
    # x[:8] và x_recon[:8] để plot side-by-side

# 2. Generation từ z ~ N(0, I)
with torch.no_grad():
    z = torch.randn(16, 20, device=device)
    samples = model.decoder(z).view(-1, 1, 28, 28)

# 3. Interpolation giữa 2 ảnh
with torch.no_grad():
    x1 = train_ds[0][0].view(1, -1).to(device)
    x2 = train_ds[1][0].view(1, -1).to(device)
    mu1, _ = model.encode(x1)
    mu2, _ = model.encode(x2)
    ts = torch.linspace(0, 1, 10, device=device).unsqueeze(1)
    z_interp = (1 - ts) * mu1 + ts * mu2
    imgs = model.decoder(z_interp).view(-1, 1, 28, 28)

Trên MNIST 30 epoch với cấu hình này:

• Loss per-sample ~85–95 (sum trên 784 pixel + 20 latent dim).
• Reconstruction: hình dạng số đúng, hơi mờ.
• Generation: 14–15/16 sample đọc được thành 1 chữ số.
• Interpolation: chuỗi 10 ảnh chuyển dần mượt.

Tutorial reference: PyTorch examples có vae chính thức (file main.py) — chạy được trực tiếp, kết quả tương đương.

1. Train VAE MNIST với latent_dim = 2, 10, 20. Với latent_dim = 2, plot \( \mu \) của test set lên 2D, tô màu theo class. Quan sát: 10 cụm cách nhau hay xen lẫn? KL có collapse chiều nào?
2. Sample 25 ảnh từ \( z \sim \mathcal{N}(0, I) \) cho mỗi latent_dim ở bài 1. Đếm tỷ lệ ảnh "đọc được thành 1 chữ số". Latent dim nào cho quality tốt nhất?
3. Implement β-VAE với beta = 1, 4, 10 (chỉ nhân KL term với beta trong loss). So sánh: (a) reconstruction loss cuối epoch; (b) chất lượng ảnh sinh; (c) plot 10 ảnh khi thay đổi 1 chiều \( z_j \) (giữ các chiều khác = 0) — chiều nào thay đổi rõ?
4. Latent interpolation: chọn 2 ảnh thuộc 2 class khác nhau (vd "3" và "8"). Plot 10 ảnh nội suy bằng (a) linear interpolation; (b) slerp. So sánh visual smoothness.
5. Compute KL closed-form thủ công cho 1 mẫu: \( \mu = [0.5, -0.3, 1.0] \), \( \log \sigma^2 = [0.0, -0.5, 0.2] \). Áp dụng \( D_{KL} = -\frac{1}{2} \sum (1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2) \). Kết quả là số dương?
6. Implement CVAE trên MNIST: concat one-hot label (10 chiều) vào input encoder và decoder. Sau train, generate 10 ảnh của mỗi class. So sánh với VAE vanilla.
7. Anomaly detection: train VAE chỉ trên số "1" của MNIST. Test trên tất cả class, plot histogram reconstruction loss theo class. Số "1" có loss thấp nhất? Threshold ở đâu để detect 90% non-"1"?

• Autoencoder standard không phải generative — latent space "lung tung", sample ngẫu nhiên decode ra nhiễu.
• VAE (Kingma & Welling, 2013) là AE probabilistic: encoder output \( q(z \mid x) = \mathcal{N}(\mu(x), \sigma^2(x)) \), ép gần prior \( p(z) = \mathcal{N}(0, I) \) qua KL term.
• Encoder thực tế dự đoán \( \mu \) và \( \log \sigma^2 \) (logvar) — numerical stable, range không ràng buộc.
• Reparameterization trick: \( z = \mu + \sigma \cdot \epsilon \) với \( \epsilon \sim \mathcal{N}(0, I) \). Cho phép backprop gradient qua bước sampling.
• Loss VAE = ELBO ngược dấu = reconstruction \( - \mathbb{E}_{q}[\log p(x \mid z)] \) cộng \( D_{KL}(q(z \mid x) \,\|\, p(z)) \).
• Reconstruction: BCE nếu pixel \( \in [0, 1] \) (Bernoulli), MSE nếu Gaussian. Dùng reduction="sum" để cân scale với KL.
• KL Gaussian closed-form: \( -\frac{1}{2} \sum_j (1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2) \). Đạt 0 đúng tại prior.
• Trade-off: recon nặng → latent giống AE, sample kém; KL nặng → posterior collapse, ảnh blur. β-VAE điều chỉnh tường minh qua \( \beta \).
• Generation: sample \( z \sim \mathcal{N}(0, I) \), decode → ảnh mới. Latent interpolation mượt (khác AE), test bằng \( z_t = (1-t) z_1 + t z_2 \).
• VAE vs GAN: VAE blur + ổn định + có likelihood + có encoder; GAN sharp + train khó + không likelihood. Diffusion là best of both.
• CVAE: concat điều kiện \( y \) vào encoder & decoder → conditional generation. Tiền đề cho text-to-image.
• Use case: anomaly detection, drug discovery, semi-supervised, representation learning. Limitation: blur, Gaussian assumption, posterior collapse.
• Stable Diffusion (Rombach 2022) dùng VAE để compress về latent space, diffusion chạy trên latent — giảm 8× chi phí compute, vẫn giữ chất lượng.
• PyTorch skeleton ~50 dòng, train MNIST 30 epoch là đủ thấy 3 hiệu ứng: reconstruction tốt, generation hợp lý, interpolation smooth.