[Paper Review] Probabilistic Weather Forecasting with Hierarchical Graph Neural Networks

Joel Oskarsson, Tomas Landelius+2 more|arXiv (Cornell University)|Jun 7, 2024

Hydrological Forecasting Using AI5 citations

TL;DR

Introduces Graph-EFM, a probabilistic weather forecasting model using a hierarchical graph neural network to produce spatially coherent ensemble forecasts with only a single forward pass per time step.

ABSTRACT

In recent years, machine learning has established itself as a powerful tool for high-resolution weather forecasting. While most current machine learning models focus on deterministic forecasts, accurately capturing the uncertainty in the chaotic weather system calls for probabilistic modeling. We propose a probabilistic weather forecasting model called Graph-EFM, combining a flexible latent-variable formulation with the successful graph-based forecasting framework. The use of a hierarchical graph construction allows for efficient sampling of spatially coherent forecasts. Requiring only a single forward pass per time step, Graph-EFM allows for fast generation of arbitrarily large ensembles. We experiment with the model on both global and limited area forecasting. Ensemble forecasts from Graph-EFM achieve equivalent or lower errors than comparable deterministic models, with the added benefit of accurately capturing forecast uncertainty.

Motivation & Objective

• Motivate the need for probabilistic forecasting to capture uncertainty in chaotic weather systems.
• Develop a hierarchical graph neural network framework to enable spatially coherent forecasts.
• Introduce Graph-EFM, a latent-variable probabilistic model enabling efficient ensemble sampling.
• Demonstrate Graph-EFM on global (1.5°) and limited-area (10 km) forecasting and assess calibration.
• Provide training objectives and improved GNN layers tailored for probabilistic MLWP.

Proposed method

• Propose a hierarchical mesh graph with multiple spatial scales to propagate information and impose spatial coherence.
• Formulate Graph-EFM as a latent-variable model where Zt is top-level latent representing uncertainty, with Xt conditioned on Zt and past states.
• Define a two-part single-step model: a latent map p(Zt|Xt−2:t−1, Ft) and a predictor p(Xt|Zt, Xt−2:t−1, Ft) realized as a deterministic mapping with skip connections.
• Train with a variational objective (ELBO) using a learned q(Zt|Xt−2:t−1, Xt, Ft) and include a CRPS term for ensemble calibration.
• Use Propagation Networks (a variant of Interaction Networks) to improve information flow from grid to latent top-level and back.
• Optionally evaluate deterministic Graph-FM as a baseline and GraphCast/GraphCast+SWAG as additional baselines.

Experimental results

Research questions

• RQ1 Can a hierarchical GNN with latent variables accurately model the distribution of future weather states?
• RQ2 Does Graph-EFM produce calibrated ensemble forecasts with spatially coherent fields at global and regional scales?
• RQ3 How does the probabilistic Graph-EFM compare to deterministic Graph-FM and existing baselines in RMSE, CRPS, and SpSkR across lead times?
• RQ4 What is the impact of the hierarchical graph design on sampling efficiency and ensemble diversity?
• RQ5 How well can the model handle limited-area modeling with boundary conditions (LAM) using the same framework?

Key findings

• Graph-EFM achieves lower CRPS than baselines across several variables and lead times, indicating better distributional accuracy.
• Ensemble mean from Graph-EFM often improves RMSE over deterministic models, especially at longer lead times.
• Without perturbing initial states, Graph-EFM attains SpSkR near 1, indicating well-calibrated uncertainty; SWAG-based ensembles can be poorly calibrated.
• For LAM (Nordic MEPS data), Graph-EFM provides spatially coherent ensembles, though short-lead RMSE gains are modest.
• The hierarchical graph structure enables efficient sampling of large ensembles (e.g., 80 members in global forecasting in ~200 seconds on a single GPU).
• Extreme weather case studies (e.g., hurricane Laura) illustrate the model’s ability to capture location and wind uncertainty in ensemble forecasts.