Equivariant Graph Network Approximations of High-Degree Polynomials for Force Field Prediction
Authors: Zhao Xu
Haiyang Yu
Montgomery Bohde
Shuiwang Ji
Venue: Accepted by TMLR ()

Abstract: Recent advancements in equivariant deep models have shown promise in accurately predicting atomic potentials and force fields in molecular dynamics simulations. Using spherical harmonics (SH) and tensor products (TP), these equivariant networks gain enhanced physical understanding, like symmetries and many-body interactions. Beyond encoding physical insights, SH and TP are also crucial to represent equivariant polynomial functions. In this work, we analyze the equivariant polynomial functions for the equivariant architecture, and introduce a novel equivariant network, named PACE. The proposed PACE utilizes edge booster and the Atomic Cluster Expansion (ACE) technique to approximate a greater number of $SE(3) \times S_n$ equivariant polynomial functions with enhanced degrees. As experimented in commonly used benchmarks, PACE demonstrates state-of-the-art performance in predicting atomic energy and force fields, with robust generalization capability across various geometric distributions under molecular dynamics (MD) across different temperature conditions. Our code is publicly available as part of the AIRS library \url{https://github.com/divelab/AIRS/}.

Published in Transactions on Machine Learning Research (08/2024)
Equivariant Graph Network Approximations of High-Degree Polynomials for Force Field Prediction
Zhao Xu"
zhaom Catomutttaa...om edu
Deportment of Computer Science & Engineering Teaas ABM University
Haiyang Yu"
haiyang Stama edu
Department of Computer Science & Engineering Texas ABM University
Montgomery Bohde
mbohde2015@tamme edu
Department of Computer Science & Engineering Texas ABM University
Shuiwang Ji
sji Sjt edu
Deportment of Computer Science & Engineering Texas A6M Universily
Reviewed on OpenReview: ht tp o/ en revi ew .ne /f or um 36 d= 7D AF up OV ne
Abstract
Recent advancements in equivariant deep models have shown promisc in accurately predicting atomic potentials and force fields in molecular dynamics simulations. Using spherical harmonics (SH) and tensor products (TP), these equivariant networks gain enhanced physical understanding, like symmetries and many-body interactions. Beyond encoding physical insights, SH and TP are also crucial to represent equivariant polynomial functions. In this work, we analyze the equivariant polynomial functions for the equivariant, architecture, and introduce novel equivariant network, named PACE. The proposed PACE utilizes edge booster and the Atomic Cluster Expansion (ACE) technique to approximate a greater number of SE(3) x S, equivariant polynomial functions with enhanced degrees. As experimented in commonly used benchmarks, PACE demonstrates state-of-the-art performance in predicting atomic energy and force fields, with robust generalization capability across various geometric distributions under molecular dynamics (MD) across different temperature conditions. Our code is publicly available as part of the AIRS library /ts
1 Introduction
Deep learning has led to notable progress in computational quantum chemistry lasks, such as predicting atomic potentials and force fields in molecular systems (Zhang et al.) 2023). Fast and accurate prediction of energy and force is desired, as plays crucial roles advanced applications such as material design and drug discovery. However, is insufficient rely solely on learning from data, as there are physics challenges that must be taken into consideration. For example, to better consider symmetries inherent in 3D molecular structures, equivariant graph neural networks (GNNs) have been developed in recent years. By using equivariant features and equivariant operations, -3u3v-equu GNNs ensure cquivariance of permutation, translation and rotation. Thus, their internal features and predictions transform accordingly as the molecule is rotated or translated. Existing equivariant GNNs can specialize in handling features with
either rotation order (Schütt et al. 2021 Jing et al. 2021
Satorras et al. 2021 Du et al. 2022 2023
"Equal contribution.

---

Published in Transactions on Machine Learning Research (08/2024)
Thölke & Fabritiis 2022) or higher rotation order { >>> 1 Thomas et al.
2018
Fuchs et al.
2020
Liao &
Smid
2023
Batzner et al., 2022 Batatia et. al.) 2022a b Unke et al. 2021a
Yu et al.
2023b
Li et al.
2022 while invariant methods only consider rotation order 8 = 0 (Schütt et al. 2017;
Smith et al. 2017
Chmicla cl ct al.) 2017;
Zhang et al. 20182 b Schütt ct al.)
2018 Ying et al.] 2021; Luo et al. 2023 Gasteiger
ett al. 2020 Liu et al.) 2022; Gasteiger el al. 2021). Generally, methods with higher rotation order exhibit
improved performance but at the cost of higher computational complexity. In addition to rotation order, some
equivariant methods (Musaelian al.] 2023; Batatia et al.) 2022a b) also consider many-body interactions in their model design. These approaches follow traditional principles (Brown et al., 2004 Braams & Bowman 2009] of decompasing the potential energy surface (PES) as linear combination of body-ordered functions. In contrast to standard message passing (Gilmer et al. 2017) that considers interactions between two atoms in each message, many-body methods aim to incorporate the interactions of multiple atoms surrounding the central node.
Interestingly, equivariant neural networks leveraging spherical harmonics as edge features can be interpreted from another perspective of approximating equivariant polynomial functions. This stems from the ability of spherical harmonics to represent both invariant and equivariant polynomial functions via tensor products. Additionally, the tensor contraction operation introduced by Atomic Cluster Expansion (ACE) (Drautz 2019:
Dusson et al.) 2022 Kovács et al. 2021 broaden the model's 's scope, covering a wider range of polynomial functions. Thus, the primary aim of this work is to develop a model that exhibits enhanced expressiveness, particularly for predicting atomic energy and force fields.
In the present study, we expand the scope of equivariant polynomial function analysis, transitioning from point cloud networks to equivariant networks that aim to predict symmetric physical properties, such atomic energy. This analytical extension informs the development of our novel equivariant network, PACE. Our proposed network effectively integrates an edge booster and many-body interactions through the ACE technique, which demonstrates an advanced capacity for approximating SE(3) x Sn equivariant functions, with higher polynomial degrecs. Our method is termed PACE as it is based on polynomial function approximation and ACE. To evaluate the efficacy of our approach, we assess our model on three molecular dynamics simulation datasets, including rMD17, 3BPA and AcÂc, and obtain consistent performance enhancements. Notably, our model not only achieves state-of-tho-art performance in energy and force predictions but also exhibits strong generalization capabilities in geometric out-of-distribution settings under MD across different temperatures.
2 Symmetries and Polynomial Functions
2.1 Symmetries
Incorporating physical symmetries into machine learning models is pivotal for tackling quantum chemistry challenges. This is because numerous quantum properties of molecules inherently exhibit equivariance or invariance to symmetry transformations. For instance, if we rotate a molecule in 3D space, forces acting on atoms rotate accordingly while the total energy of the molecule remains invariant. Besides rotation, the permutation of atoms within a molecule represents another type of symmetry. In such scenarios, the forces on atoms obey permutation oquivariance, while the molccular energy remains invariant to atomic permutations. To summarize, stomic forces can be characterized as (E(3) S--equivariant to atomic permutations, reflecting the respective symmetry groups G associated with these symmetry operations. Concurrently, molecular energy exhibits -invariine and S,-invariance.
to rotational transformation ions and
Mathemalically (Bronstein et al., 2021 given a group G and group action *, we say a function / mapping from source domain Q to target domain Y is G-equivariant if f(g * q) = g * f(q), Vq E Q, Vg E G. Similarly, if f(g q) = f() holds, we say f G-invariant. Due to the intrinsic SE(3) and S, oquivariant and invariant symmetries in quantum chemistry, it natural to encode these symmetries directly into the model architectures for effectively approximating the function f.
2

---

Published in Transactions on Machine Learning Research (08/2024)
2.2 Polynomial Functions in Equivariant GNNs
Building on this concept of learning powerful invariant or equivariant function f. various equivariant networks have been developed. These networks (Thomas et al. 2018 Batatia et al. 2022b; Yu et al. 2023b) are designed for mapping 3D molecular structures to inherently symmetric properties, such as energy, force, and the Hamillonian matrix. Typically, these networks capture directions and distances F by using real spherical harmonics Y (f) combined with radial basis functions (rbf). Owing to their equivariance property, spherical harmonics are advantagcous for encoding geometric information and predicting the physical properties of molecules.
It is noteworthy that spherical harmonics, partially discussed in Dusson et al. (2022) can also serve proper features for representing SE(3) invariant and cquivariant polynomial functions with the help of tensor products. Here, we illustrate how spherical harmonics represent polynomial functions using a specific example. We start with a polynomial function without SF(3) invariance and equivariance constrain. In this polynomial function denoted as PD (I, y, x) R R2, D is the highest degree of the polynomial terms, and f = (z, y, z). By observing Equations 1 to 33. we can find polynomial terms of p (ax, y, v)2) with D S 2 Y
and y ,==0,
yt=1 6 y where denotes tensor product, C is the coefficients of the spherical harmonics, and y (f) represents the tensor product of y==1(f)) (f) with itself, repeated t times. That's to say, the polynomial functions (1, can be represented linear combination of entries in Y yt=l, and y & y
Next, let's consider the SE(3) equivariant polynomial function PD23 (I, v, R R and we use P [k] to denote the k-th output of this function, where < k I K. Note that cach output of this function is still polynomial function, which is denoted P [k] (E, S 2). Therefore, there exists a linear function mapping from (f), ((=) ( (f) into PDE2 [ [1, v,,, K, which then comprise ( (x, y, 2).. Given that both the input features and polynomial functions are SE((3 equivariant, the corresponding linear mapping also retains this equivariance. From this, we can conclude that for any polynomial function P)S(()) (I, there exists an equivariant linear function mapping from y==0 (f). (f), and 1==1(()) (f)@2 into this function. This principle of linear universality is also discussed in
Dym & Maron (2020)
(f)= Cr [ 1
(1)
y/
y(=l(f) = Cimi
2
(2)
x
y²²
4/
yz
( y(= CC= =C
zy
Z
(3)
2
As previously elucidated, spherical harmonics and their tensor products provide way to approximate the polynomial functions, ensuring the preservation of (quuvviequv In existing equivariant graph
networks (Batzner et al. 2022 Batatia et al. 2022b Liao & Smidt 2023), spherical harmonics usually serve as
components of edge messages to provide (E (3) equivariant features. For S, permutation equivariance, Graph
Neural Networks (GNNs) Kippf & Welling 2017 Veličković et al.) 2018 Gao & Ji 2019 Liu et al. 2020 Cai
et al. 2021 Chen et al. 2024) employ a message passing scheme to aggregate neighboring messages for each central node. Therefore, by adhering to the aggregation of this message passing scheme, Sqqiivaaun features can successfully attain Sn equivariance. Specifically, edge messages are constructed using spherical harmonics (fa)) leading to aggregated equivariant features that are analogous to
fx, Σ y y(=l((u)) (Pt))
(4)
JEM,
3

---

Published in Transactions on Machine Learning Research (08/2024)
Equivariant GNNs like akin to
Anderson al. (2019) and Batatia al. (002bb) further implement architectures
1x Σ y (fin)
Σ $2EA
y (Pija..
55)
JIEN
using tensor product within their update modules. Both the basic aggregated message fx; and its updated form fx qualify as SE(3) x Sn equivariant polynomial functions with the highest degree D < 2. However, they do not cover all SB(3) x Sn equivariant polynomial functions with the highest degree D. For example,
fx.
Σ (t=)))))) JEN;
is also a SE(3) x Sn equivariant polynomial function with D 2, but it is not covered by fxx as demonstrated in appendix D Therefore, developing a model architecture capable of covering a broader range of equivariant polynomial functions crucial for enhancing its expressiveness.
2.3 Atomic Energy with Local Environment for Equivariant Polynomial Functions
The approximation of polynomial functions in point cloud networks has been explored previously, as detailed in Section 3.1 (Maron et al. 2019 Keriven & Peyré 2019 Dym Maron 2020). In this series of analyses, the architecture's 's capacity to approximate invariant or equivariant polynomial functions is demonstrated, using 3D coordinates (P1. P2, ... Pn)) as inputs. However, the analysis for equivariant networks in predicting symmetric physical properties is still underexplored. In energy and force prediction tasks, atomic cluster expansion (ACE) is widely used technique for approximating atomic energy. ACE decomposes the total energy Eg into a sum of individual atomic energies E, for each atom i, and then uses local environments to learn these atomic energies, formulated as
F, ΣΣ
50 1d (riz)+
Σ
Σ
Caro Dor (rim) ) Our (rijz) +
2
j1d2 U1
Σ
Σ
eieze dvs (reja) Dog (Tijz) Pos (rass)
3
jij2ja P1922s
(7)
where 0, = (Fej.."' ,iix)) denotes edges the atomic environment, Fiy denotes edge direction and distance from atom to atom j, 6 denotes the single edge basis function, and represents the coefficients. Therefore, in the context of approximating atomic energy, opposed to their use in point cloud networks, the polynomial function of node features fx, here focuses on the N-edge system 0 rather than atomic coordinates.
The research by Dym & Maron 2020) introduces the concept of D-spanning function sets to denote the capability of covering SE(3) Sn equivariant polynomial function with the highest degree D. Expanding upon their work, D-spanning function for approximating atomic energy within an N-bond system is defined
Σ jidamdd
0t
(8)
1j2
where K 2 D, and t (t,.... .x). Then, D-spanning set is defined as
Q = (t) (.) |||||| s D).
(9)
where 1 denotes an equivariant linear function mapping. In Appendix A.5 we provide the mathematical definition of D-spanning, and clucidate the relationship between D-spanning function sets and the model's expressiveness.
4

