SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Emails of the corresponding authors are on the first page of the main paper.
- Is there an erratum? If so, please provide a link or other access point. No. Zenodo allows versioning.
If there are changes to the dataset, these will be made and documented on Zenodo.
Versioning will be managed by Zenodo.
The issue system on the project Guthub can be used to report bugs and suggest dataset extensions. Using custom datasets on top of the provided ones is already discussed in the paper and our codebase seamlessly allows for it (see "notebooks" in GitHub repository).
B Solver validation
A crucial part of numerical solver development is the validation. This is the process of reproducing the dynamics of systems either 1) with an analytical solution or 2) with a high-quality reference solution from another numerical simulation. We chose two example systems for validation, namely the Poiseuille flow and the dam break systems.
B.1 Poiseuille flow
Here, we reproduce the Poiseuille flow case described in Morris et al. [45]. The only technical difference is that we first non-dimensionalize the problem before we run the simulation, but this procedure does not influence the Reynolds number and thus the behavior of the system. For more details on non-dimensionalization, we refer to Bezgin et al. [8]. The reference values we use are the density  $\rho_{ref} = 1000kg / m^3$ , length  $L_{ref} = 0.001m$ , and velocity  $V_{ref} = 1\mathrm{e} - 5m / s$ . Our dimensionless results almost perfectly match the analytical series expansion solution with 60 particles spanning the width of the channel, see Figure 3.
Figure 3: Comparison of SPH (○) and series solutions (-) for Poiseuille flow at  $Re = 0.0125$ .
B.2 Dam break
We recover the dam break simulation very closely to the reference result in Adami et al. [1]. Here, we include a quantitative validation of the pressure on the right wall at height 0.2 (see Figure 4), as well as a qualitative comparison of the pressure distribution at time instances  $t = 2, 5.7, 6.2$ , and 7.4 (see Figure 5).
C Dataset details
Here we give more details on the dataset-generation process by continuing Table 1 with Table 4. The properties we present are 1) the boundary conditions type of the computational domain, 2) the magnitude of external force  $F_{mag}$ , 3) the quantities from the equation of state in Section 2 being  $c_0$ ,  $\rho_0$ , and  $p_{bg}$ , 4) the length of one trajectory of a given type in dimensionless time units, 5) the viscosity, which is proportional to  $Re^{-1}$ , and 6) the size of each dataset in GB.
(a) Our results.
(b) Results from Adami et al. [1]
Figure 4: Temporal pressure evolution at  $y / H = 0.2$  on the right wall.
(b) Results from Adami et al. [1]
Figure 5: Pressure distribution along dam break simulation at  $t = 2, 5.7, 6.2$ , and 7.4.
D Dataset visualizations
Figures 6, 7, 8, 9 and Figures 10, 11, 12 show the time evolution during the rollout of 2D and 3D datasets, respectively. Each dot represents a particle, and a warmer color represents a higher velocity magnitude. In the cases of 6, 7, 8 the (approximately) correct scale is kept.
Table 4: Datasets overview. Either the trajectory length or the number of trajectories is used for splitting into training/validation/testing datasets.
Dataset	periodic domain	\(F_{mag}\)	\(c_0\)	\(\rho_0\)	\(p_{bg}\)	SPH runtime 100 steps [ms]	viscosity	size [GB]
2D TGV	yes	0	10	1	0	39.1	0.01	0.43
2D RPF	yes	1	10	1	5	43.0	0.1	0.88
2D LDC	only lid	0	10	1	1	51.2	0.01	0.34
2D DAM	no	1	14.14	1	0	89.4	5e-5	1.4
3D TGV	yes	0	10	1	0	388	0.02	2.0
3D RPF	yes	1	10	1	2	424	0.1	1.6
3D LDC	lid and in z	0	10	1	1	678	0.01	1.4
Figure 6: Taylor Green vortex velocity snapshot at step  $= 10,60$  and 113.
Figure 7: Reverse Poiseuille flow velocity snapshot at  $step = 10,9600$ , and 18000.
Figure 8: Lid driven cavity velocity snapshot at  $step = 10,4800$ , and 9000.
Figure 9: Dam break velocity snapshot at step  $= 10,192$  , and 360.
Figure 10: Taylor Green vortex 3D velocity snapshot at  $step = 10,29$ , and 54.
Figure 11: Reverse Poiseuille flow velocity snapshot at  $step = 10,4800$ , and 9000.
Figure 12: Lid driven cavity 3D velocity snapshot at step = 10, 4800, and 9000.
E Neighbor search implementations
We integrate three neighbor list routines in our codebase: original jax-md implementation (jaxmd_vmap), serialized version thereof (jaxmd_scan), and a wrapper around the matscripy neighbor list algorithm (matscripy).
jaxmd_vmap refers to using the original cell list-based implementation from the JAX-MD<sup>3</sup> library. This one can be seamlessly integrated into our JAX-based API and is the preferred choice for systems with a constant number of particles.
jaxmd_scan refers to using a more memory-efficient implementation of the JAX-MD function. We achieve this by partitioning the search over potential neighbors from the cell list-based candidate neighbors into  $M$  chunks. We need to define three variables to explain how our implementation works:
-  $X \in \mathbb{R}^{N \times d}$  - the particle coordinates of  $N$  particles in  $d$  dimensions.
-  $h \in \mathbb{N}^N$  - the list specifying to which cell a particle belongs.
-  $L \in \mathbb{N}^{C \times \text{cand}}$  - list specifying which particles are potential candidates to a particle in cell  $c \in [1, \dots, C]$ . The number of potential candidates  $\text{cand}$  is the product of the fixed cell capacity and the number of reachable cells, e.g. 27 in 3D.
The jaxmd_vmap implementation essentially instantiates all possible connections by creating an object of size  $N \cdot \text{cand}$ , and only after all distances between potential neighbors have been computed. The edge list is pruned to its actual size being roughly  $6 \times$  smaller in 3D. This factor comes from the fact that the cell size is approximately equal to the cutoff radius and if we split a unit cube into  $3^3$  cells, then the volume of a sphere with  $r = 1/3$  will be around  $1/6$  the volume of the cube. By splitting  $X$  and  $h$  into  $M$  parts and iterating over  $L$  with a jax.lax.scan loop, we can remove  $5/6$  of the edges before putting them together into one list.
matskipy is used to enable computations over systems with a variable number of particles, for which none of the above implementations can be used. We essentially wrap the neighbor search routine from matskipy.neighbours.neighbour_list<sup>4</sup> with a jax pure_callback. This is again a cell list-based algorithm, however only available on CPU. Our wrapper essentially mimics the behavior of the JAX-MD function, but pads all non-existing particles to the maximal number of particles in the dataset.
E.1 Performance
To demonstrate the improvement we get from jaxmd_scan relative to jaxmd_vmap, we compare the largest number of particles whose neighbor list computation fits into memory. We ran the comparison on an A6000 GPU with 48GB memory and observed that the default vectorized implementation can handle up to 1M particles before running out of memory, while our serialized implementation reaches 3.3M. This happens at almost no additional time cost and holds for both allocating a system and updating it after compilation.
The matscripy implementation itself is very fast for small systems (10k particles) and doesn't take any GPU memory for the construction of the edge list, but due to copying data between CPU and GPU, we observed a slowdown in training a SEGNN on 2D DAM in the order of  $+50\%$ . In addition, it seems that matscripy uses a single CPU computation and is therefore limited to smaller systems (see this timing comparison https://github.com/arturtoshev/pbc_neighbors_benchmark).
Overall, we observe reasonable performance from each of these implementations on systems with up to  $10\mathrm{k}$  particles, but more investigations need to be conducted toward comparing these algorithms on larger systems.
3 https://github.com/jax-md/jax-md
4https://github.com/libAtoms/matscipy
F JAX vs PyTorch GNN performance comparison
We compare the speed of PyTroch vs JAX on all four models (GNS, SEGNN, EGNN, PaiNN). We note that all models have existing implementations in PyTroch and all were reimplemented by us in JAX. For each model, we consider datasets and experiments from the original paper, namely:
- GNS - we use the WaterDropSample dataset from Sanchez-Gonzalez et al. [56] (available under github.com/deepmind/deepmind-research/tree/master/learning_to_simulate), which is a subset of a 2D fluid dataset with up to 1000 particles. As a reference PyTorch implementation, we use github.com/wu375/simple-physics-simulator-pytorch-geometry, and we average the forward evaluation time for one particular system with 803 particles over 5 rollouts of 1000 steps each. The neighbor search with LagrangeBench uses matscripy.
- EGNN - we evaluate runtime on the  $NBody_{(5)}$  dataset [57], a custom 3D particle dataset with 5 electrical charges in free space. The reference PyTorch implementation is the official EGNN github.com/vgsatorras/egnn. Results are reported on batches of 100 systems (500 particles per batch).
- SEGNN - similarly to EGNN, we measure SEGNN runtime on  $NBody_{(5)}$  and  $NBody_{(100)}$ , which are 5 and 100 particle systems respectively [57, 12]. The reference PyTorch implementation is the official SEGNN github.com/RobDHess/Steerable-E3-GNN. Inference is over a batch of 100 samples on the respective dataset (500 and 10000 particles per batch).
- PaiNN - we evaluate runtime on QM9 [43]. For the Jax implementation, batches must be padded to the worst possible size (because of static shapes). The reference PyTorch implementation as well as QM9 data access is adapted from the official implementation from SchnetPack github.com/atomistic-machine-learning/schnetpack. Results are reported on batches of 100 systems (variable number of nodes).
Table 5: Times [ms] per step for JAX and PyTorch inference.
GNS	SEGNN	EGNN	PaiNN
WaterDropSample	NBody(5)	NBody(100)	NBody(5)	QM9
Torch	9.2	21.22	60.55	8.27	163.23
Jax	8.0	3.77	41.72	0.94	8.42
G Benchmarking hyperparameters
The exact model configurations for SEGNN, GNS, EGNN, and PaiNN can be found in Table 6. In particular, the hyperparameters are split into 4 groups, namely:
1. Model architecture. Describe the model structure.
-  $MP$  layers. Message passing steps.
- Latent size. Hidden layer dimension.
- MLP blocks. Number of "linear-activation" blocks in the update-message functions.
-  $L_{\text{max}}$  hidden. Cutoff tensor product level in the hidden layers. SEGNN only.
-  $L_{max}$  attributes. Max spherical harmonics level for attributes. SEGNN only.
2. Optimization. Learning rate scheduler settings.
- Weight decay. Weight decay regularization.
-  $LR$  initial. Starting learning rate.
-  $LR$  final. Minimum learning rate.
- LR decay steps. Number of steps to decay from LR initial to LR final
3. Training strategies. Noise and push-forward.
- Noise std. Gaussian random walk noise standard deviation.
-  $PF$  steps. Number of push-forward steps.
-  $PF$  probs. Probability of sampling each push-forward step number.
4. Normalization. Network and input normalization
- Normalization. Network normalization type.
- Isotropic inputs. Whether to normalize all input channels with the same statistics.
Training the baselines takes several days, especially the larger instances of SEGNN. In order to reduce the amount of training runs we prioritized architecture/model-related values over arguably less impactful optimization hyperparameters, such as  $LR$  decay steps and the Weight decay. Increasing or decreasing the number of MLP blocks would respectively considerably grow the parameter count or reduce the nonlinearity of the network (since the last update layer is usually linear). We investigated in-depth non-isotropic input normalization and found it advantageous since normalizing the input dimensions independently leads to a more favorable data distribution. However, because independent rescaling of vector dimensions via normalization breaks equivariance, all equivariant models use isotropic normalization by default.
Table 6: Hyperparameter configuration for SEGNN, GNS, EGNN, and PaiNN. Push-forward settings are the same for all models. Noise levels depend on the dataset.
SEGNN-5-64	SEGNN-10-64	GNS-5-64	GNS-10-128	EGNN-5-128	PairNN-5-128
MP layers	5	10	5	10	5	5
Latent size	64	64	64	128	128	128
MLP blocks	2	2	2	2	1 or 2	1 or 2
Lmaxhidden	1	1	NA	NA	NA	NA
Lmaxattributes	1	1	NA	NA	NA	NA
Params	183K	360K	161K	1.2M	663K	1.0M
Weight decay	1e-8	1e-8	1e-8	1e-8	1e-8	1e-8
LR initial	5e-4	5e-4	5e-4	1e-4	5e-4	5e-4
LR final	1e-6	1e-6	1e-6	1e-6	1e-6	1e-6
LR decay steps	1e5	1e5	1e5	1e5	5e4	5e4
Noise std	data-dependent. ∈ [3e-4, 1e-3]
PF steps	up to 3
PF probst	[0.8, 0.1, 0.05, 0.05]
Normalization	No	No	Layer	Layer	No	No
Isotropic inputs	Yes	Yes	No	No	Yes	Yes
Table 6 only shows a single configuration for EGNN and PaiNN. For EGNN, we found that going beyond 5 message passing steps resulted in unstable gradients. The poor stability of EGNN on larger systems has already been documented in other works [12]. For PaiNN, following the approach of the original paper [59], we preferred a wider model with fewer message passing layers and a large (cutoff) interaction radius.
It is worth noting that Table 6 does not capture the full extent of EGNN and PaiNN experiments. Both models were evaluated and experimented in conceptually different setups, namely:
- EGNN was experimented on 1) original N-body formulation with position updates, and 2) directly working with velocities without updating positions. However, we found that only the first approach behaved remotely reasonably.
- Different PaiNN setups were also evaluated: with/without cutoff functions or radial embeddings, including different features in the nodes, and with/without our general vector input change. We observed that a (cosine) cutoff of 1.5 times the average inter-particle distance, together with a learnable Gaussian RBF behaved the best.
H More baseline results
The most relevant results from our experiment runs are in Table 7. Only  $\mathrm{MSE}_5$  and  $\mathrm{MSE}_{20}$  metrics are reported for the sake of simplicity and runtime. Wherever a cell is marked with "unstable" it means that the model either constantly got large errors or that the metric was not stable during training.
Table 7: Benchmarking results from all experiments. The results included in Table 2 are marked in bold.
Dataset	Model	MSE5	MSE20	Inference [ms]
TGV 2D	GNS-5-64	6.4e-7	9.6e-6	1.43
GNS-10-128	3.9e-7	6.6e-6	3.43
SEGNN-5-64L=1	3.8e-7	6.5e-6	9.78
SEGNN-10-64L=1	2.4e-7	4.4e-6	20.24
RPF 2D	GNS-5-64	4.0e-7	9.8e-6	2.05
GNS-10-128	1.1e-7	3.3e-6	3.89
SEGNN-5-64L=1	1.3e-7	4.0e-6	15.11
SEGNN-10-64L=1	1.3e-7	4.0e-6	29.65
EGNN-5-128	unstable	unstable	60.77
PaiNN-5-128	3.0e-6	7.2e-5	9.09
LDC 2D	GNS-5-64	2.0e-6	1.7e-5	1.43
GNS-10-128	6.4e-7	1.4e-5	3.43
SEGNN-5-64L=1	9.9e-7	1.7e-5	9.78
SEGNN-10-64L=1	1.4e-6	2.5e-5	20.24
DAM 2D	GNS-5-64	2.1e-6	6.3e-5	3.94
GNS-10-128	1.3e-6	3.3e-5	9.68
SEGNN-5-64L=1	2.3e-6	7.5e-5	28.82
SEGNN-10-64L=1	1.6e-6	4.1e-5	59.18
TGV 3D	GNS-5-64	3.8e-4	8.3e-3	8.41
GNS-10-128	2.1e-4	5.8e-3	30.48
SEGNN-5-64L=1	3.1e-4	7.7e-3	79.42
SEGNN-10-64L=1	1.7e-4	5.2e-3	154.30
RPF 3D	GNS-5-64	1.3e-6	5.2e-5	8.41
GNS-10-128	3.3e-7	1.9e-5	30.48
SEGNN-5-64L=1	6.6e-7	3.1e-5	79.42
SEGNN-10-64L=1	3.0e-7	1.8e-5	154.30
EGNN-5-128	unstable	unstable	250.67
PaiNN-5-128	1.8e-5	3.6e-4	43.98
LDC 3D	GNS-5-64	1.7e-6	5.7e-5	8.63
GNS-10-128	7.4e-7	4.0e-5	31.98
SEGNN-5-64L=1	1.2e-6	4.8e-5	81.21
SEGNN-10-64L=1	9.4e-7	4.4e-5	161.23
I Training strategies
We explored the impact of the training tricks by training a 10-layer 128-dim latent space GNS model on the 2D RPF dataset and selecting the best model in terms of the 20-step validation MSE loss. We ran this setup in four configurations exploring the impact of random-walk noise (Noise) and push-forward (PF).
Table 8: Comparison of training with and without training strategies.
Noise	PF	MSE5	MSE20	Sinkhorn	MSEEkin
yes	yes	1.1e-7	3.3e-6	1.4e-7	1.7e-5
no	yes	1.0e-7	6.1e-6	1.1e-6	4.8e-5
yes	no	1.2e-7	3.8e-6	1.1e-7	2.8e-5
no	no	1.4e-7	8.3e-6	1.3e-6	2.4e-5
We observe that training without noise or push-forward results in significantly worse errors by all four measures. We also observe that combining the two strategies leads to the lowest  $\mathrm{MSE}_{20}$ . Push-forward smoothens the training curves, while the additive noise prevents overfitting by adding fluctuations. It is, therefore, reasonable to expect that the training with up to 3 steps push-forward leads to the best performance on  $\mathrm{MSE}_5$ , but it seems that the random noise is more important to recover the long-term behavior and the right particle distribution (lower Sinkhorn distance). Regarding the kinetic energy, it is hard to interpret the results, but it seems like a combination of both training strategies leads to the lowest error. Overall, including both training tricks seems to have a positive impact.