基于注意力等变几何扩散模型的完整分子生成

Complete Molecular Generation Based on Attention-Equivariant Geometric Diffusion Models

摘要

针对分子生成任务，本文提出了一种注意力等变几何扩散模型（AEGDM），该模型应用离散扩散与连续扩散混合生成扩散框架，实现原子坐标、原子种类及化学键的完整生成。在 AEGDM 模型中，注意力等变图神经网络（Attention Equivariant Graph Neural Network，AEGNN）与残差增强神经网络（Residual Boosting Neural Network，RBNN）分别用以提高模型的表征能力及加快模型训练速度。在 QM9 数据集上的实验结果表明，该模型在生成分子的唯一性与新颖性方面均取得显著提升，唯一性达到 98.2%，新颖性达到 70.9%。这一性能提升表明 AEGDM 在化学空间探索方面具有较强的表征能力，能够有效推动创新性候选分子的发展。

Abstract

This paper proposes an attention-equivariant diffusion model (comprising AEGNN and RBNN) for molecular generation tasks. The AEGNN leverages multi-head self-attention to jointly update edge features, atomic coordinates, and node features within a molecular graph. Subject to strict rotational and translational equivariance, it progressively reconstructs atom types, three-dimensional structures, and chemical bonds via a reverse diffusion process. The RBNN further improves generative performance through co-training two submodels, both built by stacking identical basic modules (denoted as AEM in this paper). The Knowledge Generator (KG) adopts a deeper layer stack to capture complex nonlinear relationships, producing high-precision initial predictions that align with the ground truth. By contrast, the Residual Refiner (RR) adopts a shallower architecture, focusing on fitting the residual between the KG’s output and the true target values. This design cuts computational overhead while boosting the model’s error correction capacity. The two modules are connected in a cascade structure: the output of the KG acts as the input to the RR, and the final prediction is obtained by adding the residual correction from the RR to the initial output of the KG. Experiments conducted on the QM9 dataset show substantial improvements in the uniqueness and novelty of generated molecules, demonstrating that the AEGDM can efficiently explore chemical space and aid the discovery of structurally innovative candidate molecules. Moreover, the RBNN mechanism speeds up experimental iteration and elevates overall model performance, offering critical technical support for the iterative optimization of molecular generation frameworks.

关键词

扩散模型 / 图神经网络 / 分子生成 / 变压器 / 集成学习

Key words

diffusion model / graph neural network / molecular generation / transformer / ensemble learning

参考文献

[1] CORSO G, STÄRK H, JING B, et al. Diffdock: Diffusion steps, twists, and turns for molecular docking[EB/OL]. (2022-10-04). https://arxiv.org/abs/2210.01776.

[2] 仝建波, 李云飞, 刘淑玲, 等. 一种新型拓扑分子表征方法的应用研究 [J].环境化学, 2010(6):1167-1171.

[3] DE VIVO M, MASETTI M, BOTTEGONI G, et al. Role of molecular dynamics and related methods in drug discovery[J].Journal of Medicinal Chemistry, 2016, 59(9):4035-4061.

[4] MANSIMOV E, MAHMOOD O, KANG S, et al. Molecular geometry prediction using a deep generative graph neural network[J].Scientific Reports, 2019, 9(1):20381.

[5] GILMER J, SCHOENHOLZ S S, RILEY P F, et al. Neural message passing for quantum chemistry[C]//International conference on machine learning. Sydney, AUS: ICML, 2017:1263-1272.

[6] SCHÜTT K, KINDERMANS P J, SAUCEDA FELIX H E, et al. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions[C]//Advances in Neural Information Processing Systems. Long Beach, USA: [s.n.],2017:992-1002.

[7] GASTEIGER J, GROSS J, GUENNEMANN S. Directional Message Passing for Molecular Graphs[EB/OL]. (2020-03-06). https://arxiv.org/abs/2003.03123

[8] FUCHS F, WORRALL D, FISCHER V, et al. SE (3)-transformers: 3D roto-translation equivariant attention networks[J].Advances in Neural Information Processing Systems, 2020, 33:1970-1981.

[9] BILODEAU C, JIN W, JAAKKOLA T, et al. Generative models for molecular discovery: Recent advances and challenges[J].Wiley Interdisciplinary Reviews: Computational Molecular Science, 2022, 12(5):e1608.

[10] KINGMA D P, WELLING M. Auto-encoding variational bayes[EB/OL]. (2013-12-20). https://arxiv.org/abs/1312.6114.

[11] HO J, JAIN A, ABBEEL P. Denoising diffusion probabilistic models[J].Advances in Neural Information Processing Systems, 2020, 33:6840-6851.

[12] DINH L, SOHL-DICKSTEIN J, BENGIO S. Density estimation using real NVP[EB/OL]. (2016-05-27). https://arxiv.org/abs/1605.08803.

[13] XU M, YU L, SONG Y, et al. Geodiff: A geometric diffusion model for molecular conformation generation[EB/OL]. (2022-03-06). https://arxiv.org/abs/2203.02923

[14] SATORRAS V G, HOOGEBOOM E, WELLING M. E(n) equivariant graph neural networks[C]//International Conference on Machine Learning. Vienna, AUT: ICML, 2021:9323-9332.

[15] HOOGEBOOM E, SATORRAS V G, VIGNAC C, et al. Equivariant diffusion for molecule generation in 3d[C]//International Conference on Machine Learning. Maryland, USA: ICML, 2022:8867-8887.

[16] XU M, POWERS A S, DROR R O, et al. Geometric latent diffusion models for 3D molecule generation[C]//International Conference on Machine Learning. Hawaii, USA: ICML, 2023:38592-38610.

[17] XU C, WANG H, WANG W, et al. Geometric-facilitated denoising diffusion model for 3D molecule generation[J].Proceedings of the AAAI Conference on Artificial Intelligence, 2024, 38(1):338-346.

[18] VIGNAC C, OSMAN N, TONI L, et al. MiDi: Mixed graph and 3D denoising diffusion for molecule generation[C]//Joint European Conference on Machine Learning and Knowledge Discovery in Databases. Cham: Springer Nature Switzerland, 2023:560-576.

[19] HUA C, LUAN S, XU M, et al. Mudiff: Unified diffusion for complete molecule generation[EB/OL]. (2023-04-28). https://arxiv.org/abs/2304.14621.

[20] PENG X, GUAN J, LIU Q, et al. MolDiff: Addressing the atom-bond inconsistency problem in 3D molecule diffusion generation[C]//International Conference on Machine Learning. Hawaii, USA: ICML, 2023:27611-27629.

[21] AUSTIN J, JOHNSON D D, HO J, et al. Structured denoising diffusion models in discrete state-spaces[J].Advances in Neural Information Processing Systems, 2021, 34:17981-17993.

[22] CORSO G, CAVALLERI L, BEAINI D, et al. Principal neighbourhood aggregation for graph nets[J].Advances in Neural Information Processing Systems, 2020, 33:13260-13271.

[23] NICHOL A Q, DHARIWAL P. Improved denoising diffusion probabilistic models[C]//International conference on machine learning. Vienna, AUT: ICML, 2021:8162-8171.

[24] ZHOU Z H. Ensemble Methods: Foundations and Algorithms[M]. 2nd ed. New York: Chapman and Hall/CRC,2025.