FreeCG: Free the Design Space of Clebsch-Gordan Transform for Machine Learning Force Fields

## Abstract

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. This constraint results in reduced expressiveness of the model while simultaneously increasing computational demands.
To overcome this challenge, we first implement the CG transform layer on the permutation-invariant abstract edges generated from real edge information. We show that this approach allows complete freedom in the design of the layer without compromising the crucial symmetry. 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 % percent \% % and the maximum beyond 20 % percent \% % . The extensive real-world applications showcase high practicality. FreeCG introduces a novel paradigm for carrying out efficient and expressive CG transform in future geometric neural network designs. To demonstrate this, the recent SOTA, QuinNet, is also enhanced under our paradigm. Code will be publicly available.

## 1 Introduction

Machine Learning Force Fields (MLFFs) are of great importance for drug development (Chen et al., ), materials science (Liu et al., ), chemical reaction kinetics (Meuwly, ), nanotechnology (Wang et al., ), among others. It offers a satisfactory trade-off between accuracy and efficiency, which is expected to perform as powerful as Density Functional Theory (DFT) (Kohn & Sham, ) or other high accuracy references (Martin, ; Ceperley & Alder, ; Bartlett & Musiał, ), but with orders-of-magnitude speedup (Cui et al., ; Wang et al., ; ; Musaelian et al., ; Batzner et al., ; Drautz, ; Batatia et al., ; Thölke & De Fabritiis, ; Schütt et al., ; Chmiela et al., ).

Graph Neural Networks (GNNs) perform SOTA on several MLFFs benchmarks (Schütt et al., ; ). Group and group representation theory play important roles in the design of GNNs for MLFFs. Rotation invariance is generally required in these works, as we naturally require the potential energy unchanged w.r.t. rotations of the input molecule. A recent design trend is to maintain rotation, reflection, and translation equivariance in the design of neural networks. They hope the internal features can move with respect to the input molecule, enabling higher expressive power. GNNs that obey this property are called Equivariant Graph Neural Networks (EGNNs) (Thomas et al., ; Satorras et al., ; Gasteiger et al., ; Vaswani et al., ; Fuchs et al., ; Liao & Smidt, ).
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 degree irreps or tensors, showing significant performance boost (Batatia et al., ; Batzner et al., ; Musaelian et al., ; Gasteiger et al., ; Thomas et al., ; Simeon & De Fabritiis, ). 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 vectors, 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 equivariance, 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 equivariance, 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 we cannot naïvely 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 to produce the weights for each computation between the central atom and each neighboring one (Batzner et al., ; Musaelian et al., )).

In this work, to confront this challenge, we propose FreeCG. The model generates and refines geometric 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 is also permutation invariant, regardless of designs, and does not affect 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 to 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 et al., ; Musaelian et al., ), 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)-equivariance while being more efficient than keeping SE(3)-equivariance. The abstract edges shuffling, inspired by (Zhang et al., ), is 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 propose the concept of abstract edges, by which we design several new modules to resolve a 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 molecular property datasets QM9 reveal the SOTA performance of FreeCG, with several improvements beyond 15%. The real-world applications further indicate the practicality.

4) This work presents a new paradigm for CG transform in future research, extending beyond the design presented here. We further enhance QuinNet under this paradigm to demonstrate this point.

## 2 Related Works

Maintaining E(3)-/SE(3)-equivariance has become a popular trend in the design of neural networks for MLFFs. The challenge lies in how to properly construct geometric objects that are better capable of modeling the atomic environment. In general, the methods can be categorized into the following two lines:

Methods with geometric vectors. Several methods directly work on regular geometric vectors (Schütt et al., ; ; Gasteiger et al., ; ). The starting point of this line of works is about covering necessary information for MLFFs. For instance, SchNet (Schütt et al., ) first introduces continuous-filter to MLFFs, but only distance information is covered. DimeNet (Gasteiger et al., ) 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, ; Wang et al., ; ). This is commonly done via calculating different angles between atoms or surfaces (e.g., torsion angle, improper angle (Wang et al., ), dihedral angle (Wang et al., )).
The irreducible representations (irreps) and Clebsch-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.

Methods equipped with irreps and CG transform. Applying irreps and CG transform to Equivariant Graph Neural Networks (EGNNs) was first proposed as more of a conceptual framework, known as Tensor Field Network (Thomas et al., ). Several works have been proposed on top of this foundation. For instance, NequIP (Batzner et al., ) has implemented the idea to construct high-order irreps, and shows state-of-the-art (SOTA) results for MLFFs. Allegro (Musaelian et al., ) resolves the challenges of the scaling issue, which makes it possible to run parallelly on a large number of GPUs. SE(3)-Transformer (Fuchs et al., ) first proposes equivariant dot-product for generating self-attention. Equiformer (Liao & Smidt, ) further achieves E(3)-equivariance combining MLP attention and non-linear messages. MACE (Batatia et al., ) extends classic body order expansion methods, Atomic Cluster Expansion (ACE) (Drautz, ), to a hierarchical framework.
CG transform is a fundamental building block in these works, but with a limited design space, affecting both performance and efficiency. In this work, we completely free the design space of CG transform and propose a novel model, FreeCG, performing strong SOTA and showing high efficiency for MLFFs, which is also a new design paradigm for future works.

Figure: Figure 1: The architecture of a single layer of FreeCG. The cross-attention mechanism generates abstract edges through a permutation-invariant process. The abstract edges are also used to enhance the quality of the attention score, denoted as 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 O(3) equivariance. Abstract edges shuffling improves the information exchange between different irreps. The details for sparse path and abstract edges shuffling can be referred to Fig. 2.
Refer to caption: x1.png

## 3 Methods

### 3.1 Background

Group, equivariance and invariance. Permutation, rotation, and translation form different groups in group theory. Formally, a set with a binary operation $(G,*)$ is said to be a group if and only if the following conditions hold: 1) $g_{1}*g_{2}\in G,$ for any $g_{1},g_{2}\in G$ (closure) 2) $(g_{1}*g_{2})*g_{3}=g_{1}*(g_{2}*g_{3}),$ for any $g_{1},g_{2},g_{3}\in G$ (associativity) 3) There exists a group element $e\in G$, such that $g*e=e*g=g,$ for any $g\in G$. ($e$ identity element) 4) There is a group element $g^{\prime}$ w.r.t. $g$, such that $g*g^{\prime}=g^{\prime}*g=e$, for each $g\in G$ ($g^{\prime}$ inverse element). The group elements $g\in G$, according to the representation theory, can be represented as linear transformations $\mathcal{P}_{V}(g)\in GL(V)$ on vector space $V$. Given a function $f:X\to Y$, where $X$ and $Y$ are vector spaces. It is said to be $G$-equivariant if and only if $f(\mathcal{P}_{X}(g)x)=\mathcal{P}_{Y}(g)f(x)$, for any $g\in G$. $G$-invariance is a special case when $\mathcal{P}_{Y}(g)$ is an identity matrix. Permutation equivariance and ${\text{E}}(3)$-equivariance are two properties each layer of our model obeys. Permutation equivariance means the index of node or edge features will be consistent when passing a layer. ${\text{E}}(3)$-equivariance covers rotation, translation, and reflection, where the translation is explicitly guaranteed via only considering the relative distances between atoms, thus we consider ${\text{O}}(3)$-equivariance where translations are omitted. It is intuitive to correspondingly change directional features when the whole molecule rotates or reflects.

