[论文解读] Equivariant Diffusion for Molecule Generation in 3D

Emiel Hoogeboom, Victor García Satorras+2 more|arXiv (Cornell University)|Mar 31, 2022

Machine Learning in Materials Science被引用 118

一句话总结

介绍一个 E(3) 具等变扩散模型（EDM），它联合对3D坐标和原子类型特征去噪以生成分子，具有概率似然框架，在先前的3D分子生成器上取得显著经验提升。

ABSTRACT

This work introduces a diffusion model for molecule generation in 3D that is equivariant to Euclidean transformations. Our E(3) Equivariant Diffusion Model (EDM) learns to denoise a diffusion process with an equivariant network that jointly operates on both continuous (atom coordinates) and categorical features (atom types). In addition, we provide a probabilistic analysis which admits likelihood computation of molecules using our model. Experimentally, the proposed method significantly outperforms previous 3D molecular generative methods regarding the quality of generated samples and efficiency at training time.

研究动机与目标

• 说明需要尊重欧几里得对称性的3D分子生成（E(3)）。
• 提出一个 EDM，在扩散框架中联合建模连续坐标和离散原子特征。
• 使在所提出模型下生成分子的似然性计算可行。
• 在质量和效率方面展示相对于先前的3D分子生成方法的经验改进。

提出的方法

• 定义在坐标 x 和特征 h 上操作的等变去噪扩散模型。
• 使用基于 EGNN 的动力学 phi 以等变方式预测去噪噪声。
• 提出带有 L_t 项的似然目标以及用于预测 epsilon 的实用噪声参数化。
• 通过独热编码处理离散原子类型并为 x 和 h 引入适当的似然项。
• 为 x 引入零均值的重力约束并对特征进行尺度化以稳定训练。
• 在去噪网络输入中加入属性 c 以实现有条件生成。

实验结果

研究问题

• RQ1 一个 E(3) 具等变扩散模型是否能够在不固定原子顺序的情况下生成3D分子构象和原子类型？
• RQ2 在 QM9 上，EDM 是否在对数似然和分子稳定性方面优于先前的等变方法（如 E-NF、G-Schnet）？
• RQ3 该模型是否可以高效训练并扩展到更大、药物级数据集，同时提供可行的似然性？
• RQ4 离散原子特征（类型和电荷）如何整合入与连续坐标的联合扩散过程？
• RQ5 对分子性质的条件化对定向分子生成的影响如何？

主要发现

指标	E-NF	G-Schnet	GDM	GDM-aug	EDM（ours）
NLL	-59.7	N.A	-94.7	-92.5	-110.7 ±1.5
Atom stable (%)	85.0	95.7	97.0	97.6	98.7 ±0.1
Mol stable (%)	4.9	68.1	63.2	71.6	82.0 ±0.4

• EDM 在负对数似然（-110.7 ± 1.5）方面优于基线（E-NF -59.7，G-Schnet N.A.，GDM -94.7，GDM-aug -92.5）。
• 原子稳定性对 EDM 为 98.7 ± 0.1% 高于基线（85.0%、95.7%、97.0%、97.6%）。
• 分子稳定性对 EDM 为 82.0 ± 0.4% 高于基线（4.9%、68.1%、63.2%、71.6%）。
• EDM 在3D生成的有效且唯一分子数量上优于竞争模型（表2结果：97.5 ± 0.2% 有效，94.3 ± 0.2% 有效且唯一，包含 H 时）。
• EDM 可以更高效地训练并扩展到如 GEOM-Drugs 这样的更大数据集，优于先前的3D扩散和正则化流方法。