Conditional diffusion models for downscaling and bias correction of Earth system model precipitation

TL;DR

The paper tackles the challenge of obtaining high-resolution, bias-corrected precipitation fields from coarse-resolution Earth System Models. It introduces a conditional diffusion framework that first maps observations and ESM data into a shared embedding space, then learns an inverse mapping with a conditional diffusion model trained exclusively on observational data to downscale and correct ESM fields while preserving large-scale climate signals. The approach yields sharper small-scale structure, better representation of extremes, and calibrated uncertainty, demonstrated across regions and multiple CMIP6 models, including future scenarios under SSP5-8.5, without retraining for each new ESM. This data-efficient method offers a controllable alternative to GANs and unconditional diffusion for climate impact assessments and supports robust regional risk analysis under changing climates.

Abstract

Climate change exacerbates extreme weather events like heavy rainfall and flooding. As these events cause severe socioeconomic damage, accurate high-resolution simulation of precipitation is imperative. However, existing Earth System Models (ESMs) struggle to resolve small-scale dynamics and suffer from biases. Traditional statistical bias correction and downscaling methods fall short in improving spatial structure, while recent deep learning methods lack controllability and suffer from unstable training. Here, we propose a machine learning framework for simultaneous bias correction and downscaling. We first map observational and ESM data to a shared embedding space, where both are unbiased towards each other, and then train a conditional diffusion model to reverse the mapping. Only observational data is used for the training, so that the diffusion model can be employed to correct and downscale any ESM field without need for retraining. Our approach ensures statistical fidelity and preserves spatial patterns larger than a chosen spatial correction scale. We demonstrate that our approach outperforms existing statistical and deep learning methods especially regarding extreme events.

Figures (34)

• Figure 1: Schematic overview of our approach. (A) Bias correction and downscaling can be formulated as a mapping $\omega$ from the ESM data space to the data space of observations (OBS) used for training. We first map both datasets to a shared embedding space and then learn the inverse of the mapping $f$ with a DM. We achieve a correction of the ESM data by applying $DM \circ g$. (B) Our framework allows to train a single model for bias correction and downscaling in a supervised way despite the unpaired nature of OBS and ESM fields. We construct functions $f,g$ that map $OBS \in \mathbf{V^{obs}}$ and $ESM \in \mathbf{V^{esm}}$ fields to a shared embedding space $\mathbf{V^{emb}}$. Note that this embedding space does not enforce pairing between individual fields, but a similar distribution between the embedded fields. By inverting $f$, we can rewrite $\omega$ as $\omega = f^{-1} \circ g$. We learn the inverse $f^{-1}$ with a conditional diffusion model. This model is trained (blue arrow) on pairs of observational data to approximate the map from $f(OBS)$ to OBS. Because $f(OBS)$ and $g(ESM)$ share the embedding space (and are identically distributed by construction), we can evaluate (green arrow) the $DM$ on the embedded ESM data $g(ESM)$ and thereby approximate the bias correction and downscaling function $\omega = f^{-1} \circ g \approx DM \circ g$, without the need of paired data between OBS and ESM. The indices $i,j$ highlight that the two exemplary fields $ESM_j$ and $ESM_i$ are not paired. (C) Left: Training process of the conditional DM $DM\approx f^{-1}$. Note that the individual samples of the input OBS and their embeddings $f(OBS)$, as well as the embeddings $f(OBS)$ and the output of $DM\approx f^{-1}$ are paired, respectively. Right: Inference process of $DM\approx f^{-1}$. In this case, the individual samples of the input ESM, their embeddings $g(ESM)$, and the output of $DM\approx f^{-1}$ are paired, respectively. It is not necessary for the training embedding samples to be paired with the inference embedding samples. See fig. \ref{['fig:diffusion_model']} for details.
• Figure 2: Power spectral densities (PSDs) for different choices noising scale of the diffusion model. The noising scale $s$ (dashed line) is a hyperparameter that can be chosen depending on the ESM and observational datasets, as well as on the specific task. For the maximal choice of $s$(A) all information in the observations (ERA5) and model simulations (GFDL) is noised and thereby destroyed. Conditioning on pure noise makes the task equivalent to unconditional image generation. The diffusion model will learn to generate observational fields with no relation to the ESM fields. When $s$ is chosen to be minimal, there will be no noising and the conditional generation will directly replicate the condition, i.e. the ESM field. In (B) we chose $s$ as the point where the PSDs of the observational and simulated datasets intersect. We then apply sufficiently many forward diffusion noising steps to both datasets, destroying small-scale structure until they agree in the PSD. We call scales smaller than $s$ small scales and scales larger than $s$ large scales. In (C) and (D), the effects of choosing a smaller noise scale $s$ are shown. Prior knowledge about the ESM or its accuracy can also guide the choice of $s$.
• Figure 3: Comparative visualization of individual randomly selected samples. Each row presents three samples of the same dataset. The top row shows GFDL ESM4 data, bilinearly upsampled to 0.25° to match the other fields. The second row shows QM-corrected and the third row diffusion model-corrected GFDL fields. The bottom row shows samples of the original ERA5 data, which are unpaired to the GFDL fields above. Visual inspection shows that the diffusion model correction greatly improves upon the QM correction in terms of producing realistic spatial patterns, since the QM-corrected fields remain way too blurry compared to the HR ERA5 data. The overall large-scale patterns are preserved by the DM. There is no visual difference between the details and sharpness of diffusion model-corrected GFDL fields compared to ERA5.
• Figure 4: Comparison of climatologies and model biases. The first row shows the climatology of (A) the diffusion model-corrected GFDL at 0.25°, (B) the GFDL ESM4 model, upsampled to 0.25° and (C) the 0.25° ERA5 data. The second row shows the bias of the GFDL and the QM- and diffusion model-corrections, defined as the difference between long-term temporal averages of all validation samples. Specifically, the temporally averaged bias fields with respect to ERA5 are shown for (D) the diffusion model correction, (E) the uncorrected GFDL and (F) the QM correction. Results indicate a substantial improvement of our diffusion model (A) and the benchmark (C) over just upsampling GFDL to 0.25°. The absolute bias on top of each panel is given by the mean absolute value of the differences over the spatial and temporal dimension with respect to ERA5.
• Figure 5: Evaluation of our diffusion model's performance for downscaling and bias correction. Comparison of GFDL (bilinearly upsampled to 0.25°) (orange) and ERA5 (black) to diffusion model-corrected GFDL (magenta) and QM-corrected GFDL fields (blue) as our benchmark. The Power spectral density (PSD) plot (A) shows that the diffusion model corrects the small-scale spatial details far better than our benchmark. The spectrum aligns very well with the high resolution ERA5 target data. The histograms (B) as well as the latitude (C) and longitude (D) profiles show substantial improvements compared to the uncorrected GFDL data.