Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Authors: Christian Pedersen
Laure Zanna
Joan Bruna
Venue: ICML 2025 poster (2025)
Keywords: Diffusion models
Physics
Fluid Dynamics
Turbulence
Emulators
Surrogate Models
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: Diffusion models can mitigate the error accumulation that occurs over many autoregressive steps of an approximate model for fluid flows, allowing for significantly longer stable rollouts

Abstract: Autoregressive surrogate models (or *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 *thermalization*. Thermalizing an emulator rollout is shown to extend the time horizon of stable predictions by two orders of magnitude in complex systems exhibiting turbulent and chaotic behavior, opening up a novel application of diffusion models in the context of neural emulation.

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Christian Pedersen Laure Zanna 12
Joan Bruna" 12
Abstract
cost of numerical methods becomes extremely large. Tur- bulent fluid flows, for example, involve dynamical coupling across length scales spanning orders of magnitude, and accu- rately modelling the full dynamical range is often intractable for many important applications.
Autoregressive surrogate models (or emulators) of spatiotemporal systems provide an avenue for fast, approximate predictions, with broad applica- tions 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 tra- jectories. In essence, emulators operate out of distribution, and controlling the online distribu- tion quickly becomes intractable in large-scale settings. To address this fundamental issue, and focusing on sime-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 autore- gressive emulator rollouts by applying on-the-fly denoising during inference, process we call ther- malization. Thermalizing an emulator rollout is shown to extend the time horizon of stable pre- dictions by two orders of magnitude in complex systems exhibiting turbulent and chaotic behav- ior, opening up a novel application of diffusion models in the context of neural emulation.
Deep learning (DL) methods have been applied in vari- ous ways to mitigate the cost of obtaining PDE solutions. One approach is to couple a coarse-resolution numerical solver with a DL. component, either as a learned correction (Kochkov et al., 2021) or as an additional forcing term in the PDE, formulated using the large-eddy simulation framework (Duraisamy al., 2019). DL models have also been applied surrogate models, emulators, replacing the numerical scheme entirely, and leveraging the speed of graphical pro- cessing units (GPUs) to provide fast approximate solutions Baanchez-Gonzalee et al., 2020; Stachenfeld et al., 2021). Such emulators are being applied in the context of weather (Pathak et al., 2022: Lam et al., 2023; Bi al., 2023) and climate modelling (Kochkov et al., 2024; Watt-Meyer et al., 2023; Subel & Zanna, 2024), with computational speedups of up to five orders of magnitude (Kurth et al., 2023).
The standard approach to constructing a neural emulator is to predict the state of the system at some future time as a function of the current timestep. Simulated trajecto- ries are then generated by a rollout of many autoregressive passes, where the predicted state is fed back into the emu- lator. This framework suffers from instabilities over long timescales, accumulation of errors leads to a drift of the emulator trajectory away from the truth (Chattopadhyay et al., 2023; Bonavita, 2024; Parthipan et al., 2024; Bach et al., 2024). Eventually, the drift becomes large enough that the state is out of distribution, and trajectories exponentially diverge. Some proposed modifications to improving the stability of rollouts include the addition of noise train- ing data (Stachenfeld et al., 2021), training over multiple consecutive timesteps (Vlachas et al., 2020; Keisler, 2022), predicting multiple timesteps at once (Brandstetter et al., 2022), including additional input channels representing ex- ternal forcings -Wate et al., 2023; Subel Zanna, 2024), and adding iterative refinement steps focusing on different frequency components (Lippe et al., 2023).
1. Introduction
The modeling of dynamical systems is a cornerstone task in the physical sciences and engineering, with applications across weather (Pathak et al., 2022; Watt-Meyer et al., 2023; Lam et al., 2023; Lang et al., 2024) and climate modeling (Kochkov et al., 2024; Subel & Zanna, 2024; Lupin-Jimenez et al., 2025). The standard approach to modeling these sys- tems to solve the underlying partial differential equation (PDE) describing the system using numerical methods. In the case of high-dimensional systems, the computational
'Courant Institute of Mathematical Sciences, New York Univer- sity, USA *Center for Data Science, New York University, USA. Correspondence to: Joan Bruna <bruna@ cims.nyu.edu>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
The situation is further complicated by the fact that the sys- tems that are being modelled are generally chaotic. Indeed,

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
t- 0s, step-0
t-1s 1s, step-500
t- Ss, step step-2500
t= stap-3500
t-10s, step-5000
t-20s, stap-9999
Numerical model
Emulator
Thermalized
10
10
Einulattr
Thermalised
10P 161
16!
16P
10F
10F
107
10F
107
16F
16
Kinetic energy spectra from different trajectories
Figure Top rows: Vorticity fields from numerical model, emulated, thermalized trajectories Kolmogorov flow over trajectory of 10, 000 steps. Bottom row: Radially averaged kinetic energy spectra each timestep, for 40 randomly initialized trajectories.
for such systems with intrinsically diverging trajectories, recurrent architectures are unable to capture the dynamics while maintaining stability during training (Mikhaeil et al., 2021). Recent work has attempted to incorporate knowledge of the chaotic system's 's invariant measure into training to stabilize predictions and improve the modeling of chaotic dynamics (Li et al., 2022; Jiang al., 2023; Schiff al., 2024).
emulated trajectory from diverging away from the training distribution, thereby maintaining stability over arbitrarily long time horizons; see Figure 2. We refer to this process as thermalization, as the system is relaxed back towards a stationary measure. One of the key advantages of our scheme relative to prior work the modularity of training: we only require pretrained black box emulator and a sep- arately trained diffusion model for the stationary measure of the dynamics. These two models inference time, and, as shown in the numerical experiments, require only small overhead in inference computational cost.
only combined at
In this work, we propose an alternative approach, based on diffusion models, framework developed for generative modeling of images (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020), to stabilize autoregresssive emulator rollouts of chaotic systems. Diffusion models im- plement reversible measure transport between the data dis- tribution and reference distribution, typically the standard Gaussian measure, by learning one-parameter family of de- noising operators, resulting in an efficient non-equilibrium sampling scheme. As errors accumulate during emulator rollouts, the state of the system moves to regions of low probability, which are then transported back to equilibrium by querying the reverse diffusion models with an appropri- ately tuned noise level. Importantly, the reverse diffusion defines coupling between corrupted and equilibrium dis- tributions that respects the temporal dynamics. This can be thought of as a form of on- on-the-ff the-fly fly denoising that prevents the
Main Contributions. We present novel application of diffusion models in the context of using neural networks to approximate the solutions to chaotic PDEs. We demon- strate the thermalizer method on two high-dimensional tur- bulent systems show that arbitrarily long rollouts can be achieved when an unstable neural emulator is coupled with thermalizer.
2. Preliminaries
Problem Setup. We consider data of the form (xx) (eN as the time discretization of an underlying dynamical system (xt) X2 Rd.. solving 0,x, = F(x;), where F
2

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
to be insufficient in practice due inherent instability; the key scientific question of this work is, thus, how to over- come this instability using scalable training and inference algorithms.
Unstability of autoregressive modeling, Given observed trajectories {(x],.... .x,,,,,)))))) xx,) Ki)) where (xi..... .f..)...))))) ~ TL are independently drawn, one can estimate model Q of the one-step transitions Q. A standard choice is to as- sume a Gaussian transition kernel of the form 2(-ly) (y)((y), La), where Pg is a generic parametric map, e.g. neural network, so that the associated MLE corresponds to the least-square objective ming A
n E Ej<z || Do (o(x)))l This Gaussian model thus corresponds to a general (non-linear) diffusion in continuous time, where dx, F(x) + v2odB,, which in the limit of 0 recovers deterministic dynamics. Note that, in practice, one can replace the one-step transitions with s-step transitions without loss of generality so that the emulated model ad- vances k times faster than the numerical timestep.
Figure 2. Schematic view of the thermalisation: in dashed grey, two typical trajectories from Ty, that preserve the equilibrium distribution js, here idealised in terms of its typical set. In black, an emulated trajectory, that eventually drifts away from the typical In purple, the reverse diffusion, which defines approximate transport between jis. the law of Xt. and μe, resulting in the red "thermalized" trajectory.
R - R is an unknown, lime-independent, and generally non-linear operator.
Once trained, one can sample arbitrarily long trajectories by querying Q in an autoregressive fashion: for each T 0, we denote by कr the joint law of (xo..... xT),
where Xp ~ To
To account for the uncertainty in initial conditions, the time and space discretization errors, as well as the chaotic na- ture of some dynamics, it is convenient to adopt prob- abilistic description of the problem: we view (Xx)LeN high-dimensional Markovian stochastic process и. By the time-homogeneous Markovian assumption, its joint xz) satisfies the recurrence equation law := T x(+1) = @(x(+1 (x,. x²) for any t>o 0, where Q is the Markov kernel encoding the condi- tional distribution of x,+1 given X₃. We denote by Pe the marginal distribution of X2. We also assume that this ker- nel is irreducible, which implies the existence of unique stationary measure := Poo E P (()d) such that μ μQ
and x,-1 |x ((x(-1/xx). We easily verify that
KL ||T))
Σ 1ST
By KL (2(+)))1((1))) (I)
For large T, since μe + by ergodicity, the above identity directly yields
KI(TT| = By KT. (o(
(2)
This confirms the intuition that the law of the trajectories will drift apart, as measured with cross-entropy, at rate which is linear in the horizon. Note that RHS is precisely the test error of the regression objective - in accordance with Girsanov's theorem in the continuous-timg limit.
Faced with this Markov structural assumption, it is thus tempting to learn the system by focusing solely on the short- time conditional kernel Q, given its small memory foot- print, and the fact that it is sufficient reconstruct the full joint probability distribution. In other words, leveraging the Markov kernel enables arbitrarily long generations at inference time with a constant training budget.
While the above KL control measures whether typical sam- ples from ग. are typical under our model म.. the generative setting, where we sample from â,, corresponds instead to the reverse KI. divergence. Assuming that Q is also irre- ducible, with invariant measure fú, inverting the roles of Q and @ now yields
In this work we are particularly concerned with the long- time behavior of the system: given a possibly large time horizon T' (c.g., of the order of several decades in climate models compared to days for weather forecasts), our goal is to build scalable generative model for trajectories, that is, to efficiently sample a typical trajectory from nT, in such a way that training and inference budget remains decou- pled; ie, we can produce arbitrarily long trajectories with a training cost independent of T. As one could anticipate, solely relying on the local Markovian structure turns out
KL(FrlIaT) ~ TEy-RKL ( (Q(1)))1(()))
(3)
Crucially, we now have shift between the test error distribu- tion ğ and the training distribution μ4, leading to dramatic impact in the generalisation guarantees of the model and consistent with the observed instability of emulator rollouts. Note though that the ergodicity assumption on Q is merely for convenience: in practice, it may be that Q is not irre- ducible, e.g., when d² = O, failing to converge towards
3

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
stationary measure, or, using again the continuous-time formulation, when the associated Fokker-Plank equation 0,p4 V (-Fu,) + Aur admits no equilibrium, due to the fact that the irrotational component of F does not have enough decay at infinity.
transport equation can be formally reversed: if ig solves 0,ē, = V ((-x 2V log =)=)));) + Ae,, with in vs. then i/5 = Mo. In other words, the reverse diffusion
dX, = (X, + 2V us-
v (X (d) + V2dB,
(5)
with Xo ~ Ya. satisfies Xs ~ vo up to an exponentially small error (due to the difference [ hr) = O(e ))).
Mitigating distribution shifts. In the face of the previous discussion, it is thus tempting to control the extent of distri- bution shift by adding a regularisation term in the training objective. Given a model Qg for the one-step transitions, and assuming again irreducibility, it has an associated in- variant measure 140, characterized as the (unique) Perron eigenfunction of Qa. Provided one can access (or estimate) the true invariant measure şl of the system, one can thus consider a learning objective that combines both sources of error:
The above procedure defines a non-equilibrium sampling scheme that requires access to the scores V log Vs along the OU semigroup. As it turns out, such scores can be efficiently estimated from original samples of μ via a denoising ob- jective. Indeed, defining the Denoising Oracle D a) := E[xly], where y = oz, X ~ μ and z ~ Yas Tweedie's formula (Robbins, 1992), relates the score to the Denois- ing oracle via an explicit affine relationship, V Vlog (x())) = Dx, as)) for explicit values of a, and 3.. Finally. D(-, or) can be efficiently estimated from samples
(x -x +
T(0)= = Fyy- D(g(-)y)) (()(())) + AD(H,
(4)
using the MSE variational characterisation of the posterior
Here, we keep the presentation informal, and consider generic (and possibly distinct) probability metrics D, D, as long as they admit Monte-Carlo estimators. The important aspect of Eq (4), however, is the presence of pe, which is only known implicitly as the invariant measure of the kernel Qo. While some prior works, e.g. (Schiff et al., 2024), have explored such learning objective in physical applications, it presents an important computational challenge, since it requires estimating the invariant measure μs each time 0 is updated. In challenging situations where Q2 does not have a substantial spectral gap, obtaining (1g from Qo may become prohibitively expensive. In the next section, we introduce alternative framework to address such distri- bution shift, which also relies on having a model for the invariant measure μ of the true system but crucially avoids the computation of μıg.
mean: D arg miR I Ex- [| F(X + t [)] Here a, and B, are as defined in (Ho et al., 2020).
The standard Denoising Diffusion Probabilistic Model (DDPM) framework (Ho et al., 2020) introduced an effi- cient implementation of the above scheme, by considering appropriate discretization of the noise levels, and where the denoiser D D is represented by neural network with parameters db In the generative modeling case, the noise scheduling spans the full range from the unperturbed data distribution at s = 0, to pure Gaussian noise at s = S (where is generally in the range z 1000). As we shall see next, we will introduce a different use of diffusion models. Using the available data, we can estimate a diffusion model for the equilibrium distribution μu, that we will use to build coupling between j. a perturbation of the equilibrium, and μa. This coupling can be interleaved with timestepping performed by the emulator, constraining the emulated tra- jectory to stay in regions of state space consistent with the training data, a process we call thermalization. Compared to the generative modeling approach, we are only interested in small amounts of denoising at the very low noise end of the scheduler, i.e. at low
Estimating the Invariant Measure using Diffusion Mod- els. Diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) have been tremendously successful in the field of generative modeling. The frame- work is based on modeling the inverse of an iterative noising process, which is best thought of as the discretisation of underlying continuous diffusion process. Let us recall the main ideas which will be needed to define our method.
3. Related work
Given Nu ~ Po= == 14, we consider the Ornstein-Ulhen Ulhenbeck (OU) process ddX, = -X,ds + B2, where B, is the standard Brownian motion. The law v of X, solves the associated Fokker-Plank equation 0,v₂ = V (IV:) + Av.. which can also be written the transport equation 0,v,- V ((x + V log v.)..). It is well known that the law Ms converges exponentially fast (in KL) to the standard Gaus- sian measure Ta := (0, Ia), so Vi s a for 5 large enough. Now, by changing the sign of the velocity field and introducing again the diffusive Laplacian term, this
Large-scale autoregressive models: In the context of cli- mate and weather modelling, several large scale autoregres- sive models have been developed (Pathak et al., 2022: Lam et al., 2023; Bi et al., 2023; Lang et al., 2024; Kochkov et al., 2024), and are indeed being deployed for commercial and public use. Despite the tremendous success of these models over short timescales, long-term instability due to error accumulation over large numbers of timesteps is still a common problem.
4

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Diffusion models for PDEs: In (Lippe et al., 2023), an au- toregressive model is used to predict the next timestep, and then small number of "refinement" steps are applied to the prediction. These refinement steps are trained on denoising loss, and can therefore be interpreted as DDPM denoising steps, analogous to our thermalization steps. Therefore our work is similar in nature, however with the key differences that their denoiser is a conditional model, and that our de- noising model is constructed separately to the emulator and applied adaptively at inference time. Additionally, we in- vestigate rollouts over many more autoregressive passes, up to Le steps. Diffusion models have also been applied to generating realizations of turbulent fields (Lienen et al., 2023), and emulation in the form of autoregressive condi- tional models (Kohl al., 2023; Price et al., 2023), where the probabilistic nature of predictions is able to improve modeling of chaotic dynamics. To reduce the computa- tional cost of generating next step predictions, (Gao et al., 2024) performs the conditional generation in learned, low- dimensional latent space, and (Shehata et al., 2025) develops a more efficient conditional generation algorithm. Diffusion models have also been applied to large-scale weather mod- elling, both for the purposes of emulation (Price et al., 2023; Rühling Cachay et al., 2024) and for statistical downscaling (Mardani et al., 2023; Wan et al., 2023).
neural network parameters / are optimized by minimizing MSE loss function:
. La
11 (0((()) - (x(+1 - x() +nn 11?, (6)
N R= 1=
where trajectories start from a simulation snapshot (xo- Xo), L represents the length of training trajectory (through- out this paper, we use L = 4), k indexes an individual training trajectory, and the loss is averaged over a total of N training trajectories. Gradients are backpropagated through the full L timesteps, as done in (Brandstetter et al., 2022; List et al., 2024). We note that the separation be- tween training snapshots, At, is in general larger than the numerical timestep. The neural network $ is implemented using two choices: a U-Net (Unet) style architecture (Ron- neberger et al., 2015), where we set the number of filters and downsampling layers such that the model has 48M pa- rameters, and a Dilated ResNet (DRN) (Stachenfeld et al., 2021) with 0.5M parameters. Further details on architecture and optimization can be found in Appendix B. To build a training and test set, we use numerical simulations to gen- erate total of N = 500, 000 trajectories. We use 450, 000 of these for training and the remaining 50, 000 for vali- dation and testing. The resulting network thus defines a point-estimate, associated with the local transition kernel ((xly) = 8(x - (y + i degenerate version of the Gaussian transition kernel discussed in Section 2. Addition- ally, we implement stochastic emulator by including small amounts of noise at each timestep, during both training and inference. The solution manifold for dynamical systems is commonly a fractal, leading to singular invariant measure. The problem of learning such a solution manifold can be facilitated by applying small perturbations to the training data (Zeeman, 1988: Gritsun & Branstator, 2007; Baldovin et al., 2022). Note that this is naturally achieved in the case of the diffusion model via use of the denoising score loss.
Stabilizing rollouts: A central problem in fully learned au- toregressive simulators is the accumulation of errors and re- sultant instabilities. Several regularisation terms have been proposed (Chattopadhyay et al., 2023; Schiff et al., 2024: Guan et al., 2024). (Brandstetter et al., 2022) used training over multiple consecutive timesteps to improve stability, while the addition of training noise has also been demon- stated to improve stability when using a single timestep prediction (Stachenfeld et al., 2021). The study in (List et al., 2024) investigated an approach to multi-timestep training without propagating gradients. The inclusion of forcing from external systems has been shown to stabilize extremely long rollouts in emulation of ocean models (Watt- Meyer et al., 2023; Subel & Zanna, 2024). Decomposing the system into linear and non-linear components was also demonstrated to improve stability in simple 1D non-linear systems (Linot et al., 2023).
4.2. Thermalizer
We now introduce a procedure stabilize trajectories that operates inference time, assuming a pretrained model for local dynamics @ given by the previous emulator model, and pretrained diffusion model for μ. consisting of the denoiser D
4. Methodology
4.1. Emulator
For each time i, let fúe denote the law of the current gen- erated field x,. Unstability is to be expected as soon as this law drifts sufficiently far away from p., that we re- call corresponds to the distribution of y used to train the model for local dynamics Q(-ly). For instance, one can use the 2- 2-Wssersteein distance W2 μ) to quantify this drift. This distance is interesting because it also preseribes an optimal coupling I, between fue and μ minimizing the transportation cost: W2( = Exxxy)), ||x y
We first describe the framework for the construction of the emulator. We represent the neural-network model as &... We found the best performance when using an emulator to model the residuals between two timesteps, i.e., a predicted state can be denoted & = & + (x,) + 7n where the final term represents the stochastic component, so n ~ A (0, I) and we set = le-5 throughout the paper. The
5

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
1-0 days, stepep
ta 0d days, step=500
1=208 days, step=1000
312 days, step- 1500
t=d17 days, Step 2200=2000
t=22003 days. step-9999
Num modal
Emnzlatoo
Thermalized
101
161 Numerica model
Emulator
10 Thermalized
Kinetic energy spectra from different trajectories
Figure Top rows: QG trajectories over 10, 000 steps numerical model (top), neural network emulator (second row), and the thermalized neural emulator (third row). We show the potential vorticity, in the top layer of the 2-layer QG system. Bottom Kinetic energy spectra for each timestep. for 20 different trajectories. We show the layer thickness weikeess-weighed average of the spectra in both the upper and lower layers.
Kolmogorow: Unel
Kolmogorow: DRN
Kolmogoro
Unet DRN
10
MSE
10"
10
101
10:1
Thermalized
8000
10000
12000
Unat
10a
10a
MSE
10a
10-1
10a
10a
2000
4000
6000
se00
10000
2000
4000
6000 sooa 10000
2000
6000 8000
Emulator step
Emulator step
Thermalization steps
Figure Mean squared error between neural network trajectories and numerical model trajectories for different baseline architectures (Unet and Dilated ResNet), starting from the same initial conditions. We show 40 Kolmogorov flow trajectories in the top two pancls, and 20 QG flows in the lower two panels. Thermalized flows remain around a constant MSE long after non-thermalized flows have become unstable. The region where vectors decorrelated in time, but stable, is shown in shaded pink. On the right, show histograms for the total number thermalization steps along the full trajectories.
6

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
( Bxxy)-r x ylI where we recall that coupling I' E I1(1, (d2) is a joint distribution in R x gd such that its marginals are respectively μ1 and Pzz. More- over, under mild assumptions Chewi et al. (2024, Section 1.4), this optimal coupling is associated with optimal transport map Tz R - RR, such that (x, (x)) ~ r, In other words, the current field x, ~ fi can be mapped 'back' to the equilibrium distribution via (xx)). and this map min- imizes the expected distortion ||x: - Ti(x:)] )|² introduced in the trajectory.
Algorithm 1 Algorithm for thermalized trajectories
for = 1 to N do n -O x x,-1 + (x (x(x + rn
Spred = arg max D (xx)
if Spred > Sinit then
~ (0,
X, = Vannacs Xx +
a Sprced c
for s = Spred to Sstop do
(O,
This defines idealised stabilization method, which un- surprisingly is unfeasible to implement. First, it rests the ability to estimate the law B. from single realization X₃. While we have samples from μ in the training set, gen- erating independent samples from Re requires running the emulator multiple times at inference, which can be quickly prohibitively expensive. Even then, well-known that optimal transport suffers from the curse of dimensionality Weed & Bach (2019).
1 x, Jas
1-a. (1)
(xi)) + vB22
X1 VI-a, 1
end for end if end for
Langevin map indeed maps any initial measure to as + 0o, it suffers from two important limitations: first, its time to relaxation can be exponentially large in pres- ence of meta-stability and multiple modes, even for initial distributions jJ close equilibrium, affecting the resulting transportation cost. Next, while the score V log μ corre- sponds to the zero-noise limit of the DDPM objective, this score is only accurately learnt nearby the data distribution v, and we found unstable for our purposes. One can view our proposed transport scheme as non-equilibrium counterpart of the Langevin map that provides improved robustness.
Instead, we will leverage the diffusion model to provide an efficient approximation to these optimal transport maps. Our approximated transport will consist of running the re- verse diffusion (5) starting at a noise level = s(), that we adjust function of Xg, and X, = X,, resulting in the coupling (Xo, X.) ~ r., The rationale behind this choice of transport map is the following: one can view the samples X; ~ fe as a perturbation of 'correct' equilibrium states coming from μe; if these perturbations were isotropic Gaus- sian, then the law be would agree with V. for some noise level 8 = (f), and the reverse diffusion would implement a valid transport towards μ.. While this transport is not the op- timal transport, one has some control over the suboptimality gap Albergo et al. (2023b, Prop 3.1).
It remains to be discussed how to find suitable noise level s(),, given the current emulator state. To tackle this ques- tion, we modify the standard implementation of the DDPM framework, by adding classifier output such that the noise level s is predicted by the network, instead of being passed as an input parameter. We denote these two outputs using a superscript, where D
Such guarantees are no longer valid for generic, non- Gaussian perturbations. Instead, if one assumes an ideal choice of the noise level, given by st) = arg min, KL(selllss, and jr is the law of Xo conditioned on X., Xx ~ Pe. by the data patt-srocessing inequality one has
(1) (xx) E R represents the denoised field, and D (2) (xx) E RS represents the predicted categori- cal distribution over S noise levels. The details of how this additional classifier head are added to the Unet are given in Appendix C. The loss function we minimise
N
(L(A) & KL(Aa))us)) S KL(A, (2)) = KL(AallA). (7)
Ls
=
Σ N
1
ll s&c,
(x, + s,E,) 11?
Thus, assuming that the noise level is correctly estimated, and an accurate denoising model, the thermalizer step con- tracts towards equilibrium. Moreover, for nearly Gaussian perturbations, the LHS can be substantially smaller than the original emror. It is important to emphasize that even approx- imate transports can go long way to stabilize rollouts, as long as they can mitigate the distribution shift.
Σ Us log D (2) (x; ssci) (8)
where s is uniformly sampled from [1, x, is the training data, and ys E R is one-hot vector encoding the true noise level. To implement the thermalizer, we use the same Unet architecture as for the emulator, where the network now has 53M parameters, with the additional parameters constituting the clase-classifiing head. To train the thermalizer, we use the same N 450, 000 training set
is interesting to compare this scheme with the Langevin map, that runs the Fokker-Plank equation 0,2 = V (VVogg)) u))) + Av starting from un = f.. While the
7

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
as for the emulator, except we take only the first snapshot of each short trajectory to avoid training the thermalizer redundant data samples.
Autocorrelation
1.0
Kolmogorov
QG
Numerical model
0.8
Emilator
During an emulator rollout, at each timestep we run a for- ward pass of the classifie-compponett of the network to predict the noise level, such that for a given state vector x, the predicted noise level is Spred = arg max D %/) (x). If the predicted noise level exceeds a threshold, which we call Sinit, we then apply Stherm = Spred - Sstop steps of denoising to the state vector, where Spred the estimated noise level in the state vector, and Sstop is some minimum acceptable noise level, below which we do not thermalize. The steps for this process are detailed in Algorithm I.
Thermalized
0.6-
0.2
0.0
200
600
1000
Emulator step
5. Experimeenss
Figure 5. Autocorrelation over time for numerical model, and neu- ral network trajectories.
Additional visualizations including videos are available here.
represents a stratified fluid in rotating frame. This kind of turbulent dynamics is prevalent in climate science, and plays a central role in both atmospheric and ocean dynamics (Majda & Qi, 2018). The dataset construction procedure is the same as with Kolmogorov flow, using dataset sizes of N = 500, 000 fluid snapshots for both the emulator and thermalizer, with the only change being that now a fluid snapshot has 2 layers, and the emulator training timestep is decreased to At = 5öt. The model architectures are kept the same, except with an additional input channel to accom- modate the additional fluid layer. When training over L = 4 timesteps, we found that the Unet emulator trajectories went unstable in the region of ~ 2e to 5e timesteps, so to bring the emulator in line with our other baselines, we reduced the number of training timesteps to 2. For the DRN, we keep I, 4 with the Kolmogorov emulators. Details the numerical solver and configuration of the QG system are given in Appendix A.2.
5.1. Kolmogorov flow
To test our framework on turbulent flows, we use a fluid flow described by 2D incompressible Navier-Stokes equations with sinusoidal forcing, often referred as Kolmogorov flow, used in many works on DI. for turbulence modelling Kochkov et al., 2021; Lippe al., 2023; Schiff et al., 2024). The baseline emulator is con- structed using the loss in equation 6, where we set A/- = 251, where is the numerical model timestep, and the number of recurrent passes used during training L 4.
(Boffetta & Ecke, 2012:
In Figure 1 we compare flow trajectories for three models: in the top row, direct numerical simulation, in the second row, neural network emulator constructed as outlined in section 1, and in the third row, flow trajectory using the same emulator, but including applications of the thermalizer model described in 4.2. Each trajectory is initialized at a state obtained from the numerical model. Whilst initially, the flow fields are all consistent after 3500 passes, errors begin to accumulate in the emulator field, and instabilities become visible. By 5000 steps, the instabilities dominate and the trajectory has entirely diverged. The thermalized flow, however, remains stable and consistent with the nu- merical model, as a result of the small corrections applied dynamically during the emulator rollout. In the lowest row of Figure we show the kinetic energy spectra at each timestep, for 40 different flow trajectories, demonstrating that the stabilizing effect of the thermalizer is consistent across all randomly initialized trajectories.
In Figure 3, we compare trajectories for the numerical model, emulator and thermalized trajectories. As with the Kolmogorov flow, after approximately 4000 timesteps the un-thermalized trajectories accumulate sufficient error di- verge, but the thermalized flows remain stable. In the lower panel we show the radially averaged kinetic energy spectra for 20 trajectories at each timestep. This is the weighted average of the kinetic energy spectra in the upper and lower layers, weighted by the layer thickness.
In Figure 4, we show the MSE with respect the numer- ical model, for different neural network models and flow configurations. We observe that initially the MSE increases as trajectories diverge over the first few hundred steps, and settle to approximately unity MSE. In this regime, shaded in pink, the state vectors are stable but decorrelated. States with MSE>> > are unstable. Over x 2000 timesteps emu- lator trajectories all go unstable (with the exception of the
5.2. Quasi-geostrophic turbulence
As an additional test case, repeat the experiment 2- layer quasi-geostrophic (QG) turbulent system. This system
8

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
KE spectra after 100,000 steps
6. Conclusion
10'
Kinetic Energy
Numerical model
We introduced the thermalizer, an algorithm for stabilis- ing autoregressive surrogate models leveraging a pretrained diffusion model of the stationary data distribution. This diffusion model provides family of transport maps from Gaussian perturbations of the stationary measure back to its equilibrium, which are shown to be robust to non-Gaussian errors. As a result, this diffusion model can be deployed at inference time to constrain trajectories to stay along the solution manifold, mitigating the accumulation of error due to autoregression of an approximate timestepping model. An appealing aspect of the model is the lightweight training: we treat the emulator the diffusion model of the equilibrium separately, using standard pipelines.
Thermalized
10f
10a
Kolmogorov
10
10
10
Kinetic Energy
10t-
QG
10
10-5
10-1
pre-trained black-box, and train
Wavenumberk k
Figure 6. Radially-averaged kinetic energy for Kolmogorov (top)
and QG (bottom) trajectories from a thermalized Unet emulator after lee timesteps.
A crucial component of the thermalizer algorithm is the adaptive denoising steps, which serve to minimize both the computational cost of the algorithm, and disruption of the temporal dynamics of the flow trajectories. We demon- strate this approach on two high-dimensional turbulent sys- tems, forced 2D Navier-Stokes flow, and 2-layer quasi- geostrophic turbulent flow, enabling stable predictions over les emulator steps.
Kolmogorov DRN, where many trajectories stay stable for much longer). However all thermalized trajectories stay stable for the full 10, 000 steps, due to the corrective steps applied during the rollout. In the right panels, we show the number of thermalization steps applied along the entire trajectory for all models and flow configurations. In the case of the QG Unet, the thermalization is highly efficient, with just ~ 140 thermalization steps required on average to stabilise each 10, 000 step trajectory. In terms of computa- tional cost, we run a forward pass of the blase-classifying head of the thermalizer at every timestep, which is 60% the cost of Unet emulator step, as we are only activat- ing the downsampling and noise-classffying components of the equivalent-sizee thermalizer Unet. Additional forward passes of the full thermalizer network are only run in the case where Spred > Smitt so given an accurate noise clas- sifier, the algorithm is inherently efficient and only applies corrective denoising steps when necessary.
Limitations and Future Work. The main limitation of our current framework is the underlying (isotropic) Gaus- sianity assumption on the perturbations. Even though our experiments demonstrate some form of robustness to such misspecification, this comes at the expense of careful choice of hyperparameters. In particular, performance is strongly dependent on the setting of Sinit and Sstop- which must be individually tuned for each combination of emu- lator and thermalizer via brute force search. Also indeed some configurations remain stable, however with imperfect kinetic energy spectra at large timestep, as shown in Fig- ure 6. We suspect that such limitations could be addressed by replacing Gaussian diffusion models with a base mea- more adapted the errors introduced by the emulator; stochastic interpolants aka flow matching (Albergo et al., 2023a; Lipman et al., 2022) provide such flexibility. In that respect, an interesting theoretical question is to understand the robustness of the OU-based coupling to non non-Gaussian perturbations, beyond the data-processing inequality. Next, we focused in the time-homogeneous setting, owing to its importance in applications and the particularly lightweight implementation of the thermalizer. The next natural step is to consider autonomous systems, eg with a forcing term having some temporal periodicity. Finally, another impor- tant future direction is to consider other domains where long autoregressive rollouts are often used - perhaps even I.L.s??
Given the fact that thermalizing steps induce some small distortion in the field, it is important consider the effect the temporal consistency of the flow trajectories. In Figure 5, we show the autocorrelation for both Kolmogorov and QG fields, averaged across all 40 and 20 trajectories respectively. We see no significant change in the autocorrelation between the numerical model, and the thermalized flow, indicating that the temporal disruption caused by the thermalization steps minimal. Finally, in figure 6, we show kinetic energy spectra for Kolmogorov and QG thermalized flows after 5 steps, where the spectra demonstrate that all trajectories are successfully stabilised over this long rollout. However we note that there is some offset in the KE spectra of the QG fields, which show consistently less kinetic energy than in the numerical model.

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Acknowledgements
Brandstetter, J., Worrall, D., and Welling, M. Message pass- ing neural pde solvers. arXiv preprint arXiv:2202.03376. 2022.
The authors would like to thank Fabrizio Falasca, Pavel Perezhogin, Stephane Mallat, Etienne Lempereur, Freddy Bouchet, Eric Vanden-Eijnden and Edouard Oyallon for many valuable discussions. This project is supported by Schmidt Sciences, LLC.
Chattopadhyay, A., Sun, Y. Q., and Hassanzadeh, Chal- lenges learning multi-scale dynamics with Al weather models: Implications for stability and one solution. arXiv e-prints, arXiv:2304.0702) April 2023. doi: 10448500/axxv.22440070999
Impact Statement
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
Chewi, S., Niles-Weed, J., and Rigollet, P. Statistical opti- mal transport. arXiv preprint wxxiv:2407.88663 2024.
Dresdner, G., Kochkov, D., Norgaard, P., Zepeda-Núñezz, L.,,
Smith, J. A., Brenner, M. P., and Hoyer, S. Learning to correct spectral methods for simulating turbulent flows. 2022. doi: 10.48550/ARXIV:2227.. 207.000 URL https: // //arxiv. org.iiiorg//bae wttpivv.rrggbb//222//// 00555..
References
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 2023a.
Duraisamy, K., laccarino, G., and Xiao, H. Tur- bulence modeling in the age of data. Annual Review of Fluid Mechanics, 51(1):357-31) Jan- uary 2019. ISSN 1545-4479. doi: 10.1146/ annurev-fluide hitiiv--fiiid-00011800065477 URL http://dxx... doi. org/10. 1146/ htt4://aammurr----------011-0-00
Albergo, M. S., Goldstein, M., Boffi, N. M., Ranganath, R., and Vanden-Eijnden, E. Stochastic interpolants with data- dependent couplings. arXiv preprint arXiv:2310.03725, 2023b.
Gao, H., Kaltenbach, S., and Koumoutsakos, P. Gen- erative learning for forecasting the dynamics of high- dimensional complex systems. Nature Communicaa- tions, 15(1), October 2024. ISSN 2041-1723. doi: 10. 10.1038//411477-044433166633 URL http::/dp.dde org/10. 103./ 110.1///1111110-2-----55-------1
Bach, E., Crisan, D., and Ghil, M. Forecast error growth: A dynamic-stochant model. arXiv e-prints, art. arXiv:2411 1.06623,, November 2024. doi: 10.48550/ arXiv.2411.06623.
Baldovin, M., Caprini, L., Puglisi, A., Sarracino, A., and Vulpiani, A. The Many Faces of Fluctuation Disisipann-Dssiiiiiiinn Relations Out of Equilibrium, pp. 29-57. Springer Inter- national Publishing, 2022. ISBN 9783031044588 doi: 00107997888-031-4445888 3. URL /ttt://dxx.dt org/10 10. 1007 1007/00077770 077888-31--0458--53
Gritsun, A. and Branstator, G. Climate response us- ing a three-dimensional operator based on the fluctua- tion-dissipation theorem. Journal of the Atmospheric Sciences, 64(7):2558-255 2575, July 2007. ISSN 0022-4928. doi: 10.1175/jas3943.11 http:///ddd doi org/10. 1175/JA83943.1.
Bi, K., Xie, I.., Zhang. H., Chen, X., Gu, X., and Tian, Q Accurate medium-range global weather fore- casting with 3d neural networks. Nature, 619(7970): 533-538, July 2023. ISSN 1476-4687. doi: 10.1038/ 441586-023-0618553 URL http: http:///xxxoo...o///// (038/s41588-023-0615
Guan, H., Arcomano, T., Chattopadhyay, A., and Maulik, R. LUCIE: A Lightweight Uncoupled Cllmate Emu- lator with long-term stability and physical consistency for 0(1000) ensembles. arXiv e-prints, art. arXiv:2405.1622 May 2024. doi: 10.485500arxiv24405. 16297.
Boffetta, G. and Ecke, R. E. Two-dimensional turbu- lence. Annual Review of Fluid Mechanics, 44(1):
Gupta, J. K. and Brandstetter, J. Towards multi- spatiotemporal-scakee generalized PDE modeling. Trans- actions on Machine Learming Research, 2023. ISSN 2835-8856. URL https:// wwt.r/iiwwwee/..nnniieee id=dPSTDbCtBY.
427-451, January 2012. ISSN 1545-4479. doi: 10.1146/ annurev- fww.d--lluud-11211110 101240. URL http:://ds doi. org/10. 1146/ htth///nnuuvvviiie fwwio-llluddda 1207100-110000001200
Bonavita, M. On some limitations of current machine learning weather prediction models. Geophysical Re- search Letters, 51()), June 2024. ISSN 1944-8007. doi: 10.1029/2023311073777 URL http: /xxx doi. www.doooooo..../ 1029/2023011103777.
Ho, J., Jain, A., and Abbeel, P. Denoising diffu- sion probabilistic models. ArXiv, abs/2006.11239 2020. URL Tetto http//ppp..emanicccchoollee org/ :orpusID:21990 2199556633
10

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Jiang, R., Lu, P. Y., Orlova, E., and Willett R. M. Training neural operators to preserve invariant measures of chaotic attractors. fbs/2306.01187, 2023. URL https://api. semanticscholes www.a/ooo...rrrrooooeeooooo Crrg/CorpuIID... : 259063696.
Li, Z.., Liu-Schiaffini, M., Kovachki, N. B., Azizzadeneshelli
K., Liu, B., Bhattacharya, K., Stuart, A., and Anandku- mar, A. Learning chaotic dynamics in dissipative systems. In Oh, A. H., Agarwal, A., Belgrave. D., and Cho, K. (eds.), Advances in Neural Information Processing Sys- tems, 2022. URL httttp:://openreveeww net/forum? id=1C36tF2n7aR,
Keisler, R. Forecasting Global Weather with Graph Neural Networks. arXiv e-prints, art. arXiv:2202.07575. Febru- ary 2022. doi: 10448500a2x...2222..7555
Lienen, M., Ludke, D., Hansen- J., and Gunne- mann, S. From zero to turbulence: Generative mod- eling for 3d flow simulation. In International Confer- ence on Learning Representations, 2023, URL https: //arxiv. org/ httg/b///23000 01776.
Kochkov, D., Smith, J. A., Alieva, A., Wang, Q., Bren- ner, M. P., and Hoyer, S. Machine learning- accellraeee computational fluid dynamics. Proceedings of the Na- tional Academy of Sciences, 118(21), May 2021. ISSN URL http:
1091-6490. doi: 10. htw0:///paa..111188444. //dx.dos oeeii.org/1000 1073/0073spss htw.://///aaa..11711111111 2101784118.
Linot, A. J., Burby, J. W. Tang, Q., Balaprakash, P., Gra- ham, M. D., and Maulik, R. Stabilized neural ordinary differential equations for long-time forecasting of dy- namical systems. Journal of Computational Physics, 474:1 1838, February 2023, ISSN 0021-9991. doi: 10. 1...016j..ep..0222.. 111838. hRL doi.org/ 10. (0.1016/j.jom 202.:/22..p2222...... 20t..2222..111.0 111838.
Kochkov, D., Yuval, J., Langmore, I., Norgaard, P., Smith, J., Mooers, G., Klöwer, M., Lottes, J., Rasp. S., Düben, P., Hatfield, S., Battaglia, P., Sanahez- A., Will- son, M., Brenner, M. P., and Hoyer, S. Neural gen- eral circulation models for weather and climate. Na- lure, 632 10000027) July 2024. ISSN 1476- 4687. doi: 10. 10.10388/411366-024-007744 //dx. httpa/////...10////1188-----0000 020-118886--2--774------
URL http:
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint arXiv: 20.222200022744 2022.
Kohl, G., Chen, L.-W., and Thuerey, N. Benchmarking Au- toregressive Conditional Diffusion Models for Turbulent Flow Simulation. arXiv e-prints, art. arXiv:2309.0148 September 2023. doi: (0.48550//rXiv22309.00
Lippe, P., Veeling, B. S., Perdikaris, P., Turner, R. E., and Brandstetter, J. Pde-refiner: Achieving accurate long roll- outs with neural pde solvers. ArXiv, abs/2308.05732. 2023. URL https://appp /ta/ppppssmmaicccccollaaaaa JorpusID:260775600
Kurth, T., Subramanian, S., Harrington, P., Pathak, J., Mardani, M., Hall, D., Miele, A., Kashinath, K., and Anandkumar, A. Fourcastnet: Accelerating global high- resolution weather forecasting using adaptive fourier neu- ral operators. In Proceedings of the Platform for Ad- vanced Scientific Computing Conference, PASC '23, New York, NY, USA, 2023. Association for Computing Ma- chinery. ISBN 9798400701900. doi: 10. 10.1145/3592979 3593412. URL https:////ooi orgo 10. 1145/3592979. 3593412.
List, B., Chen, L.-W., Bali, K. and Thuerey, N. Dif- ferentiability in Unrolled Training of Neural Physics Simulators on Transient Dynamies. arXiv e-prints, art. arXiv:2402.12971, February 2024. doi: 10.48550/arXiv. 2402.
List, B., Chen, L.-W., Bali, K., and Thuerey, N. Differen- tiability in unrolled training of neural physics simulators transient dynamics. Computer Methods in Applied Mechanics and Engineering, 433:1 17441, January 2025. doi: 1.100666..cma.22220
Lam, R., Gonzaez- A., Willson, M., Wirns- berger, P., Fortunato, M., Alet, F,, Ravuri, S., Ewalds, T., Raton-Rosen, Z., Hu, W., Merose, A., Hoyer, S., Holland, G., Vinyals, O., Stott, J., Pritzel, A., Learning skillful
Lumin-Jimenez, L., Darman, M., Hazarika, S., Wu, T. Gray, M., He, R., Wong, A., and Chattopadhyay. A. Si- multaneous emulation and downscaling with physically- consistent deep learning-basee regional ocean emulators. arXiv e-prints, art. arXiv:2501 100005500.00088, January 2025. doi: 10. 0044850//rxi..2551.0555
Mohamed, S., and Battaglia, medium-range global weather forecasting. Science, 382(6677):1416-11) 2023. doi: 10.1/26/cceence. adi2336. 10. 10.112/sscience adi2336.
org/ /tw//boi////////////
Lang. S., Alexe, M., Chantry, M., Dramsch, J., Pinault, F,, Raoult, B., Clare, M. C. A., Lessig, C., Maier-Gerber, M., Magnusson, L., Bouallègue, Z. B., Nemesio, A. P., Dueben, D., Brown, A., Pappenberger, F., and Rabier, F. Aifs ecmwf's data-driven forecasting system, 2024. URL pettps://arx.te http:///riivv.o/////////0/ 01465.
Majda, A. J. and Qi, D. Strategies for reduced-ordee models for predicting the statistical responses and uncertainty quantification in complex turbulent dynamical systems. SIAM Review, 60(3):491-549 January 2018. ISSN 1095- 7200. doi: 10. 10.1137/16ml 104664.
URL http: //dx:
doi. doi.org/10 10. 1

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Mardani, M., Brenowitz, N., Cohen, Y., Pathak, J., Chen, C.-Y., Liu, C.-C., Vahdat, A., Kashinath, K., Kautz, J., and Pritchard, M. Residual Diffusion Modeling for Km- scale Amoosheric Downscaling. arXiv e-prints, art. arXiv:2309. 15214, September 2023. doi: 10.48550/arXiv. 2309.15214.
complex physies with graph networks. In International
Conference on Machine Learning, 2020.
Schiff, Y., Wan, Z. Y., Parker, J. B., Hoyer, S., Kuleshov, V., Sha, F., and Zepeda-Núñee L. DySLIM: Dynamics Stable Learning by Invariant Measure for Chaotic Sys- tems. arXiv e-prints, art. arXiv:2402.04467, February 2024. doi: (0.48550/axxiv22402..00
Mikhaeil, J. M., Monfared, Z. and Durstewitz, D. On the difficulty of learning chaotic dynamics with rnns. In Neural Information Processing Systems, 2021. URL https: /tpp/ppp...... semppaiiaaannicchhll.a. org/ CorpusID: 250626473.
Shehata, Y., Holzschuh, B., and Thuerey, N. Improved sampling of diffusion models in fluid dynamics with tweedie's formula. In The Thirteenth International Conference on Learning Representations, 2025. URL httppp//pppuuueessee..eeeerrtrra wwwp:eewwwwweeeeoooeeeeeeeeeeeee wtt..///ooorrii1i/////2////0090
Parthipan, R., Anand, M., Christensen, H. M., Hosking, J. S., and Wischik, D. J. Defining error accumula- tion in MI. atmospheric simulators. arXiv e-prints, art. 2Xiv:24055144144 May 2024. doi: 10.48550 14714.
Sohl-Dickstein, J. N., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. ArXiv, abs/1503.03585, 2015. URL https://apl www.ppsssaaatccccssrr...oo org/ Corpus 14888175.
Pathak, J., Subramanian, S., Harrington, P., Raja, S., Chat topadhyay, A., Mardani, M., Kurth, T., Hall, D., Li, Z., Azizzadenesheli, K., Hassanzadeh, P., Kashinath, K., and Anandkumar, A. FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators. arXiv e-prints, arXiv:2202.1 1114, February 2022. doi: 00448550/aXii.22222.1222
Song, Y. and Ermon, S. Generative modeling by estimat- ing gradients of the data distribution. In Neural Infor- mation Processing Systems, 2019. URL atttps://api semanticscholar: oorg org Serg//orppusDDD : 196470871.
Stachenfeld, K., Fielding, D. B., Kochkov, D., Cranmer, M., Pfaff, T.,, Godwin, J., Cui, C., Ho, S., Battaglia, P., and Banchez-Gonzalez. A. Learned Coarse Models for Efficient Turbulence Simulation. arXiv e-prints, art. arXiv:2112.15277 December 2021. doi: 10.48550/arriv. 2112.15275.
Price, I., Ganaaez, A., Alet, F., Andersson, R., El-Kadi, A., Masters, D., Ewalds, T., Stott, J., Mohamed, S., Battaglia, P., Lam, R., and Willson, M. GenCast: curfusion-basee ensemble forecasting for medium-range weather. arXiv e-prints, art arXiv:2312.15799 December 2023. doi: 104485502axxi..3312.5777
Subel, A. and Zanna, I.. Building Ocean Climate Emulators. arXiv e-prints, art. arXiv:2402.04342 February 2024. doi: 00448500arXxv.2242..00
Robbins, H. E. An empirical bayes approach statistics. In Breakthroughs in Statistics: Foundations and basic theory, pp. 388-3888 388-394.90 Springer, 1992.
Vlachas, P., Pathak, J., Hunt, B., Sapsis, T., Girvan, M., Ott, E., and Koumoutsakos, P. Backpropagation algorithms and reservoir computing in recurrent neural networks for the forecasting of complex spatiotemporal dynamics. Neural Networks, 126: 191- June 2020. ISSN 0893- 6080, doi: 10.1016/..ccmme22200.....m / //dx doi. org/ 10. 1016////..000/........ www.ent.tt....m..mmm.. 2020. 02.00..006...
Ronneherger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. ArXiv, abs/ abs/15505.0459, 2015. URL https:// api peg..smmnnticscholarr wrg/CorpusID:3373500 37112281.
URL http:
Rühling Cachay, S., Henn, B., Wett-Meyer, O., Brether- ton, C. S., and Yu, R. Probablistic emulation of global climate model with spherical dyffusion. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.), Ad- vances in Neural Information Processing Systems, volume 37, pp. 127610-127644. Curran Asso- ciates, Inc., 2024. URL. https://proceedings cttripss.c/papeer http:///...pppp..teee/////////// 86111661880011132299885559999999 pdf.
Wan, Z. Y., Baptista, R., Chen, Y.-f., Anderson, J., Boral, A., Sha, F., and Zepeda-Núñez. L. Debias Coarsely, Sample Conditionally: Statistical Downscaling through Optimal Transport and Probabilistic Diffusion Models. arXiv e-prints, art, 1.t.::333..16118 May 2023. doi: 10.485550/2Xiv.2305.5660
Meyer, O., Dresdner, G., McGibbon, J., Clark, S. K., Henn, B., Duncan, J., Brenowitz, N. D., Kashinath, K., Pritchard, M. S., Bonev, B., Peters, M. E., and Brether- ton, C. S. ACE: A fast, skillful learned global atmo- spheric model for climate prediction. arXiv e-prints, art.
Suuchez-Gonzales A., Godwin, J., Pfaff, T., Ying, R. Leskovee, J., and Battaglia, P. W. Learning to simulate
12

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
arXiv:2310.02074, October 2023, doi: 10.48550/arXiv. 2310.02074.
Weed, J. and Bach, F. Sharp asymptotic and finite-sample rates of convergence of empirical measures wasserstein distance, 2019.
Zeeman, E. C. Stability of dynamical systems. Nonlinearity, 1(1):115-155, February 1988. ISSN 1361-6544. doi: 10.
1088/09511-711511110005. URL http: /ttp//xxxxx//o//////o///////// org/ 10.1088/0911-771/11///06
13

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
A. Dynamical systems
We evaluate our method on two high-diimesiional dynamical systems, exhibiting the kinds of chaotic, turbulent dynamics that are prevalent in systems of practical interest.
A.1. Kolmogorov flow
common fluid flow to test ML surrogates, is to use a forced variant of incompressible Navier-Stokes:
0.u (u u) v vp+ff
where = [M,, Uy]. is the kinematic viscosity, p is the fluid density, p is the pressure, and f is an external forcing term. Following (Kochkov et al., 2021), we use constant sinusoidal forcing function. We use p and v 0.001, corresponding to a Reynolds number Re = 10, 000.
We numerically solve the equations using the pseudospectral method with periodic boundary conditions from the publicly available code jxx-fd cfd (Dresdner et al., 2022). We use a numerical timestep of öt = 0.001s, and simulate the system with ng = n, = 512 spatial gridsteps. The fields are then downsampled to 64 x 64 for training and testing of the emulator. We use the vorticity, w V
R64 as our state vector, unlike in some previous works which trained their emulators on the velocity vector fields (e.g. Lippe al. (2023)). All state vectors are normalized to unit variance before passing the neural network, and results in the main text are shown in normalised units for simplicity.
A.2. Quasi-geostrophic turbulence
We consider two-layer, guost-geooiioppiicccccccccooccii system, where the prognostic variable is the potential vorticity, given by
4m vi (-1)m fö
H'Hm (v) v22). m E {1,2].
(10)
where 1 denotes the upper layer, m denotes the lower layer, Hm is the depth of the layer, &b is the streamfunction, which is related to the fluid velocity by um (ulm: Um) = -. U2)m), and fo is the Coriolis frequency. The time evolution of the system is given by
Oegon + V (um + vm + Um Sm. V + ssd Tun
(11)
where Um is the mean flow the x (zonal) direction, Bm 3 /+ (-1)m+1 L fá, - U2), Ten the bottom drag coefficient, and 8,, 2 is the Kronecker delta.
g gH
We numerically solve these equations using pseudo-spectral method, with Adams-Bashforth 3rd order timestepper in the Fourier domain and periodic boundary conditions. This was implemented in PyTorch, the code for which will be made publicly available upon de- anonymization. Numerical simulations are run at resolution of 256 x 256, and the potential vorticity fields are downsampled to 64 x 64 for training and testing of the neural network models. We use the 2-layer potential vorticity, q E 2xx64464 as our state vector representing a QG system, so all neural networks used in QG experiments have an additional input channel compared to the architectures used in Kolmogorov flow. As with the Kolmogorov experiments, all state vectors are normalized to unit variance before passing to the neural network, and results in the main text are shown in normalised units for simplicity.
B. Baseline emulator
B.1. Loss function
Here we describe the construction of the baseline neural emulator. Training over multiple timesteps has been comprehensively shown to improve long term stability (Brandstetter et al., 2022; Gupta & Brandstetter, 2023; List et al., 2024; 2025), and one can imagine three different ways to build a neural network and loss function do this. First, we can simply predict the state some future timestep, as function of the current timestep, i.c. 43 X, Xt+1, with MSE loss:
N Ls
1
ΣΣ
(x++ - xi+1 + 7n 11?
(12)
N k
14

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
where k indexes training trajectory (of total of N training trajectories, each of length L snapshots), X +x((,) +n denotes predicted state (including stochastic component), and &n- Xp such that we starting from simulation snapshot.
An alternative approach is to instead emulate the residual between two snapshots, ie.. Xe+1 X, (g(&x) + rn. This can be achieved using two loss functions either by evaluating the loss on the predicted state as in equation 12:
1 4-1 N
Ls
ΣΣ
Σ
(0,(()))
( xt) - x + Tn 1|2
(13)
k=
t=
where all terms are the same as previous. Alternatively, one can evaluate the loss on the residual:
1 N
CB
Σ
( (xf++ xx)) rn rn
(14)
After experimenting with all 3 loss functions, we found significantly better performance from the residual emulators, and particularly from the residual loss function in equation 14, in terms of validation loss and MSE accuracy over 200 step rollouts.
B.2. Architecture
We experiment with two architectures for the underlying emulator. First, we use Unet (Ronneberger et al., 2015) style image-to-image architecture, adapted from attps://gitt giggtt....... com/ppearena// waaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa (Gupta & Brandstetter, 2023) (referred to as Modern Unet in that work). Our Unet configuration consists of downsampling layers, where the number of convolutional filters is doubled at each stage of downsampling. We set the number of convolutional filters in the initial layer to 64, which leads to total number of 48M parameters. We include residual connections between each stage of downsampling, and use GeLU activations throughout the architecture. We experimented with and without batch and group normalization, and found that these had impact emulator performance, we use normalization in the results presented in this paper. Given that both of the PDE solution data we are training and testing on were generated with periodic boundary conditions, all convolutional layers use circular padding.
MSE w.r.t numerical model: Kolmogorov flow
10a
48M param Unet
0.5M param DRN
10a
10a
10 -2
1000
2000
3000
4000
5000
6000
7000
Emulator step
Figure 7. MSE with respect numerical model for emulated Kolmogorov flow trajectories for the Unet and DRN models. The DRN exhibits better long-term stability than the Unet, despite the much smaller model having a higher initial MSE
Secondly, we also use Dilated ResNet (DRN) architecture (Stachenfeld et al., 2021). Our implementation for the DRN is again adapted from (Gupta& & Brandstetter, 2023). A DRN is composed of Dilated Blocks, each with 2 stacks of 7 consecutive convolutional layers. The filters of these convolutional layers have dilations of [1, 11,2,4 4, 8, 4, 4,2, 1] consecutively. Our baseline DRN is composed of 2 layers of convolutional preprocessing layers with stride 1, and then 4 Dilated Blocks, before finally passing through 2 final convolutional layers, again with stride 1. After each layer, we use GeLU activations, and we include residual connections over all Dilated Blocks. Each convolutional layer consists of 32 convolutional filters, and we do not use any batch or group normalization. The total number of parameters for the DRN 0.5M.
15

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
B.3. Optimization
The network weights are optimized using the AdamM optimizer with momenta 31 = 0.9 and B3 0.999, and a learning rate of 5e 4. In practice, equation 6 is broken up into mini-batcher of size 32. We train the model for 12 epochs, (longer training runs were experimented but we found the loss curves did not improve with more epochs). The time horizon At = 2 was chosen by experimenting with a range of different values in the range [1, 20), and selecting At based on emulator performance in terms of MSE predictions over a time horizon of 200 step rollouts. Consistent with previous works, we found that composing multiple passes of a network trained over shorter training time horizons produced better MSE than training over larger timesteps (Lippe et al., 2023).
Importantly, we observed that improvements in validation loss often had little impact on long term stability. This can be seen in Figure 7, where compare the short term MSE for the Unet and DRN models. The larger Unet model has significantly lower MSE over short number of steps, but all trajectories quickly go unstable. Indeed, we observed that long term stability varied significantly with network weight initialization, even for fixed dataset, model architecture and optimization algorithm.
C. Diffusion model thermalizer
C.1. Architecture
To implement the thermalizer as a diffusion model, we use the same core architecture the Unet used for the emulator. As discussed in section 4.2, key component of the thermalizer is some prediction of the noise level. To augment this standard Unet implementation with noise-classifying head, we take the lowest-dimensional representation from the Unet, and add an additional 2 convolutional layers, with kernel size of 3 and stride of 1. We then vectorize the output of the last convolutional filter into a length 4096 vector. This is then passed through more linear layers of size 1000. This final vector of length 1000 represents the predicted categorical distribution over noise levels for the input image. The total number of learnable parameters for this network is 53M We implement a method to rum forward pass through the clais-classifying component of the network only, for computational efficiency during rollouts. Also note that since our network does not include the noise level an input scalar, our diffusion model does not need have additional timestep embedding components.
C.2. Optimization
To optimise the network, equation 8 is broken up into mini-batches of size 64. Again use the Adam'l optimizer, with learning rate of 2e, and train the thermalizer for 35 epochs. Given that at inference time, we generally only use low noise levels of approximately X 20, we experimented with training the model only on these low noise levels. However we found that performance degraded significantly, so sample noise levels s uniformly across the full range during training, Noise levels are set using a cosine variance scheduler. All model architecture, training and inference codes can be found at https: http::///tttatt httggttttuu.... cwm/Coriiiceeeeecceeeecccca wwweederree/thermaalla.
C.3. Thermalizer settings: Sinit and Sstop
A key component of the algorithm is the adaptive nature of thermalization - which allows for stabilisation with a minimal interruption of temporal dynamics. There are two free parameters here- - Sinis which determines the estimated noise level at which the thermalizer starts to apply corrections to the state vector, and 8stop, which sets the lowest noise level to which we run the denoising process. To find the optimum settings, we run a simple grid search on Switt E [10, 6] and Sstop E [5, 2] (note that need Sstart > Sstop- For each pair of Sint and 8stop- run a thermalized trajectory for 1e4 steps. We chose the combination of Sinit and Satop in which the kinetic energy spectra of the thermalized trajectories best matches the kinetic energy in the numerical model. This experiment is repeated independently for each flow configuration, and each emulator so this procedure was run a total of 4 times. We found Sini 7 and Satop = 4 optimal for the Kolmogorov Unet, Sinkt = 10 and Ssoop = 5 optimal for the Kolmogorov DRN. For QG we use sinit = 9 and Sstop = 4 for the Unet, and Smit = 6 and Sstop = 3 for the DRN emulator.
D. Inspecting a single thermalization step
It is possible to investigate the effect of a single thermalization step on rollout. We randomly select snapshot from the thermalized trajectories show in Figure 3, and take emulator steps from this snapshot until the predicted noise level reaches
16

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Step
Step 1000
Step 2000
Step 3000
Step 4000
Withouf therm step
2000
250000
-50000
-2
-10
With themm step
10000
2
0
50
10
250000
0.05
Dift
0.00
-0.05
250000
Figure 8. Inspecting the effect of single thermalization step on the stability of QG rollout. Top row shows an emulator rollout starting from snapshot noise level spred = 10. The middle row shows the same rollout, except starting from the same snapshot after applying a thermalization step. The development of instability delayed by this single thermalization step.
the Sinit 10. We then take 8therm = Sinit Astop = thermalizing steps on this snapshot. The corresponding fields before and after thermalizing are shown in the leftmost panel of Figure where the total effect of the thermalization step shown in the lowest left panel.
We can then continue running the emulator on the snapshots with (middle row) and without (top row) applying this single thermalization step, and observe that the eventual instability develops sooner in the non-thermalized (top) row. This can be seen both visually in the potential vorticity fields as accumulation of high-frequency signal, and the growing field magnitudes shown by the colorbars. This experiment illustrates the effect that the subtle thermalization steps have on the stability of rollouts. The same procedure is repeated for other randomly selected snapshots Figures 9 and 10.
E. Additional results and figures
In Figure 11 we show schematic illustrating the principle of the thermalizer. In Figure 12, we show Spred along the full 10000 step trajectory, for the Unet emulator in both Kolmogorov and QG flows. In gray lines we show Spred for the emulator trajectories, thermalized trajectories are shown in blue, and the numerical model in red. Initially, Spred is low for all three different models, as the state vector has not had a chance to accumulate significant error. We see that 8pred increases as the un- -n-thermalized trajectories accumulate error eventually going unstable, which point our noise-classifyin chasss-issiiyiingg output becomes unreliable, leading the vertical gray lines. In contrast, the blue lines are constrained stay low noise level, never exceeding 5pred 16. This is the consequence of the thermalizer keeping states in-distribution. Finally the red lines for the numerical model are shown as sanity check that our noise-classine classiiiasi head is able to consistently identify that the numerical model states are all in-distribution.
In Figure 13, we inspect the number of thermalization steps along a 500 step segment of the trajectories show in Figures I and 15. For each of the 40 trajectories, and each of the 500 steps between timesteps 5, 000 and 5, 500, we show the number of thermalization steps applied to the state vector. We see that the amount of thermalization applies varies strongly along a trajectory, with steps ranging from 0 16. This indicates that the flow oscillates between periods of no correction, and periods where significant correction is applied.
Next we show the vorticity fields for 25 random samples from the 40 trajectories shown in Figure 1, at the end of the 10, 000 steps. In In Figure 16, we show results for the numerical model. Above each panel, we show the predicted noise level for each flow state. In Figures 17 and 18, we show the equivalent for the emulator and thermalized trajectories. Note that in the emulator case, after 10, 000 steps the fields have gone so far out of distribution that the predicted noise levels are unreliable.
17

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Step
Step 1000
Step 2000
Step 3000
Step 4000
Without therm step
10
2000
-2
10
-000
With therm step
22
2
2.5
200
-2
225
-200
0.1
10
2000
DiT
0.0
0.0
-25
110
2000
Figure 9, Same Figure 8, except with different initial condition.
In Figures 19, 20 and 21, we show the upper layer potential vorticity for all 20 QG trajectories after 10, 000 steps, for the numerical, emulator, and thermalized models respectively.
18

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Step 0
Step 1000
Step 2000
Step 3000
Step 4000
10000
25
0
Witho
-2
-25
-5000
With step
2
00
10000
2.5
25
5000
Diff
0.00
0.0
0.05
-2.5
-5000
-5.0
Figure 10. Same Figures 8 and except with different initial condition.
y
sts(t)
T
M(Rd)
Figure 11. Schema for the thermalization in the space of marginal distributions P(EE). In yellow, the OU diffuusion path is defining transport between equilibrium μ and Gaussian measure 7. In green, current law A x,. In purple, the (stochastic) transport map F defined by the reverse diffusion. While by construction it transports v, back to when applied to fi. it produces the corrected measure f.. reducing the error as per (7).
19

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Predicted noise level during rollouts
Spradd
60
Kolmogorov
0
2000
4000
6000
8000
10000
Numerical model
80
Emulator Thermalized
2000
4000
6000
10000
Emulator step
Figure 12. Predicted noise level, Speed during rollouts for the numerical model, emulator, and thermalized trajectories. In the top panel show results for the 40 Kolmogorov trajectories using the Unet emulator shown in figure 1, and in the bottom panel we show the 20 QG trajectories shown figure 3.
Number of thermalization steps during rollout, for Ko Volmogorov flow
14g
20
DRN
Trajectrry
20
200p
5000
5100
5200
5300
5400
5500
Timestep
Figure Number thermalization steps along thermalized trajectories. We inspect number of thermalization steps each of the 40 trajectories at each timestep, between timestep 5000 and 5500, for both the Unet and DRN emulators on Kolmogorov flow.
20

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
1= stop=0
t= step= 1500
5= slep=2500
t= 10s, step=4750
t= 15s, slep= 7500
1=20s, step=
Numertcal model
Envulator
Thermalized
107
302
10
Kinetic energy spectra from 40 different trajectories
Figure 14. Equivalent figure as in Figure 1, but with the Dilated ResNet (DRN) emulator.
1-0 davs, stop +0
t- 52 days, 250p-250
1-208 dlass, stop- 1000
t=521 days, step=2500
t-1042 days, step-5000
t=2083 days, step L0000
Numerical model
Emmlator
Thermalized
101
10
10
Numerical model
104
Enulator
Thermalioed
10-
Kinetio energy spectra from 20 different trajectories
Figure 15. Equivalent figure as in Figure 3, but with the Dilated ResNet (DRN) emulator.
21

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Numerical model, step
255
5.0
25
0.0
-2.5
-25
-5.0
5.0
2.5
2,5
0.0
0.0
2.5
-5.0
0
Figure 16. Kolmogorov vorticity snapshots after 10, 000 steps, from the numerical model. The predicted noise level is shown above each fluid state.
22

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Emulator step
93
le19
93
93
-1.0
1.20
1e20
162
1e20
1e20
1.0
0.5
-0.5
-5
-1.0
1620
1c20
1621
1c20
1c20
0.5
-2.5
-05
-5.0
-1
1620
1e22
1620
93
1c20
93
le1s
225
5.0
93
93
93
93
93
-1
Figure 17. Kolmogorov vorticity snapshots after 10, 000 steps, from the neural network emulator. The predicted noise level is shown above each fluid state.
23

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
Thermalized states, step
0.0
nn
B
5.0
25
0.0
2.5
2.5
-5.0
2.5
2
-25
-5.0
2
25
-25
Figure 18, Kolmogorov vorticity snapshots after 10, 000 steps, for the thermalized trajectorics. The predicted noise level is shown above each fluid state.
24

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
QG Numerical model, step 9999
2.5
0.0
-25
2
2.5
-2
-2.5
-2
-2
2
2
2.5
2
0
-2
-2
-2.5
-4
2
2.5
2
2
0.0
0
2.5
22
-2
Figure 19. QG upper layer potential vorticity snapshots after 10, 000 steps, from the numerical model. The predicted noise level is shown above each fluid state.
25

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
QG Emulator only, stop 9999
0.10
0.10
0.10
0.10
0.10
0.05
0.05
0.05
0.00
0.00
0.00
0.00
0.00
-0.05
-0.05
-0.05
-0.05
-0.05
0.10
0.05
0.05
0.05
0.05
0.05
0.00
0.00
0.00
0.00
0.00
-0.05
-0.05
-0.05
-0.05
-0.05
00.10
0.10
-0.10
0.10
0.05
0.05
0.05
0.05
0.05
0.00
0.00
0.00
0.00
0.00
0.05
0.05
0.05
00.10
0.10
-0.0
-0100
0.100
0
0.10
0.10
0.10
0.10
0.10
0.05
0.05
0.05
0.05
0.05
0.00
0.00
0.00
0.00
0.00
-0.05
-0.05
-0.10
-0.10
-0.0
-0.10
-0.10
Figure 20. QG upper layer potential vorticity snapshots after 10, 000 steps, from the neural network emulator, The predicted noise level is shown above each fluid state.
26

---

Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos
QG Thermalized states, slep 9999
2
2.5
2
0
0.0
0
-2
-2.5
12
12
5.0
2
2.5
2.5
0.0
-2
-2
-2.5
-25
12
5.0
5.0
2.5
2.5
-0.0
-2
-2.5
2
225
-2.5
12
12
13
12
2.5
2.5
2
0.0
0.0
-0
-2.5
-2.5
-2
2
Figure 21. QG upper layer potential vorticity snapshots after 10, 000 steps, for the thermalized trajectories. The predicted noise level is shown above each fluid state.
27


Reviews:
[{"review_id": "ZIY0hmE46I", "reviewer": "Reviewer_GCvA", "summary": "The paper proposes a method to stabilize predictions of an autoregressive surrogate model over long-term rollouts. for that, it proposes to learn the invariant measure with a diffusion model and perform denoising steps at inference with a noise level that is guessed by a classifier. The paper claims that arbitrarily long rollouts can be achieved with their method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments.\n\n- This is a great point which we discussed internally - indeed a diffusion model can produce realistic samples of the flow fields, so it's important to verify temporal consistency. In Figure 4 we show the autocorrelation over time for all different models, and demonstrate that the correlation between thermalized snapshots of different temporal separations is consistent with that of both the numerical model and the emulator. If one were just generating random snapshots, we would expect an autocorrelation of 0. And indeed if we were “over-thermalizing”, by introducing significant noise into the system, this autocorrelation would decrease faster than the emulator and numerical models. However we see from this figure that temporal consistency is preserved, as the number of corrective steps is miminised by the noise-classifying component of the algorithm. This ensures that only states out of distribution are denoised, and are denoised at the minimal amount to return to the invariant measure.\n\n- Additionally, at this anonymous link: https://drive.google.com/drive/folders/1b1DHvR-LJvIZtlpvr6hIRQzxn3JKEuwQ?usp=sharing please find two animations of 6 trajectories for Kolmogorov and QG, where we compare the 3 models, and the temporal flow of the thermalizer trajectories is consistent with the numerical model long after the emulator trajectories have gone unstable.In the bottom row of the QG animation, we show the number of thermalization steps over the past 100 steps, during the rollout, such that the adaptive nature of thermalization is shown. A de-anonymised link to animations will be included in the final version of the paper, as the preservation of this temporal flow is a crucial success of our algorithm.\n\n- The implementation of the denoising process is similar - indeed we also add noise before performing the denoising (see Algorithm 1) - we experimented without this component and found better performance when this forward process is included. The main difference in our work with respect to the PDE-refiner, is that their denoising is still conditioned on the previous (clean) timestep, and therefore their model is still exposed to accumulation of error and distribution shift as states wander out of distribution. We ran experiments with the PDErefiner code, but found them to be as unstable as the regular emulator in our testing conditions, so we decided not to report them. It may be that we did not find the correct hyper-parameter settings, so we decided not to report them. but we note that we are not the only ones who had difficulties setting up PDErefiner (see eg https://openreview.net/forum?id=0FbzC7B9xI section E.1, and https://arxiv.org/abs/2309.01745 Figure 8 and section 4.2). We are happy to include this remark in the updated text. \n\n- With regard to your comment on the sampling rate - this is an interesting suggestion. We do not study this degree of freedom explicitly, however in Figure 4, we see that the decorrelation time is significantly faster between QG and Kolmogorov, indicating that our emulator step size for QG is significantly larger than for Kolmogorov, with respect to the temporal dynamics of the systems. Yet in both cases, the thermalizer is able to stabilize the flow on multiple emulators. Your suggestion is nonetheless very interesting, and we will certainly explore it. \n\n- Thank you for both of your suggestions on improving the clarity of the figure 3 - we agree with these modifications and will incorporate these into the revised manuscript. We have also included an additional figure explicitly showing the number of thermalization steps over time during a rollout (similar to the bottom row in the QG animation linked above), such that the amount of thermalization is more clearly shown.", "claims_and_evidence": "The main claim of the paper is well supported by experiments, but should be better precised and tested (see suggestions below).", "methods_and_evaluation_criteria": "The datasets chosen are challenging enough, and the experiments conducted to assess the stability of the rollouts are convincing but should be extended (see suggestions below).", "theoretical_claims": "The paper makes no theoretical claims, in the sense that the authors do not prove a new theorem.", "experimental_designs_or_analyses": "Visualizing the rollout trajectories and the kinetic energy spectrum over time makes sense for analyzing the divergence of a predicted trajectory (figure 1, 2, 5). Tracking the mse between the stabilized and numerical model is important, as well as the number of thermalization steps (figure 3, 4). i think figure 3 should be complemented with an additional important measure (see suggestions below).", "relation_to_broader_scientific_literature": "This paper tackles an important problem, which is the stabilization of rollouts of an autoregressive model for evolving trajectories of a spatiotemporal system. Compared to [1], it possesses the key advantage of being constructed separately from the predictive model that we are trying to stabilize.\n\n[1] Pde-refiner: Achieving accurate long rollouts with neural pde solvers. Lippe et al. 2023", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths\n\n- The paper is well written, making it simple to understand the method.\n- The formalism of the problem (e.g., optimal transport) sheds an interesting perspective on the problem that can inspire future work on this crucial topic.\n- The experiments are compelling.\n\nWeakness\n\n- It seems to me that the paper implicitly relies on the following assumption: the trajectories must be available at a high sampling rate over time. if this were not the case, then the \"emulator\" would have a harder time predicting next steps and thus make much larger errors at each step, which the \"thermalizer\" may struggle to correct. It could still project back to the invariant measure but would not preserve the temporal dynamics. The paper would have benefited from exploring an example with trajectories that are subsampled in time.", "other_comments_or_suggestions": "- Figure 3: the mse converges to a flat level. I believe this is the average mse between an independently generated state from the diffusion model and the observed trajectory. i think you should plot such a constant level on each graph.\n- Figure 3: I would suggest plotting histograms of the relative proportion of thermalization steps vs. emulator steps rather than thermalization steps only, which would be more compelling to understand how much \"thermalization\" is involved.\n- l097: \"arbtitrarily\"", "questions_for_authors": "- related to the first comment above, you claim that your method can achieve \"arbitrarily long rollouts\". However, one very simple way to achieve this is to just generate as many independent realizations with your diffusion model as you want in the future. This gives trajectories that have a negligible probability being an actual trajectory. Even though there is technically no \"rollout\" in this naive \"solution\" to the problem, I see no guarantees in your method that the \"arbitrarily\" long sequence you generate is actually close to a true trajectory (in the sense that it has the correct time dynamics). Could you comment on that? In particular, I think the claim of your paper should be changed since it may lead the reader to think you solve the (very) hard problem of obtaining arbitrarily long realistic trajectories.\n- another difference with the work [1] you mention is in how the refinement is done. [1] adds noise to the diverged sample and denoises it back (following the standard generation process with diffusion models), while your method does a direct denoising from a diverged sample.\nI guess in your terms, this means the \"transport\" back to the invariant measure is quite different. Could you comment on that?\nCould you actually also show results with [1] on your dataset and models? I would particularly be interested to see how one does on figure 4.\n\n[1] Pde-refiner: Achieving accurate long rollouts with neural pde solvers. Lippe et al. 2023", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JTE1tQyUy1", "reviewer": "Reviewer_7AYf", "summary": "The goal of this paper is to address the problem of unstable long rollouts by autoregressive neural PDE surrogate models (also called emulators). The core idea is to combine an autoregressive emulator model with an independently trained diffusion model. The role of the autoregressive emulator is then to make an initial prediction of the next system state, whereas the diffusion model 'corrects' this initial prediction to push it towards a data manifold that is in-distribution, thus preventing instabilities. This is implemented through the combination of a denoising diffusion model together with a classifier head that predicts the noise level of the initial classifier prediction, so as to arrive at an appropriate amount of denoising steps. The results, obtained using two model architectures and two datasets, demonstrate that the proposed method maintains stability for long rollouts, whereas the baseline variants diverge.\n\n## Update after rebuttal\n\nUpdated recommendation during rebuttal phase.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments.\n\n- Thank you very much for bringing [7] to our attention, that we indeed missed. It is a very interesting application of diffusion models for fluid dynamics, and it has components related to our method. That said, we want to point out what we believe are fundamental differences between the methodology and the problem setup. In essence, (i) [7] considers parallel sequence generation, where several future frames are generated given the current state, and (ii) it is based on conditional diffusion models, which are inherently exposed to distribution shift as one applies them in an auto-regressive fashion to produce long rollouts (longer than those used for training). Needless to say these will be included in the related work section of the updated version. \n\n- The setup of [7] is concerned with diffusion models for general sequence generation, and introduces two novel schemes to improve the numerical efficiency by reducing the number of function evaluations. In that respect, their main motivation is to speed up existing methods based on conditional diffusion. The main setup they consider is not the single-step autoregressive setting, but rather the parallel sampling setting, where several future states are sampled, conditionally on the current state (fig 6). \n\n- As far as we can tell, all the models considered are conditional, similarly as the PDErefiner, where one conditions on the last available state. As such, they are exposed to distribution shifts as the horizon of the rollout at inference time increases. While the numerical results reported by the authors are indeed impressive, we note that the temporal horizon considered is roughly an order of magnitude smaller than our setting (compare our Figure 3, where we report MSE, with their figure 8, where they report correlation). As emphasized in our text, our main insight is precisely to use an unconditional diffusion model to tame such distribution shifts, by exploiting the ergodicity of the dynamical system. That said, this paper introduces interesting insights (e.g. directly using the tweedie /miyazawa estimator rather than reversing the full diffusion path) which could be interesting to combine in our setting. \n\n- Thanks for pointing out the typo in eq 7 – you are correct! \n\n- Comparison with diffusion-based baselines: While we agree with the reviewer that it would be an interesting addition to the experimental evaluation, from our previous discussion we believe that our setup makes this comparison less critical: while these diffusion-based predictions can indeed improve the stability of point-estimate emulators, they are ultimately going to suffer from long-time instabilities as they drift out of distribution. Moreover, even if they were to become stable to arbitrarily long rollouts (as in our setting), they would incur substantially higher training and inference costs. Finally, as mentioned to reviewer GCvA, we ran experiments with the PDErefiner code, but found them to be as unstable as the regular emulator in our testing conditions, so we decided not to report them. It may be that we did not find the correct hyper-parameter settings, so we decided not to report them. but we note that we are not the only ones who had difficulties setting up PDErefiner (see eg https://openreview.net/forum?id=0FbzC7B9xI section E.1, and https://arxiv.org/abs/2309.01745 Figure 8 and section 4.2).\n\n- Additionally, at this anonymous link: https://drive.google.com/drive/folders/1b1DHvR-LJvIZtlpvr6hIRQzxn3JKEuwQ?usp=sharing please find two animations of 6 trajectories for Kolmogorov and QG, where we compare the 3 models, and the temporal flow of the thermalizer trajectories is consistent with the numerical model long after the emulator trajectories have gone unstable.In the bottom row of the QG animation, we show the number of thermalization steps over the past 100 steps, during the rollout, such that the adaptive nature of thermalization is shown. A de-anonymised link to animations will be included in the final version of the paper, as the preservation of this temporal flow is a crucial success of our algorithm.", "claims_and_evidence": "See the below fields.", "methods_and_evaluation_criteria": "The method is appropriate for the problem at hand, and has the advantage that the autoregressive emulator can be used in combination with an independently (pre-)trained diffusion model.\n\nThe benchmark datasets are suitable: the authors demonstrate that baseline models on this dataset suffer from the rollout instability problem, motivating the approach for these scenarios.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "* The considered metrics and statistics consist of traditional ML metrics (like MSE), physics-based metrics (energy spectra), as well as qualitative results that clearly demonstrate the divergence of the baseline models and stability of the thermalized rollouts, which makes an appropriate analysis strategy for the goal of the paper.\n\n* The considered baseline models are relatively straightforward deterministic emulator models. Although these should definitely be included in the experimental analysis, recently diffusion-based neural PDE emulators have gained significant attention. The method should be explicitly compared against such baselines to establish whether there is any benefit of the approach over existing diffusion-based methods. Consider e.g. (a selection of) the methods studied in [1-4].\n\n[1] Lippe et al. (2023). PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers.\n\n[2] Kohl et al. (2023). Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation.\n\n[3] Shysheya et al. (2024). On conditional diffusion models for PDE simulations. \n\n[4] Cachay et al. (2023). DYffusion: A Dynamics-informed Diffusion Model for Spatiotemporal Forecasting.", "relation_to_broader_scientific_literature": "The paper relates to general research on autoregressive neural PDE emulator models. Although these models have shown promise for short-term forecasts, they can become unstable over long rollouts. Early ideas to alleviate this issue included augmenting the input data with noise [5] or performing training over longer horizons [6]. More recently, diffusion-based approaches have shown promise in mitigating this issue as well, e.g. [1-4]. This paper is most closely related to such methods, but relies on independently trained autoregressive and diffusion components, and the diffusion model is applied for an adaptive number of steps at inference time based on a separate classifier head that predicts the noise level of the emulator's initial prediction. \n\n[5] Stachenfeld et al. (2022). Learned Coarse Models for Efficient Turbulence Simulation.\n\n[6] Brandstetter et al. (2022). Message Passing Neural PDE Solvers.", "essential_references_not_discussed": "A recent work [7], accepted for publication in ICLR, proposes 'iterative refinement', a method of which the essence is similar to the method proposed in this paper. As far as I can tell, the major differences lie in that [7] uses Tweedie's formula to iteratively denoise the prediction, and that it decides on a denoising schedule using greedy optimization rather than a separate noise level classifier. The work of [7] was developed concurrently, and I do not doubt that the authors independently arrived at similar ideas. Still, it would be beneficial to inform the reader on the differences and similarities in an updated version of the paper.\n\n[7] Shehata et al. (2025). Improved Sampling Of Diffusion Models In Fluid Dynamics With Tweedie's Formula. https://openreview.net/forum?id=0FbzC7B9xI", "other_strengths_and_weaknesses": "Strengths:\n\n* The results convincingly demonstrate the method's effectiveness over baseline emulator models in preserving long rollout stability.\n* The paper is clearly written and include extensive background and related work sections explaining the problem statement and diffusion models.\n\nWeaknesses:\n\n* My main concern lies in the lack of comparison against diffusion-based baselines (see experimental design). Although the results convincingly demonstrate that Thermalizer improves long rollout stability over vanilla autoregressive emulators, the empirical results do not demonstrate whether it addresses this issue more effectively than established diffusion-based methods.\n\n* It seems that the long rollout stability comes at the cost of slightly worse short-term forecasting performance.", "other_comments_or_suggestions": "Minor comments:\n\n* line 267: \"Gradients are backpropagated\nthrough the full L timesteps, as done in (Brandstetter et al.,\n2022; List et al., 2024).\" -- If I'm not mistaken, the method of Brandstetter et al. (2022) only backpropagates the gradients by a single step after a rollout during the training process.", "questions_for_authors": "* Can you please read [7], and explain the similarities/discrepancies/advantages/drawbacks of your method compared to iterative refinement?\n\n* Equation 7: In the first line of the MSE objective, should $D_\\phi^{(1)}$ not be trained against $\\epsilon_i$ as opposed to $s_i \\epsilon_i$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0DyPId7HoK", "reviewer": "Reviewer_ZXw4", "summary": "The authors proposed using their method, thermalizer,  to make autoregressive emulator rollouts of chaotic systems more stable. This method relies on a diffusion model, stabilising the emulator's overall predictions during the inference phase. UNet and Dilated ResNet were used as emulators. To verify the quality of the proposed algorithm, the authors used two turbulent systems, Kolmogorov flow and Quasi-geostrophic turbulence. As a result, the thermalized method shows the stability of the predictions over 1e5 timesteps.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for taking the time to assess our work, and for the insightful comments:\n\n- “Am I correct that the Thermalizer Algorithm can be interpreted as an extension of the DDPM framework for turbulent systems, with a diffusion model used to predict the noise level?”\nNot really. Thank you for bringing up this important point, that perhaps was not clearly presented. The DDPM framework provides a generative model of an underlying probability distribution $\\pi(x)$ from available samples $x_i \\sim \\pi(x)$. In the context of sequence generation, all existing DDPM applications have considered *conditional* generative models, of the form $\\pi(x_{t+\\delta} | x_{t})$, in order to restore temporal consistency of the trajectories. \n The novelty in our work is that we use a DDPM model only for the *invariant measure* of the system, to stabilise autoregressive rollouts in an adaptive and efficient way; ie we model simply $\\pi(x_t)$ for $t$ large. This model just happens to be implemented via a modified DDPM algorithm in this formulation, and it is *unconditional*. Crucially, our method considers separate, independent training components for the emulator and the invariant measure, and we view this as a major advantage over more complex alternatives that require jointly training the components. As we explained in the text, the unconditional aspect of our diffusion model is what guarantees long-time stability (in ergodic systems). Additionally, another contribution of our work is the formulation of the problem of autoregressive error accumulation, which is largely missing from the literature despite extensive empirical studies.\n\n- Comparison with prior work, e.g DySLIM. We agree with the reviewer that this would be an interesting addition to our experimental setup. We are however unable to complete these experiments during this rebuttal period, but will make sure to include them in a later iteration. We emphasize though that such techniques will lead to substantially higher training costs, since the invariant measure needs to be re-estimated every time a parameter is updated. \n\n\n- $s_\\text{init}$ and $s_\\text{stop}$ were found by running a grid search, much like a hyperparameter search. We ran a search in the range [15,5] where $s_\\text{init} > s_\\text{stop}$, and ran each combination for 10,000 steps. We chose the best performing run in terms of comparison between the thermalized and numerical model kinetic energy spectra, averaged across all 40 trajectories for Kolmogorov, and 20 trajectories for QG. We have expanded the description of this search in the appendix in the revised manuscript.\n\n- The $\\alpha$ and $\\beta$ coefficients are the noise-scaling coefficients as presented in the DDPM paper (Ho et al. 2020). Thanks for notifying us that this connection is not explicitly made in the text - we have updated the manuscriptt. As described in the appendix, we use a cosine variance scheduler for our implementation, which uniquely defines alpha and beta for our 1000 noise levels.\n\n- We use a different numerical scheme for the Kolmogorov and QG flows, each are described and referenced in appendices A.1 and A.2 . The QG numerical scheme is our own pytorch implementation, and a github link will be included on de-anonymisation, but a complete description of the method is given in Appendix A2.\n\n- We would like to address your comments on the choice of experiments. In terms of systems to study, we considered Lorenz and Kuramoto-Sivashinky, however these chaotic systems are significantly lower dimensional and less complex than the fluid flows we experiment on. So once the framework was demonstrated on Kolmogorov and QGs, we considered the experiments on Lorenz and Kuramoto-Sivashinsky to be redundant.\n\n- In terms of weather and climate models - indeed, a longer term direction is to apply this framework to full-scale weather and climate models. However this presents a significant data and computational challenge, which is beyond the scope of an initial methodological work. Such an effort would need to be motivated by some initial study and presentation of the new algorithm, which we submit here.\n\n- Additionally, at this anonymous link: https://drive.google.com/drive/folders/1b1DHvR-LJvIZtlpvr6hIRQzxn3JKEuwQ?usp=sharing please find two animations of 6 trajectories for Kolmogorov and QG, where we compare the 3 models, and the temporal flow of the thermalizer trajectories is consistent with the numerical model long after the emulator trajectories have gone unstable.In the bottom row of the QG animation, we show the number of thermalization steps over the past 100 steps, during the rollout, such that the adaptive nature of thermalization is shown. A de-anonymised link to animations will be included in the final version of the paper, as the preservation of this temporal flow is a crucial success of our algorithm.", "claims_and_evidence": "The authors proposed a Thermalizer algorithm and showed its effectiveness on Kolmogorov flow and Quasi-geostrophic turbulence datasets. The proposed algorithm is based on the Denoising Diffusion Probabilistic Model (DDPM). In particular, authors modified “the standard implementation of the DDPM framework, by adding a classifier output such that the noise level s is predicted by the network, instead of being passed as an input parameter”.", "methods_and_evaluation_criteria": "Proposed methods and datasets for evaluation are quite limited. \n\n1. Limited novelty. It seems to me that the main novelty is incremental. As I can see, the authors slightly modified a DDPM algorithm “by adding a classifier output such that the noise\nlevel s is predicted by the network, instead of being passed\nas an input parameter.” The overall algorithm idea is strongly based on the DDPM algorithm. Can the authors provide more arguments to justify their novelty?\n\n2. Comparison with previous approaches. Authors propose an algorithm that is “an alternative approach” “to stabilize predictions and improve the modelling of chaotic dynamics (Li et al., 2022; Jiang et al., 2023; Schiff et al.,2024)”. However, there seems to be a lack of comparison between the Thermalizer and those alternatives. I would suggest comparing Thermalizer with two alternative approaches from the literature review. \n\n3. Datasets. The authors used Kolmogorov flow and Quasi-geostrophic turbulence to verify the proposed method as evaluation datasets. Continuing from the previous comment, I recommend comparing Thermalizer on current datasets (Kolmogorov flow and Quasi-geostrophic turbulence) and from the papers mentioned above. For example, authors can use the chaotic Lorenz-63 system and the Kuramoto-Sivashinsky equation from (Li et al., 2022).\n\n4. Real-life applications. It is always interesting to see real-life applications of the proposed method. However, only experiments on the synthetic data are present in the paper. I suggest checking for long-term instability improvement using Thermalizer on climate and weather modelling tasks.", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "The experiments were carried out correctly and showed the proposed algorithm's effectiveness. Nevertheless, I would additionally recommend trying different emulators.", "relation_to_broader_scientific_literature": "The proposed approach continues DDPM (Ho et al., 2020) and is an alternate to the previous works (Li et al., 2022; Jiang et al., 2023; Schiff et al., 2024).", "essential_references_not_discussed": "Seems like all essential references were discussed in the article.", "other_strengths_and_weaknesses": "Abstract and introduction are carefully written. Pictures are high quality.", "other_comments_or_suggestions": "It seems that authors use the term Thermalizer in several meanings - diffusion model (“To implement the thermalizer as a diffusion model”) and an algorithm to stabilize trajectories (“We introduced the thermalizer, an algorithm for stabilising autoregressive surrogate models leveraging a pretrained diffusion model of the stationary data distribution”). I would suggest to choose only one meaning for the Thermalizer term.", "questions_for_authors": "1. Am I correct that the Thermalizer Algorithm can be interpreted as an extension of the DDPM framework for turbulent systems, with a diffusion model used to predict the noise level?\n2. How do you initialize sinit in Algorithm 1?\n3. Can you please clarify what the alpha and beta coefficients in Algorithm 1 stand for?\n4. The authors compare the thermalized method with the \"numerical model\". However, the reference for the \"numerical model\" seems to be absent. Can the authors clarify this issue?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qJV3n7WQgT", "reviewer": "Reviewer_vKpF", "summary": "This paper introduces \"thermalization,\" a novel inference-time stabilization method for autoregressive emulators of chaotic spatiotemporal systems. It leverages diffusion models, trained separately on the system's invariant measure, to denoise the emulator rollouts during inference time, pulling trajectories back to the equilibrium distribution and preventing divergence. Thermalization is modular, requiring separate training of the emulator and diffusion model, offering an advantage over methods modifying emulator training for stability. Experiments on 2D Kolmogorov flow and quasi-geostrophic turbulence demonstrate significantly extended stable prediction horizons (over 100K emulator steps).\n\n## Update after rebuttal\n\nMy score remains the same.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the feedback and insightful comments.\n\n> “Dependence on Diffusion Model Quality. The effectiveness of thermalization relies on the quality of the trained diffusion model and its ability to capture the invariant measure. It's not clear whether such a time-stationary invariant measure exists for many practical systems (such as weather and climate), or whether diffusion model can capture the invariant measure. \n\nThis is absolutely correct - the ability of the diffusion model to learn the invariant measure is the limiting component of this framework.\n\n> How does the computational cost of training the diffusion model compare to the cost of training the emulator?\n\nFor the Unet emulator (which is the same architecture as we use to implement the diffusion model), the emulator training is slightly longer; 36 hours vs 24 hours for the diffusion model, on an A100. However, for the DRN emulators, training takes in the region 10 hours due to the much smaller number of parameters.\n\n> Can the diffusion model be adapted to consider non-stationary dynamics, or slowly varying invariant measure -- such as those in climate systems?\n\nIndeed this is the target for future work - conditioning the diffusion model on some time-varying system by some forcing or time-varying parameters, to allow this framework to be applied to non-stationary systems.\n\n- Additionally, at this anonymous link: https://drive.google.com/drive/folders/1b1DHvR-LJvIZtlpvr6hIRQzxn3JKEuwQ?usp=sharing please find two animations of 6 trajectories for Kolmogorov and QG, where we compare the 3 models, and the temporal flow of the thermalizer trajectories is consistent with the numerical model long after the emulator trajectories have gone unstable.In the bottom row of the QG animation, we show the number of thermalization steps over the past 100 steps, during the rollout, such that the adaptive nature of thermalization is shown. A de-anonymised link to animations will be included in the final version of the paper, as the preservation of this temporal flow is a crucial success of our algorithm.", "claims_and_evidence": "Central Claim: Thermalization extends the stable prediction horizon of autoregressive emulators for chaotic systems.\nEvidence is validation on experiments on 2D Kolmogorov flow and QG turbulence, which are highly chaotic and complex PDEs.", "methods_and_evaluation_criteria": "Evaluation metrics are kinetic energy spectra and mean Squared Error (MSE), which are appropriate for Navier Stokes and QG systems.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "The paper directly addresses the problem of long-term instability in autoregressive neural emulators, which is a well-recognized issue in the field of neural network emulators of PDEs, especially for chaotic systems.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n* Novel and Effective Method: Thermalization is a conceptually novel and empirically effective approach to stabilize autoregressive emulators.\n* Modular Training: The separation of emulator and diffusion model training is a significant strength, simplifying the training process.\n* Strong Empirical Validation: The method is rigorously tested on two challenging turbulent flow problems with convincing results.\n\nWeaknesses:\n* Dependence on Diffusion Model Quality: The effectiveness of thermalization relies on the quality of the trained diffusion model and its ability to capture the invariant measure. It's not clear whether such a time-stationary invariant measure exists for many practical systems (such as weather and climate), or whether diffusion model can capture the invariant measure. However, it is a promising direction and worth investigating further.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- How does the computational cost of training the diffusion model compare to the cost of training the emulator?\n- Can the diffusion model be adapted to consider non-stationary dynamics, or slowly varying invariant measure -- such as those in climate systems?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work addresses the issues wherein running surrogate models for long time-frames might cause degenerate behavior, requiring corrective measures. The authors propose a diffusion-model based model and technique, termed "themalization" that prevents the deterioration of the surrogate-model simulation by acting as a "prior" for the autoregressive model. The method is described and applied to two high dimensional dynamical systems. The reviewers generally appreciated the work as novel and performing well on the demonstrated examples. There were a few concerns, including a few missing references and questions about the short-term unrolling accuracy, however the authors responded and I believe that the paper is correct and significant enough for inclusion in ICML.
Decision: Accept (poster)