SlowGuess/NeurIPS_2023_OCR huggingface dataset:
While, in theory, our procedure to traverse the perception distortion tradeoff should only include values of  $\alpha$  in the range [0, 1] (see eq. (4)), we also tried to use values outside of this range. As shown in fig. 3, with  $\alpha \in [-1,0] \cup [1,2]$  we can obtain even better PD curves, and sometimes improve the perceptual quality and/or the distortion of the methods even further. For instance, the PD curve of Swin2SR obtained using  $\alpha \in [1,2]$  is strictly better than the one obtained using  $\alpha \in [0,1]$ . This deviation from the theory can be explained by the implementation choices discussed in section 4; We perform the transport in the latent space - not the pixel space. Additionally, we use FID as perceptual index to measure visual quality when the theory presented in section 3.2 only talks about the Wasserstein-2 distance. In practice, the sharpened details that appear in  $\hat{\mathbf{x}}_0$  and not in  $\mathbf{x}^*$  are either amplified when  $\alpha \in [-1,0]$  or subtracted (instead of being added) to  $\mathbf{x}^*$  when  $\alpha \in [1,2]$ . We leave the formal analysis of this interesting phenomenon for future research.
Choosing the right value of  $\alpha$ : Like any other hyper-parameter,  $\alpha$  can improve the performance with some tuning when approaching a new task or a new data set (refer to table 1). We argue that the few-shot nature of our algorithm makes this tuning actually practical:  $\alpha$  does not need to be set before performing some expensive training. Once  $\hat{\mathbf{x}}_{\alpha=0}$  is computed, any  $\hat{\mathbf{x}}_{\alpha}$  can be obtained thanks to eq. (4) without additional cost. In any case, as reported in table 1,  $\alpha = 0$  consistently improves perceptual quality for all the tasks and models considered (as expected from the theory). We consider it to be a satisfying default choice, so manually adjusting  $\alpha$  is not a great concern.
5.2 Qualitative results
Qualitative results on arbitrary image sizes and aspect ratios are shown in fig. 5. Using our method, we observe a consistent improvement of photo-realism when transporting existing restoration algorithms using our method. Hence, the qualitative results align with the quantitative perceptual performance gains.
5.3 Training details & ablation study
All the results presented in figs. 3 and 5 and table 1 were obtained using the same hyper-parameters. We used the "f8-ft-MSE" fine-tuned version of stable-diffusion's VAE from Hugging-Face's diffusers library [37]. For the training stage we use 20 randomly-drawn images from the ImageNet training set (10 images which we use as the natural images set, and 10 images which we degrade and then restore with the estimator). We used a patch-size of  $p = 3$  in the latent space.
Thanks to its simplicity, for each restoration task, our few-shot algorithm requires just a single GPU, and a few seconds for both training and inference.
We turn to detail some considerations about practical aspects of our algorithm which we empirically evaluate on the popular  $\mathrm{SISR}_{\times 4}$  task for the ESRGAN estimator.
Patch size: We experiment with increasing patch-sizes when unfolding the latent image (see appendix B.2).  $p = \{3,5\}$  yielded the best PSNR and FID. Smaller patch size ( $p = 1$ ) resulted in worse FID and bigger size  $7 \leq p \leq 15$  yielded slightly worse PSNR.
Training size: As discussed in appendix B.4, each image contributes thousands of samples to the computation of the OT operator. Still, we expect the empirical statistics estimation to benefit from a larger sample size  $S$ . To confirm this, we repeated the visual enhancement experiments while
Table 1: Using eq. (4), our algorithm can trade-off perception and distortion at inference time on any predictor [3, 11-13, 15, 16] and image restoration task. For each task, we report the performance of  $\hat{\mathbf{x}}_0$ . It consistently improves perceptual metrics on all taks and models (aside of NLM). We also report other interesting choices of  $\alpha$  that optimize perception and distortion (for more details about this choice refer to section 5.1).
Task	Signal	Distortion	Perception
PSNR ↑	SSIM ↑	LPIPS ↓	FID ↓	IS ↑	KID×103↓
x	∞	1	0	0	240.53±4.42	0
D(E(x))	27.10	0.81	0.13	0.24	234.71±4.04	0.02±0.07
SISR×4	SwinIR [11]	28.10	0.84	0.24	2.54	201.52±4.85	1.24±0.24
x0.9	28.15	0.84	0.24	2.80	198.69±2.97	1.38±0.24
x-0.2	25.08	0.77	0.25	1.19	216.74±4.26	0.38±0.89
x0	25.48	0.78	0.23	1.39	214.63±5.50	0.69±0.23
JPEGq=10	SwinIR [11]	29.68	0.86	0.30	8.95	161.73±3.36	6.52±0.77
x1.1	29.58	0.86	0.30	8.36	166.50±3.12	6.08±0.75
x-0.2	23.74	0.76	0.31	7.56	166.65±3.58	5.68±0.83
x0	24.84	0.78	0.30	8.14	163.14±3.93	6.15±0.77
AWGNσ=50	Restormer [12]	30.18	0.86	0.26	5.21	178.62±2.83	3.29±0.56
x1.1	30.09	0.86	0.25	4.63	183.36±3.20	2.61±1.53
x1.7	27.26	0.82	0.25	2.73	198.93±5.13	1.76±1.58
x0	25.31	0.78	0.27	4.42	182.86±2.21	2.93±1.62
SR×4JPEGq=10	Swin2SR [13]	19.75	0.55	0.53	205.00	5.95±0.49	40.68±3.34
x0.8	19.81	0.55	0.53	209.82	5.91±0.69	43.28±3.86
x1.9	18.44	0.49	0.51	168.12	6.36±0.69	19.95±2.84
x0	18.45	0.48	0.51	183.80	6.55±0.61	29.07±3.58
SISR×4	ESRGAN [16]	26.77	0.80	0.21	1.06	221.68±3.06	0.43±0.14
x0.7	27.00	0.81	0.21	1.51	215.87±3.64	0.56±0.21
x-0.2	24.84	0.74	0.23	0.80	221.89±2.53	0.30±0.20
x0	25.33	0.74	0.22	0.89	220.96±3.19	0.34±0.18
SR×4AWGNσ=50	DDRM [3]	26.10	0.75	0.34	36.44	43.52±3.33	5.09
x1.2	25.91	0.75	0.33	33.68	44.90±4.06	3.88
x1.7	24.48	0.70	0.35	29.05	47.91±2.69	1.47
x0	23.19	0.69	0.35	29.71	46.36±4.18	1.91
AWGNσ=50	NLM [15]	26.09	0.71	0.44	12.84	148.71±3.75	8.73±0.91
x0.8	26.24	0.72	0.43	12.46	148.78±2.49	8.60±0.95
x0	24.80	0.71	0.42	14.88	140.65±2.10	10.90±1.15
varying the number of training samples. Surprisingly, we observe no change in the performance of the evaluated metrics for  $S = \{10^5, 10^4, 10^3, 10^2\}$ , i.e., approximating the distribution statistics with 100 samples is as good as using 100,000 samples, and this is true regardless of the chosen patch size. Moreover, when With  $S = 10$ , then only for small patch sizes of  $p \leq 5$  we observe no performance drop compared to using a larger sample size. This suggests that our method can be successfully deployed in few-shot settings, where the number of available samples is small.
Paired vs. unpaired samples: Surprisingly, using paired images to compute the distribution parameters yielded better PSNR but worse FID. We suspect that using paired updates induces a bias which results in worse covariance estimation.
Transporting the degraded measurement directly: Applying our algorithm on the degraded input directly led to insufficient results as we see in appendix B.7.
Re-applying the algorithm another time on  $\hat{\mathbf{x}}_0$ : This is actually an interesting idea we tested on super-resolution when conducting our evaluations. As a matter of fact, the performance does not improve (it even degrades a bit) when applying the algorithm another time. The explanation is quite simple: After transporting once the test images using the VAE, their latent distribution aligns with that of the natural images. Hence, transporting another time does nothing (the transport operator is
Figure 5: Our method (third column from the left) notably improves the results of several benchmark predictors (second column from the left) on various degradations.
the identity matrix). We are only left with the reconstruction error introduced by the encoding and decoding of the images, which deteriorates the MSE performance.
Does the selection on the training data have an impact on the performance of restoration?: Our experiments showed that the class of images does not have a significant impact on the performance (e.g. one could use images of cars to improve images of dogs). However the resolution of images does play a significant role in attaining the best performance. I.e., to transport 512x512 images, it is best to use training images of the same resolution. This drawback is somewhat mitigated by the few-shot nature of the algorithm.
Limitations: The pre-trained VAE used for the purpose of our experiments exhibits a rate of  $R = 48$  on  $512^2$  images which inevitably translates into sub-optimal distortion performance [38]. Thus, the distortion performance of our estimates are bounded by that of the pre-trained VAE. I.e., even encoding and decoding a completely clean and natural image does not yield result in perfect
$\mathbf{x}^*$  (MMSE)
$\hat{\mathbf{x}}_0$  (ours)
$\mathbf{D}(\mathbf{E}(\mathbf{x}))$
x (original)
Figure 6: Our method's reconstruction capabilities are bounded by that of the VAE. Our algorithm is not able to preserve complex visual structures such as face identity (top row) or text (middle row).
reconstruction. Most notably, the VAE sometimes fails to reconstruct human faces, as well as text images, and such a weaknesses affects our algorithm as well (see fig. 6).
Finally, it has been recently shown that the posterior sampler is the only estimator attaining perfect perceptual quality while producing outputs that are perfectly consistent with the degraded input [39]. As such,  $\mathbf{D}_{\mathrm{max}}$  cannot hope for consistent restorations.
Potential impact: Instead of using sophisticated and data-hungry generative models, we show it is possible to obtain photo-realistic results using simple tools like MMSE estimators and VAEs. We hope our few-shot algorithm will inspire other simple and practical image restoration methods.
Potential misuse: Our algorithm aims at improving the perceptual quality of existing algorithms. However, when using a biased training set, this could potentially cause bias in the enhanced restoration as well. This could potentially harm the results of medical image diagnosis, for example.
This research was partially supported by the Council For Higher Education - Planning and Budgeting Committee.
Deep Optimal Transport: A Practical Algorithm for Photo-realistic Image Restoration - Supplementary Material
A Background and extensions
A.1 Numerical Example
Figure 7: 2D Gaussian mixture denoising. Source samples are shown in blue. The MMSE estimator  $(\mathbf{x}^*$ , orange) attains the best  $MSE$  but the worst perceptual index  $\mathcal{W}_2$ . The posterior samples  $(\mathbf{x}|\mathbf{y}$ , purple) attain the best perceptual index but half of the optimal  $MSE$  performance. The  $\mathbf{D}_{\mathrm{max}}$  estimator  $(\hat{\mathbf{x}}_0, \mathrm{green})$  maintains the  $MSE$  of  $\mathbf{x}^*$  while attaining a perceptual quality close to  $\mathbf{x}|\mathbf{y}$ . The DP curve is obtained by interpolating  $\hat{\mathbf{x}}_0$  and  $\mathbf{x}^*$  using eq. (4).
To guide the reader in understanding the MMSE transport paradigm, we showcase our method on a 2-dimensional denoising problem. To avoid a too trivial uni-modal example, we draw the clean signal from a 4-components Gaussian mixture with non-trivial covariances. We derive linear MMSE and posterior estimators from [40] and proceed by applying the closed-form transport operator introduced in eq. (3).
Note that to avoid deviating from our actual method, we refrain from using more advanced transport operators better suited for multi-modal data. Indeed, those are not a practical solution for real-world image datasets, as they require much more samples than actually available.
We summarize the experiment results in fig. 7. We observe that we obtain the best perceptual quality by sampling from the posterior distribution. However, we witness a significant decrease in  $MSE$  performance as predicted by [2]. In contrast, the  $\mathbf{D}_{\mathrm{max}}$  estimator enjoys a good perceptual index while maintaining a close-to-optimal distortion performance.
A.2 Stochastic transport operator
Throughout our experiments, we found out that increasing the patch-size  $p$  can result in numerical instabilities. Recall that the linear transport operator presented in eq. (3) uses the inverse square root of the source covariance matrix  $\Sigma_{\mathbf{x}_1}$ . When  $p$  is large, (typically  $p \geq 7$ ), we obtain ill-conditioned covariance matrices. When the smallest singular value is still positive, we add a small stability constant to the matrix diagonal to ensure it is strictly positive definite. However, the numerical errors sometimes adds up to negative eigenvalues. In this case, we clamp the negative eigenvalues to zero and use the stochastic (one-to-many) transport operator proposed by [1],
$$
\mathrm {T} _ {p _ {\mathbf {x} _ {1}} \longrightarrow p _ {\mathbf {x} _ {2}}} ^ {\text {s t o c h a s t i c}} (x _ {1}) = \Sigma_ {\mathbf {x} _ {2}} ^ {\frac {1}{2}} \left(\Sigma_ {\mathbf {x} _ {2}} ^ {\frac {1}{2}} \Sigma_ {\mathbf {x} _ {1}} \Sigma_ {\mathbf {x} _ {2}} ^ {\frac {1}{2}}\right) ^ {\frac {1}{2}} \Sigma_ {\mathbf {x} _ {2}} ^ {- \frac {1}{2}} \Sigma_ {\mathbf {x} _ {1}} ^ {\dagger} (x _ {1} - \mu_ {\mathbf {x} _ {1}}) + \mu_ {\mathbf {x} _ {2}} + w, \tag {5}
$$
when  $\Sigma_{\mathbf{x}_1}^{\dagger}$  denotes the pseudo-inverse of  $\Sigma_{\mathbf{x}_1}$  (after negative eigenvalues where clamped) and  $w \sim \mathcal{N}(0, \Sigma_{\mathbf{x}_2}^{\frac{1}{2}}(I - \Sigma_{\mathbf{x}_2}^{\frac{1}{2}}\mathrm{T}^*\Sigma_{\mathbf{x}_1}^{\dagger}\mathrm{T}^*\Sigma_{\mathbf{x}_2}^{\frac{1}{2}})^{\frac{1}{2}}\Sigma_{\mathbf{x}_2}^{\frac{1}{2}})$ , with  $\mathrm{T}^* = \Sigma_{\mathbf{x}_2}^{-\frac{1}{2}}\left(\Sigma_{\mathbf{x}_2}^{\frac{1}{2}}\Sigma_{\mathbf{x}_1}\Sigma_{\mathbf{x}_2}^{\frac{1}{2}}\right)^{\frac{1}{2}}\Sigma_{\mathbf{x}_2}^{-\frac{1}{2}}$ .
2We tried to avoid overflow when summing over the images by using 64 bit precision
B Practical choices and considerations in our algorithm
B.1 Working in latent space
We adopt the latent transport approach where the images are embedded into the latent space of a pre-trained auto-encoder. Let  $\mathbf{E}(\cdot)$ ,  $\mathbf{D}(\cdot)$  denote the encoder and decoder, respectively. Even if  $\mathbf{D}(\mathbf{E}(t)) = t$ , it is likely that  $\mathbf{E}(\cdot)$  "deforms" the space, i.e.,  $\| \mathbf{E}(s) - \mathbf{E}(t) \| \neq \| s - t \|$ , which means that the optimal transport plan in the latent space could be different than the plan we seek in the pixel space (the cost function in eq. (3) has changed). We can address this by modifying the latent cost function to account for the deformation via the following change of variables
$$
\mathbb {E} \left[ \| \hat {\mathbf {x}} - \mathbf {x} \| ^ {2} \right] = \mathbb {E} \left[ \frac {\| \mathbf {E} (\hat {\mathbf {x}}) - \mathbf {E} (\mathbf {x}) \| ^ {2}}{| \frac {\partial \mathbf {E} (\mathbf {x})}{\partial \mathbf {x}} | \cdot | \frac {\partial \mathbf {E} (\hat {\mathbf {x}})}{\partial \hat {\mathbf {x}}} |} \right], \tag {6}
$$
where  $\left|\frac{\partial\mathbf{E}(\mathbf{x})}{\partial\mathbf{x}}\right|$  is the determinant of the Jacobian matrix of  $\mathbf{E}(\cdot)$ . However it is not a practical solution since we lose access to the closed-form solution eq. (2). Note that the latent MSE approximation is usually desirable when dealing with natural images (e.g. to elaborate image quality measure [41], perceptual quality metrics [25]). It is also true in our case but it means we can no longer claim we obtain the  $\mathbf{D}_{\mathrm{max}}$  estimator.
With that, we argue that switching to a latent cost is actually a strength rather than a weakness of our method. Indeed, using the MSE between deep latent variables has shown to be a better fit to compare natural images than directly working in the pixel space [31]. The authors of [7] trained their VAE (which is used in our experiments) to remove "imperceptible details" from the latent representation, in order to better focus on higher level image semantics. In section 5.1 we validate this claim by showing that our algorithm maintains the "perceptual" discrepancy performance of the original estimator (e.g., LPIPS).
B.2 Overlapping patches extraction strategy
For Convolutional Neural Network (CNN) encoders  $^3$ , let  $(c, H_e, W_e)$  denote the shape of the latent representation (CNN encoders produce 3-dimensional encoded tensors), where  $H_e, W_e$  the spatial extent and  $c$  is the number of channels (i.e., the number of convolution kernels in the last convolution layer). The covariance matrices  $\Sigma_{\hat{\mathbf{x}}_e}$ ,  $\Sigma_{\mathbf{x}_e}$  contain  $\frac{(c, H_e, W_e)^2}{2}$  parameters, which may require a large amount of samples for large latent images with  $H_e, W_e \gg 1$ . To mitigate the quadratic dependency on  $H_e \cdot W_e$ , we assume that the latent pixels depend only on the pixels in their close neighborhood. In practice, we unfold the latent representation, extracting all overlapping patches of shape  $(c, p, p)$ . A similar approximation exists in the style-transfer literature [8, 9], where instead of patches, only the pixels are considered (i.e., this is a private case of our approach with  $p = 1$ ). In section 5.3 we empirically show that increasing  $p$  improves the perceptual quality at the expense of MSE performance, given that enough training samples are available.
B.3 Shared distribution
When dealing with natural image scenes, it is beneficial to suppose that overlapping patches share common statistical attributes [41, 42]. In the case of a CNN encoded image, this approximation remains satisfying because we ultimately look at filter activations which are spatial-invariant with each latent patch having the same receptive field. Therefore, we assume that the overlapping patches are all samples from the same distribution. This approach dramatically reduces the number of estimated parameters, and also multiplies the number of samples at our disposal by  $H_{e} \cdot W_{e}$ , which alleviates the curse of dimensionality. We demonstrate these practical benefits in section 5.3. In practice, given  $N$  images, we "flatten" all the extracted patches to vectors  $\underline{v}_{cp^2 \times 1}$  which we stack into a sample matrix  $\underline{X}$ . We then aggregate the samples to compute the MVG statistics:  $\mu = X^T \mathbf{1}$ ,  $\Sigma = \frac{NH_eW_e}{NH_eW_e - 1} (X - \mu)(X - \mu)^T$ . As  $NH_{e}W_{e}$  may be very large, we perform all computations in double precision. When training, this process is done twice; once for the natural image samples, and once for the restored samples we wish to transport.
This methodology can easily be extrapolated to other encoder architectures.
B.4 Size of the latent representation
When increasing the capacity of models with a fixed encoding rate, deepening is preferable than widening. Indeed, increasing  $c$  makes the covariance estimation dramatically harder while increasing  $H_{e}$ ,  $W_{e}$  enlarges the sample pool. Therefore, the VAE from [7] with  $c = 4$  and  $H_{e}$ ,  $W_{e} \gg 1$  is a particularly good candidate for our method. For  $p = 3$  for instance, the covariance matrix admits only 1296 parameters while each  $512^{2}$  image contributes 4096 samples to its estimation. As we see next, this greatly contributes to reducing the number of training samples needed to estimate the covariance matrices and allows to compute the transport operator in a few-shot manner.
B.5 Transport
In a single pass on a data set of natural images and a (possibly different) data set of restored samples, we compute  $\mathrm{T}_{p_{\hat{\mathbf{x}}_e}}^{\mathrm{MVG}} \longrightarrow p_{\mathbf{x}_e}$  (see eq. (2)). Note that each latent distribution could sometimes be degenerate, especially for severe degradations. Fortunately, the classical MVG transport operator can be generalized to ill-posed settings where  $\Sigma_{\hat{\mathbf{x}}}$  is a singular matrix (see appendix A.2).
B.6 Decoding
Since the transported patches overlap, we "fold" them back into a latent image  $\hat{\mathbf{x}}_{0,\text{latent}}$  by averaging. The latent image is then decoded back to the pixel space, i.e.  $\hat{\mathbf{x}}_0 = \mathbf{D}(\hat{\mathbf{x}}_{0,\text{latent}})$ . Since  $\mathbf{E}(\cdot)$  is not invertible, the decoder  $\mathbf{D}(\cdot)$  is used as a convenient approximation in the training domain of the auto-encoder. A corollary of this approximation is that the auto-encoder should in theory be trained on the image distribution we aim to transport, which weakens our claim to a fully blind algorithm.
All the steps described above are summarized in fig. 4.
B.7 Transporting the degraded measurement
We tried applying our algorithm on the degraded measurement directly. Indeed we observe qualitatively and quantitatively that transporting the degraded measurement  $\mathbf{y}$  amplifies the degradation (refer to fig. 8).
Figure 8: Transporting the degraded measurement  $(\mathrm{JPEG}_{q=10})$  directly is not enough to restore the image. It can sometimes even exacerbate the degradation. Quantitatively, the degraded sample  $\mathbf{y}$  has better PSNR and FID than its transported version (respectively 27.26 dB and 13.88 FID vs. 23.69 dB and 15.88 FID).