Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos

Authors: Christian Pedersen, Laure Zanna, Joan Bruna

ICML 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental We demonstrate this approach on two high-dimensional turbulent systems, a forced 2D Navier-Stokes flow, and a 2-layer quasigeostrophic turbulent flow, enabling stable predictions over 1e5 emulator steps.
Researcher Affiliation Academia 1Courant Institute of Mathematical Sciences, New York University, USA 2Center for Data Science, New York University, USA.
Pseudocode Yes Algorithm 1 Algorithm for thermalized trajectories
Open Source Code Yes All model architecture, training and inference codes can be found at https://github.com/Chris-Pedersen/thermalizer.
Open Datasets No To build a training and test set, we use numerical simulations to generate a total of N = 500,000 trajectories. We use 450,000 of these for training and the remaining 50,000 for validation and testing.
Dataset Splits Yes We use 450,000 of these for training and the remaining 50,000 for validation and testing.
Hardware Specification No leveraging the speed of graphical processing units (GPUs) to provide fast approximate solutions.
Software Dependencies No We numerically solve the equations using the pseudospectral method with periodic boundary conditions from the publicly available code jax-cfd (Dresdner et al., 2022). This was implemented in PyTorch, the code for which will be made publicly available upon de-anonymization.
Experiment Setup Yes The network weights are optimized using the Adam W optimizer with momenta β1 = 0.9 and β2 = 0.999, and a learning rate of 5e-4. In practice, equation 6 is broken up into mini-batches of size 32. We train the model for 12 epochs... For the thermalizer... To optimise the network, equation 8 is broken up into mini-batches of size 64. Again we use the Adam W optimizer, with a learning rate of 2e-5, and train the thermalizer for 35 epochs.