SEEDS:  Emulation of Weather Forecast Ensembles with Diffusion Models

Lizao Li

Google Research

Robert CarverContributed equally

Google Research

Ignacio Lopez-Gomez^*
Google Research

Fei Shato whom correspondence should be addressed

Google Research

John Anderson

Google Research

Uncertainty quantification is crucial to decision-making. A prominent example is probabilistic forecasting in numerical weather prediction. The dominant approach to representing uncertainty in weather forecasting is to generate an ensemble of forecasts. This is done by running many physics-based simulations under different conditions, which is a computationally costly process. We propose to amortize the computational cost by emulating these forecasts with deep generative diffusion models learned from historical data.  The learned models are highly scalable with respect to high-performance computing accelerators and can sample hundreds to tens of thousands of realistic weather forecasts at low cost. When designed to emulate operational ensemble forecasts, the generated ones are similar to physics-based ensembles in important statistical properties and predictive skill.  When designed to  correct biases present in the operational forecasting system, the generated ensembles show improved probabilistic forecast metrics. They are more reliable and forecast probabilities of extreme weather events more accurately. While this work demonstrates the utility of the methodology by focusing on weather forecasting, the generative artificial intelligence methodology can be extended for uncertainty quantification in climate modeling, where we believe the generation of very large ensembles of climate projections will play an increasingly important role in climate risk assessment.

§ INTRODUCTION

Weather is inherently uncertain, making the assessment of forecast uncertainty a vital component of operational weather forecasting [Bauer et al., 2015, Zhu et al., 2002]. Given a numerical weather prediction (NWP) model, the standard way to quantify this uncertainty is to stochastically perturb the model's initial conditions and its representation of small-scale physical processes to create an ensemble of weather trajectories [Palmer, 2019]. These trajectories are then regarded as Monte Carlo samples of the underlying probability distribution of weather states.

Given the computational cost of generating each ensemble member, weather forecasting centers can only afford to generate 10 to 50 members for each forecast cycle [ECMWF, 2021, Leutbecher, 2019, Zhou et al., 2022].
This limitation is particularly problematic for users concerned with the likelihood of high-impact extreme or rare weather events, which typically requires much larger ensembles to assess [Palmer, 2002, Palmer, 2017, Fischer et al., 2023]. For instance, one would need a 10,000-member calibrated ensemble
to forecast events with $1\%$ probability of occurrence with a relative error less than $10\%$.
Large ensembles are even more necessary for forecasting compound extreme events [Bevacqua et al., 2023, Leutbecher, 2019].
Besides relying on increases in available computational power to generate larger ensembles in the future, it is imperative to explore more efficient approaches for generating ensemble forecasts.

In this context, recent advances in generative artificial intelligence (GAI) offer a potential path towards massive reductions in the cost of ensemble forecasting. GAI models extract statistical priors from datasets, and enable conditional and unconditional sampling from the learned probability distributions. Through this mechanism, GAI techniques reduce the cost of ensemble forecast generation: once learning is complete, the sampling process is far more computationally efficient than time-stepping a physics-based NWP model.

In this work, we propose a technique that is based on probabilistic diffusion models, which have recently revolutionized GAI use cases such as image and video generation [Dhariwal and Nichol, 2021, Ho et al., 2022, Rombach et al., 2022]. Our () can generate an arbitrarily large ensemble conditioned on as few as one or two forecasts from an operational NWP system. We compare the generated ensembles to ground-truth ensembles from the operational systems, and to ERA5 reanalysis [Hersbach et al., 2020]. The generated ensembles not only yield weather-like forecasts but also match or exceed physics-based ensembles in skill metrics such as the rank histogram, the root-mean-squared error (RMSE) and continuous ranked probability score (CRPS). In particular, the generated ensembles assign more accurate likelihoods to the tail of the distribution, such as $\pm 2\sigma$ and $\pm 3\sigma$ weather events. Most importantly, the computational cost of the model is negligible; it has a throughput of 256 ensemble members (at $2^\circ$ resolution) per 3 minutes
on Google Cloud TPUv3-32 instances and can easily scale to higher throughput by deploying more accelerators.
We apply our methodology to uncertainty quantification in weather forecasting due to the wealth of data available and the ability to validate models on reanalysis. Nevertheless, the same approach could be used to augment climate projection ensembles.

Previous work leveraging artificial intelligence to augment and post-process ensemble or deterministic forecasts has focused on improving the aggregate output statistics of the prediction system. Convolutional neural networks have been used to learn a global measure of forecast uncertainty given a single deterministic forecast, trained using as labels either the error of previous forecasts or the spread of an ensemble system [Scher and Messori, 2018]. This approach has been generalized to predict the ensemble spread at each location of the input deterministic forecast; over both small regions using fully connected networks [Sacco et al., 2022], or over the entire globe using conditional generative adversarial networks [Brecht and Bihlo, 2023] based on the pix2pix architecture [Isola et al., 2017]. Deep learning has also proved effective in calibrating limited-size ensembles. For instance, self-attentive transformers can be used to calibrate the ensemble spread [Finn, 2021]. More related to our work, deep learning models have been successfully used to correct the probabilistic forecasts of ensemble prediction systems such that their final skill exceeds that of pure physics-based ensembles with at least double the number of members [Grönquist et al., 2021].

Our work differs from all previous studies in that our probabilistic generative model outputs high-dimensional weather-like samples from the target forecast distribution, akin to generative precipitation downscaling models [Harris et al., 2022]. Thus, our approach offers added value beyond improved estimates of the ensemble mean and spread: the drawn samples can be used to characterize spatial patterns associated with weather extremes [Scher et al., 2021], or as input to targeted weather applications that depend on variable and spatial correlations [Palmer, 2002].

§ METHOD

We start by framing the learning tasks. We then outline the data and neural network learning algorithm we use.  Details, including background, data processing and preparation, and learning architectures and procedures, are presented in Supplementary Information <ref> and <ref>.

§.§ Setup

To address the computational challenge of generating large weather forecast ensembles, we consider two learning tasks: generative ensemble emulation and generative post-processing. In both tasks, we are given as inputs a few examples sampled from a probability distribution $p(\vv)$, where $\vv$ stands for the atmospheric state variables. In our case, these examples represent physics-based weather forecasts. We seek to generate additional samples that either approximate the same distribution, or a related desired distribution. The central theme of statistical modeling for both tasks is to construct a computationally fast and scalable sampler for the target distributions.

Generative ensemble emulation leverages $K$ input samples to conditionally generate $N >  K$ samples such that they approximate the original distribution $p(\vv)$ from which the input samples are drawn. Its main purpose is to augment the ensemble size inexpensively without the need to compute and issue more than $K$ physics-based forecasts.

In generative post-processing, the sampler generates $N>K$ samples such that they approximate a mixture distribution where $p(\vv)$ is just one of the components. We consider the case where the target distribution is $\alpha p(\vv) + (1-\alpha) p'(\vv)$, with $\alpha \in [0, 1)$ being the mixture weight and $p'(\vv)$ a different distribution. The generative post-processing task aims not only to augment the ensemble size, but also to bias the new samples towards $p'(\vv)$, which we take to be a distribution that more closely resembles actual weather. The underlying goal is to generate ensembles that are less biased than those provided by the physics-based model, while still quantifying the forecast uncertainty captured by $p(\vv)$. We emphasize that while this task has the flavor and also achieves the effect of debiasing to some degree,  we focus on generating samples instead of minimizing the difference between their mean and a given reanalysis or observations. In both the emulation and post-processing tasks, the smaller the value of $K$ is, the greater the computational savings.

Illustration of the target distributions of generative ensemble emulation () and post-processing (Mixture). Shown are the histograms (bars: frequencies with 12 shared bins, curves: Gaussian kernel density estimators fit to the bars), , the empirical distributions of the surface temperature near Mountain View, CA on 2021/07/04 in the GEFS and ERA5 ensembles.  The goal common to both tasks is to generate additional ensemble members to capture the statistics of the desired distribution conditioned on a few GEFS samples. Note the small “bump” at the temperature of 287K in the mixture distribution.

Figure <ref> illustrates the concepts behind these two tasks. There, $p(\vv)$ is the distribution of the  surface temperature near Mountain View, CA on 2021/07/04 as predicted by the GEFS 13-day forecast ensemble [Zhou et al., 2022], and $p'(\vv)$ the corresponding ERA5 reanalysis ensemble [Hersbach et al., 2020]. While the GEFS ensemble has 31 members, our goal is to use $K\ll31$ GEFS ensemble members to steer our samplers to generate additional forecast members that are consistent with either GEFS's statistics or the mixture distribution's statistics. Inspired by terminology from natural language understanding and computer vision, we refer to those $K$ input examples from $p(\vv)$ as “seeds”. The desirability to have a small $K$ is in spirit similar to few-shot learning setups in those works.

We stress that the primary goal of both tasks is to improve the computational efficiency of ensemble weather forecasting, not to replace physics-based models. The generated samples should be not only consistent with the underlying distribution of atmospheric states (each sample is “weather-like”), but also validated by standard forecast verification metrics. In this work, we examine the generated ensembles by comparing them to other physics-based ensembles using the rank histogram, the anomaly correlation coefficient (ACC), RMSE, CRPS, and rare event classification metrics, as defined in <ref>.

§.§ Data for Learning and Evaluation

We target the GEFS (version 12) ensemble forecasting system for the generative ensemble emulation task [Zhou et al., 2022]. We use 20 years of GEFS 5-member reforecasts [Guan et al., 2022], denoted hereafter as , to learn $p(\vv)$. Generative post-processing attempts to remove systematic biases of the original forecasting system from the learned emulator. To this end, we take the ERA5 10-member Reanalysis Ensemble [Hersbach et al., 2020], denoted  as , to represent $p'(\vv)$ in the target mixture distribution. We also use ERA5 HRES reanalysis as a proxy for real historical observations when evaluating the skill of our generated ensemble predictions.

