Probabilistic Weather Forecasting with Hierarchical Graph Neural Networks

Authors: joel  oskarsson
tomas  landelius
marc peter deisenroth
fredrik  lindsten
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 14:13:32

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.

Main paper: 
 introduction: 
 Forecasting the dynamics of Earth's atmosphere is a scientific problem of utmost importance. Society is dependent on fast and informative weather forecasts for planning in areas such as transportation and agriculture and for balancing the energy system [3] . Especially important is the use of forecasts to issue warnings for extreme weather events [1] . Recent advances in Machine-Learning-based Weather Prediction (MLWP) have enabled models that produce accurate forecasts in a fraction of the time of traditional physics-based systems [37, 4, 23] . So far these developments have largely been focused on deterministic modeling. However, forecasting only one likely weather scenario ignores the many uncertainties in predicting future weather.Weather is a chaotic system, resulting in high forecast uncertainty [52] . This uncertainty comes from both imperfect representations of initial states and inaccurate descriptions of the function mapping from one time step to the next [26] . Accurately modeling this uncertainty significantly increases the value of weather forecasts. Such uncertainty can be communicated to end-users to improve decision making or be used in downstream products, for example to compute a distribution over solar power generation. Capturing the full forecast uncertainty requires us to predict not just a single likely state trajectory, but a collection of possible future weather states. Due to the complexity and dimensionality of the weather system the feasible way to achieve this is by generating samples from a modeled distribution. Such ensemble forecasting is today performed using physics-based methods, where a number of ensemble members are simulated as samples from this distribution. The computational cost of this is however massive, often limiting the spatial resolution or size of the ensemble [3] .MLWP is a promising approach for addressing this limitation and enabling large ensemble forecasts. However, for the ensemble to add value the machine learning model needs to accurately represent the distribution. Initial attempts at MLWP ensemble forecasting either rely on ad-hoc initial state perturbations [10, 37, 4] or have not been scaled to spatial resolutions of interest [19] . Also diffusion models [18] have been applied to the problem, but sampling forecasts from these is computationally expensive and can be prohibitively slow [39] . We propose a Graph-based Ensemble Forecasting Model (Graph-EFM), enabling efficient sampling of ensemble members with only one forward-pass per time step. The method builds on graph-based MLWP [20, 23] , which is a flexible framework that can be adapted to different geometries and state grid representations [24] . By combining a latent-variable formulation with a hierarchical Graph Neural Network (GNN) the distribution is modeled in a lower-dimensional space and sampled forecasts are spatially coherent.MLWP models are typically trained for and evaluated on global weather forecasting [40, 23, 4] . Another common forecasting setup in practice is the use of Limited Area Models (LAMs) to produce high-resolution regional forecasts [11] . Such LAMs are for example used by local weather services in order to provide forecasts tailored to the geographical properties and societal needs of the region [38, 44, 7, 32] . These high-resolution models are also invaluable to various industrial sectors, including energy forecasters, who rely on precise weather predictions to manage supply and demand. This motivates research into also constructing MLWP LAMs, which brings new challenges related to the high resolution and boundary conditions of the limited area. In this work we experiment not just with global forecasting, but consider also how probabilistic LAMs can be trained to produce forecasts for the Nordic region.Our main contributions are: 1) We develop a hierarchical GNN framework for both deterministic and probabilistic MLWP. The hierarchical construction encourages spatially coherent fields in forecasts. 2) We use this framework to define the probabilistic weather forecasting model Graph-EFM, capable of efficient sampling of arbitrarily large ensemble forecasts. 3) We develop a training method targeting both forecast quality and ensemble calibration. 4) We experiment with both global forecasting on 1.5°resolution and a novel limited-area modeling task at 10 km resolution. 
related work: 
 Deterministic MLWP Multiple machine learning methods have been successfully applied to large-scale weather forecasting. These include graph-based models [20, 23, 24] , transformers [4, 8, 10, 33, 25, 34] and neural operators [37, 5] . While large neural network models learn weather dynamics purely from data, there are also parallel developments in building hybrid physics-MLWP models [22, 50] .Ensembles from perturbations Most existing methods for MLWP ensemble forecasting follow closely the physics-based methods, where initial states and model parameters are perturbed to create ensemble diversity. A number of MLWP works create ensembles by ad-hoc perturbing initial states with random noise [10, 37, 4, 16, 6] . More informed perturbations have been re-used from physics based ensembles [39, 6] and created based on model-informed singular vectors [43] . Others try to perturb the forecast model itself, rolling out ensemble members using different neural network parameters [51, 43] . Such multi-model approaches require training, or at least fine-tuning, a predefined number of MLWP models. Graubner et al. [16] use the SWAG method [30] to allow for constructing multi-model ensembles of arbitrary size.Generative modeling Probabilistic machine learning approaches aim to directly learn generative models producing ensemble members. Similar to our approach, the SwinVRNN model [19] uses a latent variable formulation, but combined with a Swin Transformer architecture [29] . SwinVRNN is developed for global forecasting at 5°resolution and scales poorly to higher spatial resolutions. Also building on the graph-based framework, Price et al. [39] train a diffusion model [18, 46] to sample each time step. Their Gencast model produces ensemble forecasts of 0.25°global data with 12 h time steps. Diffusion models produce realistic-looking samples, but typically require solving an ordinary differential equation involving multiple passes through the neural network to sample each time step. For GenCast, this results in a sampling time of 8 minutes for a single 15 day forecast on a TPUv5 device [39] . Other works use diffusion models to increase the size of physics-based ensembles [27] or stochastically downscale deterministic forecasts [9, 31] .Hierarchical GNNs Motivated by capturing multiple spatial scales, hierarchical GNNs have been used for modeling general partial differential equations [14, 28] . The overall hierarchical framework shares much of its structure with the popular U-Net architecture [41] for computer vision tasks, but extended to a general graph setting. 
problem definition: 
 The weather forecasting problem can be summarized as mapping from a set of initial states X -1:0 = (X -1 , X 0 ) to the sequence of future states X 1:T = (X 1 , . . . , X T ). A table of notation is provided in appendix A. Each weather state X t ∈ R N ×dx here contains d x variables modeled at N different locations. Geospatial data is often represented as regular grids, in which case these locations correspond to the grid cells. The d x variables can include both atmospheric variables, modeled at multiple vertical levels, and surface variables. As is common in MLWP we assume the initial states to consist of two time steps, which allows for capturing first-order state dynamics. To produce a forecast, a set of forcing inputs F 1:T are also available. These contain known quantities, such as the time of day. There are also static features associated with the grid cells, such as the orography, which we here consider part of the forcing.Many variables impact the chaotic weather system, all of which are not fully captured in initial states represented on finite grids. This induces forecast uncertainty, which we view as a distribution p X 1:T |X -1:0 , F 1:T . In deterministic forecasting we seek a model that minimizes the Mean Squared Error (MSE) to the future weather states [23, 33, 34] . This is equivalent to modeling only the mean of the distribution. In probabilistic forecasting we instead aim to model the full distribution. Note that we here specifically model the conditional distribution p X 1:T |X -1:0 , F 1:T , rather than p X 1:T F 1:T . Hence we do not marginalize over uncertainty in initial states.Graph-based MLWP models use an autoregressive mapping Xt = f (X t-2:t-1 , F t ) consisting of a sequence of GNNs [20, 23, 24] . Starting from the initial states, this mapping can be iteratively applied to roll out a full forecast X 1:T . Central to the graph-based framework is the idea of mapping from the original N grid locations to a mesh graph G M = (V M , E M ). In the graph-context we refer to the grid locations as a set V G of grid nodes. By choosing |V M | < |V G | = N it becomes efficient to perform the majority of computations on the mesh. Such a mesh graph can also be tailored to the forecasting setting, for example to respect the spherical geometry in global forecasting [20] . The mapping f realizes a single-step prediction by passing X t-2:t-1 and F t through a series of GNN layers. In sequence, these layers: 1) map grid inputs to representations on the mesh graph; 2) perform a number of processing steps on the mesh; 3) map back to the grid to produce the prediction for X t . Steps 1 and 3 use bipartite graphsG G2M = (V G ∪ V M , E G2M ) and G M2G = (V G ∪ V M , E M2G) with edges connecting the grid and mesh nodes. The GNN layers in each step compute updates for node representations H ∈ R |V|×dz and edge representations E ∈ R |E|×dz in the graphs. For simplicity all representation vectors have dimensionality d z .The specific GNN layers used in previous works are Interaction Networks [2, 23] . The layers in these networks pass messages from a set of sender nodes along directed graph edges to a set of receiver nodes. Based on these messages the edge and receiver node representations are then updated. For a graph G = (V, E) let e α→β ∈ R dz be the row of E corresponding to the edge (α, β) ∈ E. Let H S be the matrix with rows containing sender node representations and H R the corresponding matrix for receiver nodes. Interaction Networks then implement the representation updateEQUATIONwhere Ne(β) = {α : (α, β) ∈ E} are the incoming neighbors of node β. Parameters in Multi-Layer Perceptrons (MLPs) are shared across nodes and edges in the graph, but not between GNN layers. Global mesh graphs Keisler [20] proposed to construct a mesh graph for global MLWP as an icosahedral grid covering the globe. This approach was extended in the GraphCast model [23] by introducing a multi-scale mesh graph with edges of varying length. Such multi-scale edges are capable of propagating information and capturing statistical dependencies both locally and over long distances in the graph. The multi-scale mesh graph is created by sequentially splitting the faces of an icosahedron into a sequence of graphs G L , . . . , G 1 with node sets satisfyingV L ⊂ • • • ⊂ V 1 by construction.The original icosahedron G L has the longest edges E L , stretching far across the globe, whereas the final graph G 1 has short edges E 1 only connecting nodes locally. The final multi-scale mesh graph is constructed asG MS = (V 1 , E L ∪ • • • ∪ E 1 ), taking the nodes from the final graph but connecting these using edges of all different lengths [23] . 
