Coarse-grained graph architectures for all-atom force predictions

TL;DR

This paper presents CGAA-FF, a coarse-grained machine learning framework for all-atom force predictions that improves efficiency and can be integrated into various architectures for soft-matter simulations.

Contribution

It introduces a novel coarse-grained message passing approach within an all-atom force field, enhancing efficiency and flexibility of molecular simulations.

Findings

Achieves 0.201 and 0.253 eV/Å accuracy on tested systems.

Provides 10-fold and 5-fold higher computational speed and memory efficiency.

Can be integrated into any equivariant architecture for soft-matter systems.

Abstract

We introduce a machine-learning framework termed coarse-grained all-atom force field (CGAA-FF), which incorporates coarse-grained message passing within an all-atom force field using equivariant nature of graph models. The CGAA-FF model employs grain embedding to encode atomistic coordinates into nodes representing grains rather than individual atoms, enabling predictions of both grain-level energies and atom-level forces. Tested on EC/EMC organic electrolytes and RDX crystalline and disordered phases, CGAA-FF achieves 0.201 and 0.253 eV A-1, respectively, while providing about 10-fold and 5-fold higher computational speed and memory efficiency, respectively, than conventional MLIPs. Since this CGAA framework can be integrated into any equivariant architecture, we believe this work opens the door to efficient all-atom simulations of soft-matter systems.