All data are derived from the publicly available sources listed in Table <ref>. Table <ref> lists the atmospheric state variables that are considered by our models. They are extracted and spatially regridded to the same cubed sphere mesh of size $6\times 48\times 48$ ($\approx2^\circ$ resolution) using inverse distance weighting with 4 neighbors [Ronchi et al., 1996]. We only retain the 00h-UTC time snapshots of the fields in Table <ref> for each day.

The climatology is computed from the ERA5 HRES dataset, using the reference period 1990-2020. The daily climatological mean and standard deviation are obtained by smoothing these two time series with a 15-day centered window over the year with periodic boundary conditions. The mean and standard deviation for February 29th is the average of those for February 28th and March 1st.

Our models take as inputs and produce as outputs the standardized climatological anomalies of variables in Table <ref>, defined as the standardized anomalies using the aforementioned climatological mean and standard deviation for the day of year and location, which facilitates learning [Dabernig et al., 2017, Lopez-Gomez et al., 2023, Qian et al., 2021].  The  outputs are converted back to raw values for evaluation.

For each unique pair of forecast lead time and number of seeds $K$, we train a diffusion model for the generative ensemble emulation task. For each unique triplet of lead time, $K$ and mixture weight $\alpha$, we train a model for the generative post-processing task. We provide results for lead times of $\{1, 4, 7, 10, 13, 16\}$ days, $K=2$ seeds, and generated ensembles with $N=512$ members. For the post-processing task, we consider the mixing ratio $\alpha=0.5$. The sensitivity to $K$, $N$, and $\alpha$ is explored in <ref>.

We evaluate our models against the operational GEFS 31-member ensemble [Zhou et al., 2022] () and the ERA5 HRES reanalysis. Note that we can do so because the and datasets are considered to have similar distributions — the reforecasts are reruns of the operational GEFS model using historical initial conditions [Guan et al., 2022]. We use the 20 years from 2000 to 2019 for training, year 2020 and 2021 for validation, and year 2022 for evaluation. In particular, to accommodate the longest lead time of 16 days, we evaluate using the forecasts initialized from 2022/01/01 to 2022/12/15 (349 days in total) and the ERA5 HRES data aligned with the corresponding days.

Data Used for Training and Evaluation

Name     Date Range     Ensemble size     Citation

ERA5-HRES     1959/01/01 – 2022/12/31     1     [Hersbach et al., 2020]

ERA5-Ensemble     1959/01/01 – 2021/12/31     10     [Hersbach et al., 2020]

GEFS     2020/09/23 – 2022/12/31     31     [Zhou et al., 2022]

GEFS-Reforecast     2000/01/01 - 2019/12/31     5     [Guan et al., 2022]

List of Atmospheric State Variables That Are Modeled

Quantity     Processed Units

Mean sea level pressure     $Pa$

Temperature at 2 meters     $K$

Eastward wind speed at 850hPa     $m/s$

Northward wind speed at 850hPa     $m/s$

Geopotential at 500hPa     $m^2/s^2$

Temperature at 850hPa     $K$

Total column water vapour     $kg/m^2$

Specific humidity at 500 hPa     $kg/kg$

§.§ Learning Method and Architecture

The use of probabilistic diffusion models to parameterize the target distributions, conditioned on a few “seeds”, is at the core of our statistical modeling algorithm for both tasks.

Probabilistic diffusion models are generative models of data.  The generative process follows a Markov chain. It starts with a random draw from an initial noise distribution – often an isotropic multivariate Gaussian. Then it iteratively transforms and denoises the sample until it resembles a random draw from the data distribution [Ho et al., 2020]. The iteration steps advance the diffusion time, which is independent from the real-world time. The denoising operation relies on the instantiation of a diffusion-time-dependent score function, which is the Jacobian of the log-likelihood of the data at a given diffusion time [Song et al., 2020]. Score functions often take the form of deep learning architectures whose parameters are learned from training data.

Typically, the score is a function of the noisy sample and the diffusion time. In this case, the resulting data distribution is a model of the unconditional distribution of the training data.  When additional inputs are passed to the score function, such as $K$ seeding forecasts in our setting, the sampler constructs the  distribution conditioned on these inputs.

In this work, our choice of the score function is inspired by the Vision Transformer (ViT), which has been successfully applied to a range of computer vision tasks [Dosovitskiy et al., 2021]. It is intuitive to view atmospheric data as a temporal sequence of snapshots, which are in turn viewed as “images”. Each snapshot is formed by “pixels” covering the globe with “color” channels. In this case, the channels correspond to the collection of atmospheric variables at different vertical levels. These can easily exceed in number the very few color channels of a typical image, e.g. 3 in the case of an RGB image. Due to this, we use a variant of ViT via axial attention [Ho et al., 2019], so that the model remains moderate in size and can be trained efficiently.

Irrespective of the lead times and the number of seeds, all the models share the same architecture and  have about 114M trainable parameters.  They are trained with a batch size of 128 for 200K steps.  The training of each model takes slightly less than 18 hours on a $2\times 2 \times 4$ TPUv4 cluster. Inference (namely, ensemble generation) runs at batch size 512 on a $4\times 8$ TPUv3 cluster at less than 3 minutes per batch. It is thus very efficient and easily scalable to generate thousands of members.

§ RESULTS

Using the methodology, we have developed two generative models. The model learns to emulate the distribution of the U.S. operational ensemble NWP system, the Global Ensemble Forecast System (GEFS) Version 12 [Zhou et al., 2022]. The model learns to emulate a blended distribution that combines the GEFS ensemble with historical data from the ERA5 reanalysis of the European Centre for Medium-Range Weather Forecasts (ECMWF), aiming to correct underlying biases in the operational GEFS system (, post-processing).

is trained using 20 years of GEFS 5-member retrospective forecasts [Guan et al., 2022], and additionally learns from ECMWF's ERA5 10-member Reanalysis Ensemble over the same period [Hersbach et al., 2020]. Once learned, both models take as inputs a few randomly selected member forecasts from the operational GEFS ensemble, which has 31 members. We refer to the selected members as the seeding forecasts. These seeds provide the physical basis used by the generative models to conditionally sample additional plausible weather states. Both and can be used to generate ensembles with a significantly larger number of forecasts than operational physics-based systems, easily reaching hundreds to tens of thousands of members.

Figure <ref> compares samples from the GEFS operational system, the ERA5 reanalysis, and the generative emulator . We also assess the quality of the generated ensembles in terms of multiple important characteristics of useful ensemble prediction systems. First, we analyze whether the forecasts in the generative ensembles display spatial coherence, multivariate correlation structures, and wavenumber spectra consistent with actual weather states. Second, we compare the pointwise predictive skill of the generative ensembles and the full operational physics-based GEFS ensemble, measured against the ERA5 high resolution (HRES) reanalysis [Hersbach et al., 2020].

We report results on a subset of field variables: the mean sea level pressure, the temperature $2~\mathrm{m}$ above the surface, and the zonal wind speed at pressure level $850~\mathrm{hPa}$. Results for all modeled fields, listed in Table <ref>, are presented in the Supplementary Information (SI). We use to refer to the full 31-member GEFS ensemble, and to an ensemble made of $2$ randomly selected seeding forecasts. Unless noted, our generated ensembles have 512 members.

Maps of total column vertically-integrated water vapor ($kg/m^2$) for 2022/07/14, as captured by (top left) the ERA5 reanalysis, (top right and middle row) 5 members of the forecast issued with a 7-day lead time, and (bottom) 3 samples from . The top 2 GEFS forecasts were used to seed the sampler.

§.§ Generated Weather Forecasts Are Plausible Weather Maps

Ensemble forecasting systems are most useful when individual weather forecasts resemble real weather maps [Lorenz, 1982]. This is because for many applications, such as ship routing, energy forecasting, or compound extreme event forecasting, capturing cross-field and spatial correlations is fundamental [Palmer, 2002, Scher et al., 2021, Worsnop et al., 2018].

Visualization of spatial coherence in forecasted weather charts for 2022/07/14,  with a 7-day lead time. The contours are for mean sea level pressure (dashed lines mark isobars below 1010 hPa) while the heatmap depicts the geopotential height at the 500 hPa pressure level. Row 1: ERA5 reanalysis, then  2 forecast members from used as seeds to our model.  Row 2–3: Other  forecast members from .    Row 4–5: 8 samples drawn from . Row 6:   Samples from a pointwise Gaussian model parameterized by the ensemble mean and variance.

To investigate this aspect of weather forecasts, we compare the covariance structure of the generated samples to those from the ERA5 Reanalysis and GEFS through a stamp map over Europe for a date during the 2022 European heatwave in Figure <ref> [Witze, 2022]. The global atmospheric context of a few of these samples is shown in Figure <ref> for reference. We also present in Figure <ref> weather samples obtained from a Gaussian model that predicts the univariate mean and standard deviation of each atmospheric field at each location, such as the data-driven model proposed in [Brecht and Bihlo, 2023]. This Gaussian model is meant to characterize the output of pointwise post-processing [Scher and Messori, 2018, Sacco et al., 2022, Brecht and Bihlo, 2023], which ignore correlations and treat each grid point as an independent random variable.

captures well both the spatial covariance and the correlation between midtropospheric geopotential and mean sea level pressure, since it directly models the joint distribution of the atmospheric state. The generative samples display a geopotential trough west of Portugal with spatial structure similar to that found in samples from or the reanalysis. They also depict realistic correlations between geopotential and sea level pressure anomalies. Although the Gaussian model predicts the marginal univariate distributions adequately, it fails to capture cross-field or spatial correlations. This hinders the assessment of the effects that these anomalies may have on hot air intrusions from North Africa, which can exacerbate heatwaves over Europe [Sánchez-Benítez et al., 2018].

The energy spectra of several global atmospheric variables for January of 2022 from the ERA5 reanalysis (thick black), members of the 7-day forecast (orange), and samples from (green).
The forecasts for each day are re-gridded to a latitude-longitude rectangular grid of the same angular resolution prior to computing the spectra. The computed spectra are averaged over the entire month. Each ensemble member is plotted separately.

