[論文レビュー] 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

ひとこと要約

本論文は、球面調和関数を2次元フーリエ基底に再表現することにより、ガウント係数を介して不変表現（irrep）のテンソル積を再定式化することで、3次元ディープラーニングにおけるE(3)同変演算を高速化するためのガウントテンソル積を導入する。この手法により、O(L⁶)からO(L³)への計算複雑度の低減が達成され、高次数irrepの効率的処理が可能となり、分子データセットにおけるモデル性能の向上が実現される。

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の使用を制限する要因を解消すること。
• クリーブシュ・ゴルダン係数とガウント係数の間の数学的関連を確立し、irrepテンソル積を球関数の積として再解釈することを可能にすること。
• 球面調和関数から2次元フーリエ基底への変換を施し、畳み込み定理と高速フーリエ変換（FFT）を活用して、関数間の積算を効率的に計算すること。
• 多様なアーキテクチャと応用分野に適用可能な汎用的で効率的なビルディングブロック「ガウントテンソル積」を構築すること。
• オープンカタリシスプロジェクトや3BPAデータセットを含む、実世界の3次元幾何学的学習ベンチマークにおいて、訓練効率とモデル性能の両方を向上させることを示すこと。

提案手法

• クリーブシュ・ゴルダン係数とガウント係数の数学的関連を確立し、これらが3つの球面調和関数の積分として定義されることを示す。これにより、irrepテンソル積を球関数の積として再解釈可能となる。
• 球関数を2次元フーリエ基底に再表現することで、関数間の積算に畳み込み定理を適用し、計算を高速化することが可能になる。
• 高速フーリエ変換（FFT）を用いて、フーリエ領域における畳み込みを効率的に計算し、irrepテンソル積全体の複雑度を低減する。
• 標準的なテンソル積演算を置き換える新しい微分可能レイヤーとしてのガウントテンソル積を導出。同変性を維持しつつ、計算効率を向上させる。
• EquiformerV2 や MACE などのモデルアーキテクチャにこの手法を実装し、ベンチマークデータセットでの公平な比較を実現するため、標準的な訓練プロトコルを用いる。
• irrepに基づく同変ネットワークの普遍的近似能力を保持しつつ、計算コストを著しく削減することを保証する。

実験結果

リサーチクエスチョン

• RQ1 E(3)同変ネットワークにおける不変表現のテンソル積は、フーリエ領域に問題を再定式化することで高速化可能か？
• RQ2 クリーブシュ・ゴルダン係数とガウント係数の間には、この高速化を可能にする数学的に厳密な関連性があるか？
• RQ3 球面調和関数から2次元フーリエ基底への基底変換により、FFTを用いた効率的計算が可能となり、同変性が保持されるか？
• RQ4 提案されたガウントテンソル積は、高次次数irrep設定において、訓練効率とモデル性能をどの程度向上させるか？
• RQ5 この手法は、3次元幾何学的学習分野におけるさまざまな同変モデルアーキテクチャとベンチマークデータセットに一般化可能か？

主な発見

• ガウント係数と2次元フーリエ基底変換を活用することで、irrepの完全テンソル積の計算複雑度がO(L⁶)からO(L³)に低減された。
• ガウントテンソル積により、EquiformerV2では高次次数irrep（L=6まで）の使用が可能となり、推論効率が著しく向上し、メモリ使用量も削減された。
• オープンカタリシスプロジェクトデータセットでは、ガウントテンソル積を用いたモデルが、ベースラインと比較して最先端の性能を達成し、訓練および推論時間が短縮された。
• 3BPAデータセットでは、分布外および二面角スライステストセットにおいて優れた一般化性能を示し、滑らかでより正確なポテンシャルエネルギー表面を実現した。
• 計算コストを削減しつつ、モデル精度を維持または向上させ、EquiformerV2 や MACE などの複数のアーキテクチャで一貫した性能向上が得られた。
• ガウントテンソル積は、既存の同変モデル設計と互換性があり、標準的なテンソル積レイヤーの即時置換としてシームレスに統合可能である。