Probabilistic Weather Forecasting with Hierarchical Graph Neural Networks

We thank reviewer PLfK for useful comments. See our response below:

1. Scores for GraphCast*

As we state clearly when introducing this baseline, this is a version of GraphCast trained on the same 1.5° dataset as the other models. It thus has different scores than the original GraphCast model evaluated in Weatherbench. As we emphasize in the paper, we focus here on a fair comparison of models operating at the same resolution and trained on the same data. In the paper we refer to appendix H for comparisons with other state-of-the-art models, trained on higher resolution data.

2. Comparison to conventional models

Results for IFS-ENS are included in appendix H. We realize that this might not be mentioned clearly in the main paper and will clarify. As the paper is mainly about probabilistic modeling, we find it more relevant to include IFS-ENS rather than IFS.

Since the limited area modeling experiment is a pure surrogate modeling task, with data coming from MEPS forecasts, it does not make sense to compare to the MEPS model itself.

3. Comparison to GenCast

See the global author rebuttal.

4. Graph-FM parameter count

When instantiating models of similar size we use the dimensionality $d_z$ of representation vectors as the main scaling factor for model size. The hierarchical architecture in Graph-FM is what enables more flexibility at the same value of $d_z$, resulting also in more parameters. Note that this flexibility does not introduce a significant increase in computational cost, which would be the case if one for example added 6 times as many layers to GraphCast.

5. U-Net connection

We disagree that this connection is overlooked and refer the reviewer to the related work section where we clearly note this similarity. It should however be emphasized that our GNN model can not be seen as equivalent to a U-net. While some hierarchical structure is shared, the GNN layers can not be reduced to simple convolutions, even in the LAM setting. For example, these feature edge representation updates that have no correspondence in convolutional layers.

We would also like to emphasize that we present a general framework that is applicable to different forecasting regions and graph constructions. The exact similarity to a U-net model will depend on these choices. In the specific LAM case from section 5.2 this connection is stronger (although the models are not equivalent), whereas in the global forecasting setting this connection is more conceptual.

6. Physical meaningfulness of VAE implementation

We would like to clarify that while the latent variables are associated with the top level of the graph, this does not mean that they only influence the coarsest spatial scales. As the GNN layers of the predictor are applied through the graph hierarchy the randomness in the latent variables can spread to all parts of the graph and introduce variations on different scales. However, explicitly constraining different latent variables to control variation at different spatial scales is an interesting idea for future work.

We have found that introducing the latent variables closer to the grid points in the model can lead to less spatially coherent forecasts. Graph-EFM (ms) can act as an ablation of this, as the latent variables are there directly associated with the nodes of the single level multiscale graph. We found Graph-EFM (ms) to be harder to train and produce worse forecasts.

7. Reference to workshop paper

By the NeurIPS policy, workshop papers not published in proceedings are not to be counted as publications. Therefore we did not explicitly cite the mentioned workshop paper. However, if it is deemed advisable to do so we would not mind including a reference to it. Perhaps the AC can help clarify the situation.

8. GraphCast* + SWAG baseline

This baseline was not picked arbitrarily, but as described it was inspired by the usage of SWAG in Graubner et al. The point of this baseline is to include a comparison to multi-model ensembles, which we see can produce too little variability. While there are indeed other ways to achieve this, we wanted to focus on a method that has been proposed in the ML weather prediction literature, rather than choosing a method arbitrarily.

9. Motivation for LAM

We currently give some motivation for this in section 5.2, but we do agree that this can be expanded on already in the introduction of the paper. We will change this.

10. Extreme weather case study

We have now included such a case study for Hurricane Laura, that can be found in the global author rebuttal.

11. Fixed variance of latent variable for global predictions

There might be some misunderstanding here, as the variance of the latent variable (Eq. 4) is fixed for both the global and LAM models. However, the variance $\sigma^2_{\alpha, j}$ of the likelihood term (Eq. 6) is fixed in the global setting while output by the model in LAM. See Appendix C for more discussion about this.

12. Robustness for longer rollouts

Note that all models are unrolled for much fewer time steps during training than at inference time. In the global case for example, Graph-EFM is unrolled 40 steps for evaluation at 10 days lead time, but only trained by unrolling up to 12 steps. As we focus here on weather forecasting rather than climate modeling, the stability of unrolling the models for months or years is of limited interest in this work.

13. Visual artefacts in Graph-EFM (ms)

We would like to emphasize that when considering the same setting as in the GraphCast paper (global, deterministic modeling) we do not see any artefacts. The artefacts in Graph-EFM (ms) are mainly observed in the LAM setting (e.g. Figure 6.d) and for the global probabilistic model, none of which were considered in the GraphCast paper. Note however that the final global Graph-EFM (ms) model does not have these artefacts, as the artefacts mentioned as limitations only appear when training with a poor choice of $\lambda_{\text{CRPS}}$.

We thank reviewer 568u for useful comments. See our response below:

1. I need to some figures during extreme events, e.g cyclones like Yaku.

We have now included such a case study for Hurricane Laura, that can be found in the global author rebuttal.

2. About higher resolution data and ERA6

There is indeed nothing that technically prevents us from applying these methods to higher resolution data. Some minor adaptations have to be done to the exact graph structures used, but the overall framework is directly applicable. The method should scale well to even higher resolutions, although naturally with an increasing memory requirement. Our choice to focus on 1.5° data is mainly due to the computational needs.

Assuming that ERA6 will follow a similar format as ERA5, although higher resolution, there is no reason to believe our methods would not be applicable also to that dataset.

Regarding higher resolution localized data, note that the MEPS data is exactly an example of this. The 10 km spatial resolution of this data is far higher than the 1.5° ERA5 data for this area (spatial resolution $\approx$ 167 x 76 km) and even higher than 0.25° ERA5 (spatial resolution $\approx$ 28 x 13 km). Also for limited area modeling we expect the method to be readily applicable also to data at even higher resolutions.

3. Do the results change a lot if it's only trained from 1980 onwards?

While we have not had the time and resources to train such a model now, we would not expect substantial differences in the results from using only data from 1980 onwards. We opted for using as much data as possible, to make sure the model was trained on a set of weather scenarios as diverse as possible. However, in the GraphCast paper it was reported that including data from before 1980 had only minor impact on model performance [1]. Due to variations in climate one might even expect an improvement in performance when considering only more recent data, as the climate in 2020 would be more similar to that of 1980-2017 than that of 1959-1979. Using only more recent years would thus create a smaller shift in the data distribution between training and testing.

[1] R. Lam et al. Learning skillful medium-range global weather forecasting. Science, 2023.

We thank reviewer i1dL for useful comments. See our response below:

1. The authors should replace Table 1 with a line graph figure instead, as it allows comparison across different variables and lead times.

Given the limited space in the main paper we did not find a way to fit line plots for all metrics, multiple variables and both datasets. We found that shrinking the line plots to make this feasible made them hard to read, so instead opted for Table 1 and keeping all line plots in the appendix. If we have the space in the camera ready version we would be happy to bring back some line plots also to the main paper.

2. How important do you think the architecture is to the performance versus the objective function?

Both of these contributions are indeed important. The objective function, including the CRPS fine-tuning, is important for learning a calibrated probabilistic model. The hierarchical architecture does also fill an important role in Graph-EFM, spreading out the stochasticity from the latent random variable over the forecast region. Since Graph-EFM (ms) uses the same objective function, but not the hierarchical architecture, it can be viewed as an ablation of this part of the model. The poor performance of Graph-EFM (ms) confirms the importance of both of these components in order to learn a useful probabilistic model.

3. What is the reason to believe latent variable models are the way to go for probabilistic weather forecasting? (as opposed to diffusion models)

While we believe there is a place also for diffusion models in weather forecasting, one strong reason to explore latent variable models is the difference in sampling speed. Being able to quickly sample large ensembles is a desirable property of probabilistic weather forecasting models. As diffusion models require multiple forward passes through a neural network to produce one sample, this can be a slow process. This issue is aggravated by the common practice of producing forecasts auto-regressively, meaning that each forecast time step must be sampled sequentially. In Graph-EFM this only requires one pass through the network per time step. However, speeding up diffusion model sampling is a very active area of research. We view this as another useful research direction also for weather forecasting, complementary to our approach.

4. Why does the paper compare with Graphcast+SWAG but not the perturbed version of Graphcast and Gencast?

Regarding GenCast comparison see the global author rebuttal.

We assume here that the perturbed version of GraphCast refers to the one considered in the GenCast paper, where GraphCast is initialized from perturbed initial conditions. We generally consider ensemble forecasting methods that start from multiple (possibly perturbed) initial conditions to be outside of our scope, and thus do not include perturbed GraphCast as a baseline. As outlines in section 3.1 we define our problem setting as modeling the distribution of future weather states conditioned on one specific set of initial conditions. The reason for this is that creating ensembles only based on multiple initial conditions typically limits the ensemble size or requires the use of ad-hoc perturbations without any guarantees of modeling the correct distribution.

5. How does the performance vary w.r.t. the number of ensemble samples?

See the “global” author rebuttal for an investigation regarding the impact of the ensemble size.

6. Why in LAM, the Graphcast architecture is doing better than the proposed architecture?

There is some more nuance when comparing GraphCast* and Graph-FM than saying one of these is strictly doing better in any of the experiments. Which out of these two that performs the best depends on what variable and lead time is considered, as can be seen from the full results in Appendix H and I. However, it is true that for longer lead times Graph-FM performs overall worse then GraphCast* in LAM. This is likely related to the structure of the problem in the LAM setting, that features the boundary forcing and other length scales than the global case. We would like to emphasize that we consider the probabilistic Graph-EFM model the main contribution and focus of the paper, and Graph-FM mainly a step in getting there. While further investigation into the possible failure modes of Graph-FM would be interesting, we prefer this to not be the focus of this work.

7. Is there an explanation why Graphcast is better than Graph-FM, but Graph-EFM is better than Graph-EFM (ms)?

As noted above, the difference between Graphcast* and Graph-FM has some nuance. Still, what should be remembered here is that the graph partly fills a different purpose in the deterministic and probabilistic models. While in the deterministic model it is used to spread information from the input and produce the best possible forecast for the next step, in the probabilistic models the graph has the key role of augmenting the sampled latent variables into a coherent realization of the weather state. In particular, as we associate the latent variable with the mesh nodes, the stochasticity enters in different ways when using the multiscale and hierarchical graphs. Since the multiscale graph does not offer any further dimensionality reduction than mapping from the grid to the mesh, it also means that the latent variable $Z^t$ will have a higher total dimensionality ($|\mathcal{V}_1| \times d_z$) in Graph-EFM (ms) than when using the hierarchical graph in Graph-EFM ($|\mathcal{V}_L| \times d_z$). For these reasons one should not assume that any performance differences between using the multiscale or hierarchical graphs in deterministic models straightforwardly translates to the probabilistic case.

We thank reviewer LjpL for useful comments. See our response below:

1. In Figure 3, it seems like the selected ensemble members vary a lot, and how close is your forecast to the ground truth?

Note that Figure 3 shows the forecasts for 10 days in the future. At such lead times there is indeed a lot of uncertainty in any forecast that can be produced. In Figure 3 we want to show that the ensemble forecast does indeed capture this variability by producing varying ensemble members that each still represent a realistic weather state. The closest forecast (in terms of RMSE) to the ground truth is instead achieved by the ensemble mean. But as expected, due to the large variability at 10 days lead time, this ensemble mean is blurry and not a realistic weather state. We can however still see that some of the large-scale features present in the ground truth are present also in this prediction.

2. The authors could better explain the underlying meaning of measures, RMSE, CRPS, and SpSkR, for the results of weather forecast.

We agree that the presentation of these metrics in the main part of the paper is quite short. Due to space restrictions we chose to move the full definitions and descriptions of these metrics to appendix D. We encourage anyone unfamiliar with these metrics in the context of weather forecasting to read also this appendix. To give a short, more high-level explanation of these metrics in the weather forecasting context:

RMSE measures the average deviation of the forecast from the ground truth data. For probabilistic models, that produce forecasts as samples from a distribution, we compute the RMSE using a forecast representing the mean of this distribution. This is because RMSE is minimized by predicting the mean.

CRPS is a probabilistic metric, and measures how well the ground truth value is captured by the distribution specified by the model. In this case this corresponds to how well the ground truth weather is captured by the probability distribution over possible future weather states specified by the model. One useful interpretation of the CRPS is that it measures the difference between the Cumulative Distribution Function (CDF) of the model distribution and the CDF of a Dirac distribution centered at the ground truth value. One can compute a version of CRPS also for deterministic models, in which case the model distribution also becomes a Dirac distribution, but centered at the predicted value. The CRPS then reduces to the Mean Absolute Error (MAE), and this is indeed what we report for the deterministic models in our paper. CRPS is thus useful as a metric that can compare both deterministic and probabilistic forecasts. See [1] for more about CRPS.

SpSkR measures how calibrated the uncertainty of the model is. At SpSkR = 1 the ensemble forecasts are exchangeable with the ground truth, meaning that they exhibit similar levels of variability. This means that we can accurately use the ensemble spread as an indicator for forecast uncertainty. Each forecast will necessarily have some error, and a SpSkR close to 1 indicates that the model expresses the uncertainty about this error correctly.

All the metrics are computed for each variable and lead time separately, but spatially averaged for all grid points.

[1] T. Gneiting and A. E. Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, 2007.