PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kristof T. Schütt, Oliver T. Unke, Michael Gastegger

arXiv:2102.03150v4cs.LGphysics.chem-ph

TL;DR

Molecular simulations and spectra are costly, and invariant representations can fail to preserve directional information while neural networks may be less data-efficient than kernel methods. The paper introduces rotationally equivariant message passing through PaiNN, which improves molecular benchmarks, reduces computational costs, and enables tensorial-property prediction for spectra. Molecular-spectrum simulations are accelerated by 4-5 orders of magnitude compared with electronic-structure references.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Accurate molecular simulation is computationally demanding, while invariant message-passing models can lose directional information and remain less data-efficient than kernel methods. The paper introduces rotationally equivariant message passing and PaiNN to address these limitations and accelerate tensorial-property and spectrum prediction.

  • Accurate ab initio molecular dynamics can require solving the Schrödinger equation for millions of time steps, making large-system simulations and spectra computationally prohibitive.Source: §I · p. 1
  • Message-passing neural networks handle large training sets and chemical-property prediction, but kernel methods can perform better with small training sets.Source: §I · p. 1
  • Rotationally invariant message passing can lose relevant directional information because molecules interact as particles in continuous 3d space.Source: §I · p. 1
  • The paper proposes rotationally equivariant message passing and the polarizable atom interaction neural network, PaiNN.Source: §I · p. 1
  • PaiNN improves common molecular benchmarks, reduces model size and inference time, predicts tensorial properties, and accelerates molecular-spectrum simulation by 4-5 orders of magnitude.Source: §I · p. 1

A. Message passing for 3d-embedded graphs

Message passing on 3d molecular graphs exchanges information over local geometric neighborhoods. Equivariant scalar-vector representations preserve directional information while supporting invariant molecular predictions.

C. Limits of rotationally invariant representations

Invariant representations based on distances and angles can fail to distinguish local structures or propagate directional information efficiently. Equivariant messages resolve angular information with linear neighbor complexity, although neither representation family is complete.

IV. POLARIZABLE ATOM INTERACTION NEURAL NETWORK (PAINN)

PaiNN encodes molecular environments with symmetry-aware atomwise representations and couples scalar and vector information. These representations support energy modeling and prediction of tensorial molecular properties.

  • Molecular potential-energy surfaces are invariant to atom permutations, rotations, and translations, so neural potentials should encode these symmetries.Source: §IV · pp. 3–4
  • A common neural-network potential decomposes total energy into atomwise contributions predicted from chemical-environment embeddings.Source: §IV · p. 4
  • Chemical properties can be equivariant tensorial quantities, including dipoles and quadrupoles in an electron-density multipole expansion.Source: §IV · p. 4
  • PaiNN interprets invariant and equivariant atomwise representations as local charges and dipole moments, enabling prediction of tensorial molecular properties.Source: §IV · p. 4

A. Representation

PaiNN initializes invariant atom features from learned atom-type embeddings and directional features at zero, then couples scalar and equivariant vector representations through message and update blocks.

  • PaiNN initializes scalar atom representations from learned embeddings of atom types and sets equivariant representations to zero because no initial directional information is available.Source: A. Representation · p. 4
  • Residual message and update blocks interchange throughout the network, producing coupled scalar and vectorial representations.Source: A. Representation · p. 4
  • Scalar messages use feature-wise continuous-filter convolutions with rotationally invariant radial filters and a cosine cutoff.Source: A. Representation · p. 4
  • Equivariant messages combine invariant filters with scaled vector features and equivariant filters to propagate directional information between neighboring atoms.Source: A. Representation · p. 4
  • Atomwise update blocks couple scalar features to contracted equivariant features through learned scaling functions and scalar products of vector-feature combinations.Source: A. Representation · p. 5

B. Prediction of tensorial properties

PaiNN extends scalar prediction with equivariant tensor outputs, using atomwise scalar and vector components to represent dipoles and polarizabilities while incorporating molecular geometry.

V. RESULTS

The reported training procedure uses Adam with squared loss and validation-based learning-rate decay, while exponential smoothing reduces fluctuations in validation loss.

  • PaiNN models use the Adam optimizer, squared loss, weight decay λ = 0.01, and two output layers unless stated otherwise.Source: §V · p. 5
  • The learning rate is reduced by a factor of 0.5 when validation loss plateaus.Source: §V · p. 5
  • Exponential smoothing with factor 0.9 reduces validation-loss fluctuations, particularly when training with both energies and forces.Source: §V · p. 6

