Scaling Transferable Coarse-graining with Mean Force Matching

TL;DR

It is demonstrated that it is possible to scale machine learning architectures for coarse-graining, enabling highly accurate and transferable models, and the advantages of mean force matching are shown both theoretically and through exhaustive benchmarking using thermodynamic consistency as the primary metric of accuracy.

Abstract

Coarse-grained molecular dynamics often sacrifices accuracy and transferability for computational efficiency, but the use of machine learned potentials is helping coarse-grained models attain performance on par with atomistic molecular dynamics. Nevertheless, developing representations of the coarse-grained potential energy surface faces severe scaling challenges due to the extreme data demands of widely used "bottom-up" coarse-graining objectives. In this work, we show that mean force matching, a strategy for training thermodynamically consistent coarse-grained models, requires 50x fewer training samples and 87% less total atomistic simulation time, while obtaining better accuracy on the potential of mean force for unseen proteins compared to other commonly used objectives. By systematically removing noise from the objective function, we demonstrate that it is possible to scale machine learning architectures for coarse-graining, enabling highly accurate and transferable models. We show the advantages of mean force matching both theoretically and through exhaustive benchmarking using thermodynamic consistency as the primary metric of accuracy.

Figures (12)

• Figure 1: Mean Force Matching (MFM) enables efficient training of transferable coarse-grained protein force fields with enhanced scalability and zero-shot accuracy. (Left) Distilled Dataset: mean forces extracted from constrained atomistic MD simulations reduce noise compared to instantaneous forces, enabling efficient MLIP training. (Center) Enhanced Scalability: test loss scaling of MFM as a function of computational budget for various model sizes. (Right) Zero-Shot Accuracy: MFM CG model produces free energy surface that closely resembles atomistic reference for Trp-cage
• Figure 2: Data efficiency comparison of MFM and FM (a) Test loss scaling of MFM and FM on randomly sampled subsets of the full dataset (b) Test loss for MFM and FM as a function of total simulation time required for the full dataset. All models use the MACE architecture and are evaluated on mean force estimates from a held-out test of 50 CATH domains.
• Figure 3: Computational costs of training and inference (a) Wall-clock time per training epoch of MFM, MFM 100K, FM, and SM objectives across SchNet, MACE, and eSEN architectures (b) Inference time as a function of protein sequence length for models trained via MFM 100K.
• Figure 4: Free energy surface (FES) of Trp-cage (a) Atomistic reference FES derived from over 30 $\mu$s of explicit solvent MD at 300 K with representative structures from each metastable state. Dihedral of the last four alpha carbons (shown in red) was used for the first CV and fraction of native contacts as the second. (b) FES produced via umbrella sampling by CG models across different training objectives (MFM, FM, and SM) and model architectures (SchNet, MACE, eSEN)
• Figure 5: RMSD to native conformation of folding trajectory of Trp-cage starting from unfolded configuration