SlowGuess/NeurIPS_2023_OCR huggingface dataset:
First, the forward model independently embeds each node's mass to an invariant node embedding  $h^t \in \mathbb{R}^{128}$ , 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 5 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 in 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 a useful E(n)-invariant node embedding.
Conditional prior:  $\left\{h^{t}\in \mathbb{R}^{n\times 128},e^{t}\in \mathbb{R}^{n^{2}\times 2}\right\} \to p_{\theta}(z|h^{t})\to \{z\in \mathbb{R}^{n\times 16}\}$
As we chose a conditional prior that is permutation equivariant, the conditional prior architecture is a message passing GNN with 5 layers that takes as input the  $E(n)$ -invariant node embeddings produced by the forward model. A linear layer maps the node embeddings to the mean of the conditional prior distribution, while a 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:
$$
\left\{h ^ {t} \in \mathbb {R} ^ {n \times 1 2 8}, e ^ {t} \in \mathbb {R} ^ {n ^ {2} \times 2}, e ^ {t + 1} \in \mathbb {R} ^ {n ^ {2} \times 2} \right\}\rightarrow q _ {\phi} (z | h ^ {t}, x ^ {t + 1}) \rightarrow \left\{z \in \mathbb {R} ^ {1 6} \right\}
$$
The architecture of the approximate posterior is identical to the conditional prior, with the exception that it accepts additional edge features as input. These are the same E(n)-invariant features that are also given as input to the forward model, but calculated for the system at time  $t + 1$ , instead of time  $t$ .
Decoder:
$$
\left\{h ^ {t} \in \mathbb {R} ^ {n \times 1 2 8}, z \in \mathbb {R} ^ {1 6}, e ^ {t} \in \mathbb {R} ^ {n ^ {2} \times 2}, \mathbf {p} ^ {t} \in \mathbb {R} ^ {n \times 3}, \mathbf {v} ^ {t} \in \mathbb {R} ^ {n \times 3} \right\}\rightarrow p _ {\theta} (x ^ {t + 1} | h ^ {t}, z) \rightarrow
$$
$$
\left\{\mathbf {p} ^ {t + 1} \in \mathbb {R} ^ {n \times 3}, \mathbf {v} ^ {t + 1} \in \mathbb {R} ^ {n \times 3} \right\}
$$
The decoder is an FA-GNN with three layers. As input, it gets the invariant node embeddings  $h^t$  produced by the forward model, the invariant  $z$  sampled from the conditional prior (during generation) or approximate posterior (during training), and the position and velocity vectors  $\mathbf{p}^t$  and  $\mathbf{v}^t$ . The decoder equivariantly maps this input to the mean and scale parameters of the Gaussian output distribution. More specifically, the model outputs vectors  $\mu^{\Delta p} \in \mathbb{R}^{n \times 3}$  and  $\mu^{\Delta v} \in \mathbb{R}^{n \times 3}$ . Furthermore, the resulting  $\mathrm{E}(\mathrm{n})$ -invariant node embeddings are also independently processed by a linear layer with Softplus activation to map to  $\sigma^p$  and  $\sigma^v$ , and by a two layer MLP mapping to a quantity  $\Delta_v \in \mathbb{R}$  which is used to scale the velocity's effect on the node position update. Using the quantities described above, the decoder's output distribution is parameterized as follows:
$$
\boldsymbol {\mu} ^ {v} = \mathbf {v} ^ {t} + \boldsymbol {\mu} ^ {\Delta v}
$$
$$
\boldsymbol {\mu} ^ {p} = \mathbf {p} ^ {t} + \Delta_ {v} \cdot \boldsymbol {\mu} ^ {v} + \boldsymbol {\mu} ^ {\Delta p}
$$
$$
p _ {\theta} \left(\mathbf {p} ^ {t + 1} \mid h ^ {t}, z\right) = \mathcal {N} \left(\boldsymbol {\mu} ^ {p}, \boldsymbol {\sigma} ^ {p}\right)
$$
$$
p _ {\theta} \left(\mathbf {v} ^ {\mathbf {t} + \mathbf {1}} \mid h ^ {t}, z\right) = \mathcal {N} \left(\boldsymbol {\mu} ^ {v}, \boldsymbol {\sigma} ^ {v}\right)
$$
The reason for using  $\pmb{\mu}^{\Delta p}$  is that we take a relatively large stepsize of  $\Delta t = 0.01$ , and as such  $\pmb{\mu}^{v}$  might not align with the difference between the coordinates  $\mathbf{p}^{t + 1} - \mathbf{p}^t$ . Accordingly,  $\pmb{\mu}^{\Delta p}$  can correct for this difference if necessary.
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 [21] with a learning rate equal to  $10^{-4}$  and weight decay of  $10^{-4}$ . For the KL annealing schedule, we increase the coefficient  $\beta$  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 as input a 2D grid of dimensions  $h \times w$  with two channels: the first channel contains integer values between 0 and  $c$ , indicating the index of each cell, where 0 stands for background and  $c = |C|$  is the number of cells. The second channel contains integer values between 0 and  $\tau$  that indicate the cell type, where 0 is again background, and  $\tau \leq c$  is 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 we describe the SpatialConv-GNN architecture for this model in more detail.
Forward model:  $\{x^{t}\in \mathbb{Z}^{h\times w\times 2}\} \to f_{\theta}\to \{h^{t}\in \mathbb{R}^{c\times h\times w\times 32}\}$
The forward model first one-hot encodes both the index and type channels of the input  $x^t$ , yielding two grids of shape  $h \times w \times c$  and  $h \times w \times \tau$ . Then, each of the channels in the one-hot encoded cell index grid is seen as a separate node, and the one-hot encoded type channel tensor is 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  $c$  grids of shape  $h \times w \times (\tau + 1)$ , which we simply store as one grid of shape  $c \times h \times w \times (\tau + 1)$ . Subsequently, a linear convolution embeds the grid to an embedding dimension of 32, resulting in a tensor of shape  $c \times h \times w \times 32$ . Then, a SpatialConv-GNN with a single message passing layer processes this tensor, again resulting in a grid of shape  $c \times h \times w \times 32$ .
Conditional prior:  $\{h^t\in \mathbb{R}^{c\times h\times w\times 32}\} \to p_\theta (z|h^t)\to \{z\in \mathbb{R}^{c\times 64}\}$
The conditional prior independently maps each cell to a latent variable  $z$  by applying a convolutional layer with ReLU activation, followed by three repetitions of  $\{\text{convolution} \to \text{ReLU} \to \text{MaxPool}\}$ . All convolution layers have a kernel size of  $9 \times 9$ , and we perform  $2 \times 2$  maxpooling. Then, global meanpooling is performed over the spatial dimensions. Subsequently, a linear layer followed by ReLU activation is applied. The resulting vector is 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.
$$
\left\{h ^ {t} \in \mathbb {R} ^ {c \times h \times w \times 3 2}, x ^ {t + 1} \in \mathbb {Z} ^ {h \times w \times 2} \right\}\rightarrow q _ {\phi} (z | h ^ {t}, x ^ {t + 1}) \rightarrow \left\{z \in \mathbb {R} ^ {c \times 6 4} \right\}
$$
The architecture of the approximate posterior is almost identical to the conditional prior. The only exception is that first,  $x^{t + 1}$  is encoded to nodes in the same way as described in the forward model, resulting in a binary tensor  $x^{t + 1} \in \{0,1\}^{c \times h \times w \times (\tau + 1)}$ . 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.
$$
\left\{h ^ {t} \in \mathbb {R} ^ {c \times h \times w \times 3 2}, z \in \mathbb {R} ^ {c \times 6 4} \right\}\rightarrow p _ {\theta} \left(x ^ {t + 1} \mid h ^ {t}, z\right)\rightarrow \left\{x ^ {t + 1} \in \mathbb {Z} ^ {h \times w \times 2} \right\}
$$
For each cell, the sampled latent variable  $z$  is concatenated along the channel axis for all pixels in the grid. The decoder then consists of a SpatialConv-GNN with a single message passing layer, followed by a linear convolutional layer with kernel size  $1 \times 1$  that maps to a single channel. The resulting tensor has shape  $c \times h \times w \times 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 a
(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^{t}$ .
SpatialConv-GNN: Recall the SpatialConv-GNN message passing layer from Equation 9:
$$
h _ {i} ^ {l + 1} = \psi^ {l} \left(h _ {i} ^ {l}, \bigoplus_ {j \in \mathcal {N} (i)} \phi^ {l} \left(h _ {j} ^ {l}\right)\right).
$$
The core components that define the SpatialConv-GNN are  $\phi, \bigoplus$ , the selection of the neighbors  $\mathcal{N}(i)$ , and  $\psi$ . We describe the details of each of the components below:
-  $\phi$ : First,  $\phi$  performs a strided convolution with kernel size  $2 \times 2$  and a stride of 2 to downsample each grid  $h_j$ . Then, three repetitions of  $\{\text{convolution} \to \text{ReLU}\}$  are applied, where the convolutions have a kernel size of  $9 \times 9$ .
-  $\bigoplus$ : We use elementwise mean as permutation-invariant aggregation function.
-  $\mathcal{N}(i)$ : We opted for a fully-connected graph, meaning  $\mathcal{N}(i) = C$  for all  $i \in C$ . The reasons for this are twofold: on the one hand, the convolutional nature of  $\phi$  and  $\psi$  already takes spatial locality into account. On the other hand, using all cells as neighbors for a single cell  $i$ , including  $i$  itself, means that we only have to compute  $\bigoplus_{j \in \mathcal{N}(i)} \phi^l(h_j^l)$  once and can re-use it for all  $i \in C$ , substantially reducing memory requirements.
-  $\psi$ :  $\psi$  first upsamples the aggregated messages using a transposed convolution to the original grid size. Then, the result is concatenated along the channel axis of  $h_i$ . Subsequently we apply a simple UNet, using three downsampling and upsampling blocks. The convolutional layers in these blocks have a kernel size of  $5 \times 5$ . We also tried to parameterize  $\psi$  with only ordinary convolutional layers, but found that the UNet'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  $\psi$  remains convolutional, and consequently it still enjoys the benefits of 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 [21] with a learning rate equal to  $10^{-4}$ , a weight decay of  $10^{-4}$ ,  $\beta_{1} = \beta_{2} = 0.9$ , and  $\varepsilon = 10^{-6}$ . For the KL annealing schedule, we increase the coefficient  $\beta$  that weighs the KL term of the loss by 0.04 every epoch, until a maximum of  $\beta = 1$ . Furthermore, we use the free bits modification of the ELBO to prevent posterior collapse [22]:
$$
\begin{array}{l} \mathrm {E L B O} _ {\text {f r e e b i t s}} = \\ \mathbb {E} _ {q _ {\phi} (z | x ^ {t}, x ^ {t + 1})} \left[ p _ {\theta} \left(x ^ {t + 1} \mid z, x ^ {t}\right) \right] - \beta \sum_ {j = 1} ^ {| z |} \operatorname * {m a x i m u m} \left(\lambda , K L \left[ q _ {\phi} \left(z _ {j} \mid x ^ {t}, x ^ {t + 1}\right) \right| \left| p _ {\theta} \left(z _ {j} \mid x ^ {t}\right) \right]\right), \\ \end{array}
$$
with  $\lambda = 0.1175$ . In the above objective function, the KL terms for each  $z_{j}$  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 is identical to the EPNS model described in Appendix C.1, with a few exceptions. Instead of an equivariant FA-GNN architecture for the forward
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.
iGPODE For the iGPODE baseline, we used the code from the original paper [59].<sup>2</sup> In terms of hyperparameter optimization, we 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.
NSDE We train the NSDE as a latent SDE, like [27] but without the adjoint sensitivity method. The model consists six networks: a context encoder, f-net, h-net, g-net, initial encoder, and a 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, 27], 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 a 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-autoregressive 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 a Gaussian distribution with diagonal covariance. We also experimented with parameterizing the covariance and keeping it 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  $\gamma = 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 backbone, which has shown state-of-the-art performance in modeling spatiotemporal dynamics on a 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, as opposed to 13 million for EPNS. We use ReLU activations and a hidden dimension of 128 throughout the model.
First, the input is one-hot encoded the same way as done for EPNS (see Appendix C.2). Then, it is processed by the model components which are defined as follows:
- **Forward model:** a modern UNet architecture with wide ResNet blocks, spatial attention, downsampling layers, and group normalization, as described in [14]. The architecture consists of three downsampling and upsampling steps, with attention only being applied in
Available at https://github.com/boschresearch/igpode at the time of writing.
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 a 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 conditional prior, with the exception that it accepts the one-hot encoding of  $x^{t+1}$  as additional input channels.
- Decoder: the decoder consists of 4 convolutional layers with kernel size  $3 \times 3$ , followed by Softmax activation at the end to map to pixel-wise probabilities over the cell indices.
- Training procedure: as for EPNS, we train PNS with multi-step training for 180 epochs, with the same KL annealing schedule. We use a free bits parameter of  $\lambda = 0.1$  to prevent posterior collapse. For optimization, we use Adam with a learning rate of  $10^{-4}$  and a weight decay of  $10^{-4}$ .
ODE²VAE Our ODE²VAE implementation is adapted from the PyTorch implementation of [58].<sup>3</sup> 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 a batch size of 8 to keep the computational requirements reasonable. The best-performing model in terms of validation loss has a 64 dimensional latent space, 32 filters in the initial convolutional layer, and 4 convolutional layers in the encoder and decoder, which has around 21 million trainable parameters.
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 a latent space but rather produces  $p_{\theta}(x^{t + 1}|x^t)$  directly, giving insight into the need of the latent variable approach. As opposed to the main model, where we sample  $z$  but construct  $x^{t + 1}$  such that  $p_{\theta}(x^{t + 1}|z,h^t)$  is maximized as is commonly done in VAEs, here we sample from the decoder distribution  $p(x^{t + 1}|h^t)$ , which does not depend on  $z$ ;
- 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  $D_{\mathrm{KS}}$ . 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 is trained using single step training, which purely optimizes the one-step transition probability.
Available at https://github.com/cagatayildiz/ODE2VAE at the time of writing.
Table 6: Results for variants of the EPNS framework on the celestial dynamics application.
PNS	EPNS	EPNS-invariant latents	EPNS-no latents	EPNS-deterministic	PNS-MLP
↑ LL (·103)	10.9±0.4	10.8±0.1	11.6±0.1	13.2±0.1	N/A	5.9±0.3
↓ DKS(KE)t=50	0.61±0.18	0.14±0.03	0.39±0.12	0.25±0.12	0.63±0.08	0.52±0.21
↓ DKS(KE)t=100	0.20±0.06	0.14±0.04	0.18±0.02	0.19±0.05	0.7±0.05	0.44±0.11
Table 7: Results for variants of the EPNS framework on the cellular dynamics application.
PNS	EPNS	EPNS-invariant latents	EPNS-no latents	EPNS-deterministic
↑LL (·104)	-16.4±0.3	-5.9±0.1	-27.5±15.4	-9.9±0.0	N/A
↓DKS(#clusters) t=30	0.70±0.10	0.58±0.09	0.60±0.22	1.00±0.00	0.93±0.08
↓DKS(#clusters) t=45	0.77±0.05	0.58±0.05	0.55±0.25	1.00±0.00	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 it to produce better rollouts - see also the  $D_{\mathrm{KS}}$  values in Table 6 - but at the cost of not directly optimizing for LL. The second reason is that the step size of  $\Delta t = 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  $D_{\mathrm{KS}}$ , EPNS with global, permutation-invariant latent variables performs competitively 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, Figure 7a shows that the main EPNS configuration with permutation equivariant latents tends to 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 an MLP backbone has poor rollout stability, which is in line with its poor performance across the other metrics.