Scaling transformer neural networks for skillful and reliable medium-range weather forecasting

2023/12/06 by Tung Nguyen, Rohan Shah, Nguyen, Tung +16 · 1 voice · 30 citations
Computer Science · Earth and Planetary Sciences · Environmental Science · Physics and Astronomy · #Artificial Intelligence (cs.AI) #Atmospheric and Oceanic Physics (physics.ao-ph) #Climate variability and models #FOS: Computer and information sciences #FOS: Physical sciences #Hydrological Forecasting Using AI #Machine Learning (cs.LG) #Meteorological Phenomena and Simulations #cs.AI #cs.LG #physics.ao-ph

paper · pdf · doi:10.48550/arxiv.2312.03876

openalex publication_date 2023/12/06 · arxiv published 2023/12/06 · arxiv updated 2024/10/22 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Weather forecasting is a fundamental problem for anticipating and mitigating the impacts of climate change. Recently, data-driven approaches for weather forecasting based on deep learning have shown great promise, achieving accuracies that are competitive with operational systems. However, those methods often employ complex, customized architectures without sufficient ablation analysis, making it difficult to understand what truly contributes to their success. Here we introduce Stormer, a simple transformer model that achieves state-of-the-art performance on weather forecasting with minimal changes to the standard transformer backbone. We identify the key components of Stormer through careful empirical analyses, including weather-specific embedding, randomized dynamics forecast, and pressure-weighted loss. At the core of Stormer is a randomized forecasting objective that trains the model to forecast the weather dynamics over varying time intervals. During inference, this allows us to produce multiple forecasts for a target lead time and combine them to obtain better forecast accuracy. On WeatherBench 2, Stormer performs competitively at short to medium-range forecasts and outperforms current methods beyond 7 days, while requiring orders-of-magnitude less training data and compute. Additionally, we demonstrate Stormer's favorable scaling properties, showing consistent improvements in forecast accuracy with increases in model size and training tokens. Code and checkpoints are available at https://github.com/tung-nd/stormer.

Cited by

• Long-Range Distillation: Distilling 10,000 Years of Simulated Climate into Long Timestep AI Weather Models
• Experiences Building Enterprise-Level Privacy-Preserving Federated Learning to Power AI for Science
• AI-Based Regional Emulation for Kilometer-Scale Dynamical Downscaling
• Signature Kernel Scoring Rule: A Spatio-Temporal Diagnostic for Probabilistic Weather Forecasting
• OmniCast: A Masked Latent Diffusion Model for Weather Forecasting Across Time Scales
• ARROW: An Adaptive Rollout and Routing Method for Global Weather Forecasting
• Zephyrus: An Agentic Framework for Weather Science
• Swift: An Autoregressive Consistency Model for Efficient Weather Forecasting
• STCast: Adaptive Boundary Alignment for Global and Regional Weather Forecasting
• FlowCast-ODE: Continuous Hourly Weather Forecasting with Dynamic Flow Matching and ODE Solver
• AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions
• Breaking the Statistical Similarity Trap in Extreme Convection Detection
• IndiaWeatherBench: A Dataset and Benchmark for Data-Driven Regional Weather Forecasting over India
• Surya: Foundation Model for Heliophysics
• AtmosMJ: Revisiting Gating Mechanism for AI Weather Forecasting Beyond the Year Scale
• LaDCast: A Latent Diffusion Model for Medium-Range Ensemble Weather Forecasting
• From General to Specialized: The Need for Foundational Models in Agriculture
• Accurate Mediterranean Sea forecasting via graph-based deep learning
• Elucidated Rolling Diffusion Models for Probabilistic Forecasting of Complex Dynamics
• AutomataGPT: Forecasting and Ruleset Inference for Two-Dimensional Cellular Automata
• Physics-Informed Teleconnection-Aware Transformer for Global Subseasonal-to-Seasonal Forecasting
• Arnoldi Singular Vector perturbations for machine learning weather prediction
• Predictability-Aware Compression and Decompression Framework for Multichannel Time Series Data
• NeuralOM: Neural Ocean Model for Subseasonal-to-Seasonal Simulation
• PEAR: Equal Area Weather Forecasting on the Sphere
• Mixture of Low Rank Adaptation with Partial Parameter Sharing for Time Series Forecasting
• Utilizing Strategic Pre-training to Reduce Overfitting: Baguan -- A Pre-trained Weather Forecasting Model
• Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency
• LatticeVision: Image to Image Networks for Modeling Non-Stationary Spatial Data
• Physics-Assisted and Topology-Informed Deep Learning for Weather Prediction
• U-Cast: A Surprisingly Simple and Efficient Frontier Probabilistic AI Weather Forecaster
• Cellular Development Follows the Path of Minimum Action

Discussions

• Scaling Transformers for skillful and reliable medium-range weather forecasting [hn, 2 points, 0 comments]

Related

• Improving Medium Range Severe Weather Prediction through Transformer Post-processing of AI Weather Forecasts
• CaFA: Global Weather Forecasting with Factorized Attention on Sphere
• Gridded Transformer Neural Processes for Large Unstructured Spatio-Temporal Data
• TENT: Tensorized Encoder Transformer for Temperature Forecasting
• OmniCast: A Masked Latent Diffusion Model for Weather Forecasting Across Time Scales