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

Authors: Shihao Shao, Haoran Geng, Zun Wang, Qinghua Cui
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Machine Learning Force Fields, Equivariant Graph Neural Network, Clebsch-Gordan Transform

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729596772428

thread_type:
official_review

<thread_object>
creator: reviewer_h6dk

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

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730270708566

thread_type:
official_review

<thread_object>
creator: reviewer_o4xz

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

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730486928437

thread_type:
official_review

<thread_object>
creator: reviewer_dhau

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

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.

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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680707861

thread_type:
official_review

<thread_object>
creator: reviewer_mcwj

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

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.

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.

questions:
- Please refer to the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731989028968

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
Dear Reviewers and ACs,

We would like to sincerely appreciate the time and effort you devoted to our paper. We are glad that all reviewers acknowledge the novelty and efficacy of our FreeCG for machine learning force fields. The extensiveness of the experiments that took months for us to complete are also widely acknowledged (Reviewer dHAu, o4xZ, and h6Dk). 

We have carefully addressed each of reviewer's concerns, and added several experiments and results suggested by reviewers at this stage to our revised paper. Here are the following details:

 loss weight of 0.9\*force MAE+0.1\*energy MAE
|         |        | UNiTE | GemNet | NequIP |  MACE  | Allegro | BOTNet | VisNet | QuinNet | ACE$^{g}$  | P$\Theta$NITA |   FreeCG   |
|---------|--------|:-----:|:------:|:------:|:------:|:-------:|:------:|:------:|:-------:|:----------:|:-------:|:----------:|
| Aspirin | Energy | 0.055 |    -   | 0.0530 | 0.0507 |  0.0530 | 0.0530 | 0.0445 |  0.0486 | 0.0392 | 0.0392 | **0.0374** |
|         | Force  | 0.175 | 0.2191 | 0.1891 | 0.1522 |  0.1684 | 0.1900 | 0.1520 |  0.1429 | 0.1407  | 0.1338 | **0.1225** | 

Ablation on sparse path
|                   | 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          |

QuinNet equipped with FreeCG trained for 1500 epochs
| Model | Metric | 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** |



A major issue about the equivariance of FreeCG has been addressed in the response to Reviewer o4xZ who suspected that irreps under the same order $l$ cannot be weighted sum. We have clarified that it was a misunderstanding and presented rigorous proof and experimental results to show the point. That being said, it is common to suspect something or make mistake (we deeply apologize if irreps under same $l$ cannot be weighted summed up, but each of the authors believe it can be and there are so many papers could be taken as evidences) as we did sometimes, and we still greatly appreciate the other comments on the writing details from Reviewer o4xZ and carefully addressed them in one-to-one response. We also agree that potentiality for chemical weapon is worth mentioning, as we added the content in the revised paper. Here we provide another snippet to show the equivariance:

```
from e3nn import o3
import torch

rot = o3.wigner_D(torch.tensor(2),torch.tensor(0.1),torch.tensor(0.2),torch.tensor(0.3))

a = torch.rand(5)
b = torch.rand(5)

2.5*(rot@a) + 1.4*(rot@b) - rot@(2.5*(a) + 1.4*(b))
```
The results would be something like:
```
tensor([-2.3842e-08,  0.0000e+00,  0.0000e+00,  1.1921e-08,  1.7881e-08])
```
One could arbitrarily change 2.5 and 1.4 as the weights. We pick a non-trivial order $l=2$. The error in this case is in 10$^{-8}$ level, which is enough for claiming the equivariance.


Overall, we humbly suggest FreeCG is a substantial contribution to the MLFFs community. It is with reasonable motivation, novel technical contribution, and significant experimental SOTA on extensive test benchmarks. We respectfully think we have solved major concerns raised by reviewers, and hope that one could see the merits it has.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.