OmniCast: A Masked Latent Diffusion Model for Weather Forecasting Across Time Scales
Tung Nguyen, Tuan Pham, Troy Arcomano, Rao Kotamarthi, Ian T. Foster, Sandeep Madireddy, Aditya Grover
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 10× to 20× 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 .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 4d04c457-08da-471b-af61-0a6cc7a02fc9Cited by top-tier papers1
Ask how each one uses itBuilds on9
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn et al.ICLR 2021 · 21,477 citations
- Zero-Shot Text-to-Image GenerationAditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray et al.ICML 2021 · 6,356 citations
- Autoregressive Image Generation without Vector QuantizationTianhong Li, Yonglong Tian, He Li, Mingyang Deng et al.NeurIPS 2024 · 758 citations
- ClimaX: A foundation model for weather and climateTung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta et al.ICML 2023 · 426 citations
- Scaling transformer neural networks for skillful and reliable medium-range weather forecastingTung Nguyen, Rohan Shah, Hritik Bansal, Troy Arcomano et al.NeurIPS 2024 · 165 citations
Related papers
- Probabilistic Transformer For Time Series AnalysisBinh Tang, David S. MattesonNeurIPS 2021 · 150 citations
- TianQuan-S2S: A Subseasonal-to-Seasonal Global Weather Model via Incorporate Climatology StateGuowen Li, Xintong Liu, Yang Liu, Mengxuan Chen et al.ICLR 2026 · 3 citations
- ClimateAR: Multi-Scale Autoregressive Generative Modeling for Climate ForecastingYue Yu, Weiqi Chen, Binqing Wu, Dongliang Cui et al.ICML 2026
- Omni-Weather: A Unified Multimodal Model for Weather Radar Understanding and GenerationZhiwang Zhou, Yuandong Pu, Xuming He, Yidi Liu et al.ICLR 2026
- SutraNets: Sub-series Autoregressive Networks for Long-Sequence, Probabilistic ForecastingShane Bergsma, Timothy Zeyl, Lei GuoNeurIPS 2023 · 14 citations
