Skip to content

[ICLR 2024 Spotlight] Official Implementation of "Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products"

License

Notifications You must be signed in to change notification settings

lsj2408/Gaunt-Tensor-Product

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

8 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Gaunt-Tensor-Product

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)

News

  • [24-01-19] 👀 Our paper is updated on arXiv.
  • [24-01-15] 🔥 Our work is accepted at ICLR 2024 as a Spotlight presentation!

Contents

TODO

  • Release the arXiv version.
  • Release the code of efficiency comparison.
  • Clean up and release the code of force field modeling.

Introduction

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:

Main results

Efficiency Comparison

Force Field Modeling -- Open Catalyst Project

Force Field Modeling -- 3BPA

Quick Start

  • Efficiency Comparison Experiment: See efficiency-comparison/README.md
  • Force Field Modeling
    • 3BPA: See force-field-modeling-3bpa/README.md

Citation

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}
}

Acknowledgments

We would like to thank the following open source efforts for pushing the boundary of equivariant networks for the Euclidean group:

About

[ICLR 2024 Spotlight] Official Implementation of "Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products"

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published