RainDiff: End-to-end Precipitation Nowcasting Via Token-wise Attention Diffusion

TL;DR

Precipitation nowcasting is challenged by chaotic spatio-temporal dynamics and the limitations of latent autoencoder–based or purely pixel-space diffusion models. The authors introduce RainDiff, an end-to-end diffusion framework that integrates Token-wise Attention across all spatial resolutions and a Post-attention mechanism within a cascaded spatio-temporal encoder, removing the need for latent autoencoders and improving long-horizon fidelity. The model decomposes forecasting into a deterministic mean $oldsymbol{bc}$ plus a stochastic residual $oldsymbol{r}$ modeled by a diffusion process conditioned on global context $h$, with residuals partitioned into autoregressive segments $s_m$ and fused as $oldsymbol{ y}=oldsymbol{bc}+oldsymbol{r}$ at inference. Across four datasets (Shanghai Radar, SEVIR, MeteoNet, CIKM), RainDiff achieves superior localization, perceptual quality, and robustness compared to both deterministic and probabilistic baselines, demonstrating scalability and practical potential for real-time nowcasting.

Abstract

Precipitation nowcasting, predicting future radar echo sequences from current observations, is a critical yet challenging task due to the inherently chaotic and tightly coupled spatio-temporal dynamics of the atmosphere. While recent advances in diffusion-based models attempt to capture both large-scale motion and fine-grained stochastic variability, they often suffer from scalability issues: latent-space approaches require a separately trained autoencoder, adding complexity and limiting generalization, while pixel-space approaches are computationally intensive and often omit attention mechanisms, reducing their ability to model long-range spatio-temporal dependencies. To address these limitations, we propose a Token-wise Attention integrated into not only the U-Net diffusion model but also the spatio-temporal encoder that dynamically captures multi-scale spatial interactions and temporal evolution. Unlike prior approaches, our method natively integrates attention into the architecture without incurring the high resource cost typical of pixel-space diffusion, thereby eliminating the need for separate latent modules. Our extensive experiments and visual evaluations across diverse datasets demonstrate that the proposed method significantly outperforms state-of-the-art approaches, yielding superior local fidelity, generalization, and robustness in complex precipitation forecasting scenarios.

Figures (9)

• Figure 1: A visualization from the SEVIR dataset shows that, at the longest forecast horizon, RainDiff avoids oversmoothed outputs and better preserves weather fronts compared to the state-of-the-art DiffCast yu2024DiffCast, resulting in closer alignment with the ground truth.
• Figure 2: Overall architecture of our precipitation nowcasting framework RainDiff. Given an input sequence $X_{0}$, a deterministic predictor $\mathcal{F}_{\theta_1}$ outputs a coarse prediction $\mu$. The concatenation of $X_{0}$ and $\mu$ is encoded by a cascaded spatio-temporal encoder $\mathcal{F}_{\theta_3}$ to yield conditioning features $h$, refined by Post-attention. A diffusion-based stochastic module $\mathcal{F}_{\theta_2}$ equipped with Token-wise Attention at all resolutions in pixel space predicts residual segments $\hat{r}$ autoregressively, where the denoising process is conditioned on $h$ and the predicted segments. This design captures rich contextual relationships and inter-frame dependencies in the radar field while keeping computation efficient.
• Figure 2: Ablation studies of RainDiff on the Shanghai Radar dataset: (a) individual components and (b) attention mechanisms in the spatio-temporal encoder.
• Figure 3: Qualitative comparison with existing works on the Shanghai Radar dataset, where the reflectivity range is on the top right.
• Figure 4: Frame-wise CSI and HSS for various methods on the Shanghai Radar dataset.