PaperlayerView exact revision on arXiv

Source-linked AI summary

Skillful and realistic precipitation forecasting with reinforced diffusion

Ozan Sener, Anand Avati, David Hafner, Jierui Lin, Laura Manduchi, Xiaoying Pang, Langwen Huang, Torsten Hoefler, Russ Schumacher, Vladlen Koltun

arXiv:2610.38775v1cs.LGphysics.ao-ph

TL;DR

Precipitation forecasting must reconcile accurate pointwise predictions with realistic structure, particularly for localized extremes. FitzRoy combines a latent-state forecasting framework with operational data and reports strong, structurally realistic forecasts through 24 hours, including advance signals of major floods.

  • Problem

    Physics-based systems struggle to resolve localized mesoscale rainfall, while precipitation forecasts have lacked a combination of high skill and realistic structure.

    Source: Introduction · p. 1, Discussion · p. 10

  • Method

    FitzRoy is a machine-learning framework that encodes atmospheric inputs into latent states and rolls them forward recurrently for precipitation forecasting.

    Source: Introduction · p. 3, Learning the latent atmospheric… · p. 4

  • Results

    FitzRoy matches or outperforms strong physics-based and ML baselines through 24 hours while preserving realistic structure and signaling extreme flood accumulations ahead of state-of-the-art models.

    Source: Introduction · p. 4, Discussion · p. 10

  • Takeaways & Limitations

    Low-latency inference provides actionable high-fidelity forecasts in seconds, extending actionable lead time for emergency management beyond NWP alone.

    Source: Discussion · p. 11

    GraphCast and GenCast comparisons use non-operational ERA5 initializations, so the results are an upper bound rather than a direct operational comparison.

    Source: Evaluation of forecast skill · p. 6

Abstract

from arXiv · show

Accurate precipitation forecasting remains elusive as physics-based models struggle to resolve mesoscale convection and machine learning often yields physically inconsistent, blurred results that consistently underestimate extreme events. We present FitzRoy, the first system that simultaneously achieves state-of-the-art predictive skill and physical realism over the continental United States. Through a reinforced diffusion framework, where a conditional discrete diffusion model is refined via reinforcement learning, FitzRoy effectively resolves the long-standing trade-off between pointwise accuracy and structural fidelity. Across lead times up to 24 hours, FitzRoy consistently matches or exceeds state-of-the-art models while preserving structural realism and capturing extreme events. Its low-latency inference further enables deployment in real-time operational settings. In case studies such as the 2025 Texas Hill Country floods and the 2026 Missouri floods, FitzRoy signals record-breaking accumulations that existing models largely missed. By bridging the gap between realism and skill, FitzRoy establishes a new frontier for high-fidelity precipitation forecasting.

Introduction

Precipitation forecasts must balance pointwise skill with realistic, sharp rainfall structure, especially for localized extremes. FitzRoy combines reinforced diffusion with operational data and reports strong forecasts through 24 hours, including advance signals of two major floods.

Introduction

  • The field lacked a system combining high forecast skill with high-fidelity precipitation structure, as standard ML often blurs rain fields and underestimates extremes.Physics-based models also struggle to resolve sharp gradients and the stochastic nature of mesoscale convection.Source: Introduction · pp. 1, 3
  • FitzRoy combines conditional discrete diffusion with reinforcement learning to forecast CONUS precipitation up to 24 hours while maintaining structural realism.The diffusion model learns plausible precipitation structures; policy gradients then bias generations toward higher-skill forecasts.Source: Introduction · pp. 3–4
  • FitzRoy uses heterogeneous operational data and produces high-resolution forecasts within seconds, unlike recent models that rely on operationally unavailable retrospective reanalysis.The system reports matching or outperforming strong physics-based and ML baselines for precipitation probability, rate, and accumulation.Source: Introduction · pp. 2, 4
  • FitzRoy signaled extreme accumulations 18–20 hours ahead in the Texas and Missouri floods, while state-of-the-art models underpredicted magnitude and misplaced the spatial footprint.The Texas figure reports a 19-hour signal; the Missouri figure describes consistent identification across initialization times.Source: Figure 1 · p. 2, Figure 2 · p. 3, Introduction · p. 4

