Enhancing Atomic Systems Energy Prediction with Graph Neural Networks

Date of Award


Document Type


Degree Name

Master of Science in Machine Learning


Machine Learning

First Advisor

Dr. Martin Takac

Second Advisor

Dr. Bin Gu


This thesis investigates the application of Graph Neural Networks (GNNs) with inte-grated attention mechanisms to enhance predictions of energy and forces in atomic sys- tems, focusing on the Polarizable Atom Interaction Neural Network (PaiNN) model [20]. By incorporating novel parameters that account for atomic distances and bond types, our approach aims to refine the model’s accuracy and interpretability in predicting molecular dynamics and properties. Utilizing the Open Catalyst Dataset (OC20), we explore several PaiNN configurations that include Principal Neighborhood Aggregation (PNA) [2] for dynamic attention-based interaction processing. Our results show that the attention- augmented PaiNN model outperforms baseline GNN models in predicting scalar and tensorial properties of molecules, resulting in a significant increase in data efficiency and predictive accuracy. This study combines computational chemistry and machine learning to present a novel way for modeling molecule energy estimates. The incorporation of attention processes into GNNs marks a significant advancement in molecular simulations, paving the path for novel applications in material science, drug development, and beyond.


Thesis submitted to the Deanship of Graduate and Postdoctoral Studies

In partial fulfilment of the requirements for the M.Sc degree in Machine Learning

Advisors: Dr. Martin Takac, Bin Gu

Online access available for MBZUAI patrons