Elucidated Rolling Diffusion Models for Probabilistic Forecasting of Complex Dynamics

TL;DR

ERDM introduces Elucidated Rolling Diffusion Models which fuse rolling forecast concepts with the EDM diffusion framework to model uncertainty growth in complex dynamics. By crafting a snapshot-aware noise schedule, per-snapshot preconditioning, a temporal 3D denoiser, and an uncertainty-aware loss, ERDM achieves superior long-range probabilistic forecasts on Navier–Stokes and ERA5 data, with competitive spectral realism and calibration versus operational models. The approach outperforms corresponding EDM baselines in CRPS and calibration, while maintaining practical training efficiency and scalable inference via a rolling window. Limitations include higher memory demands of the 3D denoiser and some short-range weaknesses relative to operational physics models, pointing to future work in latent-space variants and broader applications.

Abstract

Diffusion models are a powerful tool for probabilistic forecasting, yet most applications in high-dimensional complex systems predict future states individually. This approach struggles to model complex temporal dependencies and fails to explicitly account for the progressive growth of uncertainty inherent to the systems. While rolling diffusion frameworks, which apply increasing noise to forecasts at longer lead times, have been proposed to address this, their integration with state-of-the-art, high-fidelity diffusion techniques remains a significant challenge. We tackle this problem by introducing Elucidated Rolling Diffusion Models (ERDM), the first framework to successfully unify a rolling forecast structure with the principled, performant design of Elucidated Diffusion Models (EDM). To do this, we adapt the core EDM components-its noise schedule, network preconditioning, and Heun sampler-to the rolling forecast setting. The success of this integration is driven by three key contributions: (i) a novel loss weighting scheme that focuses model capacity on the mid-range forecast horizons where determinism gives way to stochasticity; (ii) an efficient initialization strategy using a pre-trained EDM for the initial window; and (iii) a bespoke hybrid sequence architecture for robust spatiotemporal feature extraction under progressive denoising. On 2D Navier-Stokes simulations and ERA5 global weather forecasting at 1.5-degree resolution, ERDM consistently outperforms key diffusion-based baselines, including conditional autoregressive EDM. ERDM offers a flexible and powerful general framework for tackling diffusion-based dynamics forecasting problems where modeling uncertainty propagation is paramount.

Figures (16)

• Figure 1: ERDM sampling with a window, highlighted in bold purple, of size $W=4$. Top row: ERDM starts at diffusion time ${t}=0$ with snapshots, ${\bm{x}}_1, \dots, {\bm{x}}_W$, corrupted by progressively larger noise levels, $\bar{\sigma}_1(t)<\dots<\bar{\sigma}_W(t)=\sigma_{\text{max}}$. Middle row: After $N=2$ joint denoising steps, the sequence reaches lower noise levels at ${t}=1$ such that $\bar{\sigma}_1(1)=\sigma_{\text{min}}$ and $\bar{\sigma}_{w}(1) = \bar{\sigma}_{w-1}(0)$ for $w>1$, as illustrated in the right-hand panel. The now fully denoised first snapshot, ${\bm{x}}_1$, is returned. Bottom row: The rest of the sequence is shifted one slot to the right, and a fresh pure-noise snapshot is appended to the new window. The cycle then repeats.
• Figure 2: Noise schedule comparison for a sequence length $W=6$, yielding 6 visualized segments $[\bar{\sigma}_w(1), \bar{\sigma}_w(0)]$. The schedules differ by their curvature parameter $\rho$: one using $\rho=7$ (default EDM) and the other $\rho=-10$ (ERDM). Color gradients illustrate segment progression.
• Figure 3: Sketch of one 2D U-Net block and one temporal attention block in our hybrid U-Net topology with noise embedding to both spatial and temporal paths. Dimensions b, c, t, h, w refer to batch, channel, window, height, and width. For simplicity, c, h, w do not change in the sketch.
• Figure 4: Navier-Stokes test rollout over 64 time steps with 50 ensemble members. ERDM superior performance in both CRPS and calibration compared to single- and multi-step EDM baselines, except for the initial 3 timesteps. Beyond timestep 15, ERDM consistently delivers an approximate $50\%$ improvement in CRPS over the next best model (EDM $W=4$), demonstrating particular strength in long-range forecasting scenarios.
• Figure 5: Relative CRPS (lower is better) over single-step EDM baseline as a function of lead time, up to 15 days, for 8 selected variables. ERDM consistently outperforms EDM across most variables and lead times by up to $10\%$. Furthermore, it performs competitively against the state-of-the-art operational physics-based model IFS ENS and the hybrid model NeuralGCM ENS, especially for long lead times, while being more efficient.