weather forecasting with hierarchical graph neural networks: 
 Two great challenges in weather forecasting is to accurately capture processes unfolding over different spatial scales and modeling the uncertainty in the chaotic system [52] . To tackle these challenges, we propose to construct a hierarchical mesh graph, working with different length scales at each level in the hierarchy. We use a sequence G 1 , . . . , G L of graphs as the different levels in the hierarchy, additionally adding connections between the nodes of adjacent levels. This construction is also highly suitable as a basis for building probabilistic forecasting models, as discussed below. Figure 1 shows an overview of the hierarchical mesh used in our model. See figs. 12 and 14 in the appendix for illustrations of how this differs from the multi-scale graph.There are multiple benefits to such a hierarchical mesh construction for MLWP. By keeping the graphs at different levels separate, we can define GNN layers with independent parametrizations at each level. This adds flexibility by allowing the model to learn different representation updates for edges of different spatial scales. A hierarchical mesh graph also offers a natural, spatially-aware dimensionality reduction, as the state in the grid is encoded into a few nodes at the top level. Such a representation can capture the general structure of each weather state, with finer details added as this is propagated down through the hierarchy. We leverage this property to construct a probabilistic model by imposing a distribution over these lower-dimensional representations at the top level. This allows for efficiently drawing spatially coherent samples from the distribution of future weather states.Our hierarchical mesh graph consists of L graph levels G 1 , . . . , G L with G l = (V l , E l ). Only level 1 of the hierarchy is connected to the grid, so we re-defineG G2M = (V G ∪ V 1 , E G2M ) and G M2G = (V G ∪ V 1 , E M2G ).The number of nodes |V l | decreases with the level l. The smallest set of nodes are found at the top level L.To pass information between the levels of the hierarchy we introduce additional graphs connecting the different levels. Let G l,l+1 = (V l ∪ V l+1 , E l,l+1 ) be a graph containing directed edges from mesh level l to level l + 1. We make use of a graph sequence G 1,2 , . . . , G L-1,L to propagate information up through the hierarchy and similarly a sequence G L,L-1 , . . . , G 2,1 in the downward direction. The exact layout of nodes and edges at and in-between levels are design choices that should be tailored to the specific forecasting setting. Examples for global and limited-area forecasting are given in section 5.The hierarchical graph allows for defining GNN layers both on and in-between the different levels. By sequentially updating node and edge representations at different levels in the hierarchy, information can be propagated up from the grid to the different levels. As these levels have edges of different lengths, the processing at each level happens on different spatial scales. Note that this differs from the multi-scale graph approach, where information processing over all different spatial scales happen in the same GNN layer [23] . As a step towards our probabilistic model, we define an alternative deterministic Graph-based Forecasting Model Graph-FM1 , operating on the hierarchical graph.In Graph-FM one processing step on the mesh graph is defined as a complete sweep through the hierarchy. GNNs are applied sequentially to the inter-level and intra-level graphs in the order To capture the uncertainty in the chaotic weather system we next aim to construct a probabilistic model from the ground up to capture the full distribution p X 1:T |X -1:0 , F 1:T . We start by assuming the weather system to satisfy a second-order Markov assumption, decomposingG 1 , G 1,2 , G 2 , . . . , G L-1,L , G L ,EQUATION)Factoring the distribution over time steps allows us to work with forecasts of varying length. Specifying the model for single-step prediction avoids having to learn separate parameters for different lead times. Next, we seek a flexible, but computationally efficient parametrization for the distribution p X t X t-2:t-1 , F t . This can be achieved by introducing a latent random variable Z t , and lettingEQUATIONHere the stochasticity in Z t should capture the uncertainty over X t at each time step. The corresponding graphical model is shown in fig. 2 . We impose a spatial structure over the latent variable by letting Z t be |V L | × d z matrix-valued, with each row a d z -dimensional vector associated with one node in the top level G L of the mesh graph.The single-step model consists of two components, a latent map p Z t X t-2:t-1 , F t and predictor p X t |Z t , X t-2:t-1 , F t . The latent map is parametrized using GNNs, mapping the conditioning variables to parameters of a Gaussian distribution. We consider the predictor to be concentrated around its mean, and realize p X t |Z t , X t-2:t-1 , F t as a deterministic mapping of a similar form as Graph-FM. By sampling Z t and passing this through the predictor we can draw a sample of X t from eq. ( 3). This sample can then be conditioned on at the next time step, continuing this sampling process to roll out a forecast following eq. ( 2). This forecast constitutes one ensemble member, and the process can be repeated to sample an ensemble of arbitrary size. We call our Graph-based Ensemble Forecasting Model Graph-EFM. Full details about the model are given in appendix C.We let the latent map be an isotropic GaussianEQUATIONwith the mean as a function of the conditioning variables. The variance is fixed, imposing a fixed scale for the learned latent space. The mean function µ Z consists of a sequence of GNNs. These take the inputs at the grid, propagate representations up through the hierarchical mesh graph, and finally predicts the mean of Z t α at each node α at level L. In appendix L.2 we verify empirically the importance of using the latent map over a static distribution for Z t .Predictor The predictor is a deterministic mapping Xt = g Z t , X t-2:t-1 , F t = X t-1 + g Z t , X t-2:t-1 , F t .(With the small time steps used in MLWP, X t does not change dramatically in a single step. We thus follow the common practice of including a skip connection to the previous state [23, 5, 19] . The predictor takes both inputs X t-2:t-1 , F t at the grid and Z t at the top of the mesh graph. To incorporate both we design g similar to Graph-FM, performing sweeps up and down through the mesh hierarchy. At the top of the hierarchy Z t is added to node representations H L through the residual connections in the GNN layers. A sampled value of Z t then affects the prediction Xt through the downward sweep. While multiple such sweeps are possible, we found one to be sufficient in practice.We want each sample of X t to contain spatially coherent atmospheric fields. One approach would be to impose spatial dependencies in the joint distribution over Z t . However, learning and sampling from such a distribution typically comes with computational challenges [19] . Instead, we impose spatial dependencies by integrating the latent variable formulation with the hierarchical graph. We argue that as the independent components of Z t are propagated down through the mesh graph, gradually increasing the spatial resolution, spatial dependencies are introduced by the model in the GNN layers. The hierarchical graph is key to this property, as the stochasticity in Z t is necessarily spread out over the forecast region, rather than only affecting the output locally.Deterministic forecasting models can be straightforwardly trained by minimizing a weighted MSE [23] or Negative Log-Likelihood (NLL) loss [8] for rolled out forecasts. To train Graph-EFM we instead leverage the fact that the single-step model has a structure similar to a (conditional) Variational AutoEncoder (VAE) [21, 45] , allowing us to use a variational objective. We introduce a variational approximation q Z t X t-2:t-1 , X t , F t at each time step, approximating the true posterior p Z t X t-2:t-1 , X t , F t over Z t . This variational distribution is parametrized in a similar way as the latent map, with GNN layers mapping to a Gaussian over Z t . Note however that q also depends on X t , since it approximates the posterior. Using q, we can then defineEQUATIONwhich is equal to the (negative) Evidence Lower Bound (ELBO) when the weighting is λ KL = 1.While the predictor g is a deterministic mapping, we introduce a Gaussian likelihood in eq. ( 6) to get a well-defined learning problem. This setup corresponds to the common practice in VAEs of assuming Gaussian observation noise, but not adding this to samples from the model [42] . The standard deviation σ α,j can either be a second output from the predictor or manually chosen (see appendix D for details). As with deterministic models [23, 20, 8, 34] , we found it crucial to fine-tune on rolled out forecasts of multiple time steps. This improves stability and performance for longer lead times. In the final fine-tuning we include also a Continuous Ranked Probability Score (CRPS) loss term L CRPS [15, 22] . The full objective function is then L = L Var + λ CRPS L CRPS , with λ CRPS a weighting hyperparameter. Including this CRPS loss improves the calibration of ensemble forecasts.In Graph-EFM there is a large amount of information that needs to be propagated between the grid and Z t . However, the Interaction Network GNNs are biased towards keeping old representations of receiver nodes, rather than updating this with new information from incoming edges. Note in eq. ( 1) that if the MLPs are initialized to give outputs close to 0, there will be no change to e α→β and H R β . In practice the model has a hard time learning to propagate useful information up from the grid to Z t . Even when trained purely as an auto-encoder (λ KL = 0), Z t easily ends up being ignored. To remedy this we propose an alternative GNN formulation that we call Propagation Network, defined byEQUATIONEQUATIONFor MLPs initialized with outputs close to 0, Propagation Networks reduce to averaging the values of neighboring nodes. This encourages the propagation of information from H S to H R by construction. Propagation Networks were found to perform better also in the deterministic model (see comparison in appendix L.1), so we employ these in both Graph-FM and Graph-EFM. 
experiments: 
 To evaluate our models we conduct experiments on both global and limited area forecasting. The models are implemented2 in PyTorch and trained on 8 A100 80 GB GPUs in a data-parallel configuration.Training takes 700-1400 total GPU-hours for the global models, and around half of that for the limited area models. The computational demands prevent us from re-training multiple models for statistical analysis. Once trained, sampling from Graph-EFM is highly efficient. Using batched sampling on a single GPU, 80 ensemble members are produced in 200 s (2.5 s per member) for global forecasting.Metrics We measure the skill of deterministic models by Root Mean Squared Error (RMSE). For probabilistic models we compute the RMSE for the ensemble mean. Good skill in terms of RMSE is however not enough for ensemble forecasts, where we want to capture the full distribution. For these we also assess the ensemble calibration by computing the Spread-Skill-Ratio (SpSkR). Calibrated uncertainty corresponds to SpSkR ≈ 1 [13] . We additionally use CRPS to measure how well the marginal distributions of the model matches the data. For deterministic models the CRPS reduces to Mean Absolute Error (MAE). Complete definitions of all metrics are given in appendix E.Models Achieving a fair comparisons of the actual machine learning methodology in MLWP is challenging due to models using different spatial resolution, variables and initial states. We here train an illustrative set of models on the same data and with comparable training setups. Our full Graph-EFM model is compared to: 1) Graph-EFM (ms), a version of Graph-EFM using a multiscale mesh graph instead of the hierarchical one. 2) Graph-FM, our deterministic model using the hierarchical graph. 3) GraphCast*, a reimplementation of GraphCast [23] , adapted and trained on our datasets. 4) GraphCast*+SWAG, a multi-model ensemble created by applying Stochastic Weight Averaging Gaussian (SWAG) [30] to GraphCast*. Inspired by Graubner et al. [16] , this represents a simple way to augment a deterministic model to perform ensemble forecasting. Further details about the baseline models are given in appendix C.5. For ensemble models we sample 80 members for the global experiments and 100 members for limited area forecasting. In appendix L.3 we investigate the impact of ensemble size on the evaluation. We find that improvements in metric values quickly saturate when increasing the ensemble size. This shows that sampling even more members would have negligible impact on the results of our experiments.Data and graphs We experiment on global weather forecasting up to 10 days with 6 h time steps. The dataset used for training and evaluation is a 1.5°version of the global ERA5 reanalysis3 [17] , provided through the WeatherBench 2 benchmark [40] . The models forecast d x = 83 different variables in total, including both surface-level variables and atmospheric variables at 13 different pressure levels. We use the years 1959-2017 for training, 2018-2019 for validation and 2020 as a test set. Forecasts are always started from initial conditions taken directly from ERA5, both during training and evaluation. For global forecasting we use the graph generation process from GraphCast [23] . The multi-scale graph G MS is created by refining the icosahedron 4 times. The hierarchical graph contains 4 levels of such icosahedral grids. More details on the global experiments are given in appendix H.As the models forecast many different variables we present only a selection of results in the main paper. Metric values for geopotential (z500) and 2 m temperature (2t) are listed in table 1 and results for mean sea level pressure (msl) plotted in fig. 3 . Line plots for all metrics and a large number of variables are given in appendix J.1. In the appendix we also show comparisons to additional models from the literature, trained on different data, as well as the physics-based IFS-ENS model [12] . The Extreme weather case study An important use case for ensemble forecasting is modeling extreme weather events. While higher resolutions than 1.5°are generally desirable for accurately capturing such extremes, we conduct one case study on using Graph-EFM for forecasting hurricane Laura. The full case study with visualized forecasts is available in appendix F. For this example we show that there exists ensemble members accurately predicting the landfall location of the hurricane at 7 days lead time, while the deterministic models still show no sign of the hurricane in the region. Closer to the landfall event the ensemble forecast from Graph-EFM indicates uncertainties associated with the landfall location and wind intensity. This demonstrates the added value of a probabilistic forecasting model.In LAMs weather forecasts are produced for a bounded region of the globe. LAM forecasting allows for higher resolution modeling and regionally tailored model configurations [11] , properties that can be inherited by MLWP models by training on LAM data. To model weather over a limited domain, boundary conditions need to be taken into account. In physics-based LAMs these are typically given by a global forecast [38, 44, 7, 32] . We adapt a similar approach for MLWP LAMs, by taking boundary conditions as additional forcing along the boundary of the forecast area. The problem of LAM forecasting is thus about simulating physics not just based on the initial state, but also consistent with these boundary inputs. In the models we introduce N b additional grid nodes along the area boundary, for the boundary forcing B t ∈ R N b ×dx . Boundary forcing B t is always fed together with X t to the model. Grid nodes on the boundary and within the area are treated identically by the GNN layers. We perform this adaptation to all models in our experiment. weather variables, some repeated on multiple vertical levels. Forecasts are rolled out with 3 h time steps up to lead time 57 h. In this experiment we take also the boundary forcing directly from the MEPS dataset. We define the boundary as the outermost 10 grid positions. Using the same dataset for the area and boundary allows us to investigate the modeling choices in a controlled experimental setup. In an operational scenario the boundary forcing would instead come from a re-gridded global forecast. In the LAM setting we define our graphs as regular quadrilateral meshes covering the MEPS forecasting area, but with far fewer nodes than the original grid. The graph hierarchy G 1 , . . . , G L is created by constructing such meshes at different resolutions. By placing each node in G l at the center of 3 × 3 nodes in G l-1 , we can merge 4 such graph levels to create G MS . In the hierarchical graph we instead introduce edges from each node in G l to the 3 × 3 nodes in the level below. More details about the MEPS data and experiment can be found in appendix I.Results A selection of metrics are shown in table 2 and full results given in appendix K. At these shorter lead times there is no clear benefit of probabilistic modeling in terms of RMSE. Still, as exemplified by the standard-deviation plotted in fig. 5 , probabilistic modeling provides useful information about the forecast uncertainty. Comparing the ensemble members in fig. 5 highlights the improved spatial coherency of the hierarchical graph in Graph-EFM. In contrast, the Graph-EFM (ms) forecast looks patchy and lacks physically intuitive features. There are also clear visual artifacts, that can be traced to the multi-scale graph structure. We discuss this more in-depth in appendix G. In the LAM setting all models are under-dispersed, with SpSkR < 1. One explanation for this is that the boundary forcing constrains the space of plausible forecasts, hindering the ensemble spread. 
