Continuous Ensemble Weather Forecasting with Diffusion models

Authors: Martin Andrae, Tomas Landelius, Joel Oskarsson, Fredrik Lindsten
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: weather forecasting, diffusion, ensemble forecasting

Abstract:
Weather forecasting has seen a shift in methods from numerical simulations to data-driven systems. While initial research in the area focused on deterministic forecasting, recent works have used diffusion models to produce skillful ensemble forecasts. These models are trained on a single forecasting step and rolled out autoregressively. However, they are computationally expensive and accumulate errors for high temporal resolution due to the many rollout steps. We address these limitations with Continuous Ensemble Forecasting, a novel and flexible method for sampling ensemble forecasts in diffusion models. The method can generate temporally consistent ensemble trajectories completely in parallel, with no autoregressive steps. Continuous Ensemble Forecasting can also be combined with autoregressive rollouts to yield forecasts at an arbitrary fine temporal resolution without sacrificing accuracy. We demonstrate that the method achieves competitive results for global weather forecasting with good probabilistic properties.

Decision: accept (poster)
Reviewer scores: 5, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1729525622960

thread_type:
official_review

<thread_object>
creator: reviewer_pyfz

<content>
summary:
This paper introduces Continuous Ensemble Forecasting for global weather forecasting.

strengths:
1. Weather forecasting is an important problem.

2.The method introduced in the paper seems to exhibit a certain level of novelty.

3.Some of the figures in the paper look good.

weaknesses:
1.The paper's presentation seems somewhat disorganized. The introduction does not clearly articulate the limitations of current weather forecasting methods or how the proposed method addresses these limitations. Figure 1, being the first figure, lacks sufficient information; ideally, a single figure should clearly highlight the differences between the proposed method and previous methods, as well as explain how the proposed method overcomes the limitations of prior approaches. In the Methods section, there is a tendency to mix descriptions of existing methods with the new method, which can obscure the exact nature of the contribution made by the proposed method. Figure 2 should also provide a clearer depiction of the methodology.

2.Based on the unclear description provided in the paper, the contributions appear insufficient, and the contributions listed in the introduction seem to be something that any weather forecasting model could achieve.

3.The baseline models compared in the paper are somewhat limited and do not appear to be the latest SOTA models. Moreover, according to Table 1, the ARCI-24/6h model proposed in the paper (if I understand correctly) is not the optimal model under most settings.

questions:
1.Were the limitations of previous methods that they were computationally expensive? How does the method proposed in the paper address the limitations of previous approaches?

2.What specific contributions does the paper make that differentiate it from other weather forecasting models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730152013651

thread_type:
official_review

<thread_object>
creator: reviewer_1nba

<content>
summary:
This work proposes an ODE-based method to address the limitations of previous a single forecasting step and rolled out autoregressively methods for weather prediction. It can generate long-range weather predictions without iteration and forecast arbitrary lead times.

strengths:
- It is interesting to use probability flow ODE of the noise in diffusion models. Leveraging a lead-time-dependent deterministic ODE-solver to generate a smooth and continuous trajectory is meaningful compared to autoregressive methods.
- Authors provide intensive experiments on their method.

weaknesses:
- Need more descriptions on how probability flow ODE works.
- Need more comparisons with other SOTAs. 
- No available codes to help readers better understand the paper.

questions:
**Major:**
- It is a good idea to use the probability flow ODE from the work proposed from https://arxiv.org/pdf/2206.00364 in diffusion models. But what is the assumption behind the customized ODE in Eq. (2)? Why is there a linear noise level $sigma$ in Eq. (2)? What if $\sigma^2$? More explanation would help. 
- From the paper title and section 4 name, "CONTINUOUS" seems to be used to describe $ensemble forecasting$. However, what the authors really want to convey is the temporal continuity, i.e., any arbitrary fine temporal resolution. You may need to think over the rephrasing. Otherwise, it is confusing for readers (at least for me).
- In Figure 1, do you need to include the 1-hour prediction as the input for the 6-hour prediction? I saw you have explanations in sections 3 and 4. But Figure 1 deserves a more precise description within its title since it is the main framework of your method.
- It would be better to provide more comparisons with previous SOTA ML models. Even though I know not all models did not release source codes to the public, but some did, e.g., FourCastNet and ClimaX. 
- The authors mentioned the limitation of diffusion models: computationally expensive. I suggest adding more analysis of computational efficiency.

