PaperlayerView exact revision on arXiv

Source-linked AI summary

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos

Chris Pedersen, Laure Zanna, Joan Bruna

arXiv:2503.18731v2cs.LGstat.ML

TL;DR

Autoregressive emulators struggle with long rollouts because accumulated errors cause trajectories to diverge from the training distribution. The paper introduces thermalization, which applies adaptive diffusion-based denoising from a separately trained model of the stationary measure during inference. The approach enables stable predictions over 1e5 emulator steps in two high-dimensional turbulent systems, though it depends on Gaussian perturbation assumptions and individually tuned hyperparameters.

  • Problem

    Autoregressive emulators accumulate errors during long chaotic rollouts, causing trajectories to drift out of distribution and become unstable.

    Source: §1 · pp. 1–2

  • Method

    Thermalization couples a pretrained black-box emulator with a separately trained diffusion model of the stationary distribution and applies adaptive denoising during inference.

    Source: §1 · p. 2, §2 · p. 4, §4.2 · p. 7

  • Results

    Stable predictions were demonstrated over 1e5 emulator steps for forced 2D Navier–Stokes and 2-layer quasi-geostrophic turbulent flows.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Diffusion models can stabilize autoregressive emulators by constraining rollout states toward the equilibrium solution manifold with lightweight separate training.

    Source: §6 · p. 9

    The framework assumes isotropic Gaussian perturbations, and performance depends strongly on brute-force tuning of sinit and sstop for each emulator–thermalizer pair.

    Source: §2 · p. 4, §6 · p. 9

Abstract

from arXiv · show

Autoregressive surrogate models (or \textit{emulators}) of spatiotemporal systems provide an avenue for fast, approximate predictions, with broad applications across science and engineering. At inference time, however, these models are generally unable to provide predictions over long time rollouts due to accumulation of errors leading to diverging trajectories. In essence, emulators operate out of distribution, and controlling the online distribution quickly becomes intractable in large-scale settings. To address this fundamental issue, and focusing on time-stationary systems admitting an invariant measure, we leverage diffusion models to obtain an implicit estimator of the score of this invariant measure. We show that this model of the score function can be used to stabilize autoregressive emulator rollouts by applying on-the-fly denoising during inference, a process we call \textit{thermalization}. Thermalizing an emulator rollout is shown to extend the time horizon of stable predictions by an order of magnitude in complex systems exhibiting turbulent and chaotic behavior, opening up a novel application of diffusion models in the context of neural emulation.

1. Introduction

High-dimensional turbulent systems are expensive to simulate, while autoregressive neural emulators accumulate errors and drift out of distribution over long chaotic rollouts. The paper proposes thermalization, which uses a separately trained diffusion model of the stationary distribution to stabilize inference.

  • Turbulent flows couple length scales across orders of magnitude, making accurate full-range numerical modeling often intractable.Source: §1 · p. 1
  • Autoregressive emulators feed predictions back into the model, so accumulated errors eventually drive trajectories away from the truth.Source: §1 · p. 1
  • Chaotic dynamics further challenge stability, motivating methods that incorporate knowledge of the system’s invariant measure.Source: §1 · p. 2
  • Thermalization applies diffusion-model denoising during inference to relax an emulated trajectory toward a stationary measure.The emulator and diffusion model are trained separately and combined only at inference time, with small additional computational cost.Source: §1 · p. 2
  • The method achieves arbitrarily long rollouts when an unstable neural emulator is coupled with a thermalizer on two high-dimensional turbulent systems.Source: §1 · p. 2

2. Preliminaries

The paper formulates long-horizon neural emulation as sampling trajectories from a Markov process with a stationary distribution, then analyzes why autoregressive errors create distribution shift. It introduces diffusion-based transport as an alternative that avoids repeatedly estimating the emulator’s invariant measure.

  • Thermalization starts reverse diffusion at an adaptively selected noise level for the current emulator state, treating the state as a perturbation of equilibrium.Source: Figure 2 · p. 3
  • The goal is to generate typical trajectories over arbitrarily large horizons while keeping training and inference budgets decoupled from the horizon.Relying only on local Markov structure is insufficient because of inherent instability.Source: §2 · p. 3
  • Autoregressive modeling uses a learned one-step kernel to generate long trajectories, but its trajectory law can drift from the true law at a rate linear in the horizon.Source: §2 · p. 3
  • The generated emulator distribution can shift from the training distribution, undermining generalization and producing unstable rollouts.The discussion also notes that an emulator need not be irreducible or converge to a stationary measure in practice.Source: §2 · pp. 3–4
  • Adding invariant-measure regularization to training requires estimating the model’s invariant measure repeatedly, which can become prohibitively expensive when the spectral gap is small.Source: §2 · p. 4
  • Diffusion models learn reverse transport from noisy versions of the equilibrium distribution toward that distribution using estimated scores.The OU process moves the equilibrium law toward a Gaussian reference, while reverse diffusion approximately reconstructs the equilibrium samples.Source: §2 · p. 4

3. Related work

