PreDiff: Precipitation Nowcasting with Latent Diffusion Models

Zhihan Gao; Xingjian Shi; Boran Han; Hao Wang; Xiaoyong Jin; Danielle Maddix; Yi Zhu; Mu Li; Yuyang Wang

Source-linked AI summary

PreDiff: Precipitation Nowcasting with Latent Diffusion Models

Zhihan Gao, Xingjian Shi, Boran Han, Hao Wang, Xiaoyong Jin, Danielle Maddix, Yi Zhu, Mu Li, Yuyang Wang

arXiv:2307.10422v2Jul 19, 2023cs.LGcs.AIcs.CV

TL;DR

PreDiff targets probabilistic precipitation nowcasting when point forecasts blur multiple plausible futures and may disregard physical constraints. It combines conditional latent diffusion with plug-in knowledge alignment that adjusts denoising transitions, achieving state-of-the-art results across N-body MNIST and SEVIR while retaining important evaluation caveats.

Problem
Data-driven Earth-system forecasters often produce blurry point predictions or physically implausible samples because they struggle with uncertainty and lack domain constraints.
Method
PreDiff uses a conditional latent diffusion model and a separate knowledge alignment mechanism that estimates constraint deviations at each denoising step and adjusts the transition distribution.
Results
PreDiff achieves state-of-the-art performance on N-body MNIST and state-of-the-art perceptual quality on SEVIR, with energy-aligned forecasts reaching an FVD score of 4.063.
Takeaways & Limitations
The results support knowledge-aligned probabilistic forecasting for chaotic synthetic dynamics and real-world precipitation nowcasting.
Takeaways & Limitations
Benchmark datasets and evaluation metrics remain immature, and their alignment with expert judgment remains an open challenge.

Abstract

from arXiv · show

Earth system forecasting has traditionally relied on complex physical models that are computationally expensive and require significant domain expertise. In the past decade, the unprecedented increase in spatiotemporal Earth observation data has enabled data-driven forecasting models using deep learning techniques. These models have shown promise for diverse Earth system forecasting tasks but either struggle with handling uncertainty or neglect domain-specific prior knowledge, resulting in averaging possible futures to blurred forecasts or generating physically implausible predictions. To address these limitations, we propose a two-stage pipeline for probabilistic spatiotemporal forecasting: 1) We develop PreDiff, a conditional latent diffusion model capable of probabilistic forecasts. 2) We incorporate an explicit knowledge alignment mechanism to align forecasts with domain-specific physical constraints. This is achieved by estimating the deviation from imposed constraints at each denoising step and adjusting the transition distribution accordingly. We conduct empirical studies on two datasets: N-body MNIST, a synthetic dataset with chaotic behavior, and SEVIR, a real-world precipitation nowcasting dataset. Specifically, we impose the law of conservation of energy in N-body MNIST and anticipated precipitation intensity in SEVIR. Experiments demonstrate the effectiveness of PreDiff in handling uncertainty, incorporating domain-specific prior knowledge, and generating forecasts that exhibit high operational utility.

1 Introduction

Precipitation nowcasting is important but difficult because Earth-system variability creates multiple plausible futures that point-estimation methods often blur. PreDiff addresses this gap with probabilistic diffusion forecasting and knowledge alignment for domain-consistent predictions.

Motivation: Precipitation nowcasting supports near-term decision-making, but accurately forecasting future rainfall remains challenging for data-driven algorithms.The task concerns rainfall forecasts over horizons such as 0–6 hours.
Limitations of existing methods: Point-estimation methods trained with pixel-wise losses cannot capture multiple plausible futures, producing blurry forecasts under complex climatic variability.Small initial-condition differences can lead to substantially divergent outcomes.
Diffusion models: Diffusion models provide likelihood-based generative frameworks that better handle uncertainty without GAN-style mode collapse or training instability.They also offer advantages over autoregressive image generation in speed, resolution, and quality.
Knowledge-aware forecasting: Existing diffusion applications to precipitation nowcasting remain early, while purely data-centric approaches lack domain constraints and physical controls.Prior approaches incorporate knowledge through architectural changes or additional training losses.
Contributions: PreDiff combines a latent diffusion model based on Earthformer-UNet with a second-stage knowledge alignment mechanism that guides sampling without modifying the trained model.The pipeline first captures data semantics, then injects Earth-system prior knowledge during sampling.
Contributions: PreDiff achieves state-of-the-art performance on N-body MNIST and state-of-the-art perceptual quality on SEVIR.These results support its use for probabilistic forecasting with domain-specific guidance.

2 Method

PreDiff represents conditional future distributions in a lower-dimensional latent space and generates forecasts by iterative denoising conditioned on observed frames. A separate knowledge alignment network estimates constraint violations and adjusts denoising transitions to incorporate domain-specific prior knowledge.

