Two-stage Rainfall-Forecasting Diffusion Model

TL;DR

The paper tackles the challenge of accurate long-term rainfall forecasting and spatial realism by proposing a two-stage diffusion framework (TRDM) that separates temporal evolution and spatial reconstruction. The first stage uses a 3D diffusion model to predict a low-resolution rainfall sequence from four context frames, producing 16 frames, and the second stage applies two diffusion-based super-resolution paths (SSR in pixel space and LSR in latent space) to generate high-resolution outputs. Evaluations on Swedish radar and MRMS datasets show state-of-the-art performance compared to DGMRskillful and PySTEPS, with strongest performance emerging after roughly 20 to 80 minutes, especially in medium-to-long horizons. The work demonstrates that decomposing rainfall forecasting into temporal diffusion and spatial reconstruction yields improved accuracy and realism, and provides open-source code for reproducibility.

Abstract

Deep neural networks have made great achievements in rainfall prediction.However, the current forecasting methods have certain limitations, such as with blurry generated images and incorrect spatial positions. To overcome these challenges, we propose a Two-stage Rainfall-Forecasting Diffusion Model (TRDM) aimed at improving the accuracy of long-term rainfall forecasts and addressing the imbalance in performance between temporal and spatial modeling. TRDM is a two-stage method for rainfall prediction tasks. The task of the first stage is to capture robust temporal information while preserving spatial information under low-resolution conditions. The task of the second stage is to reconstruct the low-resolution images generated in the first stage into high-resolution images. We demonstrate state-of-the-art results on the MRMS and Swedish radar datasets. Our project is open source and available on GitHub at: \href{https://github.com/clearlyzerolxd/TRDM}{https://github.com/clearlyzerolxd/TRDM}.

Figures (5)

• Figure 1: The framework of TRDM. The first part involves a low-resolution diffusion model for prediction, which is comprised of a 3D denoising diffusion model; it uses the input of four conditional frames (20 min) to predict the low-resolution state ($32 \times 32$) for the next 16 frames(80 min). The second part deals with a super-resolution diffusion model, which is constructed using a 2D denoising diffusion model, it aims to reconstruct low-resolution frames to high-resolution images ($256 \times 256$).
• Figure 2: Illustration of SSR inference processing.The diffusion process is performed in pixel space, and the dimension of $X_i$ are kept consistent with the resolution of the image.
• Figure 3: Illustration of LSR inference processing.The diffusion process is performed in the latent space, and the dimension($4\times32\times32$) of $Z_i$ is consistent with the latent variable resolution saved using the encoder.
• Figure 4: Schematic diagram of training prediction denoising network (3D-Unet)
• Figure 5: Challenging Precipitation Event Case Study start 2021-01-10 at 03:55 Sweden. Individual predictions at lead times of T+5, T+20,T+40,T+60 and T+80 minutes for different models.The continuous ranked probability score (CRPS) and Field Skill Score (FSS) for each moment displayed in the bottom-left corner.