discussion: 
 In this paper we have explored MLWP ensemble weather forecasting using graph-based latent variable models. Our Graph-EFM model is capable of efficiently producing accurate ensemble forecasts. This paves the way for large-scale MLWP ensemble forecasting both in operational use and research settings. In appendix B we further discuss the societal impact of this research. With this work we hope to emphasize that MLWP models are not just deterministic mappings, but parametrize distributions of weather states. It follows that ensemble forecasting should not be achieved by perturbing models, but by directly modeling the distribution of interest.The training process comes with some complications in terms of choosing a training schedule and hyperparameters λ KL and λ CRPS . While the CRPS fine-tuning is an important training step, we have found that choosing a too high λ CRPS can introduce visual artifacts, especially for the Graph-EFM (ms) model (see appendix G). While Graph-EFM produces diverse and physically plausible ensemble members, the forecasts still suffer from some of the blurriness common to deterministic models [23, 40] . We here trade off some of the visual fidelity achieved for example by diffusion models [39] for more efficient sampling of ensemble members.Future work Interesting avenues for future work include learning probabilistic weather models based on other types of autoencoders [48, 49] , or by directly optimizing scoring rules [36, 22] . Another approach for achieving efficient ensemble forecasting is to explore techniques for speeding up diffusion model sampling [47] .A Extreme weather Due to climate change the prevalence and severity of extreme weather events is expected to increase substantially, endangering both property and human life [56] . These events are also getting harder to predict and the cost of damages per event has increased nearly 77% over the past five decades [62] . In order to detect extreme events, there is a need for ensemble forecasting to better capture the full distribution of possible weather states. While ensemble forecasting historically has been limited by computational costs [3] , efficient MLWP ensemble models have the potential to vastly improve our ability to model extreme weather. In appendix F we present a case study showing how our Graph-EFM could be used for forecasting Hurricane Laura. More extensive evaluation of the abilities of MLWP models to capture extreme events is a complex, but important issue [54, 55] . We view these capabilities as one of the main motivations for further developments of MLWP ensemble forecasting models.Forecast failures There will always be cases where weather forecasting systems fail, and produce inaccurate predictions for future weather. Depending on the weather event, such forecast errors can have disastrous consequences. In these cases the lack of interpretability of black-box MLWP systems can be a problem, making it hard to understand why the model was wrong. Traditional physics-based systems can also be tough to interpret due to their complexity, but at their core are physical equations understood by researchers. With the rapid progress of MLWP, it is likely that we will soon see a landscape where the physical models take the back seat to a plethora of skillful MLWP forecasting systems. In such a scenario the question of how to investigate forecast failures becomes pressing. It seems desirable to be able to fall back on physical models for understanding impactful events poorly H G ← MLP(X t-2:t-1 , F t ) ▷ Embedd grid inputs to d z -dimensional vectors 2: H M , • ← GNN(G G2M , H G , E G2M , H M )▷ Map grid representation to mesh 3: for all mesh processing steps do 4:H M , E M2M ← GNN(G M , H M , E M2M , H M )▷ Update mesh representations 5: end for 6:H G ← H G + MLP(H G ) 7: H G , • ← GNN(G M2G , H M , E M2G , H G )▷ Map mesh representation back to grid 8: if outputs σ then 9:[Y, U ] ← MLP H G 10: return X t-1 + Y, Softplus(U )▷ Return prediction of X t and σ for loss 11: else 12:return X t-1 + MLP(H G )▷ Return prediction of X t 13: end if forecast by MLWP. However, to allow this we can not do away with all infrastructure and expertise related to physical modeling. These are important considerations as weather forecasting moves into an era of operational MLWP.The production of renewable energy, such as solar and wind power, can be highly volatile [60] . This creates challenges for including these sources in the larger energy system. Accurate weather forecasts, translated into forecasts of energy production, fill an important role as enablers of these energy sources by making their output predictable. Detailed probability estimates from ensemble forecasting can additionally allow for improved cost-loss decision making in these systems.The energy footprint of weather forecasting Traditional weather forecasting systems utilize massive computing clusters [11] , resulting in a substantial energy footprint of the forecasting process. In comparison, MLWP models are highly energy efficient, even when taking into account their initial training process. The total energy required to train the large global MLWP model FourCastNet is comparable to running a 10 day forecast with 50 ensemble members using a traditional forecasting system [37] . Producing a forecast using the same model uses four orders of magnitude less energy than a physics-based model. A similar reduction in energy footprint is to be expected for our models, if applied at the same scale. For ensemble forecasting the total energy saving is even greater, as it is multiplied by the number of ensemble members. However, one has to beware of rebound effects, where efficiency improvements result in more extensive use of resources. If ensemble forecasting becomes 1000 times more energy efficient and we run 1000 times as many ensemble members there is no energy saved in the end.We here give more details about the different models discussed in the main paper.C.1 Deterministic graph-based models [23, 20] Deterministic graph-based MLWP models represent the single-step prediction function Xt = f (X t-2:t-1 , F t ) as a sequence of MLPs and GNN layers. Algorithm 1 describes the full prediction process. Note that for some of the GNN layers the edge representations are not updated, as these edges are not used in later steps of the process. We denote this by H, • ← GNN(. . . ) in algorithm 1. MLPs always act on the last dimension of its concatenated inputs. Mesh node and edge representations are initialized by embedding related static features using additional MLPs. These static features contain information about node and edge positions (see appendices H.1 and I.1). The number of processing steps to run on the mesh is a hyperparameter. Each such step contains one GNN layer, and these GNNs do not share parameters. The mapping f can optionally output also standard deviations σ, which we use when training with NLL loss (see appendix D.1). We restrict σ to be positive by applying the Softplus function.Algorithm 2 Single-step prediction f in Graph-FM 1: H G ← MLP(X t-2:t-1 , F t ) ▷ Embedd grid inputs to d z -dimensional vectors 2: H 1 , • ← GNN(G G2M , H G , E G2M , H 1 )▷ Map from the grid . . . 3: for l = 2, . . . , L do ▷ . . . and up through the whole mesh hierarchy 4:H l , E l-1,l ← GNN(G l-1,l , H l-1 , E l-1,l , H l ) 5: end for 6: for number of mesh processing steps / 2 do ▷ Each iteration contains two steps, down and up 7:H L , E L ← GNN(G L , H L , E L , H L ) 8:for l = (L -1), . . . , 1 do ▷ Downward pass of sweep 9:H l , E l+1,l ← GNN(G l+1,l , H l+1 , E l+1,l , H l )10:H l , E l ← GNN(G l , H l , E l , H l ) 11:end for12:H 1 , E 1 ← GNN(G 1 , H 1 , E 1 , H 1 )13:for l = 2, . . . , L do ▷ Upward pass of sweep 14:H l , E l-1,l ← GNN(G l-1,l , H l-1 , E l-1,l , H l )15:H l , E l ← GNN(G l , H l , E l , H l ) 16:end for 17: end for 18: for l = (L -1), . . . , 1 do 19:H l , • ← GNN(G l+1,l , H l+1 , E l+1,l , H l )▷ Map down through the mesh hierarchy . . . 20: end for 21:H G ← H G + MLP(H G ) 22: H G , • ← GNN(G M2G , H 1 , E M2G , H G )▷ . . . and finally to the grid 23: if outputs σ then 24:[Y, U ] ← MLP H G 25: return X t-1 + Y, Softplus(U )▷ Return prediction of X t and σ for loss 26: else 27:return X t-1 + MLP(H G ) ▷ Return prediction of X t 28: end if C.2 Graph-FMIn the hierarchical graph there are node representations associated with every level and edge representations associated with each subset of edges. We let H l be the representations associated with nodes V l at level l in the mesh hierarchy. Similarly, let E l contain representations of intra-level edges E l , E l,l+1 representations of upwards edges E l,l+1 , and E l,l-1 representations of downward edges E l,l-1 . As in the non-hierarchical case, all representations associated with the mesh graph are initialized by MLPs applied to static features. Independent MLPs are used for the different levels of the hierarchy.Forecasting using Graph-FM follows a similar structure as previous works, but replaces the mappings between the grid and mesh to encompass all levels and changes the processing steps into sweeps through the hierarchy. Recall that for the hierarchical graph we re-defineG G2M = (V G ∪V 1 , E G2M ) and G M2G = (V G ∪ V 1 , E M2G ).The full prediction mapping f for Graph-FM is described in algorithm 2. We define one processing step in Graph-FM as one update to all node and edge representations in the graph, keeping a similar interpretation as processing steps in non-hierarchical graphs. This means that one sweep down and up (lines 7 to 16 in algorithm 2) is counted as two processing steps. It follows that only even numbers of processing steps are reasonable in Graph-FM. Note also here that GNNs only share parameters across nodes and edges in the specific graph they operate on. There is no parameter sharing across processing steps nor levels in the hierarchy. We include Propagation Networks in Graph-FM, but not all parts of the model can be expected to benefit from this change. For some steps of algorithm 2 we want to keep the inductive bias of Interaction Networks to retain information. Specifically, we use Propagation Networks for the mappings directly between the grid and lowest mesh level (lines 2 and 22) and for the upward processing steps (line 14). Algorithm 3 Sampling of ensemble member from p X 1:T |X -1:0 , F 1:T 1: X-1:0 ← X -1:0 2: for t = 1, . . . , T do 3:Z t ∼ p Z t Xt-2:t-1 , F t ▷ Sample Z t from latent map 4:Xt ← Xt-1 + g Z t , Xt-2:t-1 , F t ▷ Compute next state using predictor 5: end for 6: return X1:TAn overview of the Graph-EFM model with global graphs is shown in fig. 6 . At each tim step Graph-EFM takes as input X t-2:t-1 , F t and produces a sample of X t . By feeding the sample from the model back at the next time step, an ensemble member can be rolled out. This process is described in algorithm 3, where we let Xt be an initial state for t < 1 and then assign it sampled predictions for the actual forecast time steps 1, . . . , T .When creating an ensemble of size K from deterministic initial conditions some care has to be taken as to how many members to sample. If we at each time t would sample K values of X t for each realization of X t-2:t-1 we would end up with K T members at time T . Instead, we sample K realizations of X 1 given the initial conditions and after that only one realization of X t for each of the members. This is equivalent to doing K independent runs of algorithm 3. Note that since we can always draw new samples of Z t we are never restricted in the ensemble size.The latent map is defined asEQUATIONwhere the mean function µ Z is realized as a sequence of GNN layers mapping up through the hierarchy. The exact process is described in algorithm 4. Since only one upward pass is performed we do not update any of the edge representations. The final MLP at line 8 maps from node representations at mesh level L to the mean of p Z t X t-2:t-1 , F t . In the latent map we use exclusively Propagation Networks, encouraging the model to encode useful information from the grid into Z t .Predictor The predictor p X t |Z t , X t-2:t-1 , F t is chosen as a dirac measure, with all probability mass concentrated in one point. It thus takes the form of a deterministic mappingEQUATIONThis choice emphasizes that all randomness in p X t X t-2:t-1 , F t should come from the latent variable Z t . An alternative would be to parametrize the predictor as a diagonal Gaussian. We found this inferior in practice, as that would require also sampling from this Gaussian when rolling out Algorithm 4 Latent map mean function µ Z 1: H G ← MLP(X t-2:t-1 , F t ) ▷ Embedd grid inputs to d z -dimensional vectors 2: H 1 , • ← GNN(G G2M , H G , E G2M , H 1 ) 3: H 1 , • ← GNN(G 1 , H 1 , EH l , • ← GNN(G l-1,l , H l-1 , E l-1,l , H l ) 6: H l , • ← GNN(G l , H l , E l , H l ) 7: end for 8: return MLP H L ▷ Return mean of p Z t X t-2:t-1 , F t Algorithm 5 Predictor function g 1: H G ← MLP(X t-2:t-1 , F t ) ▷ Embedd grid inputs to d z -dimensional vectors 2: H 1 , • ← GNN(G G2M , H G , E G2M , H 1 ) 3: for l = 1, . . . , L -1 do ▷ Upward pass of sweep 4: H l , E l ← GNN(G l , H l , E l , H l ) 5: if l = L -1 then 6: H L , • ← GNN(G L-1,L , H L-1 , E L-1,L , Z t ) ▷ Incorporate Z t through GNN 7:else 8:H l+1 , • ← GNN(G l,l+1 , H l , E l,l+1 , H l+1 ) 9:end if 10: end for11: H L , • ← GNN(G L , H L , E L , H L ) 12: for l = (L -1), . . . , 1 do ▷ Downward pass of sweep 13: H l , • ← GNN(G l+1,l , H l+1 , E l+1,l , H l ) 14: H l , • ← GNN(G l , H l , E l , H l ) 15: end for 16: H G ← H G + MLP(H G ) 17: H G , • ← GNN(G M2G , H 1 , E M2G , H G ) 18: if outputs σ then 19: [Y, U ] ← MLP H G 20: return X t-1 + Y, Softplus(U )▷ Return prediction of X t and σ for loss 21: else 22:return X t-1 + MLP(H G )▷ Return prediction of X t 23: end if forecasts. This sampling would simply entail adding independent Gaussian noise to the forecasts, which only degrades their quality. We still re-introduce the Gaussian form in the variational objective (eq. ( 6)), as some choice of likelihood is necessary for a well-defined learning problem. The interpretation of this is that the forecast is rolled out as a latent process with all randomness coming from Z t . Observation noise is then assumed added to this forecast independently at each time step. We use the latent, noise-free forecast as our prediction. Connecting back to the VAE analogue of our model, this setup corresponds to the common practice in the VAE literature of using a Gaussian likelihood but treating the mean as the generated sample [42] .The exact form of the predictor function is described in algorithm 5. It has a similar form as Graph-FM, doing sweeps up and down the mesh hierarchy. We found it sufficient to perform one sweep up and down in the predictor, but it can easily be generalized to multiple sweeps as in Graph-FM. The predictor function optionally returns also the σ to be used in our variational objective. In the predictor we use Interaction Networks in the upward direction and Propagation Networks in the downward direction. This is to guarantee that randomness and useful information encoded in Z t at the top level reaches all the way down to the grid. Note that on line 6 Z t is combined with H L-1 through the Interaction Network layer mapping from level L -1 to L. For the details of how this combination happens, see eq. ( 1), where Z t takes the role of the receiver node representations H R . Note that edge and node representations updated on line 4 in the upward sweep then re-appear in the downward pass. This constitutes residual connections between the upward and downward passes, as illustrated in fig. 6 . Algorithm 6 Parameter mappings µ q , σ q of variational approximation 1: H G ← MLP(X t-2:t-1 , X t , F t ) ▷ Embedd grid inputs (including X t ) to d z -dim. vectors 2: H 1 , • ← GNN(G G2M , H G , E G2M , H 1 ) 3: H 1 , • ← GNN(G 1 , H 1 , EH l , • ← GNN(G l-1,l , H l-1 , E l-1,l , H l ) 6: H l , • ← GNN(G l , H l , E l , H l ) 7: end for 8: [Y, U ] ← MLP H L 9: return Y, Softplus(U )▷ Return mean and standard deviation of q Z t X t-2:t-1 , X t , F tVariational approximation Defining the ELBO for the model requires a variational approximation q Z t X t-2:t-1 , X t , F t at each time step, approximating the true posterior p Z t X t-2:t-1 , X t , F t over Z t . Due to the autoregressive structure of the model, it is sufficient to condition q on states at time steps t -2, t -1 and t. This conditioning removes all dependence between Z t and other variables (see fig. 2 ), simplifying the parametrization and training of q. We choose the variational approximation to have a similar Gaussian form as the latent mapEQUATIONFor added flexibility we also model the variance of q. Here both µ q and σ q are functions mapping from the grid inputs to |V L | × d z matrices. These functions are implemented jointly as described in algorithm 6. This mapping is similar to the latent map, with the differences being that X t is taken as an input and both mean and standard deviation returned. Similarly to the latent map, all GNNs used in the variational approximation are Propagation Networks.Following Lam et al. [23] , all MLPs have one hidden layer with Swish activation functions [59] . With the exception of MLPs outputting predictions or distribution parameters, we apply LayerNorm [53] to all other MLP outputs.While MLPs in different GNN layers never share parameters, we use some parameter sharing across the embedding MLPs in the different components of Graph-EFM. Specifically, the latent map and predictor share parameters for the MLP embedding the grid input (line 1 in algorithm 4 and line 1 in algorithm 5). The MLPs that embedd static graph features, for initializing graph representation vectors, are also shared across all components of the Graph-EFM model.GraphCast* GraphCast* is a reimplementation of the GraphCast model [23] in our codebase. This was done to allow for a fair comparison by using the same data and multi-scale graphs as other models in our experiments.The Graph-EFM (ms) model is a version of Graph-EFM, but using a multiscale mesh graph G MS [23] , instead of the hierarchical one. We replace the sweeps up and down the hierarchy with multiple processing steps performed on the same mesh graph, similar to the deterministic case described in algorithm 1. In the predictor of this model, the latent Z t is integrated already in the mapping from the grid to the mesh asEQUATIONNote that as Z t here is associated with the nodes of G MS , it has shape |V 1 | × d z as opposed to |V L | × d z for Graph-EFM. As V 1 has many more nodes, the total dimensionality of Z t is higher for Graph-EFM (ms). This is a consequence of the multi-scale graph offering no further dimensionality reduction beyond mapping from the grid to the mesh.GraphCast*+SWAG One simple way to create an MLWP ensemble is by re-training multiple deterministic models from random initializations. Training many large weather models from scratch is however infeasible in practice. It also limits the number of ensemble members to the number of models trained. Inspired by Graubner et al. [16] , we create a multi-model ensemble baseline using SWAG [30] . SWAG is a technique for approximate Bayesian inference based on running Stochastic Gradient Descent (SGD) training with a constant high learning rate. At regular intervals throughout this optimization process the model parameters are saved. The full SWAG process includes estimating a high-dimensional Gaussian over these parameters and drawing samples to create the ensemble. This gets around the limitation of needing to decide the maximum ensemble size when training the model. As we never require more ensemble members than saved during the SGD training we skip estimating this Gaussian and resampling. Instead, we directly use a subset of the saved model parameters to create our ensemble members. Graubner et al. [16] combine SWAG with initial state perturbations, but as we model the conditional distribution p X 1:T |X -1:0 , F 1:T we do not include any such perturbations to the initial states.To create our GraphCast*+SWAG ensemble we start from the trained GraphCast* model and run SGD training for T = 8 unrolled steps and with a fixed learning rate (10 -3 for the global experiment and 5 × 10 -5 for MEPS). Even higher learning rates led to numerical issues or the model training diverging. Model parameters are then saved every 100 gradient descent steps. During evaluation these parameters are loaded one after another and used to produce the forecast for each ensemble member.The error and spread of the GraphCast*+SWAG ensemble is highly dependent on the correlation of consecutively saved model parameters. Given the poor spread of GraphCast*+SWAG we tested also saving parameters with 1000 steps in-between, but this led to similar results.We here give detailed definitions of the objectives used to train our models.Weighted MSE Deterministic forecasting models can be trained by minimizing a weighted MSE [23] of rolled-out forecastsEQUATIONwhere we recall that N = |V G | and:• w α is the normalized area weighting for grid cell α. These are computed as described by Rasp et al. [40] .• λ j is the inverse variance of time differences for variable j. These can be computed directly from iterating over the data.• ω j is a weight associated with the vertical level of variable j. For global experiments we use the same weights ω j as Lam et al. [23] .Negative Log-Likelihood loss For the MEPS data however, choosing weights ω j is challenging due to the many different variables and their irregular vertical levels. In such cases, an alternative approach is to use an NLL loss, where the model itself determines the trade-offs between different variables [8] . As the MEPS area has grid cells of approximately equal size we can also simply set w α = 1. For the MEPS experiments we thus useEQUATIONThe standard deviation σ α,j is here a second output from the mapping f . When unrolling forecasts we do not sample from this Gaussian, but feed the mean Xt to the next time steps. Note that eq. ( 13) decomposes into weighted squared error terms and log-determinant terms preventing too large σ outputs. The weighted MSE in eq. ( 12) can therefore also be viewed as a special case of the NLL loss, up to an additive constant. This relationship corresponds to eq. ( 13) with constant variances given byEQUATIONFor the probabilistic model we can not directly apply the training objectives above, as they only seek to match the first moments of the model distribution to the data. To train the probabilistic model we instead leverage the fact that the single-step model has a structure similar to a (conditional) VAE [21, 45] . We can therefore optimize a variational objective, the ELBO, to match the distribution of the data.Variational objective For a single time step, our variational objective isEQUATIONWe include a weight λ KL in front of the KL-regularizer to allow for some tuning between the two parts of the objective function [61] . In practice we found this useful to make sure the model does not collapse to pure deterministic predictions. Note that when λ KL = 1 this is exactly equal to the (negative) ELBO. The NLL from eq. ( 13) here shows up again in the second term. Similarly to the deterministic models, σ α,j can either be output by the model (the predictor in this case) or chosen as a constant. For global experiments we choose σ α,j as in eq. ( 14), reducing the term to a weighted MSE loss. For the MEPS data we again let each σ α,j be output by the model. Equation ( 15) can in practice be computed efficiently by using a single-sample Monte Carlo estimate for the expectation and the reparametrization trick [21] . The KL-divergence is available in closed form as both distributions are Gaussian. As with deterministic models [23, 20, 8, 34] , it is crucial to fine-tune on rolled out forecasts of multiple time steps. This induces stability and improves performance for longer lead times. We thus define our fine-tuning objective asEQUATION)where Xt is an initial state from the dataset for t < 1 and otherwise sampled using the variatonal approximation for Z t . In eq. ( 16) X t is always from the training data.CRPS fine-tuning In practice we found that fine-tuning the model with only L Var tends to result in underdispersed ensemble forecasts, underestimating the variance of the distribution. To remedy this, we include an additional fine-tuning objective based on the CRPS [15, 63, 22] . The CRPS measures how well a univariate distribution matches the observed data and is defined asEQUATIONwhere x, x ′ are independent copies of a random variables with distribution p(x) and y the observed data. The CRPS is minimized only when the predicted distribution matches that of the data [15] , making it suitable for calibrating our ensemble forecasting model. We define our CRPS fine-tuning loss as an unbiased two-sample estimator, summed over all dimensions of the predictive distributionEQUATIONwith Xt and Xt coming from two independent ensemble members sampled from the model.Note that the CRPS objective targets only the marginal distributions, while spatio-temporally coherent samples requires matching the joint distribution of the data. To avoid degenerate solutions we weight L CRPS by λ CRPS and combine it with the variational objective for the final fine-tuning loss L = L Var + λ CRPS L CRPS . Computing L requires rolling out three forecasts, one for L Var (using samples from q) and two for L CRPS (using samples from the latent map). This can have a substantial memory cost. It is however sufficient to include the CRPS term only in the final steps of the training process, making this less of an issue in practice.Empirically we found that including the CRPS objective improves both ensemble calibration and forecast accuracy. The accuracy improvement can be attributed to the MAE terms in eq. ( 18). Note that in L Var forecasts are rolled out by sampling Z t from the variational approximation, while in L CRPS these samples are from the latent map. This matches how forecasting is carried out at test time.As a consequence the CRPS loss has an additional effect of bridging the gap between the variational training and the final forecasting.We here give full definitions for the metrics used to evaluate our models.We define evaluation metrics based on set of S forecasts, each a sequence X1,(m) , . . . , XT,(m) of predicted weather states. Metrics are computed per lead-time and variable. For deterministic forecasting we define the RMSE of variable j at lead time t asEQUATION)In all metrics we include area weighting w α to handle the fact that the grid cells for the global data has different area. These weights are computed as described in Rasp et al. [40] . For the MEPS data these are all set to w α = 1. Note that the square root is applied after sample averaging, following standard convention and the WeatherBench 2 benchmark [40] .In the probabilistic setting we evaluate an ensemble forecast of K members. Let Xt,(m),(k) be the prediction of ensemble member k. For ensemble forecasts we compute the RMSE of the ensemble mean asEQUATIONTo measure the calibration of uncertainty expressed by the ensemble we use the (bias-corrected) Spread-Skill-RatioEQUATIONEQUATION)If the ensemble members represent realistic forecasts, and observe exchangeability with the ground truth, SpSkR t,j ≈ 1 [13] .As a probabilistic metric we also use the CRPS [15] , here computed as a finite sample estimate [63] over all the ensemble membersEQUATION)Note that this is a spatial average over marginal CRPS values and does not take any covariance structure of the data or model distributions into account. To avoid the quadratic sum in eq. ( 23) we in practice use the idea of Zamo and Naveau [63] to compute this with linear memory by sorting the ensemble members. For deterministic models the CRPS reduces to MAE. See this for example from eq. ( 23) by letting all ensemble members be the same predicted value.One key motivations for probabilistic weather forecasting is to allow for better predictions of extreme events and improved estimates of uncertainties associated with these. To exemplify this we include here a case study looking at forecasts for Hurricane Laura, using the global models.In August of 2020 Laura developed in the Atlantic Ocean and reached hurricane levels in the Gulf of Mexico, eventually making landfall in Louisiana and causing major damages [58] . We study here forecasts for 2020-08-27T12 UTC, which is 6 hours after the hurricane hit land. All forecasts are for this exact time point, but initialized a varying number of days before. We here run 50 member ensemble forecasts using Graph-EFM and deterministic forecasts using the GraphCast* and Graph-FM models. Figure 7 shows 10 m wind speeds in ERA5 and the forecasts. For Graph-EFM we plot both forecasts from randomly sampled ensemble members and a cherry-picked best member that was deemed to most closely match ERA5. Note that the 1.5°resolution that we work with makes determining similarity or exact positions somewhat challenging. The resolution also puts some limitations on how well the most extreme wind speeds can be captured.We see in fig. 7 that at 7 days lead time there is great uncertainty, and the deterministic models do not show the hurricane at all. In the ensemble forecast from Graph-EFM there exists however already members indicating the possibility of a hurricane making landfall a week ahead (see for example the cherry-picked "Best member"). While the ensemble includes many possible scenarios, a total of 7 members show the development of a hurricane in the area. Having information of such possible scenarios a long time ahead allows for planning and readying disaster response efforts that might be needed. Note that discovering these scenarios is only possible through an ensemble forecast, as the deterministic models do not indicate such an event. At 5 days lead time all models are indicating the development of a hurricane. While the deterministic models do a good job here, they are indicating a landfall location slightly too far eastward. The ensemble members from Graph-EFM however show a range of different positions for the hurricane, indicating the uncertainty in the landfall location. At 3 days ahead 42 out of 50 ensemble members show the hurricane making landfall, but still with some uncertainty about the exact location. At 1 day lead time all models give an accurate forecast of the position of the hurricane. Apart from position, it is also interesting to consider how the models capture the intensity in terms of wind speeds. At 1 day ahead the deterministic models are somewhat underestimating the wind speed. The Graph-EFM ensemble shows a range of possible values, indicating the uncertainty in the exact wind intensity. Overall this study exemplifies how ensemble forecasts from a machine learning model can be used to discover possible extreme weather scenarios at long lead times and uncover the uncertainties associated with them.Apart from giving a low forecast error, it is also important for forecasts to look realistic. Specifically, we want forecasts to be physical, containing features that are possible under laws of atmospheric physics. This is necessary for forecasts to be interpretable by meteorologists and builds trust in the MLWP model. One key property for realistic forecasts is that of spatial coherency. Forecasts should take realistic values not just locally, in each grid cell, but show larger-scale features that are consistent over the forecasting area. This is important since the weather system contains spatial dependencies on multiple different scales. Note that in the probabilistic setting this connects closely to capturing the correct joint distribution spatially, rather than just the marginal distributions in each grid cell. As the metrics commonly used to evaluate MLWP models only consider the marginal distributions, it is important to also inspect sample forecasts to gauge the spatial coherency.This modeling aspect is especially relevant for probabilistic models, since being able to visualize and interpret individual ensemble members is an important capability of ensemble forecasting. Other uncertainty quantification techniques can be applied to estimate forecast uncertainty [6] , but these For the final Graph-EFM model forecasts are generally spatially coherent and represent physically plausible scenarios. Examples of this is shown in appendices J.2 and K.2. In this appendix we contrast this with some of the shortcomings of baseline models. We also discuss some of the challenges that we have encountered with achieving spatially coherent forecasts and ways to tackle these.Lead TimeIn the MEPS experiment the benefits of the hierarchical graph structure become the most clear when visually inspecting forecasts. Figure 8 shows a comparison between a few forecasts from Graph-EFM and Graph-EFM (ms). In the Graph-EFM (ms) model samples tend to be more patchy and poorly reproduce spatial features. We connect this to the fact that the randomness in Z t is more local, as it is associated with the nodes in G MS which directly connect to the grid. While the randomness introduced from sampling Z t can spread to the full forecast area using the multi-scale graph, this is something the model has to explicitly learn. For matching marginal distributions in grid cells it can be enough for Graph-EFM (ms) to just have the randomness in each mesh node impact the connected grid nodes locally. We can contrast this to Graph-EFM, where Z t is associated with nodes at the top of the hierarchical graph. For this randomness to affect the prediction it must necessarily be propagated down the hierarchy, which disperses it spatially by construction.We have observed that the CRPS fine-tuning (see appendix D.2) can be central to the performance of Graph-EFM. This aligns the model distribution with the distribution of the data. However, since this objective only encourages matching of the marginal distributions it is not sufficient for making the model capture the full joint distribution. This has to be achieved by constraints to the model or through additional parts of the training objective. We have observed that the weighting λ CRPS used for the CRPS term can have a large impact on the spatial coherency of forecasts. When λ CRPS is chosen too high the models trade off the ability to generate meaningful large-scale patterns for capturing local variations. This is especially severe for the Graph-EFM (ms) model, forcing us to use lower λ CRPS values in order to still get physical-looking forecasts. Figure 9 shows an example of the type of local noise-patterns that can appear when training Graph-EFM (ms) with a too high λ CRPS .Another problem observed with the multi-scale graph G MS in the LAM setting is the appearance of circular artifacts over the forecasting area. These appear in samples from Graph-EFM (ms) (see figs. 8b and 8d ), but can also be found in deterministic forecasts from GraphCast*. These artifacts can be traced to the heterogeneous structure of the multi-scale graph. Due to how G MS is constructed, mesh nodes can have different number of neighbors. The artifact positions match the positions of mesh nodes with many neighbors. Note that the hierarchical graph does not have this issue, as G 1 (that connects to the grid) has a more uniform structure. We noticed that these artifacts are even more prevalent for smaller models, as shown in fig. 10 . This points to the fact that with increased capacity, the GraphCast* model can to some extent learn to compensate for this problem.We have found it less challenging to achieve spatially coherent forecasts in the global setting. Both Graph-EFM and Graph-EFM (ms) generally produce samples with realistic physical features. Early in the training process we observe some hexagonal patterns for longer lead times (also noted by Keisler [20] ), but these disappear as we train the models on longer rollouts.One challenge that we encountered with the global probabilistic models was the appearance of small areas of instabilities when rolling out the models to longer lead times. In fig. 11 we show some examples of this in forecasts from preliminary models. These unphysical artifacts typically covered Figure 11 : Examples of artifacts in preliminary version of the global models unrolled to 10 days lead time. Small patches of deviating values appear and impact multiple variables. While these were more prevalent in Graph-EFM (ms), we did also observe this issue in Graph-EFM. only a few grid-cells, but are of course undesirable in a forecasting model. To remedy this for Graph-EFM we found it sufficient to train on longer rollouts. The stabilizing effect of unrolling up to T = 12 time steps during training did suppress these issues. This was however not enough to solve the problem for Graph-EFM (ms), for which we additionally needed to further lower λ CRPS .In this appendix we give further detail about the graphs, data and experimental setup used in our global forecasting experiment.Multi-scale graph The global multi-scale graph is created in the same way as in GraphCast [23] , by recursively splitting the faces of an icosahedron. As we work with data on a coarser resolution, we perform only 4 steps of such splitting, resulting in the graphs G 5 , . . . , G 1 . These are all then merged to create G MS , used by GraphCast* and Graph-EFM (ms). By splitting 4 times we end up with a ratio of N |V M | ≈ 11 between grid nodes and mesh nodes. This can be compared with the 6 splitting steps of GraphCast, resulting in a ratio N |V M | ≈ 25 for their 0.25°data. We initially experimented also with splitting only 3 times (resulting in N |V M | ≈ 45), but models using these graphs showed inferior performance.Hierarchical graph As the original icosahedron G 5 contains very few nodes, we do not use this for the hierarchical mesh graph, but rather construct the hierarchy using G 4 , . . . , G 1 . Edges E l,l+1 between levels are constructed by connecting each mesh node at level l with nodes at level l + 1 within a distance of 1.1 times the edge length in G l . This guarantees that each node at level l has 1 or 2 connection to the level above (see fig. 12g ). Downward edges E l+1,l are created by simply flipping the edge directions of E l,l+1 . We visualize the global mesh graphs in fig. 12 and list the exact number of nodes and edges in table 4. Connecting the grid and mesh Edges connecting the grid and mesh graph are also constructed following Lam et al. [23] . The edges E G2M are created by connecting each grid node to mesh nodes within a distance of 0.6 times the edge length in G 1 . Edges E M2G are constructed by for each grid node finding the triangle in G 1 containing it, and adding three edges from the corner mesh nodes of that triangle. Note that the edge sets E G2M and E M2G are identical for the multi-scale and hierarchical mesh graphs, since G MS and G 1 contain the same nodes. Since the grid nodes are laid out in a latitude-longitude grid and the mesh is icosahedral all mesh nodes will not be connected to the same number of grid nodes. Specifically, mesh nodes close to the poles will have many more grid connections than mesh nodes around the equator.All mesh nodes are associated with static features related to their position, specifically cos(Latitude), sin(Longitude) and cos(Longitude). Edges (in the mesh, E G2M and E M2G ) have static features containing the edge length and vector difference between its endpoints (see Lam et al. [23] for details). Edge features are normalized by the length of the longest edge.We use a version of ERA5 [17] re-gridded to 1.5°latitude-longitude gridding, provided through the WeatherBench 2 benchmark [40] . Data from the period 1959-01-01T12 to 2017-12-31T12 is used for training, 2017-12-31T18 to 2019-12-31T12 for validation and 2019-12-31T18 to 2021-01-10T18 for testing. These exact time stamps guarantee that there is no overlap between the subsets. For evaluation we consider forecasts initialized at times 00 and 12 UTC each day of 2020, following WeatherBench 2. We perform 10 day forecasts using 6 h time steps. During training we also start forecasts at times 06 and 18 UTC.The exact set of forecast variables, forcing and static fields is listed in table 5 . We use the same inputs as Lam et al. [23] . The forcing is windowed over three consecutive time steps, meaning that each F t contains forcing from times t -1, t and t + 1 as well as the the static fields. For training we rescale the values of each variable to zero mean and unit variance.We train all models using the AdamW optimizer [57] and utilize BFloat16 mixed precision to save GPU memory. The training costs for the models makes extensive hyperparameter tuning unfeasible. We choose hyperparameters based on initial experimentation with smaller models. For Graph-EFM the important weightings λ KL and λ CRPS in L can be chosen based on monitoring the model behavior during training. This was done by plotting example forecasts throughout the training process and monitoring metrics on the validation set. The weight λ KL was tuned to prevent the model from collapsing to deterministic predictions, and to still achieve useful predictions when Z t was sampled In this appendix we give further detail about the graphs, data and experimental setup used in our experiments with the MEPS data.Multi-scale graph In this limited-area setting we construct graphs as regular quadrilateral meshes covering the rectangular MEPS forecasting area. To create these we lay out mesh nodes in regular rows and columns over the area. Each node is then connected with bidirectional edges to its neighbors horizontally, vertically and diagonally (see fig. 13a ). This results in all nodes (except those at the edge of the area) having 8 neighbors. The procedure is repeated at 4 different resolutions, tripling the distance between nodes at each resolution. This means that a node at resolution level l is positioned at the center of a group of 3 × 3 nodes at resolution level l -1, sharing its exact position with the center node of the group (illustrated in fig. 13b ). To create the multi-scale mesh graph G MS we then merge the graphs at different resolutions, combining any nodes that sit at the same coordinates into one node. Note that this is possible due to how nodes align across the resolution levels.Hierarchical graph For the hierarchical model the graphs of different resolution are not merged, but used as the different levels of the hierarchy. We include only the 3 finest meshes G 1 , G 2 and G 3 , as G 4 contains only 9 nodes. Additional inter-level edge sets E l,l+1 and E l+1,l are then created for all adjacent levels. Each set E l,l+1 of upwards edges is created by connecting each node on level l with the closest node on level l + 1. This means that each node at levels l > 1 will have 9 incoming edges from the level below. The downward edges E l+1,l are a copy of E l,l+1 with the direction of each edge flipped. The mesh graphs used for the MEPS experiment are visualized in fig. 14 and the number of nodes and edges in each graph listed in table 9 .Connecting the grid and mesh To form E G2M , each grid node is connected to mesh nodes closer than 0.67 times the distance between nodes in G 1 . All distances are here 2-dimensional euclidean, computed in the MEPS Lambert projection coordinates. The set E M2G is constructed by iterating over the grid nodes, and at each creating edges from the 4 closest mesh nodes to the node in the grid. The static features associated with mesh nodes are their 2-dimensional coordinates in the MEPS Lambert projection, normalized by the maximum coordinate value. All edges have static features including the length of the edge and the vector difference between the source and target nodes (using the projection coordinates). The edge features are normalized by the length of the longest edge in the whole mesh graph.Dataset The MEPS dataset consists of archived forecasts from the operational MEPS system during the period April 2021 -March 2023. This period was chosen due to the system configuration being reasonably stable, preventing distributional shifts within the dataset. From the chosen period we extract the forecasts started at 00 and 12 UTC each day. At each initialization time there are 5 ensemble forecasts, started from slightly different initial conditions. This results in 10 forecasts per day (assuming no ensemble members are missing due to operational issues). When retrieving the data we additionally downsample the spatial resolution from the original 2.5 km to 10 km. This results in a dataset of 6069 forecasts of length 66 h with 1 h time steps. We split the forecasts into training, validation and test sets according to fig. 15 . The specific validation months were chosen to reasonably cover the seasonal variations. Note that the dataset contains around 46 years of individual time steps. However, since there are obvious correlations between ensemble members and successively started forecasts the actual information content is far lower than in a 46 year reanalysis dataset.Variables and forcing At each grid cell we model 17 weather variables, including a broad range of different quantities and different vertical levels in the atmosphere. All variables, forcing and static fields are described in table 10. The particular choice of variables was motivated by a combination of meteorological relevance, data availability and striving for a diverse set of variables to evaluate the model on. We use the same type of windowing for the forcing and standardization of variables as in the global experiment. For solar radiation (nlwrs and nswrs) we consider the net flux at ground level, aggregated over the past 3 hours (since the last time step). Apart from the solar radiation all other variables are instantaneous. For the MEPS data we use the fraction of open water in the grid cell as a forcing input. We assume this to be constant over the forecast period and take the value from the time of the initial state. In the global experiment the land-sea mask is static, but treating this as forcing could be useful for taking into account seasonal fluctuations of the ice cover in the Nordic region. The boundary forcing B t consists of the same variables as listed in table 10 . We include a static binary indicator variable describing if a node is in the boundary or forecast area. for unrolling the model. The combination of 1) the 3 h time steps, 2) using the last two weather states as model inputs, 3) including forcing from multiple times as input, means that we reduce the effective length of our ground truth forecasts in pre-processing. This explains why we predict for 57 h rather than the original 66 h. Note however that nothing prevents us from unrolling the models to create forecasts for 66 h or beyond. Although this is possible, we do not have ground truth data to compare against past 57 h and therefore view such experiments to be of limited interest. There is still no reason to expect the model performance to become drastically worse specifically past 57 h, as all models are anyhow fine-tuned on shorter rollouts than this.In all models we use boundary forcing B t in order to include information about the surrounding area. At each time step boundary forcing B t-2:t-1 is passed to the model together with X t-2:t-1 , F t , as described in fig. 16 . There is a separate set of grid nodes used for the boundary input. These nodes are treated identically as grid nodes within the forecasting area by the GNN layers. As the boundary forcing is only a change to the model inputs this does not require any substantial re-design and we can adapt all models in our experiments to this LAM setting. Note that we use a boundary area that lays inside the original MEPS area, allowing us to use parts of the ground truth forecasts as boundary forcing. We specifically define the boundary as the 10 grid cells closest to the edge of the limited area. In the operational system the boundary area is defined along the edge outside the MEPS area.There is no major conceptual difference between these and one could easily re-define the different areas to match the operational MEPS system. In this appendix we present additional results for global forecasting.Comparisons between existing MLWP models are challenging, due to the many factors that impact forecast quality. While most global models proposed in the literature are trained on ERA5, the exact choice of train/test split, spatial resolution and forecasted variables can vary (compare for example Keisler [20] , Pathak et al. [37] and Hu et al. [19] ). Models that operate on a higher spatial resolution and with more variables get more information in initial states, and should therefore be expected to produce better forecasts. When it comes to MLWP ensemble forecasting, an even greater challenge is that different initial conditions are used [39, 6, 16] . Given this situation, it is hard to disentangle the proposed machine learning methods from their surrounding design choices and make fair comparisons of model architectures.We approach these complications by re-training a smaller set of models on the same data and experimental setup. Our focus is on graph-based, non-hybrid MLWP models, as this is the regime of Graph-FM and Graph-EFM. The baseline models are described in detail in appendix C.5. In this appendix we also include additional metrics taken directly from the WeatherBench 2 Benchmark [40] . These are for the original GraphCast model [23] , KeislerNet [20] and the IFS-ENS operational ensemble from the European Centre for Medium-Range Weather Forecasts (ECMWF) [12] . The evaluation setup for these results match ours (evaluation against 2020 from ERA5, identical metrics computed on 1.5°data). However, as discussed above there are differences in the exact training data, variables and operating resolutions compared to our models.Metric values from our global experiments are presented in figs. 17 to 19. We here showcase values for surface variables and atmospheric variables at pressure levels 500, 700 and 850 hPa. The results taken directly from WeatherBench 2 are shown in gray, to emphasize that any comparison with these comes with caveats. Note that results for all variables are not available for all models.In fig. 20 we showcase example forecasts from different models at lead time 10 days. Note that forecasting this far into the future is challenging, and sampled ensemble members show only one possible scenario predicted by the models.Figure 21 shows example ensemble forecasts from Graph-EFM for surface variables, atmospheric variables at 700 hPa, z500 and t850. All forecasts are for 10 days lead time. For practical reasons we use only 20 members to estimate the ensemble mean and standard deviation in these plots. (i) t500 Lead time (Days) Lead time (Days) (n) v850 (q) z1000 Given the usefulness of Propagation Networks in Graph-EFM it is reasonable to ask if these could be beneficial to use also in the deterministic Graph-FM model. Also Graph-FM uses a hierarchical mesh graph, so the propagation of information between levels is important also for this model. We test this empirically, by training versions of Graph-FM using Interaction and Propagation networks. As the retention of information is also important for some parts of the architecture, we do not replace all GNNs with Propagation Networks (see appendix C.2). We compare models on both global and LAM forecasting, using the same experimental setups as described in appendices H and I.Global forecasting with ERA5 Figure 27 shows RMSEs for highlighted variables from ERA5 and Graph-FM using Interaction and Propagation Networks. The models have very similar errors for almost all variables. However, for variables in the upper atmosphere the Propagation Network model shows lower errors.Limited area modeling with MEPS data In fig. 28 we compare RMSEs for Graph-FM models using the different GNN layers on MEPS data. Here we see a greater advantage of the Propagation Networks. We hypothesize that this relates to the boundary forcing. Using Propagation Networks the information from boundary nodes should easier reach the top graph level, where it can faster spread throughout the forecasting area. Motivated by these results we use Propagation Networks in our final Graph-FM architecture, both for global and LAM forecasting.In Graph-EFM we define the latent map p Z t X t-2:t-1 , F t using a neural network with explicit dependence on X t-2:t-1 , F t . Given that also the predictor takes X t-2:t-1 , F t as inputs, it is not immediately clear that conditioning on this also in the latent map is necessary. Indeed, due to the ability of deep neural networks to well approximate arbitrary functions, the predictor network should internally be able to transform a simple Z t ∼ N (0, I) to introduce the dependence on X t-2:t-1 , F t . This does however assume infinite flexibility in the predictor, which might be far from the situation in practice.To investigate the importance of a learnable latent map we compare our Graph-EFM model on the MEPS dataset with one where Z t is sampled from a static distribution N (0, I). This experiment was carried out using an earlier version of Graph-EFM, with the same architecture but a slightly different training schedule. To save on computations we here only sample 16 ensemble members from each model. RMSE, CRPS and SpSkR are shown in figs. 29 to 31. In terms of RMSE and CRPS there is a clear benefit to letting the distribution over Z t be a learnable mapping. The SpSkR in fig. 31 shows no clear trends between the models. For practical network architectures the latent map does add flexibility, changing the mean of Z t that enters the predictor. The learnable latent map should also simplify the inference problem solved when optimizing our variational objective. With the dependence on previous states we can expect a smaller discrepancy between p Z t X t-2:t-1 , F t and p Z t X t-2:t-1 , t , F t , simplifying the optimization of the variational approximation q Z t X t-2:t-1 , X t , F t . We believe that this is an important aspect of the observed empirical benefit of a using the latent map.We here study how the performance of the model, in terms of metric values, varies when sampling different numbers of ensemble members. To investigate this we ran the evaluation of Graph-EFM with 5-80 members for the global model and 5-100 for the LAM model. Results for a selection of variables are shown in fig. 32 (global, ERA5) and fig. 33 (LAM, MEPS). As expected the RMSE of the ensemble mean decreases when sampling more members. However, already when sampling 20 or 25 members the results are fairly close to the full ensemble. For SpSkR the differences are even smaller. As the CRPS is a property of the distribution of the model forecast, its true value does not depend on the number of samples drawn. In practice we compute CRPS using an unbiased estimator, and the variance of this estimator decreases with ensemble size. When averaged spatially and overThe deterministic Graph-FM model was first proposed in a preliminary version of this work[35], but there only for the LAM setting under the name Hi-LAM.Our code is available at https://github.com/mllam/neural-lam/tree/prob_model_global (global forecasting) and https://github.com/mllam/neural-lam/tree/prob_model_lam (LAM).Provided by the Copernicus Climate Change Service under the ECMWF Copernicus License. 
 
 Appendix:  


