FreeCG: Free the Design Space of Clebsch-Gordan Transform for Machine Learning Force Fields
Authors: Shihao Shao
Haoran Geng
Zun Wang
Qinghua Cui
Venue: ICLR 2025 Poster (2025)
Keywords: Machine Learning Force Fields
Equivariant Graph Neural Network
Clebsch-Gordan Transform
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present a novel SoTA MLFFs model with high practicality through freeing the design space of CG transform, together with a new paradigm for guiding the future designs of EGNNs.

Abstract: Machine Learning Force Fields (MLFFs) are of great importance for chemistry, physics, materials science, and many other related fields. The Clebsch–Gordan transform (CG transform) effectively encodes many-body interactions and is thus an important building block for many models of MLFFs. However, the permutation-equivariance requirement of MLFFs limits the design space of CG transform, that is, intensive CG transform has to be conducted for each neighboring edge and the operations should be performed in the same manner for all edges. Freeing up the design space can greatly improve the model's expressiveness while simultaneously decreasing computational demands. To reach this goal, we utilize a mathematical proposition, invariance transitivity, to show that implementing the CG transform layer on the permutation-invariant abstract edges allows complete freedom in the design of the layer without compromising the overall permutation equivariance. Developing on this free design space, we further propose group CG transform with sparse path, abstract edges shuffling, and attention enhancer to form a powerful and efficient CG transform layer. Our method, known as FreeCG, achieves state-of-the-art (SOTA) results in force prediction for MD17, rMD17, MD22, and is well extended to property prediction in QM9 datasets with several improvements greater than 15% and the maximum beyond 20%. The extensive real-world applications showcase high practicality. FreeCG introduces a novel paradigm for carrying out efficient and expressive CG transform in future geometric network designs. To demonstrate this, the recent SOTA, QuinNet, is also enhanced under our paradigm. Code: https://github.com/ShihaoShao-GH/FreeCG.

Published as a conference paper at ICLR 2025
FREECG: FREE THE DESIGN SPACE OF CLEB- SCH-GORDAN TRANSFORM FOR MACHINE LEARNING FORCE FIELDS
Shihao Shao"! Haoran Geng? Zun Wang & Qinghua Cui
ABSTRACT
Machine Learning Force Fields (MLFFs) are of great importance for chemistry, physics, materials science, and many other related fields. The Clebsch-Gordan transform (CG transform) effectively encodes many-body interactions and is thus an important building block for many models of MLFFs. However, the permutation-qquivariaee requirement of MLFFs limits the design space of CG transform, that is, intensive CG transform has to conducted for each neighbor- ing edge and the operations should be performed in the same manner for all edges. Freeing up the design space can greatly improve the model's expressiveness while simultaneously decreasing computational demands. To reach this goal, we utilize a mathematical proposition, invariance transitivity, to show that implementing the abstract edges allows complete
CG transform layer on the permutation-invattee freedom in the design of the layer without compromising the overall permutation equivariance. Developing on this free design space, we further propose group CG transform with sparse path, abstract edges shuffling, and attention enhancer to form a powerful and efficient CG transform layer. Our method, known FreeCG. achieves state-of-the-ar (SOTA) results in force prediction for MD17, rMD17, MD22, and is well extended to property prediction in QM9 datasets with several improvements greater than 15% and the maximum beyond 20%. The extensive real-world applications showcase high practicality. FreeCG introduces a novel paradigm for carrying out efficient and expressive CG transform in future geomet- ric network designs. To demonstrate this, the recent SOTA, QuinNet, also en- hanced under our paradigm. Code: http://gthhbbboomo/h//hhhaaaaaap
INTRODUCTION
Machine Learning Force Fields (MLFFs) are of great importance for drug development (Chen al.
2024). materials science (Liu et al. 2024), chemical reaction kinetics (Meuwly) 2021). nanotech-
nology (Wang et al.) 2023b). among others. It offers satisfactory trade-off between accuracy and efficiency, which is expected to perform as powerful as Density Functional Theory (DFT) (Kohn &
Sham 1965) or other high accuracy references (Martin 2020 Ceperley & Alder 1980 Bartlett &
Musia
2007). but with
speedup (Cui et al.] 2024 Wang et al. 2023c 2024]
Musaclaa et al. 2023 Batzner et 2022
Drautz 2019 Batatia et al. 2022b Thölke & De Fab.
nttis 2021 Schütt et al. 2018 Chmiela et al. 2017
Graph Neural Networks (GNNs) perform SOTA on several MLFFs benchmarks (Schütt al. 2018 Group and group representation theory play important roles in the design of GNNs for MLFFs. For instance, rotation invariance generally required these works, as we naturally re- quire the potential energy unchanged w.r.t. rotations of the input molecule. A recent design trend is to maintain rotation, reflection, and translation equivariance the design of neural networks. They hope the internal features move with respect to the input molecule, enabling higher expres- sive power. GNNs that obey this property are called Equivariant Graph Neural Networks (EGNNs)
2021).
"Conesponding authors (shaoshihao www.@pkuuuuccrooee cuiqinghua @ hssc disc.pku.edu.co) School of Basic Medical Sciences, Peking University EECS, UC Berkeley AA4SSiiece,, Mierosoft

---

