[论文解读] Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation

Tuan Le, Julian Cremer+3 more|arXiv (Cornell University)|Sep 29, 2023

Machine Learning in Materials Science被引用 6

一句话总结

该论文提出EQGAT-diff，一种新颖的E(3)等变扩散模型，用于3D从头分子生成，通过利用连续原子坐标、分类的原子/键特征以及时间依赖的损失加权策略，在QM9和GEOM-Drugs数据集上实现了最先进性能，且训练和推理速度更快。该模型还通过在PubChem3D上进行预训练展示了良好的迁移能力，并通过引入杂化态等化学信息特征提升了生成分子的化学有效性。

ABSTRACT

Deep generative diffusion models are a promising avenue for 3D de novo molecular design in materials science and drug discovery. However, their utility is still limited by suboptimal performance on large molecular structures and limited training data. To address this gap, we explore the design space of E(3)-equivariant diffusion models, focusing on previously unexplored areas. Our extensive comparative analysis evaluates the interplay between continuous and discrete state spaces. From this investigation, we present the EQGAT-diff model, which consistently outperforms established models for the QM9 and GEOM-Drugs datasets. Significantly, EQGAT-diff takes continuous atom positions, while chemical elements and bond types are categorical and uses time-dependent loss weighting, substantially increasing training convergence, the quality of generated samples, and inference time. We also showcase that including chemically motivated additional features like hybridization states in the diffusion process enhances the validity of generated molecules. To further strengthen the applicability of diffusion models to limited training data, we investigate the transferability of EQGAT-diff trained on the large PubChem3D dataset with implicit hydrogen atoms to target different data distributions. Fine-tuning EQGAT-diff for just a few iterations shows an efficient distribution shift, further improving performance throughout data sets. Finally, we test our model on the Crossdocked data set for structure-based de novo ligand generation, underlining the importance of our findings showing state-of-the-art performance on Vina docking scores.

研究动机与目标

• 为解决现有扩散模型在3D分子生成中的局限性，特别是对大分子性能不佳以及数据可用性有限的问题。
• 系统性地探索E(3)等变扩散模型的设计空间，重点关注连续与离散状态空间以及损失加权策略。
• 通过时间依赖的损失加权和增强的特征工程，提升训练收敛性、推理速度和样本质量。
• 通过在大规模PubChem3D上进行预训练（隐式氢原子），实现在小样本数据集上的高效微调。
• 通过在扩散过程中引入杂化态和芳香性等化学动机特征，提升分子的有效性。

提出的方法

• 提出EQGAT-diff，一种E(3)等变图注意力网络，在去噪扩散框架中建模连续3D原子坐标和分类的原子/键类型。
• 基于信噪比（SNR）采用时间依赖的损失加权，以加速训练收敛并提升样本质量。
• 提出一种混合数据模态方法：连续坐标、离散原子类型和离散键类型，每种模态配备独立的去噪头。
• 采用DDIM采样并结合修改后的离散更新策略，在保持分类特征保真度的同时加速推理。
• 在PubChem3D数据集上进行无条件预训练，隐式处理氢原子，以实现下游任务的高效微调。
• 将杂化态和芳香性等化学上有意义的特征融入输入节点和边特征中，以提升化学有效性。

实验结果

研究问题

• RQ1 在3D分子扩散模型中，时间依赖的损失加权（如基于SNR）与均匀或离散损失加权相比，在训练收敛性和样本质量方面表现如何？
• RQ2 使用连续原子坐标与离散表示对模型性能和推理速度有何影响？
• RQ3 在大规模3D分子数据集（如PubChem3D）上预训练的模型，是否能在小规模且更复杂的数据集上通过极少的微调实现有效适应？
• RQ4 在基于扩散的生成中，杂化态和芳香性等化学信息特征如何影响生成分子的有效性和多样性？
• RQ5 参数化方式的选择（如$x_0$与$x_t$）在多大程度上影响3D分子扩散模型的性能和稳定性？

主要发现

• EQGAT-diff在QM9和GEOM-Drugs上达到最先进性能，QM9上有效率96.8%，成功率96.6%，优于先前模型。
• 采用截断SNR(t)的时间依赖损失加权显著提升训练收敛性和样本质量，当仅使用100个时间步时，推理时间最多减少5倍。
• 使用100个时间步和SNR加权训练的模型，性能优于使用500个时间步且均匀加权的模型，证明了自适应损失调度的有效性。
• 在PubChem3D上预训练后，对Crossdocked等目标数据集进行微调，仅需少量微调迭代即可实现优越性能，支持高效的分布偏移适应。
• 引入杂化态和芳香性特征后，分子有效性和多样性显著提升，GEOM-Drugs上的多样性得分为42.2%，远高于基线模型。
• 当采样步数低于500时，DDIM采样未提升样本质量，但采用100个时间步和SNR加权训练的模型实现了5倍加速推理，同时保持高性能。