Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yujie Zeng, Wenlong He, Ihor Vasyltsov, Jiaxin Wei, Ying Zhang, Lin Chen, Yuehua Dai

TL;DR

The paper tackles instability in MD simulations using GNN-based force fields, especially on out-of-distribution data. It introduces a correlation-based training objective that minimizes edge-feature correlation, combined with a dynamic cosine-like scheduler to balance stability with energy/force accuracy, and an empirical MD-stability metric to quantify robustness. Across NequIP, Allegro, and GemNet-T on Hafnium Oxide datasets, the approach improves long-time MD stability with modest overhead, at times extending stable simulation times from fractions of a picosecond to tens of picoseconds, while providing a more reliable evaluation of stability than traditional metrics. The method is architecture-agnostic and particularly benefits edge-feature representations, offering a practical path to more robust ML-based force fields for semiconductor materials, albeit with some accuracy trade-offs on certain models.

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.

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

TL;DR

The paper tackles instability in MD simulations using GNN-based force fields, especially on out-of-distribution data. It introduces a correlation-based training objective that minimizes edge-feature correlation, combined with a dynamic cosine-like scheduler to balance stability with energy/force accuracy, and an empirical MD-stability metric to quantify robustness. Across NequIP, Allegro, and GemNet-T on Hafnium Oxide datasets, the approach improves long-time MD stability with modest overhead, at times extending stable simulation times from fractions of a picosecond to tens of picoseconds, while providing a more reliable evaluation of stability than traditional metrics. The method is architecture-agnostic and particularly benefits edge-feature representations, offering a practical path to more robust ML-based force fields for semiconductor materials, albeit with some accuracy trade-offs on certain models.

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.

Paper Structure

This paper contains 21 sections, 8 equations, 9 figures, 11 tables.

Table of Contents

  1. Introduction
  2. Motivation and key contributions
  3. Related Work
  4. GNNFF models
  5. MD simulation stability
  6. Methodology
  7. Overview
  8. Feature correlation calculation
  9. Correlation loss function
  10. Dynamic coefficient scheduler
  11. Empirical metric to evaluate MD stability
  12. Experimental validation
  13. Evaluation on GNNFF models
  14. Accuracy
  15. Stability
  16. ...and 6 more sections

Figures (9)

  • Figure 1: MD simulation result with baseline Allegro model and our method
  • Figure 2: General flow of the proposed method
  • Figure 3: Feature correlation and MD simulation stability of GNNFF models of different layers
  • Figure 4: Feature correlation and MD simulation stability of GNNFF models with our method
  • Figure 5: Simulation result of Hf:O =1:1.1 with baseline Allegro model and optimized Allegro model. (a)-(c) shows the simulation status of step 20, 70 and 150 with baseline Allegro model. (d)-(f) shows the simulation status of step 20, 70 and 150 with optimized Allegro model.
  • ...and 4 more figures