SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhong Yi Wan*
Mountain View, CA 94043, USA
wanzy@google.com
Yi-fan Chen
yifanchen@google.com
Anudhyan Boral
anudhyan@google.com
Ricardo Baptista*
Pasadena, CA 91106, USA
rsb@caltech.edu
John Roberts Anderson
janders@google.com
Fei Sha
fsha@google.com
Leonardo Zepeda-Núñez
lzepedanunez@google.com
We introduce a two-stage probabilistic framework for statistical downscaling using unpaired data. Statistical downscaling seeks a probabilistic map to transform low-resolution data from a biased coarse-grained numerical scheme to high-resolution data that is consistent with a high-fidelity scheme. Our framework tackles the problem by composing two transformations: (i) a debiasing step via an optimal transport map, and (ii) an upsampling step achieved by a probabilistic diffusion model with a posteriori conditional sampling. This approach characterizes a conditional distribution without needing paired data, and faithfully recovers relevant physical statistics from biased samples. We demonstrate the utility of the proposed approach on one- and two-dimensional fluid flow problems, which are representative of the core difficulties present in numerical simulations of weather and climate. Our method produces realistic high-resolution outputs from low-resolution inputs, by upsampling resolutions of  $8 \times$  and  $16 \times$ . Moreover, our procedure correctly matches the statistics of physical quantities, even when the low-frequency content of the inputs and outputs do not match, a crucial but difficult-to-satisfy assumption needed by current state-of-the-art alternatives. Code for this work is available at: https://github.com/google-research/swirl-dynamics/tree/main/swirl_dynamics/projects/probablistic_diffusion.
Statistical downscaling is crucial to understanding and correlating simulations of complex dynamical systems at multiple resolutions. For example, in climate modeling, the computational complexity of general circulation models (GCMs) [4] grows rapidly with resolution. This severely limits the resolution of long-running climate simulations. Consequently, accurate predictions (as in forecasting localized, regional and short-term weather conditions) need to be downscaled from coarser lower-resolution models' outputs. This is a challenging task: coarser models do not resolve small-scale dynamics, thus creating bias [16, 69, 84]. They also lack the necessary physical details (for instance, regional weather depends heavily on local topography) to be of practical use for regional or local climate impact studies [33, 36], such as the prediction or risk assessment of extreme flooding [35, 44], heat waves [59], or wildfires [1].
At the most abstract level, statistical downscaling [81, 82] learns a map from low- to high-resolution data. However, it has several unique challenges. First, unlike supervised machine learning (ML), there is no natural pairing of samples from the low-resolution model (such as climate models [23]) with samples from higher-resolution ones (such as weather models that assimilate observations [40]). Even in simplified cases of idealized fluids problems, one cannot naively align the simulations in time, due to the chaotic behavior of the models: two simulations with very close initial conditions will diverge rapidly. Several recent studies in climate sciences have relied on synthetically generated paired datasets. The synthesis process, however, requires accessing both low- and high-resolution models and either (re)running costly high-resolution models while respecting the physical quantities in the low-resolution simulations [25, 43] or (re)running low-resolution models with additional terms nudging the outputs towards high-resolution trajectories [12]. In short, requiring data in correspondence for training severely limits the potential applicability of supervised ML methodologies in practice, despite their promising results [37, 39, 60, 66, 38].
Second, unlike the setting of (image) super-resolution [26], in which an ML model learns the (pseudo) inverse of a downsampling operator [13, 78], downscaling additionally needs to correct the bias. This difference is depicted in Fig. 1(a). Super-resolution can be recast as frequency extrapolation [10], in which the model reconstructs high-frequency contents, while matching the low-frequency contents of a low-resolution input. However, the restriction of the target high-resolution data may not match the distribution of the low-resolution data in Fourier space [49]. Therefore, debiasing is necessary to correct the Fourier spectrum of the low-resolution input to render it admissible for the target distribution (moving solid red to solid blue lines with the dashed blue extrapolation in Fig. 1). Debiasing allows us to address the crucial yet challenging prerequisite of aligning the low-frequency statistics between the low- and high-resolution datasets.
Given these two difficulties, statistical downscaling should be more naturally framed as matching two probability distributions linked by an unknown map; such a map emerges from both distributions representing the same underlying physical system, albeit with different characterizations of the system's statistics at multiple spatial and temporal resolutions. The core challenge is then: how do we structure the downscaling map so that the (probabilistic) matching can effectively remediate the bias introduced by the coarser, i.e., the low-resolution, data distribution?
Thus, the main idea behind our work is to introduce a debiasing step so that the debiased (yet, still coarser) distribution is closer to the target distribution of the high-resolution data. This step results in an intermediate representation for the data that preserves the correct statistics needed in the follow-up step of upsampling to yield the high-resolution distribution. In contrast to recent works on distribution matching for unpaired image-to-image translation [86] and climate modeling [32], the additional structure our work imposes on learning the mapping prevents the bias in the low-resolution data from polluting the upsampling step. We review those approaches in §2 and compare to them in §4.
Concretely, we propose a new probabilistic formulation for the downscaling problem that handles unpaired data directly, based on a factorization of the unknown map linking both low- and high-resolution distributions. This factorization is depicted in Fig. 1(b). By appropriately restricting the maps in the factorization, we rewrite the downscaling map as the composition of two procedures: a debiasing step performed using an optimal transport map [21], which couples the data distributions and corrects the biases of the low-resolution snapshots; followed by an upsampling step performed using conditional probabilistic diffusion models, which have produced state-of-the-art results for image synthesis and flow construction [5, 52, 71, 73].
(b)
Figure 1: (a) Upsampling (super-resolution) as frequency extrapolation in the Fourier domain. The model extrapolates low-frequency content to higher-frequencies (dashed blue). The debiasing map corrects the biased low-frequency content (solid red). (b) Illustration of the proposed framework where  $\mathcal{X}$  is the space of high-resolution data,  $\mathcal{V}$  is the space of low-resolution data,  $C$  is an unknown nonlinear map linking  $\mathcal{X}$  and  $\mathcal{V}$ ,  $C'$  is a known linear downsampling map,  $\mathcal{V}'$  is an intermediate (low-resolution) space induced by the image of  $C'$ , and  $T$  is an invertible debiasing map such that  $C$  can be factorized as  $T^{-1} \circ C'$ . The conditional probabilities  $p(x|C'x = y')$  are used for the probabilistic upsampling procedure.
We showcase the performance of our framework on idealized fluids problems that exhibit the same core difficulty present in atmospheric flows. We show that our framework is able to generate realistic snapshots that are faithful to the physical statistics, while outperforming several baselines.
The most direct approach to upsampling low-resolution data is to learn a low- to high-resolution mapping via paired data when it is possible to collect such data. For complex dynamical systems, several methods carefully manipulate high- and low-resolution models, either by nudging or by enforcing boundary conditions, to produce paired data without introducing spectral biases [12, 25]. Alternatively, if one has strong prior knowledge about the process of downsampling, optimization methods can solve an inverse problem to directly estimate the high-resolution data, leveraging prior assumptions such as sparsity in compressive sensing [9, 10] or translation invariance [42].
In our setting, there is no straightforward way to obtain paired data due to the nature of the problem (i.e., turbulent flows, with characteristically different statistics across a large span of spatio-temporal scales). In the weather and climate literature (see [79] for an extensive overview), prior knowledge can be exploited to downscale specific variables [81]. One of the most predominant methods of this type is bias-correction spatial disaggregation (BCSD), which combines traditional spline interpolation with a quantile matching bias correction [56], and linear models [41]. Recently, several studies have used ML to downscale physical quantities such as precipitation [78], but without quantifying the prediction uncertainty. Yet, a generally applicable method to downscale arbitrary variables is lacking.
Another difficulty is to remove the bias in the low resolution data. This is an instance of domain adaptation, a topic popularly studied in computer vision. Recent work has used generative models such as GANs and diffusion models to bridge the gap between two domains [5, 7, 14, 58, 60, 62, 68, 74, 83, 85]. A popular domain alignment method that was used in [32] for downscaling weather data is AlignFlow [34]. This approach learns normalizing flows for source and target data of the same dimension, and uses their common latent space to move across domains. The advantage of those methods is that they do not require training data from two domains in correspondence. Many of those approaches are related to optimal transport (OT), a rigorous mathematical framework for learning maps between two domains without paired data [80]. Recent computational advances in OT for discrete (i.e., empirical) measures [21, 64] have resulted in a wide set of methods for domain adaptation [20, 31]. Despite their empirical success with careful choices of regularization, their use alone for high-dimensional images has remained limited [61].
Our work uses diffusion models to perform upsampling after a debiasing step implemented with OT. We avoid common issues from GANs [75] and flow-based methods [54], which include oversmoothing, mode collapse and large model footprints [24, 52]. Also, due to the debiasing map, which matches the low-frequency content in distribution (see Fig. 1(a)), we do not need to explicitly impose that the low-frequency power spectra of the two datasets match like some competing methods do [5]. Compared to formulations that perform upsampling and debiasing simultaneously [5, 78], our
framework performs these two tasks separately, by only training (and independently validating) a single probabilistic diffusion model for the high-resolution data once. This allows us to quickly assess different modeling choices, such as the linear downsampling map, by combining the diffusion model with different debiasing maps. Lastly, in comparison to other two-stage approaches [5, 32], debiasing is conducted at low-resolutions, which is less expensive as it is performed on a much smaller space, and more efficient as it is not hampered from spurious biases introduced by interpolation techniques.
Setup We consider two spaces: the high-fidelity, high-resolution space  $\mathcal{X} = \mathbb{R}^d$  and the low-fidelity, low-resolution space  $\mathcal{Y} = \mathbb{R}^{d'}$ , where we suppose that  $d > d'$ . We model the elements  $X \in \mathcal{X}$  and  $Y \in \mathcal{Y}$  as random variables with marginal distributions,  $\mu_X$  and  $\mu_Y$ , respectively. In addition, we suppose there is a statistical model relating the  $X$  and  $Y$  variables via  $C \colon \mathcal{X} \to \mathcal{Y}$ , an unknown and possibly nonlinear, downsampling map. See Fig. 1(b) for a diagram.
Given an observed realization  $\bar{y} \in \mathcal{V}$ , which we refer to as a snapshot, we formulate downscaling as the problem of sampling from the conditional probability distribution  $p(x|E_{\bar{y}})$  for the event  $E_{\bar{y}} := \{x \in \mathcal{X} | C(x) = \bar{y}\}$ , which we denote by  $p(x|C(x) = \bar{y})$ . Our objective is to sample this distribution given only access to marginal samples of  $X$  and  $Y$ .
Main idea In general, downscaling is an ill-posed problem given that the joint distribution of  $X$  and  $Y$  is not prescribed by a known statistical model. Therefore, we seek an approximation to  $C$  so the statistical properties of  $X$  are preserved given samples of  $\mu_{Y}$ . In particular, such a map should satisfy  $C_{\sharp} \mu_{X} = \mu_{Y}$ , where  $C_{\sharp} \mu_{X}$  denotes the push-forward measure of  $\mu_{X}$  through  $C$ .
In this work, we impose a structured ansatz to approximate  $C$ . Specifically, we factorize the map  $C$  as the composition of a known and linear downsampling map  $C'$ , and an invertible debiasing map  $T$ :
$$
C = T ^ {- 1} \circ C ^ {\prime}, \quad \text {s u c h t h a t} \quad \left(T ^ {- 1} \circ C ^ {\prime}\right) _ {\sharp} \mu_ {X} = \mu_ {Y}, \tag {1}
$$
or alternatively,  $C_{\sharp}^{\prime}\mu_{X} = T_{\sharp}\mu_{Y}$ . This factorization decouples and explicitly addresses two entangled goals in downscaling: debiasing and upsampling. We discuss the advantage of such factorization, after sketching how  $C^\prime$  and  $T$  are implemented.
The range of the downsampling map  $C^\prime \colon \mathcal{X}\to \mathcal{Y}^\prime$  defines an intermediate space  $\mathcal{V}' = \mathbb{R}^{d'}$  of high-fidelity low-resolution samples with measure  $\mu_{Y^{\prime}}$ . Moreover, the joint space  $\mathcal{X}\times \mathcal{Y}'$  is built by projecting samples of  $X$  into  $\mathcal{V}'$ , i.e.,  $(x,y^{\prime}) = (x,C^{\prime}x)\in \mathcal{X}\times \mathcal{Y}^{\prime}$ ; see Fig. 1(b). Using these spaces, we decompose the domain adaptation problem into the following three sub-problems:
1. High-resolution prior: Estimate the marginal density  $p(x)$ ;
2. Conditional modeling: For the joint variables  $X \times Y'$ , approximate  $p(x|C'x = y')$ ;
3. Debiasing: Compute a transport map such that  $T_{\sharp}\mu_{Y} = \bar{C}_{\sharp}^{\prime}\mu_{X}$ .
For the first sub-problem, we train an unconditional model to approximate  $\mu_{X}$ , or  $p(x)$ , as explained in §3.1. For the second sub-problem, we leverage the prior model and  $y' \in \mathcal{Y}'$  to build a model for a posteriori conditional sampling of  $p(x|C'x = y')$ , which allows us to upsample snapshots from  $\mathcal{Y}'$  to  $\mathcal{X}$ , as explained in §3.2. For the third sub-problem, we use domain adaptation to shift the resulting model from the source domain  $\mathcal{X} \times \mathcal{Y}'$  to the target domain  $\mathcal{X} \times \mathcal{Y}$ , for which there is no labeled data. For such a task, we build a transport map  $T: \mathcal{Y} \to \mathcal{Y}'$  satisfying the condition that  $T_{\sharp}\mu_{Y} = \mu_{Y'} = C_{\sharp}'\mu_{X}$ . This map is found by solving an optimal transport problem, which we explain in §3.3.
Lastly, we merge the solutions to the sub-problems to arrive at our core downscaling methodology, which is summarized in Alg. 1. In particular, given a low-fidelity and low-resolution sample  $\overline{y}$ , we use the optimal transport map  $T$  to project the sample to the high-fidelity space  $\overline{y}' = T(\overline{y})$  and use the conditional model to sample  $p(x|C'x = \overline{y}')$ . The resulting samples are contained in the high-fidelity and high-resolution space.
The factorization in Eq. (1) has several advantages. We do not require a cycle-consistency type of loss [34, 86]: the consistency condition is automatically enforced by Eq. (1) and the conditional sampling. By using a linear downsampling map  $C'$ , it is trivial to create the intermediate space  $\mathcal{Y}'$ , while rendering the conditional sampling tractable: conditional sampling with a nonlinear map is
Algorithm 1:Downscaling from  $\bar{y}\in \mathcal{V}$  to  $\bar{x} (\bar{y})\in \mathcal{X}$
Input: low-fidelity, low-resolution sample:  $y$
1. Compute the debiased term  $\bar{y}^{\prime} = T_{\gamma}(\bar{y})$  using the barycentric approximation in Eq. (11).
2. Modify the scoring function in Eq. (6) using the conditional denoiser in Eq. (7).
3. Solve the reverse SDE in Eq. (5), and obtain  $\bar{x}$ .
Output: high-fidelity, high-resolution sample :  $\bar{x}$
often more expensive and it requires more involved tuning [17, 18]. The factorization also allows us to compute the debiasing map in a considerably lower dimensional space, which conveniently requires less data to cover the full distribution, and fewer iterations to find the optimal map [21].
3.1 High-resolution prior
To approximate the prior of the high-resolution snapshots we use a probabilistic diffusion model, which is known to avoid several drawbacks of other generative models used for super-resolution [52], while providing greater flexibility for a posteriori conditioning [17, 30, 46, 47].
Intuitively, diffusion-based generative models involves iteratively transforming samples from an initial noise distribution  $p_T$  into ones from the target data distribution  $p_0 = p_{\mathrm{data}}$ . Noise is removed sequentially such that samples follow a family of marginal distributions  $p_t(x_t; \sigma_t)$  for decreasing diffusion times  $t$  and noise levels  $\sigma_t$ . Conveniently, such distributions are given by a forward noisng process that is described by the stochastic differential equation (SDE) [45, 73]
$$
d x _ {t} = f \left(x _ {t}, t\right) d t + g \left(x _ {t}, t\right) d W _ {t}, \tag {2}
$$
with drift  $f$ , diffusion coefficient  $g$ , and the standard Wiener process  $W_{t}$ . Following [45], we set
$$
f \left(x _ {t}, t\right) = f (t) x _ {t} := \frac {\dot {s} _ {t}}{s _ {t}} x _ {t}, \quad \text {a n d} \quad g \left(x _ {t}, t\right) = g (t) := s _ {t} \sqrt {2 \dot {\sigma} _ {t} \sigma_ {t}}. \tag {3}
$$
Solving the SDE in Eq. (2) forward in time with an initial condition  $x_0$  leads to the Gaussian perturbation kernel  $p(x_{t}|x_{0}) = \mathcal{N}(x_{t};s_{t}x_{0},s_{t}^{2}\sigma_{t}^{2}\mathbf{I})$ . Integrating the kernel over the data distribution  $p_0(x_0) = p_{\mathrm{data}}$ , we obtain the marginal distribution  $p_t(x_t)$  at any  $t$ . As such, one may prescribe the profiles of  $s_t$  and  $\sigma_t$  so that  $p_0 = p_{\mathrm{data}}$  (with  $s_0 = 1, \sigma_0 = 0$ ), and more importantly
$$
p _ {T} \left(x _ {T}\right) \approx \mathcal {N} \left(x _ {T}; 0, s _ {T} ^ {2} \sigma_ {T} ^ {2} \mathbf {I}\right), \tag {4}
$$
i.e., the distribution at the terminal time  $T$  becomes indistinguishable from an isotropic, zero-mean Gaussian. To sample from  $p_{\mathrm{data}}$ , we utilize the fact that the reverse-time SDE
$$
d x _ {t} = \left[ f (t) x _ {t} - g (t) ^ {2} \nabla_ {x _ {t}} \log p _ {t} \left(x _ {t}\right) \right] d t + g (t) d W _ {t}, \tag {5}
$$
has the same marginals as Eq. (2). Thus, by solving Eq. (5) backwards using Eq. (4) as the final condition at time  $T$ , we obtain samples from  $p_{\mathrm{data}}$  at  $t = 0$ .
Therefore, the problem is reduced to estimating the score function  $\nabla_{x_t}\log p_t(x_t)$  resulting from  $p_{\mathrm{data}}$  and the prescribed diffusion schedule  $(s_t,\sigma_t)$ . We adopt the denoising formulation in [45] and learn a neural network  $D_{\theta}(x_0 + \varepsilon_t,\sigma_t)$ , where  $\theta$  denotes the network parameters. The learning seeks to minimize the  $L_{2}$ -error in predicting the true sample  $x_0$  given a noise level  $\sigma_t$  and the sample noised with  $\varepsilon_t = \sigma_t\varepsilon$  where  $\varepsilon$  is drawn from a standard Gaussian. The score can then be readily obtained from the denoiser  $D_{\theta}$  via the asymptotic relation (i.e., Tweedie's formula [29])
$$
\nabla_ {x _ {t}} \log p _ {t} \left(x _ {t}\right) \approx \frac {D _ {\theta} \left(\hat {x} _ {t} , \sigma_ {t}\right) - \hat {x} _ {t}}{s _ {t} \sigma_ {t} ^ {2}}, \quad \hat {x} _ {t} = x _ {t} / s _ {t}. \tag {6}
$$
3.2 A posteriori conditioning via post-processed denoiser
We seek to super-resolve a low-resolution snapshot  $\bar{y}^{\prime}\in \mathcal{V}^{\prime}$  to a high-resolution one by leveraging the high-resolution prior modeled by the diffusion model introduced above. Abstractly, our goal is to sample from  $p(x_0|E_{\bar{y} '})$  , where  $E_{\bar{y} '}^{\prime} = \{x_0:C'x_0 = \bar{y} '\}$  . Following [30], this may be approximated by modifying the learned denoiser  $D_{\theta}$  at inference time (see Appendix A for more details):
$$
\tilde {D} _ {\theta} \left(\hat {x} _ {t}, \sigma_ {t}\right) = \left(C ^ {\prime}\right) ^ {\dagger} \bar {y} ^ {\prime} + \left(I - V V ^ {T}\right) \left[ D _ {\theta} \left(\hat {x} _ {t}, \sigma_ {t}\right) - \alpha \nabla_ {\hat {x} _ {t}} \| C ^ {\prime} D _ {\theta} \left(\hat {x} _ {t}, \sigma_ {t}\right) - \bar {y} ^ {\prime} \| ^ {2} \right], \tag {7}
$$
where  $(C')^\dagger = V\Sigma^{-1}U^T$  is the pseudo-inverse of  $C'$  based on its singular value decomposition (SVD)  $C' = U\Sigma V^T$ , and  $\alpha$  is a hyperparameter that is empirically tuned. The  $\tilde{D}_{\theta}$  defined in Eq. (7) directly replaces  $D_{\theta}$  in Eq. (6) to construct a conditional score function  $\nabla_{x_t}\log p_t(x_t|E_{\tilde{y}'}')$  that facilitates the sampling of  $p(x_0|E_{\tilde{y}'}')$  using the reverse-time SDE in Eq. (5).
3.3 Debiasing via optimal transport
In order to upsample a biased low-resolution data  $\overline{y} \in \mathcal{V}$ , we first seek to find a mapping  $T$  such that  $\overline{y}' = T(\overline{y}) \in \mathcal{Y}'$  is a representative sample from the distribution of unbiased low-resolution data. Among the infinitely many maps that satisfy this condition, the framework of optimal transport (OT) selects a map by minimizing an integrated transportation distance based on the cost function  $c: \mathcal{Y} \times \mathcal{Y}' \to \mathbb{R}^+$ . The function  $c(y, y')$  defines the cost of moving one unit of probability mass from  $y'$  to  $y$ . By treating  $Y, Y'$  as random variables on  $\mathcal{Y}, \mathcal{Y}'$  with measures  $\mu_Y, \mu_{Y'}$ , respectively, the OT map is given by the solution to the Monge problem
$$
\min  _ {T} \left\{\int c (y, T (y)) d \mu_ {Y} (y): T _ {\sharp} \mu_ {Y} = \mu_ {Y ^ {\prime}} \right\}. \tag {8}
$$
In practice, directly solving the Monge problem is hard and may not even admit a solution [80]. One common relaxation of Eq. (8) is to seek a joint distribution, known as a coupling or transport plan, which relates the underlying random variables [80]. A valid plan is a probability measure  $\gamma$  on  $\mathcal{Y} \times \mathcal{Y}'$  with marginals  $\mu_{Y}$  and  $\mu_{Y'}$ . To efficiently estimate the plan when the  $c$  is the quadratic cost (i.e.,  $c(y, y') = \frac{1}{2} \|y - y'\|^2$ ), we solve the entropy regularized problem
$$
\inf  _ {\gamma \in \Pi \left(\mu_ {Y}, \mu_ {Y ^ {\prime}}\right)} \int \frac {1}{2} \| y - y ^ {\prime} \| ^ {2} d \gamma (y, y ^ {\prime}) + \epsilon D _ {\mathrm {K L}} (\gamma | | \mu_ {Y ^ {\prime}} \otimes \mu_ {Y}), \tag {9}
$$
where  $D_{\mathrm{KL}}$  denotes the KL divergence, and  $\epsilon > 0$  is a small regularization parameter, using the Sinkhorn's algorithm [21], which leverages the structure of the optimal plan to solve Eq. (9) with small runtime complexity [3]. The solution to Eq. (9) is the transport plan  $\gamma_{\epsilon} \in \Pi(\mu, \nu)$  given by