**Minor:**
- You are using the letter 'k' for different representations (see Eq. 1 and Alg. 1).
- It is customary to denote a time point with the superscript or subscript. It looks like a function if using $X[k+1]$ or $X(\tau)$.
- In the line under Eq. (2), why is $z(1)$ for $X[k+1]$? Isn't it the number of diffusion steps or the noise levels?
- What is $v_i^{k}$ in Algorithm 2?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730218001404

thread_type:
official_review

<thread_object>
creator: reviewer_c3dx

<content>
summary:
The "Continuous Ensemble Weather Forecasting with Diffusion Models" paper introduces an innovative framework for probabilistic weather forecasting using diffusion models. It focuses on enhancing computational efficiency by removing the autoregressive nature of sampling from diffusion models.

strengths:
The model developed in this paper presents an approach that eliminates traditional autoregressive steps during sampling from diffusion models, allowing it to achieve high temporal resolution and efficiency in forecasting. This is an exciting approach for improving computational efficiency, a feature that has been challenging to achieve in weather prediction. 

The paper's application is exciting, as diffusion models are typically used in areas like image generation rather than meteorological forecasting.

weaknesses:
The paper lacks references to closely related methods, such as Zheng et al. (2023) on fast sampling of diffusion models via operator learning, which shares similarities with the proposed model’s sampling approach. It would be beneficial to clarify how the proposed sampling method diverges from these approaches.

Although not directly used for weather data, recent models like Dyfusion, Alternator, and Mamba (with s-mamba used specifically for weather forecasting) have shown promise in similar domains, such as sea surface temperature forecasting. Including these models as baselines or discussing their relevance could offer a more comprehensive comparison, illustrating where this model stands relative to other contemporary approaches.

The paper mentions “regularity conditions” (e.g., Line 234) without specifying them rigorously. These conditions are fundamental to claims about trajectory continuity and temporal stability of forecasts. A formal mathematical description of these conditions, alongside proof or empirical validation, would be necessary to substantiate these claims.

The claims regarding fixed noise across time steps and ODE-based interpolation are under-supported. No mathematical evidence or ablation study is presented to show the impact of fixing noise on forecasting performance. Including such a study could clarify the advantages of this design choice.

The proposed Algorithm 2 seems tailored to a specific diffusion model setup. It would be valuable to discuss whether this algorithm generalizes to other diffusion models or is fundamentally limited to the framework outlined in the paper. A broader discussion on this topic could enhance the paper’s applicability to related models.

The results lack completeness, as only two of the five variables used for forecasting are presented in Table 2. Including comparisons for all five variables would provide a clearer picture of the model's strengths and weaknesses across different atmospheric variables.

The absence of error bars due to computational limitations weakens the robustness of the reported results. Given the probabilistic nature of the model, showcasing variability over multiple runs would offer greater confidence in the findings. Including error bars or confidence intervals for at least a subset of key results would be helpful, even if computational constraints prevent a complete assessment.

Numerous claims about model performance, such as trajectory continuity, lack rigorous mathematical proof and appear to rely on empirical observations. Addressing this would lend greater theoretical robustness to the proposed model.

There is no mathematical derivation for the training objective and loss function. A step-by-step derivation of these, with explanations of any modifications relative to standard diffusion model objectives and the rationale for these modifications in the weather forecasting context, would clarify the theoretical foundation of the approach.

questions:
Could you clarify the rationale behind using fixed noise across time steps? Additionally, has any sensitivity analysis or ablation study been conducted to validate this design choice? If so, please provide these findings.

Could you further explain how interpolation using an ODE solver contributes to the model’s performance, and how it compares to alternative interpolation methods?

Only two of the five forecasting variables are included in Table 2. Could you explain the rationale for this selective presentation and consider including results for all five variables to enable a comprehensive evaluation?

Could you provide formal proofs or derivations for key claims, particularly those regarding trajectory continuity and temporal stability?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648798859

thread_type:
official_review

<thread_object>
creator: reviewer_kodm

<content>
summary:
This paper uses the diffusion model for continuous ensemble weather forecasting without requiring autoregressive steps to produce future values. As introduced in Section 4.2, the main idea is to replace the fixed Z with a stochastic process Z(T) with continuous sample trajectories.

