CSU-PCAST: A Dual-Branch Transformer Framework for medium-range ensemble Precipitation Forecasting

TL;DR

CSU-PCAST addresses the challenge of medium-range precipitation forecasting by leveraging a patch-based Swin Transformer backbone with temporal embeddings and a noise-augmented CLN to produce an autoregressive, 58-channel forecast at $0.25^\circ$ resolution. The model jointly predicts 57 non-precipitation variables and one precipitation channel via a dual-branch decoder, trained through staged optimization that combines $CRPS$ and $L_{log1pMSE}$ losses to capture uncertainty and magnitude. Compared to GEFS, CSU-PCAST delivers higher deterministic skill (CSI and RMSE) and superior probabilistic performance (lower $CRPS$ and $BS$) across lead times and rainfall thresholds, with pronounced advantages for moderate-to-heavy rainfall and extended forecasts. The results indicate that this deep-learning ensemble approach can complement, and in some cases surpass, traditional NWPs, offering sharper, more reliable probabilistic guidance suitable for hydrometeorological risk management.

Abstract

Accurate medium-range precipitation forecasting is crucial for hydrometeorological risk management and disaster mitigation, yet remains challenging for current numerical weather prediction (NWP) systems. Traditional ensemble systems such as the Global Ensemble Forecast System (GEFS) struggle to maintain high skill, especially for moderate and heavy rainfall at extended lead times. This study develops a deep learning-based ensemble framework for multi-step precipitation prediction through joint modeling of a comprehensive set of atmospheric variables. The model is trained on ERA5 reanalysis data at 0.25$^{\circ}$ spatial resolution, with precipitation labels from NASA's Integrated Multi-satellite Retrievals for Global Precipitation Measurement (GPM) constellation (IMERG), incorporating 57 input variables, including upper-air and surface predictors. The architecture employs a patch-based Swin Transformer backbone with periodic convolutions to handle longitudinal continuity and integrates time and noise embeddings through conditional layer normalization. A dual-branch decoder predicts total precipitation and other variables, with targeted freezing of encoder-decoder pathways for specialized training. Training minimizes a hybrid loss combining the Continuous Ranked Probability Score (CRPS) and weighted log1p mean squared error (log1pMSE), balancing probabilistic accuracy and magnitude fidelity. During inference, the model ingests real-time Global Forecast System (GFS) initial conditions to generate 15-day forecasts autoregressively. Evaluation against GEFS using IMERG data demonstrates higher Critical Success Index (CSI) scores at precipitation thresholds of 0.1 mm, 1 mm, 10 mm, and 20 mm, highlighting improved performance for moderate to heavy rainfall.

Figures (12)

• Figure 1: CSI scores of the precipitation forecast (6 hr) from CSU-PCAST model and GEFS at different lead times and precipitation intensities during Jan 2023 and Jul 2023: (a) Jan CSU-PCAST; (b) Jan GEFS; (c) Jan CSI difference between CSU-PCAST and GEFS; (d) Jul CSU-PCAST; (e) Jul GEFS; (f) Jul CSI difference between CSU-PCAST and GEFS
• Figure 2: RMSE of the precipitation forecast (6 hr) from the CSU-PCAST model and GEFS at different lead times， evaluated against IMERG precipitation: (a) RMSE between CSU-PCAST and GEFS during January 2023; (b) Relative RMSE difference between CSU-PCAST and GEFS during January 2023; (c) RMSE between CSU-PCAST and GEFS during July 2023; (d) Relative RMSE difference between CSU-PCAST and GEFS during July 2023
• Figure 3: Precipitation forecasts initialized at 2023-07-06 00UTC with a forecast lead time of 60 hours. The top row shows forecasts from the CSU-PCAST model (left) and GEFS (right), while the bottom panel shows IMERG ground-truth observations.
• Figure 4: Brier Score (BS) differences and CRPS for precipitation forecasts in January and July 2023, against IMERG. Panels (a.1–e.1) show BS differences at thresholds of 0.1, 1, 5, 10, and 20 mm during January 2023, and panels (a.2-e.2) show BS differences during Jul 2023, respectively. The blue horizontal line denotes the GEFS baseline (0), while the red curves represent the relative BS of the CSU-PCAST model compared to GEFS. Panel (f.1, f.2) shows the CRPS, where CSU-PCAST consistently outperforms GEFS across the full 15-day forecast horizon. Together, these results highlight the enhanced probabilistic skill of CSU-PCAST across precipitation thresholds and lead times.
• Figure 5: Panel a: overall architecture of CSU-PCAST. The model is composed of patch embedding, a U-Transformer backbone, and fully connected layers. Panel b: the details of Swin layers, the transformer block has 4 Swin layers, each Swin layer contains 12 Swin blocks.