Tensor, irreps and CG transform. Tensors are high-dimensional generalizations of scalars, vectors, 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 product 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 ${\rm SO}(3)$ (omit reflection compared to ${\text{O}}(3)$) on 9-space (from tensor product of two $3\times 3$ rotation matrix) can be decomposed into $1\times 1$ ($l=0$), $3\times 3$ ($l=1$), and $5\times 5$ ($l=2$) irreps, which are called Wigner-D matrices. In EGNNs, we often project the distance vector between atoms onto the unit sphere $S^{2}$ with the central atom as the center of sphere. Actually, $S^{2}$ is homomorphic to the quotient group SO(3)$/$SO(2), thus it also has its own irreps, e.g., $l=0$ scalar and $l=1$ vector. $S^{2}$ irreps are the main features we maintain in our model, where irreps with degree $l$ has $2l+1$ 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 is known as CG transform, which utilizes CG coefficients to perform transformations. For instance, $A^{1,l_{2}l_{3}\mapsto l_{1}}_{m_{1}}=\sum_{m_{2},m_{3}}C^{l_{1}l_{2}l_{3}}_{m
_{1}m_{2}m_{3}}A^{2,l_{2}}_{m_{2}}A^{3,l_{3}}_{m_{3}}$, where $A^{l}$ are $S^{2}$ irreps, $m$ denotes the elements of irreps, and $C$ the CG coefficient. To satisfy ${\text{O}}(3)$, we consider an additional variable, parity $p$, which takes the values of $1$ or $-1$. Irreps with $p=-1$ will be inverse when the space is reflected, and $p=1$ unchanged. The above formula of CG transform becomes:

$$ $\displaystyle A^{1,l_{2}p_{2}l_{3}p_{3}\mapsto l_{1}p_{1}}_{m_{1}}=\mathbbm{1} _{(p_{1}=p_{2}p_{3})}\sum_{m_{2},m_{3}}C^{l_{1}l_{2}l_{3}}_{m_{1}m_{2}m_{3}}A^ {2,l_{2}p_{2}}_{m_{2}}A^{3,l_{3}p_{3}}_{m_{3}}$ (1) $$

where $\mathbbm{1}_{(expression)}$ is the indicator function, outputting $1$ if $expression$ is true, and 0 0 otherwise. Given a vector ($l=1$ $S^{2}$ irreps), we can lift it to irreps with arbitrary degree $l$ and $p=(-1)^{l}$, via a series of real spherical harmonics $(Y^{l}_{m=1},...,Y^{l}_{m=2l+1})$. For further details about group theory, we refer interested readers to related books and papers (Zee, ; Raczka & Barut, ; Thomas et al., ; Jeevanjee, ; Cohen et al., ).

Figure: 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 $l=1$ and $l=2$. Right: The shuffling strategy is to add a constant $k$ for the index of each abstract edge. The shuffled result is then added by $\hat{E}^{L}_{i}$, and get the final added value $d\overline{E}^{L+1}_{i}$.
Refer to caption: x2.png

### 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 $\{\bm{X},\bm{Z}\}$, the neural network $f_{\theta}$ with parameter $\theta$ aims to predict the energy, and by which it derives the predicted force on each atom. In each layer of NequIP (Batzner et al., ), messages from neighboring atoms are aggregated and combined with the features of the 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 NequIP:

$$ $\displaystyle\mathcal{L}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{acm_{o}}( \bm{X},\bm{N})=\mathbbm{1}_{(p_{o}=p_{e}p_{n})}\sum_{m_{e}m_{n}}C^{l_{o}l_{e}l _{n}}_{m_{o}m_{e}m_{n}}$ (2) $\displaystyle\sum_{b\in\mathcal{N}(a)}(R(\lVert\vec{r}_{ab}\rVert)^{l_{o}l_{e} l_{n}}_{c})Y^{l_{e}}_{m_{e}}(\frac{\vec{r}_{ab}}{\lVert\vec{r}_{ab}\rVert})N^{ l_{n}p_{n}}_{bcm_{n}}$ $$

