扩散模型在3D分子生成（等变扩散）

可视化演示

展示前向扩散加噪、逆向去噪生成及SE(3)等变性原理，切换模式查看扩散过程不同阶段。

扩散模型基础：从乱泥到陶胚

清晨的陶艺工坊里，陶艺师先做一件看似“破坏”的事：把昨日的小陶罐浸入水盆，揉捏、甩动，直到它变成看不出形状的泥团。这个“越搅越乱”的过程，对应扩散模型的正向过程。陶艺师真正的功夫，在于接下来用指尖反复推挤、回挑，一圈一圈地让泥团重新显现出器物的轮廓。这个“从乱到整”的过程，就是反向过程。扩散模型做的正是这件事：学会在每一轮微小的去噪中，把混乱的空间位形推回有理有据的分子构象。

正向过程不是一步到位的，而是像把陶罐泡在水里慢慢变软，每个时间步只加入一点高斯噪声。当 $t$ 从 0 走到 $T$，原本清晰的原子坐标逐渐被噪声淹没，最终变成一坨完全随机的三维散点。反向过程则训练一个神经网络，从任意噪声点出发，一步步“预测该往哪修”，最终得到像真实分子一样的三位原子排列。因为每一步都是小修补，训练和采样都建立在马尔可夫链的假设上。

数学上，正向过程可以一步写出闭式解：

陶艺师修坯不是一次到位，而是每一圈只微调一点；扩散模型同样依靠无数微小的去噪步骤逼近精准结构。学会它，本质上是让机器记住“每次被搅乱后，该如何把泥推回正轨”。

3D分子表示：陶土的质地与胚体

陶艺师看一个器物，从来不只看表面文字描述，而是同时关注两件事：陶土的质地（是瓷泥、粗陶还是紫砂）和胚体的三维形状（罐口多宽、腹部多鼓、底部多厚）。分子也是如此：一个分子由一组原子构成，每个原子既有自己的元素类型，又有确定的三维空间坐标。元素类型是离散标签，坐标是连续向量，两者共同决定分子的一切化学行为。

传统方法常把分子写成 SMILES 字符串，这相当于把立体雕塑压扁成一张平面说明书——虽然能描述原子连接关系，却丢失了真正的三维构象。3D 分子生成模型直接输出每个原子的坐标 $x_i$ 和元素类型 $z_i$，像同时给出陶土配方和胚体模型，无需再靠预测算法去反推空间结构。

有了原子坐标后，原子间是否存在化学键可以由几何关系判断：距离足够近的原子对形成共价键，距离过远则没有。因此图神经网络常用的边特征就是原子对之间的欧氏距离：

回到陶艺：两个陶土部件能否黏合，取决于它们之间的相对距离和贴合角度，而不是它们放在工作台的哪个角落。三维分子表示正是把这些“相对几何”作为模型的核心语言。

等变原理：转盘上的稳准手艺

陶艺工坊里最奇妙的是旋转陶轮。无论轮盘转到哪个角度，陶艺师的手只要跟随器物当前朝向精准发力，最终器形都不会变。换句话说：先旋转分子再做去噪处理，与先做去噪处理再旋转分子，结果完全一致。这个性质就是“等变”（Equivariance）。物理规律不依赖于你站在哪里、面向哪个方向，分子生成模型也必须尊重这一点。

如果用数学语言刻画，设 $R$ 是旋转矩阵，$t$ 是平移向量，$f$ 是去噪网络，那么等变要求：

普通数据增强可以给模型“看”很多旋转后的样本，但那只算“练多了眼熟”；等变网络结构则像师傅从骨子里明白“轮盘旋转只是换了视角，分子本质没有变”。于是生成出的三维坐标，天然不依赖初始朝向。

训练目标：用损失函数校准手感

陶艺师培养手感的方式很简单：把泥搅乱，感受偏差，反复修正。扩散模型的训练目标也遵循这一逻辑。正向过程为每个样本制造一个“标准答案”——真实噪声 $\epsilon$，网络则根据当前带噪坐标 $x_t$ 和时间步 $t$ 预测噪声 $\hat{\epsilon}$。训练时不断缩小预测噪声与真实噪声之间的差距。

在 3D 分子生成中，坐标是连续量、原子类型是离散量，因此损失函数由两部分组成：坐标部分用均方误差（MSE），原子类型部分用交叉熵。其中坐标去噪损失可写成：

采样时，模型并不直接预测陶罐最终形态，而是预测“该往哪个方向推泥”。于是生成新分子变成：从纯噪声坐标出发，反复使用网络给出的修正方向，逐步去噪，直到坐标清晰成形。

模型架构：等变图神经网络

陶艺师塑形靠十根手指：每根手指不仅感知自己接触的那一小块陶土，还通过附近手指的相对位置判断局部弧度。等变图神经网络（Equivariant GNN）也是如此。分子被建模成一张图：原子是节点，近邻原子之间连边，每一层网络都让节点从邻居那里收集“消息”，再更新自己的特征和坐标。

关键在于，网络不能把绝对坐标当作普通标量特征直接传递，因为绝对坐标会随整体旋转而改变。正确的做法是把两个原子的相对位移向量 $(x_i - x_j)$ 和距离 $r_{ij}$ 作为消息来源，这样消息本身就会随分子旋转而同步旋转。典型的一层坐标更新可以写成：

