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 O(L6) to O(L3), 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.Comment: 36 pages; ICLR 2024 (Spotlight Presentation); Code:
https://github.com/lsj2408/Gaunt-Tensor-Produc