GemNet: Universal Directional Graph Neural Networks for Molecules

TL;DR

GemNet advances molecular GNNs by proving universality of spherical representations for rotationally invariant and equivariant predictions, then deriving a practical edge-based directed two-hop message passing scheme. It introduces geometric message passing with symmetric updates and an efficient bilinear layer, culminating in GemNet with strong empirical gains on COLL, MD17, and OC20 (up to 34%, 41%, and 20% improvements, respectively). The method faithfully captures full geometric information—distances, angles, and dihedrals—and provides scalable variants (GemNet-Q and GemNet-T) along with direct force prediction options. Overall, GemNet delivers substantial accuracy improvements for molecular dynamics tasks while maintaining competitive computational costs and offering broad applicability to challenging, non-planar molecular geometries.

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.

Figures (4)

• Figure 1: Angles used in geometric message passing. The dihedral angle $\theta_{cabd}$ becomes visible when rotating the molecule so that atoms $a$ and $b$ lie on top of each other (right).
• Figure 2: The GemNet architecture (comprehensive version in \ref{['app:gemnet']}). $\square$ denotes the layer's input, $\|$ concatenation, and $\sigma$ a non-linearity. Directional embeddings ${\bm{m}}_{ca}$ are updated using three forms of interaction: Two-hop geometric message passing (Q-MP), one-hop geometric message passing (T-MP), and atom self-interactions. Differences between Q-MP and T-MP are denoted by colors and dashed lines.
• Figure 3: Layer-wise activation variance. GemNet's variance varies strongly between layers and increases significantly after each block without scaling factors (top). Introducing scaling factors successfully stabilizes the variance (bottom).
• Figure 4: The full GemNet architecture. $\square$ denotes the layer's input, $\|$ concatenation, $\sigma$ a non-linearity (we use SiLU in this work elfwing_sigmoid-weighted_2018), and orange a layer with weights shared across interaction blocks. Differences between two-hop message passing (Q-MP) and one-hop message passing (T-MP) are denoted by dashed lines. Numbers next to connecting lines denote embedding sizes.

Theorems & Definitions (7)

• Theorem 1: dym_universality_2021
• Theorem 2
• Theorem 3
• Lemma 1
• Proposition A: villar_scalars_2021
• Lemma A
• Lemma B