ICLR 2024 | LDM-3DG：在何处扩散比如何扩散更重要

生成3D分子结构是AI药物发现的核心挑战。当前的扩散模型大多直接在高维的原子坐标和连接关系上运行，但这真的是最优选择吗？来自德州农工大学团队的研究提出了一个根本性问题：我们应该在什么空间进行扩散？他们的答案是 LDM-3DG，一个在精心设计的潜空间中运行的3D图扩散框架。通过理论分析和精巧的级联自编码器设计，LDM-3DG 不仅在生成质量上超越现有方法，训练速度更是提升了10-34倍。这项工作揭示了一个深刻洞察：选对空间，比优化算法本身更重要。

自监督潜3D图扩散流水线概览及级联自编码器架构

自监督潜3D图扩散流水线概览及级联自编码器架构

一、空间选择：一个被忽视的根本问题

直觉的陷阱

当我们谈论分子生成时，最直观的做法是直接在分子的自然表示上工作——原子坐标构成的3D空间，以及原子间的化学键构成的图结构。这正是 EDM、GeoDiff 等先驱模型的选择。毕竟，这个空间天然具备物理意义，对称性结构清晰可见。

但直觉往往具有欺骗性。想象一个包含30个原子的小分子，在笛卡尔坐标系中就对应90维空间（每个原子3个坐标）。如果考虑原子类型、化学键类型等离散特征，维度进一步爆炸。扩散模型需要在这个高维空间中学习数据分布，每增加一个维度，所需的训练样本和计算成本都呈指数增长。

更微妙的问题在于，真实的分子并不均匀分布在这个高维空间中。它们遵循量子力学的约束、化学键的规则、热力学的稳定性要求，实际上栖息在一个受限的低维流形上。在高维空间中进行扩散，就像用大炮打蚊子——虽然最终能击中目标，但大部分计算都浪费在探索那些永远不会被访问的无意义区域。

理论的指引

LDM-3DG 的团队用数学语言阐明了这个直觉。他们证明，3D图扩散模型的性能上界取决于三个因素：

性能界限 ≲ 重建误差 + 对称性偏差 × 数据维度

这个简洁的不等式蕴含了深刻含义。首先，降低维度能直接收紧性能上界——这是潜空间扩散的数学动机。其次，重建质量至关重要，糟糕的自编码器会毁掉整个流程。最后，对称性必须被保留，否则生成的分子可能违反基本的物理不变性。

这个理论框架将问题从工程直觉提升到了数学原理层面。它告诉我们：只要能构建一个低维、高保真、对称性保持的潜空间，在其中进行扩散理论上就优于原始空间。

二、级联设计：分而治之的智慧

解耦的艺术

3D分子同时具备两种对称性：置换不变性（原子编号的顺序不重要）和SE(3)等变性（分子可以任意旋转平移）。这两种对称性在原始空间中相互交织，给模型优化带来巨大挑战。

LDM-3DG 的核心创新是级联自编码器架构，将3D图分解为两个独立的视图：

• 拓扑视图（G）：仅包含原子类型和化学键连接关系的2D图，处理置换对称性
• 几何视图（C）：原子的3D空间坐标，处理SE(3)对称性

这种分解看似简单，实则精妙。通过先编码拓扑再编码几何，模型能分别在不同的对称性约束下优化，避免了两种对称性同时作用带来的优化景观复杂化。这就像解决一个复杂方程组时，先分离变量再逐个求解，而不是直接硬解耦合的多元方程。

等变3D对象条件下的潜扩散流水线及处理复杂几何的自编码器架构

等变3D对象条件下的潜扩散流水线及处理复杂几何的自编码器架构

细节中的魔鬼

拓扑自编码器采用标准的图神经网络架构，通过消息传递聚合邻居信息。关键设计在于使用全局池化获得置换不变的图表示，压缩到低维潜向量 z_G。解码器则使用自回归Transformer逐步生成原子和化学键序列。

几何自编码器更具挑战性。为了保证SE(3)等变性，编码器使用等变图神经网络（如SchNet或E(n) Equivariant GNN），将3D坐标映射到潜向量 z_C。解码器同样保持等变性，输出的坐标会随输入的旋转平移而同步变换。

这里有一个微妙但关键的设计：潜空间 z_G 和 z_C 本身是不变的（invariant），而非等变的。这意味着无论分子如何旋转，潜表示保持不变。等变性只在解码阶段恢复，通过重新注入参考坐标来确定绝对位置。这个设计使得扩散过程可以在一个稳定、无冗余的表示空间中进行。

三、实验洞察：速度与质量的双赢

一个数量级的飞跃

