Directional Message Passing for Molecular Graphs

TL;DR

The paper presents directional message passing to incorporate angular information into molecular GNNs, introducing direction-specific embeddings and a joint 2D basis based on spherical Bessel functions and spherical harmonics. The Directional Message Passing Neural Network (DimeNet) achieves state-of-the-art performance on QM9 and MD17 while producing differentiable energy and force predictions suitable for molecular dynamics. Ablation studies confirm that directional information and the 2D basis are key to performance gains, and the approach provides a parameter-efficient framework for capturing hard angular degrees of freedom in molecular systems. This work lays the groundwork for accurate, physics-informed MD simulations with scalable graph-based models and suggests extending the method to full energy terms in future work.

Abstract

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.

Figures (8)

• Figure 1: Aggregation scheme for message embeddings.
• Figure 2: 2D spherical Fourier-Bessel basis $\tilde{a}_{\text{SBF}, ln}(d, \alpha)$.
• Figure 3: Radial Bessel basis for $N_\text{RBF} = 5$.
• Figure 4: The DimeNet architecture. $\square$ denotes the layer's input and $\|$ denotes concatenation. The distances $d_{ji}$ are represented using spherical Bessel functions and the distances $d_{kj}$ and angles $\alpha_{(kj,ji)}$ are jointly represented using a 2D spherical Fourier-Bessel basis. An embedding block generates the inital message embeddings ${\bm{m}}_{ji}$. These embeddings are updated in multiple interaction blocks via directional message passing, which uses the neighboring messages ${\bm{m}}_{kj}, k \in \mathcal{N}_j \setminus \{i\}$, the 2D representations ${\bm{a}}_\text{SBF}^{(kj,ji)}$, and the distance representations ${\bm{e}}_\text{RBF}^{(ji)}$. Each block passes the resulting embeddings to an output block, which transforms them using the radial basis ${\bm{e}}_\text{RBF}^{(ji)}$ and sums them up per atom. Finally, the outputs of all layers are summed up to generate the prediction.
• Figure 5: MAE on MD17 using 1000.0 training samples (energies in [per-mode=fraction], forces in [per-mode=fraction]). DimeNet outperforms SchNet by a large margin and performs roughly on par with sGDML.