变分自编码器(VAE)与ELBO损失

好，咱们开始聊VAE。说实话，在分子生成这个领域，VAE是绕不开的基石。我最早接触VAE是在做SMILES字符串生成的时候，那会儿就想：能不能把分子编码成一个连续向量，然后在这个向量空间里做插值、优化？VAE就是干这个的。

但VAE和普通自编码器最大的区别在于——它不学确定性的隐向量，而是学一个概率分布。你想想看，分子空间是离散的、高维的，直接学一个确定映射很容易过拟合，生成出来的分子要么无效，要么全是训练集里的老面孔。VAE通过引入隐变量的先验分布，让模型学会“平滑”的隐空间，新分子可以从这个分布里采样得到。

嗯，这里有个关键前提：VAE的损失函数不是简单的MSE，而是ELBO(Evidence Lower Bound)。我当年第一次推导ELBO的时候卡了好几天，后来发现其实就是把“对数似然”拆成两项——重建误差和KL散度。下面我带你一步步拆解。

VAE的编码-解码结构

先画张图，把整体流程先立住。

重参数化技巧：为什么需要它？

咱们从采样这步出发。隐变量z是从一个分布中随机采样的——但采样操作是不可导的，没法反向传播。我记得刚开始写VAE的时候，直接在图上加了个torch.distributions.Normal().sample()，训练直接崩掉。后来才发现必须用重参数化：

# 错误做法：直接采样，梯度断掉
z = torch.distributions.Normal(mu, sigma).sample()

# 正确做法：重参数化
epsilon = torch.randn_like(sigma)   # 从标准正态采样
z = mu + sigma * epsilon            # 线性变换，梯度保留

你仔细想想，这里的奥秘在于：我们把随机性从模型参数中剥离出去了。epsilon是外部噪声源，不参与梯度计算；而mu和sigma是网络输出的确定性函数，梯度可以畅通无阻地流回去。这就是重参数化技巧最核心的工程价值。

ELBO下界的推导

推导ELBO其实就是最大化对数似然的下界。直接写目标：

log p(x) ≥ E_z~q(z|x)[log p(x|z)] − KL(q(z|x) || p(z))

右边就是ELBO。具体推导过程我写个卡片：

log p(x) = log ∫ p(x|z) p(z) dz
        = log ∫ p(x|z) p(z) * q(z|x)/q(z|x) dz
        ≥ ∫ q(z|x) log[ p(x|z) p(z) / q(z|x) ] dz   (Jensen不等式)
        = E[log p(x|z)] − KL(q(z|x) || p(z))

这里用到了Jensen不等式，把log算子移到期望内部。注意我们取的是下界，优化ELBO等价于同时增大log p(x)和降低KL散度。我在项目中遇到过一个问题：KL散度项有时候会占主导，导致模型完全忽略重建损失，生成的分子都是随机噪声。后来我把KL权重从1降到0.1，效果才正常。

项	含义	实现方式
重建损失	衡量解码器能否还原输入	对于分子图常用BCE/CE，对于SMILES用交叉熵
KL散度	约束隐变量分布接近标准正态	封闭解：-0.5 * sum(1 + log(σ²) − μ² − σ²)

ELBO损失函数的PyTorch实现

下面给出一份我常用的VAE损失代码，融合了分子生成场景下的常见处理：

class VAELoss(nn.Module):
    def __init__(self, kl_weight=0.1):
        super().__init__()
        self.kl_weight = kl_weight

    def forward(self, recon_x, x, mu, logvar):
        # 重建损失：这里假设解码器输出是logits，x是二值图
        recon_loss = F.binary_cross_entropy_with_logits(
            recon_x, x, reduction='sum'
        )
        # 0.5 * sum(1 + log(σ²) - μ² - exp(logvar))
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        return recon_loss + self.kl_weight * kl_loss

这段代码里有几个点值得注意：

• 重建损失用了sum而不是mean，对batch内所有元素求和，搭配KL项也求sum，两个损失可以自然加权。
• kl_weight我习惯设0.1～0.5之间，分子生成任务尤其敏感，你必须调参。
• 如果解码器输出的是分子原子类型one-hot，记得先过softmax再算CE。

KL散度计算细节

KL散度封闭形式推导其实不难，但很多人一写就错。我习惯用这个公式：

# 假设 p(z) = N(0, I), q(z|x) = N(μ, σ²)
KL = 0.5 * (μ² + σ² - 1 - log(σ²))

# 等价写法（用logvar）：
KL = 0.5 * (μ² + exp(logvar) - logvar - 1)

为什么不用sigma直接用logvar？因为有数值稳定性考虑——logvar可以取负无穷，而sigma必须≥0。网络输出logvar全空间自由，你想想看，是不是省去了手动加softplus的麻烦。

实际中还会遇到隐变量空间的维度选择。做分子生成时，我试过从64维到512维，发现256维左右对大多数数据集效果不错。太大会导致隐变量稀疏，KL项容易降到0退出；太小又信息瓶颈，重建质量差。

标准化流程

1. 编码器输出两个分支：一个生成mu，一个生成logvar（最后不用激活函数，或者用tanh限制范围）。
2. 重参数化采样：z = mu + exp(0.5 * logvar) * eps。
3. 解码器：接受z，输出重建概率（可以是分子图邻接矩阵的概率，或SMILES字符概率）。
4. 损失计算：重建损失 + KL损失 × 权重。

这个流程虽然基础，但百分之九十的分子生成VAE变体都是在它上面做改进——比如β-VAE增加KL惩罚、InfoVAE加互信息项、或者用Wasserstein距离替换KL。我建议你先把标准VAE跑通，再尝试其他花样。