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

The Clebsch-Gordan Transform (CG transform) effectively encodes many-body interactions. Many studies have proven its accuracy in depicting atomic environments, although this comes with high computational needs. The computational burden of this challenge is hard to reduce due to the need for permutation equivariance, which limits the design space of the CG transform layer. We show that, implementing the CG transform layer on permutation-invariant inputs allows complete freedom in the design of this layer without affecting symmetry. Developing further on this premise, our idea is to create a CG transform layer that operates on permutation-invariant abstract edges generated from real edge information. We bring in 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 property prediction in QM9 datasets with notable enhancement. The extensibility to other models is also examined. Molecular dynamics simulations are carried out on MD17 and other periodic systems, including water and LiPS, showcasing the capacity for real-world applications of FreeCG. It introduces a novel paradigm for carrying out efficient and expressive CG transform in future geometric neural network designs.