Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

Abstract: Accurate weather forecasts are essential across various domains and are safety-critical in extreme weather conditions. Compared to simulation-based forecasting, data-driven approaches show greater efficiency, enabling short-term, high-resolution nowcasting. In particular, diffusion models proved effective in weather nowcasting due to their strong probabilistic foundation. However, existing methods rely on deterministic compression to reduce the complexity of high-dimensional weather data, limiting their ability to capture uncertainty in the decoding process. In this work, we introduce $\textbf{FREUD}$, a $\textbf{Fr}$ame-wise $\textbf{E}$ncoder and $\textbf{U}$nited $\textbf{D}$ecoder model based on rectified flow transformers for efficient compression of spatio-temporal weather data. Frame-wise encoding enables continuous forecast updates, while the unified video decoder ensures temporal consistency. Our uncertainty-preserving first stage allows us to capture aleatoric uncertainty via ensembling, which is particularly beneficial for extreme weather events with high decoding variability. We achieve state-of-the-art performance in precipitation nowcasting with a compact latent-space rectified flow transformer on the SEVIR benchmark and show further performance gains by model and test-time scaling. Code available here: https://github.com/CompVis/weather-rf

• The paper presents FREUD, a transformer-based model integrating rectified flow transformations to enhance uncertainty quantification in precipitation nowcasting.
• The paper details a novel hierarchical decoder that jointly processes sequences, achieving superior CRPS, SSIM, and reliability on the SEVIR benchmark.
• The paper demonstrates scalable, efficient training and inference with stochastic tanh latent regularization, improving performance on extreme weather forecasting.

Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

Introduction and Motivation

The paper "Probabilistic Precipitation Nowcasting with Rectified Flow Transformers" (2605.31204) addresses the challenges of short-term, high-resolution weather forecasting—nowcasting—where uncertainty quantification is critical for safety-relevant applications. Conventional deterministic ML models suffer from blurry predictions due to mode averaging and lack intrinsic uncertainty estimates. Probabilistic generative models, especially diffusion models, have advanced the field by enabling ensemble forecasting and calibrated uncertainty estimates but are bottlenecked by lossy compression stages and inefficient high-dimensional sampling.

This work introduces the FREUD framework: a transformer-based, uncertainty-preserving compression-stage tailored for weather data. The architecture uses a frame-wise transformer encoder and a unified generative rectified flow decoder supporting ensembling, with bounded and regularized latents, to solve the key problems of scalability, uncertainty quantification, and operational robustness.

FREUD Architecture and Training Methodology

The FREUD model merges recent advances in transformer modeling, flow-based generative methods, and latent-space regularization. The encoder operates on individual frames, allowing flexible conditioning, sensor fault robustness, and incremental updates, while the hierarchical decoder processes the entire sequence jointly for strict temporal consistency.

Figure 1: The FREUD architecture combines a lightweight frame-wise transformer encoder with a hierarchical rectified flow transformer decoder conditioned on encoder latents.

Rectified flows generalize diffusion and normalizing flow models by constructing a continuous velocity field bridging Gaussian noise and data. The decoder is trained directly with a rectified flow loss, removing the need for adversarial or perceptual losses and resolving the instability and bias issues of latent autoencoders in prior works. Critical to generative modeling, the FREUD latent representation is regularized using a stochastic $\tanh$ function, ensuring bounded, smooth, high-density latents that outperform traditional KL-based regularization in both reconstruction fidelity and downstream generative ability.

Figure 2: FREUD achieves more parameter-efficient compression and decoding than previous designs, accelerating both training and inference.

Uncertainty Quantification in Compression and Generative Modeling

A central innovation is the ability to sample multiple reconstructions from a single latent conditioning, directly estimating aleatoric uncertainty in the decoding process. The ensemble variance reveals localized uncertainty tightly correlated with precipitation intensity, especially for extreme or chaotic events.

Figure 3: Quantitative correlation between precipitation intensity and decoding ensemble variance, demonstrating that uncertainty increases with meteorological extremeness.

Qualitative evaluation on normal and extreme cases further underscores the decoder's capacity to track uncertainty where it matters most for operational decision-making.

Figure 4: Ensemble distribution reconstructions across precipitation regimes; variance grows for intense rainfall, capturing ground truth while deterministic baselines miss uncertainty.

Probabilistic Latent-Space Nowcasting and Forecast Performance

Downstream forecasting is executed by a scalable latent-space rectified flow transformer, trained with masking-based variable-length conditioning. This approach permits flexible inference—robust to missing or corrupt frames—and enables rapid online weather updates. The combined FREUD+LSM (Latent Space Model) achieves state-of-the-art scores in Continuous Ranked Probability Score (CRPS), structural similarity (SSIM), and reliability metrics on the SEVIR benchmark.

Figure 5: Latent-space inference pipeline: forecast generation in FREUD latent space, followed by ensemble decoding to pixel space.

The ensemble forecasting method produces well-calibrated uncertainty estimates, as demonstrated by flat rank histograms and improved reliability indices relative to previous cascaded models.

Figure 6: Ensemble rank histograms: FREUD-based pipeline exhibits superior calibration compared to CasCast.

Scaling studies confirm that performance increases with model capacity and ensemble size, establishing the transformer architecture as robustly scalable for both inference and operational deployment.

Ablations and Comparative Analysis

Architectural and regularization ablations show that the stochastic $\tanh$ latent regularization yields the most compact, structured latent space, resulting in higher downstream forecast skill and sharper reconstructions. Joint decoding eliminates flickering and temporal incoherence seen in frame-wise models. Masking-based training and diffusion forcing support flexible input lengths, mitigating error accumulation common in autoregressive approaches.

Deterministic conditioning improves localization but negatively impacts distributional coverage and calibration, highlighting an inherent trade-off.

Handling Extreme Events and Practical Robustness

FREUD's ensemble variance localizes uncertainty and detects out-of-distribution anomalies (e.g., artificially injected blobs), confirming its reliability for abnormal event detection. On rare, severe weather events (tornado, flood, hurricane), FREUD delivers superior reconstruction and probabilistic forecasting compared to previous cascaded diffusion architectures.

Figure 7: Forecasts for a hurricane event; FREUD captures circular dynamics, aligning ensemble with ground-truth motion patterns.

Implications and Future Directions

The FREUD pipeline advances probabilistic forecasting for safety-critical meteorology, providing fast, scalable, calibrated nowcasts with direct uncertainty quantification. The obsolescence of adversarial objectives, combined with transformer scalability and robust latent-space regularization, sets a new standard for applied forecasting. Practical deployments will benefit from resilience to sensor faults, computation efficiency, and accurate risk assessment for extreme weather scenarios.

Theoretically, the rectified flow paradigm and stochastic latent regularization open further lines of inquiry in generative modeling for high-dimensional, structured time series. Future research should address improved calibration in rare extremes, explore alternatives to classifier-free guidance for deterministic alignment, and focus on generalized nowcasting beyond radar-limited domains using satellite or multimodal inputs.

Conclusion

FREUD represents a significant advance in operational, uncertainty-aware precipitation nowcasting. Its transformer-based, uncertainty-preserving compression and scalable generative forecasting yield superior calibration, improved forecast quality, and practical robustness. These results support further exploration of rectified flows, stochastic regularization, and masking paradigms in generative models for spatiotemporal forecasting tasks.