Title: Deep Optimal Transport: A Practical Algorithm for Photo-realistic Image Restoration
Venue: NeurIPS 2023
Accepted: yes
Keywords: Computer Vision

Image Restoration

Deep Learning

Perceptual Quality
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

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 the source data. Thus, to improve the perceptual quality of a predictor that was originally trained to minimize MSE, we approximate the optimal transport by a linear transformation in the latent space of a variational auto-encoder, which we compute in closed-form using empirical means and covariances. Going beyond the theory, we find that applying the same procedure on models that were initially trained to achieve high perceptual quality, typically improves their perceptual quality even further. And by interpolating the results with the original output of the model, we can improve their MSE on the expense of perceptual quality. We illustrate our method on a variety of degradations applied to general content images of arbitrary dimensions.

1 Introduction

Many image restoration algorithms aim to recover a clean source image from its degraded version. The performance of such algorithms is often evaluated in terms of their average distortion , which measures the discrepancy between restored images and their corresponding clean sources, as well as perceptual quality , which refers to the extent to which restored images resemble natural images. The work in [2] exposed a fundamental trade-off between distortion and perceptual quality, where the latter is measured using a perceptual index that quantifies the statistical divergence between the distribution of restored images and the distribution of natural images. The

PD-function

x

n

σ

x

∗

↓

4

+

DRM)

=25

n

x

↓

4

+

n

σ

=25

x

∗

(DDRM)

ˆ x

0

(ours)

x

(original)

x

↓

4

+

n

σ

=25

x

∗

x

↓

4

+

n

σ

=25

x

∗

(DDRM)

σ

=25

x

(DDRM)

ˆ x

0

(ours)

x

(DDRM)

ˆ x

0

ˆ x

0

(ours)

x

ˆ x

0

(ours)

(DDRM)

x

∗

(DDRM)

∗

(ours)

(original)

x

trade-off curve reveals the predictor that achieves the lowest possible distortion, denoted as D max , while maintaining perfect perceptual quality (refer to fig. 1).

Following the methodology introduced by [2], it has become common practice to compare restoration methods on the perception-distortion (PD) plane, with many methods aiming to reach the elusive D max point. In this paper, we present a practical approach to approximate the D max predictor where distortion is measured using the MSE and perceptual quality is measured by the Wasserstein-2 distance ( W 2 ) between the distributions of restored and real images. Our approach is based on the recent work [1] which demonstrated that the D max predictor can be obtained using optimal transport (OT) from the output distribution of the MMSE predictor to the distribution of natural images. By applying an optimal transport plan to an MMSE restoration resulting from a degraded image, we can produce D max estimations by transporting the MMSE restored estimate.

Although progress has been made in finding OT plans between image distributions [4-6], it remains challenging task, particularly for high-dimensional distributions. Therefore, we propose an approximation method for the D max estimator by performing transportation in the latent space of a pre-trained auto-encoder. A similar strategy was successfully employed in the context of image generation in [7], showing effectiveness in reducing complexity and preserving details.

Inspired by the style transfer literature [8-10], we assume that the latent representations follow a Multivariate Gaussian (MVG) distribution. Thus, by considering the first and second-order statistics of the embedded MMSE estimates and embedded natural images, we can compute the well-known closed form solution of the OT operator between two Gaussians. To further reduce complexity, we make additional assumptions about the structure of the latent covariance matrices, enabling the computation of the OT operator with as few as 10 unpaired MMSE restored and clean samples. This approach leads to a few-shot algorithm that significantly enhances visual quality.

