TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials
arXiv:2306.06482
Abstract
The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce TensorNet, an innovative O(3)-equivariant message-passing neural network architecture that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Furthermore, the cost-effective decomposition of these tensors into rotation group irreducible representations allows for the separate processing of scalars, vectors, and tensors when necessary. Compared to higher-rank spherical tensor models, TensorNet demonstrates state-of-the-art performance with significantly fewer parameters. For small molecule potential energies, this can be achieved even with a single interaction layer. As a result of all these properties, the model's computational cost is substantially decreased. Moreover, the accurate prediction of vector and tensor molecular quantities on top of potential energies and forces is possible. In summary, TensorNet's framework opens up a new space for the design of state-of-the-art equivariant models.
NeurIPS 2023, camera-ready version
Cited by in corpus (6)
- Evaluation of the MACE Force Field Architecture: from Medicinal Chemistry to Materials Science
- On the design space between molecular mechanics and machine learning force fields
- Probing the effects of broken symmetries in machine learning
- Accelerating the discovery of high-performance nonlinear optical materials using active learning and high-throughput screening
- A Cartesian Encoding Graph Neural Network for Crystal Structures Property Prediction: Application to Thermal Ellipsoid Estimation
- Assessing zero-shot generalisation behaviour in graph-neural-network interatomic potentials