[Paper Review] SwinVRNN: A Data-Driven Ensemble Forecasting Model via Learned Distribution Perturbation

Yuan Hu, Lei Chen+2 more|arXiv (Cornell University)|May 26, 2022

Meteorological Phenomena and Simulations4 citations

TL;DR

SwinVRNN is a data-driven stochastic weather forecasting model that combines a Swin Transformer-based recurrent neural network (SwinRNN) with a variational inference-based perturbation module to learn spatially adaptive, multivariate Gaussian noise distributions for ensemble forecasting. It achieves state-of-the-art performance on WeatherBench, surpassing the operational ECMWF IFS model in forecast accuracy for 2-m temperature and 6-hourly total precipitation across all lead times up to five days.

ABSTRACT

Data-driven approaches for medium-range weather forecasting are recently shown extraordinarily promising for ensemble forecasting for their fast inference speed compared to traditional numerical weather prediction (NWP) models, but their forecast accuracy can hardly match the state-of-the-art operational ECMWF Integrated Forecasting System (IFS) model. Previous data-driven attempts achieve ensemble forecast using some simple perturbation methods, like initial condition perturbation and Monte Carlo dropout. However, they mostly suffer unsatisfactory ensemble performance, which is arguably attributed to the sub-optimal ways of applying perturbation. We propose a Swin Transformer-based Variational Recurrent Neural Network (SwinVRNN), which is a stochastic weather forecasting model combining a SwinRNN predictor with a perturbation module. SwinRNN is designed as a Swin Transformer-based recurrent neural network, which predicts future states deterministically. Furthermore, to model the stochasticity in prediction, we design a perturbation module following the Variational Auto-Encoder paradigm to learn multivariate Gaussian distributions of a time-variant stochastic latent variable from data. Ensemble forecasting can be easily achieved by perturbing the model features leveraging noise sampled from the learned distribution. We also compare four categories of perturbation methods for ensemble forecasting, i.e. fixed distribution perturbation, learned distribution perturbation, MC dropout, and multi model ensemble. Comparisons on WeatherBench dataset show the learned distribution perturbation method using our SwinVRNN model achieves superior forecast accuracy and reasonable ensemble spread due to joint optimization of the two targets. More notably, SwinVRNN surpasses operational IFS on surface variables of 2-m temperature and 6-hourly total precipitation at all lead times up to five days.

Motivation & Objective

• To address the suboptimal ensemble performance of existing data-driven weather forecasting models, which often rely on fixed or heuristic perturbation methods.
• To develop a stochastic weather forecasting model that jointly optimizes deterministic prediction and uncertainty quantification through end-to-end training.
• To investigate whether learned, spatially adaptive perturbations can improve ensemble spread and forecast accuracy compared to fixed distribution or Monte Carlo dropout methods.
• To demonstrate that a single deep learning model with a variational perturbation module can surpass operational NWP systems like ECMWF IFS in accuracy for key surface variables.

Proposed method

• The model uses a SwinRNN backbone based on Swin Transformer blocks with multi-scale decoder and residual prediction to enable long-range temporal and spatial modeling.
• A variational autoencoder (VAE)-inspired perturbation module learns a time-variant multivariate Gaussian distribution over stochastic latent variables from historical data.
• Ensemble forecasts are generated by sampling noise from the learned latent distribution and injecting it into the hidden features of the SwinRNN at inference time.
• Four perturbation strategies are compared: fixed distribution (standard normal), learned distribution (SwinVRNN), Monte Carlo dropout, and multi-model ensemble.
• The perturbation module jointly optimizes reconstruction loss and KL divergence to ensure both accurate prediction and meaningful uncertainty representation.
• The model is trained and evaluated on the WeatherBench dataset using standard metrics for medium-range weather forecasting.

Experimental results

Research questions

• RQ1 Can a data-driven model with a learned, spatially adaptive perturbation strategy outperform fixed-distribution and MC dropout methods in ensemble forecasting accuracy?
• RQ2 Does joint optimization of deterministic prediction and stochastic uncertainty modeling lead to better forecast skill and more reliable ensemble spread?
• RQ3 Can a single deep learning model with a variational perturbation module surpass the operational ECMWF IFS model in forecast accuracy for key surface variables?
• RQ4 How does the inclusion of spatial covariance in the learned perturbation distribution affect forecast performance compared to isotropic or diagonal noise?

Key findings

• SwinVRNN achieves superior forecast accuracy and more reasonable ensemble spread compared to fixed distribution and MC dropout methods, due to joint optimization of prediction and uncertainty modeling.
• The learned distribution perturbation method significantly improves ensemble performance, particularly in capturing spatial variability of atmospheric uncertainty across different regions and variables.
• SwinVRNN outperforms the operational ECMWF IFS model in 2-m temperature (T2M) and 6-hourly total precipitation (TP) forecasts at all lead times up to five days.
• The multi-model ensemble of SwinVRNN further enhances forecast skill, demonstrating the scalability of the approach.
• The model shows consistent improvements across all evaluated lead times, including within the training period, indicating robust generalization.
• The perturbation module effectively models complex spatial correlations in atmospheric variables, as evidenced by the learned covariance matrix's structure and performance gains.