Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products
Authors: Shengjie Luo
Tianlang Chen
Aditi S. Krishnapriyan
Venue: ICLR 2024 spotlight (2024)
Keywords: Equivariant Operations; Tensor Product; Change of basis; Spherical Harmonics; Fourier Basis; equivariant neural networks
Primary area: learning on graphs and other geometries & topologies
TL;DR: We develop a method to change the tensor product formulation of equivariant operations in equivariant neural networks, providing improved accuracy and efficiency.

Abstract: Developing equivariant neural networks for the E(3) group plays an important role in modeling 3D data across real-world applications. Enforcing this equivariance primarily involves the tensor products of irreducible representations (irreps). However, the computational complexity of such operations increases significantly as higher-order tensors are used. In this work, we propose a systematic approach to substantially accelerate the computation of the tensor products of irreps. We mathematically connect the commonly used Clebsch-Gordan coefficients to the Gaunt coefficients, which are integrals of products of three spherical harmonics. Through Gaunt coefficients, the tensor product of irreps becomes equivalent to the multiplication between spherical functions represented by spherical harmonics. This perspective further allows us to change the basis for the equivariant operations from spherical harmonics to a 2D Fourier basis. Consequently, the multiplication between spherical functions represented by a 2D Fourier basis can be efficiently computed via the convolution theorem and Fast Fourier Transforms. This transformation reduces the complexity of full tensor products of irreps from $\mathcal{O}(L^6)$ to $\mathcal{O}(L^3)$, where $L$ is the max degree of irreps. Leveraging this approach, we introduce the Gaunt Tensor Product, which serves as a new method to construct efficient equivariant operations across different model architectures. Our experiments on the Open Catalyst Project and 3BPA datasets demonstrate both the increased efficiency and improved performance of our approach. The code and models will be made publicly available at https://github.com/lsj2408/Gaunt-Tensor-Product.

Published as a conference paper at ICLR 2024
ENABLING EFFICIENT EQUIVARIANT OPERATIONS IN THE FOURIER BASIS VIA GAUNT TENSOR PRODUCTS
Shengjie Luo Peking University
Tianlang Chen
Aditi S. Krishnapriyan" 'Project lead
University of California, Berkeley
ABSTRACT
Developing equivariant neural networks for the E(3) group plays an important role in modeling 3D data across real-world applications. Enforcing this equivariance primarily involves the tensor products of irreducible representations (irreps). However, the computational complexity of such operations increases significantly as higher-order tensors are used. In this work, we propose a systematic approach to substantially accelerate the computation of the tensor products of irreps. We mathematically connect the commonly used Clebsch-Gordee Gaunt coefficients, which integrals of products of three spherical harmonies. Through Gaunt coefficients, the tensor product irreps becomes equivalent to the multiplication between spherical functions represented by spherical harmonies. This perspective further allows us to change the basis for the equivariant operations from spherical harmonics a 2D Fourier basis. Consequently, the multiplication between spherical functions represented by a 2D Fourier basis can be efficiently computed via the convolution theorem and Fast Fourier Transforms. This transformation reduces the complexity of full tensor products of irreps from O(L)) to (LL)) where L is the max degree of irreps. Leveraging this approach, we introduce the Gaunt Tensor Product, which serves as new method to construct efficient equivariant operations across different model architectures. Our experiments the Open Catalyst Project and 3BPA datasets demonstrate both the increased efficiency and improved performance of our approach.
coefficients to the
INTRODUCTION
The imposition of physical priors, corresponding to specific constraints, has played an important role in deep learning (Jumper et al., 2021). To address the curse of dimensionality, prior works exploited the geometry inherent in data, such pixel grids in images atomic positions in molecules, along with their associated symmetries. Among the constraints introduced by these physical priors, equivariance has emerged as important concept in deep learning (Bronstein et al., 2021).
function mapping with the equivariance constraint transforms the output predictably in response to transformations the input. One typical class of such transformations includes translations, rotations, and reflections in the Fuclidean space, collectively termed the Euclidean group (Rotman, 1995). Numerous fundamental real-world problems exhibit symmetries the Fuclidean group, including molecular modeling (Atz et al., 2021), protein biology (Jumper et al., 2021), and 3D vision (Weiler al., 2018). Many works have developed a rich family of equivariant operations for the Euclidean group (Han et al., 2022), demonstrating promise in challenging prediction tasks through both improved data efficiency and generalization.
Among the various approaches for enforcing equivariant operations (Cohen & Welling, 2016; 2017; Coors et al., 2018; Finzi et al., 2020; Satorras et al., 2021, see Appendix for more advances), em- ploying tensor products of irreducible representations (irreps) has recently emerged the dominant choice for the Euclidean group (Thomas et al., 2018; Fuchs et al., 2020; Brandstetter et al., 2022; Batzner et al., 2022; Liao & Smidt, 2023; Musaclian et al., 2023). Rooted in the mathematical founda- tion of representation theory, such operations allow for the precise modeling of equivariant mappings between any 27 + 1-dimensional irreps space via Clebsch-Gordee Gordan coefficients (Wigner, 2012), and are proven to have universal approximation capability (Dym & Maron, 2021). These operations have driven state-of-the-arr performance across variety of real-world applications (Ramakrishnan et al., 2014; Chmiela et al., 2017; Chanussot et al., 2021), demonstrating strong empirical results.
"Equal contribution. Correspondence to Tianlang Chen <t1 < chen@pku edu. en>, Shengjie Luo uos jestu. pku. edu. en> and Aditi S. Krishnapriyan <iditiki k @berkeley eduu

---

Published as a conference paper at ICLR 2024
The computational complexity inherent to computing the tensor products of irreps is well- acknowledged hurdle in achieving better efficiency, as well as greater accuracy (see comments in Fuchs et al., 2020; Schütt et al., 2021: Satorras et al., 2021; Frank et al., 2022: Batatia et al., 2022b; Brandstetter et al., 2022; Musaelian et al., 2023; Passaro & Zitnick, 2023; Liao et al., 2024; Simeon & Fabritiis, 2023; Wang & Chodera, 2023; Han et al., 2022; Zhang et al., 2023b). Recent advancements suggest that the empirical performance of state-of-the-an models based on such operations can be consistently boosted by using higher degrees of irreps (Batzner et al., 2022; Batatia al., 2022b; Brandstetter et al., 2022; Passaro & Zitmick, 2023; Liao et al., 2024; Musaelian et al., 2023). However, the full tensor product of irreps up to degree L have an (L)) complexity, which significantly limits L to 2 or 3 in practice, hindering efficient inference on large systems. Given this, there is strong need to improve the efficiency the tensor products of irreps.
In this work, we introduce systematic methodology to accelerate the computation of the tensor products of irreps. Our key innovation established on the observation that the Clebsch- Gordan coefficients are mathematically related to the integral of products of three spherical harmonics, known as the Gaunt coefficients (Wigner, 2012). We use this to show that the tensor product of irreps via the Gaunt coefficients is equivalent to multiplication between spherical functions represented by linear combinations of spherical harmonic bases. This understanding motivates us to employ the 2D Fourier basis to equivalently express the same functions via change of basis. Multiplications between functions represented by 2D Fourier bases can be significantly accelerated through the convolution theorem and Fast Fourier Transforms (FFT). Using these insights, the computational cost of full tensor products of irreps up to degree L can be substantially reduced to ((L)).
Building upon this approach, we introduce the Gaunt Tensor Product, which serves new method for efficient equivariant operations across various model designs. We provide a comprehensive study on major operation classes that are widely used in equivariant models for the Euclidean group, demonstrating the generality of our method and how it can be used to design efficient operations:
Equivariant Feature Interactions. Operations this category calculate tensor products be- tween two equivariant features, enabling the interaction between irreps from different dimen- sional spaces and objects. These operations have been widely used in recent advances (Unke et al., 2021; Frank et al., 2022; Yu al., 2023b), c.g., interacting atomic features and pre- dicting a Hamiltonian matrix. Employing our approach in this scenario is straightforward.
Equivariant Convolutions. Tensor products in this category are computed between equiv- ariant features and spherical harmonic filters, which are the building blocks for equivariant message passing (Thomas et al., 2018; Fuchs et al., 2020; Batzner et al., 2022: Brandstetter et al., 2022). In this case, we demonstrate that additional speed-up can be achieved via the insight from Passaro & Zitnick (2023), which introduces sparsity to the filters via a rotation. We use this to further accelerate the conversion from spherical harmonics to 2D Fourier bases.
Equivariant Many-Body Interactions. Modeling many-body interaction effects has been shown be important for advancing machine learning force fields (Drautz, 2019; Dusson et al., 2022; Batatia et al., 2022a;b; Nigam et al., 2022). Capturing these interactions requires performing multiple tensor products among equivariant features. Our connection between Flebsch-Gordar coefficients and Gaunt coefficients is naturally generalized to this setting, and we develop a divide-andd-ooquer approach to parallelize the computation of multiple 2D convolutions, demonstrating better efficiency compared to original implementations.
We conduct extensive experiments to verify the efficiency and generality of our approach. For each of the aforementioned operation classes, provide efficiency comparisons between our Gaunt Tensor Product and implementations via the e3nn library (Geiger & Smidt, 2022). We also apply our Gaunt Tensor Product to state-of-the-art models on benchmarks across task types and scales, e.g., force field modeling on OC20 (Chanussot et al., 2021) and 3BPA (Kovács et al., 2021) benchmarks. All results show both the better efficiency and improved performance of our approach.
BACKGROUND
We review the necessary background knowledge for our method, including concepts of the Euclidean group and equivariance, irreducible representations (irreps), spherical harmonics, and the tensor product of irreps. We include all the details of the mathematies used in this work in Appendix A.
2

---

