Developing equivariant neural networks for the E(3) group plays an important role in modeling 3D data across real-world applications. Enforcing this equivariance primarily involves the tensor products of irreducible representations (irreps). However, the computational complexity of such operations increases significantly as higher-order tensors are used. In this work, we propose a systematic approach to substantially accelerate the computation of the tensor products of irreps. We mathematically connect the commonly used Clebsch-Gordan coefficients to the Gaunt coefficients, which are integrals of products of three spherical harmonics. Through Gaunt coefficients, the tensor product of irreps becomes equivalent to the multiplication between spherical functions represented by spherical harmonics. This perspective further allows us to change the basis for the equivariant operations from spherical harmonics to a 2D Fourier basis. Consequently, the multiplication between spherical functions represented by a 2D Fourier basis can be efficiently computed via the convolution theorem and Fast Fourier Transforms. This transformation reduces the complexity of full tensor products of irreps from $\mathcal{O}(L^6)$ to $\mathcal{O}(L^3)$, where $L$ is the max degree of irreps. Leveraging this approach, we introduce the Gaunt Tensor Product, which serves as a new method to construct efficient equivariant operations across different model architectures. Our experiments on the Open Catalyst Project and 3BPA datasets demonstrate both the increased efficiency and improved performance of our approach.
翻译:开发面向E(3)群的等变神经网络在三维数据建模的实际应用中具有重要作用。实现这种等变性主要涉及不可约表示的张量积运算。然而,随着高阶张量的使用,此类运算的计算复杂度显著增加。本研究提出了一种系统性方法,能够大幅加速不可约表示张量积的计算过程。我们从数学上建立了常用的克利布施-戈丹系数与高特系数之间的关联——后者是三个球谐函数乘积的积分。通过高特系数,不可约表示的张量积等价于球谐函数表示的球面函数之间的乘法运算。该视角进一步允许我们将等变运算的基函数从球谐函数变换为二维傅里叶基。由此,基于卷积定理和快速傅里叶变换,可高效计算二维傅里叶基表示的球面函数乘法。这种变换将不可约表示全张量积的复杂度从$\mathcal{O}(L^6)$降低至$\mathcal{O}(L^3)$,其中$L$为不可约表示的最高阶数。基于该方法,我们提出高特张量积,作为跨不同模型架构构建高效等变运算的新方案。在Open Catalyst Project和3BPA数据集上的实验表明,该方法在提升效率的同时也改进了模型性能。