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)群开发等变神经网络在真实世界应用中建模三维数据方面发挥着重要作用。实现这种等变性主要涉及不可约表示的张量积。然而,随着使用更高阶张量,此类操作的计算复杂度显著增加。在本工作中,我们提出了一种系统方法来大幅加速不可约表示张量积的计算。我们将常用的Clebsch-Gordan系数在数学上与Gaunt系数联系起来,后者是三个球谐函数乘积的积分。通过Gaunt系数,不可约表示的张量积等价于由球谐函数表示的球面函数之间的乘法。这一视角进一步使我们能够将等变操作的基从球谐函数转换为二维傅里叶基。因此,由二维傅里叶基表示的球面函数之间的乘法可以通过卷积定理和快速傅里叶变换高效计算。这一变换将完整不可约表示张量积的复杂度从$\mathcal{O}(L^6)$降低到$\mathcal{O}(L^3)$,其中$L$为不可约表示的最大阶数。利用这一方法,我们引入了Gaunt张量积,它作为一种新方法,可在不同模型架构中构建高效的等变操作。我们在Open Catalyst Project和3BPA数据集上的实验证明了我们方法在提升效率和改进性能两方面的优势。