Deep Optimal Transport: A Practical Algorithm for Photo-realistic Image Restoration
Authors: Theo Joseph Adrai
Guy Ohayon
Michael Elad
Tomer Michaeli
Venue: NeurIPS 2023 poster (2023)
Keywords: Computer Vision
Image Restoration
Deep Learning
Perceptual Quality
TL;DR: We present a few-shot image restoration enhancement algorithm capable of improving  the visual quality of any restoration model on any task at test time.

Abstract: 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 that of 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 with arbitrary dimensions.

Deep Optimal Transport: A Practical Algorithm for Photo-realistic Image Restoration
Theo J. Adrai
Guy Ohayon
Technion- Israel Institute of Technology Computer Science
Technion-Israel Institute of Technology
Computer Science
ac.il
ohayonguy@cs .erissisttcciiiooo ac il
Michael Elad
Tomer Michaeli
Technion-Israel Institute of Technology
Technion-Israel Institute Technology
Computer Science eladêcs www.nien..c.......
Electrical Engineering
tomer. mener.m@ee.. technion.ac.
Abstract
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 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 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 predictor that was originally trained to minimize MSE, we approximate the optimal transport by linear transformation in the latent space 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 the model, illustrate our method on variety of degradations applied to general content images of arbitrary dimensions.
improve their MSE on the expense of perceptual quality. We
Introduction
Many image restoration algorithms aim to recover 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 per- ceptual quality, which refers to the extent which restored images resemble natural images. The work in [2] exposed fundamental trade-off between distortion and perceptual qual- ity, where the latter is measured using perceptual index that quantifies the statistical divergence between the distribution of restored images and the distribution of natural images. The
Wa
(MMSE estim.)
Paly (Peaterior sumplar)
to (Dmxx eetim.)
Dnreschable
Drix SMDMS
MSB
Figure 1: The Wa-MSS trade-off [
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

x,4 + no x" (DDRM)
Ey (ours)
x (original)
Figure Our few-shot algorithm improves the visual quality of any estimator at test time. For example, can improve the photo mhoto-realism of DDRM even further.
trade-off curve reveals the predictor that achieves the lowest possible distortion, denoted as D Dmax. while maintaining perfect perceptual quality (refer fig. I.
Following the methodology introduced by (2] has become common practice to compare restoration methods on the perception-distortion (PD) plane, with many methods aiming to reach the elusive this paper, we present a practical approach to approximate the max predictor where distortion is measured using the MSE and perceptual quality is measured by the Wasserstein-2 distance (W2) ) between the distributions of restored and real images. Our approach is based on the recent work [1] which demonstrated that the D max predictor be obtained using optimal transport (OT) from the output distribution of the MMSE predictor to the distribution natural images. By applying optimal transport plan to MMSE restoration resulting from degraded image, can produce D Dmax estimations by transporting the MMSE restored estimate.
Dmax point.
Although progress has been made in finding OT plans between image distributions 4. it re- mains challenging task, particularly for high-dimensional distributions. Therefore, we propose an approximation method for the Diaax estimator by performing transportation in the latent space of a pre-trained auto-encoder. A similar strategy was successfully employed in context of image generation in 77) showing effectiveness in reducing complexity and preserving details.
we assume that the latent representations follow Multivariate Gaussian (MVG) distribution. Thus, by considering the first and second-orde 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 few as 10 unpaired MMSE restored and clean samples. This approach leads to few-shot algorithm that significantly enhances visual quality.
Inspired by the style transfer literature [8-1]]
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. [ Furthermore, by adjusting a single interpolation parameter, we can trade off perception for distortion, resulting in marginal improvements in the distortion performance We demonstrate the improved photo-realism of our approach on a variety of tasks and models, including GAN and diffusion-basee based methods, using high-resolution (e.g., 5122 px) general-content images with arbitrary aspect ratios.
some regression models that were trained to prioritize source fidelity,
2 Related Work
Throughout the paper, we distinguish between two kinds of restoration algorithms: distortion and perception focused. The former category includes traditional methods that minimize distortion (c.g..
OOur code is publicly available at att.pp///ggtttb...c//taaaaaa///s

---

МSE) 15.
The latter category includes more recent works that usually involve generative models like Generative Adversarial Networks (GANs) [6. or diffusion- basssssoooooddddddd techniques B 19
This paper searches for the theoretical D max estimator which minimizes the MSE under a perfect perceptual quality constraint. However, our method is not 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 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 image latent representation method (e.g., auto-encoder), our algorithm applies 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 W2 distance between discrete, low or medium dimensional densities [2] 22]. is far more challenging 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 optimize the W2 distance but with an additional ordering constraint on the generator, which proved to be a challenging task when dealing with real-world data sets [46).. An attempt to sidestep this difficulty would be to use the Gelbrich distance [ [24], which lower bounds the W2 distance and depends only on the first two moments of the distributions. Nevertheless, it is only good estimate when the support of the distributions has elliptical level sets. To address this, one can find a lowddimensional embedding where (i) high dimensionality is no longer issue, (ii) the distributions 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 [8H10, and is also widely used by tools that compare image distributions like the Fréchet Inception Distance (FID) [25] Both use
not degenerated, and (iii) the Gelbrich
convolutional encoder and consider the pixels of the latent embedding (the vectors that span across the channel dimension) as MVG distribution.
3 Background
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 μ and be probability measures Ipn and TR"", respectively. The of all transport plans, which are probability measures æ on Ih' x IR" with marginals ys and v, denoted by I( v). The Wasserstein-2 distance between y and v is defined follows:
inf Enyyo xell(s.))
|| v1(2
(I)
A transport plan that achieves this infimum is called an optimal transport plan between μ and v. If has density (i.e. is absolutely continuous w.r.t. the Lesbegue measure), there exists a measurable function T Tę-iv R - R" such that, if x1 ~ and X, ~ are two random variables, then 0.8 Ty (x,)- We refer Tw as the optimal transport operator between l and Like in 5). we also abuse this notation even when π is non-degenerate, in which case Tn-+ represents a one-to-many (stochastic) mapping.
X2
Additionally, when considering two Multivariate Gaussians (MVGs) X1 and X2 with X1 ~ ((xx Σx,) and xy N NV Σx,), respectively, and assuming that Ex, and Ex, are non- singular, there exists a closed-form solution for the optimal transport operator, which is deterministic and linear:
MVG P (I1)
Exg E2)
Ex (#1 - ) + Mx2g
where symmetric and positive definite square root of the matrices is chosen.
3

---

3.2 Wasserstein-2 MSE tradeoff
We build upon the problem setting introduced in 2 establish our analysis. We consider the following scenario: E IR" represents a source natural image, E R represents its degraded version, and we assume that the posterior Pxly (-ly) non-degenerate for almost any y. Our objective is to construct an estimator & that predicts given A valid estimator & should be independent of x given y. Finally, Px, Px- and Ps., denote the probability distributions associated with the random variables x, x" and 8o, respectively.
Let x E[x y denote the MMSE estimator that achieves the minimal MSE, i.e. MSE(x, x*) Dmin- Additionally, let & denote the Dax estimator, which among all estimators satisfying W2(Px Ps.) = O. attains the minimal MSE, namely, MSE(x, &n) = Dmax (refer fig. []. Notably, as discussed in [T]. these estimators have compelling property: their joint distribution P is an optimal transport plan between x" and x, characterized by the following optimization problem:
arg min
||x]
x2
Pai (P
In other words, finding &o is equivalent to finding an optimal transport plan from x. Then, the D max estimator is simply &o -xx- (x*).
This estimator is particularly useful it allows obtain any point on the perception-distorie function through naive linear interpolation with the MMSE estimator. Specifically, we can define the interpolated estimator & follows:
& (1- + ax",
(4)
where 0 S a < 1 is an interpolation constant [ that depends on the perceptual index of x' and the desired perceptual index P (2 xp) (refer to fig.
SISR
SISR
AWGD mm
1.8
6.5
(SwinIR)
Restormer)
-1.0
ESRGAN)
20..5
-1.0
$2.0
1.2
0.8
2.7
x 1PPG//GGG1
SR AWGD, 25
1.5
59
2.0
4.4
1.0
MSR(X10 3)
8.1
209.0
(Swin28R)
36.8
(DDRM)
a [-1,0]
FTD
E [1,2]
E [o, 1]
-1.0
30
(predletor)
80 (Duax))
166.5
10.5
MSE(x10
22.0
MSB(X10 10))
Figure Trading perception and distortion using out-of-the-box predictors, wrapped with our method. Using eq. with a E [0, 1] we interpolate given predictor (orange) and our improved Dmax estimation (green), to approximate the PD FID-MSE function (blue curve). With [-1,0] U [1, 2] we extrapolate outside of the PD curve (light gray), beyond the theory-inspired area, to further improve performance.
Method
We start by describing the general flow of our proposed algorithm, and then move to elaborate each its components.
Theoretically speaking, our algorithm, combined with any given MMSE estimator, is approxima- tion of the Dmax estimator. practice, however, can be combined with any type of estimator and potentially improve its perceptual quality. L.c,, it can be combined with an estimator that optimizes

---

Space
Spare
Training
flamenn esch and asgrogsi:
unfsié
Hioe
Pacc
201 (
infole
Fig
JVVG
(7.4 IT
(M.
(WiR) -
Inferonce
A
faties cach pach
>1/77 xl.
A givven estimate
untodd
Our auproved Intirationn
foid
uniluation
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 given restored sample, which, after decoding, increases the visual quality of the restored sample. For a fully detailed explanation of the algorithm, see section
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 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 to approximate the optimal transport plan between ps and Px namely, Tpx- Tp-- where is given estimator. Theoretically, with such an operator, one could optimally transport such that, with minimal loss in MSE performance, the transported estimator would attain perfect perceptual quality. Computing T pp-- p2-pp- for high dimensional distributions is difficult task, involving complex (often adversarial) optimization (see discussion in section (). To solve this, we perform several assumptions and approximations that allow us to efficiently compute closed form transport operator that approximates Tp - The flow the algorithm is presented in fig. and goes follows:
Encoding: In the training stage we encode natural images (2(NNN and N restored samples ($0)) (unpaired) into their latent representations, {-r( (6) N and ( Ji=1 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. He. We). the inference stage we perform the same process but only on single estimate æ, resulting again in latent representation of size (c, He, We)
Unfold: From each latent representation we extract all the overlapping patches size (c, p), where p is the height and width of each patch.
Flatten each patch and aggregate: We flatten the extracted patches to obtain I-dimensional vectors of size cp², which we assume to come from 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 MVG
using eq. 2

---

Matmul and unflatten: We apply the pre-computed transport operator using simple matrix- multiplication on the flattened version of the patches extracted from the latent representation á.. We then reshape each vector to the original patch size (c, (c,P,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 Lo. decode the transported latent image back to the pixel space using the decoder the VAE.
Together, the inference steps form an end-to-end approximation of the desired transport operator Tp -xx In appendix B we elaborate on the choices and practical considerations of our algorithm.
5 Experiments
In all of our experiments we use the encoder and decoder of the VAE 27 from stable-diffusion
The pre-trained models we evaluate: We apply our latent transport method (described in section 4 on SwinIR [T]] Restormer [1] and Swin2SR [13], all of which attempted to minimize average pixel distortion using supervised regression loss 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]. diffusion- based method. Beyond our original goal to improve perceptual quality, demonstrate that we can also traverse the W22 -MSE tradcoff using any restoration model (e.g., SwinIR, ESRGAN). To do so, we pick any of the aforementioned algorithms and apply method improve its perceptual quality, leading to a new estimator. We then interpolate the original algorithm and its improved version using eq. (4). adjusting a to traverse the tradeoff. To clarify, we plug the original algorithm as x" in eq. (4) (instead of the theoretical MMSE estimator), and plug our improved version as &o (instead of the theoretical D max estimator).
Restoration tasks: We showcase our algorithm on Single-Imase Image Super-Resolutin (SISR), denoising of Additive White Gaussian Noise (AWGN), JPEG (28) decompression, Noisy Super-Resolution (NSR) and Compressed Super- Resolutinn (CSR). Training and inference of our algorithm are per- formed 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 used to approximate the predictor's latent statistics (Jxx.. Ex,) ): 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 statisties (ux,, Σx.) )
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) B1, both of which suit better for natural image comparison [1].
Nonetheless, LPIPS remains by definition distortion (full-reference) measure: it is non- non-negative and zero when the two images are identical (2]. Interestingly, the original perception-distorte paper [2] already classified the VGG loss (32) the ancestor of I.PIP- to be 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: is impractical to perform a serious quantitative evaluation of the tradeoff on real-world datasets (e.g., SIDD, DND, RealSR), which have too few samples to compute FID. Hence, for all models except DDRM and Swin2SR [3] we report the performance on
distortion
the 50,000 validation samples of ImageNet 2 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 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 512 512 pixels following the me-processing procedure of DDRM R Finally, we conduct the qualitative evaluation on popular samples from DIV2K or Set14.
5.1 Quantitative results
As reported in [ 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 's PSNR, even of regression models like SwinIR and Swin2SR. When using S O, systematically improve the predictor's perceptual performance (FID, KID, IS), even for estimators which were designed to achieve photo- realism in the first place, e.g., ESRGAN and DDRM. On Non-Local-Meam (NLM) [1]. an older, non deep-learning denoising algorithm, our method marginally improves all metrics.
While, in theory, our procedure to traverse the perception distortion tradeoff should only include values a the range [0. 1] (see eq. 4.. we also tried to use values outside of this range. As shown in fig. with [-1. 0] U [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
[1. 2] is strictly better than the one obtained using [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 332]only talks about the Wasserstein-2 distance. In practice, the sharpened details that appear in &o and in x" are either amplified when a E [-1, 0] or subtracted (instead of being added) to x* when a E [1, 2]. We leave the formal analysis of this interesting phenomenon for future research.
Choosing the right value of a: Like any other syper-parametes a can improve the performance with some tuning when approaching new task new data set (refer to table [T. We argue that the few-shot nature our algorithm makes this tuning actually practical: a does not need to set before performing some expensive training. Once &a=0 is computed, any &. to eq. (4) without additional cost. In any case, as reported in table [ a = 0 consistently improves perceptual quality for all the tasks and models considered (as expected from the theory). We consider to be a satisfying default choice, so manually adjusting a is not a great concern.
be obtained thanks
5.2 Qualitative results
Qualitative results on arbitrary image sizes and aspect ratios are shown in fig. Using our method, we observe 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 figs. 3 and 5 and table were obtained using the same hyper-parametent We used the "f8-ft-MSE" fine-tune tne-tuned version of stable-diffusion's VAE from Hugging- Face's library B7]. For the training stage we use 20 randomly-dran (10 images which with the estimator). We used a patch-size of
Face's ' diffusers images from the ImageNet training set the natural images set, and 10 images which we degrade and then restore
in the latent space.
Thanks to its simplicity, for each restoration task, our few-shot algorithm requires just single GPU, and 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 SISRx4 task for the ESRGAN estimator.
Patch size: We experiment with increasing patch-sizes when unfolding the latent image (see ap- pendix B.2 p (3.5) yielded the best PSNR and FID. Smaller patch size (p 1) resulted in worse FID and bigger size S $ 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 statisties estimation to benefit from larger sample size S. To confirm this, we repeated the visual enhancement experiments while
7

---

Table 1: Using eq. (4) our algorithm can trade-off perception and distortion at inference time on any predictor I
and image restoration task. For each task, we report the performance of &g. It consistently improves perceptual metrics on all taks and models (aside of NLM). We also report other interesting choices of a that optimize perception and distortion (for more details about this choice refer to section 5.1
Distortion
Perception
Signal
PSNR
SSIM
LPIPS 4
FID 4
IS f
KIDx10' +
240.53
0
Task
D(E(x))
27.10
0.81
0.13
0.24
234.7114.00
0.02
SwinIR
28.10
0.84
0.24
2.54
201.52+:.85
1.24
SISRx4
No.9
28.15
0.84
0.24
2.80
198.69-22.99
1.38::0.2 0.24
&-0.2
25.08
0.77
0.25
1.19
216.74 +4226
0.38 +0.89
80
25.48
0.78
0.23
1.39
214.63 +5.50
0.69+0.23
SwinIR III
29.68
0,86
0.30
8.95
161.73
6.52
JPFG,=10
&1.1
29.58
0,86
0,30
8.36
166.50=3.00
6.08
&-0.2
23.74
0.76
0.31
7.56
166.65 13.58
5,68 10.83
%
24.84
0.78
0.30
8.14
163.14
3.98
6.15++0.77
Restormer
30.18
0.86
0.26
5.21
178.62 12.83
3.29 3.29.10.56
AWGN 50
&1.1
30.09
0,86
0.25
4.63
183.36 13.20
2.61 16.11
&1.7
27.26
0.82
0.25
2.73
198.93 L5.13
1.76.1
&D
25.31
0.78
0.27
4.42
182.86
=2.21
2.93+:1.62
Swin2SR [13]
19.75
0.55
0.53
205.00
5.95
40.68
SR aRxa
&0.8
19.81
0.55
0.53
209.82
5.91
43.28
11.
18.44
0.49
0.51
168.12
6.36
19.95-12.80
&D
18.45
0.48
0.51
183.80
6.55 +..61
29.07
ESRGAN [16
26.77
0.80
0.21
1.06
221.68
0.43::0.14
SISR x4
$0.7
27.00
0.81
0.21
1.51
215.87 #3664
0.56-H04
*-0.2
24.84
0.74
0.23
0.80
221.89 -2.53
0.30
25.33
0.74
0.22
0.89
220.96 13119
0.34
DDRM
26.10
0.75
36.44
43.52
5.09
ARGN 0=50
&1.2
25.91
0.75
0.33
33.68
44.90
3,88
$1.7
24.48
0.70
0.35
29.05
47.91
1.47
&
23.19
0.69
0.35
29.71
46.36
1.91
NLM 15
26.09
0.71
0.44
12.84
148.71 .3.75
8.79
AWGN, G0)
&0.8
26.24
0.72
0.43
12.46
148.78 12.49
8.69
XD
24.80
0.71
0.42
14.88
140.65
#2110
1..90=
varying the number of training samples. Surprisingly, we observe no change in the performance of the evaluated metrics for {10", 10*, 103, 102), 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
10, then only for small patch sizes of p S 5 we observe no performance drop compared to using larger sample size. This suggests that our method can be successfully deployed in few-shot settings, where 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 bias which results in worse covariance estimation.
Transporting the degraded measurement directly: Applying algorithm on the degraded input directly led to insufficient results as we see in appendix B.7
Re-applying the algorithm another time on kn: This is actually an interesting idea we tested on super-resolution when conducting our evaluations. As 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 VAF, their latent distribution aligns with that of the natural images. Hence, transporting another time does nothing (the transport operator is

---

y (degraded) x (benchmark prediction)
X (ours)
(original)
SwinIR
JPEG,10
Restormer
AWGN 050
ESRGAN
S
AWGNN2 SR JPEG,10
DDRM Swin2SR
Figure Our method (third column from the left) notably improves the results of several benchmark predictors (second column from the left) 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 impact the performance of restoration?: Our experiments showed that the class of images does not have a sigmificant impact on the performance (e.g. one could use images of does play significant role in attaining the best performance. L.c., to transport 512x512 images, best use training images of the same resolution. This drawback is somewhat mitigated by the few-shot nature of the algorithm.
improve images of dogs). However the resolution of images
6 Discussion
The pre-trained VAE used for purpose of our experiments exhibits rate of R = 48 on 5122 images which inevitably translates into sub-optimal distortion performance (38] Thus, the distortion performance our estimates are bounded by that of the pre-trained VAE, Lee. encoding and decoding completely clean and natural image does not yield result in perfect
Limitations:

---

(MMSE)
XS (ours)
D(E(x))
(original)
OPa
094
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 face identity (top row) 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 well (see fig. 66.
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, max cannot hope for consistent restorations.
Potential impact: Instead of using sophisticated and data-hungry generative models, we show is possible obtain photo-realistic results using simple tools like MMSE estimators and VAEs. We hope our fhw-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 well. This could potentially harm the results of medical image diagnosis, for example.
Acknoledgements
This research was partially supported by the Council For Higher Education Planning and Budgeting Committee.
10

---

References
[1] D. Freirich, T. Michacli, and R. Meir, "A theory of the distortion perception tradcoff in wasserstein space," in Advances in Neural Information Processing Systems, 2021.
distortion-percepiion
[2] Y. Blau and T. Michaeli, "The penception-distoriion tradeoff," in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
[3] B. Kawar, M. Elad, S. Ermon, and J. Song. "Denoising diffusion restoration models," in Advances in Neural Information Processing Systems, 2022.
[4] A. Korotin, V. Egiazarian. A. Asadulaev, A. Safin, and E. Burnaev, "Wasserstein-2 generative generrate networks." in International Conference on Learning Representations, 2021.
[5] A. Korotin, D. Selikhanovych, and E. Burnaev, "Neural optimal transport," in The Eleventh International Conference on Learning Representations, 2023.
[6] A. Makkuva, A. Taghvaei, S. Oh, and J. Lee, "Optimal transport mapping via input convex neural networks," in Proceedings of the 37th International Conference Muchine Learning, 2020.
[7] R. Rombach, A. Blattmann, D. Lorenz. P. Esser, and B. Ommer. "High-resolutio image synthesis with latent diffusion models," in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[8] Y. Li, C. Fang, J. Yang, 7. Wang, X. Lu, and M.-H. Yang, "Universal style transfer via feature transforms," in Advances in Neural Information Processing Systems, 2017.
[9] M. Lu, H. Zhao, A. Yao, Y. Chen, F. Xu, and L. Zhang, "A closed-form solution to universal style transfer," in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2019.
[10] Y. Mroueh, "Wasserstein style transfer," 2019.
[11] Liang, Cao. G. Sun, Zhang, L. Van Gool, R. Timofte, "Swinir: Image restoration using swin transformer," Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV) Workshops, 2021.
[12] S. Zamir, Arora, S. Khan, M. Hayat, Khan, and M.-H. Yang. "Restormer: Efficient transformer for high-resolution image restoration." in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[13] M. V. Conde, U.-J Choi, M. Burchi, and R. Timofte, "Swin2SR: Swinv2 transformer for compressed image super-resolution and restoration," in Proceedings of the European Conference on Computer Vision (ECCV) Workshops, 2022.
[14] K. Dahov, A. Foi, V. Katkovnik, and K. Egiazarian, "Image denoising hy sparse 3-d transform-domain collaborative filtering," IEEE Transactions on Image Processing, 2007.
[15] A. Buades, B. Coll, and J.-M. Morel, "A non-local algorithm for image denoising," in 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR (55), 2005.
[16] X. Wang, K. Yu, S. WU, J. Gu, Y. Liu, C. Dong, Y. Qiao, and C. Change Loy, "Esrgan: Enhanced super- resolution generative adversarial networks," in Proceedings of the European Conference on Computer Vision (ECCV) Workshops, 2018
[17] G. Ohayon, T. Adrai, G. Vaksman, M. Flad, and P. Milanfar, "High perceptual quality image denoising with posterior sampling cgan," in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops, 2021.
[18] B. Kawar, G. Vaksman, and M. Elad, "Stochastic image denoising by sampling from the posterior distribution." in Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV) Workshops, 2021.
[19] B. Kawar, G. Vaksman, and M. Elad, "SNIPS: Solving noisy inverse problems stochastically." in Advances in Neural Information Processing Systems, 2021.
[20] M. Aharon, M. Elad, and A. Bruckstein, "K-svd: An algorithm for designing overcomplete dictionaries for sparse representation," IEEE Transactions on Signal Processing, 2006.
[21] M. Cuturi, "Sinkhorn distances: Lightspeed computation of optimal transport," in Advances in Neural Information Processing Systems, 2013.
11

---

[22] G. Peyré and M. Cuturi, "Computational optimal transport," Foundations and Trends in Muchine Learning. 2019.
[23] M. Arjovsky, S. Chintala, and I.. Botton, "Wasserstein generative adversarial networks," in Proceedings of the 34th International Conference Machine Learning, 2017.
[24] V. Panaretos and Y. Zemel, An Invitation to Statistics in Wasserstein Space. Creative Media Partners, LLC, 2020.
[25] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, "Gans trained by two time-scale update rule converge local nash equilibrium." in Advances in Neural Information Processing Systems, 2017.
[26] C. Villani, Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften, Springer Berlin Heidelberg. 2008.
[27] D. P. Kingma and M. Welling, "Auto-Encoding Variational Bayes," in 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
[28] G. Wallace, "The JPEG still picture compression standard," IEEE Transactions on Consumer Electronics. 1992.
[29] Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and Fei-Fei, "Imagenet: A large-scale hierarchical image database," in 2009 IEEE Conference on Computer Vision and Puttern Recognition, 2009.
[30] Z. Wang, A. Bovik, H. Sheikh, and E. Simoncelli, "Image quality assessment: from error visibility to structural similarity," IEEE Transactions on Image Processing, 2004.
[31] R. Zhang, Isola, A. A. Efros, F. Shechtman, and O. Wang, "The unreasonable effectiveness of deep features as perceptual metrie," in Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
[32] J. Johnson, A. Alahi, and L. Fei-Fei, "Perceptual losses for real-time style transfer and super-resolution, in European Conference on Computer Vision (ECCV), 2016.
[33] T. Salimans, I. Goodfellow, W. Zaremba, V. Cheung, A. Radford, X. Chen, and X. Chen, "Improved techniques for training gans," in Advances in Neural Information Processing Systems, 2016.
[34] M. Binkkowkii, D. Sutherland, M. Arbel, and A. Gretton, "Demystifying MMD GANs," in International Conference on Learning Representations, 2018.
[35] C. Saharia, J. Ho, W. Chan, T. Salimans, D. Flect, and M. Norouzi, "Image super-resolution via iterative refinement," in IEEE Transactions on Puttern Analysis and Machine Intelligence, 2021.
[36] F. Agustsson and R. Timofte, "Ntire 2017 challenge on single image super-resolution: in The IEFE Conference Computer Vision and Pattern Recognition (CVPR) Workshops, 2017.
Dataset and study,"
[37] P. von Platen, S. Patil, A. Lozhkov, P. Cuenca, N. Lambert, K. Rasul, M. Davaadorj, and T. Wolf, "Diffusers: State-of-the-art diffusion models." httpp:: wtw.:/hwwbb.ooooggggggngggneeeee htte//fffffferrr 2022.
con/
[38] A. Alemi, B. Poole, I. Fischer, J. Dillon, R. A. Saurous, and K. Murphy, "Fixing a broken ELBO," in Proceedings of the 35th International Conference on Machine Learning, 2018.
[39] G. Ohayon, T. Adrai, M. Flad, and T. Michaeli, "Reasons for the superiority of stochastic estimators over deterministic ones: Robustness, consistency and perceptual quality." 2022.
[40] J. T. Flam, S. Chatterjee, K. Kansanen, and T. Ekman, "On mmse estimation: A linear model under gaussian mixture statistics," IEEE Transactions on Signal Processing, 2012.
[41] I. Kligvasser, T. Shaham, Y. Bahat, and T. Michaeli, "Deep self-dissimilarities as powerful visual finger- prints," Advances in Neural Information Processing Systems, 2021.
[42] T. Rott Shaham, T. Dekel, and T. Michaeli, "Singan: Learning generative model from a single natural image." in Computer Vision (ICCV), IEEE International Conference 2019.
12

---

Deep Optimal Transport: A Practical Algorithm for Photo-realistic Image Restoration Supplementary Material
A Background and extensions
A.1 Numerical Example
y=x+n
x'
y
Pxly
80
MSE(W2) function
12.4
X
7.1
X6
MSF
6.5
MSE 6.5
MSE 12.4
MSE : 6.9
0.6
W22
1.0
W2 1.0
W2 0.5
W2 0.6
Figure 2D Gaussian mixture denoising. Source samples are shown in blue. The MMSE estimator (x', orange) attains the best MSE but the worst perceptual index W2 The posterior samples (xly. purple) attain the best perceptual index but half the optimal MSE performance. The D max estimator (&g, green) maintains the MSE of x while attaining a perceptual quality close to xly. The DP curve is obtained by interpolating & and x" using eq. 4
To guide the reader in understanding the MMSE transport paradigm,
showcase method 2 2-dimensional denoising problem. To avoid a too trivial uni-modal example, we draw the clean signal from 4-components Gaussian mixture with non- non-trivial covariances. We derive linear MMSE and posterior estimators from [40 and proceed by applying the closed-form transport operator introduced eq.
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 not practical solution for 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 predicted by (2). In contrast, the D maxx estimator enjoys good perceptual index while maintaining a closc-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 Σx, When p is large, (typically p 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 emors sometimes adds up to negative cigenvalues. In this case, we clamp the negative eigenvalues to zero and use the stochastic (one-to-many) transport operator proposed by [I].
otochastic - (x1) Ei
E Ex, I2.
Mxx + Mxx W, Ex st xxx (x1
55)
when Dt. denotes the pseudo-inverse of Ex, (after negative eigenvalues where clamped) and NO T T' ) with T - ( Exx
Exa
Exs
We tried to avoid overflow when summing the images by using 64 bit precision
13

---

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 E(-). D(-) denote the encoder and decoder, respectively. Even (EE())) f, it is likely that E(:) "deforms" the space, I.e., [[(s)) E( A ||s i|I. which means that the optimal transport plan in the latent space could be differení 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
(E() 0E(x
E(x) 0E(
E
|
x11
E
(6)
0R
where [ is the determinant of the Jacobian matrix of E(-) However it is 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. 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 D max estimator.
With that, we argue that switching to a latent cost is actually a strength rather than a weakness of method. Indeed, using the MSE between deep latent variables has shown to be better fit to compare natural images than directly working in the pixel space [31. The authors of ( 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.] 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] let (c, He, We) denote the shape of the latent representation (CNN encoders produce 3-dimensional encoded tensors), where He, W. the spatial extent and e is the number of channels (i.e., the number of convolution kernels in the last convolution layer). The covariance matrices Σx,, Σx, contain (w. parameters, which may require large amount of samples for large latent images with H., W. >> 1. To mitigate the quadratic dependency He W., 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.p). . similar approximation exists in the style-transfer literature g 9. where instead of patches, only the pixels are considered (i.e., this is private case of our approach with 1). In section 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 [4] 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, 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 He W.. 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 Mepax1 which we stack into sample
assume that the overlapping patches are
matrix X NH. Xep?" We then aggregate the samples to compute the MVG statistics: = x N.H.W.
Σ = NH. NH.W.-T(C
=)(X - =²"". As N N.W 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 wish to transport.
This methodology easily be extrapolated other encoder architectures.
14

---

B.4 Size of the latent representation
When increasing the capacity of models with a fixed encoding rate, deepening is preferable than wideniig. Indeed, increasing makes the covariance estimation dramatically harder while increasing He. We enlarges the sample pool. Therefore, the VAE from 7 with 4 and He. We > 1 is a particularly good candidate for our method. For p = 3 for instance, the covariance matrix admits only 1296 parameters while each 5122 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 few-shot manner.
B.5 Transport
In a single pass on data set of natural images and (possibly different) data set of restored samples, we compute TMvG Pa. -p. (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 Σx is a singular matrix (see appendix A.2).
B.6 Decoding
Since the transported patches overlap, we "fold" them back into latent image Ko,tatent by averaging. The latent image is then decoded back to the pixel space, i.e. &0 = (28o,tatent)) Since E(-) is not invertible, the decoder D(-) is used 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 fully blind algorithm.
All the steps described above are summarized in fig.
B.7 Transporting the degraded measurement
We tried applying our algorithm on the degraded measurement directly. Indeed we observe quali- tatively and quantitatively that transporting the degraded measurement y amplifies the degradation (refer to fig. (
y (degraded)
x (SwinIR)
T Tp pp (v)
&0
x (original)
Figure 8: Transporting the degraded measurement (JPEG, directly is not enough to restore the image. It can sometimes even exacerbate the degradation. Quantitatively, the degraded sample y has better PSNR and FID than its transported version (respectively 27. 77.26 dB and 13.88 FID v.s. 23.69 dB and 15.88 FID).
15


Reviews:
[{"review_id": "BTVRLZApke", "reviewer": "Reviewer_tvbs", "summary": "In this paper, the authors propose a few-shot algorithm to obtain higher quality restored images of the given model like VAEs and diffusion models. Specifically, the optimal transport map in the latent space is computed through the representations of the real images and the reconstructed images. By applying the OT map, a better restoration can be obtained. Experiments show that the proposed method is effecitive.", "questions": "Please see the weakness part.", "limitations": "None", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper successfully applying the theory propose by [1] in image restoration and can generate high quality reconstructed images.", "weaknesses": "* The presentation of the paper is not good and some concepts are unclear. For example,\n    * In line 116-117, what's the meaning of $x^*$ and $\\hat{x_0}$\n    * Is $x^*$ sampled from $p(x|z)$ that achieves the minimal MSE between the reconstruction and the input? Similarly, how to define $\\hat{x_0}$ and what's the meaning of max MSE error?\n    * How to define $p_x, p_x^*$?\n    * Line 121, OT plan should be between two distributions, not to data points. \n\n* Why does the unpaired dataset give better performance instead of paired dataset? It doesn't seem to make sense and the necessary explanation is needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> In line `116-117`, what's the meaning of $x^*$  and $\\hat{x}_0$ ?\n\n$x^*$ is the MMSE estimate, being the posterior mean. This solution gives the best MSE distortion performance, but at the cost of being of poor visual quality. $\\hat{x}_0$ on the other hand, is the Dmax solution - being of perfect perceptual quality while giving the smallest distortion possible. Note that these are two extremes in the Perception-Distortion curve, which we refer to in **Figure 1** of the paper. We will update the text in line `116-107` to include these more detailed explanations.\n\n> Is $x^*$  sampled from $p(x|z)$  that achieves the minimal MSE between the reconstruction and the input? Similarly, how to define $\\hat{x}_0$ and what's the meaning of max MSE error ?\n\nAs $x^*$ is the MMSE, it cannot be obtained as a single sample from the posterior. Note, however, that $x^*$ can be approximated by drawing many samples from $p(x|y)$ and averaging them, as it is the posterior mean. This property is not being used in the paper. As for $\\hat{x}_0$, it is defined in line `117` as the estimator attaining the *minimal* MSE while having perfect perceptual quality. This is not to be confused with the *notation* Dmax, referring to this MSE quantitiy, introduced by [(Blau & Michaeli, 2018)](https://openaccess.thecvf.com/content_cvpr_2018/html/Blau_The_Perception-Distortion_Tradeoff_CVPR_2018_paper.html), to refer to the maximal distortion on the perception-distortion curve (see **Figure 1**).\nWe will clarify this nuance in the revised version of the paper.\n\n> How to define $p_x$, $p_{x^*}$ ? \n\n$p_x$ and $p_{x^*}$ are the probability distributions of the random variables $x$ and $x^*$. The reviewer is correct to note that the formal definition is lacking in the paper. As we state in **Section 3.2**, $x$ and $x^*$ are random variables defined over $\\mathbb{R}^n$, so their definitions of $p_x$ and $p_x^*$ are often omitted for conciseness, like in [(Freirich et. al, 2021)](https://proceedings.neurips.cc/paper/2021/hash/d77e68596c15c53c2a33ad143739902d-Abstract.html). We will add these definitions for completeness. \n\n> Line `121`, OT plan should be between two distributions, not to data points.\n\nThe reviewer is correct, and the OT plan is indeed performed between $p_x$ and $p_{x^*}$ as the notation $T_{p_{x^*} \\longrightarrow p_x}$ suggests and as we define in `101-102` following [(C. Villani, 2008)](https://cedricvillani.org/sites/dev/files/old_images/2012/08/preprint-1.pdf). We will clarify this better in the revised version of the paper.\n\n> Why does the unpaired dataset give better performance instead of paired dataset? It doesn't seem to make sense and the necessary explanation is needed.\n\nAs explained in line `252-253`, it appears that using paired updates actually diverges from the algorithm introduced in **Section 4.1** and might introduce a statistical bias which hinders the covariance matrix estimation. As emphasized in the text, it is only a *hypothesis* and this phenomenon requires further research. Unfortunately, it is not the primary result of our algorithm and its formal explanation is out of the paper’s scope."}, {"review_id": "artezv0lBk", "reviewer": "Reviewer_YhDf", "summary": "The paper proposes 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.", "questions": "Please see the weaknesses. I am willing to improve the score if the concerns are addressed well.", "limitations": "The limitations have been described in the paper. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The method is plug-and-play, requires only a few samples, and does not require further training.\n2. The method approximates the optimal transport by a linear transformation in the latent space of a variational auto-encoder, which is somewhat novel.", "weaknesses": "1. The main concern is in the experiments section.\n(1) It would be better to evaluate the method on real-world datasets (e.g., SIDD, DND, and RealSR datasets), which may make more sense.\n(2) Quantitative results of $\\hat{x}_{0}$ should be given in Table 1.\n(3) Quantitative results of some important ablation experiments (e.g., paired vs. unpaired samples, and transporting the degraded measurement directly) can be given.\n(4) LPIPS is generally regarded as a perception metric in image restoration tasks. Since FID, IS and KID are not very stable, they are generally not used in image restoration tasks. Just looking at PSNR, SSIM and LPIPS, the method doesn't seem to achieve a good distortion-perception tradeoff.\n2. The interpolation constant $a$ seems to be adjusted manually, which may be inflexible.\n3. Projected distribution loss [1] and sliced Wasserstein loss  [2]  also show better distortion-perception tradeoff in image restoration tasks, although it needs to be used directly for training. I don't know if the authors have tried this way. Further elaboration may be needed on how the proposed method relates to and differs from this loss.\n\n[1] Projected Distribution Loss for Image Enhancement. ICCP 2021.\n[2] Self-supervised learning for real-world super-resolution from dual zoomed observations. ECCV 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> (4) LPIPS is generally regarded as a perception metric in image restoration tasks. Since FID, IS and KID are not very stable, they are generally not used in image restoration tasks. Just looking at PSNR, SSIM and LPIPS, the method doesn't seem to achieve a good distortion-perception tradeoff.\n\nPlease note that, *by definition*, LPIPS is a distortion metric, as it is evaluated on pairs of images. Interestingly, the original perception-distortion paper [(Blau & Michaeli, 2018)](https://openaccess.thecvf.com/content_cvpr_2018/html/Blau_The_Perception-Distortion_Tradeoff_CVPR_2018_paper.html) already classified the VGG loss - the ancestor of LPIPS - to be a distortion, on which the tradeoff exists. Also, the reviewer is correct to point out that FID, IS and KID are not stable on common restoration datasets, **but this is largely caused by these datasets' small size**.\n\n> (1) It would be better to evaluate the method on real-world datasets (e.g., SIDD, DND, and RealSR datasets), which may make more sense.\n\nFollowing the comment above, it is why we perform our evaluation on the 50,000 Imagenet validation samples, following very popular image restoration papers like ([Saharia et. al, 2021](https://arxiv.org/abs/2104.07636),[Rombach et. al, 2021](https://arxiv.org/abs/2112.10752)). This is also why 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. LPIPS is a distortion, not a perceptual quality metric.\nFinally, we do conduct the qualitative evaluation on popular samples from DIV2K or Set14. This is actually not a trivial question that was raised by another reviewer. We shall include these clarifications after line 194.\n\n> (2) Quantitative results of $\\hat{x}_0$ should be given in Table 1. (3) Quantitative results of some important ablation experiments (e.g., paired vs. unpaired samples, and transporting the degraded measurement directly) can be given. \n\nIn **Table 1**, we preferred not overloading the reader with yet additional rows but rather have them focus on the way we can trade perception over distortion with different values of $\\alpha$, which constitute the main result of the paper. However, concerns around the reported values of $\\alpha$ were raised by other reviewers and we understand the importance of these results. Note that quantitative results of $\\hat{x}_0$ are visible in **Figure 3**, but we shall add the exact performance in **Table 1**. We will also add the ablation figures in a complementary table in the appendix.\n\n> The interpolation constant $\\alpha$ seems to be adjusted manually, which may be inflexible.\n\nLike any other hyper-parameter, $\\alpha$ can improve performance with some tuning when approaching a new task or dataset. We argue that the few-shot nature of our algorithm makes this tuning actually practical (it does not need to be set before performing some expensive training). In any case, $\\alpha=0$ consistently improves perceptual quality for all tasks and models considered (as expected from the theory). We consider it to be a satisfying default choice, such that manually adjusting $\\alpha$ is not too great of a concern. As other reviewers pointed out the lack of consistency in selecting the reported values of $\\alpha$, we will make sure to clarify this in the experiment section, and additionally report the performance for $\\alpha=0$ in **Table 1**.\n\n> Projected distribution loss [1] and sliced Wasserstein loss [2] also show better distortion-perception tradeoff in image restoration tasks, although it needs to be used directly for training. I don't know if the authors have tried this way. Further elaboration may be needed on how the proposed method relates to and differs from this loss.\n\nThe paper is interested in the MSE-W2 tradeoff following the transport theorem introduced by [(Freirich et. al, 2021)](https://proceedings.neurips.cc/paper/2021/hash/d77e68596c15c53c2a33ad143739902d-Abstract.html). From a theoretical standpoint, it is not clear at all how we can obtain the Dmax estimator $\\hat{x}_0$ using projected losses. In practice, we carefully designed our algorithm to improve existing restoration models at test time. Therefore, the aforementioned losses are not obvious candidate to boost performance. With that said, we absolutely agree with the reviewer that approximating the Dmax estimator in other perception-distortion planes is a most interesting topic and we are determined to investigate this in future research."}, {"review_id": "VJMfuoKsAv", "reviewer": "Reviewer_HWzm", "summary": "This paper presents an image restoration algorithm that builds upon a trained network to further minimize MSE. \nTo achieve this goal, an optimal transport was approximated by a linear transformation in the latent space. Visual results show clear improvement by using the proposed approach.", "questions": "By applying the proposed algorithm another time, will the results be further improved?", "limitations": "None.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea is straightforward as optimizing delta is common in traditional image restoration\n- The visual results are convincing", "weaknesses": "- The uncertain of the proposed idea is low so relatively the novelty is less.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> `[Summary]` This paper presents an image restoration algorithm that builds upon a trained network to further minimize MSE\n\nWe would like to clarify that the main goal of our algorithm is actually to improve perceptual quality. As a side-effect, we discovered empirically that we could extend the theory introduced in [(Freirich et. al)](https://proceedings.neurips.cc/paper/2021/hash/d77e68596c15c53c2a33ad143739902d-Abstract.html) to also improve MSE, but this is not the main contribution of our paper.\n\n> The uncertain of the proposed idea is low so relatively the novelty is less.\n\nThe main contribution of this work is a proposed algorithm that improves the perceptual quality of *any* restoration algorithm and *any* inverse problem task with few unpaired examples to train on. As far as we know, these achievements are unprecedented in the literature, and thus we believe that the novelty in this work is significant. \n\n> By applying the proposed algorithm another time, will the results be further improved?\n\nThis 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 the Identity matrix). We are only left with the reconstruction error introduced by the encoding and decoding of the images, which degrades the MSE performance. We agree with the reviewer that this experiment adds to the reader's understanding and helps draw the limitations of our algorithm. We shall add it in the supplementary material.\n"}, {"review_id": "GTQE6skXxM", "reviewer": "Reviewer_Md85", "summary": "This paper presents an image restoration algorithm targeting at further restore the processed images that have been restored by pre-trained restoration models. As most of the restoration tasks use MSE as the main criteria for restoration, the restored images tend to be blurred to achieve better PSNR performance. This work takes advantage of optimal transport from source to target distribution to improve the perceptual quality, where the optimal transport operator is computed from distribution mapping from the restored image to original images. ", "questions": "What is the meaning of the term “D(E(x))” in table1 and figure6? Does it refer to the decoding results obtained from the encoded latent representation of x? ", "limitations": "Authors have addressed the limitations of the method in the paper.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+ This work proposes an image enhancement method after the degraded images restored by pre-trained restoration models. \n+ The tradeoff between MSE and perceptual quality is discussed and the balancing process between two criteria are formulated and visualized on multiple tasks.\n+ Only a few unpaired images are required for the optimal transportation computing.\n", "weaknesses": "- As the proposed method uses a few images for optimal transport computing. How to guarantee that the distribution of test data is aligned with the training data. Does the selection on the training data have an impact on the performance of restoration?\n- In table 1, the numerical results of restoration with different alpha settings are given. For different restoration tasks, the alpha is set to different values while comparison. What criteria are used to select the parameter? \n- The experiments are conducted on ImageNet dataset. But for most restoration tasks, the restored images have relatively large resolution. Images in ImageNet dataset have considerable small resolution and its original images have been compressed with lossy compression. Why authors evaluate the method upon the task-specific datasets.\n- The method has similar mechanism with flow-based image deblurring or image sharpening. The methods itself has limited restoration capabilities on degraded image restoration. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> As the proposed method uses a few images for optimal transport computing. How to guarantee that the distribution of test data is aligned with the training data. Does the selection on the training data have an impact on the performance of restoration?\n\nOur 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. We thank the reviewer for drawing our attention to this topic, and we will make sure to include this discussion in the revised version of the paper.\n\n>In table 1, the numerical results of restoration with different alpha settings are given. For different restoration tasks, the alpha is set to different values while comparison. What criteria are used to select the parameter?\n\nWe selected $\\alpha$ by observing the interpolation & extrapolation curves in **Figure 3** and chose the values with the most significant effect on performance. Like any other hyper-parameter, $\\alpha$ can improve performance with some tuning when approaching a new task or dataset. We argue that the few-shot nature of our algorithm makes this tuning actually practical (it does not need to be set before performing some expensive training). In any case, $\\alpha=0$ consistently improves perceptual quality for all tasks and models considered (as expected from the theory). We consider it to be a satisfying default choice, such that manually adjusting $\\alpha$ is not too great of a concern. As other reviewers pointed out the lack of consistency in selecting the reported values of $\\alpha$, we will make sure to clarify this in the experiment section, and additionally report the performance for $\\alpha=0$ in **Table 1**.\n\n> The experiments are conducted on ImageNet dataset. But for most restoration tasks, the restored images have relatively large resolution. Images in ImageNet dataset have considerable small resolution and its original images have been compressed with lossy compression. Why authors evaluate the method upon the task-specific datasets.\n\nAs pointed out by another reviewer, the perception metrics FID, IS and KID are not stable on common restoration datasets with small-to-medium size. Hence, we perform our evaluation on the 50,000 Imagenet validation samples, following very popular image restoration papers like ([Saharia et. al, 2021](https://arxiv.org/abs/2104.07636),[Rombach et. al, 2021](https://arxiv.org/abs/2112.10752)). This is also why 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. Finally, we conduct the qualitative evaluation on popular samples from DIV2K or Set14. This is actually not a trivial question that was raised by another reviewer. We shall include these clarifications after line `194` in the paper.\n\n> The method has similar mechanism with flow-based image deblurring or image sharpening. The methods itself has limited restoration capabilities on degraded image restoration.\n\nThere is actually much similarity with even older, classical works which apply a carefully chosen linear transformation on all overlapping patches of the degraded image. We consider this as an advantage of our method rather than a drawback. Applying linear transformations in the latent space could prove to be a powerful yet simple approach with interesting properties like few-shot learning and robustness.\nOur algorithm is not designed to directly restore degraded images but rather improve the perceptual quality of **any** restoration model at test time. In this regard, we demonstrate significant performance gains on a wide range of models (regression, GANs, diffusion) and tasks (super-resolution, denoising, JPEG artifact removal etc.).\n"}, {"review_id": "hX2HVY27th", "reviewer": "Reviewer_XzV9", "summary": "This paper considers the high perceptual image restoration problem. The authors propose a method to control the tradeoff between the perceptual quality and distortion (e.g. MSE) of a pretrained restoration model. The mehtod is developed based on a recent theory [1] on the tradeoff between MSE distortion and perception quality measured by Wasserstein-2 distance, which indicates that minimum MSE restoration under perfect perception constraint can be achieved by an optimal transport from an minimum MSE estimator to a perceptual estimator. Meanwhile, in theory the perception-distortion tradeoff, with perception quality measured by Wasserstein-2 distance and distortion measured by MSE, can be controlled by a linear interpolation between an minimum MSE estimator and a perfect perceptual estimator.  Based on this theory, the authors propose to transport the output of a pretrained restoration model to improve the perceptual quality. The optimal transport is performed in the latent space under a Gaussian distribution assumption, with which closed-form optimal transport can be derived and, finally, the perception-distortion tradeoff is controlled by linearly transforming the first- and second-order statistics (means and covariances) in latent space of a pre-trained model. Experimental results on various image restoration tasks have been provided to demonstrate the effectiveness of the proposed method.\n", "questions": "The authors mentioned the limitation in restoring human faces and text images, of which the pixels are highly correlated with each other. Is it because the patch size set to $p={3,5}$ in this paper is relative small and fail to extract mutual information while transporting?\n\nAnother question is that the authors demonstrate that a larger patch size $7≤p≤15$ can yield slightly worse PSNR. Is it because the distribution is only close to MVG with small $p$? The assumption might be far from valid when $p$ gets larger, especially for face and text images, which are highly structured.\n", "limitations": "The authors have adequately addressed the limitations.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed method is interesting, which can be viewed as a post-processing method that, given a pre-trained restoration model, it can achieve tradeoff between perception quality (measured by Wasserstein-2 distance) and distortion (measured by MSE) by a simple training stage computing empirical fisrt- and second-order statistics (means and covariances) in the latent space under a Gaussian distribution assumption.\n2. This paper is clearly written and easy to read.\n2. Experimental results verified the effectiveness of the proposed method.\n", "weaknesses": "1. Lack of theorectical novelty. The proposed method is based on the theory in [1] (see Section 3) and appears to be more of an extended evaluation of the results in [1]. Besides, the adopted appraoch that performs transport in the latent space is also not new and borrowed from existing work.\n\n2. While the proposed method is interesting, its post-processing nature may limit its practical use since it requires a test-time training procedure. In comparison, one-stage restoration models would be preferred in practical applications. Although the test-time training only uses a dozen images restored by the pretrained model, retraining may be necessary for it to perform well when the degradation strength changes.\n\n3. The assumption that the latent representations follow a Multivariate Gaussian distribution is relatively strong. The authors claim that their method can be applied to any pre-trained model. I’m afraid this assumption might limit its application as it requires the latent representation distribution to be (close to) Multivariate Gaussian.\n\n4. Using optimal transport for image restoration is not new, e.g. for image denoising and super-resolution, which has demonstrated the capability of yielding high perception quality. A natural question is how does the proposed method compare with such one stage optimal transport based restoration approaches.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Lack of theorectical novelty. The proposed method is based on the theory in [1] (see Section 3) and appears to be more of an extended evaluation of the results in [1]. Besides, the adopted appraoch that performs transport in the latent space is also not new and borrowed from existing work.\n\nThe reviewer is correct to point out that we extend the theory introduced by [(Freirich et. al)](https://proceedings.neurips.cc/paper/2021/hash/d77e68596c15c53c2a33ad143739902d-Abstract.html). Note however, that this work is only theoretical and does not propose a practical way to achieve the Dmax estimator.\nSimilarly, we do not hide the fact that our latent transport approach is inspired by prior works (see lines 50-53).\nNonetheless, we introduce several key-improvements to the naive channel-wise transport that enable our algorithm to improve perceptual quality of *any* restoration algorithm and *any* inverse problem task with few unpaired examples to train on. As far as we know, these achievements are unprecedented in the literature, and thus we are confident the novelty in our work is significant.\nWe will make sure to emphasize this in the **related-work Section**.\n\n> While the proposed method is interesting, its post-processing nature may limit its practical use since it requires a test-time training procedure. In comparison, one-stage restoration models would be preferred in practical applications.\n\nThe reviewer's standpoint is interesting because it contrasts with that of the other reviewers, which identified the plug-and-play property of our approach as a strength rather than a weakness. In the vast majority of applications, it seems sound to assume that we can obtain a dozen of unpaired images to apply the test-time procedure, in which case all the one-stage restoration models we considered were improved by our algorithm.\n\n> Although the test-time training only uses a dozen images restored by the pretrained model, retraining may be necessary for it to perform well when the degradation strength changes.\n\nAs noted by the reviewer, we apply the algorithm at test-time, on the output of a given model. Interestingly, retraining is necessary only if the pre-trained model itself needs retraining, in which case the few-shot procedure is likely to consume a negligible fraction of the computation and time resources allocated to retrain the one-stage model. We will make sure to clarify this in the revised version of the paper.\n\n> The assumption that the latent representations follow a Multivariate Gaussian distribution is relatively strong. The authors claim that their method can be applied to any pre-trained model. I’m afraid this assumption might limit its application as it requires the latent representation distribution to be (close to) Multivariate Gaussian.\n\nActually, we perform the transport in the latent space of a Variational Auto-Encoder, which is precisely trained to achieve a normally distributed latent representation. Additionally and as pointed out by the reviewer, Gaussian close-form transport in the latent space is not new, and has proven its stability over a wide range of application.\n\n> Using optimal transport for image restoration is not new, e.g. for image denoising and super-resolution, which has demonstrated the capability of yielding high perception quality. A natural question is how does the proposed method compare with such one stage optimal transport based restoration approaches.\n\nSince our algorithm is conceived to improve an already existing method, it cannot be compared toe-to-toe with other one-stage approaches (whether they are transport based or not). Nonetheless, we can apply our algorithm to improve further the performance of these methods.\n\n> The authors mentioned the limitation in restoring human faces and text images, of which the pixels are highly correlated with each other. Is it because the patch size set to $p=3,5$ in this paper is relative small and fail to extract mutual information while transporting?\n\nThis is an interesting idea. Note however, that the Encoder used to embed the images admits a large receptive field (typically 64 pixels depending on the configuration). Hence, it is safe to assume that the spatial correlations are conserved even with a small latent patch-size. On the other hand, we show in the **Limitation Section** that human faces and text are already largely distorted by the Variational Auto-Encoder (VAE). I.e., encoding and then decoding a clean image containing text (without any transport) typically results in deteriorated text. Therefore, we deem more plausible that the problem resides in the VAE rather than the patch-size. This is an important clarification we will add to **Section 6**. \n\n> Another question is that the authors demonstrate that a larger patch size  can yield slightly worse PSNR. Is it because the distribution is only close to MVG with small $p$ ? The assumption might be far from valid when $7 \n\\leq p \\leq 15$ gets larger, especially for face and text images, which are highly structured.\n\nAs discussed earlier, we use the latent representation of a VAE, therefore latent images' prior follows a MVG. The reviewer is correct to point out that larger patch-size typically result in a deterioration in the approximated distribution. As explained in **Appendices B.3-B.4**, The dimensionality of the prior grows quadratically with the patch-size, and therefore the number of covariance matrix parameters is proportional to $p^4$. With that, the number of samples available to compute these parameters shrinks quadratically with $p$. In this regard, it is much more challenging to accurately approximate the transport operator for $7 \n\\leq p \\leq 15$ compared to $p=3,5$."}]

Author rebuttal:

Meta-review: The final ratings for the rebuttal after the paper were three accepts (two borderline and one weak), with the fourth reviewer recommending reject. However, the negative reviewer did not elaborate on which of their concerns remained and why they were unsatisfied with the rebuttal.

Given the mix of scores, the AC read the paper as well as the reviews and discussions carefully. Based on this reading, the AC believes that the weaknesses listed by the negative reviewer have either been addressed by the rebuttal or are not a bar for acceptance.

Moreover, the AC finds this to be an interesting paper with a novel idea, and with both good theory and practical empirical results. The paper would be of considerable interest to the NeurIPS audience --- beyond offering a new useful algorithm for restoration tasks, the ideas in the paper have the potential to be adapted by other researchers in different settings. Therefore, the AC is happy to recommend acceptance.
Decision: Accept (poster)