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

AI summary

Overview

• Research area: Data-driven weather forecasting, specifically probabilistic machine learning models that span medium-range and subseasonal-to-seasonal (S2S) prediction, combining variational autoencoders, masked generative modeling, and diffusion.
• Technical level: Intermediate — the paper assumes familiarity with VAEs, transformers, and diffusion models, though the core idea (generate all future states at once instead of one step at a time) is explainable without that background.
• Scope: The paper introduces OmniCast, a two-stage latent diffusion model evaluated on WeatherBench2 (medium range) and ChaosBench (S2S), and reports accuracy, physics-based, probabilistic, and efficiency results against deep-learning and numerical baselines.

What This Paper Is About

Deep-learning weather models do well up to roughly two weeks but degrade at longer S2S horizons of two to six weeks, largely because they predict in a loop — feed each short-interval prediction back as input and repeat — which compounds errors and makes multi-step fine-tuning computationally prohibitive. OmniCast instead learns the distribution of the entire sequence of future weather states at once, given the initial condition, by compressing weather fields into a continuous latent space and then generating all future latent tokens together through an iterative unmasking process.

Key Contributions

1. A unified architecture for weather forecasting across timescales. OmniCast pairs a continuous VAE that compresses raw weather data into a lower-dimensional latent space with a masked generative transformer that generates a sequence of future latent tokens conditioned on initial tokens.
2. A continuous-token masked generative framework with a per-token diffusion head. Rather than discretizing latent tokens and training with cross-entropy (which the paper argues is poorly suited to weather data with many physical variables), OmniCast models each token's continuous distribution with a small MLP denoising network operating on the transformer's output vector.
3. A deterministic auxiliary objective applied selectively. An MSE head predicts the first 10 future frames with exponentially decreasing per-frame weights, while masked-token prediction is applied across the whole sequence; the paper shows this combination outperforms both removing the MSE loss and applying it to all frames.
4. Efficiency and stability demonstrations. OmniCast is reported as 10× to 20× faster than leading probabilistic methods at medium range, requires a single forward pass through the transformer backbone per sampling iteration (subsequent diffusion steps run only through a compact MLP head), and produces stable rollouts up to 100 years ahead.

Main Findings

• S2S accuracy at longer leads: On ChaosBench, OmniCast performs slightly worse than baselines on RMSE and SSIM at short lead times, but its relative performance improves with lead time and it matches ECMWF-ENS as one of the top two methods beyond day 10.
• Bias: OmniCast shows the lowest bias among all baselines, maintaining near-zero bias across T850, Z500, and Q700.
• Physical consistency: On Spectral Divergence (SDIV) and Spectral Residual (SRES), OmniCast achieves substantially better physical consistency than other deep-learning methods and often outperforms all baselines.
• Probabilistic skill: On CRPS and Spread/Skill Ratio, OmniCast and ECMWF-ENS are the two leading methods; OmniCast trails ECMWF-ENS at shorter leads but outperforms it beyond day 15.
• Medium-range competitiveness: Against Gencast and IFS-ENS, OmniCast performs comparably with IFS-ENS across ensemble RMSE, CRPS, and SSR, and is only slightly behind Gencast.
• Training cost: OmniCast was trained for 4 days on 32 NVIDIA A100 GPUs in a single stage, versus 5 days on 32 TPUv5e devices for Gencast (which uses a two-stage pipeline, pretraining at 1.0° then finetuning at 0.25°) and 10 days on 128 TPUv5e devices for NeuralGCM.
• Inference speed: For a 15-day forecast at 0.25° resolution, Gencast requires 480 seconds on TPUv5 while OmniCast takes 29 seconds on an A100; at 1.0°, OmniCast takes 11 seconds versus 224 seconds for Gencast on the same hardware.
• Deterministic objective ablation: Removing the MSE objective degrades both RMSE and CRPS, especially at short lead times; applying MSE to all future frames is counterproductive; applying it only to the first 10 frames gives the best RMSE and CRPS.
• Training sequence length ablation: Models trained on shorter sequences or smaller intervals excel at short- and medium-range forecasting but underperform at S2S scales because of error accumulation.
• Unmasking order ablation: Fully randomized unmasking across space and time achieves the best SSR scores; autoregressive and random-framewise alternatives produce under-dispersive ensembles.
• Diffusion temperature ablation: τ below 1 produces under-dispersive ensembles; τ = 1.5 pushes samples away from the mean and hurts RMSE and CRPS; τ = 1.3 was identified as the best balance and used in the main experiments.

