The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products
Authors: YuQing Xie
Ameya Daigavane
Mit Kotak
Tess Smidt
Venue: ICML 2025 poster (2025)
Keywords: equivariance
tensor product
spherical harmonics
gaunt
benchmarking
asymptotics
expressivity
systems
Primary area: deep_learning->graph_neural_networks
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

The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products
YuQing Xie
Ameya Daigavane
Mit Kotak
Tess Smidt
Abstract
1. Introduction
Many complex physical systems possess inherent spatial symmetries, and incorporating these symmetries into mod- els has been shown to significantly improve both learning ef- ficiency and robustness (Batzner et al., 2022; Rackers et al., 2023; Frey et al., 2023; Owen et al., 2024). To address the specific symmetries present in 3D systems, consider- able effort has been dedicated to the development of E(3)- equivariant neural networks (E(3)NNs) (Thomas al., 2018; Weiler et al., 2018; Kondor, 2018; Kondor et al., 2018). E(3)NNs have delivered strong performance across wide range of scientific applications, including molecular force fields (Batzner et al., 2022: Musaelian et al., 2023; Batatia et al., 2022), catalyst discovery (Liao & Smidt, 2023), gen- erative models (Hoogeboom et al., 2022), charge density prediction (Fu et al., 2024), and protein structure prediction (Lee al., 2022: Jumper et al., 2021).
(33)-qquivariant neural networks have demon- strated success across a wide range of 3D mod- elling tasks. A fundamental operation these networks is the tensor product, which interacts two geometric features in an equivariant manner to create new features. Due to the high compu- tational complexity of the tensor product, signifi- cant effort has been invested to optimize the run- time of this operation. For example, Luo et al. (2024) recently proposed the Gaunt tensor prod- uct (GTP) which promises a significant speedup. In this work, we provide careful, systematic analysis of number of tensor product opera- tions. In particular, we emphasize that different tensor products are not performing the same op- eration. 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 po- tential by 30%. Finally, we provide the first sys- tematic microbenchmarks of the various tensor product operations. We find that the theoretical runtime guarantees can differ wildly from em- pirical performance, demonstrating the need for careful application-specifie benchmarking. Our code available at https: httt:::://gthub..co atomm ca rch
The group E(3) consists of all rotations, translations and reflections in 3 dimensions. A function f X + Y is
E(3)-cquivariant if it satisfies:
f(g x) g Ege
(1)
where the group action may differ on the input space X and output space Y.
E33 -qquiiaaivr neural networks work with features that transform as irreducible representations of O()), termed 'irreps', described in Appendix B. How these irreps transform under 3D rotations (elements of SO(3)) is defined by positive integer L, which can intuitively be thought of as an angular frequency. As described in Section 2. tensor products are used to construct equivariant bilincarities, a key component of E33) networks.
The only true tensor product is the Clebsch-Gordan ten- sor product (CGTP) which uses the well-studied Clebsch- Gordan (Varshalovich et al., 1988) coefficients. However, it has time complexity of (LL)) as we show in Appendix E, which can quickly become expensive for larger L. This scal- ing limits the direct application of ((3)-equivariant neural
'Equal contribution 'Department of EECS, Massachusetts Institute of Technology, Cambridge MA 02139, USA DDppartment of CSE, Massachuselts Institute of Technol- ogy, Cambridge MA 02139, USA. Correspondence to: YuQing Xie <xyuqing@@minn @ii. Tess Smidt Dismidt@@it.cdu..
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Note that Passaro & Zitnick (2023) claims a runtime of (IT)) for this tensor product. In Appendix E, show that this runtime can be reduced to (T.)).

---

The Price of Freedom
networks to larger systems. Indeed, a number of works suggest that incorporating higher I. features can be crucial to improving model performance (Cen et al.; Frank et al., 2022; Aykent & Xia). Hence, there is significant interest in optimizing the tensor product.
2. Analyzing Tensor Product Operations
2.1. Motivating Irreps
We begin by motivating why many equivariant architectures (Geiger & Smidt, 2022: Unke & Maennel, 2024) use irreps to represent features. In particular, it is much easier to enforce equivariance when our features are in terms of irreps because they facilitate parameterization of the most general equivariant linear maps.
One optimization was identified by Passaro & Zitnick (2023) in the special case where one input is a projection of a single vector onto spherical harmonics, an operation commonly used in 3D graph convolutions. Under suitable rotation, these irreps become sparse, allowing for a runtime of O(L)) However, the extreme sparsity is not generally true for arbi- trary irrep values.
First, while there are an infinite number of ways to choose finite dimensional representations, we can always decom- pose them into a direct sum of irreps. Irreps are hence the fundamental building blocks of arbitrary representations and they are well studied for many groups. Further, assuming we are working over an algebraically closed field, linear layers between a pair of irreps must either be 0 if they are inequivalent or multiples of the identity if they are equiva- lent. This is known as Schur's 's lemma (Dresselhaus et al., 2007). Hence once we rewrite representations explicitly as a direct sum of irreps, parameterizing equivariant linear layers becomes trivial. This is why identifying features in terms of irreps is so prevalent in equivariant architectures.
For arbitrary irrep values, Luo et al. (2024) proposed the Gaunt Tensor Product (GTP) which they show has com- plexity O(L)). Further, Unke & Maennel (2024) introduced another ((L)) operation which we call matrix tensor prod- uct (MTP). While this represents exciting progress, it raises an important question: What is fundamentally different be- tween these new tensor products and the general CGTP with )LL) complexity?
In this paper, we provide a framework for systematically analyzing these new operations. Importantly, we clarify that most new proposed 'tensor products' are not tensor products in the mathematical sense, and propose to call them lensor product operations (TPOs) instead. We orga- nize this paper by first motivating TPOs in Section 2. We then define a measure of expressivity and interactability for TPOs. In Section 3, we introduce a number of TPOs, and we summarize the asymptotic runtimes and expressivities of various TPO implementations in Section 4. Finally, in Sec- tion 5 we benchmark the various tensor products showing that asymptotics do not always correspond with practical performance.
However, Schur's 's lemma explicitly prevents linear mixing of inequivalent irreps. Hence we need appropriate nonlinear interactions. The simplest possibility is to use multiplication, motivating the tensor product.
2.2. Tensor Products and Bilinearities
Here we provide a unified framing for understanding why tensor products are expensive and how other works attempt to remedy this problem. In particular we want to empha- size the connection to constructing bilinearities and that other 'tensor products' are fundamentally different opera- tions from the formal mathematical tensor product.
We summarize our core contributions as follows:
A measure for expressivity and interactability of TPOs.
Given two vector spaces X, Y, one can construct a tensor product space X @ Y with an associated bilinear mapping T : X x Y - X & Y. This bilinear mapping often referred as the tensor product. In practice, if we have a basis 21,..., Im and y...... Un of X, Y respectively, then we usually use basis T( y₁) (often written as x; @ yj) for the tensor product space X @ Y.
A simpler implementation of GTP using projection onto the sphere S², which has the same asymptoties but is faster in practice.
Now suppose we have a space Z. Then in practice, we typically parameterize a linear layer Lin : X & Y + Z. The composition Lin T : X x Y + Z is a bilinearity X x Y - Z. Note that this construction is universal, as all possible bilincarities can be written as Lin T for some choice of linearity Lin.
A comprchensive analysis of asymptotic runtimes and expressivity of different classes of TPOs.
Rigorous benchmarks of various TPO implementations.
We assume familiarity with group representations of SO()). For a brief introduction to representations (reps) and ir- reducible representations (irreps) of (3(3), we refer the reader to Appendix B.
While tensor products allow interactions between differ- ent irreps and simple construction of all possible bilineari- ties, they are expensive. First, the dimension of X x Y is O[X]]] where [X] and |Y| | are the dimensions of X and
2

---

