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

arXiv: https://arxiv.org/abs/2502.12548 · HTML: https://mlsanity.com/paper/2502.12548/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 15 of its 44 references that are in the map.
- First posted on arXiv 18 Feb 2025.

## Builds on (most cited first, 15 of 15)

- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations](https://mlsanity.com/paper/2210.07237/) (2022)
- [Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-Scale Graph Networks](https://mlsanity.com/paper/2204.10348/) (2022)
- [Feature Overcorrelation in Deep Graph Neural Networks: A New Perspective](https://mlsanity.com/paper/2206.07743/) (2022)
- [EGraFFBench: Evaluation of Equivariant Graph Neural Network Force Fields for Atomistic Simulations](https://mlsanity.com/paper/2310.02428/) (2023)
- [Fast Neural Network Approach for Direct Covariant Forces Prediction in Complex Multi-Element Extended Systems](https://mlsanity.com/paper/1905.02791/) (2019)
- [Data efficiency and extrapolation trends in neural network interatomic potentials](https://mlsanity.com/paper/2302.05823/) (2023)

## Abstract (arXiv)

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.