分子潜在扩散中的暗区平滑化

# 平滑分子潜在扩散中的暗区  
来源：https://arxiv.org/html/2606.13955  

王曦¹ 李嘉涵¹ 夏宇轩¹ 吴英成² 郑少义¹ 王胜杰¹  
¹纽约大学 ²斯坦福大学  

###### 摘要  

潜在扩散是可扩展的三维分子生成的一个有前景的框架，但它要求潜在空间在事后样本之外保持平滑、有效且可导航。然而，现有的分子VAE通常通过基于重建的目标来学习，这并不能保证这样的潜在空间。我们表明，这会导致**暗区**：在扩散采样过程中可达但解码为不连通或化学无效的分子的潜在空间区域。与图像生成不同，分子解码要求严格的结构和化学精度，即使是微小的潜在扰动也可能导致灾难性失败。因此，我们提出**TopVAE**，一种拓扑优化的VAE，通过在训练过程中让解码器内化结构和化学约束来减少暗区，从而消除测试时化学校正的需要。TopVAE大幅提升了离事后鲁棒性，当与标准DiT搭配时，在QM9上实现了77%更低的FCD3D、最高的V&C，在GEOM-Drugs上实现了52%更低的FCD3D，并在零样本支架修补中生成了1.29倍更稳定且连通的分子。  

## 1 引言  

潜在扩散模型最初为图像生成而开发(Rombach et al., 2022; Peebles and Xie, 2023)，现已扩展到三维分子生成(Xu et al., 2023; Chen et al., 2025)。这一设定更具挑战性，因为有效的分子必须同时满足结构、几何和化学约束(You et al., 2024; Luo et al., 2025)。因此，潜在变量公式特别有吸引力：如果这些复杂性可以被解码器吸收，扩散可以在更干净的潜在空间中运行，从而简化建模和采样。  

分子潜在空间通常通过重建目标加上正则化来学习，如VAE或VQ-VAE框架(Luo et al., 2025)。然而，强大的重建并不能保证潜在的适合扩散。最近的研究表明，重建质量不能很好地预测，有时甚至与下游生成性能相冲突(Yao et al., 2025; Skorokhodov et al., 2025; Xu et al., 2026)，因为扩散需要平滑、可插值的潜在流形，而不仅仅是在事后样本上的正确性。虽然潜在正则化旨在扩大有效潜在的区域，但这对于分子来说尤其困难：分子解码要求严格的结构和化学精度，并且与图像不同，它缺乏使附近潜在点自然解码的强架构偏置。因此，即使是事后潜在周围的微小扰动也可能产生严重损坏的分子，包括不连通的结构或无法通过化学净化的图(图1)。我们将这些失败区域称为**暗区**：解码器失效的潜在邻域，使得周围的流形不可插值且难以被扩散导航。  

暗区主要来自两个来源：**拓扑不连通**（解码图断裂成孤立组件）和**化学无效性**（预测的键违反价键或兼容性规则）。先前的工作通常通过在生成过程中强制执行化学规则来解决这些问题(Jin et al., 2018; Liu et al., 2018; Ma et al., 2018; Krenn et al., 2022)。虽然这样的约束可以修复输出，但它们不能消除底层的潜在空间脆弱性，并且可能使生成产生偏差。相反，我们旨在将这些约束内化到VAE训练过程中，使有效分子占据更鲁棒的潜在邻域，而无需在推理时进行化学约束优化(ChemCO)。为此，我们提出**TopVAE**（拓扑优化VAE），它通过基于BFS的邻接精化促进连通性，通过在训练过程中展开原始-对偶优化来强制价键和键类型约束，并有选择地将ChemCO的校正注入解码器(Oh et al., 2018)，从而支持无ChemCO的推理。  

总之，我们的贡献如下：  

1. 1.**分子潜在空间中的暗区。** 我们识别并形式化了暗区：扩散采样可达但解码为化学无效或不连通分子的潜在区域。  
2. 2.**TopVAE**包含三个组件：a) TopoBridge通过邻接精化保证连通性；b) 固有ChemCO通过展开的原始-对偶优化提供化学约束；c) 优势门控约束学习(AGCL)有选择地将基于约束的校正信号注入解码器训练。  
3. 3.**暗区闭合的最先进生成。** TopVAE与标准DiT搭配，在QM9上实现了77%更低的FCD3D、最高的V&C，在GEOM-Drugs上实现了52%更低的FCD3D，并在零样本支架修补中生成了1.29倍更稳定且连通的分子。  