strengths:
The proposed method is presented clearly.

weaknesses:
W1: there have been exsiting works for nonautoregressive sequence forecasting by diffusion models. However, realted investigation and discussions are missing, e.g., 

[1] NeurIPS'21 CSDI Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation

[2] ICML'23 Non-autoregressive Conditional Diffusion Models for Time Series Prediction 

W2: moreover, some questions and details are not provided.
- it would be better to explain what ensemble forecasting is formally
- it is still not very convincing why a stochastic process is better than random Gaussian noises for nonautoregressive generation. Because we can simply represent the forecast outputs as a vector of a length N. Then we can use a Gaussian vector to yield forecasts. 
- is there a significant difference between weather forecasting and general time series forecasting?
- what is the input size (the number of variables) of the model in the experiments? It seems there are many variables in the weather forecasting task. Did you perform univariate forecasting tasks for each variable?   
- how to fuse $n_{ens}$ trajectories to get the final predictions?

questions:
In the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731964721069

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response about proof of trajectory continuity

comment:
In this paper we have presented a novel method for generating forecast trajectories in a nonsequential fashion, leveraging a stochastic process in the Gaussian noise. Empirically, we show in figure 5 that this generates temporally consistent and continuous forecast trajectories. This can also be seen in the supplementary animations.

However, as pointed out by the reviewers, this claim lacks substantial theoretical proof. In section 4.1 we provide some intuition for why the trajectories would be continuous with enough data and training. This analysis relies heavily on the assumption that the solution to eq. 2 preserves this continuity. In the paper, we claim that this will hold under some regularity conditions but leave this for future work. As this was highlighted by the reviewers as something important, we have spent some time investigating this and provided a proof of this property. **The specific assumptions and proof of desired properties is given in appendix E of the revised paper**.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731965055866

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response about baselines

comment:
We'd like to stress the **importance of using ensemble forecasting models as baselines**. Many of the proposed baselines (FourCastNet [1], ClimaX [2] etc) are deterministic models that can not generate ensemble forecasts without perturbing initial conditions¹. We have included a deterministic baseline trained using the same backbone U-Net as the diffusion model, which we outperform significantly. In fact, the contribution of our paper is not to provide a new backbone architecture, but rather the development of a diffusion-model-based framework for enabling time-continuous ensemble predictions. We can in principle use another backbone architecture in our framework, including GNN-based (GraphCast [3]), Transformer-based (ClimaX), or Neural-operator-based (FourCastNet),  but focused on a U-Net for simplicity. 

We believe latent-variable models are an interesting class of models since they have computational benefits compared to diffusion models. Because of this, we have chosen to compare to Graph-EFM [4], a SOTA latent-variable model for probabilistic weather forecasting. This model is chosen because it's 1) specifically developed for weather, 2) set to appear as a [spotlight at NeurIPS](https://nips.cc/virtual/2024/poster/93149) and crucially 3) have publicly available code. Their model outperforms the GraphCast model which in turn has shown to be superior to other models such as Fourcastnet and ClimaX. Thus we believe that outperforming Graph-EFM is a great achievement. 

We would also like to highlight the connection between our model and the SOTA weather forecasting model GenCast [5]. The AR models used as baselines in the paper is the exact forecasting setup from [5]. The difference is that it's implemented with a U-net backbone instead of a graph neural network. The reason for this is that we want to use the same backbone architecture for the different diffusion model variants that we consider in the paper for a fair comparison, and to avoid confounding the effect of the backbone architecture with the effect of different diffusion model setups. In the paper we show that the ARCI model outperforms the GenCast setup for high temporal resolution, making this a SOTA model for global weather forecasting at 1h resolution. We have updated the paper to clarify this connection between the AR model and GenCast.

Howevever, as suggested by reviewer c3dX, there are other similar models that haven't been applied to weather forecasting. DYffusion [6] is a related model that we discuss in some detail in the related work section. We will try to implement DYffusion in the weather forecasting setup and hope to share the results with you during the discussion period.

[1] Kurth, Thorsten, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, and Anima Anandkumar. “FourCastNet: Accelerating Global High-Resolution Weather Forecasting Using Adaptive Fourier Neural Operators.” In Proceedings of the Platform for Advanced Scientific Computing Conference, 1–11. PASC ’23. New York, NY, USA: Association for Computing Machinery, 2023. https://doi.org/10.1145/3592979.3593412.

