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 considers the problem of estimating the rate-distortion function R(D) for arbitrary sources with unknown distribution. An upper bound is established using a variational distribution and is estimated using an iterative coordinate descent algorithm inspired by the Blahut-Arimoto algorithm. A lower bound is established using a parameterization of the dual form of the rate-distortion function by Cisszar. These are then evaluated using neural network encoder/decoders on Gaussian and Banana sources, and tested on natural images suggesting that there is room for improved compression rates achieved by current state-of-the-art image compression algorithms.