Figure <ref> contrasts the energy spectra of forecasts with that of ERA5 and . The large overlap between samples from both forecast systems and the reanalysis demonstrates that the two ensembles have similar spatial structure. Small systematic differences can be observed in some variables like the zonal wind in the low troposphere, but for most variables the differences between and are similar to the differences between the operational system and the ERA5 reanalysis.

Generated ensembles provide better statistical coverage of the extreme heat event over Portugal. Each plot displays 16,384 generated forecasts from our method, extrapolating from the two seeding forecasts randomly taken from the operational forecasts. Contour curves of iso-probability are also shown. The first row is from and the second from .  characterizes the event best. Most notably, in the two rightmost plots of the bottom row, is able to generate well-dispersed forecast envelopes that cover the extreme event, despite the two seeding ones deviating substantially from the observed event.

In addition to examining the coherence of regional structures and the global spectra of the generative samples, we also examine the multivariate correlation structure of generative samples locally.
Figure <ref> depicts the joint distributions of temperature at 2 meters and total column water vapour at the grid point near Lisbon during the extreme heat event on 2022/07/14. We used the 7-day forecasts made on 2022/07/07. For each plot, we generate  16,384-member ensembles. The observed weather event from ERA5 is denoted by the star. The operational ensemble, denoted by the squares (also used as the seeding forecasts) and triangles (the rest of the GEFS ensemble), fails to predict the intensity of the extreme temperature event. This highlights that the observed event was so unlikely 7 days prior that none of the 31 forecast members from attained near-surface temperatures as warm as those observed. In contrast, the generated ensembles are able to extrapolate from the two seeding forecasts, providing an envelope of possible weather states with much better coverage of the event. This allows quantifying the probability of the event taking place (see Figure <ref> and <ref>). Specifically, our highly scalable generative approach enables the creation of very large ensembles that can capture the likelihood of very rare events that would be characterized with a null probability by limited-size ensembles. Moreover, we observe that the distributions of the generated ensembles do not depend critically on the (positioning of the) seeding forecasts. This suggests that the generative approach is plausibly learning the intrinsic dynamical structure, i.e., the attractor of the atmosphere, in order to expand the envelopes of the phase of the dynamical systems to include extreme events that deviate strongly from the seeds.

§.§ Forecast Reliability and Predictive Skills

An important characteristic of ensemble forecast systems is their ability to adequately capture the full distribution of plausible weather states. This characteristic is known as forecast calibration or reliability [Wilks, 2019]. Forecast reliability can be characterized for a given lead time in terms of the rank histogram [Anderson, 1996, Talagrand et al., 1997]. Deviations from flatness of this histogram indicate systematic differences between the ensemble forecast distribution and the true weather distribution. Rank histograms for 7-day forecasts from , , , and over California and Nevada are shown in Figure <ref>. The GEFS ensembles display systematic negative biases in mean sea level pressure and near-surface temperature over the region, as well as an underestimation of near-surface temperature uncertainty. Our model ensembles are more reliable than and , due in part to the larger number of members that can be effortlessly generated. shows the highest reliability of all, validating generative post-processing as a useful debiasing methodology. In particular, Figure <ref> shows how substantially reduces the ensemble under-dispersion for 2-meter temperature forecasts.

The reliability information contained in the rank histogram can be further summarized in terms of its bulk deviation from flatness, which we measure using the unreliability metric $\delta$ introduced by Candille and Talagrand [Candille and Talagrand, 2005]. Higher values of $\delta$ indicate higher deviations from flatness, or a lower reliability of the forecasts. Figure <ref> confirms that the generated ensembles are on a global average more reliable than for all lead times. In addition, the refined calibration of is more noticeable in the first forecast week.

Top: Rank histograms from 7-day forecasts for grid points in the region bounded by parallels 34N and 42N, and meridians 124W and 114W, for the year 2022. This region roughly encompasses California and Nevada, USA. To compare the histograms of ensembles of different size, the $x$ axis is normalized to quantiles instead of ranks, and the $y$ axis shows the difference to the uniform distribution. A perfectly calibrated ensemble forecast should have a flat curve at $0$. Bottom: Unreliability parameter $\delta$ [Candille and Talagrand, 2005] as a function of lead time, computed for the same year and averaged globally.

The predictive skill of the generated ensembles is measured in terms of the root-mean-squared-error () and the anomaly correlation coefficient () of the ensemble mean, as well as the continuous ranked probability score (), treating the ERA5 HRES reanalsyis as the reference ground-truth.
These metrics are computed and averaged over the grid points every forecast day in the test set and then aggregate over the test days.  <ref> details how these metrics are defined.

Figure <ref> reports these metrics for 3 atmospheric fields: the mean sea level pressure, the temperature 2 meters above the ground, and the eastward wind speed at 850hPa. Both or perform significantly better than the seeding ensemble across all metrics. The emulator shows similar but slightly lower skill than across all metrics and variables. Our generative post-processing is noticeably better than the physics-based at predicting near-surface temperature, roughly matching its skill for the other two fields. Intuitively, the potential benefits of statistical blending with a corrective data source are determined by the variable-dependent biases of the emulated forecast model. In this case, the GEFS model is known to have a cold bias near the surface [Zhou et al., 2022].

Metrics of point-wise skill (, and ) of the generative and physics-based ensemble forecasts, measured against the ERA5 HRES reanalysis as ground-truth. Shown are results for mean sea level pressure (left), $2$-meter temperature (center), and zonal velocity at $850$ hPa (right). A detailed description of these metrics is included in  <ref>.

Binary classification skill of the different ensembles regarding extreme events ($\pm 2\sigma$ from climatology) in mean slea level pressure, $2$-m temperature, and zonal velocity at $850$ hPa, using ERA5 HRES as the ground-truth. Skill is measured in terms of the cross-entropy; lower values are indicative of higher skill. First row: Brier score for $+2\sigma$. Second row: Brier score for $-2\sigma$.

A particularly challenging but important task of ensemble forecasts is being able to forecast extreme events and assign meaningful likelihoods to them [Palmer, 2002].   Figure <ref> compares the skill of the same 4 ensembles in predicting events deviating at least $\pm 2\sigma$ from the mean climatology. We measure binary classification skill by computing the Brier score of occurrence using ERA5 HRES as the binary reference, and assigning a probability of occurrence to the ensemble forecasts equal to the fraction of occurrences within the ensemble.

We observe that is comparable in skill to the full ensemble and far exceeds the skill of the seeding forecast ensemble .  In the forecast of 2-meter temperature, performs noticeably better than the other ensembles. For other variables, despite the less apparent advantage, remains the best extreme forecast system for most lead times and variables. This highlights the relevance of our generative approach for forecasting tasks focused on extremes.

§.§ Hallucination or In-filling?

One frequently cited issue of generative AI technology is its tendency to “hallucinate information”. We conclude this section by exploring the nature of the distribution information that the generative ensembles are able to represent, beyond what is present in the two seeding forecasts from the GEFS full ensemble. As shown previously, the generated ensembles outperform the seeding forecast ensembles in all metrics and often match or improve over the physics-based full ensemble.

Figure <ref> measures the correlation of the generative ensembles (and ), the seeding ensemble , and the GEFS model climatology, with respect to the ensemble forecasts.  While comparing full joint distributions remains infeasible, we compute how well the spread of each ensemble forecast correlates with that of the full physics-based ensemble .  The plots show that at long-lead times ($\ge 10$ days),  all ensembles but converge to high correlations ($\ge 95\%$) with .  This is also true for the model climatology.  However, in the medium range (more than 4 days but less than 10 days ahead), the generative ensembles display a higher correlation with the than both the model climatology and . This suggests that the generative models are indeed able to generate information about forecast uncertainty beyond the two seeding forecasts. In addition, the fact that generative ensembles can capture a higher correlation with than the model climatology in the short and medium range shows that the diffusion models are learning to emulate dynamically-relevant features beyond model biases; they have resolution beyond climatology.

Thus, we put forward a reasonable hypothesis that the generated ensembles in-fill probability density gaps in the small seeding ensembles.  They also extend the (tails of the) envelopes of the full ensembles such that extreme events are well represented in the envelops.

Comparing the ensembles and the model climatology to in terms of how the ensemble spreads are correlated with those from .  The plots show that in medium-range between 4 to 10 days, the model has leveraged the  two seeding forecasts to generate different, yet informative, ensemble members to represent uncertainty.

§ RELATED WORK

Previous work leveraging artificial intelligence to augment and post-process ensemble forecasts has focused on improving the aggregate output statistics of the prediction system. [Scher and Messori, 2018] trained a convolutional neural network to quantify forecast uncertainty given a single deterministic forecast [Scher and Messori, 2018]. They learned a global measure of uncertainty in a supervised setting, using as labels either the error of previous forecasts or the spread of an ensemble system. [Brecht and Bihlo, 2023] generalized this approach by predicting the ensemble spread at each forecast location, given a deterministic forecast [Brecht and Bihlo, 2023]. For this task, they used a conditional generative adversarial network based on the pix2pix architecture [Isola et al., 2017]. [Grönquist et al., 2021] trained a deep learning system to post-process a 5-member ensemble forecast, resulting in a lower CRPS than a 10-member ensemble from the same operational system [Grönquist et al., 2021]. [Sacco et al., 2022] extended this work to build a system capable of predicting the ensemble mean and spread over a limited domain [Sacco et al., 2022].

Our work differs from that of [Brecht and Bihlo, 2023], [Grönquist et al., 2021], and [Sacco et al., 2022] in that our probabilistic generative model outputs actual samples from the target forecast distribution. Thus, our approach offers added value beyond the ensemble mean and spread: the drawn samples can be used to characterize spatial patterns associated with weather extremes [Scher et al., 2021], or as input to targeted weather applications that depend on variable and spatial correlations [Palmer, 2002].

