Equivariant message passing for the prediction of tensorial properties and molecular spectra

TL;DR

This work addresses data efficiency and the ability to predict tensorial molecular properties with graph-based models. It introduces PaiNN, a rotationally equivariant message-passing neural network that uses coupled scalar and vector features to propagate directional information efficiently, enabling accurate predictions of scalar properties and tensorial quantities like dipole moments and polarizabilities. The approach achieves state-of-the-art or competitive performance on QM9 and MD17 with smaller models, and demonstrates dramatic speedups in simulating infrared and Raman spectra via RPMD, reducing runtimes from years to hours. Overall, equivariant message passing with PaiNN offers significant gains in both predictive accuracy and computational efficiency for large-scale molecular simulations and spectral analyses.

Abstract

Message passing neural networks have become a method of choice for learning on graphs, in particular the prediction of chemical properties and the acceleration of molecular dynamics studies. While they readily scale to large training data sets, previous approaches have proven to be less data efficient than kernel methods. We identify limitations of invariant representations as a major reason and extend the message passing formulation to rotationally equivariant representations. On this basis, we propose the polarizable atom interaction neural network (PaiNN) and improve on common molecule benchmarks over previous networks, while reducing model size and inference time. We leverage the equivariant atomwise representations obtained by PaiNN for the prediction of tensorial properties. Finally, we apply this to the simulation of molecular spectra, achieving speedups of 4-5 orders of magnitude compared to the electronic structure reference.

Figures (5)

• Figure 1: Illustration of message passing using angles and directions for two structures. All edges within the cutoff range (dashed lines) have equal length. The representations of the blue and red node are the same using angles (left), while directions allow to distinguish both structures (right).
• Figure 2: The architecture of PaiNN with the full architecture (a) as well as the message (b) and update blocks (c) of the equivariant message passing. In all experiments, we use 128 features for $\mathbf{s}_i$ and $\vec{\mathbf{v}}_i$ throughout the architecture. Other layer sizes are annotated in grey.
• Figure 3: Gated equivariant block.
• Figure 4: Rotational energy profile of substituted ferrocene obtained by varying the rotation angle $\theta$ while keeping all bond distances fixed at their equilibrium values (H: white, C: black, F: purple, Fe: grey). For small cutoff $r_{\rm cut}$ (dashed grey circles), MPNNs with scalar feature representations are unable to represent information about the rotation angle $\theta$.
• Figure 5: IR (top) and Raman (bottom) spectra of ethanol and aspirin. Spectra calculated with the reference method using the harmonic oscillator approximation are shown in black (QM harmonic). The inset table shows the mean absolute errors on the respective test set.