Latent diffusion: PreDiff trains a frame-wise VAE and then a conditional diffusion model in its lower-dimensional latent representation for more efficient forecasting.The encoder maps pixel-space frames to latents, while the decoder reconstructs predicted frames.
Latent diffusion: The conditional latent diffusion model factorizes future generation into denoising transitions conditioned on the encoded observation context.Initial latent noise is progressively transformed into the future latent sequence.
Latent diffusion: Earthformer-UNet models the high-dimensional spatiotemporal transition distribution using hierarchical cuboid attention and concatenated context and noisy future latents.The architecture is derived from Earthformer and is intended to capture extensive Earth-system dependencies.
Knowledge alignment: Knowledge alignment constrains forecasts through a domain-defined function F and quantifies violations by comparing F(bx,y) with the desired value F0(y).For isolated physical systems, this formulation can enforce conservation of energy relative to the last observation.
Knowledge alignment: A knowledge alignment network estimates constraint values from intermediate noisy latents, and each denoising transition is adjusted to reduce expected constraint violations.The adjustment shifts the predicted transition mean using a guidance scale and the gradient of the estimated violation.
Knowledge alignment: Knowledge alignment is trained independently from the latent diffusion model and applied as a plug-in during inference without changing the trained VAE or LDM.No inference sampling process is required to train the alignment network.

3 Experiments

Experiments on N-body MNIST and SEVIR show that PreDiff supports probabilistic spatiotemporal forecasting, while knowledge alignment improves compliance with energy or intensity guidance.

N-body MNIST: PreDiff outperforms nine deterministic and probabilistic baselines on N-body MNIST across conventional frame metrics and FVD.Its Earthformer-UNet backbone improves spatiotemporal modeling, producing sharp digits in accurate positions.
Knowledge Alignment: Energy Conservation: PreDiff-KA substantially improves E.MSE and E.MAE over PreDiff and all baselines while retaining FVD = 4.063.Lower energy errors indicate better compliance with conservation of energy.
Knowledge Alignment: Energy Conservation: PreDiff-KA achieves lower energy errors than even the target sequences, overcoming defects caused by noise in the simulated training data.The target data may violate conservation because of discretization errors, whereas purely data-oriented methods retain larger errors.
Knowledge Alignment: Energy Conservation: Pixel-wise MSE, MAE, and SSIM are less informative for PreDiff-KA because energy correction can shift digit positions while preserving sharp visual quality.Small position changes can create large pixel-wise errors even when the generated digit remains clear.
SEVIR: On SEVIR, PreDiff is strongest among probabilistic methods, with better FVD and CSI-pool16 despite weaker per-pixel CSI scores.The authors attribute this pattern to better capture of local data patterns and distributions, which can support weather-event recognition.
Knowledge Alignment: Anticipated Average Intensity: Knowledge alignment guides SEVIR forecasts toward anticipated precipitation intensities, enabling qualitative generation of rainstorm and drought scenarios.PreDiff-KA preserves distributional fidelity, with FVD 34.18 versus 33.05 for PreDiff.

4 Conclusions and Broader Impacts

PreDiff combines latent diffusion with knowledge alignment for Earth system forecasting, while acknowledging unresolved challenges in evaluation, scientific integration, and data scarcity. The paper positions this approach as a promising advance but calls for further collaboration and robustness improvements.

PreDiff introduces a latent diffusion model and knowledge-alignment mechanism for Earth system forecasting, with reported state-of-the-art performance on N-body MNIST and SEVIR.The knowledge-alignment mechanism guides forecasts toward domain-specific prior knowledge without modifying the trained PreDiff model.
Earth system forecasting research still faces maturing benchmarks and metrics, incomplete integration of physical principles, and scarcity of high-quality data.
Future progress requires improved benchmarking, stronger incorporation of scientific knowledge, and collaboration between AI researchers and domain experts.

B Implementation Details

All experiments run on NVIDIA A10G GPUs, and PreDiff, knowledge-alignment networks, and baselines fit on a single GPU without gradient checkpointing or model parallelization.

All models fit on a single NVIDIA A10G GPU without gradient checkpointing or model parallelization.The GPUs provide 24GB of memory.

B.1 PreDiff

PreDiff uses frame-wise variational autoencoders, an Earthformer-UNet latent diffusion backbone, and Earthformer-based knowledge-alignment networks. The implementation includes dataset-specific latent dimensions, denoising-step embeddings, adversarially trained autoencoders, and separately configured optimizers and baselines.