where $\mathcal{N}(a)$ is the set of neighboring atoms of atom $a$. $R$ is a MLP. $\lVert*\rVert$ is Euclidean norm. $N_{b}$ is the features of node $b$. $\vec{r}_{ab}$ is the vector pointing from atom $a$ to $b$. Consider the vector function form of Eq. 2:
$\bm{\mathcal{L}}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{cm_{o}}=(\mathcal{L
}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{1cm_{o}},\mathcal{L}^{l_{e}p_{e}l_
{n}p_{n}\mapsto l_{o}p_{o}}_{2cm_{o}},...)$, which is permutation equivariant w.r.t. permutation operations acting on $\bm{X}$ and $\bm{N}$. Formally, it means $\bm{\mathcal{L}}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{cm_{o}}(\mathcal{P}
_{\bm{X}}\bm{X},\mathcal{P}_{\bm{N}}\bm{N})=\mathcal{P}_{\bm{\mathcal{L}}}\bm{
\mathcal{L}}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{cm_{o}}(\bm{X},\bm{N})$. Put simply, if we exchange the indexes of two atoms, for example, 1 and 2, and feed them into function $\bm{\mathcal{L}}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{cm_{o}}$, it equals to that we directly change the index 1 and 2 of the output of function $\bm{\mathcal{L}}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{cm_{o}}$, which is $(\mathcal{L}^{l_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{2cm_{o}},\mathcal{L}^{l
_{e}p_{e}l_{n}p_{n}\mapsto l_{o}p_{o}}_{1cm_{o}},...)$. This property is simple and very important for the molecular neural networks, as the properties of a molecule should not depend on the order in which these atoms are arranged.

Most works take this property for granted. However, the permutation equivariance is actually important but vulnerable. It 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 as $\mathcal{O}(\max\limits_{i}{\rm card}(\mathcal{N}(i)))$, where ${\rm card}(X)$ is the number of elements in set $X$. One cannot arbitrarily remove calculations for a specific neighboring atom because it 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 1 poses heavy computation challenges, as 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.5. 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 structures can be designed to maintain permutation equivariance. To address these problems, we aim to 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:

$$ $h(x)=h^{{}^{\prime}}(h_{1}(x),h_{2}(x),...)$ (3) $$

if $h_{*}(x)$ are all $G$-invariant, then, regardless of how we design $h^{{}^{\prime}}$, the overall function $h$ must be $G$-invariant as well. The proof is simple, as:

$$ $\displaystyle h^{{}^{\prime}}(h_{1}(\mathcal{P}_{X}(g)x),h_{2}(\mathcal{P}_{X} (g)x),...)=\mathcal{P}_{h}(e)h^{{}^{\prime}}(h_{1}(x),h_{2}(x),...)$ (4) $$

Here, the invariance we care about is w.r.t. the internal features of a given atoms and the neighbours. Specifically, it is the sum $\sum_{b\in\mathcal{N}(a)}(R(\lVert\vec{r}_{ab}\rVert)^{l_{o}l_{e}l_{n}}_{c})Y^
{l_{e}}_{m_{e}}(\frac{\vec{r}_{ab}}{\lVert\vec{r}_{ab}\rVert})N^{l_{n}p_{n}}_{
bcm_{n}}$, and the term for each $b$. Invariant components will not affect the equivariance of the layer. Thus, we first obtain a set of invariant functions $h_{*}$. Then, we can freely design the function $h^{{}^{\prime}}$ based on these invariant functions.

###### Problem 1

###### Problem 2

### 3.3 FreeCG

Abstract edges. The above proposition presents an elegant way to solve problem 1 and 2. The idea is to put CG transform inside the function $h^{{}^{\prime}}$, and by the conclusion, we can 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_{*}$’s to be the aggregation of edge features. We call $h_{*}$’s abstract edges. For the concrete design, we take the transformer architecture in ViSNet (Wang et al., ) as an efficient tool to construct abstract edges. The detailed information of the implementation is in the Sec. A.2. In ViSNet, each edge maintains high-degree features $E_{ij}=E^{l=1}_{ij}\oplus E^{l=2}_{ij}$ consisting of irreps $E^{l}_{ij}=Y^{l}(\vec{r}_{ij}/\lVert\vec{r}_{ij}\rVert)$. The above features are invariant to layer index $L$. The computed attention $a^{L,t}_{ij}$ is multiplied to each edge. The sum of them $\hat{E}^{L}_{i,t}=\sum_{ij\in\mathcal{E}(i)}a^{L}_{ij,t}E^{L}_{ij}$ forms an temporary abstract edge, where we omit the degree $l$, and $L$ the index of the layer. $t$ denotes the index of the $t$-th abstract edges $(\hat{E}^{L}_{i,t=1},\hat{E}^{L}_{i,t=2},...)$.
In the original ViSNet, it was used to update the geometric feature $d\overline{E}^{L+1}_{i}=\hat{E}^{L}_{i}+o^{L,1}_{i}\cdot{\rm Linear}(\overline
{E}^{L}_{i})$, where ${\rm Linear}$ is a fully-connected linear operation, which performs across the dimension of $t$, thus does not break the equivariance, and $o^{L,1}_{i}$ is a variable generated from the node feature, as we will introduce in Sec. A.2. We leverage the fact that $\hat{E}^{L}$ fits the requirements of $h_{*}$ in 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 $h_{*}$, 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 in Problem 1 is controlled to be constant. The proposition above gives us enough freedom to construct the CG transform function $h$, expanding the design space to maximum, alleviating Problem 2. The idea is to use CG transform to replace the updating mechanism of $\overline{E}$ in ViSNet. A naïve attempt is to directly take the CG transform between $\overline{E}^{L}$ and $\hat{E}^{L}$ to acquire $\overline{E}^{L+1}$. However, we want to further decrease the $O(T^{2})$ time complexity for the CG transform, where $T$ is the number of abstract edges, even though it is a constant number. Leveraging the unlimited freedom in constructing $h$, and taking inspiration of group convolution (Krizhevsky et al., ), we propose group CG transform (distinct from the group in group theory). We first split the abstract edges of $\overline{E}^{L}$ and $\hat{E}^{L}$ into groups, where each index of abstract edge belongs to some group $U_{g}$ , the integer $g$ ranges from $1$ to $G$, and $G$ a hyper-parameter for the number of total groups. Then a group CG transform acts as:

$$ $\displaystyle d\overline{E}^{\prime L+1,l_{o},p_{o}}_{i,t_{o}m_{o}}=\mathbbm{1 }_{(p_{o}=p_{1}p_{2})}o^{L,1}\sum_{l_{1},l_{2}}\sum_{m_{1},m_{2}}C^{l_{o},l_{1 },l_{2}}_{m_{o}m_{1}m_{2}}$ (5) $\displaystyle\sum_{t_{1},t_{2}\in U_{g}}W^{l_{o},l_{1},l_{2}}_{t_{o}t_{1}t_{2} }\rm{Linear}(\overline{E}^{L}_{i})^{l_{1},p_{1}}_{t_{1}m_{1}}\hat{E}^{L,l_{2}, p_{2}}_{i,t_{2}m_{2}}$ $$

where $t_{o}\in U_{g}$. The group CG transform decreases the time complexity to $O(T^{2}/G)$. Here, the parameters $W$ for CG transform are also worth emphasizing. They are not necessary to be kept the same across different abstract edges $t$ 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 to enhance the model expressive ability. In contrast to 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 O(3) equivariance. It is because we often need to consider both $p=1$ and $p=-1$ for a single $l$ for O(3), thus the total computation is quadrupled, and memory usage is doubled. Here we propose a method to keep O(3) while being more efficient than SO(3). We only keep $(l=1,p=-1)$ and $(l=2,p=1)$, which is same as the order of directly using spherical harmonics. In such way, It suffices that each output irreps containing information from both input irreps through CG transform, as $(l=1,p=-1)*(l=2,p=1)\mapsto(l=1,p=-1)$, $(l=1,p=-1)*(l=1,p=-1)\mapsto(l=2,p=1)$, $(l=2,p=1)*(l=2,p=1)\mapsto(l=2,p=1)$, and $(l=1,p=-1)*(l=1,p=-1)\mapsto(l=2,p=1)$. There are only 4 path in contrast to 8 path for SO(3), being O(3) equivariant but more efficient than being SO(3) equivariant, illustrated in Fig. 2.

Abstract edges shuffling. Inspired by ShuffleNet (Zhang et al., ), we can also shuffle the abstract edges to make the information exchanged comprehensively. 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 1 again. Theoretically, the shuffling strategy can be arbitrary as 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 Results section.

Abstract edges enhance cross-attention. The Transformer integrates neighboring atoms information in molecular tasks through cross-attention mechanism, which aims to capture relations for those atoms exhibiting strong interatomic correlations. In order to better utilize abstract edge information, we leverage it to augment the generation of attention scores. 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_{ij}=\mathop{\rm SiLU}(Q_{i}\odot K_{j})$, where $\odot$ represents dot product, and SiLU is the activation function. Note that ViSNet is different from other transformer-based models where $A_{ij}$ is scaled by the SiLU instead of Softmax across different $j$. We integrate the information of abstract edges by:

$$ $A_{ij}=\mathop{\rm SiLU}(Q_{i}\odot K_{j}+\max\limits_{t}(\overline{E}^{L}_{j, t}\odot E_{ij}))$ (6) $$

where $\max\limits t(\overline{E}^{L}_{j,t}\odot E{ij})$ denotes the maximum value of the dot product between each abstract edge $\overline{E}^{L}_{j,t}$ and the real edge features $E_{ij}$ across different abstract edges. $E_{ij}$ does not have an $L$ 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 $i$ and $j$. The generation of abstract edges and the operations performed on them are illustrated in Fig. 1, respectively.

**Table 1: Performances on MD17 dataset. The results are reported in mean abosolute error (MAE). The energies and forces are measured in kcal/mol and kcal/mol/Å, respectively. The best numbers are marked in bold.**
| Molecule | SchNet | DimeNet | PaiNN | SpookeyNet | ET | GemNet | NequIP | SO3KRATES | 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 |
| Ethanol | 0.08 | 0.064 | 0.064 | 0.052 | 0.052 | - | 0.051 | 0.052 | 0.051 | 0.050 | 0.049 |
| Malondialdehyde | 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 |
| Toluene | 0.12 | 0.102 | 0.095 | 0.094 | 0.074 | - | 0.092 | 0.095 | 0.074 | 0.080 | 0.076 |
| Uracil | 0.14 | 0.115 | 0.106 | 0.105 | 0.095 | - | 0.104 | 0.103 | 0.095 | 0.096 | 0.097 |
| Force Prediction |  |  |  |  |  |  |  |  |  |  |  |
| Aspirin | 1.35 | 0.499 | 0.338 | 0.258 | 0.253 | 0.217 | 0.184 | 0.236 | 0.155 | 0.145 | 0.122 |
| Ethanol | 0.39 | 0.230 | 0.224 | 0.094 | 0.109 | 0.085 | 0.071 | 0.096 | 0.060 | 0.060 | 0.053 |
| Malondialdehyde | 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 |

**Table 2: Performances on rMD17 dataset. The results are reported in MAE. The energies and forces are measured in kcal/mol and kcal/mol/Å, respectively. The best numbers are marked in bold.**
| Molecule | UNiTE | GemNet | NequIP | MACE | Allergo | BOTNet | ViSNet | QuinNet | FreeCG |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Energy Prediction |  |  |  |  |  |  |  |  |  |
| Aspirin | 0.055 | - | 0.0530 | 0.0507 | 0.0530 | 0.0530 | 0.0445 | 0.0486 | 0.0530 |
| Azobenzene | 0.025 | - | 0.0161 | 0.0277 | 0.0277 | 0.0161 | 0.0156 | 0.0394 | 0.0217 |
| Benzene | 0.002 | - | 0.0009 | 0.0092 | 0.0069 | 0.0007 | 0.0007 | 0.0096 | 0.0107 |
| Ethanol | 0.014 | - | 0.0092 | 0.0032 | 0.0092 | 0.0092 | 0.0078 | 0.0096 | 0.0087 |
| Malonaldehyde | 0.025 | - | 0.0184 | 0.0185 | 0.0138 | 0.185 | 0.0132 | 0.0168 | 0.0146 |
| Naphthalene | 0.011 | - | 0.0046 | 0.1153 | 0.0046 | 0.0046 | 0.0057 | 0.0174 | 0.0118 |
| Paracetamol | 0.044 | - | 0.0323 | 0.0300 | 0.0346 | 0.0300 | 0.0258 | 0.0362 | 0.0392 |
| Salicylic acid | 0.017 | - | 0.0161 | 0.0208 | 0.0208 | 0.0185 | 0.0161 | 0.033 | 0.0233 |
| Toluene | 0.010 | - | 0.0069 | 0.0115 | 0.0092 | 0.0069 | 0.0059 | 0.0139 | 0.0334 |
| Uracil | 0.013 | - | 0.0092 | 0.0115 | 0.0138 | 0.0092 | 0.0069 | 0.0149 | 0.0116 |
| Force Prediction |  |  |  |  |  |  |  |  |  |
| Aspirin | 0.175 | 0.2191 | 0.1891 | 0.1522 | 0.1684 | 0.1900 | 0.1520 | 0.1429 | 0.1212 |
| Azobenzene | 0.097 | - | 0.0669 | 0.0692 | 0.0600 | 0.0761 | 0.0585 | 0.0513 | 0.0486 |
| Benzene | 0.017 | 0.0115 | 0.0069 | 0.0069 | 0.0046 | 0.0069 | 0.0056 | 0.0047 | 0.0056 |
| Ethanol | 0.085 | 0.083 | 0.0646 | 0.0484 | 0.0484 | 0.0738 | 0.0522 | 0.0516 | 0.0438 |
| Malonaldehyde | 0.152 | 0.1522 | 0.0118 | 0.0946 | 0.0830 | 0.1338 | 0.0893 | 0.0875 | 0.0802 |
| Naphthalene | 0.060 | 0.0438 | 0.0300 | 0.0369 | 0.0208 | 0.0415 | 0.0291 | 0.0242 | 0.0228 |
| Paracetamol | 0.164 | - | 0.1361 | 0.1107 | 0.1130 | 0.1338 | 0.1029 | 0.0979 | 0.0840 |
| Salicylic acid | 0.088 | 0.1222 | 0.0922 | 0.0715 | 0.0669 | 0.0992 | 0.0795 | 0.0771 | 0.0648 |
| Toluene | 0.058 | 0.0507 | 0.0369 | 0.0350 | 0.0415 | 0.0438 | 0.0264 | 0.0244 | 0.0239 |
| Uracil | 0.088 | 0.0876 | 0.0669 | 0.0484 | 0.0415 | 0.0738 | 0.0495 | 0.0487 | 0.0446 |

## 4 Experiments

To evaluate the performance of FreeCG, we collect molecular force field datasets, on which we compare our methods with other SOTA MLFFs, which include small molecules dataset MD17 (Chmiela et al., ) with its revised version, rMD17 (Christensen & Von Lilienfeld, ), and large molecules dataset MD22 (Chmiela et al., ). To test the generalization capacity of the proposed FreeCG, we also evaluate the performance of FreeCG on a standard molecule property prediction dataset, QM9 (Ruddigkeit et al., ; Ramakrishnan et al., ). We take popular SOTA models into the comparison, including sGDML (Chmiela et al., ), SchNet (Schütt et al., ), DimeNet (Gasteiger et al., ), SphereNet (Liu et al., ), PaxNet (Zhang et al., ), PaiNN (Schütt et al., ), SpookyNet (Unke et al., ), ForceNet (Hu et al., ), ET (Thölke & De Fabritiis, ), GemNet (Gasteiger et al., ), ComENet (Wang et al., ), NequIP (Batzner et al., ), UniTE (Qiao et al., ), SO3KRATES (Frank et al., ), MACE (Batatia et al., ), Allegro (Musaelian et al., ), BOTNet (Batatia et al., ), ViSNet (Wang et al., ), ViSNet-LSRM (Li et al., ), and QuinNet (Wang et al., ). To asses the practicality of FreeCG on real-world tasks, Molecular dynamics simulations are run for MD17 molecules and two periodic systems, water (Fu et al., ; Wu et al., ) and LiPS (Batzner et al., ) under PBCs. FreeCG is also evaluated on the conformation space of a 166-atom mini-protein, Chignolin (Wang et al., ). The results reveal that FreeCG is capable to make accurate predictions on force and energy efficiently, and it also exhibits strong practicality on real-world applications. The ablation on each component and corresponding hyper-parameters of FreeCG are also presented. The common settings and extra experiments are reported in the Sec. A.1.

### 4.1 Comparison with state-of-the-arts for MLFFs

Force field on small molecules and periodic systems. MD17 is a famous molecular dynamics benchmark for small molecules. FreeCG outperforms others in all force prediction tasks. It also significantly decreases the force prediction errors for the most hard-to-predict molecule in this datasets, aspirin, by 15%. Remarkably, FreeCG also decreases the MAE by over 10% for ethanol, naphthalene, and salicylic acid. FreeCG does not have a particular preference for the size of molecules. It demonstrates strong performance for aspirin (180.2 g/mol) and excels on ethanol (46.1 g/mol). The energy prediction is also competitive when compared to other SOTA methods. rMD17 is the revised version of MD17. It recomputed the trajectories of each atom with higher accuracy. The force prediction accuracy of FreeCG is still leading in majority of the molecules. It improves the force results compared to the baseline model, ViSNet, in all molecules except for the 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. 1 and 2, respectively. The accuracy of force prediction on periodic systems, water and LiPS, are also evaluated. Fig. 8 and 9 show that, FreeCG achieves the best performance across other methods, and decrease the MAE value of the second best by 50% for water.

Force field on large molecules. MD22 is a large molecules benchmark adopted by several studies (Wang et al., ; ; Li et al., ). As shown in Tab. 3, it reveals that FreeCG also performs well for large scale data. 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 on Ac-Ala3-NHMe are both around 20%. The performances for the other models are not consistent 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, as a classical kernel method.

**Table 3: Performances on MD22 dataset. The results are reported in MAE. The energies and forces are measured in kcal/mol and kcal/mol/Å, 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. (), unlike Wang et al. (); Chmiela et al. (; ), 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-CG-CG | 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-NHMe | 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 |
| Double-walled nanotube | 0.520 | 0.362 | 0.292 | 0.3391 | 0.2767 | 0.2473 | 0.2449 |

### 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. ViSNet also performs the second best in most measures. Although these two methods are proposed to be MLFFs, but they are even more comparable than others in molecular property prediction tasks. The results are in Tab. 4.

Figure: Figure 3: The speed and memory occupation of FreeCG compared with other SoTA models. Numbers are reported based on a single chignolin molecule. The right three models are based on high-order irreps and CG transform.
Refer to caption: x3.png

### 4.3 Efficiency benchmarking

AIMD-Chig dataset is taken as a benchmark for testing the memory usage and inference speed. We compare the inference speed and memory usage of FreeCG with ViSNet, NequIP, and Allegro. The results are shown in Fig. 3. FreeCG adds little extra time and memory cost, compared to the baseline model, ViSNet. It is also the most efficient one for both memory and speed, compared to the other two CG transform-based methods, NequIP and Allegro. The overall results prove the effciency of FreeCG. The number of groups in group CG transform also impacts the inference speed. Fig. 4 shows the theoretical number of paths and the actual inference time for different group numbers. A computation analysis for CG transform can be referred to Sec. A.5.

### 4.4 Ablation study

We conduct ablations on different modules we propose, as well as the strategies for abstract edges shuffling. The results are shown in Tab. 5. 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/G$. Here we also study the influence of the shuffling strategies. We adopt $0.5*T/G$, $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 a large number of groups appears to be a good choice.

**Table 4: Molecular property prediction on QM9 dataset. The results are reported in MAE. The best numbers are marked in bold.**
| Target |  | SchNet | EGNN | DimeNet++ | PaiNN | SphereNet | PaxNet | ET | ComENet | ViSNet | FreeCG |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\mu$ | mD | 33 | 29 | 29.7 | 12 | 24.5 | 10.8 | 11 | 24.5 | 9.5 | 11.4 |
| $\alpha$ | m$a^{3}_{0}$ | 235 | 71 | 43.5 | 45 | 44.9 | 44.7 | 59 | 45.2 | 41.1 | 38.2 |
| $\epsilon_{HOMO}$ | meV | 41 | 29 | 24.6 | 27.6 | 22.8 | 22.8 | 20.3 | 23.1 | 17.3 | 16.6 |
| $\epsilon_{LUMO}$ | meV | 34 | 25 | 19.5 | 20.4 | 18.9 | 19.2 | 17.5 | 19.8 | 14.8 | 13.5 |
| $\Delta\epsilon$ | meV | 63 | 48 | 32.6 | 45.7 | 31.1 | 31 | 36.1 | 32.4 | 31.7 | 31.5 |
| $\langle R^{2}\rangle$ | m$a^{2}_{0}$ | 73 | 106 | 331 | 66 | 268 | 93 | 33 | 259 | 29.8 | 82.1 |
| $ZPVE$ | meV | 1.7 | 1.55 | 1.21 | 1.28 | 1.12 | 1.17 | 1.84 | 1.2 | 1.56 | 1.10 |
| $U_{0}$ | meV | 14 | 11 | 6.32 | 5.85 | 6.26 | 5.9 | 6.15 | 6.59 | 4.23 | 4.11 |
| $U$ | meV | 19 | 12 | 6.28 | 5.83 | 6.36 | 5.92 | 6.38 | 6.82 | 4.25 | 4.51 |
| $H$ | meV | 14 | 12 | 6.53 | 5.98 | 6.33 | 6.04 | 6.16 | 6.86 | 4.52 | 4.13 |
| $G$ | meV | 14 | 12 | 7.56 | 7.35 | 7.78 | 7.14 | 7.62 | 7.98 | 5.86 | 5.65 |
| $C_{v}$ | mcal/mol K | 33 | 31 | 23 | 24 | 22 | 23.1 | 26 | 24 | 23 | 20.4 |

**Table 5: Ablation on different modules. Abstract edges shuffling and Attention enhancer are added upon the best choices of the above modules, with respect to the validation loss.**
| Method | Aspirin |  |  |
| --- | --- | --- | --- |
| Val loss | Energy | Force |  |
| ViSNet | - | 0.116 | 0.155 |
| + Group CG transform |  |  |  |
| 8 groups | 0.0509 | 0.123 | 0.144 |
| 32 groups | 0.0416 | 0.112 | 0.129 |
| + Abstract edges shuffling |  |  |  |
| 1-group shuffle | 0.0401 | 0.112 | 0.128 |
| 0.5-group shuffle | 0.0396 | 0.110 | 0.128 |
| 1.5-group shuffle | 0.0384 | 0.111 | 0.125 |
| + Attention enhancer | 0.0345 | 0.110 | 0.122 |

### 4.5 Extension to other models

The free CG transform proposed in this work also presents a paradigm to enhance other geometric models. To demonstrate that the FreeCG concept is applicable to other models, we use the former SOTA model, QuinNet, as an example to illustrate how effectively FreeCG can be extended to other architectures. 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 train both the vanilla QuinNet and QuinNet equipped with FreeCG modules for 1000 epochs, evaluating on the test set every 200 epochs. As shown in Fig. 10 and 11, 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.

## 5 Conclusion

This work proposes FreeCG, an equivariant neural network that frees the design space of CG transform. 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-protein Chignolin. As we show FreeCG helps improve QuinNet, it also introduces a new paradigm for expressive and efficient CG transform-based neural network design in the future.

## Appendix 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 a single GPU. Pytorch 1.10.0 is used as the basic machine learning python library. For the CG transform operations, we adopt e3nn 0.5.1. Matplotlib 3.0.3 is utilized for plotting. The details can be referred to Tab. 6. We report the hyperparameters used in Tab. 7. For training/validation/test splits, we follow previous works (Wang et al., ; ; Fu et al., ). We pick up the model for evaluating on 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, Exponential Moving Average (EMA) is adopted to generate the model weights. The detailed training configurations are shown in Tab. 7.

**Table 6: Hardware and software settings.**
| Hardware | Software |  |  |  |
| --- | --- | --- | --- | --- |
| CPU | GPU | Neural Network | Equivariance | Plotting |
| Intel^® Xeon^®<br>Gold 6330 CPU @ 2.00GHz | NVIDIA Tesla A100 | Pytorch 1.10.0 | e3nn 0.5.1 | Matplotlib 3.0.3 |

**Table 7: Hyperparameters for each dataset.**
| Hyperparameter | MD17 | rMD17 | MD22 | QM9 | Water-1k | LiPS | Chignolin |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Initial learning rate | 4e-4, 2e-4 | 2e-4 | 2e-4, 1e-4 | 1e-4 | 5e-4 | 1e-3 | 2e-4 |
| Learning rate decay factor | 0.8 |  |  |  |  |  |  |
| Learning rate decay patience | 30 | 30 | 30 | 15 | 5 | 5 | 10 |
| Learning rate warmup step | 1000 | 1000 | 1000 | 10000 | 1000 | 1000 | 1000 |
| Optimizer | AdamW ($\beta(0.9,0.999)$) |  |  |  |  |  |  |
| Epoch | 3000 | 3000 | 3000 | 1500 | 1000 | 100 | 3000 |
| Batch size | 4 | 4 | 4 | 32 | 1 | 1 | 4 |
| Number of layers | 9 | 9 | 9 | 9 | 9 | 9 | 6 |
| Cutoff | 5.0, 4.0 | 5.0 | 5.0, 4.0 | 5.0 | 6.0 | 6.0 | 5.0 |
| 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 |
| Dimension of latent feature | 256 | 256 | 256 | 512 | 256 | 256 | 128 |
| Number of groups | 8 |  |  |  |  |  |  |
| Output head | Equivariant/Scalar |  |  |  |  |  |  |
| EMA rate | 0.999 |  |  |  |  |  |  |

### 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 $\{\bm{X}=\vec{r}_{0},\vec{r}_{1},\vec{r}_{2},...,\vec{r}_{N}),\bm{Z}=(z_{1},z_
{2},...,z_{n})\}$, where $\vec{r}\in\mathbbm{R}^{3}$ the Cartesian coordinates of atom, and $z$ 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 $h_{i}={\rm embedding}(z_{i})\in\mathbbm{R}^{C}$. $C$ is the dimension of the latent space. For each atom, we only consider neighboring atoms within a given radius $\mathcal{N}(i)$, where we maintain the distance vector from the central atom to the neighboring atoms, and lift them to $(l=1,p=-1)$ and $(l=2,p=1)$ irreps $E_{ij}\in\mathbbm{R}^{3+5}$ via real spherical harmonics applied on the unit vector $E_{ij}=Y^{l=1}(\vec{r}_{ij}/\lVert\vec{r}_{ij}\rVert)\oplus Y^{l=2}(\vec{r}_{
ij}/\lVert\vec{r}_{ij}\rVert)$, where we also calculate the corresponding Euclidean norm $\lVert\vec{r}_{ij}\rVert$. The Euclidean norms of vectors are then converted to high-dimension scalar features (edge attributes) $f_{ij}={\rm RBF}(\vec{r}_{ij})\in\mathbbm{R}^{C}$ by radial basis functions (RBFs). We also maintain zero-initialized abstract edges $\overline{E}^{L=0}_{i}={\bm{0}}$ for each node to be updated in the following layers. We assign the same number of abstract edges as the dimension of the latent features, such that additional operations to align the dimension numbers are not required.