Published as a conference paper at ICLR 2025
Thomas et al., [2018] Satorras et al. 2021 Gasteiger al. 2020b; Vaswani al. 2017 Fuchs et al.
2020 Liao & Smidt 2022). To better model many-body interactions, irreducible representations (irreps) are adopted to represent high-order geometric objects. In this context, the Clebsch-Gordan
(CG) transform is used to translate between different irreps. Several works leverage such high de- gree irreps tensors, showing significant performance boost (Batatia al., 2022bb |Bazner al.
2022 Musaclian et al. 2023; Gasteiger et al.] 2021] Thomas et al.) 2018 Simeon & De Fabritiis
2024 However, the benefit of high degree irreps and CG transform performed on them is at the cost of heavy computational overhead. The reason is, tensors are extensions of scalars and vec- tors, and in this way CG transform also extends the dot product. Thus, the higher the degree of irreps for the CG transform, the greater the computational demands. The requirements for being permutation equivariant make this burden hard to alleviate. Unlike rotation or translation equivari- ance, permutation equivariance is often implicitly guaranteed in EGNNs, which means the order of internal features should changes according to it of input atoms. To maintain permutation equiv- ariance, EGNNs require each node to receive information from neighboring atoms together with the edges linking them, where the heavy computation of CG transform occurs for each neighboring atom and edge. This means cannot naively remove some neighbor computations, as it will break permutation equivariance. Moreover, the narrowness of the design space prevents us from freely constructing the CG transform layer, and thus limits the expressivity of models. For instance, we need to operate on each neighboring atom in an equal way (e.g., the predecessors typically assign a same MLP operating on scalar features of the edge produce the weights for each computation between the central atom and each neighboring one (Batzner et al. 2022 Musaelian et al. 2023)
In this work, to confront this challenge, we propose FreeCG. The model generates and refines ge- ometric features from the surrounding edges near each atom. We call the different aggregated edge geometric features abstract edges, which are permutation invariant when we consider the internal features maintained by a given atom w.r.t. the neighbouring atoms and edges. By the invariance transitivity, we show that CG transform on these abstract edges also permutation invariant, re- gardless of designs, and does not affecc the permutation equivariance of the layer, thus being free of the burdens above. Furthermore, the abstract edges are constructed from different real edges, so they contain refined features of them for better model expressive power. The invariance nature of abstract edges allows us to assign different weights to different edges, instead of weights computed by the same MLP. The free design space allows us do more. We put abstract edges into groups, and operate on each group individually, to further decrease the computation demands. Previous works that keep E(3)-equivariance are more expensive (Batzner al., 2022 Musaelian al. 2023). as they require an extra parity argument being 1 or -1. and thus the number of irreps is doubled. Instead, we select an efficient set of paths for CG transform so that we maintain E(3)-qquivariance while being more efficient than keeping SE(3) -quivaraave The abstract edges shuffling, inspired by (Zhang et al., 2018), also implied for combination of irreps features. The abstract edges are then plugged back into the cross-attention calculation to improve the quality of the attention scores. The operations above are available thanks to the invariance properties of the abstract edges.
The contributions are summarized as follows:
1) We utilize the invariance transitivity with permutation-invarian abstract edges to resolve major challenge in the current EGNNs: the narrowness design space of CG transform, which results in reduced expressivity and high computation overhead.
2) We propose FreeCG, comprising of three main components: Group CG transform with sparse path, abstract edges shuffling, and Attention enhancer. These contribute to an informative and efficient model with high-order irreps and CG transform.
3) Experiments on small molecule datasets MD17, rMD17, large molecules ones MD22, and molec- ular property datasets QM9 reveal the SOTA performance of FreeCG, with several improvements beyond 15%. The extensive real-world applications further indicate the practicality.
4) This work presents new paradigm for CG transform in future research, extending beyond the design presented here. We further enhance QuinNet (Wang et al. 2023c) under this paradigm to demonstrate this point.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORKS
has become a popular trend in the design of neural networks for MLFFs. In general, the methods can be categorized into two lines:
Maintaining
Methods with geometric vectors. Several methods directly work on regular geometric vectors (Schütt et al. 2021: 2018) Gasteiger et al. 2020b[a). The starting point of this line of works about covering necessary information for MLFFs. For instance. SchNet (Schütt et al., 2018) first intro- duces continuous- filter to MLFFs, but only distance information is covered. DimeNet (Gasteiger et al. 2020b) further considers bond angle information to obtain higher capacity. To make models capable of capturing many-body interactions, most of the following models in this line explicitly encode such types of interactions (Thölke & De Fabritiis, 2021] Wang al.] 2024 (023c). This is commonly done via calculating different angles between atoms or surfaces (e.g., torsion angle, improper angle (Wang et al. 2024), dihedral angle (Wang al.) 2023c)). The irreducible repre- sentations (irreps) and Seebsch-Gordan (CG) transform, on the other hand, can implicitly encode many-body interactions in well-defined mathematical objects. This work presents FreeCG with highly efficient and expressive CG transform layers, significantly outperforming previous works by maximum margins and with minor computational load.
Applying irreps and CG transform to Equiv- ariant Graph Neural Networks (EGNNs) was first proposed as more of a conceptual framework, known Tensor Field Network (Thomas et al. 20118) Several works have been proposed on top of this foundation. For instance, NequlP (Batzner et al. 2022) has implemented the idea to con- struct high-order irreps, and shows state-of-the-aa (SOTA) results for MLFFs. Allegro (Musaclian cet al. 2023) resolves the challenges of the scaling issue, which makes it possible to run parallelly on large number of GPUs Furthermore, Allegro constructs high order geometric features with pair-wise edge CG transform, which is similiar to FreeCG at this point. However, Allegro requires to compute the same CG transform for each pair of edges in the same manner to maintain permu- tation equivariance, while FreeCG is capable of designing different CG transform for each pair of abstract edges because of their permutation invariance. The atom-centric style and the rich rep- resentation of abstract edges together makes stronger expressivity of FreeCG.
Methods equipped with irreps and CG transform.
(uuhhs et al., 2020) first proposes equivariant dot-product for generating self-attention. Equiformer [iao & Smidt 2022) further achieves equivvariee combining MIP attention and non-linear messages. MACE (Batatia et al. 20226 extends classic body order expansion methods, Atomic Cluster Expansion (ACE) (Drautz) 2019). to hierarchical framework. CG transform is funda- mental building block in these works, but with a limited design space, affecting both performance and efficiency. In this work, completely free the design space of CG transform and propose novel model, FreeCG, performing strong SOTA and showing high efficiency for MLFFs, which is also a new design paradigm for future works.
3 METHODS
3.1 BACKGROUND
Group, equivariance and invariance. Permutation, rotation, and translation form different groups
in group theory. Formally, set with binary operation (G, +) is said to be a group if and only if the following conditions hold: 1) * 92 E G, for any 91, 92 E G (closure) 2) (g) * g2) * ga = gi (92 ?ga), for any g 92, 93 E G (associativityy) 3) There exists group element e E G. such that g*te = e*g g, for any E G. (e identity element) 4) There is a group element g' w.r.f. g, such that g "g' = g' * g = e, for each g G (g' inverse element). The group elements g E G, according to the representation theory, can be represented as linear transformations Pv(gg) E GL(V) on vector space V. Given function f X -> Y, where X and Y are vector spaces. equivariant if and only if (g)xx) = (g) for any 9 E G. G-invariance is special case when Py (g) is an identity matrix. Permutation equivariance and B(3) e(u)-aquiva are two properties each layer of our model obeys. Permutation equivariance means the index of node edge features will be consistent when passing a layer. E(3) equiivarivee covers rotation, translation, and reflection, where the translation is explicitly guaranteed via only considering the relative distances between atoms, thus we consider ((3)-equivariance where translations are omitted. It is intuitive to correspondingly change directional features when the whole molecule rotates reflects.
said to be G-
3

---

Published as a conference paper at ICLR 2025
f
R;
fu
FTNI
daj
103
120
100
To
group DG transform
sporse path
Amonnt
100
f1y
oigU oupale
shuffle
E
df.
100
100
abstract edges shuff
det amotuct
digg
VisNet
CB vandfor with path
The architecture of single layer of FreeCG. The cross-attention mechanism generates abstract edges through a inouuutaionninvariin process. The abstract edges are also used to enhance the quality of the attention score, denoted Attention Enhancer. In the right part, the Group CG transform organizes abstract edges into groups and performs the CG transform on each group. We adopt sparse path for CG transform, enabling lower computation demands while maintaining 0(3) equivariance. Abstract edges shuffling improves the information exchange between different irreps. Details for sparse path and abstract edges shuffling can be referred to Fig. Better viewed in color.
Figure 1:
Tensor, irreps and CG transform. Tensors are high-dimensional generalizations of scalars, vec- tors, and matrices. Scalars and vectors are both special cases of Cartesian tensors. Tensor product can generate high-rank tensors from low-rank ones. Formally, tensors are the results of tensor prod- uct of several vectors and covectors. In our context, it is not essential to distinguish between vectors and covectors. Tensors representing groups can be further decomposed to the direct sum of irreps. For example, tensors of SO(3) (omit reflection compared O(3)) 9-space (from tensor product of two 3 x 3 rotation matrix) can be decomposed into x 1 (1 = 0), 3 x 3 (1 = 1), and x called Wigner-D matrices. In EGNNs,
(1= 2) irreps, which between atoms onto the unit sphere 5 with the central atom as the center of sphere. Actually, 52 is homomorphic the quotient group (0(3)/SO(2) thus it also has its own irreps, e.g., 0 scalar and = 1 vector. S irreps are the main features we maintain in our model, where irreps with degree 7 has 21 elements, which are often indexed by m. To combine these features, we can calculate the tensor product between them, and the results can, again, be decomposed to irreps. This process known as CG transform, which utilizes CG coefficients perform taansfomattios. For instance, A 11/211+11 = E Cmimana 42. A3.1s, where A' are S irreps, denotes the elements of ir- reps, and C the CG coefficient. To satisfy O(3), we consider an additional variable, parity p. which takes the values or -1. Irreps with p = will be inverse when the space reflected, and p unchanged. The above formula of CG transform becomes:
often project the distance vector
A A.ppalapo =1(p:=pppa)
Σ
Chilas 2ddaps 33,9ps mimggm A3,1ps
ma m9
where ((espression) is the indicator function, outputting if expression is true, and 0 otherwise.
Given a vector (1 - S irreps), we can lift it to irreps with arbitrary degree I and = (-1).. via series of real spherical harmonics (Y'-=,,..,, yt 1 ). For further details about group theory,
we refer interested readers to related books and papers (Zee 2016 Raczka & Barut 1986 Thomas
al. 2018 Jeevanjee 2011 Cohen et al. 2018).
4

---

where ((a) is the set of neighboring atoms of atom a. R a MLP. ||*|| is Euclidean norm. No is the features of node b. Fab is the vector pointing from atom a to b. Consider the vec- tor function form of Eq. 2 C Clapeaaprooooo permutation equivariant w.r.t. permutation operations acting on X and Z. Formally, it means PC PePepaaaaot.. (X, Z). Put simply, if we exchange the in- dexes of two atoms, for example, and 2, and feed them into function Lem. claoeoaa-olaoo it equals Cl which is Published as a conference paper at ICLR 2025
p=1-
1
2041
Linear(f))
p=1
|=1 p=-1
8
8
Irdex
p=1
p=1.p=-1
|=2, p=1
d
Figure 2: Details on sparse path and abstract edges shuffling. Left: The sparse path holds two useful
properties: 1) The number of paths is less than the weaker SO(3) equivariance (4 vs. 8). 2) Each
output irreps contains the information from input ones with both degree / = 1 and I = 2. Right: The shuffling strategy to add a constant & for the index of each abstract edge. The shuffled result
is then added by E/, and get the final added value dE, 1 Better viewed in color.
3.2 PROBLEM ANALYSIS
The task of force field prediction can be formalised as follows: Given a set of atoms with their positions and atom types {XX 2). the neural network fo with parameter I aims to predict the energy, and by which it derives the predicted force on each atom. In each layer of NequlP (Batzner et al. 2022). messages from neighboring atoms central atom. The messages are created via CG transform between the irreps. Here, we revisit the critical step constructing messages to a central atom a in NequlP:
aggregated and combined with the features of the
P..pehaan-taa (X,N) 11 (po=p.PP) gemo
Σ mc
Chataha Momemma
(2)
(R((IFabll) Ioaatan Jy BEA((a
Σ
Fab
Nrhaaa
Fab
bemn
(ChPPetPP Pt.Po pleDetnDa-ttal ,...). which is _L2em,
Poo (Px P2Pz) ena
to that we directly change the index and 2 of the output of function cm.
(cl.pe.....oap. clapollp..too.. ....). This property is simple and very important for the molecular Lemm. neural networks, as the indexes of atomic features should change correspondingly w.r.t. the changes of atom indexes, and the feature values for each atom should stay unchanged.
2emm
Most works take this property for granted. However, the permutation equivariance is actually im- portant but vulnerable. limits the design space to a very small scope, and make the network poorly scalable when the number of neighbors arises. Specifically, it brings the following issues:
Problem 1 The CG transform layer scales OCE, (i)d(N(i))), where card(X) is the number of elements set X. One cannot arbitrurily remove calculations for a specific neighboring atom because would break the permutation equivariance.
Problem 2 The design space is limited for maintaining permutation equivariance. For example, in Eq. 2 the formulation and the parameters of R should be the same across different neighboring atoms, thus forbidding the design for complicated CG transform layers.
Problem poses heavy computation challenges, the CG transform itself is very time-consuming, compared to dot product and element- wise multiplication. We provide a detailed analysis for the efficiency of CG transform in Sec. A.6 On the other hand, the narrowness for design space brought by problem 2 makes it hard to design a high expressive CG transform layer, as only limited struc- tures can be designed maintain permutation equivariance. To address these problems, aim to
5

---

Published as a conference paper at ICLR 2025
free the CG transform in messages transmissions from the constraints of permutation equivariance without compromising the overall equivariance of the network. Here, we leverage a simple and useful mathematical property. Consider a function h that can be written as:
M(x) h (I11
(3)
if (x:) are all G-invariant, then, regardless of how design h', the overall function h must be G-invariant as well. The proof is simple, as:
h (1e (g)xr. h2 (Px(())))) Pa (11 (g())))
(4)
Here, the invariance study is about the internal features of given atom w.r.t. the neighbours. Specifically, is the sum (ve)(a)) (R( (l) Yr ( Fub N and the term for each b. (R(
Faoo
Such invariant components guarantee the equivariance of the layer mentioned above. Thus, we can freely design the function h once we have these invariant functions hT,
3.3 FREECG
Abstract edges. The above proposition presents an elegant way to solve problem Looblem[ande and put CG transform inside the function h and by the conclusion, we completely free the design space of the CG transform. The first step is to construct the permutation invariant function h,, To emphasize the geometric information, we want these h, to be the aggregation of edge features. We call h., 's abstract edges. For the concrete design, we take the transformer architecture in ViSNet efficient tool construct abstract edges. The detailed information of the complete implementation is in the Sec.
The idea is
Wang et al. 2024)
A.2 In ViSNet, each edge maintains high-degree features Y' (1) The above features are invariant to is multiplied to each edge. The sum of them E/;, Enjer(o) a Faj forms a temporary abstract edge, where we omit the degree and I, the index of the layer. t denotes the index of the t-th abstract edges (É, ( In the original ViSNet,
Egg ij E consisting of irreps
layer index L. The computed attention Gij
L,C
was used to update the geometric feature dE = E + L Linear(E where Linear is fully-connected linear operation, which performs across the dimension of t, thus does not break the equivariance, and E represents propagated abstract edges in contrast to the temporary one É, where the former updated by dE in each layer, derived from the temporary abstract edge É. o L,1 is variable generated from the node feature, as we will introduce in Sec. A.2 We utilize that RL and F. fits the requirements of h. our proposition, take them as abstract edges, and propose methods to construct CG transform function h upon it. The proof that each abstract edge meets the requirement for ht, namely, it is permutation invariant, is in Sec. A.4
Group CG transform. The number of abstract edges is decided by us, so the complexity for computing CG transform Problem [ controlled to be constant. The proposition above gives us complete freedom to construct the CG transform function hi, expanding the design space to maximum, alleviating Problem The idea is to use CG transform function to replace the updating mechanism of E in ViSNet. A naïve attempt is to directly take the CG transform between E' and EL to acquire =++ transform, where 7' is the number of abstract edges, even though is constant number. Leveraging the unlimited freedom in constructing h, and taking inspiration of group convolution (Krizhevsky al. 2012), we propose group CG transform (distinet from the group group theory). We first split the abstract edges of E and ÉL into groups, where each index of abstract edge belongs to some group Uy the integer g ranges from to G. and groups. Then group CG transform acts as:
However, we want to further decrease the ( time complexity for the CG
for the number total
Notldoso d.
Σ Σ
dlodid momimma
i,tomo
hodo mi
(5)
Σ hu.tEU
Wood tottl
Tinear
Iyh.Ps AL. Piagmg

---

Published as a conference paper at ICLR 2025
Table 1: Performances on MD17 dataset. The results reported in mean abosolute enror (MAE). The energies and forces are measured in kcal/mol and kcal/mol/A, respectively. The best numbers marked in bold.
Molecule
SchNet
DimeNet
PaiNN
SponkeyNet
ET
GemNet
NequIP
SOSKRATES
ViSNet
QuinNet
FreeCG
Energy Prediction
Aspirin
0.37
0.204
0.167
0.151
0.123
0.131
0.139
0.116
0.119
0.110
Fthanol
0.08
0.064
0.064
0.052
0.052
0.051
0.052
0.051
0.050
0.049
Malondiakdehyde
0.13
0.104
0.091
0.079
0.077
0.076
0.077
0.075
0.078
0.094
Naphthalene
0.16
0.122
0.116
0.116
0.085
0.113
0.115
0.085
0.101
0.083
Salicylic acid
0.20
0.134
0.116
0.114
0.093
0.106
0.016
0.092
0.101
0.090
Toloene
0.12
0.102
0.095
0.094
0.074
0,092
0.095
0.074
0.080
Uracil
0.14
0.115
0.095
0.104
0.103
0.095
Prediction
Aspirin
1.35
0.499
0.338
0.258
0.253
0.217
0.184
0.236
0.155
0.145
0.122
Ethunol
0.39
0.230
0.224
0.094
0.109
0.085
0.071
0.096
0.060
0.060
0.053
Malondiakdehydee
0.66
0.383
0.319
0.167
0.169
0.155
0.129
0.147
0.100
0.097
0.095
Naphthalene
0.58
0.215
0.077
0.089
0.061
0.051
0.039
0.074
0.039
0.039
0.034
Salicylic acid
0.85
0.374
0.195
0.180
0.129
0.125
0.090
0.145
0.084
0.080
0.070
Toluene
0.57
0.216
0.094
0.087
0.067
0.060
0.046
0.073
0.039
0.039
0.035
Uracil
0.56
0.301
0.139
0.119
0.095
0.097
0.076
0.111
0.062
0.062
0.059
where Ug. The group CG transform decreases the time complexity OGT² Here, the parameters W for CG transform are also worth emphasizing. They are not necessary to be kept the same across different abstract edges & to keep permutation equivariance, and do not need to adopt the same MLP for each edge to calculate weights. Thus, we directly assign different weights W for different abstract edges simplify the model design. In contrast previous methods, We save the computational cost for calculating weights for each edge.
Sparse path. Typically, ensuring SO(3) equivariance is considered more efficent than ensuring 0(3) equivariance. It is because we often need to consider both p = 1 and = -1 for a single I for 0(3), thus the total computation is quadrupled, and memory usage doubled. Here we propose method to keep O(3) while being more efficient than SO(3). We only keep (7 1, = -1) and (1 = 2,p = 1). which is same as the order of directly using spherical harmonics. In such way, sufficess that each output irreps containing information from both input irreps through CG transform, as illustrated in Fig. 2 There are only 4 path in contrast to 8 path for SO(3), being 0(3) equivariant but more efficient than being SO(3) equivariant.
A Abstract edges shuffling. Inspired by ShuffleNet (Zhang et al. 2018), we can also shuffle the ab- stract edges to make the information exchanged comprehensively. If we omit the shuffling operation, the information in the abstract edges would only be exchanged within the group, which decreases the capacity of the model. We shuffle all the abstract edges. Specifically, we increase the indices of all irreps by [1.5 T/G]. If the index exceeds T, we start counting from again. Theoretically, the shuffling strategy can be arbitrary long as maintaining the same strategy for each layer during every inference. This process is also depicted in Fig. 2 The ablation on different strategies is shown in Sec. 4.4
Abstract edges enhance cross-attention. The transformer integrates neighboring atoms informa- tion in the model through cross-attention mechanism, which aims to capture relations for those atoms exhibiting strong interatomic correlations. Thanks to the informative abstract edge, we utilize it to augment the generation of attention scores. the original design, to calculate the cross-attention, the node scalar features are processed to generate query Q, key K, and value V for each atom, respectively. Then, the self attention is computed as a;j = SiLU (SutofCtfff((] kkkk)) Note that ViSNet is different from other transformer based models where azj is scaled by the SiLU instead of Softmax across different j. We integrate the information of abstract edges by:
ag SiLU ( + AttEnhancer (Esg. E3))
(6)
where AttEnhancer (E.j. E,) max, CE,, @ Egj.. that is, for each real edge feature Eij, compute the dot product with all abstract edges F, and take the maximum value across different
7

---

Published as a conference paper at ICLR 2025
Table 2: Performances on rMD17 dataset. The results are reported in MAF. The energies and forces are measured in kcal/mol and kcal/mol/A, respectively. The best numbers are marked in bold.
Molecule
UNiTE
GemNet NequlP
MACE
Allegro
BOTNet
ViSNet
QuinNet
ACE"
PONITA
FreeCG
Energy Prediction
Aspirin
0.055
0.0530
0.0507
0.0530
0.0530
0.0445
0.0486
0.0:992
0.0392
0.0530
Azobenzene
0.025
0.0161
0.0277
0.0277
0.0161
0.0156
0.0394
0.0115
0.0161
0.0217
Benzene
0.002
0.0009
0.0092
0.0069
0.0007
0.0007
0.0096
0.0002
0.0039
0.0107
Ethamol
0.014
0.0092
0.0032
0.0092
0.0092
0.0078
0.0096
0.0046
0.0592
0.0087
Malonaldehyde
0.025
0.0184
0.0185
0.0138
0.185
0,0132
0.0168
0.0115
0.0138
0.0146
Naphthalene
0.011
0.0046
0.1153
0.0046
0.0046
0.0057
0.0174
0.0046
0.0069
0.0118
Paracetamol
0.044
0.0323
0.0300
0.0346
0.0300
0.0258
0.0362
0.0208
0.0254
0.0392
Sulieylic acid
0,017
0.0161
0,0208
0.0208
0,0185
0,0161
0.033
0.0115
0.0161
0,0233
Toluene
0.010
0.0069
0.0115
0.0092
0.0069
0.0059
0.0139
0.0046
0.0069
0.0334
Uracil
0.013
0.0092
0.0115
0.0138
0.0092
0.0069
0.0149
0.0046
0.0092
0,0116
Force Prediction
Aspirin
0.175
0.2191
0.1891
0.1522
0.1684
0.1900
0,1520
0.1429
0.1407
0.1338
0.1212
Azobenzone
0.097
0.0669
0.0692
0.0600
0.0761
0.0585
0.0513
0.0507
0.0530
0.0486
Benzene
0,017
0.0115
0.0069
0.0069
0.0046
0,0069
0.0056
0.0047
0.0023
0.0069
0,0056
Ethanol
0.085
0.083
0.0646
0.0484
0.0484
0.0738
0.0522
0.0516
0.0322
0.0577
0.0438
Malonaldehyde
0.152
0.1522
0.0118
0.0946
0.0830
0.1338
0.0893
0.0875
0.0715
0.0922
0,0802
Naphthalene
0.060
0.0438
0.0300
0.0369
0.0208
0.0415
0.0291
0.0242
0.0208
0.0300
0.0228
Paracetamol
0.164
0.1361
0.1107
0.1130
0.1338
0.1029
0.0979
0.0922
0.0922
0,0840
Salicylic acid
0.088
0.1222
0.0922
0.0715
0.0669
0.0992
0.0795
0.0771
0.0623
0.0761
0.0648
Toluene
0.058
0.0507
0.0369
0.0350
0.0415
0.0438
0.0264
0.0244
0.0254
0.0300
0.0239
Uracil
0.088
0.0876
0.0669
0.0484
0.0415
0.0738
0.0495
0.0487
0.0392
0.0553
0.0446
abstract edges. Eij does not have an I. superscript because these features remain constant across different layers. This as an additional contribution to the cross-attention, as it quantifies how well the abstract edges capture the information of the edge linking atoms and j. The detailed imple- mentation, including the Cutoff function and dk:, in the above formula, is introduced in Sec. A.2
Equivariance. It is important to maintain equivariance for physics-informed tasks. Here, we sum- marize the overall idea of ensuring equivariance for FreeCG. There are two equivariances maintained in this work, namely O(3) (3D rotation and reflection) and permutation equivariance. We maintain 0()) equivariance via leveraging spherical irreps, CG transform, and the linear combinations on them. Each of the objects and operations above is -33--qquivrrian,, and so does the combination of them. Specifically, the proof of linear combination is in Sec. A.5 The core challenge of this work lies in how to maintain permutation equivariance while generating abstract edges. It is shown, in Sec. 3.3 that we leverage invariance transitivity to maintain the permutation equivariance.
4 EXPERIMENTS
To evaluate the performance of FreeCG, we collect force field datasets, on which we compare our methods with other SOTA MLFFs. These datasets include small molecules dataset MD17
(Chmiela et al. 2017) with its revised version, rMD17 (Christensen & Von Lilienfeld 2020), and large molecules dataset MD22 (Chmiela et al. 2023). To test the generalizability of FreeCG, we
also evaluate the performance of FreeCG on standard molecule property prediction dataset, QM9 (Ruddigkeit et al., 2012) Ramakrishnan et al. 2014). We take popular SOTA models into the com-
| et al. 2020b). SphereNet (Liu et al.) 2021). PaxNet (Zhang et al. 2022p. PaiNN (Schütt et al. 2021
parison, including sGDML (Chmiela et al. (017)). SchNet (Schütt et al.) 2018). DimeNet (Gasteiger)
SpookyNet (Unke et al.
2021). ForceNet (Hu et al.) (021). ET (Thölke & De Fabritis 2021
GemNet (Gasteiger et al.
2021). ComENer (Wang et al. (202).. NequlP (Batzner et al. 2022
UniTE (Qiao et al. 2022). , SO3KRATES (Frank et al. (022). MACE (Batatia et al. 2022b). Allegro
MMusaelian et al. 2023). BOTNet (Batatia et al. 2022a) ViSNet (Wang et al.) 20244, ViSNet-LSRM
[Li et al. 2023) QuinNet (Wang et al. (023cc) ) ACE" Bochkarev et al. 2024). PONITA (Bekkers)
et al. 2024),and EquiformerV2 V2 (Liao et al. 2024). To asses the practicality on weal-world tasks,
Molecular dynamics simulations are run for MDI7 molecules and two periodic systems, water (Fu
et al. 2022 Wu et al. 2006) and LiPS (Batzner et al. 2022) under Periodic Boundary Conditions
8

---

Published as a conference paper at ICLR 2025
Table 3: Performances MD22 dataset. The results reported in MAE. The energies and forces are measured in kcal/mol and kcal/mol/A, respectively. The best numbers are marked in bold. Note that the energy MAE is calculated without being divided by the total number of atoms as Wang et al. (2024), unlike Wang et al. (2023c); Chmiela et al. (2017) 2023) which does not affect the comparison.
Molecule
sGDML
ViSNet
ViSNet-Improper
ViSNet-LSRM
MACE
QuinNet FreeCG
Energy Prediction
Ac-Ala3-NHMe
0.391
0.0636
0.0546
0.0673
0.0631
0.0840
0.0507
AT-AT
0.720
0.0708
0.0668
0.0780
0.108
0.144
0.0665
AT-AT-C--GG
1.42
0.196
0.197
0.118
0.154
0.379
0.254
DHA
1.29
0.0741
0.0700
0.0897
0.135
0.118
0.0761
Buckyball catcher
1.17
0.508
0.537
0.319
0,489
0.563
0.512
Stachyose
4.00
0.0915
0,0882
0.104
0.122
0.226
0.183
Double-walled nanotube
4.00
0.800
0.601
1.81
1.67
1.81
0.543
Force Prediction
Ac-Ala3-NMMe
0.790
0.0830
0.0709
0.0942
0.0876
0.0681
0.0531
AT-AT
0.690
0,0812
0,0776
0.0781
0.0992
0.0687
0.0634
AT-AT-CG-CG
0.700
0.148
0.139
0.1064
0.1153
0.1273
0.1252
DHA
0.750
0.0598
0.0554
0.0598
0.0646
0.0515
0.0507
Buckyball catcher
0.680
0.184
0.201
0.1026
0.0853
0.1091
0.1783
Stachyose
0.680
0.0879
0.0802
0.0767
0.0876
0.0543
0.612
Bouble-walled nanotube
0.520
0.362
0.292
0.3391
0.2767
0.2473
0.2449
(PBCs). FreeCG is also evaluated on the conformation space of 166-atom mini-protein, Chignolin Wang et al. 2023a). Mean Squared Error (MSE) loss is adopted for all experiments. The ablation on each component and eyperpparametess of FreeCG are also presented. The common settings and extra experiments are reported the Sec. A.T
4.1 COMPARISON WITH SHATE-OF-THE-ARE FOR MLFFs
Force field on small molecules and periodic systems. MD17 is famous molecular dynamics benchmark for small molecules. FreeCG outperforms others all force prediction tasks. It also sig- nificantly decreases the force prediction errors for the hardest-to-pre to-rrddcpttttte molecule in this datasets, aspirin, by 15%. Remarkably, FreeCG also decreases the MAE by over 10% for ethanol, naphtha- lene, and salicylic acid. FreeCG does not have particular preference for the size of molecules. It demonstrates strong performance for aspirin (180.2 g/mol) and excels on ethanol (46. g/mol). The energy prediction is also competitive when compared to other SOTA methods. rMD17 is the revised version of MD17. recomputed the trajectories of each atom with higher accuracy. The force pre- diction accuracy of FreeCG still leading in majority of the molecules. It improves the force results compared to the baseline, ViSNet, in all the molecules except for benzene, and performs SOTA on more atoms. Note that the results on benzene is already extreme high with previous models. The results for MD17 and rMD17 can be referred to Tab. and 2 prediction on periodic systems, water and LiPS, are also evaluated. Fig. and show that, FreeCG achieves the best performance across other methods, and decrease the MAE value of the second best by 50% for water. Note that the energy prediction on rMD17 does not show SOTA results. The main focus of our work is for force prediction, and the downstream task is mainly for molecular dynamics and related ones that requires precise force prediction. That being said, we conduct experiments where the loss weight is adjusted to 0.9*force MAE 0.1 *eneryy MAE (different the number ported Tab. 5). The checkpoint in 2500 epochs shows that it performs SOTA on energy prediction for aspirin in rMD17, shown in Tab.
respectively. The accuracy of force
Force field on large molecules. MD22 is large molecules benchmark adopted by several studies Li et al. 2023). As shown in Tab. Bl it reveals that FreeCG also performs well for large molecules. It leads in most tracks for force prediction, and shows comparable results for energy prediction. Remarkably, The decreasing in MAE for energy and force prediction
Wang et al. 2024 2023c

---

Published as a conference paper at ICLR 2025
Ac-Ala3-NHMe both around 20%. The performances for the other models are not consistently well for force prediction, while ViSNet-LSRM exhibits strong performance for energy prediction. It is also reasonable that all modern deep neural network-based methods outperform sGDML, classical kernel method.
4.2 COMPARISON WITH STATE-OF-THE-ARTS FOR MOLECULAR PROPERTIES PREDICTIONS
To examine the generalization power on molecular property prediction of FreeCG, we collect QM9 as a standard benchmark for this task. FreeCG performs the best for most properties. Although FreeCG proposed to be MLFFs, but it even more comparable than others in molecular property prediction tasks. The results in Tab. T demoosttrae strong generalization capabilities of FreeCG.
4.3 EFFICIENCY BENCHMARKING
Except for the conformation exploration in Sec. A.7 AIMD-Chig dataset Wang et al.
(023a) is
taken for testing the memory usage and inference speed. It contains mini-protein which is more challenging for inference speed and memory usage. We compare the inference speed and memory usage of FreeCG with ViSNet, NequlP, and Allegro. The results are shown in Fig. 5 The training time and the numbers of parameters are also shown in Tab. [2 FreeCG adds little extra time and memory cost, compared the baseline model. ViSNet. memory and speed, compared to the other two CG transform based methods, NequlP and Allegro. The overall results prove the effciency of FreeCG. The number of groups group CG transform also impacts the inference speed. Fig. shows the number of paths and the actual inference time for different group numbers. A computation analysis for CG transform can be referred to Sec. A.6
also the most efficient one for both
4.4 ABLATION STUDY
We conduct ablations on different modules we propose, well the strategies for abstract edges shuffling. The results are shown in Tab. 7 It reveals that each of our module contributes to the final score of FreeCG. In the final implementation of abstract edges shuffling, we add the index of each abstract edge by [1.5 * T/GJ. Here we also study the influence of the shuffling strategies. We adopt [0.5 T/GJ. 1.0 T/G], and [1.5 T/G] for comparing the performance. We can see from the result that [1.5 T/G] works the best. The group numbers are also evaluated and large number of groups appears be a good choice. The ablation on sparse path module reported in Tab. X reveals that sparse path exhibits satisfying trade-off between performance and efficiency. Full path only decreases minor energy and force prediction MAE, while brings huge overhead in training. inference speed (see Fig. 3). and memory cost.
4.5 EXTENSION TO OTHER MODELS
FreeCG also presents a paradigm to enhance other geometric models. We use the former SOTA model. QuinNet, as an example to illustrate how effectively FreeCG can be extended other archi- tectures. QuinNet has a transformer architecture inside, so we construct abstract edges the same way FreeCG does, and adopts all FreeCG modules upon those abstract edges. We evaluate the training curve of 1000 epochs training for both QuinNet+FreeCGG and vanilla QuinNet. As shown in Fig. and [T] when equipped with FreeCG, QuinNet significantly gets improved for both energy and force prediction. This trend becomes more pronounced with longer training periods, as evidenced by the results at the 1000th epoch. We also report the results of training for 1500 epochs (see Fig. which shows that QuinNet+FreeCG significantly outperforms fully-trained QuinNet.
5 CONCLUSION
This work proposes FreeCG, an equivariant neural network that frees the design space of CG trans- form. It achieves SOTA performance in force prediction for MD17, rMD17, and MD22 datasets, as well as in molecular properties prediction for QM9, with only minor computational overhead. The practicality of FreeCG for conducting molecular dynamics simulations is thoroughly examined across periodic systems, small molecules in MD17, and the mini mini-protein Chignolin. As show that FreeCG helps improve QuinNet, it also introduces a new paradigm for expressive and efficient CG transform-basse based neural network design in the future.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
is crucial to ensure that the use of MLFFs is strictly regulated and controlled to prevent their misuse for illegal purposes, such the development and deployment of chemical weapons. Furthermore, global cooperation and information sharing between the academic community and industry are key to identifying and mitigating potential threats timely manner. We require that further research published top of this work could adhere the code of ethics, avoiding the misuse of MLFFs. Meanwhile, the data used in our work does not contain any human subjects. All the data we used freely available the internet.
REFERENCES
Rodney J Bartlett and Monika Musiat. Coupled-cluster theory in quantum chemistry. Reviews of Modern Physics, 79(1):291, 2007.
Ilyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor NC Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi. The design space of (3)-equivariant atom- centered interatomic potentials. arXiv preprint arXiv:2205. 06643, 2022a.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:1 1423-1 1436, 2022b.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Ko- rnbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):
2453, 2022.
Erik J Bekkers, Sharvaree Vadgama, Rob Hesselink, Putri A Van der Linden, and David W. Romero. Fast, expressive $ mathrm (SE) (n)S equivariant networks through weight-sharing in position- orientation space. In ICLR, 2024.
Anton Bochkarev, Yury Lysogorskiy, and Ralf Drautz. Graph atomic cluster expansion for semilocal interactions beyond equivariant message passing. Physical Review X, 14(2):021036, 2024.
David M Ceperley and Berni J Alder. Ground state of the electron gas by a stochastic method. Physical review letters, 45(7):566, 1980.
Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras- Derras-Domingo, Caleb Ho, Weihua Hu, al. Open catalyst 2020 (oc20) dataset and community challenges. Acs Catalysis, 11(10):0059-609) 6072, 2021.
Sijie Chen, Tong Lin, Ruchira Basu, Jeremy Ritchey, Shen Wang. Yichuan Luo, Xingcan Li, Dehua Pei, Levent Burak Kara, and Xiaolin Cheng. Design of target specific peptide inhibitors using generative deep learning and molecular dynamics simulations. Nature Communications, 15(1): 1611, 2024.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning accurate energy-conserving molecular force fields. Science advances, 365)ee6033 2017. 2017.
Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9())eaddf0873, 2023.
Anders Christensen and O Anatole Von Lilienfeld. On the role of gradients machine learning of molecular energies and forces. Machine Learning: Science and Technology, 144):045018, 2020.
Taco S Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical cnns. arXiv preprint arXiv: 1801. 10130, 2018.
11

---

Published as a conference paper at ICLR 2025
Taoyong Cui, Chenyu Tang, Mao Su, Shufei Zhang, Yuqiang Li, Lei Bai, Yuhan Dong, Xingao Gong, and Wanli Ouyang. Geometry-enhance pretraining on interatomic potentials. Nature Machine Intelligence, pp. 1-9, 2024.
Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review B, 99(1):014104, 2019.
Thorben Frank, Oliver Unke, and Rabur- Müller. So3krates: tions arbitrary length-scales in molecular systems. Advances in Neural Information Processing Systems, 35:29400-2999 29413, 2022.
Equivariant attention for interac-
Xiang Fu, Zhenghao Wu, Wujic Wang. Tian Xie, Sinan Keten, Rafael
and
Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. arXiv preprint arXiv:2210.07237, 2022.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se trn-transformmrs: 3d roto-
translation equivariant attention networks. Advances in neural information processing systems, 33:1970-1981, 2020.
Johannes Gasteiger, Shankari Giri, Johannes Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.141155 2020a.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molec- ular graphs. arXiv preprint urXiv:2003.03123 2020b.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790- 6802, 2021.
Weihua Hu, Muhammed Shuaibi, Abhishek Das. Siddharth Goyal, Anuroop Sriram, Jure Leskovec, Devi Parikh, and C Lawrence Zitnick. Forcenet: A graph neural network for large-scale quantum calculations. arXiv preprint arXiv:2103.01436 2021.
Nadir Jeevanjee. An introduction to tensors and group theory for physicists. Springer, 2011.
Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A113, 1965.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Yunyang Li, Yusong Wang, Lin Huang. Han Yang, Xinran Wei, Jia Zhang, Tong Wang. Zun Wang, Bin Shao, and Tie-Yan Liu. Long-short-range message-passinng physics-informed frame- work capture non-local interaction for scalable molecular dynamics simulation. arXiv preprint arXiv:2304. 13542, 2023.
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206. 11990, 2022.
Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivari- ant transformer for scaling to higher-degree representations. In ICLR, 2024.
Mingfeng Liu, Jiantao Wang, Junwei Hu, Peitao Liu, Haiyang Niu, Xuexi Yan, Jiangxu Li, Haile Yan, Bo Yang, Yan Sun, al. Layer-by-layer phase transformation in ti3o5 revealed by machine- learning molecular dynamics simulations. Nature Communications, 15(1):3079, 2024.
Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d graph networks. arXiv preprint arXiv:2102.05011 2021.
Richard M Martin. Electronic structure: basic theory and practical methods. Cambridge university press, 2020.
Markus Meuwly. Machine learning for chemical reactions. Chemical Reviews, 121(16):10218- 10239, 2021.
12

---

Published as a conference paper at ICLR 2025
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Ko- mbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Zhuoran Qiao, Anders Christensen, Matthew Welbor,, Frederick R Manby, Anima Anandku- mar, and Thomas F Miller III. Informing geometric deep learning with electronic interactions to accelerate quantum chemistry. Proceedings of the National Academy of Sciences, 119(31): e2205221119, 2022.
Ryszard Raczka and Asim Orhan Barut, Theory of group representations and applications. World Scientific Publishing Company, 1986.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. 17. Journal of chemical information and modeling, 52(1 11:22)2 2875, 2012.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference on machine learning, pp. 9323 22332.3.. PMLR, 2021.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-9388. PMLR, 2021.
Kristof T Schütt, Huziel Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018.
Guillem Simeon and Gianni De Fabritiis. Tensomeet: Cartesian tensor representations for efficient learning of molecular potentials. Advances in Neural Information Processing Systems, 36, 2024.
Philipp Thölke and Gianni De Fabritis. Equivariant transformers for neural network based molecu- lar potentials. In International Conference on Learning Representations 2021.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivarian neural networks for 3d point 2018.
clouds. arXiv preprint /
Oliver T Unke, Stefan Chmiela, Michael Gastegger, Kristof Schüitt, Huziel Sauceda, and Klaus- Robert Müller. Spookynet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nature communications, 12(1))7273, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Limei Wang. Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. Comenet: Towards complete and efficient message passing for 3d molecular graphs. Advances in Neural Information Processing Systems, 35:650-664, 2022.
Tong Wang, Xinheng He, Mingyu Li, Bin Shao, and Tie- Yan Liu, Aimd-chig: Exploring the con- formational space of a 166-atom protein chignolin with ab initio molecular dynamics. Scientific Data, 10(1):549, 2023a.
Yusong Wang, Tong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing. Nature Communications, 15(1):313, 2024.
Zhoujie Wang, Jianzhong Zhu, and Songyan Li. Novel strategy for reducing the minimum mis- cible pressure in a co2-oil system using nonionic surfactant: Insights from molecular dynamics simulations. Applied Energy, 352: 121912, 2023b.
13

---

Published as a conference paper at ICLR 2025
Zun Wang, Guoqing Liu, Yichi Zhou, Tong Wang, and Bin Shao. Quinnet: efficiently incorporat- ing quintuple interactions into geometric deep learning force fields. In Proceedings of the 37th International Conference Neural Information Processing Systems, pp- 77043-77055, 2023c.
Yujie Wu, Harald Tepper, and Gregory A Voth. Flexible simple point-charge water model with improved liquid-state properties. The Journal of chemical physics, 124(2). 2006.
Anthony Zee. Group theory in nutshell for physicists, volume 17. Princeton University Press, 2016.
Shuo Zhang, Yang Liu, and Lei Xic. Efficient and accurate physics-aware multiplex graph neural networks for 3d small molecules and macromolecule complexes. arXiv preprint arXiv:2206.02789, 2022.
Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun. Shufflenet: An extremely efficient convolutional neural network for mobile devices. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6848-6856, 2018.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 EXPERIMENTAL SETTINGS
We conduct all the experiments under the same software and hardware settings. The machine is equipped with an Intel Xeon Gold 6330 CPU @ 2.00GHz, with NVIDIA Tesla A100 80G GPUs. We run the experiments for each molecule on single GPU. Pytorch 1.10.0 is used as the basic ma- chine learning python library. For the CG transform operations, we adopt e3nn 0.5. Matplotlib 3.0.3 is utilized for plotting. The details can be referred Tab. 4 We report the hyperparameters used in Tab. 5 For training/validationnness splits, we follow previous works (Wang et al., 2024]] 2023c Fu et al. 2022). We pick up the model for evaluating test set based on the performance on the validation set. If the model does not improve for a given number of epochs, we will terminate the training and select the checkpoint with the best validation score. As previous works, Expo- nential Moving Average (EMA) is adopted to generate the model weights. The detailed training configurations are shown in Tab.
Table 4: Hardware and software settings.
Hardware
Software
CPU
GPU
Neural Network
Equivariance
Plotting
Intel Xeon R Gold 6330 CPU @ 2.00GHZz
NVIDIA Tesla A100
Pytorch 1.10.0
c3nn 0.5.1
Matplotlib 3.0.3
Table 5: Hyperparameters for each dataset.
Hyperparameter
MD17
rMD17
MD22
QM9
Water-Ik
LiPS
Chignolin
OC20
Initial learning rate
4e-4, 2e-4
2e-4
2e-4, le-4
le-4
Se-4
1e-3
2e-4
le-4
Learning rate decay factor
0.8
Learning decay patience
30
30
30
15
5
Learning rate warmup step
1000
1000
1000
10000
1000
1000
1000
6000
Loss
Mcan Squared Error (MSE) loss
Optimizer
Adam W (8,0.9, 0.999))
Epoch
3000
3000
3000
1500
1000
100
3000
12
Batch size
4
32
1
4
4
Number of layers
9
9
9
9
Cutoff
5.0, 4.0
5.0
5.0, 4.0
5.0
6.0
5.0
14
Force/Energy loss weights
0,95/0.05
0.95/0.05
0.95/0.05
1.0/0
1.0/0
0.95/0.05
0.99/0.1
Dimension of latent feature Number of groups
256
256
256
512
256
256
256
Output head
Bquivarian//callar 0.999
FMA rate
A.2 MODEL IMPLEMENTATION
Here we show how FreeCG is built upon ViSNet. This section provides detailed explanations of the implementation details, ensuring FreeCG can be replicated effectively.
Input layer. Given the atom coordinates and types (X = Fo, Fi, F2..... FN). 2 (21, ..., 2n } where 7 E IR3 the Cartesian coordinates of atom, and the atom type (atom numbers). First we embed the atom types to the latent space, and take them as our first layer's node features h4 embedding(z))
R.C. C. is the dimension of the latent space. For each atom, we only consider neighboring atoms within a given radius Ni)). where we maintain the distance vector from the central atom the neighboring atoms, and lift them
1,p -1) and (1 = 2.p = 1) irreps Eij RR355 via real spherical harmonics applied on the unit vector Eg = y - (F,s//(/,,)))) + (Fag//F
where we also calculate the corresponding Euclidean norm || The Euclidean norms of vectors are then converted to high-dimension scalar features (edge attributes) fúz RBF(7,) E IR. by radial basis functions (RBFs). We also maintain zero- initialized abstract edges -T O for each node be updated in the following layers. We assign
15

---

Published as a conference paper at ICLR 2025
Table 6: Molecular property prediction QM9 dataset. The results reported in MAE, The best numbers are marked in bold.
Target
SchiNet
EGNN
DimeNet++
PaiNN
SphereNet
PaxNet
ComENet
ViSNet
EquiformerV2
FreeCG
g
24.5
10
11.4
43.5
GROMO
meV
41
29
24.6
27.6
22.8
22.8
20.3
23.1
17.3
14
16.6
mcV
34
25
19.5
20.4
18.9
19.2
17.5
19.8
14.8
13
13.5
A
meV
63
48
32.6
45.7
31.1
31
36.1
32.4
31.7
29
31.5
(R))
mig
73
106
331
66
268
93
33
259
29.8
186
82.1
ZPVE
meV
1.7
1.55
1,21
1,28
1.12
1.17
1.84
1.2
1.56
1.47
1.10
5,92
mcV
6.53
6.33
6.04
6.16
4.52
6.22
4.13
meV
12
7.56
7.35
7.78
7,14
7.62
7,98
5,86
7,57
5.65
G
meal/mol K
33
31
23
24
22
23.1
26
24
23
23
20.4
Table 7: Ablation on different modules. Abstract edges shuffling and Attention enhancer are added upon the best choices of the above modules, with respect the validation loss.
Aspirin
Method
Val loss
Energy
Force
ViSNet
0.116
0.155
Group CG transform
groups
0.0509
0.123
0.144
32 groups
0.0416
0.112
0.129
Abstract edges shuffling
g-roup shuffle
0,0401
0.112
0.128
0.5-group shuffle
0.0396
0.110
0.128
1.5 -rouu shuflle
0.0384
0.111
0.125
Attention enhancer
0.0345
0.110
0.122
the same number of abstract edges the dimension of the latent features, such that additional oper- ations to align the dimension numbers are not required.
Intermediate layers. Here, are in. The sesssage-passingg between atoms is implemented by a transformer architecture. For each atom the neighboring atoms j (i)) will send messages to i, and the messages to update the information of i. The query, key, and value of the node features are first calculated, respectively: qi (h,),, ky = fa (h,). Bj fog)). The edge attributes are also converted to auxiliary terms dk, = fax (fe) and duj = (fa) to modulate keys and and values of atoms. Here functions f's are all fully-connected linear operations. Then we calculate the cross-attention between and j. which is
superscript to denote the index of layer that the features
aggregated
ai SiLU Cutoff ( + AttEnhancer(E.j.. Eyy)
(7)
where (-toff(-) is cosine cutoff function, and AttEnhancer(() the proposed attention enhancer module, as we will formulate its details. First, recall the dimension of E/ C+8 and Ei E R Each of the C abstract edges will undergo dot product with Eij. The highest value among them will be the output of AuEnhancer. In other word,
AttEnhaaceer(E, E.y) = max(E, o Enj)
(8)
as we introduce in Eq. Then the values are multiplied with de and attention.
+ー 02 duy. Oaj
It then undergoes two different fully- fully-connected operations to generate two coefficients s and They are used generate the abstract edges:
EJ E 81 + Eig 82
(10)
This variable, together with منرز are aggregated by sum:
B+ = Σ A
(11)
JEN())
16

---

Published as a conference paper at ICLR 2025
Table 8: Ablation on Sparse Path. The energy and force prediction errors are reported in mean abosolute error (MAE). The unit of training speed is sterations/eccon.. The mini-batch size is set to 32.
Force MAE
Energy MAE
Train. speed
Num. of param.
Sparse path
0.122
0.110
3.62
10.4M
Full path
0.118
0.106
0.79
43.8M
6f
(12)
JEN(G)
/ then converts to three variables for further operation:
1.2 = Linear(0!))
(13)
Ef E R+8 and E E RC8 are used for the following group CG transform and abstract edges shuffling. First E RC undergoes fully- tully-connected operation along C dimension, and multiply with 1,1 which means o of Linear(E.))
It, together with E RC*5 are then divided into G groups along C' dimension, where we get Eiiee, R R8-8 and (o,-) Linear(EE)) ))EG, R 8 Then, perform CG transform between two variables in fully connected form with learnable weights, and concatenate the results gen- erate d.F.L+l before shuffling, as shown in Eq. 5 For the shuffling strategies, we add [2G] to each 3C
index of the abstract edges F=L-11 Then, it is added with EL to form a residual structure, as we show here:
dE, L+1 shuffle(dE E
(14)
where dE, dE added to E to obtain E, El+11 Next, we update h and f. We first show the update for
dh dif++ h + [Lincar, (E @ Lincarz(E (EE))
22 L.3 o.
+ +
(15)
To update f. we follow ViSNet to leverage rejection of vectors:
afttt flj + RejCaleun .
(16)
RejCal (E Fb)) (E), SiLU(Linear(//])
where rejection calculation module RejCalc is:
RejCalCCootee b) = a (Linearmode (a) bb
(17)
The updated L+1 n,L11, and p4+1 are fed into the next layer.
Output layers are different with respect to the task our model performs on. We introduce the details for each task.
Force field prediction. FreeCG is based veegrvivoneevvvtive field, which means derive the force from the predicted potential energy. Following ViSNet (Wang et al.) 2024) and PaiNN (Schütt) al. 202I), predict the potential energy of the molecule via equivariant gated module.
L+1 h
L+1 = MLP Con ||Linear (BF)II) II)
(18)
where MLP is an 1-hidden layer multi layer preceptor. There is one more step to update E _b+1_
El Linear (E u, 1++1
(19)
These calculations then repeated twice in succession. There is also alternative head design for force field prediction which uses only scalar features h. Under this setting, Eq. and and (19 replaced by (L denotes the last layer):
are
1 = MLPP(L--
(20)
17

---

Published as a conference paper at ICLR 2025
Theoritical complexxit
Actual running time
Sparse DG tranatom transfor
Spane CG transtor Cu tranaform
Number of pottt
Duaninng Simne (195)
614
6.6
Sroup
proup
groups
The number of paths for CG transform under different group numbers, where the numbers of irreps are the same. Right: The actual running time for CG transform for different group numers. Here we adopt sparse path strategy for computing 512 irreps (before grouping) for each I. Full CG transform denotes not using sparse path.
Figure 3: Efficiency analysis of group CG transform. Left:
In Tab. where we introduce our parammee choices, the scalar version of output head is denoted as Scalar, and the other one Equivariant. Finally, the total energy of the molecule is the sum of the last-layer node features h RR:
Σ 4
(21)
and the force is the negative gradients of the total energy:
F.
(22)
Property prediction. The calculations for properties in QM9 follow the same procedure as energy prediction in force field prediction, with the exception of molecular dipole and electronic spatial extent. We first need to calculate the center of mass Fo, which is:
Fe E, m. F Σ, m4
(23)
For molecular dipole, the formula is:
Σ E + 14 Fo)
(24)
and for electronic spatial extent:
(
Σ 1 1 Fall
(25)
The output head can be easily adapted for different tasks, providing flexibility in property prediction.
A.3 DATASETS DETAILS
MD17 and rMD17. They are both molecular dynamics datasets for small molecules. MD17 (Cmieela et al. 2017). proposed by Chmiela, S., et al. contains ab-initio level molecular dynamics trajectories. Four types of data are included the dataset: atomic numbers, atomic positions, molec- ular energy, and the force acting on each atom. To alleivate the noise during the trajectory computa- tion, Christensen, A. et al. also propose revised MD17 (rMD17) (Christensen & Von Lilienfeld] 2020). where molecular trajectories calculated at the PBE/def2-S level of theory. The preci- sion of the calculated trajectories upheld by the tight SCF convergence and dense DFT integration grid.
MD22 consists of larger molecules with atoms numbering from 42 to 370, in contrast to MD17 and rMD17. The trajectories sampled between 400K and 500K Ifs resolution. The energy and
18

---

Published as a conference paper at ICLR 2025
Table 9: The basic operation number for each type of CG transform. to denotes the output degree. The column and row numbers represent the degrees of two input irreps, respectively. The cyan blocks represent the operations in regular neural networks, while the others are for high-order CG transform.
4
2
0
2
6 =
0 1 2
2 0
0
1
3
0
1
3
6
9
1
9
12
2
5 2
9
12
2
5
12
19
Performances on MD17 dataset under loss weight being 0.9* MAE+0.1 MAE. The results are reported in mean abosolute error (MAE). The energies and forces are measured in kcal/mol and kcal/mol/A, respectively. The best numbers are marked in bold.
Table 10:
UNITE
GemNet
NequlP
MACE
Allergo
BOTNet
VisNet
QuinNet
ACE"
PeNITA
FreeCG
Energy
0.055
0.0530
0.0507
0.0530
0.0530
0.0445
0.0486
0.0392
0.0392
0.0374
Aspirin
Force
0.175
0.2191
0.1891
0.1522
0.1684
0.1900
0.1520
0.1429
0.1407
0.1338
0.1225
Energy
0.025
0.0161
0.0277
0.0277
0.0161
0.0156
0.0394
0.0115
0.0161
0.0181
Azobenzene
Force
0.097
0.0669
0.0692
0.0600
0.0761
0.0585
0.0513
0.0507
0.0530
0.0475
Energy
0.002
0,0009
0.0092
0.0069
0.0007
0.0007
0.0096
0.0002
0.0039
0,0084
Benzene
Force
0.017
0.01 15
0.0069
0.0069
0.0046
0.0069
0.0056
0.0047
0.0023
0.0069
0.0057
Energy
0.014
0,0092
0.M032
0.0092
0.0092
0.0078
0.0096
0.0046
0.0892
0.0072
Fthanol
Force
0.085
0.083
0,0646
0.0484
0.0484
0.0738
0.0522
0,0516
0.0322
0,0577
0.0450
force labels are obtained at the PBE+MBD level of theory. The root mean squared test error of force prediction is controlled to be around kcal/mol/Ā in the original paper (Chmiela al. 2023). Thus, the training data sizes for different molecules vary. Generally, the larger the molecules, the smaller the training data size.
QM9 consists of around 130,000 molecules with 12 properties regression tasks. is subset of the GDB-17 database (Ruddigkeit et al.) 2012)). The data is calculated at B3LYP/6-31G(2dEp)) based DFT level of accuracy. Since the attributes vary for different properties, we use distinct output head for each, as discussed Sec. A.2
Chignolin. The AIMD-Chig dataset (Wang al. 2023a) comprises of two million conformations of the 166-atom protein chignolin, obtained through sampling at the M06-2X/6-31 G* based DFT level. There approximately 10,000 different conformations, including folded, unfolded, and metastable states. We report the performances of FreeCG on different parts of the energy landscape, and adopt this dataset benchmark efficiency, following (Wang et al. 2024).
Periodic systems. Periodic Boundary Conditions (PBCs) are vital in molecular dynamics simula- tions of periodic systems they eliminate surface effects, enhance statistical sampling, and provide a realistic representation of bulk properties. Here we focus on two typical molecules, water and LiPS. The water dataset are generated by the flexible version of the Extended Simple Point Charge water model (SPC/E-fw) (Wu et al.) 2006) in (Fu et al.) 2022). The authors provides with several sizes of training sets, including II, 1Ok, 90k. Here we adopt the 1k version, where 950 samples are for training and the rest for validation. LiPS is an important solid-state materials for battery de- velopment. It can help predict key performance metrics such as capacity, energy density, and cycle life, aiding in the development of next- next-generation lithium-ion batteries. We follow the same train/val split as (Pu et al. 2022). Some 3D structures of the data can be referred to Fig.
A.4 PROOF OF THE PERMUTATION INVARIANCE OF ABSTRACT EDGES
According to Sec.
A.2
we first recall the last step for generating abstract edges:
Efp) Card(P)
Σ Joz
Σ
(26)
(E)
PEP
JEN(i)
19

