🤖 machine learning

Graph Mamba Operator: A Latent Simulator for Interacting Particle Systems

The paper introduces the Graph Mamba Operator (GraMO), a latent-space simulator that uniquely couples graph-based spatial interactions with state-space temporal dynamics in a single linear recurrence to overcome the error accumulation and limited context of existing methods, achieving superior long-horizon prediction performance across diverse interacting particle systems.

Original authors: Karn Tiwari, Niladri Dutta, N M Anoop Krishnan, Prathosh A P

Published 2026-06-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Karn Tiwari, Niladri Dutta, N M Anoop Krishnan, Prathosh A P

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict how a flock of birds, a group of dancers, or a bunch of bouncing balls will move in the future. This is a tough job because every object is influenced by two things at once: who is near them right now (spatial interaction) and how they have been moving for a long time (temporal memory).

The paper introduces a new AI tool called GraMO (Graph Mamba Operator) that acts like a "super-simulator" for these moving groups. Here is how it works, explained simply:

The Problem: The "Domino Effect" of Mistakes

Previous AI models tried to predict the future by taking one step at a time, like a person guessing the next move in a game of chess.

The Flaw: If the AI makes a tiny mistake predicting the next second, that mistake gets carried over to the next second, and then the one after that. By the time it tries to predict 20 seconds into the future, the prediction is a mess of accumulated errors.
The Limitation: Older models also treated "space" (who is touching whom) and "time" (how things move over seconds) as two separate problems, solving them one after the other. This is like trying to learn to drive a car by first studying the steering wheel for an hour, then the pedals for an hour, and never practicing them together.

The Solution: GraMO (The "Living Blueprint")

GraMO solves this by creating a single, unified "living blueprint" (called a latent state) that holds the entire history of the system in its memory.

Think of GraMO not as a person guessing step-by-step, but as a highly skilled conductor leading an orchestra:

The Graph (The Sheet Music): The "Graph" part of the name means the AI knows the relationships between the particles (like knowing which violinist is next to the cellist). It understands that if one moves, the neighbor might react.
The Mamba (The Memory): The "Mamba" part comes from a type of AI known for having a long, perfect memory. It remembers the entire history of the movement without forgetting the beginning.
The Magic Trick (The Coupling): Instead of checking the sheet music then checking the memory, GraMO does both at the exact same time. It updates the "blueprint" of the system in one smooth motion.

How It Works: The "Smart Sponge" Analogy

Imagine the system's state is a smart sponge.

Old Models: They would squeeze the sponge, guess the shape, squeeze it again, guess again, and so on. Every time they squeezed it, they lost a little bit of the original shape (error accumulation).
GraMO: It treats the sponge as a single object that changes shape based on two things happening simultaneously:
1. Pressure from neighbors: If a neighbor pushes the sponge, it deforms (Graph interaction).
2. The flow of time: The sponge has an internal rhythm that remembers how it was stretched 10 seconds ago (Temporal memory).

Crucially, the "rules" for how the sponge deforms change depending on what is happening right now. If the particles are moving slowly, the sponge stretches gently. If they crash into each other (a "regime change"), the sponge instantly tightens its rules to handle the chaos. This is called input-dependent dynamics.

What They Tested It On

The authors tested this "smart sponge" on three very different types of moving systems:

N-Body Systems: Simulating charged particles, springs, and gravity (like planets or magnets).
Robotics: Controlling a flexible rope and a soft, squishy robot that moves in water.
Motion Capture: Predicting human walking and running movements.

The Results

In every test, GraMO was the most accurate.

Long-Term Accuracy: It was especially good at predicting far into the future without the "messy errors" that plagued other models.
Generalization: It could predict the movement of a rope with 9 knots even if it was only trained on ropes with 5, 7, or 8 knots. It learned the rules of the rope, not just the specific rope it saw.
Efficiency: It was faster to train than many complex models while being more accurate.

The Bottom Line

GraMO is a new way to teach computers to simulate moving things. Instead of guessing the future one step at a time and getting tired of mistakes, it builds a single, evolving memory that understands both who is touching whom and how time is passing all at once. This allows it to predict complex movements—like dancing, bouncing balls, or swimming robots—with much higher precision and stability.

Technical Summary: Graph Mamba Operator (GraMO)

Problem Statement

Modeling interacting dynamical systems—such as N-body systems, articulated bodies, and deformable structures—requires capturing both spatial interactions and long-range temporal dependencies. Existing approaches face two primary limitations:

Autoregressive Error Accumulation: Graph Neural Networks (GNNs) typically rely on autoregressive rollouts where spatial and temporal dynamics are treated separately. This leads to compounding prediction errors over long horizons.
Limited Scope: Current methods often focus on local interactions and short temporal contexts, failing to capture multi-hop dependencies and global structural dynamics effectively.

While State-Space Models (SSMs) offer a mechanism for long-range temporal memory with linear cost, applying them to graphs often requires imposing artificial sequential orderings (e.g., node sorting or random walks), which disrupts the intrinsic topology governing particle interactions. Conversely, standard SSMs treat inputs as unstructured sequences, ignoring the graph structure essential for physical systems.

Methodology: The Graph Mamba Operator (GraMO)

The authors propose GraMO, a latent-space simulator that integrates state-space models with graph-based interaction learning within a single recurrence. Unlike prior work that sequences nodes or applies spatial and temporal updates in separate stages, GraMO couples graph propagation and temporal state evolution simultaneously.

