[论文解读] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation

Clément Vignac, Nagham Osman+2 more|arXiv (Cornell University)|Feb 17, 2023

Machine Learning in Materials Science被引用 10

一句话总结

MiDi 通过端到端扩散联合生成分子图及其 3D 构象，相较于现有方法，在 GEOM-DRUGS 上达到高比例的稳定、有效分子。

ABSTRACT

This work introduces MiDi, a novel diffusion model for jointly generating molecular graphs and their corresponding 3D arrangement of atoms. Unlike existing methods that rely on predefined rules to determine molecular bonds based on the 3D conformation, MiDi offers an end-to-end differentiable approach that streamlines the molecule generation process. Our experimental results demonstrate the effectiveness of this approach. On the challenging GEOM-DRUGS dataset, MiDi generates 92% of stable molecules, against 6% for the previous EDM model that uses interatomic distances for bond prediction, and 40% using EDM followed by an algorithm that directly optimize bond orders for validity. Our code is available at github.com/cvignac/MiDi.

研究动机与目标

• 激发端到端生成能够同时捕捉图结构和 3D 构象的分子。
• 提出一个扩散模型，同时对图和坐标进行去噪。
• 通过新颖的 rEGNN 增强 Transformer 架构，提升 3D 等变处理。
• 在 GEOM-DRUGS 和 QM9 数据集上证明无条件生成性能的提升。

提出的方法

• 将分子表示为嵌入在 3D 中的图，具有节点特征（原子类型、电荷）和边特征（键类型）。
• 使用混合高斯（坐标）和离散（原子类型、电荷、键）扩散过程，并配以分量特定的噪声时间表。
• 引入自适应噪声时间表，在原子类型和电荷之前优先考虑坐标和键类型。
• 开发一个放松的等变图神经网络（rEGNN）层，扩展 EGNN 以包含非平移不变特征。
• 将 rEGNN 集成到基于 Transformer 的去噪网络中，保持 SE(3) 等变性。
• 训练时使用一个损失函数，结合坐标回归和离散特征的交叉熵。

Figure 1 : Samples from our model. MiDi generates simultaneously a 2D graph structure and a 3D conformation that is consistent with this structure.

实验结果

研究问题

• RQ1 扩散模型是否能够以端到端可微的方式同时学习分子图和 3D 坐标？
• RQ2 自适应分量级别噪声时间表是否改善生成分子的质量和稳定性？
• RQ3 放松的 SE(3) 等变架构（rEGNN）是否优于标准 EGNN 在图和构象生成方面的表现？
• RQ4 在 QM9 和 GEOM-DRUGS 上，与基于 OpenBabel 的后处理以及先前的扩散方法相比，MiDi 在无条件分子生成方面的表现如何？

主要发现

• MiDi 在 GEOM-DRUGS 上实现了高比例的稳定分子，超过基于 EDM 的方法和 Open Babel 后处理。 (Adaptive MiDi: Mol stable 91.6%, At stable 99.8%, Validity 77.8%, Connected 100.0%).
• 在 QM9 上，MiDi 在基于图的指标上优于基础 EDM 模型，尽管 Open Babel 优化在简单数据集上仍然表现强劲。
• MiDi 实现了对 2D 图结构和 3D 构象的端到端可微生成，避免了单独的键预测步骤。
• 自适应噪声时间表提高了样本质量，引导模型在原子类型和电荷之前优先考虑坐标和键类型。
• MiDi 可扩展到大型药物样数据集（GEOM-DRUGS），并在端到端性能方面优于依赖事后键预测的方法。

Figure 2 : The noise schedule is tuned separately for each component. Atom coordinates and bond types are denoised earlier during sampling, while atom types and formal charges are updated later in the process. Experimentally, the adaptive schedule allows to obtain better 3D conformers and more stabl