Cross-Domain Lossy Compression as Optimal Transport with an Entropy Bottleneck

We study an extension of lossy compression where the reconstruction distribution is different from the source distribution in order to account for distributional shift due to processing . We formulate this as a generalization of optimal transport with an entropy bottleneck to account for the rate constraint due to compression . We provide expressions for the tradeoff between compression rate and the achievable distortion with and without shared common randomness between the encoder and decoder . demonstrate We study the examples of binary , uniform and Gaussian sources ( in an asymptotic setting ) in detail and demonstrate that shared randomness can strictly improve the tradeoff . For the case without common randomness and squared-Euclidean distortion , we show that the optimal solution partially decouples into the problem of optimal compression and transport and also characterize the penalty associated with fully decoupling them . We provide experimental results by training deep learning end-to-end compression systems for performing denoising on SVHN and super-resolution on MNIST , and demonstrate consistency with our theoretical results . 1 INTRODUCTION . Using deep neural networks for lossy image compression has proven to be effective , with ratedistortion performance capable of dominating general-purpose image codecs like JPEG , WebP or BPG ( Rippel & Bourdev , 2017 ; Agustsson et al. , 2017 ; Mentzer et al. , 2018 ) . More recently , many of these works have included generative aspects within the compression to synthesize realistic elements when the rate is otherwise too low to represent fine-grained details ( Tschannen et al. , 2018 ; Agustsson et al. , 2019 ; Mentzer et al. , 2020 ) . Though this has been found to deteriorate ratedistortion performance , it has generally resulted in more perceptually-pleasing image reconstruction by reducing artifacts such as pixelation and blur . Using a distributional constraint as a proxy for perceptual measure , several works have subsequently formalized this in a mathematical framework known as the rate-distortion-perception tradeoff ( Blau & Michaeli , 2018 ; 2019 ; Matsumoto , 2018 ; 2019 ; Theis & Wagner , 2021 ; Yan et al. , 2021 ; Zhang et al. , 2021 ) . As is conventional in lossy compression , these works address the scenario in which both low distortion , whereby each individual image reconstruction resembles the ground truth image , and closeness in distribution in which it is not easy to discriminate between image samples from the data-generating distribution and reconstruction distribution , are desirable . The underlying ideal in conventional compression systems is to have perfect reconstruction with respect to some ground truth input . However this is not the case in applications such as denoising , deblurring , or super-resolution ( SR ) , which require restoration from a degraded input image . In fact , in these cases a ground truth may not even be available . In such applications naturally the reconstruction distribution must match the original source rather than the degraded input distribution . A large body of literature has been devoted to various image restoration tasks , including several methods based on deep learning including both supervised ( e.g. , ( Blau & Michaeli , 2018 ) ) and unsupervised learning methods ( e.g. , ( Wang et al. , 2021b ) ) . Although most of the literature exclusively treats compression and restoration separately , in many application they can co-occur . For example , the encoder which records a degraded image may not be co-located with the decoder , but must transmit a compressed version of the image over a digital network . In turn , the decoder must perform both decompression and restoration simultaneously . To that end , we study an extension of lossy compression in which the reconstruction distribution is different than the source distribution to account for distributional shift due to processing . The problem can be described as a transformation from some source domain to a new target domain under a rate constraint , which generalizes optimal transport . This readily extends other works which view image restoration under the perception-distortion tradeoff ( Blau & Michaeli , 2018 ) or under optimal transport ( Wang et al. , 2021b ) . It also provides a generalization of the rate-distortion-perception problem ( Blau & Michaeli , 2019 ) where the reconstruction distribution must be close to the input distribution . Following ( Theis & Agustsson , 2021 ; Theis & Wagner , 2021 ) , we also utilize common randomness as a tool for compression in our setting . Our results are summarized as follows : • We provide a formulation of lossy compression with distribution shift as a generalization of optimal transport with an entropy constraint and identify the tradeoff between the compression rate and minimum achievable distortion both with and without common randomness at the encoder and decoder . We identify conditions under which the structure of the optimal solution partially decouples the problems of compression and transport , and discuss their architectural implications . We study the examples of binary , uniform and Gaussian sources ( in asymptotic regime ) in detail and demonstrate the utility of our theoretical bounds . • We train deep learning end-to-end compression systems for performing super-resolution on MNIST and denoising on SVHN . Our setup is unsupervised and to the best of our knowledge the first to integrate both compression and restoration at once using deep learning . We first demonstrate that by having common randomness at the encoder and decoder the achievable distortion-rate tradeoffs are lower than when such randomness is not present . Furthermore , we provide experimental validation of the architectural principle suggested by our theoretical analysis . 2 THEORETICAL FORMULATION We consider a setting where an input X ∼ pX is observed at the encoder , which is a degraded ( e.g. , noisy , lower resolution , etc ) version of the original source . It must be restored to an output Y ∼ pY at the decoder , where pY denotes the target distribution of interest . For example , if X denotes a noise-corrupted image and Y denotes the associated clear reconstruction , then pY can be selected to match the distribution of the original source . We will assume pX and pY are probability distributions over X , Y ⊆ Rn and require X and Y to be close with respect to some fidelity metric , which will be measured using a non-negative cost function d ( x , y ) over X × Y . We will refer to d ( · , · ) as the distortion measure and assume that it satisfies d ( x , y ) = 0 if and only if x = y . We further assume that X can not be directly revealed to the decoder , but instead must be transmitted over a bit interface with an average rate constraint . Such a scenario occurs naturally in many practical systems when the encoder and decoder are not co-located such as communication systems or storage systems . As one potential application , when aerial photographs are produced for remote sensing purposes , blurs are introduced by atmospheric turbulence , aberrations in the optical system and relative motion between camera and ground . In such scenarios unsupervised restoration is preferred as it is often intractable to accurately model such degradation processes and collection of paired training data can be time consuming or require significant human intervention . Unsupervised image restoration has been studied recently in Zhang et al . ( 2017 ) ; Pan et al . ( 2021 ) ; Wang et al . ( 2021b ) ; Menon et al . ( 2020 ) . These works also fix the reconstruction distribution Y ∼ pY and propose to minimize a distortion metric between the output and the degraded input as in our present work , but do not consider compression . 2.1 OPTIMAL TRANSPORT AND EXTENSIONS . Definition 1 ( Optimal Transport ) . Let Γ ( pX , pY ) be the set of all joint distributions pX , Y with marginals pX and pY . The classical optimal transport problem is defined as D ( pX , pY ) = inf pX , Y ∈Γ ( pX , pY ) E [ d ( X , Y ) ] , ( 1 ) where we refer to each pX , Y ∈ Γ ( pX , pY ) as a transport plan . Operationally the optimal transport plan in ( 1 ) minimizes the average distortion between the input and output while keeping the output distribution fixed to pY . This may generate a transport plan with potentially unbounded entropy , which may not be amenable in a rate-constrained setting . We therefore suggest a generalization to Definition 1 which constrains the entropy of the transport plan . It turns out that having common randomness at the encoder and decoder can help in this setting , so we will distinguish between when it is available and unavailable . Definition 2 ( Optimal Transport with Entropy Bottleneck — no Common Randomness ) . Let Mncr ( pX , pY ) denote the set of joint distributions pX , Z , Y compatible with the given marginal distributions pX , pY satisfying pX , Z , Y = pXpZ|XpY |Z . The optimal transport from pX to pY with an entropy bottleneck of R and without common randomness is defined as Dncr ( pX , pY , R ) , inf pX , Z , Y ∈Mncr ( pX , pY ) E [ d ( X , Y ) ] s.t . H ( Z ) ≤ R , ( 2 ) where H ( · ) denotes the Shannon entropy of a random variable . We note that when the rate constraintR is sufficiently large such that one can select Z = X or Z = Y in ( 2 ) , then Dncr ( pX , pY , R ) = D ( pX , pY ) in ( 1 ) . More generally , D ( pX , pY ) serves as a lower bound for Dncr ( pX , pY , R ) for any R > 0 . Definition 2 also has a natural operational interpretation in our setting . We can view the encoder as implementing the conditional distribution pZ|X to output a representation Z given the input X , and the decoder as implementing the conditional distribution pY |Z to output the reconstruction Y given the representation Z . The entropy constraint H ( Z ) ≤ R essentially guarantees that the representation Z can be losslessly transmitted at a rate close to R1 . It turns out that when we specialize to the squared Euclidean distance , we can without loss of optimality impose a more structured architecture for implementing the encoder and the decoder . Let W 22 ( · , · ) be the squared quadratic Wasserstein distance , by setting d ( X , Y ) = ||X − Y ||2 in Definition 1 . Theorem 1 . Let Dmse ( pX , pY , R ) , inf pX̂|X , pŶ |Y E [ ‖X − X̂‖2 ] + E [ ‖Y − Ŷ ‖2 ] +W 22 ( pX̂ , pŶ ) s.t . E [ X|X̂ ] = X̂ , E [ Y |Ŷ ] = Ŷ , H ( X̂ ) ≤ R , H ( Ŷ ) ≤ R , ( 3 ) and Dmse ( pX , R ) , inf pX̂|X E [ ‖X − X̂‖2 ] s.t . H ( X̂ ) ≤ R. ( 4 ) Moreover , let . Dncr ( pX , pY , R ) , Dmse ( pX , R ) +Dmse ( pY , R ) +W 2 2 ( pX̂∗ , pŶ ∗ ) , ( 5 ) Dncr ( pX , pY , R ) , Dmse ( pX , R ) +Dmse ( pY , R ) , ( 6 ) where pX̂∗ and pŶ ∗ are the marginal distributions induced by the minimizers pX̂∗|X and pŶ ∗|Y that attain Dmse ( pX , R ) and Dmse ( pY , R ) , respectively ( assuming the existence of such minimizers ) . Then under the squared Eucledian distortion measure , Dncr ( pX , pY , R ) = Dmse ( pX , pY , R ) . ( 7 ) In addition , we have Dncr ( pX , pY , R ) ≥ Dncr ( pX , pY , R ) ≥ Dncr ( pX , pY , R ) , ( 8 ) and both inequalities are tight when pX = pY . 1The source coding theorem guarantees that any discrete random variable Z can be losslessly compressed using a variable length code with average length of no more than H ( Z ) + 1 bits . Theorem 1 deconstructs Z into the quantizations X̂ of X and Ŷ of Y , and decomposes the overall distortion in ( 2 ) in terms of the losses due to quantization , transport , and dequantization in ( 3 ) . It also suggests a natural architecture that partially decouples compression and transport without loss of optimality . First , the sender uses the distribution pX̂|X to produce the compressed representation X̂ from X . This is then passed through a “ converter ” pŶ |X̂ to transform X̂ to an optimal representation Ŷ of Y . Finally , the receiver maps Ŷ back to Y using the conditional distribution pY |Ŷ . This is illustrated in Figure 1 . The entropy constraint H ( X̂ ) ≤ R in ( 2 ) essentially guarantees that X̂ can be losslessly transmitted to the decoder where the converter can be applied to map X̂ to Ŷ before outputting Y . Alternately the constraint H ( Ŷ ) ≤ R guarantees that the converter could also be implemented at the encoder and then Ŷ can be compressed and transmitted to the decoder . Finally note that our proposed architecture is symmetric2 with respect to the encoder and the decoder and in particular the procedure to transport Y to X would simply be the inverse of transporting X to Y , and indeed the distortion incurred by dequantizing pY |Ŷ is the same as the distortion incurred by quantizing pŶ |Y . For the special case of same source and target distribution , we have Dmse ( pX , pX , R ) = 2Dmse ( pX , R ) , implying that the rate required to achieve distortion D under no output distribution constraint ( and with the output alphabet relaxed to Rn ) achieves distortion 2D under the constraint that Y equals X in distribution . This recovers the result of Theorem 2 in Yan et al . ( 2021 ) for the one-shot setting . More generally , ( 8 ) shows that we may lower bound Dmse ( pX , pY , R ) by the distortion incurred when compressing X and Y individually , each at rate R , through ignoring the cost of transport . On the other hand , the upper bound corresponds to choosing the optimal rate-distortion representations X̂∗ , Ŷ ∗ for X , Y , then considering transport between them . The advantage of this approach is that knowledge of the other respective distribution is not necessary for design . Although not optimal in general , we will , in fact , provide an example where this is optimal in Section 2.2 . Finally , the following result implies that under mild regularity conditions , the optimal converter pŶ |X̂ can be realized as a ( deterministic ) bijection , and in the scalar case it can basically only take the form as illustrated in Figure 1 . Theorem 2 . Assume that Dncr ( pX , pY , R ) is a strictly decreasing function in a neighborhood of R = R∗ and Dncr ( pX , pY , R∗ ) is attained by pX , Z , Y . Let X̂ , E [ X|Z ] and Ŷ , E [ Y |Z ] . Then H ( X̂ ) = H ( Ŷ ) = R∗ , ( 9 ) E [ ‖X̂ − Ŷ ‖2 ] = W 22 ( pX̂ , pŶ ) , ( 10 ) and there is a bijection between X̂ and Ŷ . We remark that in general computing the optimal transport map is not straightforward . For the case of binary sources we can compute an exact characterization for Dncr as discussed in Section 2.2 . Furthermore as discussed in Appendix A.6 , W 22 ( pX̂ , pŶ ) can be computed in closed form when X̂ and Ŷ are scalar valued , which can be used to obtain upper bounds on Dncr . In our experimental results in Section 3 we use deep learning based methods to learn approximately optimal mappings . So far we have focused on the setting when there is no shared common randomness between the encoder and the decoder . We will now consider the setting when a shared random variable denoted by U is present at the encoder and decoder . We assume that the variable U is independent of the input X so that the decoder has no apriori information of the input . In practice the sender and receiver can agree on a pseudo-random number generator ahead of time and some kind of seed could be transmitted , after which both sides can generate the same U . We further discuss how shared randomness is used in practice in the experimental section . Definition 3 ( Optimal Transport with Entropy Bottleneck — with Common Randomness ) . Let Mcr ( pX , pY ) denote the set of joint distributions pU , X , Z , Y compatible with the given marginal 2We say that the problem is symmetric if it is invariant under reversing pX , pY with a new distortion measure defined by reversing the arguments of d ( · , · ) . distributions pX , pY and satisfying pU , X , Z , Y = pUpXpZ|X , UpY |Z , U , where pU represents the distribution of shared randomness . The optimal transport from pX to pY with entropy bottleneck R and common randomness is defined as Dcr ( pX , pY , R ) , inf pU , X , Z , Y ∈Mcr ( pX , pY ) E [ d ( X , Y ) ] s.t . H ( Z|U ) ≤ R. ( 11 ) Note that we optimize over pU ( the distribution associated with shared randomness ) , in addition to pZ|X , U and pY |Z , U in ( 11 ) . Furthermore , Dcr ( pX , pY , R ) ≤ Dncr ( pX , pY , R ) in general , as we do not have access to shared randomness in Definition 2 . Also from the same argument that was made following Definition 2 , we have that Dcr ( pX , pY , R ) ≥ D ( pX , pY ) in Definition 1 . As with Definition 2 , we can also provide a natural operational interpretation . In particular , given the input X and common randomness U the encoder can output a compressed representation Z using the conditional distribution pZ|X , U . The representation Z can be losslessly compressed approximately to an average rate of R again by exploiting the shared randomness U . Finally the decoder , given Z and U can output the reconstruction Y using the conditional distribution pY |Z , U . An interesting difference with Definition 2 is that the setup is no longer symmetric between encoder and decoder , as X is independent of U but Y is not . The following result provides a simplification to the architecture in Definition 3 . Theorem 3 . Let Qcr ( pX , pY ) denote the set of joint distributions pU , X , Y compatible with the given marginals pX , pY satisfying pU , X , Y = pUpXpY |U , X as well as H ( Y |U , X ) = 0 . Then Dcr ( pX , pY , R ) = inf pU , X , Y ∈Qcr ( pX , pY ) E [ d ( X , Y ) ] s.t . H ( Y |U ) ≤ R. ( 12 ) Before discussing the implications of Theorem 3 we remark on a technical point . Because the Shannon entropy is defined only for discrete random variables , U must be chosen in a way such that Y |U = u is discrete for each u , even for continuous ( X , Y ) . This is known to be possible , e.g. , Li & El Gamal ( 2018 ) have provided a general construction for a U with this property , with additional quantitative guarantees to ensure that U is informative of Y . In the finite alphabet case we show in Appendix A.3 that optimization of U can be formulated as a linear program . We next discuss the implication of Theorem 3 . First note that the problem can be modelled with only pY |U , X producing a reconstruction Y without the need for the intermediate representation Z , much like the conventional optimal transport in Definition 1 . The condition H ( Y |U , X ) = 0 also implies that the transport plan is deterministic when conditioned on the shared randomness , which plays the role of stochasticity . Furthermore in this architecture the encoder should compute the representation Y given the sourceX and the shared random-variable U ( which corresponds to the transport problem ) and then compress it losslessly at a rate close to H ( Y |U ) ( which corresponds to the compression problem ) . The receiver only needs to decompress and reconstruct Y . This is in contrast to the case without common randomness in Theorem 1 where the reconstruction Y must be generated at the decoder .

Summary: The paper explores the problem of learned compression scheme for the specific setup where the source and target distributions are not the same (such as in restoring a degraded image). The authors formulate the problem as an optimal transport problem with entropy constraints to control the compression rate. They further argue that the use of common randomness in the encoder and decoder improves the rate-distortion tradeoff and dcoument it on denoising and super-resolution experiments.