Decision: The authors propose an autoencoder version of GraphCast and UNet for probabilistic global medium-range weather forecasting.

- **Reviews**: While the paper received (very) high scores from some reviewers, some reviewers were initially more critical.
- **Rebuttal**: The authors successfully addressed the concerns of the more critical reviewers by discussing computational challenges and justifying their approach through experiments.
- **Decision**: I recommend acceptance.

Reviews: review: paper_summary: paper_summary: The paper proposes Graph-EFM, a method that combines a hierarchical multi-scale graph neural network with a variational objective for probabilistic weather forecasting. The method performs on par with Graphcast on deterministic metrics with the extra benefit of uncertainty estimation.
strength_weakness: strength_weakness: - The authors should replace Table 1 with a line graph figure instead, as it allows comparison across different variables and lead times.
- Please see my questions below.
questions: questions: - How important do you think the architecture is to the performance versus the objective function? The proposed architecture has an intuition similar to UNet, i.e., multi-scale features and the lowest layer can be used to parameterize the hidden variable.
- Diffusion models are considered the best family of models for generative modeling, surpassing GANs and latent variable models for other fields such as computer vision. What is the reason to believe latent variable models are the way to go for probabilistic weather forecasting?
- Why does the paper compare with Graphcast+SWAG but not the perturbed version of Graphcast and Gencast?
- How does the performance vary w.r.t. the number of ensemble samples? Given that sampling from a latent variable is fast, have the authors tried using more ensemble members?
- Is there an explanation why Graphcast is better than Graph-FM, but Graph-EFM is better than Graph-EFM (ms)?
- Why in LAM, the Graphcast architecture is doing better than the proposed architecture?
limitations: limitations: See above.
paper_summary: paper_summary: The authors propose a graph-based ensemble forecasting model (Graph-EFM) to provide weather prediction with a hirearchical GNN framework. They used a hierarchical mesh graph to handle the challenges of capturing processes unfolding over different spatial scales and modeling the uncertainty in the chaotic system. The Graph-EFM provides a probabilistic weather forecasting with the benefit of capturing forecast uncertainity. The experiment results show the effectiveness and advantages of Graph-EFM compared to other deterministic models.
strength_weakness: strength_weakness: In Figure 3, it seems like the selected ensemble members vary a lot, and how close is your forecast to the ground truth. Possibly, explaining a little bit of the underlying meaning of the measures in table 1 & 2 in the paper.
questions: questions: 1. The authors could better explain the underlying meaning of meaures, RMSE, CRPS, and SpSkR, for the results of weather forecast. Basically, I would like to ask authors to show how close their Graph-EFM's forecast is to the ground truth weather.
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_summary: This paper introduces a new method for predicting weather using advanced deep learning models. The approach, called Graph-EFM, improves accuracy and better handles uncertainties in weather forecasts. It uses a 1.5 degree version of ERA5 and making weather predictions more reliable and useful for real-world applications.
strength_weakness: strength_weakness: What happened if 0.25 degree ERA5 is used?
questions: questions: I need to some figures during extreme events, e.g cyclones like Yaku.

