Knowledge Distillation of Noisy Force Labels for Improved Coarse-Grained Force Fields

TL;DR

This paper introduces a knowledge distillation approach to improve coarse-grained molecular force fields by denoising force labels and training student models with ensemble teacher predictions.

Contribution

It presents a novel framework that enhances coarse-grained force field accuracy by leveraging teacher models trained on noisy forces and distilling their knowledge into student models.

Findings

Training on ensemble teacher forces improves force field quality.

Denoising force labels enhances stability of coarse-grained models.

Framework validated on complex molecular fluid with improved property predictions.

Abstract

Molecular dynamics simulations are an integral tool for studying the atomistic behavior of materials under diverse conditions. However, they can be computationally demanding in wall-clock time, especially for large systems, which limits the time and length scales accessible. Coarse-grained (CG) models reduce computational expense by grouping atoms into simplified representations commonly called beads, but sacrifice atomic detail and introduce mapping noise, complicating the training of machine-learned surrogates. Moreover, because CG models inherently include entropic contributions, they cannot be fit directly to all-atom energies, leaving instantaneous, noisy forces as the only state-specific quantities available for training. Here, we apply a knowledge distillation framework by first training an initial CG neural network potential (the teacher) solely on AA-to-CG mapped forces to denoise those labels, then distill its force and energy predictions to train refined CG models (the student) in both single- and ensemble-training setups while exploring different force and energy target combinations. We validate this framework on a complex molecular fluid, a deep eutectic solvent, by evaluating two-, three-, and many-body properties and compare the CG and all-atom results. Our findings demonstrate that training a student model on ensemble teacher-predicted forces and per-bead energies improves the quality and stability of CG force fields.