Methodology in Plain English

OmniCast is trained in two stages. First, a VAE compresses each weather state into a smaller continuous representation — for S2S, a state of shape 69×128×256 becomes a latent map of 1024×8×16, a 16× spatial reduction. The paper argues for a continuous latent space rather than a discrete, vector-quantized one because weather data with many variables (the paper's illustrative example uses 100 variables at 32 bits each) would otherwise require an extremely aggressive compression: a discrete vocabulary of 2^13 = 8192 with 4× spatial reduction gives a compression ratio of roughly 3938, whereas a continuous latent of dimension D = 16 gives a ratio of 100.

Second, a transformer learns to generate future latent tokens. During training, a random subset of future tokens is replaced with a learnable mask token, and the model must predict those masked tokens from the initial condition and the visible tokens. Because masked tokens are continuous, a small diffusion MLP head models each token's distribution using the transformer's output vector at that position as conditioning; a separate deterministic MLP head simultaneously predicts the first 10 frames with exponentially decreasing weights.

At inference, generation starts from a fully masked future sequence and unmasks random subsets of tokens over a cosine schedule until all tokens are produced, then the VAE decoder maps them back to weather fields. For S2S, the model forecasts 44 future states at 24-hour intervals (lead times 1 to 44 days) using 44 iterations, a diffusion temperature of τ = 1.3, and an ensemble of 50 sequences per initial condition. The transformer backbone follows the MAE encoder-decoder design with 16 layers, 16 attention heads, hidden dimension 1024, and dropout 0.1; the diffusion head has six residual blocks of width 2048 with AdaLN conditioning, trained with a 1000-step linear noise schedule resampled to 100 steps at inference.

For medium-range forecasting, the same recipe is applied at 0.25° resolution (a 69×721×1440 state compressed to 256×45×90), with two-step-ahead prediction at 12-hour intervals and autoregressive sampling, one sampling iteration per frame, τ = 1.0, and 50 ensemble members.

Why This Matters

Impact on research. The paper challenges the dominant autoregressive paradigm in data-driven weather forecasting, showing that joint generation across space and time can extend skill into the S2S range where prior deep-learning methods degrade. It also provides evidence about design trade-offs — sequence length, masking order, auxiliary deterministic losses, and sampling temperature — that other generative weather models can reuse.

Real-world applications:

• Subseasonal outlooks for drought, flood, and heatwave risk, which the paper identifies as a key motivation for S2S prediction.
• Resource management and planning decisions that depend on conditions two to six weeks out.
• Ensemble forecasting at scales where running large numerical ensembles is computationally constrained.
• Operational forecasting pipelines that need fast medium-range ensemble products — the reported 29-second 15-day forecast at 0.25° is framed as evidence of practicality.

Industry relevance. Energy, agriculture, insurance, logistics, and disaster-response organizations rely on forecast ensembles for decisions; the paper's efficiency claims (10× to 20× faster at medium range, single-stage training, far lower inference cost than Gencast, NeuralGCM, and IFS-ENS) target the cost barrier that has limited large ensembles in operational settings.

Future Directions

• VAE capacity versus transformer capacity: The paper explicitly names the fundamental trade-off between VAE reconstruction quality and transformer modeling capacity as future work; the VAE imposes an upper bound on forecasting performance.
• Richer generative frameworks: The authors suggest exploring more sophisticated generative frameworks to enhance the diffusion objective.
• Reconstruction limits at low resolution: The VAE trained on 1.40625° data performs worse than the one trained on 0.25° data at the same compression ratio and latent dimension, and the latent dimension was not increased beyond 1024 because of training difficulties with the diffusion objective — leaving headroom unexplored.
• Baseline coverage: Gencast and NeuralGCM could not be run for S2S due to computational demands, and Fuxi-S2S was excluded because it forecasts daily averages rather than point-in-time states, so the S2S comparison set remains incomplete.

Target Audience

Machine learning researchers working on generative models and scientific forecasting; meteorologists and climate scientists interested in data-driven S2S prediction and physical-consistency metrics; and practitioners in operational forecasting or weather-dependent industries who care about ensemble skill, inference cost, and training efficiency.

Authors’ 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.