---

Published as a conference paper at ICLR 2025
a
(1)
637
FreeOG
ViSNet
FreeCG
ViSNet
(3)
(4)
PlturtAtt Stree 1982)
0.7MM
FreeCG
ViSNet
FreeCG
ViSNet
FreeCG
ViSNet
(6)
C
orgon daspat
0.244
FreeCG
ViSNet
Sir.baiion Time (p1)
330
Figure 4: Applications for the 166-atom mini-protein, Chignolin. a. The energy landscape of Chignolin was sampled using Replica Exchange Molecular Dynamics (REMD). This landscape is characterized by two key distance parameters: the x-axis represents the distance between the car- bonyl oxygen on the D3 backbone and the nitrogen the G7 backbone, while the y-axis depicts the distance between the carbonyl oxygen on the E5 backbone and the nitrogen on the T8 backbone. These two distance metrics collectively illustrate the conformational states of Chignolin across its energy landscape. The left and right energy basins are corresponded to folded and unfolded states, respectively. b. The six conformations landscape. The force and energy performances (kcal/mol) are reported, with a comparison made to ViSNet. These six conformations cover both folded and unfolded states. The RMSD (A) during the molecular dynamics simulation. The shaded area denotes the values of standard derivations. The RMSD values obtained by taking average of 10 trajectories. Better viewed in color.
sampled at the localization highlighted in the energy
where P is the set for all permutation operations, here we omit the subscript of P for specific spaces work on. This equation indicates that, performing our aggregation (the middle part) is equivalent to performing the last part (which guarantees the operation is permutation invariant). Here, note we are proving that the abstract edges for each atom are permutation invariant, and we can freely design CG transform per alom, thus the permutation is applied to j but not i. It sums over all the permutation operations, and thus the last step permutation invariant. Then, suffices to show that each of the previous step are all least permutation equivariant. It also suffices to show they permutation equivariant w.r.t single index switch operation, as each permutation operation can be made by several switches. If we exchange, without loss of generality, index and y. then those di that z or y shows up the subscript for j will exchange with each other, and do Ú and ÉT. Thus, the rest steps are equivariant w.r.f. single switch, and so they are permutation equivariant. Therefore, conclude our proof that abstract edges E's permutation invariant. For E, it is also permutation invariant, as it is generated by permutation invariant function h in the previous layer. /-=1 in the first layer is permutation invariant too, it is set to be fixed zeros.
20