The Price of Freedom
Y respectively. This is in fact lower bound on asymptotic runtime. Further, even if X,Y Y are explicitly written in a ba- sis which is direct sum of irreps, the most natural basis of X & Y is generally not. Hence we must perform change of basis of X & Y into one which explicitly is a direct of irreps. The change of basis matrix is known as the Clebsch-Gore Gordan coefficients. Naively, it requires O [+) time to per- form this change of basis. The tensor product whose output is explicitly written in an irrep basis is the (lebsch-Gordan tensor product.
We can now naturally define the expressivity of given tensor product by how many equivariant bilinearities it can parametrize.
Definition 2.2 (Expressivity of TPOs). Given TPO T X' x X + Z', the space Xr., = (Br,x,y.z. 4A E ) is the set of all bilinear maps we can construct by insert- ing equivariant linear layers in the inputs (Lingx Lingy) ) and outputs (Ling,) ) in Equation 2. We define the expressivity of I with respect to X, Y, Z as the dimension of B X,Y..
In Section 4, we apply this expressivity measure to various tensor product operations.
To optimize the tensor product, one can leverage the struc- ture of the Clebsch- Gorddn coefficients (such as sparsity) to reduce the (((x))|X|²) cost. Still, such approaches can never beat the lower bound of (( In typical analysis, | are (L2)) making the lower bound as O(L)). However, other 'tensor products' claim asymp- totic times lower than ((T.A) while still demonstrating good performance (Luo et al., 2024; Unke & Maennel, 2024). Fundamentally, what these operations do is re- place T X x Y + X @ Y with some other bilinearity T X' x Y + Z'. While these bilinearities are often called 'tensor products', they are not truly tensor products in the mathematical sense. In this paper, we will refer to these bilinearities as tensor product operations (TPOs). Definition 2.1 (Tensor product operations). Let X', Y', z' be vector spaces equipped with actions of G. We refer to any equivariant bilinear map T: : X' x Y + Z' as a tensor product operation.
2.4. Interactability and Selection Rules
A natural question ask about the constructed bilinearities is which inputs can affect which outputs. Since equivariant linear layers can mix irreps of the same type together, an- swering this question for bilinearities between single irress suffices. Using our definition of expressivity, we provide a natural definition for interactability.
Definition 2.3 (Interactability). Let X' XY' - Z' be a TPO. Let (41, Eg, E3) be a triple of irrep types. Define vector spaces V V and V which transform under irrep types f l2 and la respectively. We say [4,, l2, fa] is interactable if the expressivity of T with respect to Vd,, V V is nonzero.
We can analyze interactability by directly analyzing the TPO T. Suppose the irreps within X', Y', Z' are labeled by the tuples (fx, ex). cy) and C22, where & in- dicates the irrep type and is index over its multiplic- ity. A selection rule for T is condition on the labels (l7, ez) which must be satisfied for : x'(xxiex) x y((y,,o) + Z(tzzoz)
2.3. Expressivity and Construction of Bilincarities
It turns out that the runtime savings of alternative TPOs can be understood as coming from a reduction in expressivity. Motivated by the use of tensor product operations con- struct bilinearities, we use the dimension of constructible bilinearities as a proxy for expressivity.
T Tleex.xx)(tv.vr...
to be nonzero. Importantly, note that satisfying selec- tion rule is necessary but not sufficient condition for the restricted bilincarity to be nonzero.
To do so, we must first understand how to use TPOs to construct a bilinearity. Similar to tensor products, we use equivariant linear layers; however, now we can also add linear layers for the inputs'. Hence, we define:
The reason selection rules are useful is that if there is no choice of (ex, ey, such that cx), (fr, and (42, satisfy the selection rules for a given T,, then (2, (z) is not interactable under T. Hence, analyzing se- lection rules helps formally characterize which interactions are excluded by a specific TPO.
Linex: X X' Liney: Y
Linez= Z' Z
parameterized by 0 = (0x, 02) O. Composing these operations, we obtain bilinearity BT.XYY XxY + Z defined as:
Br.xY.ze(xx,)) = LLiiLinnnLiinn,.. T'Linn x. Ling, y).
(2)
2.5. Generalizations to z-fold tensor products and non-irrep basis
UUsullly X transforms as the direct sum of all irreps up to some L. As each irrep has dimension 2L + 1. the dimension of X |X|= -/ 21 + (L + 2 O(L?). "In the case of the CGTP, adding cquivariant linear layers for the inputs is redundant, as the weights can be absorbed into Lineg- Hence, we only use equivariant linear layer Lingz for
We would like to also mention that there have been some re- cent works dealing with a Cartesian tensor basis (Shao et al., 2024; Zaverkin et al.). Specifying a fully expressive equiv- ariant linear layer is more difficult for such representations, however, the corresponding tensor product is usually simpler
the outputs in Section 2.2.
3

---

The Price of Freedom
Input Lincar
Lino,
lrep Decompostion
Ling 2
Octput Linear
x
Line, X
T(Ling x, Lingy y)
Tonsor Product Operation
Figure 1. Overall schematic of an equivariant bilinearity, where the two inputs x and y are passed through linear layers and combined using the tensor product operation 7 to form z. The output irreps are passed through final linear layer. In the case of the CGTP, this would consist of elementwise multiplication of x and y followed by contraction with the irreps Each irrep of particular type is denoted by its color. Linear layers Linex Liney Line can only map between irreps of the same type (indicated by arrows of the corresponding color). However, T can create irreps of a different type (shown in yellow).
coefficients to form output
to compute. In particular, computing multiple v-fold tensor products in succession seems to be asymptotically more efficient in this basis (Zaverkin et al.). However, most archi- tectures only use 2-fold tensor products for which Cartesian based approaches have poor asymptotic scaling (Zaverkin al.). Hence, analysis focuses on the 2-fold case with an irrep basis.
studied. The corresponding selection rules are well known (Varshalovich er al., 1988). For simplicity, we describe them assuming single copy of each irrep in the inputs.
Proposition 3.1 (Selection rule for CGTP). Suppose have irrep labels (f1, 1), 1), and (&s, cz). We must have la I to le for all choices of distinct a, b, E {1, and cz = (21.
In principle, our definitions of expressivity and interactabil- ity can generalize for y-fold tensor products as well. We simply replace the TPO which a fixed bilinearity with fixed u-linearity in Definitions 2.2 and 2.3 and similarly consider attaching equivariant linear layers the inputs and outputs of the v-linearity.
Importantly, interactions which can form a & irrep are placed in separate channels, which means that there can be multiple E irreps created as a result of the CGTP We refer to any triple [ei, , f,, fa] for which (E,, 1), (f2, 1), (fs, (f,, E2) satisfies the selection rule for CGTP valid path.
For alternative choices of basis, we may incur additional costs in the construction of equivariant linear layers. In the case of irrep basis, this construction cheap aaand so we focus on the asymptotic runtime of the TPO. However for other basis choices, this may not be true and one must take care in analyzing the cost of constructing equivariant linear layers for fair comparison.
3.2. Gaunt Tensor Product
The Gaunt tensor product operation (GTP) as introduced by Luo et al. (2024) a bilinearity
(0 +
+ L) x (0
+ + L) + (0 @ ... 2L).
It uses the intimate connection between spherical harmon- ies, SO(3) irreps, and spherical signals, as defined more precisely in Appendix C. The idea is that any rep of form (0...., L) can be interpreted as coefficients for the spherical harmonics, and hence, corresponds to function on S2.
3. Tensor Product Operations
3.1. Clebsch-Gordan Tensor Product
The Clebsch-Gordan Tensor Product (CGTP) is the only TPO considered which is actually tensor product. The change of basis of tensor product reps into irreps is well
In particular, given two (0, 1,... L) reps x and y, let fx ToSphere(x) and fy = ToSphere(y) be the associated signals on 52, Taking the pointwise product of fx and fy

---

The Price of Freedom
rep using Clebsch-Gordan coefficients. After doing so, matrix multiply the tensor product reps. Finally, decompose the resulting tensor product rep back into irreps. Details are provided in Section E.3.
xm
(4)
y
Similar to GTP, MTP only outputs one copy of each possi- ble output irrep. Hence, the outputs the same irrep type get weighted and summed together. However, in contrast to GTP, MTP is not symmetric operation (because matrix multiplication not commutative), symmetric tensor product terms. The selection rules are inherited from CGTP.
Figure Schematic GTP. interpret input irreps as scalar SH coefficients to create spherical signals. We then take pointwise products of the two signals to create a new signal which we de- compose back into scalar SH coefficients.
have anti-
on S gives us a new function fx fy: also on S²_ Then, converting back to irreps gives us the Gaunt tensor product:
Proposition 3.3 (Selection rule for MTP). Suppose we have irrep labels (f1, (4,,,, 1),(E,, 1). We must have la S E + le for all distinct choices of a, b, E (1,2,3}.
x &GTP FromSphere(ff fy)
(3)
Selection rules for GTP can be derived from the Gaunt coefficients (Gaunt, 1929). Note that there is necessarily only a single copy of each irrep type, even in the output. Proposition 3.2 (Selection rules for GTP). For irrep la- bels (e1. 1), (lg, 1), (lg, 1). the corresponding interaction is nonzero only if the following are satisfied:
4. Summary of Asymptotic Runtimes
In Appendix E, analyze the asymptotic runtimes, and in Appendix F, we analyze the expressivity of implementations of the various tensor product operations. The results are summarized in Table For expressivity, we assume we are using the various tensor products to construct bilinearities from spaces X, Y transforming as (0 ... . L) space Z transforming (0 ... 2L).
la S lo + le for any distinet b, c. (1,
e + 6 + la is even.
We would like to highlight the following findings specific implementations. Details are found in Appendix E.
Note particular a cross product corresponds to the [1, path which fails rule 2. This formalizes the intuition that GTP is a symmetric operation and unable to perform anti- symmetric interactions. In Section 6.2, we show that GTP is incapable of solving simple task of classifying chiral 3D structures. However, the impact of the loss of antisymmetric interactions on real datasets remains unexplored.
The often cited O(L)) runtime for CGTP does not leverage all the sparsity in the coefficients (Passaro & Zitnick, 2023). Using all sparsity, we can obtain ((L)) runtime as noted by (Cobb et al.).
We provide and analyze novel implementation of GTP which represents spherical signals using a grid on the sphere as opposed to the original 2D Fourier basis (Luo et al., 2024). Our implementation has the same asymptotics and is 30% faster practice. Futher, our spherical grid implementation unlocks the use of 52 fast Fourier transforms (S2FFT) (Healy et al., 2003) algorithms which leads to asymptotically faster ver- sion of GTP.
3.3. Matrix Tensor Product
Figure 3. Schematic of the process in taking a matrix tensor prod- uct. We embed input irreps into tensor product rep. We interact using matrix multiplication before decomposing the result- ing tensor product rep back into direct sum irreps.
From this table, faster TPOs comes from loss of expressivity. In particular, when normalizing for our expressivity measure, the only true asymptotic speedup comes from implementations lever- aging the fast algorithm for spherical harmonic transforms.
that asymptotic speedups in the
We also analyze another interaction introduced in the new e3x framework in the FusedTensor class (Unke & Maennel, 2024; Maennel et al., 2024). The key idea is that tensor product rep is a matrix, and we can interact two tensor product reps through matrix multiplication.
5. Microbenchnarking Tensor Product Implementations
Matrix tensor products (MTP) first takes each input and embeds the irreps in a single large enough tensor product
For the machine- leamnng practitioner, the fundamental ques- tion of which tensor product choose for their network
5

---

The Price of Freedom
challenging do a fair comparison between different tensor products.
Table 1. Asymptotic runtimes and expressivity of various TPO implementations. Note that when normalized for expressivity, most TPOs have the same asymptoties as sparse CGTP. The only true speed up comes from a fast spherical transform algorithm by Healy et al. (2003).
To test the analysis in Table 1, we implemented all of the tensor products in e e built on top of the JAX (Brad- bury et al., 2018) framework. The just-in-time compilation of JAX (via jax. jit) automatically fuses multiple opera- tions to reduce memory transfer and kernel launch overhead. We also include an unweighted implementation of the matrix tensor product from e3x and more GPU-friendly imple- mentation of Clebsch-Gordan (Sparse) Algorithm 1 in our benchmarking suite. Through this we ensure that all tensor product implementations benefit equally from the various pattern- -ased fusion strategies implemented within JAX's XI.A compiler (Sabne, 2020; Snider & Liang, 2023).
TPO
Expressivity
Runtime
Runtime / Expressivity
CGTP (Naive)
O(1")
O
O
CGTP (Sparse)
O
O
O(L2)
GTP (Fourier)
O(L)
((L*)
O(12)
GTP (Grid)
O(L)
O!L*)
O(L")
GTP (S2FFT)
O(L)
O log' L)
O(L
MTP (Naive)
O(L)
OiL
O(13)
MTP (Sparse)
O(L)
O(L*)
O(LA)
often heavily dependent on the resulting wall-clock time of their training runs. In this section, we show that the practical runtimes of different tensor products is more mu- anced than the theoretical guarantees of Table 1. We focus on the tensor products since these operations are usually the bottleneck in equivariant neural networks. We perform careful microbenchmarking on CPU (some details about the hardware here) and on GPU (specifically, the NVIDIA RTX A5500 and A A00),, and report hardware-agnostic FLOPs, derdddee-dependent GPU utilization, and overall wall-clock time for each tensor product.
However, these fusion strategies and heuristics are still largely centered around dense linear algebra workloads over multi-dimensional arrays. This makes it harder to see out- of-the-box performance gains for operations that don't fit this paradigm (Barham & Isard, 2019). Thus, to analyze how efficiently a given tensor product is being executed on the GPU, we also report average GPU utilization.
Benchmark Settings: We define our benchmarking metrics by counting instructions executed within each kernel run, DRAM read and write memory accesses and throughput. All metrics were collected through NVIDIA's Nsight Com- pute toolkit. This approach has been successfully used in FourCastNet (Kurth et al., 2023) and takes inspiration from the Empirical Roofline Toolkit (Yang, 2020).
To the best of our knowledge, this is the most rigorous benchmarking of such tensor product operations performed till date.
5.1. Summary of Findings
All of the experiments were performed an NVIDIA A5500 with 24 GB of off-chip GPU memory. Our inputs are randomly generated and batch size refers to number of sam- ples used at once. Benchmarks on other hardware (A100, CPU) can be found in Appendix M. Additional details about hardware counters can be found in Appendix L.
We find that the Clebsch- Carschn Tensor Products incur low GPU utilization and higher wall-clock time when compared to the Gaunt Tensor Product and Matrix Tensor Product. However, if we normalize for expressivity (Appendix F), the wall-clock times for the Clebsch-Gordan Tensor Products much lower.
5.3. Results
Overall, our benchmarks highlight the need for carefully balancing expressivity, asymptotics, FLOPs and GPU uti- lization while designing and implementing tensor products.
Asymptotics A FLOPs. The first trend we report is the discrepancy between asymptotics in Table I and compute FLOPs. While the GTPs have lower asymptotic complexity, they end up having comparable (or even higher) FLOPPs than the CGTPs. This is due to different constant scaling factors causing TPOs with similar asymptotics to scale differently with Lmax-
5.2. Methodology
While wall-clock times of tensor products correlate well with downstream inference/training times, they depend on specific code implementations (eg. compiled vs uncompiled code, the use of custom kernels) or GPU compute capability (number of tensor cores, memory bandwidth). This makes it hard to extend benchmarking claims beyond the specific execution environment. While using FLOPP can help get around these limitations (Brehmer et al., 2024), they may not necessarily correspond to the actual compute needed for using equivariant operations given their poor GPU utiliza- tion. The difference in GPU utilization makes it particularly
FLOPs A Wall-clock time. Next, we report discrepancy be- tween FLOPs and wall-clock times. While GTP has higher FLOPs than the CGTPs, their wall-clock time is lower due to their high GPU utilization. For Gaunt (Fourier), code does not leverage sparsity when transforming to a 2D Fourier basis, potentially causing the slowdown. We tried implementing a version that leveraged the sparsity through jax.expe r immnttal... al. sparse but it was slower than
6

---

The Price of Freedom
Total GFLOPSs ba000
Walltime batch=10000
10
Average Throughput for betch=10000
10
10a
10
Waltine (mp)
GFLOPIS
10
10P
10a
10
10
10a
10
1
2
5
6
8
9
10
4
5
6
7
8
10
3
4
5
7
8
9
10
Lmax
Lmax
Lmux
Clebsch-Gordan (Naive)
Gaunt (Grid)
Clebsch Gordan (Sparse)
Mordshh-Gordan
Gaunt (Fourier)
Matrix-TP
Total Normalized GRLOPs batch- 10000
Total Normalized Wa (GPU) for batch=10000
10
10a
1973
G2 W111
10-1
10
10-
10
10
Liaa
Lmax
Figure 4. Top: Analysis of tensor products compute scaling RTX A5500 GPU: Total GFLOPs (Left), Total walltime (Middle), and Average throughput GFLOPs/s (Right). Bottom: Analysis of tensor products compute scaling per path on RTX A5500 GPU: (Left) Total Walltime Expressivity, (Right) Total GFLOPs/ Expressivity. Batch refers to the number of tensor products performed in parallel.
the version report suggesting that its not yet well op- timized. We also report low GPU utilization for Clebsch- Gordan Tensor Product despite having a more GPU-friendly implementation Algorithm 1.
days to profile all the kernels being executed in single model inference call. This is the biggest bottleneck when profiling model runs.
By reporting the GPU utilization and FLOPs these tensor products, we hope to highlight the performance gains rela- tive the device peak that are still left on the table. Recent works (NVIDIA, 2024; Bharadwaj al., 2025; Firoz et al., 2025; Tan et al., 2025) show this is promising direction.
GPU utilization * Wall-clock time. Incase of MTP, we show that high GPU utilization does not always necessarily lead to the lowest wall-clock time due the high compute FLOPs.
Wall-clock time / Expressivity. After normalizing against expressivity defined in Section 4, we find that the Clebsch- Gordan tensor products were the fastest both terms of wall-clock time and FLOPs.
6. Experiments
6.1. Force and Energy Prediction on Organic Molecules
Luo et al. (2024) demonstrates the speedups obtained by GTP in real-world application by replacing the sym- metrized tensor product the MACE (Batatia al., 2022) model (a popular machine learning interatomic potential) with GTP. We redo their experiments the 3BPA (Kováes et al., 2021) and revised MD17 (Christensen & von Lilien- feld, 2020) datasets with our S2Grid implementation of GTP. Our implementation is drop-in replacement for the original GTP implementation using Fourier transforms. While sig- nificantly simpler due the use of pre-existing primitives in e3nn, find that GTP implementation using S2Grid is also ss 30% faster in practice, shown in Figure 5.
5.4. Limitations and Open Questions
Here, we have effectively only benchmarked the 'forward pass' of the tensor product operation in an equivariant neural network. In practice, depending on the dataset and the task, may out the 'optimal' tensor product here may not result in significant performance gains. For example, the lack of expressivity in some of these tensor products may be captured in other operations of the neural network, or may not be important to the task at hand.
Since Nsight Compute does multiple replays on single kernel in order gather performance metrics, can take

---

The Price of Freedom
3BPA
GTP Grid) GTP (Original)
RMSE (Energy)
10-1
(a)
10-7
Accuracy on 3D Tetris
250
500 Epoch
750
1000
1.0
Cletscn Gardon
RMD17
0.8
Gount
GTP (52 Grid)
Accuracy
0.6
10-1
GTP (Original)
RMSE (Energy)
0.4
Lriie
Lrsx =1
0.2
Lreas
Lyax =
10
Loiss =3
Lrwx #3
0.0
Leeas
Lyax
100
200
300
400
Training Steps
250
500
750
1000
Epoch
Figure 5. Training the MACE interatomic force field with the orig- inal GTP (Luo al., 2024) and our faster implementation GTP (Grid) for the same number of GPU hours on the 3BPA and revised MD17 datasets.
Figure 6. (a) The 8 different 3D Tetris pieces, with the first two pieces being mirror images of each other. (b) Training curves networks trained with different tensor products on the 3D Tetris task. The maximum is varied from 4. All of the CGTP networks attain 100% accuracy while none of the GTP networks do.
6.2. Classifying 3D Tetris Pieces
We consider a simple task of classifying 8 different 3D Tetris-like pieces, shown in 6a. Note that the first two pieces are non-superimposable mirror reflections of each other; they are chiral. Given a randomly oriented 3D structure, the network needs to predict which of the tetris pieces it corresponds to.
While there is much focus on improving how runtimes scale with L, there has been little work analyzing how expressivity is affected. Hence we provided the first systematic analysis of the distinctions between different ((3) cquivariant TPOs that exist the literature.
On the theoretical side, we clarify that many new TPOs not formally tensor products. We then introduced measure of expressivity and interactability based on how TPOs are used in practice. We applied this to CGTP, GTP, and MTP analyzing their asymptotic runtimes, expressivities, and in- teractability. Interactability lets us formalize the intuition that GTP unable to represent antisymmetric interactions. In addition, we asymptotic speedups in runtime typically come from losses in expressivity.
We use simple message- message-passing neural network, described in Appendix I, using either the Gaunt and Clebsch-Gordaa Gordan tensor products. Our network architecture is almost identical to that of NequlP (Batzner et al., 2022).
As shown Figure 6, the network very easily able to solve this task with CGTP, but the same network parametrized with GTP is unable distinguish between the two chiral pieces. Adding more channels incorporating the pseudo- spherical harmonics (which have the opposite parity of the spherical harmonics under reflection) did not help. The fundamental failure is the inability to create the le term via lo @ lo + Ie because this the cross product, an antisymmetric operation. Indeed, there is no way create pseudoscalar using the GTP in this setting.
For implementations. we highlight our novel and simpler implementation of GTP which directly uses a spherical grid. This implementation has the same asymptotic cost as the original GTP and perfroms 20% faster practice. Further, we identify that leveraging fast spherical harmonic trans- forms can allow faster asymptotic runtimes. Finally, clarify that the commonly cited OLL)) runtime for CGTP does not use all sparsity in the Clebsch- Gaoscan coefficients and we can in fact achieve O(L5) runtime (Cobb et al.).
7. Conclusion
This work was inspired by the observation that specific anti- symmetric paths were missing in GTP and that paths which result in the same output irrep type are merged together.
We then provide the first microbenchmarks of implemen- tations of CGTP, GTP, and MTP. We observe asymptotic
8

---

The Price of Freedom
gains do not always correspond performance gains in practice. In fact, CGTP leveraging sparsity has the slowest walltime despite having the fewest FLOPs. As expected, we observe better walltime performance for GTP compared with CGTP. In addition, our grid implementation of GTP performs better than the original 2D Fourier implementation. However, normalizing for our expressivity measure, we see that standard CGTP has the best performance, highlighting these new TPOs do not truly speed up tensor products but rather remove degrees of freedom.
References
Aykent, S. and Xia, T. Gotennet: Rethinking efficient 3d equivariant graph neural networks. In The Thirteenth International Conference on Learning Representations.
Barham, P. and Isard, M. Machine learning systems are stuck in rut. In Proceedings of the Workshop on Hot Topics in Operating Systems, HotOS '19, pp. 177-183, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367271. doi: 10. 10.1145/3317550.3321444 URL https: patttp:///oi.org:: 10.1145/3317550.3222
Finally, we replicated the experiments originally performed in Luo et al. (2024). We directly replaced their GTP imple- mentation with our grid implementation and observe a 30% speed increase. Additionally, we demonstrate the lack of antisymmetric operations means classifying 3D tetris pieces using solely GTP is impossible.
Batatia, I., Kovacs, D. P., Simm, G. N. C., Ortner, C., and Csanyi, G. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: /tpp//eeeee.wee......../// forum?i Dorum?id==xpssnE--20
We conclude that while it is important to improve runtimes, it is equally important to properly analyze where the savings come from and the impact of these new operations both in theory and in practice. We believe there are many opportu- nities for creative design of new TPOs. We hope this work can provide a foundation for analyzing these operations.
Batzner, S., Musaelian, A., Sun, L.,, Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozin- sky. B. E(3)-equivariant graph neural networks for data- efficient and accurate interatomic potentials. 13, May 2022. URL httt:s///ooooorgg1.000////1111.. 022-29939-5.
Acknowledgments
YuQing Xie, Ameya Daigavane, and Mit Kotak were sup- ported by the NSF Graduate Research Fellowship program under Grant No. 1743022 Tess Smidt was supported by DOE ICDI grant DE-SC0022215. Mit Kotak and Tess Smidt would like additionally acknowledge support from the National Science Foundation under Cooperative Agree- ment PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interactions).
Beentjes, C. H. Quadrature on spherical surface. Working note available on the website http://people. maths. ox. ac. uk/beemijes/issayyy 2015.
Bharadwaj, V., Glover, A., Buluc, A., and Demmel, J. An efficient sparse kernel generator for o(3)-equivariant deep networks, 2025, URL https: /tt./xixiorrrr///////..// abs 25001.1396..
Impact Statement
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPP programs, 2018. URL http: http://tthhub.ooooo jaae/aaxxaaaa ml jax.
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Brehmer, J., Behrends, S., de Haan, P., and Cohen, T. Does equivariance matter at scale?, 2024. URL https:// rrxiv.org/abs/2410.22
Cen, J., Li, A., Lin, N., Ren, Y., Wang, 7., and Huang, W. Are high-degree representations really unnecessary in equivariant graph neural networks? In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
Christensen, A. S. and von Lilienfeld, O. A. On the role gradients for machine learning of molecular energies and forces, 2020. URL https: http////a//vv//o////////// /bbs/ 2007.09593.
9

---

The Price of Freedom
Cobb, O., Wallis, C. G., Mavor-Parker, A. N., Marignier, A., Price, M. A., d' Avezac, M., and McEwen, J. Efficient generalized spherical cnns. In International Conference on Learning Representations.
P., and Hassabis, D. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589) 2021.
Kondor, R. N-body networks: a covariant hierarchi- cal neural network architecture for learning atomic po- tentials, 2018. URL https: httpp:///axxx.org///aa abs/ 1803.01588.
Dresselhaus, M. S., Dresselhaus, G., and Jorio, A. Group theory: application the physics of condensed matter. Springer Science & Business Media, 2007.
Firoz, J., Pellegrini, F., Geiger, M., Hsu, D., Bilbrey, A., Chou, H.-Y., Stadler, M., Hochnerbach, M., Wang. T., Lin, D.. Kucukbenli, E., Sprueill, H. W. Batatia, I,, Xantheas, S. S., Lee, M., Mundy, C., Csanyi, G., Smith, J. S., Sadayappan, P., and Choudhury, S. Optimizing data distribution and kernel performance for efficient training of chemistry foundation models: case study with mace, 2025. URL https: www.i.orggg/oaa..t.tooo...ooooor abs/ 2504.10700.
Kondor, R.. Lin, Z., and Trivedi, Clebsch gorscan nets: a fully fourier space spherical convolutional neural network. Advances in Neural Information Processing Systems, 31, 2018.
Kovács, D. P., Oord, C. v. d., Kucera, J., Allen, A. E. A., Cole, D. J., Ortner, C., and Csányi, G. Linear atomic cluster expansion force fields for organic molecules: Beyond rmse. Journal of Chemical Theory and Com- putation, 17(12):7696-7711 Dec 2021. ISSN 1549- 9618. doi: 10. 10.1021/acs,jccc 1c00447 URL https: / //bbb./rgg10100/aaass//ttte/eeee
Frank, T., Unke, T., and Muller, K. R. So3krates: Equiv- ariant attention for interactions on arbitrary length-scales in molecular systems. In Advances Neural Information Processing Systems, 2022.
Kurth, T., Subramanian, S., Harrington, P., Pathak, J., Mardani, M.. Hall, D., Miele, A., Kashinath, K., and Anandkumar, A. Fourcastnet: Accelerating global high- resolution weather forecasting using adaptive fourier neural operators. In Proceedings of the Platform for Advanced Scientific Computing Conference, PASC '23 New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400701900. doi: 10. 1 N0.11457359297933593422 URL pestps://doi.oom 10.1155/392999.359523
Frey, N. C., Soklaski, R., Axelrod, S., Samsi, S., Gómez- Bombarelli, R., Coley, C. W., and Gadepally, V. Neural scaling of deep chemical models. Nature Machine Intelli- gence, 5(11): 1297-1355, October 2023.
Fu, X., Rosen, A., Bystrom, K., Wang, R., Musaelian, A., Kozinsky, B., Smidt, T., and Jaakkola, T. A recipe for charge density prediction, 2024.
Gaunt, J. A. Iv. the triplets of helium. Philosophical Trans- actions of the Royal Society of London. Series A, Con- laining Papers of Mathematical or Physical Character, (1116696699:1111111
Lebedev, I. Quadratures on a sphere. USSR Compu- tational Mathematics and Mathematical Physics, 16(2): 10-24, 1976.
1929.
Geiger, M. and Smidt, T. e3nn: Euclidean neural networks. arXiv preprint arXiv:2207.09453, 2022.
Lee, J. H., Yadollahpour, P., Watkins, A., Frey, N. C., Leaver-Fay, A., Ra, S., Cho, K., Gligorijevic, V., Regev, A., and Bonneau, R. Equifold: Protein structure predic- tion with a novel coarse- coarse-grained structure representation. bioRxiv, 2022. doi: 10. 10.000002222......1111133 URL
Healy, D. M., Rockmore, D. N., Kostelec, P. J., and Moore, S. FFTs for the 2-sphere-imroovements and variations. Journal of Fourier analysis and applications, 9:341-385, 2003.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d, 2022.
2022 0112//0008//2222222
Liao, Y.-L. and Smidt, T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference Learning Repre- sentations, 2023, URL https: /ppp/eeveee..ee..... Forum?id=Kwrn Pf
Jumper, J.. Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zidek, A., Potapenko, A., Bridgland, A., Meyer, C., Kohl, S. A. A., Ballard, A. J., Cowie, A., Romera-Paredes, B., Nikolov, S., Jain, R., Adler, J., Back, T., Petersen, S., Reiman, D., Clancy, E., Zielinski, M., Steinegger, M., Pacholska, M., Berghammer, T., Bodenstein, S., Silver, D., Vinyals, O., Senior, W., Kavukcuoglu, K., Kohli,
Luo, S., Chen, T., and Krishnapriyan, A. S. Enabling effi- cient equivariant operations in the fourier basis via gaunt tensor products, arXiv preprint arXiv:2401. /02/6, 2024.
10

---

The Price of Freedom
Maennel, H., Unke, O. T., and Müller, K.-R. Complete and efficient covariants for 3d point configurations with appli- cation to learning molecular quantum properties. arXiv preprint PrXii:2409.027:000 2024.
Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation- and translation-equivarianl neural networks for 3d point clouds, 2018. URL htttt:t/aaxiii..oo abs / 1112008222...
McLaren, A. D. Optimal numerical integration on a sphere. Mathematics of Computation, 17(84):361-385 1963.
Musaelian, A., Batzner, S., Johansson, A., Sun, L., Owen, C. J., Kornbluth, M., and Kozinsky, B. Learning local equivariant representations for large-scale atomistic dy- namics. Nature Commanications, 14(1):579, 2023.
Unke, O. T. and Maennel, E3x: F(3)--quivvriant deep learning made easy. arXiv preprint arXiv:2401.07595 2024.
URL https: / //docs.nvidi..com/... cuda///uu/////aaaa// vareeuuivviiance/// ndex.html.
NVIDIA.
cuequivariance, 2024.
Varshalovich, D. A., Moskalev, A. N., and Khersonskii, V. Quantum theory of ungular momentum. World Scientific, 1988.
Owen, C. J., Torrisi, S. B., Xie, Y., Batzner, S., Bystrom, K., Coulter, J., Musaelian, A., Sun, L.,, and Kozinsky. B. Complexity of many-body interactions transition metals via machine-learned force fields from the TM23 data set. Npj Comput. Mater. 10(1), May 2024.
Weiler, M., Geiger, M., Welling, M., Boomsma, W., and Cohen, T. 3d steerable cnns: Learning rotationally equivariant features in volumetric data, 2018. URL https: http:://axxvvrr/////oo abs httpr//aas//1100000 abbs/180..2257.
Xin, H., Zhou, Z., An, D., Yan, L.-Q., Xu, K., Hu, S.-M., and Yau, S.-T. Fast and accurate spherical harmonics products. ACM Trans. Graph., 40(6):280-1, 2021.
Passaro, S. and Zitnick, C. L.. Reducing so (3) convolu- tions to so (2) for efficient equivariant gnns. In Inter- national Conference on Machine Learning, pp. 27420- 27438. PMLR, 2023.
Yang, C.
Hierarchical roofline analysis: How to col- lect data using performance tools on intel cpus and nvidia gpus, 2020. URL
Rackers, J. A., Tecot, L., Geiger, M., and Smidt, T. E. A recipe for cracking the quantum scaling limit with ma- chine learned electron densities. Mach. Learn: Sci. Tech- nol., 4(1):015027, February 2023.
Yang, C., Wang. Y., Farrell, S., Kurth, T., and Williams, S. Hierarchical roofline performance analysis for deep learn- ing applications. 2020. URL petttss//arxxv.ortt abs /b0./200.05257.
Sabne, A. Xla Compiling machine learning for peak performance, 2020.
Shao, S., Li, Y., Lin, Z., and Cui, Q. High-rank irreducible cartesian tensor decomposition and bases of equivariant spaces. arXiv preprint arXiv:2412.18263 2024.
Zaverkin, V., Alesiani, F., Maruyama, T., Errica, F., Chris- tiansen, H., Takamoto, M., Weber, N., and Niepert, M. Higher-rank irreducible cartesian tensors for equivariant message passing. In The Thirty-eighth Annual Confer- ence on Neural Information Processing Systems.
Snider, D. and Liang, R. Operator fusion in xla: Analysis and evaluation, 2023. URL https: /atpos////xxvvorr// abs/ /bo//201.3306..
Tan, C. W. Descoteaux, M. L., Kotak, M., de Mi- randa Nascimento, G., Kavanagh, S. R., Zichi, L., Wang, M., Saluja, A., Hu, Y. R., Smidt, T., Johans- son, A., Witt, W. C., Kozinsky, B., and Musaelian, A. High-performance training and inference for deep equivariant interatomic potentials, 2025. URL ht https: //aaxxivorrgg/e///e/////...o..ot abs/ 250 4.16668.
TensorCoreWeight. Why the compute through- puts value is different from the actual per- formance (peak performance)?, 2023. URL s/Mtyythhe-comptteettruuuuuu value-is malle-is-differenee frrnttt-rromt-hheeeee

---

The Price of Freedom
A. Notation
Here, we present the notation we use throughout this paper and the typical variable names.
Table 2. Notation used throughout this paper
SO(n)
Group of rotations in n n-dimensional space
((n)
Group rotations and inversion in n--dimensional space
Sa
The 2-sphere, surface defined by x² + y² + x²
Spherical harmonic function of degree l and order m
Y
The collection of spherical harmonic functions of degree for all orders m
Denotes direct
@ Denotes tensor product
X
Denotes a Cartesian product of spaces and also denotes cross products
O
Big O notation
ToSphere
Function which takes in spherical harmonic coefficients consisting of single copies of irrep up to some cutoff I. and converts it into a spherical signal f : S² R
FromSphere
Function takes in spherical signal f S + and converts spherical harmonic coefficients consisting of single copies of irrep up some cutoff L
Table 3. Commonly used meanings of symbols
G
Denotes group
pla)
Representation of a group
A First input space of a constructed bilinearity
Second input space of constructed bilinearity
Output space of constructed bilinearity
V First input space of a tensor product operation (fixed equivariant bilinearity)
W Second input space of tensor product operation (fixed equivariant bilinearity)
V Output space of tensor product operation (fixed equivariant bilinearity)
T Tensor product operation (fixed equivariant bilinearity V x W + 2)
f Typically used to denote irrep type for SO(3). For spherical signals, used instead to denote spherical harmonic degree which naturally indexes multiplicities of irrep types for VSTP/ISTPs
c Indexes multiplicities an irrep type in given space (channel)
j Used instead of to denote irrep type for VSTP and general ISTPs
$ Denotes irrep type of spherical signal (ie. our signal a map f S2 + R2s+1)
B. Irreducible Representations of E(3)
representation of a group G maps each group element g to a bijective linear transformation p(g) E GL(V), where some vector space. Representations must preserve the group multiplication property:
plg h) p(g) p(h) Vg, G
(4)
Thus, the representation p defines group action on vector space V. The dimension of the representation p simply defined as the dimension of the vector space V.
There may be subspaces W C V which are left invariant under actions of p(g) for all g E G. If this is the case, then restricting to W also gives representation p/w (g) GL.(W). If there is no nontrivial W, then we say the representation is an irreducible representation (irrep).
To build E ((3)-qquivariant neural networks, the irreducible representations of E(3) play a key role. Because E(3) is not a compact group, the usual approach has been consider irreducible representations of the group 50(3) of 3D rotations, and
12

---

The Price of Freedom
compose them with the representation in which translations act the identity:
PRRTT) (R)
(5)
This is why translations are often handled in E(3) equivaaria neural networks by centering the system or only using relative
vectors.
The 'scalar' representation Psealar representation of 50(3) is defined as:
Pscalar R
id
YR E 50(3)
(6)
and is of dimension 1 over V = R. Elements of R are unchanged by any rotation R. We call such elements 'scalars' to indicate that they transform under the 'scalar' representation of SO(3). An example of a 'scalar' element could be mass of an object, which does not change under rotation of coordinate frames.
Let (R) E Rax3 be the rotation matrix corresponding to rotation E SO(3). Then, the 'vector' representation of SO(3) is defined as:
Pvector R) = T(R) AR E 8O(3)
and is of dimension over V= TR The name arises from the way vectors in R transform under rotation of the coordinate frame. We call such elements 'vectors' indicate that they transform under the 'vector' representation of SO(3). For example, the velocity and position of object in certain coordinate frame are 'vectors'.
Weyl's theorem for the Lie group SO(3) states that all dinee-aimeeiional representations of SO(3) are equivalent to direct sums of irreducible representations. The irreducible representations of SO(3) are indexed by an integer é 2 O, with dimension 26 + 1. é = corresponds to the 'scalar' representation, while = 1 corresponds to the "vector' representation above. We will often use m, where -6 S m S to index of each of the 2 + 1 components.
We say that quantity x E 26+ is irrep. if transforms as the irreducible representation ('irrep') of 50(3) indexed by l. If Xı is l irrep and X2 an & irrep, we say that (x1, x2) is direct sum of & and &a irreps, which we call (f1, f2) 'rep'. Weyl's theorem states that all reps multiplicity of an irrep in rep is exactly the number of repeats.
direct sum of l irreps, possibly with repeats over 4; xt,)(f), The
An important lemma for constructing equivariant linear layer is Schur's lemma (Dresselhaus et al., 2007).
Lemma B.1 (Schur's Lemma). Suppose Vi, V/2 are irreps of a Lie group any algebraically closed field (such SO(3)). Let é: Vi + V be an equivariant linear map. Then á is either 0 or an isomorphism. Further, if V, = V/2 then & is multiple of identity. Finally, for any two Φ1 Ø2 Vi + 1/2 we must have @1 A62-
This tells us that to construct equivariant linear layers between reps written as direct sum of irreps, only have weights between input and output irreps of the same type and that those weights must be tied together so they give multiples of the identity transformation.
C. Spherical Harmonics
The spherical harmonics are intimately connected to the representations of SO(3) and play a key role in the Gaunt tensor product.
We define the spherical coordinates (r, 0, sp) as:
rrsin cos yp) rsin f sin yp cos
y
(8)
for 0 E 10,T), E [0,
- The spherical harmonics Yem are a set of functions S2 + IR indexed by (6, m), where again 2 0, - S m S Here, 22 ((r, 0, b) = =} denotes the unit sphere.
13

---

The Price of Freedom
Indeed, as suggested by the notation, the spherical harmonics closely related to the irreducible representations of SO(3). Let Ye be the concatenation of all Ye, over all m for given f:
Yi, (0,-)(0,0)) Ye, (0.-0+1(0.0)
()
(9)
Ye,(0,,)
When we transform the inputs Ye(0, o). the output transforms as a irrep.
The spherical harmonics satisfy orthogonality conditions:
J 52
Vegim d52 = 2mim
(10)
where:
J f ds²
/ f(0, (00,)))(0, (0) sin 0d8dp
-- p=0
2m
(11)
The orthogonality property allows us to treat spherical harmonics as basis for functions on S². We can linearly combine the spherical harmonies using irreps approximate arbitrary functions on the sphere. Given (0, 1.... L) rep x())), we associate the function fx S2 R as:
(x(0), x(1)
L
5x
Σ
X (0,m(00)))))
(12)
2=0
m==
The function fx is uniquely determined by In particular, by the orthogonality of the spherical harmonics (Equation 10), we can recover the x!! component:
- fx Y.m ds²
(13)
S2
Thus, can define the operations ToSphere and FromSphere:
ToSphere
FromSphere
X
fx
X
(14)
D. Clrdsnh Tensor Product Details
The most natural map is V x W > V @ W constructed by taking an outer product of the inputs. If the inputs are explicitly written as direct sum of irreps, we can write the tensor product as
2 y
(x((1) ((2))
(15)
x X y'2'yyy
new basis which is the sum of tensor product reps.
The key idea of Srebsch-Gordan tensor product is we can explicitly reduce the tensor product reps back into direct sum of irreps with change of basis. This change of basis is the definition of the Clebsch-Gordan coefficients, giving
(1) @ (f2)
(f1) 8oG ((2))
(16)
where
(x)
OGG y (f2) (Ea)
)
damma
Cyami &2.22 Xm ylal)
(E2)
(17)
82
14

---

The Price of Freedom
Therefore the original tensor product can also be rewritten direct of irreps. This defines the Clebsch-Gordan tensor product (CGTP)
OCG y=
(x GGG y(t)).
(18)
xx
y y(yy
Note that full CGTP really just change of basis from a sum of tensor product reps sum irreps. Hence, we do not lose any information.
E. Runtime Analysis
Here, we provide a detailed asymptotic analysis of runtimes for different tensor products. We consider different settings.
Single Input, Single Output (SISO):
Here we are computing only one path [6,, E2, Eg] where & E(L).
& x 82 - Es
Single Input, Multiple Output (SIMO):
Here we fix l,, Ez but allow all possible irreps generated by the respective tensor products.
4 l2 +
Multiple Input, Multiple Output (MIMO):
Here we only bound the L that the tensor products use but allow full capacity for the input and output irreps. In the case of CGTP, we can have an arbitrary number of copies of each irrep but we assume we only use single copies of each irrep in the input.
Z x W +
In the SISO and SIMO settings, the asymptotic runtimes of different tensor products are directly comparable. However, in the MIMO setting, we lose expressivity in some tensor products. This is discussed more in Appendix F. Note the MIMO setting what one would typically want use in practice. Hence the runtimes reported in Table are for the MIMO setting.
F. Clebsch-Gordan Tensor Product
The tensor product operation defined as:
(+))
OeG x (()))((s) mms
Σ
(g.ma) (2x) x(2))(Et) da.mi
X X
(19)
my=
where C' denotes the Clebsch-Gordan (CG) coefficients which can be precomputed.
E. NAIVE RUNTIME
Let max(f1, E2, Cz). From Equation 19, for each m3, we would need to sum over m,, m2 which range from -( to f and -e to & respectively. Hence, expect O(L)) operations. To compute the values for all m which range from -63 f3. that computing single & & la tensor product requires ( operations.
E. 1.2. OPTIMIZED RUNTIME WITH SPARSITY
f2.mi.22.m is nonzero only if m; + mmg = mg. Transforming to the real basis for the irreps, this condition becomes +-,, + m2 = my. either case for a fixed m, and My, we only ever need to sum over a constant number of m2 rather than (LL) of them as naively expected. Therefore an implementation taking this sparsity into account gives us a runtime of ( This optimization was noted in Cobb et al..
However, the CG coefficients are sparse. In the complex basis for the irreps, (f)..m)
15

---

The Price of Freedom
E.2. Gaunt Tensor Product
The Gaunt Tensor Product (GTP) is based on the decomposition of product of spherical harmonic functions back into spherical harmonics (Luo et al., 2024). In particular, suppose of our inputs x(f1) transforms direct of irreps up to some cutoff L (ic. & ranges from O..... L). We can view these irreps as coefficients of spherical harmonics which gives Xma (f2)) W Ya (0, ạ).
spherical signal Fi F1 (p) Eei,,, X Yt, (6, :p). We similarly construct F2(0, up) Etamm
Taking the product of these spherical signals gives decomposed into spherical harmonics which use define the GTP. This results in
signal F2 (p) Fi (0, (2) F2 (0, sp). This signal can be
Fo s)) = Σ (x((t)
GGTP ( mm
Yy (0,6))
(20)
Eg.ma
E.2.1. 2D FOURIER BASIS
Luo et al. (2024) describe an implementation which decomposes spherical harmonics into 2D Fourier basis in their original paper introducing GTP. This also turns out to be the same implementation in Xin et al. (2021). We describe their procedure here
Note that for any l < L we can always write the spherical harmonics in the 2D Fourier basis:
Ye. (0, (0:
E I.S.
Va,o e((((+up)
(21)
for some coefficients vf.m Hence, any signal Xm can be encoded as
F1(0,0)
Σ
X Vin c(ullto))
Σ
Σ
Σ
x 19 Viiu
eltuatup),
(22)
t=0 m= LSu,es
-LSuNEL
1== m=-i
We identify the encoding
I.
Xa,v Σ
Σ m=-
(()
c.m.
(23)
One observe that the vir sparse and only nonzero when tv. Therefore, finding Xa,v if we have set of irreps is (L) and it is O(1) if we only want one irrep. Because there are O(L²) possible values for u, v, encoding into the 2D Fourier is O(L)) if we encode all irreps up to L or O(L?) if encoding a single irrep.
For functions of 0, yo encoded using a 2D Fourier basis x x2,0 we can compute their product using a standard 2D FFT in Ol log L) time. This gives some output encoded as Ya.v where now u, e range from -2L, 2L to capture all information.
Finally, we decode the resulting function in the 2D Fourier basis back into spherical harmonic basis to extract the output irreps. Suppose L < u L. We can always write
v(u Fol, ()
Σ E=0
Σ
L.m 20y0 Yim (0,5p
(24)
where F.. (0, sp) is some function in the space orthogonal to that spanned by the spherical harmonics. By construction, our output signal is always in the space spanned by the spherical harmonics so the orthogonal parts cancel. Hence we can write
cuupeuub =
Σ
Vu,
Lmy y. Zu, t.m (0, sp)
(25)
-2L50,U52L
22L5,
t= m=-t
1
Σ
Yu C.
(0.))
(26)
t=0
m=-!
-2LSuv52L
16

---

The Price of Freedom
Hence identify:
ym
Σ
Nu,
(27)
Once again, we can note that 2/, must be sparse and only nonzero when t =m. Hence, evaluating the above takes (I.) time since sum (L) values of paired with constant number of v's. If we only extract irrep, then range over (L) values of m giving O(I2) runtime. If we extract all irreps up to 2L this becomes (
E.2.2. GRID TENSOR PRODUCT
Rather than use a 2D Fourier basis, we can instead represent the signal by directly giving its value for set of points on the sphere. Quadrature on the sphere is a well-studied topic (Beentjes, 2015; Lebedev, 1976); in general, O(L?) points are needed to exactly integrate spherical harmonics upto degree L (McLaren, 1963). For this section, consider product grid on the sphere formed by the Cartesian produet of two 1D grids for and y with (L)) points each, for total of O points.
We write:
1.
Fi(0,,98)
Σ Σ
( Xm
Ymm (0,
Σ
Σ
Xm Ne (csm) (Ppk)
(28)
É=0 m=-t
t=0 m=-i
where am is some normalization factor, Pr are the associated Legendre polynomials, and
sin([m/{) 0
m (8)( =
m
(29)
cos(mip)
We note that can first evaluate
1
9m (0,)
Σ t=
x() Ne. Prr (cos(0,))
(30)
where we set P.m 0 m f. we have a set of irreps up to then we the summation and this takes ((L) time. If only have one irrep encode then this takes (1) time. But we also have (L) values of 0, on the grid and O(L) values of to evaluate. This gives ((L3) runtime to encode onto the grid for irreps up L and O )) for a single irrep. Finally evaluating
Fi(0),PP))
Σ m==
Im (000))cs (vpk)
(31)
for set of fp can be done through FFT in (I log L.) time for each 0, giving (IL2 log I.) total. Hence we see encoding onto the sphere takes OLL)) time for irreps up and OLL log L) time for single irrep.
For the multiplication of signals, just have elementwise multiplication F3(Ok, (Px) Fi (p)) F2 spx). Since
there are O(I2) grid points this takes (I.2) time.
Finally, we decode the signal back into irreps. To do so we use the fact that
(())
Σ aj (e), (0,, j.k
(32)
for some coefficients aj. This is essentially performing numerical integration of our signal against spherical harmonic. Once again using the factorization of the spherical harmonics we get
(()) m
=
Σ
Σ
F(O;, 4a (pp)) aj NyNe,mP PC (cos(0,)).
(33)
17

---

The Price of Freedom
The inner sum in parentheses can be computed in OLL) time and need compute for O(L²) values of 0,, m pairs giving a runtime of (T.)). Of course, we note that really is just sines and cosines afternatively we takes OLL² log L) total. Computing the outer sum takes (L) since we sum over (L)) values of j. For single irrep there are (L) values of j giving O(L²) for the outer sum. For irreps up to & there are O(L²) pairs of E, m giving ((L2) runtime for the outer sum. In total, we see going from the grid to the coefficients takes ( log L) for a single irrep and (T,)) for
FFT which
all irreps.
However, turns out that the associated Legendre polynomials have recurrence properties which can be exploited make transforming a set of irreps up L to the grid and a set of irreps up to L back from the grid asymptotically more efficient (Healy et al., 2003). The runtime for this algorithm which will call S2FFT is OL? log" L).
E.3. Matrix Tensor Product
Here describe and analyze the time complexity of matrix tensor product. Let I1, L2 be the max l's of the inputs and La be the max l of the outputs. We pick some = [max [1, I2, Lz)/ so that E @ E when decomposed into irreps can contain all irreps of the inputs and outputs. Note in principle we could always choose larger L.
In the following runtime analysis, we assume L1 = L2 L, i = L, and L3 2L.
E. NAIVE RUNTIME
The first step MTP is to convert our input irreps into tensor product rep using Clebsch- coefficients as
X mi mi
Σ
Cla.ma
x(e)
(34)
f.ms E.m,,Emm E,mz
my=-t
Y 1222
Σ
Cla.mss &mmg
(E) y
(35)
ma
Naively we sum over (L) values of my and need to do the computation for O(L?) possible pairs of mi, m22 This gives (LL) naive runtime for converting a single irrep into tensor product rep. To do so for all irreps up to L the takes ((L))
time.
We can then sum over tensor product reps to create
x Σ x
Y= Σr Y(
(36)
There are OLL) matrices to sum over if we have irreps up to L. Summing matrices takes O(L?) time since our matrices are size (LL) (L)) Hence, this takes (II3) time if we have irreps up to L. If we have single irrep then we do not need to do anything.
We then multiply the matrices giving Z XY. Using the naive matrix multiplication algorithm requires O(L)) runtime.
Finally we can use Clebsch-Gordan to extract individual irreps giving
(x SFTP ()(
Zmi m2
(37)
12
Again, naively sum over O(L²) pairs of mi, m2 and need to evaluate O(L) values of my for O(L*) conversion for single irrep. If we want all irreps up to 2T. then we need (T.4).
E.3.2. OPTIMIZED RUNTIME WITH SPARSITY
Similar to the CGTP, we can take sparsity of the Clebsch-Gordan coefficients into account. We have nonzero values only if +m, + m2 = My. Hence in the encoding step, for fixed m, m2 we only need to sum over constant number of ma instead of (L). This gives a reduction of in encoding to tensor product rep. Similarly in the decoding step, we see for fixed m3 we only need sum over (..) pairs of my, m2. This gives reduction of as well decoding back into irreps.
18

---

The Price of Freedom
E.4. Asymptotic runtimes in different settings
Table Asymptotic runtimes various tensor products for different output settings. The performing tensor products for each oulput settings are highlighted in green. In the MIMO setting, the Gerdah tensor products are highlighted in red to indicate that they can output irreps with multiplicity > 1, unlike the Gaunt tensor products.
Tensor Product
SISO
SIMO
MIMO
Clehsch-ordan (Naive)
O
O(L.)
O(E))
Clebsch- Gordan (Sparse)
(1(2)
O
((19)
Gaunt (Original)
O log
O(27)
O(2))
Gaunt (Naive Grid)
O log 4)
((29)
O(L
Gaunt (S2FFT Grid)
O(17 log L))
O(1.)og L)
O log LA)
Matrix (Naive)
0(29)
O(1))
O(L4)
(Spurse)
0(49)
O(19)
O(13)
F. Expressivity
Here, we analyze the expressivity, as defined Definition 2.2, of the various tensor products. In this case, we assume we use tensor product construct bilinear maps
B (00 (00.............................
2L)
by inserting equivariant linear layers before and after the tensor product. This choice of input and output irreps for our bilinearity is inspired by what is commonly found in practice. is often the case that we have the same number of copies of each irrep type for our features. Since multiplicity of rep only affects expressivity by scaling factor, we focus on the case where there is a single copy of each irrep type up some cutoff L.
By Schur's ' lemma, the total number of inputs and output irreps to our tensor product gives the degrees of freedom for paramterizing the linear layers from 0 @ @ L and to 0 @ @ 2L. There is an additional 2-fold redundancy in overall scaling so #Input irreps + #Ouput irreps gives an upper bound on expressivity.
can only linear maps between irreps of the same type and these maps must be identity. Hence,
F.1 Clebsch-Gordan tensor product
In the case of CGTP, we assume input which is single copy of each irrep up to order L for ((L) irreps in the input. In general, tensor products of single pairs of irreps gives (L) output irreps. There are O(L2) pairs for a total of O(L3) output irreps.
F.2 Gaunt tensor product
In the case of GTP, we note that coefficients of spherical harmonics corresponds to single copies of each irrep. Hence, we have OL)) input irreps. Similarly, in the output there is only one copy each irrep we obtain from the spherical harmonic coefficients. By selection rules, the highest order harmonic we could obtain of order 2L. Hence the number of output irreps also (L).
F.3. Matrix tensor product
the case of FTP, we encode single copies irreps into a tensor product rep (L/2) @ (L/2). Hence there are O(L) inputs. We then perform matrix multiplication which results into (L/2) & (L/2) tensor product rep. But this decomposes into @ @ L giving (IL) output irreps.
G. CGTP Sparse Algorithm
While leveraging sparsity of the Clebsch-Gordee coefficients will improve asymptotic runtime, in practice we would like an implementation which GPU friendly. Here we present an algorithm which uses the sparsity to create a constant number of generalized convolution operations.
19

