This repo is the official implementation of ICLR 2024 Spotlight Presentation: Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products.
Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products
Shengjie Luo*, Tianlang Chen*, Aditi S. Krishnapriyan
Contact: Shengjie Luo (luosj@stu.pku.edu.cn), Tianlang Chen (tlchen@pku.edu.cn), Aditi S. Krishnapriyan (aditik1@berkeley.edu)
- [24-01-19] 👀 Our paper is updated on arXiv.
- [24-01-15] 🔥 Our work is accepted at ICLR 2024 as a Spotlight presentation!
- Release the arXiv version.
- Release the code of efficiency comparison.
- Clean up and release the code of force field modeling.
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(L^6) to O(L^3), where L is the max degree of irreps.
Building upon this approach, we introduce the Gaunt Tensor Product, which serves as a new method for efficient equivariant operations across various model designs. We provide a comprehensive study on major operation classes that are widely used in equivariant models for the Euclidean group, demonstrating the generality of our method and how it can be used to design efficient operations:
- Equivariant Feature Interactions:
- Equivariant Convolutions:
- Equivariant Many-Body Interactions:
- TODO: Please stay tuned.
If you find this work useful, please kindly consider citing our work:
@misc{luo2024enabling,
title={Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products},
author={Shengjie Luo and Tianlang Chen and Aditi S. Krishnapriyan},
year={2024},
eprint={2401.10216},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
We would like to thank the following open source efforts for pushing the boundary of equivariant networks for the Euclidean group: