A Probabilistic U-Net Approach to Downscaling Climate Simulations

AI summary

Authors: Maryam Alipourhajiagha, Pierre-Louis Lemaire, Youssef Diouane, Julie Carreau (Polytechnique Montréal) arXiv: 2511.03197v1 [cs.LG], 05 Nov 2025

Overview

• Research area: Machine learning for climate science — statistical downscaling of climate simulations using deep generative models.
• Technical level: Intermediate. Readers should be comfortable with U-Net architectures, latent-variable models, and standard verification metrics such as CRPS, MAE, and power spectral density.
• Scope: The paper adapts the probabilistic U-Net to climate downscaling and compares four training objectives — three settings of a WMSE–MS-SSIM loss and one afCRPS loss — for producing ensembles of high-resolution precipitation and temperature fields from data 16×16 times coarser.

What This Paper Is About

Climate models are computationally expensive and therefore produce outputs at coarse spatial resolution, while climate change impact studies typically need much finer scales. The authors take the probabilistic U-Net — a generative model originally designed for medical image segmentation — and repurpose it as a statistical downscaling emulator that turns coarse fields into ensembles of plausible high-resolution fields. The central question is which training objective best preserves extreme events and fine-scale spatial variability, since conventional MSE training tends to produce overly smoothed fields that miss extremes.

Key Contributions

1. First application of the probabilistic U-Net to climate downscaling, wrapping a deterministic U-Net backbone in a prior/posterior latent-variable framework to obtain probabilistic high-resolution ensembles.
2. A systematic comparison of four training objectives — WMSE–MS-SSIM with λ = 1 (WMSE only), λ = 0 (MS-SSIM only), and the tuned λ = 0.158, plus afCRPS — evaluated qualitatively and quantitatively.
3. An optimized training objective (the tuned WMSE–MS-SSIM setting) that better captures precipitation extremes, alongside the finding that afCRPS better reproduces small-scale spatial variability.
4. A publicly released implementation at github.com/MaryamAlipourH/prob-unet-climate-downscaling.

Main Findings

• No single loss dominates: The authors report that no variant outperforms all others on every metric; different losses emphasize different aspects of skill.
• Precipitation extremes: Among the three WMSE–MS-SSIM variants, the tuned setting (λ = 0.158) performs best on return-level curves. MS-SSIM (λ = 0) closely reproduces observed tail behavior in log-frequency histograms. The WMSE-only variant (λ = 1) substantially underestimates high-intensity precipitation and fails to capture extremes.
• afCRPS overshoots extremes: The afCRPS model tends to overestimate extreme precipitation events in both the return-level analysis and the log-frequency histograms.
• Spatial variability: In power spectral density analysis, the λ = 1 variant exhibits spectral smoothing and underestimates variance at higher radial wavenumbers; MS-SSIM (λ = 0) better recovers small scales; afCRPS provides the closest match to the observed spectrum across scales.
• Temperature is insensitive to the loss choice: Minimum and maximum temperature histograms and spectra align well with ground truth across all four variants, with only minor deviations at the extremes.
• Quantitative metrics (Table 1): afCRPS achieves the lowest CRPS overall — 0.94 ± 0.74 mm/day for precipitation, 0.68 ± 0.20 °C for T_min, and 0.62 ± 0.12 °C for T_max — and the best temperature MAE (0.90 ± 0.28 °C for T_min; 0.75 ± 0.17 °C for T_max). The λ = 1 variant has the best temperature CRPS for T_max (0.59 ± 0.14 °C) and the best precipitation MAE (1.19 ± 0.94 mm/day). Nearest-neighbour interpolation, used as a baseline, has the highest MAE of all (1.51 ± 1.14 mm/day for precipitation; 1.76 ± 0.60 °C for T_min; 1.30 ± 0.30 °C for T_max); CRPS is not reported for the NN baseline.
• Complementarity: The authors conclude that combining afCRPS with MS-SSIM may offer a more balanced solution, since afCRPS gives the best CRPS and accuracy while MS-SSIM variants better represent precipitation extremes despite higher aggregate error.
• Ensemble diversity: Qualitative examples show the model generating diverse but physically consistent realizations — large-scale precipitation patterns are reproduced across predictions while variability appears in higher-intensity regions, enabling representation of uncertainty in extremes.

Methodology in Plain English

• The backbone: Downscaling is cast as supervised image-to-image translation using a four-level U-Net patterned on a StyleGAN/EDM backbone. The encoder halves spatial resolution four times and doubles channels from 64 to 256; the decoder mirrors this with nearest-neighbour up-sampling followed by 3×3 convolutions. Encoder levels use two residual blocks each and decoder levels use three, with skip connections linking matching scales.
• Handling scale mismatch: Because a U-Net needs matching input and output resolution, low-resolution fields are first upsampled with nearest-neighbour interpolation (chosen to avoid smoothing and artifacts), and the network predicts the residual between this interpolated field and the true high-resolution target.
• Adding probability: The deterministic backbone is wrapped in the probabilistic U-Net framework. A prior network produces P(z|X) from the input alone; a posterior network produces Q(z|X,Y) when the high-resolution target is available. Both are axis-aligned Gaussians. During training, z is sampled from Q(z|X,Y), broadcast to a feature map, concatenated to the final U-Net activations, and passed through three 1×1 convolutions to produce the prediction. The loss is cross-entropy plus a weighted KL divergence term, with the weight γ adapted after a short warm-up phase. At inference, latent vectors are sampled from P(z|X) to generate an ensemble.
• Physics constraints: Two physical constraints are enforced through re-parametrization: precipitation is kept non-negative with a softplus function, and T_max ≥ T_min is ensured by applying the constraint to T_max − T_min. All metrics and losses are computed after converting predictions back to physical units.
• Data: One member of the ClimEx ensemble of dynamically downscaled simulations over southern Quebec and the Canadian Maritimes at 0.11° (≈ 12 km). The high-resolution domain is 128×128 grid cells; low-resolution data are obtained by averaging 16×16 blocks, yielding an 8×8 grid. Training uses 1960–1990, validation 1990–1997, and testing 1998–2005, avoiding the period when the RCP8.5 scenario begins. For return-level estimation, the test set is extended to 30 years to capture a complete cycle of climate variability.
• Training settings: 10 epochs, batch size 32, latent space dimension 16.
• Evaluation: Return-level curves with 95% confidence bands built by fitting a Generalized Extreme Value distribution to annual maxima at each grid cell and applying a parametric bootstrap; empirical return levels from 5 predictions are superimposed, with a match considered good if empirical levels fall within the confidence bands for at least 95% of points. Log-frequency histograms assess distributional fidelity across all pixels, and azimuthally averaged power spectral density assesses variance across spatial scales. Hyperparameters were α = 0.007, β = 0.048, λ = 0.158, and η = 0.95.

Why This Matters

The paper shows that deep generative emulators can produce ensembles of fine-scale climate fields at a fraction of the cost of dynamic downscaling, and that the choice of training loss directly shapes which aspects of the climate signal are captured. It matters because impact studies need both realistic extremes and realistic spatial detail, and these two goals pull in different directions depending on the loss.

Real-world applications:

• Hydrology and flood risk: Precipitation extremes drive flood assessment, and realistic tail behavior in downscaled fields feeds directly into hydrological models.
• Heatwave and temperature impact studies: Reliable minimum and maximum temperature fields underpin assessments of heat-related health and energy demand.
• Infrastructure planning: Fine-scale projections inform design standards for drainage, buildings, and other climate-sensitive assets.
• Agricultural and water-resource management: Local-scale variability and extremes affect crop stress and water availability.

Industry relevance: The probabilistic U-Net produces ensembles rather than single deterministic outputs, giving uncertainty quantification and latent-space interpretability that deterministic emulators lack. This is relevant to reinsurance and catastrophe modelling, climate risk consultancies, and energy utilities — sectors that consume large ensembles of high-resolution projections but are constrained by the cost of regional climate model runs.

Future Directions

• Combining losses: The authors explicitly suggest that combining afCRPS with MS-SSIM may yield a more balanced solution, since the two are complementary — afCRPS for spectral fidelity and overall CRPS, MS-SSIM for extremes.
• Extending beyond precipitation: Temperature distributions proved relatively insensitive to the loss choice, leaving open whether the findings generalize to other variables and regions.
• Better extreme reproduction with afCRPS: The tendency of afCRPS to overestimate extremes is unresolved and would need correction before operational use.
• Ensemble and calibration work: The model produces ensembles with substantial spread; the paper does not report formal calibration diagnostics beyond return-level coverage, leaving room for further verification.

Target Audience

Climate scientists and machine learning researchers working on statistical downscaling and climate emulation; practitioners in hydrology and climate risk who need high-resolution ensembles for impact studies; and graduate students or applied researchers interested in adapting generative segmentation architectures such as the probabilistic U-Net to geoscientific problems. Familiarity with U-Nets, latent-variable generative models, and standard forecast verification metrics is assumed.

Authors’ abstract

Climate models are limited by heavy computational costs, often producing outputs at coarse spatial resolutions, while many climate change impact studies require finer scales. Statistical downscaling bridges this gap, and we adapt the probabilistic U-Net for this task, combining a deterministic U-Net backbone with a variational latent space to capture aleatoric uncertainty. We evaluate four training objectives, afCRPS and WMSE-MS-SSIM with three settings for downscaling precipitation and temperature from $16\times$ coarser resolution. Our main finding is that WMSE-MS-SSIM performs well for extremes under certain settings, whereas afCRPS better captures spatial variability across scales.