Interestingly, our method can even improve the visual quality of generative models that were trained to achieve high perceptual quality in the first place (see fig. 2). Furthermore, by adjusting a single interpolation parameter, we can trade off perception for distortion, resulting in marginal improvements in the distortion performance of some regression models that were trained to prioritize source fidelity. We demonstrate the improved photo-realism of our approach on a variety of tasks and models, including GAN and diffusion-based methods, using high-resolution ( e.g. , 512 2 px) general-content images with arbitrary aspect ratios. 1

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 (e.g.,

∗

(origin

ˆ x

(DDR

0

x

(ours

MSE) [11-15]. The latter category includes more recent works that usually involve generative models like Generative Adversarial Networks (GANs) [16-18], or diffusion-based techniques [3, 19].

This paper searches for the theoretical D max estimator which minimizes the MSE under a perfect perceptual quality constraint. However, our method is not a stand-alone restoration algorithm, i.e., its input is not the degraded image. Rather, it can be applied on top of any existing predictor. We provide a new way to potentially improve the performance (either MSE or perceptual) of any given estimator in a few-shot, plug-and-play fashion.

Provided with an image latent representation method (e.g., an auto-encoder), our algorithm applies a linear transformation on all the overlapping patches of its input (after encoding). In this regard, its functioning is not far from classical image restoration methods [14, 15, 20].

2.1 Wasserstein-2 transport

While many successful approaches exist to compute the W 2 distance between discrete, low or medium dimensional densities [21, 22], it is far more challenging to determine an optimal transport plan in the continuous, high-dimensional setting. In fact, the task of even computing the Wasserstein distance (without its optimal plan) on empirical distributions has drawn significant attention with WGANs [23]. Thus, computing the transport operator requires to optimize the W 2 distance but with an additional ordering constraint on the generator, which proved to be a challenging task when dealing with real-world data sets [4-6]. An attempt to sidestep this difficulty would be to use the Gelbrich distance [1, 24], which lower bounds the W 2 distance and depends only on the first two moments of the distributions. Nevertheless, it is only a good estimate when the support of the distributions has elliptical level sets. To address this, one can find a low-dimensional embedding where (i) high dimensionality is no longer an issue, (ii) the distributions are not degenerated, and (iii) the Gelbrich distance equals the Wasserstein-2 distance. A possible option would be to use the bottleneck of an auto-encoder. This approach was adopted by style transfer works [8-10], and is also widely used by tools that compare image distributions like the Fréchet Inception Distance (FID) [25]. Both use a convolutional encoder and consider the pixels of the latent embedding (the vectors that span across the channel dimension) as a MVG distribution.

3 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 on R m and R n , respectively. The set of all transport plans, which are probability measures π on R m × R n with marginals µ and ν , is denoted by Π( µ, ν ) . The Wasserstein-2 distance between µ and ν is defined as follows:

A transport plan that achieves this infimum is called an optimal transport plan between µ and ν . If µ has a density ( i.e. is absolutely continuous w.r.t. the Lesbegue measure), there exists a measurable function T µ -→ ν : R m -→ R n such that, if x 1 ∼ µ and x 2 ∼ ν are two random variables, then x 2 a.s = T µ -→ ν ( x 1 ) . We refer to T µ -→ ν as the optimal transport operator between µ and ν . Like in [5], we also abuse this notation even when π is non-degenerate, in which case T µ -→ ν represents a one-to-many (stochastic) mapping.

Additionally, when considering two Multivariate Gaussians (MVGs) x 1 and x 2 with x 1 ∼ N ( µ x 1 , Σ x 1 ) and x 2 ∼ N ( µ x 2 , Σ x 2 ) , respectively, and assuming that Σ x 1 and Σ x 2 are nonsingular, there exists a closed-form solution for the optimal transport operator, which is deterministic and linear:

where a symmetric and positive definite square root of the matrices is chosen.

3.2 Wasserstein-2 MSE tradeoff

We build upon the problem setting introduced in [1, 2] to establish our analysis. We consider the following scenario: x ∈ R n represents a source natural image, y ∈ R m represents its degraded version, and we assume that the posterior p x | y ( ·| y ) is non-degenerate for almost any y . Our objective is to construct an estimator ˆ x that predicts x given y . A valid estimator ˆ x should be independent of x given y . Finally, p x , p x ∗ and p ˆ x 0 denote the probability distributions associated with the random variables x , x ∗ and ˆ x 0 , respectively.

Let x ∗ = E [ x | y ] denote the MMSE estimator that achieves the minimal MSE, i.e. , MSE ( x , x ∗ ) = D min . Additionally, let ˆ x 0 denote the D max estimator, which among all estimators satisfying W 2 ( p x , p ˆ x 0 ) = 0 , attains the minimal MSE, namely, MSE ( x , ˆ x 0 ) = D max (refer to fig. 1). Notably, as discussed in [1], these estimators have a compelling property: their joint distribution p ˆ x 0 , x ∗ is an optimal transport plan between x ∗ and x , characterized by the following optimization problem:

In other words, finding ˆ x 0 is equivalent to finding an optimal transport plan from x ∗ to x . Then, the D max estimator is simply ˆ x 0 = T p x ∗ -→ p x ( x ∗ ) .

This estimator is particularly useful as it allows to obtain any point on the perception-distortion function through a naive linear interpolation with the MMSE estimator. Specifically, we can define the interpolated estimator ˆ x P as follows:

where 0 ≤ α ≤ 1 is an interpolation constant [1] that depends on the perceptual index of x ∗ and the desired perceptual index 0 ≤ P = W 2 ( x , ˆ x P ) (refer to fig. 1).

1.5

5.9

1.2

FID

3.4

ˆ

x

-

1

.

0

ˆ

x

2

.

0

x

∗

(

SwinIR

)

ˆ

x

0

SISR

×

4

2.0

4.4

0.8

1.8

ˆ

x

-

0

.

5

ˆ

x

1

.

5

x

∗

(

ESRGAN

)

ˆ

x

0

SISR

×

4

1.0

MSE

(

×

10

-

3

)

8.1

2.7

6.5

ˆ

x

-

1

.

0

ˆ

x

2

.

0

x

∗

(

Restormer

)

ˆ

x

0

AWGD

σ

=50

10.5

MSE

(

×

10

-

3

)

22.0

166.5

FID

209.0

ˆ

x

-

1

.

0

ˆ

x

2

.

0

x

∗

(

Swin2SR

)

ˆ

x

0

SR

×

4

JPEG

q

=10

2.5

MSE

(

×

10

-

3

)

10.2

29.0

36.8

ˆ

x

-

1

.

0

ˆ

x

1

.

9

x

∗

(

DDRM

)

ˆ

x

0

SR

×

4

AWGD

σ

=25

α

∈

[

-

1

,

0]

α

∈

[1

,

2]

α

∈

[0

,

1]

x

∗

(

predictor

)

ˆ

x

0

(

D

)

4 Method

We start by describing the general flow of our proposed algorithm, and then move to elaborate on each of its components.

Theoretically speaking, our algorithm, combined with any given MMSE estimator, is an approximation of the D max estimator. In practice, however, it can be combined with any type of estimator and potentially improve its perceptual quality. I.e., it can be combined with an estimator that optimizes

flatten each patch and

aggregate

Our improved

estimation

Training

Inference

flatten each patch

unflatten

unfold

unfold

unfold

fold

A given estimate

distortion (e.g., SwinIR [11]) and improve its perceptual quality at the expense of distortion, or it can be combined with an estimator that optimizes perceptual quality (e.g., DDRM [3]) and improve its perceptual quality even further. As a result, our algorithm is agnostic to the type of degradation. To clarify, our algorithm is not really a restoration algorithm by itself, but rather a wrapper which can potentially improve the perceptual quality of any given estimator.

4.1 The algorithm

The main goal of our algorithm is to approximate the optimal transport plan between p ˆ x and p x , namely, T p ˆ x -→ p x , where ˆ x is a given estimator. Theoretically, with such an operator, one could optimally transport ˆ x such that, with minimal loss in MSE performance, the transported estimator would attain perfect perceptual quality. Computing T p ˆ x -→ p x for high dimensional distributions is a difficult task, involving complex (often adversarial) optimization (see discussion in section 2). To solve this, we perform several assumptions and approximations that allow us to efficiently compute a closed form transport operator that approximates T p ˆ x -→ p x . The flow of the algorithm is presented in fig. 4, and goes as follows:

Encoding : In the training stage we encode N natural images { x ( i ) } N i =1 and N restored samples { ˆ x ( i ) } N i =1 (unpaired) into their latent representations, { x ( i ) e } N i =1 and { ˆ x ( i ) e } N i =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, H e , W e ) . In the inference stage we perform the same process but only on a single estimate ˆ x , resulting again in a latent representation of size ( c, H e , W e )

Unfold : From each latent representation we extract all the overlapping patches of size ( c, p, p ) , where p is the height and width of each patch.

Flatten each patch and aggregate : We flatten all the extracted patches to obtain 1-dimensional vectors of size cp 2 , which we assume to come from a MVG distribution. In the training stage we compute their empirical mean and covariance matrix (aggregating over the N dimension), and compute the optimal transport in closed-form T MVG p ˆ x e -→ p x e using eq. (2).

Matmul and unflatten : We apply the pre-computed transport operator using a simple matrix-vector multiplication on the flattened version of the patches extracted from the latent representation ˆ x e . We then reshape each vector to the original patch size ( c, p, p ) (unflatten).

Fold : We rearrange the transported patches back to the original size of the latent representation (reversing the unfold operation). Since the patches overlap, we simply average the shared pixels.

Decoding : To produce our final enhanced estimation ˆ x 0 , we decode the transported latent image back to the pixel space using the decoder of the V AE.

Together, the inference steps form an end-to-end approximation of the desired transport operator T p ˆ x -→ p x . 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 V AE [27] from stable-diffusion [7].

The pre-trained models we evaluate: We apply our latent transport method (described in section 4) on SwinIR [11], Restormer [12] and Swin2SR [13], all of which attempted to minimize average pixel distortion using a supervised regression loss on paired image samples. Additionally, we apply our algorithm on models that are trained to achieve high perceptual quality, and show that we can improve their visual quality even further. As such, we tested two benchmark models in high perceptual quality image restoration: ESRGAN [16], a GAN-based method, and DDRM [3], a diffusion-based method. Beyond our original goal to improve perceptual quality, we demonstrate that we can also traverse the W 2 -MSE tradeoff using any restoration model (e.g., SwinIR, ESRGAN). To do so, we pick any of the aforementioned algorithms and apply our method to improve its perceptual quality, leading to a new estimator. We then interpolate the original algorithm and its improved version using eq. (4), adjusting α 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 ˆ x 0 (instead of the theoretical D max estimator).

Restoration tasks: We showcase our algorithm on Single-Image Super-Resolution (SISR), denoising of Additive White Gaussian Noise (AWGN), JPEG [28] decompression, Noisy Super-Resolution (NSR) and Compressed Super-Resolution (CSR). Training and inference of our algorithm are performed on each restoration model separately, and the evaluation is performed on the restoration task that corresponds to the given model.

Transport operator computation: The transport operator is computed using two disjoint sets of 10 randomly picked images from the ImageNet [29] dataset train split. The first set is used to approximate the predictor's latent statistics ( µ ˆ x e , Σ ˆ x e ): we degrade each image according to the restoration task the predictor is intended to solve, compute the 10 restored outputs, embed the results and compute the embeddings' statistics. The second set is embedded into the latent representation without further modification and serves to approximate the natural image latent statistics ( µ x e , Σ x e ).

Metrics: In addition to Peak Signal-to-Noise Ratio (PSNR), we evaluate distortion performance with Structural Similarity Index Measure (SSIM) [30] and Learned Perceptual Image Patch Similarity (LPIPS) [31], both of which suit better for natural image comparison [31].

Nonetheless, LPIPS remains by definition a distortion (full-reference) measure: it is non-negative and zero when the two images are identical [2]. Interestingly, the original perception-distortion paper [2] already classified the VGG loss [32] - the ancestor of LPIPS - to be a distortion, on which the tradeoff exists (but is less severe).

Therefore, to evaluate perceptual quality, we use the Inception Score (IS) [33], the Fréchet Inception Distance (FID) [25] and the Kernel Inception Distance (KID) [34] following popular image restoration papers [3, 7, 35].

Data sets: It is impractical to perform a serious quantitative evaluation of the perception-distortion tradeoff on real-world datasets (e.g., SIDD, DND, RealSR), which have too few samples to compute FID. Hence, for all models except DDRM [3] and Swin2SR [13], we report the performance on the 50,000 validation samples of ImageNet [29] following [7, 35]. Because of its computational complexity, DDRM [3] reported its performance on a subset of a 1000 ImageNet [29] validation samples. For Swin2SR [13], we use the official DIV2K [36] restored samples provided by the authors.

Although our algorithm can be applied to images with arbitrary aspect ratios, all the tested models were trained on square images. Thus, we resize the samples to 512 × 512 pixels following the pre-processing procedure of DDRM [3]. Finally, we conduct the qualitative evaluation on popular samples from DIV2K or Set14.

5.1 Quantitative results

As reported in table 1, our algorithm can trade perceptual quality for distortion (and vice versa) at test time. We sometimes even manage to improve the pre-trained predictor's PSNR, even of regression models like SwinIR and Swin2SR. When using α ≤ 0 , we systematically improve the predictor's perceptual performance (FID, KID, IS), even for estimators which were designed to achieve photorealism in the first place, e.g., ESRGAN and DDRM. On Non-Local-Means (NLM) [15], an older, non deep-learning denoising algorithm, our method marginally improves all metrics.

While, in theory, our procedure to traverse the perception distortion tradeoff should only include values of α in the range [0 , 1] (see eq. (4)), we also tried to use values outside of this range. As shown in fig. 3, with α ∈ [ -1 , 0] ∪ [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 3.2 only talks about the Wasserstein-2 distance. In practice, the sharpened details that appear in ˆ x 0 and not in x ∗ are either amplified when α ∈ [ -1 , 0] or subtracted (instead of being added) to x ∗ when α ∈ [1 , 2] . We leave the formal analysis of this interesting phenomenon for future research.

Choosing the right value of α : Like any other hyper-parameter, α 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: α does not need to be set before performing some expensive training. Once ˆ x α =0 is computed, any ˆ x α can be obtained thanks to eq. (4) without additional cost. In any case, as reported in table 1, α = 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 α 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 V AE 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 SISR × 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 ≤ p ≤ 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

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 ≤ 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 ˆ 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 V AE, their latent distribution aligns with that of the natural images. Hence, transporting another time does nothing (the transport operator is

Restor

ESRGAN

Swin2SR

σ

AWGN

4

×

SISR

10

q

JPEG

4

×

SR

x

(original)

SwinIR

JPEG

q

10

Restormer

AWGN

σ

50

ESRGAN

SISR

×

4

Swin2SR

SR

×

4

JPEG

q

10

DDRM

SR

×

4

AWGN

σ

25

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.

6 Discussion

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 V AE. I.e., even encoding and decoding a completely clean and natural image does not yield result in perfect

ˆ x

MSE)

0

∗

x

(MMSE)

ˆ x

0

(ours)

ˆ x

0

(ours)

D

(ours)

D

(

E

(

x

))

D

(

E

(

x

))

x

(original)

x

∗

(MMSE)

ˆ x

0

(ours)

D

(

E

(

x

))

x

(original)

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, D 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.

Acknoledgements

This research was partially supported by the Council For Higher Education - Planning and Budgeting Committee.

(

E

(

))

x

x

(

References

D. Freirich, T. Michaeli, and R. Meir, 'A theory of the distortion-perception tradeoff in wasserstein space,' in Advances in Neural Information Processing Systems , 2021.

Y. Blau and T. Michaeli, 'The perception-distortion tradeoff,' in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

B. Kawar, M. Elad, S. Ermon, and J. Song, 'Denoising diffusion restoration models,' in Advances in Neural Information Processing Systems , 2022.

A. Korotin, V. Egiazarian, A. Asadulaev, A. Safin, and E. Burnaev, 'Wasserstein-2 generative networks,' in International Conference on Learning Representations , 2021.

A. Korotin, D. Selikhanovych, and E. Burnaev, 'Neural optimal transport,' in The Eleventh International Conference on Learning Representations , 2023.

A. Makkuva, A. Taghvaei, S. Oh, and J. Lee, 'Optimal transport mapping via input convex neural networks,' in Proceedings of the 37th International Conference on Machine Learning , 2020.

R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, 'High-resolution image synthesis with latent diffusion models,' in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , 2022.

Y. Li, C. Fang, J. Yang, Z. Wang, X. Lu, and M.-H. Yang, 'Universal style transfer via feature transforms,' in Advances in Neural Information Processing Systems , 2017.

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.

Y. Mroueh, 'Wasserstein style transfer,' 2019.

J. Liang, J. Cao, G. Sun, K. Zhang, L. Van Gool, and R. Timofte, 'Swinir: Image restoration using swin transformer,' in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops , 2021.

S. W. Zamir, A. Arora, S. Khan, M. Hayat, F. S. 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.

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.

K. Dabov, A. Foi, V. Katkovnik, and K. Egiazarian, 'Image denoising by sparse 3-d transform-domain collaborative filtering,' IEEE Transactions on Image Processing , 2007.

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'05) , 2005.

X. Wang, K. Yu, S. Wu, J. Gu, Y. Liu, C. Dong, Y. Qiao, and C. Change Loy, 'Esrgan: Enhanced superresolution generative adversarial networks,' in Proceedings of the European Conference on Computer Vision (ECCV) Workshops , 2018.

G. Ohayon, T. Adrai, G. Vaksman, M. Elad, and P. Milanfar, 'High perceptual quality image denoising with a posterior sampling cgan,' in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops , 2021.

B. Kawar, G. Vaksman, and M. Elad, 'Stochastic image denoising by sampling from the posterior distribution,' in Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops , 2021.

B. Kawar, G. Vaksman, and M. Elad, 'SNIPS: Solving noisy inverse problems stochastically,' in Advances in Neural Information Processing Systems , 2021.

M. Aharon, M. Elad, and A. Bruckstein, 'K-svd: An algorithm for designing overcomplete dictionaries for sparse representation,' IEEE Transactions on Signal Processing , 2006.

M. Cuturi, 'Sinkhorn distances: Lightspeed computation of optimal transport,' in Advances in Neural Information Processing Systems , 2013.

G. Peyré and M. Cuturi, 'Computational optimal transport,' Foundations and Trends in Machine Learning , 2019.

M. Arjovsky, S. Chintala, and L. Bottou, 'Wasserstein generative adversarial networks,' in Proceedings of the 34th International Conference on Machine Learning , 2017.

V. Panaretos and Y. Zemel, An Invitation to Statistics in Wasserstein Space . Creative Media Partners, LLC, 2020.

M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, 'Gans trained by a two time-scale update rule converge to a local nash equilibrium,' in Advances in Neural Information Processing Systems , 2017.

C. Villani, Optimal Transport: Old and New . Grundlehren der mathematischen Wissenschaften, Springer Berlin Heidelberg, 2008.

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.

G. Wallace, 'The JPEG still picture compression standard,' IEEE Transactions on Consumer Electronics , 1992.

J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei, 'Imagenet: A large-scale hierarchical image database,' in 2009 IEEE Conference on Computer Vision and Pattern Recognition , 2009.

Z. Wang, A. Bovik, H. Sheikh, and E. Simoncelli, 'Image quality assessment: from error visibility to structural similarity,' IEEE Transactions on Image Processing , 2004.

R. Zhang, P. Isola, A. A. Efros, E. Shechtman, and O. Wang, 'The unreasonable effectiveness of deep features as a perceptual metric,' in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

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.

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.

M. Bi´ nkowski, D. J. Sutherland, M. Arbel, and A. Gretton, 'Demystifying MMD GANs,' in International Conference on Learning Representations , 2018.

C. Saharia, J. Ho, W. Chan, T. Salimans, D. Fleet, and M. Norouzi, 'Image super-resolution via iterative refinement,' in IEEE Transactions on Pattern Analysis and Machine Intelligence , 2021.

E. Agustsson and R. Timofte, 'Ntire 2017 challenge on single image super-resolution: Dataset and study,' in The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops , 2017.

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.' https://github.com/huggingface/diffusers , 2022.

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.

G. Ohayon, T. Adrai, M. Elad, and T. Michaeli, 'Reasons for the superiority of stochastic estimators over deterministic ones: Robustness, consistency and perceptual quality,' 2022.

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.

I. Kligvasser, T. Shaham, Y. Bahat, and T. Michaeli, 'Deep self-dissimilarities as powerful visual fingerprints,' Advances in Neural Information Processing Systems , 2021.

T. Rott Shaham, T. Dekel, and T. Michaeli, 'Singan: Learning a generative model from a single natural image,' in Computer Vision (ICCV), IEEE International Conference on , 2019.

Deep Optimal Transport: A Practical Algorithm for Photo-realistic Image Restoration - Supplementary Material

A Background and extensions

A.1 Numerical Example

y

=

x

+

n

MSE

: 6

.

5

W

2

: 1

.

0

x

∗

=

E

[

x

|

y

]

MSE

: 12

.

4

W

2

: 0

.

5

ˆ x

∼

p

x

|

y

MSE

: 6

.

9

W

2

: 0

.

6

ˆ x

0

0.6

W

2

1.0

6.5

MSE

7.1

ˆ x

0

x

∗

12.4

p

x

|

y

MSE

(

W

2

)

function

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 D 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 Σ x 1 . 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 errors sometimes adds up to negative eigenvalues. 2 In this case, we clamp the negative eigenvalues to zero and use the stochastic (one-to-many) transport operator proposed by [1],

when Σ † x 1 denotes the pseudo-inverse of Σ x 1 (after negative eigenvalues where clamped) and w ∼ N (0 , Σ 1 2 x 2 ( I -Σ 1 2 x 2 T ∗ Σ † x 1 T ∗ Σ 1 2 x 2 ) 1 2 Σ 1 2 x 2 ) , with T ∗ = Σ -1 2 x 2 ( Σ 1 2 x 2 Σ x 1 Σ 1 2 x 2 ) 1 2 Σ -1 2 x 2 .

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 if D ( E ( t )) = t , it is likely that E ( · ) 'deforms' the space, I.e., ∥ E ( s ) -E ( t ) ∥ ̸ = ∥ 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

where | ∂ E ( x ) ∂ x | is the determinant of the Jacobian matrix of E ( · ) 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 D 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 Σ ˆ x e , Σ x e contain ( 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 ≫ 1 . To mitigate the quadratic dependency on H e · 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 · 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 v cp 2 × 1 which we stack into a sample matrix X NH e W e × cp 2 . We then aggregate the samples to compute the MVG statistics: µ = X T 1 , Σ = NH e W e NH e W e -1 ( X -µ )( X -µ ) 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.

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 ≫ 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 T MVG p ˆ x e -→ p 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 Σ ˆ 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 ˆ x 0 ,latent by averaging. The latent image is then decoded back to the pixel space, i.e. ˆ x 0 = D ( ˆ x 0 ,latent ) . Since E ( · ) is not invertible, the decoder D ( · ) 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 y amplifies the degradation (refer to fig. 8).

y

(degraded)

x

∗

(SwinIR)

T

p

y

-→

p

x

(

y

)

ˆ x

0

x

(original)

®

O

®

?

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\nThe 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\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n+ 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\nWEAKNESSES:\n- 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\nQUESTIONS:\nWhat 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?\n\nLIMITATIONS:\nAuthors have addressed the limitations of the method in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper presents an image restoration algorithm that builds upon a trained network to further minimize MSE. 
To 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.

STRENGTHS:
- The idea is straightforward as optimizing delta is common in traditional image restoration
- The visual results are convincing

WEAKNESSES:
- The uncertain of the proposed idea is low so relatively the novelty is less.

QUESTIONS:
By applying the proposed algorithm another time, will the results be further improved?

LIMITATIONS:
None.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nPlease see the weaknesses. I am willing to improve the score if the concerns are addressed well.\n\nLIMITATIONS:\nThe limitations have been described in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

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,
    * In line 116-117, what's the meaning of $x^*$ and $\hat{x_0}$
    * 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?
    * How to define $p_x, p_x^*$?
    * Line 121, OT plan should be between two distributions, not to data points. 

* 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.

QUESTIONS:
Please see the weakness part.

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10