Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics
Authors: Koen Minartz
Yoeri Poels
Simon Martinus Koop
Vlado Menkovski
Venue: NeurIPS 2023 poster (2023)
Keywords: stochastic simulation
equivariance
dynamical systems
probabilistic simulation
generative models
TL;DR: We propose a method for equivariant simulation of stochastic spatiotemporal dynamics

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
Koen Minartz'
Yoeri Poels Simon Koop'
Vlado Menkovski
Data and Artificial Intelligence Cluster, Eindhoven University of Technology Swiss Plasma Center, École Polytechnique Fédérale de Lausanne [k minartz, y.r.j. j-polls, s s.m.koop, v. v.menkovski @uuue... ml
Abstract
Neural networks emerging tool for scalable data-driven simulation of high- dimensional dynamical systems, especially in settings where numerical methods are infeasible computationally expensive. Notably, it has been shown that incor- porating domain symmetries deterministic neural simulators can substantially improve their accuracy, sample efficiency, and parameter efficiency. However, to in- corporate symmetries in probabilistic neural simulators that can simulate stochastic phenomena, we need model that produces equivariant distributions over trajecto- ries, rather than equivariant function approximations In this paper, we propose Equivariant Probabilistic Neural Simulation (EPNS), a framework for autoregres- sive probabilistic modeling of equivariant distributions over system evolutions. We use EPNS to design models for 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 promising method for efficient and effective data-driven probabilistic simulation diverse range domains.
Introduction
The advent of fast and powerful computers has led to numerical simulations becoming key tool in the natural sciences. For example, in computational biology, models based on proba- bilistic cellular automata are used to effectively reproduce stochastic cell migration dynamics [2] 13 26 [51), and in physics, numerical meth- ods for solving differential equations are used in diverse domains such weather forecast-
Pola (aa)
Simurata & sleps
Permuee
indices
ing 4 32], thermonuclcar fusion T6. 55], and computational astrophysics 24 52) amongst
Do (P "PP'
others. Recently, data-driven methods driven by neural networks have emerged complemen- tary paradigm to simulation. These neural sim- ulators particularly useful when numerical methods are infeasible or impractical, for exam- ple due computationally expensive simulation protocols or lack of an accurate mathematical model
stops
Figure 1: Schematic illustration of EPNS applied stochastic cellular dynamics. EPNS models distributions over trajectories that are equivari- permutations P of the cell indices, so that pe (P (Px*) pa (x (l*)).
30 37 42
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Notably, most neural simulators are deterministic: given some initial condition, they predict 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 [0), all of which imperative for high-quality neural simulator. As such, both the abilities 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-awant simulation. For probabilistic spatiotemporal simulation, methods like neural stochastic differential equations (NSDEs) 20 27 neural stochastic partial differential equa- tions (NSPDEs) 40, Bayesian neural networks 14 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 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 wwo-dimensional data with rotation symmetry.
are limited to the deterministic case 7 I8 19.31
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 simula- tion of dynamical systems. this work, we propose 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 prob- abilistic 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 grid domain. For the first problem, incorporate an existing equivariant architecture in EPNS. For the second, we propose a novel graph neural network that is equivariant permutations of cell indices, as illustrated in Figure [
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.
2 Background and related work
2.1 Equivariance
For given group of transformations G, function : X + X is G-equivariant if
=
holds for all g E G, where p(g) denotes the group action of g on x E X 12) In the context of probabilistic models, a conditional probability distribution p(aly,, =) is defined to be G-equivariant with respect to y if the following equality holds for all g G:
p(rly, =) = p()))))))(s))..).
The learning of such equivariant probability distributions has been employed for modeling symmetri- cal density functions, for example for generating molecules in 3D (Euclidean symmetry) [25 43 57
and generating sets (permutation symmetry) G 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 [7] 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 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 3 II I8] 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 [3 53 [4], climate science 77. 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 (58) simulate dynamics by probabilistically evolving (stochastic) differential equation in latent space. Alternatively, methods like Bayesian neural networks (BNNs) and Gaussian processes (GPs) have also been applied probabilistic spatiotemporal simulation 33 55]. 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-agen 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 Method
3.1 Modeling framework
Problem formulation. At time 6, the system's 's state is denoted as x", and a sequence of states denoted as o. We presuppose that we have samples from a ground-truth distribution p (2 for training, and for inference we assume that x is given work, we assume Markovian dynamics, although our approach can be extended to non-Markee non-Markovian dynamics as well. Consequently, the modeling task boils down to maximizing Equation 3
starting point for the simulation. In this
Eg Ee EE log pa (x ++1 (x=))
As such, we can approximate p" (a: (=T))) by learning pe pa(x=*1 [x²) and applying it autoregressively. Furthermore, when we know that the distribution p* (a.1:T (2:)) is equivariant to p(g). want to guarantee that po (x 1:T (a=0) is equivariant to these transformations as well. This means that
po(x (T (20) = po(p(g) (e(ll))xxT((())))) (((g)>>)
should hold, where p1g)x117 1:T simply means applying p(g) to all states in the sequence 1:T
Generative model. The model pa ((r++ should adhere to three requirements. First, since our goal is to model equivariant distributions over trajectories diverse domains, the type of model used for should be amenable to incorporating general equivariance constraints. Second, we argue that generative models with itcrative sampling procedures, although capable of producing high-quality samples [15 46). are impractical in this setting 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.
need to repeat the sampling procedure over
As such, we model po (x*+1|zz (2²) as Je (xr**111 1/gee+1|2, 12, (o dz by using latent variable 2, following the conditional Variational Autoencoder (CVAE) framework (47. Specifically, to generate x++1 I' is first processed by a forward model fe, producing embedding h' fo(z"). ). If there is little
3

---

ooward Model
Forward Model
Decader
fa
h'
(a
s.
h
Pi (a j.
Conditional Prior
Approximate Posterior
2
He (zilf -
2
(at
Pa
2-1-1
(a) Sampling.
(h) Training.
Figure EPNS model overview.
variance in p" (x*+1 [2=), ), h contains almost all information for generating x'+1_ On the other hand, if the system contains implicit information about possible continuations of the trajectories. To enable the model to simulate the latter kind of behavior, the embedding h is subsequently processed by the conditional prior distribution. Then, 2 ~ po(2)h² is processed in tandem with by the decoder, producing a distribution Pa (2 |z, h²) over the next state. The sampling process is summarized in Figure 2a
point where random effects result in strongly bifurcating trajectories. h' only
Incorporating symmetries. Using the approach described above, we can design models that can simulate trajectories by iteratively sampling from po(x:++12=") Now, the question how we can ensure that po (x: 1: | (x:)) respects the relevant symmetries. In Appendix A we prove Lemma] |
Lemma 1. Assume we model po 1-T po:x1-T)
(x:") as described in Section 3.1 Then, po (x)) is equiv- ariant linear transformations p(g) of a symmetry group the sense of Definition if:
(a) The forward model is G-equivariant: fo(p(g)x) plg)f((x);
(b) The conditional prior is G-invariant G-equivariant: po(z|p(g)h) po(2[h'), or po (()))" =oo po(z)h();
(c) The decoder is G-equivariant with respect to he (For an invariant conditional prior), G-equivariant with respect to both h'' and (for equivariant con- ditional prior): po pu(x+11/2 (2, h²)
or po(x 1, 22,
po (((g))
The consequence of Lemma is that, given equivariant neural network layer, we can stack these layers carefully parameterize 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 and p(g)æ* First, applying fo to xf yields ht, and since fe is G'eqquivarian applying it to p(g)x* yields p(g)h Second, due to the invariance of po(z(h'), both ht' and p(g) lead to exactly the same distribution over z. Third, since ~ Po (=(h1) is used parameterize the decoder pox h²), which is equivariant with respect to h², the model's one-step distribution po (x** ++ ) is equivariant as well. Finally, autoregressively sampling from the model inductively leads to equivariant distribution over the entire trajectory.
Invariant versus equivariant latents. When po(2(h²) is equivariant, the latent variables are typ- ically more local in nature, whereas they are generally more global in the case of an invariant conditional prior. For example, for graph domain, we may design Po(z[h*) to first perform commutaton-invariaat node aggregation to obtain global latent variable, alternatively equivari- antly map each node embedding to a node-wise distribution over to obtain local latent variables. Still, both design choices result in an equivariant model distribution po (x.1: (x0).
3.2 Training procedure
Following the VAE framework 23
approximate posterior distribution qs { xf+1) is 39
learned to approximate the true posterior pa (2|æ² 2.)1). During training, go is used to optimize the

---

Evidence Lower Bound (ELBO) on the one-step log- likelihood:
(5) In our case, 9 takes both the current and next timestep as imput to infer a posterior distribution over the latent variable z. Naturally, parameterize 4a be invariant equívariant with respect to both xf and x/+1 for an invariant equivariant conditional prior po (2/x²) respectively.
log po (x'=11x') (as(s [po(='t||,,) |2, x²]] KL [4s (=| ,ot , (po((())))
well-known challenge for autoregressive neural simulators concerns the accumulation of error over long rollouts 19 48] imperfections in the model lead to its output deviating further from the ground-trun -ruth for increasing rollout length k. producing feedback loop that results in exponentially increasing errors. In the stochastic setting, we observe similar phenomenon: pe p= does not fully match p" (a*++|2>" |z²),, resulting in a slightly biased one-step distribution. Empirically, we see that these biases accumulate, leading to poor coverage of p" p> (x²) for > 1, and miscalibrated uncertainty estimates. Fortunately, we found that 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 po (z -+++1 |x²)) can lead to the simulated trajectory diverging from the ground-trute sample. Although this behavior desired for stochastic simulation model, it ruins the training signal. Instead, we iteratively sample reconstruction a++1 from the model, where &++1 Po(x*+1j0 | 2) and (s(2|2*+1 .*) 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
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 traditionally modeled as 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 In this setting, the dynamics of the system governed by the
noise the forcing terms [5 33 45)
following SDE: dx, vddl
Vi E [n]]
Gmj(# xi) ||., x4
dv; Σ
dt + o(r)dW,
Vi E [m]
1#1
where (xx) R is 1% of the magnitude of the deterministic part of the acceleration of the i'th body. We approximate Equation with the Fuler-Maruyama method to generate the data, Because of the chaotic nature of n-body systems, the accumulation of the noising term a (x:) (x)dWW 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) E(n)-qquivariance. We use a frame averaging GNN model as detailed Section 3.3 of (38) as backbone architecture, denoted as FA-GNN, since has demonstrated state-of-thee tooftthe-art art performance in n-body dynamics prediction [8 Using FA FA-GNN, can parameterize functions that are invariant or equivariant to Euclidean symmetries. We also tried the EGNN [44), but during preliminary experimentation found that training 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 Ba
Forward model fo: FA-GNN CNNuuiivariin (x²)
Conditional prior pe (p(22 (z)ht); N ((otio prior (h²))
* Decoder
(= (u ht)): N (ndecoderzz, 2, h(',,, o deccdder ht"))
'However, just in 99. the gradients only backpropagated through single autoregressive step.
5

---

po(=/k²)
P=(=/h)
Vodinsso convolutons spatinl aggregation
Plo Par"
2,
FA-GHN PA-GiAiunvecne
FA- Dillppicnia
SpatialCoavy-Or
(a) Model structure for celestial dynamics.
(b) Model structure for cellular dynamics.
Figure 3: Schematic overviews of model designs.
As fa needs to be equivariant, we parameterize it with the equivariant version of the FA-GNN, denoted PA-GGNepuivarinn-- Recall that po (b4) can be either invariant equivariant. We choose for node- wise latent variables that invariant to Euclidean transformations. As such, we parameterize pariorr and orior with FA-GNN Finally, Ps (x:"_1 2, h²) needs to be equivariant, which we achieve with normal distribution where the parameters y de and deesder are modeled with FA- FA-GGNeeuivariant and FA FA-GNN respectively. More details can be found in Appendix C.I
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 regular lattice L, set of cells C, and a time-evolving function x - C' [13. Furthermore, each cell has an associated type (C).. See also Figures and 3b for illustrations. Specifically, we consider extension of the cell sorting simulation, proposed [3]. which is a prototypical application of the CPM. In this setting, the system is initialized as 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
J (x(1,,,(())) (1 52(1).0(1)) + Σ Av -Vc))- AL)EN(()
Σ
(8)
contact energy
volume constraint
where A(L) the set of all adjacent lattice sites L, (x(44, x(1,)) is the contact energy between cells a1,4) and sr(2z), and dz.y is the Kronecker delta. Furthermore, C' is the set of all cells in the system, Vec is the volume of cell c, V*(c) the target volume of cell e, and Xy multiplier. To evolve the system, an MCMC algorithm is used, which randomly selects an I E T., and proposes to change x(4) to I()), where (4, (L). The change is accepted with probability min(1, e AI/TT) where is the temperature parameter of the model. The values of the model parameters be found in Appendix B
Lagrange
Symmetries and backbone architecture. As the domain L is a regular lattice, one relevant symmetry shift equivariance. Further, since x maps set of cells the second symmetry we consider is equivariance to permutations of C; see Figure [ for an intuitive schematic. To respect both symmetries, we propose novel message passing graph neural network architecture. We first transform the input to one-hot encoded version, such that x maps all lattice sites to one-hot vector of cell indices along the channel axis. We then consider each channel as a separate node h; Subsequently, message passing layers operate according to the following update equation:
1 vb
h// JEN(O)
o' (h1))
(9)

---

where + is ietiuuttttiiiiinvviinn aggregation function, performed over neighboring nodes. Usually, v/ and ø² would be parameterized with MIPs, but since Ay lives on a grid domain, we choose both functions be convolutional neural nets. We refer this architecture SpatialConv-GNW
Model design.
We again start by laying out the overall model structure, also shown in Figure 3b
Forward model fe: PpatiaaCConvGGMM(c
Conditional prior: pa p=(=)("): N (netior (f'), o
Decoder: Po (1**1111
+ |2,
h²:: C
(orecoderr (2,
Since fe is parameterized by At any point time, each cell is only affected by itself and the directly adjacent cells, so we design pe(=lh") to be equivariant with respect to permutations, but invariant with respect to translations; concretely, we perform alternating convolutions and spatial pooling operations for each node h; to infer distributions over a latent variable for each cell. Finally, the decoder is parameterized by (2, )²), SpatialConv categorical distribution over the set of cells C. More details can be found in Appendix C.2
it is equivariant to translations and permutations.
which equivariantly maps the parameters of pixel-wise
5 Experiments
5.1 Experiment setup
Objectives. Our experiments aim to a) assess the effectiveness of EPNS in terms of simulation faithfulness, calibration, and stability; b) empirically verify the equivariance property of EPNS, and investigate the benefits of equivariance in the context of probabilistic simulation; and c) compare EPNS existing methods for probabilistic simulation of dynamical systems. We report the results conveying the main insights in this section; further results and ablations are reported in Appendix E
Data. For both problem settings, we generate training sets with 800 trajectories and validation and test sets each consisting of 100 trajectories. The length of each trajectory is 1000 for the celestial dynamics case, saved temporal resolution of dt 0.01 units per step. We train the models to predict a time step of At = 0.1 ahead, resulting in an effective rollout length of 100 steps for a full trajectory. Similar earlier works 38 44 59], we choose n = 5 bodies for our experiments. For the cellular dynamics case, we use trajectories with |C| = 64 cells. The length of each trajectory is 59 steps, where we train the models to predict a single step ahead. Additionally, we generate test sets of 100 trajectories starting from the same initial condition, which are used to assess performance regarding uncertainty quantification.
Implementation and baselines. The EPNS implementations are based on the designs explained in Section 4 The maximum rollout lengths used for multi-step training are 16 steps (celestial dynamics) and 14 steps (cellular dynamics). We use linear KL-annealing schedule [8], as well as the "free bits' modification of the ELBO as proposed in (2 for the cellular dynamics data. For celestial dynamics, we compare to neural SDEs (NSDE) [27] and interacting Gaussian Process ODEs (iGPODE) [59]. as these methods can probabilistically model dynamics of (interacting) particles. For cellular dynamics, we compare to ODE*VAD (58] it probabilistically models dynamics on grid domain. For each baseline, we search over the most relevant hyperparameters, considering at least seven configurations per baseline. We also compare to non-equivariant counterpart of EPNS, which we will refer t to as PNS. All models are implemented in Py Torch (35) and trained on a single NVIDIA A100 GPU, Further details on the bascline models can be found in Appendix D Our code is available at https:// thub. com/ km kt////matttttt//SSS
5.2 Results
Qualitative results. Figure 4 shows a sample from the test set, as well as simulations produced by neural simulators starting from the same initial condition 20 for both datasets. More samples are provided in Appendix E For celestial dynamics, both EPNS and iGPODE produce a simulation that looks plausible, as all objects follow orbiting trajectories with reasonable velocities. As expected, the simulations start to deviate further from the ground truth time proceeds, due to the accumulation
4igure
7

---

Ground truth
Ground truth
EPNS
EPNS
IGPODE
ODE
0
t= 20
40
0
t=
t= 10
15
Figure 4: Qualitative results for celestial dynamics (left) and cellular dynamics (right).
Table I:
Comparison of EPNS to related work. For celestial dynamics, the kinetic energy of the system is used to calculate Dks, while for celullar dynamics, the same type. We report the standaad error over three independent training runs.
the number of cell clusters of
Celestial dynamics
Cellular dynamics
LL
Des (KE)
LL -10
(xs(Hcclustes) L
Model
Model
-10
1=50
t=100
1=30
t=45
NSDE
-5.0:0.00
0.8011.01
0.65=10.33
ODE'VAE 58
-30.1+0.00
0.98+0.01
0.96++0.04
iGPODE 59
-8.510.1
0.42:0.07
0.57:0.02
PNS (ours)
10.91:.4
0.6110.18
0.2011.066
PNS (ours)
-16.410.3
0.70
0.77=10.05
EPNS (ours)
10.8:0.11
0.1410.03
0.1410.04
EPNS (ours)
-5.910.1
0.58
0.58
of randomness. In the cellular dynamics case, the dynamics produced by EPNS look similar to the ground truth, as cells of the same type tend to cluster together. In contrast, ODE2VAL fails to capture this behavior. Although cells located roughly the right location for the initial part of the trajectory, the clustering and highly dynamic nature of the cells' movement and shape absent.
Comparison to related work. We consider two metrics to compare to baselines: the first is the ELBO on the log-likelihood on the test set, denoted LL. LL is widely used metric in the generative modeling community for evaluating probabilistic models, expressing goodness-of- on unseen data, see for example [15 [2 23 59]. Still, the calculation of this bound differs per model type, and it does not measure performance in terms of new sample quality. This is why the second metric, the Dolmogorov-Sminoo (KS) test statistic DKs. compares empirical distributions of newly sampled simulations to ground truth samples. Specifically, Dks is defined as Dxs (t) max F (y) Fo
the largest absolute difference between the ground truth and model empirical cumulative distribution functions F* and Fg. Dks 0 corresponds identical samples, whereas DRs corresponds to no distribution overlap at all. We opt for this metric since we have explicit access to the ground-truth simulator, and such we can calculate DKs over 100 simulations starting from fixed x order to assess how well EPNS matches the ground-truth distribution over possible trajectories. The results are shown in Table D indicating distribution that has high peaks around the ground truth samples. In terms of DRs. EPNS performs best, meaning that simulations sampled from EPNS also match better to the ground truth in terms of the examined properties of Table
EPNS performs comparative to or better than the other methods in terms of LL,
Uncertainty quantification. In addition to Tables I and 3 we further investigate how distributions produced by EPNS, PNS, and EPNS trained with one step training (EPNS-on one step) align with the ground truth. Figure 5 shows how properties of 100 simulations, all starting from the same z0, evolve over time. In the case of celestial dynamics, the distribution over potential energy produced by EPNS closely matches the ground truth, as shown by the overlapping shaded areas. In contrast, EPNS--oe step effectively collapses deterministic model. Although PNS produces distribution that overlaps with the ground truth, it doesn't match as closely EPNS. For the cellular dynamics

---

EPNS
EPNS-one one step
PNS
EPNS
ground truth
BPNS
step
ground wuth
PNS
ground trath
eoergy
Poterital eneegy
Poleniaal energy
-75
-75
-75
-100
-100
-100
-126
-125
122
26
60
75
100
0
60
103
50
76
100
Time step
Time step
Time step
EPNS
grouna Inuh
EPNS-
use viep
grondd Inch
PNS
grrand Inuh
# clossers
20
cell clusters
20
# cell choserrs
20
10
Time
Time step
Time step
Figure 5: Distributions of the potential energy in system (top row, celestial dynamies) and the
number of cell clusters of the same type (bottom row, cellular dynamics) over time. The shaded area indicates the (0.1, (.99-quanntil. interval of the observations; the solid line indicates the median value.
case, the distribution produced by EPNS is slightly off compared to the ground truth, but still matches considerably better than EPNS-one step and PNS. These results indicate that both equivariance and multi-step training improve uncertainty quantification for these applications.
Verification of equivariance.
Table 2: KS-test results for equivariance verification.
To empirically verify the equivari- ance property of EPNS, we sam-
Celestial Dynamics
Cellular Dynamics
ple ten simulations for each x in
x-coordinate y-coordinate z-coordinate Model Des p-value Des p-value Des p-value Des
distance traveled
the test set, leading to 1000 simu- lations in total, and apply a trans- formation p(g) the final outputs x² pf) is a random rotation and random permutation for celestial dynamics and cellular dynamics respectively. Let us refer to the transformed outcomes as distribution We then repeat the procedure, but now transform the inputs x by the same transformations before applying EPNS, yielding outcome distribution 2. For an equivariant model, distributions I and 2 should be statistically indistinguishable. We test this by comparing attributes of a single randomly chosen body or cell using two-sample KS-test. We also show results of PNS to investigate equivariance is learned by non-equivariant model. The results are shown in Table 2 In all cases, distributions and 2 are statistically indistinguishable for EPNS (p >> 0.05), but not for PNS (p < 0.05). As such, EPNS indeed models equivariant distributions over trajectories, while PNS does not learn to be equivariant.
p-value
PNS 0.09 0.00
0.22
0.00
0.09
0.00
0.503
0.00
EPNS 0.04
0.37
0.02
0.99
0.02
0.98 0.036
0.536
Data efficiency. We now examine the effect of equivariance on data efficiency. Table 3 shows the ELBO of both EPNS and PNS when trained on varying training set sizes. In the celestial dynamics case, the ELBO of EPNS degrades more gracefully when shrinking the training set compared to PNS. For cellular dynamics, the ELBO of EPNS does not worsen much at all when reducing the dataset to only 80 trajectories, while the ELBO of PNS deteriorates rapidly. These results demonstrate that EPNS is more data efficient due to its richer inductive biases.
Table 3: Test set ELBO values (-10²) for varying
amounts of training samples.
Training Celestial Dynamics Cellular Dynamics samples
EPNS
PNS
EPNS
PNS
80
9.9
7.7
-59.6
-326.5
400
10.2
10.0
-61.0
-173.9
800
10.8
10.9
-59.4
-163.9
Rollout stability. As we cannot measure stability by tracking error metric over time for stochastic dynamics, we resort to domain-specific stability criteria. For celestial dynamics, we define a run be unstable when the energy in the system jumps by more than 20 units, which we also used to ensure

---

Fraction of stable runs
1.00
1.00
0.75
Fraction of stuble
0.75
FPNS
iGODDR
0.50
0.50
PNS
NSDE
ODE VAE
0.25
0.00
0.00
200
400
600
800
1000
0
10
20
30
40
50
60
Time step
Time step
(a) Celestial dynamics.
(b) Cellular dynamics.
Figure 6: Fraction of simulations that remain stable over long rollouts.
the quality of the training samples. For cellular dynamics, we define a run to be unstable soon as 20% of the cells have volume outside the range of volumes in the training set. Although these are not the only possible criteria, they provide reasonable proxy to the overall notion of stability for these applications. The fraction of stable simulations over time, starting from the initial conditions of the test set, is depicted in Figure o For both datasets, EPNS generally produces simulations that remain stable for longer than other models, suggesting that equivariance improves rollout stability. Although stability still decreases over time, most of EPNS' simulations remain stable for substantially longer than the rollout lengths seen during training.
Conclusion
Conclusions. In this work, we propose EPNS, generally applicable framework for equivariant probabilistic spatiotemporal simulation. We evaluate EPNS in the domains of stochastic celestial dynamics and stochastic cellular dynamics. Our results demonstrate that EPNS outperforms existing methods for these applications. Furthermore, we observe that embedding equivariance constraints in EPNS improves data efficiency, stability, and uncertainty estimates. In conclusion, we demonstrate the value of incorporating symmetries in probabilistic neural simulators, and show that EPNS provides the means achieve this for diverse applications.
Limitations. The autoregressive nature of EPNS results in part of the simulations becoming unstable as the rollout length increases. This a known limitation of autoregressive neural simulation models, and we believe that further advances in this field can transfer to our method [9 SpatialConv-
Second, the relatively large memory requirements, and would interesting to improve the architecture to be more memory efficient. Third, we noticed that training can be unstable, either due to exploding loss values or due to getting stuck in bad local minima. Our efforts to mitigate this issue resulted in approximately two out of three training runs converging to good local optimum, and we believe this could be improved by further tweaking of the training procedure. Finally, our experimental evaluation has shown the benefits of equivariant probabilistic simulation using relatively small scale systems with stylized dynamics. An interesting avenue for future work to apply EPNS to large-scale systems with more complex interactions that of significant interest to domain experts various scientific fields. Examples such applications could be the simulation Langevin dynamics or cancer cell migration.
48 49
7 Acknowledgements
The authors would like to thank Jakub Tomczak and Tim d'Hondt for the insightful discussions and valuable feedback. This work used the Dutch national e-infrastructure with the support of the SURF Cooperative using grant no. EINF-3935 and grant no. EINF-3557.
10

---

References
[1] Kelsey R Allen, Yulia Rubanova, Tatiana LopezaCuevara, William F Whitmey, Alvaro Sanchez- Gonzalez, Peter Battaglia, and Tobias Pfaff. Learning rigid dynamics with face interaction graph networks. In International Conference on Learning Representations, 2023.
[2] Ariel Balter, Roeland M. H. Merks, Nikodem J. Popławski. Maciej Swat, and James A. Glazier. The Blazier-Graner-Hoggeww Model: Extensions, Future Directions, and Opportunities for Further Study, pages 151-167. Birkhäuser Basel, Basel, 2007.
[3] Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and koray kavukcuoglu. Interaction networks for learning about objects, relations and physics. In Advances in Neural Information Processing Systems, 2016.
[4] Peter Bauer, Alan Thorpe, and Gilbert Brunet. The quiet revolution numerical weather prediction. Nature, 42557567):47-55, September 2015.
[5] M. M. Bierbaum, R. I. Joseph, R. L. Fry, and J. B. Nelson. A fokker-planck model for a two-body problem. AIP Conference Proceedings. 617(1):340-37,, 2002.
[6] Marin Biloš and Stephan Günnemann. Scalable normalizing flows for permutation invariant densities. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[7] Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Modelling atmospheric dynamics with spherical fourier neural operators. In ICLR 2023 Workshop on Tackling Climate Change with Machine Learning. 2023.
[8] Samuel R. Bowman, Luke Vilnis, Oriol Vinyals, Andrew M. Dai. Rafal Józefowicz, and Samy Bengio. Generating sentences from a continuous space. In Proceedings of the 20th SIGNLL Conference Computational Natural Language Learning, pages 10-21. ACL,, 2016
[9] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, 2022.
[10] Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges, 2021.
[11] Michael Chang, Tomer Ullman, Antonio Torralba, and Joshua Tenenbaum. A compositional object-based approach to learning physical dynamics. In International Conference on Learning Representations, 2017.
[12] Taco Cohen and Max Welling. Group equivariant convolutional networks. Proceedings of The 33rd International Conference Machine Learning, 2016.
[13] François Graner and James A. Glazier. Simulation of biological cell sorting using a two- dimensional extended potts model. Phys. Rev. Lett., 09:2013-2015 Sep 1992.
[14] Jayesh K. Gupta and Johannes Brandstetter. Towards multi saatiisppppottmprra-sslaal generalized pde modeling. arXiv preprint, 2022.
[15] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[16] M. Hoelzl, G.T.A. Huijsmans, S.J.P. Pamela, M. Bécoulet, E. Nardon, F.J. Artola, B. Nkonga, C.V. Atanasiu, V. Bandaru, A. Bhole, D. Bonfiglio, A. Cathey, O. Czarny, A. Dvorovaa, T. Fehér, A. Fil, E. Franck, S. Futatani, M. Gruca, H. Guillard, J.W. Haverkort, I. Holod, D. Hu, S.K. Kim, S.Q. Korving, L. Kos, Krebs, Kripner, G. Latu, F. Liu, P. Merkel, D. Meshcheriakov, V. Mit- terauer, S. Moochalkky,, J.A. Morales, R. Nies, N. Nikulsin, F. Orain, J. Pratt, R. Ramasamy, Ramet, C. Reux, Särkimäki, N. Schwarz, P. Singh Verma, S.F. Smith, Sommariva, E. Strumberger, D.C. van Vugt, M. Verbeek, E. Westerhof, F. Wieschollek, and J. Zielinski. The JOREK non-linear extended MHD code and applications to large-scale instabilities and their control in magnetically confined fusion plasmas. Nuclear Fusion, (166000,,,,, May 2021.

---

[17] Peter Holderrieth, Michael Hutchinson, and Yee Whye Teh. Equivariant learning of stochastic fields:
Gaussian processes and steerable conditional neural processes. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 4297-4307. PMIR, 18-24
Jul 2021.
[18) Masanobu Horie and Naoto Mitsume. Physics-embedded neural networks: Graph neural pde solvers with mixed boundary conditions. In Advances in Neural Information Processing Systems, 2022.
[19] Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equiv- ariant graph mechanics networks with constraints. In International Conference on Learning Representations, 2022.
[20] Patrick Kidger, James Foster, Xuechen Li, and Terry J Lyons. Neural sdes as infinite-dimensioaal gans. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[21] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[22] Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, lya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems, 2016.
[23] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. International Conference Leurning Representations, 2014.
[24] Anatoly A. Klypin, Sebastian Mujillo-Gomez, and Joel Primack. Dark matter halos in the standard cosmological model: Results from the bolshoi simulation. The Astrophysical Journal, 740(2):102, October 2011.
[25] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: Exact likelihood generative learning for symmetric densities. In Proceedings of the 37th International Conference on Machine Learning, 2020.
[26] Sandeep Kumar, Aastha Kapoor, Sejal Desai, Mandar M. Inamdar, and Shamik Sen. Prote- olytic and non-proteolytic regulation of collective cell invasion: tuning by eem density and organization. Scientific Reports, 6(1):19905, Feb 2016.
[27] Xuechen Li, Ting-Kan Kam Leonard Wong, Ricky T. Q. Chen, and David K. Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In Proceedings of The 2nd Symposium on Advances in Approximate Bayesian Inference, 2020.
[28] Yang Li, Haidong Yi, Christopher Bender, Siyuan Shan, and Junier Oliva. Exchangeable neural ode for set modeling. In Advances in Neural Information Processing Systems, 2020.
[29] Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids. In International Conference Learning Representations, 2019.
[30] Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021.
[31] Julia Ling, Andrew Kurzawski, and Jeremy Templeton. Reynolds averaged turbulence mod- elling using deep neural networks with embedded invariance. Journal of Fluid Mechanics, 807:155-166, 2016.
[32] Peter Lynch. The origins of computer weather prediction and climate modeling. Journal of Computational Physics, 227(7):34313444 2008
[33] Matteo Manzi and Massimiliano Vasile. Analysis of stochastic nearly-ittegrable dynamical systems using polynomial chaos expansions. In 2020 AAS/AIAA Astrodynamics Specialist Conference. Univelt Inc/ American Astronautical Society, USA, June 2020.
12

---

[34] Maximilian Mueller, Robin Greif, Frank Jenko, and Nils Thuerey. Leveraging stochastic predictions of bayesian neural networks for fluid simulations. In NeurIPS 2022 Workshop on Machine Learning and the Physical Sciences, 2022.
[35] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems 32, 2019.
[36] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzale,, and Peter Battaglia. Learning mesh- based simulation with graph networks. In International Conference on Learning Representations, 2021.
[37] Yoeri Poels, Gijs Derks, Egbert Westerhof, Koen Minartz, Sven Wiesen, and Vlado Menkovski. Fast dynamic 1D simulation of divertor plasmas with neural PDE surrogates. Nuclear Fusion, 63(12) 63(12):126012, sep 2023.
[38] Omri Puny, Matan Atzmon, Edward J. Smith, Ishan Misra, Aditya Grover, Heli Ben-Hamu, and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022.
[39] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning, 2014.
[40] Cristopher Salvi, Maud Lemercier, and Andris Gerasimovies. Neural stochastic PDEs: Resolution-invariint learning of continuous spatiotemporal dynamics. In Advances in Neural Information Processing Systems, 2022.
[41] Alvaro Sanchez-Gonzale Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. In NeurIPS 2019 Workshop on Machine Learning and the Physical Sciences, 2019.
[42] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In Proceedings of the 37th International Conference on Machine Learning, 2020.
[43] Victor Garcia Satorras, Emiel Hoogeboom, Fabian Bernd Fuchs, Ingmar Posner, and Max Welling. B(n) equivariant normalizing flows. In Advances in Neural Information Processing Systems, 2021.
[44] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In Proceedings of the 38th International Conference on Machine Learning, 2021.
[45] Shambhu N Sharma and H Parthasarathy. Dynamics of a stochastically perturbed two-body problem. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 463(2080):99--00) 2007.
[46] Jascha Sohi-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Conference on Machine Learning, 2015.
[47] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in Neural Information Processing Systems, 2015.
[48] Kim Stachenfeld, Drummond Buschman Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez-Gonzakee Learned coarse models for efficient turbulence simulation. International Conference on Learning Representations, 2022.
13

---

[49] Jingtong Su, Julia Kempe. Drummond Fielding, Nikolaos Tsilivis, Miles Cranmer, and Shirley Ho. Adversarial noise injection for learned turbulence simulations. In NeurIPS 2022 Workshop on Machine Learning and the Physical Sciences, 2022.
[50] Sophia Huiwen Sun, Robin Walters, Jinxi Li, and Rose Yu. Probabilistic symmetry for multi- agent dynamics. In Nikolai Matni, Manfred Morari, and George J. Pappas, editors, Proceedings of The 5th Annual Learning for Dynamics and Control Conference, volume 211 of Proceedings of Machine Learning Research, pages 1231-1244. PMLR, 15-16 Jun 2023.
[51] Thomas Thenard, Anita Catapano, Michel Mesnard, and Rachele Allena. A cellular potts energy-based approach analyse the influence of the surface topography single cell motility. Journal of Theoretical Biology. 509:1 1 10487, September 2020.
[52] Mark Vogelsberger, Shy Genel, Volker Springel, Paul Torrey, Debora Sijacki, Dandan Xu, Greg Snyder, Dylan Nelson, and Lars Hernquist. Introducing the illustris project: simulating the coevolution of dark and visible matter in the universe. Monthly Notices of the Royal Astronomical Society, 444(2):1518-
August 2014.
[53] Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. In International Conference Learning Representations, 2021.
[54] Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. Proceedings of the 39th International Conference on Machine Learning. 2022.
[55] S. Wiesen, Reiter, Kotov, M. Baclmans, W. Dekeyser, A.S. Kukushkin, S.W. Lisgo, R.A. Pitts. V. Rozhansky, G. Saibene, I. Veselova, and S. Voskoboynikov. The new SOLPS-ITER code package. Journal of Nuclear Materials, 463:480-484, August 2015.
[56) Tailin Wu, Takashi Maruyama, Qingqing Zhao, Gordon Wetzstein, and Jure Leskovec. Learning controllable adaptive simulation for multi-resolution physics. In International Conference on Learning Representations, 2023.
[57] Minkai Xu, Lantao Yu. Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[58] Cagatay Yildiz, Markus Heinonen, and Harri Lahdesmaki. Ode2vae: Deep generative second order odes with bayesian neural networks. In Advances in Neural Information Processing Systems, 2019.
[59] Cagatay Yildiz, Melih Kandemir, and Barbara Rakitsch, Learning interacting dynamical systems with latent gaussian process ODEs. Advances Neural Information Processing Systems, 2022.
14

---

A Proof of equivariance
Here, we prove Lemma
T which we repeat below:
Lemma 1. Assume we model Po(z ( as described in Section 3.1 that is, po(x 1:7 (0))
Io po (z") and po(x*+1(2") (xx) po(x<+11,, h³ po(z))")) where ht So (xx") Then, po(x"" 1:7 ( is equivariant to linear transformations p(g) of a symmetry group G in the sense of definition iif
(a) The forward model G-equivariant: fo(p(g)=) p(g)fo(>);
(b) The conditional prior is G-invariant or G-equivariant: pa(=lp(o)h)) po(=)h"), or po po(p(g) (())))) = po (2(
(c) The decoder is G-equivariant with respect to h' (for an invariant conditional prior), or G-equivariant with respect to both and z (for an equivariant con- ditional prior): pa p ht)
or po
po po(pg))= (p()))==
To support the proof, we first show that, if G-equivariant distribution p(y|æ) exists, i.e. p(y|z) p(p(g)))p(g)zt) then the absolute value of the determinant p(g) must equal 1. To this end, let yg pggyy and Bg p(g)z
plya ( ddya
Uu
[ plylz)dd, Ns
{ is equivariant }
J. pla)()) |det p(a)l dy
(substitute yo - -((g)}
p(و) J p(ulx)dy
dee
|eet p(g)
For the lemma to hold for first-order Markovian data-generating process, it suffices to have an autoregressive model of which the one-step transition distribution po (x( is equivariant. To show this, let r = p(g)x'
po(p(g)x! (p(g)=))
= po(xz (p())))
(129)
po ( (2n 2
(Markov property]
II Pe (af)(x" 1
[Equivariant one-step pa}
[Ineerse Markov property] }
As such, it remains to show that (o (1]++111ppp (xx) is equivariant when conditions (a)-(c) of the lemma hold. We show here the case where Pg (=]h²) ) equivariant, the invariant case analogous, where acts as 2 Iz. Again, let x p(g)x* and let 2g p(g)2:
15

---

po(p(g)=*+1 (p(g)==))))))))>>>))))))))))))") - = (m1g)
=po -=o(att11ng (61)
{equivariant fo
/ po po ( - hy)po ( (=g (Egltin) )i
po 11, hookppp(2 (2)
fequivariant po)
po ( 1z, (2 (=(h²) ddet p(a) dz
{ -[()==
po 1 (2o
{|det p(9)p = 1}
po (c'*t((x))
B Parameters of the Cellular Potts model
Recall that the Hamiltonian for the cell sorting simulation is defined as follows:
= Σ (()()) (1- L.ENNL
Σ Av (V())-V(())
(10)
CEC
For explanation the interpretation of this equation, please consult Section 4.2 values used in Equation can can be found Tables and below:
The parameter
Table 4: Parameter values of the Cellular Potts simulation.
Table 5: Contact energy J (x(4), xx(4)) lookup table.
Cell type
Medium
Type
Type 2
Type 3
Type
Parameter
Value
Medium
0
16
16
16
16
xx(4))((4))
See Table
Type
11
11
Av
Type 2
16
11
8
11
11
(c(
25
Type 3
11
11
11
T
Type 4
16
11
11
11
These parameters were not chosen to be biologically plausible, but rather to result in behavior in which there is clear tendency of clustering of same-type cells time proceeds, and to strike balance between the stochastic fluctuations and clustering tendencies in the system.
C
EPNS model details
C.1 Celestial dynamics
Before going into the specifics, first discuss some general remarks on the EPNS model imple- mentation. Recall that the data modeled as geometric graph consisting of n nodes. The celestial dynamics EPNS model distinguishes between E(n)-invariant features (mass ant features (position pI e 3 and velocity vi E R$)) for each node We provide the Euclidean distance (llpı' 2l2))) as well the magnitude of the difference between the velocity vectors (1|v;' vj" (]2) between all nodes as edge features et, R² to the model. ReLU is used throughout model as activation function. Our implementation do not construct and invert the frame at each message passing layer. Instead, we construct the frame once, then apply all message passing layers, and then invert the frame again to arrive invariant or equivariant output. Overall, the model has about 1.8 million trainable parameters.
IR-)) and equivari-
FA-GNN differs slightly from (38), as we
16

---

Forward model: m E R , E xx²2 -p' E Rux3 v$ E R + fo
h² E R2x ERRxx2 p' g v E [R=3]
First, the forward model independently embeds each node's mass to an invariant node embedding h' R22 independently of position and velocity attributes. The number of hidden dimensions is kept constant throughout the remainder of the model for the node embeddings. Then, we apply the FA-GNN to the graph for message passing layers to update the invariant embedding vectors the input features are the invariant node embedding as well as the coordinate and velocity vectors of each node. The forward model does not yet update the coordinate and velocity vectors, but simply applies the identity function to the input coordinates and velocities, as we found that it worked better to update these vectors the decoder only. Note that this trivially implies an equivariant forward model. As such, in this stage, the spatial information is only used to construct useful E(n) -nnvaraat node embedding.
Conditional prior:
{ E R2x122 e' E [R p=(2)h)) (= [
As we chose a conditional prior that is permutation equivariant, the conditional prior architecture is a message passing GNN with 5 layers that takes input the Ei()-invariiat node embeddings produced by the forward model. A linear layer maps the node embeddings to the mean of the conditional prior distribution, while linear layer followed by Softplus activation maps the embeddings to the scale parameter. We use a 16-dimensional latent space for each node.
Approximate posterior:
et1 I E 128 et E TRn ett E m=rx2]] (o(")"" x/t)) - {a E R
The architecture of the approximate posterior is identical to the conditional prior, with the exception that it accepts additional edge features input. These the same E(n)-invariaat features that are also given as input to the forward model, but calculated for the system at time t + 1, instead time t.
Decoder:
h1 E R2x128 E BEREE I21,, e R=2x2 ,p'''' E R v E R - 16
(p't1 E Rax3 v41 E R"
The decoder is FA-GNN with three layers. As input, it gets the invariant node embeddings h' produced by the forward model, the invariant & sampled from the conditional prior (during generation) approximate posterior (during training), and the position and velocity vectors p' and v'. The decoder cquivariantly maps this input to the mean and scale parameters of the Gaussian output distribution. More specifically, the model outputs vectors م R and át R Furthermore, the resulting E(n)-invariant node embeddings are also independently processed by a linear layer with Softplus activation to map to dº and σ" and by two layer MLP mapping to quantity A E IR which is used to scale the velocity's effect on the node position update. Using the quantities described above, the decoder's output distribution parameterized as follows:
μ¹ = v + μ nº p' + A. 4² + Ap
po 1 2) N (
Po po 1 14',) =V (n", σ")
The reason for using AP is that we take relatively large stepsize of A/ = 0.01, and such Ju" might align with the difference between the coordinates p [+1 p'. Accordingly, AA can correct for this difference if necessary.
17

---

Training: For each epoch, we select one starting point uniformly at random from each trajectory in the training set for the multi-step training. We train the model for 40000 epochs with the Adam optimizer [2i] with learning rate equal to 10-1 and weight decay of 10-2. For the KI. annealing schedule, we increase the coefficient 9 that weighs the KL. term of the loss by 0.005 every 200 epochs. We use a batch size of 64. We do not apply the free bits variant of the ELBO for celestial dynamics.
C.2 Cellular dynamics
Again, we start with general model remarks before going into the specific components. For the cellular dynamics EPNS model, we get input a 2D grid of dimensions h X w with two channels: the first channel contains integer values between 0 and c, indicating the index of each cell, where O stands for background and |C| is the number of cells. The second channel contains integer values between 0 and that indicate the cell type, where O is again background, and S e the number of different cell types. We again use ReLU activation throughout the model as nonlinearity. Overall, the model has around 13 million trainable parameters.
We first describe the overall model structure in terms of the forward model, conditional prior, and decoder, after which describe the SpatialConvvGMi architecture for this model in more detail.
Forward model: {af E ] - fo - {Tr E [Rxxhxx332]
The forward model first one-hot encodes both the index and type channels of the input æ', yielding and h x w X Then, each of the channels in the one-hot encoded cell
two grids of shape h x index grid is seen as a separate node, and the one-hot encoded type channel tensor concatenated to each node along the channel axis. Note that we do not assume permutation equivariance of the cell types, but only of the cell indices. This means we now have grids of shape h x w x (T + 1), (T + 1). Subsequently, a linear convolution
which we simply store as grid of shape X h x embeds the grid to an embedding dimension of 32, resulting in a tensor of shape e x w X 32. with single message passing layer processes this tensor, again resulting
Then, a apatialConv-- in grid of shape x h x w x 32.
Conditional prior: (h' E Rcxhxux32] + p=(=11") + f= E Rex61]
The conditional prior independently maps each cell latent variable by applying a convolutional layer with ReLU activation, followed by three repetitions of {convolution - ReLU - MaxPool]. All convolution layers have kernel size of 9 9, and we perform meanpooling is performed over the spatial dimensions. Subsequently, a linear layer followed by ReLU activation applied. The resulting vector mapped by a linear layer to the mean of the conditional prior distribution, while a linear layer followed by Softplus activation maps the same vector to the standard deviation of the conditional prior.
maxpooling. Then, global
Approximate posterior:
(h' E IR Rexxxxxx32 x4
++1 E zhxwx2) + 9o(2[h', + (2 R
The architecture of the approximate posterior is almost identical to the conditional prior. The only exception is that first, z++1 is encoded to nodes the same way as described in the forward model, resulting in binary tensor r1+1 E {0, 1] (=+1x For each cell, this tensor is concatenated to the input along the channel axis, and is then processed by a model with the same architecture as the conditional prior.
Decoder:
(h' E Rexhxxxx22 22 E Rex64) - (n fatt1 E zhxwx2]
For each cell, the sampled latent variable 2 concatenated along the channel axis for all pixels in the grid. The decoder then consists of a SpatialConv-GNT with a single message passing layer, followed by a linear convolutional layer with kernel size 1 x that maps to a single channel. The resulting tensor has shape x h x x 1. We then apply Softmax activation along the cell axis to get the pixel-wise probabilities over cell indices. We only optimize the likelihood for the cell index, not for the cell type, as the cell type can be determined by the cell index. Specifically, to get
18

---

(discretized) output where the cell type is included, we simply take the cell index with the highest probability and match it with the corresponding cell type, which we can extract from the known input x²
SpatialConv-OMW:
Recall the SpatialConv-ooe message passing layer from Equation
vbl
h
o' (11)) JEN(
The core components that define the (i), and v.. We describe the details of each of the components below:
are 6, A the selection of the neighbors
First, à performs strided convolution with kernel size 2 x 2 and stride of to downsample each grid hy. Then, three repetitions of {convolution + ReLU) are applied, where the convolutions have a kernel size of
We use elementwise mean permutation inniaattationnnvaiiant aggregation function.
(i): We opted for fully-connected graph, meaning Nii) C' for all E C. The reasons for this are twofold: on the one hand, the convolutional nature of á and vi already takes spatial locality into account. On the other hand, using all cells as neighbors for single cell
i, including itself, means that we only have to compute
+ o' (h)) once and can re-use
JEN(()
it for all i E C, substantially reducing memory requirements.
vi: vo first upsamples the aggregated messages using a transposed convolution to the original grid size. Then, the result concatenated along the channel axis of h4 Subsequently we apply simple UNet, using three downsampling and upsampling blocks. The convolutional layers in these blocks have kernel size of 5. We also tried to parameterize with only ordinary convolutional layers, but found that the UNet's s ability to capture dynamics at different scales substantially helped performance. However, the aliasing effects introduced from the downsampling and upsampling might result in translation equivariance not actually being achieved. Nevertheless, permutation equivariance is the most crucial symmetry to respect in this setting. Furthermore, the architecture of di remains convolutional, and consequently it still enjoys the benefits inductive biases related to locality and approximate translation equivariance.
Training: For each epoch, we select one random starting point uniformly at random from each trajectory in the training set for the multi-step training. We train the model for 180 epochs with the Adam optimizer [2I] with learning rate equal to 10-4, a weight decay of 10-4, 81 = .3 = 0.9, and = 10-6, For the KL annealing schedule, we increase the coefficient 3 that weighs the KL term of the loss by 0.04 every epoch, until a maximum of 3 = 1. Furthermore, we use the free bits modification of the ELBO to prevent posterior collapse (22
ELBO free bits
=1
E Bootsle w0t11) (p 12,x0)) 3 maximum (A, KL [go(2z |r x x++1) allpllloooo [x)]),
x²)]
with À = 0.1175. In the above objective function, the KL terms for each 2j are summed over the cells and averaged over a minibatch. We use a batch size of 8.
D Baseline model details
D.1
Celestial dynamics
PNS The PNS model for celestial dynamics identical to the EPNS model described Ap- pendix
C.I with few exceptions. Instead of an equivariant FA-GNN architecture for the forward
19

---

model and decoder, we use an ordinary message passing GNN. Consequently, we do not distinguish between E(n)-equivariant and invariant features anymore, but simply concatenate all node features to one node feature vector as input. The training procedure is identical to the EPNS model. Similar to the EPNS model, this model has approximately 1.8 million trainable parameters.
terms of hyperparameter optimization, consider combinations of: {100, 500, 2000} inducing points for the sparse Gaussian process; whether or not to use an initial value encoder; and Softplus versus ReLU activations. We selected the model with the best validation loss, which had 100 inducing points, an initial value encoder, and ReLU activations.
iGPODE
For the iGPODE baseline, we used the code from the original paper [5
NSDE We train the NSDE as a latent SDE, like 27 but without the adjoint sensitivity method. The model consists six networks: context encoder, f-net, h-net, g-net, initial encoder, and decoder. The context encoder consists of a bidirectional GRU, and encodes the data into a context vector at each time step to be used by the drift of the encoding distribution. The f-net parameterizes the drift of the encoding distribution and the h-net that of the prior distribution. The g-net parameterizes the diffusion of both the encoding and prior distribution. For this we tried both the 'diagonal' and 'general' noise options in TorchSDE (20, 271, but settled on 'diagonal" noise as it performed significantly better. Contrary to [27]. the initial encoder does not depend on the context vector at the initial time step because we did not want it to suffer from the GRU performing poorly on sequences of one time step in the sampling experiments. Instead, it uses a fully connected network with skip connections to process the data at the first time step. Because the data comes from a Markov process and the full state of the system is available to the network, this should suffice to parameterize the initial distribution. The prior against which this is measured is learned mixture of multivariate Gaussians to allow the model a large degree of flexibility for the initial distribution. This prior plays no role in the evaluation of the model, however.
For the decoding distribution, we experimented with both an autoregressive distribution and with a non- aon-autoregresivee distribution. The non-autoregressive distribution worked significantly better than the autoregressive one. In both cases, we use a fully connected neural network with skip connections to parameterize Gaussian distribution with diagonal covariance. We also experimented with parameterizing the covariance and keeping fixed. The fixed value for the variance worked better.
The initial encoder and prior have their own, higher, learning rate separately from the parameters of the other networks. Both sets of parameters were optimized using Adam [21]. During hyperparameter optimization we found that a learning rate of 0.001 for the initial part of the model and one of 0.0005 for the other networks, combined with exponential learning rate scheduling with 7 = 0.9999 for both sets of parameters resulted in consistent and stable training. Moreover, we used weight decay with value 0.0001. The latent size and context size we settled on are 64 and 256 respectively. The hidden size of the f- g-, and h-net are 128 (with two hidden layers), and that of the decoding distribution is 140 with three hidden layers. Finally, we restricted the length of the sequences on which we trained to 20 and used batches of size 256.
D.2 Cellular Dynamics
PNS As non-equivariant baseline, PNS models the cellular dynamics using a modern UNet back- bone, which has shown state-of state-of-the-ar performance in modeling spatiotemporal dynamics grid [14]. Notably, the PNS model for cellular dynamics has about an order of magnitude more trainable parameters than its EPNS counterpart: around 127 million parameters, opposed to 13 million for EPNS. We use ReI.U activations and a hidden dimension of 128 throughout the model.
First, the input is one-hot encoded the same way done for EPNS (see Appendix C.2 Then, it is processed by the model components which are defined as follows:
Forward model: modern UNet architecture with wide ResNet blocks, spatial attention, downsampling layers, and group normalization, as described in [4. The architecture consists of three downsampling and upsampling steps, with attention only being applied in
? Available http://gihabbboom//ooosseer////p at the time of writing.
20

---

the lowest layer of the UNet. The number of channels is 128 at the top of the UNet, and doubled after each spatial downsampling step to maximum of 512 channels.
Conditional prior: the conditional prior architecture is identical to the one used in EPNS, except that the number of channels used in the convolutional layers is 128 as opposed to 32.
Approximate posterior: The approximate posterior architecture is the same as the condi- tional prior, with the exception that it accepts the one-hot encoding of x++1 as additional input channels.
Decoder: the decoder consists of 4 convolutional layers with kernel size x 3, followed by Softmax activation at the end to map pixel-wise probabilities over the cell indices.
Training procedure: as for EPNS, we train PNS with multi- training for 180 epochs, with the same KI. annealing schedule. We use a free bits parameter of 0.1 to prevent posterior collapse. For optimization, we use Adam with learning rate of 10-2 and weight decay of 10-4
VDE Our ODE implementation is adapted from the PyTorch implementation of [58] Specifically, we changed the expected shapes of the convolutional layers to match with the shape of the cellular dynamics data, and changed the decoding distribution to a categorical distribution instead of a Bernoulli distribution. The hyperparameters we search over are: (4. 5} convolutional layers in the encoder and decoder; {16, 32} filters in the initial encoder layer and last decoder layer (which are doubled in each subsequent convolutional layer); and {32, 64} dimensions for the latent space. We disregarded model configurations that did not fit within the 40GB of GPU VRAM with batch size to keep the computational requirements reasonable. The best best-performing model in terms of validation loss has a 64 dimensional latent space, 32 filters in the initial convolutional layer, and convolutional layers in the encoder and decoder, which has around 21 million trainable parameters.
E
Additional results
E.1
Ablation studies
We investigate the following ablations:
Changing the design choice of using permutation invariant or permutation equivariant latent variables, giving insight into the importance of this decision;
A deterministic variant of the EPNS framework, which does not have a latent space but autoregressively produces point predictions. This gives insight into the need of a probabilistic model in the stochastic setting;
An ablation of the EPNS framework which does not have latent space but rather produces po (1*++(l*") directly, giving insight into the need of the latent variable approach. As opposed to the main model, where sample 2 but construct = such that (= |2, h²) is maximized as commonly done in VAEs, here we sample from the decoder distribution p(="+1[1") which does not depend 2;
A non-equivariant PNS model with MLP backbone architecture, investigating the need for permutation equivariance in n-body dynamics.
Since the multi-step training heuristic depends on steering the trajectory using the latent space, the deterministic and no-latent ablations were trained with one-step training. All other parameters and methods were kept constant.
Overall, Table 6 shows that the main EPNS configuration with equivariant latent variables either outperforms or performs competitively with its ablations, in terms of LL and DRs. More specifically. a few observations stand out:
The EPNS variant without latent variables outperforms EPNS in terms of LL for celestial dynamics. This can be explained by the following two reasons: first, EPNS-no latents trained using single step training, which purely optimizes the one-step transition probability,
Available http:///hhhhbccoom/pppptyl//bbla at the time of writing.
21

---

Table 6: Results for variants of the EPNS framework on the celestial dynamics application.
EPNS- invariant latents
EPNS- no latents
PNS- deterministic
EPNS-
PNS
EPNS
MLP
LL (.10%)
10.9=0.4
10.8+0.11
11.6=+0.1
13.2+0.1
N/A
5.9+0.3
DKs 5=
0.61=0.18 0.20=0.06
0.14
0.39=0.12
0.25=0.12 0.19+0.05
0.63#0.08 0.7+0.05
0.52+0.21
(K 1=10
0.14+0.04
0.18
0.44+0.11
Table 7: Results for variants of the EPNS framework the cellular dynamics application.
EPNS-
EPNS-
EPNS-
PNS
EPNS
invariant latents
no latents
deterministic
1 LL (-104)
-16.4+0.3
-5.9+0.1
15,
-9.9+0.0
N/A
4 D (#cluste)) 1=30
0.70+0.10
0.58+0.09
0.60:1.22
1.00+0.00 1.00+0.00
0.93+0.08
4 (xs(Hclusters)) t=
0.77+0.05
0.58+0.05
0.55+0.25
0.97=0.03
Since the log-likelihood is calculated as the sum of one-step transition probabilities for an autoregressive model, this is perfectly aligned with the LL metric. On the other hand, EPNS is trained with multi-step training, which enables to produce better rollouts - see also the DKs values in Table but at the cost of not directly optimizing for LL. The second reason is that the step size of At = 0.1 is relatively small. This means that the assumption of a Gaussian one-step distribution made by EPNS-no latents is reasonable.
In terms of Dks. EPNS with global, cermuutationinvariin latent variables performs compet- itively with EPNS when taking the average over training runs, as seen in Table 7 However, the variance is much higher. Consequently, training EPNS with invariant latents is much more unstable than EPNS with equivariant latents. Moreover, the LL metric is always better for EPNS with equivariant latents in this setting.
Further, we also investigate the rollout stability of the various ablations see Section 5.2 for the details on how this is calculated. For the celestial dynamics problem, (alure/ssssowss that the main EPNS configuration with permutation equivariant latents tends produce simulations that remain stable for slightly longer than its counterpart with invariant latents. Moreover, both the permutation equivariant and permutation invariant alternatives of EPNS remain stable for substantially longer than variants without any latent variables, or a deterministic variant. Further, Figure 7b shows that PNS with MLP backbone has poor rollout stability, which is in line with its poor performance across the other metrics.
For the cellular dynamics problem, shown in Figure & the deterministic variant of EPNS appears to produce more stable simulations at first glance. However, upon further inspection, this is only the case because the deterministic model tends to produce 'frozen' dynamics, where cells move much slower speeds than expected. Specifically, the average velocities with which cells move, expressed in grid sites per time step, are 1.82 for the ground-truth dynamics, 1.35 for EPNS, and 0.13 for the deterministic ablation of EPNS, As such, the high stability values for CPNN-deterministie shown in Figure 8 can be considered an artifact of the stability metric, rather than the result of truly stable and realistic simulations. In fact, a model that simply propagates the initial state x. would retain 100% stability indefinitely. Consequently, EPNS with permutation equivariant latents produces the most stable simulations where cells move at reasonable speeds, albeit still slower than in the ground-trus -tuth..
E.2 Phase space coverage for celestial dynamics
We also investigated the phase space that explored by the ground truth and EPNS, PNS and PNS with an MLP backbone respectively. To this end, we generate histograms of the values of the positions and velocities taken by the object with index 1, both for the ground truth and the model. Note that these histograms show empirical densities that have been aggregated over time and over 100 runs all starting from the same initial condition, with different random seeds. The results are shown in Figure 9 (EPNS), Figure 10 (PNS) and Figure II (PNS-MLP). Visually, the phase space explored by
22

---

Fraction of stable runs
1.00
1.00
0.75
Fraction of stuble
0.75
PNS-
0.50
EPNS
MLP
0.25
EPNS-
SPNS- latent
0.25
deterministic
0.00
200
400
600
800
1000
0
200
400
600
800
1000
Time step
Time step
(a) EPNS and ablations.
(b) EPNS and INS-MLE
Figure 7: Fraction of simulations that remain stable long rollouts for celestial dynamics.
Fraction of stable runs
1.0
EPNS
EPNS-
deterninistic
0.5
EPNS-
no lateot
IPNS. invuriant latent
Time step
Figure Fraction of simulations that remain stable over long rollouts for cellular dynamics.
both EPNS and PNS appears overlap well with the ground truth, and their performance appears similar. contrast, the coordinates and velocities explored by PNS-MLP do not overlap well with the ground truth. This suggests that the permutation symmetry. respected by both PNS and EPNS, is vital in designing a good model for these systems, even when they consist of only five bodies.
Density
0.2
Density
1.0
Density
0.2
0.5
0.0
0.0
0.0
5.0
Coordinate value Ur
Coordinate value dy
Coordinate value
U/s
Density
Density
1.0
Density
0.2
0.1
0.0
0.0
0.0
Coordinate value
Coordinate value
Coordinate value
Figure 9: Phase-space coverage plots of the n-body system for EPNS. The plots show histograms of phase-space coordinates of the body with index aggregated over time and 100 samples starting from the same initial condition. The solid blue line indicates the histograms obtained from the ground-truth simulator, while the dashed orange line indicates the histograms obtained from the model.
23

---

1.0
Density
Density
Density
0.6
0.1
-5.0
-2.5
0.0
2.5
5.0
-5.0
-2.5
0.0
2.5
Coordinate value
Coordinate value
Coordinate value
0.2
Density
Density
Density
0.2
0.5
-5
55
Coordinate value
Coortinate value
Coordinate value
Figure 10: Phase-space coverage plots of the n-body system similar to Figure but with PNS as model.
Density
Density
Density
0.1
-10
-5
-2
-5.0
225
0.0
2.5
Coordinate value
Coordinate value
Coordinate value
Density
Density
Density
0.2
0.5
-5
10
-5
Coordinate value
Coordinate value
Coordinate value
Figure 11: Phase-space coverage plots of the n-body system similar to Figure 9 but with PNS with MLP backbone as model.
24

---

E.3 20-body system results
In this section, we report results of applying EPNS 20-body system, as opposed 5-body system in Section Although this system is not as high- sigh-dimessiona as most challenging potential real-life applications. the 20-body system can help assess the potential of scaling EPNS such systems. Notably, for these experiments, we do not perform any hyperparameter optimization relative to the 5-hody EPNS model.
Table shows the quantitative results of applying EPNS and PNS to a 20 body system. Similar to the 5-body system, we see that EPNS outperforms PNS on all metrics. Note that LL is calculated by summing over all dimensions, and since we have 20 bodies here, LL is higher than the 5-body model. If we were calculate the per-body LL instead, it would be lower, as expected for a more challenging problem setup without any specific hyperparameter optimization. Surprisingly, the Dks scores EPNS are similar those of the 5-body problem, indicating that EPNS has the potential to generalize to larger systems.
Further, Figure 12 shows the distribution of the kinetic energy time, both of the ground truth and of EPNS and PNS, calculated over 100 simulations starting from the same initial conditions, similar to Figure 5 Again, observe that the distribution generated by EPNS overlaps very well with the ground-truth, and notably better than the distribution generated by PNS.
we plot the empirical density of the phase space coordinates of the body with index 1, aggregated over time and over 100 samples starting from the same initial condition but with different random seeds. Again, the density produced by EPNS is plotted in dashed orange lines, while the ground truth is plotted in solid blue. Similar to Figure [ both densities overlap very well, suggesting that the distribution over trajectories from fixed initial condition produced by EPNS matches the ground truth, also in this higher-dimensional setting.
Finally, similar to Section
we observe that
Table 8: Results for scaling EPNS to 20-body system.
PNS
EPNS
t LL (-103)
15.4+001 0.22+0003
15.6+0.0
1 D (KE( 1=30
0.14 0.21
1 (KE) 145
0.14
groud
ground
Kinetic energy
Kinetic energy
400
400
200
200
20
60
100
40
100
Time step
Time step
(a) EPNS.
(b) PNS.
Figure 12: Ground-truth distribution over kinetic energy for a 20-body system compared to EPNS and PNS.
E.4
Additional samples for celestial dynamics
We provide additional qualitative results for celestial dynamics in Figure 14 To do so, we select four samples from the test set uniformly random. Overall, we observe that EPNS, iGPODE and PNS generally produce simulations that look visually plausible. NSDE does not generate plausible new simulations. We attribute this primarily to the relatively high KL divergence component (around 1900) of the NSDE model BLOO values (shown in Table T). Recall that, following the same procedure as for the other baselines, we selected the NSDE hyperparameters that resulted in the best validation ELBO value, which might not have been optimal in terms of new sample quality. Notably, all models
25

---

Density
0.1
Density
0.2
Density
0.2
0.1
0.0
0.0
10
10
10
Coordinate value
Coordinate value
Coordinate value
Density
Density
Density
6.1
0.1
-10
5
10
-6
10
Coordinate value
Coordinate value
Coordinate value
Figure 13: Phase-space coverage plots of the n-body system similar to Figure but with EPNS trained on a celestial dynamics with 20 bodies.
can struggle with handling close encounters, and body may get 'stuck' another body, for example shown by the red body in the bottom right trajectory. Generally, EPNS seemed to struggle less with handling such close encounters, although the top left trajectory suggests that such behavior is not completely eliminated from the model.
E.5 Additional samples for cellular dynamics
Additional qualitative results for cellular dynamics shown in Figure 15 four samples from the test set uniformly at random. EPNS produces samples that look realistic, as the cells generally have reasonable sizes, show dynamic movement and cell shapes, and tend to cluster together with other cells of the same type over time. Still, sometimes the volumes of some cells become unreasonably small after many rollout steps, which is the main cause for its stability decreasing, as shown in Figure 6b In the case of PNS, although it does exhibit clustering behavior, cells often quickly evolve have unrealistic volumes, also reflected in the rapidly decaying stability shown in Figure 6b ODE generally produces simulations in which cells look 'frozen' and lack clustering behavior. We attribute this the fact that ODE models the dynamics in latent space, and struggles to propagate the geometric structure of the data as time proceeds result.
Again, selected
26

---

Ground
trutth
Ground
truth
EPNS
EPNS
IGPODE
iIRODE
NSDE
NSDE
PNS
PNS
PNS MLP
PNS MLP
20
40
60
20
f
40
60
Ground truth
Ground truth
EPNS
EPNS
IGPODE
IGPODE
NSDE
NSDE
PNV
PNS
PNS MLP
MUP
(.9%
PNS
t= 0
t= 20
t= 40
t
t=
t 20
t = 40
60
Figure 14: Additional qualitative results for celestial dynamics simulation.
27

---

Ground truth
Ground truth
EPNS
EPNS
ODE VAE
ODE
PNS
PNS
t
5
10
15
t 10
15
Ground truth
Ground truth
EPNS
EPNS
ODE*VAE
VAE
PNS
PNS
t
t 5
t=
1=
5
f
10
15
Figure 15: Additional qualitative results for cellular dynamics simulation.
28


Reviews:
[{"review_id": "9LPvUZiNio", "reviewer": "Reviewer_AV4g", "summary": "The paper presents EPNS, a framework for equivariance in autoregressive probabilistic multidimensional dynamics modeling. The framework is general such that one can use existing equivariant NN architectures; for the task of cellular dynamics forecasting the authors also introduce a ConvNet GNN. The models are shown to outperform current SOTA in both accuracy and data efficiency.", "questions": "1. The paper mentions no work has considered equivariant probabilistic modeling of spatiotemporal dynamics. I can think of the following papers that tackles similar settings as in this paper:\n\n- Köhler, Jonas, Leon Klein, and Frank Noé. \"Equivariant flows: sampling configurations for multi-body systems with symmetric energies.\" arXiv preprint arXiv:1910.00753 (2019).\n- Garcia Satorras, Victor, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. \"E (n) equivariant normalizing flows.\" Advances in Neural Information Processing Systems 34 (2021): 4181-4192.\n- Sun, Sophia Huiwen, Robin Walters, Jinxi Li, and Rose Yu. \"Probabilistic Symmetry for Multi-Agent Dynamics.\" In Learning for Dynamics and Control Conference, pp. 1231-1244. PMLR, 2023.\n\nI believe your framework is more general w.r.t. symmetry groups, but it would be good to cite these works as well to provide context.\n\n2. How are the p-values in table 2 calculated, and how should we interpret them? \n\n3. Some minor questions on writing:\n-  Is it a typo on line 143? equivariant with regards to $h^t$?\n- Line 277: \"meaning that new simulations also match better to the ground truth in terms of the examined properties\" <- what does 'new simulation' and 'examined properties' refer to here?\n\n\n", "limitations": "The author provides a detailed discussion of limitations in the conclusion section of the paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The framework and algorithm proposed is sound and general, which I can see being widely used in the spatiotemporal forecasting community. The authors did a great job outlining and explaining the EPNS algorithm step by step. The illustrations are great and helpful as well. The method even addressed the long rollout deviation problem common to autoregressive models. Overall, the algorithm is convincingly effective and practical. \n", "weaknesses": "- It was a little misleading when the authors say that they theoretically proved that EPNS produces equivariant distributions over trajectories. From my understanding, the equivariance comes from the paper's formulation by construction. Maybe a better way to word the contribution is something like \"Our framework is able to construct equivariant distributions over the entire trajectory autoregressively from equivariant single-step predictions\"?\n \n- There is no substantial contributions on model architecture. I personally found the framework and training setup to be insightful enough to make an impact on the forecasting community, but one could argue it's just a new way to use existing models on new datasets. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review and specific comments and pointers. We are glad that you found our work to be sound and general, convincingly effective and practical, and well-explained. Please find our responses to your specific comments below:\n\n**W1 - wording of contribution:** thank you for pointing out that this formulation can be confusing. We will change the formulation to ‘We demonstrate that EPNS produces single-step distributions that are equivariant to transformations of the relevant symmetry group, and that it employs these distributions to autoregressively construct equivariant distributions over the entire trajectory’.\n\n**W2 - contributions on model architecture:** While the primary focus of this paper is the general framework for creating and learning equivariant distributions over trajectories, we also introduced the novel SpatialConv-GNN architecture. This new GNN architecture was necessary in order to respect both the shift and permutation symmetries of the cellular dynamics problem, and is described in Section 4.2 and Appendix B.2. In our opinion, this constitutes a substantial architectural innovation, since it exemplifies designing an EPNS model tailored to a combination of data symmetries - in this case permutational and shift symmetry on a grid domain - that is previously unexplored. In other words, the value does not only come from being a new layer, but also from demonstrating how application-specific architectures can be designed that greatly outperform generic baselines using the EPNS framework.  \n\n**Q1 - related literature:** Thank you for pointing us to these relevant works. While we briefly mentioned the first two papers in Section 2.1 (due to space constraints), we will extend the discussion on them in the updated version of the paper. Both of these works also use equivariant models in the parameterization of probability distributions. A crucial difference is that their time parameter is internal to the model and used to define a continuous flow from a base distribution to a target distribution, whereas we explicitly consider time as a variable in the data. Furthermore, they only consider modeling E(n)-invariant distributions, as you already mentioned.\nThe third paper is more closely related to our work, and it was published in PMLR after the NeurIPS submission deadline. Nevertheless, we will discuss it in our updated related work section. The major differences are twofold: \n1. generality of symmetries: the method in the third paper considers only planar rotations (SO(2)), while we consider general linear symmetry groups. Concretely, this means that their method is not applicable to e.g. the cellular dynamics problem or the three-dimensional n-body problem of our experimental evaluation; \n2. model expressivity: the one-step distribution of their method is restricted to being a multivariate normal distribution. In contrast, by introducing a latent variable into the model architecture, our method can model more general one-step distributions. This should be especially relevant in the case the one-step distribution can no longer be assumed to be unimodal, for example due to coarse temporal resolution and the presence of bifurcations in the dynamics underlying the data. \n\n**Q2 - p-values:** the p-values in table 2 are calculated using the two-sample Kolmogorov-Smirnov (KS) test, a statistical test of equality of two probability distributions. The null hypothesis of this test is that the two sample sets under consideration are drawn from the same probability distribution. To test the null hypothesis, it compares the empirical cumulative distribution functions (CDFs) of the two sample sets under consideration, finds the maximum distance between them, and calculates a p-value based on this statistic. This p-value basically expresses the probability of finding a distance that is at least as large as the actual measured distance, under the assumption that the null hypothesis is true. As such, if the p-value is small (say, less than 0.05), it is very unlikely that we encounter this distance (or an even larger one) under the null hypothesis, and thus we reject it. For EPNS, the large p-values in Table 2 of our paper indicate that the KS test cannot distinguish between the distributions from which the two sets are sampled, meaning that the null hypothesis that the underlying distributions are statistically equivalent cannot be refuted. In contrast, for PNS, the p-values are very small, meaning that according to the KS test it is very unlikely that the underlying distributions are identical, providing strong evidence that the equivariance property is violated.\n\n**Q3 - typo:** Thank you for spotting this typo. Indeed that should have been “equivariant with regards to”.\n\n**Q3 - line 277:**  We now see that this formulation can be a bit confusing. With ‘new simulations’, we mean sampled trajectories from the ML model starting from the given initial condition x^0. This is in contrast to the log-likelihood, which is measured on ground-truth trajectories from the test set, i.e. no sampling from the model is involved.  With examined properties, we mean the properties of the system mentioned in the caption of Table 1, i.e. kinetic energy for the n-body system and the number of cell clusters of the same type for cellular dynamics, which are aggregate properties that are relevant to these specific domains.\n"}, {"review_id": "qsHfUPQ53W", "reviewer": "Reviewer_QKcR", "summary": "This work presents a new framework for autoregressive probabilistic modeling called Equivariant Probabilistic Neural Simulation (EPNS) that produces equivariant distributions over system evolutions. The authors demonstrate that this method outperforms existing neural network-based methods for probabilistic simulation, improving simulation quality, data efficiency, rollout stability, and uncertainty quantification. The paper also discusses the benefits of incorporating domain symmetries in deterministic neural simulators and provides specific examples of domains where EPNS could be applied for efficient and effective data-driven probabilistic simulation.", "questions": "There are several concerns regarding the work as outlined below.\n1. There are several equivariant architectures including the EGNN authors have mentioned. The choice of the specific GNN which is not truly state-of-the-art for dynamics simulation is not clear. Authors mentioned that EGNN is not stable while training. EGNN has one of the simplest architectures among the equivariant GNNs and it is not clear why such an architecture would be unstable. Authors should clarify this.\n2. The reason for the choice of a graph for 3-body system is not clear. Why not an MLP for such a simple system?\n3. Can the graph architecture be generalized to larger systems, say 30 or 300 particles? It is suggested to demonstrate this ability. Otherwise, the approach remains applicable only to toy systems.\n4. What are the characteristics of the learned dynamics? Can the drift and diffusion term of the SDE be decoupled? If yes, how well is the variance of the drift term captured by the simulator? If not, how can the ability of the model to realistically capture the simulation be established.\n5. Authors discuss about the stability of the simulation. Indeed, autoregressive models are known to be unstable. Can the authors discuss about the stability of the learned simulation with respect to the length of the training trajectory or inference trajectory? That is, can they say that the inference trajectory will remain stable for the same length as the training trajectory or lower or more $x$ times more. This would be helpful to establish the usefulness of the ML-model\n6. What about the data-efficiency of the approach? Does it continue to become better with increasing data. If yes, by how much? Have authors investigated the data efficiency of the model?\n7. Finally, the model is not interpretable and not generalizable to larger system sizes. For instance, the model cannot give any insights about the physics or dynamics of the system in terms of the drift or the diffusion term. This means it can simulate the same system on which it is trained. This would be useful only if the model is computationally much more economical than the ground truth. Is that the case? Authors have mentioned that the GNN is computationally expensive. In that case, where do the authors think the approach would be useful practically?", "limitations": "While some of the limitations are briefly mentioned, there are a large number of limitations for the present work as discussed in the weaknesses and questions section. Authors should either address those or include those in the manuscript as limitations.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Authors attempt to solve a challenging problem of simulating systems following a probabilistic simulation. They present an autoregressive approach toward probabilistic simulations combining equivariance. This can be used for several problems such as Langevin dynamics, Brownian motion to name a few.", "weaknesses": "There are several weaknesses for the work as outlined below.\n\n1. **Simple systems:** Authors have considered very simple toy systems such as 3-body gravitational system, and a simplified cellular dynamics, a lattice based system. There are several examples of realistic probabilistic simulations such as Langevin or Brownian dynamics; these are not considered. This limits the application to toy systems and validity of such approach for larger realistic systems are questionable.\n2. **Generalizability:** The work is demonstrated on a 3-body gravitational system. The GNN nature of the approach allows natural zero-generalizability to larger system sizes. These are not explored making the application of the approach to real-world problems limited. Similarly, due to the data-driven approach, it is not clear whether the approach can be generalized to unseen noise conditions for the same system. That is, for the gravitational system, if the interactions with the dust particles increased or for the cellular the systems, if the temperature increased, will the approach be applicable is not clear.\n3. **Metrics:** This approach is presented as a data-driven approach to learn the dynamics. However, the ability of the model to preserve the physical quantities meaningfully is not explored. For instance, the forces on each particle in the rollout, how do they compare with those in the ground truth in a statistical fashion? Similarly, the phase space explored by the rollout, do they remain same ensuring statistical similarity of the trajectory. Indeed, the authors have shown the $LL$ and $D_{KS}$, but more elaborate error metrics can provide insights into the nature of learned simulations.\n4. **Baselines:** It is not clear why a graph architecture is chosen especially when there is no generalizability to larger system sizes explored. To this extent, an additional MLP-based baseline can give insights into whether a GNN architecture provides superior results or not.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and the concrete comments and suggestions. Please find below our response to your points:\n\n**W1:** To avoid misunderstandings, let us clarify that the n-body experiments were conducted on a 5-body problem, rather than a 3-body problem. This means the data comes from a 30-dimensional SDE - it has three spatial coordinates and three velocity components for each of the five bodies. Moreover, similar 5-body system dynamics prediction problems are commonly used for evaluating new machine learning methods, see [1, 2, 3] for recent examples. \n\nFor cellular dynamics, the data forms a system of 80x80 lattice sites, for each of which we need to produce a 64-dimensional probability vector for the cell index, as there are 64 cells which serve as nodes in the EPNS model. So, this system is extremely high-dimensional, and its complexity is evident from the poor performance of baseline models. \n\n**W2:** Let us first address generalizability to larger systems, also in response to Q3 and Q7. For cellular dynamics the system is already large, as there are n=64 cells on a 80x80 grid. For the n-body dynamics, we have now run additional experiments with n=20 bodies to demonstrate EPNS’ capability to generalize to larger system sizes – the ground-truth simulator struggled to remain stable for larger n. Results can be found in Table R3 and Figures 2 and 3a in the general response. The key observations are that performance for 20-body dynamics case is similar to the 5-body case, as evidenced by good match with the ground truth in the coordinate exploration plots (Figure 2, compare bottom and top rows), and the LL and D_KS scores (Table R3), which together with Figure 3a indicate the large overlap between EPNS’ trajectory distribution and the ground truth. In terms of generalizing to e.g. different noise conditions and temperatures, EPNS could easily be extended to take such variables as a conditioning signal, similar to e.g. [4], but this is left out of scope for this paper.\n\n**W3:** For the 5-body system (resp. for the cellular dynamics), we have investigated the behavior of physical quantities such as kinetic energy and potential energy (resp. the number of cell clusters and cell volume), as shown in Figures 5, 6 and Table 1 of our original paper and the accompanying text. However, the quantities you suggested are interesting to investigate as well. In terms of forces, due to the specific parameterization of our model (Appendix B.1), the forces are not directly modeled, so this is not possible. For the phase space exploration, we have added plots showing the distribution over the spatial coordinates for a body in Fig. 2 of the general response PDF – we will add more plots in the appendix of the updated version of the paper where we have more space. The probability distribution over coordinate values between the ground truth and EPNS overlaps very well, both for 5-body and 20-body systems (top and bottom rows of Fig. 2). \n\n**W4:** This work is about the incorporation of equivariances into stochastic simulation models. The 5-body system has a permutation symmetry - it does not matter for the physics in which order we label the planets - and a graph neural network is equivariant to precisely this symmetry, whereas an MLP is not. To stress the importance of the permutation symmetry, we have now run experiments with an MLP baseline. As can be seen in Figures 1b, 2, and Table R1 in the general response, this performs much worse than EPNS. \n\n**Q1:** The FA-GNN is SOTA for n-body dynamics prediction, see Section 5.3 of [3] for a comparison to EGNN. The instability of EGNN occurred in combination with the multi-step training heuristic, whereas the FA-GNN suffered significantly less from this problem. So, the choice for FA-GNN in this specific EPNS model is motivated both from a performance and practical perspective.\n\n**Q2:** See W4.\n\n**Q3:** See W2.\n\n**Q4:** Our method is not restricted to modeling SDEs - see for example the cellular dynamics application - and as such we do not model a separate drift and diffusion term, since those are only applicable to SDEs. The ability to realistically capture the dynamics can be established by looking at domain-specific quantities of interest such as kinetic and potential energy for celestial dynamics, the number of clusters for cellular dynamics, the phase space distributions you suggested, and common metrics like log-likelihood.\n\n**Q5:** The maximum rollout length for EPNS during training was 16 steps for the 5-body problem and 14 steps for the cellular dynamics. As can be seen from Figure 6 in the paper, EPNS generally remains stable for much longer than the rollouts seen during training.\n\n**Q6:** We have investigated data efficiency, see Table 3 in the paper and the accompanying text.\n\n**Q7:** See W2 and Q4 for our comments on larger systems and interpretability. Further, our primary goal was to demonstrate the potential of our method for effectively learning stochastic dynamics from data. Nevertheless, the n-body EPNS model is very economical in terms of simulation time: around 20 times faster than the ground truth simulator. Further, as can be seen in Table 3 of our paper, our method is data efficient, meaning that with few samples, a model can be trained that can simulate from many initial conditions. \nMore importantly, this work constitutes a step in the direction of data-driven simulations of complex stochastic systems, which can be useful in many real-world settings where the true laws governing the dynamical system are unknown.\n\n[1] V. G. Satorras et al. “E(n) equivariant graph neural networks”, ICML 2021\n\n[2] C. Yildiz et al. “Learning interacting dynamical systems with latent gaussian process ODEs”, NeurIPS 2022\n\n[3] O. Puny et al. “Frame averaging for invariant and equivariant network design”, ICLR 2022\n\n[4] J. Brandstetter et al. “Message passing PDE solvers”, ICLR 2022\n"}, {"review_id": "4DPNEeYpu7", "reviewer": "Reviewer_n81N", "summary": "This paper introduces a general framework named Equivariant Probabilistic Neural Simulation (EPNS) for equivariant stochastic simulation with neural simulators. The framework incorporates domain equivariances and utilizes an autoregressive probabilistic approach. The authors prove that, under equivariance assumptions on the neural network layers, the probability distribution over trajectories remains equivariant. The proposed framework is applied to two scenarios: a stochastic n-body system with celestial dynamics and stochastic cellular dynamics on a grid domain. Additionally, the authors propose a novel message passing graph neural architecture that maintains equivariance to permutations of cell indices, motivated by stochastic cellular dynamics problem. Experiments showcase multiple facets of the benefits of this framework.\n", "questions": "Q1.  While I understand how equivariance can help in maintaining cell clusters and representing the true number of clusters, I am unclear about its connection to kinetic energy in celestial dynamics. Could the authors provide further insight into how equivariance relates to kinetic energy in this context?\n\nQ2. The authors did not cite any papers when selecting the metrics presented. Could the authors explain the rationale behind not basing their metric selection on previous work in neural simulation?\n\nQ3. How can EPNS be extended to handle non-Markovian dynamics? What would be the potential benefits of such an extension?\n\nQ4. Please comment to W2.", "limitations": "Limitations were mentioned. See W2 for another potential limitation. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1. The paper presents a method for learning stochastic equivariant distributions over trajectories, which constitutes a valuable contribution to neural simulation. Additionally, the authors introduce a novel message passing graph neural network that enables equivariant embeddings in domains containing lattices and sets.\n\nS2. The experiments are comprehensive and demonstrate several benefits of the framework, including improved data efficiency, uncertainty quantification, and rollout stability.\n\nS3. The paper's presentation is excellent, with a detailed appendix and code that facilitates reproduction of results. The results are comprehensive, and the insights are well-described.", "weaknesses": "W1. The authors state the benefits of learning equivariant distributions over equivariant function approximations. However, it is unclear whether any of their experiments support this claim beyond enabling uncertainty quantification.\n\nW2. The authors did not explicitly state whether the simulations used for the experiments are toy examples or scenarios that are of significant interest to domain experts. It would be beneficial for the paper to provide clarification on this aspect. Specifying the nature and relevance of the simulations would help us assess the potential impact of the developed methods on scientific applications. This information is crucial for understanding the current progress in developing methods that have practical implications in the scientific community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful review. We are happy to hear that you believe that EPNS is a valuable contribution to the field of neural simulation, that you find our experiments to be comprehensive, and that you liked the presentation of the paper. Please find our specific responses to your points below:\n\n**W1:** Initially, we took the need for a probabilistic model that is capable of sampling as a requirement to our problem setting of stochastic simulation, and aimed to demonstrate the benefits of equivariant distributions over non-equivariant counterparts. We have now run additional experiments with deterministic versions of the EPNS models on our datasets to investigate whether there are other benefits to using probabilistic models on these problem settings. The results can be found in Tables R1 and R2 and Figures 1a, 1c and 3b in the general response. Interestingly, using EPNS instead of its deterministic variant improves rollout stability for celestial dynamics; for cellular dynamics, stability for the deterministic model appears better at first glance, but this is only because the model produces ‘frozen’ dynamics, as shown in Figure 3b, which is not the case for the ground truth. To quantify this, we calculated the mean cell velocity, which is shown in Table R4. Of course, the D_KS metric, indicating calibration of model-sampled trajectories, is better in the probabilistic case. \n\n**W2:** Thank you for making us aware of this unclarity. We did not design our experiments with the aim of trying to solve open problems in two distinct domains. However, this does not mean that the experimental settings are easy from a machine learning modeling perspective, demonstrated by the clear performance difference between EPNS and the baseline methods for both scenarios. Additionally, for the n-body case, the complexity of the data (5 bodies, corresponding to a 30-dimensional SDE) is very similar to experiments in recent machine learning papers [1, 2, 3]. For the cellular dynamics case, the high dimensionality and level of complexity of the dynamics is comparable to computational biology studies that use the ground-truth simulation algorithm to investigate the migration of cancer [4, 5]. So, although these specific settings are not *directly* of specific interest in their relevant application domains, they are of representative complexity of problems studied in ML literature (for n-body dynamics) and in computational biology (for cellular dynamics simulations). \nWe will update our paper to include an explicit discussion on the scope, relevance and limitations of the experimental setup.\n\n**Q1:** The E(n) symmetries introduced into the model architecture do not directly contribute to better modeling of quantities like kinetic energy in the n-body experiment. Nevertheless, for the n-body case, we observe that EPNS performs better than the non-equivariant models in terms of the evolution of the kinetic/potential energy of the system. The explanation for this is that including symmetries enables more effective and efficient learning of the dynamics in general, which is reflected in aggregate quantities like kinetic energy.\n\n**Q2:** Log-likelihood is a common metric for evaluating probabilistic models, and has also been used in previous works (e.g. [2]) – we will make sure to add relevant citations to emphasize this. In addition to evaluating goodness-of-fit, we want to evaluate whether our simulations are physically realistic. To do so, we look at characteristics of the underlying dynamics, such as various types of energy for the n-body system and number of cell clusters in the cellular dynamics, and compare the distributions of these to the ground truth. What characteristics are relevant strongly depends on the system at hand, meaning there is no one metric of physicality applicable to all datasets or dynamics. However, the metric we use for comparing the distributions of these characteristics, the Kolmogorov-Smirnov test statistic, is well understood and widely used. Note that in practice, when running simulations, one is often particularly interested in these kinds of characteristics, so they are crucial in assessing the performance of a simulation model.  \n\n**Q3:** In order to extend EPNS to non-Markovian dynamics, we need to be able to pass information (possibly arbitrarily many) timesteps into the future. One way to do this would be to introduce a memory component into the model, for example similar to LSTMs. However, we need to make sure that this does not break the equivariance properties of EPNS. As such, the updating of this component should be done in a similar way in which the latent variables are constructed in EPNS, e.g. by having equivariant memory updates, and a decoder which is equivariant with respect to this memory cell as well. The proof of Lemma 1 would then remain similar. Other approaches could be investigated as well, e.g. using transformers to extract information from the prior sequence. Further research is needed to investigate what works best empirically. \nThe benefits of such extensions would be that this will make EPNS suitable for modeling systems where the Markov property does not hold, e.g. in case of long-range dependencies in the data-generating process or when not all relevant information of the system is observed.\n\n**Q4:** See W2.\n\n[1] V. G. Satorras et al. “E (n) equivariant graph neural networks”, ICML, 2021\n\n[2] C. Yildiz et al. “Learning interacting dynamical systems\nwith latent gaussian process ODEs”, NeurIPS, 2022\n\n[3] O. Puny et al. “Frame averaging for invariant and equivariant network design”, ICLR, 2022\n\n[4] S. Kumar et al. “Proteolytic and non-proteolytic regulation of collective cell invasion: tuning by ECM density and organization” Scientific Reports 6, 2016.\n\n[5] Y. Zhang et al. “Computational Modeling to Determine the Effect of Phenotypic Heterogeneity in Tumors on the Collective Tumor–Immune Interactions”. Bull Math Biol 85, 2023.\n"}, {"review_id": "imDBuNgbDS", "reviewer": "Reviewer_yRZx", "summary": "This paper proposed to incorporate symmetries into probabilistic models to simulate stochastic dynamics over trajectories. In particular, the authors developed EPNS, an equivariant version of sequential CDVAE for the generation of the state transition. The results on two tasks including stochastic n-body system and stochastic cellular dynamics, show that EPNS considerably outperforms existing neural network-based methods for probabilistic simulation in terms of simulation quality, data efficiency, rollout stability, and uncertainty quantification.", "questions": "See the weakness part above.", "limitations": "The authors have adequately discussed the limitations. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "\n1.\tIt is valuable to design equivariant probabilistic neural simulators against existing deterministic counterparts. It is true that random effects exist in common systems and should be characterized. \n\n2.\tLemma 1 provides convincing derivations for the formulation of equivariant CDVAE. It is interesting to see that there are two different ways to derive equivariant models: invariant latent and equivariant latent. Besides, severe deviation over long rollout simulations has always been a tricky problem. To mitigate this phenomenon, the paper adopts a heuristic method by using the ground truth of the next state to steer the sampling progress， which is inspiring.\n\n3.\tThis paper is well-organized, first introducing the general formulation and then detailing the specific architectures for two interesting applications. I enjoy the writing. The experimental evaluations, although insufficient, is appropriately demonstrated to support certain claims of the authors.\n", "weaknesses": "\n1.\tMy biggest concern is that one important reference [A] is not cited and discussed. The authors claim that “current probabilistic methods do not incorporate general domain symmetries into their model architectures”, which is incorrect if checking [A]. The work by [A] has developed equivariant Gaussian Process (GP) to model stochastic fields, including scalar, vector and stacked fields. Equivariant GP is applicable to the task of modeling system evolutions discussed in this paper. We can directly consider p(x_{t+1}|x_t) as the prediction of vector fields studied in [A]. More importantly, when referring to [A], many claims of the authors in Introduction (Lines 47-51, Lines 55-57) and related works (Line 95-97) are no longer correct, and [A] should be at least considered as an indispensable baseline in the experiments. It is unclear how much novelty remains and how many extra efforts should be made to reflect the modification of discussing [A].\n\n2.\tFor the experimental evaluations, several crucial contribution points are NOT justified sufficiently. First of all, the equivariant deterministic methods should be tested, by which we can observe why we need stochastic modeling. Second, besides the reference [A] mentioned above, the authors are suggested to compare against typical probabilistic models (such as ODEVAE) with trivial equivariant modification, which is able to better support the benefit of the proposed designed model. Third, in comparison with using the latent model z, one can directly model the transition probability p(x_{t+1}|x_t) straightforwardly via multiple equivariant layers. How does this baseline perform?  \n\n3.\tThe authors proposed to use invariant latent for Celestial dynamics, and equivariant latent for cellular dynamics. Certain ablations are suggested to provide to support this design. \n\n4.\tThe conclusion of Lemma 1 is interesting. It provides sufficient conditions to maintain the symmetries, but are these conditions also necessary?\n\n5.\tWould you please explain more on why EGNN is unstable for the Celestial dynamics task?\n\n6.\tIs it possible to combine invariant latent and equivariant latent into one single model. In this case, does it mean that invariant latent corresponds to invariant features in EGNN and equivariant latent corresponds to equivariant features in EGNN? \n\n7.\tHow to generate stochastic ground truth trajectories in the experiments? How to compute quantile interval of the observations in Figure 5. \n\n\n[A] Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes, ICML 2021.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review. We appreciate that you find equivariant probabilistic neural simulation a valuable contribution, that you find our methodology convincing and inspiring, and that you liked the presentation of the paper. Please find our responses to the your specific comments below:\n\n**W1:** Thank you for bringing this paper to our attention. We agree that [A] provides a method for equivariant modeling stochastic phenomena, and we will update our related work discussion accordingly. However, the main method described in [A], SteerCNPs, is not a natural fit to the problem setting considered in our work, as it is a method for spatial interpolation rather than temporal dynamics simulation. More specifically, as described in 5.3 of [A], it requires a context set Z to encode, which in their experiments are, for example, certain pixel values of MNIST data that have been observed, or wind speeds that have been recorded only at certain locations. In our setting, such a context set is not available. Additionally, from Section 5.1 in [A] it is clear that the SteerCNP method is exponential in n, which makes it infeasible for the problems we consider: for the 5-body system, n would be 30. In fact, [A] poses the use of point-cloud based techniques to overcome this limitation for scenarios like ours as an interesting avenue for future work.\n\nTo directly model the transition probability as a neural field distributed according to a Gaussian Processes, as you suggest, would be more akin to learning an equivariant GP prior, rather than the posterior [A] focusses on. Although [A] provides sufficient and necessary conditions for a GP prior to be invariant (theorem 1), these conditions need modification for the GP to be equivariant. Moreover, [A] does not provide methods for parameterizing the kernel and mean of a GP prior in a way that it is learnable, flexible, and satisfies the relevant conditions for in- or equivariance given a transformation group G.\n\nHowever, we see that without context the specific lines in our paper that you pointed out appear as if we claim to be the first paper to model equivariant stochastic phenomena in general. Besides updating the related work section to discuss this work, we will also update these text segments to stress that our claims concern modeling dynamical systems using equivariant distributions over trajectories with regard to general symmetry groups specifically.\n\n**W2 & W3:** Thank you for suggesting these additional experimental evaluations. We have now run additional experiments with the following ablations: a deterministic variant of EPNS, a variant that models the transition probability directly, and an ablation on using in-/equivariant latents. The main results can be found in Tables R1 and R2 and Figures 1a, 1c and 3b in the general response. For celestial dynamics, it turned out that a model with permutation equivariant latents worked even better than our original configuration with invariant latents in terms of stability and the D_KS scores. For cellular dynamics, the original configuration with equivariant latents substantially outperforms the other variants of the model. \n\nAs for comparing against modifications of probabilistic models, the architectural adaptations required to make the baseline models equivariant to the relevant transformation group are not trivial. For example, for the n-body system, introducing E(3)-equivariance in the iGPODE and Neural SDE models would require a complete overhaul of the model structure. For the cellular dynamics, introducing permutation and translation equivariance in ODE2VAE would require the integration of our SpatialConv-GNN, which is practically infeasible since both already have high memory requirements by themselves. Of course, the question of how to build equivariant architectures into various generative models is an interesting avenue of research. Our paper provides a promising step in this direction.\n\n**W4:** The conditions in this lemma are sufficient but not necessary: for example, if the decoder ignores the latent state entirely, the entire model can be G-equivariant even if the conditional prior is neither invariant nor equivariant. However, doing so considerably limits the class of distributions that can be learned.\n\n**W5:** We did preliminary experiments with EGNN for the architecture, but the multi-step training would often become unstable with the EGNN. As such, we switched to the FA-GNN for practical reasons, as instability issues seemed to occur less frequently. Furthermore, the FA-GNN has also been shown to outperform the EGNN for n-body dynamics prediction (see Section 5.3 of [B]), so the choice for the FA-GNN is also motivated from a performance perspective.\n\n**W6:** It is possible to combine both invariant and equivariant latent variables within the EPNS model. Whether doing so effectively separates invariant and equivariant features seems plausible, but further research is necessary to confirm if this is the case.\n\n**W7:** We will update the appendix with the equations and algorithms that are used for simulating the ground truth trajectories. Note that the code for generating the ground-truth dynamics can be found in the supplemental material. The quantile intervals in Figure 5 have been computed by running the ground-truth simulator and the ML model 100 times from the same initial state, and measuring the potential energy (resp. number of cell clusters) of the system at each point in time. We then calculate the (0.1, 0.9) quantiles of these values at each time point.\n\n[A] P. Holderrieth et al. “Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes”, ICML 2021.\n\n[B] O. Puny et al. “Frame Averaging for Invariant and Equivariant Network Design”, ICLR 2022\n"}]

Author rebuttal:
# General response to reviewers

We would like to thank all the reviewers for their insightful and detailed comments! We are happy to read that the reviewers generally agree that our work tackles a valuable problem, that we present an interesting and effective framework and that the paper is well-written. 

Two overarching points were mentioned multiple times, which are also addressed in more detail in the individual responses:
* The need for extra baselines/ablations to verify the utility of the stochastic modeling framework and the architectural design choices. We address this by providing new experiments with: 
    + Deterministic equivariant methods (verifying the utility of stochastic modeling); 
    + Modeling the transition probability directly without a latent variable (verifying the utility of the latent-variable approach); 
    + Using equivariant (resp. invariant) latents for the celestial dynamics (resp. cellular dynamics) rather than the setting used in the paper (ablation study of design choices); 
    + Using a non-equivariant baseline (MLP) for celestial dynamics (validating the equivariant/GNN approach).
* The complexity of the considered systems, and a discussion thereof. We address this by elaborating on the systems in more detail, and we kindly wish to highlight recent works addressing very similar systems [1, 2, 3]. Further, the complexity is also demonstrated through the baseline methods struggling in some scenarios. Finally, to better demonstrate the scalability of our approach for the n-body case, we have performed additional experiments on a system of 20 bodies, as opposed to 5 bodies in the paper.

In addition, we address individual comments with respect to our work in the context of other related works, and explain how the paper will be updated to include a discussion on these works. We thank the reviewers for pointing us to these references. Further, we have clarified some specific pieces of the paper that could be confusing or misinterpreted. Finally, we have addressed the remaining specific individual questions. 

The additional results can be found in the tables at the bottom of this general response and in the figures in the attached pdf file, and are further explained in the individual responses.

We would like to again thank the reviewers for their detailed comments as it allowed us to substantially improve our work!

[1] V. G. Satorras et al. “E(n) equivariant graph neural networks”, ICML 2021

[2] C. Yildiz et al. “Learning interacting dynamical systems with latent gaussian process ODEs”, NeurIPS 2022

[3] O. Puny et al. “Frame averaging for invariant and equivariant network design”, ICLR 2022

---

Table R1: additional results of ablations for celestial dynamics.
|  | NSDE | iGPODE | PNS | EPNS |\||  EPNS-equivariant latents | EPNS-no latents | EPNS-deterministic | PNS-MLP |
|-|:-:|:-:|:-:|:-:|-|:-:|:-:|:-:|:-:|
| $\uparrow$ LL ($\cdot 10^3$) |-5.0±0.0|-8.5±0.1|11.3±0.3|11.6±0.1|\||10.9±0.1|**13.2±0.1**|N/A|5.9±0.3|
| $\downarrow$ $D_{ks}(KE)$ t=50 |0.80±0.01|0.42±0.07|0.55±0.13|0.39±0.12|\||**0.13±0.03**|0.25±0.12|0.63±0.08|0.52±0.21|
| $\downarrow$ $D_{ks}(KE)$ t=100 |0.65±0.03|0.57±0.02|0.30±0.10|**0.18±0.02**|\||**0.18±0.02**|**0.19±0.05**|0.7±0.05|0.44±0.11|

---

Table R2: additional results for ablations for cellular dynamics
|  | ODE$^2$VAE| PNS | EPNS |\||  EPNS-invariant latent | EPNS-no latents | EPNS-deterministic |
|-|:-:|:-:|:-:|-|:-:|:-:|:-:|
| $\uparrow$ LL ($\cdot 10^4$) |-30.1±0.0|-16.4±0.3|**-5.9±0.1**|\||-58.2|-9.77|N/A|
| $\downarrow$ $D_{ks}$(#clusters) t=30 |0.98±0.01|0.70±0.10|**0.58±0.09**|\||1|1|0.77|
| $\downarrow$ $D_{ks}$(#clusters) t=45|0.96±0.04|0.77±0.05|**0.58±0.05**|\||1|1|0.9|

---

Table R3: results for celestial dynamics with 20 bodies.
|  | PNS | EPNS |
|-|:-:|:-:|
| $\uparrow$ LL ($\cdot 10^3$) |15.3±0.2|15.7±0.0|
| $\downarrow$ $D_{ks}(KE)$; t=50 |0.54±0.29|0.28±0.01|
| $\downarrow$ $D_{ks}(KE)$; t=100 |0.19±0.02|0.26±0.12|

---

Table R4: mean velocity of cells (in #grid sites per step)
|Ground truth  | EPNS | EPNS - deterministic |
|-|:-:|:-:|
| 1.821 |1.35|0.125|




Meta-review: The paper proposes a model for equivariance in distributions of trajectories rather than function approximators directly, the idea being that with stochastic simulators probabilistic models and distributions are the natural fit. All but one reviewer have a positive outlook to the paper. They appreciate the innovation, the generality, and good results in a wide array of benchmarks.

One of the reviewers voiced concerns regarding the scalability and generalizability of the method on higher-dimensional settings. I agree with the reviewer: the experiments are on rather small settings, and I do not find very convincing the argument that the 80x80x64 cellular dynamics problem constitutes a high-dimensional problem, because the dimensions here is due to the way the problem is parameterized for the sake of the approach. The question is whether the model can handle large systems (which have large dimensionalities), not whether they can handle any large dimensionality problem (one can also append 0s or random noise dimensions to trivially make very larger dimensional problems, this does not increase the intrinsic dimensionality, however).

That said, I agree with the sentiment of the authors, that there has been sufficient empirical justification and even they tried their method in a 20-body problem. Larger systems are for sure of interest, however, I think for this paper the novelty and validation is ok. I ask from the authors however to incorporate all the criticisms regarding the absence of truly high-dimensional systems and other types of dynamics (eg Langevin) in their limitations.
Decision: Accept (poster)