---

Published in Transactions on Machine Learning Research (08/2024)
2.4 Motivation of PACE
Inspired by the goal of approximating SE(3) x Sp equivariant functions with full spanning and higher polynomial degrees for the atomic energy based on the local atomic environment, in this work, we propose an equivariant network PACE following the message passing scheme. To briefty summarize using the same example from Equation
4 PACE incorporates a novel edge booster that leverages the tensor product of spherical harmonics from & S max construct edge messages analogous to rij S mmax N boost. Hence,
the aggregated equivariant features in PACE become
fx = Σ
ig L Lmax honst
(10)
JEAG
Then, PACE utilizes an update module with the symmetric contraction module from MACE to conduct many-body interaction with correlation v and the corresponding equivariant polynomial function defined
fx =
Σ
@h in
x ... @ @temax t1. tt S maxx * N boost.
(11)
Jjmmax
J.E,
vmax times
Consequently, the updated equivariant features fx. can effectively cover SE(3) Sn equivariant functions with formulation Equation 11 and achieve D-spanning with D = min (fmax * N boost; Umax
3 Related Works
3.1 Universality Analysis
Universality is powerful property for neural networks that can approximate arbitrary functions. While Zaheer
et al. [2017]; Maron et al. (2019); Keriven & Peyrê (2019) study the universality of permutation invariant
networks, several works have recently studied the rotational equivariant networks. Dym & Maron (2020) takes
use of the proposed tensor representation to build D-spanning family and shows that Tensor Field Networks
(TFN) (Thomas al.] 2018)) proved to be universal equivariant network capable of approximating arbitrary equivariant functions defined the point coordinates of point cloud data. Furthermore, GemNet (Gasteiger
et al. 2021 uses the conclusion in Dym & Maron 20020) and is proved to be universal GNN with directed edge embeddings and two-hop message passing.
3.2 Equivariant Graph Neural Networks
In recent years, equivariant graph neural networks have been developed for 3D molecular representation learn- ing, as they are capable of effectively incorporating the symmetries required by the specific task (Ruddigkeit let al. 2012 Chmiela et al. 2017 Yu et al. 2023a Khrabrov et al., 2022). Existing equivariant 3D GNNs can be broadly classified into two categorics, depending on whether they utilize order & equivariant features or higher order >>>>> equivariant. features. Methods belonging to the first calegory (Satorras et al. 2021 Schütt al., 2021 Deng et al.) 2021 Jing et al. [002 Thölke & Fabritiis 2022) achieve equivariance by applying constrained operations order 1 vectors, such vector scaling, summation, linear transformation.
vector product, and scalar product. The second category of methods (Thomas et al. 2018 Fuchs al. 2020
Liao & Smidt
2023 Batzner et al. 2022 Batatia et al. 2022a b Brandstetter et al. 2021 predominantly
employs tensor products (TP) to preserve higher-order equivariant features, with some works (Luo et al.
2024
focusing on accelerating tensor product computations.
3.3
Atomic Cluster Expansion
Molecular potential and force field are crucial physical properties in molecular analysis. To approximate these properties, the atomic cluster expansion (ACE) (Drautz) 2019 Kovács et al. 2021 is used to approximate the atomic potential. Recently, several neural networks aiming to predict atomic potential and forces have been developed to consider many-bot interactions by incorporating ACE into their model architectures. Specifically, BOTNet (Batatia. et. al. 2022a) takes multiple message passing layers to encode the many-body
5

---

For a MACE layer, Published in Transactions on Machine Learning Research (08/2024)
interaction and analyzes the body order for various message passing schemes. MACE (Batatia et al. 2022b takes generalized Clebsch-Golden coefficients to couple the aggregated message to incorporate higher-order interactions. Allegro (Musaelian et al. 2023) uses a series of tensor product layers to calculate the equivariant representations without using message passing, learning many-body interactions.
3.4 Comparison to Existing Works
A NequlP [Batzner et al. 2022) layer builds its message using the original spherical harmonics, and its output irreducible representations can cover equivariant polynomials defined as Ejev, Fij : S max which can form D-spanning family with the highest degree D
the output irreducible representations can cover the equivariant polynomials defined as QK &t &tomax 41," 11,'0 bumax fmax- which can form a D-spanning family with Lt @ ..... @ Jomax U1
En temax EN.
Umax
degree = max {lmax: Umax) J. For the proposed PACE layer, the output irreducible representations cover @temax 11, Lumax S Eaxx * N honst, spanning D-spanning fam-
K= - En Ivmax CN;
34 2 ... & Jjmmax
Umnxx times
ily with D = max {fmax N boost, Umax as explained in Section 2.4 During our experiments, set fmax 3 and Umax = 3 following previous work (Batatia et al. 2022b Musaclian et al. 2023) for fair comparsion. Although both MACE and PACE can form a D-spanning family with D = 3, our PACE model can still approximate a greater number of polynomial functions with D > in this scenario. Moreover, by using additional self-interactions, PACE needs fewer channels to approximate the same amount of positions in D-spanning functions.
4 The Proposed PACE
The node features x² are initialized through a linear transformation applied to its atomic type. Edges constructed based on cutoff distance have orientations fgj denoted by spherical harmonics Y (fy)), and pairwise distances Fs embedded using a learnable radial basis function R. The proposed PACE comprises two distinct message passing layers followed by an output layer. The illustration of PACE architecture is provided in Figure [ In this section, we will introduce the architecture our proposed PACE model as well as the corresponding equivariant polynomial functions from the perspective of irreducible representations. Detailed information about the irreducible representations used in equivariant GNNs can be found in Appendix A.I As shown in Appendix B.1 the irreducible representations can also represent D-spanning family.
Theorem 4.1. For any D-spanning function (0.) oppeared in Q and for any position P (P1,P2,"* .Pk) in tensor representation, where pk E R3 denotes the element position, if there exists w1 and irreps, such that Φ (10((()))) (O.) Etm wft irrepst" and the set of irreps, forms D-spanning family.
Definition 4.2. If the set of irreps, forms D-spanning family ((+), we can say that the irreducible representation irreps, represents Q/ (0.).
4.1 The First Layer
Following the message passing neural network, the first step is to build edge messages from neighboring nodes the central node. As discussed in Section 2.4, the construction of these edge messages for node pairs equivariant neural networks generally requires a tensor product combine edge spherical harmonics and node features at each layer. Diverging from this conventional design, the first layer of PACE uses the edge booster two consecutive tensor products, which means Nboost = 2, to construct the edge message
m ((is)) Own,y MLP(xx || xx)"),
(12)
m Y (fiz) @w2, m where the learnable weights wi, and 22,6 applied to each tensor product are obtained by
(13)
wi,y = MTP (R(:))
(14)
6

---

Published in Transactions on Machine Learning Research (08/2024)
Layer
E
Output
rbf43
MLP
Salf
inier rereectioon
MLP
V11.5
022.3
Many body interaction
X
Sphcs
Sphiy
! E
B
Layer
AE
Self-
M3.8j
MLP
x
m
E
rbfe
Many-body interaction
MLP
C
ACE 4-body interaction
D
Polynomial many-body interaction
J1
ja
Self-interaction
Air
Tensor contraction
Fest
A
Aing
A
Self-interaction
Aivz
Ateg @ Asa
Edn (Fig))
Σ Do (Fija)
Dea (Fija)
Self-interaction
Ates
Airs
An architecture overview of PACE. A: The first message passing layer. The initial node features x! and x are concatenated and transformed by an MLP. Then, a tensor product applied to the transformed node features and edge spherical harmonics Sph, with an RBF and MLP transformed edge distance as learnable weights. Next, output is further combined with Sph,y using another tensor product, and the weights involve the initial node features and edge distance. These operations in the yellow box comprise the edge booster, aiming to enhance the model's expressiveness. The obtained edge-boosted messages are aggregated to form the atomic base. Then, an update module comprising polynomial many-body interaction module and skip connection is used to update the fcatures of central nodes. B: The second message passing layer. A tensor product is applied to the edge message and updated node features that are obtained from the first layer. The learnable weights of the tensor product are based on the initial node features and distance each neighboring pair. Then, edge messages are aggregated and node features are updated similarly those in the first layer. C: An example of 4-body interaction in ACE. We aim fit E Φeg (Faji ) Σ Bey (rija) s9 Σ &ug (Fija) using Aiv, @ Ain, Aivg: where {Ø} denotes the atomic base in ACF. D: Polynomial many-body interaction module. The atomic base A; is fed to multiple self-interaction layers separately to produce different Açe. Then, tensor contraction is performed to produce ă, E: Output. The invariant part of node features produced by both layers are transformed and summed to predict the deviation from the total molecular energy to its average.
Figure 1:
22,9 MLP(R(Fg)) + MLP(2) xx]),
(15)
considering the atomic types for node pairs and their pairwise distances. Then, we obtain the edge message by concatenating these two together, denoted as
m (m m
(16)
7

---

Published in Transactions on Machine Learning Research (08/2024)
As shown in Appendix B.1 the edges features from spherical harmonics represent rij t maxx, with mmx = 3.
Then the second tensor product increases it to Ki I N boost max with N boost = 2 using Lemma B.2
After obtaining edge messages via the edge booster, we aggregate neighboring messages to the central node
4 = [(O)
Σ
m.j.
(17)
JEN(()
where [i) is the number of neighbor nodes. Here, A, analogous the atomic base of the central node in
the original ACE method (Drautz 2019: Dusson et al. 2022 Kovács ct al. 2021 Mcanwhile, the equivariant
polynomial form of A, is ΣJEN, 5 6.
Next, A is processed by a polynomial many-body interaction module to produce features a,. The corresponding 8ts K=a (0.) Σj Ja F & r @ Tijs t1, 12, tg 6 for (0)) with D < 3 as shown in Appendix B.2.2 These features are then combined with the initial node features to update for cach node. Details of the polynomial many-body interaction module are described in Section and the design of the first layer is illustrated in Figure TA.
equivariant polynomial form of à is Q
4.2 The Second Layer
The second layer of PACE (Figure (B) involves only one lensor product to construct the message. However, unlike NeuqulP and MACE, which leverage the original edge spherical harmonics, we instead use the odge-boosted message m obtained from the first layer. Specifically,
mi ml & x)
(18)
where x/ denotes the updated node features from the first layer and the learnable weights are
W3,i MLP((((F))) MLP(( MMP()) x[])-
(19)
Next, messages are aggregated for the central nodes. Finally, another polynomial many-body interaction module is applied to update node features. The I= 0 invariant features outputted by both layers are further used for energy prediction.
4.3 Polynomial Many-Body Interaction Module
The polynomial many-body interaction module playing important role in both PACE layers is proposed to incorporate many-body interactions by mixing the atomic base A,. As shown in Figure [P, we first use Umax different self-interactions to map the input atomic base A₁ to different bases Ajo following
W AC +b t=0
(20)
Af A%
e>
where f is the rotation order of irreps, is the channel index for A; and d' is the channel index for AG Then, use tensor contraction (Batatia al.] 2022b) with generalized Clebsch-Golden to fuse multiple atomic bases. Tensor contraction is illustrated in Appendix A.4
As shown in Equation 25 in Appendix, in the case of two irreps, Clebsch-Gordan coefficients Chams tyma.emm are used to maintain equivariance when fusing two irreps with rotation orders & and l2 to the output f3, and the triplet (f1, f2, (a) is defined as a path. When fusing N irreps, the generalized Clobsch-Gordan cocfficients used to maintain the equivariance can be defined as
C[NNMMNN]
L. M.a Ma .amz Lymy
MN
(21)
Cam,,
where [[N] Lz, LN) with [4-1 4 < L, S 1,-1 + G/. L, E N, Vi 2 2, E N+, and the path is shown as n/N]= (ei, l2, Toz, la Tay.... 1N-11 LN -
LN ToN).
Then, the output irreps is contracted one by
8

---

Published in Transactions on Machine Learning Research (08/2024)
one to consider the coupled many-body interactions shown as
MN 1(v--),,bu 11,bu- Mo-1
Ao
E (Wenr nivl
+ t Lo
(22)
where W is the path weight, ä is the intermediate irreps, and N+ starts from N to to incorporate N-body interactions. Note that C = 1 and Lo = O when v = 1, and ā = 0 when == N.
In contrast to MACE (Batatia et al. 2022b) applying tensor contraction to the atomic base A,, our PACE first takes multiple self-interactions to distinguish the atomic base from the same A; to different Air for different body orders before applying the contraction. According to Darby et al. (2023), the out of tensor contraction in MACE is defined as the tensor decomposed product basis while the one obtained in PACE is called tensor sketched product basis. Although Darby et al. (2023) has demonstrated that the lensor-decomposed basis can consistently reconstruct the tensor-sketched basis with factor of Umax times channels, we implement the latter one in PACE to improve the model expressiveness without increasing the number of channels.
4.4 Output
As illustrated in Figure IE to compute the total energy of the molecule, we extract and transform the invariant part of node features produced by each PACE layer
E E= EMMLP(<) + MXP(xx
(23)
where E' is the averaged total energy of the training set, and N represents the number of atoms in a molecule. Once the total energy is predicted, we then use f; = 0p. to calculate the force acting on cach atom, as it ensures energy conservation.
4.5 Summary of Key Contributions
The first key contribution of our work is the edge booster, which is computed by Equa- tions 12) to (16 and illustrated as the yellow-boxed operations in Figure [ Our edge booster leverages consccutive tensor products on spherical harmonics which helps enhance an equivariant modcl's ability
Edge Booster
approximate higher-degree polynomial functions. This idea is shown Equation (10) and the overall polynomial degree is shown Appendix B.2.2, There might be various implementations to build an edge booster following the idea of Equation (10). However, we choose to design our edge booster as described by Equations 12 to (16] for two reasons. First, edge-boosted message only considers spherical harmonics and atomic types of nodes. Note that the computation of our edge booster does not involve node features that are updated layer by layer. Such a design follows our motivation to create an independent edge booster module that can produce an edge-boosted message to replace the original edge spherical harmonics used in each layer. Based on this design, we only need to compute the edge booster once in the first layer and we can reuse the odgo-boosted message all other layers. As shown in Figure AA&B, the odgo-boosted message m! is used in both layers. Hence, the design of our edge booster also considers computations ional costs. In summary, the result of our edge booster can be used in each layer to help approximate more polynomial functions with higher degree in an efficient manner. A detailed cost analysis of edge booster is provided in Appendix D.2.3
Additional Self-interaction Layers The second enhancement comes from the additional self-interaction layers used in the polynomial many-body interaction module. In Scction 4.3 we use the tensor contraction proposed in MACE (Batatia el al. 2022bb to approximate many-body interactions, but we apply additional self-interactions (SI) to the atomic base A, so that we use different inputs of the tensor contraction. In MACE, the atomic base A, directly fed into the tensor contraction, while the atomic base A, our design first fed to multiple inferreteratnonn layers separately to produce different atomic bases Age. The functionality of these self-interaction is to increase the ability to approximate more positions in D-spanning functions.

---

Published in Transactions on Machine Learning Research (08/2024)
Table 1: Performance on the rMD17 dataset. Mean absolute errors (MAE) are reported for both energy (E) and force (F) prodictions, with mcV and meV/À as units, respectively. Bold numbers highlight the best performance.
ACE
FCHIL
CAP
ANI
GemNet (T/Q)
NoqulP
BOTNet
Allegro
MACE
Ours
E
6.1
6.2
17.7
16.6
2.3
2.3
2.3
2.2
1.7
Aspirin
F
17.9
20.9
44.9
40.6
9.5
8.2
8.3
7.3
6.6
5.8
3.6
2.8
8.5
15.9
0.7
1.2
1.2
0.5
Azobenzene
10.9
10.8
24.5
35.4
3.3
2.6
3.0
E
0.04
0.35
0.75
3.3
0.04
0.03
0.4
0.4
0.02
Benzene
F
0.5
2.6
6.0
10.0
0.5
0.3
0.3
0.2
0.3
0.2
1.2
3.5
2.5
0.4
0.4
0.4
0.3
Ethanol
7.3
18.1
13.4
3.2
1.7
1.5
4.8
4.6
0.8
0.8
0.6
0.8
0.5
Malonaldehyde
F
11.1
10.3
26.4
24.5
6.6
5.1
5.8
3.6
4.1
3.6
0.9
3.8
11.3
0.2
0.2
0.5
0.2
Naphthalene
5.1
16.5
29.2
1.8
0.9
1.6
4.0
2.9
8.5
11.5
1.4
1.3
1.5
1.3
0.9
Paracetamol
F
12.7
12.3
28.9
30.4
5.9
5.8
4.9
4.8
4.0
F
1.8
1.8
5.6
9.2
0.7
0.8
0.9
0.9
0.5
Salicylic acid
F
9.3
24.7
29.7
3.1
4.3
2.9
1.1
1.7
4.0
7.7
0.3
0.3
0.4
0.5
0.2
Toluene
F
6.5
8.8
17.8
24.3
2.2
1.6
1.9
1.8
1.5
1.1
F
1.1
0.6
3.0
5.1
0.4
0.4
0.6
0.5
0.3
Uracil
F
6.6
4.2
17.6
21.4
3.8
3.1
3.2
1.8
2.1
2.0
Specifically, to approximate n different positions in D-spanning function, PACE without SI needs Umax channels, while PACE with ST only needs channels. This is because the ST in the polynomial many-body interaction module enables every single channel in PACE to approximate
position in D-spanning functions. The corresponding proof is shown in Appendix B.2.2 The total number of self-interactions equals the max correlation order Umax, where Umax corresponds to (body order - 1). In our implementations, we use three self-interactions because the highest body order we consider in each layer is 4. Although Umax? channels without self-interactions can achieve the same expressiveness theoretically, we find adding these additional self ief-iinteractons can significantly enhance the empirical performance.
5 Experiments
We conduct experiments on three molecular dynamics simulation datasets, the revised MD17 (rMD17), 3BPA and AcAc datasots. The proposed PACE is trainod using these datasets predict both the invariant energy of the entire molecule and the equivariant forces acting individual atoms. Among our baselines (ovácc
et al.
2021
Christensen et al. 2020 Bartók et al.
2010
Smith et al. 2017 Gasteiger et al. 2021 Batzner
et al.
2022
Batatia et al. 2022a Musaelian et al.
2023
Batatia al. 2022b), NequlP, BOTNet, Allegro
and MACE
are all equivariant graph neural networks (GNNs) with rotation order & > 1. In particular,
BOTNet, Allegro, and MACE incorporate many-body interactions, while ACE is a parameterized physical model that does not belong to the class of neural networks. Our experiments are implemented with PyToh 1.11.0 Paszke al. 2019), PyTorch Geometric 2.1.0 (Fey & Lenssen 2019), and e3nn 0.5.1 (Geiger & Smidt, 2022 In experiments, we train models on a single 11GB Nvidia GeForce RTX 2080Ti GPU and Intel Xeon
Gold 6248 CPU.
5.1
The rMD17 Dataset
Datasct. The rMD17 (Christensen Von Lilienfeld, 2020) is a benchmark datasct that compriscs ten small organic molecular systems. Each molecule in the dataset is accompanied by 1000 3D structures, which were generated through meticulously accurate ab initio molecular dynamic simulations employing density
10

---

Published in Transactions on Machine Learning Research (08/2024)
Table 2: Performance on the 3BPA dataset. Root-mean-square errors (RMSE) are reported for both energy (E) and force (F) predictions, with meV and meV/A as units, respectively. Standard deviations are calculatod over three runs with different seeds. Bold numbers highlight the best performance.
ACE
NoqulP
BOTNet
Allegro
MACE
Ours
E
3.28 (0.12)
3.1 (0.13)
3.84 (0.10)
3.0 (0.2)
2.4 (0.1)
300K
F
27.1
10.77 (0.28)
11.0 (0.14)
12.98 (0.20)
8.8 (0.3)
9.1 (0.1)
E 24.0
11.16 (0.17)
11.5 (0.6)
12.07 (0.55)
9.7 (0.5)
7.9 (0.2)
600K
F
64.3
26.37 (0.11)
26.7 (0.29)
29.11 (0.27)
21.8 (0.6)
21.4 (0.3)
E
85.3
38.52 (2.00)
39.1 (1.1)
42.57 (1.79)
29.8 (1.0)
29.6 (0.4)
1200K
F 187.0
76.18 (1.36)
81.1 (1.5)
82.96 (2.17)
62.0 (0.7)
60.7 (2.0)
E
23.3
16.3 (1.5)
7.9 (0.6)
7.6 (0.4)
Dihedral Slices
F
23.1
20.0 (1.2)
16.5 (1.7)
16.0 (0.5)
Table 3: Performance on the AcAc dataset. Root-mean-square errors (RMSE) are reported for both energy (E) and force (F) predictions, with meV and meV/À as units, respectively. Standard deviations are calculated over three runs with different seeds. Bold numbers highlight the best performance.
BOTNet
NequlP
MACE
Ours
0.89 (0.0)
0.81 (0.04)
0.9 (0.03)
0.73 (0.02)
300K
F
6.3 (0.0)
5.90 (0.38)
5.1 (0.1)
4.2 (0.1)
6.2 (1.1)
6.04 (1.26)
4.6 (0.3)
3.8 (0.6)
600K
F
29.8 (1.0)
27.8 (3.29)
22.4 (0.9)
16.9 (0.3)
functional theory (DFT). These structures capture the diverse conformational space of the molecules and are valuable for studying their quantum properties.
Setup. In our experiments, we use officially provided random splits. Next, we use the same splitting seed as MACE to further divide the training set into a training set comprising 950 structures and a validation set comprising 50 structures. Then, perform our evaluations on the test set with 1000 structures. Training details are provided in Appendix C.2
Results. Table summarizes the performance of our proposed method in comparison to baselines on all ten molecules in the rMD17 dataset. Mean absolute errors (MAE) are employed as the evaluation metric for both energy and force predictions. is worth noting that our PACE demonstrates state-of-ar of-rtt performance in energy prediction across all molecules. Specifically, we achieved significant improvements of 33.3%, 33.3% and 30.8% on Benzene, Toluene, and Paracctamol, respectively. In terms of force prodiction, PACE achieves stale-of-the-art performance on nine out of the (en molecules, exhibiting substantial improvements of 26.7 76,7% 16.7% and 15.4% on Toluene, Paracetamol, and Azobenzene, respectively. Besides, we attain the second-best on Uracil. The comparison of algorithm efficiency and the affect from rotation order can be found in Appendix D.2 and D.3
5.2 The 3BPA & AcAc Datasets
Dataset. The 3BPA dataset (Koáács et al. 2021 also generated through molecular dynamic simulations employing Density Functional Theory (DFT). Unlike rMD17, this dataset is specifically focused on single flexible molecule, namely the 3BPA molecule. 3BPA is characterized by three freely rotating angles, which primarily induce structural changes at varying temperatures. The AcAc dataset (Batatia et al. 2022a) is also generated through molccular dynamic simulation using Density Functional Theory (DFT). Similar to 3BPA, this dataset specifically focuses on a single flexible molecule, Acely Acetylacetone. Given that many real-world applications involve temperature fluctuations, evaluating a model's robustness and extrapolation capabilities under varying temperature conditions is crucial for predicting molccular bohavior accurately. 3BPA and AcAc datasets are frequently employed assess a model's out-of-distribation (OOD) generalization capability MD geometric distributions across different temperatures for the same molecule.
11

---

Published in Transactions on Machine Learning Research (08/2024)
Table 4: Results of ablation experiments. Compared to PACE, PACE (-EB) removes edge booster in the first. layer, and PACE (-ST) uses no additional sclf-intcraction in the polynomial many-body interaction module. Mean absolute errors (MAB) are reported for both energy (B) and force (F) predictions, with meV and meV/ Á units, respectively. Bold numbers highlight the best performance.
AcAc 300K
AcAc 600K
E
0.73 (0.02)
3.8 (0.6)
PACE
4.2 (0.1)
16.9 (0.3)
E
0.84 (0.04)
3.8 (0.1)
PACE (-EB)
4.9 (0.1)
18.0 (0.3)
E
0.95 (0.01)
4.0 (0.2)
PACE (-SI)
5.7 (0.1)
21.3 (1.0)
Setup. For both 3BPA and AcAc, We follow MACE to train our model using a training set consisting of 450 structures and a validation set comprising 50 structures. Both the training and validation sets were sampled at 300K. Both 3BPA and AcAc have two test sets sampled 300K and 600K. 3BPA has another test set sampled at 1200K. Moreover, 3BPA has an additional test set, including optimized conformations with dihedral slices, where two dihedral angles are fixed and the third angle varies from O to 360 degrees. This test set probes regions of the potential energy surface (PES) that are distant from the training set. We provide training details in Appendix C.2
Results. Table 2 and Table summarize the performance of our proposed method in the 3BPA and AcAc datascts, respectively. Here, noot-eean-square error (RMSE) is used as the evaluation metric. On 3BPA dataset, the proposed PACE shows comparable performance to MACE, while outperforming other baseline methods significantly. On AcAc dataset, our PACE demonstrates significant performance improvement from all basclincs. Morcover, we use the model trained on 3BPA to conduct molecular dynamic simulations and the comparison of PACE-generatie trajectories and the ground-truth trajectories is shown in Appendix D.1
5.3 Ablation Studies
The first ablation study aims to demonstrate the effectiveness of the proposed edge booster, designed to facilitate a higher polynomial degree. In this experiment, we remove the second tensor product, namely the odge booster, from the first layer of PACE. As indicated in Table 4 this comparison reveals that the edge booster, which enhances expressiveness, indeed imrpoves the model's performance in force field prediction.
The second ablation study is for demonstrating the effectiveness of using different atomic bases Ase for different body orders in many-body interaction. In this experiment, we remove additional self-interaction operations and directly use A; in symmetric contraction in the polynomial many-body interaction module. The comparison shown in Table Sable-justin justifies that different alomic bases A, produced by additional self-interactions can improve the expressiveness and model performance without increasing the number of channels.
6 Conclusion
In this work, we introduced PACE, a new equivariant network for atomic potential and force field predictions. Our PACE is designed from two perspectives, increasing the ability to approximate more positions in D- spanning functions, as well expanding the space cover greater number of higher degree polynomial equivariant functions. The integration of the edge booster and Atomic Cluster Expansion (ACE) technique with additional self-interactions allows PACE to approximate SE((3 x Sn equivariant polynomial functions with higher degrees, thereby improving the accuracy in prediction. The comprehensive experimental results and analyses provide compelling evidence for the efficacy of the proposed PACE model. In the future, we aim to advance the development of models capable of approximating polynomial functions of even higher degrees while simultaneously maintaining high computational efficiency.
12

---

Published in Transactions on Machine Learning Research (08/2024)
Acknowledgements
This work was supported in part by National Science Foundation grant IIS-2243850 and National Institutes of Health grant U01AG070112.
References
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neurol information processing systems, 32, 2019.
Kendall Atkinson and Weimin Han. Spherical harmonics and approcimations on the unil sphere: introduction, volume 2044. Springer Science & Business Media, 2012.
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.
Tyes Batalia, Simon Batzner, Dávid Péter Kováes, Albert Musaelian, Gregor NC Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi. The design space of -quivaaran atom-centered interatomic potentials. ar preprint ar orXi:2205.06644 2022a.
llyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order cquivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:11422 1423 11436, 2022b.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess F. Smidt, and Boris Kozinsky. E. (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communiccations, 13(1)22453, 2022.
Bastiaan Braams and Joel M Bowman, Permutationally invariant potential energy surfaces in high dimensionality. International Reviews in Physical Chemistry, 28(4):577-606, 2009.
Johannes Brandstetter, Rob Hesselink, Flise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve (3) equivariant message passing. arXiv preprint erXiv:211000200 02905, 2021.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.1 13478, 2021.
Alex Brown, Anne B McCoy, Bastiaan J Braams, Zhong Jin, and Joel M Bowman. Quantum and classical studics of vibrational motion of ch 5+ on a global potential energy surface obtained from a novel ab initio direct dynamics approach. The Journal of chemical physies, 121 (9):4105 4116, 2004.
Lei Cai, Jundong Li, Jie Wang, and Shuiwang Ji. Line graph neural networks for link prediction. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(9):5103 --11,, 2021.
Pei Chen, Sounajyoti Sarkar, Leonard Lausen, Balasubrananiam Srinivasan, Sheng Zha, Ruihong Huang, enhanced tabular data representation learning. In Advances in
and George Karypis. Hytrel: Neural Information Processing Systems, 2024.
Stefan Chmiela, Alexandre Tkatchenko, Huziel F. Sauceda, Tgor Poltavsky, Kristof T Schütt, and Robert Müller. Machine learning of accurate onergy-consorving molecular force fields. Science advances, 3(5): e1603015, 2017.
Anders S Christensen and O Anatole Von Lilienfeld. On the role of gradients for machine learning of molecular energies and forces. Machine Learning: Science and Technology, 144):045018, 2020.
Anders S Christensen, Lars A Bratholm, Felix A Faber, and O Anatole von Lilienfeld. Fchl revisited: Faster and more accurate quantum machine learning. The Journal of chemical physics, 152(4):044107, 2020.
James Darby, Dávid P Kovács, Ilyes Batatia, Miguel Caro, Gus LW Hart, Christoph Ortner, and Gábor Csányi. Tensor-reduced atomic density representations. Physical Review Letters, 131 (2):028001, 2023.
13

---

Published in Transactions on Machine Learning Research (08/2024)
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J Guibas. Vector neurons: A general framework for squiiaaqnn networks. In Proceedings of the IEEE/CV International Conference on Computer Vision, pp. 12200 12209, 2021.
Ralf Drautz. Atomic cluster expansion for accurate and transfcrable interatomic potentials. Physical Review B, (1100144,, 2019.
Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and TieYYan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning, pp. 5583-5608 PMLR, 2022.
Weitao Du, Yuanqi Du, Limei Wang, Dieqiao Feng, Guifeng Wang, Shuiwang Ji, Carla Gomes, and Zhi-Ming Ma. A new perspective on building efficient and expressive 3d equivariant graph neural networks. ar preprint ar arXiu:2304.00 04777, 2023.
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 Computa- tional Physics, 454:110946, 2022.
Nadav Dym and Haggai Maron. On the universality of rotation equivariant point cloud networks. In International Conference on Learning Representations, 2020.
Matthias Fey and Jan E. Lenssen. Fast Graph Representation Learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling, SE(3)-transformmrs:: 3D roto-translation oquivariant attention networks. Advances in Neural Information Processing Systems, 33:1970-1981, 2020.
Hongyang Gao and Shuiwang Ji. Graph u-nets. In international conference on machine learning, pp. PMLR, 2019.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs In International Conference on Learning Represenlations, 2020.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal directional graph neural
networks for molecules. Advances in Neural Information Processing Systems, 34:6790-688 6802, 2021.
Mario Geiger and Tess Smidt. c3nn: Euclidean Neural Networks, 2022. URL
09453
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
David J Griffiths and Darrell F Schroeter. Introduction to Quontum Mechanics. Cambridge University Press, 2018.
Kurt Hornik. Approximation capabilities of multilayer feodforward networks. Neural networks, 4(2):251-277, 1991.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359 366, 1989.
Bowen Jing, Stephan Fismann, 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 https::/ net ntt.fforum//..-1..J........
Nicolas Keriven and Gabriel Peyré. Universal invariant and equivariant graph neural networks. Advances in Newral Information Processing Systems, 32, 2019.
14

---

Published in Transactions on Machine Learning Research (08/2024)
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, 25(42):25853-285 25863, 2022.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
Dávid Péter Kovács, Cas van der Oord, Jiri Kucera, Alice EA Allen, Daniel J 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):7696 7711, 2021.
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-struccuue calculation. Nature Computational Science, 2(6))367- 377, 2022.
Yi-Lum Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Fleventh International Conference on Learning Represenlations, 2023. URL https: //openreviee ew
Meng Liu, Hongyang Gao, and Shuiwang Ji. Towards deeper graph neural networks. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 338 348, 2020.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representations, 2022.
Shengjie Luo, Tianlang Chen, Yixian Xu, Shuxin Zheng, Hin Liu, Liwei Wang, and Di He. One transformer can understand both 2D & 3D molecular data. In International Conference on Learning Representations, 2023. URL http::://ppeeeerreee..//rrrrrror wtt.:///orrmmm.ddd////ooo....
Shengjie Luo, Tianlang Chen, and Aditi $ Krishnapriyan. Enabling efficient equivariant operations in the fourier basis via gaunt tensor products. In International Conference on Learning Representations, 2024.
Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality of invariant networks. In International conference on machine learning, pp. 1363-4371. PMLR, 2019.
Albert. Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit. Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, Perfomeero Deep Learning Library. Advances in Neural Information Processing Systems, 32, 2019.
Lars Ruddigkeit, Rund Van Deursen, Lorenz C Blum, and Jcan-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(11):2864 2275, 2012.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(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, Farhad Arbabzadah, Stefan Chmicla, Klaus R Müller, and Alexandre Tkatchenko. Quantum- chemical insights from deep tensor neural networks. Nature commanications, 8(1):13890, 2017.
15

---

Published in Transactions on Machine Learning Research (08/2024)
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. SchNet-a deep learning architecture for molecules and materials. The Journal of Chemical Plyyicc, 142(22))241722, 2018.
Nimrod Segol and Yaron Lipman. On universal equivariant set networks. In International Conference Learning Representations, 2019.
Justin Smith, Olexandr Isayev, and Adrian E Roitberg. ANI-1: an extensible neural network potential with DFT accuracy at force field computational cost. Chemical science, 844):3192-3203, 2017.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Represenlations, 2022.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and transsation--qquiiariint neural nerworks for 3D point clouds. orXiv preprinl wwXXi:1802.08811 2018.
Oliver Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus-Robent Müller. SE (3)-cquivariant prediction of molccular wavefunctions and electronic densitics. Advances in Neural Information Processing Systems, 34:14434 14447, 2021a.
Oliver T Unke, Stefan Chmiela, Michael Gastegger, Kristof T Schütt, Huziel Sauceda, and Klaus-Robert Müller. Spookynet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nature communications, 12(1):7273, 2021b.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tic-Ya Yan Liu. Do transformers really perform badly for graph representation? Advances in Newral Information Processing Systems, 34:28877-28888, 2021.
Haiyang Yu, Meng Liu, Youzhi Luo, Alex Strasser, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. QH9: A quantum hamiltonian prediction benchmark for QM9 molecules. arXio preprint ar arXix:33060096 09549, 2023a.
Haiyang Yu, Zhao Xu, Xiaofeng Qian, Xiaoning Qian, and Shuiwang Ji. Efficient and equivariant graph networks for predicting quantum Hamiltonian. In Proceedings of the 40th International Conference Machine Learning, 2023b.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and EJPRL Weinan. Deep potential molccular dynamics: a scalable model with the accuracy of quantum mechanics. Physical review lellers, 120(14):
143001, 2018a.
Linfeng Zhang, Jiequn Han, Han Wang, Wissam Saidi, Roberto Car, et al. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. Advances in neural information processing systems, 31, 2018b.
Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xic, ... and Shuiwang Ji. Artificial intelligence for science in quantum, alomistic, and continuum systems. ar Xie preprinl or Xiv:22337 Xtw:2307.08422, 2023.
16

---

Published in Transactions on Machine Learning Research (08/2024)
A Background and Related Work
A.1 Spherical Harmonics and Irreducible Representation
To encode geometric information molecules into
features, spherical harmonics are used for its equivariance properly. Specifically, we use real spherical harmonic basis functions Y to encode an orientation fyy between node pair. If the molecule is rotated by rotation matrix T in 3D coordinate system, then have:
(YY'(TF.)) ((D((T)')X"(())))
(24)
where L [0, L] is the degree, of with size 28 + denotes coefficients of spherical harmonics, and the Wigner D-matrix (T) with size (2l 1) x (2 + 1) specifies the corresponding rotation acting coeflicients q'. In practice, equivariant features are often represented by irreducible representations (irreps), which correspond to the coefficients of real spherical harmonics. Concretely, irreducible representations are formod by the concalenation of multiple lype-& vectors, where E [0, L]. More detailed explanation about irreducible representations can be found in Liao & Smidt (2023
A.2 Tensor Product
Equivariant networks using higher order > equivariant features (Thomas al. 2018 Fuchs al. 2020
Liao & Smidt 2023 Batzner al. 2022 Batatia et al., 2022a b) predominantly employs tensor products
(TP) to preserve higher-order equivariant features. Tensor product operates irreducible representations u of rotation order & and of rotation order f,, yielding new irreducible representation of order l as
(u x a mg
Σ
(Bs.ms) fi ((a.m) mmm
(25)
m22
where C' denotes the Clebsch-Gordan (CG) coefficients (Grilfiths & Schroeter 2018 and E N denoles the m-th element in the irreducible representation. Here, la satisfies 14 l2 S 43 S l + l2. and 4,, ls E N. High-order equivariant 3D GNNs commonly use tensor products on the irreducible reprosentations of noighbor nodes and edges to construct messages as
mij
(aleettaa (To)) 8xf.
(26)
EuRly
where 14, 56 Sh++l; Xy denotes features of node j, and R is learnable non-linear function that takes the embedding of pairwise distance Fiy as input.
A.3 Atomic Cluster Expansion
Molecular potential and force field are crucial physical properties in molecular analysis. To approximate these properties, the alomic cluster expansion (ACE) (Drautz 2019 the atomic potential denoted as
Kováes el al. 2021 is used to approximate
E, = ΣΣ ① d (1T)
+ ΣΣ (2) do, (rij. Poa (Tija)
22
31J2 01
ΣΣ
Vi02nng Pey (rig,) Pe, (raja) Dvg (rija) + ...,,
(3)
(27)
31
J1ja2a PiPgi
where 0, = (r.j..".. ,Fix)) denotes the N bonds in the atomic environment, ó is the single bond basis function and c is the coefficients. The computations ional complexity of modeling many-body potential increases exponentially with number of neighbors. To reduce the complexity, ACE further makes use of density trick to calculate the atomic energy via atomic base A = E, (n(rr))) which has a linear complexity with the number of neighbors, denoted as
VT 200
v1 0120220 2.
E,(0,) Σ
(1)
A +
Σ 1y02
0102 Awn Awy +
(2)
Σ
(3)
Crrgua Air Aivz Ales
(28)
17

---

Published in Transactions on Machine Learning Research (08/2024)
In this case, the computational complexity of modeling many-body interactions decreases to linear growth with the number of neighbors. With the reduced linearly complexity, many equivariant networks are designed to learn the many-body interactions. In these networks, spherical harmonic functions Y (ft;)) combined with radial functions R(Fiy) are usually taken as the single bond basis function (r Then the aggrogated atomic bases Awe typically represented as equivariant irreducible representations in these net works, are combined through tensor product operations to encode the many-body interactions while maintaining equivariance. Note that 2,,, distinguishes among different basis.
A.4 Tensor Contraction in Polynomial Many-Body Interaction Module
In the polynomial many-body interaction module of PACE, we apply tensor contraction different atomic bases for encoding the interaction of multiple nodes around the central node. The tensor contraction which is originally proposed in (Batatia et al. 2022b). However, Figure 2 illustrates the computation of tensor contraction described in Equation [22]. This figure demonstrates an example of 4-body interactions with the max correlation order Umax = 3, where different atomic bases An, Ai2, Aig are fed as inpuls. Intuitively, this example is the computation similar to A, 8 A2 @ Aa + A, @ A2 + As1, where An @ Ag @ Al3 approximates 4-body interactions, An @ A2 approximates 3-body intcractions, and A₁ approximates 2-body interactions. For casier understanding, this example only has hidden channel, and only shows one output LM, that is Lg = 0, M3 0. Generally, the complete final output irreducible representation is the concatenation of all required L3
Let's delve deeper into the "contract weights" box and "contract features" box at the top of this figure. The path 1/33 here denotes a coupled path tuple (41, f2, L22, f3. Lz). The purpose of using such coupled paths is for computing consecutive tensor products as a whole instead of applying tensor products one by one. For each path, there is a generalized CG coefficient matrix with shape of (Emax 2 emaxx mmxx 2, emax' 2 2 x Lg + 1). Note that L3 in this example. The top "contrast weights" box executes the first step to take a weighted sum of the generalized CG matrices with learnable weights over the paths. Then, the "contract features" box below applies a dot product on this summed generalized CG matrix and atomic base Aia along the dimension Lg
Ms denoted as fg. m3. The output of this step is denoted as L2 Then, L2,,,, is further contracted with A2 and An to obtain the final result with similar "contract features" operations. This flow considers the 4-body interactions. To further consider 3-body and 2-body interactions, similar "contract weights" operations are applied over n/22 and n/1]. Then, the results of contracted weights are added to aLa.Ms 2,, and Ma M, (represented by arrows pointing from right to left). The summation of generalized CG and the M.. addition of contracted weights to M,, is for efficient computation of various paths and various body order interactions simultaneously.
As mentioned this example is computing something similar to An X A2z 8 As + An @ A₂ + An for many-body interactions. However, they are not exactly equivalent, because two tensor products in An @ Azz @ As consider the coupled path n with generalized CG, which is beyond the original tensor product.
A.5 Defination
Definition A.I. (D-spanning). For D E N+, let Fyeat be subset. of Co (R33N W.N. feat). We say that Freat is D-spanning, if there cxist f1.... fk E Ffeat: such that every polynomial R - R N of dogree D which is invariant to translations and equivariant to permutations, can be written as p(X) Σ K Åu (fa(X)), where x= A : W feat > R are all linear functionals, and Åg Weatt + R are the functions defined by element-wise applications of Ag.
When discussing the expressiveness of equivariant networks, equívariant polynomial functions are commonly used to analyze these networks (Dym & Maron 2020; Segol & Lipman, 2019). According to Lemma 1 from Dym & Maron (2020), any continuous G-equivariant function can be uniformly approximated on compact sets by G-equivariant polynomials. Consequently, the ability of a model to approximate equivariant functions increases with its capacity to cover more equivariant polynomials. In this context, a D-spanning family refers to a set of functions that span the polynomial function space up to degree D. This implies that any S,, permutation equivariant and rotation invariant polynomial function up to degree D can be represented as
18

---

Published in Transactions on Machine Learning Research (08/2024)
Contract weights
X
Scalar weight
Generalized CG
Contract features
X
Contract weights
Σ
Scalar weight
6/3]
Generalized CG
Contract features
Ag
Contract weights
Σ
X
Generalized
weight
Contract features
Σ
X
Features An
Alass La
Figure 2: Illustration of lensor contraction in the polynomial many-body interaction module. In this figure, it demonstrates an example of 4-body interactions with fmax = 3 and final Lg = 0, M3 = 0. Note that the contract weights operation learns weighted summation over all paths r/v],, where n/v] (E1, l2 Lo, ... lys Lo).
linear combination of elements within the D-spanning family. Therefore, networks that can encompass the D-spanning family with higher D exhibit greater expressiveness.
19

---

Published in Transactions on Machine Learning Research (08/2024)
B Theoretical Proof
B.1 Proof of Theorem
Proof. Since QK is D-spanning family, then there exists fi,... fx 0 with ||t||| D, that each polynomial function with the highest degree no more than D can be represented as the linear combination with linear pooling function AĄ on them
p =
Σ wk * (n(Jk) = Σ wx EuAALQQ
(29)
Σ wk Σ W77k (O.)(P))))
(30)
where P= (p1,p2,... .PK) is the position of the entry, and Wp is the corresponding weight.
Since (0:))() = -ee irreps""""
fm the can be represented as
p
Wp
Σ
fm With pirrepst.a..
fm
Σ k
Σ tmP
etm W L.t, pirreps
(31)
P
tm
where the irrepsfmm.. can be obtained by various channels and =dm bm Wit, W. Therefore, the set of
Wi,t,, P
irreps, is D-spanning.
B.2
Proof of D-spanning irreps in PACE
B.2.1 Lemmas
Lemma B.1. The linear combination of elements for spherical harmonics with radial basis function R (Fis) (F.)) m (fi)), where 0 S E lmax E N and -2 S m S f,m E 2, can represent any polynomial function Pimae (I:z).
Proof. When the radial basis function set to (FTg) = riy, N4, real spherical harmonics with radial basis function R (Fag) Y i (faz) represent homogencous polynomials P(fy) with degree l. Morcover, by using multi-layer perceptrons (MLPs) R2 (Faj), 1 it can approximate any polynomial with respect to the pairwise distance Fij (Hornik, 1991: Hornik et al. 1989). We first define the space of homogeneous polynomials H, (IR",, and space of harmonic homogeneous polynomials Y,, (R*). When we constrain the function to be on $2, Y, = Y(??) and is composed of spherical harmonies. From the Lemma 4.1 in Atkinson & Han 2012) homogeneous polynomial p(x) E H, can be decomposed as p(x) h(x) x (xx), where h W and g TH. n
for 2 2. When considering polynomial functions constrained to $2, have p ( ( + a( { ) where is in the space spanned by spherical harmonics of degree n. The dimension of homogeneous polynomials is [++1](a+2) For n = 1, the dimension of homogeneous polynomials space is 3, and the dimension of spherical harmonics space is also 3. For n the dimension of homogencous polynomial space is and the dimension of spherical harmonic space is 5. Then, with the constrain x²2 + y² 2² 1 demoting the input (x, y, 2) on $ which is independent with spherical harmonics y (f). Therefore, homogeneous polynomials on 8 can be linear combination of spherical harmonics for 2. Thus, ( can be expressed a linear combination of spherical harmonics when H,, 1,2. Then, we deduce that homogeneous polynomial function g( [ can be represented as linear combination of spherical harmonics in S Since a homogeneous polynomial function of degree n can be represented by q(x) = | where g E H,, it can represented by spherical harmonics with radial function to encode the distance. Above all, spherical harmonics with radial basis R (Fi)) (Pig) can represent the polynomial function pf (riz)-
2
Lemma B.2. If irreps, can represent Q*1 (0:), irreps, can represent (+,( and their tensor product output irreps can represent (it..tz((O.)..
20

---

Published in Transactions on Machine Learning Research (08/2024)
Fij 8h for tensor representation format, and the Proof. With the aggregated node features with Q E
value at position P1 = (P11,"' ,Pit) can be represented by linear combination of the elements in irreps shown as
(²)(P:) = Σ
" f irross
(32)
fim
When the tensor product of irreps, and irreps, is irrepss; the representation of irreps is denoted as
irreps (2)
Coma &am, irreps irrepps
irreps Cama
(33)
When the (22( can be linear combination of elements irreps2, then
((t,,)) = E
wf.n irrepsi
fime )
Cimi
w2
irreps
(2m2)
fimi fama irreps, wi wh L
(34)
6, E2m2
Then when the linear combination of the irreps is shown as f2)m
Q' Pp) =
Wg
La
irreps
fa)
(2)mm
Σ
Σ
Chamns
irreps fim inreps? Aam
W3
lams
(a)mm
timi Ezm2
irreps
fimi irreps famz
Σ
Clams domg W3
fa(11.fa)ms
(35)
& Bame
When Wz
fs(f)
EE Eeyms 82m2
Clams Camm
&2mp w
fimi u2
we have fama
Σ
Clams Camilema" W/
f 22ma
Σ
Clama
Σ
Clamy
Cymi wh fifm mm w waam
fama
Σ
fims w
fama
Σ
(trmi fama
Eyra, Pamz
f1mi,
Ea(f)
Σ
w
fim mm
U1
(36)
fama
Therefore,
(P),..t))(P,.P)))
Σ
irreps f1 mi
irreps
wf fimm w fama
41m2.82m2
fimi.aam &amma
Σ
irreps 4im irreps
mym w
Eims fuma
wz
6 8ama
(37)
Above all, the irreps, can represent (2t (P,,
B.2.2 Equivariant Polynomial Functions for Polynomial Symmetric Contraction
Considering tensor product, the path can be represented as n/2] (61, l2, Lz), we have irreps, and irreps, represent Q² (P:) and (P²( with ||t]]] & u, ||t|]]l & v, respectively. Then with Lemma B.2 the output.
21

---

Published in Transactions on Machine Learning Research (08/2024)
irreps can represent ((t,, P2). Note that there might be multiple channels for the same rotation order I2, and we use (f1,,, f2) to distinguish them. Meanwhile, the weighted sum over these equivariant irreducible representations can also achieve representativity.
For the path n/o] = (61, l2 L2, .. ,,: Lu), the irreps can represent Qit to 1), and the irreps, can represent Qtr
irreps3
M,
Σ
Cly Lun, Mo- M..m, irreps 1/-1 My
irreps
(38)
My-1
where L (nju 11. (v) and -1-11 L 2), (1)) Then, in this case, we extend the proof of Lemma B.2 and prove the
Q ( P2, Pp. Po)
Σ L,ma
W3 1.
iirepss
M.M.
Σ
W/
M,
Σ
Cdama The Medoo
irreps
L. Mr
irreps h.m.
L.M.
E Mosam
E
irreps irreps
1/ M
Σ
L. Mo, W3
L.M.
(39)
1 Muluma
1y Mr
Then we take W3
M.M, Ee Ui Mo-i.tmm
Ml
1y Ma whom, and with similar procedure to Equa- wy
tion 37 we can derive that
Q' da Qut aw)
Pa.........
Σ
Mo 'irreps,
1/ M. Lom: irreps t.mo
W1
W2
U, I Molomo dol
Po-1)Qt
(40)
Above all, the output of equivariant base can represent Q(t,, .w)) an) (0.) with ||t, S S Vjj E [o]. Since it can S. v, the constructed D-spanning family Q with K = and =. v. Then, the set of the output equivariant features also a D-spanning family with D =v..
select any (t1, "2, t.)) within [[(t1, t,, "t..",,,,,)))) tt,)ll
Specifically, in the update module the first layer, it uses symmetric contraction with correlation order Umax = 3. Since the aggrogated messages follow the form Qi = Σ, @h with t I 6max * N Nbost rj
with Noosst = and Umax = 3. Then, the updated fcatures after the symmetric contraction is Q% 8t2 Sts
EnnrassE,, r Fij @ Fije @ iijs f1, f2, ty S 6. Therefore, it covers all the possible t with ||t|| S 3, and then represent the D-spanning set with 3.
C Implementation Details
C.1 Data Preprocessing
The raw data and splits for the rMD17 dataset are downloaded from https://figshare, https/saarr.comm/rrrccc////////a taset/Revised MD17 dwtasattee.....mm Batiastt_____1____ rMD17 wttpir/DD211M111 /12672088 The raw data is provided in Numpy .npz format, with energies in units of keal and forces in units of kcal/Å. We convert the units for energies and forces to eV and eV/,, and save the data in .xyz format. The 3BPA and AcAc datascts, already in xyz format with enorgies in eV and forces in eV/A, are directly downloaded from https: http://ggtttub...t /ttp://////ooooa////////eeeeaaa a/tree/main
Next, we extract muclcar charges, conformation coordinates, energies, and forces from the .xyz files and encapsulate all this information into PyTorch Geometric dataset. Subsequently, we perform another pre-processing step to obtain one-hot encoding for each atom and construct molecular graph using radius
22

---

Published in Transactions on Machine Learning Research (08/2024)
cutoff. It is important to note that train a separate model for each molecule. Consequently, the dimension of the one-hot encoding for nodes depends on the total number of atomic types present in the molecule.
Following this, we compute several statistics of the dataset. First, we calculate the average energy using the training set. Then, we determine the mean per-atom energy errors. Second, we calculate the standard deviation of per-atom forces. Third, we compute the average number of neighboring nodes. The processed training dataset, along with the pre-computed statistics, is then used to train our model. This pre-processing workflow primarily follows the baseline method, MACE, without incorporating additional features.
C.2 Experiment Details
For all molecules in, we use 2 GNN layers with 256 hidden channels and fmax 3. For multi-layer perceptions (MLPs), SiLU is used as nonlinearity. During training, we train each model for 5000 epochs using a batch size of 5 and learning rate of 0.01. Besides, we use Adam-AMSGrad optimizer with default parameters of 10-8, and without weight decay. Moreover, exponential moving average with a weight of 0.99 is used in training.
81 0.9, 82 = 0.999,
For rMD17, edges are built for pairwise distances within a radius cutoff. Edge features are initiated by either Exponential Bernstein radial basis functions (Unke et al. 2021b) or bessel functions with a smoothed polynomial cutoff (Gasteiger et al., 2020). Table 5 shows our options for molccules rMD17 dataset. The training loss considers both energy and force with weight ratio of 9: 1000. In addition, we use a rotation order of for the hidden irreducible representations outputted by PACE's polynomial many-body interaction module. For both 3BPA and AcAc, edges are built using radius cutoff of 5 A, and edge features are initiated with 4 Bessel basis functions. The ratio of energy and force in loss is 15:1000. Moreover, the rotation order of the hidden irreducible representations is 2.
Table 5: Model architectural hyperparameters for rMD17.
Aspirin
Azobenzene
Beuzuue
Ethanol
Malonaldekyde
Naplthalene
Paracetamol
Salicyclic acid
Toluene
Uracil
Edge embeddinn
Besel
Bessel
EBRadial
EBRadial
Besel
Besel
EBRadial
Bessel
Bessel
# of basis Radius (A)
4
20
12
8
4
10
4
5
5
5
5
5
5
6
C.3 Pseudocode of PACE
In this subsection, we provide pseudocode for implementing our PACE. The initial node features are embedded based on atomic types, edge directions are encoded by spherical harmonics, and edge lengths are encoded by radial basis functions. These node and edge features are then fed into the first message passing layer. In this layer, an edge booster first applied to produce boosted messages along the edges. These boosted messages are then aggregated around the central node. Next, the aggregated message and initial node features are processed by polynomial many-body interaction module, followed by skip connection with ceff-interaction to update the node features. In the second message passing layer, a tensor product is applied to the updated node features and the boosted messages obtained from the first layer to update edge messages. As in the first layer, edge messages are aggregated and further processed by the polynomial many- interaction module to update the node features. Finally, the total energy and forces acting on atoms are calculated based on the updated node features and atom coordinates.
The first key contribution of our method is the use of an edge booster to increase the number of the digrer-degree polynomial functions that model can approximate. The edge booster is implemented in the first PACE layer using two consecutive tensor products. Initially, tensor product is applicd to the edge spherical harmonics and the initial node features to generate the boosted message with N Whonst = 1. Subsequently, another tensor product is applied to the edge spherical harmonics and the previously obtained boosted message, resulting a boosted message with N boost 2. The final boostod message obtained by concatenating these boosted messages with Nbocst and Nboost 2. This boosted message is used in both PACE layers without repetitive computation, taking computational costs into consideration. The second key
23

---

Published in Transactions on Machine Learning Research (08/2024)
33PA 300K
600K
3BPA 1200K
PNCE
PACE
2500
PACE
Ground-truth
Ground-truth
G Ground-truth
2500
2000
2000
1200
R0R
1500
1000
1000
2000
Distance CA)
Distance LA)
Distance (AI
Figure 3: Illustration of radial distribution functions (RDF) of MD trajectories. Values are averaged over five MD simulations with five initial molecular structures. The shell thickness dr = 0.05 is used.
contribution our method is the incorporation of additional self-interactions in the polynomial many-body interaction module. These self-interactios produce different atomic bases, which serve as inputs for symmet to enhance the model's ability approximate more
contraction. The function of these positions in D-spanning functions.
In our implementation, the ML.P applied to invariant features consists of a 2-layer non-linear transformation with SiLu activation. Detailed information about the tensor product and self-interaction applied to equivariant features can be found in the TFN (Thomas et al. 2018) and c3nn (Geiger & Smidt 2022) libraries. For a more comprehensive explanation of the algorithm used for symmetric contraction, please refer to MACE (atattia
ct al. 2022b).
D Additional Analysis
D.1 Molecular Dynamic Simulation
To further assess the ability of PACE to simulate realistic structures and dynamics, conduct Molecular Dynamics (MD) simulations on the 3BPA dataset using PACE. These simulations are implemented with the Langevin dynamics provided by the ASE library, using 1000 timesteps with 1 Is for each timestep. We randomly select an initial structure from the testing set to generate MD trajectories with PACE. Correspondingly, we extract the ground truth MD trajectory, starting from the same initial structure, from the testing set. The total radial distribution functions (RDF) of the ground truth MD trajectories and PACE-generated MD trajectories are visualized in Figure 3 The results of the MD simulations further a affirm that our PACE model not only achieves state-of-the-art (SOTA) performance in force field prediction but also holds practical value in realistic simulations.
D.2 Algorithm Efficiency
D.2.1 Time Complexity of PACE
Table G shows the time complexity for several key components of existing equivariant model architectures. In this table, C' is the number of channels, L is the maximum rotation order of the equivariant fcatures, and v denotes the correlation order. Our PACE model requires three lensor products (TPs) to obtain the edge features and two polynomial many-body interaction modules on the aggregated node features. Each tensor product has time complexity of OCLL)). Compared to the many-body interaction module used in MACE with a time complexity of OCC2 + CL co+), our polynomial many-body interaction module has time complexity of O(NCT, FCC.1++2), where N denotes the number of nodes, E. denotes the number of edges.
24

---

Published in Transactions on Machine Learning Research (08/2024)
Algorithm 1 Code sketch of PACE network. The PACE architecture comprises embedding layers, two distinct message passing (MPNN) layers, and an output layer. The proposed edge booster is integrated into the first MPNN layer to generate the boosted edge message m In the polynomial many-body interaction module, we utilize additional self-interactions to produce different atomic bases Aie-
Require: Zi (i)), Ps, Tej, E One-hot encoding of node by atomic types Zi, the number of neighbor nodes [W(i)), node positions Pi, edge orientation vector Pij. edge distance Tij, the averaged total energy of the training set E.
function POLNOOHIALMANYYOOVVYEEETTTOOOOOD x!)
A₁ < SELFINTERACTIONN(AA))
> Different atomic bases
4₃2 +- SELFINTERACTION(AA]]
Ag + SELFIINTERACTION(AA))
+ SYAMETHICCOOYTRACTIONN,, A12, A13xx))
return à
end function
function MPNNLAYER 1SSphig. rbJig, xl, x9)
Wi,6 < (MPP(bff))
22.j < -MLP(rbfg) MLP(x(1 x()
m j.1 -Sphij @ aan.
MLP(X x1x
Boosted message with N Nboust = 1.
1.2 -Shhi Swn.y
Boosted message with N Noost 2
m -mmbaa m.).))
Concatenated boosted message
A < [V( EjENTO ml
b Aggregated message
d POLNOMALAAABBOOYYYTRRRAACCOOO
x à! SELFINTERACTION(E))
Skip connection
return m!.
end function
function MPNNLAYER 2(, rbfij, X,, X x!,
W3,ij < MLP(rbfig) + MLP((x) xx))
m m. Sws, x).
Tensor product
A3 Val EjEN()) m
Aggregated message
a3 POLYOUIALAAYYOOYYYYEERRYYYYAAAA
x² t a? + SELFINNTERACTON((S))
Skip connection
return x²
end function
xP + NODELMBEDDING(22)
Embedding of node via atomic type Embodding of edge orientation Embedding of edge length
Sph:, + SPHERACALEAAROOOVCEEEEEENNNCCCC
rbfij RADIALBBSSSEEAARDDDNNNEEE
x m MPPNNLLVER ((ppaa rb fijj X
x + MPNNLAVER MPNNLAVER_2(m) rbfuj. x). x
2
E - E+ (MLP(x)) .m==00) MLP ( (m==00))
Molccule-Ievel energy
f, 0p:
> Atom-level force
D.2.2 Training Time and Memory of PACE
Table presents comparative analysis of training time and memory consumption between proposed PACE method and the bascline methods. Results are reported in seconds per epoch and MB as units. A consistent batch size of 5 is used for each method, with average times calculated over 10 epochs, where validation occurs once every 2 epochs. In these experiments, Allegro configured with 5 layers and a rotation order of 3, and NequIP with 3 layers and rotation order of 3. MACE is configured with layers and a rotation order of 2. PACE is configured as reported in Appendix C.2 The results show that compared MACE, PACE offers enhanced expressiveness and superior performance, albeit with justifiable increase in
25

---

Published in Transactions on Machine Learning Research (08/2024)
Table 6: Time complexity of components used in equivariant model architectures. Here, C is the number of channels, L is the maximum rotation order of the equivariant features, and denotes the correlation order.
Model component
Time Complexity
Self-interaction
O O(CL9)
Tensor product
Polynomial many-body Interaction
+ O CL
MACE many-body interaction
Table 7: Training time and memory consumption, with scc/epoch and MB as units, respectively.
NequlP
Allegro
MACE
Ours
Training time
130
50.3
29.5
45
Paracetamol
Memory
3511
10211
3239
4406
Training time
110
35.8
23.5
35
Toluene
Memory
3507
3595
2540
3394
computational costs. In contrast to NequIP and Allegro, PACE stands out for its combined advantages in both performance and efficiency.
D.2.3 Cost Analysis of Edge Booster
As summarized in Section 4.5 the proposed edge booster (EB) module enhances our PACE model's ability to approximate higher-degree polynomial functions. Instead of integrating the EB module into each PACE layer, we employ it solely in the first layer of our PACE network. The edgeeboosted message generated by the EB is then utilized to replace the original spherical harmonics in every PACE layer. Moreover, the EB module only takes spherical harmonics and atomic types of nodes inputs without considering the updated node features. This design choice aims to minimize the computational cost associated with approximating more higher-degree polynomial functions. Accordingly, in Table g we provide a quantitalive analysis of the computational cost incurred by our edge booster.
This analysis of runtime and memory consumption is conducted using the AcAc dataset. Results are reported in seconds per epoch and megabytes (MB). A batch size of used, and the average time is calculated over 10 epochs, with validation occurring once every epochs. Our PACE model comprises 2 layers. In the first experiment, we remove the EB module and use the original spherical harmonics for tensor products in both PACE layers. The second experiment employs the EB module only in the first layer, which corresponds to the architecture of our PACE. In the third experiment, we add an additional EB module in the second layer. Here, the second-layer EB module uses the updated node features output by the first layer, instead of the node embeddings based on atomic types. The results demonstrate that the EB module with two consecutive tensor product operations is computationally expensive. Compared incorporating the EB module into both layers, the design of using only EB module in the first layer, and without considering updated node features, can elficiently reduce the lime and memory cost by 30.5% and 9.7%, respectively.
D.3 Rotation Order
As described in Appendix C.2 for the rMD17 dataset, PACE utilizes a rotation order of 3 for the hidden irreducible representations outputted by its polynomial many-body interaction module, contrasting with MACE's use of a rotation order of 2 for encoding many-body interactions. Hence, we conduct further experiments to determine if the enhancod performance of PACE is solcly due to this increased rotation order. The results in Table indicate that while the higher rotation order of hidden representations does contribute to improved model performance, it is not the sole factor in PACE's superior performance.
26

---

Published in Transactions on Machine Learning Research (08/2024)
Table 8: Analysis of computational cost. introduced by the edge booster (EB) module on the AcAc dataset.
Training time and memory consumption are reported with scc/epoch and MB as units, respectively.
No EB
EB in first layer EB in both layers
Training time
12.1
17.8
25.6
Memory 2080
3076
3408
Table 9: Results of experiments on rotation order of hidden irreducible representation of polynomial many- body interaction module. PACE uses mmax = for rMD17 dataset. Mcan absolute errors (MAE) are reported for both energy (E) and lorce (F) predictions, with meV and meV/ Á as units, respectively. Bold numbers highlight the best performance.
Ethanol
Toluene
PACE
F
1.8
1.1
PACE (hidden <max 2)
2,1
1.1
D.4 Hyperparameter Exploration
In this subsection, we conduct experiments on the AcAc dataset to investigate our model's sensitivity to different hyperparameter settings. Figure presents the performance curves corresponding to varions hyperparameter choices. From top to bottom, the results pertain four hyperparameters: radius cutoff, number of channels, number of radial basis functions (RBFs), and maximum rotation order. Each data point three runs with different random seeds.
represents the average performance
It is widely believed that the radius cutoff can impact model performance, as it determines the connectivity within the 3D molecular graph. Our results indicate that cutoff of approximately 4-5 Á yields the best performance for our model. Next, the results in the second row indicate that the model performance remains relatively stable with the number of channels ranging from 100 to 300. Having too few channels reduces the model's capacity, while having too many channels does not. provide additional benefits. The results in the third row demonstrate that the model's performance remains relatively stable, with optimal performance occurring when using around radial basis functions. Using few RBFs will reduce the impact of edge distance. In contrast, the last row shows that the maximum rotation order for spherical harmonics and irreducible representations has significant impact on model performance, which is consistent with trends observed in previous studies (Batzner et al. 2022 Batatia et al. 2022b). However, it is important to note that higher rotation orders introduce greater computational costs. Consequently, existing models typically use a maximum rotation order of three. The experimental results shown in Figure
f demonstrate that our PACE model exhibits reasonable stability and robustness to variations hyperparameter choices.
E
Examples
Considering [N|| = 2, ri,j=1 01,,812211 21) and r,j = (x2, U2; 22), then
vî + vī
11 + 1/222
71 + 1222
fx. Σ
(fg) -C
219/1 + 228/2
2 22
21 + 2252
(41)
JEN.
213/1 + Z23/2
Ti + 22
V1 + 3/2
1/1 + 3/2
fx Σ yt=l((ng) 8 Σ y(=)(An))))) Ct= JIEN
21 + 22
3
21 + 2/
(42)
JEN.
I + 322
3 + 322
For fx. it has clement yi + y2. while fx only contains elements like yf y2 + 2y1 Thus, fx. can not be represented linear combination of elements in Lx.-
27

---

Published in Transactions on Machine Learning Research (08/2024)
RMSS: 30601/11
NUSS Incok
1962
19022
Nalko
300K, E
300K
000K
BOOK
13
1.-
1962 (m8))
0.0
RASE 3mel:/A
BA52 (mea))
300
302
100
200
302
300
130
290
630
200
400
900
600
chamnees
300K, E
300K
60OK, E
600K
13-
0.0
RVSE mneerAA
20-
300K_E
300K_F
600K_E
600K_F
RMSE meto (Ab
D00LE (mee))
Figure 4: Exploration of hyperparameter space.
28


Reviews:
[{"review_id": "NV2v1fBFeU", "reviewer": "Reviewer_h3gn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The introduction of the PACE network, which integrates edge booster and Atomic Cluster Expansion (ACE) techniques, represents a novel approach in the field of molecular dynamics simulations.\n\n2. The paper provides an in-depth analysis of equivariant polynomial functions, transitioning from point cloud networks to equivariant networks, which adds valuable theoretical insights.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This submission introduces a novel equivariant network called PACE, which aims to improve the prediction of atomic energy and force fields in molecular dynamics simulations by integrating edge booster and Atomic Cluster Expansion (ACE) techniques. PACE extends the analysis of equivariant polynomial functions, approximating high-degree SE(3) × Sn equivariant polynomial functions. It demonstrates state-of-the-art performance and generalization capabilities on datasets like rMD17, 3BPA, and AcAc. By leveraging spherical harmonics and tensor products, PACE effectively incorporates rotational and permutational equivariance, significantly enhancing the prediction accuracy of atomic potentials and force fields.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "1. The paper would benefit from the inclusion of an algorithm table that outlines the steps of the PACE network. This would help clarify the implementation details and make it easier for readers to understand and reproduce the methodology.\n\n2. The authors need to discuss the proposed network's time complexity. Analyzing the computational complexity of PACE, not only empirical evaluations but also theoretical discussions, would provide valuable insights into its efficiency and practicality for large-scale applications.\n\n3. The paper currently lacks a thorough exploration of the hyperparameter space. Investigating the sensitivity of the model to different hyperparameter settings could provide insights into the stability and robustness of the network.", "broader_impact_concerns": "Not involved.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WbL9z2W2vi", "reviewer": "Reviewer_dhJf", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- This paper is well-written and well-organized.\n- The experimental results are good, which surpass the baselines by a large margin.\n\nWeaknesses:\n- Though using polynomial approximation and ACE is innovative, what makes PACE shows better performance than existing methods is still not clearly explained. The claims of achieving state-of-the-art performance are not strongly justified. It is suggested to provide clear intuition about why PACE differs from or improves upon previous works.\n- The theoretical justifications provided for the use of polynomial approximations and spherical harmonics are not adequately rigorous. The paper could benefit from a more detailed and formal mathematical exposition. Additionally, some key concepts, such as the D-spanning function sets and their relevance to the model's expressiveness, are not fully explained, leaving gaps in the reader's understanding\n- While the paper claims strong generalization capabilities in out-of-distribution settings, the evidence provided is insufficient. The experiments lack diversity in molecular complexity and size, which raises questions about the model's robustness and applicability to a wider range of molecular systems. More diverse and challenging test cases should be included to convincingly demonstrate the effectiveness of out-of-distribution generalization.\n- The paper does not provide enough information to ensure reproducibility. Details on the implementation, such as specific software versions, data preprocessing steps, and code availability, are not adequately provided.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Summary:\nThis paper studies the equivariant graph network approximation, which largely relies on the effectiveness of spherical harmonics (SH) and tensor products (TP). This paper studies the equivariant polynomial functions for the equivariant architecture by proposing a novel framework called PACE which utilizes edge booster and Atomic Cluster Expansion (ACE) to approximate equivariant polynomial functions. By conducting extensive experimental analysis, the learning performance and generalization effectiveness of PACE is carefully validated on typical benchmark datasets in predicting atomic energy and force fields.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the weakness part.", "broader_impact_concerns": "No ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0Ks42EJubF", "reviewer": "Reviewer_QPwE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. This paper is well-written with clear logic plus theoretical and empirical analysis.\n2. The authors innovatively utilize Eq. (12) to (16) to construct their edge booster, which considers only spherical harmonics and atomic types without involving node features. Such design cleverly increases reusability and reduces computational costs.\n3. The inclusion of additional self-interaction layers deepens the network's symmetry comprehension and expands its capacity to approximate a broader array of positions in D-spanning functions. This innovation is both novel and practical.\n4. The authors present thorough empirical studies, the results of which are clearly SOTA.\n\n## Weaknesses\n\nThis work is generally neatly presented. I only have several minor suggestions.\n\n1. The authors discuss computational cost reductions in Sec.3.5. I recommend providing quantitative details on these reductions and any associated efficiency gains.\n2. Chap.4 related works could be positioned before the introduction of the PACE architecture. This change would help readers better appreciate the advancements made with PACE.\n3. I am confused by the color scheme in Fig.1. For example, the \"Self-interaction\" block is in dark blue in part A and B while light purple in part D. However, the block \"Many-body interaction\" is in the same light purple color in part B. Please consider a unified color scheme.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel equivariant network named PACE to better predict atomic potentials and force fields. PACE is based on polynomial function approximation and Atomic Cluster Expansion (ACE). First, it incorporates edge booster in the first layer to build edge messages from neighboring nodes. Next, the additional self-interaction (SI)layers are introduced in the polynomial interaction module. Empirical results validate the network's efficacy, demonstrating its robust performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to *Weaknesses*.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the equivariant graph network approximation, which largely relies on the effectiveness of spherical harmonics (SH) and tensor products (TP). This paper studies the equivariant polynomial functions for the equivariant architecture by proposing a novel framework called PACE which utilizes edge booster and Atomic Cluster Expansion (ACE) to approximate equivariant polynomial functions. By conducting extensive experimental analysis, the learning performance and generalization effectiveness of PACE is carefully validated on typical benchmark datasets in predicting atomic energy and force fields. This paper is well-written and well-organized. The experimental results are good, which surpass the baselines by a large margin. The authors address reviewers' concerns well in their rebuttal. Thus, I would like to recommend accept as is.
Decision: Accept (poster)