[2] Nguyen, Tung, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. “ClimaX: A Foundation Model for Weather and Climate.” In Proceedings of the 40th International Conference on Machine Learning, 25904–38. PMLR, 2023. https://proceedings.mlr.press/v202/nguyen23a.html.

[3] Lam, Remi, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, et al. “Learning Skillful Medium-Range Global Weather Forecasting.” Science 382, no. 6677 (December 22, 2023): 1416–21. https://doi.org/10.1126/science.adi2336.

[4] Oskarsson, Joel, Tomas Landelius, Marc Peter Deisenroth, and Fredrik Lindsten. “Probabilistic Weather Forecasting with Hierarchical Graph Neural Networks.” arXiv, October 26, 2024. https://doi.org/10.48550/arXiv.2406.04759.

[5] Price, Ilan, Alvaro Sanchez-Gonzalez, Ferran Alet, Tom R. Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, et al. “GenCast: Diffusion-Based Ensemble Forecasting for Medium-Range Weather.” arXiv, May 1, 2024. https://doi.org/10.48550/arXiv.2312.15796.

[6] Cachay, Salva Rühling, Bo Zhao, Hailey James, and Rose Yu. “DYffusion: A Dynamics-Informed Diffusion Model for Spatiotemporal Forecasting,” 2023. https://openreview.net/forum?id=WRGldGm5Hz&noteId=jDWiuqrB98.

¹ They can produce a type of ensemble forecasts by perturbing initial conditions, but this captures a different type of uncertainty than what we are focusing on, namely the uncertainty in the dynamics/forecasts even for fixed initial conditions. Note that our method can also be used with initial condition perturbation for capturing both types of uncertainties, but we have not explored this combination to avoid confounding the results.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732536819773

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Gentle Reminder

comment:
Dear Reviewers, Area Chair,
The discussion period ends in two days. We believe that we have addressed all your concerns. See the updated draft, our general reply, and the specific answers we sent you. We would love to hear your thoughts on these updates and provide further clarifications if needed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732793989424

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Message to reviewers

comment:
Dear reviewers,

We have updated the draft with **highlights of all the changes we made to the original draft**, making it easier for you to find the relevant changes. This includes a calculation of error bars from multiple re-trained models in appendix C, as suggested by reviewer c3dx. We have also added a more detailed discussion on the limitations of the DYffusion model which was mentioned by reviewer c3dx as a related model. We are still working on a full retraining and evaluation of DYffusion in the weather forecasting setup and hope to share the results with you before the discussion period ends.

We believe that we have addressed all your concerns. See the updated draft, our general replies, and the specific answers we sent you. We would love to hear your thoughts on these updates and provide further clarifications if needed.

Regards,

The authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733218499453

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Implementation of DYffusion baseline

comment:
Dear Reviewers, Area Chair,

Almost all current probabilistic spatio-temporal models rely on autoregressive rollouts that struggle at high temporal resolution. This includes the current SOTA GenCast model [1]. In this paper we have presented a novel method that scales well to arbitrarily fine temporal resolutions. Another alternative to autoregressive rollouts is the DYffusion framework [2], where stochastic interpolation and deterministic forecasting is combined into a diffusion-like model. The method allows for forecasting at arbitrary temporal resolution, but still requires sequential computations for sampling the prediction. **As promised in the answers to reviewer c3dx**, we have worked on a full retraining and evaluation of DYffusion in the weather forecasting setup. This is the result of these efforts.

**We have adapted the DYffusion codebase to work with weather forecasting using ERA5 data.** This application was not considered in the original DYffusion papers which focused on sea-surface forecasting [2] and climate modeling [3]. To ensure a fair evaluation, we have used the same backbone U-Net as our diffusion model, and a lat-lon weighted mean-squared-error loss. We have trained it with a horizon of 24h, with 1h time-steps, making it directly comparable to ARCI-24/1h. For the hyperparameters, we were not able to do an extensive search due to the time constraints, and instead used values similar to those proposed in the original paper. Below we evaluate DYffusion on the 2018 test data and compare it to the ARCI-24/1h model. 

