cs.LGOct 28, 2026
SaveSkillful and realistic precipitation forecasting with reinforced diffusion
Organizations: Apple · ETH Zurich · Colorado State University
Abstract
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.
Figures & tables
Figure 1: Flooding in the Texas Hill Country. ( A ) 6-hour observed accumulated precipitation from radar between 03:00 and 09:00 UTC on 4 July 2025 around Camp Mystic (center of red box). ( B ) Area-averaged observed and forecasted accumulations over the red box shown in panel A (Lat: – , Lon: – ) for different initialization times. ( C ) 6-hour accumulation forecasts for different initialization times (lead times). While physics-based systems such as HRRR struggle even at 6-hour lead times, FitzRoy consistently predicts extreme precipitation 19 hours in advance. FitzRoy slightly overpredicts the rain, but for an extreme event this overprediction is preferable to underprediction for emergency alertness. FitzRoy also maintains consistency across initialization times, whereas HRRR degrades as the lead time shortens and the deadly flooding approaches.
Figure 2: Missouri flood. ( A ) 9-hour observed accumulated precipitation from radar between 03:00 and 12:00 UTC on 10 July 2026. ( B ) Area-averaged observed and forecasted accumulations over the red box shown in panel A (Lat: – , Lon: – ) for different initialization times. ( C ) 9-hour accumulation forecasts for different initialization times (lead times). FitzRoy consistently identifies the extreme accumulation magnitude across multiple initialization times, whereas operational baselines exhibit significant underprediction.
Figure 3: Overview of FitzRoy. ( A ) An encoder with cross-attention ingests heterogeneous inputs and produces a latent atmospheric state. A pointwise probability decoder queries this state at specific locations to generate probabilistic forecasts recurrently for each lead time, while a gridwise intensity decoder maps the latent state to a spatially coherent precipitation-intensity field over the full grid via discrete diffusion. For clarity, only a single lead time is visualized. ( B ) To represent high-fidelity precipitation structure faithfully, we learn a tokenizer that maps precipitation fields to a discrete vocabulary, and train the joint decoder in this token space. ( C ) Diffusion objectives encourage realism but do not directly optimize forecast skill. We therefore fine-tune with reinforcement learning to improve skill while preserving realism: we sample multiple precipitation fields from the diffusion model, rank them, and apply policy gradients to bias sampling towards high-skill forecasts.
Figure 4: Evaluation of deterministic and probabilistic forecast skill via scorecards. ( A ) Precipitation accumulation. ( B ) Instantaneous precipitation rate. ( C ) Probability of rain. Each scorecard displays a collection of metrics over all forecast lead times, where hatched cells indicate statistically not significant differences (two-sided test at a 5% significance level following Hamill (1999) ). Across nearly all tasks, metrics, and lead times, FitzRoy matches or outperforms the physics-based and ML baselines within their intended scope of operation.
Figure 5: Evaluation of forecast realism. ( A ) Frequency bias as a function of lead time. FitzRoy is essentially unbiased regarding rain vs. no-rain occurrence and reproduces event frequencies accurately, showing only slight overprediction for heavy rain. FitzRoy demonstrates superior frequency bias compared to NBM and ENS-CF, both of which significantly underpredict heavier rainfall events. ( B ) Calibration curves for FitzRoy and probabilistic baselines at four precipitation thresholds. The x-axis shows forecast probability and the y-axis the observed frequency of events exceeding the threshold. Perfect calibration follows the diagonal; points below (above) indicate overconfidence (underconfidence). Only bins with at least 1,000 samples are shown. FitzRoy is well-calibrated across all intensity thresholds. ( C ) Top: mean power spectral density of hourly precipitation accumulation as a function of spatial frequency. Bottom: relative difference with respect to MRMS. Curves closer to MRMS indicate better performance. FitzRoy closely matches the MRMS power spectrum, whereas all other methods overpredict power at low spatial frequencies. ( D ) Hurricane Milton case study (12:00 UTC October 9th, 2024). FitzRoy and HRRR both capture the hurricane structure and eye, whereas all other evaluated methods blur these features. FitzRoy localizes the storm more accurately than HRRR, which misplaces the eye of the hurricane, especially for longer lead times.
Appendix figures & tables28 assets
Supplementary material from the paper’s appendix.
Appendix
Figure S1 : Latency-aware evaluation of intensity forecasts. ( A ) In a real-time operational setting, model latency (red) must be accounted for when producing effective forecasts up to 24 hours ahead (pink). In addition, non-hourly model runs must be taken into account when providing forecasts for every hour of the day (brown). Earlier model initialization times lead to reduced effective forecast quality at a given valid time. This highlights the practical advantages of low-latency ML models such as FitzRoy in real-time operations, where a lower latency translates into even higher relative skill compared to baseline models. We achieve a sub-hour latency of FitzRoy in a real-time setting by using “MultiSensor_QPE_01H_Pass1” with a 20 min latency ( NOAA National Severe Storms Laboratory, 2025 ) , together with the most recent available NBM and GFS runs as inputs. Requested lead times are adjusted according to the time differences between the current time and the respective model initialization times. The encoding stage of FitzRoy requires approximately 2 seconds, whereas the diffusion process for a given lead time requires approximately 20 seconds on a single A100 GPU. Diffusion processes for different lead times can be executed in parallel. ( B ) Elapsed hours since model initialization time for each hour of the day. For example, at 04:00 UTC, the most recent available HREF-PMM output is 4 hours old; thus, the 5–28 hour forecasts from that model run are deployed as effective real-time operational 1–24 hour forecasts. ( C ) Scorecards for hourly accumulated precipitation and instantaneous precipitation rate; cf. Supplementary Fig. S19 – S21 . Hatched cells indicate statistically not significant differences (using a two-sided test with a 5% significance level ( Hamill, 1999 ) ).
Figure S2 : Analysis of model components. ( A ) FitzRoy compared to various model variants. Reinforcement learning provides the largest performance gains. The inclusion of NBM is the second-largest source of improvement, particularly for heavier rain categories. The results are robust to the choice of NBM fusion weight. Using a perceiver-based architecture rather than a grid-based one and discrete rather than continuous diffusion yields marginal but statistically significant gains. Using cross-entropy instead of regression performs similarly but cross-entropy enables probabilistic reasoning beyond regression. ( B ) Robustness to the RL reward definition. We evaluate the model across intensity thresholds and area sizes not used in the RL reward. FitzRoy generalizes strongly, showing even larger gains for several unseen categories, especially for heavier rain. ( C ) FitzRoy without NBM input compared to baseline approaches. The model without NBM performs better than NBM across all metrics except 2-mm CSI and FSS. Furthermore, FitzRoy without NBM outperforms ENS-CF on all metrics, and it surpasses or matches HRRR on all metrics. In all scorecards, hatched cells indicate differences that are not statistically significant (two-sided test, 5% significance level ( Hamill, 1999 ) ).
Figure S3 : Effect of reinforcement learning on forecast realism. ( A ) Frequency bias as a function of lead time. Both FitzRoy with and without RL exhibit low frequency bias across all lead times. ( B ) Top: Mean power spectral density of hourly precipitation accumulation as a function of spatial frequency. Bottom: Relative difference with respect to MRMS, where curves closer to zero indicate better performance. FitzRoy yields similar spectral characteristics with and without RL. ( C ) Hurricane Milton case study (12:00 UTC October 9th, 2024). Both model variants produce spatially realistic precipitation fields, while our final RL-fine-tuned model better captures several structures present in the observations.
Figure S4 : Long Island, New York flooding. ( A ) 6-hour observed accumulated precipitation from radar between 00:00 and 06:00 UTC on 19 August 2024 in the northeastern US. ( B ) Area-averaged observed and forecasted accumulations over the red box shown in panel A (Lat: – , Lon: – ) for different initialization times. ( C ) 6-hour accumulation forecasts for different initialization times (lead times). Although no models had perfect forecasts of these storms, FitzRoy provided earlier indication for the potential of large rainfall accumulations on Long Island, where significant flooding occurred.
Figure S5 : Remnants of Tropical Storm Chantal. ( A ) 9-hour observed accumulated precipitation from radar between 18:00 and 03:00 UTC on 6–7 July 2025 in North Carolina. ( B ) Area-averaged observed and forecasted accumulations over the red box shown in panel A (Lat: – , Lon: – ) for different initialization times. ( C ) 9-hour accumulation forecasts for different initialization times (lead times). FitzRoy consistently identifies the extreme accumulation magnitude across multiple initialization times, whereas operational baselines exhibit significant underprediction and higher run-to-run variability.
Figure S6 : Ruidoso, New Mexico flooding. ( A ) 3-hour observed accumulated precipitation from radar between 21:00 and 00:00 UTC on 8 July 2025 in the state of New Mexico. ( B ) Area-averaged observed and forecasted accumulations over the red box shown in panel A (Lat: – , Lon: – ) for different initialization times. ( C ) 3-hour accumulation forecasts for different initialization times (lead times). FitzRoy’s forecasts of this very localized cluster of rainstorms were a closer match to the observed precipitation.
Figure S7 : Northern US Derecho. MRMS and forecasted precipitation rate (Lat: – , Lon: – ) for different forecast lead times initialized at 22:00 UTC June 20th, 2025 (latency-adjusted). FitzRoy successfully forecasts the squall line and structure of future observations.
Figure S8 : Missouri flood. MRMS and forecasted precipitation rate (Lat: – , Lon: – ) for different forecast lead times initialized at 21:00 UTC July 9th, 2026 (latency-adjusted).
Figure S9 : Evaluation of model components. For each model variant, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC.
Figure S10 : Additional evaluation of the effect of reinforcement learning. For each model variant, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC.
Figure S11 : Evaluation of a model variant without NBM. For each model, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC.
Figure S12 : Attempted comparison with MetNet-3. Since neither source code nor operational forecasts from MetNet-3 ( Andrychowicz et al., 2023 ) are publicly available, we follow the evaluation protocol described in Andrychowicz et al. (2023) (solid lines) and compare against the values reported in Andrychowicz et al. (2023) (dashed lines). Reported values for HRRR (dashed, cyan) are substantially higher than those obtained in our reproduction attempt (solid, cyan), which complicates a direct comparison between MetNet-3 (dashed, pink) and FitzRoy (solid, blue). Despite those discrepancies, which persisted even after discussion with the authors of Andrychowicz et al. (2023) , the FitzRoy curves (solid, blue) show higher CSI values at medium to longer lead times. We expect this performance advantage to increase further under a fully consistent evaluation procedure. ( A ) Hourly accumulated precipitation. ( B ) Instantaneous precipitation rate.
Figure S13 : Evaluation of hourly probability of rain. For each model, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC.
Figure S14 : Evaluation of instantaneous precipitation rate. For each model, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC.
Figure S15 : Evaluation of instantaneous precipitation rate w.r.t. StormCast ( Pathak et al., 2026b ) . For each model, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC. Since StormCast only operates on a 1,536 km 1,920 km CONUS subregion ( Pathak et al., 2026b ) , we use the same subregion here.
Figure S16 : Evaluation of hourly precipitation accumulation. For each model, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC.
Figure S17 : Evaluation of 6-hour precipitation accumulation compared to existing ML prediction models. For each model, we evaluate four runs per day for the given evaluation time range: 00, 06, 12, and 18 UTC, against current state-of-the-art ML weather prediction models AIFS-single and GraphCast. FitzRoy outperforms these models across all metrics and intensities analyzed.
Figure S18 : Evaluation of 12-hour precipitation accumulation compared to GenCast. For each model, we evaluate two runs per day for the given evaluation time range: 00 and 12 UTC against the GenCast model. FitzRoy outperforms GenCast across all metrics and intensities analyzed.
Figure S19 : Latency-aware evaluation of hourly precipitation accumulation. For each model, we evaluate on every hour for the given evaluation time range where we take into account the operational latencies of the models (see Fig. S1 ).
Figure S20 : Latency-aware evaluation of instantaneous precipitation rate. For each model, we evaluate on every hour for the given evaluation time range where we take into account the operational latencies of the models (see Fig. S1 ).
Figure S21 : Latency-aware evaluation of instantaneous precipitation rate w.r.t. StormCast. For each model, we evaluate on every hour for the given evaluation time range where we take into account the operational latencies of the models. Since StormCast only operates on a 1,536 km 1,920 km CONUS subregion ( Pathak et al., 2026b ) , we evaluate here on the same subregion.
| Training stage | Hourly records | Equivalent passes | Hardware | Wall time | GPU-hours |
| Encoder + probability decoder | 302K | 20 | 64 A100 | 2 days | 3,072 |
| Tokenizer | 70K | 71 | 32 A100 | 1 day | 768 |
| MaskGIT pretraining | 167K | 48 | 64 A100 | 6 days | 9,216 |
| MaskGIT RL fine-tuning | 39K | 0.35 | 48 B200 | 1 day | 1,152 |
Table S1 : Computational cost of the FitzRoy training pipeline. We report the number of unique hourly records used by each stage, the corresponding number of equivalent passes through those records, the hardware configuration, wall-clock training time, and total GPU-hours. Equivalent passes are used instead of conventional epochs because initialization times are sampled with replacement and paired with dynamically sampled forecast lead times.
| Model | Type | Region | Spatial Resolution | Forecast Frequency | Precipitation Accumulation | Precipitation Rate | Probability of Precipitation |
| HREF | NWP | CONUS | 3 km | 6 hours | hourly | — | hourly |
| NBM | NWP | CONUS | 2.5 km | 1 hour | hourly | — | hourly (rain/no-rain) |
| ENS | NWP | global | 0.25° | 6 hours | hourly | hourly | hourly |
| HRRR | NWP | CONUS | 3 km | 1 hour | hourly | hourly | — |
| AIFS-single | ML | global | 0.25° | 6 hours | 6 hours | — | — |
| GraphCast | ML | global | 0.25° | 6 hours | 6 hours | — | — |
Table S2 : Output characteristics of FitzRoy and baseline models for precipitation forecasts up to 24 hours. StormCast operates on a 1,536 km 1,920 km CONUS subregion ( Pathak et al., 2026b ) . (NWP: Numerical Weather Prediction, ML: Machine Learning, CONUS: Continental United States.)
| Input Source | #Channels |
| ERA5 atmospheric variables | 21 |
| MRMS hourly accumulation | 1 |
| MRMS precipitation type | 1 |
| NBM hourly accumulation forecast | 1 |
| NBM precipitation type forecast | 1 |
Table S3 : Training inputs for the encoder and pointwise probability decoder.
| Target Variable | Loss Function |
| ERA5 atmospheric variables | Mean squared error |
| ERA5 hourly accumulation | Cross entropy (401 bins) |
| MRMS hourly accumulation | Cross entropy (401 bins) |
| MRMS precipitation type | Cross entropy (4 classes) |
Table S4 : Training targets and loss functions for the encoder and probability decoder.
| Variable | Level |
| geopotential | 1000 hPa |
| u_component_of_wind | 1000 hPa |
| v_component_of_wind | 1000 hPa |
| geopotential | 850 hPa |
| relative_humidity | 850 hPa |
| temperature | 850 hPa |
Table S5 : ERA5 variables at 0.25° resolution, cropped to CONUS with a 7.5° padding.
| Data Source | Variable | Description |
| MRMS | MultiSensor_QPE_01H_Pass1 | Multi-sensor accumulation 1-hour (20-minute latency) |
| MultiSensor_QPE_01H_Pass2 | Multi-sensor accumulation 1-hour (1-hour latency) | |
| PrecipRate | Radar precipitation rate | |
| PrecipFlag | Surface Precipitation Type | |
| NBM | QPF01 | Quantitative precipitation forecast |
| PType | Probability of Snow, Rain, Sleet, or Freezing Rain |
Table S6 : Precipitation data sources and variables used by FitzRoy. MRMS and NBM are gridded data. GHCN provides point observations.
| Encoder and probability decoder | Loss weight (ERA5 channels) | 0.1 |
| Loss weight (ERA5 hourly accumulation) | 0.1 | |
| Loss weight (MRMS hourly accumulation) | 1.0 | |
| Loss weight (MRMS precipitation type) | 0.1 | |
| Latent dimension ( ) | 1024 | |
| Number of query tokens ( ) | 1024 | |
| MRMS quality clamp | 0.01 |
Table S7 : Hyperparameters for training FitzRoy.
More in this topic
Precipitation Forecasting
Precipitation forecasting predicts where, when, and how much rain or snow will fall, from local radar nowcasts to global multi-day outlooks.
See all 56 papers