AGDIFF: Attention-Enhanced Diffusion for Molecular Geometry Prediction

Abstract

Accurate prediction of molecular geometries is crucial for drug discovery and materials science. Existing fast conformer prediction algorithms often rely on approximate empirical energy functions, resulting in low accuracy. More accurate methods like ab initio molecular dynamics and Markov chain Monte Carlo can be computationally expensive due to the need for evaluating quantum mechanical energy functions. To address this, we introduce AGDIFF, a novel machine learning framework that utilizes diffusion models for efficient and accurate molecular structure prediction. AGDIFF extends previous models (such as GeoDiff) by enhancing the global, local, and edge encoders with attention mechanisms, an improved SchNet architecture, batch normalization, and feature expansion techniques. AGDIFF outperforms GeoDiff on both the GEOM-QM9 and GEOM-Drugs data sets. For GEOM-QM9, with a threshold (δ) of 0.5 Å, AGDIFF achieves a mean COV-R of 93.08% and a mean MAT-R of 0.1965 Å. On the more complex GEOM-Drugs data set, using δ = 1.25 Å, AGDIFF attains a median COV-R of 100.00% and a mean MAT-R of 0.8237 Å. These findings demonstrate AGDIFF’s potential to advance molecular modeling techniques, enabling more efficient and accurate prediction of molecular geometries, thus contributing to computational chemistry, drug discovery, and materials design. https://github.com/ADicksonLab/AGDIFF