Intermediate layers. Here, we use a superscript $L$ to denote the index of layer that the features are in. The message-passing between atoms is implemented by a transformer architecture. For each atom $i$, the neighboring atoms $j\in\mathcal{N}(i)$ will send messages to $i$, and the messages are aggregated to update the information of $i$. The query, key, and value of the node features are first calculated, respectively: $q_{i}=f_{q}(h_{i})$, $k_{j}=f_{k}(h_{j})$, $v_{j}=f_{v}(h_{j})$. The edge attributes are also converted to auxiliary terms $dk_{j}=f_{dk}(f_{ij})$ and $dv_{j}=f_{dv}(f_{ij})$ to modulate keys and and values of atoms. Here functions $f$ are all fully-connected linear operations. Then we calculate the cross-attention between $i$ and $j$, which is

$$ $\displaystyle a_{ij}={\rm SiLU}\bigg{(}{\rm Cutoff}(\lVert\vec{r}_{ij}\rVert)q _{i}k_{j}dk_{j}+$ (7) $\displaystyle{\rm AttEnhancer}(\vec{r}_{ij},\overline{E}^{L}_{j})\bigg{)}$ $$

where ${\rm Cutoff}(\cdot)$ is a cosine cutoff function, and ${\rm AttEnhancer}(\cdot)$ the proposed attention enhancer module, as we will formulate its details. First, recall the dimension of $\overline{E}^{L}_{i}\in\mathbbm{R}^{C*8}$ and $\vec{r}_{ij}\in\mathbbm{R}^{8}$. Each of the $C$ abstract edges will undergo a dot product with $\vec{r}_{ij}$. The highest value among them will be the output of ${\rm AttEnhancer}$. In other word,

