SwinVRNN: A Data‐Driven Ensemble Forecasting Model via Learned Distribution Perturbation

Abstract The data‐driven approaches for medium‐range weather forecasting are recently shown to be extraordinarily promising for ensemble forecasting due to their fast inference speed compared to the traditional numerical weather prediction models. However, their forecast accuracy can hardly match the state‐of‐the‐art operational ECMWF Integrated Forecasting System (IFS) model. Previous data‐driven approaches perform ensemble forecasting using some simple perturbation methods, like the initial condition perturbation and the Monte Carlo dropout. However, their ensemble performance is often limited arguably by 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 the future states deterministically. Furthermore, to model the stochasticity in the prediction, we design a perturbation module following the Variational Auto‐Encoder paradigm to learn the multivariate Gaussian distributions of a time‐variant stochastic latent variable from the data. Ensemble forecasting can be easily performed by perturbing the model features leveraging the noise sampled from the learned distribution. We also compare four categories of perturbation methods for ensemble forecasting, that is, fixed distribution perturbation, learned distribution perturbation, MC dropout, and multi model ensemble. Comparisons on the WeatherBench data set show that the learned distribution perturbation method using our SwinVRNN model achieves remarkably improved forecasting accuracy and reasonable ensemble spread due to the joint optimization of the two targets. More notably, SwinVRNN surpasses operational IFS on the surface variables of the 2‐m temperature and the 6‐hourly total precipitation at all lead times up to 5 days (Code is available at https://github.com/tpys/wwprediction ).