Equivariant Neural Simulators for Stochastic Spatiotemporal Dynamics
Authors: Koen Minartz, Yoeri Poels, Simon Koop, Vlado Menkovski
NeurIPS 2023 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | 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 evaluate EPNS on two diverse problems: an n-body system with stochastic forcing, and stochastic cellular dynamics. For both problem settings, we generate training sets with 800 trajectories and validation and test sets each consisting of 100 trajectories. |
| Researcher Affiliation | Academia | 1Data and Artificial Intelligence Cluster, Eindhoven University of Technology 2Swiss Plasma Center, École Polytechnique Fédérale de Lausanne |
| Pseudocode | No | The paper describes the model architecture and training procedure in detail, including mathematical formulations and schematic figures (e.g., Figure 2, Figure 3), but it does not provide any explicitly labeled pseudocode or algorithm blocks. |
| Open Source Code | Yes | Our code is available at https://github.com/kminartz/EPNS. |
| Open Datasets | No | For both problem settings, we generate training sets with 800 trajectories and validation and test sets each consisting of 100 trajectories. The paper states they generated the data rather than using a publicly available dataset, and no access information for the generated data is provided. |
| Dataset Splits | Yes | For both problem settings, we generate training sets with 800 trajectories and validation and test sets each consisting of 100 trajectories. |
| Hardware Specification | Yes | All models are implemented in Py Torch [35] and trained on a single NVIDIA A100 GPU. |
| Software Dependencies | No | The paper mentions 'implemented in Py Torch [35]' but does not provide a specific version number for PyTorch or any other software libraries or dependencies with version numbers. |
| Experiment Setup | Yes | The maximum rollout lengths used for multi-step training are 16 steps (celestial dynamics) and 14 steps (cellular dynamics). We use a linear KL-annealing schedule [8], as well as the free bits modification of the ELBO as proposed in [22] for the cellular dynamics data. 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 β that weighs the KL term of the loss by 0.005 every 200 epochs. We use a batch size of 64. For cellular dynamics, we train for 180 epochs with a learning rate of 10 4, a weight decay of 10 4, β1 = β2 = 0.9, and ε = 10 6. The free bits parameter λ = 0.1175 and batch size of 8 are also provided. |