$$ ${\rm AttEnhancer}(\overline{E}^{L}_{i},\vec{r}_{ij})=\max_{C}(\overline{E}^{L} _{i}\odot\vec{r}_{ij})$ (8) $$

as we introduce in Eq. 6. Then the values are multiplied with $dv$ and attention.

$$ $\hat{v}_{j\mapsto i}^{L}=v_{j}\cdot dv_{j}\cdot a_{ij}$ (9) $$

It then undergoes two different fully-connected operations to generate two coefficients $s_{1}$ and $s_{2}$. They are used to generate the abstract edges:

$$ $\hat{E}^{L}_{j\mapsto i}=\overline{E}^{L}_{i}\cdot s_{1}+E_{ij}\cdot s_{2}$ (10) $$

This variable, together with $\hat{v}_{j\mapsto i}$, are aggregated by sum:

$$ $\hat{E}^{L}_{i}=\sum_{j\in\mathcal{N}(i)}\hat{E}^{L}_{j\mapsto i}$ (11) $$

$$ $\hat{v}^{L}_{i}=\sum_{j\in\mathcal{N}(i)}\hat{v}^{L}_{j\mapsto i}$ (12) $$

$\hat{v}^{L}_{i}$ then converts to three variables for further operation:

$$ $o^{L,1}_{i},o^{L,2}_{i},o^{L,3}_{i}={\rm Linear}(\hat{v}^{L}_{i})$ (13) $$

