STLDM: Spatio-Temporal Latent Diffusion Model for Precipitation Nowcasting

TL;DR

STLDM addresses precipitation nowcasting by coupling a deterministic Translator with a latent diffusion module that operates in the VAE latent space, formulating the task as Forecasting followed by Enhancement and training all components end-to-end. The approach yields high-precision, sharp forecasts while maintaining temporal coherence and achieving faster inference than pixel-space diffusion models, demonstrated across SEVIR, HKO-7, and MeteoNet. The key innovations include the end-to-end three-module architecture (VAE, Conditioning Network, Latent Denoising Network), the constraint loss $\mathcal{L}_{C}$ to align the first estimation with ground truth, and the use of CFG in latent space with a spatio-temporal attention-based translator. These properties enable accurate probabilistic nowcasting with practical real-time applicability for meteorological operations, especially when generating large ensembles.

Abstract

Precipitation nowcasting is a critical spatio-temporal prediction task for society to prevent severe damage owing to extreme weather events. Despite the advances in this field, the complex and stochastic nature of this task still poses challenges to existing approaches. Specifically, deterministic models tend to produce blurry predictions while generative models often struggle with poor accuracy. In this paper, we present a simple yet effective model architecture termed STLDM, a diffusion-based model that learns the latent representation from end to end alongside both the Variational Autoencoder and the conditioning network. STLDM decomposes this task into two stages: a deterministic forecasting stage handled by the conditioning network, and an enhancement stage performed by the latent diffusion model. Experimental results on multiple radar datasets demonstrate that STLDM achieves superior performance compared to the state of the art, while also improving inference efficiency. The code is available in https://github.com/sqfoo/stldm_official.

Figures (16)

• Figure 1: This demonstrates that deterministic models result in blurry predictions while generative models suffer from the issue of inaccurate predictions. Our proposed STLDM is capable of forecasting accurate predictions while maintaining a nice appearance.
• Figure 2: Top: Model architecture of the proposed STLDM, consisting of a Variational Autoencoder, ($\{\mathcal{E}, \mathcal{D}\}$), a Conditioning Network (aka Translator), $\Psi_{\theta}$, and a Spatio-Temporal Latent Denoising Network, $D_{\theta}$. Input radar frames are denoted as $X_{1:M}$; while the decoded outputs of both the final prediction after denoising from pure Gaussian Noise, $z_{1:N}^{T}$, and the first estimation, $\bar{z}_{1:N}$, are denoted as $\hat{Y}_{1:N}$ and $\overline{Y}_{1:N}$ respectively. Bottom: Overall architecture of $D_{\theta}$ (Yellow Box) and the details of its sub-modules (Green Box). "L-Spatial Attention" stands for Linearized Spatial Attention.
• Figure 3: A set of sample predictions on the SEVIR test set. From top to bottom: Input, Ground truth, SimVP, PreDiff, DiffCast, and STLDM. The red region of the last prediction frame is zoomed in for a clearer comparison.
• Figure 4: Illustration of the implemented encoder, $\mathcal{E}$, decoder, $\mathcal{D}$ and conditioning network, $\Psi_{\theta}$.
• Figure 5: Gated Spatio-Temporal Attention (gSTA) module