Frame-wise autoencoder: The frame-wise VAE encodes N-body MNIST frames from 64×64×1 to 16×16×3 and SEVIR frames from 128×128×1 to 16×16×4.
Latent diffusion model: The Earthformer-UNet concatenates encoded context and noisy future latents along the temporal dimension, then predicts the previous denoising latent.The denoising-step embedding is added before each cuboid-attention block.
Knowledge alignment networks: The knowledge-alignment network uses an Earthformer encoder with pooling to predict constraint-related quantities from the denoising step and concatenated context and noisy latent.
Optimization: Frame-wise VAEs use Adam, while the latent Earthformer-UNet and knowledge-alignment networks use AdamW with detailed configurations provided separately.
Baselines: Baseline implementations follow official configurations with dataset-specific modifications and tuning of learning rates, schedulers, and working resolutions.

C Derivation of the Approximation to Knowledge Alignment Guidance

The derivation rewrites knowledge-alignment guidance with a normalization constant and approximates the guided denoising transition using a Taylor expansion around the unguided mean. The resulting Gaussian transition retains the original covariance while shifting its mean according to the guidance term.

The derivation rewrites the guided transition using a normalization constant Z before taking logarithms and applying the approximation results.
For brevity, µθ and Σθ denote the conditional transition mean and covariance, while C_i denotes constants for i = 1, ..., 7.
A Taylor expansion around z_t = µθ provides the approximation in the limit of infinitely many diffusion steps.The limiting condition is expressed as ||Σθ||→0.
The guided transition is approximated by a Gaussian whose mean is shifted by Σθ relative to unguided denoising.This preserves the Gaussian form while incorporating knowledge-alignment guidance into the transition mean.

D.1 Quantitative Analysis of BIAS on SEVIR

The SEVIR BIAS analysis evaluates whether forecasts favor false alarms or misses across precipitation thresholds. PreDiff achieves the best threshold-wise BIAS, with scores consistently closest to 1, indicating less intensity bias than the compared models.

BIAS is computed from hits, misses, and false alarms after binarizing predictions at thresholds [16, 74, 133, 160, 181, 219].The metric assesses whether predictions favor false alarms or misses.
PreDiff achieves the best BIAS across thresholds, with scores consistently closest to 1 among the evaluated models.A BIAS score closer to 1 indicates less bias toward false alarms or misses.
Deterministic forecasting models tend to produce lower-intensity predictions by avoiding high-intensity signals under uncertainty.
Probabilistic baselines predict more possible high-intensity signals, even when this leads to lower CSI scores.

D.2 CSI at Varying Thresholds on SEVIR

The SEVIR CSI analysis compares deterministic and probabilistic forecasting methods across precipitation thresholds and spatial pooling scales. PreDiff remains competitive on standard CSI and outperforms baselines on pooled CSI, particularly at higher thresholds.

PreDiff outperforms baselines on CSI-pool4 and CSI-pool16, particularly at higher precipitation thresholds.The comparisons cover pooling scales of 4 × 4 and 16 × 16.
PreDiff achieves competitive CSI scores across thresholds when compared with ConvLSTM, Earthformer, and probabilistic baselines.
The evaluation reports CSI, CSI-pool4, and CSI-pool16 at thresholds [16, 74, 133, 160, 181, 219].
CSI is computed from averaged samples for each model, whereas the other metrics are averaged over individual sample scores.

E More Qualitative Results on N-body MNIST

The N-body MNIST qualitative results present example prediction sequences for several forecasting methods, including PreDiff and PreDiff-KA. The figures also compare predicted total energy with the final context-frame energy through E.MSE.

Figures 6–13 compare context and target sequences with predictions from ConvLSTM, Earthformer, VideoGPT, LDM, PreDiff, and PreDiff-KA.
E.MSE measures the average error between each prediction’s total energy and the total energy of the last context frame.Total energy is defined as the sum of kinetic and potential energy.
The qualitative examples span test sequences whose final frames contain digits 0, 1, 4, 7, and 8.
Red dashed lines mark the position of the indicated digit in the last frame to support visual judgment of the predictions.

F More Qualitative Results on SEVIR

Figures 14–19 present qualitative SEVIR test-set predictions comparing PreDiff and baselines, alongside PreDiff-KA forecasts guided by anticipated future intensity.

F More Qualitative Results on SEVIR: Figures 14–19 compare ConvLSTM, Earthformer, VideoGPT, LDM, PreDiff, and PreDiff-KA predictions against the context and target sequences.The comparison appears in subfigure (a) of each figure.
F More Qualitative Results on SEVIR: Subfigure (b) shows PreDiff-KA predictions conditioned on anticipated average future intensity values from µτ + nστ, with n ranging from 4 to −4.The displayed guidance levels are n = 4, 2, 0, −2, and −4.
F More Qualitative Results on SEVIR: Each figure contains example predictions from the SEVIR test set, with context and target sequences included for visual reference.Figures 14–19 provide several prediction sets rather than a numerical summary.