[论文解读] Geometry-Complete Diffusion for 3D Molecule Generation and Optimization

Alex Morehead, Jianlin Cheng|arXiv (Cornell University)|Feb 8, 2023

Analytical Chemistry and Chromatography被引用 7

一句话总结

本文提出几何完备扩散模型（GCDM），一种新颖的3D分子生成框架，利用SE(3)-等变图神经网络与几何完备的消息传递机制，在去噪过程中学习几何与物理先验。GCDM在QM9和GEOM-Drugs数据集上达到最先进性能，能够生成稳定且逼真的大分子，并实现对现有分子的直接优化以满足目标性质，而无需微调或重新训练。

ABSTRACT

Denoising diffusion probabilistic models (DDPMs) have pioneered new state-of-the-art results in disciplines such as computer vision and computational biology for diverse tasks ranging from text-guided image generation to structure-guided protein design. Along this latter line of research, methods have recently been proposed for generating 3D molecules using equivariant graph neural networks (GNNs) within a DDPM framework. However, such methods are unable to learn important geometric properties of 3D molecules, as they adopt molecule-agnostic and non-geometric GNNs as their 3D graph denoising networks, which notably hinders their ability to generate valid large 3D molecules. In this work, we address these gaps by introducing the Geometry-Complete Diffusion Model (GCDM) for 3D molecule generation, which outperforms existing 3D molecular diffusion models by significant margins across conditional and unconditional settings for the QM9 dataset and the larger GEOM-Drugs dataset, respectively, and generates more novel and unique unconditional 3D molecules for the QM9 dataset compared to previous methods. Importantly, we demonstrate that the geometry-complete denoising process of GCDM learned for 3D molecule generation enables the model to generate a significant proportion of valid and energetically-stable large molecules at the scale of GEOM-Drugs, whereas previous methods fail to do so with the features they learn. Additionally, we show that extensions of GCDM can not only effectively design 3D molecules for specific protein pockets but also that GCDM's geometric features can be repurposed to consistently optimize the geometry and chemical composition of existing 3D molecules for molecular stability and property specificity, demonstrating new versatility of molecular diffusion models. Our source code and data are freely available at https://github.com/BioinfoMachineLearning/Bio-Diffusion.

研究动机与目标

• 解决现有3D分子扩散模型因采用非几何、分子无关的GNN而无法学习关键几何与物理性质的局限性。
• 开发一种可有效扩展至大尺寸3D分子（如GEOM-Drugs数据集中分子）的生成模型，解决以往模型在该尺度下的困难。
• 探究几何扩散模型是否可被重新用于分子优化任务，而无需微调或重新训练。
• 通过将手性等物理归纳偏置与不变注意力机制整合进扩散框架，建立3D分子生成的新最先进水平。

提出的方法

• GCDM采用基于GCPNet++的几何完备消息传递框架，对带有向量特征的3D分子图进行去噪，确保SE(3)等变性。
• 模型采用去噪扩散概率框架，通过迭代去除噪声从噪声3D分子图中生成逼真结构。
• 通过引入如不变图注意力和分子手性等物理归纳偏置，提升生成分子的几何保真度与有效性。
• 训练目标为标准DDPM损失，采用加权MSE衡量预测噪声与真实噪声之间的差异，其中w(t) = 1以实现最优生成效果。
• 通过统一的图去噪过程，模型实现无条件与条件3D分子生成，预测原子类型、电荷与3D坐标。
• GCDM通过利用学习到的几何特征，无需重新训练即可迭代优化现有分子以提升特定性质，实现直接分子优化。

Figure 1: A framework overview for our proposed Geometry-Complete Diffusion Model ( GCDM ). Our framework consists of ( i. ) a graph (topology) definition process, ( ii. ) a GCPNet -based graph neural network for 3D graph representation learning, ( iii. ) denoising of 3D input graphs using GCPNet++

实验结果

研究问题

• RQ1 结合SE(3)-等变、几何完备消息传递的扩散模型是否能比非几何模型生成更真实、更稳定的3D分子？
• RQ2 引入手性与不变注意力等物理归纳偏置是否能提升生成3D分子的有效性与质量？
• RQ3 单一预训练的几何扩散模型是否可有效复用于分子优化任务，而无需微调或重新训练？
• RQ4 与以往在该尺度上失败的模型相比，GCDM在大尺寸3D分子（如GEOM-Drugs数据集中的分子）上的表现如何？
• RQ5 GCDM是否能在蛋白质口袋条件下的3D分子生成任务及直接分子优化任务中超越非几何扩散模型？

主要发现

• 在QM9数据集上，GCDM在无条件与条件3D分子生成任务中均达到最先进性能，其在有效性、稳定性与性质特异性方面优于先前方法。
• 在GEOM-Drugs数据集上，GCDM成功生成了逼真且稳定的大型3D分子，而以往方法因几何学习不足而无法实现此目标。
• GCDM首次实现无需重新训练即可对现有3D分子进行特定性质的优化，这是几何扩散模型中的开创性成果。
• 该模型的几何完备去噪过程使其能够学习并保留关键物理与几何性质，如手性与键角。
• 在蛋白质口袋条件下的3D分子生成任务中，GCDM表现具有竞争力，展现出在药物设计中的强大潜力。
• GCPNet++中采用的SE(3)-等变、基于向量的消息传递相比以往模型的仅标量消息传递，显著提升了几何泛化能力。