WIND: Weather Inverse Diffusion for Zero-Shot Atmospheric Modeling

TL;DR

WIND introduces a diffusion-forcing, probabilistic foundation model for atmospheric dynamics that learns a task-agnostic prior and solves downstream weather and climate problems at inference time via MMPS-based posterior sampling. By framing probabilistic forecasting, downscaling, sparse reconstruction, conservation enforcement, and counterfactual storytelling as inverse problems with a single pre-trained backbone, WIND forgoes task-specific fine-tuning while preserving physical coherence and uncertainty quantification. Empirical results on ERA5 data demonstrate spectral fidelity, stable long-rollouts, and competitive performance across diverse tasks, with explicit mechanisms to enforce dry air mass conservation and to generate warming-era storylines. This approach offers a computationally efficient, flexible alternative to specialized models, with potential impact on real-time assimilation, risk assessment, and climate storytelling under future forcings.

Abstract

Deep learning has revolutionized weather and climate modeling, yet the current landscape remains fragmented: highly specialized models are typically trained individually for distinct tasks. To unify this landscape, we introduce WIND, a single pre-trained foundation model capable of replacing specialized baselines across a vast array of tasks. Crucially, in contrast to previous atmospheric foundation models, we achieve this without any task-specific fine-tuning. To learn a robust, task-agnostic prior of the atmosphere, we pre-train WIND with a self-supervised video reconstruction objective, utilizing an unconditional video diffusion model to iteratively reconstruct atmospheric dynamics from a noisy state. At inference, we frame diverse domain-specific problems strictly as inverse problems and solve them via posterior sampling. This unified approach allows us to tackle highly relevant weather and climate problems, including probabilistic forecasting, spatial and temporal downscaling, sparse reconstruction and enforcing conservation laws purely with our pre-trained model. We further demonstrate the model's capacity to generate physically consistent counterfactual storylines of extreme weather events under global warming scenarios. By combining generative video modeling with inverse problem solving, WIND offers a computationally efficient paradigm shift in AI-based atmospheric modeling.

Figures (20)

• Figure 1: Training setup and inference capabilities of WIND(a) Training: We apply an independent noise level $\sigma_k$ to each frame in the sequence. The model is trained without explicit noise level information. WIND learns to jointly denoise the sequence, enabling it to handle arbitrary combinations of clean and noisy context frames. (b) Inference:WIND addresses various climate- and weather- related questions by framing them as inverse problems: generate fields $y$ under the constraint $A(x) = y$. We demonstrate how to formulate a task-specific operator $A$ for each question in (b).
• Figure 2: Probabilistic forecast performance We evaluate the 14-day forecast skill using the CRPS (lower is better) averaged over 100 initializations in 2021. We compare the unconstrained WIND baseline against the performance with enforced dry air mass conservation. The large overlap demonstrates that the physics constraint does not degrade the probabilistic forecast skill. WIND outperforms the autoregressive AR-UViT baseline after very few days.
• Figure 3: Power spectra for spatial downscaling. We compare the PSD of the ERA5 ground truth, a specialized FNO and UViT model, and WIND. WIND closely tracks the energy spectrum of ERA5 across all scales, preserving high-frequency details. In contrast, the deterministic FNO baseline exhibit spectral drop-off at high frequencies. While UViT performs on par with our method for surface variables, it struggles with the atmospheric variables $Q$ and $TP$.
• Figure 4: Qualitative comparison of sparse reconstruction.First column: ERA5 ground truth (2m Temperature) and 10% sparse input mask. Top row: ERA5 ground truth (2m Temperature) and the corresponding predictions from WIND, UViT, and Kriging. Bottom row:10% sparse input mask and the prediction error. While WIND and UViT recover physically coherent fields with realistic gradients, Kriging yields overly smooth interpolations that miss fine-grained patterns.
• Figure 5: Long-term stability of dry air mass. The DAM of WIND without constraint drifts around 200 forecast days. ERA5 ground truth shows a seasonal cycle. WIND with DAM guidance strictly enforces conservation for the entire 4-year rollout.