FitzRoy

FitzRoy uses a shared latent atmospheric representation with separate point-probability and gridwise intensity decoders. Its intensity forecasts combine discrete tokenization and conditional diffusion with reinforcement learning to balance structural realism and forecast skill.

FitzRoy

  • FitzRoy uses probability and intensity decoders in a two-stage curriculum to predict point-level probabilities and spatial precipitation fields from a shared atmospheric state.The framework targets accumulated precipitation, instantaneous rates, and exceedance probabilities, and accommodates heterogeneous inputs.Source: FitzRoy · p. 4, Figure 3 · p. 5

Learning the latent atmospheric representation

  • Cross-attention maps heterogeneous observations into latent states that are recurrently rolled forward to the desired lead time, without fixed-grid or Markovian temporal constraints.The probability decoder queries latent states at target coordinates and trains against observations for pointwise accuracy.Source: Learning the latent atmospheric… · pp. 4–5

Intensity forecasts via reinforced diffusion

  • The intensity decoder tokenizes precipitation fields, generates coherent accumulation and rate fields with conditional discrete diffusion, then uses RL to bias samples toward higher-skill forecasts.RL ranks candidate realizations using a composite reward based on fractions skill score and frequency bias at multiple intensity thresholds.Source: Intensity forecasts via reinforced… · pp. 5–6, Figure 3 · p. 5

Evaluation of forecast skill

Across deterministic and probabilistic precipitation tasks, FitzRoy is consistently competitive with or better than the evaluated baselines, with gains varying by model, metric, threshold, and lead time. It also produces sharper, calibrated distributions and can issue forecasts with sub-hour latency.

Evaluation of forecast skill

  • FitzRoy is evaluated on a held-out test period beginning in August 2024 against operational and ML baselines, with GraphCast and GenCast results interpreted as upper bounds.Those two models use non-operational ERA5 reanalysis, unlike FitzRoy's operational inputs.Source: Evaluation of forecast skill · p. 6
  • Across deterministic and probabilistic tasks, FitzRoy is consistently competitive with or better than all baselines, with improvements varying by model, metric, threshold, and lead time.Against HREF, improvements are generally more modest and some comparisons are statistically indistinguishable; FitzRoy nevertheless matches or exceeds it in the evaluated regimes.Source: Evaluation of forecast skill · p. 7, Figure 4 · p. 7
  • FitzRoy shows its most significant gains for rain intensities and improves discrimination between rainy and non-rainy events and skill for heavy-precipitation thresholds.Probabilistic components produce sharper distributions that remain calibrated in the tails.Source: Evaluation of forecast skill · p. 8
  • FitzRoy can ingest recent observations and produce forecasts with sub-hour latency, improving its advantage in a latency-aware evaluation over NWP systems whose available forecasts may be hours old.Pure forecasting-quality comparisons omit data-delivery and model-runtime latency; supplementary evaluation adjusts effective lead time for initialization, data availability, and runtime.Source: Evaluation of forecast skill · pp. 8–9
  • Supplementary analyses examine the contributions of NBM forecasts, reinforcement learning, and architectural and modeling choices.These analyses assess individual components of FitzRoy's performance.Source: Evaluation of forecast skill · p. 9

Evaluation of forecast realism

FitzRoy's probabilistic and intensity forecasts are well calibrated, and its precipitation fields preserve observed spatial structure across scales. In the Hurricane Milton case study, it captures the storm eye and rainbands while localizing the storm more accurately than HRRR at longer lead times.

Evaluation of forecast realism

  • FitzRoy is well calibrated across four intensity thresholds and is unbiased for rain/no-rain events, with event frequencies consistent with observations for light and heavy precipitation.Its frequency-bias performance shows only slight heavy-rain overprediction in the figure description.Source: Evaluation of forecast realism · p. 9
  • FitzRoy closely matches the observed MRMS precipitation spectrum across spatial scales, preserving variance from convective-scale features to mesoscale organization.At some scales its spectra are closer to observations than those of HRRR and HREF.Source: Evaluation of forecast realism · p. 9
  • In the Hurricane Milton case study, FitzRoy captures the eye and spiral rainbands, while HRRR's position errors grow at 12–18-hour lead times and HREF blurs these features.The comparison illustrates how the aggregate realism diagnostics relate to depictions of an extreme weather event.Source: Evaluation of forecast realism · p. 9

