SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\gamma_ {\epsilon} \left(y, y ^ {\prime}\right) = \exp \left(\left(f _ {\epsilon} (y) + g _ {\epsilon} \left(y ^ {\prime}\right) - \frac {1}{2} \| y - y ^ {\prime} \| ^ {2}\right) / \epsilon\right) d \mu_ {Y} (y) d \mu_ {Y ^ {\prime}} \left(y ^ {\prime}\right), \tag {10}
$$
in terms of potential functions  $f_{\epsilon}, g_{\epsilon}$  that are chosen to satisfy the marginal constraints. After finding these potentials, we can approximate the transport map using the barycentric projection  $T_{\gamma}(y) = \mathbb{E}_{\gamma}[Y' | Y = y]$ , for a plan  $\gamma \in \Pi(\mu_Y, \mu_{Y'})$  [2]. For the plan in Eq. (10), the map is given by
$$
T _ {\gamma_ {\epsilon}} (y) = \frac {\int y ^ {\prime} e ^ {\left(g _ {\epsilon} \left(y ^ {\prime}\right) - \frac {1}{2} \| y - y ^ {\prime} \| ^ {2}\right) / \epsilon} d \mu_ {Y} \left(y ^ {\prime}\right)}{\int e ^ {\left(g _ {\epsilon} \left(y ^ {\prime}\right) - \frac {1}{2} \| y - y ^ {\prime} \| ^ {2}\right) / \epsilon} d \mu_ {Y} \left(y ^ {\prime}\right)}. \tag {11}
$$
In this work, we estimate the potential functions  $f_{\epsilon}, g_{\epsilon}$  from samples, i.e., empirical approximations of the measures  $\mu_{Y}, \mu_{Y'}$ . Plugging in the estimated potentials in Eq. (11) defines an approximate transport map to push forward samples of  $\mu_{Y}$  to  $\mu_{Y'}$ . More details on the estimation of the OT map are provided in Appendix H.
A core advantage of this methodology is that it provides us with the flexibility of changing the cost function  $c$  in Eq. (8), and embed it with structural biases that one wishes to preserve in the push-forward distribution. Such direction is left for future work.
4.1 Data and setup
We showcase the efficacy and performance of the proposed approach on one- and two-dimensional fluid flow problems that are representative of the core difficulties present in numerical simulations of weather and climate. We consider the one-dimensional Kuramoto-Sivashinski (KS) equation and the two-dimensional Navier-Stokes (NS) equation under Kolmogorov forcing (details in Appendix F) in periodic domains. The low-fidelity (LF), low-resolution (LR) data  $(\mathcal{V}$  in Fig. 1(b)) is generated using a finite volume discretization in space [51] and a fractional discretization in time, while the high-fidelity
(HF), high-resolution (HR) data  $(\mathcal{X}$  in Fig. 1(b)) is simulated using a spectral discretization in space with an implicit-explicit scheme in time. Both schemes are implemented with jax-cfd and its finite-volume and spectral toolboxes [28, 48] respectively. After generating the HF data in HR, we run the LF solver using a spatial discretization that is  $8\times$  coarser (in each dimension) with permissible time steps. For NS, we additionally create a  $16\times$  coarser LFLR dataset by further downsampling by a factor of two the  $8\times$  LFLR data. See Appendix F for further details.
For both systems, the datasets consist of long trajectories generated with random initial conditions<sup>2</sup>, which are sufficiently downsampled in time to ensure that consecutive samples are decorrelated. We stress once more that even when the grids and time stamps of both methods are aligned, there is no pointwise correspondence between elements of  $\mathcal{X}$  and  $\mathcal{Y}$ . This arises from the different modeling biases inherent to the LF and HF solvers, which inevitably disrupt any short-term correspondence over the long time horizon in a strongly nonlinear dynamical setting.
Finally, we create the intermediate space  $\mathcal{Y}'$  in Fig. 1(b) by downsampling the HFHR data with a simple selection mask $^3$  (i.e., the map  $C'$ ). This creates the new HFLR dataset  $\mathcal{Y}'$  with the same resolution as  $\mathcal{Y}$ , but with the low-frequency bias structure of  $\mathcal{X}$  induced by the push-forward of  $C'$ .
Baselines and definitions. We define the following ablating variants of our proposed method
- Unconditional diffusion sampling (UncondDfn).
- Diffusion sampling conditioned on LFLR data without OT correction (Raw cDfn).
- [Main] Diffusion sampling conditioned on OT-corrected (HFLR) data  $(OT + cDfn)$ .
We additionally consider the following baselines to benchmark our method:
- Cubic interpolation approximating HR target using local third-order splines (Cubic).
- Vision transformer (ViT) [27] based deterministic super-resolution model.
- Bias correction and statistical disaggregation (BCSD), involving upsampling with cubic interpolation, followed by a quantile-matching debiasing step.
- CycleGAN, which is adapted from [86] to enable learning transformations between spaces of different dimensions (cycGAN).
- ClimAlign (adapted from [32]), in which the input is upsampled using cubic interpolation, and the debiasing step is performed using AlignFlow [34] (ClimAlign).
The first two baselines require paired data and, therefore, learn the upsampling map  $\mathcal{Y}' \to \mathcal{X}$  (i.e., HFLR to HFHR) and are composed with OT debiasing as factorized baselines. BCSD is a common approach used in the climate literature. The last two baselines present end-to-end alternatives and are trained directly on unpaired LFLR and HFHR samples. Further information about the implemented baselines can be found in Appendix D.
OT training. To learn the transport map in Eq. (11), we solve the entropic OT problem in Eq. (9) with  $\epsilon = 0.001$  using a Sinkhorn [21] iteration with Anderson acceleration and parallel updates. We use 90,000 i.i.d. samples of  $Y\in \mathcal{V}$  and  $Y^{\prime}\in \mathcal{V}^{\prime}$ , and perform 5000 iterations. Implementations are based on the ott-jax library [22].
Denoiser training and conditional sampling. The denoiser  $D_{\theta}$  is parametrized with a standard U-Net architecture similar to the one used in [67]. We additionally incorporate the preconditioning technique proposed in [45]. For  $s_t$  and  $\sigma_t$  schedules, we employ the variance-preserving (VP) scheme originally introduced in [73]. Furthermore, we adopt a data augmentation procedure to increase the effective training data size by taking advantage of the translation symmetries in the studied systems.
Samples are generated by solving the SDE based on the post-processed denoiser  $\tilde{D}_{\theta}$  using the Euler-Maruyama scheme with exponential time steps, i.e.,  $\{t_i\}$  is set such that  $\sigma(t_i) = \sigma_{\max}(\sigma_{\min} / \sigma_{\max})^{i/N}$  for  $i = \{0, \dots, N\}$ . The number of steps used,  $N$ , vary between systems and downscaling factors. More details regarding denoiser training and sampling are included in Appendix B.
Metrics. To quantitatively assess the quality of the resulting snapshots we compare a number of physical and statistical properties of the snapshots: (i) the energy spectrum, which measures the
2The presence of global attractors in both systems renders the exact initial conditions unimportant. It also guarantees sufficient coverage of the target distributions sampling from long trajectories.
<sup>3</sup>It is worth noting that careful consideration should be given to the choice of  $C'$  to avoid introducing aliasing, as this can potentially make the downscaling task more challenging.
(a) KS, conditional samples
Figure 2: (a) KS samples generated with diffusion model conditioned on LR information with and without OT correction applied, (b) empirical probability density function for relevant LR and HR samples in KS and (c) mode-wise log energy ratios with respect to the true samples (Eq. (13) without weighted sum) at  $8 \times$  downscaling for NS.
(b) KS, marginal PDFs
(c) NS  $8\times$  ,log energy ratios
Table 1: Metrics of the LFLR source and OT-corrected samples for KS and NS. The precise metric definitions are provided in Appendix C.
KS 8×	NS 8×	NS 16×
Metric	LFLR	OT-corrected	LFLR	OT-corrected	LFLR	OT-corrected
covRMSE ↓	0.343	0.081	0.458	0.083	0.477	0.079
MELRu ↓	0.201	0.020	1.254	0.013	0.600	0.016
MELRw ↓	0.144	0.020	0.196	0.026	0.200	0.025
KLD ↓	1.464	0.018	29.30	0.033	12.26	0.017
energy in each Fourier mode and thereby providing insights into the similarity between the generated and reference samples, (ii) a spatial covariance metric, which characterizes the spatial correlations within the snapshots, (iii) the KL-divergence (KLD) of the kernel density estimation for each point, which serves as a measure for the local structures (iv) the maximum mean discrepancy (MMD), and (v) the empirical Wasserstein-1 metric (Wass1). We present (i) below and leave the rest described in Appendix C as they are commonly used in the context of probabilistic modeling.
The energy spectrum is defined as
$$
E (k) = \sum_ {| \underline {{k}} | = k} | \hat {u} (\underline {{k}}) | ^ {2} = \sum_ {| \underline {{k}} | = k} \left| \sum_ {i} u \left(x _ {i}\right) \exp \left(- j 2 \pi \underline {{k}} \cdot x _ {i} / L\right) \right| ^ {2} \tag {12}
$$
where  $u$  is a snapshot system state, and  $k$  is the magnitude of the wave-number (wave-vector in 2D)  $\underline{k}$ . To assess the overall consistency of the spectrum between the generated and reference samples using a single scalar measure, we consider the mean energy log ratio (MELR):
$$
\operatorname {M E L R} = \sum_ {k} w _ {k} \left| \log \left(E _ {\text {p r e d}} (k) / E _ {\text {r e f}} (k)\right) \right|, \tag {13}
$$
where  $w_{k}$  represents the weight assigned to each  $k$ . We further define  $w_{k}^{\mathrm{unweighted}} = 1 / \mathrm{card}(k)$  and  $w_{k}^{\mathrm{weighted}} = E_{\mathrm{ref}}(k) / \sum_{k}E_{\mathrm{ref}}(k)$ . The latter skews more towards high-energy/low-frequency modes.
Effective debiasing via optimal transport. Table 1 shows that the OT map effectively corrects the statistical biases in the LF snapshots for all three experiments considered. Significant improvements are observed across all metrics, demonstrating that the OT map approximately achieves  $C_{\sharp}^{\prime}\mu_{X} \approx T_{\sharp}\mu_{Y}$  as elaborated in §3 (extra comparisons are included in Appendix H).
Indeed, the OT correction proves crucial for the success of our subsequent conditional sampling procedure: the unconditional diffusion samples may not have the correct energy spectrum (see
This definition is applied to each sample and averaged to obtain the metric (same for MELR below).
(d)
Figure 3: Example showing the vorticity field of samples debiased and super-resolved using different techniques at  $8 \times$  (top row) and  $16 \times$  (bottom row) downscaling factors. From left to right: (a) LR snapshots produced by the low-fidelity solver (input  $\bar{y}$  of Alg. 1), (b) OT-corrected snapshots ( $\bar{y}'$  in line 1 of Alg. 1), (c) BCSD applied to LR snapshots, (d) snapshots downscaled with cycle-GAN directly from LR snapshots, (e) ClimAlign applied to LR snapshots, (f) cubic interpolation of the OT-corrected snapshots, (g) deterministic upsample of the OT-corrected snapshots with ViT, (h) diffusion sample conditioned on the OT-corrected snapshots (output  $\bar{x}$  in Alg. 1, ours), and (i) two true HR samples in the training data with the closest Euclidean distance to the OT-corrected generated sample. The  $16 \times$  source is the same as the  $8 \times$  source but further downsampled by a factor of two. OT maps are computed independently between resolutions.
UncondDfn in Fig. 2(c), i.e. suffering from color shifts - a known problem for score-based diffusion models [72, 15]. The conditioning on OT corrected data serves as a sparse anchor which draws the diffusion trajectories to the correct statistics at sampling time. In fact, when conditioned on uncorrected data, the bias effectively pollutes the statistics of the samples (Raw cDfn in Table 2). Fig. 2(b) shows that the same pollution is present for the KS case, despite the unconditional sampler being unbiased.
In Appendix E, we present additional ablation studies that demonstrate the importance of OT correction in the factorized benchmarks.
Comparison vs. factorized alternatives. Fig. 3 displays NS samples generated by all benchmarked methods. Qualitatively, our method is able to provide highly realistic small-scale features. In comparison, we observe that Cubic expectedly yields the lowest quality results; the deterministic ViT produces samples with color shift and excessive smoothing, especially at  $16\times$  downscaling factor.
Quantitatively, our method outperforms all competitors in terms of MELR and KLD metrics in the NS tasks, while demonstrating consistently good performance in both  $8 \times$  and  $16 \times$  downscaling, despite the lack of recognizable features in the uncorrected LR data (Fig. 3(a) bottom) in the latter case. Other baselines, on the other hand, experience a significant performance drop. This showcases the value of having an unconditional prior to rely on when the conditioning provides limited information.
Comparison vs. end-to-end downscaling. Although the cycGAN baseline is capable of generating high-quality samples at  $8 \times$  downscaling (albeit with some smoothing) reflecting competitive metrics, we encountered persistent stability issues during training, particularly in the  $16 \times$  downscaling case.
Diffusion samples exhibit ample variability. Due to the probabilistic nature of our approach, we can observe from Table 2 that the OT-conditioned diffusion model provides some variability in the downscaling task, which increases when the downscaling factor increases. This variability provides a measure of uncertainty quantification in the generated snapshots as a result of the consistent formulation of our approach on probability spaces.
We introduced a two-stage probabilistic framework for the statistical downscaling problem. The framework performs a debiasing step to correct the low-frequency statistics, followed by an upsampling step using a conditional diffusion model. We demonstrate that when applied to idealized physical fluids, our method provides high-resolution samples whose statistics are physically correct, even when there is a mismatch in the low-frequency energy spectra between the low- and high
Table 2: Evaluation of downscaling methods for NS. The best metric values are highlighted in bold. Precise metric definitions (except MELR, given by Eq. (13)) are included in Appendix C.
Model	Var	covRMSE↓	MELRu↓	MELRw↓	KLD↓	Wass1↓	MMD↓
8× downscale
BCSD	0	0.31	0.67	0.25	2.19	0.23	0.10
cycGAN	0	0.15	0.08	0.05	1.62	0.32	0.08
ClimAlign	0	2.19	0.64	0.45	64.37	2.77	0.53
Raw+cDfn	0.27	0.46	0.79	0.37	73.16	1.04	0.42
OT+Cubic	0	0.12	0.52	0.06	1.46	0.42	0.10
OT+ViT	0	0.43	0.38	0.18	1.72	1.11	0.31
(ours) OT+cDfn	0.36	0.12	0.06	0.02	1.40	0.26	0.07
16× downscale
BCSD	0	0.34	0.67	0.25	2.17	0.21	0.11
cycGAN	0	0.32	1.14	0.28	2.05	0.48	0.13
ClimAlign	0	2.53	0.81	0.50	77.51	3.15	0.55
Raw+cDfn	1.07	0.46	0.54	0.30	93.87	0.99	0.39
OT+Cubic	0	0.25	0.55	0.13	7.30	0.85	0.20
OT+ViT	0	0.14	1.38	0.09	1.67	0.32	0.07
(ours) OT+cDfn	1.56	0.12	0.05	0.02	0.83	0.29	0.07
resolution data distributions. We have shown that our method is competitive and outperforms several commonly used alternative methods.
Future work will consider fine-tuning transport maps by adapting the map to the goal of conditional sampling, and introducing physically-motivated cost functions in the debiasing map. Moreover, we will address current limitations of the methodology, such as the high-computational complexity of learning OT-maps that scales quadratically with the size of the training set, and investigate the model's robustness to added noise in the collected samples as is found in weather and climate datasets. We will also further develop this methodology to cover other downscaling setups such as perfect prognosis [57] and spatio-temporal downscaling.
Statistical downscaling is important to weather and climate modeling. In this work, we propose a new method for improving the accuracy of high-resolution forecasts (on which risk assessment would be made) from low resolution climate modeling. Weather and climate research and other scientific communities in computational fluid dynamics will benefit from this work for its potential to reduce computational costs. We do not believe this research will disadvantage anyone.
The authors would like to sincerely thank Toby Bischoff, Katherine Deck, Nikola Kovachki, Andrew Stuart and Hongkai Zheng for many insightful and inspiring discussions that were vital to this work. RB gratefully acknowledges support from the Air Force Office of Scientific Research MURI on "Machine Learning and Physics-Based Modeling and Simulation" (award FA9550-20-1-0358), and a Department of Defense (DoD) Vannevar Bush Faculty Fellowship (award N00014-22-1-2790).
[1] J. T. Abatzoglou and T. J. Brown. A comparison of statistical downscaling methods suited for wildfire applications. International Journal of Climatology, 32(5):772-780, 2012. doi: https://doi.org/10.1002/joc.2312. URL https://rmets.onlinelibrary.wiley.com/doi/abs/10.1002/joc.2312.