Beyond Force Metrics: Pre-Training MLFFs for Stable MD Simulations

TL;DR

Problem: ML force fields can achieve low force prediction errors yet yield unstable MD trajectories, undermining long-time simulations. Approach: compare direct MD17 aspirin training to OC20 pretraining followed by MD17 fine-tuning using GemNet-T, evaluating stability, latent structure, and local force behavior. Key findings: pre-training yields dramatically longer stable trajectories, more structured latent representations, smoother local force responses, and better local force-difference consistency, with only modest gains in force MAE. Significance: demonstrates that large-scale pretraining enhances robustness and generalization of MLFFs for MD, guiding evaluation metrics beyond force error alone.

Abstract

Machine-learning force fields (MLFFs) have emerged as a promising solution for speeding up ab initio molecular dynamics (MD) simulations, where accurate force predictions are critical but often computationally expensive. In this work, we employ GemNet-T, a graph neural network model, as an MLFF and investigate two training strategies: (1) direct training on MD17 (10K samples) without pre-training, and (2) pre-training on the large-scale OC20 dataset followed by fine-tuning on MD17 (10K). While both approaches achieve low force mean absolute errors (MAEs), reaching 5 meV/A per atom, we find that lower force errors do not necessarily guarantee stable MD simulations. Notably, the pre-trained GemNet-T model yields significantly improved simulation stability, sustaining trajectories up to three times longer than the model trained from scratch. By analyzing local properties of the learned force fields, we find that pre-training produces more structured latent representations, smoother force responses to local geometric changes, and more consistent force differences between nearby configurations, all of which contribute to more stable and reliable MD simulations. These findings underscore the value of pre-training on large, diverse datasets to capture complex molecular interactions and highlight that force MAE alone is not always a sufficient metric of MD simulation stability.

Figures (7)

• Figure 1: Overview of the MLFF framework for MD simulations. a. GemNet-T is trained on molecular (MD17) and catalyst (OC20) datasets, then used as a force predictor in MD simulations. b. Atomic positions are iteratively updated based on forces predicted by the trained GemNet-T model. c. Simplified illustration of the GemNet-T architecture, which incorporates multi-hop geometric message passing and directional embeddings to model atomic interactions.
• Figure 2: Impact of pre-training on force accuracy and simulation stability. a. Force MAE (meV/Å) over epochs for the GemNet-T model trained without pre-training (MD17 only) versus with pre-training (MD17 pre-training followed by OC20 finetuning). b. Instability onset time over epochs. Error bars represent results from three independent simulations.
• Figure 3: Pair-distance distribution comparisons for Aspirin. Each plot shows normalized interatomic distance distributions $h(r)$ from reference ab initio MD (blue dashed) and MLFF-based MD (orange dashed) as a function of distance $r$ (Å). Subplots a. and b. show results for pre-trained models at epochs 50 and 300; c. and d. show non-pre-trained models at the same epochs. Red labels indicate instability onset time (ps).
• Figure 4: 2D projections of the learned embeddings for the eight MD17 molecular systems and OC20. a. t-SNE projection for the pre-trained model. b. t-SNE projection for the non-pre-trained model. c. UMAP projection for the pre-trained model. d. UMAP projection for the non-pre-trained model.
• Figure 5: Silhouette score statistics for the 2D projections in Figure \ref{['fig1']}. a. Distribution of per-sample silhouette scores for the t-SNE projection of the pre-trained model. b. Distribution for the t-SNE projection of the non-pre-trained model. c. Distribution for the UMAP projection of the pre-trained model. d. Distribution for the UMAP projection of the non-pre-trained model.