Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos

Thank you very much for the feedback and insightful comments.

“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.

This is absolutely correct - the ability of the diffusion model to learn the invariant measure is the limiting component of this framework.

How does the computational cost of training the diffusion model compare to the cost of training the emulator?

For 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.

Can the diffusion model be adapted to consider non-stationary dynamics, or slowly varying invariant measure -- such as those in climate systems?

Indeed 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.

• 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.

Thank you very much for taking the time to assess our work, and for the insightful comments:

• “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?” Not 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. 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.

• 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.

• $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.

• 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.

• 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.

• 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.

• 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.

• 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.

We thank the reviewer for their insightful comments.

• 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.

• 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).

• 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.

• Thanks for pointing out the typo in eq 7 – you are correct!

• 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).

• 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.

We thank the reviewer for their insightful comments.

• 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.

• 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.

• 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.

• 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.

• 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.