Research map / Probabilistic and scientific modeling / Molecules and proteins
OmniCast: A Masked Latent Diffusion Model for Weather Forecasting Across Time Scales
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes OmniCast, a continuous-latent masked generative transformer with per-token diffusion heads that jointly samples future weather sequences for medium-range and S2S forecasting.
Problem
Autoregressive data-driven weather models accumulate errors and are computationally expensive for subseasonal-to-seasonal (S2S) forecasting, and fail to capture boundary conditions at those horizons.
Approach
Combine a continuous per-frame VAE latent space with a masked generative transformer that predicts masked future tokens via a small per-token diffusion head, training on full future sequences and sampling by iterative unmasking so space-time tokens are generated jointly rather than autoregressively.
Key results
- On ChaosBench S2S (1-44 days), OmniCast matches ECMWF-ENS as one of the top two methods beyond day 10 and outperforms it on CRPS beyond day 15, with lowest bias among baselines.
- At 15-day forecast runtime, OmniCast takes 29 seconds on an A100 at 0.25 degrees versus 480 seconds for GenCast on TPUv5, and 11 seconds vs 224 seconds at 1.0 degree.
- At medium range, OmniCast performs comparably with IFS-ENS and is only slightly behind GenCast across variables and metrics.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use OmniCast instead of Gencast for 15-day inference at 0.25° | 15-day forecasting · 15-day forecast, 0.25° | runtime: 480 seconds -> 29 seconds | Helped |
| Use OmniCast instead of Gencast for 15-day inference at 1.0° | 15-day forecasting · 15-day forecast, 1.0° | runtime: 224 seconds -> 11 seconds | Helped |
| Increase VAE latent dimension (ablation) | Continuous VAE · VAE reconstruction across variables at 1.40625° | reconstruction error: 0.96 -> 0.80 -> 0.71 (T2m); 0.65 -> 0.51 -> 0.43 (U10); 0.62 -> 0.48 -> 0.40 (V10); 48.72 -> 35.42 -> 27.34 (Z500); 0.77 | Helped |
| Remove the auxiliary MSE objective (ablation) | OmniCast · ChaosBench S2S, T850 | RMSE and CRPS | Hurt |
| Apply MSE to all future frames instead of only the first 10 (ablation) | OmniCast · ChaosBench S2S, T850 | RMSE and CRPS | Hurt |
| Train on shorter sequences or smaller forecast intervals rather than the full sequence (ablation) | OmniCast · ChaosBench S2S, T850 | forecast performance | Mixed |
| Use autoregressive framewise unmasking (ablation) | OmniCast · ChaosBench S2S, T850 | SSR | Hurt |
| Use random framewise unmasking rather than fully randomized token unmasking (ablation) | OmniCast · ChaosBench S2S, T850 | SSR | Hurt |
| Set diffusion temperature below 1 (ablation) | OmniCast · ChaosBench S2S, T850 | SSR and other forecast metrics | Hurt |
| Increase diffusion temperature, but not as high as 1.5 (ablation) | OmniCast · ChaosBench S2S, T850 | SSR, RMSE, and CRPS | Mixed |
| Use VQ-VAE instead of the continuous VAE (ablation) | OmniCast VAE · Weather-state reconstruction | reconstruction error: 2 to 3 times higher | Hurt |
| Compress across both spatial and temporal dimensions with Video VAE instead of per-frame VAE (ablation) | OmniCast VAE · Weather-state reconstruction at equivalent effective compression ratio | reconstruction quality | Hurt |
Surprising findings
- Naively applying the deterministic MSE loss to all future frames hurts S2S performance, so restricting it to the first 10 frames works best.
- Models trained on shorter sequences or smaller intervals excel at short range but underperform at S2S horizons due to error accumulation.
Limitations
- Short-lead RMSE and SSIM are slightly worse than baselines, since the model targets full sequences rather than short-range accuracy.
- Medium-range performance is only comparable to, not better than, GenCast and IFS-ENS.
- VAE reconstruction error imposes an upper bound on forecasting performance; no ablation of GenCast or NeuralGCM at S2S due to compute.
Experimental setup
Datasets: ERA5 WeatherBench2 ChaosBench
Benchmarks: ChaosBench WeatherBench2
Metrics: RMSE absolute bias SSIM spectral divergence spectral residual CRPS spread/skill ratio
Baselines: PanguWeather GraphCast UKMO-ENS NCEP-ENS CMA-ENS ECMWF-ENS ClimaX Stormer Gencast IFS-ENS
Builds on: MaskGIT Masked Autoencoder denoising diffusion GenCast PDEArena UNet
Closest papers by meaning
- Scaling transformer neural networks for skillful and reliable medium-range weather forecasting (NeurIPS 2024 · similarity 0.87)
- DiffDA: a Diffusion model for weather-scale Data Assimilation (ICML 2024 · similarity 0.87)
- ABC: Any-Subset Autoregression via Non-Markovian Diffusion Bridges in Continuous Time and Space (NeurIPS 2026 · similarity 0.86)
- Probablistic Emulation of a Global Climate Model with Spherical DYffusion (NeurIPS 2024 · similarity 0.86)
- Latent Diffusion Transformer for Probabilistic Time Series Forecasting (AAAI 2024 · similarity 0.86)
- Elucidated Rolling Diffusion Models for Probabilistic Forecasting of Complex Dynamics (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.