VAE 变分自编码器

普通自编码器（AE）学到了数据的压缩表示，但它的潜在空间是"混乱"的 — 不同数据的编码可能散落在潜在空间的任意位置，彼此之间没有规律。

问题：如果你想生成新数据，你需要从潜在空间中随机采样一个点，然后解码。但 AE 的潜在空间不连续、不规则，随机采样很可能落在"空白区域"，解码出来的东西毫无意义。

VAE 的核心思想：强制让潜在空间服从正态分布 N(0,1)。这样潜在空间是连续、平滑的，任意采样一个点都能解码出有意义的结果。

KL 散度公式

KL 散度衡量"编码器输出的分布 q(z|x)"与"标准正态分布 p(z) = N(0,1)"之间的差异。 对于对角高斯分布，KL 散度有解析解：

import torch
import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=256, latent_dim=2):
        super().__init__()
        # 编码器：输入 → 隐藏层 → μ 和 log(σ²)
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
        )
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)      # 输出 μ
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)   # 输出 log(σ²)

        # 解码器：z → 隐藏层 → 重建
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid(),
        )

    def reparameterize(self, mu, logvar):
        """重参数化技巧：z = μ + σ × ε"""
        std = torch.exp(0.5 * logvar)   # σ = exp(0.5 * log(σ²))
        eps = torch.randn_like(std)        # ε ~ N(0, 1)
        return mu + std * eps              # z = μ + σ × ε

    def forward(self, x):
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decoder(z)
        return x_recon, mu, logvar

def vae_loss(x_recon, x, mu, logvar):
    # 重建损失：二元交叉熵（或 MSE）
    recon_loss = nn.functional.binary_cross_entropy(
        x_recon, x, reduction='sum'
    )
    # KL 散度损失：-0.5 × Σ(1 + log(σ²) - μ² - σ²)
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_loss