[Paper Review] Equivariant message passing for the prediction of tensorial properties and molecular spectra

Kristof T. Schütt, Oliver T. Unke+1 more|arXiv (Cornell University)|Feb 5, 2021

Machine Learning in Materials Science265 citations

TL;DR

The paper introduces PaiNN, a rotationally equivariant message passing neural network for molecular graphs that predicts scalar and tensorial properties and enables efficient molecular spectra simulations, achieving state-of-the-art results with smaller models and faster inference.

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.

Motivation & Objective

• Motivate the limitations of rotationally invariant representations in propagating directional information in molecular graphs.
• Propose rotationally equivariant message passing and the PaiNN architecture to improve data efficiency and expressiveness.
• Enable prediction of tensorial properties (e.g., dipole moments, polarizabilities) using equivariant representations.
• Demonstrate applications to molecular spectra via RPMD simulations with large speedups over electronic structure references.

Proposed method

• Develop equivariant message passing where scalar and vector features are updated with rotationally equivariant functions.
• Use continuous-filter convolutions with radial basis function filters for invariant scalar messages and vector-aware convolutions for equivariant messages.
• Couple scalar and vector updates through a residual architecture to propagate directional information.
• Predict tensorial properties through a rank-1 tensor decomposition combining atomwise latent features and local vector features.
• Train end-to-end on energies, forces, dipole moments, and polarizabilities using Adam with weight decay and learning-rate scheduling.

Experimental results

Research questions

• RQ1 Can rotationally equivariant message passing outperform invariant GNNs on standard molecular-property benchmarks at various data regimes?
• RQ2 Do equivariant representations enable accurate prediction of tensorial properties such as dipoles and polarizabilities without excessive model complexity?
• RQ3 Can PaiNN accelerate molecular spectra simulations (IR/Raman) via fast, accurate trajectories compared to ab initio methods?

Key findings

• PaiNN achieves state-of-the-art or competitive results on QM9 across multiple scalar properties, with a smaller parameter count (~0.6M) than some baselines like DimeNet++ (~1.8M).
• PaiNN reduces inference time by over 70% versus DimeNet++ on QM9 (13 ms vs 45 ms per batch of 50 molecules).
• Ablation studies show all components of the equivariant architecture contribute to accuracy, with vector-feature convolutions and scalar–vector coupling providing notable gains.
• Equivariant representations enable accurate prediction of tensorial properties such as dipole moments and polarizabilities, leveraging atom-centered polarization terms and global geometry.
• PaiNN enables fast, RPMD-based simulations of infrared and Raman spectra, achieving substantial speedups over electronic-structure references and producing spectra in close agreement with experiment for ethanol (RPMD) and aspirin datasets.
• Equivariant features improve the model’s ability to capture directional information with small cutoffs, enabling efficient propagation beyond local neighborhoods.