CoDiCast: Conditional Diffusion Model for Global Weather Forecasting with Uncertainty Quantification

1 Introduction

Weather prediction describes how the weather states evolve by mapping the current weather states to future weather states [Palmer, 2012]. Accurate weather forecasting is crucial for a wide range of societal activities, from daily planning to disaster preparedness [Shi et al., 2025]. For example, governments, organizations, and individuals rely heavily on weather forecasts to make informed decisions that can significantly impact safety, economic efficiency, and overall wellbeing. However, weather predictions are intrinsically uncertain largely due to the complex and chaotic nature of atmospheric processes [Slingo and Palmer, 2011]. Therefore, assessing the range of probable weather scenarios is significant, enabling informed decision-making.

Traditional numerical weather prediction (NWP) methods achieve weather forecasting by approximately solving the differential equations representing the integrated system between the atmosphere, land, and ocean [Price et al., May 2024; Nguyen et al., 2023]. However, running such an NWP model can produce only one possibility of the forecast, which ignores the weather uncertainty. To solve this problem, Ensemble forecast 1 of multiple models is often employed to model the probability distribution of different future weather scenarios [Palmer, 2019; Leinonen et al., 2023]. While such NWP-based ensemble forecasts effectively model the weather uncertainty, they have two primary limitations: physics-based models inherently make restrictive assumptions of atmospheric dynamics, and running multiple NWP models requires extreme computational costs [Rodwell and Palmer, 2007].