---

The Price of Freedom
Algorithm 1 CGTP sparse
Require: Irrep 1 (x()), Irrep 2 (())
Clebsch-Gordan coefficients C
Cla.ma
f2.m2
for my = -33..... Ex do:
for my -l...., l do:
A f, 22
fi,,,, ma +my
0
f2, +ma
for mg = -83. -3.............................. L3 do:
for my 1 l do:
mB
mi. .2
fm.mm.mmmmm mg
Ckamms
f2 ms
0
for m3 = -3:.......
Eg do:
for my l
do:
Clams
m,,22
Ei,m,, 1 +
1, ma +mg
0
for mg -l:.... La do:
for my
E,, l do:
D ,l
mn -
my
Chs.ms
f1.m1.12 m2
0
for mg --23.... ................................ Lg do
for m1 11.... l do
(Ea) (1a) mss
(61) (f2)
Zms +- Zmss
mm l2
X y mi +m
for mg = ll...... Eg do
for m1 = ......
4 do
(fa) +- Zms
(1a) B
ma B..m
(f1) (f2)
Zmm
Xmi y m
for mg lg..... ...... Ey do
for m l1 L1 do (fa)
(1a)
Csams (f1) (F2)
Zmg +- Zin
& Xmi y -mm +mag
for m3 = ......., Ey do
for ml = E1 . do
Zmg + Zun
(Ea)
(2s)
+ D m2 Xim --m-- ms
(81) (E2)
return z (fa)
H. Simulating the Fully- Clebsch-Gordan Tensor Product with Gaunt Tensor Products
One way to increase the expressivity of GTP is to first reweight the inputs x,y. That is, we first create
x (e) axx(1) (2)
(38)
(0) (t)
(39)
where at and by are learnable weights. We then perform GTP after this reweighting and extract some output irrep(s) Ez.
That is we get
(x' @GTP (ts).
(40)
The analogous operation is fully connected CGTP. There may be multiple pairs of irreps which give Ey output. We can always weight and sum these to get
Σ
Wee (x) 8cG y ((s))()))
(41)
20

