Research map / Probabilistic and scientific modeling / Dynamical systems

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos

ICML 2025 · Probabilistic and scientific modeling › Dynamical systems · approach: Data assimilation from sparse observations

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes the thermalizer, which couples a pretrained black-box autoregressive emulator with a separately trained diffusion-model denoiser of the invariant measure, applying adaptive on-the-fly denoising at inference to stabilize long chaotic rollouts.

Problem

Autoregressive neural emulators of spatiotemporal chaotic systems accumulate errors and drift out of distribution, causing trajectories to diverge over long rollouts.

Approach

Use a diffusion model trained on samples of the stationary measure to estimate its score, then at inference time detect (via a noise-level classifier head) when the emulator state drifts out of distribution and run partial reverse-diffusion denoising steps to pull it back toward the invariant measure, without retraining the emulator.

Key results

• Thermalized emulator rollouts remain stable for 1e5 steps on Kolmogorov flow and QG, extending stable prediction horizon by about two orders of magnitude versus un-thermalized emulators that diverge around 2000-5000 steps.
• On QG with the Unet emulator, only about 140 thermalization steps on average are needed to stabilize each 10,000-step trajectory.
• The noise-classifying head costs about 60% of a Unet emulator step per timestep, with full denoising passes triggered only when the predicted noise level exceeds s_init.

Measured effects reported in the paper

Surprising findings

• Lower validation loss and short-term MSE do not reliably translate into long-term stability; the larger Unet had lower short-term MSE but all trajectories went unstable quickly, while the smaller DRN was more stable long-term.
• Training the thermalizer only on low noise levels (relevant at inference) degraded performance significantly, so noise levels were sampled across the full range.

Limitations

• Relies on isotropic Gaussian perturbation assumption of the diffusion model, which may not match emulator errors.
• Performance depends strongly on hyperparameters s_init and s_stop, tuned by brute-force grid search per emulator-thermalizer pair.
• Some stable configurations show imperfect kinetic energy spectra at long timesteps; QG shows an offset in kinetic energy spectra versus the numerical model.

Experimental setup

Datasets: Kolmogorov flow simulations (jax-cfd, 64x64 vorticity) 2-layer quasi-geostrophic turbulence simulations (64x64 potential vorticity)

Benchmarks: Kolmogorov flow rollouts up to 1e5 steps 2-layer quasi-geostrophic turbulence rollouts up to 1e5 steps

Metrics: MSE versus numerical model radially averaged kinetic energy spectra autocorrelation number of thermalization steps stability over rollout horizon

Baselines: Unet emulator without thermalization Dilated ResNet emulator without thermalization numerical solver (jax-cfd / PyTorch pseudo-spectral)

Builds on: DDPM Tweedie's formula reverse diffusion sampling Ornstein-Uhlenbeck process

Closest papers by meaning

• Hierarchical Implicit Neural Emulators (NeurIPS 2025 · similarity 0.88)
• The Dynamic-Probabilistic Consistency Gap in Chaotic Surrogate Modeling (NeurIPS 2026 · similarity 0.87)
• APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs (NeurIPS 2024 · similarity 0.86)
• ABC: Any-Subset Autoregression via Non-Markovian Diffusion Bridges in Continuous Time and Space (NeurIPS 2026 · similarity 0.86)
• Sparse Diffusion Autoencoder for Test-time Adapting Prediction of Complex Systems (NeurIPS 2025 · similarity 0.86)
• Predicting partially observable dynamical systems via diffusion models with a multiscale inference scheme (NeurIPS 2025 · similarity 0.86)