TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials (2306.06482v2)

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.

• The paper introduces TensorNet, which uses Cartesian tensor embeddings and matrix product operations to efficiently learn molecular potentials.
• It employs O(3)-equivariant message passing and isolates scalars and vectors, significantly cutting computational overhead.
• Experiments on QM9 and rMD17 demonstrate that TensorNet attains high accuracy with only 23% of the parameters used in comparable models.

TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials

The paper introduces TensorNet, a machine learning model employing $\mathrm{O}(3)$-equivariant message-passing neural networks that utilize Cartesian tensor representations to efficiently learn molecular potentials. This approach demonstrates a significant stride in managing the complexity and computational demands typically associated with modelling molecular systems.

Key Contributions and Methodology

TensorNet innovates by using Cartesian tensors, specifically rank-2 tensors represented as 3x3 matrices, to simplify feature mixing via matrix product operations. This contrasts with traditional higher-rank spherical tensor models. The decomposition of these tensors into rotation group irreducible representations enables isolated processing of scalars, vectors, and tensors, which results in substantial reductions in computational cost while maintaining state-of-the-art performance with fewer parameters.

The architecture of TensorNet is structured around several key mechanisms:

1. Cartesian Tensor Embeddings: Cartesian rank-2 tensors are used as atomic embeddings. These embeddings allow efficient feature mixing through matrix operations.
2. Equivariant Message-Passing: TensorNet employs an $\mathrm{O}(3)$-equivariant framework for message passing, ensuring that transformations (like rotations) of the input do not affect the model's output.
3. Efficient Decomposition: The model leverages the decomposition of Cartesian tensors into scalar, vector, and other forms, enabling separate processing of these components as dictated by their unique transformation properties.
4. Computational Efficiency: By minimizing the number of parameters and message-passing steps needed, TensorNet achieves high accuracy with less computational overhead compared to spherical tensor methodologies.

Numerical Results and Performance

TensorNet exhibits impressive quantitative metrics, reflecting in its outperformance of several contemporary models on datasets like QM9 and rMD17. With a reduction in parameter count to merely 23% of models like Allegro, it still leads in accuracy for molecular property prediction. The model's efficacy is especially notable in predicting molecular energy-related properties, where it matches or surpasses state-of-the-art performances.

TensorNet's architecture also enables the straightforward prediction of vector and tensor molecular quantities and potential energies. For example, in QM9, TensorNet achieves mean absolute errors as low as 3.9 meV for certain molecular properties, outperforming notable models like DimeNet++ and PaiNN.

Implications and Future Directions

The development of TensorNet marks a notable shift in how machine learning models approach molecular systems by emphasizing computational efficiency without sacrificing performance. This work broadens the horizon for the design of efficient equivariant models, providing a template for future research in designing models that incorporate physical system symmetries more effectively.

The usage of Cartesian tensors, rather than more complex spherical ones, implies that future models can increasingly focus on reducing computational burdens, a critical advantage for scaling machine learning methods to larger systems or more demanding simulations. Future directions could explore further optimizations, incorporation into broader chemical simulations, or extensions to handle higher-rank quantities when needed.

Conclusion

Overall, TensorNet provides a significant contribution by harmonizing computational efficiency and accuracy in molecular potential learning, highlighting the potential to explore and design state-of-the-art equivariant models with practical implications for scientific research and computational chemistry.