§ DISCUSSION

The () proposed in this work leverages the power of generative artificial intelligence to produce ensemble forecasts comparable to those from the operational GEFS system at accelerated pace – the results reported in this paper need only 2 seeding forecasts from the operational system, which generates 31 forecasts in its current version [Zhou et al., 2022]. This leads to a hybrid forecasting system where a few weather trajectories computed with a physics-based model are used to seed a diffusion model that can generate additional forecasts much more efficiently. This methodology provides an alternative to the current operational weather forecasting paradigm, where the computational resources saved by the statistical emulator could be allocated to increasing the resolution of the physics-based model [Ma et al., 2012], or issuing forecasts more frequently.

is trained on historical retrospective forecasts (, reforecasts) issued with the operational physics-based model, which are already required for post-processing in the current paradigm [Hagedorn et al., 2008]. Our framework is also flexible enough to enable direct generation of debiased ensembles when the generative post-processing task is considered during training; the only additional requirement is access to historical reanalysis for the reforecast period.

For future work, we will conduct case studies of high-impact weather events to further evaluate ' performance, and consider specific ensemble forecast applications such as tropical and extratropical cyclone tracking [Froude et al., 2007, Lin et al., 2020]. We will also explore more deeply the statistical modeling mechanisms that such models employ to extract information from weather data and in-fill the ensemble forecast distribution. It is our belief that our application of generative AI to weather forecast emulation represents just one way of many that will accelerate progress in operational NWP in coming years. Additionally, we hope the established utility of generative AI technology for weather forecast emulation and post-processing will spur its application in research areas such as climate risk assessment, where generating a large number of ensembles of climate projections is crucial to accurately quantifying the uncertainty about future climate [Deser et al., 2020].

Our colleagues at Google Research have provided invaluable advice.  Among them, we thank Stephan Rasp, Stephan Hoyer, and Tapio Schneider for their inputs and useful discussion on the manuscript. We thank Carla Bromberg and Tyler Russell for technical program management, as well as Alex Merose for data coordination and support. We also thank Cenk Gazen, Shreya Agrawal and Jason Hickey for discussions with them in the early stage of this work.

[Anderson, 1996]
Jeffrey L Anderson.
A method for producing and evaluating probabilistic forecasts from
ensemble model integrations.
J. Clim., 9:0 1518–1530, 1996.
URL <http://www.jstor.org/stable/26201352>.

[Bauer et al., 2015]
Peter Bauer, Alan Thorpe, and Gilbert Brunet.
The quiet revolution of numerical weather prediction.
Nature, 525:0 47–55, 2015.
URL <https://doi.org/10.1038/nature14956>.

[Benedetti, 2010]
Riccardo Benedetti.
Scoring rules for forecast verification.
Mon. Weather Rev., 1380 (1):0 203–211,
January 2010.

[Bevacqua et al., 2023]
Emanuele Bevacqua, Laura Suarez-Gutierrez, Aglaé Jézéquel, Flavio Lehner,
Mathieu Vrac, Pascal Yiou, and Jakob Zscheischler.
Advancing research on compound weather and climate events via large
ensemble model simulations.
Nat. Commun., 14:0 2145, 2023.
URL <https://doi.org/10.1038/s41467-023-37847-5>.

[Brecht and Bihlo, 2023]
Rüdiger Brecht and Alex Bihlo.
Computing the ensemble spread from deterministic weather predictions
using conditional generative adversarial networks.
Geophys. Res. Lett., 50:0 e2022GL101452, 1 2023.
URL <https://doi.org/10.1029/2022GL101452>.

[Brier, 1950]
Glenn W Brier.
Verification of forecasts expressed in terms of probability.
Mon. Weather Rev., 78:0 1–3, 1950.

[Candille and Talagrand, 2005]
G Candille and O Talagrand.
Evaluation of probabilistic prediction systems for a scalar variable.
Q. J. R. Meteorol. Soc., 1310 (609):0
2131–2150, July 2005.

[Dabernig et al., 2017]
Markus Dabernig, Georg J Mayr, Jakob W Messner, and Achim Zeileis.
Spatial ensemble post-processing with standardized anomalies.
Q. J. R. Meteorol. Soc., 143:0 909–916, 1 2017.
URL <https://doi.org/10.1002/qj.2975>.

[Deser et al., 2020]
C Deser, F Lehner, K B Rodgers, T Ault, T L Delworth, P N DiNezio, A Fiore,
C Frankignoul, J C Fyfe, D E Horton, J E Kay, R Knutti, N S Lovenduski,
J Marotzke, K A McKinnon, S Minobe, J Randerson, J A Screen, I R Simpson, and
M Ting.
Insights from earth system model initial-condition large ensembles
and future prospects.
Nature Climate Change, 10:0 277–286, 2020.
ISSN 1758-6798.
URL <https://doi.org/10.1038/s41558-020-0731-2>.

[Dhariwal and Nichol, 2021]
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat GANs on image synthesis.
In Advances in Neural Information Processing Systems,
volume 34, pages 8780–8794, 2021.

[Dosovitskiy et al., 2021]
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn,
Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg
Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.
An image is worth 16x16 words: Transformers for image recognition at
In International Conference on Learning Representations, 2021.
URL <https://openreview.net/forum?id=YicbFdNTTy>.

[ECMWF, 2021]
IFS Documentation CY47R3 - Part V Ensemble prediction system.
Number 5. ECMWF, 09/2021 2021.
URL <https://www.ecmwf.int/node/20199>.

[Finn, 2021]
Tobias Sebastian Finn.
Self-Attentive ensemble transformer: Representing ensemble
interactions in neural networks for earth system models.
June 2021.

[Fischer et al., 2023]
E M Fischer, U Beyerle, L Bloin-Wibe, C Gessner, V Humphrey, F Lehner, A G
Pendergrass, S Sippel, J Zeder, and R Knutti.
Storylines for unprecedented heatwaves based on ensemble boosting.
Nature Communications, 14:0 4643, 2023.
ISSN 2041-1723.
URL <https://doi.org/10.1038/s41467-023-40112-4>.

[Froude et al., 2007]
Lizzie S R Froude, Lennart Bengtsson, and Kevin I Hodges.
The prediction of extratropical storm tracks by the ECMWF and
NCEP ensemble prediction systems.
Mon. Weather Rev., 1350 (7):0 2545–2567, July

[Gneiting and Raftery, 2007]
Tilmann Gneiting and Adrian E Raftery.
Strictly proper scoring rules, prediction, and estimation.
J. Am. Stat. Assoc., 1020 (477):0 359–378,
March 2007.

[Grönquist et al., 2021]
Peter Grönquist, Chengyuan Yao, Tal Ben-Nun, Nikoli Dryden, Peter Dueben,
Shigang Li, and Torsten Hoefler.
Deep learning for post-processing ensemble weather forecasts.
Philos. Trans. R. Soc. A, 3790 (2194):0
20200092, April 2021.

[Guan et al., 2022]
Hong Guan, Yuejian Zhu, Eric Sinsky, Bing Fu, Wei Li, Xiaqiong Zhou, Xianwu
Xue, Dingchen Hou, Jiayi Peng, M M Nageswararao, Vijay Tallapragada, Thomas M
Hamill, Jeffrey S Whitaker, Gary Bates, Philip Pegion, Sherrie Frederick,
Matthew Rosencrans, and Arun Kumar.
GEFSv12 reforecast dataset for supporting subseasonal and
hydrometeorological applications.
Mon. Weather Rev., 1500 (3):0 647–665, March

[Hagedorn et al., 2008]
Renate Hagedorn, Thomas M Hamill, and Jeffrey S Whitaker.
Probabilistic forecast calibration using ECMWF and GFS ensemble
reforecasts. part I: Two-meter temperatures.
Mon. Weather Rev., 136:0 2608–2619, 2008.

[Harris et al., 2022]
Lucy Harris, Andrew T T McRae, Matthew Chantry, Peter D Dueben, and Tim N
A generative deep learning approach to stochastic downscaling of
precipitation forecasts.
J. Adv. Model. Earth Syst., 140 (10):0
e2022MS003120, April 2022.

[Hersbach et al., 2020]
Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András
Horányi, Joaquín Muñoz-Sabater, Julien Nicolas, Carole Peubey,
Raluca Radu, Dinand Schepers, Adrian Simmons, Cornel Soci, Saleh Abdalla,
Xavier Abellan, Gianpaolo Balsamo, Peter Bechtold, Gionata Biavati, Jean
Bidlot, Massimo Bonavita, Giovanna De Chiara, Per Dahlgren, Dick Dee, Michail
Diamantakis, Rossana Dragani, Johannes Flemming, Richard Forbes, Manuel
Fuentes, Alan Geer, Leo Haimberger, Sean Healy, Robin J Hogan, Elías
Hólm, Marta Janisková, Sarah Keeley, Patrick Laloyaux, Philippe
Lopez, Cristina Lupu, Gabor Radnoti, Patricia de Rosnay, Iryna Rozum, Freja
Vamborg, Sebastien Villaume, and Jean-Noël Thépaut.
The ERA5 global reanalysis.
Q. J. R. Meteorol. Soc., 1460 (730):0
1999–2049, 2020.

[Ho et al., 2019]
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans.
Axial attention in multidimensional transformers, 2019.

[Ho et al., 2020]
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin,
editors, Advances in Neural Information Processing Systems, volume 33,
pages 6840–6851. Curran Associates, Inc., 2020.

[Ho et al., 2022]
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey
Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet,
et al.
Imagen video: High definition video generation with diffusion models.
arXiv preprint arXiv:2210.02303, 2022.

[Isola et al., 2017]
P Isola, J Zhu, T Zhou, and A A Efros.
Image-to-image translation with conditional adversarial networks.
In 2017 IEEE Conference on Computer Vision and Pattern
Recognition (CVPR), pages 5967–5976, 2017.
URL <http://doi.ieeecomputersociety.org/10.1109/CVPR.2017.632>.

