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

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.

• The paper introduces dual strategies for structure-only initialization using a differentiable MD engine and a GNN cascade to generate dynamic context from a single static structure.
• It employs an inference-time physics-based optimization framework that refines predictions, ensuring physical fidelity and robust out-of-distribution generalization.
• The study demonstrates high accuracy (R² up to 0.98) in simulating diverse disordered elastic networks, bolstering inverse design applications.

Enabling Structure-Only Initialization and OOD Generalization in GNN-based Molecular Dynamics Simulators

Introduction and Motivation

Machine learning-based molecular dynamics (MD) simulators, particularly those built on graph neural networks (GNNs), have achieved compelling accuracy and computational efficiency across a range of systems. However, these models historically require temporal context—multiple past system states (e.g., velocities, accelerations)—to deliver stable, physically realistic rollouts. This constraint undermines their compatibility with inverse design workflows, where simulations must be initialized from a single, static structure, and also limits their ability to extrapolate to novel out-of-distribution (OOD) regimes inherent in the design of mechanical metamaterials and disordered solids.

The paper "Enabling Structure-Only Initialization and Out-of-Distribution Generalization in GNN-based Molecular Dynamics Simulators" (2605.09495) addresses these deficiencies through a set of algorithmic innovations and architectural strategies. It introduces two complementary approaches for structure-only initialization (SOI), a physics-aware inference-time optimization framework, and a differentiable GNN-based barostat that collectively deliver (i) stable "cold-start" dynamics from single-structure inputs and (ii) robust OOD generalization on disordered elastic networks (DENs) with highly varying Poisson ratios and network geometries.

Disordered Elastic Networks: Testbed and Dataset

As a testbed, the work focuses on 2D DENs—particle (bead) networks with random topology and tunable bond stiffness, where the macroscopic Poisson ratio can span both positive (non-auxetic) and highly negative (auxetic) values. DENs provide a rigorous benchmark due to the diversity of their dynamical behaviors and geometric complexity. Examples of DEN topologies, including both non-auxetic and auxetic variants generated by global node optimization and bond stiffness tailoring, illustrate the diversity of network structures modeled in the study.

Figure 1: Disordered elastic network structures ranging from high Poisson ratio (a) to auxetic (b,c) configurations with complex topology.

Structure-Only Initialization: Bootstrapping Dynamics from Single Configurations

Problem Statement

