[Paper Review] GemNet: Universal Directional Graph Neural Networks for Molecules

Johannes Gasteiger, Florian Becker+1 more|arXiv (Cornell University)|Jun 2, 2021

Computational Drug Discovery MethodsComputer Science127 citations

TL;DR

GemNet establishes universality of spherical representations for rotationally invariant/equivariant molecular predictions, introduces geometric two-hop message passing with directional edge embeddings, and achieves state-of-the-art force prediction on multiple MD datasets.

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.

Motivation & Objective

• Address expressiveness gaps in regular GNNs for molecular graphs and enable rotationally invariant and equivariant predictions.
• Prove universality of spherical representations for rotationally invariant and equivariant tasks.
• Develop a practical, directional, edge-based two-hop message passing framework with geometric information to improve accuracy and sample efficiency.

Proposed method

• Prove universality of spherical representations for SO(3)-invariant and permutation-equivariant functions.
• Discretize spherical representations into directed edge embeddings and model two-hop message passing between edges.
• Introduce geometric message passing with interatomic directions and angles (including φ and θ) using spherical Fourier-Bessel bases.
• Incorporate symmetric message passing and efficient bilinear layers to improve stability and computational efficiency.
• Stabilize activation variance with predetermined scaling factors instead of standard normalization.
• Develop GemNet architecture with three interaction forms (two-hop geometric message passing, one-hop geometric message passing, and atom self-interactions).

Experimental results

Research questions

• RQ1 Can spherical representations achieve universality for rotationally invariant and equivariant predictions in molecular GNNs?
• RQ2 How can directional (edge-based) representations be discretized to realize a practical, expressive GNN for molecules?
• RQ3 What architectural innovations (symmetric messaging, bilinear layers, variance stabilization) most improve predictive accuracy for energies and forces in molecular dynamics datasets?
• RQ4 How does GemNet compare to state-of-the-art models on COLL, MD17, and OC20 in force and energy predictions?

Key findings

• GemNet achieves substantial error reductions across datasets, outperforming prior models on COLL, MD17, and OC20 by notable margins.
• Two-hop geometric message passing with directional edge embeddings significantly improves performance, especially on challenging, non-planar molecules.
• Symmetric message passing, bilinear layers, and activation variance stabilization contribute meaningfully to accuracy with ablation studies confirming each component’s benefit.
• Direct force prediction can accelerate training and inference, with trade-offs in MAE depending on dataset and task.
• GemNet-T and GemNet-Q variants offer different trade-offs between computational cost and accuracy, with GemNet-Q providing larger gains on more complex combined datasets.