Source-linked AI summary
Forecasting Global Weather with Graph Neural Networks
Ryan Keisler
TL;DR
The paper addresses global weather forecasting with data-driven methods that could complement numerical weather prediction. It builds a graph-neural-network model that advances a dense 3D atmospheric state in six-hour steps and chains those steps into multi-day forecasts. Performance improves on previous data-driven approaches and is comparable to operational physical models under specified evaluation conditions.
Problem
Machine learning is being explored to improve numerical weather prediction through better, faster, or larger-ensemble forecasts.
Method
A message-passing graph neural network advances a dense 3D atmospheric state by six hours, with autoregressive chaining for longer forecasts and training on ERA5 or GFS data.
Results
Forecast performance improves upon previous data-driven approaches and is comparable to operational GFS and ECMWF models at 1-degree scales with reanalysis initial conditions.
Takeaways & Limitations
The model supports skillful multi-day global forecasts and can be connected to live GFS forecasts to form a hybrid physics+ML system.
Takeaways & Limitations
Comparisons with operational forecasts are affected because the data-driven forecasts use ERA5 initial conditions that may contain information from up to 12 hours into the future.
Abstract
from arXiv · showhide
We present a data-driven approach for forecasting global weather using graph neural networks. The system learns to step forward the current 3D atmospheric state by six hours, and multiple steps are chained together to produce skillful forecasts going out several days into the future. The underlying model is trained on reanalysis data from ERA5 or forecast data from GFS. Test performance on metrics such as Z500 (geopotential height) and T850 (temperature) improves upon previous data-driven approaches and is comparable to operational, full-resolution, physical models from GFS and ECMWF, at least when evaluated on 1-degree scales and when using reanalysis initial conditions. We also show results from connecting this data-driven model to live, operational forecasts from GFS.
1 Introduction
The paper develops a data-driven global weather forecasting model motivated by improving successful numerical weather prediction with faster, more accurate, or more numerous forecasts. It uses dense atmospheric representations and graph neural networks rather than conventional convolutional networks.
- Machine learning is pursued to improve numerical weather prediction through better, faster, or larger-ensemble forecasts.
- The paper presents a data-driven machine-learning model for global weather forecasting, building on earlier efforts to emulate global numerical weather prediction.
- The model forecasts a substantially denser physical system than earlier data-driven approaches.Each step predicts six variables across 13 pressure levels on a 1-degree grid, representing approximately 5,000,000 physical quantities.
- The denser representation is intended to move closer to the physical system modeled by traditional numerical weather prediction.The authors frame the model as nonlinear interpolation between previously observed spatiotemporal patterns.
- The approach uses message-passing graph neural networks instead of convolutional neural networks.The stated motivation is that GNNs offer a more general architecture and flexible spatial graphs suited to spherical geometry.
2 Data
The model is trained and validated primarily on ERA5 reanalysis data, with additional experiments using GFS forecasts. Both datasets use the same atmospheric variables and 1-degree, 13-level grid, but differences prevent straightforward interoperability.
- The study trains and validates on ERA5 reanalysis data or a subset of NOAA GFS forecasts.
- ERA5 provides hourly 3D atmospheric-state snapshots from 1950 to the present using decades of observations and 4D-Var data assimilation.
- The atmospheric inputs use six physical variables on a 1-degree latitude/longitude grid and 13 pressure levels.The variables are temperature, geopotential height, specific humidity, and three wind components.
- The GFS experiment uses 712 forecasts issued between 2021-04-23 and 2021-10-18, retaining three-hourly outputs through 249 forecast hours.
- ERA5 and GFS share the nominal variables and grid but differ in interpolation details, preventing straightforward interoperability.One cited difference concerns values below the surface for pressure levels exceeding surface pressure.
- Additional model inputs include ERA5 land-sea mask and orography data plus analytic top-of-atmosphere solar radiation.
3 Model
The forecasting system is a machine-learned message-passing graph neural network trained on historical atmospheric data. Its design follows and substantially draws on prior GNN-based physical simulation work.
- The model uses message-passing graph neural networks whose parameters are learned from historical training data.The implementation uses JAX, Jraph, and Haiku.
3.1 Architecture
The architecture maps atmospheric data from a latitude/longitude grid into an icosahedron latent space, processes it with graph message passing, and decodes a six-hour atmospheric-state change back to the original grid.
- 3.1 Architecture: The Encoder, Processor, and Decoder map between physical latitude/longitude data and abstract icosahedron-grid features.
- 3.1 Architecture: The icosahedron intermediate grid has 5,882 nodes versus 65,160 original latitude/longitude nodes, with approximately 3-degree spacing.
- 3.1 Architecture: Each component is implemented as a message-passing GNN operating on nodes and directional edges with learned node and edge updates.
- Encoder: The Encoder gathers local spatial and channel information from nearby latitude/longitude nodes into each icosahedron node.
- Processor: The Processor performs nine rounds of message passing among each icosahedron node and its immediate neighbors, with residual connections.
- 3.1 Architecture: Figure 2 depicts local bipartite encoding and decoding connections alongside local icosahedron-grid processing.
- Decoder: The Decoder maps processor outputs back to the latitude/longitude grid, using a skip connection from the original 78 atmospheric variables.It predicts the six-hour change, which is added to the initial state; six hours balances rollout iterations against dynamical complexity.
3.2 Model Size and Latency
The model is compact and fast enough to generate multi-day forecasts rapidly, while showing no evidence of overfitting on the training dataset.
- 6.7M parameters occupy 27 MB using float32 representation.
- 80,000X smaller than the 2.1 TB of ERA5 data used for training.
- Validation and test losses are indistinguishable from training loss, with no evidence of overfitting.
- 0.04 seconds produces one 6-hour model step on an NVIDIA A100 GPU.
- 0.8 seconds creates a 5-day forecast after model-loading and compilation overhead.
3.3 Training
Training combines staged optimization, multi-resolution data, multi-step rollouts, and physically motivated loss normalization to improve forecast training. The final procedure also includes density and area-based rescaling, while Figure 3 illustrates predicted six-hour atmospheric changes.
- Training schedule: 5.5 days of training on one NVIDIA A100 GPU cost approximately $370.The schedule used 3.5 days at lr=3e-4, 1 day at lr=3e-5, and 1 day at lr=3e-6.
- Multi-resolution training: Training switched from 2-degree to 1-degree data across three rounds, producing better results than training on 1-degree data throughout.
- Multi-step loss: Multi-step loss rolled out the model for approximately 10 steps and accumulated loss at each 6-hour step.The training rounds used progressively larger 4-, 8-, and 12-step losses, corresponding to 1-, 2-, and 3-day rollouts.
- Loss normalization: Loss normalization rescales each physical variable to unit variance in its 3-hour temporal difference before calculating MSE.The normalization targets atmospheric dynamics and gives a physically meaningful error comparable across pressure levels.
- Additional rescaling: The final model also rescales data by nominal static air density at each pressure level and weights grid losses by pixel area using cos(lat).Density rescaling had no strong performance impact.
- Qualitative example: Figure 3 compares ERA5 six-hour changes with model predictions for geopotential height, temperature, and humidity.The caption states that predictions use only the initial state and that raster figures use a 1-degree latitude/longitude grid.
4 Results
The model accurately predicts six-hour atmospheric changes and maintains useful behavior across multi-step rollouts, although smoothing and instabilities emerge at longer horizons. It outperforms earlier data-driven systems and compares favorably with operational models under the paper’s evaluation conditions, while also supporting rapid live-GFS extensions.
- 4.1 Single-step results: The model learns to predict the atmospheric state’s change over the next 6 hours from the current state.Example wind-component predictions accurately match ERA5 six-hour changes using only the initial state.
- 4.2 Multi-step results: A 12-step, 3-day autoregressive rollout generally tracks ERA5’s large-scale 850-hPa specific-humidity flows, though predictions become smoother over time.Each predicted state becomes the input to the next six-hour step.
- 4.3 Stability: The forecast remains numerically stable for approximately 6 days, but hexagon-patterned instabilities begin beyond that horizon.The instability is associated with the icosahedron-grid processor and was not observed in a simpler latitude/longitude architecture.
- 4.4 Comparison to data-driven models: The model outperforms the previous data-driven approaches of Rasp and Thuerey [2021] and Weyn et al. [2020] on globally evaluated 2016 metrics.The comparison uses a 1.0-degree grid with FWHM=1.5-degree spherical Gaussian smoothing.
- 4.5 Comparison to operational NWP models: The data-driven model generally outperforms GFS v15.2 and is comparable to ECMWF on 2020 extratropical metrics under matched coarse-scale evaluation conditions.The comparison accounts approximately for ERA5’s assimilation-window advantage by shifting the data-driven metrics 6 hours earlier.
- 4.6 Connecting to live GFS runs: Connected to live GFS forecasts, the model produced a high-correlation approximation of forecast-to-forecast temperature and humidity changes 22 minutes before publication.The example extended a forecast by 3 days, corresponding to the time GFS required to publish that forecast output.
5 Discussion
The discussion attributes the model’s performance to dense atmospheric simulation, graph-based processing, hardware-enabled scale, and a temporal-difference loss heuristic, while acknowledging limited causal understanding and optimization constraints.
- Interpretation and limitations: The study optimized forecast performance through many small experiments rather than clean ablation studies, leaving its explanation of why the model works partly subjective.The authors explicitly prioritize experiment-informed interpretation over clarity of causal understanding.
- GNNs: Message-passing GNNs support physically uniform neighborhood aggregation on the sphere, with mild early evidence of improvement over latitude-dependent CNN neighborhoods.The authors emphasize that this architectural choice was motivated by prior success in physical-system simulation, but they did not fully test its generality.
- Simulating a dense physical system: The model evolves key prognostic variables with short time steps on a relatively dense 3D grid, bringing it closer to traditional NWP than prior data-driven systems.Despite this design, the model remains roughly 10X coarser in each spatial dimension and 100X coarser in time than operational NWP simulations.
- GPU hardware and memory management: A 40-GB NVIDIA A100 GPU and gradient checkpointing enabled a larger model and losses accumulated over many rollout steps.Checkpointing reduced instantaneous GPU memory usage, allowing multi-step training within the available hardware budget.
- Loss: Temporal-difference rescaling performed significantly better than standard-deviation rescaling when reducing multidimensional MSE values into the training loss.The loss aggregates errors across latitude, longitude, pressure level, physical variable, and rollout time step.
- Time step: A 6-hour time step outperformed a 12-hour step, presumably because shorter intervals have simpler dynamics to model.The authors selected 6 hours as the final time step.
- Loss: The loss design retains human judgment about the target metric, while directly optimizing across many rollout steps may be ineffective.This limits the appeal of replacing the heuristic with a fully end-to-end-learned alternative.
6 Conclusion
The paper presents a graph-neural-network system that advances global atmospheric states in six-hour increments and chains those steps into multi-day forecasts. Its performance improves on earlier data-driven approaches, is comparable to operational physical models at 1-degree scales with reanalysis initialization, and supports a hybrid system connected to live GFS forecasts.
- Conclusion: The system uses local information to advance the current 3D atmospheric state by six hours, chaining steps into forecasts several days ahead.The conclusion describes the approach as data-driven global weather forecasting with graph neural networks.
- Conclusion: Forecast performance improves upon previous data-driven approaches and is comparable to GFS and ECMWF physical models at 1-degree scales with reanalysis initial conditions.The comparison is explicitly limited to the stated evaluation scale and initialization setting.
- Conclusion: The model can connect to live GFS forecasts to form a hybrid physics+ML system that anticipates the physics-based model’s output.The conclusion identifies this live operational connection as an additional demonstrated capability.
- Future work: Future work includes finer spatial resolution, adaptive mesh refinement, large ensemble generation, and additional data-driven forecasting applications.The authors specifically mention moving toward 0.25-degree resolution and using the model to generate large ensembles cheaply.
- Conclusion: The authors view the work as evidence that data-driven techniques can model physical-system dynamics and motivate further machine-learning research for weather forecasting.This conclusion is framed as a direction for follow-up efforts rather than a claim of universal superiority.