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

TL;DR

FreeCG introduces a flexible, efficient, and expressive approach to the Clebsch-Gordan Transform for machine learning force fields, significantly improving accuracy and computational efficiency across multiple datasets.

Contribution

It proposes a novel free design space for the CG transform, enabling more expressive and efficient models without compromising symmetry, and extends this to enhance existing SOTA models.

Findings

Achieves state-of-the-art results on MD17, rMD17, MD22 datasets.

Improves property prediction accuracy on QM9 by over 15%.

Enhances the SOTA QuinNet model using the proposed paradigm.

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% 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 neural network designs. To demonstrate this, the recent SOTA, QuinNet, is also enhanced under our paradigm. Code will be publicly available.