[Karras et al., 2022]
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.
Elucidating the design space of diffusion-based generative models.
In Proc. NeurIPS, 2022.

[Leutbecher, 2019]
Martin Leutbecher.
Ensemble size: How suboptimal is less than infinity?
Q. J. R. Meteorol. Soc., 145:0 107–128, 9 2019.
URL <https://doi.org/10.1002/qj.3387>.

[Lin et al., 2020]
Jonathan Lin, Kerry Emanuel, and Jonathan L Vigh.
Forecasts of hurricanes using Large-Ensemble outputs.
Weather Forecast., 350 (5):0 1713–1731,
October 2020.

[Lopez-Gomez et al., 2023]
Ignacio Lopez-Gomez, Amy McGovern, Shreya Agrawal, and Jason Hickey.
Global extreme heat forecasting using neural weather models.
Artif. Intell. Earth Syst., 2:0 e220035, 2023.

[Lorenz, 1982]
E N Lorenz.
Atmospheric predictability experiments with a large numerical model.
Tellus A, 340 (6):0 505–513, 1982.

[Ma et al., 2012]
Juhui Ma, Yuejian Zhu, Richard Wobus, and Panxing Wang.
An effective configuration of ensemble size and horizontal resolution
for the NCEP GEFS.
Adv. Atmos. Sci., 29:0 782–794, 2012.
URL <https://doi.org/10.1007/s00376-012-1249-y>.

[Palmer, 2002]
Tim Palmer.
The economic value of ensemble forecasts as a tool for risk
assessment: From days to decades.
Q. J. R. Meteorol. Soc., 128:0 747–774, 4 2002.
URL <https://doi.org/10.1256/0035900021643593>.

[Palmer, 2017]
Tim Palmer.
The primacy of doubt: Evolution of numerical weather prediction from
determinism to probability.
J. Adv. Model. Earth Syst., 9:0 730–734, 6 2017.
URL <https://doi.org/10.1002/2017MS000999>.

[Palmer, 2019]
Tim Palmer.
The ECMWF ensemble prediction system: Looking back (more than) 25
years and projecting forward 25 years.
Q. J. R. Meteorol. Soc., 145:0 12–24, 9 2019.
URL <https://doi.org/10.1002/qj.3383>.

[Qian et al., 2021]
Weihong Qian, Jun Du, and Yang Ai.
A review: Anomaly-based versus full-field-based weather analysis and
Bull. Am. Meteorol. Soc., 102:0 E849–E870, 2021.

[Rombach et al., 2022]
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn
High-resolution image synthesis with latent diffusion models.
In Proceedings of the IEEE Conference on Computer Vision and
Pattern Recognition (CVPR), 2022.

[Ronchi et al., 1996]
C Ronchi, R Iacono, and P S Paolucci.
The “cubed sphere”: A new method for the solution of partial
differential equations in spherical geometry.
J. Comput. Phys., 124:0 93–114, 1996.

[Sacco et al., 2022]
Maximiliano A Sacco, Juan J Ruiz, Manuel Pulido, and Pierre Tandeo.
Evaluation of machine learning techniques for forecast uncertainty
Q. J. R. Meteorol. Soc., 1480 (749):0
3470–3490, October 2022.

[Scher and Messori, 2018]
Sebastian Scher and Gabriele Messori.
Predicting weather forecast uncertainty with machine learning.
Q. J. R. Meteorol. Soc., 144:0 2830–2841, 10
URL <https://doi.org/10.1002/qj.3410>.

[Scher and Messori, 2018]
Sebastian Scher and Gabriele Messori.
Predicting weather forecast uncertainty with machine learning.
Q. J. R. Meteorol. Soc., 1440 (717):0
2830–2841, 2018b.

[Scher et al., 2021]
Sebastian Scher, Stephen Jewson, and Gabriele Messori.
Robust worst-case scenarios from ensemble forecasts.
Weather Forecast., 36:0 1357–1373, 2021.

[Slingo and Palmer, 2011]
Julia Slingo and Tim Palmer.
Uncertainty in weather and climate prediction.
Philos. Trans. R. Soc. A, 369:0 4751–4767, 12 2011.
URL <https://doi.org/10.1098/rsta.2011.0161>.
doi: 10.1098/rsta.2011.0161.

[Song et al., 2020]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano
Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential
In International Conference on Learning Representations, 2020.

[Sánchez-Benítez et al., 2018]
A Sánchez-Benítez, R García-Herrera, D Barriopedro, P M Sousa, and R M
June 2017: The earliest european summer mega-heatwave of reanalysis
Geophys. Res. Lett., 45:0 1955–1962, 2 2018.
URL <https://doi.org/10.1002/2018GL077253>.

[Talagrand et al., 1997]
O. Talagrand, R. Vautard, and B Strauss.
Evaluation of probabilistic prediction systems.
In Workshop on Predictability, 20-22 October 1997, pages
1–26, Shinfield Park, Reading, 1997. ECMWF.

[Tancik et al., 2020]
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin
Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng.
Fourier features let networks learn high frequency functions in low
dimensional domains.
In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin,
editors, Advances in Neural Information Processing Systems, volume 33,
pages 7537–7547. Curran Associates, Inc., 2020.

[Wilks, 2019]
Daniel S Wilks.
Statistical Methods in the Atmospheric Sciences.
Elsevier, 4th edition, 2019.

[Witze, 2022]
Alexandra Witze.
Extreme heatwaves: Surprising lessons from the record warmth.
Nature, 608:0 464–465, 2022.

[Worsnop et al., 2018]
R. P. Worsnop, M. Scheuerer, T. M. Hamill, and J. K. Lundquist.
Generating wind power scenarios for probabilistic ramp event
prediction using multivariate statistical post-processing.
Wind Energy Sci., 30 (1):0 371–393, 2018.
URL <https://wes.copernicus.org/articles/3/371/2018/>.

[Zhou et al., 2022]
Xiaqiong Zhou, Yuejian Zhu, Dingchen Hou, Bing Fu, Wei Li, Hong Guan, Eric
Sinsky, Walter Kolczynski, Xianwu Xue, Yan Luo, Jiayi Peng, Bo Yang, Vijay
Tallapragada, and Philip Pegion.
The development of the NCEP global ensemble forecast system version
Weather Forecast., 370 (6):0 1069–1084, June

[Zhu et al., 2002]
Yuejian Zhu, Zoltan Toth, Richard Wobus, David Richardson, and Kenneth Mylne.
The economic value of ensemble-based weather forecasts.
Bull. Am. Meteorol. Soc., 83:0 73–84, 2002.

§ PROBABILISTIC DIFFUSION MODELS

We give a brief introduction to probabilistic diffusion models. Diffusion models are powerful methods for learning distributions from data. They have recently become one of the most popular approaches for image generation and video synthesis [Ho et al., 2022]. For a detailed description, see [Song et al., 2020].

Consider a multivariate random variable $\vV$ with an underlying distribution $p_{\text{data}}(\vv)$.  Intuitively, diffusion-based generative models iteratively transform samples from an initial noise distribution $p_\mathcal{T}$ into samples from the target data distribution $p_{\text{data}}$ through a denoising operator. By convention, $\mathcal{T}$ is set to 1 as it is just nominal and does not correspond to the real time in the physical world[For notation clarity, we will be using $\tau$ and its capitalized version $\mathcal{T}$ to denote diffusion times, which are different from the physical times.]. Additionally, as it will be clear from the explanation below, this initial distribution is a multivariate Gaussian distribution.

Noise is removed  such that the samples follow a family of diffusion-time-dependent marginal distributions $p_\tau(\vv_\tau; \sigma_\tau)$ for decreasing diffusion times $\tau$ and noise levels $\sigma_\tau$ [Ho et al., 2020]. The distributions are given by a forward blurring process that is described by the following stochastic differential equation (SDE) [Karras et al., 2022, Song et al., 2020]

\begin{equation}
\label{eq:fwd_sde}
d\vV_\tau  = f(\vV_\tau, \tau) d\tau +g(\vV_\tau, \tau) dW_\tau,
\end{equation}

with drift $f$, diffusion coefficient $g$, and the standard Wiener process $W_\tau$. Following [Karras et al., 2022], we set

\begin{equation}
f(\vV_\tau, \tau) = f(\tau) \vV_\tau := \frac{\dot{s}_\tau}{s_\tau}\vV_\tau,  \qquad \text{and} \qquad
g(\vV_\tau, \tau) = g(\tau) := s_\tau\sqrt{2\dot{\sigma}_\tau\sigma_\tau},
\end{equation}

where the overhead dot denotes the time derivative.
Solving the SDE in (<ref>) forward in time with an initial condition $\vv_0$ leads to the Gaussian perturbation kernel $p_\tau(\vV_\tau|\vv_0) = \mathcal{N}( s_\tau \vv_0, s_\tau^2\sigma_\tau^2\textbf{I})$. Integrating the kernel over the data distribution $p_0(\vv_0) = p_{\text{data}}$, we obtain the marginal distribution $p_\tau(\vv_\tau)$ at any $\tau$. As such, one may prescribe the profiles of $s_\tau$ and $\sigma_\tau$ so that $p_0 = p_{\text{data}}$ (with $s_0 = 1, \sigma_0 = 0$),
and more importantly

\begin{equation}
\label{eq:dist_t}
p_\mathcal{T}(\vV_\mathcal{T}) \approx \mathcal{N}( 0, s_\mathcal{T}^2\sigma_\mathcal{T}^2\mathbf{I}),
\end{equation}

i.e., the distribution at the (nominal) terminal time $\mathcal{T}=1$ becomes indistinguishable from an isotropic, zero-mean Gaussian. To sample from $p_{\text{data}}$, we utilize the fact that the  reverse-time SDE

\begin{equation}
\label{eq:bwd_sde}
d\vV_\tau = \big[f(\tau) \vV_\tau - g(\tau)^2\nabla_{\vV_\tau} \log p_\tau(\vV_\tau)\big]d\tau +g(\tau) dW_\tau
\end{equation}

