Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shengjie Luo, Tianlang Chen, Aditi S. Krishnapriyan

TL;DR

The Gaunt Tensor Product is introduced, which serves as a new method to construct efficient equivariant operations across different model architectures and 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.

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.

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

TL;DR

The Gaunt Tensor Product is introduced, which serves as a new method to construct efficient equivariant operations across different model architectures and reduces the complexity of full tensor products of irreps from to , where L$ is the max degree of irreps.

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 to , where 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.
Paper Structure (73 sections, 2 theorems, 66 equations, 1 figure, 2 tables)

This paper contains 73 sections, 2 theorems, 66 equations, 1 figure, 2 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Euclidean group and Equivariance.
  4. Irreducible representations.
  5. Spherical harmonics.
  6. Tensor products of irreducible representations.
  7. Gaunt Tensor Product
  8. A New Perspective on Tensor Products of Irreps via Gaunt Coefficients
  9. Fast Tensor Product using Convolution Theorem & Fast Fourier Transform
  10. From spherical harmonics to 2D Fourier bases.
  11. Performing 2D convolution via Fast Fourier Transform.
  12. From 2D Fourier bases to spherical harmonics.
  13. Computational complexity analysis.
  14. Discussion.
  15. Designing Efficient Equivariant Operations via the Gaunt Tensor Product
  16. ...and 58 more sections

Key Result

Theorem 3.1

Let $j,l,J \in \mathbb{N}_{\geq0}$ and let $\boldsymbol{T}^{(j)}$ be a spherical tensor operator of rank $j$. Then there is a unique complex number, the reduced matrix element$\lambda \in \mathbb{C}$ (often written $\langle J \| \boldsymbol{T}^{(j)} \| l \rangle \in \mathbb{C}$), that completely det

Figures (1)

  • Figure 1: Results on efficiency comparisons and sanity check. We comprehensively compare our Gaunt Tensor Product with implementations of e3nn, eSCN, and MACE on corresponding equivariant operation classes. In all settings, our approach achieves significant speedups. Our Gaunt Tensor Product parameterization further passes the sanity check with SEGNN on the N-body simulation task.

Theorems & Definitions (2)

  • Theorem 3.1: Wigner-Eckart theorem for Spherical Tensor Operators, from jeevanjee2011introduction
  • Theorem A.1: Wigner-Eckart theorem for Spherical Tensor Operators, from jeevanjee2011introduction