---

The Price of Freedom
where we,e are learnable weights.
However, even if we only care about symmetric tensor products, the weighted GTP operation is strictly less expressive than fully connected CGTP.
More concretely, suppose we have nontrivial e == 2 and l = 4 data in our inputs. From CGTP and the selection rules we see that
(x(2) GG (()))(2) (xx44) Oca y(2))(2)
(x(2) DeG y(4))(2)
(42)
(x(4) GcG y y())
(43)
are all nonzero. In particular, is possible to create a output of
(x GcG y (2)) 2) + (x 8cg (()))
with a fully connected CGTP. However, GTP instead gives single & = 2 output of form
(x((12 Geg (2) + 22 (x OcG (2) 2 + (x Ocg y (2)) + o2 (x1(4) Eog y(a))((2) (44)
where the c's are nonzero coefficients. Note that in order to have nonzero (x(2) 8cg y2)) and (xlt) Oca y(4))(2
contributions, ag, b2, a4, by must all be nonzero. However, that means we must have nonzero (x(2) 8oG y(d) y(2) and (x(4) Oeg y(2))(2) contributions. Therefore weighted GTP not expressive enough output (x(2) Oco y (2)))))) + (x(4) &oc (2)) it will necessarily mix additional terms.
I. Details of Messane-Passing Network
Algorithm 2 LEARN..LLTTINNOPPPOOII
Require: Tensor Product 6,, Number of Channels C (for Gaunt tensor product). procedure LEARNABLETP(X,
if ② = &cG then
return LINEAR(XI &CG X2)
= ØGTP then
for
1,2,..., C do
x LINEAR (9) (x1)
( LINEAR, (1)
(x2)
LINEAR (1) (x{ @GTP X2
(1)
return CONCATRNATR( (x! (1) l (1.2...., C}))
return LearnableTP
In Algorithm 2, create learnable (ie, parametrized) variants of the purely functional tensor products. For the Clebsch- Gordan tensor product @cG. we simply add linear layer to its output. For the Gaunt tensor product @GTP, we create multiple channels, perform the tensor product channel-wise and then concatenate all irreps. This allows the output to have irreps of multiplicity > 1, even with the Gaunt tensor product. We set the number of channels C' as 4 in all experiments with the Gaunt tensor product.
In Algorithm 3, we use these learnable tensor products in simple message-passingg network, very similar to NequlP (Batzner al., 2022).
J. Tetris Experiment Details
The pieces are normalized such that the side length of each cube is 1. When represented as a graph, the center of each cube is a node. We instantiate the network with dmax = 1.1 so that the centers are connected only to its immediately adjacent centers. The network finally outputs x De + x Oo irreps. (As reminder, Oe are scalars and Oo are pseudoscalars).
21

---

The Price of Freedom
Algorithm 3 Operation of our Message Passing Neural Network
Require: Graph G, Message Passing Iterations T, Cutoff dmmax Spherical Harmonic Degree e, Tensor Product Compute neighbor lists for each node in G:
(u, E - | rull S dana
Create LEARMBLLETENSORPROOUU from @. for E V do:
(0) t [1]
for t = 1,,,.... . do:
for v E V do:
[NO EVEN(Y) MLP([ rell) x LEARVABLETENNSOOPPOUUCUC.. (t-1) Ye(r, re))
(1) < GATE(HE))
((0) CONCATENATE(1M5 (1-)) 1)
hv < LINBAR(h?))
(T) return {ht LVEV
The logits and predicted probabilities then computed by:
to x (0o) x (0n)o
(45)
1 x(00) x x (00))
(46)
1. = x (0c), for i 2 2
(47)
P: softmax(4))
(48)
is clear that defining the logits in this manner preserves the rotational and reflection symmetries. The predictions are clearly invariant under rotations (as they are irreps), and under reflections: (0o) -> (000)) but (0e): x (0e).
We set the number of message-passing steps 7' be 3, to allow the interactions lo 1o le and then le 1o 0o, so the pseudoscalar can be created. The degree of spherical harmonics is kept as l 4. The irreps of the hidden layers restricted to some cutoff' L, which is varied from 1 to 4 vary the expressivity of the network. We train the model with the Adam optimizer with learning rate 0.01 to minimize the standard cross-entropy loss to one-hot encoded labels for the 8 pieces.
K. Proof of Theorem 3.2
Proof. It is known (Gaunt, 1929) that:
Y
Ymm x Σ
0,0002000 Jom Y.m
j
Hence, we see that
(x) @GTP y (O2))(0) x f.4. Jimm Ja (1)) 8eG (1))(())
Ct.0
For selection rule 1, this inherited from the selection rules of CGTP, For selection rule 2, this follows from the selection rules for the 2,20 0,0,0 term which is nonzero only when E + lg + is even.
L. Details of Benchmarking Setup
L.1. GPU
Wall-Clock Time: Total wall-clock time reported as the sum of gpu_.time executed for that function call. We confirmed that this closely matches the jax wall-clock time with some profiling overhead due to kernel replays.
uppu_ttmmedduratiin www.adiiioo..... sum metric for all the kernels
22

---

The Price of Freedom
We obtain the list of all FLOPS generating instructions that are executed within each kernel on the GPU. We then count all instructions and multiply them by their specific weighing factors (the number of FLOPS in each instruction). The weighing factor for a multiplication addition is 1 and the weight for multiply-and-add (FMA) is 2. For Tensor Core instructions on an Ampere GPU, the peak performance is 1024 FMA/cycle/SM and since FMA = 2 FLOPS, we get 2048 CensorCooeWeiihtn 2023). We check the consistency of this weighing factor by running a matrix-multipication benchmark. Further details can be found in Empirical Roofline Toolkit (Yang et al., 2020).
FLOPs:
Average Throughput (FLOPs/s): We first calculate the Tensor Cores and CUDA Cores Throughput for every kernel by dividing the FLOPs for that kernel by the time taken execute that kernel measured using gpu..t aameduuraiion sum. We report the average throughput for range of kemels launched for given tensor product function call.
Average DRAM Bandwidth (GB/s): For every kernel, we take the sum of dram bytes read. sum and dram.bytes.wr wwweee.wwweeeee..ee..e...ee.ee.e sum and divide it by the kernel execution time. We report the average bandwidth for a range of kernels launched for a given tensor product function call
GPU: We gathered the GPU plots on an NVIDIA RTX A5500 with default JAX precision (F32), running the CUDA driver version 550.90.07 and CUDA toolkit version 12.5. We used Nsight Compute 2024.2.0.0 build 34181891 and JAX version 0.4.30. JAX automatically switches to TF32 wherever appropriate.
Precision
Metrics
Weight Factor
FP64
ptt. sum
Sm sass w Sm m...asshhheeccnnttcecrrrrre.pppp
sm hr..aaahhhaaannnncccccceeeeeeeee
FP32
sm sass ott.aaaaaaaaaaaaaaaaaaaaaooaaaoa sm hmg.aaas.hhcccc..........rreooe.
sm hmw.ssasshhhaaa.....c....u.ooaas
FP16
sm sua.saasshhaaaallac.c..e.eeeeee. sm ot..saashhhaaa.nnttccceeeeeeeoer
Tensor Core
sm sn__instexxcateedppppe.ccce
2048
Table 5. FLOPS weights for various Nsight Compute metrics.
L.2. CPU
Wall-Clock Time: The elapsed time after compiling using jax. jit. To enable accurate measurements, calculate the mean wall-clock time for 100 rounds after performing 10 warmup rounds.
We use Linux's 's perf profiler to directly count the number of FLOPs through fp.ret.sseeavx.ops... all metric (specific to AMD).
FLOPs:
We had to skip throughput and bandwidth because we could not colleft hardware-counters for them using perf
CPU: The CPU plots were gathered an AMD EPYC 7313 16-core 3.7 GHz processor with default JAX precision.
M. Additional Benchmarks
23