$\hat{E}^{L}_{i}\in\mathbbm{R}^{C*8}$ and $\overline{E}^{L}_{i}\in\mathbbm{R}^{C*8}$ are used for the following group CG transform and abstract edges shuffling. First $\overline{E}^{L}_{i}\in\mathbbm{R}^{C*8}$ undergoes a fully-connected operation along $C$ dimension, and multiply with $o^{L,1}_{i}$, which means we get $o^{L,1}_{i}\cdot{\rm Linear}(\overline{E}^{L}_{i})$.
It, together with $\hat{E}^{L}_{i}\in\mathbbm{R}^{C*8}$, are then divided into $G$ groups along $C$ dimension, where we get $\hat{E}^{L}_{i,t\in G_{g}}\in\mathbbm{R}^{\frac{C}{G}*8}$, and $(o^{L,1}_{i}\cdot{\rm Linear}(\overline{E}^{L}_{i}))_{t\in G_{g}}\in\mathbbm{R
}^{\frac{C}{G}*8}$. Then, we perform CG transform between two variables in a fully connected form with learnable weights, and concatenate the results to generate $d\overline{E}^{\prime L+1}_{i}$ before shuffling, as shown in Eq. 5. For the shuffling strategies, we add $\frac{3C}{2G}$ to each index of the abstract edges $\overline{E}^{\prime L+1}_{i}$. Then, it is added with $\hat{E}^{L}_{i}$ to form a residual structure, as we show here:

$$ $d\overline{E}^{L+1}_{i}={\rm shuffle}(d\overline{E}^{\prime L+1}_{i})+\hat{E}^ {L}_{i}$ (14) $$