Published as a conference paper at ICLR 2024
Let S denote a set of objects located in the three-dimensions Euclidean space. We use X E R to denote the objects with features, where n is the number of objects, and d is the feature dimen- sion. Given object i, let r, E RS denote its Cartesian coordinate. We define S = (X, R), where R = (rı,... }. These data structures are inherent in a variety of real-world systems, such as molec- ular systems and point clouds. Problems defined on S generally exhibit symmetries to transformations on the objects. For instance, given a molecular system, if a rotation is applied to each atom position r,, the direction of the force on each atom will also rotate. Mathematically, such transformations and sym- metries are directly related to the concepts of groups and equivariance that we will introduce below.
Euclidean group and Equivariance. Formally, a group G is non-empty set, together with a binary operation G x G + G satisfying group axioms that describe how the group elements interact with each other. The class of transformations including translations, rotations, and reflections in the three-dimensional Euclidean space is of great interest, as it describes many real-world scenarios. These transformations form a group structure known as the Euclidean group, and denoted as E(3).
To describe how the group elements act on vector spaces, group representations are used. Formally, the representation a group G on a vector space V is a mapping p : G + GL(V), where GL(V) is the set of invertible matrices on V. For example, representations of the rotation group SO(3) on R2 can be parameterized by orthogonal transformation matrices R E R3x3 det(R) = 1, which are applied to x E RS via matrix multiplication Rx.
We now introduce the concept of Equivariance. Let o + y denote a function mapping between vector spaces. Given group G. let xK and pd denote its representations on X' and y respectively. A function Φ X - y is said to be equivariant if it satisfies the following condition x² (g) (g)[ø(z)] = ((p*(g))
for all g E G, x E X. When p = It (identity transformation) for G, it is a special case of equivariance, known as Invariance.
The equivariance constraint is a commonly exploited physical prior in machine learning models. For tasks defined on the data structure S. it is typically important to ensure operations maintain -(3)-qquivariaace. The E(3) group a semi-direct product of the group of translations 123 with the group of orthogonal transformations ((3), i.c., E'3) = R (3).. In general, achieving translation equivariance straightforward through convolution- pype operations. Therefore, we primarily focus on the orthogonal group (3)) which consists of rotations and reflections.
Irreducible representations. We use D(g) to denote the representation of group element g in the matrix form. For vector spaces of dimension 27 + 1, E N, there exists a collection of matrix representations indexed with for each g E SO(3), which are called Wigner-D matrices. They are denoted as D (*) (g). Each Wigner-D matrix (3) (g) is the irreducible representation of SO(3) on a subspace with dimension 2/ 1. Any matrix representation D(g) of SO(3) on vector space V can be reduced to an equivalent block diagonal matrix representation with Wigner-D matrices along the
(g)
diagonal: D(g) (D (1) (g) @ (g)
+ ....
Q-1
(()a
Q. Here,
Q is the change of basis that makes D(g) and the block diagonal matrix representation equivalent.
This equivalence provides an alternative view that the representations that act on the vector space V are made up of each Wigner-D matrix D(4,,, which only acts on a 24, + 1-dimension subspace Vi, of Then, the vector space V can be factorized into Vi, @ Vi, @ We call the space V/2 type-l irrep space, and (1) E Vi2 a type-l irrep. We have already introduced the Wigner-D matrix in the type-1 space, i.e., (1)) = R. E 3333 (c)((R)= = 1. The Wigner-D matrix representations can be naturally extended to the O()) group by simply including the reflections via direct product.
Spherical harmonics. Spherical harmonics are used to map from R33 to any type-l irrep space. We use YI(0 S2 + R, -1 - m S 1,7 E N to denote the spherical harmonics, where ( and m denote their de- grees and orders respectively. For any x E R.3, (0
x )]" is type-l [ | .... y()) |
|
irrep. This mapping is equivariant to the 0(3) group. i.e., given F E S2 .g 50(3) and its representa- tion on R3 as rotation matrix R, we have (RF) = D (g) (f) and (() = (-))
Tensor products of irreducible representations. Given ((),, ((aa) in type-11 and type-12 irrep space respectively, their tensor product yields (())) @ (l)), which lies in the Vi, @ Vz space with
3

---

Published as a conference paper at ICLR 2024
(24 1) x (2lz + 1) dimension. This operation preserves equivariance to the 0(3) group, ie., for all g E O(3). As previously stated, the tensor product of irreducible representations D (g) ② (g) can be further decomposed into irreducible pieces: (()(gg)
(D @ (g)) @ x((a)>= (g(h)) @
4 1/2 -11 D (g). The resulting irrep of the tensor product can be organized via change of basis into parts that are in different type-/ irrep spaces, i.c., x ((1) @ x((a) E V = Vit-12l @ @ Via+
(g) @ D (g)
x(1:) @ x
The Clebsch- Gordan tensor product provides an explicit way to calculate the resulting type-/ irrep after the direct sum decomposition. Given x(11) [R and x(aa E IR2 the m-th element of the resulting type-l irrep of the tensor product between x(11) and x((2) is given by:
(x Gog
x
m
Σ
(llam) (4amm))aamma)) 3mi "mn
(())
(1)
11
where (4am))0smma)) m()))) are the Clebsch-Gordae coefficients (Wigner, 2012). We use x(4)
(x(0); x(0:: E TR((++)² to denote features containing irreps of up to degree I. in practice, and the full tensor product given by (x(L.) Big x(L2)) ! = Σhi - Beg (. Retaining
all output irreps of the Clebsch- tensor product requires ((*) 3D matrix multiplications, i.e.,
an ((L)) cost, which becomes computationally unfeasible when scaling irreps to higher L degrees.
3 GAUNT TENSOR PRODUCT
We now introduce our Gaunt Tensor Product, which serves as new method for efficient F(3)-- equivariant operations. First, we propose new perspective on the tensor products of irreps (Sec- tion 3.1), which connects this operation to multiplication between functions defined on the sphere via Gaunt coefficients (Wigner, 2012). Our perspective inspires the development of an efficient approach to accelerate the tensor product computation by using the convolution theorem and Fast Fourier Transforms (Section 3.2), which achieves substantial speed-ups compared prior implementations. Finally, we provide a comprehensive study on major equivariant operation classes in the literature, and show how design efficient counterparts via our Gaunt Tensor Product (Section 3.3), which opens new opportunities for pushing the frontier of -qqivvvaiva models.
3.1 A NEW PERSPECTIVE ON TENSOR PRODUCTS OF IRREPS VIA GAUNT COEFFICIENTS
As shown in Eqn. (I), the essence of the tensor product operation is in the Clebsch-Gordan co- efficients, which are used to perform the explicit direct-sum decomposition. The CG coefficients commonly appear in quantum mechanies, and serve as the expansion coefficients of total angular mo- mentum eigenstates in an uncoupled tensor product basis. Interestingly, spherical harmonics actually correspond to the eigenstates of angular momentum, indicating that there exist underlying relations be- tween these two concepts. In physics, the Wigner-Eckart theorem in quantum mechanics reveals these relationships an elegant manner (refer to Appendix A.7 for the details of these physical concepts):
Theorem 3.1 (Wigner- Eckart theorem for Spherical Tensor Operators, from Jeevanjee (2011)). Let J. E N20 and let T') be spherical tensor operator of rank j. Then there is unique complex number, the reduced matrix element À E C (often written (J(TT() [[7) E C) that completely determines any of the (2.J + 1) x (2j + 1) x (2l + 1) matrix elements (J. M T(? I),)):
Wiiger-Ecckart
M.MITO T() 17, m) X
Here, n)) (J, M denote eigenstates for angular momentum operators with |n|st, |M|S. and re- late to spherical harmonies via (0, v) (0, v), with spherical coordinate 0E(0, m] $E (0, 2T). The Wigner-Eckan theorem describes how the matrix elements (JJM M| |4, n), |m| 5 of the spherical tensor operator (T(() in the basis of angular momentum eigenstates are decomposed into the Clebsch-Gordan coefficients and the reduced representation ((][TO))))) If we instantiate TÚ) as the spherical harmonics y() and apply the Wigner- -ckkrt theorem, we obtain the following relation (J, M y-U) m) = (J C(J.A)) which can be rewritten as, (l.))).m))'
|0
Yills) (4,)(0000) + mg (0, Y (0, +) sin 0d0dg' 8(0) (..mm)) d2) (lsmmm)))
(3)
4

---

Published as a conference paper at ICLR 2024
where 6(0)
(2) is real constant only determined by 11 12 and From Eqn. (3), the Clebsch-Gordan coefficients are mathematically related to the integrals of products of three spherical harmonies,
which are known as the Gaunt coefficients (Wigner, 2012). We denote them as (1amm))(22mma))
((l.)
If we use the Gaunt coefficients in the tensor product operation instead, we have the following result on the tensor product of irreps:
(x(11) Gainns x
Jm
Σ
(1) GGautt x12)) )m
11=0 1==
b
12
(1 (1gmm)))g.
Tmii Ting
(11) (12)
11 =0 12
-1 mg m
(4)
L2
2
(2)
I Yuu (0, ()) (B v) sin Od0des
Mm Jin
1
1a
mi
1 mig 12
b
y (1) (0, wt)
(12 Lmz
(0, 2/) r 1 (0, +) sin 0d0dy. x
Tm m0
The Y'''s form an orthonormal basis set for functions defined on the unit sphere S² i... any square- integrable function +)),v) 5 + can be expanded as Ei- Em--, p y!!! (0, v). and the coefficients fll can be calculated by J²t Jo f(0. b)Y we Imi Iigg
! (6, v)) sin Adodw.. If
(11) (12)
to for 11>1,22222 we actually obtain
spherical functions, (12) Ima (iaa))((a))))
(0 =i= En = 2m01
(l))(())(B (0, and b) Ei
(0, tb).
Thus, (x 8 Gaint x x( m calculates the coefficients of function (s = (0, v) F2 v) This perspective connects tensor products of irreps multiplication between spherical functions, which creates new ways to accelerate the tensor product operation, as we will discuss further.
3.2 FAST TENSOR PRODUCT USING CONVOLUTION THEOREM & FAST FOURIER TRANSFORM
For multiplication between functions defined on the unit sphere S2, a natural question is the existence of efficient computational methods. We can leverage the change of basis to achieve this goal: instead of using spherical harmonics, the same functions defined on the unit sphere S can be equivalently represented using 2D Fourier basis, which is known to have favorable numerical properties.
Given F(B, Φ) : S2 + R, can be represented by linear combination of 2D Fourier bases as E.E. E., (())+)))))) u, v E Z with complex coefficients fū,, to make a real function. To keep the convention of spherical harmonics, we also describe the 2D Fourier basis in terms of (up to) degrees I I u S L, -I. S L. The multiplication between Fi (0, v) and F2 v) can thus be expressed as:
(2) = ΣΣ plr Meentton)) )
Fi -(0,,))FF0000))
ΣΣ Par (ez2)tsw)))))
(5)
ΣΣΣΣ 41 42
Son.e Fen Almy (lons
By comparing the 2D Fourier bases, the coefficients of (0( (0, (2 4) can be obtained by fäx = Eu, +wa=us E +01 f/ f2 which is exactly the form of 2D convolution. It well-known that these 2D convolutions can be significantly accelcrated by using the convolution theorem and Fast Fourier Transforms (FFT) (Proakis, 2007). We can then use this to develop an efficient approach to compute the tensor product of irreps in Eqn. (4).
From spherical harmonics to 2D Fourier bases. Since both spherical harmonics and 2D Fourier bases form an orthonormal basis set for spherical functions, we have the following equivalence: for spherical harmonics Y(!!) (0, v), 0 < l S L, -l < < L, there exist coefficients s.m such that y(1) (0, v) = Et= Ele-p y.m. e((ud+vv)) These coefficients y.m. are sparse due to the orthogonality of such basis sets. That is, I.m.* are non-zero only when to.
5

---

Published as a conference paper at ICLR 2024
To compute x((a) @ Goint (T(=2) in Eqn. (4), the first step is to convert the coefficients of spherical harmonics to coefficients of the 2D Fourier basis. From the above equivalence, we have the following conversion rule from spherical harmonics to 2D Fourier bases:
Σ E
! Lim. LLSLL LS I'm
(6)
I V0O
Performing 2D convolution via Fast Fourier Transform. After converting coefficients from spherical harmonics to 2D Fourier bases, we can use the convolution theorem and Fast Fourier Transform (FFT) to efficiently compute the 2D convolution from Eqn. (5). Both ===² and v are transformed from the spatial domain to the frequency domain via FFT. In the frequency domain, the 2D convolution can be simply computed by element wise multiplication. Then, the results of multiplication in the frequency domain are transformed back to the spatial domain via FFT again.
From 2D Fourier bases to spherical harmonics. Similarly, for 2D Fourier bases c((u(+v))) where- L S u S and -L S S L, there exist coefficients s.m., such that ed(u(+oo)
レ Et - Z..0" 2.m.+y/l (0, vb). and z.m,.* are also sparse, i.e., =.m.. are non-zero only when Y = tv. The last step is to convert the resulting coefficients of multiplication between spherical functions back to coefficients of spherical harmonics, which can be calculated as:
Σ
Σ L
"i."" ,0SIS L. -s SL
(7)
Computational complexity analysis. We provide analysis on the computational complexity of the original implementation of full tensor products of irreps of degree up to L. As can be seen from Eqn. (4), to calculate the m-th element of the output type-/ irrep, there exist ( (2) + combinations of input irreps. Each combination requires a summation of OLL²) multiplications. This means it requires ((L*) operations for the m-th element of the output type-l irrep. Then, retaining all m-elements of all output type-l irreps with 0 S / < L has an (LL)) complexity.
In our approach, we first convert coefficients of spherical harmonics to coefficients of 2D Fourier bases (see Eqn. (6)). For each zu,us since v'oe.'* are non-zero only when m = Σv, it requires O(L) operations. Thus, the first step of our approach requires ( operations in total. For the second step, the 2D convolution in Eqn. (5) has an (T, log I.) complexity using FFT. Similarly, the last step of our approach also only requires ((L)), due to the sparsity of 2/,". Therefore, our approach has an (LL) complexity compute the full tensor products of irreps of degrees up to L, which achieves substantial acceleration compared to the original implementation.
3.3 DESIGNING EFFICIENT EQUIVARIANT OPERATIONS VIA THE GAUNT TENSOR PRODUCT
Based on the approach in Section 3.2, we introduce the Gaunt Tensor Product, a technique that efficiently computes tensor products of irreps utilizing Gaunt coefficients. Our Gaunt Tensor Product can be applied across different model architectures in the literature, and serves as new method for efficient equivariant operations. We provide a comprehensive study on major equivariant operation classes (Equivariant Feature Interactions, Equivariant Convolutions, and Equivariant Many-body Interactions) and show how to design efficient counterparts via our Gaunt Tensor Product. We present additional details in Appendix C.
Equivariant Feature Interactions. Given two features (1) (((aa) containing irreps of up to degree L and L2 respectively, the operations in this class are commonly in the following form: 1.1
(x(41) 2x (La))(()) Σi== Elg w/, (x(11) Sep (l))))(1), where w/, IR are the learnable weights for each (4, 12) + l combination. Compared to Eqn. (4), our Gaunt Tensor Product is exactly in the same form, except for the learnable weights wf Our approach can be naturally extended by reparameterizing w//
Wi,- wi, wy, where wi, W1,, wi E R. The weighted combinations can be equivalently achieved in Eqn. (4) by separately multiplying Wi,, W11 wy with x((1), (1aa and (x(1 @Gaunt x(12)(() (1) correspondingly. in which case the tensor product of irreps can be still efficiently computed by using the approach in Section 3.2.
These operations facilitate interactions between irreps from different dimensional spaces and objects, which have been widely used in recent work. For example, So3krates (Frank al., 2022) developed

---

Published as a conference paper at ICLR 2024
the atomic interaction layer, where both x(L.) and ý(L2) are the equivariant features x,(D) of atom i. Additionally, in the Quantum Hamiltonian prediction task, such operations are commonly used to interact features for constructing Hamiltonian matrices (Unke et al., 2021: Yu et al., 2023b). For instance, QHNet (Yu et al., 2023b) instantiates x(()) and ý(lz) equivariant features of atom pairs (&,,) Xz) and uses the output features as the input of the non-diagonal matrix elements prediction block.
Equivariant Convolutions. One special case of Equivariant Feature Interaction is to instantiate y((22) as spherical harmonics filters, i.c. Σi, -1z=0 hi (xx 8og y(l)( [ (F)-)))))) where TG E I are positions of objects, ((2a), S2 + 222 is the type-12 spherical harmonic, and M E R are learnable weights calculated based on the relative distance and types of objects i,j- Similarly, the h/g.12 can be reparameterized for our approach. Passaro & Zitmick (2023) provide an interesting observation that if we select a rotation g E SO(3) with ((
P. TE.-E = (0, 1, 0), then dm i.e., is non-zero only if = 0. This sparsification further propagates to the coefficients of the 2D Fourier basis, which brings additional acceleration in Eqn. (6).
Ym (D (g) Fs r -.
Equivariant Convolutions are the basic operations for equivariant message passing (Thomas et al., 2018; Fuchs et al., 2020; Brandstetter et al., 2022; Batzner et al., 2022: Liao & Smidt, 2023; Passaro & Zitnick, 2023). The spherical harmonics filters encode the geometric relations between objects and interact with the object features an equivariant manner, which plays an essential role in learning semantics and geometric structures under equivariance constraints.
Equivariant Many-body Interactions. Operations in this class perform multiple tensor products among equivariant features, i.e., x((1) Bcg X2 ((22)) Beg Bcg Xn (Lm)) Such operations are commonly used for constructing many-body interaction effects (Drautz, 2019; Batatia al., 2022a;b; Nigam et al., 2022; Kovács et al., 2023), which can improve performance for ML force fields. For example,
given equivariant features A44²) of atom i, MACE (Batatia et al., 2022b) constructs the many-body features by En, Eim C.L.M. mm I A which is the explicit form of performing tensor products with A!4)) itself v times, i.e., B = 1(55 Big ... Big A(4)) Similarly, such operations can be equivalently viewed as multiplications of multiple spherical functions. Using this, we can use divide-and-conquer method to efficiently compute multiple 2D convolutions for further speedups.
(1e)
Discussion. It is common in models to maintain C channels of equivariant features, i.e., x(2), ,0 S c C contains irreps of degrees up to L. This extension can be naturally included in our approach by defining the combination rules of features in different channels, e.g., channel-wise (O(C) ops) or mixxng (O(C2) ops). Note that our approach is not limited to full tensor products only. In fact, given x(1)) Beg y (12) only, the approach presented in Section 3.2 can still be used, which reduces (TT,3) (i.² log L.) with T. = l 122 Overall, we demonstrate the generality of our Gaunt Tensor Product, and we believe that it is a starting point with more possibilities to explore in the future.
4 EXPERIMENTS
We empirically study the efficiency and effectiveness of our Gaunt Tensor Product. First, provide comprehensive efficiency comparisons between our Gaunt Tensor Product and the original implementations of the operation classes mentioned in Section 3.3. Second, we perform a sanity check to investigate the effects brought by the different parameterizations of our Gaunt Tensor Product versus the (ebschh-Gordan Tensor Product. Finally, we present results applying our Gaunt Tensor Product to various model architectures and task types. Experiment details are presented in Appendix E. The code will be released at https: / github. com 1s 12 00////48/// Gaantt Tensor-Productt
Efficiency Comparisons. (1) For the Equivariant Feature Interaction operation, we use the e3nn implementation (Geiger Smidt, 2022) as the baseline. We randomly sample 10 pairs of features con- taining irreps of up to degree L with 128 channels. The average inference time is reported. (2) For the Equivariant Convolution operation, we use the eSCN implementation (Passaro & Zitnick, 2023) as the baseline, due to its enhanced efficiency over e3nn for this operation. In this setting, pairs of features and spherical harmonics filters are sampled instead. We implement our Gaunt Tensor Product by using the insights of eSCN for further acceleration (Section 3.3). (3) For the Equivariant Many-body Interac- tion operation, we compare our approach with both the e3nn and the MACE implementations (Batatia et al., 2022b). We evaluate various degrees of I. and the number of tensor product operands v.
7

---

Published as a conference paper at ICLR 2024
Equivariant
Equivariant Convolution
Equivariant Many body Interaction (I)
2000
eSnn
80
#SCN
65nn
Timn (100)
1500
OLLSS
CurS
Time (mss
2000
1000
1030
500
0-
10
12
14
Degrae
Degree L
Degree
Equivariant Many- Interaction (I))
Equivariant Ma body Interaction (III)
Fix v=3
Fix L=2
Sanity Check IN body Simulation)
MACE
MICCE
Tensor Product
7.5
Giaunt Tensor Product fours)
Tirne (mos)
Time (mas)
20
Test MSE
5.0
10
SEGNN La-1
SEGNN La-2
Degrac
Number Tonsor Opcrands
Figure 1: Results on efficiency comparisons and sanity check. We comprehensively compare our Gaunt Tensor Product with implementations of e3nn, eSCN, and MACE on corresponding equivariant operation classes. In all settings, our approach achieves significant speedups. Our Gaunt Tensor Product parameterization further passes the sanity check with SEGNN on the N-body simulation task.
OC20 2M Validation
S2EF Force MAE [meV/A]
Energy MAE
Force Cos
EFwT
Model
[meV] L
[%]] 1
SchNet (Schiitt et al., 2018)
1400
78.3
0.109
0.00
DimeNet++ (Gasteiger et al., 2020a)
805
65.7
0.217
0.01
SpinConv (Shuaibi el al., 2021)
0.479
0.13
GemNet-dT (Gasteiger al., 2021)
358
29.5
0.557
0.61
Ge (Gasteiger et al., 2022)
286
25.7
0.598
1.06
L
# layers
# batch
SCN 1-tap I-band (Zitnick al., 2022)
6
12
64
299
24.3
0.605
0.98
SCN 4-tup 2-band (Zitnick et al., 2022)
6
12
64
279
22.2
0.643
1.41
eSCN (Passaro & Zitnick, 2023)
4
12
96
291
22.2
0.637
1.39
cSCN (Passaro & Zitnick, 2023)
6
12
96
294
21.3
0.653
1.45
EquiformerV2 (Liao et al., 2024)
4
12
64
284
21.4
0.657
1.51
EquiformerV2 (Liao et al., 2024)
6
12
64
285
20.5
0.663
1.67
EquiformerV2 Gaunt-Selfmix (ours)
4
12
64
277
20.9
0.661
1.82
EquiformerV2 Gaunt-Selfmix (ours)
64
20.1
0.669
1.95
Table 1: Results on the OC20 S2EF task. The results are averaged across the four OC20 Validation set splits. Bolded values denote the best performance.
The results presented in Fig.1. Our Gaunt Tensor Product achieves substantial acceleration pared to the e3nn implementation, e.g., multiple orders of magnitude speed-up when L>7. Compared the efficient eSCN implementation. our approach has further acceleration, due to the combination of these two methods. Finally, our Gaunt Tensor Product consistently outperforms the MACE implemen- tation across different and v. These results demonstrate the improved efficiency of our approach.
Sanity Check. As shown in Eqn. (3), the Gaunt coefficients are proportional to the Glebsch-Gordan coefficients for the constants, which leads different parameterizations of our Gaunt Tensor Product and the (lebsch-Gordanl Tensor Product. To investigate the underlying effects, we conduct sanity check by applying the parameterization of the Gaunt Tensor Product the SEGNN (Brandstetter et al., 2022) model the N-body simulation task (Satorras et al., 2021). which requires the model to forecast the positions set of particles modeled by predefined interaction rules. Following Brandstetter al. (2022), we employ 4-layer SEGNN and keep all hyperparameters the same to compare the performance of different parameterizations.
In the last panel of Figure 1, tively compared to the Clebsch- Gorddan Tensor Product implementation, which demonstrates that our parameterization does not hurt the performance of equivariant operations.
that SEGNN with our Gaunt Tensor Product performs competi-
8

---

Published as a conference paper at ICLR 2024
Table 2: Results on the 3BPA dataset. Energy (E, meV) and force (F, meV/A) errors are evaluated.
Standard deviations over three runs are reported if available. Bold values denote the best efficiency.
ACE
sGDML
Allegro
NequlP
BOTNet
MACE
MACE-Gaunt
7.1
3.84 (0.08)
(0.1)
3.1 (0.13)
(.2)
(.91
300 K
F
27.1
46.2
12.98 (0.17)
10.8 (0.2)
(0.14)
8.8 (0.3)
9.2
E
24.0
484.8
12.07 (0.45)
(0.6) (0.1)
(0.5)
(0.6
600 K
64.3
439.2
29.17 (0.22)
26.4 (0.1)
26.7 (0.29)
21.8 (0.6)
(0.2)
E
85,3
774.5
42.57 (1.46)
38.5 (1.6)
29.0 (1.1)
30.4
1200 K
187.0
711.1
82.96 (1.77)
76.2
(1.5)
(0.7)
(1.2)
Dihedral Slices
16.3 (1.5)
7.8 (0.6)
(0.3)
F
(1.0
16.5 (1.7)
(1.7
Speed ups (v3. e33n)
33.2x
43.7x
Memory costs (V.S. c3nn)
32.8%
5,8%
OC20 S2EF performance. Following Gasteiger et al. (2021); Zitnick et al. (2022); Passaro & Zitnick (2023), we evaluate our approach on the large-scale Open Catalyst 2020 (OC20) dataset, which consists of 1.2M DFT relaxations. Each data instance in OC20 has adsorbate molecule placed on a catalyst surface. The core task is Structure-to- Lotooo-nnn--rrreeeeeeeee (S2EF), which requires the model to predict the energy and per-atom forces of the ddsorbate-catalyye complex. Recently, the EquiformerV2 model (Liao et al., 2024) achieved sfte-of-thhe-rr performance on the OC20 S2EF task, which uses the efficient eSCN implementation to build scalable Transformer backbone. However, this implementation restricts the available tensor product operations to Equivariant Convolutions only.
We apply our Gaunt Tensor Product construct an efficient Equivariant Feature Interaction operation, Selfmix, and add it to each layer of the EquiformerV V2 model. Our approach retains efficiency even with such modifications. Following Liao al. (2024), we train 12-layer EquiformerV2 models and keep all configurations the same. We use the OC20 2M subset for training. The results are presented in Table 1. We show that EquiformerV2 with our Gaunt Selfmix operation achieves better performance on the S2EF task, i.e., 16.8% relative improvement on the EFwT metric with L=6.
3BPA performance. Following Batzner et al. (2022); Musaclian et al. (2023); Batatia et al. (2022b), we use the 3BPA dataset (Kovács et al., 2021) to benchmark the Equivariant Many-body Interaction operations. The training set of 3BPA contains 500 geometries sampled from the 300 molecular dynamics simulation of the large and flexible drug-like like molecule 3- ((benyylxyypyrriinn222 (3BPA). Both in- (300 K) and out-of-distribution test sets (600 K, 1200 K, dihedral slices) are used to evaluate the model's ability to predict the energy and forces of the geometrical structure. We apply our Gaunt Tensor Product to the MACE architecture, which uses equivariant features containing irreps of up degree L = 2. The Equivariant Many-body Interaction operation in MACE performs two tensor products among three equivariant features. We keep all training configurations the same as in Batatia al. (2022b). The details of baselines and settings are presented in Appendix E.4.
The results are presented in Table 2. MACE, integrated with our Gaunt Tensor Product, performs competitively compared its original implementation, which has the best performance among com- petitive baselines. This affirms that our parameterization in the Equivariant Many-body Interaction does not hurt accuracy. MACE with our Gaunt Tensor Product achieves significant acceleration compared to both the e3nn implementation and the original MACE implementation, i.c., 43.7x speed-ups compared the baselines. Note that the MACE implementation trades space for speed and incurs additional memory costs, while our approach has low memory costs that are friendly for deployment, e-g., reducing 82.3% relative memory costs versus MACE.
CONCLUSION
We introduce a systematic approach, the Gaunt Tensor Product, to significantly accelerate the compu- tation of the tensor products of irreps. Our Gaunt Tensor Product is established through the connection between Stebsch-Gordaal coefficients and Gaunt coefficients, which enables the equivalent expression of the tensor product operation as the multiplication between spherical functions represented by spherical harmonics. This viewpoint allows us to change the basis for the equivariant operations to 2D Fourier basis. Multiplication between functions in this basis can thus be accelerated via the convolution theorem and Fast Fourier Transforms. Our experiments demonstrate both the improved efficiency and effectiveness of our approach. Overall, our Gaunt Tensor Product provides new opportunities for efficient equivariant operations, opening avenues for exploration in the future.

---

Published as a conference paper at ICLR 2024
Acknowledgements. This work acknowledges support from Laboratory Directed Research and Development (LDRD) funding under Contract Number DE-ACO2-05CHI123) We thank Nithin Chalapathi, Di He, Eric Qu, Danny Reidenbach, and Bohang Zhang for helpful discussions and feedback. We also thank all reviewers for their valuable suggestions.
REFERENCES
Kenneth Atz, Francesca Grisoni, and Gishert Schneider. Geometric deep learning on molecular representations. Nature Machine Intelligence, 3(12):1023-1033 2021.
Albert P Bartók, Mike C Payne, Risi Kondor, and Gábor Csányi. Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons. Physical review letters, 104(13): 136403, 2010.
Ilyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor NC Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi The design space of e (3)-equivariant atom-centered interatomic potentials. arXiv preprint arXiv:2205. 06643, 2022a.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Informaionn Processing Systems, 35:1 1121-112241411 11436, 2022b.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korblluhh Nicola Molinari, Tess Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-energy surfaces. Physical review letters, 98 (8(14): 164001, 2007.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling. Ge- ometric and physical quantities improve e(3) equivariant message passing, In International /
Conference on Learning Representations, 2022. URL https::t//// http///ppeeeee...eeeeeeee net/ fformm?
id= [d=_xx:8gOBeV1.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2 13478, 2021.
Gerald Burns.
Introduction to group theory with applications: materials science and technology. Academic Press, 2014.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier feras-Doming Caleb Ho. Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi, Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 11 11(10):0099-677, 2021.
Tianlang Chen, Shengjie Luo, Di He, Shuxin Zheng, Tie- Yan Liu, and Liwei Wang. GeoMFormer: A general architecture for geometric molecular representation learning. In NeurIPS 2023 Al for Science Workshop, 2023. URL https: / opwnrevieeeeeeeeeeeeeeeeee.oe..e net/ forum?i.:=:0UNNuuuqtt
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, 3(5):e1603015, 2017.
Stefan Chmiela, Huziel E Sauceda, Igor Poltavsky, Klaus- Müller, and Alexandre Tkatchenko. sgdml: Constructing accurate and data efficient molecular force fields using machine learning. Computer Physics Communications, 240:38 2019.
Taco Cohen and Max Welling, Group equivariant convolutional networks. In International conference on machine learning, pp. 2990-2999. PMLR, 2016.
Taco S. Cohen and Max Welling. Steerable CNNs. International Conference on Learning Representations, 2017. URL htt.httst:::/oerer//// net/ ftee:/ooormmm..rrgrr.....
10

---

Published as a conference paper at ICLR 2024
Eugene D Commins. Modern quantum mechanics, revised edition, 1995.
Benjamin Coors, Alexandru Paul Condurache, and Andreas Geiger. Spherenet: Learning spherical representations for detection and classification in omnidirectional images. In Proceedings of the European conf. on computer vision (ECCV). 2018.
John F Cornwell. Group theory in physics: An introduction. Academic press, 1997.
Albert Cotton.
Chemical applications of group theory. John Wiley & Sons, 1991.
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector neurons: A general framework for so (3)-equivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 12200-12209, 2021.
Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review B, 99(1):0,4 2019.
Genevieve Dusson, Markus Bachmayr, Gábor Csányi, Ralf Drautz, Simon Etter, Cas van der Oord, and Christoph Ortner. Atomic cluster expansion: Completeness, efficiency and stability, Journal of Computational Physics, 454:1 10946, 2022.
Alexandre Duval, Simon Mathis, Chaitanya Joshi, Victor Schmidt, Santiago Miret, Fragkiskos Malliaros, Taco Cohen, Pietro Lio, Yoshua Bengio, and Michael Bronstein. A hitchhiker's guide to geometric gnns for 3d atomic systems. arXiv preprint arXiv:2312.05511,222 2023.
Nadav Dym and Haggai Maron. On the universality of rotation equivariant point cloud networks. In International Conference Learning Representations, 2021. URL https: http::::////eeeeveeeeeee hett/forum? d= 6NF BvW1 RXaG.
Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional
neural networks for equivariance to lie groups on arbitrary continuous data. International Conference on Machine Learning, pp. 3165-3176. PMLR, 2020.
Thorben Frank, Oliver Unke, and Klaus-Robert Müller. So3krates: Equivariant attention for interac- tions on arbitrary length-scales in molecular systems. Advances in Neural Information Processing Systems, 35:29400-2090 29410, 2022.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (r)-tsfsfformr:s: 3d roto- translation equivariant attention networks. Advances in neural information processing systems, 33: 1970-1981, 2020.
Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In NeurlPS-W, 2020a.
Johannes Gasteiger, Janek Groß, and Stephan Günnemann, Directional message passing for molecular graphs. In International Conference Learning Representations, 2020b. URL https:// openreview net/
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34: 6790-6802, 2021.
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C Lawrence Zitnick, and Abhishek Das. Gemnet-oc: developing graph neural networks for large and diverse molecular simulation datasets. Transactions on Machine Learning Research, 2022.
Mario Geiger and Tess Smidt. e3nn: Fuclidean neural networks. arXiv preprint arXiv:2207.09453. 2022.
Xiaoxum Gong, He Li, Nianlong Zou, Runzhang Xu, Wenhui Duan, and Yong Xu. General framework for (3)-equivariant neural network representation of density functional theory hamiltonian. Nature Communications, 14(1):2848, 2023.
11

---

Published as a conference paper at ICLR 2024
Mojtaba Haghighatlari, Jie Li, Xingyi Guan, Oufan Zhang, Akshaya Das, Christopher Stein, Farnaz Heidar Heidar-Zadeh, Meili Liu, Martin Head-Gordon, Luke Bertels, et al. Newtonnet: A newtonian message passing network for deep learning of interatomic potentials and forces. Digital Discovery, 133):333-34,, 2022.
BHLB Hammer, Lars Bruno Hansen, and Jens Kehlet Nørskov. Improved adsorption energetics within censity-functions theory using revised bardew-burke-erzeeee functionals. Physical review B, 59(1 1):7413, 1999.
Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: A survey, arXiv preprint XxXv:2202.07233000 2022.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In Computer Vision-ECCY 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pp. 646-661. Springer, 2016.
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lictransformer: Equivariant self-attention for lie groups. In International Conference on Machine Learning, pp. 4533-4543 PMLR, 2021.
Sigeru Huzinaga. Molecular integrals. Progress of Theoretical Physies Supplement, 40:52-77, 1967.
Nadir Jeevanjee. An introduction to tensors and group theory for physicists. Springer, 2011.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2021. URL httpt::/://oo htttt:::://ppenzeiieee new /ortfform?
1d=1YLJDvSxx
Chaitanya K. Joshi, Cristian Bodnar, Simon V Mathis, Taco Cohen, and Pietro Lio. On the expressive power of geometric graph neural networks. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 15330-15333 15355. PMLR, 23-29 Jul 2023. URL httph:://///tc...tt http::ss/ppcceedinngg mlr. press/v202 joshi23a html.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 59667877):583-5588 2021.
Kuzma Khrabrov, Ilya Shenbin, Alexander Ryabov, Artem Tsypin, Alexander Telepov, Anton Alekseev, Alexander Grishin, Pavel Strashnov, Petr Zhilyaev, Sergey Nikolenko, et al. nabladft: Large-scale conformational energy and hamiltonian prediction benchmark and dataset. Physical Chemistry Chemical Physics, 24(42):25553-25588 2022.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Yvette Kosmann-Sehwarzbaal et al. Groups and symmetries. Springer, 2010.
Dávid Péter Kovács, Cas van der Oord, Jiri Kucera, Alice EA Allen, Daniel Cole, Christoph Ortner, and Gábor Csányi Linear atomic cluster expansion force fields for organic molecules: beyond rmse. Journal of chemical theory and computation, 17(12):766-7711,,2 2021.
Dávid Péter Kovács, Ilyes Batatia, Eszter Sára Arany, and Gábor Csányi. Evaluation of the MACE force field architecture; From medicinal chemistry to materials science. The Journal of Chemical Physics, 14444), 07 2023. ISSN 0021-9606. doi: 10.1063/5.015532 URL https: //doi.org////oo// / 0. 10006.56000 0000/.00555.............1111115. 0155322.
Janice Lan, Aini Palizhati, Muhammed Shuaibi, Brandon M Wood, Brook Wander, Abhishek Das, Matt Uyttendaele, Lawrence Zitnick, and Zachary W Ulissi. Adsorbml: Accelerating adsorption energy calculations with machine learning. arXiv preprint arXiv:221 1,,46666,6 2022.
12

---

Published as a conference paper at ICLR 2024
He Li, Zun Wang, Nianlong Zou, Meng Ye, Runzhang Xu, Xiaoxun Gong, Wenhui Duan, and Yong Xu. Deep-learning density functional theory hamiltonian for efficient ab initio electronic-structio calculation. Nature Computational Science, 2(6):367-377, 2022.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023. URL
https: http:://oppereeeee /ppp//peeeeeww.e////e/ee net Forum?id=KKin EARR---O
Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: /ppeneeeeeeeeeeee net/ / forum? iid= mCCBKZmrZD.
Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Tie- Liu, Liwei Wang, and Di He. One transformer can understand both 2d & 3d molecular data. In The Eleventh International Conference Learning Representations, 2023, URL https: https///opppeeeeeeeeee net forum? id=v2Tp1oPV3PC.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Korbbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Jigyasa Nigam, Sergey Pozdnyakov, Guillaume Fraux, and Michele Ceriotti. Unified theory of atom-centered representations and message-passing machine-leaming schemes. The Journal of Chemical Physics, 156(20), 2022.
Saro Passaro and C. Lawrence Zitnick. Reducing SO(3) convolutions to SO(2) for efficient equivariant GNNs. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learn- ing, volume 202 of Proceedings of Machine Learning Research, pp. 27420-27430 PMLR, 23-29 Jul 2023. URL https: httptttsssppppeeeeeeeeett http:///ooceeeenseee/ee/ mlr press 222
html.
Richard Pio. Euler angle transformations. IEEE Transactions on automatic control, 11(4):707-715, 1966.
John Proakis. Digital signal processing: principles, algorithms, and applications, 4/E. Pearson Education India, 2007.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):-7,, 2014.
Joseph Rotman. An introduction to the theory of groups Graduate Texts in Mathematics, 1995.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323-9332, PMLR, 2021.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-9388, PMLR, 2021.
Kristof T Schütt, Huziel Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):24122, 2018.
Alexander V Shapeev. Moment tensor potentials: class of systematically improvable interatomic potentials. Multiscale Modeling & Simulation, 14(3):11531113,, 2016.
Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tic- Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv preprint ar arXiv:2203.0483 2022.
Muhammed Shuaibi, Adcesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C Lawrence Zitncc.. Rotation invariant graph neural networks using spin convolutions. arXiv preprint arXiv: (06.09575 2021.
13

---

Published as a conference paper at ICLR 2024
Guillem Simeon and Gianni De Fabritiis. Tensomet: Cartesian tensor representations for efficient learning of molecular potentials. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
Armold Sommerfeld. Partial differential equations in physics. Academic press, 1949.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference Learning Representations, 2022. URL https: // //openreview....////////e .nt. 2 zeteummid-zNiizqee
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-an and translatonn-cqiivaraan neural networks for 3d point clouds. arXiv preprint arXiv: 1X2::8802.08219 2018.
Aidan P Thompson, Laura P Swiler, Christian R Trott, Stephen M Foiles, and Garritt J Tucker. Spectral neighbor analysis method for automated generation quantum-accurate interatomic potentials. Journal of Computational Physics, 18-:316-330, 2015.
Oliver Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus-Robers Müller. Se (3)-euvivariant prediction of molecular wavefunctions and electronic densities. Advances in Neural Information Processing Systems, 34: 34:14434-14447, 2021.
Mikaela Angelina Uy, Quang-Hies Pham, Binh-Son Hua, Thanh Nguyen, and Sai-Kit Yeung. Revisiting point cloud classification: A new benchmark dataset and classification model real- world data. In Proceedings of the IEEF/CVE international conference on computer vision, pp. 1588-15977 2019.
Dmitriui Aleksandrovich Varshalovich, Anatolij Nikolaevic Moskalev, and Valerii Kel manovich Khersonskii. Quantum theory of angular momentum. World Scientific, 1988.
Yuanqing Wang and John Chodera. Spatial attention kinetic networks with <qnivvvuivai In The Eleventh International Conference on Learning Representations, 2023. URL https: ///pppnreeiew..e/////// net wit.:////ooooddddIIdd0d03
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data, Advances Neural Information Processing Systems, 31, 2018.
EP Wigner. On the matrices which reduce the kronecker products of representations of sr groups. manuscript (1940), published in lc biedenharn and h. van dam (eds.). quantum theory of angular momentum, 87-133, 1965.
Eugene Wigner. Group theory: and its application the quantum mechanics of atomic spectra, volume 5. Elsevier, 2012.
Zhirong Wu. Shuran Song. Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1912-1920, 2015.
Chengxuan Ying. Tianle Cai, Shengjic Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances Neural Information Processing Systems, 34:28877-288888 2021a.
Chengxuan Ying, Mingqi Yang, Shuxin Zheng, Guolin Ke, Shengjie Luo, Tianle Cai, Chenglin Wu, Yuxin Wang, Yanming Shen, and Di He. First place solution of kdd cup 2021 ogb large-scale challenge graph prediction track. arXiv preprint 2rXiv:2106.082 2021b.
Haiyang Yu, Meng Liu, Youzhi Luo, Alex Strasser, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. QH9: A quantum hamiltonian prediction benchmark for QM9 molecules. In Thinth-eeentth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023a. URL https: /ttp///ppp//e//ee ew wwt.ee...eeeeeeeeeeeeeeeeeeeeeee forum? d1 uR 9N39A.
14

---

Published as a conference paper at ICLR 2024
Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Efficient and equiv- ariant graph networks for predicting quantum Hamiltonian. In Andreas Krause, Emma Brun- skill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Pro- ceedings of Machine Learning Research, pp. 40412 40424. PMLR, 23-29 Jul 2023b. URL https: http:://ppoceedinn...... wtt.sss.mm.ppppeeee.. / v2 222222y22233 htppyyu22......
Bohang Zhang, Shengjie Luo, Liwei Wang, and Di He. Rethinking the expressive power of GNNs via graph biconnectivity. In The Eleventh International Conference on Learning Representations, 2023a. URL htwpt:ttt//..o...... wpw.reveee.eeeee.eeeeeeeem.e...e net/ forum? d= d 6KoT3.
Xuan Zhang, Limei Wang. Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin. Zhao Xu, Keqiang Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423, 2023b.
Ya A Granovskinand AS Zhedanov al. Nature of the symmetry group of the 6j-symbol. Zh. Eksp. Teor. Fiz., 94:49, 1988.
C. Lawrence Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ward Ulissi, and Brandon M Wood. Spherical channels for modeling atomic interactions. In Advances in Neural Information Processing Systems, 2022.
15

---

Published as a conference paper at ICLR 2024
A BACKGROUND
A.1 GROUP THRORY
Groups. Let G' denote group, which is a non-empty set together with a binary operation denoted
as o: G x G - G. The group operation o satisfies the following properties known as group axioms:
Associativity: For all a, b, E G, we have (b c) = (a b) o c a boc.
Identity: For all a E G, there exists an element e G such that a o coo a = a.
Inverse: For each
G, there exists an element a -1 E G such that a Closure: For all a, b E G. a abbe E G.
aoaal
For brevity, the group operation o can be omitted and a o can be abbreviated as ab. Groups can be finite or infinite, countable or uncountable, compact or non-compact. The Groups are algebraic structures that are widely used across physics (Cornwell, 1997), chemistry (Cotton, 1991), materials science (Burns, 2014), 4), and other scientific areas.
Group Representations. As introduced above, the group operation o describes how the group elements interact with each other. To describe how the group elements act on vector spaces, group representations are used. Formally, representation of a group G on a vector space V is mapping p² G GL. Here, GLV denotes the general linear group on vector space V, which the set of invertible matrices together with the binary operation of ordinary matrix multiplication. V is called the representation space and the dimension mapping pK (g) parameterized by the group element g E G is thus an invertible linear transformation on V. which follows the group structure (also known as group homomorphism') that p(g192)e ((2)) for all 01.92 E G, V.
V is called the dimension of the representation. The
Assuming is the matrices in ddd Let ((g) denote the representation of g in the matrix form. Representations of the rotation group in the three-dimensional Fuclidean space are matrices R E R333 deee()) = 1, and the representations act on Rd via matrix multiplications D(g)x= Rx.
real vector space Rd, the group representations are invertible
For any two representations, we say D(g) and D'(g) are equivalent if there exists a similarity transform such that:
D'(g) - DDD(D)
(8)
That is. both representations (g) and ((g) represent the same group element g with different basis, and the change of basis is carried out by Q.
Examples of Groups. While the Group is an abstract mathematical concept, there are numerous classes of objects and transformations that have group structures. In this paper, we are interested in the group of translations, rotations, and reflections in the three-dimensional Euclidean space, which is known as the 3D Euclidean group and denoted as E(3). For each element g in the E33) group, its representation on R can be parameterized by pairs of translation vectors E R3 and orthogonal transformation matrices R E R det(R) =1, i.e., p*² (g) (t, R). Given a vector x E 3R we have p² (g) [g) :=
= Rx + t.. The group operation and inverse of F((3 are defined by g gog g = (Rt' + t. RR' and g 1 (R-1x, R 1),, where g = (t, R). g = (t', R`) E E(3).
The E(3) group is semi-direct product of the group of translations R.3 with the group of orthogonal transformations 0(3), i.e., E(3) R x O()). Thus, the group elements in E(3) can be decomposed an ( (r)))raanfoormatoonn (rotation and/or reflection) followed by a translation. Note that translations casy to deal with in modern neural networks. Therefore, focus mainly on the orthogonal group (3), which consists of rotations and reflections. For each clement g in the (()) group, its represen- tation on 3 is thus parameterized by orthogonal transformation matrices R 3x3 .det(R) ie.. g = R. Similarly, the group operation o and inverse of O(3) are defined by g o g' == RR' and g = R-1, where g R, R' E O(3). The special orthogonal group SO(3) has the same group operation and inverse, but excludes reflections. Thus, each element = R E (0( has (R) 1.
+1,
Given two groups (G, og) and (H, on) function mapping G + H is said to be group homomor- phism if for all g, h E G we have T(g) on (h) <g oa h)
16

---

Published as a conference paper at ICLR 2024
A.2 EQUIVARIANCE
Formally, let Ø X' + y denote a function mapping between vector spaces. Given a group G. let p* and p² denote its group representations X' and y respectively. function @ X + y is said to be equivariant it satisfies the following condition:
(=(())((=)] (()(())==)) for all G, e X.
(9)
When p² = =3' for all g G, it is special case of equivariance which is known as invariance, i.e., function Φ + said to be invariant if satisfies the following condition:
(xx) = ó(p [x), for all E G, E X.
(10)
Intuitively, an equivariant function mapping transforms the output predictably in response to trans- formations the input, whereas invariant function mapping produces output that remains unchanged by transformations applied to the input. For example, the total energy of a molecular system is invariant the rotations, translations, and reflections acted on each position of atoms, i.e., the energy function with respect to the atom positions is an invariant function to the E(3) group. If a rotation is applied to each position of atoms, the direction of the force on each atom will also rotate, i... the force function is an equivariant function to the SO(3) group.
A.3 IRREDUCIBLE REPRESENTATIONS
Subrepresentations and Irreducible Representations. Let W be subspace of V. When W is invariant under the group G, i.e., pd (g)xr W for all x E W and g E G. the representation p of the group on the vector space W called a subrepresentation of p The subspace W also called G-invariant. For any representation p of a group G, it has two trivial subrepresentations: (1) pV itself; (2) the zero representation, both of which satisfy the above definition.
After defining the concept of subrepresentations, sentations, which plays an essential role in modern rotationally equivariant neural networks. A representation pV G GL(V) is said to be irreducible if it has only trivial subrepresentations. If there is a proper nontrivial G-invariant subspace, p is said to be reducible.
ready to introduce irreducible repre-
Wigner-D matrices. We introduce the Wigner-ID matrices, which serve as the irreducible repre- sentations of the rotation group SO(3). For vector spaces of dimension 2/ + 1, 1, = 0, 1,2... there exists collection of representations indexed with for each element in the SO(3) group, which are called Wigner-D matrices and denoted as D (g) for g E (3))). Each Wigner-D matrix ( is the irreducible representation of the (O(3) group on subspace with dimension 21 + 1, and any matrix representation D(g) of SO(3) on vector space V can be reduced to an equivalent block diagonal matrix representation with Wigner-D matrices along the diagonal:
D (
D(g) ( Q- + (()) + Q..) --
D
Q
(1I)
Here, Q is the change of basis that makes D(g) and the block diagonal matrix representation equivalent. This equivalence provides an alternative perspective that the representations that act on the vector space are made up of each Wigner-D matrix D(4), which only acts on subspace Vi, with dimension 21, + 1 of V. Then, the vector space V can be factorized into VI, @ V1, o .... By convention, the space Vi, is called a rype-l, vector space and x E Vi, called type-l; vector, which transformed by the Wigner-D matrix ((4)), In Appendix A.1, we have already introduced the Wigner-D matrix in the type-1 space, i.e., ( R 33x3 det (R) = The type-O space is RR R.
and is invariant to the SO(3) group, i.e., D (0) (g) ar x for all g E G and
As previously mentioned, we are interested in the F(3) group consisting of translations, rotations, and reflections. The Wigner-D matrix representations can be naturally extended the (3)) group by further including the reflections via a direct product, while the translations can be casily dealt with in modern equivariant networks. We refer the readers Commins (1995); Wigner (2012) for the explicit forms of the Wigner-D matrices and the e3nn² codebase (Geiger & Smidt, 2022) for the numerical implementation.
Pttpps:// hom.bccmm//mm
17

---

Published as a conference paper at ICLR 2024
A.4 SPHERICAL HARMONICS
Spherical harmonics are a class of functions defined on the unit sphere S These functions form an orthonormal basis set, and thus every function defined on S can expressed a series of spherical harmonics, which is similar to the Fourier series. By convention, we use Y : 5² - R, -I S m S 1,1=0, = 0, 0, ... to denote the real spherical harmonics, where I and m denote the degrees and orders of spherical harmonics respectively. Both spherical coordinates (0, v)) 0 E [0, TI, E [0, 2z) and Cartesian coordinates the input variables of the spherical harmonics. Next, we will introduce several important properties of the spherical harmonics used in equivariant networks, and their connections to the irreducible representations of the SO(3) group.
RR ||||| are widely used to express the points on the sphere S as
Spherical harmonics expansion. Formally, on the unit sphere S any qquare-integrable function f: S - R can be expanded as a linear combination of spherical harmonics:
S(0, v))
yyy m (0, v)
(12)
I=0 m=-!
This expansion holds in the sense of mean-square convergence:
2x ('" (5(0 -
lim
Σ
y (0, ii sin OdOdy
(13)
Similar to the Fourier series, the spherical harmonic coefficients fm can be obtained by using the orthonormal property of spherical harmonics:
==/ [ f(0, (0 v) sin 0d0dy.
2
p0)
(14)
Equivariance of spherical harmonics. The spherical harmonics have various intriguing properties, especially with respect to the equivariance of some groups consisting of transformations of interest, e.g., rotations and reflections:
Rotations. Given rotation g SO(3) and its representation on R33 as p (g) Rax3 dct(R.) = 1, the unit vector r E R3, ||||| = 1 is rotated to r' = Rr. Then we have:
R
y (Rr) = Y (r') = D m (1) (r).
(15)
Let (r) = (r), (r), y(u)) (Rr) D ( ())) (r), for all SSO(3,, S?,
r, (('((r))] (T)]T. then the above property can be abbreviated as,
(16)
where D'(g) denotes the matrix representation of the SO(3) group element y on the 2/ + 1 vector space, and D (1)
mm (g) corresponds to the (m, m')-th matrix element.
Reflections. The reflection transformation (also called inversion) is represented by the operator (T(r) (--r). Then, we have:
m (-r) = (-1) (r).
(17)
Connections between spherical harmonies and irreducible representations of SO The equivariance of spherical harmonics indicates that there are strong connections between this function class and group representation theory. The yi(")s of degree 1 form basis set of functions for the irreducible representation of the S0(3) group on the 2/ + vector space Messmann-Schwa//hie et al., 2010). Each irreducible representation of SO(3) can be realized a finite-dimensional Hilbert space of functions on the sphere, the restrictions of harmonic homogeneous polynomials' of a given
GGiven multivariate polynomial P over field, if the polynomial p satisfies: (1) AP 0; (2)
P(A Ar.) = P((1 -.. Jn); is called harmonic homogenous polynomial.
18

---

Published as a conference paper at ICLR 2024
degree, and such representation is unitary. The orthonormal basis of this space is exactly the introduced spherical harmonics. From this perspective, the spherical harmonics are proportional to the Wigner-D matrix elements: Ym (r) D mo((s) (g) given that the rotation g is parameterized by Fuler angles (Pio, 1966) 8,,8,7), and the sphere point r is obtained by using the rotation g' parameterized by (a, 3, 0) on the z-unit vector. By definition, the spherical harmonics are invariant to the r-rotatonn, and the relations to the Wigner-D matrix elements show the connections to the irreducible representations of the SO(3) group.
A.5 TENSOR PRODUCT OF REPRESENTATIONS
In the previous subsections, we defined the group representations on vector spaces. To map the representations between vector spaces, the tensor product operation provides a principled approach, with useful properties for constructing equivariant networks. Formally, given two vectors x = (21,,-.,, Fdi
Rd2, the tensor product between and y follows x @ y xy E RR xd= Given two representations T1 : G + (V, and 72 G - GT((V) on vector spaces and 1/2 respectively, the tensor product of representations given by क @ m2 G GL(V & V22, i.e., गı @ (g) (g) @ (g))), where (g) @ тe (g(g) is the tensor product of linear maps previously defined. This operation maps representations on spaces V1, V22 to the V12 & V/2 space with dimension being (21, + 1) x (212 + 1). Then, the equivariance of the tensor product operation
Rd: y = (/,, Vdg
the (3) group natural show:
(D (1 (g) & (a) (g))(x @ x²) = (D (g)x) @ (D (22 ( (g)xx)) for all x1 E V1, X2 E V2, (18)
where (g) & (g) is matrix representation of g E SS(() in the vector space V1 s V/2
Clebsch-Gordan Tensor Product. As the tensor product of representations is still a representation on the corresponding vector space, we can further decompose it into a direct sum of irreducible representations, shown in Eqn. (11). Then, the resulting vector of the tensor product can be organized via a change of basis into parts that are individually transformed by Wigner-D matrices of different types, i.c. X1 @ X2
Vo V1 ... Vi + ..., with Vi being the type-/ vector space.
The Clebsch-Gordan Tensor Product provides a way to explicitly calculate the resulting type-l vector after the direct sum decomposition. Given ((1) E 2 and x(a2) E IR22++1, the m-th element of the resulting type-/ vector of the tensor product between (()) and (!a) is given by:
4 Σ
(x Beg x
(.)) (12,m)) 3 aa
(19)
MLL1 11
12
where (l..m)
(fi.mi)(a.m) are the Clebsch-Gordan coefficients. Interestingly, the (lebsch-Gordas coeffi- cients are also closely related to angular momentum theory in quantum mechanics (Varshalovich et al., 1988), and have various intriguing properties that have underlying relationships to other important concepts used in this work.
Properties of Clebsch-Gordan coefficients. The Clebsch-Gorde coefficients have the following (l.ma1)(2.mm) 0f < 141 12| or | > l + +12 That is, type-0 vector and a type-1 vector cannot produce type-2 vector via the (lebsch-Gordan tensor product. The Clebsch-Gordan coefficients have the following orthogonality and symmetry properties:
selection rules: (l.m)
Orthogonality:
Σ
(1.m) (1.ma)(2aama)) (1.ma)llamma)) δi 8mm'
(1m))
(20)
Σ I.
(Lm) (la.m)
(m)))
(1: Smymg 8 Omggin
Symmetry: (l.m)
= (-1)) (((()
(h,
(1,-
(21)
(-1) (2) (2.mm))ls.m1))
19

---

Published as a conference paper at ICLR 2024
(-1)4-m)
21 + 1
(1))=ma)
(-1))a++
21 + 1
212
(11,,))(,mm)
V
2/1 (1-m)((2,mma)
(-1)k--m
21 + 1
(llammg)
((.m))1:maaa) (-1)katma
21 + 1
212 +
212 (2-=ma)(..m)
A.6 RELATIONSHIP BETWEEN CLEBSCH- WIGNER 3-J AND GAUNT COFFFICIENTS.
We are motivated by the orthogonality and symmetry of the Clebsch- Gordann coefficients to explore further relations two other concepts, Wigner 3-j symbols and Gaunt coefficients (Wigner, 1965). These concepts are widely used in quantum mechanics, and inspire the approaches presented in the main body of this work.
In quantum mechanics, the Wigner 3-j symbols are an alternative to Clebsch-Gordan coefficients, with the purpose of adding angular momenta. Mathematically, the Wigner 3-j symbols are related to Clebsch- coefficients as:
cll.m (0s.m)))
(-1)-1+++
,12 m1
l2
1
(22)
22 +1
m2
m
(1
li mi
l2
1 m
where
denotes the Wigner 3-j symbol that has the following explicit expression:
ma
l2
(m= + m2 m. (-1)) (-1)41-12-
mi
ma
(11 1 1)1(11 ()!(
+1++++++1)
(11 - my)! (f1 + m²) (12 ma)! (12 + m2)! m): m)!
(-1)*
21 mi k)! (12 my k:]! (1 1/ my 1:J! - my + k)t
(3
where K and N are as K max -0,, - N min -12,2211,1111111 m, 1 m2)
and factorials of negative numbers are conventionally taken equal to zero. The Wigner 3-j symbols are highly symmetric. They are also invariant under an even permutation of their columns, and enumerate phase factor under an odd permutation or changing the sign of the numbers. The Wigner 3-j symbols also have the so-called Regge symmetries (Zhedanov et al., 1988).
The Wigner 3-j symbols are related to spherical harmonics via the Gaunt coefficients. which are the integrals of the products of three spherical harmonies:
2
((l (H, ())r (0, (i)
ab) sin OdAds
(24)
(2 + (212 + 1) (21 + 1)
4
12 000
1
11 )
4 mi
12 m2
4T
The Gaunt coefficients are used quantum mechanical computations, e.g., multi-center molecular integrals (Huzinaga, 1967). In this work, we use their relations to the Clebsch-Gordan coefficients to accelerate the tensor product of representations, thereby improving the efficiency of different classes of equivariant operations. From the above relations, conclude this subsection with the following equivalence:
2nT
y (1, +) (2) (0, ( +) sin ddldg: = a(l)) (m.)) (1,,2) (ly (4amm))) mm)))))
(25)
(1) where (5,))) is constant only determined by 41 l2, and l. This equivalence inspires the development of the Gaunt Tensor Product presented in the main body of this work.
We note one difference between employing Clebsch-Gordan coefficients and Gaunt coefficients in the tensor product computation, namely that the Gaunt coefficients impose more symmetry
20

---

Published as a conference paper at ICLR 2024
constraints. For instance, the tensor product with Gaunt coefficients excludes pseudovectors from the output. However, we empirically observe that this does not hurt performance (See the sanity check experiments in Section 4). We also note that several existing models that achieved strong performance in real-world applications (Brandstetter et al., 2022; Batzner et al., 2022; Batatia et al., 2022b) that consider the parity of irreps do not use irreps of even parities, e.g., pseudovectors, in computing cquivariant representations. In these cases, the tensor product with Clebsch-Gordan coefficients is equivalent to the tensor product with Gaunt coefficients.
A.7 ANGULAR MOMENTUM, TENSOR OPERATORS AND WIGNER-ECKART THEOREM
In the previous subsection, we introduced the relations between the Clebsch-Gordan coefficients and the Gaunt coefficients, which serve as the key to motivate the new perspective tensor products of irreps via Gaunt coefficients in Section 3.1. For completeness, we thoroughly elaborate on the necessary knowledge to derive this equivalence in this subsection. First, leverage the lens of angular momentum in quantum mechanics to describe spherical harmonics and Clebsch-Gordan coefficients. We then elaborate on the concepts of tensor operators and the Wigner-Fckart theorem, which we use to relate Clebsch-Gordan coefficients and Gaunt coefficients, as introduced below.
A.7.1 PRELIMINARY
We introduce several mathematical concepts that are used in quantum mechanics.
Hilbert space and Dirac notation ("bra-ket") notation). A Hilbert space H. is abstract vector space satisfying the following conditions: (1) it is equipped with the inner product operation; (2) it is complete metric space with respect to the distance function induced by the inner product. A typical example of a Hilbert space is the three-dimensional Euclidean vector space R which is equipped with inner product operation known as the dot product R3 x R R. As general and abstract concept, Hilbert spaces allow mathematical tools like linear algebra for (finite-dimeng) dimensiaim) Euclidean vector spaces spaces that may be infinit-ddimensiona c.g., function spaces frequently used in physics.
In Hilbert space H. we instead use bra-ket notation to describe our system. Vectors in the Hilbert space H. are written as kets: la) E H. The operational rules from Euclidean vector space are also generalized as,
1. Addition: (a) + 13) 19) + la) (commutative), (a) + ((9) + In) = (Ia) + 13)) + In)
(associative);
2. Scalar product: (c1 + c2)la) eala) + o2)a), (o)la)) = (c1e2)la), c(la) + 1B})
cla) + cla), laa) = a). where the scalars are complex c E C (as we will deal exclusively
with complex Hilbert spaces for quantum mechanics);
3. Null ket [0): |a) + (0) = Ja). 0/a) = 10);
4. Inverse: any ket la) has a corresponding inverse - a). and we have la) +1 a) = 10). aa) = lla. For an N V-dimensional Hilbert space, there exists set of N kets which constitute basis, and any ket can be expressed as a sum of the basis kets, i.e., (a) =
Σn a,lA,),, az E C.
We further define the inner product in a Hilbert space, i.e., a map which takes two kets and returns a scalar (complex number): (:-)) 7 H + C. With this definition, we now define the bra notation, which is written as a backwards-facing ket (or). We have the one-to-one correspondence from kets to bras and vice-versa:
1. la) (al;
2. cala) col.3 + c c3 c3)
where c cf E C are the complex conjugates of ca. C3 C. By using the bra-ket notation, the inner product of two kets in a Hilbert space can be written as a product of a bra and a ket: (Ia), 13)) = (a)3). In this notation, the inner product has the following properties:
21

---

Published as a conference paper at ICLR 2024
1. linearity (in the bras): (c(al)]8) c(a)3) and ((a1) (a2)) (0]]8) + (a2)3);
2. linearity (in the kets): (al(c[3)) c(a and (al (a]())) (82) = -(o(8)) + (al
3. conjugate symmetry: (als) (Bla)"
4. positive comi-definitenes:: (ala) 2 0.
From (3), we can see that the order of the kets matters in the product. In fact, bra denotes linear map that maps kets in H. to C, i.e. (la), -) = (a). We say that two kets |a) and (3) are orthogonal their inner product is zero, i... (a]3) = (3)a) = 0 = [a)_ Again, analogous to Euclidean vector space, we use ||a||| = v a) to denote the norm of the ket (a).
As previously mentioned, the three-dimensional Fuclidean vector space is one example of the Hilbert space with finite dimensions. Here, we provide an example of diinite-dimensione Hilbert spaces for illustration. Consider a function y(z) = >> On sin naw ), I R: each ket in this space represents a function, and the basis kets are given by the infinite set of functions [n) sin( J. (we can use n in the ket notation for brevity, because it the identity index in this Hilbert space). In this space, the inner product is defined between functions as (y1 y2) = y/2 (2) dz.
naw
Operators. To describe the mapping between kets in Hilbert space, we introduce the concept of operators, i.e., O: : H - H. In general, properties of operators include:
1. Equality: À = B if Ála) Bla) for all la);
2. Null: there exists special operator Ô such that 0/a) (0) for all |a);
3. Identity: there exists another special operator I such that l/a) = la) for all Ja);
4. Addition: A+ B B + Á and À + (B + + - (A B) +C;
5. Scalar multiplication: (cΛ)lo) A(cla));
6. Multiplication: (ÁB)Ja) A(Bla)))
linear operator has the distributive law, i.e., A(cala) + cg(3)) = ca Alo) + ccg Using R3 as an example, linear operators are thus matrices 3R3x3
We use A-1 1 to denote the inverse operator of Å, and Å Ala) = ila). Equipped with the inner product of Hilbert space, we can define the adjoint of an operator denoted with dagger by (Ia),A|3) Á|33) = (Á) [aa,;,{)))))) for any choice of (a) and (8). Combined with the operations, the adjoint has the following properties: (1) (cA)t e" Åt; (2) (Ā + B)t At + B"; (3) (AB)t = Btt Ät. One important class operators are those that are self-adjoint, i.e., A Á, which are also known as Hermitian operators. If the adjoint gives the inverse operator, i.c., Øt = --11 these are called unitary operators, and act like rotation in the Hilbert space.
Similar to matrices, the order of the operators matters when we multiply them, and the difference between the two orderings has special symbol called the commutator: [Å, B) = ÁB - BA. If the commutator of two operators zero, say they commute. Properties of commutators include:
1. [ = 0;
2. [Á, ß] - 1B, A):
3. [Ā+ T.BB ĆJ = (A, C1 |B, ČJ:
4. [A, = [À. + B[Á,C]:
5. [ÁB, A[B, 0 + +(A,C]B:
6. [Á, CJJ + [8, [C, Ál + [Á, = 0
One intriguing property of two Hermitian operators that commute with each other is that we can simultaneously find their eigenkets. This property relates to the measurement in quantum mechanics.
22

---

Published as a conference paper at ICLR 2024
Eigenkets, basis kets and matrix elements of operators. Analogous to linear algebra in the Euclidean vector space, define the eigenkets of operator Å as the kets which are left invariant up scalar multiplication, i.e., Ala) ala). where the scalar is the eigenvalue associated with eigenket (a) (by convention, we use eigenvalues as indexes for eigenkets). Note that eigenkets are also called eigenstates in the context of quantum mechanies.
Recall that each ket in a Hilbert space can be expressed as a linear combination of the basis kets. In fact, there always exists an orthonormal basis, i.c., set of basis kets which are mutually orthogonal and have norm 1. Let us use (en) to denote such an orthonormal basis. Then we have (em aa) =
En (m = am, and |a) = (en))))))))) In fact, we introduce a special operator called projection operator, i.e., A, = (en) (e,) which project (a) onto one of the basis kets (ex).
By using the basis kets, we can now represent any ket column vector consisting of the coefficients
corresponding to the basis kets: lab) = Σn (n(en) = [[v1, .... Wn...." and the associated bra (v8) the conjugate transpose, i.e., a row vector whose entries are the complex conjugates of the tha
(vl >n v (enl i ..... Now, we have (xle² = E, X ths and (( (xleey vix..
Specially, we have (ei) being almost zero matrix except for the (1, 1)-th entry being 1.
Thus, any operator Å can be similarly represented as matrix form:
(ei [lee) (eilAle2) (e2 Ale (e2lAle2)
Ā = ΣΣ (em (len)(eal
(26)
The entries of the matrix (cm (lle)) are known as matrix elements of operators. There also exists a relation between the matrix elements of Å and its adjoint Åt: (emlAlen) = (enlA' lem)".
A.7.2 A BRIEF OVERVIEW OF ANGULAR MOMENTUM QUANTUM MECHANICS
Momentum Operator and Position Operator. Before we go into the concept of angular momen- tum in quantum mechanics, discuss two basic examples of operators that we will use later. Here, we use a concrete example to see how previously introduced concepts are reflected real-world systems: the ID De Broglie plane wave function for free particle with momentum p and energy E, ie.. (xx,) t) = <k(kz=wt), where = fil, E = fw, E 2m denotes the Planck Constant, and m denotes the mass of the particle. To extract the value of momentum from the wave function, can use the following operator:
p²
h 2 (z, i Oz
h 2 cilko-wt) h az
- (ik)e = (ik) ((t)),)))))))))))))))(())))))))) (z,
(x,
(27) We then define the momentum operator, p. as p- - - Moreover, we can see that (z, t) indeed lies in Hilbert functional space, and is an eigenstate of the momentum operator because V (x, 1) left invariant up to a scalar p being multiplied after the action p.
Another important operator is the position operator 2, which acts on functions of x, and gives another function of as follows:
fx(x)) = =f)))))())
(28)
In fact, both the operator à and P are related. Given some arbitrary function Φ(x), the commutator [ p) has the following equivalence:
[s. f))() = (āp -f))(())= = &pp(2z) - p)((()) = #(fp(z)) - (((((())) = (= 8б(4)
) (()((z))
Ox
ab(+) iiu dr
+ ih a
(ro(#)) -itur 05(z)
+ itu 0б(x)
+ (ibb(())) il(((x=.
Dr
dx
dr
(29)
Thus, we obtain a fundamental commutation relation in quantum mechanics: [ p1 = ih. is also straightforward to generalize both á and P to three dimensions:
0
= (pr, Py: P2)) = -- i = -ih
Ox
Oy Oz
(30)
x = (ž, 2,
23

---

Published as a conference paper at ICLR 2024
And the commutation relationships follows:
[ Pa] = ih, [Ā. fy.pyl= ih, [2, ih,
(31)
with all other commutators involving the three coordinates and momenta being zero.
Angular Momentum Operator: We are now ready to introduce the concept of angular momentum in quantum mechanics. This helps us understand the underlying relations between spherical harmonics and Clebsch-Gordan coefficients, as stated in Section 3.1. In classical physics, the angular momentum of a particle with momentum p and position r is defined by L
p. and each component of
L = (Lz, Ly Lz) is given by L = VP:- 2py. Ly = 2pv rPz: Lg rPy ypz. Similarly, the angular momentum operator L (I.,, by L.)) can be obtained follows:
L uī= E
L, 2p. 2=22
(32)
L = Tp. - DP..
From the definition, we can check that the angular momentum components are Hermitian operators: it. = Lun, it, = Ly, it = Lz. The square of L is also Hermitian, which is Í = L3 1 + L. Therefore, all the angular momentum operators have real eigenvalues.
Eigenstates and eigenvalues of Angular Momentum Operators. Recall that we can simultane- ously find the eigenstates of two Hermitian operators that commute with each other. fact, we have the following commutation relations for angular momentum operators:
[i., [yy = i [T. i..] = ih [[.,, i.]] = ihil.g
(33)
jŻ.. i = 0, !Ż.z L = 0, [
(34)
This means that we can simultaneously find the eigenstates of L. and L (also for Lo, Ly). To achieve this, let us first turn our view from the Cartesian coordinate systems to the spherical coordinate system, for convenience:
x = r sin 0 cos o
V
x
+
82
+ 22
= r sin θ sin Φ, 0 = cos -1
(35)
r cos 0 tan 1
By the definition L. 2. úP we have:
L. =
I
(36)
8y
0x
In fact, this is related
oø by the chain rule:
26
0$ Oz
04 By
0y
36
0=
Dy
y t
i.e.
L ih
Jth
θ
L ih
sin Φ
+ cot 0 cos
00
(37)
i
L, ih
cos
+ cot 0 sin
00
1
i
02
1 h22
sin 0
+
sin 0
a0
00
sin" 0 002
Since construct the simultaneous eigenstates of both i. and L the union of their eigenvalues m)) to index these simultaneous eigenstates as 11, m), and establish the following equations:
L, 11, m) mm/, m).
(38)
L U1 m) 1|4, m))
24

---

Published as a conference paper at ICLR 2024
where m, I E R because both L. and L are Hermitian operators. Let now solve the first eigenvalue equation using the coordinate system for ī,2 -ih 0 86. and denote [1, m) using its functional form Bø - 1/. From this equation,
determine the dependence of the solution and rewrite 3m (0, Φ) = e Pt.m where (t.m captures the still undetermined 0 dependence of the eigenstate vl.m. As a function of the angles, it required that fm uniquely defined: (0, 2 Vr, (0, Φ), from which we can obtain e 1. If we instead write this eigenvalue equation as Lg]4, m) = fum I. m), we thus have the fact that m E Z. Similarly, it can be shown that the possible values of are I E N and -1 S m S L, with the eigenvalue equations being written as,
Vt,m= -th Ovr.m Be mmhmm 8
L, m) = hmm]l, m)
(39)
1 m)) 111(1 1)]1,
After determining the eigenvalues of operators L, and L?2, the next step is to obtain the form of the eigenstates m). By using the functional form, we rewrite the second equation as:
a2 + sin 0 0 + 101 +
122
sin
1/ 1/ (0, Φ) = -0.
(40)
sin 00
00)
This is exactly the Laplacian equation in the spherical domain (Sommerfeld, 1949), and the solution is exactly the spherical harmonics 1 (0 Y!! (0, Φ) with -1 S m S 1,1 = 0, 1,2, ..... as introduced in Section A.4. the context quantum mechanics, it conventional to use the ket |1, m)
to denote the spherical harmonics yr(() and specify it in the spherical coordinate representation as (0, ø)(, m) Y!!! (0, Ø). In fact, the of kets {14, m3) S m (} exactly forms an orthonormal basis for the Hilbert space with 21- - dimension.
Addition of Angular Momentum Operators. In quantum mechanics, we often consider systems with two physically different angular momenta, c.g., the orbital angular momenta of two electrons. means that the angular momentum operators jointly act on two Hilbert spaces with different dimensions, which prompts to define the concept of the total angular momentum operator. Given two angular momentum operators j1 and f2 acting spaces V1 and V/2 with 211 + and 21g + 1 dimensions respectively, these two spaces are spanned by the eigenstates 141 m1). -41 S m I l and |/2, m2), -12 S m2 S l2. The total angular momentum operator is the addition of these two operators, and acts on the tensor product space Va3 = V1 @ 1/2 with (2l, + 1) x (212 + 1) dimension, which is spanned by |41,,,,4 12, m2) 1/1, m1) @ (12, m2), -4 S m1 S -12 S m2 S l2. We call 11, m1, 12, m2) the uncoupled tensor product basis. Given any angular momentum operator is defined to act on states in V/3 in the following manner:
( & 1) m1, 1,,,,,444 12, m2) 314, m1) & 12, m2). (i & 1) m1, 12, m2) = |4, mi) @ j|l2 m2).
(41)
where i denotes the identity operator. The total angular momentum operators are then defined as:
j= j i i j2
(42)
Similar to Eqn, (33), the total angular momentum operators also have the following commutative relationship:
[Ja, Jyl = ih.J (J,, = ih.jz, [J J.] = ihJj,
(43)
[J. o. 0. -,
(44)
Therefore, of coupled eigenstates also simultaneously exists for j2 and J., which we denote as 14, T., M). This has the following eigenvalue equations:
j_ l2; L, M) hM 12 M)
(45)
44, 12: M) = 122 + 1))4, 1,
where -L S M < L, 14 - 121 < L < 4 + 12, and the total number of eigenstates of the total angular momentum operator is exactly equal the dimension of the tensor product space V/::
25

---

Published as a conference paper at ICLR 2024
-, 1) = (21 +
1) x (212 1). Both the uncoupled tensor product basis and the
coupled eigenstates form basis for the tensor product space, i.c., they are related via a change of basis:
4 Σ
11,12 M)
Σ
11, (1,mi,4a 12, m2) m, m2)T, M):
(46)
11
11
where the expansion coefficients (1,m,, 12, m2 L4, M) are the Clebsch-Gordan coefficients denoted C.L,,))
as (1,mm))ls, which we previously introduced in Appendix A.5,
In summary, spherical harmonics correspond to the eigenstates of angular momentum operators, while the Clebsch-Gordas coefficients serve the expansion coefficients of coupled eigenstates of total angular momentum operators in an uncoupled tensor product basis. This indicates an underlying relationship between these concepts, which will further reveal by using the Wigner-Fckart theorem below.
A.7.3 TENSOR OPERATORS AND WIGNER-ECKART THEOREM
As we introduced in Section 2, representations of the rotation group SO(3) on R are parameterized by orthogonal transformation matrices E R3333 dei(R) = 1, which are applied to x E R$ via matrix multiplication Rx.. Each rotation matrix R. has the explicit form that is specified by the rotation axis and angle, c.g., rotation by angle & about the z-axis is given
Rotation Operator.
COS sin & cos 0
sin Φ
0
by R. (0)
0
We also use R(ñ, Φ) to denote the rotation g E SO(3) that
1
performs rotation by angle about the ń E R3, | direction. In quantum mechanies, given any rotation R(ñ, o), there exists a unitary operator Ú (ū7 Φ) that transforms ket from the unrotated coordinate system to the rotated one, i.e. (w) R = Üni, ,)w)) The previously introduced angular momentum operator exactly the generator of the rotation operator Ó:
i(L. 7
( の) exp(
(47)
Recall that given E N. the eigenstates of angular momentum operators |1, m), m I form a basis for the Hilbert space with 2/ + dimension. i.c., any ket (a) this space can be expressed as (a) = -m am L1, m). We can also calculate the matrix element of the rotation operator Ü(i, o) corresponding to g E SO(3) by using this basis:
-i(L- m)
D mmm (g) = (1, m² lexp
|1, m),
(48)
h
which is exactly the Wigner-D matrix in the shape of (27 + 1) x (21 + 1), introduced in Section 2. Hence, the rotated ket (a') can be expressed as:
la') (aa x)1a)
amil ml) ml () m)
Σ D m
(g)am mm').
(49)
me ,7
Therefore, the rotation transformation for kets in the 2/ + 1-dimensional space spanned by [1, m), -1 S m S / can also be computed through matrix multiplication:
I'm
ΣD (1) me m (g)0m.
(50)
Vector Operators and Tensor Operators. Classically, a (three-dimensiona)) vector is defined by its properties under rotation: the three components corresponding the Cartesian 2, y, a axes transform V; - Σラ R, A tensor is generalization of such vector to an object with more than one index, e.g., T.j or Tiyk (with 3 and 3 x 3 components respectively). and satisfies the condition that these components mix among themselves under rotation by each individual index following the vector rule Tiyk Σ R. R. RA,, Timn- Tensors written like this called
3
26

---

Published as a conference paper at ICLR 2024
Cartesian tensors, and the number of indexes is denoted as the rank of the Cartesian tensor, i.e., rank n Cartesian tensor has 3" components in total.
In quantum mechanics, an operator V called a vector operator the expected values of its three components in any ket transform like the components of classical 3D vector under rotation:
ÚiVO Σ RigVy
(51)
where U denotes the corresponding unitary rotation operator of the rotation matrix R. Given any ket (a) and its rotated ket la') via the rotation operator Ü, we can check the expected value of the vector operator V in the rotated system:
(o'lV/oo)) (alú - (.y((alVla)).
(52)
This exactly meets the requirement for the expected values its three components under rotation. Similarly, we further generalize Cartesian tensor operators, which satisfy the condition that the expected values of its components in any ket transform like the components of classical tensor under rotation.
Although a Cartesian tensor is easy to denote, it is usually hard to work with, especially when dealing with rotations. Instead, we further define the concept of spherical tensor operators. Given a basis ket [t, m) for the 21 1-dimensional Hilbert space, is transformed via a rotation Ü Φ) as:
Ünü ((..))).m))
= Ell. ml 1 d)) mes
Σ 1,,m') mmm (g), m
(53)
where E (O))) corresponds to the rotation Ün, o). From this equation, we can see that the rotation acts separately different 21 + 1-dimensional Hilbert spaces spanned by different of basis kets 11, m.mm I E N. By doing this, we can handle this rotation. Spherical tensor operators of rank are de- fined families T())
(I")) ...... T()) of 2/ + operators T!ll that satisfy the following constraint:
ÚtT
Σ
T:( m mmm (g).
(1)
(54)
m/=-1
Since the transformation rule of the spherical tensor operators originates from the properties of the eigenstates of angular momentum operator [1, m)) spherical harmonies Y!!!) are special class of spherical tensor operators.
Wigner-Eckart theorem. Recall that all information about an operator Å H + H is encoded by its matrix elements (eml Á Ále,) with the basis em. = I,.... From the definitions of vector operators and spherical harmonics operators, we can see their underlying symmetry constraints, indicating restricted freedom. Due to this, the matrix elements of such operators are fully specified by comparatively small number of elements, which are called reduced matrix elements in the context of quantum mechanics. This reduction for operators is exactly described by the Wigner-Eckart theorem. For clarity, we discuss this theorem in its most popular form, i.c., for spherical tensor operators.
Given spherical tensor operators of rank j. it is natural use the eigenstates of angular momentum operators [1. m) as the basis to calculate the matrix elements. By fixing j, 1, J, there are 2j + 1 components 77 of the spherical tensor operator, 2/ + 1 basis kets [1, m) and 2. + 1 basis bras (J,
Therefore, there are (2j + 1) x (21 + 1) x (2J + 11 matrix elements (J, M.MMI) (() n). The following Wigner-Eckart theorem reveals how these matrix elements are specified.
Theorem A.1 (Wigner-Ee Eekkrt Eckart theorem for Spherical Tensor Operators, from Jeevanjee (2011)). Let J. J Nzo and let TU) be sphericul tensor operator of rank j. Then there is unique complex number, the reduced matrix element X E C (often written (J (J| [[] E C), that completely determines any of the (2JJ + 1) x (2j + 1) x (2l + 1) matrix elements (J, M T() |t, n):
(J, M TU) I, X (J.M)
(55)
(l.)))m)))
27

---

Published as a conference paper at ICLR 2024
The Wigner-Eckart theorem describes how the matrix elements ((J, M ( [1, n²) of the spherical ten- operator (T()) in the basis of angular momentum eigenstates ((J, M and [1, n)) are decomposed into the Clebsch-Gordan coefficients (CLL)) ((.))..m)) and the reduced matrix elements ((][TT())))) [1]). We
can see that the (J]IT()) values do not depend on M, n. m, which are thus called reduced matrix elements. The Eckart theorem elegantly reveals the underlying connections between matrix elements, which is extremely useful for simplifying calculations in quantum mechanics.
Furthermore, if we instantiate TU) as the spherical harmonics y and apply the Wigner-Eckart theorem, we obtain the following relation (J, M y:k (1, = (J] (I.))i.m)) which can
be explicitly rewritten using the definition of inner products in the Hilbert functional space:
( J,a J," (
()l e m2 (0, y (0, vb) sin 0d0dg = 60 Cillm) (11.2) (1.m1)02.ma))
(56)
where 6( is real constant only determined by I l2 and I. From Eqn. (56), the (eessch-Gordat coefficients are mathematically related to the integrals of products of three spherical harmonics, which are known the Gaunt coefficients (Wigner, 2012), and are denoted as G.l.m) allows us obtain the relationship presented in Appendix A.6.
(2.
This
(l.mm))2mmm))"
B ADVANCES IN EQUIVARIANT NETWORKS FOR THE EUCLIDEAN GROUP
One standard way to enforce equivariant operations is through group convolutions, which define convolution filters as functions on groups (Cohen & Welling, 2016; 2017: Coors et al., 2018). However, the integral in the group convolution becomes intractable when dealing with the continuous Euclidean groups motivating the development of die-algebra-basee operations (Finzi et al., 2020; Hutchinson et al., 2021). Finzi et al. (2020) proposed LieConv, which determines the group convolution via lifting (mapping the input in vector space to group element), and discretization of the convolution integral via the PointConv trick. Following a similar idea, Hutchinson al. (2021) developed Lierransformer, which employs the self-attention strategy to dynamically re-weight the convolutional kernel. Nevertheless, such methods still induce high computational costs due to discretization and sampling, and the non-linearity of equivariant features is restricted.
Another line of work leverages vector operations involving scalar (R) and vector (R3) features to keep equivariance (Satorras et al., 2021; Jing et al., 2021; Schüit et al., 2021; Deng et al., 2021; Haghighatlari et al., 2022: Thölke Fabritiis, 2022; Wang & Chodera, 2023; Chen et al., 2023). EGNN (Satorras et al., 2021) uses the scalar-vector product to update the vector features under the equivariance constraints. GVP (Jing et al., 2021), PaiNN (Schütt et al., 2021) and TorchMD-Net (Thölke & Fabritiis, 2022) further leverage vector products to enhance the learning of scalar features. Chen et al. (2023) recently proposed a unified framework called GeoMFormer, to simultaneously learn effective scalar and vector features. The empirical performance and applicability of these models are limited, due to bounded model capacity (Joshi et al., 2023).
Recently, tensor products of irreducible representations (irreps) have become the dominant equivariant operations for the Fuclidean group (Thomas et al., 2018; Drautz, 2019; Fuchs et al., 2020; Unke et al., 2021; Brandstetter et al., 2022; Batatia et al., 2022a;b: Nigam et al., 2022; Batzner al., 2022; Frank et al., 2022; Liao & Smidt, 2023; Musaelian et al., 2023; Passaro & Zitnick, 2023; Yu et al., 2023b; Gong al., 2023; Liao al., 2024). With solid mathematical foundation from representation theory, such operations allow modeling equivariant mappings between any 2/+ 1-dimmnnsioonll irreps space via Clebsch-Gordan coefficients (Wigner, 2012), and are proven to have universal approximation capabil- ity (Dym & Maron, 2021). State-of-the the-art performance has been achieved in various real-world appli- cations via these operations (Ramakrishnan et al., 2014; Chmiela al., 2017; Chanussot et al., 2021).
Various works use pair-wise tensor products to develop equivariant neural networks. Tensor Field Network (TFN) (Thomas et al., 2018) uses tensor products of irreps to build an S SE(3) -equiaatqi convolution layer to handle data in the three-dimennsonnal Euclidean space. Following TFN, Fuchs et al. (2020) developed the 3E(3)-Transformer enhance the feature learning. NequlP (Batzner al., 2022) further extended the tensor product operations to be E33 (quuiarruv,,,, and achieved strong performance on learning interatomic potentials. Allegro (Musaelian al., 2023) further developed strictly local equivariant network using the
which incorporates the attention mechanism to
28

---

Published as a conference paper at ICLR 2024
iterated tensor products of irreps. SEGNN (Brandstetter et al., 2022) generalizes the node and edge features to include vectors or tensors, and incorporates geometric and physical information by using the tensor product operations. So3krates (Frank et al., 2022) developed a modified attention mechanism for modeling the non-local effects in molecular modeling. In particular, the atomic interaction layer in So3krates computes the full tensor products of features containing irreps of different degrees. Equiformer (Liao & Smidt, 2023) uses the tensor product operations to build scalable equivariant Transformer architecture, testing this on the baselines of the large-scale OC20 dataset (Chanussot et al.. 2021). This extended Graph Transformers ( (Ying et al., 2021a;;; Shi et al., 2022; Luo et al., 2023; Zhang et al., 2023a) to equivariant networks for the Euclidean group. More recently, eSCN (Passaro & Zitnick, 2023) presented an interesting observation on the sparsity of spherical harmonics filters and developed an equivariant network by using the SO(2) convolution instead, to scale the degree of irreps up. Following eSCN, EquiformerV2 (Liao et al., 2024) combined the Equiformer backbone with the eSCN convolution for scaling to higher-degree representations, which has sate-of-the-art performance on the OC20 dataset.
Existing works have also developed operations involving multiple tensor products among equivariant features (Equivariant Many-body Interaction). Atomic Cluster Expansion (ACE) (Drautz, 2019; Dus- son et al., 2022) developed systematic framework for constructing many-body interaction features encode the atomic environment, which involves multiple tensor products among atomic basis functions. MACE (Batatia et al., 2022b) combines the many-body interaction features with message- passing neural networks, which shows strong performance for learning machine learning force fields (Kovács et al., 2023). Nigam et al. (2022); Batatia et al. (2022a) provided a thorough analysis of the theoretical properties and design space of Equivariant Many-body Interaction operations.
Equivariant neural networks based on tensor products of irreps have been widely used in modeling 3D data across variety of real-world applications, e.g., force field modeling (Chmiela al.. 2017: Chanussot et al., 2021), physical dynamics simulation (Brandstetter et al., 2022), point cloud learning (Wu al., 2015; Uy et al., 2019), and Quantum Hamiltonian Prediction (Unke al., 2021; Li et al., 2022; Khrabrov et al., 2022; Yu et al., 2023b;a). PhiSNet developed an SE(3)- cquivariant message passing network by using the tensor product of irreps. In particular, to predict the Hamiltonian matrix elements, operations in the Equivariant Feature Interaction class used for coupling atom features themselves (for diagonal block prediction), or features of atom pairs (for non-diagonal block prediction), which models interaction of irreps from different dimensional space and objects. Similarly, DeepH-E(3) (Li et al., 2022) further uses the tensor product operations to incorporate the spin-orbit coupling effects. QHNet (Yu et al., 2023b) simplifies the PhiSNet architecture by controlling the number of tensor product operations.
Finally, for comprehensive surveys equivariant networks for the Euclidean group and their applications, we refer the interested readers to Bronstein et al. (2021); Han et al. (2022); Zhang et al. (2023b); Duval et al. (2023).
29

---

Published as a conference paper at ICLR 2024
C IMPLEMENTATION DETAILS
In this section, we provide more details on designing efficient counterparts for operations in the Equivariant Feature Interaction, Equivariant Convolution, and Equivariant Many-body Interaction categories.
Equivariant Feature Interactions. Given two features ( 5 containing irreps of up to degree L1 and Le respectively, the operations in this class are commonly in the following form:
(x(4)) 2% y(l)))() Et, -tz- w/ (x(1) Beg yl2))(0), where w 1 IR are the learnable weights for each (4, 12) + combination. By reparameterizing w as wi, Wig wy where Why W12 wi E I,, Eqn. (4) in the main body becomes:
41 La
(x(11)
Gains x222))00
W (x) @Gannt x (2))()) mm
=
1==0
Wi, Mi, (x (1) @Gaunt x (1)))
12=0
(l)) (11) (1a)
Wi, Wl,
(1)m))()).m2)
2m1 Fing
l m0 12
my
mig
L1
4
25
(12)
vill)) yulhi) (0, () m2 (0, v 1 (0, v)) sin 0d0d
Wi Mig W
T mi
1==0
mu2
L2
Wh, 2m (1))))))( (0, V)
Σ 1=
Wiy ma y(al)) (0, Y (0, v- sin Adedy
wt
11=0
(57)
Let (11)
wi, (11) and x (12) w1, w1,x(la), then Eqn. (57) actually calculates the coefficients
((h)):(()) of function (0, v) (0, 1)) (0, v²), where (0, +)= Eh, Im, mi (0, v))
((0))y(())) and (0, -) I my - Timz m² (0, w)). To this end, we can use the efficient approach presented in Section 3.2 accelerate such computation. After calculate the coefficients of F3 (0, )) we then multiply the weight wy with corresponding irreps. Thus, the setting of weighted combinations are naturally included in our Gaunt Tensor Product.
Equivariant Convolutions. One special case of Equivariant Feature Interaction is to instantiate
1,2 (12)) as spherical harmonics filters, i.e., Σ/i Eiz Mi,,, (x;
(11 Beg y(la) ( (1==) (|F=-B] (),), where
ry E R3 are positions of objects, (f)) 5 + 2| is the type-l2 spherical harmonic, W/ R are learnable weights calculated based the relative distance and types of objects j,
e.g., M,,,, = FI, 12 (1)r ryll, (1,, .ry), where Fr, are instantiated as Gaussian basis functions and IJ are the type of object j. Similarly, h4, .1, can be reparameterized for our approach.
Moreover, Passaro & Zimick (2023) provides an interesting observation that we select a rotation E SO(3) with (g) |r. |. = (0, 1,0), then ( (D ) ε <!)) i.e., being non-zero only if m O. Such sparsity further propagates to its coefficients of the 2D Fourier basis, which brings further acceleration Equ. (6): Let ym denote the elements of the spherical harmonic filter.
Due to the sparsity of (!) dil) and coefficients vl.m. the conversion rule from spherical harmonics 2D Fourier bases can be rewritten as:
Vu.o
Σ 1=0
(0) 1.0.a yo Va.d
-L L,
(58)
Vi.
if # 0.
The coefficients of the spherical harmonic filters in the 2D Fourier basis are also sparse, and such a conversion has an O(L") complexity instead.
Equivariant Many-body Interactions. Operations in this class perform multiple tensor products among equivariant features, i.e., x(41) @ Beg... x with tensor products. For exam-
30

---

Published as a conference paper at ICLR 2024
ple, given equivariant features A E (I( of atom i, MACE (Batatia et al., 2022b) constructs (1))
the many-body features by En, Σim C IT:=1 i,mg where A is the mę-th element of the type-Ig irrep of A! A!4), Im are the multi-indices (Iimi,..,lomm C. are the generalized Clebsch- Gordan coefficients, i.e., CLiM lima....hum, = Ma Lo
1am, ClaMa Im
Ma.Ma, Lx- -i ddym with L (I2,..., Lx); 11 - 12| S L2 S l + 12, Vi 2 3 L-1 -41 4/ L S Li-1 1 and M, E {m, - 4 S my I 4.), fu enumerates all possible couplings of l,...,y. by. In particular, it is the explicit form of performing tensor products with a!") itself times via the generalized (lebsch-Gordian coefficients:
C.. My
Bud A(4) Bag Seg (44))
(59)
We generalize the Gaunt coefficients to this setting. (1) M (0, v)) sin Od0deb.
Let (
(1.mm)...(l.m))
(tl)) generalized Gaunt coefficients is equivalent to:
2T lo Jo (lll (0, Y (0, Y)...
(..) (0, )
Then Eqn. (59) in the
(A!")
Gaunt ... S Goient y
12 I times
2
L' Σ 11
L'
In
Σ
A
yoll (0,
Σ
villn) (0, ()))))n sin 0d0du.
In DD
(60)
In fact, we actually compute multiple multiplications among multiple spherical functions. We can also use the parameterization tricks to similarly multiply each A with learnable weights wy to achieve weighted (Im) + combinations. The approach presented in Section 3.2 can be naturally used: after the conversion from spherical harmonics to 2D Fourier bases, we use A, V 1,... n to
denote the coefficients of 2D Fourier bases of these n spherical functions, and the 2D Convolution can be compactly ex expressed as:
A 2+
A (n-1)) 1]+ @ A
(61)
where @ denotes the convolution. Interestingly, these 2D Convolutions have associativity, and use a diviide-and-conquer approach for parallelizing the computation of these convolution operations. For example, if perform the tensor product operation among 8 functions, the coefficients A!"" can be divided into pair-wise groups for parallel computation, i.c., (1, 2) - 1x, (3, 4) - 2*, (5, + 3+, (7, 8) + 4+, Such a procedure iteratively conducted until the final results are calculated. The computational complexity of performing 2D convolutions in the Equivariant Many-body Interaction setting is (n I log L),, while recursively computing Eqn. (61) has an O 3 I log L) complexity.
Discussions. In practice, it is common in models to maintain C' channels of equivariant features, i.e., x!4) E [(++))=xC denotes the equivariant features of object which contains C' channels of irreps { 0 < << C' of degrees up to L. Our approach can be naturally extended by simply defining the combination rules of features in different channels:
Channel-wise: given X (L) x(4) E IR(++))xC, output features the e-th channel are (L) calculated based on x
x and (L) ji
only (O(C) full tensor products):
W
x(1)).
Guunt X
(62)
Gaunt
Channel-mixing: given x E ( output features in the e-th channel (L)
are calculated based weighted combinations of X C 0, 1,., C and x (L)) 0, 1,.,, C' (O()?) full tensor products);
x C2
{x(1))
(L) Wain
Gount 8 =
Σ
wee es
(L)
@N
(D) Grint j.02
(63)
ej mD D.
31

---

Published as a conference paper at ICLR 2024
Our approach is also not limited to full tensor products only. Given x(()) &Gant x(12) only, the approach presented in Section 3.2 can still be used:
(x Gaant
Cllm) (11)) ((2.m) am mm2
(4) (12)
("
(1) U2)
you (0,, (O) (0, sin Ad0de,
Dmny Tiz
2m
L (0)
daly (0, 1p) m2 m2
y (0, si Ad0dd..
Fm
2
(64)
where we set aml Im, to for in + 41,4 + l2. Our Gaunt Tensor Product between x(h) and (Ia) can still be equivalently viewed as multiplication between two spherical functions, which can be efficiently calculated by using the approach presented Section 3.2.
(Bv) (t.)
D PROOF OF THE EQUIVARIANCE OF THE GAUNT TENSOR PRODUCT
In this section, we provide formal proof of the equivariance of our Gaunt Tensor Product. As stated in Section 2, tensor products of irreps via the group. Given x(f1) and mndy((+) yf(a), have the following equation for any O()):
coefficients equivariant to the O(3)
(x)((()) Bog ( (
Ocg (D"a (1)x)))) (
(65)
Let D ( (g) denote the (m, n)-th element of the (() matrix. First, the m-th element of D (x)(( Dog (ll))(t)) can explicitly written via the Alesshh-Gordan coefficients:
D (g)(x) Dry (2))
Σ ml
D (1)
mon (x)(x Beg (
Σ m
( mm (g)
Σ
(1i.m1).(2amm)) Bny Ming
(66)
Σ
(1)
(g)
I (1)
(1) (12)
nam
(1amm))((amm)))
m2
Second, the m-th element of ((D (g)x (h)) Beg (D (g)(())))))) can also be explicitly written via the Clebsch-Gordan coefficients:
(()())))(()
(ID (1) (g)("1) @.s (D (g)x
Σ
Oflmm
D (n)
mims (g)
Σ
E (1)) mgms (1)uit
my mm
(l))
Gy (2 mmme (g)
11
(12)
(o) u wn1
mi.ing my ml
Σ
(lla) (1mmp)dlmmm)
D my
(o)
(2) m mg
(g)
m miinn
(67)
We can conclude that the Alebsch-Gordan coefficients have the following properties:
(llmm D ( ( (1) m.)...mm)))))
Σ
(llam) D ((1)) (lmm)) mn m. mlmma (0)-
Dg) (12)
(68)
mi.m/
The Gaunt coefficients also satisfy these constraints. Here, we use the Cartesian coordi- nates the input of spherical harmonics and denote the Gaunt coefficients (l.m)
(1.m)) as
32

---

Published as a conference paper at ICLR 2024
Jgz Ymn (F)Y (F) (f)di. We have the following equations:
D mm (1.m1))(2amaa)
mmm
(g)
{ yot (F) (F)) (F)dF
ma
{ ga
(
Y
E mi (o) (F) df
(69)
ma
Ym (F) m2 (F)) (D) (g)
82
E Gilmm
D (4)
(g)
(aa
mmmma (g)
(f1 (2)
m/m:
mi my
D D(l)
(g)
(9)
y (F)y ( (F) (F)dF
mime
mlmm
mi m/mmg
52
(il
(g)
Σ
D mlgma
(g) (f) y (F)dr
82
mi
mi my
Σ D
(11)
(g) (+)
D (12) mmm3
(0)
rills) (F)
y (P)dr
(70)
82
mimm
D
(11)
(g) -1 "ylll"' (F)
Σ
D
(12)
(g)
y-llaa (F)
y: (F)dF
mimm
m2mm
Del (1)
(g
" (F)
Σ
D
(12)
(g
(F)dF
mimm
mgmm
82
Y (11) mi
(D
12)
( ( ))
(
82
Due to the symmetry of the unit sphere S we have Sga Y!ll (F)Y (F) (g)) = Ige Y!lli) (g=)) (D (F)) (f)di. Therefore, the Gaunt coefficients also satisfy the properties to make our Gaunt Tensor Product equivariant the O(3) group.
33

---

Published as a conference paper at ICLR 2024
E EXPERIMENTAL DETAILS
E EFFICIENCY COMPARISONS
Settings. (1) For the Equivariant Feature Interaction operation, we use the e3nn implementa- tion (Geiger & Smidt, 2022) as the baseline. We randomly sample 10 pairs of features containing irreps of up to degree L. with 128 channels. We evaluate different degrees of and report the average inference time". (2) For the Equivariant Convolution operation, we use the eSCN imple- mentation (Passaro & Zitnick, 2023) as the baseline, due to its enhanced efficiency over e3nn. In this setting, 10 pairs of features and spherical harmonics filters are sampled instead. The number of channels is kept as 128. We also evaluate different degrees of L. We implement our Gaunt Tensor Product by using the insights of eSCN for further acceleration, as stated in Section 3.3. (3) For the Equivariant Many-body Interaction operation, we compare our approach with both the e3nn and the MACE implementations (Batatia et al., 2022b). We evaluate different degrees of L and the number of tensor product operands v: (a) fix V- = 3 and vary L, (b) fix L, = and vary v. The benchmarking results are obtained by using the same NVIDIA Tesla V100 GPU.
E.2 SANITY CHECK
Baselines. In this experiment, we choose the SEGNN (Brandstetter et al., 2022) as the baseline, which achieves state-of-ther of-tth-art performance on the N-body simulation task. In particular, SEGNN proposes generalization of equivariant GNNs such that node and edge attributes are not restricted to scalars. Based on tensor products of irreps, SEGNN injects geometric and physical quantities into node updates. To investigate the underlying effects of different parameterizations between our Gaunt Tensor Product and the flebsch-Gordal Tensor Product, we compare two architectures: (1) SEGNN with Clebsch-Gordan Tensor Product implementation; (2) SEGNN with Gaunt Tensor Product. For other configurations, we follow Brandstetter et al. (2022) to keep the same.
Settings. Following Brandstetter et al. (2022), we use the N-body simulation task to investigate the effects of different parameterizations. The N- N-body particle system consists of 5 particles that carry positive or negative charge, having initial position and velocity in a 3-dimensional space. The task is to estimate all particle positions after 1,000 timesteps. We use the 4-layer SEGNN and set the max degree of attributes (l,,) irreps in [1, 2]. The max degree of feature vectors (1y) is set to 1. The learning rate is set to 0.01 for la = 1 and 0.02 for L, 2. The learning rate was reduced using an on-plateau scheduler based on the validation loss with patience of 10 and a decay factor of 0. for La = and 0.2 for la- = 2. Other settings are kept the same the original model. The model is trained on 1 NVIDIA Tesla V100 GPU.
E.3 OC20 S2EF
Datasets. OC20 is a diverse and large-scale dataset consisting of 1.2M DFT relaxation trajectories, which are computed by using the revised Perdew- Beww-BBrrkeeeee Ernzeef (RPBE) functional (Hammer et al., 1999). Each structure in OC20 has an adsorbate molecule placed on catalyst surface, and the core task is Structure-to- centture-toEEneree -eeetty-------eeeee -Forces (S2EF), which is to predict the energy and force mean absolute error (MAB). The "All" split of OC20 contains 134M training examples. Due to limited computational resources, we adopt the widely used S2EF-2M subset, which consists of 2 million training examples in total. Four splits of validation sets are used including in-distribution (ID), adsorbates (OOD-Ads adsorbates and catalysts -oth-B The evaluation metrics include the Energy MAE and Force MAE, the cosine similarity of predicted forces and ground-truth forces (Force Cos), and the ratio of examples on which the predicted energy and forces are under a predefined threshold (EFwT).
on out-of--iittibution on catalysts (OOD-Cat) (atD. and out-of-distribution on both
Baselines. We compare our approach based on the EquiformerV2 (Liao et al., 2024) with several competitive baselines from the leaderboard of the Open Catalyst Challenge (Chanussot et al., 2021). First, we compare with several classic invariant neural networks. SchNet (Schütt al., 2018) lever- aged the interatomic distances encoded via radial basis functions, which serve as the weights of
WWe noticed that after running the same e3nn code multiple times, the subsequent executions are faster than the initial attempt, because of the code compilation. Nevertheless, our approach remains more efficient overall.
34

---

Published as a conference paper at ICLR 2024
continuous-filted convolutional layers. DimeNet++(Gasteigee al., 2020a) built upon the DimeNet model (Gasteiger et al., 2020b) which introduced the directional message passing that encodes both distance and angular information between triplets of atoms. SpinConv (Shuaibi et al., 2021) used per-edge local coordinate frame and a spin convolution over the remaining degree of freedom to build the model. GemNet (Gasteiger al., 2021) embedded all atom pairs within a given cutoff distance based on interatomic directions, and proposed three forms of interaction to update the directional em- beddings: Two-hop geometric message passing (Q-MP), one-hop geometrie message passing (T-MP), and atom self-interactions. An efficient variant named GemNet- T is proposed to use cheaper forms of interaction. Based on GemNet, Gasteiger al. (2022) comprehensively investigated the relationships between the performance and chemical diversity, system size, dataset size and domain shift. From these results, an efficient improvement to GemNet architectures was developed, named GemNet-OC.
Recently, SCN (Zitnick et al., 2022) used rotation trick to encode the atomic environment by using the irreps, while not using the tensor product operations. Note that SCN does not strictly obey the equivariance constraint. cSCN (Passaro & Zitnick, 2023) presented an interesting observation that the spherical harmonics filters have specific sparsity patterns if a proper rotation acts on the input spherical coordinates. Combined with patterns of the Clebsch-Gordan coefficients, eSCN proposed an efficient equivariant convolution to reduce SO(3) tensor products to SO(2) linear operations, achieving efficient computation. Built upon the eSCN convolution, EquiformerV2 (Liao al., 2024) developed a scalable Transformer-basee model, which achieved state-of-the-art performance on the OC20 and AdsorbML (Lan al., 2022) benchmarks.
Architectures. EquiformerV2 V2 is composed of stacked equivariant Transformer blocks. Each block consists of two layers: an Equivariant Graph Attention layer followed by feed-forward layer, with both layers having separable equivariant Layer Normalization and skip connections. EquiformerV2 can use irreps of higher degrees, e.g., 4 or 6, with the eSCN convolution. However, such implemen- tations also restrict the available equivariant operations for this scalable architecture. In this work, we proposed the Gaunt Tensor Product, which serves as a new method for efficient equivariant operations, enabling further combinations between architectures like EquiformerV2 and effective equivariant operations. As an example, we use our Gaunt Tensor Product to implement the Selfmix layer in the Equivariant Feature Interaction operation class. This layer receives equivariant features of each atom as input. For each atom, full tensor products of its irreps are performed with weighted combinations. In each EquiformerV2 block, we insert this layer with the equivariant Layer Normalization in between the Equivariant Graph Attention layer and the feed-forware layer. For this modified architecture, the original implementation would be extremely slow because to the limitation of the eSCN implemen- tation. In contrast, our Gaunt Tensor Product has high efficiency when using irreps of higher degrees.
Settings. Following Liao et al. (2024), we use 12-layer EquiformerV2 model. The max degree of irreps is set in [4. 6]. The dimension of hidden layers and feed-forward layers is set to 128. The number of attention heads is set to 8. The number of radial Basis kernels set to 600. We use AdamW as the optimizer (Kingma& & Ba, 2014). The gradient clip norm is set to 100.0. The peak learning rate is set to 2e-4. The batch size set to 64. The dropout ratios set to 0.1. The ratio of stochastic depth (Huang et al., 2016) is set to 0.001. The weight decay is set to 0.0. The model is trained for 12 epochs with a O. l-pooch warm-up stage. After the warm-up stage, the learning rate decays by using a cosine learning rate scheduler. The batch size is set to 64. Other hyperparameters are kept the same as EquiformerV2 for a fair comparison. The model is trained on 16 NVIDIA Tesla V100 GPUs.
E.4 3BPA
Following Batatia et al. (2022b), we use the 3BPA dataset to benchmark the Equivariant Many-body Interaction operations. The 3BPA dataset contains geometric structures of flexible drug- like organic molecule 3-(benzylox)ppyridng 2-amine sampled from different temperature molecular dynamics trajectories. The training set of the 3BPA dataset contains 500 geometries sampled the 300 K temperature, while the three test sets contain geometries sampled at 300K, 600 K, and 1200 K to assess in- and out-of-domain accuracy. A fourth test set includes optimized geometries, in which two dihedral angles of the molecule are fixed, while a third angle is varied between 0 and 360 degrees. This variation produces dihedral slices that traverse regions of the potential energy surface (PES), different from the training data. This test assesses the smoothness and accuracy of the potential energy surface (PES) region that dictates the conformers present in simulation. Consequently, it directly influences properties of interest, such the binding free energies associated with protein targets.
Datasets.
35

---

Published as a conference paper at ICLR 2024
Baselines. We follow (Batatia al., 2022b; Musaclian al., 2023) to choose several competitive baselines for comparison, Atomic Cluster Expansion (Drautz, 2019) provided a systematic framework for constructing high body order complete polynomial basis functions, which includes many existing atomic environment representations as special cases (Behler Parrinello, 2007; Bartók et al., 2010; Thompson al., 2015; Shapeev, 2016). sGDML. (Chmiela al., 2019) is optimized implementation of the symmetric gradient domain machine learning model, which is able to faithfully reproduce global potential energy surfaces (PES) for molecules with a few dozen atoms from limited number of user-prov ussrrprovided reference molecular conformations and the associated atomic forces. NequlP (Batzner al., 2022) extended the tensor product operations be E equiivarra Allegro (Musaelian et al., 2023) further developed strictly local equivariant network to improve scalability. Batatia et al. (2022a) provided an analysis on the design space of E(3) -33)-aqnua operations and developed BoTNee. MACE (Batatia et al., 2022b) combined the message passing neural networks with Equivariant Many-body Interaction operations, which achieved better data efficiency and improved generalization performance.
Settings. Based on the MACE architecture (Batatia et al., 2022b), we compare our Gaunt Tensor Product with the implementations of eănn and MACE. Following Batatia et al. (2022b), we two-layer MACE model. The dimension of hidden layers is set to 256. Radial features are generated using 8 Bessel basis functions and a polynomial envelope for the cutoff with p 5. The radial features are fed to an MLP of size [64, 64, 64, 1024]. using SiLU nonlinearities on the outputs of the hidden layers. The readout function of the first layer is implemented as a simple linear transformation. The readout function of the second layer is a single-layer MLP with 16 hidden dimensions. We use the AMSGrad variant of Adam as the optimizer, and set the hyper-parameter to 1e-8 and (31,32) to (0.9,0.999). The batch size is set to 5. The peak learning rate is set to 1e-2. The learning rate was reduced using an on- on-plateau scheduler based on the validation loss with patience of 50 and a decay factor of 0.8. We use exponential moving average with weight of 0.99 to evaluate the validation set as well as for the final model. We set the exponential weight decay to 5e-7. The model is trained on 1 NVIDIA A6000 GPU.
36


Reviews:
[{"review_id": "PjyEMwYqfT", "reviewer": "Reviewer_wTba", "summary": "The authors propose a systematic approach to accelerate the computation of tensor products of irreps by connecting Clebsch-Gordan coefficients to Gaunt coefficients. Based on this, the authors introduce a change of basis from spherical harmonics to a 2D Fourier basis. This transformation enables efficient computation via the convolution theorem and Fast Fourier Transforms (FFT), reducing the complexity from $O(L^6)$ to $O(L^3)$.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well-structured, featuring clear and logical derivations. The authors deserve commendation for their efforts in showing the equivariance of tensor products with Gaunt coefficients (Appendix D).\n- Interpreting the spherical tensor product as a multiplication of spherical functions and subsequently transforming it into a 2D convolution provides valuable insight into the underlying operations. In my opinion, this perspective is more important than the subsequent application of FFT, which is a logical extension.\n- The substantial reduction in time and memory cost on the 3BPA dataset serves as compelling evidence of the method's effectiveness.", "weaknesses": "Majors:\n- The current derivation and application are closely tied to $\\text{O}(3)$. Is it feasible to extend this convolutional perspective beyond $\\text{O}(3)$, e.g., to other Lie groups like $\\text{SU}(2)$? It would be beneficial to see discussions on the aspect of its generalization.\n- I would like to see the inference time of EquiformerV2 model with and without the proposed Gaunt tensor products.\n\nMinors: \n- Please adjust Equation (23)(26)(36) as they are written in the margins of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "cf5HeKURbc", "reviewer": "Reviewer_AAYb", "summary": "This work introduces a new framework for calculating the equivariant operations efficiently. The key idea is to convert the Clebasch-Gordan coefficients to Gaunt coefficients, by which the tensor product of irreps can be calculated efficiently in the Fourier domain.", "questions": "1. Could you give me more explanation about the *braket* notations used in Theorem 3.1? It seems not clearly introduced in the paper.\n2. Could you explain more or give the intuition about the irreducible representation?\n3. How to determine in practice the truncation, ie. the degree L, in the 2D Fourier expansion (above Eq. (5))?\n4. Is it possible to choose other basis systems instead of Fourier? May I know how choosing different systems affects the result of this paper?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written, and the main idea is explained clearly.\n2. It is smart and interesting that the basis transformation is applied to the Fourier domain. The usefulness of the proposed method for improving computational efficiency is verified with numerical experiments.", "weaknesses": "It would be better if more discussion could be given about the added value of this work. It will help the readers who are not familiar with this direction (like me) to easily understand the contribution and main differences compared with the existing works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "hShHs4UCLe", "reviewer": "Reviewer_1Msw", "summary": "This paper introduces the Gaunt Tensor Product, an efficient method to compute the tensor products of (high-degree) irreps as compared to the widely-used Clebsch-Gordan (CG) product. This is established through the connection between CG coefficients and Gaunt coefficients via the Wigner-Eckart theorem, and this connection implies that tensor product can be equivalently represented as multiplication of spherical functions. Through convolution theorem and FFT, the computational cost has been reduced from $O(L^6)$ to $O(L^3)$, where $L$ is the highest-degree of the irreps.", "questions": "Please refer to the weakness section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-structured and articulate. The authors have included most of the requisite background material to the appendix, thereby maintaining the flow of the main text.\n\n2. The paper's approach to enhancing computational efficiency, specifically through establishing a connection between CG coefficients, Gaunt coefficients, and the FFT, is both innovative and compelling.\n\n3. The experimental evaluations comparing the proposed tensor product with the conventionally employed CG product are both extensive and convincing.", "weaknesses": "1. One point of critique is related to the treatment of Theorem 3.1 in the manuscript. The text uses a range of specialized terms—such as \"spherical tensor operator,\" \"reduced matrix element,\" and \"total angular momentum\"—and introduces various notations that may not be readily accessible to readers unfamiliar with quantum mechanics. To enhance clarity and comprehension, I recommend that the authors elaborate on these concepts, particularly in their exposition of the Wigner-Eckart theorem.\n\n2. Another inquiry concerns the necessity of using high-degree irreps. In most E3-NN architectures, irreps with degrees no greater than L=2 are commonly employed. Consequently, I question whether the computational overhead incurred through the FFT in implementing the proposed tensor product is actually more costly than directly using the CG product for low-degree irreps.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (spotlight)