Source-linked AI summary
Scaling transformer neural networks for skillful and reliable medium-range weather forecasting
Tung Nguyen, Rohan Shah, Hritik Bansal, Troy Arcomano, Romit Maulik, Veerabhadra Kotamarthi, Ian Foster, Sandeep Madireddy, Aditya Grover
TL;DR
Weather forecasting needs accurate, efficient models, but prior deep-learning systems often use complex architectures whose contributions are difficult to isolate. Stormer uses a minimally modified transformer with randomized dynamics forecasting and related training components, achieving strong short- and medium-range performance and outperforming current methods beyond 7 days with far less data and compute.
Problem
Deep-learning weather forecasters can be accurate but often rely on complex customized architectures with insufficient ablation analysis, making their effective ingredients difficult to identify.
Method
Stormer minimally modifies a standard vision transformer with weather-specific embedding, randomized dynamics forecasting, and pressure-weighted loss.
Results
Stormer is competitive for 1–7-day forecasts, outperforms state-of-the-art methods beyond 7 days, and uses orders-of-magnitude fewer GPU hours than baselines.
Takeaways & Limitations
Randomized forecasting lets one model generate and combine multiple forecasts for each target lead time, while performance improves with model capacity and data size.
Takeaways & Limitations
The training setup uses a uniform distribution over 6-, 12-, and 24-hour intervals, and multi-step finetuning destabilized training because step losses have different magnitudes.
Abstract
from arXiv · showhide
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.
1. Introduction
Stormer shows that a simple vision-transformer architecture, paired with a carefully analyzed training recipe, can deliver strong weather-forecasting performance while using substantially fewer resources than baselines.
- 1. Introduction: Deep-learning weather models can produce forecasts in seconds rather than the hours required by typical numerical weather-prediction models.They are trained on historical weather data such as ERA5 reanalysis.
- 1. Introduction: The model combines a standard ViT backbone with weather-specific embedding, randomized dynamics forecasting, and pressure-weighted loss.These components were identified through extensive ablation studies.
- 1. Introduction: Randomized dynamics forecasting trains the model at varying intervals, then combines multiple rollouts for a target lead time.For example, a 3-day forecast can use twelve 6-hour predictions or six 12-hour predictions.
- 1. Introduction: Stormer achieves competitive forecast accuracy for 1–7 days and outperforms state-of-the-art methods beyond 7 days.The model uses more than 5× lower-resolution data and orders-of-magnitude fewer GPU hours than the baselines.
2. Background and Preliminaries
Weather forecasting predicts future atmospheric states from initial conditions across variables and spatial locations, while data-driven methods differ in how they represent forecast lead times and rollout intervals.
- 2. Background and Preliminaries: Global weather forecasting predicts future atmospheric conditions from initial conditions across variables and spatial locations.The formulation includes a target lead time, atmospheric variables such as temperature and humidity, and a spatial grid over the globe.
- 2. Background and Preliminaries: Direct forecasting trains separate outputs for target lead times, whereas continuous forecasting conditions one model on the desired lead time.Direct forecasting can be challenging for large lead times and costly when many target times are needed.
- 2. Background and Preliminaries: Figure 3 reports that different forecast intervals suit different lead times, while weather-specific embedding and adaptive layer norm outperform their standard alternatives.These preliminary comparisons motivate Stormer’s design choices.
3. Methodology
Stormer uses a simple Transformer framework with weather-specific embedding, randomized dynamics forecasting, pressure-weighted loss, and multi-step finetuning. Its inference averages forecasts from interval combinations, while the embedding models atmospheric-variable interactions more efficiently than standard patch embedding.
- Training and inference: Stormer combines randomized dynamics forecasting, pressure-weighted loss, multi-step finetuning, and forecast averaging within a simple Transformer framework.The model is trained on varying intervals, finetuned over multiple rollout steps, and evaluated using averaged forecasts from interval combinations.
- Training and inference: Randomly sampled 6-, 12-, and 24-hour intervals target both diurnal and longer synoptic-scale atmospheric dynamics.The 6- and 12-hour intervals resolve the day-night cycle, while 24 hours filters that cycle to support longer-scale dynamics.
- Loss design: Pressure weighting prioritizes near-surface variables using atmospheric pressure as a proxy for density, while latitude weighting accounts for spherical-grid non-uniformity.The pressure-weighted loss adapts a weighting introduced in GraphCast to Stormer’s architecture.
- Training and inference: A single conditioned model generates multiple forecasts for a target lead time, and averaging diverse interval combinations provides an ensembling-like inference effect.Homogeneous combinations improve efficiency, whereas best m in n selects validated combinations for greater expressivity.
- Architecture: Weather-specific embedding tokenizes each atmospheric variable independently and uses cross-attention to aggregate variables, reducing sequence length by V and modeling nonlinear interactions.The aggregation module converts the variable-expanded sequence back to hidden dimension D while reducing transformer computation.
- Architecture: Weather-specific embedding outperforms standard patch embedding across lead times from 1 to 10 days, including when trained on a single dataset.Its advantage is attributed to modeling interactions among atmospheric variables rather than treating channels like natural-image channels.
4. Experiments
Stormer is evaluated against leading deep-learning and numerical weather-forecasting systems on WeatherBench 2, with ablations and scaling analyses examining its design and efficiency. It is competitive at short lead times, substantially stronger at longer horizons, and improves with randomized forecasting, pressure-weighted loss, dynamics prediction, larger models, and more tokens.
- 4.1. Comparison with State-of-the-art models: Stormer reaches competitive short-range accuracy and longer-range superiority while using approximately 190× less data than Pangu-Weather and 90× less than GraphCast.The comparison uses 6-hourly data at 1.40625° resolution and 13 pressure levels.
- 4.1. Comparison with State-of-the-art models: At 14 days, Stormer performs 10% −20% better than GraphCast across all 9 key variables.Stormer also outperforms both Pangu-Weather and GraphCast from day 6 onwards by a large margin.
- 4.1. Comparison with State-of-the-art models: Stormer training finished in under 24 hours on 128 A100 GPUs, versus 60 days for Pangu-Weather and 28 days for GraphCast.
- 4.2. Ablation studies: Randomized forecasting consistently outperforms single-interval models, with the advantage widening at longer lead times and no computational overhead.The models share the same architecture and training duration, while multiple forecasts are combined during inference.
- 4.2. Ablation studies: Pressure-weighted loss improves performance, while forecasting weather dynamics is consistently more accurate than predicting complete weather states.The pressure weighting prioritizes variables nearer the surface, and dynamics prediction targets changes between weather conditions.
- 4.3. Scaling analysis: Forecast accuracy improves significantly with larger models and with smaller patches that increase training-token counts.The performance gap between model sizes widens as lead time increases.
5. Related Work
Prior data-driven weather-forecasting work spans deterministic models that improve long-horizon skill and probabilistic models that quantify uncertainty. These approaches often involve specialized architectures, multiple models, or additional physical or probabilistic components.
- Deterministic weather forecasting: Deterministic forecasting research includes convolutional, transformer, and graph-based models, with some improving long-lead accuracy through specialized models or architectures.
- Deterministic weather forecasting: Some existing methods require multiple models or do not disclose complete architecture and training details.
- Probabilistic weather forecasting: Probabilistic weather forecasting combines forecasting architectures with probabilistic objectives to quantify forecast uncertainty.Examples include diffusion objectives and variational autoencoder components.
6. Conclusion and Future Work
The paper presents Stormer as a simple weather-forecasting model whose training recipe yields competitive short-range accuracy and strong performance beyond 7 days. It also identifies future directions involving uncertainty, additional randomization, higher resolution, and larger models.
- Conclusion and Future Work: Stormer combines a simple architecture with randomized iterative forecasting to produce and combine forecasts at different time intervals.
- Conclusion and Future Work: Stormer is competitive at short range and performs exceptionally beyond 7 days while using significantly less data and computing resources.
- Conclusion and Future Work: Future work could use multiple forecasts to quantify uncertainty, randomize additional model components, and evaluate higher-resolution data and larger models.
- Conclusion and Future Work: Deep-learning weather models require careful checks and monitoring when predicting new or uncertain scenarios.
B.1. Stormer architecture
Stormer tokenizes and aggregates atmospheric variables, conditions its transformer blocks on the forecast interval, and predicts weather changes through a multi-phase training objective. Its evaluation reports latitude-weighted RMSE across nine variables and lead times from 1 to 14 days.
- B.1. Stormer architecture: Stormer separately tokenizes variables, aggregates them with cross-attention, embeds the interval δt, and outputs the predicted weather change Δδt.The resulting token sequence has (H/p) × (W/p) tokens after aggregation.
- B.1. Stormer architecture: Each Stormer block uses adaptive layer normalization conditioned on δt within a stack of N transformer blocks.
- B.2.2. Three-phase training: Training proceeds in three phases with rollout steps K equal to 1, 4, and 8, fine-tuning the best preceding checkpoint in phases 2 and 3.
- B.2.2. Three-phase training: Pressure-level variables receive weights proportional to pressure, while surface variables use weights 1 for T2m and 0.1 for U10, V10, and MSLP.
- B.2.5. Software and hardware stack: Stormer and baselines are evaluated using latitude-weighted RMSE for nine key variables at lead times from 1 to 14 days.
B.3. Evaluation protocol
Stormer forecasts are evaluated after regridding all model outputs to a shared 1.40625° resolution, using initial conditions at 00/12 UTC throughout 2020.
- All model forecasts are regridded to the shared 1.40625° (128 × 256) resolution before evaluation.
- Evaluation uses initial conditions at 00/12 UTC for every day in 2020.
C.1. Complete comparison with SoTA models
Stormer matches or exceeds benchmark weather models across short and medium lead times and achieves particularly strong performance at longer horizons, despite using much lower-resolution data.
- Stormer achieves lower RMSE than IFS for most variables, except near-surface temperature at initial lead times, and is only slightly worse than IFS ENS.
- Stormer is on par with or outperforms every benchmark model, with significant gains during later forecast periods.
- Stormer is the first model trained on 1.40625° data to surpass IFS.
- Stormer delivers competitive ACC from 1 to 5 days and outperforms baselines by a large margin beyond 6 days.
- Impact of multi-step fine-tuning: Multi-step fine-tuning significantly improves performance at long lead times when progressing from K=1 to K=8.
- The non-ensemble Stormer is explicitly compared with baseline methods to assess performance without ensemble inference.
C.4. Probabilistic forecasting with IC perturbations
Stormer’s inference strategies support probabilistic forecasting through initial-condition perturbations, but introduce a trade-off between probabilistic quality, deterministic accuracy, and computational efficiency.
- Stormer’s unperturbed forecasts are underdispersive and should not be used for uncertainty estimation without additional randomization.
- Adding four Gaussian noises for each Best m in n interval combination produces 128 ensemble members.
- Initial-condition perturbations significantly improve probabilistic metrics but may reduce deterministic performance at short lead times.
- Homogeneous inference performs competitively with Best m in n while requiring only 3 forward passes instead of n.
C.6. Qualitative results
Qualitative forecasts are visualized from 1- to 14-day lead times across nine key variables, with each display comparing initial conditions, truth, predictions, and bias.
- Stormer qualitative forecasts are presented at lead times ranging from 1 to 14 days for nine key variables.
- Each qualitative figure shows one lead time, with rows for variables and columns for initial condition, ground truth, forecast, and bias.