A Theory of the Distortion-Perception Tradeoff in Wasserstein Space

1 Introduction . Inverse problems that involve signal reconstruction from partial or noisy measurements , arise in numerous scientific domains . Examples range from medical imaging to tomography , microscopy , astronomy and audio enhancement . In many such problems it is desired to design an estimator that ( i ) has a small reconstruction error ( low distortion ) , and ( ii ) outputs reconstructions that can not be told apart from valid signals ( good perceptual quality ) . Interestingly , however , it has been shown that the lower the average distortion of an estimator , the more the distribution of its outputs generally deviates from the distribution of the signals it attempts to estimate [ 4 ] . In other words , low distortion generally comes at the price of poor perceptual quality , and vice versa . This phenomenon , known as the perception-distortion tradeoff , has found particular interest in the image restoration domain ( see Fig . 1 ) , where algorithms are now commonly being evaluated using both distortion measures and perception indices [ 5 ] . Unfortunately , despite the increasing popularity of performing comparisons on the perceptiondistortion plane , the minimal distortion that can be achieved under a given perception constraint ( red curve in Fig . 1 ) remains an open question . Blau and Michaeli [ 4 ] investigated several properties 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of this distortion-perception function , such as monotonicity and convexity . But beyond this rather general characterization , little is known about its precise nature . In this paper , we derive a closed form expression for the distortion-perception ( DP ) function for the special case where distortion is measured by mean squared-error ( MSE ) and perception is measured by the Wasserstein-2 distance between the probability laws of the estimate and the estimand . Our main contributions are : ( i ) We prove that the DP function is always quadratic in the perception constraint P , regardless of the underlying distribution ( Theorem 1 ) . ( ii ) We show that it is possible to construct estimators on the DP curve from the estimators at the two extremes of the tradeoff ( Theorem 3 ) : The one that globally minimizes the MSE , and a minimizer of the MSE under a perfect perceptual quality constraint . The latter can be obtained as a stochastic transformation of the former . ( iii ) In the Gaussian setting , we further provide a closed form expression for optimal estimators and for the corresponding DP curve ( Theorems 4 and 5 ) . We show this Gaussian DP curve is a lower bound on the DP curve of any distribution having the same second order statistics . Finally , we illustrate our results , numerically and visually , in a super-resolution setting in Section 5 . The proofs of all our theorems are provided in Appendix B . Our theoretical results shed light on several topics that are subject to much practical activity . Particularly , many recent works adress the task of diverse perceptual image reconstruction , by employing randomization among possible restorations [ 15 , 3 , 22 , 1 ] . Commonly , such works attempt to sample from the posterior distribution of natural images given the degraded input image . This is done , for example , using priors over image patches [ 7 ] , conditional generative models [ 18 , 21 ] , or implicit priors induced by deep denoiser networks [ 11 ] . Theoretically , posterior sampling leads to perfect perceptual quality ( the restored outputs are distributed like the prior ) . However , a fundamental question is whether this is optimal in terms of distortion . As we show in Section 3.1 , posterior sampling is often not an optimal strategy , in the sense that there exist perfect perceptual quality estimators that achieve lower distortion . Another topic of practical interest is the ability to traverse the distortion-perception tradeoff at test time , without having to train a different model for each working point . Recently , interpolation between distortion-oriented models and perception-oriented ones , has been suggested for this end . Existing methods include interpolation in pixel space [ 31 ] or in some latent space [ 26 ] , interpolation between network weights [ 31 , 32 ] , and style transfer between a low-distortion reconstruction and a high perceptual quality one [ 6 ] . In light of this plethora of approaches , it is natural to ask which strategy is optimal . In Section 3.2 we show that for the MSE–Wasserstein-2 tradeoff , linear interpolation in pixel space leads to optimal estimators . We also discuss a geometric connection between interpolation and the fact that estimators on the DP curve form a geodesic in Wasserstein space . 2 Problem setting and preliminaries . 2.1 The distortion-perception tradeoff . Let X , Y be random vectors taking values in Rnx and Rny , respectively . We consider the problem of constructing an estimator X̂ ofX based on Y . Namely , we are interested in determining a conditional distribution pX̂|Y such that X̂ constitutes a good estimate of X . For example , in the super-resolution setting shown in Fig . 1 , Y is the low resolution image ( left ) , X is the corresponding ground-truth high-resolution image ( not shown ) , and X̂ is a super-resolution reconstruction generated from Y ( e.g . the EDSR or ESRGAN estimators in the middle ) . In many practical cases , the goodness of an estimator is associated with two factors : ( i ) the degree to which X̂ is close to X on average ( low distortion ) , and ( ii ) the degree to which the distribution of X̂ is close to that of X ( good perceptual quality ) . An important question , then , is what is the minimal distortion that can be achieved under a given level of perceptual quality ? and how can we construct estimators that achieve this lower bound ? In mathematical language , we are interested in analyzing the distortion-perception ( DP ) function ( defined similarly to the perception-distortion function of [ 4 ] ) D ( P ) = min pX̂|Y { E [ d ( X , X̂ ) ] : dp ( pX , pX̂ ) ≤ P } . ( 1 ) Here , d : Rnx × Rnx → R+ ∪ { 0 } is some distortion criterion , dp ( · , · ) is some divergence between probability measures , and pX̂ is the probability measure on R nx induced by pX̂|Y and pY . The expectation is taken w.r.t . the measure pXX̂ induced by pX̂|Y and pXY , where we assume that X̂ is independent of X given Y . As discussed in [ 4 ] , the function D ( P ) is monotonically non-increasing and is convex whenever dp ( · , · ) is convex in its second argument ( which is the case for most popular divergences ) . However , without further concrete assumptions on the distortion measure d ( · , · ) and the perception index dp ( · , · ) , little can be said about the precise nature of D ( P ) . Here , we focus our attention on the squared-error distortion d ( x , x̂ ) = ‖x−x̂‖2 and the Wasserstein-2 distance dp ( pX , pX̂ ) = W2 ( pX , pX̂ ) , with which ( 1 ) reads D ( P ) = min pX̂|Y { E [ ‖X − X̂‖2 ] : W2 ( pX , pX̂ ) ≤ P } . ( 2 ) We assume that all distributions have finite first and second moments . In addition , from Theorem 3 below it will follow that the minimum is indeed attained , so that ( 2 ) is well defined . It is well known that the estimator minimizing the mean squared error ( MSE ) without any constraints , is given by X∗ = E [ X|Y ] . This implies that D ( P ) monotonically decreases until P reaches P ∗ , W2 ( pX , pX∗ ) , beyond which point D ( P ) takes the constant value D∗ , E [ ‖X −X∗‖2 ] . This is illustrated in Fig . 2 . It is also known that D ( 0 ) ≤ 2D∗ since the posterior sampling estimator pX̂|Y = pX|Y achieves W2 ( pX , pX̂ ) = 0 and E [ ‖X − X̂‖2 ] = 2D∗ [ 4 ] . However , apart from these rather general properties , the precise shape of the DP curve has not been determined to date , and neither have the estimators that achieve the optimum in ( 2 ) . This is our goal in this paper . 2.2 The Wasserstein and Gelbrich Distances . Before we present our main results , we briefly survey a few properties of the Wasserstein distance , mostly taken from [ 20 ] . The Wasserstein-p ( p ≥ 1 ) distance between measures µ and γ on a separable Banach space X with norm ‖ · ‖ is defined by W pp ( µ , γ ) , inf { E ( U , V ) ∼ν [ ‖U − V ‖p ] : ν ∈ Π ( µ , γ ) } , ( 3 ) where Π ( µ , γ ) is the set of all probabilities on X × X with marginals µ and γ . A joint probability ν achieving the optimum in ( 3 ) is often referred to as optimal plan . The Wasserstein space of probability measures is defined as Wp ( X ) , { γ : ∫ X ‖x‖pdγ < ∞ } , and Wp constitutes a metric onWp ( X ) . For any ( m1 , Σ1 ) , ( m2 , Σ2 ) ∈ Rd × Sd+ ( where Sd+ is the set of symmetric positive semidefinite matrices in Rd×d ) , the Gelbrich distance is defined as G2 ( ( m1 , Σ1 ) , ( m2 , Σ2 ) ) , ‖m1 −m2‖22 + Tr { Σ1 + Σ2 − 2 ( Σ 1 2 1 Σ2Σ 1 2 1 ) 1 2 } . ( 4 ) The root of a PSD matrix is always taken to be PSD . For any two probability measures µ1 , µ2 on Rd with means and covariances ( m1 , Σ1 ) , ( m2 , Σ2 ) , from [ 8 , Thm . 2.1 ] we have that W 22 ( µ1 , µ2 ) ≥ G2 ( ( m1 , Σ1 ) , ( m2 , Σ2 ) ) . ( 5 ) When µ1 = N ( m1 , Σ1 ) and µ2 = N ( m2 , Σ2 ) are Gaussian distributions on Rd , we have that W2 ( µ1 , µ2 ) = G ( ( m1 , Σ1 ) , ( m2 , Σ2 ) ) . This equality is obvious for non-singular measures but is true for any two Gaussian distributions [ 20 , p. 18 ] . If Σ1 and Σ2 are non-singular , then the distribution attaining the optimum in ( 3 ) corresponds to U ∼ N ( m1 , Σ1 ) , V = m2 + T1→2 ( U −m1 ) , ( 6 ) where T1→2 = Σ − 12 1 ( Σ 1 2 1 Σ2Σ 1 2 1 ) 1 2 Σ − 12 1 ( 7 ) is the optimal transformation pushing forward from N ( 0 , Σ1 ) to N ( 0 , Σ2 ) [ 12 ] . This transformation satisfies Σ2 = T1→2Σ1T1→2 . For a discussion on singular distributions , please see App . A .

Summary: The paper studies the trade-off between minimizing the mean-square error (distortion) and the Wasserstein-2 distance (perception) when constructing an estimator \hat{X} of a random variable X. In particular, the work analyzes a distortion-perception function, which given the radius P of a Wasserstein-2 ball around X returns the minimal mean-square error of \hat{X} constrained to be in that ball.   The analysis reveals that for estimators with small Wasserstein-2 distance to the true distribution, one can often obtain an estimator with much better mean-square error by a simple linear interpolation in the output space.   The theoretical results are applied to deep image super-resolution, where it is shown that the perceptual quality of an estimator with small MSE can be improved by interpolating it with one having small Wasserstein-2 distance. Remarkably, the interpolated estimator often still has low MSE, while significantly lower Wasserstein-2 distance, leading to a greatly improved perceptual quality.