Improving the Stability of GNN Force Field Models by Reducing Feature Correlation

TL;DR

This paper introduces a feature correlation reduction method to improve the long-term stability of GNN-based force field models in molecular dynamics simulations, especially for out-of-distribution data.

Contribution

It reveals the negative impact of feature correlation on GNNFF stability and proposes a loss function with a dynamic scheduler to mitigate this issue.

Findings

Significantly improves MD simulation stability for GNNFF models.

Extends stable simulation time from 0.03ps to 10ps.

Achieves these improvements with less than 3% additional computational cost.

Abstract

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.