---

Published as a conference paper at ICLR 2025
High-order irreps and CG transform
Inference time (msecs)
Size (mbytes)
8000
2000
20
FrenCG
"Modal"
Figure 5: The speed and memory occupation of FreeCG compared with other SoTA models. Num- bers reported based single chignolin molecule. The right three models based on high- order irreps and CG transform. Better viewed in color.
31prrh
ethandl
This challeeydde
FreeCG
DFT
101.02 cy sxid
tolvene
2.1
20-
0.1
Figure 6: The distributions of interatomic distances during molecular dynamic simulations of MD17. The unit of is A, and the unit of h(r) is A
A.5 PROOF OF THE EQUIVARIANCE THE LINEAR COMBINATION OF SPHERICAL IRREDUCIBLE REPRESENTATION
Let Us the spherical irreps with order 1, g E (O()), w; the weight for i-th vector in the weighted sum, and PV the group homomorphism. We need to prove (v(g) E, With Ee WaPv; (0)%. It holds we can prove the following conditions 1) pv, (g)w; = WiPy, (g)x; which is trivially true as w; is scalar; 2) pV; pV; for arbitrary i and j.
To show 2) is golden, we need to show a) V, an invariant subspace of SO(3), since the pv; would be ill-defined if it is not the case, which is also trivially true because that is how spherical spaces are generated; b) Vi = V, for arbitrary and j. Here, e are spherical irreps with same l. Spherical space of order is the same vector space spanned by spherical harmonics Y!,,s Spherical irreps with order I are the vectors in the same order spherical space, V; Vj for arbitrary and Since all V,s are the same we can simply denote each of them as V. Summarize over for both sides of pv; (g) (g)w;; wipp; (g)v,, we get pv(g) Ee wiy; Ee wipp (g)u,. This concludes our proof.
A.6 ANALYSTS ON THE EFFICIENCY OF CG TRANSFORM
The CG transform consists two steps: 1) performing a tensor product between two irreps, and 2) decomposing the resulting tensors into irreps. These transforms actually quadratic homogeneous polynomials. For the sake of convenience, we discuss SO(3) group here. Recall the CG transform formula:
21

