SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theo J. Adrai
Technion-Israel Institute of Technology
theoad@cs.technion.ac.il
Michael Elad
elad@cs.technion.ac.il
Guy Ohayon
ohayonguy@cs.technion.ac.il
Tomer Michaeli
tomer.m@ee.technion.ac.il
We propose an image restoration algorithm that can control the perceptual quality and/or the mean square error (MSE) of any pre-trained model, trading one over the other at test time. Our algorithm is few-shot: Given about a dozen images restored by the model, it can significantly improve the perceptual quality and/or the MSE of the model for newly restored images without further training. Our approach is motivated by a recent theoretical result that links between the minimum MSE (MMSE) predictor and the predictor that minimizes the MSE under a perfect perceptual quality constraint. Specifically, it has been shown that the latter can be obtained by optimally transporting the output of the former, such that its distribution matches the source data. Thus, to improve the perceptual quality of a predictor that was originally trained to minimize MSE, we approximate the optimal transport by a linear transformation in the latent space of a variational auto-encoder, which we compute in closed-form using empirical means and covariances. Going beyond the theory, we find that applying the same procedure on models that were initially trained to achieve high perceptual quality, typically improves their perceptual quality even further. And by interpolating the results with the original output of the model, we can improve their MSE on the expense of perceptual quality. We illustrate our method on a variety of degradations applied to general content images of arbitrary dimensions.
Many image restoration algorithms aim to recover a clean source image from its degraded version. The performance of such algorithms is often evaluated in terms of their average distortion, which measures the discrepancy between restored images and their corresponding clean sources, as well as perceptual quality, which refers to the extent to which restored images resemble natural images. The work in [2] exposed a fundamental trade-off between distortion and perceptual quality, where the latter is measured using a perceptual index that quantifies the statistical divergence between the distribution of restored images and the distribution of natural images. The
Figure 1: The  $\mathcal{W}_2$ -MSE trade-off [1].
$\mathbf{x}_{\downarrow 4} + \mathbf{n}_{\sigma = 25}$
$\mathbf{x}^*$  (DDRM)
$\hat{\mathbf{x}}_0$  (ours)
x (original)
Figure 2: Our few-shot algorithm improves the visual quality of any estimator at test time. For example, we can improve the photo-realism of DDRM [3] even further.
trade-off curve reveals the predictor that achieves the lowest possible distortion, denoted as  $\mathbf{D}_{\mathrm{max}}$  while maintaining perfect perceptual quality (refer to fig. 1).
Following the methodology introduced by [2], it has become common practice to compare restoration methods on the perception-distortion (PD) plane, with many methods aiming to reach the elusive  $\mathbf{D}_{\mathrm{max}}$  point. In this paper, we present a practical approach to approximate the  $\mathbf{D}_{\mathrm{max}}$  predictor where distortion is measured using the MSE and perceptual quality is measured by the Wasserstein-2 distance  $(\mathcal{W}_2)$  between the distributions of restored and real images. Our approach is based on the recent work [1] which demonstrated that the  $\mathbf{D}_{\mathrm{max}}$  predictor can be obtained using optimal transport (OT) from the output distribution of the MMSE predictor to the distribution of natural images. By applying an optimal transport plan to an MMSE restoration resulting from a degraded image, we can produce  $\mathbf{D}_{\mathrm{max}}$  estimations by transporting the MMSE restored estimate.
Although progress has been made in finding OT plans between image distributions [4-6], it remains challenging task, particularly for high-dimensional distributions. Therefore, we propose an approximation method for the  $\mathbf{D}_{\mathrm{max}}$  estimator by performing transportation in the latent space of a pre-trained auto-encoder. A similar strategy was successfully employed in the context of image generation in [7], showing effectiveness in reducing complexity and preserving details.
Inspired by the style transfer literature [8-10], we assume that the latent representations follow a Multivariate Gaussian (MVG) distribution. Thus, by considering the first and second-order statistics of the embedded MMSE estimates and embedded natural images, we can compute the well-known closed form solution of the OT operator between two Gaussians. To further reduce complexity, we make additional assumptions about the structure of the latent covariance matrices, enabling the computation of the OT operator with as few as 10 unpaired MMSE restored and clean samples. This approach leads to a few-shot algorithm that significantly enhances visual quality.
Interestingly, our method can even improve the visual quality of generative models that were trained to achieve high perceptual quality in the first place (see fig. 2). Furthermore, by adjusting a single interpolation parameter, we can trade off perception for distortion, resulting in marginal improvements in the distortion performance of some regression models that were trained to prioritize source fidelity. We demonstrate the improved photo-realism of our approach on a variety of tasks and models, including GAN and diffusion-based methods, using high-resolution (e.g.,  $512^{2}$  px) general-content images with arbitrary aspect ratios. $^{1}$
Throughout the paper, we distinguish between two kinds of restoration algorithms: distortion and perception focused. The former category includes traditional methods that minimize distortion (e.g.,
1Our code is publicly available at https://github.com/theoad/dot-dmax.
MSE) [11-15]. The latter category includes more recent works that usually involve generative models like Generative Adversarial Networks (GANs) [16-18], or diffusion-based techniques [3, 19].
This paper searches for the theoretical  $\mathbf{D}_{\mathrm{max}}$  estimator which minimizes the MSE under a perfect perceptual quality constraint. However, our method is not a stand-alone restoration algorithm, i.e., its input is not the degraded image. Rather, it can be applied on top of any existing predictor. We provide a new way to potentially improve the performance (either MSE or perceptual) of any given estimator in a few-shot, plug-and-play fashion.
Provided with an image latent representation method (e.g., an auto-encoder), our algorithm applies a linear transformation on all the overlapping patches of its input (after encoding). In this regard, its functioning is not far from classical image restoration methods [14, 15, 20].
2.1 Wasserstein-2 transport
While many successful approaches exist to compute the  $\mathcal{W}_2$  distance between discrete, low or medium dimensional densities [21, 22], it is far more challenging to determine an optimal transport plan in the continuous, high-dimensional setting. In fact, the task of even computing the Wasserstein distance (without its optimal plan) on empirical distributions has drawn significant attention with WGANs [23]. Thus, computing the transport operator requires to optimize the  $\mathcal{W}_2$  distance but with an additional ordering constraint on the generator, which proved to be a challenging task when dealing with real-world data sets [4-6]. An attempt to sidestep this difficulty would be to use the Gelbrich distance [1, 24], which lower bounds the  $\mathcal{W}_2$  distance and depends only on the first two moments of the distributions. Nevertheless, it is only a good estimate when the support of the distributions has elliptical level sets. To address this, one can find a low-dimensional embedding where (i) high dimensionality is no longer an issue, (ii) the distributions are not degenerated, and (iii) the Gelbrich distance equals the Wasserstein-2 distance. A possible option would be to use the bottleneck of an auto-encoder. This approach was adopted by style transfer works [8-10], and is also widely used by tools that compare image distributions like the Fréchet Inception Distance (FID) [25]. Both use a convolutional encoder and consider the pixels of the latent embedding (the vectors that span across the channel dimension) as a MVG distribution.
3.1 Optimal transport in Wasserstein Space
In this section, we briefly introduce key concepts of optimal transport theory that we draw from [26].
Let  $\mu$  and  $\nu$  be probability measures on  $\mathbb{R}^m$  and  $\mathbb{R}^n$ , respectively. The set of all transport plans, which are probability measures  $\pi$  on  $\mathbb{R}^m\times \mathbb{R}^n$  with marginals  $\mu$  and  $\nu$ , is denoted by  $\Pi (\mu ,\nu)$ . The Wasserstein-2 distance between  $\mu$  and  $\nu$  is defined as follows:
$$
\mathcal {W} _ {2} ^ {2} (\mu , \nu) = \inf  _ {\pi \in \Pi (\mu , \nu)} \mathbb {E} _ {x, y \sim \pi} \left[ \| x - y \| _ {2} ^ {2} \right]. \tag {1}
$$
A transport plan that achieves this infimum is called an optimal transport plan between  $\mu$  and  $\nu$ . If  $\mu$  has a density (i.e. is absolutely continuous w.r.t. the Lesbegue measure), there exists a measurable function  $\mathrm{T}_{\mu \longrightarrow \nu}:\mathbb{R}^m\longrightarrow \mathbb{R}^n$  such that, if  $\mathbf{x}_1\sim \mu$  and  $\mathbf{x}_2\sim \nu$  are two random variables, then  $\mathbf{x}_2\stackrel {a.s}{=}\mathrm{T}_{\mu \longrightarrow \nu}(\mathbf{x}_1)$ . We refer to  $\mathrm{T}_{\mu \longrightarrow \nu}$  as the optimal transport operator between  $\mu$  and  $\nu$ . Like in [5], we also abuse this notation even when  $\pi$  is non-degenerate, in which case  $\mathrm{T}_{\mu \longrightarrow \nu}$  represents a one-to-many (stochastic) mapping.
Additionally, when considering two Multivariate Gaussians (MVGs)  $\mathbf{x}_1$  and  $\mathbf{x}_2$  with  $\mathbf{x}_1 \sim \mathcal{N}(\mu_{\mathbf{x}_1}, \Sigma_{\mathbf{x}_1})$  and  $\mathbf{x}_2 \sim \mathcal{N}(\mu_{\mathbf{x}_2}, \Sigma_{\mathbf{x}_2})$ , respectively, and assuming that  $\Sigma_{\mathbf{x}_1}$  and  $\Sigma_{\mathbf{x}_2}$  are nonsingular, there exists a closed-form solution for the optimal transport operator, which is deterministic and linear:
$$
\mathrm {T} _ {p _ {\mathbf {x} _ {1}} \longrightarrow p _ {\mathbf {x} _ {2}}} ^ {\mathrm {M V G}} (x _ {1}) = \Sigma_ {\mathbf {x} _ {1}} ^ {- \frac {1}{2}} \left(\Sigma_ {\mathbf {x} _ {1}} ^ {\frac {1}{2}} \Sigma_ {\mathbf {x} _ {2}} \Sigma_ {\mathbf {x} _ {1}} ^ {\frac {1}{2}}\right) ^ {\frac {1}{2}} \Sigma_ {\mathbf {x} _ {1}} ^ {- \frac {1}{2}} \cdot (x _ {1} - \mu_ {\mathbf {x} _ {1}}) + \mu_ {\mathbf {x} _ {2}}, \tag {2}
$$
where a symmetric and positive definite square root of the matrices is chosen.
3.2 Wasserstein-2 MSE tradeoff
We build upon the problem setting introduced in [1, 2] to establish our analysis. We consider the following scenario:  $\mathbf{x} \in \mathbb{R}^n$  represents a source natural image,  $\mathbf{y} \in \mathbb{R}^m$  represents its degraded version, and we assume that the posterior  $p_{\mathbf{x}|\mathbf{y}}(\cdot | y)$  is non-degenerate for almost any  $y$ . Our objective is to construct an estimator  $\hat{\mathbf{x}}$  that predicts  $\mathbf{x}$  given  $\mathbf{y}$ . A valid estimator  $\hat{\mathbf{x}}$  should be independent of  $\mathbf{x}$  given  $\mathbf{y}$ . Finally,  $p_{\mathbf{x}}, p_{\mathbf{x}^*}$  and  $p_{\hat{\mathbf{x}}_0}$  denote the probability distributions associated with the random variables  $\mathbf{x}, \mathbf{x}^*$  and  $\hat{\mathbf{x}}_0$ , respectively.
Let  $\mathbf{x}^{*} = \mathbb{E}[\mathbf{x}|\mathbf{y}]$  denote the MMSE estimator that achieves the minimal MSE, i.e.,  $\mathrm{MSE}(\mathbf{x},\mathbf{x}^{*}) = D_{\min}$ . Additionally, let  $\hat{\mathbf{x}}_0$  denote the  $\mathbf{D}_{\max}$  estimator, which among all estimators satisfying  $\mathcal{W}_2(p_{\mathbf{x}},p_{\hat{\mathbf{x}}_0}) = 0$ , attains the minimal MSE, namely,  $\mathrm{MSE}(\mathbf{x},\hat{\mathbf{x}}_0) = D_{\max}$  (refer to fig. 1). Notably, as discussed in [1], these estimators have a compelling property: their joint distribution  $p_{\hat{\mathbf{x}}_0,\mathbf{x}^*}$  is an optimal transport plan between  $\mathbf{x}^{*}$  and  $\mathbf{x}$ , characterized by the following optimization problem:
$$
p _ {\hat {\mathbf {x}} _ {0}, \mathbf {x} ^ {*}} \in \underset {p _ {\mathbf {x} _ {1}, \mathbf {x} _ {2}} \in \boldsymbol {\Pi} \left(p _ {\mathbf {x}}, p _ {\mathbf {x} ^ {*}}\right)} {\arg \min } \mathbb {E} \left[ \| \mathbf {x} _ {1} - \mathbf {x} _ {2} \| _ {2} ^ {2} \right]. \tag {3}
$$
In other words, finding  $\hat{\mathbf{x}}_0$  is equivalent to finding an optimal transport plan from  $\mathbf{x}^*$  to  $\mathbf{x}$ . Then, the  $\mathbf{D}_{\mathrm{max}}$  estimator is simply  $\hat{\mathbf{x}}_0 = \mathrm{T}_{p_{\mathbf{x}^*}} \longrightarrow p_{\mathbf{x}}(\mathbf{x}^*)$ .
This estimator is particularly useful as it allows to obtain any point on the perception-distortion function through a naive linear interpolation with the MMSE estimator. Specifically, we can define the interpolated estimator  $\hat{\mathbf{x}}_P$  as follows:
$$
\hat {\mathbf {x}} _ {P} = (1 - \alpha) \hat {\mathbf {x}} _ {0} + \alpha \mathbf {x} ^ {*}, \tag {4}
$$
where  $0 \leq \alpha \leq 1$  is an interpolation constant [1] that depends on the perceptual index of  $\mathbf{x}^*$  and the desired perceptual index  $0 \leq P = \mathcal{W}_2(\mathbf{x},\hat{\mathbf{x}}_P)$  (refer to fig. 1).
Figure 3: Trading perception and distortion using out-of-the-box predictors, wrapped with our method. Using eq. (4) with  $\alpha \in [0,1]$  we interpolate a given predictor (orange) and our improved  $\mathbf{D}_{\mathrm{max}}$  estimation (green), to approximate the PD FID-MSE function (blue curve). With  $\alpha \in [-1,0] \cup [1,2]$  we extrapolate outside of the PD curve (light gray), beyond the theory-inspired area, to further improve performance.
We start by describing the general flow of our proposed algorithm, and then move to elaborate on each of its components.
Theoretically speaking, our algorithm, combined with any given MMSE estimator, is an approximation of the  $\mathbf{D}_{\mathrm{max}}$  estimator. In practice, however, it can be combined with any type of estimator and potentially improve its perceptual quality. I.e., it can be combined with an estimator that optimizes
Figure 4: With a pre-trained VAE, we estimate the first and second order statistics of the latent patches of natural images and the restorations of some given estimator. At inference time, we use the closed-form OT eq. (2) operator between MVG distributions to transport the latent representation of a given restored sample, which, after decoding, increases the visual quality of the restored sample. For a fully detailed explanation of the algorithm, see section 4.
distortion (e.g., SwinIR [11]) and improve its perceptual quality at the expense of distortion, or it can be combined with an estimator that optimizes perceptual quality (e.g., DDRM [3]) and improve its perceptual quality even further. As a result, our algorithm is agnostic to the type of degradation. To clarify, our algorithm is not really a restoration algorithm by itself, but rather a wrapper which can potentially improve the perceptual quality of any given estimator.
4.1 The algorithm
The main goal of our algorithm is to approximate the optimal transport plan between  $p_{\hat{\mathbf{x}}}$  and  $p_{\mathbf{x}}$ , namely,  $\mathrm{T}_{p_{\hat{\mathbf{x}}}} \longrightarrow p_{\mathbf{x}}$ , where  $\hat{\mathbf{x}}$  is a given estimator. Theoretically, with such an operator, one could optimally transport  $\hat{\mathbf{x}}$  such that, with minimal loss in MSE performance, the transported estimator would attain perfect perceptual quality. Computing  $\mathrm{T}_{p_{\hat{\mathbf{x}}}} \longrightarrow p_{\mathbf{x}}$  for high dimensional distributions is a difficult task, involving complex (often adversarial) optimization (see discussion in section 2). To solve this, we perform several assumptions and approximations that allow us to efficiently compute a closed form transport operator that approximates  $\mathrm{T}_{p_{\hat{\mathbf{x}}}} \longrightarrow p_{\mathbf{x}}$ . The flow of the algorithm is presented in fig. 4, and goes as follows:
Encoding: In the training stage we encode  $N$  natural images  $\{x^{(i)}\}_{i = 1}^N$  and  $N$  restored samples  $\{\hat{x}^{(i)}\}_{i = 1}^N$  (unpaired) into their latent representations,  $\{x_{e}^{(i)}\}_{i = 1}^{N}$  and  $\{\hat{x}_e^{(i)}\}_{i = 1}^N$ , respectively. The size of each image sample is denoted by  $(3,H,W)$ , where  $H$  and  $W$  are the height and width, respectively, and the size of their latent representation is denoted by  $(c,H_e,W_e)$ . In the inference stage we perform the same process but only on a single estimate  $\hat{x}$ , resulting again in a latent representation of size  $(c,H_e,W_e)$
Unfold: From each latent representation we extract all the overlapping patches of size  $(c, p, p)$ , where  $p$  is the height and width of each patch.
Flatten each patch and aggregate: We flatten all the extracted patches to obtain 1-dimensional vectors of size  $cp^2$ , which we assume to come from a MVG distribution. In the training stage we compute their empirical mean and covariance matrix (aggregating over the  $N$  dimension), and compute the optimal transport in closed-form  $\mathrm{T}_{p_{\mathbf{x}_e}}^{\mathrm{MVG}} \longrightarrow p_{\mathbf{x}_e}$  using eq. (2).
Matmul and unflatten: We apply the pre-computed transport operator using a simple matrix-vector multiplication on the flattened version of the patches extracted from the latent representation  $\hat{x}_e$ . We then reshape each vector to the original patch size  $(c,p,p)$  (unflatten).
**Fold:** We rearrange the transported patches back to the original size of the latent representation (reversing the unfold operation). Since the patches overlap, we simply average the shared pixels.
Decoding: To produce our final enhanced estimation  $\hat{x}_0$ , we decode the transported latent image back to the pixel space using the decoder of the VAE.
Together, the inference steps form an end-to-end approximation of the desired transport operator  $\mathrm{T}_{p_{\mathbf{x}}}\longrightarrow p_{\mathbf{x}}$ . In appendix B we elaborate on the choices and practical considerations of our algorithm.
In all of our experiments we use the encoder and decoder of the VAE [27] from stable-diffusion [7].
The pre-trained models we evaluate: We apply our latent transport method (described in section 4) on SwinIR [11], Restormer [12] and Swin2SR [13], all of which attempted to minimize average pixel distortion using a supervised regression loss on paired image samples. Additionally, we apply our algorithm on models that are trained to achieve high perceptual quality, and show that we can improve their visual quality even further. As such, we tested two benchmark models in high perceptual quality image restoration: ESRGAN [16], a GAN-based method, and DDRM [3], a diffusion-based method. Beyond our original goal to improve perceptual quality, we demonstrate that we can also traverse the  $\mathcal{W}_2$  -MSE tradeoff using any restoration model (e.g., SwinIR, ESRGAN). To do so, we pick any of the aforementioned algorithms and apply our method to improve its perceptual quality, leading to a new estimator. We then interpolate the original algorithm and its improved version using eq. (4), adjusting  $\alpha$  to traverse the tradeoff. To clarify, we plug the original algorithm as  $\mathbf{x}^*$  in eq. (4) (instead of the theoretical MMSE estimator), and plug our improved version as  $\hat{\mathbf{x}}_0$  (instead of the theoretical  $\mathbf{D}_{\mathrm{max}}$  estimator).
Restoration tasks: We showcase our algorithm on Single-Image Super-Resolution (SISR), denoising of Additive White Gaussian Noise (AWGN), JPEG [28] decompression, Noisy Super-Resolution (NSR) and Compressed Super-Resolution (CSR). Training and inference of our algorithm are performed on each restoration model separately, and the evaluation is performed on the restoration task that corresponds to the given model.
Transport operator computation: The transport operator is computed using two disjoint sets of 10 randomly picked images from the ImageNet [29] dataset train split. The first set is used to approximate the predictor's latent statistics  $(\mu_{\hat{\mathbf{x}}_e}, \Sigma_{\hat{\mathbf{x}}_e})$ : we degrade each image according to the restoration task the predictor is intended to solve, compute the 10 restored outputs, embed the results and compute the embeddings' statistics. The second set is embedded into the latent representation without further modification and serves to approximate the natural image latent statistics  $(\mu_{\mathbf{x}_e}, \Sigma_{\mathbf{x}_e})$ .
Metrics: In addition to Peak Signal-to-Noise Ratio (PSNR), we evaluate distortion performance with Structural Similarity Index Measure (SSIM) [30] and Learned Perceptual Image Patch Similarity (LPIPS) [31], both of which suit better for natural image comparison [31].
Nonetheless, LPIPS remains by definition a distortion (full-reference) measure: it is non-negative and zero when the two images are identical [2]. Interestingly, the original perception-distortion paper [2] already classified the VGG loss [32] - the ancestor of LPIPS - to be a distortion, on which the tradeoff exists (but is less severe).
Therefore, to evaluate perceptual quality, we use the Inception Score (IS) [33], the Fréchet Inception Distance (FID) [25] and the Kernel Inception Distance (KID) [34] following popular image restoration papers [3, 7, 35].
Data sets: It is impractical to perform a serious quantitative evaluation of the perception-distortion tradeoff on real-world datasets (e.g., SIDD, DND, RealSR), which have too few samples to compute FID. Hence, for all models except DDRM [3] and Swin2SR [13], we report the performance on the 50,000 validation samples of ImageNet [29] following [7, 35]. Because of its computational complexity, DDRM [3] reported its performance on a subset of a 1000 ImageNet [29] validation samples. For Swin2SR [13], we use the official DIV2K [36] restored samples provided by the authors.
Although our algorithm can be applied to images with arbitrary aspect ratios, all the tested models were trained on square images. Thus, we resize the samples to  $512 \times 512$  pixels following the pre-processing procedure of DDRM [3]. Finally, we conduct the qualitative evaluation on popular samples from DIV2K or Set14.
5.1 Quantitative results
As reported in table 1, our algorithm can trade perceptual quality for distortion (and vice versa) at test time. We sometimes even manage to improve the pre-trained predictor's PSNR, even of regression models like SwinIR and Swin2SR. When using  $\alpha \leq 0$ , we systematically improve the predictor's perceptual performance (FID, KID, IS), even for estimators which were designed to achieve photorealism in the first place, e.g., ESRGAN and DDRM. On Non-Local-Means (NLM) [15], an older, non deep-learning denoising algorithm, our method marginally improves all metrics.