Core Formulation

GraMO operates on a latent state $H(t) \in \mathbb{R}^{V \times D \times N}$ , where $V$ is the number of nodes, $D$ is the feature dimension, and $N$ is the memory dimension. The continuous-time dynamics are defined as:
$\frac{dH(t)}{dt} = A \odot (\tilde{L} H(t) W_s) + B(t) \odot X(t)$
Where:

$\tilde{L}$ is the normalized adjacency matrix (with self-loops) acting on the node dimension to propagate spatial information.
$W_s$ is a learnable spatial mixing matrix acting on the state dimension.
$A$ governs stable temporal state decay (initialized via HiPPO).
$B(t)$ is an input-dependent projection that injects current observations $X(t)$ into the latent state.
The output is obtained via a contraction $Y(t) = H(t) \cdot C(t)$ .

Key Architectural Features

Joint Recurrence: The update is linear in the latent state $H(t)$ at each step, but the operator coefficients ( $A, B, C, \Delta t$ ) are input-dependent. This allows the system to adapt to different regimes (e.g., contact, slow drift, fast transitions) without resetting the latent state, thereby avoiding the error accumulation typical of autoregressive predictors.
Discretization: The model uses a Zero-Order Hold (ZOH) discretization with a learnable timestep $\Delta t$ . The update rule becomes:
$H_{t+1} = \bar{A}_t \odot (\tilde{L} H_t W_s) + \bar{B}_t \odot X_t$
This structure unifies state-space models with graph Laplacian operators, enabling principled message passing within a continuous-time latent dynamical system.
Selectivity Mechanism: Parameters $B_t, C_t,$ and $\Delta t$ are generated dynamically from the input $X_t$ via a graph neural network. This enables content-aware gating (selecting which input information enters the state) and non-stationary adaptation (modulating temporal dynamics based on the current graph state).
Bidirectional Processing: The architecture employs a bidirectional SSM formulation, applying the graph-coupled update in both forward and backward temporal directions to capture richer trajectory context.

Theoretical Properties

The paper establishes that GraMO admits a Graph Autoregressive Moving-Average (ARMA) representation with input-dependent coefficients. Theoretical analysis confirms:

Stability: The recurrence ensures stable long-range information flow due to the bounded spectrum of the normalized Laplacian and the HiPPO-initialized temporal operator.
Linearity in Latent Space: The update remains linear in the lifted latent state, connecting the method to time-varying Koopman dynamics where the operator adapts to the current graph state.

Experimental Results

GraMO was evaluated on three distinct domains: N-body simulations, robotics, and motion capture.

Benchmarks

N-Body Systems: Charged particles, spring dynamics, and gravitational systems.
Robotics: Rope control, soft robot manipulation, and swimming soft robots.
Motion Capture: CMU dataset (Walk and Run sequences).

Performance

GraMO achieved the lowest error across all benchmarks and metrics (Average Mean Squared Error - AMSE, and Final Mean Squared Error - FMSE), outperforming strong baselines including ST-GNN, SE(3)-Transformers, Koopman-based methods, and GraphMamba.

N-Body: On the Gravity System, GraMO reduced FMSE from 6.020 (GeoTDM) to 3.720. On Spring Dynamics, it reduced FMSE from 0.063 (ESTAG) to 0.024.
Robotics: On the Rope dataset, GraMO improved FMSE from 0.5380 (GeoTDM) to 0.4120.
Motion Capture: GraMO reduced the average FMSE on Run sequences from 0.508 (NS-EGNN) to 0.161, representing a 67.3% relative error reduction.

Ablation Studies

Ablation experiments confirmed the necessity of each component:

Removing temporal modeling caused the largest performance degradation.
Removing the bidirectional SSM hurt performance, particularly in motion forecasting.
Removing input-dependent selectivity (making parameters fixed) significantly increased error, highlighting the need for adaptive dynamics.
Replacing HiPPO initialization degraded long-range memory capabilities.

Generalization and Efficiency

Zero-Shot Generalization: GraMO demonstrated the ability to generalize to unseen system sizes (e.g., training on $N \in \{5, 7, 8\}$ particles and testing on $N=9$ ), outperforming Compositional Koopman operators.
Efficiency: While simpler equivariant models train faster, they exhibit lower predictive performance. GraMO offers a favorable trade-off, achieving lower FMSE than complex baselines like EqMotion and SE(3)-Transformer with competitive training times.

Significance and Claims

The authors claim that GraMO provides a principled framework for learning complex spatiotemporal dynamics by unifying graph-based interaction learning with state-space evolution. Key claims include:

Coupling Mechanism: It is the first latent simulator to couple graph propagation and temporal state updates within a single recurrence, allowing spatial interactions and temporal memory to evolve jointly.
Error Mitigation: By evolving a persistent latent state rather than recursively predicting in observation space, GraMO avoids the error accumulation characteristic of autoregressive rollouts, leading to superior long-horizon prediction.
Topology Preservation: Unlike approaches that serialize nodes (e.g., GraphMamba), GraMO preserves the intrinsic interaction topology via Laplacian-based propagation, which is critical for modeling interacting particle dynamics.
Adaptability: The input-dependent coefficients allow the model to handle regime-switching systems (e.g., contact vs. free flight) effectively.

The paper concludes that while GraMO scales with graph size due to repeated graph-state interactions, it represents a significant step forward in modeling interacting particle systems, with future work potentially exploring graph sparsification to improve scalability.