[论文解读] Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation

Han Huang, Leilei Sun+2 more|arXiv (Cornell University)|May 21, 2023

Computational Drug Discovery Methods被引用 5

一句话总结

该论文提出JODO，一种联合2D与3D的扩散模型，通过引入一种新型扩散图Transformer来建模分子图与几何结构之间的相关性，从而生成包含原子类型、形式电荷、键类型和3D坐标在内的完整分子。该模型在QM9和GEOM-Drugs数据集上的无条件与条件生成任务中均达到最先进性能，优于基线模型，在有效性、FCD和骨架相似性方面表现更优，同时支持快速的少步采样与逆向分子设计。

ABSTRACT

Designing new molecules is essential for drug discovery and material science. Recently, deep generative models that aim to model molecule distribution have made promising progress in narrowing down the chemical research space and generating high-fidelity molecules. However, current generative models only focus on modeling either 2D bonding graphs or 3D geometries, which are two complementary descriptors for molecules. The lack of ability to jointly model both limits the improvement of generation quality and further downstream applications. In this paper, we propose a new joint 2D and 3D diffusion model (JODO) that generates complete molecules with atom types, formal charges, bond information, and 3D coordinates. To capture the correlation between molecular graphs and geometries in the diffusion process, we develop a Diffusion Graph Transformer to parameterize the data prediction model that recovers the original data from noisy data. The Diffusion Graph Transformer interacts node and edge representations based on our relational attention mechanism, while simultaneously propagating and updating scalar features and geometric vectors. Our model can also be extended for inverse molecular design targeting single or multiple quantum properties. In our comprehensive evaluation pipeline for unconditional joint generation, the results of the experiment show that JODO remarkably outperforms the baselines on the QM9 and GEOM-Drugs datasets. Furthermore, our model excels in few-step fast sampling, as well as in inverse molecule design and molecular graph generation. Our code is provided in https://github.com/GRAPH-0/JODO.

研究动机与目标

• 为解决现有生成模型仅孤立地建模2D分子图或3D几何结构的局限性，该局限性会损害生成质量并限制下游应用。
• 开发一种统一的生成框架，联合建模2D键合图与3D坐标，以捕捉互补的分子描述符。
• 通过学习分子结构与几何的联合分布，提升样本保真度并实现高效的逆向分子设计。
• 设计一种可微分的连续时间扩散过程，保持几何对称性（SE(3)-不变性）和图排列不变性。
• 实现快速、高质量的采样与量子性质条件生成，无需依赖后处理或额外模拟。

提出的方法

• 提出一种连续时间扩散模型，采用统一的噪声调度，同时对2D图和3D坐标分量进行破坏。
• 引入扩散图Transformer，通过关系注意力机制联合更新节点与边表示，整合标量特征与几何向量。
• 采用数据预测模型（而非噪声预测），并引入自条件机制，利用最后一步采样输出作为辅助条件，以提升离散结构恢复能力。
• 通过消息传递架构参数化数据预测模型，实现在分子图上传播并更新标量与3D向量特征。
• 通过等变与不变神经网络设计，强制在3D几何生成中保持SE(3)-不变性，在2D图学习中保持排列不变性。
• 通过将生成过程条件化于目标量子性质，将模型扩展至逆向设计，实现性质导向的分子生成。

实验结果

研究问题

• RQ1 联合2D与3D的扩散模型是否能比仅生成2D图或3D坐标的模型生成更高保真度的完整分子？
• RQ2 统一的扩散框架在反向去噪过程中，能否有效建模分子拓扑与3D几何之间的相关性？
• RQ3 所提出的具有关系注意力与几何向量更新的扩散图Transformer是否能提升生成质量与采样效率？
• RQ4 该模型能否在无条件与条件（性质导向）分子生成任务中均达到最先进性能？
• RQ5 与现有基线相比，该模型在少步采样与量子性质逆向设计中的表现如何？

主要发现

• 在QM9数据集上，JODO-2D的生成有效性达99.91%（未进行价键检查），优于所有基线模型，包括CDGS（98.13%）与DiGress（92.10%）。
• 在QM9数据集上，JODO-2D的FCD得分为0.472，显著低于CDGS（2.069）与DiGress（3.597），表明其分布匹配能力更优。
• 在MOSES数据集上，JODO-2D的生成有效性为88.9%，唯一性为100%，新颖性为91.0%，优于DiGress（85.7%有效性，95.0%新颖性），并匹配或超越其他强基线模型。
• 该模型在逆向分子设计方面表现优异，能够有效生成具有期望量子性质的分子。
• JODO支持快速的少步采样，输出质量高，表明其在实际应用中具有实用价值。
• 消融实验证实，所提出的带自条件机制的数据预测模型与扩散图Transformer显著提升了生成质量与结构保真度。