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.
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