has the same marginals as eq. (<ref>). Thus, by solving (<ref>) backwards using (<ref>) as the initial condition, we obtain samples from $p_{\text{data}}$ at $\tau=0$.

There are several approaches to learn the score function $\nabla_{\vV_\tau} \log p_\tau(\vV_\tau)$. The key step is to parameterize it with a neural network $s_\theta(\vv_\tau,  \sigma_\tau)$ such that

\begin{equation}
s_\theta(\vv_\tau,  \sigma_\tau) \approx \nabla _{\vv_\tau}\log p_\tau(\vv_\tau| \vv_0) = -\frac{s_\tau \vv_0 - \vv_\tau}{\sigma^2_{\tau}} =z_\tau/\sigma_\tau
\label{eScoreNetwork}
\end{equation}

for any $\vv\sim p_\tau(\vV_\tau|\vv_0 )$ and $\sigma_\tau$. In other words, the normalized score network $\sigma_\tau s_\theta(\vv_\tau,  \sigma_\tau)$ identifies the noise $z_\tau$ injected into the data $\vv_0$. The parameters $\theta$ of the score network are trained to reduce the L2 difference between  them. Once this score network is learned, it is used to denoise a noised version of the desired data $\vv_0$.   We refer to this process as learning the normalized score function. In this work, we use the hyperparameters $f(\tau) = 0$ and $g(\tau)=100^\tau$.

§ METHOD DETAILS

§.§ Problem Setup and Main Idea

Here we provide a detailed description of the two generative tasks considered. Formally, let $p(\vv)$ denote an unknown distribution of the atmospheric state $\vV$. In the task of generative ensemble emulation, we are given a few examples or seeds sampled from $p(\vv)$ and our task is to generate more samples from the same distribution.

Given $K$ samples $\mathcal{E}^K=(\vv^1, \vv^2, \cdots, \vv^K)$ from $p(\vv)$, we construct an easy-to-sample conditional distribution

\begin{equation}
\hat{p}(\vv) = p( \vv ; \mathcal{E}^K),
\label{eCondition}
\end{equation}

which approximates $p(\vv)$.  The conditional distribution $\hat{p}(\vv)$ needs to have two desiderata: it approximates well $p(\vv)$, and it is much less costly to sample than $p(\vv)$.  Note that this problem extends the typical density estimation problem, where there is no conditioning, , $K=0$.

In the task of generative post-processing, we construct an efficient sampler to approximate a mixture distribution. As another example,  let $p'(\vv)$ be the reanalysis distribution of the atmospheric state corresponding to the forecast $p(\vv)$.  We construct a conditional distribution $p( \vv ; \mathcal{E}^K)$ to approximate the following mixture of distributions:

\begin{equation}
p^\textsc{mix}(\vv) =  \alpha p(\vv) + (1-\alpha)p'(\vv).
\label{eMix}
\end{equation}

In practice, the information about $p'(\vv)$ is given to the learning algorithm in the form of a representative $K'$-member ensemble $\mathcal{E}'^{K'}$. The mixture weight $\alpha$ is then given by

