Tensor Decomposition Networks for Fast Machine Learning Interatomic Potential Computations

2025-07-01

AI Summary

To address the high computational cost (O(L⁶)) of Clebsch–Gordan (CG) tensor products in SO(3)-equivariant interatomic potential models, this paper proposes the Tensor Decomposition Network (TDN). TDN approximates CG products via CP decomposition, ensuring theoretically bounded equivariance error and universal approximation capability. It further introduces path-weight sharing—compressing O(L³) path-specific parameters into a single shared set—reducing overall complexity to O(L⁴). TDN is fully compatible with mainstream SO(3)-equivariant frameworks as a drop-in replacement. Extensive evaluation on large-scale datasets—including PubChemQC-R, OC20, and OC22—demonstrates that TDN maintains state-of-the-art prediction accuracy while significantly accelerating molecular structure relaxation. The method thus achieves an optimal trade-off between computational efficiency and physical consistency, preserving rotational equivariance without sacrificing predictive fidelity.

Technology Category

• Machine Learning: Matrix & Tensor Methods
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

$ m{SO}(3)$-equivariant networks are the dominant models for machine learning interatomic potentials (MLIPs). The key operation of such networks is the Clebsch-Gordan (CG) tensor product, which is computationally expensive. To accelerate the computation, we develop tensor decomposition networks (TDNs) as a class of approximately equivariant networks whose CG tensor products are replaced by low-rank tensor decompositions, such as the CANDECOMP/PARAFAC (CP) decomposition. With the CP decomposition, we prove (i) a uniform bound on the induced error of $ m{SO}(3)$-equivariance, and (ii) the universality of approximating any equivariant bilinear map. To further reduce the number of parameters, we propose path-weight sharing that ties all multiplicity-space weights across the $O(L^3)$ CG paths into a single path without compromising equivariance, where $L$ is the maximum angular degree. The resulting layer acts as a plug-and-play replacement for tensor products in existing networks, and the computational complexity of tensor products is reduced from $O(L^6)$ to $O(L^4)$. We evaluate TDNs on PubChemQCR, a newly curated molecular relaxation dataset containing 105 million DFT-calculated snapshots. We also use existing datasets, including OC20, and OC22. Results show that TDNs achieve competitive performance with dramatic speedup in computations.

Problem

• Accelerate expensive Clebsch-Gordan tensor products in MLIPs
• Develop low-rank tensor decomposition networks for SO(3)-equivariance
• Reduce computational complexity from O(L^6) to O(L^4)

Innovation

• Replace CG tensor products with low-rank CP decomposition
• Introduce path-weight sharing to reduce parameters
• Achieve competitive performance with dramatic speedup