参考图注  
图1：TopVAE概述。(A) 暗区诊断显示，现有VAE解码器在离事后潜在区域产生无效分子，TopVAE通过内化化学规则来闭合这些区域。(B) TopVAE结合了拓扑优先的多阶段解码器，利用TopoBridge保证连通性，并在训练过程中利用ChemCO进行化学约束学习。  

## 2 相关工作  

#### 三维分子生成。  
用于从头设计三维分子的扩散和流模型分为两类。**数据空间**模型直接在分子数据空间中生成：一些模型扩散原子类型和坐标，并事后推断键(Hoogeboom et al., 2022)，而后来的图感知或完整分子变体通过联合建模图/拓扑和几何、显式预测键变量或引入键形成感知的训练目标来减少原子-键不匹配(Peng et al., 2023; Huang et al., 2023; Vignac et al., 2023; Le et al., 2023; Reidenbach et al., 2026; Xu et al., 2024)。**潜在**模型首先将分子压缩为连续表示，然后拟合扩散模型(Xu et al., 2023; You et al., 2024; Luo et al., 2025; Joshi et al., 2025)，获得了可扩展性和可控性，但需要高质量潜在用于扩散模型。我们的暗区诊断使这一假设变得明确且可检验，从而激励了对现有潜在分子扩散模型的解码器中心补充。  

#### 约束感知的分子解码。  
化学和结构有效性已通过语法级约束(Kusner et al., 2017; Krenn et al., 2022)、结构化图解码器(Jin et al., 2018; Liu et al., 2018)、有效性导向的正则化(Ma et al., 2018)以及可微约束满足层(Wang et al., 2023; Zeng et al., 2024)来强制执行。TopVAE采取了不同的路线：我们不在推理时保留永久的约束层，而是使用AGCL将ChemCO的约束引导校正作为选择性训练信号，目标是让无约束的解码器本身内化化学规则，并支持无ChemCO的推理(第4.5节，图2)。扩展的相关工作见附录A。  

## 3 前提条件  

#### 三维分子VAE。  
三维分子可以表示为M = (A, T, B, R)，其中A ∈ {0,1}<sup>N×N</sup>是二值邻接矩阵，T表示原子类型，B表示键类型，R ∈ ℝ<sup>N×3</sup>存储原子坐标。分子VAE将M编码为潜在变量z，并通过解码器D<sub>θ</sub>重建分子：  

q<sub>φ</sub>(z | M) = 𝒩(μ<sub>φ</sub>(M), diag(σ<sub>φ</sub><sup>2</sup>(M)))， M̂ = D<sub>θ</sub>(z)， z = μ<sub>φ</sub>(M) + σ<sub>φ</sub>(M) ⊙ ε。 (1)  

这里q<sub>φ</sub>(z | M)是编码器后验，由编码器参数φ参数化，将输入分子映射到高斯潜在分布的均值和方差。VAE通过平衡重建和潜在正则化来训练：  

L<sub>VAE</sub> = 𝔼<sub>M∼𝒟</sub> 𝔼<sub>z∼q<sub>φ</sub>(z | M)</sub> [ -log p<sub>θ</sub>(M | z) ] + β KL( q<sub>φ</sub>(z | M) ∥ p(z) )。 (2)  

#### 分子潜在上的潜在扩散。  
在VAE潜在空间上训练扩散模型以生成新的潜在码。从干净的潜在z<sub>0</sub>开始，前向过程添加高斯噪声：  

z<sub>t</sub> = √(ᾱ<sub>t</sub>) z<sub>0</sub> + √(1-ᾱ<sub>t</sub>) ε， ε ∼ 𝒩(0, I)。 (3)  

去噪器ε<sub>ψ</sub>(z<sub>t</sub>, t)被训练来预测ε，通常使用MSE损失。在采样时，迭代去噪产生干净的潜在ẑ<sub>0</sub>，然后解码为分子：M̂ = D<sub>θ</sub>(ẑ<sub>0</sub>)。因此，解码器不仅要在事后潜在上工作，还要在扩散过程中访问的噪声和先验类潜在上工作。  