在 QM9 数据集上，LDM-3DG 的训练时间从 EDM 的 26.7小时 降至 2.2小时，快了12倍。在更大的 Drugs 数据集上，差距更加惊人——从 146.1小时 降至 4.3小时，快了34倍。这不是渐进式改进，而是范式转变。

速度提升的根源在于维度降低。原始空间中，一个30原子分子对应90维坐标加上离散特征；潜空间中，被压缩到仅64维的连续向量。扩散模型的每一步去噪操作，在低维空间中自然更高效。

质量的提升更令人惊讶

人们可能担心，压缩会损失信息，导致生成质量下降。但实验结果恰恰相反。在 QM9 上，LDM-3DG 生成分子的化学合法性达到100%，而 EDM 为95.6%。在药物相似性（QED）指标上，LDM-3DG 达到0.948，显著高于 EDM 的0.497。

这个看似矛盾的现象揭示了深刻真理：好的压缩不是丢弃信息，而是提炼本质。潜空间通过编码器的学习，滤除了原始空间中的噪声和冗余，保留了决定分子性质的关键特征。在这个纯净的空间中，扩散模型更容易捕捉真实数据分布。

分布外泛化的惊喜

在量子属性条件生成任务中，模型需要根据指定的能隙、极化率等物理量生成分子。当测试属性值超出训练范围（OOD设置）时，大多数模型性能急剧下降。但 LDM-3DG 通过自监督学习正则化潜空间，显著改善了OOD表现。

具体来说，在训练自编码器时，额外引入对比学习目标，迫使潜空间中距离相近的点对应性质相似的分子。这种语义有序性使得模型在面对未见过的属性值时，能通过潜空间插值合理外推，而不是盲目猜测。

实验证明，引入自监督学习后，能量最高占据轨道（εH）的MAE从6.84降至3.92，改善了43%。更有趣的是，通过可视化潜空间的同质性比率（homogeneity ratio），研究者发现改善程度与潜空间的语义聚类质量呈正相关——这为未来设计更好的潜空间正则化策略提供了方向。

随极化率变化的分子生成可视化，展示模型如何准确捕捉分子大小与属性的关系

随极化率变化的分子生成可视化，展示模型如何准确捕捉分子大小与属性的关系

四、蛋白质结合：等变性的价值

条件生成的挑战

在基于结构的药物设计中，目标是生成能与特定蛋白质口袋结合的配体分子。这个任务的特殊之处在于条件本身是3D的——蛋白质口袋的形状和位置决定了配体应该如何放置。

传统方法通常将蛋白质信息编码为固定的向量，作为扩散模型的条件。但这种做法丢失了关键信息：配体与蛋白质的相对位置必须在物理空间中对齐。如果蛋白质旋转了，生成的配体也应该同步旋转，保持相对几何关系不变。

等变解码的精妙

LDM-3DG 通过等变解码器优雅地解决了这个问题。关键技巧是：在解码几何信息时，将蛋白质口袋的原子坐标作为参考系重新输入模型。这样，解码器输出的配体坐标天然地处于与蛋白质对齐的坐标系中。

更重要的是，这个设计保持了SE(3)等变性。数学上可以证明，如果对蛋白质和生成配体同时施加旋转或平移，它们的相对位置关系保持不变。这不是通过复杂的约束实现的，而是架构设计的自然结果。

在 CrossDocked 数据集上的实验验证了这个设计的价值。LDM-3DG 生成的配体在高亲和力比例（HiAff）上达到6.84%，远高于 TargetDiff 的2.14%。Top-10%样本的Vina对接分数达到-12.34 kcal/mol，显示出极强的结合潜力。

蛋白质结合目标下的生成分子可视化，展示模型发现的新结合位点

蛋白质结合目标下的生成分子可视化，展示模型发现的新结合位点

五、消融的启示：质量的瓶颈

级联不可或缺

一个自然的问题是：能否用单个自编码器同时处理拓扑和几何？研究者尝试了这种单阶段设计，结果灾难性失败。单阶段模型生成分子的3D几何Tanimoto相似度几乎为零，说明几何信息完全丢失。

这个对比实验证实了理论分析：当两种对称性交织在一起时，优化变得极其困难。级联设计通过解耦，让模型能分别专注于不同的对称性约束，是保证重建质量的关键瓶颈。

扩散优于VAE

在潜空间中，既可以使用扩散模型，也可以使用更简单的VAE先验。实验显示，扩散模型在捕捉分子量、LogP等属性的分布（TV距离）上显著优于VAE。这印证了扩散模型在建模复杂多峰分布方面的天然优势。

几何不可省略

