Precipitation Downscaling Using Foundation Model-Conditioned Diffusion
Authors: Victor Nascimento Ribeiro, Jorge Guevara, Jorge Sebastian Moraga, Chris Lucas, Natalie Lord, Andrew Taylor, Edward Lockhart, Will Trojak, +2 more
Organizations: IBM Research · Fathom · STFC Hartree Centre
High-resolution precipitation fields are essential for hydrological impact assessment, yet global climate model outputs are too coarse and biased for direct use. AI-based statistical downscaling with diffusion models offers a promising approach, but the mechanism by which large-scale atmospheric predictors condition generation remains largely unexplored. We investigate three conditioning strategies for a denoising diffusion probabilistic model applied to daily precipitation downscaling: channel concatenation of upsampled coarse predictors, cross-attention conditioning with a learned convolutional encoder, and cross-attention conditioning with the frozen encoder of the pretrained Prithvi WxC weather foundation model. All strategies are evaluated against an unconditioned baseline under identical conditions using probabilistic, distributional, spectral, and extreme-event metrics for the Colorado River Basin. Concatenation conditioning achieves the lowest point-wise CRPS and MSE, but tends to produce over-smoothed fields that suppress high-intensity events. In contrast, cross-attention conditioning provides substantially better distributional realism and modest improvements in spectral fidelity. Improvements are greatest for extremes: the Prithvi-WxC conditioned model retains over half of >100mm/day events, although estimates are uncertain due to limited samples. When trained on the full dataset, the learned convolutional model performs similarly to the foundation model-conditioned approach while requiring lower computational resources. However, the Prithvi-WxC-conditioned model achieves comparable performance with only five years of training data. These results indicate that cross-attention conditioning offers advantages over simple concatenation for probabilistic precipitation downscaling, and that pre-trained foundation model representations may offer benefits in data-limited settings.
Figures & tables
Variable
Description
Vertical Atmospheric Variables at {500,700,850} hPa
q
Specific Humidity
t
Temperature
u
Zonal Wind Speed (West–East)
v
Meridional Wind Speed (South–North)
w
Vertical Velocity
Table 1: ERA5 atmospheric, surface, and static variables used to condition the diffusion precipitation downscaling models. Vertical atmospheric variables are provided at three pressure levels (500, 700, and 850 hPa). Static physiographic variables are provided at their native resolution, while surface and vertical variables are coarse-resolution.
Figure 1: Schematic of concatenation-based conditioning (CC) for the diffusion precipitation downscaling model. Surface, vertical, and static physiographic predictor fields are concatenated with the noise input before being passed to the UNet denoising model, which outputs the downscaled total precipitation field. Surface and vertical coarse-resolution predictors must be interpolated to the high-resolution input grid.
Figure 2: Schematic of cross-attention conditioning (CA) integrated into the UNet diffusion denoiser. Static physiographic variables are incorporated via channel concatenation at their native high resolution. Coarse-resolution surface and vertical predictors are encoded into a compact latent representation, which is supplied to cross-attention blocks at multiple spatial resolutions within the UNet encoder–decoder.
Figure 3: Modifications for the integration of the Prithvi WxC weather foundation model as a conditioning encoder within the diffusion framework. A lightweight convolutional preprocessing stage maps the coarse-resolution ERA5 predictor stack to the input format expected by the Prithvi WxC encoder. The frozen foundation model encoder produces high-dimensional spatiotemporal embeddings, which are subsequently projected to a compact feature map via convolutional post-processing layers and injected into the diffusion UNet via cross-attention.
CRPS
MSE
Abs. Bias
Abs. Rel. Bias
Raw
(x16)
Raw
(x16)
Raw
Time avg.
Raw
Time avg.
UNC
1.092
1.034
15.993
11.893
1.617
0.402
317.22
0.380
CC
0.511
0.389
4.282
1.918
0.760
0.185
45.656
0.169
CA-CE
0.581
0.444
5.497
2.825
0.876
0.107
49.494
0.112
CA-PWC
0.557
0.415
5.934
2.548
0.827
0.095
35.710
0.097
Table 2: Comparison of conditioning strategies for probabilistic accuracy and bias metrics for daily total precipitation (mm/day) over the test period (2013–2015). Bold values indicate best performance per metric. ”Raw” values correspond to mean of daily per-pixel precipitation metrics. ”x16” are calculated on 16x16 degree grid mean values. and ”Time avg” on mean values across all time steps in the test period.
Precipitation Distribution
Spatial Statistics
Distribution
Calibration
Histogram
RALSD
RAPSD
Bias
Error
CRPS
CRPS
UNC
0.358
0.028
0.150
13.107
7.011
CC
0.385
0.065
0.102
5.942
2.763
CA-CE
0.138
0.079
0.102
5.390
2.632
CA-PWC
0.065
0.050
0.102
5.539
2.598
Table 3: Comparison of conditioning strategies for precipitation distribution and spatial spectral metrics on the test period.
Figure 4: Mean and maximum daily total precipitation (mm/day) over the test period (2013–2015) for the ERA5-Land reference dataset and each conditioning strategy. Top row (a): time-mean daily precipitation fields and corresponding marginal histograms. Bottom row (b): time-maximum daily precipitation fields and corresponding marginal histograms. Red dashed lines on histograms indicate the reference mean value.
Figure 5: Spatial maps of mean relative bias (%) in daily total precipitation (mm/day) for the test period (2013–2015). Top row: pixel-wise mean relative bias across all test days. Bottom row: mean relative bias computed after spatial maximum aggregation.
Annual Maximum Precipitation
CRPS
MSE
MAE
Distribution Bias
UNC
9.234
288.551
11.869
0.359
CC
7.522
194.825
9.486
0.388
CA-CE
5.750
148.712
8.271
0.134
CA-PWC
5.730
175.959
8.854
0.062
Table 4: Comparison of conditioning strategies for annual maximum daily precipitation (mm/day) over the test period (2013–2015). Bold values indicate best performance per metric.
Figure 6: Event frequency ratios across precipitation intensity bins (mm/day) for each conditioning strategy over the test period (2013–2015). The ratio is computed as the number of predicted events divided by the number of reference events within each intensity bin, with a value of 1.0 (dashed red line) indicating a perfect match. Observation pixel-event counts: 0-0.1 mm/day (n=14.38M); 0.1-1 mm/day (n=4.68M); 1-10 mm/day (n=4.40M); 10-30 mm/day (n=569.6k); 30-50 mm/day (n=37.9k); 50-100 mm/day (n=6.7k); 100-200 mm/day (n=263). Lower counts in extreme bins indicate higher uncertainty in those ratios.
Figure 7: Radially averaged power spectral density (RAPSD) of daily total precipitation over the test period (2013–2015) for the ERA5-Land reference dataset and each conditioning strategy (top panel). The bottom panel shows the ratio of predicted to reference power spectra as a function of spatial wavelength (km). Values close to 1.0 indicate accurate reproduction of spatial variability at a given scale.
Figure 8: Fractions Skill Score (FSS) for daily total precipitation at multiple intensity thresholds (mm/day) as a function of spatial scale (km) for each conditioning strategy over the test period (2013–2015). Higher FSS values indicate better spatial skill at a given scale and threshold.
Figure 9: Data efficiency of cross-attention conditioning strategies as a function of training data volume (5–25 years) for CA-CE and CA-PWC. (a): CRPS, (b): Absolute Relative Bias, (c): RALSD.
Trainable
Total
GPU Memory
Inference
Training Time
Params
Params
Usage
Time
per Epoch
UNC
48.9 M
48.9 M
15.9 GB
0.210 s
138 s
CC
48.9 M
48.9 M
16.0 GB
0.211 s
141 s
CA-CE
52.9 M
52.9 M
27.8 GB
0.605 s
324 s
CA-PWC
85.0 M
301 M
29.4 GB
0.684 s
342 s
Table 5: Model complexity across conditioning strategies. Trainable parameters refers to the number of parameters updated during training, while total parameters includes the frozen Prithvi WxC backbone. GPU memory usage includes loading both model and data. Inference time is reported per day per ensemble for the 1095 day × 5 ensemble test set on a single NVIDIA H100 80GB HBM3.
Accurate daily precipitation forecasts are important for weather-sensitive decisions. Although recent Weather Foundation Models (WFMs) have achieved significant progress, they are trained on reanalysis data, which exhibit regional and intensity-dependent biases. Since retraining WFMs on observation-based targets is costly, we study Dynamics-Precipitation Decoupling, a plug-and-play framework that uses a separately trained generator to map non-precipitation forecasts from a frozen weather model to observation-aligned daily precipitation. Its key challenges are to capture heterogeneous rainfall distributions and align daily accumulations with atmospheric signals that evolve across space and time. To address them, we train a diffusion generator conditional on atmospheric states and satellite-gauge estimates of precipitation. A multi-scale recurrent encoder summarizes atmospheric evolution, while cross-attention aligns atmospheric and precipitation features during denoising. Extensive evaluation of various backbones, including Pangu-Weather, GraphCast, and HRES, on global and regional forecasting shows improved daily precipitation accuracy and event detection, with better representation of localized accumulation patterns. Global ensemble evaluation shows lower Brier scores than ENS and training-from-scratch GenCast at all reported lead times, with lower CRPS in most comparisons. Remarkably, a heavy-rainfall case study using 2,288 stations in China further supports improved daily accumulation patterns and event detection against ground observations.
Guowen Li, Yang Liu, Mengxuan Chen +11
Tsinghua University · The Chinese University of Hong Kong · The Hong Kong University of S & T +5
Skillful medium-range precipitation forecasting at kilometer scale remains challenging over complex terrain because precipitation arises from multiscale nonlinear processes that global models cannot explicitly resolve at affordable cost. Global AI weather models can produce skillful medium-range forecasts, but their native 0.25 degrees resolution limits direct use for local hazard applications. Statistical downscaling can help bridge this gap, yet existing approaches often struggle with state-dependent, and especially lead-time-dependent, biases in global forecasts. We introduce SwAIther-Precip, a lead-time-aware downscaling framework that converts coarse-resolution AIFS forecasts into probabilistic km-scale precipitation fields over Switzerland. First, a U-Net conditioned on lead time via feature-wise linear modulation deterministically corrects systematic biases at coarse resolution. This targeted correction enables a cheaper super-resolution stage conditioned only on corrected precipitation, allowing direct training on observations rather than on the full atmospheric state. A diffusion-based model then generates fine-scale spatial variability independently of lead time. Using AIFS forecasts and CombiPrecip radar-gauge observations, SwAIther-Precip reduces CRPS by 48% relative to raw AIFS. The generated fields reproduce observed spatial variability with spectral fidelity above 0.85 at large scales and 0.88 at small scales, corresponding to an effective resolution of approximately 4 km on a 1 km grid for lead times up to 5 days. Training across lead times further improves long-range performance, yielding a 13% CRPS reduction at 6 days relative to lead-time-specific models. These results show that explicitly correcting lead-time-dependent biases before generative super-resolution is key to efficient km-scale probabilistic downscaling of global AI precipitation forecasts.
Dan Assouline, Erwan Koch, Federico Amato +5
Swiss Data Science Center, ETH, Zürich, Switzerland · Expertise Center for Climate Extremes, University of Lausanne, Lausanne, Switzerland · Faculty of Geosciences and Environment, University of Lausanne, Lausanne, Switzerland +1
High-resolution precipitation information is essential for climate impact assessment, yet global climate models remain too coarse to resolve key small-scale processes. Existing machine learning downscaling methods often require paired low- and high-resolution data for supervised learning, are tied to fixed regions or scale factors during inference, and can be computationally expensive to train and run in physical space. Here we introduce Longwang, a zero-shot latent generative framework for global spatiotemporal precipitation downscaling. Longwang learns a context-conditioned latent generative prior and combines it with a physically informed observation operator through posterior sampling, enabling daily O(10 km) precipitation fields to be generated from monthly O(100 km) inputs. On ERA5 reanalysis, Longwang outperforms standard posterior sampling with an unconditional generative prior in reconstructing fine-scale spatial patterns, preserving temporal coherence, and recovering extreme precipitation intensities. The framework further generalizes to historical climate simulations and future climate projections under substantial distribution shift.
Yue Wang, Daniele Visioni
1*Department of Earth and Atmospheric Sciences, Cornell University, 112 Hollister Drive, Ithaca, 14853, NY, USA.