where $d\overline{E}^{L+1}_{i}$ is added to $\overline{E}^{L}_{i}$ to obtain $\overline{E}^{L+1}_{i}$. Next, we update $h$ and $f$. We first show the update for $h$:

$$ $\displaystyle dh^{L+1}_{i}=h^{L}_{i}+\bigg{(}{\rm Linear_{1}}(\overline{E}^{L} _{i})\odot{\rm Linear_{2}}(\overline{E}^{L}_{i})\bigg{)}\cdot o^{L,2}_{i}$ (15) $\displaystyle+o^{L,3}_{i}$ $$

To update $f$, we follow ViSNet to leverage rejection of vectors:

$$ $\displaystyle df^{L+1}_{ij}=f^{L}_{ij}+{\rm RejCalc_{trg}}(\overline{E}^{L}_{i },\vec{r}_{ij})\odot$ (16) $\displaystyle{\rm RejCalc_{src}}(\overline{E}^{L}_{i},\vec{r}_{ij})\cdot{\rm SiLU }({\rm Linear}(f^{L}_{ij}))$ $$

where rejection calculation module $RejCalc$ is:

$$ $RejCalc_{mode}(a,b)=a-({\rm Linear_{mode}}(a)\odot b)\cdot b$ (17) $$

The updated $\overline{E}^{L+1}$, $h^{L+1}$, and $f^{L+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 on energy-conservative field, which means we derive the force from the predicted potential energy. Following ViSNet (Wang et al., ) and PaiNN (Schütt et al., ), we predict the potential energy of the molecule via equivariant gated module.

$$ $h^{L+1}_{i},u^{L+1}_{i}={\rm MLP}\bigg{(}{\rm Concat}(h_{i}^{L},\lVert{\rm Linear _{1}}(\overline{E}^{L}_{i})\rVert)\bigg{)}$ (18) $$

where ${\rm MLP}$ is an $1$-hidden layer multi-layer preceptor. There is one more step to update $\overline{E}^{L+1}_{i}$:

$$ $\overline{E}^{L+1}_{i}={\rm Linear_{2}}(\overline{E}^{L}_{i})\cdot u^{L+1}_{i}$ (19) $$

These calculations are then repeated twice in succession. There is also an alternative head design for force field prediction which uses only scalar features $h$. Under this setting, Eq. 18 and (19) are replaced by:

$$ $h^{\bm{L}}={\rm MLP}(h^{{\bm{L}}-1})$ (20) $$

In Tab. 7 where we introduce our hyper-parameter 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_{i}^{\bm{L}}\in\mathbbm{R}$:

$$ $y=\sum_{i}h^{\bm{L}}_{i}$ (21) $$

and the force is the negative gradients of the total energy:

$$ $F_{i}=-\nabla_{\vec{r}_{i}}y$ (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 $\vec{r}_{c}$, which is:

$$ $\vec{r}_{c}=\frac{\sum_{i}m_{i}\cdot\vec{r}_{i}}{\sum_{i}m_{i}}$ (23) $$

For molecular dipole, the formula is:

$$ $\mu=\left\|\sum_{i}\overline{E}_{i}^{\bm{L}}+h_{i}^{\bm{L}}(\vec{r}_{i}-\vec{r }_{c})\right\|$ (24) $$

and for electronic spatial extent:

$$ $\langle R^{2}\rangle=\sum_{i}h_{i}^{\bm{L}}\lVert\vec{r}_{i}-\vec{r}_{c}\rVert$ (25) $$

The output head can be easily adapted for different tasks, providing flexibility in property prediction.

Figure: Figure 4: Efficiency analysis of group CG transform. Left: 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 $l$. Full CG transform denotes not using sparse path.
Refer to caption: x4.png

**Table 8: The basic operation number for each type of CG transform. $l_{o}$ 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.**
| $l_{o}=2$ | 0 | 1 | 2 |  | $l_{o}=1$ | 0 | 1 | 2 |  | $l_{o}=2$ | 0 | 1 | 2 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 0 | 1 |  |  |  | 0 | - | 3 | - |  | 0 | - | - | 5 |
| 1 |  | 3 |  |  | 1 | 3 | 6 | 9 |  | 1 | - | 9 | 12 |
| 2 |  |  | 5 |  | 2 | - | 9 | 12 |  | 2 | 5 | 12 | 19 |

### A.3 Datasets details

MD17 and rMD17. They are both molecular dynamics datasets for small molecules. MD17 (Chmiela et al., ), proposed by Chmiela, S., et al. contains ab-initio level molecular dynamics trajectories. Four types of data are included in the dataset: atomic numbers, atomic positions, molecular energy, and the force acting on each atom. To alleivate the noise during the trajectory computation, Christensen, A. S. et al. also propose revised MD17 (rMD17) (Christensen & Von Lilienfeld, ), where molecular trajectories are calculated at the PBE/def2-SVP level of theory. The precision of the calculated trajectories is 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 are sampled between 400K and 500K at 1fs resolution. The energy and 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 1 kcal/mol/Å in the original paper (Chmiela et al., ). 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. It is a subset of the GDB-17 database (Ruddigkeit et al., ). The data is calculated at B3LYP/6-31G(2df,p) based DFT level of accuracy. Since the attributes vary for different properties, we use distinct output head for each, as discussed in Sec. A.2.

Chignolin. The AIMD-Chig dataset (Wang et al., ) 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 are 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 to benchmark efficiency, following (Wang et al., ).

Periodic systems. Periodic Boundary Conditions (PBCs) are vital in molecular dynamics simulations of periodic systems as 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., ) in (Fu et al., ). The authors provides with several sizes of training sets, including -1k, -10k, 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 development. It can help predict key performance metrics such as capacity, energy density, and cycle life, aiding in the development of next-generation lithium-ion batteries. We follow the same train/val split as (Fu et al., ). Some 3D structures of the data can be referred to Fig. 7.

Figure: Figure 5: 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 carbonyl oxygen on the D3 backbone and the nitrogen on 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 are sampled at the localization highlighted in the energy 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. c. The RMSD ($\mathrm{\SIUnitSymbolAngstrom}$) during the molecular dynamics simulation. The shaded area denotes the values of standard derivations. The RMSD values are obtained by taking average of 10 trajectories.
Refer to caption: x5.png

### 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:

$$ $\hat{E}^{L}_{i}=\sum_{j\in\mathcal{N}(i)}\hat{E}^{L}_{j\mapsto i}=\sum_{p\in P }\sum_{j\in\mathcal{N}(i)}\frac{\mathcal{P}(p)\hat{E}^{L}_{j\mapsto i}}{{\rm Card }(P)}$ (26) $$

where $P$ is the set for all permutation operations, here we omit the subscript of $\mathcal{P}$ for specific spaces to work on. Here, note we are proving that the abstract edges for each atom are permutation invariant, and we can freely design CG transform per atom, thus the permutation is applied to $j$ but not $i$. It sums over all the permutation operations, and thus the last step is permutation invariant. Then, it suffices to show that each of the previous step are all at least permutation equivariant. It also suffices to show they are 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 $x$ and $y$, then those $a_{ij}$ that $x$ or $y$ shows up in the subscript for $j$ will exchange with each other, and so do $\hat{v}_{j\mapsto i}$ and $\hat{E}^{L}_{j\mapsto i}$. Thus, the rest steps are equivariant w.r.t. single switch, and so they are permutation equivariant. Therefore, we conclude our proof that abstract edges are permutation invariant.

### A.5 Analysis on the efficiency of CG transform

The CG transform consists of two steps: 1) performing a tensor product between two irreps, and 2) decomposing the resulting tensors into irreps. These transforms are actually quadratic homogeneous polynomials. For the sake of convenience, we discuss SO(3) group here. Recall the CG transform formula:

