# Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 61 references that are in the map.
- First posted on arXiv 23 May 2023.

## Builds on (most cited first, 30 of 37)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [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)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [Interaction Networks for Learning about Objects, Relations and Physics](https://mlsanity.com/paper/1612.00222/) (2016)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities](https://mlsanity.com/paper/2006.02425/) (2020)
- [Learning Particle Dynamics for Manipulating Rigid Bodies, Deformable Objects, and Fluids](https://mlsanity.com/paper/1810.01566/) (2018)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [Frame Averaging for Invariant and Equivariant Network Design](https://mlsanity.com/paper/2110.03336/) (2021)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [Incorporating Symmetry into Deep Dynamics Models for Improved Generalization](https://mlsanity.com/paper/2002.03061/) (2020)
- [A Compositional Object-Based Approach to Learning Physical Dynamics](https://mlsanity.com/paper/1612.00341/) (2016)
- [Hamiltonian Graph Networks with ODE Integrators](https://mlsanity.com/paper/1909.12790/) (2019)
- [Approximately Equivariant Networks for Imperfectly Symmetric Dynamics](https://mlsanity.com/paper/2201.11969/) (2022)
- [Learned Coarse Models for Efficient Turbulence Simulation](https://mlsanity.com/paper/2112.15275/) (2021)
- [Equivariant Graph Mechanics Networks with Constraints](https://mlsanity.com/paper/2203.06442/) (2022)
- [Neural Stochastic PDEs: Resolution-Invariant Learning of Continuous Spatiotemporal Dynamics](https://mlsanity.com/paper/2110.10249/) (2021)
- [Learning rigid dynamics with face interaction graph networks](https://mlsanity.com/paper/2212.03574/) (2022)
- [Physics-Embedded Neural Networks: Graph Neural PDE Solvers with Mixed Boundary Conditions](https://mlsanity.com/paper/2205.11912/) (2022)

## Built on by (most cited first, 6 of 6)

- [STFlow: Data-Coupled Flow Matching for Geometric Trajectory Simulation](https://mlsanity.com/paper/2505.18647/) (2025)
- [Symmetry-Based Structured Matrices for Efficient Approximately Equivariant Networks](https://mlsanity.com/paper/2409.11772/) (2024)
- [Understanding complex crowd dynamics with generative neural simulators](https://mlsanity.com/paper/2412.01491/) (2024)
- [Equivariant Flow Matching for Symmetry-Breaking Bifurcation Problems](https://mlsanity.com/paper/2509.03340/) (2025)
- [Accelerating Simulation of Two-Phase Flows with Neural PDE Surrogates](https://mlsanity.com/paper/2405.17260/) (2024)
- [Efficient Probabilistic Modeling of Crystallization at Mesoscopic Scale](https://mlsanity.com/paper/2405.16608/) (2024)

## Abstract (arXiv)

Neural networks are emerging as a tool for scalable data-driven simulation of high-dimensional dynamical systems, especially in settings where numerical methods are infeasible or computationally expensive. Notably, it has been shown that incorporating domain symmetries in deterministic neural simulators can substantially improve their accuracy, sample efficiency, and parameter efficiency. However, to incorporate symmetries in probabilistic neural simulators that can simulate stochastic phenomena, we need a model that produces equivariant distributions over trajectories, rather than equivariant function approximations. In this paper, we propose Equivariant Probabilistic Neural Simulation (EPNS), a framework for autoregressive probabilistic modeling of equivariant distributions over system evolutions. We use EPNS to design models for a stochastic n-body system and stochastic cellular dynamics. Our results show that EPNS considerably outperforms existing neural network-based methods for probabilistic simulation. More specifically, we demonstrate that incorporating equivariance in EPNS improves simulation quality, data efficiency, rollout stability, and uncertainty quantification. We conclude that EPNS is a promising method for efficient and effective data-driven probabilistic simulation in a diverse range of domains.