Towards Empirical Sandwich Bounds on the Rate-Distortion Function

1 INTRODUCTION . From storing astronomical images captured by the Hubble telescope , to delivering familiar faces and voices over video chats , data compression , i.e. , communicating the “ same ” information but with less bits , is commonplace and indispensable to our digital life , and even arguably lies at the heart of intelligence ( Mahoney , 2009 ) . While for lossless compression , there exist practical algorithms that can compress any discrete data arbitrarily close to the information theory limit ( Ziv & Lempel , 1977 ; Witten et al. , 1987 ) , no such universal algorithm has been found for lossy data compression ( Berger & Gibson , 1998 ) , and significant research efforts have dedicated to lossy compression algorithms for various data . Recently , deep learning has shown promise for learning lossy compressors from raw data examples , with continually improving compression performance often matching or exceeding traditionally engineered methods ( Minnen et al. , 2018 ; Agustsson et al. , 2020 ; Yang et al. , 2020a ) . However , there are fundamental limits to the performance of any lossy compression algorithm , due to the inevitable trade-off between rate , the average number of bits needed to represent the data , and the distortion incurred by lossy representations . This trade-off is formally described by the rate-distortion ( R-D ) function , for a given source ( i.e. , the data distribution of interest ; referred to as such in information theory ) and distortion metric . The R-D function characterizes the best theoretically achievable rate-distortion performance by any compression algorithm , which can be seen as a lossy-compression counterpart and generalization of Shannon entropy for lossless compression . Despite its fundamental importance , the R-D function is generally unknown analytically , and establishing it for general data sources , especially real world data , is a difficult problem ( Gibson , 2017 ) . The default method for computing R-D functions , the Blahut-Arimoto algorithm ( Blahut , 1972 ; Arimoto , 1972 ) , only works for discrete data with a known probability mass function and has a complexity exponential in the data dimensionality . Applying it to an unknown data source requires discretization ( if it is continuous ) and estimating the source probabilities by a histogram , both of which introduce errors and are computationally infeasible beyond a couple of dimensions . Previous work characterizing the R-D function of images and videos ( Hayes et al. , 1970 ; Gibson , 2017 ) all assumed a statistical model of the source , making the results dependent on the modeling assumptions . In this work , we make progress on this old problem in information theory using tools from machine learning , and introduce new algorithms for upper and lower bounding the R-D function of a general ( i.e. , discrete , continuous , or neither ) , unknown memoryless source . More specifically , 1 . Our upper bound draws from the deep generative modeling toolbox , and is closely related to a class of β-VAEs in learned data compression ( Ballé et al. , 2017 ) ; we clarify how these models optimize a model-independent upper bound on the source rate-distortion function . 2 . We theoretically derive a general lower bound on the R-D function that can in principle be optimized by gradient ascent . Facing the difficulty of the problem involving global optimization , we restrict to a squared error distortion and obtain an approximate algorithm . 3 . We experimentally show that our upper bound can recover the R-D function of Gaussian sources , and obtain non-trivial sandwich bounds on high-dimension data with low intrinsic dimension . In particular , our sandwich bounds on 128×128 GAN-generated images shed light on the effectiveness of learned compression approaches ( Ballé et al. , 2021 ) . 4 . Although we currently can not verify its tightness , our estimated R-D upper bound on natural images indicates possible room for improving the image compression performance of state-of-the-art methods by one dB in PSNR , at various bitrates . We begin by reviewing the prerequisite rate-distortion theory in Section 2 , then describe our upper and lower bound algorithms in Section 3 and Section 4 , respectively . We discuss related work in Section 5 , report experimental results in Section 6 , and conclude in Section 7 . 2 BACKGROUND . Rate-distortion ( R-D ) theory deals with the fundamental trade-off between the average number of bits per sample ( rate ) used to represent a data source X and the distortion incurred by the lossy representation X̂ . It asks the following question about the limit of lossy compression : for a given data source and a distortion metric ( a.k.a. , a fidelity criterion ) , what is the minimum number of bits ( per sample ) needed to represent the source at a tolerable level of distortion , regardless of the computation complexity of the compression procedure ? The answer is given by the rate-distortion function R ( D ) . To introduce it , let the source and its reproduction take values in the sets X and X̂ , conventionally called the source and reproduction alphabets , respectively . We define the data source formally by a random variable X ∈ X following a ( usually unknown ) distribution PX , and assume a distortion metric ρ : X × X̂ → [ 0 , ∞ ) has been given , such as the squared error ρ ( x , x̂ ) = ‖x − x̂‖2 . The rate-distortion function is then defined by the following constrained optimization problem , R ( D ) = inf QX̂|X : E [ ρ ( X , X̂ ) ] ≤D I ( X ; X̂ ) , ( 1 ) where we consider all random transforms QX̂|X whose expected distortion is within the given threshold D ≥ 0 , and minimize the mutual information between the source X and its reproduced representation X̂ 1 . Shannon ’ s lossy source coding theorems ( Shannon , 1948 ; 1959 ) gave operational significance to the above mathematical definition of R ( D ) , as the minimum achievable rate with which any lossy compression algorithm can code i.i.d . data samples at a distortion level within D. The R-D function thus gives the tightest lower bound on the rate-distortion performance of any lossy compression algorithm , and can inform the design and analysis of such algorithms . If the operational distortion-rate performance of an algorithm lies high above the source R ( D ) -curve ( D , R ( D ) ) , then further performance improvement may be expected ; otherwise , its performance is already close to theoretically optimal , and we may focus our attention on other aspects of the algorithm . As the R-D function does not have an analytical form in general , we propose to estimate it from data samples , making the standard assumption that various expectations w.r.t . the true data distribution PX exist and can be approximated by sample averages . For a discrete source , this assumption automatically holds , and the R-D function also provides a lower bound on the Shannon entropy of discrete data . 3 UPPER BOUND ALGORITHM . For a known discrete source , the Blahut-Arimoto ( BA ) algorithm converges to R ( D ) from above by fixed-point equations . For a general unknown source , it is not clear how this can be done . We 1Both the expected distortion and mutual information terms are defined w.r.t . the joint distribution PXQX̂|X . We formally describe the general setting of the paper , including the technical definitions , in Appendix A.1 . propose to solve the underlying variaitonal problem approximately by gradient descent . In exchange for generality and scalability , we lose the optimality guarantee of the BA algorithm , only arriving at a stochastic upper bound of R ( D ) in general . By R-D theory ( Cover & Thomas , 2006 ) , every ( distortion , rate ) pair lying above the R ( D ) -curve is in principle realizable by a ( possibly expensive ) compression algorithm ; an upper bound on R ( D ) , therefore , reveals what kind of ( improved ) R-D performance is theoretically possible , without suggesting how it can be practically achieved . Variational Formulation . We adopt the same unconstrained variational objective as the BlahutArimoto ( BA ) algorithm ( Blahut , 1972 ; Arimoto , 1972 ) , in its most general form , L ( QX̂|X , QX̂ , λ ) : = Ex∼PX [ KL ( QX̂|X=x‖QX̂ ) ] + λEPXQX̂|X [ ρ ( X , X̂ ) ] , ( 2 ) where QX̂ is an arbitrary probability measure on X̂ and KL ( ·‖· ) denotes the Kullback-Leibler ( KL ) divergence . This objective can be seen as a Lagrangian relaxation of the constrained problem defining R ( D ) , where the first ( rate ) term is a variational upper bound on the mutual information I ( X ; X̂ ) , and the second ( distortion ) term enforces the distortion tolerance constraint in Eq . 1 . For each fixed λ > 0 , a global minimizer of L yields a point ( R , D ) on the R ( D ) curve ( Csiszar , 1974 ) , where R and D are simply the two terms of L evaluated at the optimal ( QX̂|X , QX̂ ) . Repeating this minimization for various λ then traces out the R ( D ) curve . Based on this connection , the BA algorithm carries out the minimization by coordinate descent on L via fixed-point equations , each time setting QX̂|X to be optimal w.r.t . QX̂ and vice versa ; the sequence of alternating distributions can be shown to converge and yield a point on the R ( D ) curve ( Csiszar , 1974 ) . Unfortunately , the BA algorithm only applies when X and X̂ are finite , and the source distribution PX known ( or estimated from data samples ) in the form of a vector of probabilities PX ( x ) of every state x ∈ X . Moreover , the algorithm requires storage and running time exponential in the data dimensionality , since it operates on exhaustive tabular representations of PX , ρ , QX̂|X , and QX̂ . The algorithm therefore quickly becomes infeasible on data with more than a couple of dimensions , not to mention high-dimension data such as natural images . The fixed-point equations of the BA algorithm are known in general settings ( Rezaei et al. , 2006 ) , but when X or X̂ is infinite ( such as in the continuous case ) , it is not clear how to carry them out or exhaustively represent the measures QX̂ and QX̂|X=x for each x ∈ X . Proposed Method . To avoid these difficulties , we propose to apply ( stochastic ) gradient descent on L w.r.t . flexibly parameterized variational distributions QX̂|X and QX̂ . The distributions can be members from any variational family , as long as QX̂|X=x is absolutely continuous w.r.t . QX̂ for ( PX -almost ) all x ∈ X . This technical condition ensures their KL divergence is well defined . This is easily satisfied , when X̂ is discrete , by requiring the support of QX̂|X=x to be contained in that of QX̂ for all x ∈ X ; and in the continuous case , by representing both measures in terms of probability density functions ( e.g. , normalizing flows ( Kobyzev et al. , 2021 ) ) . In this work we represent QX̂ and QX̂|X by parametric distributions , and predict the parameters of each QX̂|X=x by an encoder neural network φ ( x ) as in amortized inference ( Kingma & Welling , 2014 ) . Given i.i.d . X samples , we optimize the parameters of the variational distributions by SGD on L ; at convergence , the estimates of rate and distortion terms of L yields a point that in expectation lies on an R-D upper bound RU ( D ) . The variational objective L ( Eq . 2 ) closely resembles the negative ELBO ( NELBO ) objective of a β-VAE ( Higgins et al. , 2017 ) , if we regard the reproduction alphabet X̂ as the “ latent space ” . The connection is immediate , when X̂ is continuous and a squared error ρ specifies the density of a Gaussian likelihood p ( x|x̂ ) ∝ exp ( −‖x− x̂|2 ) . However , unlike in data compression , where X̂ is determined by the application ( and often equal to X for a full-reference distortion ) , the latent space in a ( β- ) VAE typically has a lower dimension than X ; a decoder network is then used to parameterize a likelihood model in the data space . To capture this setup , we introduce a new , arbitrary latent space Z on which we define variational distributions QZ|X , QZ , and a ( possibly stochastic ) decoder function ω : Z → X̂ . This results in an extended objective ( with L being the special case of an identity ω ) , J ( QZ|X , QZ , ω , λ ) : = Ex∼PX [ KL ( QZ|X=x‖QZ ) ] + λEPXQZ|X [ ρ ( X , ω ( Z ) ) ] . ( 3 ) How does this relate to the original rate-distortion problem ? We note that the same results from rate-distortion theory apply , once we identify a new distortion function ρω ( x , z ) : = ρ ( x , ω ( z ) ) and treat Z as the reproduction alphabet . Then for each fixed decoder , we may define a ω-dependent rate-distortion function , Rω ( D ) : = infQZ|X : E [ ρω ( X , Z ) ] ≤D I ( X ; Z ) . The minimum of J w.r.t . the variational distributions produces a point on the Rω ( D ) curve . Moreover , as a consequence of the data processing inequality I ( X ; Z ) ≥ I ( X ; ω ( Z ) ) , we can prove that Rω ( D ) ≥ R ( D ) for any ω ( Theorem A.3 ) . Moreover , the inequality is tight for a bijective ω , offering some theoretical support for the use of sub-pixel instead of upsampled convolutions in the decoder of image compression autoencoders ( Theis et al. , 2017 ; Cheng et al. , 2020 ) . We can now minimize the NELBO objective Eq . 3 w.r.t . parameters of ( QZ|X , QZ , ω ) similar to training a β-VAE , knowing that we are optimizing an upper bound on the information R-D function of the data source . This can be seen as a generalization to the lossless case ( with a countable X ) , where minimizing the NELBO minimizes an upper bound on the Shannon entropy of the source ( Frey & Hinton , 1997 ) , the limit of lossless compression . The tightness of our bound depends on the choice of variational distributions . The freedom to define them over any suitable latent space Z can simplify the modeling task ( of which there are many tools ( Salakhutdinov , 2015 ; Kobyzev et al. , 2021 ) ) . e.g. , we can work with densities on a continuous Z , even if X̂ is high-dimensional and discrete . We can also treat Z as the concatenation of sub-vectors [ Z1 , Z2 , ... , ZL ] , and parameterize QZ in terms of simpler component distributions QZ = ∏L l=1QZl|Z < l ( similarly forQZ|X ) . We exploit these properties in our experiments on images .

Summary: This paper aims to provide stronger upper and lower bounds for the RD function of arbitrary sources. The authors handle unknown sources by requiring only i.i.d. samples. Specifically, the authors derive an upper bound to the RD function using a $\beta$-VAE-like generative model, which has some similarities to the Blahut-Arimoto (BA) algorithm with two distinctions: 1) they do not restrict the source to be a low-dimensional discrete data, and 2) they use gradient descent to predict the parameters of the variational distributions $Q_{\hat{X}|X}$ and $Q_{\hat{X}}$. They also derive a lower bound to the RD function using the dual form of the RD function. Finally, the authors provide experimental results to show that the proposed upper bound is, in fact, exact for random Gaussian samples. For more complex data such as banana-shaped sources and their high-dimensional projections, they show that the proposed upper and lower bounds are tight. Furthermore, the proposed upper bounds on the RD function of natural images indicate that the state-of-the-art image compression methods can still be improved by at least one dB PSNR.