---

The Price of Freedom
Total
Waltime batch=10000
Walkime
10a
10
10a
WAlinn (60)
10
100
10
10a
10a
10*
10a
10a
10
2
6
10
5
8
9
10
8
9
10
Total Normalized Walltime (CPU) for SISO batch-10000
Total Normalized Walttime (CPU) for SIMO batch=10000
Total Normalized Walltmme (CPU) for MIHO batch-10000
10a
10
10
10
Wetmm 2700
10
10
10c
10-2
109
10"
10-7
10
Lana
Total GFLLDSs (Linux Peril SISO batch 10000
Total GPLL (Linux Peril SIMO batche 10000
Total GPLOPss [Linux Perfi for MIMO batchee 10000
10a
10:
10a
10
10'
101
01.0%
10a
10
10a
10a
10
10-2
10-1
10
Lmw
Leaav
Lecsv
Nommalized CFLOPs Parf) S.50 batch=10000
Normalized GRLOPs SIMO batch-10000
Total Normalized GRLOPs (Linux MIO batch-10000
10
10
10a
107
10a
10
10a
10a
10a
10
10P
10-1
10
10
10
10
8
9
10
8
10
2
8
10
Figure 7. Analysis of SISO, SIMO and MIMO (Table 4 for input settings) performance for different tensor products AMD EPYC 7313
Total walltime (top row), Total normalized walltime (second row), Total GFLOPs (third row) and Total normalized GFLOPs (bottom row). Note that MTP only supports MIMO, We had to skip values due profiling emror..
24

---

The Price of Freedom
Walliime (GPU) for batch=10000
Walltime SIDD betch=10000
Total Walltime (GPU) batch-10000
10
10
Ole mmd
10
10
100
10"1
10-1
10
10
2
6
9
10
8
10
2
2
6
9
10
Total Normalized Walttime (GPU) for SISO batch=19930
Total Normalized Valitime (GPU) for SIMD batch=10000
Total Normalized Walltime (GPUU) for MISO batch 10000
10
10
10
Wi 1092
Waltoe twet 9177
10a
100
10
Walim
10-3
10-1
10"2
10
5
10
10
Leea
Leaa
Leoo
GFLOP% for batch-10000
GRLO% SIMD batch- 10000
Total GFLOPs banch- 10000
10a
102
10a
10:
00.00
10
Crai
10a
GPLORS
1.0
10
10-2
30-1
10"2
10
10
1mw
Leeav
Normalized GFLOPs betch-10000
Total Normalized GFLOPs for batch-10000
ized GFLOPs MINO batch-10000
10a
10'
10
10a
10a
100
10-1
10
10-2
10
10-2
10-3
5
6
8
9
10
7
8
10
2
4
5
8
9
10
10
Average Throughout for SISO for batch- 10000
10
Average Throughout for SIMG for batch=
10
Auerago rage Throughhut for NIMO for batche 10000
PRO 1900 THEER
105
10a
10'
GFLOPIS
10
GELOR!S
10*
GFLOPIS
10
10a
10a
10
10²
107
10'
10
2
8
10
10
Lnww
Lynor
Figure 8. Analysis of SISO, SIMO and MIMO (Table 4 for input settings) performance for different tensor products on RTX A5500: Total forward walltime (top row), Total forward normalized walltime (second row). Total GFLOPs (third row). Total forward normalized GFLOPs (fourth row) and Average forward throughput (bottom row). Note that MTP only supports MIMO. We had to skip values due to profiling emors.
25

---

The Price of Freedom
Walltime (GPU) batch=10000
betch=10000
Total Walltime (GPU) batch-10000
10
10a
Owl
100
10-1
10-1
2
10
8
10
2
9
10
Total Normatized Walitime (GPU) for SISO batch=19900
Total Normalized Walitime (GPU) for SIMD betch=10000
Total Normalized Walltime (SPU)) for MIMO batch 10000
10"
watooe twe 2172
10F
Weltme (1wel
10-1
10
10"2
5
10
10
Livax
Leoo
for batch-10000
GRLO% SIMO batch-
Total GRLOPs banch 10000
10'
10a
10
10
109
GraI
GPLOPS
100
10-1
10
10
Leeav
Normalzzed GFLOPs betch-10000
Total Normalized GFLOPs batch-10000
ized GFLOPs
10
10a
601090 K0
GILONI
10a
GRLPSS km
10a
10-1
10-3
10-:
10
10-2
8
9
10
7
8
10
2
4 5
8
9
10
10
Average Throughout for SISO for batch- 10000
10
Average Throughout for SIMG for batch-
10
Aue rage Throughhut for NIMO for batch 10000
reo 198 1H.O
PRO 19C0N THEER
10
10a
10'
GFLOPIS
10
10*
GFLOPIS
10
10a
10a
10a
10a
107
10'
10
10
knww
Lyna
Figure 9. Analysis of SISO, SIMO and MIMO (Table for input settings) performance for different tensor products on A100 GPU, showing Total forward walltime (top row). Total normalized forward walltime (second row), Total forward GFLOPs (third row) and Total forward normalized GFLOPs (fourth row) and Average forward throughput (bottom row). Note that MTP only supports MIMO. We had skip values due to profiling error..
26

---

The Price of Freedom
Backward Walltime (GPUI for SISO batch-10000
10
Backward Walltime (GPU) for SIMO batch-10000
10'
Bockward Walttime (GPU) for HINO batch-10000
10a
10
Wwt
10a
10a
10
10
10a
10-1
-1
Lmar
Lreay
10
10
10a
10a
109
10
10
10
6
2
10
10
4yax
Figure 10. Analysis of SISO, SIMO and MIMO (Table 4 for input settings) performance for different tensor products RTX A5500:
Total backward walltime (top row). Total backward normalized walltime (bottom row). Note that MTP only supports MIMO
27


Reviews:
[{"review_id": "IPbX0UX9Et", "reviewer": "Reviewer_a8uP", "summary": "This paper investigates tensor product operations in E(3)-equivariant neural networks, an \nimportant class of models for 3D modeling tasks, which have been recently proposed as a \nfaster alternative to the standard Clebsch-Gordan tensor product. In particular, the authors \nintroduce measures of expressivity and interactability, and analyze the runtimes (FLOPs and \nwall-clock time), GPU utilization, expressivity, and asymptotics of various operations. Finally, \nthey  provide  a  novel  implementation  of  the  Gaunt  tensor  product  without  sacrificing \nasymptotics, which is shown to be faster in practice.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough feedback on our work. We appreciate that the reviewer finds our work well-organized and presented, the experiments thorough, and provides important insights to practitioners.\n\nRegarding the weaknesses\n## Weaknesses\n1. This is a great question. The motivation is to highlight the difference in runtime improvements caused by cleverly reducing degrees of freedom from improvements caused by actually making tensor products faster. In particular, we also emphasize that the output spaces of these algorithms are different, making a direct comparison unfair. \n\n2. We provide a framework to analyze the expressivity-vs-runtime tradeoffs in popular tensor product operations (TPOs) in equivariant neural networks. These tradeoffs are non-obvious, and can help other practitioners realize that these operations are not necessarily equivalent to each other. Importantly, we highlight that our work has revealed that most existing TPOs actually only get improvements from reducing degrees of freedom. Our connection to $S^2$ fast Fourier transforms provides the first (though for now impractical) algorithm which has an asymptotic improvement on the runtime/expressivity ratio.\n\n     Our microbenchmarking efforts help highlight the different algorithmic tradeoffs. MTP focuses on maximising GPU utilization with a more matrix-multiplication friendly algorithm which ends up costing more FLOPs and hence higher overall runtime. CGTP on the other hand is FLOPs efficient but suffers from poor GPU utilization. GTP offers a balance between both.\n\n\n3. This is a great question as indeed, different implementations could have drastically different walltimes in practice and we tried to mitigate implementation differences as much as possible (e.g. implementing all of the algorithms in JAX). We would like to highlight that it is not the results themselves but the method of benchmarking that is important.\n\n     Our microbenchmarks provide valuable insight on how different implementations can be improved. In particular, the discrepancy between FLOP counts (which are invariant) and walltime indicates potential for significant acceleration (eg. custom kernels) by better utilizing the GPU.\n\n## Other comments and suggestions\n1. Good catch! Indeed we do need an algebraically closed field for the maps to be identity. In the case of $SO(3)$, the irreps over complex vector spaces are real irreps.\n2. Thanks for catching this! We meant to add 'AMD EPYC'.\n3. Thanks for catching this!\n\n## Questions\n1. As discussed above, the normalization for expressivity accounts for the fact that the output spaces of the different TPOs are not identical. We agree that there can be many ways to account for this difference.\n2. Good question! The algorithm we used for CGTP-sparse (Appendix G) while being able to exploit the sparsity, introduced a lot of overhead due to the conditional logic leading to poor GPU utilization. As is seen in recent efforts (https://developer.nvidia.com/blog/accelerate-drug-and-material-discovery-with-new-math-library-nvidia-cuequivariance/, https://arxiv.org/abs/2501.13986v2), a more GPU friendly algorithm can improve the runtime.\n3. This is an excellent question! In most of the popular datasets, we strongly suspect antisymmetric operations to have minimal impact. Loosely speaking, this is because irrep types of commonly predicted quantities such as energy or forces can be constructed purely from symmetric tensor products of the input irrep types. We plan to explore the impact of antisymmetric tensor products further in future work.", "claims_and_evidence": "The main claim in this paper is that, although several operations have been proposed to achieve E(3) equivariance efficiently, efficiency, in many cases, comes at the cost of expressivity. This claim is supported by appropriate expressivity measures, while runtime is examined by carefully designed benchmarks illustrating discrepancies from the theoretical bounds. Additionally, the improvement they proposed is shown experimentally to be indeed a faster version of the GTP tensor operation.", "methods_and_evaluation_criteria": "The evaluation criteria (expressivity and runtime) are reasonable, since it is a typical trade-off in equivariant machine learning. Additionally, the proposed implementation (although not strictly a new method) makes sense as it shows runtime improvements.", "theoretical_claims": "I have looked at the proof of Theorem 3.2, which appears correct.", "experimental_designs_or_analyses": "I found the experimental designs/analyses carefully designed, with several factors tested (asymptotic, FLOPs, GPU utilisation, wall-clock time).", "relation_to_broader_scientific_literature": "E(3)-equivariant neural networks are an important class of models for 3D modeling tasks, \nwith the tensor product being the key non-linear operation, in several such architectures. \nThese have found application in, e.g. molecular modeling and physical simulations.\nThe authors consider recently proposed several alternatives to the standard Clebsch-Gordan tensor \nproduct  in  the  literature,  which  offer  improved  runtimes, thus well-contextualising their work.", "essential_references_not_discussed": "Nothing to note.", "other_strengths_and_weaknesses": "**Strengths**\n \n1.  The paper is well-organized and well-presented (although I believe that some prior knowledge of the field is required to carefully follow it).\n2. The experimental section is thorough, clearly delivers evidence for the claims made, and illustrates the differences across different tenor product operations.\n3.  The paper provides important clarifications and clears possible misunderstandings regarding the operations that have been proposed in the literature (Table 1 is quite useful). Additionally, it provides important insights to practitioners by thoroughly analyzing the asymptotics and expressivity of various tensor product implementations. These insights can help pave the way for developing and analysing new tensor product operations. \n \n**Weaknesses**\n \n1.  Normalizing runtimes for expressivity does not seem sufficiently justified, or more precisely, does not necessarily give sufficient guidance to practitioners on what to choose. For example, generalisation or optimisation might be equally or more crucial for certain real-world tasks, and therefore it might be the case that the cheaper operation should be chosen.\n2. Although the paper makes important clarifications for the field, I am a bit concerned that it lacks novelty, since it mainly analyses existing methods without necessarily providing actionable guidelines. \n3.  It is not evident how much of a fair comparison is made in benchmarking the various tensor product implementations. Could it be the case that different implementations might (significantly) alter the experimental metrics (as with GTP)?", "other_comments_or_suggestions": "1.  In  lines  59-60,  the  authors  state  that  linear maps between equivalent irreps are \nmultiples of the identity. To my knowledge, this is not the case for real representations \nof  an  arbitrary  group.  However,  this  does  hold  by  Schur’s  lemma  for  complex \nrepresentations. \n \n2.  Typo in lines 249-250: “(some details about the hardware here)”. \n \n3.  Typo in Equation (19).", "questions_for_authors": "1.  Can you justify the contribution of normalizing for expressivity? While the defined \nmeasure  of  expressivity  and  asymptotic  runtimes  make  sense  on  their  own,  for \nexample, the conclusion about the ratio runtime/expressivity is not really clear. \n \n2.  In  Figure  4,  what  is  the  reason  for  CGTP  (Sparse) having low FLOPs but high \nwalltime? \n \n3.  Have  you  considered  any  examples  of  real-world  datasets  where  leveraging \nantisymmetric  interactions would improve performance? It would be interesting to \nevaluate the various tensor product operations on such a problem.\n\n**Note**. The paper is technically sound and provides certain valuable insights. I am currently (hesitantly) leaning towards acceptance, but I hold reservations due to my concerns about novelty, and perhaps also relevance/impact (since the improvements proposed are mostly related to implementation).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tHi2038yHD", "reviewer": "Reviewer_t24S", "summary": "This paper aims to advance the fundamental understanding of equivariant neural networks by studying different mappings from the product of vector spaces into tensor product spaces (including the well-known CG tensor product), which serve as the building blocks of expressive equivariant architectures. In particular, it introduces a straightforward but novel measure to compare the expressivity of different mapping designs. Additionally, the paper discusses implementation details and provides a comprehensive analysis of each operation in terms of FLOPs, wall-clock time, GPU utilization, and expressivity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your careful reading of our work and helpful feedback. We appreciate that the reviewer finds our measure of expressivity reasonable and our runtime evaluation rigorous and comprehensive.\n\nRegarding the weaknesses\n1. > Contribution of the new GTP implementation appears limited\n\n     We very much agree that spherical harmonic transforms (which we refer to as $S^2$ fourier transforms) have been extensively studied and we cited Healy et. al. 2003 which improves upon the seminal work by Driscoll and Healy 1994 (which we will add as a citation). However, much of the work on S2FFTs involves numerical precision issues for super high $\\ell$s when using asymptotically fast algorithms. This is beyond the scope of our work as most equivariant networks are still limited to small $\\ell$. Instead, we just point out asymptotic improvements exist if equivariant networks scale to super high $\\ell$ in the future. Importantly, we emphasize that this connection gives the first asymptotic runtime improvement over sparse CGTP that is not simply a consequence of reduced expressivity.\n\n     Our actual grid implementation uses a much simpler seminaive algorithm described in Healy et. al. 2003 and Dricoll and Healy 1994.\n\n2. Our experiments in Section 6.1 (Figure 6) mainly serve to show a drop-in replacement of Fourier GTP with Grid GTP leads to speedups. These are done on standard benchmarks of 3BPA and rMD17 and with an equivalent number of GPU hours. If the reviewer is concerned about actual training performance, comparisons against CGTP can be found in the original GTP paper (https://arxiv.org/abs/2401.10216).\n\n     Our experiment in Section 6.2 (Figure 7) is a simple demonstration that lack of “antisymmetric” interactions can have tangible limitations on model performance. If the reviewer is concerned about how the lack of such interactions affect real world datasets, which is an excellent question, we actually strongly suspect most common molecular datasets are minimally impacted. Loosely speaking, this is because irrep types of commonly predicted quantities such as energy or forces can be constructed purely from symmetric tensor products of the input irrep types. We are actively exploring the impact of antisymmetric tensor products in future work.\n\n# Other Comments or Suggestions\nYour understanding is correct, the $c_*$ are meant to separate the multiplicity of irreps. For CGTP, we wrote $c_Z$ as a tuple to try to highlight which inputs contributed, however in practice, this tuple is simply embedded as an integer. We will explain this more clearly in our revised manuscript.", "claims_and_evidence": "The experiments demonstrate the runtime and expressivity, aligning well with the theoretical analysis.", "methods_and_evaluation_criteria": "The proposed measure of expressivity for different mappings into tensor product spaces seems standard and reasonable. The runtime evaluation is rigorous and comprehensive, whereas the assessment of expressivity appears relatively adequate.", "theoretical_claims": "The theoretical claims appear to be correct overall, as well as their proofs.", "experimental_designs_or_analyses": "The experimental design is sound.", "relation_to_broader_scientific_literature": "Potentially, building expressive yet efficient equivariant neural networks using tensor product operations is a challenging and important problem. I believe this paper contributes to the development of such networks and, consequently, to their application.", "essential_references_not_discussed": "The novelty of the new GTP implementation seems weak. First, the discrete spherical harmonic transform has already been extensively studied. In particular, this transform has a well-known convolution theorem, which appears to be closely related to the proposed implementation. However, it is unclear whether the authors have adequately addressed the overlap with existing works or provided relevant references from related fields.\n\n[1] Blais, J. R. (2008, June). Discrete spherical harmonic transforms: Numerical preconditioning and optimization. In International Conference on Computational Science (pp. 638-645). Berlin, Heidelberg: Springer Berlin Heidelberg.", "other_strengths_and_weaknesses": "Weaknesses:\n- The contribution of the new GTP implementation appears limited (see my comment above). \n- The evaluation of expressivity is constrained by the simplicity and small size of the dataset used. It would be highly beneficial to see a discussion of the trade-offs when applied to a large-scale real-world dataset. I believe the implications and potential impact of this work are limited by its current application to a relatively simple dataset.", "other_comments_or_suggestions": "I believe I understand the definition of the selection rule, where $c_*$ appears to represent the multiplicity of irreducible representations in different $G$-spaces. However, in Proposition 3.1, $c_Z$ is a tuple rather than a single number. It seems that the goal is to separate the interaction outputs—essentially, different tuples correspond to different indices. If my interpretation is correct, I believe the definition and related explanations could be revised to improve clarity for a broader audience.", "questions_for_authors": "The theoretical results and discussion on runtime appear rigorous to me. However, the similarities between the new implementation and existing methods should be clarified. Additionally, the real-world trade-offs remain unclear. The experiments in Figures 6 and 7 do not cover all the different methods and do not consider the performance with normalization in terms of runtime. I understand that such trade-offs might not be consistent across different tasks, but I believe that providing more experimental results on existing real-world benchmarks comparing different methods (CGTP, GTP, MTP) with normalization would significantly enhance the contribution. I would consider raising my score if the authors address these concerns.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S7gr3xmkga", "reviewer": "Reviewer_W5ML", "summary": "The paper presents a comprehensive analysis of tensor products and tensor product operations used in $E(3)$-equivariant models based on spherical tensors, including the Clebsch-Gordan tensor product (CGTP), Gaunt tensor product (GTP), and Matrix tensor product (MTP). The authors introduce expressivity and interactability as key measures to characterize these operations, demonstrating that speedups presented in prior work often come at the cost of expressivity. Furthermore, the paper proposes a more efficient implementation of GTP using a spherical grid, demonstrating practical performance improvements.\n\n## update after rebuttal\n\nThe authors addressed all my questions. However, I maintained my score at 4 as I remain uncertain about the broader impact of the work. Overall, I recommend this work for publication.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your careful reading of our work and positive feedback. We appreciate that the reviewer finds our evaluation criteria and experiments well aligned with the problem, the theoretical claims well supported, and discussion on expressivity a clear explanation for why certain TPOs fail to capture specific interactions.\n\nRegarding the comments and suggestions\n1. We thank the reviewer for bringing our attention to recent work in the Cartesian basis. This is a very interesting body of work which we will definitely mention in our revised manuscript.\n\n     We would like to point out that our focus is on irrep based frameworks as in that case maximally expressive linear layers are easy to parameterize (as briefly explained in section 2.1). For the case of Cartesian tensors, it is easy to perform tensor products and remain a Cartesian tensor. However, especially at higher rank, it becomes difficult to parameterize fully expressive equivariant linear layers. Regarding the specific sources:\n\n     It seems the interaction in https://arxiv.org/abs/2306.06482 is limited to rank 2 tensors and is a special case of MTP. https://arxiv.org/abs/2412.18263 seems to have significant work on decomposing Cartesian tensors to irreps in order to construct equivariant linear layers.\n\n     It seems https://arxiv.org/abs/2405.14253 introduces a Cartesian tensor product which is also a true tensor product. However, this method has poor asymptotic scaling for 2-fold tensor products which was the focus of this work. Instead, their speedups are for small $\\ell$ or large $\\nu$-fold tensor products (performing multiple tensor products in succession like in MACE). We believe our framework can also be adapted for analyzing large $\\nu$ by replacing the fixed bilinearities with fixed $\\nu$-linearities and defining expressivity as the dimension of constructible $\\nu$-linearites by inserting equivariant linearities in the inputs and output.\n\n     We will add a discussion of these points in our final manuscript.\n\n2. This is a good question. We would expect the same optimizations that were accelerating the forward pass to work for the backward pass. We were able to confirm this hypothesis by running a short experiment on the different input/output irreps settings and are able to see similar trends as the forward time. The plots and experiment details are included below. In particular, we create a random $z$ and then benchmark how long it takes to compute $\\nabla_x(T(x,y)-z)^2$.\n\nhttps://anonymous.4open.science/r/PriceofFreedom-A835/benchmarking/plots/png/walltime_bwd_gpu_MIMO_10000_RTX-1.png\nhttps://anonymous.4open.science/r/PriceofFreedom-A835/benchmarking/plots/png/walltime_bwd_normalized_gpu_MIMO_10000_RTX-1.png\nhttps://anonymous.4open.science/r/PriceofFreedom-A835/benchmarking/plots/pof_legend.png\n\n3. We will definitely mention other basis choices in our revised manuscript and discuss Cartesian tensors in particular. The content in Section 2.1 motivates why we focus on an irrep basis.", "claims_and_evidence": "The paper's main claim is that tensor product operations proposed as an alternative to CGTP achieve computational efficiency at the cost of expressivity. This claim is supported through theoretical analysis and numerical experiments. The authors comprehensively compare asymptotic computational complexities and runtimes, showing that CGTP retains the highest expressivity but can be more computationally expensive. Furthermore, the paper provides empirical evidence that GTP cannot represent antisymmetric interactions using a 3D Tetris classification task.", "methods_and_evaluation_criteria": "The proposed methods, such as the spherical grid implementation of GTP, and evaluation criteria, including 3BPA and 3D Tetris data sets, are well-aligned with the problem. The authors carefully assess theoretical expressivity and runtime complexity as key metrics, including rigorous numerical experiments.", "theoretical_claims": "The theoretical claims regarding the expressivity of different tensor product operations are well-supported. The derivations of runtime complexities and selection rules for each tensor product align with established results in the literature. The discussion on expressivity offers a clear explanation for why certain tensor products fail to capture specific interactions.\n\nHowever, the statement that CGTP is the only true tensor product is arguable, especially given that the study focuses exclusively on spherical tensors. For example, an alternative tensor product can be defined in the Cartesian basis, an area with a smaller but gradually growing body of literature that deserves mention in this context; see, e.g., https://arxiv.org/abs/2306.06482, https://arxiv.org/abs/2405.14253, https://arxiv.org/abs/2412.18263, and references therein.", "experimental_designs_or_analyses": "The experimental design is sound, with well-chosen benchmarks and clear comparisons across tensor product operations. However, the benchmarks focus primarily on the forward pass, while the section title for the 3BPA data set states that atomic forces and energies were evaluated. Given the importance of gradients in applications such as machine-learned force fields, it would be beneficial to analyze backward pass performance if possible or clarify whether this analysis has already been conducted for the 3BPA data set.", "relation_to_broader_scientific_literature": "The paper fits well with the broader literature on equivariant neural networks and tensor product operations. It thoroughly discusses prior work on Clebsch-Gordan coefficients, corresponding tensor products, and equivariant architectures based on spherical tensors. However, it does not address tensor products in alternative bases, e.g., the Cartesian basis. While the spherical basis is currently the dominant choice in the community, emerging studies on Cartesian representations could provide additional context for the paper’s claims.", "essential_references_not_discussed": "The paper does not reference works discussing tensor products in bases other than the spherical one (e.g., https://arxiv.org/abs/2306.06482, https://arxiv.org/abs/2405.14253, https://arxiv.org/abs/2412.18263). Aside from this, the paper sufficiently covers related literature.", "other_strengths_and_weaknesses": "Please refer to the comments in previous sections for all strengths and weaknesses of the presented work.", "other_comments_or_suggestions": "1. The authors could clarify why the discussion is limited to the spherical basis and whether their framework extends to other bases.\n\n2. It would be helpful to discuss potential extensions of the framework to “backward pass” computations.\n\n3. Briefly mentioning how alternative bases might impact expressivity and runtime would offer a more comprehensive perspective.\n\nMinor comments:\n\n4. The caption of Figure 5 appears redundant, as it is a part of Figure 4.\n5. \"($\\mathbf{x}^{(l_1)}$\" appears to be missing in Equation 19.", "questions_for_authors": "My questions would relate to the issues or comments raised above, so addressing them would suffice to change my evaluation of the paper.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zh4KgJufIV", "reviewer": "Reviewer_Hd64", "summary": "This paper analyzes various tensor product operations in equivariant neural networks for 3D modeling. It introduces measures of expressivity and interactability and improves the Gaunt tensor product (GTP) with a spherical grid, achieving a 30% speedup. The paper also presents microbenchmarks, showing discrepancies between theoretical and empirical runtime, emphasizing the need for application-specific benchmarking.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your reading of our work and feedback. We appreciate that the reader finds our explanation of differences between TPOs clear.\n\nRegarding the weaknesses:\n## Weaknesses\n1. We would like to clarify that the main point of the paper is understanding how alternatives to CGTP are achieving their performance gains. Crucially, we want to clarify that direct speed comparisons between GTP and CGTP are unfair. Instead, **GTP is performing a different operation from CGTP**. This realization was a major motivation for our work and in fact through our work we see that GTP does not give speed ups compared to CGTP in a fairer comparison.\n\n     The key point is our framework provides a **systematic way** to analyze the differences between any TPOs (such as GTP, MTP, CGTP) and compare speedups in a fairer setting. We believe that this is extremely valuable for designing new TPOs.\n\n     We do not claim that our measure of expressivity has direct correlation with actual training performance. In fact, combining the previous observation that GTP can achieve similar performance as CGTP and our result showing it **provably** has fewer degrees of freedom is already interesting. The **true benefit of GTP should not be understood as a speedup but rather the elimination of many degrees of freedom with minimal impact**. These observations illuminate potential avenues for further improvement.\n\n2. By identifying that we can directly use a S2Grid, we highlight a possible asymptotic runtime improvement over the original implementation of GTP by using a S2FFT. This is the **first** algorithmic improvement which cannot simply be explained by a reduction in degrees of freedom. We believe this is quite interesting and can potentially have a significant impact in the future.\n\n     However, we emphasize that the current $\\ell$s used are far too small to see these asymptotic benefits from the complicated S2FFT algorithm. Despite this, using a S2Grid and a seminaive S2 Fourier transform (instead of S2FFT) gives a much simpler implementation of GTP which already sees performance gains over the original.\n\n3. This is a good point, and we did try to mitigate implementation differences as much as possible (e.g. using JAX for everything). We would like to highlight that it is not the results themselves but the method of benchmarking that is important.\n   \n     Our microbenchmarks provide valuable insight on how different implementations can be improved. In particular, the discrepancy between FLOP counts (which are invariant) and walltime indicates potential for significant acceleration (eg. custom kernels) by better utilizing the GPU.\n\n## Questions\n1. This is a great question! Our experiment in 6.2 gives a simple demonstration highlighting this failure. However, we actually strongly suspect most common molecular datasets are minimally impacted. Loosely speaking, this is because irrep types of commonly predicted quantities such as energy or forces can be constructed purely from symmetric tensor products of the input irrep types. We are actively exploring the impact of antisymmetric tensor products in future work.\n\n     Further, we have come up with a new GTP-like tensor product which does not suffer the antisymmetry issues and has the same potential asymptotic benefit from S2FFT. However, we feel it distracts from the main message of this paper and have omitted it.\n\n2. Thanks for the question! We chose the highest batch size (10,000) that we could fit onto the GPU without going out of memory. To offset any profiling overheads we chose to measure our walltime/FLOPs/throughput metrics using hardware instruction counters at the GPU driver level which is based on Roofline Toolkit Framework for Deep Learning (https://arxiv.org/abs/2009.05257) and has been successfully used in FourCastNet (https://dl.acm.org/doi/10.1145/3592979.3593412).\n\n\n3. We are providing a framework for systematically analyzing how different TPOs achieve efficiency gains. Theoretically, we provide a way to see whether these came from a clever reduction in degrees of freedom or from actual algorithmic improvements. Our microbenchmarks provide a way to see runtimes of TPOs in practice and avenues for implementation improvements.\n\n     Many previous works simply state their method is faster and achieves similar training results. Our work provides a way to analyze why their method was faster and whether there may be possible improvements or limitations. Our experiments are designed with this goal in mind. As such, the comparison of different methods in training results is not the central focus of this work (those results can be found in the original papers introducing the specific TPOs). However we understand the importance of training performance and agree that evaluation of training performance is also crucial for new TPO proposals.", "claims_and_evidence": "The content has already been provided in the subsequent subsections.", "methods_and_evaluation_criteria": "The content has already been provided in the subsequent subsections.", "theoretical_claims": "The content has already been provided in the subsequent subsections.", "experimental_designs_or_analyses": "The content has already been provided in the subsequent subsections.", "relation_to_broader_scientific_literature": "No.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper provides an extensive theoretical introduction and comparison of different types of tensor products, making it easy for readers to understand the differences between them.\n2. The paper introduces a new evaluation method for tensor products, enhancing our understanding of their relative efficiency and expressivity, which can help in selecting the appropriate type of tensor product for networks.\n\nWeaknesses:\n\n1. The problem the paper aims to solve is not entirely clear. Specifically, while a new evaluation criterion for tensor-product computations is introduced, the paper does not explain what additional benefits this provides for tensor-product-based equivariant models. For example, is the model's runtime/accuracy proportional to the tensor-product's runtime/expressivity as proposed? Can tensor product types be evaluated solely based on this criterion without model training? The paper also lacks experiments to verify this point.\n\n2. Considering the theoretical improvement made to the Gaunt Tensor Product, the 30% speedup of S2grid should be seen as marginal. Compared to other acceleration works like Gaunt Tensor Product and cuEquivariance, this improvement seems relatively modest.\n\n3. Regarding the evaluation of different tensor products, the actual performance of these tensor products also depends on the degree of optimization applied, which greatly impacts the scaling constants. It is necessary to discuss or clarify the implementation method in the paper. Otherwise, such comparisons would be unfair and could not be strictly aligned with the theoretical complexity.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "1. Based on Section 6.2, the lack of antisymmetry in GTP seems to significantly affect its performance. Should further experiments be conducted on general datasets to explore this issue? Additionally, how do the evaluation criteria proposed in the paper inspire solutions for addressing this problem?\n\n2. For single tensor products not embedded in networks, performance benchmarking often faces challenges such as very short execution times or inaccurate timing tools. How did you address these issues in your actual experiments?\n\n3. Overall, the paper gives the impression that it touches on many subtopics but provides only a superficial investigation of each, including the evaluation criteria, improvements to GTP, and the study of antisymmetry. Have you considered organizing these topics into separate papers, each with more detailed experiments, such as embedding the tensor products into networks and testing results on popular small-molecule or materials datasets?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper had 3/4 reviewers recommending acceptance. Reviewers appreciated the well-supported claims, quality of presentation, and insights about several tensor product operations used in equivariant NNs.

The paper has one contribution towards improving E(3)-equivariant models, which is a simpler and faster  implementation of the Gaunt tensor product (GTP). This was considered a small contribution by the reviewers.

The main contribution claimed is the method for analyzing and comparing different tensor product operations in the equivariant NN literature. Most reviewers found this contribution valuable and that it provides useful insights about prior work. Reviewer Hd64 questioned whether the methodology is useful in practice and raised several potential issues with the implementation details and metrics. In the back-and-forth with the authors, the reviewer raised their score but not to the point of acceptance.

I side with the majority here. Even if the method to compare and evaluate prior work is not ideal as raised by Hd64, it seems to be the first of its kind and brings useful insights to researchers and practitioners on the field. I believe it is worth being published.
Decision: Accept (poster)