Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products

要約

E(3) グループの等変ニューラル ネットワークの開発は、現実世界のアプリケーション全体で 3D データをモデリングする際に重要な役割を果たします。
この等分散性の強制には、主に既約表現のテンソル積 (irreps) が関係します。
ただし、高次のテンソルが使用されると、そのような演算の計算の複雑さは大幅に増加します。
この研究では、irreps のテンソル積の計算を大幅に高速化する体系的なアプローチを提案します。
一般的に使用されるクレブシュ・ゴルダン係数を、3 つの球面調和関数の積の積分であるゴーント係数に数学的に接続します。
Gaunt 係数により、irreps のテンソル積は、球面調和関数で表される球面関数間の乗算と等価になります。
この観点により、等変演算の基礎を球面調和関数から 2D フーリエ基礎に変更することがさらに可能になります。
その結果、2D フーリエ基底で表される球面関数間の乗算は、畳み込み定理と高速フーリエ変換によって効率的に計算できます。
この変換により、irreps の完全テンソル積の複雑さが $\mathcal{O}(L^6)$ から $\mathcal{O}(L^3)$ に軽減されます。ここで、$L$ は irreps の最大次数です。
このアプローチを活用して、異なるモデル アーキテクチャ間で効率的な等変演算を構築する新しい方法として機能する Gaunt Tensor Product を導入します。
Open Catalyst プロジェクトと 3BPA データセットに関する私たちの実験は、私たちのアプローチの効率の向上とパフォーマンスの向上の両方を実証しています。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google