Improving the Stability of GNN Force Field Models by Reducing Feature Correlation
This addresses a critical stability issue in GNNFF models for semiconductor material research, offering an incremental improvement to enhance reliability in molecular dynamics simulations.
The paper tackles the instability of Graph Neural Network Force Field (GNNFF) models in long-time molecular dynamics simulations, especially for out-of-distribution data, by proposing a feature correlation reduction method that improves stability with less than 3% computational overhead, extending stable simulation time from 0.03ps to 10ps for the Allegro model.
Recently, Graph Neural Network based Force Field (GNNFF) models are widely used in Molecular Dynamics (MD) simulation, which is one of the most cost-effective means in semiconductor material research. However, even such models provide high accuracy in energy and force Mean Absolute Error (MAE) over trained (in-distribution) datasets, they often become unstable during long-time MD simulation when used for out-of-distribution datasets. In this paper, we propose a feature correlation based method for GNNFF models to enhance the stability of MD simulation. We reveal the negative relationship between feature correlation and the stability of GNNFF models, and design a loss function with a dynamic loss coefficient scheduler to reduce edge feature correlation that can be applied in general GNNFF training. We also propose an empirical metric to evaluate the stability in MD simulation. Experiments show our method can significantly improve stability for GNNFF models especially in out-of-distribution data with less than 3% computational overhead. For example, we can ensure the stable MD simulation time from 0.03ps to 10ps for Allegro model.