A. Chemical compound space

On QM9 and MD17, PaiNN is evaluated for chemical-property prediction and combined energy–force modeling, with strong benchmark accuracy and reduced model size and inference time.

C. Advantages of equivariant features

Equivariant vector features improve PaiNN’s force accuracy and enable directional information to be propagated at small cutoffs, unlike invariant alternatives in the ferrocene test.

Ablation studies

  • Removing equivariant vector-feature components increases the aspirin force MAE beyond 1 kcal/mol/˚A despite keeping parameter count approximately constant.Source: §1 · p. 6

Propagation of directional information in substituted ferrocene

  • The ferrocene derivative is challenging because its torsional energy depends mainly on the orientation of distant fluorine atoms while energy fluctuations remain below 1 kcal/mol.Source: §2 · p. 7
  • SchNet predicts a flat energy profile for rcut ≤3.0 ˚A because scalar features cannot resolve the rotation angle θ.Source: §2 · p. 8, Figure 4 · p. 8
  • PaiNN faithfully reproduces the substituted-ferrocene torsional energy profile even with very small cutoffs.Source: §2 · pp. 7–8
  • DimeNet resolves θ at smaller cutoffs through explicit triplet-angle terms but still fails at rcut = 2.5 ˚A.Source: §2 · p. 8

D. Molecular Spectra

PaiNN enables efficient tensorial-property prediction for infrared and Raman spectra, including nuclear quantum effects through RPMD. Its simulations reproduce reference spectra closely while reducing computational time from years or hundreds of days to hours.

Simulation requirements

  • RPMD is needed for high-quality spectra because nuclear quantum effects require multiple molecular replicas and increase computational cost.The simulations use 64 replicas for RPMD.Source: D. Molecular Spectra · p. 8

Model and training

  • PaiNN jointly predicts energies, forces, dipole moments, and polarizability tensors for ethanol and aspirin models.The models use 8k ethanol conformations or 16k aspirin training examples, with additional validation and test molecules.Source: D. Molecular Spectra · p. 8

Spectral accuracy

  • RPMD spectra capture experimentally observed frequency shifts more faithfully than classical MD for ethanol and aspirin.For ethanol, RPMD improves agreement with experiment, while classical MD remains closer to the static electronic-structure spectrum; aspirin shows similar trends.Source: D. Molecular Spectra · p. 8

Computational efficiency

  • ≈140 seconds versus 14 ms: PaiNN reduces an individual ethanol reference computation to GPU evaluation time.The overall ethanol simulation decreases from 400 days to approximately one hour.Source: D. Molecular Spectra · p. 8
  • 25 years versus one hour: PaiNN reduces the estimated aspirin simulation time using 15 ms per step.The aspirin reference computation requires 3140 seconds per step.Source: D. Molecular Spectra · p. 9

Appendix A: Data

The appendix documents reference-data generation, training procedures, and molecular-dynamics settings for the reported experiments. It specifies molecule-dependent electronic-structure methods, sampling protocols, optimization details, and spectral-processing parameters.

Reference data

  • Ferrocene reference data used GFN2-xTB, with training structures generated by 300 K normal-mode sampling and random Cp-moiety rotations.The rotational profile sampled a full Cp-Fe-Cp-axis rotation in 1k steps while fixing other degrees of freedom.Source: §1 · p. 12
  • Aspirin reference computations used PBE0/def2-TZVP with ORCA, tight SCF convergence, RIJK acceleration, and 20,000 randomly selected MD17 configurations.Integration grid levels 4 and 5 were used during SCF iterations and final property calculations.Source: §2 · p. 12

Training

  • All models used Adam optimization with validation-loss-based learning-rate decay and exponential smoothing by factor 0.9.Learning rates were selected from 1e-4, 5e-4, 1e-3, and higher values until divergence boundaries were identified.Source: Appendix B · p. 12

Simulation settings

  • Classical molecular dynamics ran for 50 ps at 300 K, discarding the first 10 ps; RPMD used 64 replicas under the same simulation framework.A Nose-Hoover chain thermostat with chain length 3 and a 100 fs time constant controlled temperature.Source: Appendix C · p. 12

Spectral processing

  • Spectra used Wiener-Khinchin autocorrelation processing with 2048 fs depth, Hann windows, zero-padding, 514 nm laser frequency, and 300 K temperature for Raman calculations.These settings specify the spectral post-processing and Raman measurement conditions.Source: Appendix C · p. 12

Related papers

Loading 2102.03150v4…