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

TL;DR

This paper systematically analyzes various equivariant tensor products in 3D neural networks, highlighting tradeoffs between expressivity and runtime, and introduces simplified, faster implementations with empirical benchmarks.

Contribution

It provides a systematic analysis of tensor product operations, introduces measures for expressivity, simplifies the Gaunt tensor product, and offers comprehensive empirical benchmarks.

Findings

Different tensor products perform distinct operations.

Speedups often reduce expressivity.

Spherical grid implementation improves runtime by 30%.

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 https://github.com/atomicarchitects/PriceofFreedom.