Sharpness-Aware Minimization for Robust Molecular Dynamics Simulations

Sharpness-aware minimization (SAM) is a novel regularization technique that takes advantage of not only the training error but also the landscape geometry of model parameters to improve model robustness. Although SAM has demonstrated the state-of-the-art (SOTA) performance in image classification, its applicability to physical system is yet to be examined. An ideal testbed is neural-network quantum molecular dynamics (NNQMD) simulations that accurately predict material properties, but the stability of their trajectories is severely limited by thermal noise. In this paper, we demonstrate for the first time that SAM regularizer achieves an order-of-magnitude reduction of the out-of-sample error in potential energy prediction using several SOTA models. Comparing NNQMD datasets with distinct structural characteristics, we found that SAM consistently reduces the out-of-sample error for a crystal dataset at high temperatures with enhanced thermal noise, thus proving the concept of SAM-enhanced robust NNQMD, while no clear trend was observed with an amorphous dataset. Our result suggests a possible correlation between materials structure and model parameter landscape.