PaperlayerView exact revision on arXiv

Source-linked AI summary

GemNet: Universal Directional Graph Neural Networks for Molecules

Johannes Gasteiger, Florian Becker, Stephan Günnemann

arXiv:2106.08903v10physics.comp-phcs.LGphysics.chem-phstat.ML

TL;DR

Molecular GNNs are powerful but theoretically limited because regular architectures cannot distinguish certain molecular graphs. The paper proves spherical representations universal for the relevant symmetries, discretizes them into GemNet, and reports improved force prediction across COLL, MD17, and OC20.

  • Problem

    Regular GNNs cannot distinguish certain molecules and require many training samples to achieve good accuracy, despite enabling fast molecular-property prediction.

  • Method

    The paper proves universality for spherical representations, discretizes them with directed edge embeddings and two-hop geometric message passing, and incorporates structural improvements into GemNet.

  • Results

    GemNet substantially improves molecular-dynamics prediction error and outperforms previous models on force predictions across COLL, MD17, and OC20.

  • Takeaways & Limitations

    Spherical representations provide universal approximation guarantees for the relevant molecular symmetries without requiring full SO(3) representations.

  • Takeaways & Limitations

    GemNet is focused on molecular-simulation predictions, and two-hop message passing introduces significant computational overhead despite mitigation with down-projection and GemNet-T.

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses theoretical limitations of regular GNNs by proving universality for spherical representations and discretizing them into geometric message passing. The resulting GemNet combines these ideas with structural improvements and improves molecular force prediction across several benchmarks.

  • Regular GNNs cannot distinguish certain molecules and require many training samples for good accuracy.
  • The paper proves sufficient conditions for universal approximation under translation invariance and permutation and rotation equivariance.
  • Spherical representations preserve relative rotational information while remaining invariant to global rotations.
  • The authors discretize spherical representations using directed edge embeddings and two-hop message passing.
  • GemNet adds structural enhancements and predetermined activation-variance scaling to improve practical performance.
  • 34%, 41%, and 20% average improvements over previous models are reported for force predictions on COLL, MD17, and OC20, respectively.The largest improvements occur for molecules with dynamic, non-planar geometries.

2 Related work

Related work spans machine learning potentials, directional GNNs, expressiveness theory, and equivariant neural networks. GemNet is positioned as a directional model that studies expressiveness and proposes an improved variant.

  • Machine learning potentials: Machine learning potentials progressed from hand-fitted analytical force fields to kernel methods and end-to-end GNNs.
  • Directional GNNs: Directional GNNs encode geometric information explicitly or through angles and dihedral angles to achieve rotational equivariance and invariance.
  • Expressiveness of GNNs: Prior expressiveness research established that standard GNNs are limited by the Weisfeiler-Lehman test and studied higher-order representations as a remedy.
  • Equivariant neural networks: Equivariant neural networks use known symmetries, including SO(3), as foundational design principles for physics-related models.
  • Equivariant neural networks: Universality results for translations, rotations, and permutations characterize the symmetry group relevant to general molecules, apart from reflections.

3 Universality of spherical representations

The paper establishes that spherical representations are sufficient for universal approximation of molecular functions with the relevant translation, rotation, and permutation symmetries. It extends the result from invariant to rotationally equivariant predictions without requiring full SO(3) representations.

  • Preliminaries: Molecular point clouds consist of atomic positions and rotationally invariant features, with functions required to be translation-invariant and permutation- and rotation-equivariant.
  • Tensor field networks: Tensor field networks use equivariant feature-embedding functions followed by pooling functions that aggregate pointwise results into model outputs.
  • Spherical networks: Spherical harmonics and Clebsch-Gordan coefficients connect spherical networks to tensor field networks, establishing equivalent universality results.
  • Spherical networks: Spherical representations replace intermediate SO(3) representations with functions defined on the S2 sphere.
  • Universality results: Every continuous invariant function can be approximated uniformly on compact sets by the spherical-network function class.
  • Universality results: Theorem 3 extends the invariant construction to rotationally equivariant functions under a non-spanning condition on relative vectors.
  • Universality results: S2 representations suffice for universal approximation in R3, eliminating the need for SO(3) representations, spin-weighted harmonics, triplet embeddings, or complex-valued functions.

4 From spherical representations to directional message passing

