SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Koen Minartz<sup>1</sup> Yoeri Poels<sup>1,2</sup> Simon Koop<sup>1</sup> Vlado Menkovski<sup>1</sup>
$^{1}$ Data and Artificial Intelligence Cluster, Eindhoven University of Technology  $^{2}$ Swiss Plasma Center, École Polytechnique Fédérale de Lausanne {k.minartz, y.r.j.poels, s.m.koop, v.menkovski}@tue.nl
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.
The advent of fast and powerful computers has led to numerical simulations becoming a key tool in the natural sciences. For example, in computational biology, models based on probabilistic cellular automata are used to effectively reproduce stochastic cell migration dynamics [2, 13, 26, 51], and in physics, numerical methods for solving differential equations are used in diverse domains such as weather forecasting [4, 32], thermonuclear fusion [16, 55], and computational astrophysics [24, 52], amongst others. Recently, data-driven methods driven by neural networks have emerged as a complementary paradigm to simulation. These neural simulators are particularly useful when numerical methods are infeasible or impractical, for example due to computationally expensive simulation protocols or lack of an accurate mathematical model [9, 14, 30, 37, 41, 42].
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{Px}^{t + k}|\mathbf{Px}^t) = p_{\theta}(x^{t + k}|x^t)$ .
Notably, most neural simulators are deterministic: given some initial condition, they predict a single trajectory the system will follow. However, for many applications this formulation is insufficient. For example, random effects or unobserved exogenous variables can steer the system towards strongly diverging possible evolutions. In these scenarios, deterministic models are incentivized to predict an 'average case' trajectory, which may not be in the set of likely trajectories. Clearly, for these cases probabilistic neural simulators that can expressively model distributions over trajectories are required. Simultaneously, it has been shown that incorporating domain symmetries in model architectures can be highly beneficial for model accuracy, data efficiency, and parameter efficiency [10], all of which are imperative for a high-quality neural simulator. As such, both the abilities to perform probabilistic simulation and to incorporate symmetries are crucial to develop reliable data-driven simulators.
Accordingly, prior works have proposed methods for probabilistic simulation of dynamical systems, as well as for symmetry-aware simulation. For probabilistic spatiotemporal simulation, methods like neural stochastic differential equations (NSDEs) [20, 27], neural stochastic partial differential equations (NSPDEs) [40], Bayesian neural networks [34, 58], autoregressive probabilistic models [50], and Gaussian processes [59] have been successfully used. Still, these methods do not incorporate general domain symmetries into their model architectures, hindering their performance. In contrast, almost all works that focus on incorporating domain symmetries into neural simulation architectures are limited to the deterministic case [7, 18, 19, 31, 53, 54]. As such, these methods cannot simulate stochastic dynamics, which requires expressive modeling of distributions over system evolutions. A notable exception is [50], which proposes a method for probabilistic equivariant simulation, but only for two-dimensional data with rotation symmetry.
Although the above-mentioned works address the modeling of uncertainty or how to build equivariant temporal simulation models, they do not provide a general method for equivariant stochastic simulation of dynamical systems. In this work, we propose a framework for probabilistic neural simulation of spatiotemporal dynamics under equivariance constraints. Our main contributions are as follows:
- We propose Equivariant Probabilistic Neural Simulation (EPNS), an autoregressive probabilistic modeling framework for simulation of stochastic spatiotemporal dynamics. We mathematically and empirically demonstrate that EPNS produces single-step distributions that are equivariant to the relevant transformations, and that it employs these distributions to autoregressively construct equivariant distributions over the entire trajectory.
- We evaluate EPNS on two diverse problems: an n-body system with stochastic forcing, and stochastic cellular dynamics on a grid domain. For the first problem, we incorporate an existing equivariant architecture in EPNS. For the second, we propose a novel graph neural network that is equivariant to permutations of cell indices, as illustrated in Figure 1.
- We show that incorporating equivariance constraints in EPNS improves data efficiency, uncertainty quantification and rollout stability, and that EPNS outperforms existing methods for probabilistic simulation in the above-mentioned problem settings.
For a given group of transformations  $G$ , a function  $f: X \to X$  is  $G$ -equivariant if
$$
f (\rho (g) x) = \rho (g) f (x) \tag {1}
$$
holds for all  $g \in G$ , where  $\rho(g)$  denotes the group action of  $g$  on  $x \in X$  [10, 12]. In the context of probabilistic models, a conditional probability distribution  $p(x|y,z)$  is defined to be  $G$ -equivariant with respect to  $y$  if the following equality holds for all  $g \in G$ :
$$
p (x \mid y, z) = p (\rho (g) x \mid \rho (g) y, z). \tag {2}
$$
The learning of such equivariant probability distributions has been employed for modeling symmetrical density functions, for example for generating molecules in 3D (Euclidean symmetry) [25, 43, 57] and generating sets (permutation symmetry) [6, 28]. However, different from our work, these works typically consider time as an internal model variable, and use the equivariant probability distribution as a building block for learning invariant probability density functions over static data points by starting from a base distribution that is invariant to the relevant transformations. For example, each
3D orientation of a molecule is equally probable. Another interesting work is [17], which considers equivariant learning of stochastic fields. Given a set of observed field values, their methods produce equivariant distributions over interpolations between those values using either equivariant Gaussian Processes or Steerable Conditional Neural Processes. In contrast, our goal is to produce equivariant distributions over temporal evolutions, given only the initial state of the system.
2.2 Simulation with neural networks
Two commonly used backbone architectures for neural simulation are graph neural networks for operating on arbitrary geometries [1, 3, 9, 11, 18, 29, 36, 42, 56] and convolution-based models for regular grids [14, 30, 48]. Recently, the embedding of symmetries beyond permutation and translation equivariance has become a topic of active research, for example in the domains of fluid dynamics and turbulence modeling [31, 53, 54], climate science [7], and various other systems with Euclidean symmetries [18, 19, 44]. Although these works stress the relevance of equivariant simulation, they do not apply to the stochastic setting due to their deterministic nature.
Simultaneously, a variety of methods have been proposed for probabilistic simulation of dynamical systems. Latent neural SDEs [27] and ODE $^{2}$ VAE [58] simulate dynamics by probabilistically evolving a (stochastic) differential equation in latent space. Alternatively, methods like Bayesian neural networks (BNNs) and Gaussian processes (GPs) have also been applied to probabilistic spatiotemporal simulation [34, 59]. Further, NSPDEs are able to effectively approximate solutions of stochastic partial differential equations [40]. Due to their probabilistic nature, all of these methods are applicable to stochastic simulation of dynamical systems, but they do not incorporate general domain symmetries into their models. Finally, closely related to our work, [50] addresses the problem of learning equivariant distributions over trajectories in the context of multi-agent dynamics. However, their method is limited to two spatial dimensions, assumes a Gaussian one-step distribution, and only takes the symmetry group of planar rotations into account. In contrast, we show how general equivariance constraints can be embedded in EPNS, and that this leads to enhanced performance in two distinct domains.
3.1 Modeling framework
Problem formulation. At each time  $t$ , the system's state is denoted as  $x^t$ , and a sequence of states is denoted as  $x^{0:T}$ . We presuppose that we have samples from a ground-truth distribution  $p^*(x^{0:T})$  for training, and for inference we assume that  $x^0$  is given as a starting point for the simulation. In this work, we assume Markovian dynamics, although our approach can be extended to non-Markovian dynamics as well. Consequently, the modeling task boils down to maximizing Equation 3:
$$
\mathbb {E} _ {x ^ {0: T} \sim p ^ {*}} \mathbb {E} _ {t \sim U \{0, \dots , T - 1 \}} \log p _ {\theta} \left(x ^ {t + 1} \mid x ^ {t}\right). \tag {3}
$$
As such, we can approximate  $p^*(x^{1:T}|x^0)$  by learning  $p_\theta(x^{t+1}|x^t)$  and applying it autoregressively. Furthermore, when we know that the distribution  $p^*(x^{1:T}|x^0)$  is equivariant to  $\rho(g)$ , we want to guarantee that  $p_\theta(x^{1:T}|x^0)$  is equivariant to these transformations as well. This means that
$$
p _ {\theta} \left(x ^ {1: T} \mid x ^ {0}\right) = p _ {\theta} \left(\rho (g) x ^ {1: T} \mid \rho (g) x ^ {0}\right) \tag {4}
$$
should hold, where  $\rho(g)x^{1:T}$  simply means applying  $\rho(g)$  to all states in the sequence  $x^{1:T}$ .
Generative model. The model  $p_{\theta}(x^{t + 1}|x^t)$  should adhere to three requirements. First, since our goal is to model equivariant distributions over trajectories in diverse domains, the type of model used for  $p_{\theta}$  should be amenable to incorporating general equivariance constraints. Second, we argue that generative models with iterative sampling procedures, although capable of producing high-quality samples [15, 46], are impractical in this setting as we need to repeat the sampling procedure over potentially long trajectories. Third, the model should be able to simulate varying trajectories for the same initial state, especially in the case of sensitive chaotic systems, so good mode coverage is vital.
As such, we model  $p_{\theta}(x^{t + 1}|x^t)$  as  $\int_zp_\theta (x^{t + 1}|z,x^t)p_\theta (z|x^t)dz$  by using a latent variable  $z$ , following the conditional Variational Autoencoder (CVAE) framework [47]. Specifically, to generate  $x^{t + 1}$ ,  $x^t$  is first processed by a forward model  $f_{\theta}$ , producing an embedding  $h^t = f_{\theta}(x^t)$ . If there is little
(a) Sampling.
(b) Training.
Figure 2: EPNS model overview.
variance in  $p^*(x^{t+1}|x^t)$ ,  $h^t$  contains almost all information for generating  $x^{t+1}$ . On the other hand, if the system is at a point where random effects can result in strongly bifurcating trajectories,  $h^t$  only contains implicit information about possible continuations of the trajectories. To enable the model to simulate the latter kind of behavior, the embedding  $h^t$  is subsequently processed by the conditional prior distribution. Then,  $z \sim p_\theta(z|h^t)$  is processed in tandem with  $h^t$  by the decoder, producing a distribution  $p_\theta(x^{t+1}|z,h^t)$  over the next state. The sampling process is summarized in Figure 2a.
Incorporating symmetries. Using the approach described above, we can design models that can simulate trajectories by iteratively sampling from  $p_{\theta}(x^{t + 1}|x^t)$ . Now, the question is how we can ensure that  $p_{\theta}(x^{1:T}|x^0)$  respects the relevant symmetries. In Appendix A, we prove Lemma 1:
Lemma 1. Assume we model  $p_{\theta}(x^{1:T}|x^0)$  as described in Section 3.1. 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 2 if:
(a) The forward model is  $G$ -equivariant:  $f_{\theta}(\rho(g)x) = \rho(g)f_{\theta}(x)$ ;
(b) The conditional prior is  $G$ -invariant or  $G$ -equivariant:  $p_{\theta}(z|\rho(g)h^t) = p_{\theta}(z|h^t)$ , or  $p_{\theta}(\rho(g)z|\rho(g)h^t) = p_{\theta}(z|h^t)$ ;
(c) The decoder is  $G$ -equivariant with respect to  $h^t$  (for an invariant conditional prior), or  $G$ -equivariant with respect to both  $h^t$  and  $z$  (for an equivariant conditional prior):  $p_{\theta}(x^{t + 1}|z,h^t) = p_{\theta}(\rho (g)x^{t + 1}|z,\rho (g)h^t)$ , or  $p_{\theta}(x^{t + 1}|z,h^t) = p_{\theta}(\rho (g)x^{t + 1}|\rho (g)z,\rho (g)h^t)$ .
The consequence of Lemma 1 is that, given an equivariant neural network layer, we can stack these layers to carefully parameterize the relevant model distributions, guaranteeing equivariant distributions over the entire trajectory. To understand this, let us consider the case where the conditional prior is  $G$ -invariant, assuming two inputs  $x^t$  and  $\rho(g)x^t$ . First, applying  $f_{\theta}$  to  $x^t$  yields  $h^t$ , and since  $f_{\theta}$  is  $G$ -equivariant, applying it to  $\rho(g)x^t$  yields  $\rho(g)h^t$ . Second, due to the invariance of  $p_{\theta}(z|h^t)$ , both  $h^t$  and  $\rho(g)h^t$  lead to exactly the same distribution over  $z$ . Third, since  $z \sim p_{\theta}(z|h^t)$  is used to parameterize the decoder  $p_{\theta}(x^{t+1}|z,h^t)$ , which is equivariant with respect to  $h^t$ , the model's one-step distribution  $p_{\theta}(x^{t+1}|x^t)$  is equivariant as well. Finally, autoregressively sampling from the model inductively leads to an equivariant distribution over the entire trajectory.
Invariant versus equivariant latents. When  $p_{\theta}(z|h^t)$  is equivariant, the latent variables are typically more local in nature, whereas they are generally more global in the case of an invariant conditional prior. For example, for a graph domain, we may design  $p_{\theta}(z|h^t)$  to first perform a permutation-invariant node aggregation to obtain a global latent variable, or alternatively to equivariantly map each node embedding to a node-wise distribution over  $z$  to obtain local latent variables. Still, both design choices result in an equivariant model distribution  $p_{\theta}(x^{1:T}|x^0)$ .
3.2 Training procedure
Following the VAE framework [23, 39], an approximate posterior distribution  $q_{\phi}(z|x^t,x^{t + 1})$  is learned to approximate the true posterior  $p_{\theta}(z|x^{t},x^{t + 1})$ . During training,  $q_{\phi}$  is used to optimize the
Evidence Lower Bound (ELBO) on the one-step log-likelihood:
$$
\log p _ {\theta} \left(x ^ {t + 1} \mid x ^ {t}\right) \geq \mathbb {E} _ {q _ {\phi} \left(z \mid x ^ {t}, x ^ {t + 1}\right)} \left[ p _ {\theta} \left(x ^ {t + 1} \mid z, x ^ {t}\right) \right] - K L \left[ q _ {\phi} \left(z \mid x ^ {t}, x ^ {t + 1}\right) \right\rvert   | p _ {\theta} \left(z \mid x ^ {t}\right) ]. \tag {5}
$$
In our case,  $q_{\phi}$  takes both the current and next timestep as input to infer a posterior distribution over the latent variable  $z$ . Naturally, we parameterize  $q_{\phi}$  to be invariant or equivariant with respect to both  $x^{t}$  and  $x^{t + 1}$  for an invariant or equivariant conditional prior  $p_{\theta}(z|x^t)$  respectively.
A well-known challenge for autoregressive neural simulators concerns the accumulation of error over long rollouts [9, 48]: imperfections in the model lead to its output deviating further from the ground-truth for increasing rollout length  $k$ , producing a feedback loop that results in exponentially increasing errors. In the stochastic setting, we observe a similar phenomenon:  $p_{\theta}(x^{t + 1}|x^t)$  does not fully match  $p^* (x^{t + 1}|x^t)$ , resulting in a slightly biased one-step distribution. Empirically, we see that these biases accumulate, leading to poor coverage of  $p^* (x^{t + k}|x^t)$  for  $k \gg 1$ , and miscalibrated uncertainty estimates. Fortunately, we found that a heuristic method can help to mitigate this issue. Similar to pushforward training [9], in multi-step training we unroll the model for more than one step during training. However, iteratively sampling from  $p_{\theta}(x^{t + 1}|x^t)$  can lead to the simulated trajectory diverging from the ground-truth sample. Although this behavior is desired for a stochastic simulation model, it ruins the training signal. Instead, we iteratively sample a reconstruction  $\hat{x}^{t + 1}$  from the model, where  $\hat{x}^{t + 1} \sim p_{\theta}(x^{t + 1}|\hat{x}^t,z)$  and  $z \sim q_{\phi}(z|x^{t + 1},\hat{x}^t)$  as illustrated in Figure 2b. Using the posterior distribution over the latent space steers the model towards roughly following the ground truth. Additionally, instead of only calculating the loss at the last step of the rollout as done in [9], we calculate and backpropagate the loss at every step. $^{1}$
4 Applications and model designs
4.1 Celestial dynamics with stochastic forcing
Problem setting. For the first problem, we consider three-dimensional celestial mechanics, which are traditionally modeled as a deterministic n-body system. However, in some cases the effects of dust distributed in space need to be taken into account, which are typically modeled as additive Gaussian noise on the forcing terms [5, 33, 45]. In this setting, the dynamics of the system are governed by the following SDE:
$$
d x _ {i} = v _ {i} d t \quad \forall i \in [ n ] \tag {6}
$$
$$
d v _ {i} = \sum_ {j \neq i} \frac {G m _ {j} \left(x _ {j} - x _ {i}\right)}{\left\| x _ {j} - x _ {i} \right\| ^ {3}} d t + \sigma (x) d W _ {t} \quad \forall i \in [ n ] \tag {7}
$$
where  $\sigma(x) \in \mathbb{R}$  is  $1\%$  of the magnitude of the deterministic part of the acceleration of the  $i$ 'th body. We approximate Equation 6 with the Euler-Maruyama method to generate the data. Because of the chaotic nature of n-body systems, the accumulation of the noising term  $\sigma(x)dW_{t}$  leads to strongly diverging trajectories, posing a challenging modeling problem.
Symmetries and backbone architecture. We model the problem as a fully connected geometric graph in which the nodes' coordinates evolve over time. As such, the relevant symmetries are (1) permutation equivariance of the nodes, and (2)  $\mathrm{E}(\mathrm{n})$ -equivariance. We use a frame averaging GNN model as detailed in Section 3.3 of [38] as backbone architecture, denoted as FA-GNN, since it has demonstrated state-of-the-art performance in n-body dynamics prediction [38]. Using FA-GNN, we can parameterize functions that are invariant or equivariant to Euclidean symmetries. We also tried the EGNN [44], but during preliminary experimentation we found that training was relatively unstable in our setting, especially combined with multi-step training. The FA-GNN suffered substantially less from these instabilities.
Model design. Before explaining the individual model components, we first present the overall model structure, also depicted in Figure 3a:
- Forward model  $f_{\theta}$ : FA-GNNequivariant(xt)
- Conditional prior  $p_{\theta}(z|h^t)\colon \mathcal{N}\left(\mu^{\mathrm{prior}}(h^t),\sigma^{\mathrm{prior}}(h^t)\right)$
- Decoder  $p_{\theta}(x^{t + 1}|z,h^{t}) \colon \mathcal{N}\left(\mu^{\mathrm{decoder}}(z,h^{t}),\sigma^{\mathrm{decoder}}(z,h^{t})\right)$
<sup>1</sup> However, just as in [9], the gradients are only backpropagated through a single autoregressive step.
(a) Model structure for celestial dynamics.
Figure 3: Schematic overviews of model designs.
(b) Model structure for cellular dynamics.
As  $f_{\theta}$  needs to be equivariant, we parameterize it with the equivariant version of the FA-GNN, denoted as  $\mathrm{FA - GNN}_{\mathrm{equivariant}}$ . Recall that  $p_{\theta}(z|h^t)$  can be either invariant or equivariant. We choose for node-wise latent variables that are invariant to Euclidean transformations. As such, we parameterize  $\mu^{\mathrm{prior}}$  and  $\sigma^{\mathrm{prior}}$  with FA-GNNinvariant. Finally,  $p_{\theta}(x^{t + 1}|z,h^t)$  needs to be equivariant, which we achieve with a normal distribution where the parameters  $\mu^{\mathrm{decoder}}$  and  $\sigma^{\mathrm{decoder}}$  are modeled with FA-GNNequivariant and FA-GNNinvariant respectively. More details can be found in Appendix C.1.
4.2 Stochastic cellular dynamics
Problem setting. The second problem considers spatiotemporal migration of biological cells. The system is modeled by the Cellular Potts model (CPM), consisting of a regular lattice  $L$ , a set of cells  $C$ , and a time-evolving function  $x: L \to C$  [13]. Furthermore, each cell has an associated type  $\tau(c)$ . See also Figures 1 and 3b for illustrations. Specifically, we consider an extension of the cell sorting simulation, proposed in [13], which is a prototypical application of the CPM. In this setting, the system is initialized as a randomly configured culture of adjacent cells. The dynamics of the system are governed by the Hamiltonian, which in the cell sorting case is defined as follows:
$$
H = \underbrace {\sum_ {l _ {i} , l _ {j} \in \mathcal {N} (L)} J \left(x \left(l _ {i}\right) , x \left(l _ {j}\right)\right) \left(1 - \delta_ {x \left(l _ {i}\right) , x \left(l _ {j}\right)}\right)} _ {\text {c o n t a c t e n e r g y}} + \underbrace {\sum_ {c \in C} \lambda_ {V} \left(V (c) - V ^ {*} (c)\right) ^ {2}} _ {\text {v o l u m e c o n s t r a i n t}}, \tag {8}
$$
where  $\mathcal{N}(L)$  is the set of all adjacent lattice sites in  $L$ ,  $J(x(l_i),x(l_j))$  is the contact energy between cells  $x(l_{i})$  and  $x(l_{j})$ , and  $\delta_{x,y}$  is the Kronecker delta. Furthermore,  $C$  is the set of all cells in the system,  $V(c)$  is the volume of cell  $c$ ,  $V^{*}(c)$  is the target volume of cell  $c$ , and  $\lambda_V$  is a Lagrange multiplier. To evolve the system, an MCMC algorithm is used, which randomly selects an  $l_{i}\in L$ , and proposes to change  $x(l_{i})$  to  $x(l_{j})$ , where  $(l_i,l_j)\in \mathcal{N}(L)$ . The change is accepted with probability  $\min (1,e^{-\Delta H / T})$ , where  $T$  is the temperature parameter of the model. The values of the model parameters can be found in Appendix B.
Symmetries and backbone architecture. As the domain  $L$  is a regular lattice, one relevant symmetry is shift equivariance. Further, since  $x$  maps to a set of cells  $C$ , the second symmetry we consider is equivariance to permutations of  $C$ ; see Figure 1 for an intuitive schematic. To respect both symmetries, we propose a novel message passing graph neural network architecture. We first transform the input  $x$  to a one-hot encoded version, such that  $x$  maps all lattice sites to a one-hot vector of cell indices along the channel axis. We then consider each channel as a separate node  $h_i$ . Subsequently, message passing layers operate according to the following update equation: