Thesis: Adaptive Graph Neural Networks for Weather Forecasting

Context

Over the past decades, numerical weather prediction (NWP) models have improved their ability to make accurate weather forecasts. These models roll out forecasts based on initial conditions and complex physical equations, making them very (computationally) expensive to run.

Recently, deep learning techniques and models have been applied in the field of weather forecasting. These data-driven techniques have been able to outperform NWP models on a variety of meteorological variables, and are orders of magnitude more computationally efficient than NWP models [1]. The currently most promising model architectures for weather forecasts are graph neural networks (GNNs), which previously also proved to be well suited to model complex dynamic systems like fluid dynamics [2].

Problem statement

The state-of-the-art data-driven models (e.g. Graphcast [1], AIFS [3]) use a static (fixed) hierarchical (multiscale) graph to efficiently pass messages across larger areas of the graph [4]. The 2020 paper on MeshGraphNets, which introduced the encode-process-decode architecture for GNNs, suggested the use of an adaptive graph to increase resolution in areas with more dynamics [5]. Even though these adaptive graphs showed promising results on multiple physical simulations, they have not yet been applied to the problem of weather forecasting.

The dynamics of the weather aren’t equal everywhere; some areas have very little spatial variance (e.g. high-pressure areas), while others have a much higher degree of dynamics (e.g. cyclones), so we hypothesise that an adaptive graph can help improve accuracy in these high-variance areas.

Research options

There are multiple different ways to implement the graph adaptivity, listed here in order of increasing complexity (for the model to learn and to implement):

1. Heuristic-based: use some heuristic (e.g. the spatial or temporal gradient of some meteorological variables) to determine the node positions using some deterministic algorithm. Then perform Delaunay triangulation on the nodes.
2. Learned heuristic-based: let the model learn to construct a heuristic (instead of coming up with one ourselves, e.g. using a CNN), and then base the node positions on this heuristic as before.
3. Direct node positioning: instead of using a heuristic, let the model predict the positions of the nodes directly from the data. Then use a Delaunay triangulation to construct the edges.
4. Direct graph construction: do not only predict the node positions, but also predict which nodes are connected to each other (bypasses the use of the triangulation step).

The obvious way forward is to start with the simplest methodology, then move our way to more complex options, while constantly evaluating if this added complexity actually improves skill and/or performance.

Doing this for a multiscale graph adds the complexity of determining which of the nodes are also in coarser meshes (to connect over longer distances). This could be done by first predicting the nodes in the coarser grids, or by also predicting a label for the nodes.

Research Question

Hypotheses

H1: Adaptive graphs can improve the skill of medium-range weather forecasts, especially in high-variance areas.

H2: Adaptive graphs can improve runtime performance of medium-range weather forecasts (for an equal accuracy) by using fewer nodes.

H3: Determining the node positions as part of the training improves skill over a variance/gradient-based approach.

(H4): The use of adaptive graphs can enable us to train LAMs by limiting the training objective to a limited area.

Methodology

We first create a good baseline largely inspired by the Graphcast model on a small yet representative subset of the ERA5 dataset. We perform hyperparameter tuning to get the model performance as good as possible within the computational constraints.

The model can then be adjusted to make use of an adaptive graph system. To the extent possible, we use the same hyperparameters as for the baseline to keep the comparison as equal as possible. We compare both models on next-timestep error and on rollouts for up to +120 h.

References

1. Lam, R., Sanchez-Gonzalez, A., Willson, M., et al. (2023). Learning skillful medium-range global weather forecasting. Science, 382(6677), 1416–1421.
2. Sanchez-Gonzalez, A., Godwin, J., Pfaff, T., Ying, R., Leskovec, J., & Battaglia, P. W. (2020). Learning to simulate complex physics with graph networks. Proceedings of the 37th International Conference on Machine Learning.
3. Lang, S., Alexe, M., Chantry, M., et al. (2024). AIFS—ECMWF’s data-driven forecasting system. arXiv.
4. Fortunato, M., Pfaff, T., Wirnsberger, P., Pritzel, A., & Battaglia, P. (2022). MultiScale MeshGraphNets. 2nd AI4Science Workshop at the 39th International Conference on Machine Learning.
5. Pfaff, T., Fortunato, M., Sanchez-Gonzalez, A., & Battaglia, P. W. (2021). Learning mesh-based simulation with graph networks. International Conference on Learning Representations.