TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials
This work addresses the need for efficient and accurate machine learning models in molecular science, offering a novel approach that could impact computational chemistry and materials research.
The paper tackles the problem of efficiently representing molecular systems for machine learning by introducing TensorNet, an O(3)-equivariant message-passing neural network that uses Cartesian tensor representations, achieving state-of-the-art performance with significantly fewer parameters and reduced computational cost.
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.