Directional Message Passing for Molecular Graphs

Johannes Gasteiger; Janek Groß; Stephan Günnemann

Source-linked AI summary

Directional Message Passing for Molecular Graphs

Johannes Gasteiger, Janek Groß, Stephan Günnemann

arXiv:2003.03123v2Mar 6, 2020cs.LGphysics.comp-phstat.ML

TL;DR

Molecular GNNs typically use interatomic distances but omit directional information relevant to molecular interactions. The paper introduces directional message passing and DimeNet, which use direction-associated message embeddings and spherical basis representations. DimeNet outperforms previous GNNs on average by 76% on MD17 and 31% on QM9 while supporting molecular-dynamics requirements.

Problem
Molecular GNNs omit directional information despite its role in angular and torsional interactions in empirical molecular potentials.
Method
The paper embeds messages associated with directions, transforms them using distances and angles, and represents these quantities with spherical Bessel functions and spherical harmonics.
Results
DimeNet outperforms previous GNNs on average by 76% on MD17 and by 31% on QM9.
Takeaways & Limitations
DimeNet provides a molecular GNN suitable for predicting molecular properties and for molecular-dynamics simulations.
Takeaways & Limitations
For molecular-dynamics applications, stable simulations require twice-continuous differentiability and exclude discontinuous transformations such as ReLU nonlinearities.

Abstract

from arXiv · show

Graph neural networks have recently achieved great successes in predicting quantum mechanical properties of molecules. These models represent a molecule as a graph using only the distance between atoms (nodes). They do not, however, consider the spatial direction from one atom to another, despite directional information playing a central role in empirical potentials for molecules, e.g. in angular potentials. To alleviate this limitation we propose directional message passing, in which we embed the messages passed between atoms instead of the atoms themselves. Each message is associated with a direction in coordinate space. These directional message embeddings are rotationally equivariant since the associated directions rotate with the molecule. We propose a message passing scheme analogous to belief propagation, which uses the directional information by transforming messages based on the angle between them. Additionally, we use spherical Bessel functions and spherical harmonics to construct theoretically well-founded, orthogonal representations that achieve better performance than the currently prevalent Gaussian radial basis representations while using fewer than 1/4 of the parameters. We leverage these innovations to construct the directional message passing neural network (DimeNet). DimeNet outperforms previous GNNs on average by 76% on MD17 and by 31% on QM9. Our implementation is available online.

1 INTRODUCTION

The paper identifies a limitation of molecular GNNs: distance-based atom embeddings omit directional information important for angular and torsional interactions. It proposes directional message passing and DimeNet to incorporate this information while preserving physical symmetries.

GNNs update atom embeddings using messages based on previous embeddings and pairwise distances, without directly using directional information.
Directional message passing embeds messages associated with directions to neighboring atoms rather than embedding atoms themselves.
76% average improvement on MD17 and 31% on QM9 establish DimeNet’s reported advantage over previous GNNs.
DimeNet uses spherical Bessel functions and spherical harmonics to provide orthogonal representations for distances and angles.
The proposed framework connects directional information with equivariance, graph neural networks, belief propagation, and empirical molecular potentials.

2 RELATED WORK

Related work spans molecular machine learning, general graph neural networks, and equivariant neural networks. The paper positions DimeNet at their intersection, replacing hand-crafted molecular representations with learned directional representations.

ML for molecules: Molecular property prediction has progressed from hand-crafted atomic-neighborhood representations with Gaussian processes or neural networks to feature-learning GNNs.
Graph neural networks: GNN research has pursued architectures with greater expressive power than the 1-Weisfeiler-Lehman isomorphism test.
Equivariant neural networks: Equivariant neural networks incorporate symmetry principles across domains including spheres, molecules, volumetric data, and general manifolds.

3 REQUIREMENTS FOR MOLECULAR PREDICTIONS

Molecular prediction models must use atomic numbers and positions while respecting physical symmetries and, for molecular dynamics, force and differentiability requirements. The paper therefore emphasizes conservative force prediction, continuous differentiability, and avoiding fixed auxiliary features.

