GemNet: Universal Directional Graph Neural Networks for Molecules
This work addresses the problem of accurately predicting molecular interactions for chemical simulations, offering a novel theoretical foundation and practical gains, though it builds incrementally on existing GNN methods.
The paper tackled the theoretical limitations of graph neural networks (GNNs) in distinguishing certain graphs for molecular interactions by proposing GemNet, a universal directional GNN that achieved performance improvements of 34%, 41%, and 20% on COLL, MD17, and OC20 datasets, respectively.
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.