| | | | 5 days | | | 10 days | | 
| -- | --  | -- | -- | -- | -- | -- | -- | 
| Variable | Model      | RMSE  | CRPS  | SSR  | RMSE   | CRPS  | SSR  | 
| z500  | DYffusion     | 633.5 | 303.7 | 0.93 | 823.1  | 393.6 | 1.00 | 
| | ARCI-24/1h    | 564   | 264.4 | 0.88 | 776.1  | 369   | 0.95 | 
| t850  | DYffusion | 2.9   | 1.52  | 0.88 | 3.53   | 1.86  | 0.95 | 
| | ARCI-24/1h | 2.66  | 1.35  | 0.91 | 3.37   | 1.73  | 0.96 | 
| t2m   | DYffusion | 2.41  | 1.18  | 1.09 | 2.9    | 1.4   | 1.09 | 
| | ARCI-24/1h    | 2.07  | 0.95  | 0.91 | 2.57   | 1.18  | 0.94 | 
| u10   | DYffusion | 3.65  | 1.89  | 0.87 | 4.1    | 2.14  | 0.92 | 
| | ARCI-24/1h | 3.44  | 1.74  | 0.94 | 3.99   | 2.05  | 0.97 | 
| v10   | DYffusion | 3.72  | 1.93  | 0.88 | 4.16   | 2.17  | 0.94 | 
| | ARCI-24/1h | 3.54  | 1.8   | 0.93 | 4.1    | 2.11  | 0.98 | 
| ws10  | DYffusion | 2.51  | 1.36  | 0.91 | 2.68   | 1.47  | 0.95 | 
| | ARCI-24/1h | 2.38  | 1.28  | 0.97 | 2.62   | 1.42  | 0.99 | 

Although DYffusion achieves much better results than the 1h-GenCast setup AR-1h (Fig. 6 in paper), **it is beaten by ARCI-24/1h in all variables and lead times**. While doing this retraining and evaluation we have learned key things about DYffusion that we highlight below.

**1. DYffusion is hard to train**

DYffusion requires training two networks, one interpolator and one forecaster. To get a probabilistic model, they introduce a layer dropout term in the interpolator that they keep on during inference. This makes the performance sensitive to the dropout rate, which can not be changed without retraining both the interpolator and forecasting networks. 

For our model, we took a successful Diffusion framework for images and applied it to weather, with minimal changes to the hyperparameters. We can also tune the performance after training by changing the noise levels in Table 3.

**2. The interpolator has problems with blurring and temporal consistency**

Since all intermediate forecasts are generated using the interpolator, this sets an upper limit on the performance of how well we can expect to perform. The interpolator is trained in the same way as a deterministic models using a MSE loss, which has shown to be prone to blurring.

Further, the interpolator gives no guarantee of temporal continuity of trajectories. The stochastic interpolator (with dropout activated) might predict forecasts at 11h, 12h, that are both possible given the endpoints, but that are not consistent with each other due to sampling different dropout at the different time steps. In our paper we propose a way to correlate the forecasts at each time, overcoming this issue.

[1] Price, Ilan, Alvaro Sanchez-Gonzalez, Ferran Alet, Tom R. Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, et al. “GenCast: Diffusion-Based Ensemble Forecasting for Medium-Range Weather.” arXiv 2024. https://doi.org/10.48550/arXiv.2312.15796.

[2] Cachay, Salva Rühling, Bo Zhao, Hailey James, and Rose Yu. “DYffusion: A Dynamics-Informed Diffusion Model for Spatiotemporal Forecasting,”. Advances in Neural Information Processing Systems, 36, 2023. https://openreview.net/forum?id=WRGldGm5Hz

[3] Cachay, Salva Rühling, Brian Henn, Oliver Watt-Meyer, Christopher S. Bretherton, and Rose Yu. “Probabilistic Emulation of a Global Climate Model with Spherical DYffusion.” Advances in Neural Information Processing Systems, 37, 2024. https://openreview.net/forum?id=Ib2iHIJRTh
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a new approach for ensemble weather forecasting using diffusion modeling. The method is not autoregressive and can generate ensemble member trajectories from initial conditions without relying on iterative steps, even though it can be combined with autoregressive methods for long-range forecasts. Uncertainty quantification is provided via repeated sampling. The approach is theoretically sound.