STLDM: Spatio-Temporal Latent Diffusion Model for Precipitation Nowcasting

AI summary

Overview

• Research area: Machine learning for weather forecasting, specifically precipitation nowcasting, spatiotemporal prediction, and latent diffusion models.
• Technical level: Intermediate. Readers should be comfortable with deep learning, diffusion models, and standard forecasting metrics.
• Scope: This paper introduces STLDM, an end-to-end latent diffusion model that separates precipitation nowcasting into deterministic forecasting and generative enhancement to improve both accuracy and visual realism.

What This Paper Is About

Precipitation nowcasting predicts near-future rainfall from radar observations, but deterministic models often produce blurry forecasts while generative models can produce sharp yet inaccurate ones. The paper proposes STLDM to combine a deterministic first estimate of future precipitation with a latent diffusion refinement, aiming for accurate and realistic nowcasts with faster inference. The method is evaluated on multiple real-world radar datasets and compared against both deterministic and diffusion-based baselines.

Key Contributions

1. Reformulates precipitation nowcasting as two sequential subtasks: deterministic Forecasting and generative Enhancement.
2. Proposes STLDM, a simple latent diffusion architecture with three modules: a Variational AutoEncoder, a Translator/conditioning network, and a Latent Denoising Network.
3. Trains the latent denoising network end-to-end with the VAE and conditioning network for precipitation nowcasting, rather than training each module independently.
4. Achieves state-of-the-art or competitive results on SEVIR, HKO-7, and MeteoNet across accuracy and perceptual metrics while using fewer sampling steps and faster inference than diffusion baselines.

Main Findings

• Dual-stage design improves accuracy and realism. STLDM uses a deterministic Translator for global motion and a latent diffusion model for detail, avoiding blurry deterministic outputs and inaccurate generative outputs.
• Strong benchmark performance. It achieves best or near-best scores on most metrics, especially on HKO-7, and the best LPIPS among compared models.
• Faster inference. STLDM uses 20 sampling steps and runs about 10x faster than DiffCast on SEVIR and about 40x faster on HKO-7 and MeteoNet.
• Constraint loss is critical. Removing the constraint loss L_C causes noisy first estimates and over-prediction, hurting CSI and HSS.
• End-to-end training improves perceptual quality. Strategy C achieves the best SSIM and LPIPS, while Strategy B is marginally better on some CSI/HSS scores.
• Not universally superior. On SEVIR, STLDM is 1.64% worse than DiffCast on CSI16-m; on MeteoNet, it drops 1.40–2.16% in SSIM, CSI-m, and HSS but gains 0.39–3.79% on other metrics.

Methodology in Plain English

STLDM encodes radar frames into a compressed latent space using a VAE. A Translator, based on gated spatio-temporal attention, predicts a rough future sequence in that latent space. A conditional latent diffusion model then starts from Gaussian noise and denoises it, guided by the rough prediction, to produce a refined latent sequence. The VAE decoder converts this refined latent sequence back into radar images. Training combines reconstruction loss, KL regularization, a constraint loss that keeps the final prediction aligned with the Translator's global motion trend, a prior loss, and a diffusion loss. Classifier-free guidance is used during sampling, and linearized spatial attention reduces computation. All modules are trained end-to-end.

Why This Matters

• Research impact: Shows that end-to-end latent diffusion can balance deterministic accuracy and generative realism in spatiotemporal forecasting, and that a simple two-stage reformulation can beat more complex diffusion pipelines.
• Real-world applications:
  • Severe weather warning: faster, sharper nowcasts can improve flash flood and storm alerts.
  • Aviation and transportation: better short-term precipitation forecasts support route planning and safety.
  • Emergency management: accurate 0–2 hour predictions help allocate response resources.
  • Agriculture and water management: near-term rainfall forecasts inform irrigation and drainage decisions.
• Industry relevance: Weather services, insurance, energy, logistics, and smart-city operations can benefit from lower-latency, higher-quality nowcasting models.

Future Directions

• Extend STLDM to higher-resolution radar data and longer lead times without large computational cost.
• Improve uncertainty calibration and ensemble diversity for operational probabilistic forecasting.
• Integrate physical constraints or numerical weather prediction guidance more directly.
• Investigate why STLDM is not uniformly best on all metrics and datasets, and adapt the architecture or loss weighting accordingly.

Target Audience

Machine learning researchers working on diffusion models and spatiotemporal prediction; meteorologists and weather forecasting engineers; graduate students in AI for weather and climate; and industry practitioners building operational nowcasting systems. Intermediate technical background in deep learning is recommended.

Authors’ 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.