---

Published as a conference paper at ICLR 2025
Table 11: Performances of QuinNet equipped with FreeCG MD17. The results reported in
MAE. The energies and forces are measured in kcal/mol and kcal/mol/A, respectively. The best numbers are marked in bold.
Aspirin
Ethanol
Malonaldehyde
Naphthalene
Salicylic acid
Toluene
Uracil
QuinNet
Energy
0.119
0.050
0.078
0.101
0.101
0.080
0.096
(from paper)
Force
0.145
0.060
0.097
0.039
0,080
0.039
0.062
QuinNet
Energy
0.132
0.052
0.076
0.109
0.106
0.081
0.099
(1500 epochs)
Force
0.152
0.065
0,113
0.043
0,080
0.039
0.061
QuinNet+FreeCG
Energy
0.113
0.048
0.073
0.094
0.103
0.078
0.094
(1500 epochs)
Force
0.127
0.056
0.095
0.034
0.073
0.037
0.057
Table 12: Comparison of parameters and training speed. The mini-batch size is set to 32.
QuinNet+FreeCG QuinNet
FreeCG
VisNet
Number of parameters (M)
9.4
9.1
10.4
9.8
Iterations per second
4.19
4.45
3.62
3,95
Diaboota ma
Σ
Clatato
Ale B ma ma mib
(27)
ma m2
where m. m, = my. To illustrate, if we regard single multiplication and addition as the two basic operations, then combining two 1 irreps to form / = 2 irreps will use up 1 basic operation for ma = + 2, 3 operations for ma = 11, and 5 operations for my = O, making a total of 13 basic operations. One effective approach to understanding irreps is to view them vectors and scalars. The dot product between vectors requires only 5 basic operations, in contrast to the 13 operations mentioned earlier. Hence, the CG transform is extremely time-consuming. The table of basic operations for the CG transform between each pair is shown in Tab. g
extension of
A.7 ADDITIONAL EXPERIMENTS
Applications for full-atom proteins. AIMD-Chig dataset (Wang et al. (023a) comprises nearly 10,000 conformations of the 166-atom mini protein, Chignolin. These conformations were obtained including folded, unfolded, and metastable states. It is important to evaluate the performance of FreeCG on such real-world proteins. Following ViSNet (Wang et al. 2024), as shown in Fig. 4 we explore the energy landscape of Chignolin, where we sample six conformations located different parts of the landscape, covering folded and unfolded states. The energy and force predictions on these six conformations are compared with ViSNet. FreeCG succeeds in outperforming ViSNet for the most sampled conformations. Molecular dynamics simulations are from six different initial conformations. To assess the simulation stability, we calculate the Root Mean Square Deviation (RMSD) between each step in the trajectory and the initial conformation, shown in Fig. e). The results demonstrate satisfying performance of FreeCG on real-world proteins.
Applications for molecular dynamics simulation. We conduct molecular dynamics simulations for FreeCG on MD17 and compare the results with DFT calculations. We run a 300ps simulation for 0.5 fs, under a Nosé-Hoover thermostat 500K temperature. Like previous works (Fu et al. 2022 Wang et al. 2024] (023c), we are interested in the distribution of interatomic distances. Here, h(r) is defined as the probability density function of interatomic distances r. We plot the distribution as h(r). where h(r) are averaged along frames or predicted trajectories, and we desire the distribution is similar to the MD17 datasets calculated by DFT. The results are shown in Fig. 6 which shows that FreeCG is capable to well recover the interatomic dis- tances distribution. We also conduct molecular dynamics simulation on water (Ik version) (Fu et al.) 2022 Wu et al.] 2006) and LiPS (Batzner et al. 2022) datasets under periodic boundary conditions (PBCs) to evaluate how FreeCG performs on large molecular systems. We set the timestep to 0.25 fs and run total 200,000 steps for each type of molecules. We focus the recovery of radial distri- bution functions (RDFs) because they effectively describe structural and thermodynamic propertics. It is similarly calculated h(r) but under different constant multipliers. FreeCG finishes all the
each molecule. The time step is
22

