通过 Gaunt 张量积在 Fourier 基中实现高效的等变操作

开发针对 E(3) 群的等变神经网络在跨现实世界应用的 3D 数据建模中发挥着重要作用。强制执行这种等变性主要涉及不可约表示的张量积。然而，随着使用更高阶的张量，此类操作的计算复杂度显著增加。在这项工作中，我们提出了一种系统方法来大幅加速不可约表示张量积的计算。我们在数学上将常用的 Clebsch-Gordan 系数与 Gaunt 系数联系起来，后者是三个球谐函数乘积的积分。通过 Gaunt 系数，不可约表示的张量积等价于由球谐函数表示的球面函数之间的乘法。这一视角进一步允许我们将等变操作的基从球谐函数转换为 2D Fourier 基。因此，由 2D Fourier 基表示的球面函数之间的乘法可以通过卷积定理和快速 Fourier 变换高效计算。这种变换将不可约表示全张量积的复杂度从 $\mathcal{O}(L^6)$ 降低到 $\mathcal{O}(L^3)$，其中 $L$ 是不可约表示的最大次数。利用这种方法，我们引入了 Gaunt 张量积，作为一种在不同模型架构中构建高效等变操作的新方法。我们在 Open Catalyst Project 和 3BPA 数据集上的实验证明了我们方法效率的提升和性能的改善。