\begin{equation}
\alpha = \frac{M-K}{M-K+K'}.
\end{equation}

where $M>K $ is the number of samples we have drawn from $p(\vv)$.
In the following, we describe how to apply the technique of generative modeling introduced in <ref> to these two tasks.  We first mention how data is organized.

§.§ Data for Training and Evaluation

Each ensemble $\mathcal{E}_{tl}$ is identified by the forecast initialization day $t$ and the lead time $l$. Here, $t \in T^{\textsc{train}}$ indexes all the days in the training set and $l \in L$, where $L= \{1, 3, 6, 10, 13, 16\}$ days is the set of lead times we consider.  We group ensembles of  the same lead time $l$ as a training dataset $\mathcal{D}^\textsc{train}_l = \{\mathcal{E}_{1l}, \mathcal{E}_{2l}, \cdots, \mathcal{E}_{T^\textsc{train}l}\}$. In generative ensemble emulation, $\mathcal{D}^\textsc{train}_l$ contains samples from exclusively, since the target distribution is . In generative post-processing, $\mathcal{D}^\textsc{train}_l$ contains samples from both and  , to captue the target distribution eq. (<ref>).

Note that the $K$ seeds can be included in the generated ensembles for evaluation. However, in this work, we decide to adopt a stricter evaluation protocol where the seeds are excluded from the generated ensembles; including them would mildly improve the quality of the generated ensembles, since the number of generated forecasts is far greater than $K$.

As a shorthand, we use $\vv_t$ to denote all members in the ensemble forecast made at time $t$, and $\vv_{tm}$ to denote the $m$th forecast at the same time.

§.§.§ Generative Ensemble Emulation

In the ensemble emulation task we use $K$ samples to construct a sampler of the conditional distribution $\hat{p}(\vv)$. To this end, for each lead time $l$ and value of $K$, we learn a distinct conditional distribution model $\mathcal{M}_{lK}$. We choose $K=1,2,3,4$; and use $\mathcal{D}^\textsc{train}_l$ to learn $\mathcal{M}_{lK}$.

To train the conditional generation model $\mathcal{M}_{lK}$, we use as input data $K$ randomly chosen members from each ensemble forecast $\vv_t$ in $\mathcal{D}^\textsc{train}_l$, and another member $\vv_{tm_0}$ from the remaining $(M-K)$ members as our target. This procedure augments the data in $\mathcal{D}^\textsc{train}_l$ by $C_M^K$-fold, where $C_M^K$ is the number of $K$-combinations of the $M$ members.

The learning goal is then to optimize the score network eq. (<ref>), so that the resulting sampler gives rise to:

\begin{equation}
p_{lK}( \vv_{tm_0} |  \vv_{tm_1},   \vv_{tm_2}, \cdots,   \vv_{tm_K}),
\label{eConditionalDistribution}
\end{equation}

where $m_1, m_2, \cdots, m_K,$ and $m_0$ index the chosen ensemble members.  Specifically, the score network is given by

\begin{equation}\label{scoreEmulation}
s_\theta( \vv_{tm,\tau}, \vv_{tm_1}, \vv_{tm_2}, \cdots, \vv_{tm_K}, \mathbf{c}_t, \sigma_\tau)
\end{equation}

where $\vv_{tm,\tau}$ is a perturbed version of $\vv_{tm_0}$.  We include the climatological mean $\mathbf{c}_t$ as an input to the score function eq. (<ref>), since the atmospheric state is non-stationary with respect to the time of the year.

§.§.§ Generative Post-Processing

For the task of generative post-processing, the setup is similar except that we seek to approximate the mixture distribution of $p(\vv)$ and $p'(\vv)$ defined in eq. (<ref>). The training dataset $\mathcal{D}_l^{\textsc{train}}$ thus contains ensembles from two different distributions, $\vv_t$ and $\vv'_t$.

In this case, we train conditional distribution models $\mathcal{M}_{lKK'}$ for each triplet $(l, K, K')$ using the same input as the conditional emulator $\mathcal{M}_{lK}$, , $K$ out of $M$ members of each ensemble $\vv_t$. However, the target sample is now drawn randomly from either $p(\vv)$ or $p'(\vv)$. Consistent with eq. (<ref>), we draw a sample from the remaining $(M-K)$ members of ensemble $\vv_t$ with probability $(M-K)/(M-K+K')$, or a member from $\vv'_t$ with probability $K'/(M-K+K')$. Denoting the final selection as $\vx_{t0}$, the score function becomes

\begin{equation}\label{scorePost}
s_\theta( \vx_{t0,\tau}, \vv_{tm_1}, \vv_{tm_2}, \cdots, \vv_{tm_K}, \mathbf{c}_t, \sigma_\tau)
\end{equation}

where $\vx_{t0, \tau}$ is a perturbed version of $\vx_{t0}$.

In this task, $K'$ is a hyperparameter that controls the mixture distribution being approximated. We choose  $K'$ from 2, 3 or 4.

§.§ Model Architecture

We employ an axial variant of the vision Transformer (ViT) [Dosovitskiy et al., 2021] to model the normalized score functions $\sigma_\tau s_\theta(\cdot, \cdot)$ in eqns. (<ref>) and (<ref>), adapted to the particular characteristics of atmospheric modeling.

Each atmospheric state snapshot $\vv$ is represented as a tensor $v_{qp}$, where $q$ indexes the physical quantities and $p$ the locations on the cubed sphere. In particular, the altitude levels of the same physical quantity are also indexed by $q$. The overall model operates on sequences of atmospheric state snapshots $v_{sqp}$ where $s$ indexes the sequence position. In the context of computer vision, this is similar to a video model, where $p$ indexes the position in the image, $q$ the “RGB” channels, and $s$ the time.

The number of channels, which we use to model physical fields at different atmospheric levels, can be $q\sim\mathcal{O}(100)$ in our setting.
For this reason, we resort to axial attention, applied to $p$, $q$, and $s$ separately [Ho et al., 2019].
Each sequence dimension has its own learned position embedding instead of the usual fixed positional encodings.

§.§.§ Spatial Embedding

Schematic of the spatial attention mechanism. Patches of cubed sphere grid points are mapped to patch embeddings, and a transformer stack transforms their combination with patch position embeddings in the patch ($l$) dimension.

We call a piece of data assigned to each grid point of the cubed sphere as a slice.
Regarding spatial attention, the $6\times C\times C$ values of a slice associated with the cubed sphere grid points indexed by $p$ are partitioned into $6(C/P)^2$ square patches of size $P\times P$. We employ $C=48$ in our work, corresponding roughly to $2^\circ$ resolution. Each patch is then flattened and linearly embedded into $D$ dimensions with learned weights, leading to a tensor $h_{sqld}$, where $l=0,\ldots 6(C/P)^2-1$ indexes the patches and $d$ the in-patch dimension.
Due to the nontrivial neighborhood structure of the cubed sphere, a shared learned positional encoding $e^{\text{pos}}_{ld}$ is used for the $l$ dimension.
This is then fed into a transformer of $T_L$ layers operating with $l$ as the sequence dimension and $d$ as the embedding dimension. The output is the the  slice embedding corresponding to each patch:

\begin{equation}
h^{\text{slice}}_{sqld} = T_L(h_{sqld} + e^{\text{pos}}_{ld}).
\end{equation}

For a schematic illustration, see Figure <ref>.

Transformer stacked in the $q$ dimension

§.§.§ Atmospheric Field Embedding

The physical quantities have discrete semantic meanings, so we encode the positions in the
$q$ dimension using a sum of two shared learned embeddings $e^{\text{field}}_{qd}$ and $e^{\text{level}}_{qd}$.
We treat field labels like “mean sea-level pressure” and “eastward wind speed”, and generalized levels like “surface”, “2m”, “500hPa”, or “integratal” as strings. We then encode them with a categorical embedding.
This is then fed into a transformer of $T_F$ layers operating with $q$ as the sequence dimension and $d$ as the embedding dimension. The output is the snapshot embedding

\begin{equation}
h^{\text{snapshot}}_{sqld} = T_F(e^{\text{field}}_{qd} + e^{\text{level}}_{qd} + h^{\text{slice}}_{sqld}).
\end{equation}

For the schematic illustration, see Figure <ref>.

§.§.§ Sequence Embedding

The neural network operates on a sequence of snapshots.
We tag each snapshot by adding two embeddings: a learned categorical embedding $e^{\text{type}}_{sd}$
with values like “Denoise”, “GEFS”, “Climatology” for the type of the snapshot and a random Fourier embedding [Tancik et al., 2020] $e^{\text{time}}_{sd}$ for the relative physical time.
In particular, the label “Denoise” denotes a snapshot for the score function input $\vv$.
For examples, GEFS ensemble members are labeled as type “GEFS” and the same relative physical time embedding, so to the sequence transformer, they are exchangeable as intended.

The score function of a diffusion model requires the diffusion time $\tau$ as an additional input,
which we model by prepending to the sequence the token $e^{\tau}_{d}$, derived from a random Fourier embedding of $\tau$.
The sequence and its embeddings are then fed into a transformer of $T_S$ layers operating with $s$ as the sequence dimension and $d$ as the embedding dimension. The output embedding is

\begin{equation}
h^{\text{out}}_{sqld} = T_S(e^{\tau}_{d} \oplus_s (e^{\text{type}}_{sd} + e^{\text{time}}_{sd} + h^{\text{snapshot}}_{sqld})),
\end{equation}

where $\oplus_s$ means concatenation in the $s$ dimension.  This is illustrated in Figure <ref>.

Transformer stacked in the $s$ dimension

To obtain the score function $s_\theta(\vv, \tau)$, we take the output token $h^{\text{out}}_{1qld}$ matching the position of the denoise input and project it with learned weights and reshape back to a tensor $v^{\text{out}}_{qp}$ as the final output.

§.§.§ Hyperparameters

The most salient hyperparameters of this architecture are the patch size $P$, the model embedding dimension $D$, and the number of layers $L_L,L_F,L_S$ of transformer stacks $T_L,T_F,T_S$. In our work, we use patch size $P=12$ (number of patches $96$), embedding dimension $D=768$, and transformer layers $(T_L,T_F,T_S)=(6,4,6)$. The hidden layers in the feed-forward network have dimension $4D$ as usual.
The model has 113,777,296 trainable parameters. These parameters are learned by training over randomly sampled days from the 20-year GEFS reforecast dataset [Guan et al., 2022], using a batch size 128 for 200,000 steps.

§ EVALUATION METRICS AND MORE RESULTS

Similar to how we organize the training data (cf. <ref>), we partition all evaluation data according to the lead time $l$, and generate one evaluation dataset $\mathcal{D}^\textsc{eval}_l$ for each  training dataset $\mathcal{D}^\textsc{train}_l$. In the main text, we report results for ensembles generated by models $\mathcal{M}_{l2}$ and $\mathcal{M}_{l2K'}$, conditioned on $K=2$ ensemble forecasts.

§.§ Evaluation Metrics

We use $\vv$ (or $\vw$) to denote ensemble forecasts, which are indexed by $t$ (time), $m$ (ensemble member ID), $q$ (the atmospheric variable of interest), and $p$ (the geolocation). We use $M_{\vv}$ to denote the number of members in the ensemble, $\overline{\vv}$ to denote the ensemble mean forecast, and $s_{\vv}$ to denote the spread:

\begin{align}
\overline{\vv} & = \frac{1}{M_{\vv}} \sum_{m=1}^{M_{\vv}} \vv_m,\\
s_{\vv} & = \sqrt{\frac{1}{M_{\vv}-1}\sum_{m}(\vv_{m} -\overline{\vv})^2}.
\end{align}

We omit other indices ($t$, $q$ and $p$) to avoid excessive subscripts, assuming the operations are applied element-wise to all of them.

§.§.§ RMSE

The (spatial) root-mean-square-error (RMSE) between two (ensemble) means $\overline{\vv}$ and $\overline{\vw}$  is defined as

\begin{equation}
\rmse_t(\overline{\vv}, \overline{\vu}) = \sqrt{\frac{1}{P} \sum_{p=1}^P (\overline{\vv}_{tp} - \overline{\vw}_{tp})^2},
\end{equation}

where $p$ indexes all the geospatial locations. Again, this is computed element-wise with respect to every variable $q$. We summarize the RMSE along the temporal axis by computing the sample mean and variance,

\begin{align}
\overline{\rmse}(\overline{\vv}, \overline{\vu}) & =  \frac{1}{T}\sum_{t=1}^T  \rmse_t(\overline{\vv}, \overline{\vu}),     \label{eEnsembleMeansRMSE}\\
s^2_{\rmse(\overline{\vv}, \overline{\vu})} & = \frac{1}{T-1} \sum_{t=1}^T \left( \rmse_t(\overline{\vv}, \overline{\vu}) - \overline{\rmse}(\overline{\vv}, \overline{\vu})\right)^2.
\end{align}

Assuming the errors at different time are i.i.d., the sample variance can be used to estimate the standard error of the mean estimate eq. (<ref>).

§.§.§ Correlation Coefficients and ACC

The spatial correlation between two (ensemble) means $\overline{\vv}$ and $\overline{\vw}$  is defined as

\begin{equation}
\corr_t(\overline{\vv}, \overline{\vu}) =\frac{\sum_p (\overline{\vv}_{tp} -\overline{\overline{\vv}_{tp}}) (\overline{\vw}_{tp} -\overline{\overline{\vw}_{tp}} ) }{\sqrt{\sum_p(\overline{\vv}_{tp} -\overline{\overline{\vv}_{tp}})^2} \sqrt{\sum_p(\overline{\vw}_{tp} -\overline{\overline{\vw}_{tp}})^2} },
\label{eEnsembleMeansPearson}
\end{equation}

where $\overline{\overline{\cdot}_{tp}}$ refers to averages in space.

The centered anomaly correlation coefficient () is defined as the spatial correlation between the climatological anomalies [Wilks, 2019],

\begin{equation}
\acc_t(\overline{\vv}, \overline{\vu})= \frac{\sum_p (\overline{\vv}_{tp}' -\overline{\overline{\vv}_{tp}'}) (\overline{\vw}_{tp}' -\overline{\overline{\vw}_{tp}'} ) }{\sqrt{\sum_p(\overline{\vv}_{tp}' -\overline{\overline{\vv}_{tp}'})^2} \sqrt{\sum_p(\overline{\vw}_{tp}' -\overline{\overline{\vw}_{tp}'})^2} },
\label{eEnsembleMeansACC}
\end{equation}

where the anomalies are defined as the raw values minus the corresponding climatological mean $\vc_{tp}$,

\begin{equation}
\overline{\vv}_{tp}' = \overline{\vv}_{tp} -\vc_{tp}, \;\;\;\; \overline{\vw}_{tp}' = \overline{\vw}_{tp} -\vc_{tp}.
\label{eAnom}
\end{equation}

Both estimates eq. (<ref>) and (<ref>) are then averaged in time over the evaluation set.

§.§.§ CRPS

The continuous ranked probability score (CRPS) [Gneiting and Raftery, 2007] between an ensemble and  is given by

\begin{equation}
\crps(\vv, \era) = \frac{1}{M_{\vv}} \sum_{m=1}^{M_{\vv}} |\vv_m - \era| - \frac{1}{2M_{\vv}^2}\sum_{m=1}^{M_{\vv}}\sum_{m'=1}^{M_{\vv}} |\vv_m - \vv_{m'}|.
\label{eEnsembleCRPSAppendix}
\end{equation}

This is computed for all time and geo-spatial locations, and then averaged.

We use the traditional CRPS instead of the ensemble-adjusted CRPS$^*$ [Leutbecher, 2019]. The latter examines properties of ensembles in the theoretical  asymptotic limit where their sizes are infinite.  Since our approach aims at augmenting the physics-based ensembles, the definition eq. (<ref>) is  more operationally relevant; even perfect emulation would obviously not improve the CRPS$^*$ score.
Indeed, our model can be seen as a practical way to bridge the skill gap between a small ensemble and the asymptotic limit.

§.§.§ Rank Histogram and Reliability Metric

We assess the reliability of the ensemble forecasts in terms of their rank histograms [Anderson, 1996].
An ideal ensemble forecast should have a flat rank histogram, so deviations
from a flat rank histogram indicate unreliability.

We aggregate rank histograms for all $n_{\text{testing}}$ dates in the evaluation set of  days to obtain the average rank histogram $\{s_i\}_{i=0}^{M_{\vv}}$, where $s_i$ is the number of times the label has rank $i$ and $M_{\vv}$ is the ensemble size.
Following Candille and Talagrand [Candille and Talagrand, 2005], the unreliability $\Delta$ of the ensemble is defined
as the squared distance of this histogram to a flat histogram