The paper makes spherical representations tractable by sampling them along directions to neighboring atoms, represented as directed edge embeddings. This yields a two-hop message-passing scheme derived from a model with universal approximation guarantees.

  • Directional representations: Directional representations sample spherical functions in specific directions rather than using spherical harmonics directly.The directions are taken from neighboring atoms, forming an equivariant mesh that avoids aliasing effects from fixed grids.
  • Directional representations: The directional mesh connects spherical representations to molecular graph neural networks through directed edge embeddings.An edge embedding points toward the corresponding neighboring atom direction.
  • Directional representations: A learned filter relates input and output directions through a convolution, improving expressiveness when directional meshes differ.The input is defined only at specific directions, and Dirac deltas simplify the resulting expressions.
  • General filters: The rotationally invariant spherical filter depends only on the angle between the input direction and the edge direction.Its rotational invariance permits replacement by a general learnable filter parameterized by the relative angle without losing expressivity.
  • Directional message passing: The resulting scheme performs two-hop message passing between directional edge embeddings via an intermediate edge.Its filter products resemble Hadamard products in modern GNNs, while retaining two-hop rather than one-hop message passing.

5 Geometric message passing

Geometric message passing uses interatomic directions, distances, and angular information to update directional embeddings. Structural enhancements add symmetric updates, efficient computation, and a bottleneck intended to improve generalization.

  • Geometric representation: GemNet represents directional embeddings for atom pairs within an embedding cutoff and updates them using quadruplets of atoms.Two atoms interact while two additional atoms define the directions used in the update.
  • Geometric representation: The geometric update incorporates three angles: φcab, φabd, and the dihedral angle θcabd.These angles encode relative directional information for updating the embedding mca based on mdb.
  • Geometric representation: Spherical Fourier-Bessel bases with polynomial radial envelopes represent the relative directional information while ensuring smoothly differentiable predictions.The basis is split into three parts to incorporate the available geometric information.
  • Geometric representation: The first weight matrix in each representation part has a small output dimension, creating a bottleneck that improves generalization.The filters are obtained after transforming the basis representations through two linear layers.
  • Symmetric message passing: Symmetric message passing computes one update for both opposing edge embeddings and separates their directions with two learnable weight matrices.This avoids executing the message-passing scheme twice while preserving the distinction between mca and mac.
  • Efficient bilinear layer: Reordering linear summations can reduce memory usage by 50% even for Hadamard products.The optimization applies because basis transformation, neighbor aggregation, and the bilinear layer use linear functions.

6 GemNet: Geometric message passing neural network

GemNet refines directional message passing for molecular energy and force prediction by combining geometric interactions, directed embeddings, and variance stabilization. Its design improves accuracy while exposing a trade-off between two-hop expressiveness and computational cost.

  • Architecture: GemNet predicts molecular energy and forces from atomic positions and atomic numbers, generalizing from COLL to datasets such as MD17 without architectural changes.The architecture is based on DimeNet++ and is designed for molecular simulations.
  • Interactions: Directional embeddings are updated through two-hop geometric message passing, one-hop geometric message passing, and atom self-interactions.The one-hop interaction passes messages between directional embeddings pointing toward the same atom and provides angle-based pair interactions.
  • Variance stabilization: GemNet stabilizes activation variance with predetermined scaling factors instead of relying solely on regular normalization layers.The method uses a SiLU gain of γ = 1/0.6 and standardized weight matrices; the factors keep activation variance roughly constant.
  • Accuracy: 34 % more accurate forces on COLL and 44 % lower average force error on MD17@CCSD are reported for GemNet relative to prior methods.The COLL comparison concerns force accuracy, while the MD17@CCSD comparison concerns force MAE.
  • Computational trade-offs: Two-hop geometric message passing has higher computational cost, motivating the cheaper GemNet-T ablation.GemNet-Q uses quadruplet-based interactions, and GemNet-T removes the two-hop scheme.
  • Limitations: GemNet is scoped to molecular simulations, while two-hop message passing adds significant overhead despite mitigation through down-projection and GemNet-T.The authors report that parameter count and training or inference time are otherwise on par with previous models.

7 Experiments

The experiments evaluate GemNet on molecular-dynamics and catalyst datasets using energy and force metrics, relaxation outcomes, and energy prediction. GemNet generally outperforms prior methods, while efficiency and the value of two-hop messaging vary by dataset and task.

  • Experimental setup: The experiments cover COLL, MD17, MD17@CCSD, and OC20, with OC20 including structure-to-energy-and-forces, structure relaxation, and relaxed-energy tasks.MD17 and MD17@CCSD use separate molecules, while OC20 contains catalyst relaxation trajectories with adsorbates.
  • Benchmark results: 41 % average improvement on MD17 force MAE and 20 % average improvement across OC20 measures are reported over previous methods.The OC20 result averages across four test sets and all three tasks.
  • Benchmark results: GemNet variants perform best across the investigated molecular-dynamics datasets, including against kernel methods and models trained with substantially more samples.The largest improvements occur for challenging chain-like molecules with broad movement.
  • Ablations: Two-hop GemNet-Q improves over GemNet-T on COLL and combined multi-molecule MD17, but performs approximately on par with GemNet-T on regular MD17.The results suggest regular MD17 is too simple to reveal the benefits of two-hop message passing.
  • Computational aspects: GemNet-Q is roughly two times slower than GemNet-T, while efficient aggregation reduces memory usage for regular Hadamard products by around 50 %.The reported memory reduction is from 4.1GB to 2.2GB for a batch of 32 Toluene molecules.
  • Direct force prediction: Direct force prediction accelerates training by four times and inference by 1.6 times on average, but increases MAE on COLL and MD17 while performing better on OC20 S2EF.The authors state that choosing direct prediction depends on the dataset and application’s computational requirements.

