Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics

Koen Minartz, Yoeri Poels, Simon Koop, Vlado Menkovski

TL;DR

This work tackles probabilistic simulation of stochastic spatiotemporal dynamics under domain symmetries. It introduces Equivariant Probabilistic Neural Simulation (EPNS), an autoregressive CVAE-style framework that enforces equivariance through forward, prior, and decoder components, ensuring trajectory distributions respect the relevant symmetry group. Across stochastic n-body dynamics and stochastic cellular dynamics, EPNS shows improved data efficiency, tighter uncertainty quantification, and enhanced rollout stability compared with existing probabilistic baselines. The results highlight the practical value of embedding symmetries in probabilistic neural simulators and suggest broad applicability to scientific domains requiring robust, symmetry-aware stochastic forecasting.

Abstract

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.

Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics

TL;DR

This work tackles probabilistic simulation of stochastic spatiotemporal dynamics under domain symmetries. It introduces Equivariant Probabilistic Neural Simulation (EPNS), an autoregressive CVAE-style framework that enforces equivariance through forward, prior, and decoder components, ensuring trajectory distributions respect the relevant symmetry group. Across stochastic n-body dynamics and stochastic cellular dynamics, EPNS shows improved data efficiency, tighter uncertainty quantification, and enhanced rollout stability compared with existing probabilistic baselines. The results highlight the practical value of embedding symmetries in probabilistic neural simulators and suggest broad applicability to scientific domains requiring robust, symmetry-aware stochastic forecasting.

Abstract

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.
Paper Structure (67 sections, 1 theorem, 20 equations, 15 figures, 8 tables)

This paper contains 67 sections, 1 theorem, 20 equations, 15 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Background and related work
  3. Equivariance
  4. Simulation with neural networks
  5. Method
  6. Modeling framework
  7. Problem formulation.
  8. Generative model.
  9. Incorporating symmetries.
  10. Invariant versus equivariant latents.
  11. Training procedure
  12. Applications and model designs
  13. Celestial dynamics with stochastic forcing
  14. Problem setting.
  15. Symmetries and backbone architecture.
  16. ...and 52 more sections

Key Result

Lemma 1

Assume we model $p_\theta(x^{1:T} | x^0)$ as described in Section para:generative-model. Then, $p_\theta(x^{1:T} | x^0)$ is equivariant to linear transformations $\rho(g)$ of a symmetry group $G$ in the sense of Definition eq:def-equi-distr if:

Figures (15)

  • Figure 1: Schematic illustration of EPNS applied to stochastic cellular dynamics. EPNS models distributions over trajectories that are equivariant to permutations $\mathbf{P}$ of the cell indices, so that $p_\theta(\mathbf{P}x^{t+k} | \mathbf{P}x^t)=p_\theta(x^{t+k} | x^t)$.
  • Figure 2: EPNS model overview.
  • Figure 3: Schematic overviews of model designs.
  • Figure 4: Qualitative results for celestial dynamics (left) and cellular dynamics (right).
  • Figure 5: Distributions of the potential energy in a system (top row, celestial dynamics) and the number of cell clusters of the same type (bottom row, cellular dynamics) over time. The shaded area indicates the $(0.1, 0.9)$-quantile interval of the observations; the solid line indicates the median value.
  • ...and 10 more figures

Theorems & Definitions (1)

  • Lemma 1