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

YuQing Xie, Ameya Daigavane, Mit Kotak, Tess Smidt

TL;DR: We explore the efficiency and expressivity of various tensor product operations in equivariant neural networks both theoretically and experimentally.

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

Lay Summary:

Neural networks that understand 3D geometry are becoming increasingly important for tasks like predicting how molecules interact or how objects move in space. These networks often rely on a complex mathematical operation called the tensor product to combine 3D information. However, this operation is usually slow and uses a lot of computing power.

Recently, researchers have looked into developing faster versions. One such example is the Gaunt Tensor Product (GTP). In our work, we took a closer look at several of these fast methods to see how they really compare. We found that while these methods are faster, they do so by losing expressivity. We propose new ways to measure the expressivity/speed tradeoff for these alternative tensor product operations.

We also discovered a simpler and even faster way to implement GTP by using a spherical grid, which cuts down training time in a real-world example by 30%. Lastly, we tested all the methods in detail and found that how fast they are in theory doesn’t always match up with how they perform in practice. So it’s important to test them carefully depending on the task.

Our findings provide researchers a useful too for understanding the right balance between speed and accuracy when building powerful 3D AI models.

Link To Code: https://github.com/atomicarchitects/PriceofFreedom

Primary Area: Deep Learning->Graph Neural Networks

Keywords: equivariance, tensor product, spherical harmonics, gaunt, benchmarking, asymptotics, expressivity

Submission Number: 6290