Transformer-based bias correction & extreme-event calibration for S2S forecasts

A Space-Time Transformer post-processes subseasonal-to-seasonal ensemble forecasts — correcting bias, calibrating spread, and sharpening extremes — benchmarked with ACC, RMSE, CRPS, CRPSS and Brier skill.

GitHub repository Executable notebook Notebook source

+0.146
CRPSS vs raw (transformer)
2.34 K
test RMSE (best of all methods)
+0.24
CRPSS at week 6
71,185
trainable parameters

Test-set skill

Held-out test set (514-case dataset, 6 lead weeks on a 16×16 grid), evaluated on anomalies. Green = best in column; the transformer row is highlighted. Lower is better for RMSE/CRPS; for spread/err, 1.0 is calibrated.

methodRMSE (K)ACCCRPS (K)CRPSS vs rawCRPSS vs climspread/errBSS q90BSS q95BSS q98
RAW2.510+0.5821.468+0.000+0.0840.504-0.020-0.098-0.168
MeanVarDebias2.508+0.5801.328+0.095+0.1711.000+0.040-0.078-0.216
QuantileMapping2.426+0.6161.426+0.028+0.1100.526+0.025-0.026-0.061
Transformer2.337+0.6231.253+0.146+0.2181.240+0.087-0.006-0.088

The transformer leads on every aggregate score — deterministic (RMSE, ACC), probabilistic (CRPS, CRPSS) and moderate-extreme (BSS q90). MeanVarDebias fixes dispersion only; QuantileMapping fixes the marginal distribution but not dispersion; the transformer fixes both jointly and conditions on the large-scale state.

Results

crpss_by_lead
Probabilistic skill (CRPSS vs the raw model) by lead time. The transformer's advantage grows with lead — exactly where S2S forecasts are weakest.
rmse_by_lead
Deterministic error (RMSE) by lead time.
example_maps
An extreme case at week 4: the raw ensemble compresses the warm anomaly; the transformer restores its magnitude.
reliability
Reliability for the q90 heatwave event — proximity to the diagonal indicates calibrated probabilities.
rank_histogram
Rank histograms: the raw ensemble is U-shaped (under-dispersed); the transformer is far flatter (calibrated).
spread_error
Ensemble dispersion (spread/RMSE); 1.0 is perfectly calibrated.
extreme_bss
Extreme-event Brier skill vs climatology across thresholds.
gpd_tail
Generalised-Pareto heatwave return levels from the peaks-over-threshold fit.
training_curve
Training convergence of the NumPy transformer (tail-weighted CRPS + MSE).

Method

Each forecast case — ensemble-mean and ensemble-spread anomaly fields — is split into spatial patches; every (lead, patch) pair is a token and the model applies full space-time self-attention. The head emits a per-gridpoint mean correction and a per-token spread inflation, trained on a tail-weighted CRPS + mean-anchoring MSE objective. The reference model is implemented in pure NumPy with a hand-written autograd engine (gradient-checked against finite differences); a parity PyTorch version is provided for GPU training on real data.

Reproducibility: the figures and table above are produced on a physically-motivated synthetic S2S dataset so the pipeline runs anywhere with no downloads. The same code runs on real ERA5 + WeatherBench2 S2S forecasts (public, anonymous Zarr on Google Cloud) via s2s_bc.data.real.build_real_dataset; the NASA GEOS-S2S path is wired through the IRI SubX endpoint.