A Space-Time Transformer for Precipitation Nowcasting

AI summary

Overview

Research area: AI weather prediction (AI-WP), specifically precipitation nowcasting — computer vision (video transformers) applied to satellite meteorology.

Technical level: Advanced. The paper assumes familiarity with transformer encoders, patch tokenization, class-weighted cross-entropy, and probabilistic verification metrics such as CRPS.

Scope: One-sentence summary: the paper adapts a video transformer into a model called SaTformer for predicting cumulative rainfall from satellite imagery, treating the problem as classification rather than regression, and reports a first-place finish on the NeurIPS Weather4Cast 2025 "Cumulative Rainfall" challenge.

What This Paper Is About

Numerical weather prediction (NWP) models are computationally demanding and perform poorly at short lead times, which limits their usefulness for nowcasting. Most recent AI-weather prediction work focuses on medium-range, global forecasting, leaving precipitation nowcasting at continental and regional scales comparatively unexplored. This paper's goal is to adapt a video transformer to predict future rainfall intensities from low-resolution HRIT satellite radiances, and to work out which techniques keep such a model skillful on a heavily skewed, fat-tailed precipitation dataset.

Key Contributions

1. SaTformer, a video transformer adapted for precipitation nowcasting. The model takes one hour of 11-band HRIT satellite imagery and uses full space-time self-attention — all tokens attend to all other tokens across both time and space — with a class token aggregating information for prediction.
2. Reformulating nowcasting as classification. Instead of regressing rainfall intensity directly, the target space is partitioned into n non-overlapping bins of step size δ, and the model predicts one-hot categorical labels over those bins with a softmax over a single-layer MLP head.
3. A class-weighted cross-entropy loss. Each class is weighted by its inverse frequency in the training set, to counter the strong bias toward low/no-rain events in the target distribution.
4. Ablations and a competition result. The paper ablates loss weighting, three attention variants, and the number of target bins, and reports that its best model scored first in the Weather4Cast 2025 "Cumulative Rainfall" challenge with a CRPS of 3.135. Code and model weights are released at github.com/leharris3/satformer.

Main Findings

• First place on the challenge: The best-performing SaTformer achieves a CRPS score of 3.135 on the Weather4Cast 2025 "Cumulative Rainfall" challenge set, which the authors report as first place in that challenge.
• Loss weighting is decisive: Without class-frequency re-weighting, the model "quickly overfits to low/no-rainfall events," scoring a bin-weighted Top-3 accuracy of 0.076 and bin-weighted CRPS of 6.91. Adding the weighted loss raises these to 0.272 and 2.64, respectively.
• Full space-time attention wins: Comparing attention variants on the W4CC validation set, space-then-time gives BW-Top-3 0.250 / BW-CRPS 4.39, time-then-space gives 0.214 / 3.39, and full space-time attention gives the best result at 0.272 / 2.64.
• Bin count matters, and more bins is not monotonically better: CRPS on the challenge set is 14.181 with 4 bins, 5.987 with 8, 4.293 with 16, 3.898 with 32, 3.135 with 64, 3.610 with 128, 5.312 with 256, and 4.783 with 512. The authors note that finer bins should help in theory but in practice lead to degenerate solutions as class representation becomes sparse; 64 classes is the reported empirical middle ground.
• Minimal modification is needed: Video transformers perform surprisingly well on this nowcasting task with only modest task-specific changes, and full self-attention is affordable here because the inputs are small — 4 input frames over a 32×32 pixel area.
• The data are strongly imbalanced: The paper's Figure 2 shows the training-set target distribution is heavily skewed toward low/no-rain events, with large gaps in sample coverage at the extremes.

Methodology in Plain English

The model starts from a video-transformer design (the authors describe it as derived step-by-step from a TimeSformer baseline) and adapts it to weather data:

