Asymptotically Fast Clebsch-Gordan Tensor Products with Vector Spherical Harmonics
Authors: YuQing Xie, Ameya Daigavane, Mit Kotak, Tess Smidt
Organizations: Department of EECS, Massachusetts Institute of Technology, Cambridge, MA 02139, USA · Department of CSE, Massachusetts Institute of Technology, Cambridge, MA 02139, USA
E(3)-equivariant neural networks have proven to be effective in a wide range of 3D modeling tasks. A fundamental operation of such networks is the tensor product, which allows interaction between different feature types. Because this operation scales poorly, there has been considerable work towards accelerating this interaction. However, recently \citet{xieprice} have pointed out that most speedups come from a reduction in expressivity rather than true algorithmic improvements on computing Clebsch-Gordan tensor products. A modification of Gaunt tensor product \citep{gaunt} can give a true asymptotic speedup but is incomplete and misses many interactions. In this work, we provide the first complete algorithm which truly provides asymptotic benefits Clebsch-Gordan tensor products. For full CGTP, our algorithm brings runtime complexity from the naive O(L6) to O(L4log2L), close to the lower bound of O(L4). We first show how generalizing fast Fourier based convolution naturally leads to the previously proposed Gaunt tensor product \citep{gaunt}. To remedy antisymmetry issues, we generalize from scalar signals to irrep valued signals, giving us tensor spherical harmonics. We prove a generalized Gaunt formula for the tensor harmonics. Finally, we show that we only need up to vector valued signals to recover the missing interactions of Gaunt tensor product.
Figures & tables
Figure 1: Fourier transforms from a group theoretic perspective. An action on a group induces an action on functions on the group. These functions form a vector space and hence this defines a group representation. We can decompose this representation into irreps, giving exactly the standard Fourier basis. The Peter-Weyl theorem generalizes this idea for compact Lie groups.
Figure 2: Schematic of the process in taking a vector signal tensor product. We interpret input irreps as vector SH coefficients to create vector spherical signals. We then take pointwise cross products of the two signals to create a new signal which we decompose back into vector SH coefficients.
Tensor Product
Expressivity
Runtime
Runtime / Expressivity
CGTP (Naive)
O(L3)
O(L6)
O(L3)
CGTP (Sparse)
O(L3)
O(L5)
O(L2)
GTP (Fourier)
O(L)
O(L3)
O(L2)
GTP (Grid)
O(L)
O(L3)
O(L2)
GTP (Healy)
O(L)
O(L2log2L)
O(Llog2L)
VSTP (Grid)
O(L)
O(L3)
O(L2)
Table 1: Asymptotic runtimes and expressivity of various tensor product implementations. Note that the fast spherical transform algorithm by Healy et al. (2003) gives the best runtime/expressivity tradeoff ratios.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
SO(n)
Group of rotations in n -dimensional space
O(n)
Group of rotations and inversion in n -dimensional space
S2
The 2-sphere, surface defined by x2+y2+z2=1
Yℓm
Spherical harmonic function of degree ℓ and order m
Yℓ
The collection of spherical harmonic functions of degree ℓ for all orders m
Yj,mℓ,s
Tensor spherical harmonic function
Yjℓ,s
Collection of tensor spherical harmonic functions for all m
Appendix
Table 2: Notation used throughout this paper
G
Denotes a group
ρ(g)
Representation of a group
A
First input space of a constructed bilinearity
B
Second input space of a constructed bilinearity
C
Output space of a constructed bilinearity
X
First input space of a tensor product operation (fixed equivariant bilinearity)
Y
Second input space of a tensor product operation (fixed equivariant bilinearity)
Appendix
Table 3: Commonly used meanings of symbols
SISO
SIMO
MIMO
Encode
O((s1+s2)L2logL)
O((s1+s2)L2logL)
O((s12+s22)L2+(s1+s2)L2log2L)
Interact
O(s1s2s3L2/s~)
O(s1s2s3L2/s~)
O(s1s2s3L2/s~)
Decode
O(s3L2logL)
O(s32L2+s3L2log2L)
O(s32L2+s3L2log2L)
Total
O(s1s2s3L2/s~+s~L2logL)
O((s1+s2)L2logL+s32L2+s3L2log2L)
O(s~2L2+s~L2log2L)
Appendix
Table 4: Asymptotic runtimes of ISTPs assuming a naive grid implementation.
SISO
SIMO
MIMO
Encode
O((s1+s2)L2logL)
O((s1+s2)L2logL)
O((s12+s22)L2+(s1+s2)L2log2L)
Interact
O(s1s2s3L2/s~)
O(s1s2s3L2/s~)
O(s1s2s3L2/s~)
Decode
O(s3L2logL)
O(s32L2+s3L2log2L)
O(s32L2+s3L2log2L)
Total
O(s1s2s3L2/s~+s~L2logL)
O((s1+s2)L2logL+s32L2+s3L2log2L)
O(s~2L2+s~L2log2L)
Appendix
Table 5: Asymptotic runtimes of ISTPs assuming asymptotically fast S2FFT grid implementation.
Table 6: Asymptotic runtimes of various tensor products for different output settings. The best performing tensor products for each output settings are highlighted in green . In the MIMO setting, the Clebsch-Gordan tensor products are highlighted in red to indicate that they can output irreps with multiplicity >1 , unlike the Gaunt tensor products.
We derive integral formulas that simplify the Vector Signal Tensor Product recently introduced by Xie et al., which generalizes the Gaunt tensor product to anti-symmetric couplings. In particular, we obtain explicit closed-form expressions for the anti-symmetric analogues of the Gaunt coefficients. This enables us to simulate the Clebsch-Gordan tensor product using a single Vector Signal Tensor Product, yielding up to a 9× reduction in the required tensor product evaluations. Our results enable efficient and practical implementations of the Vector Signal Tensor Product, paving the way for applications of this generalization of Gaunt Tensor Products in SO(3)-equivariant neural networks. Moreover, we discuss how the Gaunt and the Vector Signal Tensor Products allow to control the expressivity-runtime tradeoff associated with the usual Clebsch-Gordan Tensor Products. Finally, we investigate low rank decompositions of the normalizations of the considered tensor products in view of their use in equivariant neural networks.
E(3)-equivariant networks are promising for 3D atomistic system modeling, yet their scalability is limited by the O(L6) complexity of the Clebsch-Gordan Tensor Product (CGTP). The recently proposed Gaunt Tensor Product (GTP) reduces the complexity but is unable to capture the antisymmetric paths, resulting in incomplete expressivity. In this work, we present SpinGTP, an approach to overcome the GTP incompleteness by generalizing from scalar functions to Spin-Weighted Spherical Harmonics (SWSH). By relying on the algebraic properties of SWSH, SpinGTP recovers the missing antisymmetric interactions while maintaining the asymptotic efficiency of GTP. It also allows for a more expressive equivariant basis that naturally accounts for the parity-odd components of tensor products. We evaluate SpinGTP across diverse benchmarks, including Tetris, 3BPA, SPICE-MACE-OFF, and OC20. Our results show that SpinGTP achieves accuracies comparable to full CGTP. Notably, by explicitly capturing antisymmetric paths, SpinGTP exhibits superior performance in tasks involving chiral materials and non-centrosymmetric geometries. This work provides a complete, scalable, and mathematically rigorous path toward high-order equivariance in large-scale 3D atomistic system simulations.
Chenxing Liang, Yuchao Lin, Andrii Kryvenko +5
Department of Computer Science and Engineering, Texas A&M University · Lambda, Inc. · Department of Materials Science and Engineering, Texas A&M University +3
Equivariant machine learning interatomic potentials (MLIPs) promise atomistic simulations with accuracy approaching that of quantum-mechanical calculations. However, the high computational cost of their core Clebsch--Gordan tensor product (CGTP) convolutions severely limits model scalability and efficiency. In this paper, we rethink the parallel execution of CGTP convolutions and introduce LunaTP. Built on node-centric execution, LunaTP exploits path-group parallelism and multi-level reuse to accelerate CGTP convolutions and higher-order derivative computation while preserving the underlying mathematical formulation. On an NVIDIA H100 GPU, LunaTP achieves kernel speedups of up to 250.0× over e3nn and 17.1× over OpenEquivariance. It also delivers end-to-end training throughput gains of up to 28.7× over e3nn and 3.11× over OpenEquivariance. Finally, LunaTP demonstrates scalability and robust performance across models, workload scales, and GPUs, delivering substantial acceleration for training and atomistic simulation with high-accuracy equivariant MLIPs.
Yuanchang Zhou, Mingzhen Li, Chen Wang +3
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences