[论文解读] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

Minkai Xu, Lantao Yu+4 more|arXiv (Cornell University)|Mar 6, 2022

Machine Learning in Materials Science被引用 173

一句话总结

GeoDiff 直接用一个 SE(3)-等变扩散框架对原子坐标进行建模，以生成分子构象，在构象生成任务上达到最先进的结果，尤其是在大分子上。

ABSTRACT

Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative models. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to a noise distribution, in this paper, we propose a novel generative model named GeoDiff for molecular conformation prediction. GeoDiff treats each atom as a particle and learns to directly reverse the diffusion process (i.e., transforming from a noise distribution to stable conformations) as a Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be roto-translational invariant. We theoretically show that Markov chains evolving with equivariant Markov kernels can induce an invariant distribution by design, and further propose building blocks for the Markov kernels to preserve the desirable equivariance property. The whole framework can be efficiently trained in an end-to-end fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments on multiple benchmarks show that GeoDiff is superior or comparable to existing state-of-the-art approaches, especially on large molecules.

研究动机与目标

• 激发并开发一个扩散式生成模型，直接输出分子的 3D 原子坐标。
• 通过 SE(3) 等变设计以及无质心的等不变量初始化，确保旋转平移不变性。
• 提供一个端到端可训练框架，使用为给定分子图的条件生成定制的 ELBO 目标。

提出的方法

• 将 3D 坐标 C^t 的前向扩散公式化，采用固定方差日程，从 C^T ~ N(0, I) 开始。
• 定义一个条件反向马尔可夫链 p_θ(C^{t-1}|G, C^t)，在对分子图 G 条件下逐步降噪。
• 使用 SE(3) 等变神经网络（Graph Field Network）参数化反向均值 μ_θ，确保旋转平移等变性。
• 使用无质心的等不变量初始密度和无质心的马尔可夫核以保证平移不变性；通过 μ_θ 与网络结构强制旋转等变性。
• 通过加权 ELBO 训练，简化为 ε 与 ε_θ 之间的加权 L2 损失，提供两种强制等变性的策略（对齐基和链式规则方法）。
• 提供一个采样算法，从 C^T ~ N(0, I) 开始，迭代地使用 p_θ(C^{t-1}|G, C^t) 采样 C^{t-1}。

实验结果

研究问题

• RQ1 直接在 3D 坐标上运作的扩散模型，是否能在构象体生成方面达到甚至超过依赖不变量几何中介的模型？
• RQ2 如何在训练和推理阶段确保几何生成的旋转平移不变性？
• RQ3 哪些有效的参数化与架构（如 SE(3)-等变网络）适用于建模分子构象的反向扩散步骤？
• RQ4 端到端的基于坐标的扩散生成是否可扩展到更大、药物样分子，并提供多样性高、质量高的构象？

主要发现

• GeoDiff 在 GEOM-Drugs 上持续超越最先进的 ML 基线，在 COV-R 和 MAT-R 指标上取得显著提升。
• GeoDiff 对大分子实现更高的 COV-R 更低的 MAT-R，显示出良好的覆盖率和准确性。
• GeoDiff-C 在优化方面通常略好于 GeoDiff-A，表明基于链式法则的训练更好地引导学习。
• 该方法显著提升了生成构象的多样性和质量，特别是在测试大分子药物数据集时。
• 定性样本显示 GeoDiff 能有效捕捉局部与全局结构特征。