Can the model deal with higher resolution data?

Or ERA6 when available?

Or more localised higher resolution data?

Do the results change a lot if it's only trained from 1980 onwards?
limitations: limitations: Not many, this is a fantastic piece of work
paper_summary: paper_summary: This work introduces a VAE variant of GraphCast for global medium-range weather forecasting and a VAE variant of a UNet (that is formulated as a GNN) for limited area modeling over Scandinavia. For this, they adapt GraphCast to have a similar hierarchical structure to UNets, and then treat the coarsest hierarchical layer (the bottleneck) as a latent variable representing the mean of isotropic Gaussians. The ensemble predictions from the model are similarly fast as a single deterministic prediction, achieved through batching. Their calibration can be good for some variables (e.g. global t2m for 10 day lead time has a spread/skill ratio of 0.99), while poorer for others (e.g. local wvint has a spread/skill ratio of 0.57 for 24h lead time).
strength_weakness: strength_weakness: Major points:
1. Questionable baselines: I am unsure if the chosen baselines are very strong, let me name a few reasons for this:
    - Tab 1 presents performance for GraphCast, e.g. RMSE=387 for z500, 5 day leadtime. However, if I check the headline scores in the WeatherBench 2 (https://sites.research.google/weatherbench/) for GraphCast, i see RMSE=274 for z500, 5 day leadtime, which is significantly higher and beats all models presented in this study.
    - Both Tab 1 and Tab 2 do not include scores for the conventional weather models. I would expect Tab 1 to include IFS & IFS-ENS scores and Tab 2 to include MEPS scores.
    - Since this work introduces a probabilistic weather model, i would expect comparison with other recent works on probabilistic weather models, like the ones cited in this paper (e.g. GenCast).
    - Graph-FM has almost 6x the parameters compared to GraphCast* (Tab 5) - which quite possibly could be the major reason for its improved performance, and not the introduced architectural feature of hierarchical layers.
2. Overlooked connection to UNets: The Graph-FM that was introduced for the LAM setting looks to me as equivalent to a UNet:
    - The input data comes on a regular grid with 236 x 268 pixels. Which is subsequently downsampled using 3x3 windows. Processing at each depth level is done with a locally connected layer (in other words: a local convolutional filter). A semantically simpler description of such a model would be a UNet with 3x3 pooling (learned in this case) and 3x3 conv filters at each stage. Possibly, implementing it as a UNet could also be computationally advantageous, making use of the highly optimized kernels for 2d convolutions and pooling operations, instead of GNN layers that rely on scatter sums.
    - UNets have been previously used for Weather forecasting: e.g. https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2018GL080704 & https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2020MS002203
3. Proposed VAE implementation physically not meaningful? Your VAE is implemented with a latent variable at the coarsest level. This is supposed to capture epistemic uncertainty related to the forward model (and not due to initial state). However, one may argue for atmospheric models most model uncertainty comes from the subgrid-scale parametrizations and not from the coarse-scale representation of atmospheric dynamics. Hence, to me it seems far more intuitive to introduce the stochasticity at the finest level, representing the small scales. I assume you chose the hierarchical description mostly for computational reasons, but given a lack of physical basis, i would at least expect a more thorough investigation of potential errors introduced by this, e.g. is the ensemble variability too smooth?
4. Missing reference to previously published work? A workshop paper at last years NeurIPS has introduced both the hierarchical GNN and the MEPS dataset https://arxiv.org/abs/2309.17370 , if I am not misstaken. I am not really sure about NeurIPS policy here, but even if this work is a direct extension of the previous work and the previous work is to be considered as non-archival, I still believe you should at least cite the workshop paper.

Minor points:
1. GraphCast + SWAG: This is a baseline with poor performance, that is somewhat arbitrarily picked from many possible approaches to obtain ensemble predictions from neural networks. I see two options here: Either you keep it, but also introduce many other such baselines, to make clear that you did not cherrypick a particularly weak one. Other approaches  that should not be prohibitively expensive to run could e.g. be MC-Dropout or Laplace Approximation. Or, you simply drop it, as is, it does not  add much to the paper.
2. Introduction lacks motivation for LAM: This is an ML conference that you are submitting to. It would probably be good to briefly motivate why doing LAM is even necessary (i.e., why can't we just rely on global models instead)?
3. Extreme Weather evaluation / case study: One key reason for ensemble prediction is capturing the tails, i.e. the extremes. You state in Appendix A that this is out-of-scope for the work. I would argue you are making your life too easy here. Since the presented models are likely not useful unless they display decent performance also for extreme weather, it would be important to evaluate just that. It may be enough for this paper to e.g. study a single extreme event as a case study.
questions: questions: Why does the original GraphCast paper not report the visual artifacts that you found for Graph-EFM (ms)? Could it be that your models have simply not been trained sufficiently or that there is a bug in your implementation?

Why do you use a fixed variance of the latent variable for the global predictions? It would be interesting to see this ablation.
limitations: limitations: The section on limitations is somewhat short. I believe the key limitation of this work is in the evaluation, i.e. it is unclear to me after reading the work how well the models perform. E.g., are these models robust for longer rollouts or if applied at prediction time further away from training time?
Moreover, the work does transparently communicate a limitation of their multi-scale Graph-EFM: visual artifacts. But, this is strange to me, as it could mean two things: a) the VAE formulation and multiscale edges simply don't work well together or b) since GraphCast did not observe such issues, the presented Graph-EFM (ms) is simply not trained sufficiently or suffers from a buggy implementation. While a) would be an interesting finding, i believe b) can not be ruled out given the presented results in this paper.

Comments: comment: We are happy to hear that our clarifications could address some of your concerns. Thank you again for actively engaging during this discussion period and for much valuable input on the paper!
Thanks, this partially addressed my concerns, I am thus adjusting my rating.
**Global results** (corresponding to Table 1 in paper)

|           |                   |       | 24h   |        |       | 57h   |        |
|-----------|-------------------|-------|-------|--------|-------|-------|--------|
| Variable  | Model             | RMSE  | CRPS  | SpSkR  | RMSE  | CRPS  | SpSkR  |
| z500      | Graph-EFM + ICP   | 410   | 193   | 1.34   | 707   | 319   | 1.20   |
|           | Graph-FM + ICP    | -     | -     | -      | -     | -     |        |
|           | GraphCast* + ICP  | 475   | 244   | 0.66   | 914   | 486   | 0.46   |
| 2t        | Graph-EFM + ICP   | 1.66  | 0.73  | 1.07   | 2.36  | 1.03  | 1.06   |
|           | Graph-FM + ICP    | -     | -     | -      | -     | -     | -      |
|           | GraphCast* + ICP  | 2.14  | 1.11  | 0.31   | 3.48  | 1.82  | 0.35   |

(Unfortunately the results for Graph-FM + ICP did not finish in time, and we do not want to delay our answer further.) 

**LAM results** (corresponding to Table 2 in paper)

|           |                   |       | 24h   |        |       | 57h   |        |
|-----------|-------------------|-------|-------|--------|-------|-------|--------|
| Variable  | Model             | RMSE  | CRPS  | SpSkR  | RMSE  | CRPS  | SpSkR  |
| z500      | Graph-EFM + ICP   | 174   | 90    | 0.98   | 222   | 113   | 0.72   |
|           | Graph-FM + ICP    | 233   | 124   | 0.73   | 357   | 192   | 0.43   |
|           | GraphCast* + ICP  | 156   | 82    | 0.93   | 200   | 109   | 0.42   |
| wvint     | Graph-EFM + ICP   | 1.63  | 0.8   | 0.9    | 2.11  | 1.02  | 0.83   |
|           | Graph-FM + ICP    | 1.67  | 0.89  | 0.36   | 2.50  | 1.39  | 0.22   |
|           | GraphCast* + ICP  | 1.54  | 0.82  | 0.38   | 2.10  | 1.18  | 0.20   |

We can note that the ICPs do not result in improved RMSE values for any model, instead making it worse. Visually inspecting the samples (sadly we cannot include figures here) show that the noise tends to stay in the forecast over time, and does not induce physically plausible diversity in the ensemble. The added noise does as expected lead to improved CRPS-values for GraphCast* and Graph-FM, as any added variance up to the model error would improve the CRPS. As the noise creates some spread, the SpSkR also takes values higher than 0. We note however that in these evaluations the SpSkR starts high and rapidly decreases with lead time, as can be expected since all the stochasticity comes from time 0. Adding ICPs to Graph-EFM thus also results in higher SpSkR, even when it was already close to 1. As mentioned this is a simple experiment and there are better ways to perturb ICs. However, we hope that it shows that Graph-EFM (which also uses independent noise, but on the graph) does something more than just use this noise to perturb initial conditions.
5. Regarding U-net connection 

We are happy to elaborate on our answer here, as the rebuttal left us a bit short on space to discuss this point. It is indeed possible that you can put together the components you have listed to construct something that would be basically the same as our LAM model using the specific graph constructed for the MEPS area. To get something *equivalent* to the InteractionNetwork layers would still require quite some careful tailoring of these components, and would not resemble a standard off-the-shelf U-net. Although we do agree that a comparison to an off-the-shelf U-net model (possibly also with a VAE modification) would indeed be interesting, even though we lack the time now. 

We again want to stress that this discussion about U-net similarity is fully confined to the specific graph created for the MEPS experiments. This should be interpreted as one example application of a general framework. We think that this aligns with the presentation in the paper. The global model is another example of applying the framework, that is clearly different from a U-net. The specific LAM graph construction is one example of how a LAM graph could be created. There are many options to further explore there, including triangular and irregular graphs, which have no equivalence to a U-net.  

6. Regarding clarification 6 (randomness at different spatial scales) 

Thank you for the clarification, we understand your point better now. We base our model on an assumption that the effective dimension of the dynamics is lower than the dimension of the grid. This is modelled by setting the dimension of the latent random variable to be significantly lower than the data. We agree with your point that sampling the randomness at the top level and deterministically mapping this down does limit what kind of variation that can be modeled. However, this is part of our modeling assumption and since the mapping is non-linear and the random variables contain many dimensions of randomness we do not believe it to be a significant limitation. We back up this modeling choice with our experiments using Graph-EFM (ms), where introducing the noise closer to the grid has resulted in a worse model with less realistic forecasts. 

We are not entirely sure about how you mean the uncertainty in the model would be related to ICs rather than the dynamics. As the forecast is sampled per time step, the randomness is included in the mapping from one time point to the next, rather than only between time 0 and 1. Uncertainty in initial states corresponds to a distribution $p(X^{-1:0})$ and combining initial state perturbations with a deterministic forecasting model corresponds to sampling from some approximation of this distribution and then applying deterministic dynamics (meaning that $p(X^{1:T} | X^{-1:0}, F^{1:T})$ is a dirac distribution). This is fundamentally different from what we do in Graph-EFM, where we model the distribution $p(X^{1:T} | X^{-1:0}, F^{1:T})$ as a product of random dynamics $p(X^t | X^{t-2:t-1}, F^t)$ at each time point (eq. 2).  

But to give some more clarity on this matter we have tried evaluating the models with perturbations added to the initial conditions. Within this short time frame we have only been able to do a simple experiment here, where we use our pre-trained models and add independent Gaussian noise to $X^{-1:0}$. We tune the scale of this noise on the validation set. The results for models with Initial Condition Perturbations (ICPs) are shown in the tables below. 

(See continuation below for tables)
Thank you reviewer PLfK for your careful consideration of our response. We are happy to discuss and clarify our points further! 

1. Regarding the method contribution 

We find it important to point out that we are not just applying off-the-shelf methods to weather forecasting, performing a purely empirical study. As all deep learning papers we build upon established methods, but combine these in novel ways motivated by the application. In the paper we contribute details all the way from training objectives to the internal workings of GNN layers. Hence we think our contribution is also methodological, with a focus on making the method work for weather forecasting. 

2. Regarding scaling to higher spatial resolutions 

It is true that we limit the study to 1.5° resolution data for computational reasons. As we are sure the reviewer is aware, but we would like to clarify here for completeness, the computational requirements for training models on 0.25° data are substantial. Already now we use 8 high-end GPUs for a week in order to train one model. For comparison, the original training of GraphCast on 0.25° data took 4 weeks on 32 devices. By considering this and our training times for GraphCast* and Graph-EFM, we can get a rough estimate that training a 0.25° version of Graph-EFM would take more than 6 months using the same 8 GPUs. 

In our opinion it is crucial that new ideas, models, and methods can be proposed and evaluated at scales manageable by academic researchers who do not have access to the same computational and engineering resources as a company like DeepMind. This allows for much faster advancements in machine learning methodology for the ML weather prediction area and enables for the best ideas to then be integrated and scaled up in high resolution models. Indeed, we believe that the evaluation that we present in the paper on a lower resolution is indicative of the results that one might expect on a higher resolution, and that these results are sufficient to show that the model that we propose is of interest to the research community.  

3. Clarification regarding GenCast comparison 

We would like to reiterate that for GraphCast there is code openly available, making it much more feasible to implement and retrain. While the GraphCast* model is implemented in our codebase, we still rely on the GraphCast code for parts of the global model. It is of course not impossible to reimplement and retrain GenCast, but doing this without access to the implementation or detailed documentation would risk getting important details wrong, leading to unfair comparisons.  

4. Regarding the evaluation of LAM models 

The perfect emulator of a system should be one that gives exactly the same output as the system. In this way it makes sense to evaluate a MEPS emulator by comparing to what the MEPS system outputs (the forecasts). However, we do agree with your point that this is not the full story here. As the MEPS system itself has some error w.r.t. the true weather, the error of the emulator should be considered in this context. The reported error of the emulator (compared to MEPS forecasts) can in this way be seen as an upper bound on how much worse its performance is compared to MEPS (on observations).  

While we do not have data for synoptic stations readily available at the moment to run such an evaluation, one can use the values reported in [37] to put errors into some context. Figure 9a in [37] shows MAE of 2m temperature for different months in 2014 and 2015. The corresponding MAE, when comparing Graph-EFM to the MEPS forecasts in the test set, is 0.61 (and similar for other models). Note that we have reported RMSEs rather than MAEs in the paper, but we have evaluated also the MAE values. An increased error of 0.61 is not negligible, but comparable to the difference between IFS and MEPS over Norway in the figure. Now if one considers trading off the superior MEPS forecasts for a model with a similar error to IFS over the region, but that runs in seconds and can create arbitrary large ensembles, this has some clear usefulness.  This comparison comes with a whole bag of caveats (especially since we are comparing values for different years), but the point is mainly to give some example of how the magnitude of emulation errors relates to model errors on observations. 

(See continuation in answer below)
Dear Authors,

thank you for your effort to respond to the points raised in my review.

After reading your rebuttal and the referenced sections in the paper again, I would like to ask you for some further clarifications.
I believe the merit of this work mainly consists in an experimental study of cleverly adjusting multiple known concepts from main-stream deep learning to probabilistic weather forecasting. Hence, it is most important to demonstrate two things: 1) that the new architectures are robustly advantageous compared to existing state-of-the-art and 2) that these advantages are relevant for practical use.
For computational reasons you choose to train and evaluate on ERA5 at 1.5° resolution, which is arguably not relevant for practical use, and also not the resolution at which current SOTA is trained. You try to mitigate this by re-implementing one of the SOTA methods (GraphCast) to run at 1.5°, and then show at that resolution your proposed changes lead to similar performance deterministically, but can create fast probabilistic ensembles with improved skill. The crux is, your experiments do not demonstrate that these advantages are robust to scaling.
Also you argue you are not comparing to probabilistic SOTA (e.g. GenCast), because code is not open source. But given you reimplemented GraphCast, i suspect you should be able to reimplement GenCast also. So while I understand your reason, I do not believe you can shy away so easily from comparing against SOTA probabilistic models (or also: initial condition perturbations, as other reviewers point out).
For the LAM models, thank you for clarifying that you evaluate against MEPS forecasts, I overlooked that. While this is of course valid to compare the different neural networks, it does not give an indication how good the emulation is. Now a weather service would arguably only use the emulator if they understood, what they loose in terms of performance compared to the conventional model, and more importantly, if the performance drop is not too large. For this it would be important to have an independent reference. For instance, the cited reference [37] for AROME-MetCoOp evaluated on 154 synoptic weather stations, at such stations, you could compare the performance of your neural network emulators against the original MEPS forecast, and also against the IFS-ENS. It would be important to show that your emulators are at least outperforming the IFS-ENS.