• Input and tokenization. Each frame in the input sequence is cut into non-overlapping patches, each patch is flattened and projected into a token, learnable positional embeddings are added, and a randomly initialized class token is prepended to the sequence.
• Encoder. The token sequence passes through L transformer encoder blocks implemented with multi-headed attention. In the authors' variant, attention is full over space and time, so every token can look at every other token in the sequence.
• Prediction. The class token is spliced out of the final encoder layer and passed to a single-layer MLP head, which produces a softmax probability distribution over the target classes.
• Turning rainfall into classes. The regression label provided by the organizers is converted to a bin index by rounding (y_reg − D_ymin)/δ, and the model is trained to predict that one-hot bin.
• Weighted loss. The categorical cross-entropy loss is reweighted by the inverse frequency of each class in the training set.
• Data and training. Training uses 7 region-specific datasets of paired HRIT radiances and OPERA precipitation fields. Each training step samples a random region, a random 32×32 pixel (512×512 km) HRIT crop over a random four-consecutive-frame (1 hour) window, and a 32×32 pixel (64×64 km) OPERA sample centered on the same point spanning the following 16 frames (4 hours). Inputs are normalized roughly to [0, 1] using training-set statistics. The configuration uses input height and width 32, 11 channels, 64 classes, patch size 4, hidden dimension 512, head dimension 64 with 8 heads, and L = 12 encoder blocks. Training runs for 200 epochs and 25,000 steps with Adam at a learning rate of 1e-5, an effective batch size of 128, on 1×A6000 GPU.
• Evaluation. Models are scored with continuous ranked probability score (CRPS), which compares the predicted cumulative distribution function against the ground-truth observation; checkpoints with the lowest average validation loss are evaluated.

Why This Matters

Impact on research. The paper provides evidence that video transformers transfer to precipitation nowcasting with only minimal modification, giving the AI-weather community a strong, simple baseline at regional scale where most prior AI-WP work has focused on global medium-range forecasting. It also highlights a general problem for AI-WP: fat-tailed, imbalanced precipitation data can push models toward degenerate "always predict no rain" solutions unless calibration or class-regularization techniques are used. The authors explicitly frame taming fat-tailed datasets as a critical open challenge.

Real-world applications (the paper motivates these by noting that extreme precipitation creates risk to life, property, and commerce, and that the public and decision-makers need timely, accurate rainfall forecasts at convective-allowing scales):

• Short-lead-time rainfall guidance for forecasters, complementing numerical models that struggle at short lead times.
• Input to flood and extreme-weather warning workflows, given that extreme events are predicted to increase in frequency and severity.
• Regional-scale rainfall estimation where high-resolution radar data are sparse — the competition regions are exactly those where OPERA data are sparse.
• Any downstream decision process that needs probabilistic rainfall estimates rather than a single deterministic number, since the model outputs a full distribution and is scored with CRPS.

Industry relevance. Weather services, insurers, agriculture, transportation, and disaster-response organizations all depend on nowcasting. Because AI-WP pipelines can run at sub-hour resolution and in large ensembles, and are cheaper to run than physically based NWP pipelines, models like SaTformer point toward operational, real-time guidance that is inexpensive at inference time.

Future Directions

• Dense, video-like nowcasts. The authors state plainly that significant architectural modification would be required for SaTformer to produce dense nowcasts over space and time, and suggest graph networks or diffusion-based approaches that are better suited to generative prediction.
• Sub-quadratic attention for larger inputs. Full self-attention costs quadratic time in sequence length, so extending to longer satellite videos likely requires decoupled space-time attention, sparse self-attention, or similar efficient variants — and researchers applying transformers to higher-dimensional weather data will "almost certainly" need sub-quadratic methods.
• Self-supervised compression of video inputs. The authors propose using a self-supervised encoder to compress video inputs before the tokenization stage.
• Better handling of fat-tailed data. The paper suggests data augmentation techniques or alternative task formulations as routes to improve on these results, and flags imbalanced learning in AI-WP as a continuing challenge.

Target Audience

This paper is most useful to machine learning researchers working on weather and climate applications, especially those familiar with video transformers or building entries for forecasting competitions such as Weather4Cast. It also suits meteorologists and operational forecasting groups interested in what current AI-WP methods can and cannot do at nowcasting scale, and practitioners who need a practical recipe — classification reformulation plus class-weighted loss — for training on highly imbalanced geophysical data. Readers without a background in transformer architectures, probabilistic forecast verification, or the Weather4Cast benchmark will find the method sections demanding.

Authors’ abstract

Until recently, numerical weather prediction (NWP) models have stood rivalless in operational forecasting despite a few limitations. Namely, physically-based models are computationally demanding and struggle at short lead times, reducing their utility for nowcasting. Motivated by these shortcomings, recent work proposes AI-weather prediction (AI-WP) alternatives that emulate analysis data with neural networks. While these data-driven approaches have achieved high skill for medium-range forecasting-applications of AI-WP to precipitation and to nowcasting are less explored. To these ends, this paper discusses \textit{SaTformer}: a video transformer adapted for precipitation nowcasting. To ameliorate some problems related to what is essentially a fat-tailed regression task, we find it prudent to formulate nowcasting as a classification problem and employ a frequency-weighted loss. This straightforward approach scored first on the NeurIPS Weather4Cast 2025 ``Cumulative Rainfall'' challenge. Code and model weights are available: \texttt{\href{github.com/leharris3/w4c-25}{github.com/leharris3/satformer}}.