[Paper Review] 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 Analysis4 citations

TL;DR

This paper introduces the Gaunt Tensor Product, a method to accelerate E(3)-equivariant operations in 3D deep learning by reformulating irreducible representation (irrep) tensor products via Gaunt coefficients. By re-expressing spherical harmonics in a 2D Fourier basis, the method leverages the convolution theorem and Fast Fourier Transforms to reduce computational complexity from O(L⁶) to O(L³), enabling efficient high-degree irrep operations with improved model performance on molecular datasets.

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.

Motivation & Objective

• To address the high computational complexity of full tensor products of irreducible representations (irreps) in E(3)-equivariant deep learning, which limits the use of higher-degree irreps in practice.
• To establish a mathematical link between Clebsch-Gordan coefficients and Gaunt coefficients, enabling a re-interpretation of irrep tensor products as multiplications of spherical functions.
• To enable efficient computation of these operations by transforming the basis from spherical harmonics to a 2D Fourier basis, exploiting the convolution theorem and Fast Fourier Transforms (FFTs).
• To develop a general-purpose, efficient building block—Gaunt Tensor Product—for constructing equivariant models across diverse architectures and applications.
• To demonstrate both improved efficiency and enhanced performance on real-world 3D geometric learning benchmarks, including Open Catalyst Project and 3BPA datasets.

Proposed method

• Mathematically connect Clebsch-Gordan coefficients to Gaunt coefficients, which are integrals of three spherical harmonics, enabling a re-interpretation of irrep tensor products as multiplications of spherical functions.
• Re-express spherical functions in the 2D Fourier basis, allowing the use of the convolution theorem to accelerate multiplications between functions.
• Apply Fast Fourier Transforms (FFTs) to compute the Fourier-domain convolutions efficiently, reducing the overall complexity of irrep tensor products.
• Derive the Gaunt Tensor Product as a new differentiable layer that replaces standard tensor product operations in equivariant models, maintaining equivariance while improving efficiency.
• Implement the method in model architectures such as EquiformerV2 and MACE, using standard training protocols for fair comparison on benchmark datasets.
• Ensure the method preserves the universal approximation capability of irrep-based equivariant networks while drastically reducing computational cost.

Experimental results

Research questions

• RQ1 Can the tensor product of irreducible representations in E(3)-equivariant networks be accelerated by re-expressing the problem in the Fourier domain?
• RQ2 Is there a mathematically rigorous connection between Clebsch-Gordan coefficients and Gaunt coefficients that enables this acceleration?
• RQ3 Does transforming the basis from spherical harmonics to 2D Fourier basis allow for efficient computation via FFTs while preserving equivariance?
• RQ4 To what extent does the proposed Gaunt Tensor Product improve training efficiency and model performance on high-degree irrep settings?
• RQ5 Can the method be generalized across different equivariant model architectures and benchmark datasets in 3D geometric learning?

Key findings

• The computational complexity of full tensor products of irreps is reduced from O(L⁶) to O(L³) by leveraging Gaunt coefficients and 2D Fourier basis transformations.
• The Gaunt Tensor Product enables the use of higher-degree irreps (up to L=6) in EquiformerV2 with significantly improved inference efficiency and reduced memory usage.
• On the Open Catalyst Project dataset, models using the Gaunt Tensor Product achieve state-of-the-art performance with faster training and inference times compared to baselines.
• On the 3BPA dataset, the method demonstrates superior generalization, especially on out-of-distribution and dihedral-slice test sets, indicating a smoother and more accurate potential energy surface.
• The method maintains or improves model accuracy while reducing computational cost, with consistent gains across multiple architectures including EquiformerV2 and MACE.
• The Gaunt Tensor Product is compatible with existing equivariant model designs and can be seamlessly integrated as a drop-in replacement for standard tensor product layers.