Regarding the difference to UNets, could please elaborate what you mean with "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."?
Let me elaborate on why I believe they could be basically the same: On a regular grid, a message passing layer is a nonlinear transformation of local features (both edges and nodes), followed by a (possibly weighted) sum of multiple features per grid cell. If the graph is such that the weighted sum includes features that are e.g. in 2x2 quadratic neighborhoods, then this operation is essentially a kernel. In UNets we learn these kernels. For the nonlinear transformation of local features, you can use 1x1 Conv layers. For downsampling, you just use dilation, for upsampling, linear interpolation followed by 1x1 Conv.
There would of course be a fairly simple route to alleviate my concern here: actually training a UNet-based VAE on the MEPS data and compare the performance to the proposed Graph-EFM, but I do understand there is too little time left for the rebuttal to do this. By the way, I consider this point to only influence the assessment of methodological novelty in this paper, but if you can prove that your MEPS emulators are indeed practically useful, such methodological novelty might not even be necessary at all.

Regarding your clarification 6, you write "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.", which is clear to me, but was not the point of my concern. Maybe my writing was not clear, let me restate. Because you only draw random variables at the top level, the model lacks fine-scale randomness beyond the top level resolution. This does not mean that the top level random variables can not influence finer scales, but rather that this influence is deterministic. By the way, this touches upon the point of comparing to IC perturbations as a baseline: quite possibly, the uncertainty your model outputs is related to the initial conditions, instead of the forward atmospheric model, so running Graph-FM with perturbed IC and comparing with Graph-EFM with perturbed IC could give insights here. But again, too little time left in the rebuttal..

Nonetheless, I will reconsider my rating after the rebuttal period.
Thanks!
Thanks for your input!
We agree that this could be interesting to investigate in more detail. However, we believe that it should in that case also involve a version of Graph-EFM with initial condition perturbation, to understand the effect of different sources of randomness on the final uncertainty estimates and to ensure that we compare "apples and apples". Note that there is nothing that prevents us from combining the latent variable formulation of Graph-EFM with initial condition perturbation, but in our view these mechanisms conceptually correspond to different sources of uncertainty: uncertainty about the initial conditions, and uncertainty in the modelled dynamics even in the (hypothetical) case of perfectly known initial conditions. A version of Graph-EFM with initial condition perturbations should additionally be trained with the perturbations present in order to learn the correct distribution. We will consider a more in-depth investigation into this matter for future work.
Thank you for your answering my questions. I would still love to see a comparison of Graph-EFM with perturbed Graphcast. Regardless of the formulation of the problem, IC perturbations are always a common technique to achieve uncertainty estimation, and have been used in numerical methods, hybrid ML methods like NeuralGCM, and deep learning. Comparing with such a simple baseline will help shed light on the performance gain of Graph-EFM.