OmniCast: A Masked Latent Diffusion Model for Weather Forecasting Across Time Scales

TL;DR

OmniCast addresses the challenge of accurate subseasonal-to-seasonal weather forecasting by unifying short- and long-horizon predictions through a two-stage latent diffusion framework. It uses a continuous VAE to map weather states to a compact latent space and a diffusion-based transformer to jointly model the entire future sequence via masked token unmasking, mitigating error accumulation typical of autoregressive methods. The approach achieves state-of-the-art performance on ChaosBench across deterministic, probabilistic, and physics-based metrics, while delivering substantial efficiency gains (10x–20x faster) and stable rollouts up to 100 years. This work demonstrates a scalable path toward fast, reliable probabilistic forecasts across multiple time scales using a single, end-to-end trained model.

Abstract

Accurate weather forecasting across time scales is critical for anticipating and mitigating the impacts of climate change. Recent data-driven methods based on deep learning have achieved significant success in the medium range, but struggle at longer subseasonal-to-seasonal (S2S) horizons due to error accumulation in their autoregressive approach. In this work, we propose OmniCast, a scalable and skillful probabilistic model that unifies weather forecasting across timescales. OmniCast consists of two components: a VAE model that encodes raw weather data into a continuous, lower-dimensional latent space, and a diffusion-based transformer model that generates a sequence of future latent tokens given the initial conditioning tokens. During training, we mask random future tokens and train the transformer to estimate their distribution given conditioning and visible tokens using a per-token diffusion head. During inference, the transformer generates the full sequence of future tokens by iteratively unmasking random subsets of tokens. This joint sampling across space and time mitigates compounding errors from autoregressive approaches. The low-dimensional latent space enables modeling long sequences of future latent states, allowing the transformer to learn weather dynamics beyond initial conditions. OmniCast performs competitively with leading probabilistic methods at the medium-range timescale while being 10x to 20x faster, and achieves state-of-the-art performance at the subseasonal-to-seasonal scale across accuracy, physics-based, and probabilistic metrics. Furthermore, we demonstrate that OmniCast can generate stable rollouts up to 100 years ahead. Code and model checkpoints are available at https://github.com/tung-nd/omnicast.

Figures (20)

• Figure 1: OmniCast processes the latent tokens through a transformer backbone that outputs a vector $z_i$ for each position $i$ in the sequence.
• Figure 2: The denoising network $e_\theta$ predicts the noise $\epsilon$ from $z_i$ and $x_i^s$.
• Figure 3: The deterministic network predicts directly $x_i$ from $z_i$.
• Figure 4: Deterministic performance of different methods at lead times from $1$ to $44$ days across three key variables. Solid curves are deep learning methods and dashed curves are numerical methods.
• Figure 5: Physics-based metrics of different methods at lead times from $1$ to $44$ days across three key variables. Solid curves are deep learning methods and dashed curves are numerical methods.