Technical Summary: Unified Simulation of Lagrangian Particle Dynamics via Transformer

Problem Statement

The construction of a unified simulation framework capable of modeling the full diversity of physical phenomena remains a long-standing challenge in simulation science. Traditional approaches rely on distinct discretizations (e.g., Eulerian grids for fluids, Lagrangian meshes for solids) and governing equations (e.g., Navier-Stokes, constitutive laws, atomistic potentials) for each phenomenon. While modern engines combine multiple solver backends, they still operate on a per-phenomenon, per-solver paradigm, making cross-phenomenon coupling difficult and constraining the applicability of a single pipeline.

Existing machine learning-based simulators offer a complementary approach by learning dynamics directly from data, bypassing explicit partial differential equations (PDEs). However, current neural simulators are typically designed and trained for specific phenomena or equation classes. Furthermore, reduced-order methods often tie learned bases to fixed geometries or shape families, requiring model redesign when physics or geometry changes. There is a need for a single architecture that can generalize across diverse Lagrangian particle systems without solver-specific redesign.

Methodology

The authors propose a particle-based prediction-correction transformer that unifies the simulation of cloth, elastic solids, Newtonian and non-Newtonian fluids, granular materials, and molecular dynamics within a single, fixed architecture. The model follows a prediction-correction scheme common in particle simulation (e.g., Position-Based Dynamics), separating the integration of known external forces from the learning of inter-particle interactions.

1. Prediction-Correction Framework

The simulation step is divided into two phases:

• Explicit Predictor: Integrates known external forces ($\mathbf{F}_t$) over a timestep $\Delta t$ to produce an intermediate state ($\tilde{\mathbf{X}}_t, \tilde{\mathbf{V}}_t$). This step accounts for external acceleration but ignores inter-particle interactions.
• Learned Corrector: A neural network $F_\theta$ predicts residual updates ($\Delta \mathbf{X}_t, \Delta \mathbf{V}_t$) based on the intermediate state, particle attributes ($\mathbf{C}$), topology ($\mathbf{T}$), and boundary conditions ($\mathbf{X}_b, \mathbf{C}_b$). The final state is updated as $\mathbf{X}_{t+\Delta t} = \tilde{\mathbf{X}}_t + \Delta \mathbf{X}_t$ and $\mathbf{V}_{t+\Delta t} = \tilde{\mathbf{V}}_t + \Delta \mathbf{V}_t$.

2. Architecture Components

The corrector consists of three stages designed to handle both local and global interactions efficiently:

• Particle Tokenizer: Encodes local interactions before global attention. It aggregates contributions from three neighborhoods using learnable kernels (parameterized as 3D grids with trilinear interpolation):

  • Spatial Neighborhood ($N^s_i$): Particle-particle interactions based on relative displacement and velocity.
  • Boundary Neighborhood ($N^b_i$): Particle-boundary interactions.
  • Topological Neighborhood ($N^t_i$): Interactions guided by rest-shape adjacency (for cloth/solids).
    The output is a per-particle token $\mathbf{h}_i$ concatenating these aggregated features with the particle's own state.

• Super-Token Encoder: Addresses the computational cost of global self-attention on $N$ particles. Inspired by multigrid methods, this module alternates between global self-attention and token merging.

  • Self-Attention: Contextualizes tokens at the current resolution using 3D Rotary Position Embeddings (RoPE).
  • Token Merging: Halves the token count at each layer by greedily matching and merging the most cosine-similar tokens. This creates a compact set of "super tokens" that aggregate local features, velocities, and material information.
  • This process repeats for $L_e$ layers, progressively reducing the sequence length and computational cost for deeper layers.

• Super-Token Decoder: Lifts the coarse context back to full particle resolution through $L_d$ layers.

  • Particle tokens query the super tokens via cross-attention, allowing particles to access global information encoded in the super tokens.
  • Particle tokens are then refined via self-attention and feed-forward networks.
  • A final MLP maps the refined tokens to the residual position and velocity corrections.

3. Training

The model is trained using an autoregressive rollout loss over a window of $W$ states. The loss function combines the mean squared error (MSE) of positions and velocities with optional physics-informed regularizers (e.g., a divergence-free constraint for incompressible fluids). The architecture remains fixed across all dynamics categories; only the training data distribution changes.

Key Contributions

1. Unified Architecture: A single transformer-based model capable of simulating six distinct Lagrangian dynamics categories (fluids, cloth, granular materials, elastic solids, non-Newtonian fluids, and molecular dynamics) without modifying the network structure.
2. Novel Token Mechanism: A particle tokenizer for local neighborhood interactions combined with a super-token encoder-decoder that enables global communication through progressive token merging, balancing expressiveness and computational efficiency.
3. Generalization and Applications: Demonstrated generalization to unseen materials, boundary configurations, initial conditions, and external forces. The framework supports downstream tasks including interactive control, inverse design via differentiable rollout, and learning from real-world manipulation data.

Experimental Results

The model was evaluated on six categories using data generated by established solvers (DFSPH, GIPC, MPM, VBD, OpenMM).

• Accuracy: The proposed model achieves lower rollout Mean Squared Errors (MSE) for positions and velocities compared to state-of-the-art baselines (GNS, DeepLagrangian, Neural Operators) across fluid, cloth, and sand tasks. For example, on the fluid task, the position error was $5.17 \times 10^{-1}$ (scaled), significantly outperforming baselines which exhibited artifacts like floating particles or interpenetration in long rollouts.
• Stability: The model maintains stable simulations over long horizons (e.g., 800 frames for fluids, 250 frames for cloth) where baselines often diverge or collapse.
• Generalization:
  • Unseen Conditions: The model successfully simulates unseen boundary geometries (e.g., tiger-shaped sand piles) and initial conditions.
  • Particle Density: It generalizes to particle sampling densities different from the training set (e.g., training on ~7k particles, testing on 8k–9k).
  • Long Rollouts: It remains stable beyond its training horizon (200 frames) when tested at 250 frames.
• Ablation Studies: Removing either the local tokenizer or the global super-token communication degrades performance, confirming the necessity of both local detail and global coupling. Adding a divergence-free regularizer further improves physical consistency for fluids.

Significance and Claims

The paper claims that this work represents a step toward a unified simulation framework where a single differentiable architecture replaces the need for multiple specialized solvers (FEM, SPH, MPM). By treating each phenomenon as a different training distribution over a shared input-output format, the model reduces the need for per-phenomenon solver engineering.

The authors highlight three specific downstream applications enabled by this architecture:

1. Interactive Control: Real-time manipulation of deformable objects under user-specified forces, leveraging the predictor's ability to integrate arbitrary external forces.
2. Inverse Design: Gradient-based optimization of physical parameters (e.g., friction coefficients) via backpropagation through the neural rollout.
3. Learning from Real-World Data: The ability to train directly on particle trajectories extracted from real-world observations (using techniques like 3D Gaussian Splatting) without relying on a classical simulator for supervision.

The authors acknowledge limitations, primarily memory constraints limiting the current implementation to approximately 50,000 particles on a single 80GB GPU due to the full-resolution self-attention in the first encoder layer. They suggest that sparse attention variants could extend this capacity. Additionally, while the architecture is unified, the current implementation trains separate parameter sets per dynamics category; a future goal is a single model with shared weights across all categories.