Molecular prediction models map atomic numbers and positions to scalar targets, while some prior models additionally use auxiliary features such as bond types or electronegativity.
Explicitly incorporating physical invariances can avoid duplicate weights and reduce training time and model complexity.
For molecular dynamics, forces are obtained as the negative gradient of a predicted potential, ensuring a conservative force field.
Stable simulations require the model to be twice continuously differentiable, ruling out discontinuous transformations such as ReLU nonlinearities.
Because atom positions can change arbitrarily, molecular-dynamics models cannot use pre-computed auxiliary information such as bond types.

4 DIRECTIONAL MESSAGE PASSING

Directional message passing augments distance-based molecular GNNs with equivariant directional embeddings and angle-aware interactions. Messages are updated analogously to belief propagation, enabling direct use of angular information.

Regular GNN messages use atom embeddings and pairwise distances, omitting directional information such as bond angles and rotations.
Directional message passing embeds each atom-neighbor direction as a separate message rather than embedding atoms alone.
Directional message embeddings are rotationally equivariant because their associated directions rotate with the molecule.
Neighboring messages are updated through interactions based on interatomic distances and angles between message directions, following a belief-propagation analogy.
The angle-aware interaction function enables the model to learn angular potentials directly and can distinguish molecular structures that regular cutoff-based GNNs cannot.

5 PHYSICALLY BASED REPRESENTATIONS

The paper represents molecular distances and directional geometry with orthogonal spherical Bessel and spherical Fourier-Bessel bases. Continuous cutoff treatment supports stable, differentiable predictions while using compact radial representations.

Gaussian radial bases are replaced by orthogonal representations based on spherical Bessel functions and spherical harmonics to reduce redundancy and improve parameter efficiency.
The 2D spherical Fourier-Bessel basis jointly represents an incoming-message distance and the angle between message directions.
The radial basis represents interatomic distance alone by setting the angular indices to zero, producing a purely radial Bessel representation.
The real-valued distance and angular bases are orthogonal on their domain and bound the highest-frequency components, regularizing the model and stabilizing predictions.
Six spherical radial functions and 16 radial basis functions were sufficient; 16 is 4x lower than PhysNet’s 64 and 20x lower than SchNet’s 300.
An envelope function makes the cutoff representation twice continuously differentiable by driving the basis functions and their first two derivatives to zero at the cutoff.

6 DIRECTIONAL MESSAGE PASSING NEURAL NETWORK (DIMENET)

DimeNet integrates directional message passing and spherical Fourier-Bessel representations into a molecular GNN. Its architecture produces invariant predictions while supporting differentiable, energy-conserving force prediction.

Architecture: DimeNet integrates directional message passing and spherical Fourier-Bessel representations into a streamlined PhysNet-based architecture.The model is designed for molecular-property prediction and molecular-dynamics simulations.
Differentiability and symmetries: DimeNet is invariant to atom permutations and molecular translation, rotation, and inversion, while remaining twice continuously differentiable.Its construction enables energy-conserving, rotation-equivariant atomic forces and smooth predictions under small deformations.
Embedding block: Atomic numbers receive learnable embeddings, which are combined with interatomic distances to generate initial message embeddings.The embedding block shares atom-type embeddings across molecules and applies a learnable transformation.
Performance: DimeNet sets the state of the art on 11 QM9 targets and outperforms the second-best model by 31% in mean standardized MAE.The result is reported in the QM9 comparison table.
Interaction block: Stacked interaction blocks transform each message using radial and two-dimensional distance-angle representations, bilinear layers, and residual blocks.The 2D representation is first projected before interaction with message embeddings, while radial information uses element-wise multiplication.
Output block: Message embeddings are passed through output blocks, aggregated per atom, transformed into atom-wise outputs, and summed into the final prediction.The output block uses radial basis information before atom-wise and global aggregation.

7 EXPERIMENTS