#### 分子潜在空间中的暗区。  
在推理时，潜在扩散模型可能采样到位于训练期间见过的事后码之间或之外的潜在。如果解码图(i)通过RDKit净化（要求合法价键、一致芳香性和识别的键类型）并且(ii)形成单个连通组件，我们称其为化学有效。用A<sub>LDM</sub>表示采样可达区域，用M<sub>valid</sub>表示有效分子集，我们定义**暗区**为：  

D<sub>dark</sub> = { z ∈ A<sub>LDM</sub> : D<sub>θ</sub>(z) ∉ M<sub>valid</sub> }。 (4)  

## 4 TopVAE：用于化学约束学习的拓扑优化VAE  

TopVAE针对分子潜在空间中的暗区。目标是让VAE解码器学习化学规则，从而闭合暗区。因此，TopVAE结合了结构化分子解码器和两个约束感知模块。TopoBridge将预测的邻接精化为保证连通的邻接矩阵。ChemCO展开化学约束优化，并将价键限制、键互斥性、原子对合法性和度规则转化为可微训练信号。  

### 4.1 编码器  

编码器使用关系变换器(Diao and Loynd, 2022)将分子图映射为潜在特征。我们首先嵌入原子和原子对：  

h<sub>i</sub><sup>(0)</sup> = f<sub>enc-node</sub>(T<sub>i</sub>, R<sub>i</sub>)， p<sub>ij</sub> = f<sub>enc-edge</sub>(B<sub>ij</sub>, ‖R<sub>i</sub> - R<sub>j</sub>‖<sub>2</sub>)。 (5)  

关系变换器通过边感知注意力更新原子特征，其中对特征p<sub>ij</sub>被注入到原子i和j之间的注意力中：  

H = RTrans<sub>φ</sub>( {h<sub>i</sub><sup>(0)</sup>}<sub>i=1</sub><sup>N</sup>, {p<sub>ij</sub>}<sub>i,j=1</sub><sup>N</sup> )。 (6)  

得到的原子级特征H = (h<sub>1</sub>, ..., h<sub>N</sub>)用于参数化式(1)中的潜在分布。  

### 4.2 拓扑感知的多阶段解码器  

使用第3节中的符号，TopVAE以拓扑优先的顺序解码每个潜在码z：  

p<sub>θ</sub>(A, T, B, R | z) = p<sub>θ</sub>(A | z) p<sub>θ</sub>(T | A, z) p<sub>θ</sub>(B | A, T, z) p<sub>θ</sub>(R | A, T, B, z)。 (7)  

这种分解使邻接成为第一个解码的对象。特别地，键预测受限于：A<sub>ij</sub> = 0 ⇒ B<sub>ij</sub> = 0，A<sub>ij</sub> = 1 ⇒ B<sub>ij</sub> ∈ {1, ..., K}。 (8)  

具体地，节点投影头f<sub>node</sub>首先将每个潜在token z<sub>i</sub>映射为初始解码器特征g<sub>i</sub><sup>(0)</sup>。邻接头f<sub>A</sub>然后取对特征[g<sub>i</sub><sup>(0)</sup>, g<sub>j</sub><sup>(0)</sup>]并预测边存在概率：  

g<sub>i</sub><sup>(0)</sup> = f<sub>node</sub>(z<sub>i</sub>)， P<sup>A</sup><sub>ij</sub> = σ( f<sub>A</sub>( [g<sub>i</sub><sup>(0)</sup>, g<sub>j</sub><sup>(0)</sup>] ) )。 (9)  

由于分子键是无向的且自键无效，我们对P<sup>A</sup>进行对称化并将对角线设为零；TopoBridge（记为Π<sub>TB</sub>）然后将其精化为连通的邻接矩阵：  

à = Π<sub>TB</sub>(P<sup>A</sup>)。 (10)  

精化后的邻接作为注意力偏置注入TopoFormer：  

α<sub>ij</sub> = softmax<sub>j</sub>( (q<sub>i</sub><sup>⊤</sup> k<sub>j</sub> / √d) + e<sub>A</sub>(Ã<sub>ij</sub>) )， g<sub>i</sub> ← ∑<sub>j</sub> α<sub>ij</sub> v<sub>j</sub>， (11)  

其中q<sub>i</sub>、k<sub>j</sub>、v<sub>j</sub>是查询、键、值向量。