一个有趣的消融是：能否只用拓扑信息，忽略3D几何？毕竟，化学性质主要由连接关系决定。但实验表明，缺少几何信息会导致生成分子的能量分布偏移，说明3D构象虽然不决定化学式，但对分子稳定性和性质至关重要。

这提醒我们：即使在看似2D的化学问题中，3D信息也不是可有可无的装饰，而是物理现实的反映。

六、更深的思考

表示学习的哲学

LDM-3DG 的成功揭示了一个超越技术细节的原则：在正确的表示空间中，简单的方法就足够强大。这呼应了深度学习的核心哲学——学习好的表示比设计复杂的算法更重要。

潜空间扮演的角色不仅是降维，更是一种概念提炼。它将原子坐标和化学键这些表面特征，转化为捕捉分子本质的抽象概念。在这个抽象层次上，扩散模型能更高效地学习和泛化。

对称性与约束

物理学告诉我们，自然界的对称性不是装饰，而是基本规律的体现。在机器学习中硬编码这些对称性，不是限制模型自由度，而是注入先验知识，减少需要从数据中学习的冗余自由度。

LDM-3DG 通过架构设计实现对称性保持，而非通过数据增强或正则化惩罚。这种结构性解决方案比统计性解决方案更优雅也更可靠，因为它从根本上消除了违反对称性的可能性。

泛化的本质

自监督学习改善OOD性能的现象值得深思。传统观点认为，泛化来自大量数据的统计平均。但 LDM-3DG 的经验表明，结构化的潜空间本身就是一种泛化机制。

当潜空间被组织成语义有序的流形时，插值和外推变得有意义——在两个已知分子之间的中间点，对应的是性质介于两者之间的合理分子，而不是随机的垃圾。这种几何结构蕴含了归纳偏置，使得模型能在有限数据上学到可外推的规律。

七、局限与展望

尽管 LDM-3DG 在多个方面取得突破，仍有改进空间。几何重建的RMSE虽然可接受，但对于大型复杂分子仍有提升余地。这可能需要更强大的等变神经网络架构，或者引入物理驱动的正则化项。

潜空间的显式正则化仍依赖自监督学习这一通用策略。未来可以探索更具体的先验——例如，利用量子化学计算的能量面来引导潜空间几何，或者使用药效团约束来组织拓扑潜空间。

从应用角度看，将 LDM-3DG 扩展到蛋白质设计、材料发现等更广泛的3D结构生成任务是自然的下一步。核心思想——在对称性保持的低维潜空间中扩散——具有普适性，只需根据具体对称群调整自编码器设计即可。

值得思考的问题

Q1: 为什么降维不会损失关键信息？

这涉及流形假设：高维数据实际分布在低维流形上。降维不是随机丢弃维度，而是学习这个流形的参数化。好比地球表面是2维的，虽然嵌入在3维空间中，但我们用经纬度两个参数就能准确定位。关键是降维方式要尊重数据的内在结构，这正是自编码器通过端到端学习实现的。

Q2: 潜空间扩散能否应用于其他类型的图结构？

核心思想是可迁移的，但具体实现需要根据对称性调整。社交网络没有SE(3)对称性，但有置换对称；蛋白质骨架有扭转角度的周期性；晶体结构有平移和点群对称。每种情况都需要设计相应的等变/不变编码器。挑战在于，复杂对称群的神经网络实现仍是开放问题。

Q3: 如何理解自监督学习改善OOD性能？

自监督学习通过对比目标，将潜空间组织成语义聚类：相似分子靠近，不同分子远离。这种几何结构使得插值有意义。当面对OOD属性值时，模型在潜空间中沿着已知点的连线外推，而不是随机游走。这就像在已知数据点之间拟合了一个平滑的流形，新点沿流形外推，而非跳到任意位置。

Q4: 级联设计的计算开销如何？

虽然需要训练两个自编码器，但它们是独立训练的，可以并行。总训练时间与单阶段模型相当，但收敛更快、质量更高。关键优势在推理阶段：扩散在低维空间进行，每步更新成本降低一个数量级。解码虽然额外消耗时间，但只需一次，不影响扩散的迭代效率。整体推理速度反而更快。

参考文献: You Y, Zhou R, Park J, et al. Latent 3d graph diffusion[C]. International Conference on Learning Representations (ICLR), 2024.

代码数据: https://github.com/GRAPH-0/LDM-3DG

数据集: QM9, Drugs, CrossDocked2020

欢迎在评论区分享您对潜空间生成模型的看法或在分子设计中的应用经验。如果您对AI药物发现和3D分子生成感兴趣，欢迎关注 MindDance 公众号获取更多前沿内容，也欢迎将文章分享给您的同行。