SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the cellular dynamics problem, as shown in Figure 8, 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 at 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 EPNS-deterministic 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^0$  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-truth.
E.2 Phase space coverage for celestial dynamics
We also investigated the phase space that is 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 11 (PNS-MLP). Visually, the phase space explored by
(a) EPNS and ablations.
(b) EPNS and PNS-MLP.
Figure 7: Fraction of simulations that remain stable over long rollouts for celestial dynamics.
Figure 8: Fraction of simulations that remain stable over long rollouts for cellular dynamics.
both EPNS and PNS appears to overlap well with the ground truth, and their performance appears similar. In 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.
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 1, 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.
Figure 10: Phase-space coverage plots of the n-body system similar to Figure 9, but with PNS as model.
Figure 11: Phase-space coverage plots of the n-body system similar to Figure 9, but with PNS with MLP backbone as model.
E.3 20-body system results
In this section, we report results of applying EPNS to a 20-body system, as opposed to a 5-body system in Section 5. Although this system is not as high-dimensional as most challenging potential real-life applications, the 20-body system can help to assess the potential of scaling EPNS to such systems. Notably, for these experiments, we do not perform any hyperparameter optimization relative to the 5-body EPNS model.
Table 8 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 to 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  $D_{\mathrm{KS}}$  scores of EPNS are similar to 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 over 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, we observe that the distribution generated by EPNS overlaps very well with the ground-truth, and notably better than the distribution generated by PNS.
Finally, similar to Section E.2, 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 9, we observe that both densities overlap very well, suggesting that the distribution over trajectories from a fixed initial condition produced by EPNS matches the ground truth, also in this higher-dimensional setting.
Table 8: Results for scaling EPNS to a 20-body system.
PNS	EPNS
↑ LL (·103)	15.4±0.1	15.6±0.0
↓ DKS(KE) t=30	0.22±0.03	0.14±0.03
↓ DKS(KE) t=45	0.21±0.03	0.14±0.05
(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 at 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 ELBO values (shown in Table 1). 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
Figure 13: Phase-space coverage plots of the n-body system similar to Figure 9, but with EPNS trained on a celestial dynamics with 20 bodies.
can struggle with handling close encounters, and a body may get 'stuck' to 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 are shown in Figure 15. Again, we selected 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 to have unrealistic volumes, also reflected in the rapidly decaying stability shown in Figure 6b. ODE $^2$ VAE generally produces simulations in which cells look 'frozen' and lack clustering behavior. We attribute this to the fact that ODE $^2$ VAE models the dynamics in latent space, and struggles to propagate the geometric structure of the data as time proceeds as a result.
Figure 14: Additional qualitative results for celestial dynamics simulation.
Figure 15: Additional qualitative results for cellular dynamics simulation.