Elucidated Rolling Diffusion Models for Probabilistic Weather Forecasting
Salva Rühling Cachay, Miika Aittala, Karsten Kreis, Noah Brenowitz, Arash Vahdat, Morteza Mardani, Rose Yu
While applications of diffusion models to probabilistic forecasting of high-dimensional dynamical systems have shown promising results, these applications often use limited temporal context and predict future snapshots one by one.To explicitly model the sequential nature of the data and the growth of physical uncertainty inherent to chaotic systems, we introduce Elucidated Rolling Diffusion Models (ERDM). ERDM fuses the principled design of Elucidated Diffusion Models (EDM) with a rolling diffusion framework that assigns progressively larger noise levels to farther lead times.We adapt the EDM noise schedule, network preconditioning, and Heun sampler for this specific setting and introduce three key contributions: $(i)$ a novel loss weighting scheme, that puts more weights to middle snapshots where the stochastic to deterministic demixing occurs; $(ii)$ an efficient initialization strategy using a pretrained next-step EDM for the initial window to mitigate sensitivity in nonlinear, chaotic systems; and $(iii)$ a bespoke hybrid sequence architecture for robust spatiotemporal feature extraction under progressive denoising. On 2D Navier–Stokes simulations, ERDM improves probabilistic skill scores by up to $50\%$ over prior conditional diffusion baselines.On ERA5 global weather at $1.5^\circ$ resolution, it matches the ensemble skill of leading forecasters like IFS ENS and NeuralGCM, while being more efficient, requiring only 4 H200 GPUs and 5 days for training.