The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products

2025-06-16

AI Summary

This work systematically investigates the expressivity–efficiency trade-off in tensor product operations within E(3)-equivariant neural networks. We propose quantitative metrics for expressivity and geometric interaction, and theoretically and empirically demonstrate that acceleration often comes at the cost of diminished higher-order geometric interactions. We establish, for the first time, that the Gaunt tensor product can be equivalently reformulated as a spherical grid-based implementation (GTP), achieving a 30% speedup in MACE potential training without asymptotic complexity overhead. To rigorously evaluate such operations, we introduce the first micro-benchmarking suite specifically designed for E(3) tensor products—revealing substantial discrepancies between theoretical complexity estimates and empirical runtime performance. We publicly release the complete benchmark suite, reference implementations, and evaluation tools. This work provides both theoretical foundations and practical standards for designing efficient, expressive equivariant models.

Technology Category

• Machine Learning: Matrix & Tensor Methods
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Computer Vision: 3D Computer Vision

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

$E(3)$-equivariant neural networks have demonstrated success across a wide range of 3D modelling tasks. A fundamental operation in these networks is the tensor product, which interacts two geometric features in an equivariant manner to create new features. Due to the high computational complexity of the tensor product, significant effort has been invested to optimize the runtime of this operation. For example, Luo et al. (2024) recently proposed the Gaunt tensor product (GTP) which promises a significant speedup. In this work, we provide a careful, systematic analysis of a number of tensor product operations. In particular, we emphasize that different tensor products are not performing the same operation. The reported speedups typically come at the cost of expressivity. We introduce measures of expressivity and interactability to characterize these differences. In addition, we realized the original implementation of GTP can be greatly simplified by directly using a spherical grid at no cost in asymptotic runtime. This spherical grid approach is faster on our benchmarks and in actual training of the MACE interatomic potential by 30%. Finally, we provide the first systematic microbenchmarks of the various tensor product operations. We find that the theoretical runtime guarantees can differ wildly from empirical performance, demonstrating the need for careful application-specific benchmarking. Code is available at href{https://github.com/atomicarchitects/PriceofFreedom}{https://github.com/atomicarchitects/PriceofFreedom}

Problem

• Analyzing expressivity and runtime tradeoffs in equivariant tensor products.
• Comparing different tensor product operations and their expressivity costs.
• Benchmarking empirical performance versus theoretical runtime guarantees.

Innovation

• Simplified GTP using spherical grid
• Introduced expressivity and interactability measures
• Systematic microbenchmarks for tensor products