In recent years, machine learning (ML)-based weather predictions (MLWP) have been proposed to challenge NWPbased prediction methods [Ben Bouall`egue et al., 2024; B¨ulte et al., 2024]. They have achieved enormous success with comparable accuracy and a much (usually three orders of magnitude) lower computational overhead. They are typically trained to learn weather patterns from a huge amount of historical data and predict the mean of the probable trajectories by minimizing the mean squared error (MSE) of model forecasts [Hewage et al., 2021]. Representative work includes Pangu [Bi et al., 2023], GraphCast [Lam et al., 2023], ClimaX [Nguyen et al., 2023], ForeCastNet [Pathak et al., 2022], Fuxi [Chen et al., 2023b], Fengwu and [Chen et al., 2023a]. Despite the notable achievements of these MLWP methods, most of them are deterministic [Kochkov et al., 2024], falling short in capturing the uncertainty in weather forecasts [Jaseena and Kovoor, 2022].

1 Generating a set of forecasts, each of which represents a single possible scenario.

To compute the uncertainty for ML models, two methods exist. Perturbing initial conditions [Morley et al., 2018] helps estimate the aleatoric uncertainty (data noise), while the Monte Carlo Dropout approach [Gal and Ghahramani, 2016] estimates epistemic uncertainty (model uncertainty) [Siddique et al., 2022]. However, neither approach fully captures uncertainty in both the input conditions and the evolution of weather models. Additionally, these methods require manual tuning of perturbations and dropout rates, which can negatively impact model accuracy. These limitations motivate us to explore an approach for comprehensive uncertainty quantification without the loss of accuracy.

Denoising diffusion probabilistic models (DDPMs) [Ho et al., 2020] stand out as a probabilistic type of generative model, which can generate high-quality images. By explicitly and iteratively modeling the noise additive and its removal, DDPMs can capture intricate details and textures of images. Furthermore, controllable diffusion models [Zhang et al., 2023] enable the generation process to be guided by specific attributes or conditions, e.g., class labels, textual descriptions, or other auxiliary information. By doing so, the models can generate images that adhere to the specified conditions. This inspires us to consider the weather "prediction" tasks as "generation" tasks - generating plausible weather scenarios with conditional diffusion models. Promising potentials could be the following: (1) Weather numerical data is usually a 2-D grid over latitude and longitude, sharing a similar modality with the image. Diffusion models can capture the intricate weather distribution with iterative denoising. (2) Weather states from the recent past (i.e., initial conditions) can be injected into diffusion models to guide the generation of future weather evolution. (3) More notably, the starting noise sampling from the Gaussian distribution can mimic the aleatoric uncertainty while iteratively adding and removing noise captures the epistemic uncertainty. These features prompt probabilistic diffusion models to generate a set of diverse weather scenarios rather than a single deterministic one. Our contributions are presented as follows:

• We identify the shortcomings of current weather prediction methods. NWP-based methods are limited to restrictive assumptions and computationally intensive. Moreover, a single deterministic NWP- and MLWPbased method cannot achieve uncertainty quantification.
• To address these problems, we propose CoDiCast, a conditional diffusion model for global weather prediction conditioning on observations from the recent past while probabilistically modeling the uncertainty. In addition, we use the cross-attention mechanism to effectively integrate conditions into the denoising process to guide the generation tasks.
• We conduct extensive experiments on a decade of ERA5 reanalysis data from the European Centre for MediumRange Weather Forecasts (ECMWF), and evaluate our method against several state-of-the-art models in terms of accuracy, efficiency, and uncertainty. It turns out that CoDiCast achieves an essential trade-off among these valuable properties.

2 RelatedWork

Physics-based Numerical Weather Prediction. Numerical Weather Prediction (NWP) methods achieve weather forecasts by modeling the system of the atmosphere, land, and ocean with complex differential equations [Bauer et al., 2015]. For example, High-Resolution Forecasts System (HRES) [ECMWF, 2023] forecasts possible weather evolution out to 10 days ahead. However, it is a deterministic NWP method that only provides a single forecast. To overcome the limitation of deterministic methods, the ensemble forecast suite (ENS) [Buizza, 2008] was developed as an ensemble of 51 forecasts by the European Centre for Medium-Range Weather Forecasts (ECMWF). ENS provides a range of possible future weather states, allowing for investigation of the detail and uncertainty in the forecast. Even if NWP ensemble forecasts effectively model the weather evolution, they exhibit sensitivity to structural discrepancies across models and high computational demands [Balaji et al., 2022].

ML-Based Weather Prediction. Pangu [Bi et al., 2023] employed three-dimensional transformer networks and Earthspecific priors to deal with complex patterns in weather data. GraphCast [Lam et al., 2023] achieved medium-range weather prediction by utilizing an "encode-process-decode" configuration with each part implemented by graph neural networks (GNNs). GNNs perform effectively in capturing the complex relationship between a set of surface and atmospheric variables. A similar GNN-based work is [Keisler, 2022]. Fuxi [Chen et al., 2023b] and Fengwu [Chen et al., 2023a] also employ the "encode-decode" strategy but with the transformer-based backbone. FourCastNet [Pathak et al., 2022] applied Vision Transformer (ViT) and Adaptive Fourier Neural Operators (AFNO), while ClimaX [Nguyen et al., 2023] also uses a ViT backbone but the trained model can be fine-tuned to various downstream tasks. However, these models fall short in modeling the uncertainty of weather evolution [Jaseena and Kovoor, 2022] even though perturbations are added to initial conditions [B¨ulte et al., 2024] and dropout methods [Gal and Ghahramani, 2016] are used to mimic the uncertainty. Additionally, ClimODE [Verma et al., 2024] incorporated the physical knowledge and developed a continuous-time neural advection PDE weather model.

DiffusionModels. Diffusion models [Ho et al., 2020] have shown their strong capability in computer vision tasks, including image generation [Li et al., 2022], image editing [Nichol et al., 2021], semantic segmentation [Brempong et al., 2022] and point cloud completion [Luo and Hu, 2021]. Conditional diffusion models [Ho and Salimans, 2022] were later proposed to make the generation step conditioned on the current context or situation. However, not many efforts have adopted diffusion models in global medium-range weather forecasting. More recent research has focused on precipitation nowcasting [Asperti et al., 2023; Gao et al., 2024], and is localized in its predictions. GenCast [Price et al., May 2024] is a recently proposed conditional diffusion-based ensemble forecasting for medium-range weather prediction. However, their conditioning is to directly use the observations from the recent past, which is shown to be insufficient (see the last case in the ablation study).

3 Preliminaries

In this section, we formulate the global weather prediction problem and briefly review Denoising Diffusion Probabilistic Models (DDPMs) [Ho et al., 2020].

3.1 ProblemFormulation

Deterministic Global Weather Predictions. Given the input consisting of the weather state(s), Xt ∈ RH×W×C at time t, the problem is to predict a point-valued weather state, Xt+∆t ∈ RH×W×C at a future time point t +∆t. H ×W refers to the spatial resolution of data which depends on how densely we grid the globe over latitudes and longitudes, C refers to the number of channels (i.e., weather variables), and the superscripts t and t +∆t refer to the current and future time points. The long-range multiple-step forecasts could be achieved by autoregressive modeling or direct predictions.

Probabilistic Global Weather Predictions. Unlike the deterministic models that output point-valued predictions, probabilistic methods model the probability of future weather state(s) as a distribution P(Xt+∆t | Xt ), conditioned on the state(s) from the recent past. Probabilistic predictions are appropriate for quantifying the forecast uncertainty and making informed decisions.

3.2 Denoising Diffusion Probabilistic Models

A denoising diffusion probabilistic model (DDPM) [Ho et al., 2020] generates target samples by learning a distribution pθ(x0) that approximates the target distribution q(x0). DDPM comprises a forward diffusion process and a reverse denoising process. The forward process transforms an input x0 with a data distribution of q(x0) to a Gaussian noise vector xN in N diffusion steps. It can be described as a Markov chain that gradually adds Gaussian noise to the input according to a variance schedule {β1, . . . , βN}:

where at each step n, the diffused sample xn is obtained from xn-1 as described above. Multiple steps of the forward process can be described as follows in a closed form:

In the reverse process, the denoiser network is used to recover x0 by stepwise denoising starting from the pure noise sample, xN. This process is formally defined as:

where pθ(xn) is the distribution at step n parameterized by θ. For each iteration, n ∈ [1, N], diffusion models are trained to minimize the following KL-divergence:

where q(xn-1|xn) can be computed from q(xn|xn-1) using Bayes rule and the multistep forward process equation above.

4 Methodology

This section introduces our approach for global weather prediction, CoDiCast, implemented as a conditional diffusion model. The key idea is to consider "prediction" tasks as "generation" tasks while conditioning on the context provided by recent past observation(s). An overview of the proposed CoDiCast is shown in Figure 2.

4.1 Forward Diffusion Process

At time point t, the forward diffusion process assumes that a pure noise sample Xt+1 N is generated from Xt+1 0 ∈ RH×W×C by adding noise N times (see the dotted lines in Figure 2):

where ϵ is sampled from N(0, I) with the same dimensions as Xt+1 0 , and ¯α is as described in Section 3.2).

4.2 Reverse Conditional Denoising Process

CoDiCast models the probability distribution of the future weather state, conditioned on the current and previous weather states. We exploit a pre-trained encoder to learn conditions as embedding representations of the past observations Xt-1 and Xt , which are used to control and guide the synthesis process. Working in the latent space of embeddings works better than the original space of the observations.

where Xt+1 N ∼ N(0, I), ˜ Z t-1:t is the embedding representation as shown in Eq. (6).

After prediction at the first time point is obtained, a forecast trajectory, X1:T , of length T, can be auto-regressively modeled by conditioning on the predicted "previous" states.

4.3 Pre-trained Encoder

We learn an encoder by training an autoencoder network [Baldi, 2012]. An Encoder compresses the input at each time point into a latent-space representation, while Decoder reconstructs the input from the latent representation. After the encoder, F, is trained, it can serve as a pre-trained representation learning model to project the original data into a latent embedding in Eq. (6). We provide more details on the model architecture in the arXiv version 2 .

4.4 Attention-based Denoiser Network

Our denoiser network consists of two blocks: cross-attention and U-net (as shown in Figure 4). Cross-attention mechanism [Hertz et al., 2022] is employed to capture how past observations can contribute to the generation of future states. The embedding of past observations, ˜ Z t-1:t , and the noise data Xt+1 n at diffusion step n, are projected to the same hidden dimension d with the following transformation:

2 Preprint version: https://arxiv.org/abs/2409.05975

learnable projection matrices. Then we implement the cross-attention mechanism by Attention(Q, K, V) = softmax( QKT √d )V . A visual depiction of the cross-attention mechanism is in the arXiv version.

U-Net [Ronneberger et al., 2015] is utilized to recover the data by removing the noise added at each diffusion step. The skip connection technique in U-Net concatenates feature maps from the encoder to the corresponding decoder layers, allowing the network to retain fine-grained information that might be lost during downsampling. The detailed U-Net architecture is presented in the arXiv version.

4.5 Training Process

The training procedure is shown in the arXiv version.Firstly, we pre-train an encoder to learn the condition embedding of the past observations. Subsequently, we inject it into our conditional diffusion model and train CoDiCast with the devised loss function:

where Xt+1 n = √ ¯αnXt+1 0 +√1- ¯αnϵ, cond = F(Xt-1:t ), and ϵθ is the denoiser in Figure 4.

4.6 Inference Process

We first extract the conditional embedding representations, ˜ Z t-1:t , by the pre-trained encoder, and then randomly generate a noise vector XN ∼ N(0, I). The sampled noise vector, XN, is autoregressively denoised along the reversed chain to predict the target until n equals 1, we obtain the weather prediction ˆX0 at the time t + 1. Later, multi-step prediction can be implemented autoregressively - the output from the previous time step is the input while predicting the next step, as shown in Eq. (5). We provide the pseudocode in the arXiv version.

4.7 Ensemble Forecast

To enhance the reliability of weather forecasts, ensemble forecast strategy is often employed to capture the variability among forecasts by separately running multiple deterministic models, e.g., ensemble forecast suite (ENS) [Buizza, 2008]. In our approach, since CoDiCast is a probabilistic model that can generate a distribution of future weather scenarios rather than a single prediction, following [Price et al., May 2024], we run the trained CoDiCast multiple times to get the ensemble instead. More specifically, by integrating both initial conditions and noise sampled from a Gaussian distribution, CoDiCast implements the ensemble forecast through multiple stochastic samplings during inference, capturing a range of possible forecasts for the uncertainty quantification.

5 Experiments

5.1 Dataset and Baselines

Dataset. ERA5 [Hersbach et al., 2020] is a publicly available atmospheric reanalysis dataset provided by the European Centre for Medium-Range Weather Forecasts (ECMWF). Following the existing work [Verma et al., 2024], we use the preprocessed 5.625◦ resolution (32 × 64) and 6-hour increment ERA5 dataset fromWeatherBench [Rasp et al., 2020]. We downloaded 5 variables for the globe: geopotential at 500 hPa pressure level (Z500), atmospheric temperature at 850 hPa pressure level (T850), ground temperature (T2m), 10 meter U wind component (U10) and 10 meter V wind component (V10). More details are in the arXiv version.

Baselines. We comprise the following baselines: ClimODE [Verma et al., 2024]: a spatiotemporal continuous-time model that incorporates the physical knowledge of atmospheric advection over time; ClimaX [Nguyen et al., 2023]: a state-ofthe-art vision Transformer-based method trained on the same dataset (without pre-training that is used in the original paper); FourCastNet [Pathak et al., 2022]: a global data-driven weather model using adaptive Fourier neural operators; Neural ODE [Chen et al., 2018]: an ODE network that learns the time derivatives as neural networks by solving an ordinary differential equation; Integrated Forecasting System IFS [Rasp et al., 2020]: a global numerical weather prediction (NWP) system, integrating multiple advanced physics-based models to deal with more meteorological variables across multiple altitudes. Our study focuses solely on a subset of these variables due to the limited computational resources, with IFS serving as the gold standard. For a fair comparison, all ML models use the same data set described in Section 5.1.

5.2 Experiments Design

We use data between 2006 and 2015 as the training set, data in 2016 as the validation set, and data between 2017 and 2018 as the testing set. We assess the global weather forecasting capabilities of our method CoDiCast by predicting the weather at a future time t +∆t (∆t = 6 to 144 hours) based on the past two time units. To quantify the uncertainty in weather prediction, we generate an "ensemble" forecast by running CoDiCast five times during the inference phase.

Training. We first pretrain an encoder model with the Autoencoder architecture. For the diffusion model, we used U-Net as the denoiser network with 1000 diffusion/denoising steps. The architecture is similar to that of DDPM [Ho et al., 2020] work. We employ four U-Net units for both the downsampling and upsampling processes. Before training, we apply Max-Min normalization [Ali et al., 2014] to scale the input data within the range [0, 1], mitigating potential biases stemming from varying scales [Shi et al., 2023]. Adam was used as the optimizer, where the learning rate = 2e-4 , decay steps = 10000, decay rate = 0.95. The batch size and number of epochs were set to 64 and 800, respectively. More training details and model configurations are in the arXiv version.We conduct all experiments on one NVIDIA A100 GPU with 80GB memory.

Evaluation Metrics. Following [Verma et al., 2024], we use latitude-weighted Root Mean Square Error (RMSE) and Anomaly Correlation Coefficient (ACC) as deterministic metrics. RMSE measures the average difference between values predicted by a model and the actual values. ACC is the correlation between prediction anomalies relative to climatology and ground truth anomalies relative to climatology. It is a critical metric in climate science to evaluate the model's performance in capturing unusual weather or climate events. The formulas of these metrics are in the arXiv version.

5.3 Quantitative Evaluation

Accuracy. We compare different models in forecasting five primary meteorological variables as described in Section 5.1. Table 1 shows that CoDiCast presents superior performance over other MLWP baselines, demonstrating diffusion models can capture the weather dynamics and make predictions accurately. However, there is still room for CoDiCast to be improved compared with the gold-standard IFS model, which integrates more meteorological variables.

Uncertainty. The gray-shaded error range in Table 1 represents model uncertainty, with fluctuations remaining within 10% of the ground truth scale. This indicates the robustness of ML models. However, ForeCastNet predictions fluctuate relatively larger due to the sensitive selection of the dropout rate. We provide a case study of CoDiCast forecast for 72 hours with uncertainty quantification in Figure 5. It shows the mean prediction tracks the general trend of the ground truth and the uncertainty grows as the lead time increases. Besides, we can tell that most actual values fall within the 1 or 2 standard deviations (σ) ranges, indicating these predictions are reasonably accurate.

Inference efficiency. Generally, numerical weather prediction models (e.g., IFS) require around 50 minutes for the medium-range global forecast, while deterministic ML weather prediction models take less than 1 minute [Rasp et al., 2020] but cannot model the weather uncertainty. CoDiCast needs about 3.6 minutes (see the last row in Table 2) for the global weather forecast, potentially balancing the efficiency and accuracy with essential uncertainty quantification. The efficiency also depends on the model complexity.

5.4 Qualitative Evaluation

In Figure 6, we qualitatively evaluate the performance of CoDiCast on global forecasting tasks for all target variables, Z500, T850, T2m, U10 and V10 at the lead time of 6 hours. The first row is the ground truth of the target variable, the second row is the prediction and the last row is the difference between the model prediction and the ground truth. From the scale of color bars, we can tell that the error percentage is less than 3% for variables Z500, T850, and T2m. However, error percentages over 50% exist for U10 and V10 even though only a few of them exist. We also observe that most higher errors appear in the high-latitude ocean areas, probably due to the sparse data nearby. We provide visualizations for longer lead times in the arXiv version.

5.5 Ablation Study

CoDiCast includes two important components: pre-trained encoder and cross attention. To study their effectiveness, we conduct an ablation study: (a) No-encoder directly considers past observations as conditions to the diffusion model; (b) No-cross-attention simply concatenates the embedding and the noisy sample at each denoising step; (c) No-encodercross-attention concatenates the past observations and the noisy sample at each denoising step. Figure 7 shows that the full version of CoDiCast consistently outperforms all other variants, verifying positive contributions of both components.

5.6 Parameter Study

Table 2 shows that the accuracy improves as the number of diffusion steps, N, increases when N < 1000, indicating that more intermediate steps are more effective in learning the imperceptible attributes during the denoising process. However, when 1000 < N < 2000, the accuracy remains approximately flat but the inference time keeps increasing linearly. Considering the trade-off between accuracy and efficiency, we finally set N = 1000. Additionally, we study the effect of "linear" and "quadratic" variance scheduling β, where β ∈ [0.0001, 0.02]. Figure 8 shows that the "linear" variance scheduling performs better than the "quadratic" one.

6 Conclusions

In this work, we start with analyzing the limitations of current deterministic numerical weather prediction (NWP) and machine-learning weather prediction (MLWP) approaches. They either require substantial computational costs or lack uncertainty quantification in their forecasts. We propose a conditional diffusion model, CoDiCast, to address these limitations. Experimental results demonstrate it can simultaneously complete more accurate predictions than existing MLWP-based models and a faster inference than NWP-based models while being capable of providing uncertainty quantification compared to deterministic methods. In conclusion, our model offers three key characteristics at the same time.

7 Collaborations and Broader Impacts

Machine learning (ML) and atmospheric researchers were involved in the multidisciplinary collaborations. The domain experts provide ML researchers with data sources and contextual experiment results while obtaining insights into the potential applications of diffusion models in atmospheric sciences. Impact: Our work facilitates weather warnings such as extreme heat and cold, and storm warnings, which could help the government, local agencies, and residents to take necessary preventive and protective measures to minimize losses.

Acknowledgments

This work is supported by the Institute for Geospatial Understanding through an Integrative Discovery Environment (I-GUIDE), which is funded by the National Science Foundation (NSF) under award number: 2118329. Any opinions, findings, conclusions, or recommendations expressed herein are those of the authors and do not necessarily represent the views of NSF.