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

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 57 references that are in the map.
- First posted on arXiv 10 May 2026.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [TorchMD: A deep learning framework for molecular simulations](https://mlsanity.com/paper/2012.12106/) (2020)
- [Relational inductive bias for physical construction in humans and machines](https://mlsanity.com/paper/1806.01203/) (2018)
- [Simple GNN Regularisation for 3D Molecular Property Prediction & Beyond](https://mlsanity.com/paper/2106.07971/) (2021)
- [Physical Design using Differentiable Learned Simulators](https://mlsanity.com/paper/2202.00728/) (2022)
- [Learning to Solve PDE-constrained Inverse Problems with Graph Networks](https://mlsanity.com/paper/2206.00711/) (2022)
- [Constraint-based graph network simulator](https://mlsanity.com/paper/2112.09161/) (2021)
- [Learning Physical Constraints with Neural Projections](https://mlsanity.com/paper/2006.12745/) (2020)
- [Physics-Informed Neural Networks with Hard Nonlinear Equality and Inequality Constraints](https://mlsanity.com/paper/2507.08124/) (2025)
- [GNS: A generalizable Graph Neural Network-based simulator for particulate and fluid modeling](https://mlsanity.com/paper/2211.10228/) (2022)
- [Audacity of huge: overcoming challenges of data scarcity and data quality for machine learning in computational materials discovery](https://mlsanity.com/paper/2111.01905/) (2021)
- [MeshGraphNet-Transformer: Scalable Mesh-based Learned Simulation for Solid Mechanics](https://mlsanity.com/paper/2601.23177/) (2026)
- [FeNNol: an Efficient and Flexible Library for Building Force-field-enhanced Neural Network Potentials](https://mlsanity.com/paper/2405.01491/) (2024)
- [Improving Long-Range Interactions in Graph Neural Simulators via Hamiltonian Dynamics](https://mlsanity.com/paper/2511.08185/) (2025)
- [Enforcing governing equation constraints in neural PDE solvers via training-free projections](https://mlsanity.com/paper/2511.17258/) (2025)
- [Learning to Simulate Unseen Physical Systems with Graph Neural Networks](https://mlsanity.com/paper/2201.11976/) (2022)
- [Fast, Modular, and Differentiable Framework for Machine Learning-Enhanced Molecular Simulations](https://mlsanity.com/paper/2503.20541/) (2025)
- [Scaling Kinetic Monte-Carlo Simulations of Grain Growth with Combined Convolutional and Graph Neural Networks](https://mlsanity.com/paper/2511.17848/) (2025)

## Abstract (arXiv)

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.