---

Published as a conference paper at ICLR 2025
Table 13: Force prediction results OC20 S2EF validation set. The unit of force and energy MAE are eV/A and eV, respectively. The best numbers are marked in bold.
Id
OOD Ads
OOD Cat
OOD Both
Force Prediction
VisNet
0.0599
0.0673
0.0597 0.0536
0.0821 0.0717
FreeCG
0.0535
0.0614
Energy Prediction
VisNet
0.8807
1.0278 0.7586
0.8989 0.7294
1.3175
FreeCG
0.7157
0.8938
dynamics simulation with accurate recovery of the atomic distributions. The results for molecules under PBCs are shown in Fig. a
Results on OC20 dataset, S2EF task. OC20 dataset Chanussot et al. (2021) contains 1.2 million DFT relaxations of molecular adsorption onto catalytic surfaces, serving a large-scale benchmark for three key tasks in catalysis simulations: Structure to Energy and Force (S2EF), Initial Structure Relaxed Structure (IS2RS), and Initial Structure Relaxed Energy (IS2RE). We focus on the S2EF task, which involves predicting energy and forces based on atomic positions and types under Periodic Boundary Conditions (PBCs). Unlike other MLFF datasets, OC20 incorporates both or- ganic and inorganic materials relevant to catalysis research, which is also critical field for climate and energy applications. The dataset contains variably sized training sets and validation sets with distinct out-of-domain conditions: in-domain (ID), out-of-domain adsorbates (OOD Ads), out-of- domain catalysts (OOD Cat), and combined out-of-domain (OOD Both). Using the OC20 200K training evaluated our model on all validation sets. As shown in Table [3] FreeCG demon- strates significant performance advantages over VisNet across all OC20 benchmarks, highlighting its domain adaptability.
23

