第3章：分子生成模型概览 — VAE、GAN、Flow、扩散模型、强化学习

做分子生成，说白了就是让模型学会“画”出有药效的分子。模型选对了，损失函数才能派上用场。我个人习惯先捋一遍主流生成模型的底牌——每种模型怎么生成分子，它用损失函数达成什么目标。嗯，咱们一个个来看。

3.1 变分自编码器（VAE）

VAE是分子生成里的常青树。它把分子编码成隐空间向量，再解码回来。核心损失函数由两部分组成：

• 重构损失：让解码出的分子跟原分子尽可能像。对于SMILES序列，常用交叉熵；对于图结构，常用图编辑距离或边预测损失。
• KL散度：惩罚隐变量分布偏离标准正态，说白了就是压着隐空间不要发散。

我在项目里试过VAE生成类药分子，刚开始KL散度系数设得太小，隐空间崩成一团，生成的全是乱七八糟的字符串。后来调成β-VAE的思路才稳住。

def vae_loss(recon_x, x, mu, logvar):
    recon_loss = F.cross_entropy(recon_x, x, reduction='sum')
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + beta * kl_loss

3.2 生成对抗网络（GAN）

GAN在图像生成上名声大噪，但在分子生成里有点水土不服。它包含生成器和判别器两个网络，互相博弈。

• 对抗损失：生成器想骗过判别器，判别器想分辨真假。典型的二分类交叉熵损失。
• 偶尔加辅助损失：比如合法性惩罚、药性约束，防止生成乱码分子。

你想想看，分子是离散数据（原子类型、键），GAN的梯度很难直接回传。我记得早期用GAN生成SMILES时，判别器很容易就能区分真伪，生成器死活学不会。后来有人用WGAN-GP的损失函数缓解了这个问题。

3.3 标准化流（Normalizing Flow）

Flow模型走的是另一条路：用可逆变换把简单分布映射成分子分布。损失函数很纯粹——最大化负对数似然（NLL）。

Flow的好处是能精确计算概率密度，生成质量高。但可逆变换限制了网络结构，计算量也大。我试过用Flow生成小分子，训练一天才收敛，但样本多样性的确比VAE好。

# Flow的损失函数直接是负对数似然
def flow_loss(z, log_det_jacobian):
    # z是经过flow变换后的隐变量
    nll = 0.5 * torch.sum(z ** 2, dim=1) - log_det_jacobian
    return nll.mean()

3.4 扩散模型（Diffusion Model）

最近扩散模型火得不行。它的思路是把分子逐步加噪变纯噪声，再学逆向过程去噪。损失函数通常是一个简单的噪声预测误差：

• 在每一步预测加入的噪声，用MSE跟真实噪声比较。
• 有时会用变分下界（ELBO）作为统一框架。

扩散模型生成分子质量贼高，但采样速度慢。我在优化时发现，损失函数里对时间步的加权很重要——早期步和晚期步的噪声幅度不一样，用均等权重容易顾此失彼。

def diffusion_loss(model, x_0, noise, t):
    x_t = forward_diffusion(x_0, noise, t)
    noise_pred = model(x_t, t)
    return F.mse_loss(noise_pred, noise)

3.5 强化学习（RL）

RL在分子生成里常作为“精调”阶段出现。它不直接学分布，而是根据奖励信号优化生成策略。

• 策略梯度损失：最大化期望奖励，常用REINFORCE或PPO。
• 结合预训练模型的KL惩罚：防止生成器忘了药性以外的化学合理性。

说白了，RL的损失函数不是训练生成模型的起点，而是收尾。我常用的套路是：先用VAE或扩散模型预训练一个基础生成器，再跑RL微调，损失函数里加上KL项稳住基础。

def reinforce_loss(log_probs, rewards):
    # rewards已经想过baseline
    loss = -torch.mean(log_probs * rewards)
    return loss

3.6 模型对比一览

3.7 本章知识体系

下面这张图把各模型在损失函数视角下的关系理清了。你可以看到：不同模型用不同损失函数解决“让生成分子合理”这件事。