GemNet: Universal Directional Graph Neural Networks for Molecules

Johannes Gasteiger,Florian Becker,Stephan Günnemann
2024-06-22
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.
Computational Physics,Machine Learning,Chemical Physics
What problem does this paper attempt to address?
The paper aims to address the problem of molecular interaction prediction to accelerate molecular dynamics simulations and revolutionize chemical simulation techniques. Specifically, the paper focuses on improving the performance of Graph Neural Networks (GNNs) in predicting molecular energy and other quantum mechanical properties. The main contributions of the paper are as follows: 1. **Gap between theory and practice**: Conventional GNNs have theoretical limitations that prevent them from distinguishing certain types of graph structures, which limits their performance in practical applications. This paper addresses this gap by introducing a spherical representation method and demonstrates that this representation is a universal approximator for tasks involving translation invariance, permutation equivariance, and rotation equivariance. 2. **GemNet model**: The authors propose a Geometric Message Passing Neural Network (GemNet) that captures geometric structural information of molecules through directional edge embeddings and a two-hop message passing mechanism. GemNet achieves significantly better results than previous models on multiple datasets, particularly excelling in handling complex molecules with dynamic, non-planar geometries. 3. **Theoretical proof**: The authors prove that the spherical representation method is universal for rotation equivariant predictions and show how these theoretical models can be discretized into a two-hop message passing form. Additionally, they prove that using only directional information without relying on SO(3) group representations can achieve the same expressive power. 4. **Structural improvements**: To enhance the practical performance of GemNet, the authors propose a series of structural enhancements, including variance stabilization of activation functions, efficient bilinear layers, and symmetric message passing schemes. These improvements have been validated in multiple ablation studies to significantly reduce model error. 5. **Experimental results**: GemNet outperforms existing methods on several benchmark datasets such as COLL, MD17, and OC20, achieving performance improvements of 34%, 41%, and 20% respectively in force prediction tasks. In summary, through theoretical analysis and model innovation, this paper makes significant progress in the field of molecular interaction prediction, providing an effective solution for the accurate simulation of complex molecular structures.