GemNet: Universal Directional Graph Neural Networks for Molecules

TL;DR

GemNet introduces a universal directional graph neural network that is theoretically capable of modeling complex molecular interactions and demonstrates superior performance on multiple molecular datasets.

Contribution

The paper develops GemNet, a GNN with spherical representations and structural improvements, bridging the gap between theory and practice in molecular modeling.

Findings

GemNet outperforms previous models on COLL, MD17, and OC20 datasets.

GemNet achieves 34%, 41%, and 20% improvements respectively.

The model performs especially well on challenging molecules.

Abstract

Effectively predicting molecular interactions has the potential to accelerate molecular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. However, they still appear very limited from a theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with spherical representations are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then discretize such GNNs via directed edge embeddings and two-hop message passing, and incorporate multiple structural improvements to arrive at the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms previous models on the COLL, MD17, and OC20 datasets by 34%, 41%, and 20%, respectively, and performs especially well on the most challenging molecules. Our implementation is available online.