Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics
This work addresses the problem of efficient and effective data-driven probabilistic simulation for researchers in fields like physics and biology, representing a novel method for a known bottleneck rather than an incremental improvement.
The paper tackled the challenge of incorporating domain symmetries into probabilistic neural simulators for stochastic spatiotemporal dynamics, proposing the Equivariant Probabilistic Neural Simulation (EPNS) framework, which demonstrated improved simulation quality, data efficiency, rollout stability, and uncertainty quantification compared to existing methods.
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.