Enabling Structure-Only Initialization and Out-of-Distribution Generalization in GNN-based Molecular Dynamics Simulators

2026-05-10

AI Summary

This work addresses the limitations of existing graph neural network (GNN)-based molecular dynamics simulators, which rely on temporal context for initialization and thus struggle with inverse design tasks that provide only static structures, while also exhibiting poor generalization under out-of-distribution (OOD) conditions. To overcome these challenges, the authors propose a novel framework integrating structural initialization, physics-informed optimization during inference, and a differentiable GNN-based barostat. This approach significantly enhances simulation stability and physical consistency without requiring temporal context. The method accurately tracks system volume and pressure, demonstrates exceptional rollout stability in uniaxial compression of disordered elastic networks, and successfully generalizes to complex dynamical scenarios not encountered during training.

Abstract

Machine learning-based simulators offer the potential to model the dynamics of complex systems more efficiently than classical approaches, while retaining differentiability, a key property for materials design. Graph neural network (GNN)-based simulators have shown strong performance across a range of physical domains, including molecular dynamics. However, their reliance on temporal context for accurate prediction limits their use in inverse design settings, where simulations must be initialized from a single static configuration. Moreover, inverse design requires robust out-of-distribution (OOD) generalization, as candidate structures typically lie outside the training domain. Here, we address both challenges by introducing two complementary strategies that enable stable and accurate structure-only initialization of GNN-based simulations. To directly target OOD generalization, we propose an inference-time physics-based optimization framework that constrains model predictions to remain physically consistent during rollout. In addition, we introduce a differentiable, GNN-based barostat that enables accurate tracking of system dimensions and pressure, critical for capturing macroscopic responses and supporting OOD generalization. We evaluate these approaches in the context of uniaxial compression of disordered elastic networks spanning a broad range of geometries, Poisson ratios, and microscopic behaviors. We find that, together, these methods substantially improve rollout stability and enable reliable OOD generalization, including regimes with distinct, more complex dynamics than those in the training data. These results show that, when properly initialized and constrained, GNN-based simulators can serve as efficient and generalizable tools for materials discovery and structural optimization, advancing their use in materials, molecular, and dynamical system design.

Problem

• structure-only initialization
• out-of-distribution generalization
• GNN-based simulators
• inverse design
• molecular dynamics

Innovation

• structure-only initialization
• out-of-distribution generalization
• physics-based optimization
• differentiable barostat
• GNN-based simulators