Related work addresses long-horizon instability through architectural, training, regularization, forcing, and decomposition strategies, while diffusion models have also been used for PDE refinement and turbulent-field generation. The paper distinguishes thermalization through its separately trained, adaptively applied denoiser and much longer rollouts.

  • Large-scale climate and weather autoregressive models succeed over short timescales but commonly remain unstable over long horizons because errors accumulate.Source: §3 · p. 4
  • Diffusion-based PDE work applies denoising refinement after autoregressive prediction, whereas this paper uses a denoiser trained separately from the emulator and applied adaptively at inference time.Source: §3 · p. 5
  • Prior stabilization methods include regularization, multi-timestep training, training noise, external forcing, and linear–nonlinear decomposition.Source: §3 · p. 5

4. Methodology

The method combines a pretrained autoregressive emulator with a separately trained diffusion model that adaptively denoises rollout states toward the stationary distribution. Thermalization approximates an ideal transport back to equilibrium while using predicted noise levels to decide when and how much denoising to apply.

Model construction

  • The emulator predicts residual state changes with a small stochastic perturbation, is trained through four autoregressive steps, and uses a diffusion denoiser trained on trajectory snapshots.The emulator update is ˆx_t+1 = ˆx_t + Φ_β(ˆx_t) + τn with τ = 1e−5, while the thermalizer uses the first snapshot from each training trajectory.Source: §4.1 · p. 5, §4.2 · pp. 7–8

Thermalizer

  • Thermalization couples a pretrained local-dynamics emulator with a separately trained diffusion model at inference time to stabilize autoregressive rollouts.The approach is modular: the emulator and diffusion model are trained separately and combined only during inference.Source: §4.2 · pp. 5, 7

Idealized stabilization

  • The ideal stabilization maps the emulator’s current law back to the invariant measure through an optimal transport map, but estimating that map is infeasible in high dimensions.It requires estimating the rollout law from a single state or many expensive emulator runs, and optimal transport suffers from the curse of dimensionality.Source: §4.2 · p. 7

Diffusion-based transport

  • Reverse diffusion approximates the ideal transport by starting from the current state at an adaptively chosen noise level and denoising it toward equilibrium.The construction treats rollout states as perturbations of equilibrium states and uses reverse diffusion to approximate the corresponding transport.Source: §4.2 · p. 7

Contraction guarantee

  • With correctly estimated noise and accurate denoising, the thermalizer contracts the rollout distribution toward equilibrium, with stronger contraction possible for nearly Gaussian perturbations.Even approximate transports can stabilize rollouts when they sufficiently mitigate distribution shift.Source: §4.2 · p. 7

Adaptive implementation

  • At each timestep, a classifier predicts the state’s noise level, and denoising is triggered only above a threshold using s_therm = s_pred − s_stop steps.The predicted noise level determines both whether thermalization occurs and how many denoising steps are applied.Source: §4.2 · pp. 7–8

5. Experiments

Experiments on Kolmogorov flow and quasi-geostrophic turbulence show that thermalization stabilizes autoregressive emulator rollouts while largely preserving temporal and spectral behavior.

Kolmogorov flow

  • In Kolmogorov flow, thermalized trajectories remain consistent with the numerical model after unthermalized emulator errors cause divergence by 5,000 steps.Kinetic-energy spectra across 40 randomly initialized trajectories show the stabilizing effect is consistent.Source: §5.1 · p. 8

Quasi-geostrophic turbulence

  • In quasi-geostrophic turbulence, unthermalized trajectories diverge after approximately 4,000 timesteps, while thermalized flows remain stable.The comparison uses 20 trajectories and radially averaged kinetic-energy spectra weighted across the two layers.Source: §5.2 · p. 8

Stability and efficiency

  • All thermalized trajectories remain stable for 10,000 steps, whereas nearly all emulator trajectories become unstable after about 2,000 steps.The exception is the Kolmogorov DRN, for which many trajectories remain stable longer without thermalization.Source: §5.2 · pp. 8–9
  • Approximately 140 thermalization steps stabilize each 10,000-step QG Unet trajectory on average.The thermalizer’s noise-classifying head costs 60% of an equivalent Unet emulator step, while full denoising is applied only when needed.Source: §5.2 · p. 9

Long-rollout fidelity

  • After 1e5 steps, thermalized trajectories remain stabilized, with no significant autocorrelation change but consistently lower QG kinetic energy than the numerical model.The long-rollout kinetic-energy spectra show an offset for QG fields despite successful stabilization.Source: §5.2 · p. 9, Figure 6 · p. 9

6. Conclusion

The thermalizer uses a separately trained diffusion model of the stationary data distribution to stabilize autoregressive surrogate rollouts. Across two high-dimensional turbulent systems, adaptive denoising enabled stable predictions over 1e5 emulator steps, while Gaussian perturbation assumptions and hyperparameter tuning remain limitations.

Conclusion

  • The thermalizer constrains autoregressive trajectories using transport maps from Gaussian perturbations back toward the stationary distribution, mitigating accumulated rollout error.The emulator remains a pretrained black box, while the equilibrium diffusion model is trained separately with standard pipelines.Source: §6 · p. 9
  • Stable predictions over 1e5 emulator steps were demonstrated for forced 2D Navier–Stokes and 2-layer quasi-geostrophic turbulent flows.Adaptive denoising reduces computational cost and disruption of temporal dynamics.Source: §6 · p. 9