---

Published as a conference paper at ICLR 2025
b
a
d
c
f
e
Figure 7: The 3D structures of the data considered in this work. a. Aspirin in MD17 and rMD17. b. Ac-Ala3-NHMe in MD22. e. 1-Cyclohexene- 1 I-carboxylic acid in QM9. d. Chignolin. e. A single cell of water molecules under PBCs. f. single cell of LiPS under PBCs (Note that the cell is not cubic).
24

---

Published as a conference paper at ICLR 2025
RDF 00
RDF HH
RDF HO
FreeCG
2.5
30
DFT
RDF(O)
2.0
1.5
2-
20
1.0
10
0.0
15.0
13.6
12.5
Force MAE
10.0
7.5
5.0
5.0
5.2
3.5
2.5
1.4
0.7
0.0
FreeCG
NequlP
GemNet-T
ForceNet
PaiNN
DimeNet
Figure 8: Results on water- lk. Top: The RDFs molecular dynamic simulations for each bond of
water under PBCs. The unit of is Å, and the unit of RDF(r) is A Bottom: The force MAE comparison with other methods. The unit is V/A.
25

---

Published as a conference paper at ICLR 2025
RDF
25
FreeCG
2.0
RDF(0
10
14
12.8
12
11.7
Force MAE
10
8
6
4
3.7
3.2
2
1.2
1.3
FreeCG
NequiP
GemNet-T
ForceNet
PaiNN
DimeNet
Figure 9: Results on LiPS. Top: The RDFs in molecular dynamic simulations for LiPS under PBCs.
The unit r is A, and the unit of RDF Á Bottom: The force MAE comparison with other methods. The unit is me V/À.
26

---

Published as a conference paper at ICLR 2025
aspirin force MAE
ethanol force MAE
0.40
FreeCe
0.35
0.35
0.30
025
MWE
0.5
0.0
020
0.10
0.15
9.05
200
300
400
500
600 Epochs
700
800
900
1000
200
300
400
500
600
700
800
930
1000
Epochs
malonaldehyde force
naphthalene
0.6
0.18
9.16
0.4
MAG
0.12
8.0
6.2
0.06
9.04
1000
Epochs
salicylic acid force MAE
toluene force MAE
0.11
0.10
MAE
MAC
0.08
0.06
0.10
0.04
600
1000
600 Epochs
Epochs
uracil force MAE
0.18
0.16
0.14
0.10
0.08
200
300
430
500
6.00
700
800
900
1000
Epochs
Figure 10: QuinNet force prediction performance on MD17 when equipped with modules from FreeCG. The unit of force is kcal/mol/A.
27

---

Published as a conference paper at ICLR 2025
aspirin energy MAE
ethanol energy MAE
0.6
OunNet
0.30
0.5
MAE
0.4
MAE
0.15
0.05
800
600
900
Epochs
Epochs
malonaldehyde MAE
naphthalene energy MAE
0.6
0.5
MWE
0.4
0.4
0.3
0.2
0.1
0.1
200
300
400
500
000 Epochs
700
800
900
1000
200
300
400
500
600
700
800
930
1060
Epochs
salicylic_ energy
toluene energy
0.8
MAG
0.5
1.0
0.3
0.5
0.1
600
900
Epochs
uracil energy MAE
MAE
380
400
800
800
900
Epochs
Figure 11: QuinNet energy prediction performance MD17 when equipped with modules from FreeCG. The unit of energy is kcal/mol.
28


Reviews:
[{"review_id": "zx6kdJKQ47", "reviewer": "Reviewer_mCwJ", "summary": "This paper proposes FreeCG, a method that implements the CG transform layer on the permutation-invariant abstract edges, which allows complete freedom in the design of the layer without compromising the overall permutation equivariance. This can greatly improve the model’s expressiveness and decrease the computational demands.", "questions": "- Please refer to the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- This paper uses invariance transitivity with permutation-invariant abstract edges to solve the narrowness design space of CG transform.\n- This work further proposes a FreeCG model that contains Group CG transform with sparse path, abstract edges shuffling, and Attention enhancer to improve the representation power and efficiency.\n- The model shows good performance on several small molecule datasets.", "weaknesses": "- In Table 2, FreeCG is not effective in energy prediction on the rMD17 dataset. Could authors elaborate more on this?\n- In Table 4, other competitive baselines such as equiformerV2 should be included for a comprehensive comparison. \n- The datasets used in this paper are quite small and the results are sometimes not robust to evaluate the model performance, such as QM9. I would like to see the force and energy prediction performance on a large dataset such as OC20, to better understand the effective and efficiency of the FreeCG.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZJvHBRq9lw", "reviewer": "Reviewer_dHAu", "summary": "This paper proposes an approach to address the computational inefficiency of Clebsch-Gordan (CG) transforms in rotation-translation equivariant graph neural networks (EGNNs). Leveraging the invariance transitivity property, the proposed method, FreeCG, implements the CG transform layer on permutation-invariant abstract edges, enabling a more flexible layer design without compromising permutation equivariance. Additional architectural modifications are introduced to enhance model efficiency, with extensive empirical results demonstrating FreeCG’s performance.", "questions": "1. For abstract edge shuffling, did the authors try using a linear layer to mix different channels of $\\hat{E}_i^L$? Linear mixing could enhance information exchange effectively in this case.\n2. Did the authors consider applying FreeCG to more representative tensor product-based EGNNs? This could further validate the method’s efficacy.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. FreeCG could reduce the computational overhead of CG transforms and enhance expressivity by incorporating several architectural improvements.\n2. FreeCG achieves competitive or superior performance across multiple benchmarks and demonstrates compatibility with various EGNN architectures.", "weaknesses": "1. FreeCG’s approach to freeing CG transform space with abstract edges resembles the customized tensor product mechanism used in Allegro (see Eq. 13 in [1]). Allegro applies CG transforms to geometric features of a pair (similar to FreeCG’s node features) and environment embeddings of the pair (similar to FreeCG’s abstract edges from neighbors). While FreeCG focuses on atom-wise message passing with greater interaction between atom features and abstract edges, further clarification on the advantages of FreeCG over Allegro would strengthen the technical contribution.\n2. Introducing sparse paths may come with a trade-off in expressivity. The authors should provide an ablation study to analyze the impact of sparse paths on computational efficiency and performance.\n3. Results in Table 5 are somewhat unclear. If the final FreeCG model uses the best-performing modules, it should have a group number of 32, which conflicts with the number 8 reported in Table 7. Additionally, the ablations indicate that the primary performance improvement is due to increased group numbers, suggesting that the other modules’ contributions may be limited.\n4. Certain aspects of the methodology require additional clarity: a) Some notations are under-defined, e.g., $N$ in line 227, $\\bar{E_i^L}$, and $d\\bar{E_i^{L+1}}$ in line 277. b) The relationship between $d\\bar{E_i^{L+1}}$ and $\\bar{E_i^{L+1}}$ in line 277 is unclear, and more details are needed on how $\\bar{E_i^L}$ is constructed in the first layer.\n\n[1] Musaelian, A., Batzner, S., Johansson, A. et al. Learning local equivariant representations for large-scale atomistic dynamics. Nat Commun 14, 579 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bEQotSo3QG", "reviewer": "Reviewer_o4xZ", "summary": "The paper introduces a novel approach to Machine Learning Force Fields (MLFFs) that leverages permutation-invariant hidden features for efficient computation. These permutation-invariant features are aggregated over high-irreducible (high-irrep) edge features with attention weights, forming what is referred to as the “abstract edge” in the paper. The Clebsch-Gordan (CG) transformation is applied within each layer, and these permutation-invariant hidden features serve as the input of the kernel function for the CG transformation layer.\n\nTo improve the efficiency of the CG product, the authors propose a grouped CG transformation, which operates similarly to a group convolution. Additionally, to enhance the utility of the abstract edge, they introduce abstract edge shuffling and an Attention Enhancer mechanism.", "questions": "- How does the model ensure the SO(3) equivariance?\n- Why was the AIMD-Chig dataset chosen for memory usage and inference speed benchmarks?\n- Why does abstract edge shuffling improve the model? A brief discussion on the role of shuffling is needed.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper includes extensive experiments in the MLFF domain, providing empirical solid support for the model’s performance.\n\n- By aggregating high-irreducible (high-irrep) edge features with attention-weighted, permutation-invariant hidden features (abstract edges), the model achieves robust and flexible representations.", "weaknesses": "- The paper’s organization and readability could be improved. In particular, modified margins between paragraphs, figures, titles, and subtitles reduce overall readability.\n\n- The paper does not adequately address  $\\text{SO}(3)$ -equivariance, a critical concept in equivariant GNNs. Typically, in Clebsch-Gordan (CG) transformations, a radial distance-based kernel function is used to ensure  $\\text{SO}(3)$ -equivariance. However, FreeCG may lack  $\\text{SO}(3)$ -equivariance, as the aggregation of abstract edges is based on a weighted summation over components of each irreps. While this aspect is unclear, explicitly showing how the model achieves  $\\text{SO}(3)$ -equivariance would enhance understanding. If the model does not fully satisfy  $\\text{SO}(3)$ -equivariance, it would be helpful to justify using abstract edges within the CG transformation layer.\n\n- The paper should include recent results in Table 2 to demonstrate the model’s state-of-the-art (SOTA) status. Specifically, adding comparisons with Graph ACE [1] and PONITA would strengthen the claim.\n\n[1] Bochkarev, Anton, Yury Lysogorskiy, and Ralf Drautz. \"Graph Atomic Cluster Expansion for Semilocal Interactions beyond Equivariant Message Passing.\" Physical Review X 14.2 (2024): 021036.\n\n[2] Bekkers, Erik J., et al. \"Fast, Expressive $\\mathrm {SE}(n) $ Equivariant Networks through Weight-Sharing in Position-Orientation Space.\" The Twelfth International Conference on Learning Representations.\n\n- In Equation (26), the model employs a summation over all permutations, which may constrain the model complexity to  O(N \\times N!) . To fully discuss the model’s efficiency, scalability with respect to the number of atoms should also be considered.\n\n### Minor Issues (Typos, Formatting, and Readability, etc..)\n- Information about what kinds of loss are used is omitted.\n- On page 5,  $N$  appears undefined and may be a typo; perhaps it should be  $Z$ .\n- In Figure 3, the shaded regions impair readability; using a transparent background or adjusting the shading would improve clarity.\n- Unnecessary line breaks are present in Appendix equations (7) and (15).\n- Table 8 is unnecessarily expanded to fill the entire line width.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sDzilTlNMj", "reviewer": "Reviewer_h6Dk", "summary": "The paper introduced an informative and efficient model utilizing high-order irreps and CG transform. This model achieves state-of-the-art (SOTA) results in force prediction for MD17, rMD17, and MD22. However, while the authors asserted that this work represented a new paradigm for CG transform that can be integrated with other models, such as QuinNet, to enhance performance, the experimental results do not support this claim.", "questions": "For MILP, both model execution speed and memory usage are critically important. The proposed method seems to significantly increase the computational complexity of the model. It is crucial to include a comparison of the number of parameters and training speed.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposed an informative and efficient model with high-order irreps and CG transform.\n\nThe model achieved state-of-the-art (SOTA) results in force prediction for MD17, rMD17, MD22.", "weaknesses": "Although the author claimed that this work presented a new paradigm for CG transform that can be combined with other models, such as QuinNet, to achieve better performance, the experimental results did not demonstrate this. Simply analyzing the training curves did not allow readers to determine the impact of adding or removing FreeCG on QuinNet. Furthermore, while other models were trained for 3000 epochs, QuinNet was only trained for 1000 epochs. At this point, there is no clear advantage for QuinNet+FreeCG, and readers cannot ascertain whether extended training would improve QuinNet's performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)