Elucidated Rolling Diffusion Models for Probabilistic Forecasting of Complex Dynamics

TL;DR

This paper introduces Elucidated Rolling Diffusion Models (ERDM), a novel framework that unifies rolling forecast structures with high-fidelity diffusion techniques to improve probabilistic forecasting of complex, high-dimensional systems.

Contribution

The paper presents ERDM, integrating EDM components into rolling forecasts with new loss weighting, pre-trained initialization, and a hybrid architecture for better uncertainty modeling.

Findings

ERDM outperforms diffusion-based baselines on Navier-Stokes and weather data.

The framework effectively models uncertainty propagation in complex systems.

ERDM demonstrates improved accuracy in probabilistic forecasting tasks.

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.