\begin{equation}
\Delta=\sum_{i=0}^{M_{\vv}} \left(s_i-\frac{n_{\text{testing}}}{1+M_{\vv}}\right)^2.
\end{equation}

This metric has the expectation $\Delta_0=n_{\text{testing}}\frac{M_{\vv}}{1+M_{\vv}}$ for a perfectly reliable forecast.
The unreliability metric $\delta$ is defined as

\begin{equation}
\delta=\frac{\Delta}{\Delta_0},
\end{equation}

which has the advantage that it can be used to compare ensembles of different sizes. Once $\delta$ is computed for each location, a global $\delta$ value is obtained by averaging over all locations.

§.§.§ Brier Score for Extreme Event Classification

To measure how well the ensembles can predict extreme events,  we first apply a binarization criterion $b:\mathbb{R}\rightarrow \{0,1\}$ (such as “if t2m is 2$\sigma$ away from its mean, issue an extreme heat warning.”), and verify against the classification of the same event by the ERA5 HRES reanalysis,

\begin{equation}
\end{equation}

by converting the ensemble forecast into a probabilistic (Bernoulli) prediction

\begin{equation}
b^{\vv}_{tqp}=\frac{1}{M_{\vv}}\sum_{m=1}^{M_{\vv}} b(\vv_{tmqp}).
\end{equation}

We can evaluate how  well the  probabilistic predictions align with the discrete labels by evaluating the Brier score [Brier, 1950]

\begin{equation}
\brier^b(\vv, \era)_{tq} =\frac{1}{P}\sum_{p=1}^P
(b^{\era}_{tqp} - b^{\vv}_{tqp} )^2.
\label{eBrierScore}
\end{equation}

Alternatively, the probabilistic predictions can be evaluated by their cross-entropy, also known as the logarithmic loss [Benedetti, 2010]

\begin{equation}
\logloss^b(\vv, \era)_{tq} =-\frac{1}{P}\sum_{p=1}^P
b^{\vv}_{tqp} \ln b^{\era}_{tqp} + (1 - b^{\vv}_{tqp} ) \ln (1 - b^{\era}_{tqp}).
\label{eLogLoss}
\end{equation}

It is possible for all ensemble members to predict the wrong result (especially for small ensembles), leading to the $\logloss^b = -\infty$ due to $\ln(0)$.
In practice, a small number $\epsilon=10^{-7}$ is added to both logarithms since no forecast is truly impossible [Benedetti, 2010].
The value of the metric is thus affected by the choice of $\epsilon$, a known inconvenience of this otherwise strictly proper metric.

§.§.§ GEFS Model Climatological Spread

The baseline GEFS-Climatology spread in Figure <ref> is computed from the training dataset, which is the 20-year GEFS reforecast [Guan et al., 2022]. The climatological spread is computed for each lead time independently. For a fixed lead time, the pointwise spread of the 5-member ensemble is computed for each day. Then the model climatological spread for that lead time is defined as the day-of-year average over the 20 years.

Note that we slightly abuse standard notation, since this spread accounts not only for the internal variability of the model, but also for forecast uncertainty [Slingo and Palmer, 2011].

§.§ Detailed and Additional Results

Unless noted, the results reported here are obtained from $N=512$ generated ensembles from our models, with $K=2$ seeding forecasts and $K'=3$ for .

§.§.§ RMSE, ACC, CRPS, Reliability for All Fields

We report values for (Figure <ref>), (Figure <ref>), (Figure <ref>), and $\delta$ (Figure <ref>) for all 8 modeled fields in Table <ref>. The results  are consistent with those reported in the main text: ensembles attain similar or better skill than the physics-based ensembles, with performing better than for variables that are biased in . In addition, is in general more reliable than and , particularly within the first forecast week.

RMSE of the ensemble means with ERA5 HRES as the “ground-truth” labels.

ACC of the ensemble means with ERA5 HRES as “ground-truth” labels.

CRPS with ERA5 HRES  as the “ground-truth” labels.

Unreliability $\delta$ with ERA5 as the label.

§.§.§ Brier score of all fields at various thresholds

From  Figure <ref> to Figure <ref>, we report the Brier score at different thresholds: $\pm 3\sigma,\pm 2\sigma, \text{and} \pm 1\sigma$, respectively.  For all thresholds, provides the most accurate extreme forecast ensembles among all others, for most lead times and variables. Moreover, the generative ensembles from and perform particularly well for longer lead times and more extreme prediction tasks. For the most extreme thresholds considered ($\pm 3\sigma$), the generative ensembles in general outperform .

Brier score for $-3\sigma$ with ERA5 HRES as the “ground-truth” label.

Brier score for $-2\sigma$ with ERA5 HRES as the  “ground-truth” label.

Brier score for $-1\sigma$ with ERA5 HRES as the  “ground-truth” label.

Brier score for $+1\sigma$ with ERA5 HRES as the  “ground-truth” label.

Brier score for $+2\sigma$ with ERA5 HRES as the  “ground-truth” label.

Brier score for $+3\sigma$ with ERA5 HRES as the  “ground-truth” label.

§.§ Effect of $N$, $K$ and $K'$

§.§.§ Generative ensemble emulation with varying $N$ for 7-day lead time with $K=2$

Figure <ref> shows the effect of $N$, the size of the generated ensemble, on the skill of forecasts. (Some metrics such as Brier scores for $2\sigma$ are omitted as they convey similar information.)

For most metrics, increasing $N$ offers diminishing marginal gains beyond $N=256$. Metrics which are sensitive to the sampling coverage, such as the unreliability $\delta$ and the Brier score at $-2\sigma$, still improve with larger $N >256$.

The effect of $N$ on  various metrics. Top down: RMSE, ACC, CRPS, unreliability, Brier score for $-2\sigma$ events.

§.§.§ Generative ensemble emulation with varying $K$ for 7-day lead time

Figure <ref> studies the effect of $K$, the number of seeds,  on the skill of forecasts. We find that the ensemble forecast skill of is comparable to physics-based full ensembles,  when conditioned on $K\geq2$ seeds from GEFS.  Note that is a drastically better forecast system than physics-based ensembles formed by the conditioning members ().

For all values of $K$, the affordability of large ensemble generation warranted by leads to very significant improvements in extreme event classification skill with respect to the conditioning ensemble.

It is also shown that the generative emulation approach is also feasible when $K=1$. However, the skill of the generative ensemble  is closer to physics-based ensembles with just a few members () than to . These studies suggest that $K=2$ could be a good compromise between generative ensemble performance and computational savings.

The effect of $K$ on  various metrics. Top down: RMSE, ACC, CRPS, unreliability, Brier score for $-2\sigma$ event.

§.§.§ Generative post-processing with varying $K'$ for 7-day lead time with $K=2$

Figure <ref> studies the effect of $K'$, which contains the mixture weight, on the skill of forecasts. Since the training data has ensemble size $5$, the number of available GEFS training labels is $3$. We see no forecast skill gains blending in beyond 4 reanalyses members (around 50% mixing ratio). In the main text, we report results with $K'=3$.

The effect of $K'$ on various metrics. Top down: RMSE, ACC, CRPS, unreliability, Brier score for $-2\sigma$ event.

§.§.§ Case Study:  Visualization of Generated Ensembles

Ensemble forecasts are samples drawn from multivariate high-dimensional distributions. Making comparisons among these distributions using a limited number of samples is a challenging statistical task. In this case study, we study several ensembles by reducing the dimensions and focusing on two fields at a single spatial location.

Figures <ref> and <ref> visualize the joint distributions of
temperature at 2 meters and total column water vapour at the grid point near Lisbon during the extreme heat event on 2022/07/14. We used the 7-day forecasts made on 2022/07/07. For each plot, we generate 16,384-member ensembles, represented by the orange dots using different seeds from the 31-member GEFS ensemble. The density level sets are computed by fitting a kernel density estimator to the ensemble forecasts.  The observed weather event is denoted by the black star.  The operational ensembles, denoted by blue rectangles (seeds) and orange triangles,  do not cover the extreme event well, while the generated ensembles extrapolate from the two seeding forecasts, providing a better coverage of the event.

In Figure <ref>, we compare and generated samples conditioned on 6 different randomly sampled two seeding forecasts from the 31-member GEFS ensemble. The plots are arranged in 2 pairs of rows (4 rows in total) such that all the odd rows are from and even rows are from and that on each paired rows, the seeds for above and below are exactly the same. We see that given the same seeds, the distribution is correctly more biased towards warmer temperatures while maintaining the distribution shape and spread compared to .

In Figure <ref> we compare based on 2 seeds with based on 4 seeds. We first generate samples conditioned on 4 randomly sampled seeding forecasts from the 31-member GEFS ensemble. Then we generate samples conditioned on 2 randomly selected seeding forecasts from the 4 chosen in the first round. The plots show 6 such pairs of samples, arranged in 2 pairs of rows (4 rows in total) such that all odd rows use 2 seeds and even rows use 4 seeds and that for each pair of rows, the 4 seeds below always contain the 2 seeds above. We observe that the distribution of the generated samples conditioned on 4 seeds are more similar to each other, which verifies the intuition that more seeds lead to more robustness to the sampling of the seeds. More seeds also lead to better coverage of the extreme heat event by the generated envelopes.

Temperature at 2 meters and total column water vapour of the grid point near Lisbon on 2022/07/14. SEEDS-GEE (green) and SEEDS-GPP (red) with $N=16384$ samples (dots and kernel density level sets) were conditioned on various GEFS subsamples of size $K=2$ (square).

Conditioned on 2

Conditioned on 4

Conditioned on 2

Conditioned on 4

Temperature at 2 meters and total column water vapour of the grid point near Lisbon on 2022/07/14. SEEDS-GEE with $N=16384$ samples (dots and kernel density level sets) were conditioned on various GEFS subsamples of size $K=4$ (square).