# An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks

arXiv: https://arxiv.org/abs/2501.13986 · HTML: https://mlsanity.com/paper/2501.13986/

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 36 references that are in the map.
- First posted on arXiv 23 Jan 2025.

## Builds on (most cited first, 20 of 20)

- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data](https://mlsanity.com/paper/1807.02547/) (2018)
- [DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking](https://mlsanity.com/paper/2210.01776/) (2022)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [Cormorant: Covariant Molecular Neural Networks](https://mlsanity.com/paper/1906.04015/) (2019)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network](https://mlsanity.com/paper/1806.09231/) (2018)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains](https://mlsanity.com/paper/2410.24169/) (2024)
- [Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products](https://mlsanity.com/paper/2401.10216/) (2024)
- [What is an equivariant neural network?](https://mlsanity.com/paper/2205.07362/) (2022)
- [Design Principles for Sparse Matrix Multiplication on the GPU](https://mlsanity.com/paper/1803.08601/) (2018)
- [Higher-Order Equivariant Neural Networks for Charge Density Prediction in Materials](https://mlsanity.com/paper/2312.05388/) (2023)
- [Fast evaluation of spherical harmonics with sphericart](https://mlsanity.com/paper/2302.08381/) (2023)
- [ALTO: Adaptive Linearized Storage of Sparse Tensors](https://mlsanity.com/paper/2102.10245/) (2021)
- [Deconstructing equivariant representations in molecular systems](https://mlsanity.com/paper/2410.08131/) (2024)

## Built on by (most cited first, 22 of 22)

- [The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products](https://mlsanity.com/paper/2506.13523/) (2025)
- [EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers](https://mlsanity.com/paper/2604.09130/) (2026)
- [Training a Foundation Model for Materials on a Budget](https://mlsanity.com/paper/2508.16067/) (2025)
- [PFT: Phonon Fine-tuning for Machine Learned Interatomic Potentials](https://mlsanity.com/paper/2601.07742/) (2026)
- [Quick ViTs: Speeding up Vision Transformers through Equivariance](https://mlsanity.com/paper/2505.15441/) (2025)
- [Towards A Transferable Acceleration Method for Density Functional Theory](https://mlsanity.com/paper/2509.25724/) (2025)
- [A Cartesian-3j Framework for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2512.16882/) (2025)
- [E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory](https://mlsanity.com/paper/2601.16622/) (2026)
- [Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2607.02499/) (2026)
- [Smooth Dynamic Cutoffs for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2601.21147/) (2026)
- [Edge Cluster Expansion with Radial Rotary Attention for Interatomic Potentials](https://mlsanity.com/paper/2607.10664/) (2026)
- [Sobek: Streaming Equivariant Tensor Product Convolutions](https://mlsanity.com/paper/2607.18074/) (2026)
- [The Ubiquitous Sparse Matrix-Matrix Products](https://mlsanity.com/paper/2508.04077/) (2025)
- [Perspective: Towards sustainable exploration of chemical spaces with machine learning](https://mlsanity.com/paper/2604.00069/) (2026)
- [MALOQ: Massively Accelerated Learning of Operators for Quantum Transport](https://mlsanity.com/paper/2606.28911/) (2026)
- [Transferable Implicit Solvent Machine Learning Potential for Drugs and Proteins Approaching Ab Initio Accuracy](https://mlsanity.com/paper/2607.10887/) (2026)
- [A Density-Matrix Framework for Electronic-Structure Analysis of Electrolytes for Lithium Batteries](https://mlsanity.com/paper/2607.25597/) (2026)
- [Harmonic Torsional Diffusion for Protein-Ligand Flexible Docking](https://mlsanity.com/paper/2608.20366/) (2026)
- [E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU](https://mlsanity.com/paper/2609.35099/) (2026)
- [BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials](https://mlsanity.com/paper/2610.02013/) (2026)
- [FlashCart: Fast Cartesian Tensor Products for Equivariant Interatomic Potentials](https://mlsanity.com/paper/2610.06409/) (2026)
- [LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs](https://mlsanity.com/paper/2610.18393/) (2026)

## Abstract (arXiv)

Rotation equivariant graph neural networks, i.e. networks designed to guarantee certain geometric relations between their inputs and outputs, yield state of the art performance on spatial deep learning tasks. They exhibit high data efficiency during training and significantly reduced inference time for interatomic potential calculations compared to classical approaches. Key to these models is the Clebsch-Gordon (CG) tensor product, a kernel that contracts two dense feature vectors with a highly-structured sparse tensor to produce a dense output vector. The operation, which may be repeated millions of times for typical equivariant models, is a costly and inefficient bottleneck. We introduce a GPU sparse kernel generator for the CG tensor product that provides significant speedups over the best existing open and closed-source implementations. Our implementation achieves high performance by carefully managing the limited GPU shared memory through static analysis at model compile-time, minimizing reads and writes to global memory. We break the tensor product into a series of smaller kernels with operands that fit entirely into registers, enabling us to emit long arithmetic instruction streams that maximize instruction-level parallelism. By fusing the CG tensor product with a subsequent graph convolution, we reduce both intermediate storage and global memory traffic over naive approaches that duplicate input data. We also provide optimized kernels for the gradient of the CG tensor product and a novel identity for the higher partial derivatives required to predict interatomic forces. Our kernels offer up to 1.3x speedup over NVIDIA's closed-source cuEquivariance package, as well as 10x speedup over the widely-used e3nn package. In FP64 precision, we offer up to 6.2x inference-time speedup for the MACE chemistry foundation model over the original unoptimized version.