FreeCG: Free the Design Space of Clebsch-Gordan Transform for Machine Learning Force Fields
This work addresses a bottleneck in MLFFs for chemistry, physics, and materials science by enabling more efficient and expressive geometric neural networks, though it is incremental in advancing existing CG transform methods.
The paper tackles the limitation of Clebsch-Gordan Transform (CG transform) in Machine Learning Force Fields (MLFFs) due to permutation-equivariance constraints, which reduces expressiveness and increases computational demands. It proposes FreeCG, a method that implements CG transform on permutation-invariant abstract edges, achieving state-of-the-art results with improvements up to 20% on force and property prediction datasets.
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.