[论文解读] MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation

Xingang Peng, Jiaqi Guan+2 more|arXiv (Cornell University)|May 11, 2023

Machine Learning in Materials Science被引用 11

一句话总结

MolDiff 通过在三维分子中联合扩散原子和键，使用一个 E(3)-等变框架并结合键引导，解决原子-键不一致问题，且生成质量高于先前模型。

ABSTRACT

Deep generative models have recently achieved superior performance in 3D molecule generation. Most of them first generate atoms and then add chemical bonds based on the generated atoms in a post-processing manner. However, there might be no corresponding bond solution for the temporally generated atoms as their locations are generated without considering potential bonds. We define this problem as the atom-bond inconsistency problem and claim it is the main reason for current approaches to generating unrealistic 3D molecules. To overcome this problem, we propose a new diffusion model called MolDiff which can generate atoms and bonds simultaneously while still maintaining their consistency by explicitly modeling the dependence between their relationships. We evaluated the generation ability of our proposed model and the quality of the generated molecules using criteria related to both geometry and chemical properties. The empirical studies showed that our model outperforms previous approaches, achieving a three-fold improvement in success rate and generating molecules with significantly better quality.

研究动机与目标

• 激发解决三维分子扩散生成中的原子-键不一致问题。
• 提出一个同时对原子与键进行采样的扩散模型，以保持化学一致性。
• 设计一个 E(3)-等变图神经网络，协同更新原子与键的表示。
• 引入以键为先的噪声调度，以稳定训练与生成。
• 开发能够捕捉生成分子的几何形状、药物相似性及结构合理性的评估指标。

提出的方法

• 将三维分子建模为 M={A,R,B}，其中原子类型为 A，坐标为 R，键为 B。
• 使用前向扩散过程，对原子位置以及离散的原子/键类型施加吸收型噪声来扰动。
• 使用带有 E(3)-等变神经网络的反向扩散，预测 pθ(M^{t-1}|M^{t}) 对于原子和键。
• 将原子类型和键类型视为离散变量，并采用以键为先的扩散调度，将键扩散与原子扩散分开。
• 引入一个键预测器，通过 log C 的梯度引导原子位置去噪，其中 C 为键预测器的置信度。
• 采用一个 E(3)-等变信息传递方案，更新顶点和边的表示以预测原子和键的类型（Eq. 4）。
• 训练时使用损失 L^{t-1}=L_pos^{t-1}+λ1 L_atom^{t-1}+λ2 L_bond^{t-1}，在位置去噪、原子类型去噪与键类型去噪之间取得平衡。

实验结果

研究问题

• RQ1 原子和键的联合扩散是否相对于事后添加键的做法提升了生成三维分子的真实感与有效性？
• RQ2 以键为先的扩散调度是否能减少原子-键不一致并提高生成过程中的化学拓扑？
• RQ3 同时更新原子和键的 E(3)-等变架构如何影响生成质量？
• RQ4 在基于扩散的生成中，哪些指标最能捕捉三维分子的几何形状、药物相似性和结构合理性？

主要发现

有效性	连通性	成功率	新颖性	多样性	唯一性	相似性
0.997	0.996	0.993	0.972	0.769	0.986	0.634
EDM	0.447	0.830	0.371	1.000	0.729	1.000	0.441
Predict bond (Lookup table)	0.297	0.937	0.278	1.000	0.751	0.999	0.434
Predict bond (NN Predictor)	0.956	0.978	0.935	0.984	0.768	0.999	0.592
Predict bond (Openbabel)	0.992	0.969	0.961	0.986	0.762	0.999	0.576
Conti. diffusion	0.914	0.885	0.808	0.999	0.752	1.000	0.513
Conti. diffusion (scaling)	0.972	0.934	0.907	0.997	0.759	1.000	0.511
Add bond length loss	0.998	0.959	0.958	0.986	0.752	0.998	0.590

• MolDiff 在生成性能上显著优于 EDM，实现更高的有效性、连通性，以及接近完美的成功率。
• 原子/键类型的离散扩散在多种基线上优于连续扩散和事后键分配。
• 以键为先的扩散调度通过先让键类型扩散到先验再进行原子扩散来稳定训练，从而在后续原子细化时改善键的一致性。
• 在原子位置去噪过程中引入键预测器引导的梯度，提升键长精度和整体分子几何。
• 相较于 EDM，MolDiff 提供更优的三维几何指标（更低的 RMSD，更好的键角/二面角分布）以及更逼真的键类型和环结构统计。
• MolDiff 在成功率上较基线 EDM 提高了三倍，并生成具有改进药物相似性与结构质量的分子。