$$ $D^{l_{a}l_{b}\mapsto l_{d}}_{m_{d}}=\sum_{m_{a},m_{b}}C^{l_{d}l_{a}l_{b}}_{m_{ d}m_{a}m_{b}}A^{l_{a}}_{m_{a}}B^{l_{b}}_{m_{b}}$ (27) $$

where $m_{a}+m_{b}=m_{d}$. To illustrate, if we regard single multiplication and addition as the two basic operations, then combining two $l=1$ irreps to form a $l=2$ irreps will use up 1 basic operation for $m_{d}=\pm 2$, 3 operations for $m_{d}=\pm 1$, and 5 operations for $m_{d}=0$, making a total of 13 basic operations. One effective approach to understanding irreps is to view them as an extension of 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. 8.

Figure: Figure 6: The distributions of interatomic distances $r$ during molecular dynamic simulations of MD17. The unit of $r$ is $\mathrm{\SIUnitSymbolAngstrom}$, and the unit of $h(r)$ is $$\mathrm{\SIUnitSymbolAngstrom}$^{-1}$.
Refer to caption: x6.png

### A.6 Additional experiments

Applications for full-atom proteins. AIMD-Chig dataset (Wang et al., ) 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., ), as shown in Fig. 5, we explore the energy landscape of Chignolin, where we sample six conformations located at 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 run 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. 5(c). 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 each molecule. The time step is set to 0.5 fs, under a Nosé–Hoover thermostat at 500K temperature. Like previous works (Fu et al., ; Wang et al., ; ), 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)$ w.r.t. $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 distances distribution. We also conduct molecular dynamics simulation on water (1k version) (Fu et al., ; Wu et al., ) and LiPS (Batzner et al., ) 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 on the recovery of radial distribution functions (RDFs) because they effectively describe structural and thermodynamic properties. It is similarly calculated as $h(r)$ but under different constant multipliers. FreeCG finishes all the dynamics simulation with accurate recovery of the atomic distributions. The results for molecules under PBCs are shown in Fig. 8 and 9.

Figure: Figure 7: The 3D structures of the data considered in this work. a. Aspirin in MD17 and rMD17. b. Ac-Ala3-NHMe in MD22. c. 1-Cyclohexene-1-carboxylic acid in QM9. d. Chignolin. e. A single cell of water molecules under PBCs. f. A single cell of LiPS under PBCs (Note that the cell is not cubic).
Refer to caption: x7.png

Figure: Figure 8: Results on water-1k. Top: The RDFs in molecular dynamic simulations for each bond of water under PBCs. The unit of $r$ is $\mathrm{\SIUnitSymbolAngstrom}$, and the unit of ${\rm RDF}(r)$ is $$\mathrm{\SIUnitSymbolAngstrom}$^{-1}$. Bottom: The force MAE comparison with other methods. The unit is meV/$\mathrm{\SIUnitSymbolAngstrom}$.
Refer to caption: x8.png

Figure: Figure 9: Results on LiPS. Top: The RDFs in molecular dynamic simulations for LiPS under PBCs. The unit of $r$ is $\mathrm{\SIUnitSymbolAngstrom}$, and the unit of ${\rm RDF}(r)$ is $$\mathrm{\SIUnitSymbolAngstrom}$^{-1}$. Bottom: The force MAE comparison with other methods. The unit is meV/$\mathrm{\SIUnitSymbolAngstrom}$.
Refer to caption: x9.png

Figure: Figure 10: QuinNet force prediction performance on MD17 when equipped with modules from FreeCG. The unit of force is kcal/mol/Å.
Refer to caption: x10.png

Figure: Figure 11: QuinNet energy prediction performance on MD17 when equipped with modules from FreeCG. The unit of energy is kcal/mol.
Refer to caption: x11.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Please refer to the weakness.
soundness: 2
strengths: - This paper uses invariance transitivity with permutation-invariant abstract edges to solve the narrowness design space of CG transform.
- 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.
- The model shows good performance on several small molecule datasets.
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.
weaknesses: - In Table 2, FreeCG is not effective in energy prediction on the rMD17 dataset. Could authors elaborate more on this?
- In Table 4, other competitive baselines such as equiformerV2 should be included for a comprehensive comparison. 
- 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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
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.
2. Did the authors consider applying FreeCG to more representative tensor product-based EGNNs? This could further validate the method’s efficacy.
soundness: 2
strengths: 1. FreeCG could reduce the computational overhead of CG transforms and enhance expressivity by incorporating several architectural improvements.
2. FreeCG achieves competitive or superior performance across multiple benchmarks and demonstrates compatibility with various EGNN architectures.
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.
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.
2. 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.
3. 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.
4. 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.

[1] Musaelian, A., Batzner, S., Johansson, A. et al. Learning local equivariant representations for large-scale atomistic dynamics. Nat Commun 14, 579 (2023).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - How does the model ensure the SO(3) equivariance?
- Why was the AIMD-Chig dataset chosen for memory usage and inference speed benchmarks?
- Why does abstract edge shuffling improve the model? A brief discussion on the role of shuffling is needed.
soundness: 2
strengths: - The paper includes extensive experiments in the MLFF domain, providing empirical solid support for the model’s performance.

- By aggregating high-irreducible (high-irrep) edge features with attention-weighted, permutation-invariant hidden features (abstract edges), the model achieves robust and flexible representations.
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.

To 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.
weaknesses: - The paper’s organization and readability could be improved. In particular, modified margins between paragraphs, figures, titles, and subtitles reduce overall readability.

- 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.

- 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.

[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.

[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.

- 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.

### Minor Issues (Typos, Formatting, and Readability, etc..)
- Information about what kinds of loss are used is omitted.
- On page 5,  $N$  appears undefined and may be a typo; perhaps it should be  $Z$ .
- In Figure 3, the shaded regions impair readability; using a transparent background or adjusting the shading would improve clarity.
- Unnecessary line breaks are present in Appendix equations (7) and (15).
- Table 8 is unnecessarily expanded to fill the entire line width.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 2
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.
soundness: 3
strengths: The paper proposed an informative and efficient model with high-order irreps and CG transform.

The model achieved state-of-the-art (SOTA) results in force prediction for MD17, rMD17, MD22.
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.
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.

Meta-review: This paper introduces a clever, highly expressive neural parameterization of force fields which obeys the requisite equivariances without increasing computational complexity.  The new parameterization is then applied to a suite of challenging molecular benchmarks where it achieves strong results across the board, often improving SOTA by quite a margin.

The reviewers initially criticised clarity and the scope of experiments and comparisons, but in the rebuttal phase the authors put in a massive effort to alleviate these criticisms, run many new experiments, and clarify things. There has been some concerns about SO(3) equivariance by o4xZ but it is clear that this is satisfied and the authors have argued with both derivations and numerics.

Overall the rebuttal phase made it clear that this is a solid, original contribution to the field of machine learned force fields and with the all the improvements that the authors made and the positive assessments from the reviewers, I can only recommend acceptance.

Decision: Accept (Poster)