arXiv:2506.13523cs.LGcs.AI2025-06ICML被引 19

剖析3D等变网络中张量积的表达能力与速度权衡,发现快不等于好。

The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products

  • 通过表达能力与交互性指标对比不同张量积方法
  • 提出球面网格法在保持理论复杂度下提速30%
  • 揭示理论速度与实际表现差异,强调针对性测试必要性

E(3)等变神经网络在多种3D建模任务中表现优异。其核心操作是张量积,用于以等变方式组合几何特征生成新特征。由于张量积计算复杂度高,研究者投入大量精力优化运行效率。例如,Luo等人(2024)提出的Gaunt张量积(GTP)宣称有显著加速。本文对多种张量积操作进行系统分析,指出不同方法并非执行相同操作,普遍的提速常以牺牲表达能力为代价。我们引入表达能力与交互性度量来刻画差异。此外,发现原始GTP实现可通过直接使用球面网格大幅简化,且不增加渐近时间复杂度。该球面网格方法在基准测试和MACE原子间势的实际训练中均快30%。最后,我们首次提供各类张量积操作的系统微基准测试,发现理论运行时承诺与实际性能差异巨大,凸显特定应用测试的重要性。代码已公开于https://github.com/atomicarchitects/PriceofFreedom。

原文摘要 · Abstract (English)

$E(3)$-equivariant neural networks have demonstrated success across a wide range of 3D modelling tasks. A fundamental operation in these networks is the tensor product, which interacts two geometric features in an equivariant manner to create new features. Due to the high computational complexity of the tensor product, significant effort has been invested to optimize the runtime of this operation. For example, Luo et al. (2024) recently proposed the Gaunt tensor product (GTP) which promises a significant speedup. In this work, we provide a careful, systematic analysis of a number of tensor product operations. In particular, we emphasize that different tensor products are not performing the same operation. The reported speedups typically come at the cost of expressivity. We introduce measures of expressivity and interactability to characterize these differences. In addition, we realized the original implementation of GTP can be greatly simplified by directly using a spherical grid at no cost in asymptotic runtime. This spherical grid approach is faster on our benchmarks and in actual training of the MACE interatomic potential by 30%. Finally, we provide the first systematic microbenchmarks of the various tensor product operations. We find that the theoretical runtime guarantees can differ wildly from empirical performance, demonstrating the need for careful application-specific benchmarking. Code is available at https://github.com/atomicarchitects/PriceofFreedom.

等变网络张量积3D建模

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