Case studies: forecasting catastrophic extremes

In the Texas Hill Country flood case, FitzRoy forecasts the location and magnitude of extreme accumulations about 12 hours before the heaviest rain, while HRRR underpredicts and HREF displaces the maxima. The Missouri case likewise describes a consistent extreme-precipitation signal from FitzRoy amid smaller, more variable predictions from other models.

Case studies: forecasting catastrophic extremes

  • About 12 hours before the heaviest Texas Hill Country rain, FitzRoy forecasts a coherent band aligned with the observed training convection and flood-producing region, capturing the location and magnitude of extreme accumulations.HRRR substantially underpredicts accumulation and misses the full flood-producing swath; HREF's maxima are displaced and its footprint too limited.Source: Case studies: forecasting catastrophic… · p. 10
  • Across multiple forecast cycles for the Texas event, FitzRoy's extreme-rain signal remains anchored over the impacted region, while the heaviest-rain location shifts between HREF and HRRR initializations.The authors connect this run-to-run stability with resolving mesoscale organization.Source: Case studies: forecasting catastrophic… · p. 10
  • For the Missouri flood, FitzRoy predicts a narrow band of heavy rain with intensity and structure similar to the observed storm, though with some location errors.Other operational models underpredict rainfall, place the convection band incorrectly, or both.Source: Case studies: forecasting catastrophic… · p. 10
  • In Missouri, FitzRoy maintains a consistent extreme-precipitation forecast while other models predict smaller accumulations with more run-to-run variation.The cases highlighted in the main text include the Texas Hill Country and Missouri floods; the broader case set spans several meteorological regimes.Source: Case studies: forecasting catastrophic… · pp. 9–10

Discussion

FitzRoy uses reinforced diffusion to reduce the precipitation forecasting trade-off between skill and structural realism, while a shared latent representation supports heterogeneous inputs. The authors report high skill and realistic spatial structure across 24-hour lead times, with low-latency inference and case-study performance on high-impact floods.

Discussion

  • FitzRoy achieves strong performance across 24-hour lead times while preserving the high-frequency spatial structure characteristic of observed precipitation.The paper frames this as substantially reducing the historical skill-realism trade-off.Source: Discussion · p. 10
  • A precipitation tokenizer and conditional discrete diffusion anchor forecasts in physically plausible structures, while reinforcement learning biases the diffusion process toward high-skill trajectories.The approach is described as mitigating the double-penalty problem for realistic but slightly displaced storms.Source: Discussion · p. 10
  • A shared latent representation lets FitzRoy integrate additional external forecasts and operate across forecasting, nowcasting, downscaling, and post-processing roles depending on inference-time information.The framework also achieves strong performance using initializations constructed from physics-based model states and observations.Source: Discussion · p. 10
  • In the Texas and Missouri flood cases, the authors report that operational physics-based models failed to signal event magnitude in time for effective emergency response, while FitzRoy resolved extreme convection consistently.The passage presents this as evidence of real-world utility for high-impact localized extremes.Source: Discussion · p. 10
  • FitzRoy produces high-fidelity forecasts in seconds, whereas high-resolution physics-based ensembles require hours, and the authors say this extends actionable lead time for emergency management beyond NWP alone.The claimed operational benefit is tied to low-latency inference.Source: Discussion · p. 11
  • The authors present reinforced diffusion and shared latent representations as a blueprint for operational meteorology that synthesizes precipitation skill and realism.They connect this direction to weather anticipation for resource management and disaster resilience.Source: Discussion · p. 11

Supplementary Information

The supplementary information details model and evaluation protocols, analyzes component contributions to forecast skill and probabilistic value, and presents case studies, a qualified comparison with MetNet-3, and training costs.