8 Conclusion

The conclusion presents spherical representations as universal for GNNs and GemNet as a geometrically enhanced architecture that substantially improves molecular dynamics prediction.

  • All proposed components yield significant improvements in COLL ablation studies.The studies report force MAE in meV/Å after 500 000 training steps.
  • GemNet combines spherical-representation universality, geometric message passing, symmetric message passing, and efficient bilinear layers.
  • The proposed enhancements substantially improve error across various molecular dynamics datasets.The conclusion also states that most enhancements are independently relevant to other molecular GNNs.

A Proof of Theorem 2

The proof establishes universality by relating spherical-representation models to TFNs and handling real-valued, translation-invariant, rotation-equivariant, and permutation-equivariant functions.

  • The universal approximation theorem is proved by showing equivalence between TFN and the proposed model.Complex spherical harmonics are related to Clebsch-Gordan coefficients in this argument.
  • Real-valued representations can replace complex-valued ones without changing the resulting function space.The real and imaginary components span equivalent function spaces.
  • Scalar pooling functions are limited to linear functions of the constant l = 0 component.This limitation is equivalent to integrating over the real-space spherical representation.
  • Translation invariance is handled by expressing functions through relative vectors such as x2 − x1.The proof rewrites h(x1, x2, ..., xn) using translated inputs anchored at zero.
  • Permutation-equivariant functions are represented using invariant functions under the stated non-degeneracy condition.The condition excludes cases where the relative vectors span a (d − 1)-dimensional space.
  • Activation variance depends on neighborhood size and is estimated empirically because shared weight matrices create non-zero covariance.The paper reports that neighborhood-dependent rescaling negatively affects accuracy.

F GemNet architecture

The full GemNet architecture uses stacked interaction blocks, shared weights, specified embedding sizes, and directed connections that distinguish two-hop from one-hop message passing.

  • GemNet uses 4 stacked interaction blocks with embedding size 128 throughout the model.The basis-function and bilinear-layer settings are also specified for the architecture.
  • Sharing the first weight matrix in Eq. (11) preserves validation loss while reducing training time by up to 15%.
  • The architecture diagram marks inputs, concatenation, SiLU non-linearity, shared weights, and embedding sizes.Dashed lines distinguish two-hop message passing (Q-MP) from one-hop message passing (T-MP).

G Training and hyperparameters

Training uses dataset-specific targets, losses, hyperparameters, and optimization schedules, with batch size strongly affecting MD17(@CCSD) performance.

  • The section includes a table of model and training hyperparameters.
  • MD17 energies are centered by subtracting each molecule’s mean energy, with energies in eV and forces in eV Å−1.
  • A force weighting factor of ρ = 0.999 is used in the training objective.
  • OC20 uses an adapted model with 128 Gaussian radial basis functions, three interaction blocks, and dataset-specific optimization settings.

H Additional experimental results

Additional experiments show that GemNet’s accuracy depends on architectural choices, while several variants improve efficiency or retain strong performance under demanding data and model settings. Across these tests, GemNet remains competitive with or better than prior methods, although some simplifications trade accuracy for speed.

  • Architectural and cutoff effects: Increasing the interaction cutoff to 10 Å slightly lowers COLL force MAE, whereas reducing the embedding cutoff to 3 Å significantly raises it.The results were measured after 500,000 training steps.
  • Efficiency and trade-offs: Direct force prediction accelerates training but significantly reduces accuracy, while GemNet remains more accurate than many previous models; GemNet-T is comparably fast to prior methods.GemNet-Q and GemNet-T required around six days per molecule for MD17 training, with little progress after 100 hours.
  • Architectural and cutoff effects: All independent improvements added to DimeNet++ significantly affect COLL force MAE after increasing its basis embedding size from 8 to 16.
  • Accuracy and data efficiency: GemNet outperforms SchNet on six of eight MD17 molecules while using 1,000 rather than 50,000 training samples.This corresponds to 50x fewer samples for GemNet.
  • Accuracy and data efficiency: On revised MD17, GemNet outperforms FCHL19 by 52% on average and UNiTE by 5%.UNiTE is described as a delta-machine-learning approach using quantum mechanical features.
Loading 2106.08903v10…