[Paper Review] FuXi-ENS: A machine learning model for medium-range ensemble weather forecasting

Xiaohui Zhong, Lei Chen+8 more|arXiv (Cornell University)|May 9, 2024

Meteorological Phenomena and Simulations9 citations

TL;DR

FuXi-ENS is a probabilistic ML model that delivers 6-hourly global ensemble weather forecasts up to 15 days at 0.25° resolution, outperforming ECMWF ensemble in most cases across numerous metrics.

ABSTRACT

Ensemble forecasting is crucial for improving weather predictions, especially for forecasts of extreme events. Constructing an ensemble prediction system (EPS) based on conventional NWP models is highly computationally expensive. ML models have emerged as valuable tools for deterministic weather forecasts, providing forecasts with significantly reduced computational requirements and even surpassing the forecast performance of traditional NWP models. However, challenges arise when applying ML models to ensemble forecasting. Recent ML models, such as GenCast and SEEDS model, rely on the ERA5 EDA or operational NWP ensemble members for forecast generation. Their spatial resolution is also considered too coarse for many applications. To overcome these limitations, we introduce FuXi-ENS, an advanced ML model designed to deliver 6-hourly global ensemble weather forecasts up to 15 days. This model runs at a significantly increased spatial resolution of 0.25 extdegree, incorporating 5 atmospheric variables at 13 pressure levels, along with 13 surface variables. By leveraging the inherent probabilistic nature of Variational AutoEncoder (VAE), FuXi-ENS optimizes a loss function that combines the CRPS and the KL divergence between the predicted and target distribution, facilitating the incorporation of flow-dependent perturbations in both initial conditions and forecast. This innovative approach makes FuXi-ENS an advancement over the traditional ones that use L1 loss combined with the KL loss in standard VAE models for ensemble weather forecasting. Results demonstrate that FuXi-ENS outperforms ensemble forecasts from the ECMWF, a world leading NWP model, in the CRPS of 98.1% of 360 variable and forecast lead time combinations. This achievement underscores the potential of the FuXi-ENS model to enhance ensemble weather forecasts, offering a promising direction for further development in this field.

Motivation & Objective

• Motivate reducing computational cost of ensemble forecasts while maintaining or improving forecast quality.
• Develop an ML-based ensemble forecasting system that yields high-resolution, probabilistic forecasts.
• Incorporate uncertainty quantification through a variational autoencoder framework and perturbation mechanisms.
• Demonstrate performance gains over ECMWF ensemble across multiple variables and lead times.

Proposed method

• Use ERA5 data at 0.25° to train a probabilistic FuXi-ENS model that outputs ensembles by sampling from a learned Gaussian perturbation distribution.
• Two main components: a perturbation model (VAE-based) producing a Gaussian N(Θp) and a forecasting model generating forecasts from perturbed initial conditions.
• Train with a loss combining CRPS and KL divergence to align predicted and target distributions, with CRPS shown to outperform L1 loss for ensemble forecasting.
• Generate ensemble members by sampling from the perturbation Gaussian and autoregressively forecasting, enabling 6-hourly forecasts up to 15 days.
• Employ knowledge distillation to align perturbation model distributions with target distributions during training.
• Implement an autoregressive curriculum training regime on 8 Nvidia A100 GPUs with the AdamW optimizer.

Experimental results

Research questions

• RQ1 Can an ML-based ensemble forecast at high spatial resolution (0.25°) outperform a leading NWP ensemble (ECMWF) in probabilistic metrics?
• RQ2 Does integrating CRPS with KL divergence in a VAE-based perturbation framework improve ensemble forecast skill and uncertainty representation?
• RQ3 How does FuXi-ENS perform across deterministic and ensemble metrics for mid-range forecasts up to 15 days?
• RQ4 Is FuXi-ENS computationally efficient enough for operational-like production of large ensembles?

Key findings

• FuXi-ENS outperforms the ECMWF ensemble in 360 of 360 CRPS-related comparisons (98.1%) across 60 lead times and 6 variables.
• FuXi-ENS provides superior ensemble mean performance and extreme-event forecasting for multiple variables including WS850, WS10M, MSL, and T2M.
• For Z500, FuXi-ENS surpasses ECMWF after day 4 and remains better from days 3 to 15; for T850, FuXi-ENS is better across most lead times.
• The model runs at 0.25° resolution and 6-hourly steps up to 15 days, with around 10 seconds per 15-day forecast per member on an Nvidia A100 GPU.
• FuXi-ENS yields comparable ensemble spread to ECMWF and shows underdispersion that gradually approaches unit spread-skill ratio with lead time.
• Extreme rainfall forecasts (Brier Score) show FuXi-ENS outperforming ECMWF in 96.3% of 1440 cases across percentiles, variables, and lead times.