Graph Mamba Operator for Learning the Dynamics of Particle-based Systems

ICLR 2026 Conference Submission16568 Authors

Keywords: Graph Neural Network, Neural Operator, State Space Models

TL;DR: Introduces a new graph neural operator that leverages state-space models (SSMs) for long range modelling.

Abstract: Modelling complex 3D dynamics---from molecular conformations to particle interactions and human motion---requires capturing dependencies spanning long temporal horizons and non-local spatial interactions. Graph neural networks (GNNs) have shown promise in spatio-temporal settings but often suffer from instability and degraded accuracy in long-range forecasting. We propose the Graph Mamba Operator (GraMO), a neural operator that integrates state-space models (SSMs) with geometric learning to capture spatio-temporal correlations jointly. To jointly model complex dynamics, GraMO integrates a stable, SSM-based temporal backbone with an SSM-parameterized graph update to capture long-range spatial dependencies. Unlike stepwise predictors that accumulate errors over time, GraMO learns entire trajectories in a single forward pass. Across diverse benchmarks ranging from molecular dynamics to human motion capture, GraMO shows notable improvements in trajectory fidelity, stability, and robustness over strong baselines with relative improvements of over 26.3\% on motion capture benchmarks and 25.2\% on MD17 final-state prediction. Ablation studies reveal that temporal SSM components consistently improve performance, while spatial SSM updates show task-dependent benefits---helping with long-range interactions in large molecules but potentially hindering performance on systems with primarily local dependencies. Altogether, these results suggest that selective integration of SSM components with graph neural networks can improve performance on particle-based systems, with applications in molecular simulations, articulated rigid body dynamics, and particle systems.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 16568