Research map / Learning theory and algorithms / Matrix and tensor methods
The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Systematically analyzes tensor product operations (CG, Gaunt, matrix TP) for E(3)-equivariant networks by defining expressivity and interactability, proposes a faster spherical-grid Gaunt TP implementation, and microbenchmarks TPOs.
Problem
Speedups of equivariant tensor products (Gaunt and matrix tensor products) are reported without analysis of what expressivity is lost relative to the Clebsch-Gordan tensor product, and wall-clock behavior is not well characterized.
Approach
Defines expressivity (dimension of constructible equivariant bilinearities) and interactability (nonzero selection rules) for tensor product operations, shows asymptotic speedups come with loss of expressivity, and reimplements Gaunt TP on a spherical grid (same asymptotics, ~30% faster in MACE training).
Key results
- Grid-based Gaunt TP is about 30% faster than the original Fourier-basis Gaunt TP when training MACE on 3BPA and revised MD17 for the same GPU hours.
- On the 3D Tetris task, all CGTP networks reach 100% accuracy while no GTP network does, because GTP cannot form the antisymmetric 1e cross-product term.
- Asymptotic speedups of GTP and MTP correspond to reduced expressivity; after normalizing by expressivity, sparse CGTP is fastest in wall-clock time and FLOPs, and sparse CGTP has the fewest FLOPs but the slowest walltime among the implementations measured.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace GTP (Fourier) with GTP (Grid) | GTP implementation in MACE · 3BPA and revised MD17 | speed: 30\% faster | Helped |
| Use GTP instead of CGTP | message-passing network · 3D Tetris pieces | accuracy: 100\% accuracy with CGTP; none of the GTP networks do | Hurt |
| Use a spherical grid rather than the 2D Fourier basis for GTP | GTP implementation · Tensor product microbenchmark | wall-clock time: 30\% faster | Helped |
| Use GTP instead of CGTP | tensor product operation · Asymptotic analysis (MIMO) | runtime: CGTP (Sparse) $\mathcal{O}(L^5)$; GTP (Fourier) $\mathcal{O}(L^3)$; expressivity $\mathcal{O}(L^3)$ vs $\mathcal{O}(L)$ | Mixed |
| Normalize tensor product runtime by expressivity | CGTP, GTP, and MTP implementations · Asymptotic analysis (MIMO) | runtime / expressivity: CGTP (Sparse), GTP (Fourier), GTP (Grid), and MTP (Sparse): $\mathcal{O}(L^2)$; GTP (S2FFT): $\mathcal{O}(L\log^2 L)$ | Helped |
| Use CGTP rather than GTP | message-passing network · 3D Tetris pieces | ability to distinguish chiral pieces | Helped |
| Add more channels to GTP (ablation) | GTP network · 3D Tetris pieces | ability to distinguish chiral pieces | No clear effect |
| Add pseudo-spherical harmonics to GTP (ablation) | GTP network · 3D Tetris pieces | ability to distinguish chiral pieces | No clear effect |
| Use CGTP with sparsity rather than the reported CGTP implementation | tensor product microbenchmark · RTX A5500 | wall-clock time | Hurt |
Surprising findings
- Sparse CGTP has the fewest FLOPs yet the slowest wall-clock time among the benchmarked implementations.
- Adding channels or pseudo-spherical harmonics does not allow the Gaunt TP network to solve the chiral Tetris task.
Limitations
- Only the forward pass of tensor products was benchmarked, not full training behavior across tasks.
- Nsight Compute profiling is very slow, limiting the scope of profiling.
- Expressivity measure is defined for the two-fold tensor product with irrep basis, not for general nu-fold or Cartesian bases.
Experimental setup
Datasets: 3BPA revised MD17 3D Tetris pieces (synthetic)
Benchmarks: 3BPA revised MD17 3D Tetris classification
Metrics: wall-clock time GFLOPs GPU utilization expressivity classification accuracy force/energy error
Baselines: Clebsch-Gordan tensor product Gaunt tensor product (Fourier) matrix tensor product (e3x) MACE
Builds on: Gaunt tensor product MACE NequIP e3nn
Closest papers by meaning
- Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products (ICLR 2024 · similarity 0.90)
- A Tale of Two Symmetries: Exploring the Loss Landscape of Equivariant Models (NeurIPS 2025 · similarity 0.86)
- Spin-Weighted Spherical Harmonics Enable Complete and Scalable E(3)-Equivariant Networks (NeurIPS 2026 · similarity 0.86)
- Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing (ICLR 2025 · similarity 0.86)
- On the Fourier analysis in the SO(3) space : the EquiLoPO Network (ICLR 2025 · similarity 0.86)
- Convolutions and More as Einsum: A Tensor Network Perspective with Advances for Second-Order Methods (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.