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.

  • Problem

    Accurate molecular dynamics and spectra are computationally prohibitive, while invariant representations can be insufficient for propagating directional information.

  • Method

    The paper introduces rotationally equivariant message passing and the PaiNN architecture, retaining equivariant scalar-vector representations for molecular modeling.

  • Results

    PaiNN outperforms invariant message-passing networks, matches kernel methods at small sample sizes, and accelerates molecular-spectrum simulation by 4-5 orders of magnitude.

  • Takeaways & Limitations

    Equivariant atomwise representations support tensorial-property prediction and make molecular-spectrum simulations feasible on timescales reduced from years to hours.

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.
  • Message-passing neural networks handle large training sets and chemical-property prediction, but kernel methods can perform better with small training sets.
  • Rotationally invariant message passing can lose relevant directional information because molecules interact as particles in continuous 3d space.
  • The paper proposes rotationally equivariant message passing and the polarizable atom interaction neural network, PaiNN.
  • 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.

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.

  • 3d-embedded molecular graphs connect atom i to neighbors j through relative positions within a cutoff radius.
  • Message-passing neural networks repeatedly exchange neighborhood messages and update node features.
  • More expressive local representations can use scalar and vectorial features instead of scalar neurons alone.
  • Rotational equivariance requires vector message and update functions to transform consistently under any rotation matrix.
  • Equivariant networks apply nonlinear functions to scalars, scale vectors by scalars, and linearly combine equivariant vectors.
  • Directional information remains available during message passing, while scalar products can recover rotationally invariant 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.

  • Rotationally invariant atomic-environment representations may be insufficient for successive message passes to construct complex, nonlocal representations.
  • Distances resolve changing bond lengths but not bond angles, while angles resolve angular changes but not distances; combining both increases expressiveness.
  • Equivariant messages resolve angular information with runtime complexity O(|N|), avoiding the quadratic neighbor scaling of angle-inclusive messages.
  • In an equidistant four-atom chain, invariant distances and angles cannot distinguish two arrangements, whereas equivariant representations differ in component sign.
  • Retaining equivariance allows directional information to propagate beyond the neighborhood while keeping the cutoff small and scaling linearly with neighbor count.
  • Neither many-body invariant representations nor finite-order multipole expansions are complete for representing arbitrary molecular structures or equivariant functions.
  • The proposed architecture is intended to exploit equivariant representations and overcome the discussed limitations of invariant representations.

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.
  • A common neural-network potential decomposes total energy into atomwise contributions predicted from chemical-environment embeddings.
  • Chemical properties can be equivariant tensorial quantities, including dipoles and quadrupoles in an electron-density multipole expansion.
  • PaiNN interprets invariant and equivariant atomwise representations as local charges and dipole moments, enabling prediction of tensorial molecular properties.

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.
  • Residual message and update blocks interchange throughout the network, producing coupled scalar and vectorial representations.
  • Scalar messages use feature-wise continuous-filter convolutions with rotationally invariant radial filters and a cosine cutoff.
  • Equivariant messages combine invariant filters with scaled vector features and equivariant filters to propagate directional information between neighboring atoms.
  • Atomwise update blocks couple scalar features to contracted equivariant features through learned scaling functions and scalar products of vector-feature combinations.

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.

  • PaiNN predicts scalar properties from invariant atomwise representations and equivariant tensors of order M through rank-1 tensor decomposition.
  • Gated equivariant output blocks produce the scalar and vector components used in tensorial property decompositions.
  • PaiNN keeps pairwise message functions lightweight while restricting complex transformations to atomwise updates and the tensor output network.
  • The molecular dipole is constructed from polarized atoms, combining atomic dipoles with charges multiplied by atom positions.
  • For polarizability tensors, isotropic atomwise polarizabilities are combined with anisotropic components, with atom positions incorporating global molecular structure.

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.
  • The learning rate is reduced by a factor of 0.5 when validation loss plateaus.
  • Exponential smoothing with factor 0.9 reduces validation-loss fluctuations, particularly when training with both energies and forces.

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.

  • For QM9, PaiNN is trained on 110k molecules, with 10k used for validation and the remainder for testing across three random splits.
  • PaiNN obtains validation and test MAEs of 0.045a0 for isotropic polarizability, comparable to DimeNet++.
  • PaiNN achieves state-of-the-art results for six QM9 target properties, matches DimeNet++ on two others, and ranks second on the remaining properties.
  • 600k parameters make PaiNN substantially smaller than DimeNet++ with about 1.8M parameters.
  • Inference time falls from 45 ms to 13 ms for random QM9 batches of 50 molecules, an improvement of more than 70% versus DimeNet++.
  • On MD17, PaiNN is evaluated with 1k known structures, using 950 for training and 50 for validation, with a separate model for each trajectory.

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.
  • 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.
  • Propagation of directional information in substituted ferrocene: SchNet predicts a flat energy profile for rcut ≤3.0 ˚A because scalar features cannot resolve the rotation angle θ.
  • Propagation of directional information in substituted ferrocene: PaiNN faithfully reproduces the substituted-ferrocene torsional energy profile even with very small cutoffs.
  • Propagation of directional information in substituted ferrocene: DimeNet resolves θ at smaller cutoffs through explicit triplet-angle terms but still fails at rcut = 2.5 ˚A.

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.
  • 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.
  • 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.
  • 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.
  • Computational efficiency: 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.

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.
  • Reference data: 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.
  • 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.
  • 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.
  • 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.
Loading 2102.03150v4…