Limitations and Future Work

  • The framework assumes isotropic Gaussian perturbations, and performance depends strongly on brute-force tuning of sinit and sstop for each emulator–thermalizer pair.Some configurations remain stable but produce imperfect kinetic-energy spectra at large timesteps.Source: §6 · p. 9

Conclusion

  • The evaluation covers two high-dimensional dynamical systems exhibiting chaotic, turbulent behavior relevant to practical applications.Source: A. Dynamical systems · p. 14

A.1. Kolmogorov flow

The appendix describes the forced 2D Navier–Stokes and quasi-geostrophic systems used for evaluation and the autoregressive emulator losses. Residual emulators, particularly residual loss, performed best on validation loss and 200-step rollout MSE.

Kolmogorov flow

  • The forced Navier–Stokes setup uses constant sinusoidal forcing, p = 1, ν = 0.001, and Reynolds number Re = 10,000.Source: §A.1 · p. 14
  • The Navier–Stokes simulations use 512 × 512 grids and downsample vorticity fields to 64 × 64 for emulator training and testing.Source: §A.1 · p. 14

Quasi-geostrophic flow

  • The second evaluation system is a two-layer quasi-geostrophic flow whose prognostic variable is potential vorticity.Its numerical solution uses a pseudo-spectral method, third-order Adams–Bashforth stepping, periodic boundaries, and 64 × 64 downsampled fields.Source: §A.2 · p. 14

Loss function

  • Residual emulators, especially the residual loss in equation 14, achieved better validation loss and MSE accuracy over 200-step rollouts than the other tested losses.The residual formulation predicts the next state as the current state plus an emulator output and stochastic component.Source: §B.1 · p. 15

B.2. Architecture

The emulator comparison contrasts a large Unet with a smaller Dilated ResNet. Although the Unet has lower short-term MSE, the DRN has better long-term stability.

Architecture

  • The Unet uses three downsampling layers with doubled filter counts and contains 48M parameters.Source: §B.2 · p. 15
  • The DRN exhibits better long-term stability than the Unet despite having a higher initial MSE and a much smaller model.Source: Figure 7 · p. 15
  • The DRN is built from dilated blocks using convolutional dilation pattern [1, 2, 4, 8, 4, 2, 1].Source: §B.2 · p. 15

B.3. Optimization

Short-term validation improvements do not reliably predict long-term rollout stability. In the architecture comparison, the larger Unet has lower short-term MSE, yet trajectories quickly become unstable.

Optimization

  • The emulator training uses AdamW with learning rate 5e−4, 32-sample minibatches, 12 epochs, and a Δt = 2 horizon selected using 200-step rollout MSE.Source: §B.3 · p. 16
  • The larger Unet achieves significantly lower short-term MSE, but all trajectories quickly become unstable.Long-term stability also varies substantially with network weight initialization for fixed data, architecture, and optimizer.Source: §B.3 · p. 16

C.1. Architecture

The thermalizer uses a Unet-based architecture with a noise-classifying head and is trained across the full noise range. Its adaptive denoising is controlled by thresholds that determine when corrections begin and end.

C.1. Architecture

  • The thermalizer augments a Unet with a noise-classifying head built from the lowest-dimensional representation and two additional convolutional layers.The added convolutions use kernel size 3 and stride 1.Source: §C.1 · p. 16
  • Training uses mini-batches of 64, AdamW with learning rate 2e−5, and 35 epochs.Training only on low noise levels degraded performance, so noise levels are sampled uniformly across the full range.Source: §C.2 · p. 16
  • Adaptive thermalization uses sinit to start corrections and sstop to set the lowest noise level reached by denoising.The settings are selected through a grid search over the stated threshold ranges.Source: §C.3 · p. 16

D. Inspecting a single thermalization step

A single thermalization step delays instability in a QG emulator rollout. The intervention denoises a snapshot before continuing the same autoregressive rollout, reducing the later buildup of unstable structure.

D. Inspecting a single thermalization step

  • A single thermalization step delays the development of instability in a QG emulator rollout.The comparison starts both rollouts from the same snapshot, with the thermalized case applying one correction beforehand.Source: Figure 8 · p. 17
  • The experiment applies four denoising steps after the predicted noise reaches sinit = 10, using sstop = 6.The thermalization count is stherm = sinit − sstop = 4.Source: D. Inspecting a single thermalization… · p. 17
  • Without thermalization, high-frequency signal accumulates and field magnitudes grow sooner than in the thermalized rollout.The difference is visible in both the potential-vorticity fields and their colorbars.Source: D. Inspecting a single thermalization… · p. 17

E. Additional results and figures

Additional diagnostics show that thermalization keeps long rollouts at low predicted noise while applying corrections unevenly across trajectories. End-state visualizations compare numerical, emulator, and thermalized fields for Kolmogorov and QG systems.

E. Additional results and figures

Loading 2503.18731v2…