Equivariant message passing for the prediction of tensorial properties and molecular spectra (2102.03150v4)

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.

• The paper introduces rotationally equivariant representations via the PaiNN MPNN, achieving improved accuracy in predicting tensorial properties.
• The paper's PaiNN architecture outperforms invariant models on molecular benchmarks, offering up to 4-5 orders of magnitude speedup in molecular simulations.
• The paper demonstrates that equivariant message passing reduces computational resources while maintaining high prediction reliability for diverse molecular spectra and energies.

Equivariant Message Passing for the Prediction of Tensorial Properties and Molecular Spectra

This paper presents notable advancements in the prediction of molecular properties using message passing neural networks (MPNNs), specifically through the development of the polarizable atom interaction neural network (PaiNN). The authors extend the MPNN framework to incorporate rotationally equivariant representations, addressing limitations posed by previous invariant methodologies. This work demonstrates improved predictive accuracy on molecular benchmarks while maintaining a reduced model size and faster inference times.

Key Contributions

The primary contributions of this research are:

1. Introduction of Rotationally Equivariant Representations: The paper extends message passing formulations to include rotationally equivariant representations. This enhancement allows for more expressive and data-efficient predictions, particularly benefiting the learning of tensorial properties such as dipoles and polarizabilities.
2. Development of PaiNN: A novel MPNN architecture, PaiNN, is introduced. It facilitates refined predictions across various molecular properties and outperforms previous invariant models in common molecular benchmarks. The architecture leverages equivariant representations, enhancing the model's capability to predict both scalar and tensorial molecular properties.
3. Substantial Speedups in Molecular Simulation: By employing PaiNN, the prediction of molecular spectra, specifically infrared and Raman spectra, achieves significant speed enhancements—up to 4-5 orders of magnitude faster than traditional electronic structure methods. This advancement markedly reduces computational times, making such simulations feasible for extensive analyses.

Numerical Results and Benchmarks

The performance of PaiNN is evaluated on the QM9 dataset and the MD17 molecular dynamics benchmark. The results demonstrate PaiNN's improved accuracy and computational efficiency:

• In the QM9 dataset, PaiNN achieves leading performance in various property predictions and competes effectively with models such as DimeNet++.
• On the MD17 dataset, PaiNN shows superior data efficiency, performing comparably to kernel methods with small training datasets and excelling in both energy and force predictions.

Theoretical and Practical Implications

The proposed approach offers significant theoretical implications by demonstrating the superiority of equivariant over invariant representations in extending information propagation capabilities within MPNNs. Practically, this translates to reduced computational resources and enhanced prediction reliability, making it a valuable tool in computational chemistry for accurate and efficient molecular simulations.

Future Developments

Future directions of this work could explore the integration of PaiNN in more complex chemical systems, such as enzyme catalysis or surface reactions. The equivariant framework may also find applications in other domains requiring 3d geometries, such as generative models or wavefunction predictions.

Conclusion

This paper contributes to the field of computational chemistry by addressing the limitations of invariant representations and introducing a neural network architecture capable of efficiently predicting both scalar and tensorial properties of molecules. The development of PaiNN represents a significant step in advancing the capability and efficiency of MPNNs for molecular property predictions and simulations.