Model details

  • FitzRoy encodes observations into recurrent atmospheric latents and uses probability and 1-km intensity decoders; intensity generation combines tokenization, conditional discrete diffusion, and reinforcement learning.The probability decoder predicts hourly precipitation as a 401-bin classification target, with additional weather-variable and precipitation-type supervision.Source: Model details · p. 15, Encoder · p. 15, Probability decoder · p. 16

Evaluation details

Analysis of model components

  • RL yields the largest ablation gains, improving CSI and FSS through 24 hours—especially for heavier precipitation—and generalizes to unseen thresholds and scales without degrading realism.The realism check covers calibration, frequency bias, power spectral density, and qualitative analyses.Source: Effects of reinforcement learning · p. 19
  • NBM guidance is the second-largest contributor, with particularly pronounced improvements at heavier precipitation thresholds.The study compares the full architecture with a version trained without NBM conditioning.Source: Contribution of NBM forecasts · p. 20
  • Without NBM, FitzRoy outperforms ENS-CF on all evaluated metrics, matches or outperforms HRRR, and outperforms NBM except on CSI and FSS at 2 mm.The authors frame the without-NBM and full configurations as complementary comparisons to standalone systems and ensemble guidance, respectively.Source: Contribution of NBM forecasts · p. 20
  • Changing the latent NBM integration coefficient from 0.5 to 0.25 or 0.75 has no systematic effect on forecast performance.This indicates robustness to the precise coefficient choice.Source: Contribution of NBM forecasts · p. 20
  • Discrete diffusion yields minor but statistically significant accuracy gains over continuous diffusion, while categorical training performs comparably to regression and enables probabilities above arbitrary rain thresholds.The continuous variant falls slightly below discrete diffusion with reinforcement learning, especially at higher rain thresholds; the categorical-loss latent performs comparably to the regression-loss latent.Source: Discrete vs. continuous diffusion · p. 20, Categorical vs. regression loss · p. 20, Figure S2 · p. 24
  • Replacing point-cloud attention with a regular raster slightly degrades hourly accumulation skill, while point-based encoding also supports irregular inputs and missing observations.The ablation changes only the input encoder and keeps downstream components and parameter counts comparable.Source: Point-based vs. grid-based encoder · p. 21, Figure S2 · p. 24

Additional hourly accumulation case studies

  • FitzRoy indicates large Long Island accumulations 9–12 hours ahead and predicts Tropical Storm Chantal’s North Carolina heavy-rain location and approximate intensity 24 hours ahead, earlier than other models.For Chantal, FitzRoy’s extreme-precipitation forecast remained consistent while other models predicted smaller accumulations with more run-to-run variation.Source: Additional hourly accumulation case… · p. 21, Figure S4 · p. 26, Figure S5 · p. 27
  • FitzRoy consistently indicates potential heavy precipitation near Ruidoso and closely matches accumulated rainfall there, while other forecast models predict less.The case concerns highly localized storms that caused flash flooding on 8 July 2025.Source: Additional hourly accumulation case… · p. 21, Figure S6 · p. 28

Precipitation rate case studies

  • FitzRoy’s precipitation-rate forecasts reproduce observed squall-line structure in the June 2025 northern US derecho and depict repeated, back-building intense convection in the July 2026 Missouri storm.The Missouri forecasts show high rain rates and storm development similar to the observed convective system.Source: Precipitation rate case studies · p. 21, Figure S7 · p. 29, Figure S8 · p. 30

Training cost

  • The training-cost summary reports wall-clock time and GPU-hours for four pipeline stages, using equivalent passes rather than epochs because initialization times are sampled with replacement and forecast lead times dynamically.See Table S1 for the stage-by-stage computational-cost summary.Source: Training cost · p. 22, Table S1 · p. 22

Attempted comparison with MetNet-3

  • A strict quantitative comparison with MetNet-3 remains challenging, but FitzRoy shows higher CSI at medium-to-longer lead times in a retrained evaluation that excludes the full test period.Neither MetNet-3 code nor operational forecasts are publicly available, and reported HRRR values differ substantially from the reproduction; the authors caution that their interpretation is limited.Source: Attempted comparison with MetNet-3 · p. 22, Figure S12 · p. 34
Loading 2610.38775v1…