Conventional GNN-based simulators demand temporal context (sequences of prior states) that are unavailable when initializing from new, static candidate structures in design loops. Naive position-only models, lacking velocity/acceleration input, immediately lose accuracy and destabilize over short rollouts, as evidenced by large mean square errors and poor prediction of emergent macroscopic properties (e.g., Poisson's ratio).

Figure 3: Position-only GNN simulator fails to track Poisson's ratio and incurs high position errors after 10-step rollouts, confirming instability of naïve single-step initialization.

Differentiable MD Bootstrapping

The first SOI approach introduces a custom, minimal, fully-differentiable MD engine that creates a short bootstrap trajectory from the initial configuration. This transient trajectory generates the necessary dynamic state variables (velocity, acceleration), which are then used to seed the GNN-based simulation. The pipeline ensures the GNN receives the required history for autoregressive rollout without recourse to empirical or artificially padded initializations. This approach preserves fully differentiable workflows, which is essential for inverse design pipelines based on backpropagation.

Simulator Cascade

The second approach eschews explicit MD and implements a sequential GNN cascade: a series of specialized GNN models, each trained to predict the state at a particular early timestep, thus iteratively building up the historical context necessary for the main simulator. Each component in the cascade has increasing temporal history; earlier stages reason solely from positions, later ones build upon their outputs, enabling progressive accumulation of dynamic features.

Figure 2: Overview of SOI strategies—differentiable MD bootstrapping (top), and the autoregressive simulator cascade (bottom).

Comparative Results

Both methods, evaluated on DEN compression, enable stable rollouts for up to 50 steps from static input. With in-distribution networks (Poisson ratio $\nu \geq 0.1$), the MD-bootstrapped GNN simulator attains $R^2 \approx 0.97$ for Poisson's ratio; even in OOD regimes (e.g., highly auxetic, $\nu \ll 0$), it achieves $R^2 \approx 0.71$. The simulator cascade performs even better OOD ($R^2 \approx 0.89$ down to $\nu = -0.3$, and $R^2 \approx 0.99$ in-domain).

Figure 4: Parity plots and error curves for SOI methods; MD bootstrapping (a,b) and simulator cascade (c,d) compared on long rollouts over a broad $\nu$ interval.

Inference-Time Physics-Based Optimization for Stabilization and Generalization

Model performance in OOD scenarios can degrade due to error accumulation and a lack of physically realistic trajectories, especially in regions unseen during training (e.g., networks designed for extreme auxeticity). The proposed Inference-Time Physics-based Optimization (ITPO) framework addresses this by applying an auxiliary, physics-informed optimization to the model outputs at each rollout step. Specifically, after the GNN forwards, predicted accelerations are treated as learnable parameters and refined by minimizing a loss that balances fidelity to the GNN’s prediction (anchor term) and satisfaction of explicit physical constraints (pressure/barostat, potential energy, mean-squared net force).

This procedure, applied during inference only, does not alter learned model weights and is agnostic to the overall system; it simply constrains each predicted state to remain physically plausible. ITPO regularizes the rollout, suppresses error accumulation, and enforces conservation and macroscopic consistency—critical for extrapolation in high-stiffness or negative-$\nu$ regimes.

Impact of ITPO

ITPO leads to major improvements: both MD-bootstrapped and cascade-based simulators reach $R^2 \approx 0.98$ across the whole range of $R^2 \approx 0.97$0, matching in-distribution accuracy and yielding plateau-like, low position error across long rollouts up to 200 steps (40,000 MD-equivalent steps). The error profiles become invariant to Poisson's ratio, and the stability with respect to rollout depth is markedly enhanced.

Figure 5: ITPO-corrected SOI initializations sustain high accuracy ($R^2 \approx 0.97$1) and low position error across 50-step rollouts, including in the highly auxetic, OOD regime.

Figure 6: ITPO also improves accuracy and stability as rollout length increases, maintaining physical fidelity on extended prediction horizons.

GNN-Based Barostat and Macroscopic Consistency

A differentiable barostat is embedded as part of the rollout pipeline. This algorithm tracks and adjusts the simulation box’s periodic boundaries and internal pressure, ensuring that system-level properties such as Poisson's ratio are well-conditioned even over long, coarse-grained rollouts. Accurate tracking of box dimensions and stress/strain is achieved, as demonstrated in both non-auxetic and auxetic materials. The barostat is critical for coupling rollouts to the dynamic macroscopic behavior needed for inverse design applications.

Figure 7: Barostat performance—dynamically calculated box dimensions and pressures (orange) closely match ground truth (blue) over compression trajectories.

Implications, Limitations, and Future Directions

Practical Implications

• Inverse Design: Reliable SOI in GNN simulators enables fully differentiable materials optimization, making these models practical surrogates for high-throughput screening, topology optimization, and autonomous discovery of unconventional material properties—the focus of recent AI for science efforts [Zheng 2023, Bordiga 2024, Song 2024].
• Generalization: The demonstrated OOD stability evidences that physics-informed or constraint-regularized ML can extrapolate substantially beyond the training domain without explicit re-training or data augmentation, addressing concerns about brittleness and distribution shift in scientific ML [Li 2023, Meredig 2018, Omee 2024].
• Computational Efficiency: The coarse-graining capacity of GNN simulators—enabled via MD bootstrapping or GNN cascades—achieves order-of-magnitude acceleration versus classical MD or standard differentiable simulation, facilitating large-scale optimization and design.

Theoretical Implications

• The performance of ITPO validates the use of train-time/data-driven models augmented with inference-time physical constraints—a modular, scalable strategy now gaining traction for robust scientific extrapolation [Rubanova 2021, Karniadakis 2021, Cranmer 2020].
• The efficacy and reliability of the multi-step training (MST) regime confirm its role in mitigating error accumulation in autoregressive simulators, complementing earlier proposals in the numerical and ML simulation communities [Sanchez-Gonzalez 2020, Hoang 2025].

Limitations and Future Work

• Computational Overhead: ITPO introduces additional per-step optimization cost. Efficient variants (subsampling, partial-step correction) merit exploration to reduce this burden.
• Applicability to 3D and Multiphysics Systems: Extension to 3D networks, complex loading, and multiphysics couplings (thermal, electrical) could broaden realistic deployment.
• Unified Simulation Frameworks: Integration into general-purpose differentiable simulation libraries (cf. JaxMD [Schoenholz 2021], TorchMD) would streamline adoption by the computational materials community.
• Symbolic Extraction and Interpretability: The distilled knowledge learned by GNNs could be subjected to symbolic regression to recover interpretable, closed-form models for material design, as recently demonstrated [Cranmer 2020].

Conclusion

The proposed framework decisively resolves the structure-only initialization bottleneck in GNN-based simulators and achieves robust OOD generalization through inference-time physical regularization. These advances establish GNN-based simulators, combined with differentiable MD, tailored GNN cascades, and on-demand physics optimization, as viable surrogates for autonomous, high-throughput materials design in domains characterized by structural disorder and extreme mechanical behavior. The methodology’s generality and robust accuracy open new avenues for inverse design, optimization, and scientific exploration in computational materials science and beyond.