Experiments evaluate DimeNet on QM9 and MD17, compare it with established models, and ablate its directional and basis representations. DimeNet achieves strong benchmark performance, while the ablations show that both contributions materially affect accuracy.

Experimental setup: DimeNet is compared with SchNet, PhysNet, PPGN, MEGNet-simple, Cormorant, and sGDML across molecular-property and molecular-dynamics tasks.The experiments use QM9 for molecular properties and MD17 for energies and atomic forces.
QM9: DimeNet outperforms SchNet by a large margin and performs roughly on par with sGDML on QM9.QM9 contains roughly 130,000 equilibrium molecules and evaluates multiple molecular-property targets.
Ablation studies: Replacing 16 and 6 Bessel basis functions with 64 Gaussian RBFs increases error by 10%.This indicates that the Bessel basis provides a helpful inductive bias in addition to reducing parameter count.
Ablation studies: Ignoring message angles increases DimeNet’s error by around 26%, while replacing message embeddings with node embeddings increases MAE by 68%.The node-embedding variant performs worse than SchNet, and learned filters show dependence on both distance and angle.

8 CONCLUSION

The paper introduces directional message passing and spherical Fourier-Bessel representations to make molecular GNN interactions more expressive. DimeNet performs strongly on QM9 and MD17, while future work remains to model torsional and non-bonded terms.

Conclusion: Directional message passing lets molecular GNNs use directional information in addition to interatomic distances.The approach represents messages associated with directions and targets the important hard degrees of freedom in molecular interactions.
Conclusion: Spherical Bessel and two-dimensional spherical Fourier-Bessel bases provide principled representations for distances and directional information.These representations are integrated into DimeNet for molecular-property prediction and molecular-dynamics simulations.
Conclusion: DimeNet’s QM9 and MD17 results support the directional message-passing and basis contributions as essential ingredients of its state-of-the-art performance.The conclusion explicitly attributes the demonstrated performance to these contributions.
Future work: Future work should incorporate torsional and non-bonded interaction terms beyond the first two modeled terms.The paper states that doing so could improve predictions further.

A INDISTINGUISHABLE MOLECULES

A distance-only or graph-structure-only GNN cannot distinguish certain molecules with identical local bond-length neighborhoods. Directional message passing resolves this ambiguity by incorporating bond directions.

Indistinguishable molecules: A standard non-directional GNN cannot distinguish a hexagonal molecule from two triangular molecules when their bond-length neighborhoods are identical.Cyclohexane and two slightly stretched cyclopropane molecules provide an example when the cutoff is c ≤2.5 Å.
Indistinguishable molecules: Directional message passing solves this problem by considering the direction of each bond.The added directional information separates structures that share the same bond lengths.

B EXPERIMENTAL SETUP

The setup uses optimized interaction-block hyperparameters, standardized error metrics for multi-target evaluation, and learned distance- and angle-dependent filters to inspect DimeNet’s spatial dependencies.

Model configuration: The model uses 6 interaction blocks with embedding size F = 128, while basis and bilinear dimensions are selected as sufficiently large values.The chosen values are NSHBF = 7, NSRBF = NRBF = 6, and Nbilinear = 8; performance was not very sensitive once dimensions reached at least 4.
Evaluation metrics: The evaluation reports mean standardized MAE and mean standardized logMAE across 12 targets.Standardized MAE compares average error with each target’s standard deviation, while logMAE gives equal weight to relative improvements but is sensitive to outliers.
Filter analysis: DimeNet’s learned filters are analyzed by separating distance dependence from the interaction function and examining their variation with distance and angle.The first 15 elements of the angular filter are visualized for tensor slice n = 1, with α = 0 at the top of the figure.

E MULTI-TARGET RESULTS

The multi-target QM9 results compare single-target and multi-target learning, including the effect of using separate output blocks for each target.

Multi-target learning: Single-target learning significantly improves performance on all QM9 targets compared with multi-target learning.Using a separate output block per target slightly reduces this performance difference with little impact on training time.