[论文解读] Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products

Shengjie Luo, T. Chen+1 more|arXiv (Cornell University)|Jan 18, 2024

Medical Imaging and Analysis被引用 4

一句话总结

本文提出了广义张量积（Gaunt Tensor Product），通过利用Gaunt系数重新表述不可约表示（irrep）张量积，实现了对E(3)-等变3D深度学习操作的加速。通过将球谐函数表示为二维傅里叶基，该方法利用卷积定理与快速傅里叶变换（FFT），将计算复杂度从O(L⁶)降低至O(L³)，从而实现了高阶不可约表示的高效运算，并在分子数据集上取得了更优的模型性能。

ABSTRACT

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)-等变深度学习中不可约表示（irrep）全张量积的高计算复杂度问题，该问题在实践中限制了高阶irrep的使用。
• 建立Clebsch-Gordan系数与Gaunt系数之间的数学联系，使不可约表示张量积可被重新解释为球函数的乘积。
• 通过将基从球谐函数转换为二维傅里叶基，实现高效计算，利用卷积定理与快速傅里叶变换（FFT）加速运算。
• 开发一种通用、高效的构建模块——广义张量积，用于在多种架构与应用中构建等变模型。
• 在真实世界的3D几何学习基准（包括Open Catalyst Project与3BPA数据集）上，展示其在效率与性能上的双重提升。

提出的方法

• 在数学上将Clebsch-Gordan系数与Gaunt系数关联，后者为三个球谐函数的积分，从而实现将irrep张量积重新解释为球函数的乘积。
• 将球函数重新表达为二维傅里叶基，使卷积定理可用于加速函数之间的乘法运算。
• 应用快速傅里叶变换（FFT）高效计算傅里叶域卷积，从而降低irrep张量积的整体复杂度。
• 推导出广义张量积作为新型可微分层，替代等变模型中的标准张量积操作，在保持等变性的同时提升效率。
• 在EquiformerV2与MACE等模型架构中实现该方法，采用标准训练协议以确保在基准数据集上的公平比较。
• 确保该方法在大幅降低计算成本的同时，保持基于irrep的等变网络的通用逼近能力。

实验结果

研究问题

• RQ1 能否通过在傅里叶域中重新表述问题，加速E(3)-等变网络中不可约表示的张量积？
• RQ2 Clebsch-Gordan系数与Gaunt系数之间是否存在数学上严谨的联系，从而实现该加速？
• RQ3 将基从球谐函数转换为二维傅里叶基，是否能在保持等变性的同时，通过FFT实现高效计算？
• RQ4 所提出的广义张量积在高阶irrep设置下，对训练效率与模型性能的提升程度如何？
• RQ5 该方法是否可推广至不同的等变模型架构与3D几何学习基准数据集？

主要发现

• 通过利用Gaunt系数与二维傅里叶基变换，不可约表示全张量积的计算复杂度从O(L⁶)降低至O(L³)。
• 广义张量积使EquiformerV2能够高效使用更高阶irrep（最高至L=6），显著提升推理效率并减少内存占用。
• 在Open Catalyst Project数据集上，采用广义张量积的模型实现了最先进性能，且训练与推理速度优于基线模型。
• 在3BPA数据集上，该方法展现出更优的泛化能力，尤其在分布外与二面角切片测试集上表现突出，表明其能量势面更平滑、更精确。
• 该方法在降低计算成本的同时保持或提升模型精度，在EquiformerV2与MACE等多种架构中均表现出一致的性能增益。
• 广义张量积与现有等变模型设计兼容，可作为标准张量积层的即插即用替代品无缝集成。