Technical Summary: Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

Problem Statement

Accurate short-term weather forecasting (nowcasting) is critical for safety and economic sectors like energy, agriculture, and aviation. While data-driven machine learning models offer speed advantages over traditional Numerical Weather Prediction (NWP), existing approaches face significant limitations:

1. Deterministic Blurring: Deterministic models often average over multiple plausible futures, resulting in blurry predictions that fail to capture the sharpness of extreme weather events.
2. Uncertainty Quantification: Safety-critical applications require understanding the range of plausible outcomes (aleatoric uncertainty), which deterministic models cannot provide.
3. Limitations of Current Generative Approaches:
   • GANs: Suffer from training instability and mode collapse, limiting their ability to cover rare, critical long-tail events.
   • Diffusion/Flow Models: While effective, they are computationally expensive when operating in high-dimensional pixel space.
   • Two-Stage Architectures: Current state-of-the-art methods often use a two-stage setup (compression autoencoder followed by a generative model). However, standard compression is inherently lossy; errors imperceptible in natural images can correspond to substantial shifts in precipitation. Furthermore, these compressors often rely on adversarial or perceptual losses, which can reintroduce instability and degrade the coverage of extreme events. Additionally, deterministic decoders in these pipelines fail to quantify uncertainty during the decoding process.

Methodology

The authors propose FREUD (FRamewise Encoder and United Decoder), a novel framework designed for efficient, probabilistic precipitation nowcasting. The approach consists of two main stages: a generative compression stage and a latent-space forecasting stage.

1. FREUD Architecture (First Stage)

The core innovation is a rectified flow-based encoder-decoder tailored for weather data:

• Frame-wise Encoder: Unlike sequence-level encoders that risk information leakage from future to past, FREUD uses a lightweight transformer to encode each frame independently. This ensures causal structure, robustness to missing or corrupted radar frames, and supports incremental updates.
• United Video Decoder: To prevent temporal flickering and ensure consistency, a hierarchical rectified flow transformer decodes all frames jointly. It utilizes space-time factorized attention and neighborhood attention to manage computational costs.
• Rectified Flow Training: The decoder is trained using a Rectified Flow (RF) objective (minimizing the velocity field error) rather than adversarial or perceptual losses. This simplifies training, ensures stability, and allows for sharp, realistic reconstructions.
• Stochastic Tanh Regularization (T-reg.): To ensure a smooth, bounded latent space suitable for generation without the complexity of KL-divergence weighting, the authors introduce a novel regularization. Latents are constrained to $[-1, 1]$ via a tanh activation and perturbed with Gaussian noise ($\epsilon \sim \mathcal{N}(0, \sigma^2 I)$). This architectural constraint encourages smoothness and robustness without requiring loss balancing.

2. Latent-Space Nowcasting (Second Stage)

• Masking-Based Conditioning: The latent-space model employs a masking-based training strategy (inspired by RaMViD). It learns to denoise arbitrary subsets of past frames, allowing for flexible conditioning lengths and robustness to variable input availability during inference.
• Probabilistic Sampling: The system generates forecasts by sampling from the conditional distribution $p(x_{out}|c)$ in the latent space.
• Uncertainty Quantification:
  • Decoding Uncertainty: By sampling multiple times from the generative decoder given the same latent, the model captures uncertainty inherent in the decoding process.
  • Ensemble Variance: The variance across an ensemble of reconstructions serves as an estimate of aleatoric uncertainty, which is particularly high in regions of intense, chaotic rainfall.

Key Contributions

1. A Scalable First-Stage Design: A rectified-flow encoder-decoder with stochastic tanh regularization that achieves high-fidelity reconstruction without perceptual or adversarial losses, ensuring a smooth, bounded latent space.
2. Asymmetric Architecture: A frame-wise encoder to prevent information leakage and support variable-length inputs, combined with a united video decoder for temporal consistency.
3. Decoding Uncertainty Quantification: A probabilistic decoder design that enables multiple reconstructions from the same latent, providing natural estimates of decoding uncertainty crucial for safety-critical applications.
4. Robust Latent-Space Nowcasting: State-of-the-art performance achieved through masking-based rectified-flow training, which handles variable-length conditioning and scales effectively with model size and ensemble count.

Experimental Results

The method was evaluated on the SEVIR benchmark (20,393 extreme weather events) and compared against deterministic and probabilistic baselines (e.g., ConvLSTM, PredRNN, CasCast, PreDiff).

• Reconstruction Quality: FREUD with T-reg. achieved superior metrics (RMSE, MAE, SSIM, PSNR) compared to CasCast and other variants, even with a single ensemble member. It demonstrated higher temporal consistency (lower dMAE) than frame-wise diffusion autoencoders.
• Forecasting Performance: The FREUD pipeline (specifically the LSM-L variant) achieved state-of-the-art Continuous Ranked Probability Score (CRPS) and SSIM, outperforming all existing methods in distribution coverage and perceptual quality.
• Uncertainty Calibration: The ensemble rank histograms indicated significantly better calibration (closer to uniform) compared to CasCast, which showed strong overconfidence. The variance of the ensemble correlated strongly (R = 0.997) with precipitation intensity, correctly identifying high-uncertainty regions.
• Efficiency: The transformer-based FREUD autoencoder uses fewer parameters and FLOPs than CasCast, with encoding being 96% faster and decoding 68% faster (at 5 NFE).
• Scaling: Performance consistently improved with larger model sizes (LSM-S, LSM-B, LSM-L), confirming favorable scaling properties.

Significance and Claims

The paper claims that FREUD addresses the critical gap between efficiency and reliability in weather nowcasting. By replacing deterministic compression with a probabilistic rectified flow framework, the authors demonstrate that:

• Uncertainty can be captured end-to-end: The system quantifies both compression (decoding) uncertainty and predictive uncertainty, providing valuable insights for extreme events where the range of outcomes matters more than the most likely one.
• Simplicity yields robustness: Eliminating adversarial and perceptual losses simplifies training while improving reconstruction sharpness and stability.
• Operational Viability: The frame-wise encoding and flexible conditioning make the system robust to real-world sensor faults and missing data, a common issue in operational radar networks.

The authors conclude that their approach paves the way for reliable, purely data-driven short-term precipitation forecasting without relying on deterministic physics-based predictors, offering a scalable solution that balances accuracy, speed, and uncertainty quantification.