消息权重 $m_{ij}$ 由旋转不变的距离和元素类型决定，更新方向由相对向量决定，于是整个网络“既有判断力，又有方向感”——这正是等变扩散模型能精准生成三维结构的关键。

应用与挑战：从工坊到药房

出师后的陶艺工坊不再只做摆件，而是开始为药厂定制“分子花瓶”。在药物发现中，科学家希望生成能嵌入靶蛋白口袋的小分子配体。等变扩散模型可以直接在三维空间中生成原子的坐标和元素类型，生成结果本身就是可对接的三维构象，省去了传统 2D 分子再转 3D 的繁琐步骤。

这项技术已经大量用于靶向药物先导物生成、分子构象采样、以及逆合成分析辅助。比如给定一个蛋白口袋，扩散模型在口袋内部不断加噪、去噪，生成一系列候选配体，再通过对接打分筛选出高潜力分子。又比如给定目标分子的骨架，逆合成算法可以借助扩散生成“片段连接处”的合理构象，减少合成路线搜索空间。

但远非尽善尽美。首先是采样效率瓶颈：从纯噪声到完整分子往往需要数百步去噪，速度远慢于单次前向的生成模型。其次是离散与连续混合优化：坐标是连续值，元素类型是离散标签，两类变量必须联合生成，很容易出现“碳原子长在离谱位置”的错误。最后是化学有效性：模型可能生成键角扭曲、立体化学不合理或能量过高的结构，仍需要物理场评估和分子动力学校验。

未来的方向就像陶艺师从手工拉坯进化到半自动化批量生产：更快的高阶采样器、可控的条件扩散、与口袋几何深度耦合的引导机制，都是让这座“分子陶坊”真正进入药企流水线的必经之路。

AI 应用场景

Python 演示

# Minimal equivariant diffusion for 3D molecular skeleton (numpy only)import numpy as np # Target molecule: regular hexagon (like a benzene carbon skeleton)ANG = np.linspace(0.0, 2.0 * np.pi, 6, endpoint=False)X0 = np.stack([np.cos(ANG), np.sin(ANG)], axis=1)N, D = X0.shapeT = 60 # Diffusion noise schedulebetas = np.linspace(0.005, 0.02, T, dtype=np.float64)alphas = 1.0 - betasalpha_bar = np.cumprod(alphas)sqrt_ab = np.sqrt(alpha_bar)sqrt_omab = np.sqrt(1.0 - alpha_bar)sqrt_betas = np.sqrt(betas)  def q_sample(x0, t):    """Forward noising: x_t = sqrt(alpha_bar_t) * x0 + sqrt(1-alpha_bar_t) * eps"""    eps = np.random.randn(*x0.shape)    return sqrt_ab[t] * x0 + sqrt_omab[t] * eps, eps  def denoise_predict(xt):    """Equivariant noise predictor: only uses inter-atomic distances"""    eps_hat = np.zeros_like(xt)    for i in range(N):        for j in range(i + 1, N):            diff = xt[i] - xt[j]            dist = np.linalg.norm(diff) + 1e-6            weight = dist - 1.0            eps_hat[i] -= weight * diff / dist            eps_hat[j] += weight * diff / dist    return eps_hat  def p_sample(xt, t):    """One reverse denoising step: recover x_{t-1} from x_t"""    eps_hat = denoise_predict(xt)    coef = betas[t] / sqrt_omab[t]    mean = (xt - coef * eps_hat) / np.sqrt(alphas[t])    if t == 0:        return mean    return mean + sqrt_betas[t] * np.random.randn(*xt.shape)  # Generate molecule by starting from pure noise and denoising step by steprng = np.random.default_rng(42)x = rng.standard_normal((N, D))for t in reversed(range(T)):    x = p_sample(x, t) bond_len = np.linalg.norm(np.roll(x, -1, axis=0) - x, axis=1)print("Generated skeleton coordinates:")print(np.round(x, 3))print("Mean bond length:", round(float(bond_len.mean()), 3))

1. betas = np.linspace(0.005, 0.02, T)定义扩散过程中的噪声预算表，决定每一步加噪多还是少。
2. q_sample前向加噪公式 x_t = sqrt(alpha_bar_t) * x0 + sqrt(1-alpha_bar_t) * eps，把干净骨架逐步打乱。
3. denoise_predict等变噪声预测器，只用原子间相对距离和方向，天然满足旋转与平移等变。
4. weight = dist - 1.0距离越偏离理想键长 1.0，修复力度越大，像陶艺师对塌陷处加大推挤。
5. p_sample反向去噪一步，根据预测噪声计算去噪均值，并保留少量随机性。
6. coef = betas[t] / sqrt_omab[t]DDPM 反向更新的关键系数，把预测噪声换算为坐标修正力度。
7. for t in reversed(range(T))从最高噪声步倒序回放，逐级雕塑出清晰骨架。
8. np.roll(x, -1, axis=0)将坐标按原子顺序错位一位，方便计算首尾相接的相邻键长。
9. bond_len.mean()输出平均键长，验证生成结构是否符合六元环的等边几何。

常见误区

面试题