Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos

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.

AI simulations of fluid flows have the potential to revolutionise many areas of science and engineering due to immense computational speedups. However current models tend to only be reliable over short timescales, as they build up error and the predicted state degrades over time. This is due to the fact that the AI models are approximations, and so small errors introduced at each step accumulate over many timesteps.We tackle this problem by using an algorithm from the field of image generative modelling called diffusion models, to "denoise" the AI-simulated fluid flow on the fly - a process we call thermalization. This algorithm is able to only remove the errors introduced by the AI-simulator, without disrupting the temporal dynamics of the fluid flow. We demonstrate this approach on multiple flows and multiple AI-simulators. In every case we are able extend the length of reliable rollouts by a factor of 100 when applying the thermalizer to the AI-simulated flow. While we demonstrate thermalization on fluid flows, in principle it could be applied to many other kinds of AI models, including language modelling.