arXiv:2507.01131cs.LGphysics.comp-ph2025-07NeurIPS被引 1

用低秩分解加速原子间势能计算,速度提升显著且保持精度。

Tensor Decomposition Networks for Fast Machine Learning Interatomic Potential Computations

  • 用CP分解替代昂贵的克莱布什-戈登张量积,实现近似等变性。
  • 在1.05亿分子构型上测试,性能接近现有模型,计算复杂度从O(L⁶)降至O(L⁴)。
  • 适合需要高效分子模拟的研究者,尤其关注计算速度与可扩展性。

SO(3)等变网络是机器学习原子间势能(MLIP)的主流模型,其核心操作为克莱布什-戈登(CG)张量积,计算开销大。为此,我们提出张量分解网络(TDNs),以低秩张量分解(如CP分解)近似替代CG张量积。理论证明:(i)SO(3)等变性的误差存在统一上界;(ii)可逼近任意等变双线性映射。为进一步减少参数量,提出路径权重共享机制,将O(L³)条CG路径的多重性空间权重统一为单一参数集,不破坏等变性。该层可直接替换现有网络中的张量积模块,计算复杂度由O(L⁶)降至O(L⁴)。我们在新构建的PubChemQCR数据集(含1.05亿个基于DFT计算的分子构型)及现有数据集OC20、OC22上评估,结果表明TDNs在性能相当的前提下实现显著提速。代码已开源,集成于AIRS库。

原文摘要 · Abstract (English)

$\rm{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 in which 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 $\rm{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 $\mathcal{O}(L^3)$ CG paths into a single shared parameter set 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 $\mathcal{O}(L^6)$ to $\mathcal{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. Our code is publicly available as part of the AIRS library (\href{https://github.com/divelab/AIRS/tree/main/OpenMol/TDN}{https://github.com/divelab/AIRS/}).

分子模拟张量分解等变网络加速计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。