Year: 2021

Abstract:
Rate-distortion (R-D) function, a key quantity in information theory, characterizes the fundamental limit of how much a data source can be compressed subject to a fidelity criterion, by any compression algorithm. As researchers push for everimproving compression performance, establishing the R-D function of a given data source is not only of scientific interest, but also sheds light on the possible room for improving compression algorithms. Previous work on this problem relied on distributional assumptions on the data source (Gibson, 2017) or only applied to discrete data. By contrast, this paper makes the first attempt at an algorithm for sandwiching the R-D function of a general (not necessarily discrete) source requiring only i.i.d. data samples. We estimate R-D sandwich bounds on Gaussian and high-dimension banana-shaped sources, as well as GAN-generated images. Our R-D upper bound on natural images indicates room for improving the performance of state-of-the-art image compression methods by 1 dB in PSNR at various bitrates.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 BACKGROUND
text: 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.

heading: 3 UPPER BOUND ALGORITHM
text: 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.

heading: 4 LOWER BOUND ALGORITHM
text: Without knowing the tightness of an R-D upper bound, we could be wasting time and resources trying to improve the R-D performance of a compression algorithm, when it is in fact already close to the theoretical limit. This would be avoided if we could find a matching lower bound on R(D). Unfortunately, the problem turns out to be much more difficult computationally. Indeed, every compression algorithm, or every pair of variational distributions (QX̂ , QX̂|X) yields a point above R(D). Conversely, establishing a lower bound requires disproving the existence of any compression algorithm that can conceivably operate below the R(D) curve. In this section, we derive an algorithm that can in theory produce arbitrarily tight R-D lower bounds. However, as an indication of its difficulty, the problem requires globally maximizing a family of partition functions. By restricting to a continuous reproduction alphabet and a squared error distortion, we make some progress on this problem and demonstrate useful lower bounds on data with low intrinsic dimension (see Sec. 6.2).
Dual characterization of R(D). While upper bounds on R(D) arise naturally out of its definition as a minimization problem, a variational lower bound requires expressing R(D) through a maximization problem. For this, we introduce a “conjugate” function as the optimum of the Lagrangian Eq. 2 (QX̂ is eliminated by replacing the rate upper bound with the exact mutual information I(X; X̂)):
F (λ) := inf QX̂|X
I(X; X̂) + λE[ρ(X, X̂)]. (4)
As illustrated by Fig. 1, F (λ) is the maximum R-axis intercept of a straight line with slope −λ, among all such lines that lie below or tangent to R(D); the RD curve can then be found by taking the upper envelope of lines with slope −λ and R-intercept F (λ), i.e., R(D) = maxλ≥0 F (λ)−λD. This key result is captured mathematically by Lemma A.1, and the following:
Theorem 4.1. (Csiszár, 1974) Under basic conditions (e.g., satisfied by a bounded ρ; see Appendix A.2), it holds that (all expectations below are with respect to the data source r.v. X ∼ PX )
F (λ) = max g(x) {E[− log g(X)]} (5)
where the maximization is over g : X → [0,∞) satisfying the constraint E [
exp(−λρ(X, x̂)) g(X)
] = ∫ exp(−λρ(x, x̂))
g(x) dPX(x) ≤ 1,∀x̂ ∈ X̂ (6)
In other words, every admissible g yields a lower bound of R(D), via an underestimator of the intercept E[− log g(X)] ≤ F (λ). We give the origin and history of this result in related work Sec. 5.
Proposed Unconstrained Formulation. The constraint in Eq. 6 is concerning – it is a family of possibly infinitely many constraints, one for each x̂. To make the problem easier to work with, we propose to eliminate the constraints by the following transformation. Let g be defined in terms of another function u(x) ≥ 0 and a scalar c depending on u, such that
g(x) := cu(x), where c := sup x̂∈X̂
Ψu(x̂), and Ψu(x̂) := E [
exp−λρ(X, x̂) u(X)
] . (7)
This reparameterization of g is without loss of generality, and can be shown to always satisfy the constraint in Eq. 6. While this form of g bears a superficial resemblance to an energy-based model (LeCun et al., 2006), with 1c resembling a normalizing constant, there is an important difference: c = supx̂ Ψu(x̂) is in fact the supremum of a family of “partition functions” Ψu(x̂) indexed by x̂; we thus refer to c as the sup-partition function. Although all these quantities have λ-dependence, we omit this from our notation since λ is a fixed input parameter (as in the upper bound algorithm).
Consequently, F is now the result of unconstrained maximization over all u functions, and we obtain a lower bound on it by restricting u to a subset of functions with parameters θ (e.g., neural networks),
F (λ) = max u≥0 {E[− log u(X)]− log sup x̂∈X̂ Ψu(x̂)} ≥ max θ {E[− log uθ(X)]− log sup x̂∈X̂ Ψθ(x̂)}
For convenience, define the maximization objective by
`(θ) := E[− log uθ(X)]− log c(θ), c(θ) := sup x̂∈X̂ Ψθ(x̂). (8)
Then, it can in principle be maximized by stochastic gradient ascent using samples from PX . However, computing the sup-partition function c(θ) poses serious computation challenges: even evaluating Ψθ(x̂) for a single x̂ value involves a potentially high-dimensional integral w.r.t. PX ; this is only exacerbated by the need to globally optimize w.r.t. x̂, an NP-hard problem even in one-dimension.
Proposed Method. To tackle this problem, we propose a sample-based over-estimator of the suppartition function inspired by the IWAE estimator (Burda et al., 2015). Fix uθ for now. Noting that Ψ(x̂) := E[ψ(X, x̂)] is an expectation w.r.t. PX , where ψ(x, x̂) := exp−λρ(x,x̂)u(x) , one may then naturally consider a plug-in estimator for c, replacing the expectation by a sample estimate of Ψ(x̂). Formally, given a sequence of i.i.d. random variables X1, X2, ... ∼ PX , we define the estimator Ck := supx̂ 1 k ∑ i ψ(Xi, x̂) for each k ≥ 1. We can then prove (see Theorem A.4 and proof in Appendix) that E[C1] ≥ E[C2] ≥ ... ≥ c, i.e., Ck is on average an over-estimator of the sup-partition function c; and like the Importance-Weighted ELBO (Burda et al., 2015), the bias of the estimator decreases monotonically as k → ∞, and that under continuity assumptions, Ck is asymptotically unbiased and converges to c. In light of this, we replace c by E[Ck] and obtain a k-sample under-estimator of the objective `(θ) (which in turn underestimates F (λ)):
`k(θ) := E[− log uθ(X)]− logE[Ck]; moreover, `1(θ) ≤ `2(θ) ≤ ... ≤ `(θ). Unfortunately, we still cannot apply stochastic gradient ascent to `k, as two more difficulties remain. First, Ck is still hard to compute, as it is defined through a global maximization problem. We note that by restricting to a suitable ρ and X̂ = X , the maximization objective of Ck has the form of a kernel density estimate (KDE). For a squared error distortion, this becomes a Gaussian mixture density, 1 k ∑ i ψ(xi, x̂) ∝ ∑ i πi exp(−λ‖xi − x̂‖2), with centroids defined by the samples x1, ..., xk, and mixture weights πi = (ku(xi))−1. The global mode of a Gaussian mixture can generally be found by hill-climbing from each of the k centroids, except in rare artificial examples (Carreira-Perpinan, 2000; 2020); we therefore use this procedure to compute Ck, but note that other methods exist (Lee et al., 2019; Carreira-Perpinan, 2007). The second difficulty is that even if we could estimate E[Ck] (with samples of Ck computed by global optimization), the objective `k requires an estimate of its logarithm; a naive application of Jensen’s inequality − logE[Ck] ≥ E[− logCk] results in an over-estimator (as does the IWAE estimator), whereas we require a lower bound. Following Poole et al. (2019), we underestimate − log(x) by its linearization around a scalar parameter α > 0, resulting in the following lower bound objective:
˜̀ k(θ) := E[− log uθ(X)]− E[Ck]/α− logα+ 1. (9)
`k(θ) can finally be estimated by sample averages, and yields a lower bound on the optimal intercept F (λ) by the chain of inequalities, ˜̀k(θ) ≤ `k(θ) ≤ `(θ) ≤ F (λ). A trained model uθ∗ then yields an R-D lower bound, RL(D) = −λD + ˜̀k(θ∗). We give the detailed algorithm in Appendix A.4.

heading: 5 RELATED WORK
text: Machine Learning: The past few years have seen significant progress in applying machine learning to data compression. For lossless compression, explicit likelihood models (van den Oord et al., 2016; Hoogeboom et al., 2019) directly lead themselves to entropy coding, and bits-back coding techniques are actively being developed for efficient compression with latent variable models (Townsend et al., 2019; Kingma et al., 2019; Ho et al., 2019; Ruan et al., 2021). In the lossy domain, Theis et al. (2017); Ballé et al. (2017) showed that a particular type of VAE can be trained to perform data compression using the same objective as Eq. 3. The variational distributions in such a model have shape restrictions to simulate quantization and entropy coding (Ballé et al., 2017). Our upper bound is directly inspired by this line of work, and suggests that such a model can in principle compute the source R-D function when equipped with sufficiently expressive variational distributions and a “rich enough” decoder (see explanation in Sec. 3). We note however not all compressive autoencoders admit a probabilistic formulation (Theis et al., 2017); recent work has found training with hard quantization to improve compression performance (Minnen & Singh, 2020), and methods have been developed (Agustsson & Theis, 2020; Yang et al., 2020b) to reduce the gap between approximate quantization at training time and hard quantization at test time. Departing from compressive autoencoders, Yang et al. (2020c) and Flamich et al. (2020) applied regular Gaussian VAEs to data compression to exploit the flexiblility of variable-width variational posterior distributions. The REC algorithm from Flamich et al. (2020) can in theory transmit a sample of QX̂|X with a rate close to the rate term of the NELBO-like Eq. 3, but comes with non-negligible overhead. Our experiment in Sec. 6.3 points to the theoretically possible gain in image compression performance from this approach. Agustsson & Theis (2020) proved the general difficulty of this approach without assumptions on the form of QX̂|X , and showed that the particular case of a uniform QX̂|X leads to efficient implementation based on dithered quantization.
Information theory has also broadly influenced unsupervised learning (Alemi et al., 2018; Poole et al., 2019) and representation learning (Tishby et al., 2000). The Information Bottleneck method (Tishby et al., 2000) was directly motivated as a more general R-D theory and borrows from the BA algorithm. Alemi et al. (2018) analyzed the relation between generative modeling and representation learning with a similar R-D Lagrangian to Eq. 2, but used an abstract, model-dependent distortion ρ(x̂, x) := − log p(x|x̂) with an arbitrary X̂ and without considering a data compression task. Recently, Huang et al. (2020) proposed to evaluate decoder-based generative models by computing a restricted version of Rω(D) (with QX̂ fixed to a Gaussian); our result in Sec. 3 (Rω(D) ≥ R(D)) gives a principled way to interpret and compare these model-dependent R-D curves.
Information Theory: While the BA algorithm (Blahut, 1972; Arimoto, 1972) solves the problem of computing the R(D) of a known discrete source, no tools currently exist for the general and unknown case. Riegler et al. (2018) share our goal of computing R(D) of a general source, but still require the source to be known analytically and supported on a known reference measure. Harrison & Kontoyiannis (2008) consider the same setup as ours of estimating R(D) of an unknown source from samples, but focus on purely theoretical aspects, assuming prefect optimization. They prove statistical consistency of such estimators for a general class of alphabets and distortion metrics, assuring that our stochastic bounds on R(D) optimized from data samples, when given unlimited computation and samples, can converge to the true R(D). Perhaps closest in spirit to our work is Gibson (2017), who estimate lower bounds on R(D) of speech and video using Gaussian autoregressive models of the source. However, the correctness of their bounds depends on the modeling assumptions.
A variational lower bound on R(D) was already proposed by Shannon (1959), and later extended (Berger, 1971) to the present version similar to Theorems 4.1 and A.2. In the finite-alphabet case, the maximization characterization of R(D) directly follows from taking the Lagrange dual of its standard definition in Eq. 1 (which is itself a convex optimization problem); the dual problem can then be solved by convex optimization (Chiang & Boyd, 2004), but faces the same limitations as the BA algorithm. A rigorous proof of Thereom 4.1 for general alphabets is more involved, and is based on repeated applications of information divergence inequalities (Csiszár, 1974; Gray, 2011).

heading: 6 EXPERIMENTS
text: We experiment on four types of sources. On random Gaussian sources, we show that our upper bound algorithm can converge to the exact R-D function; our lower bounds, however, become increasingly loose in higher dimensions, and our experiments with varying k offers some insight into this issue. We obtain tight sandwich bounds on a banana-shaped source (Ballé et al., 2021) and its higher-dimension projections, demonstrating that our lower bound algorithm can handle high-dimension data with a low intrinsic dimension. We further confirm this by obtaining sandwich bounds on realistic GAN-generated images with varying intrinsic dimension, and use them to assess the performance of common learned image compression methods. Finally, we estimate bounds on the R-D function of natural images. The intrinsic dimension is still too high for our lower bounds to be useful. Our upper bounds, on the other hand, suggest one dB (in PSNR) of theoretical improvement in the compression quality of state-of-the-art image compression methods, at various bitrates. We provide extensive experimental details and additional results in Appendix sections A.5 and A.6.

heading: 6.1 GAUSSIAN SOURCES
text: We start by testing our algorithms on the factorized Gaussian distribution, one of the few sources with an analytical R-D function. We randomly generate the Gaussian sources in increasing dimensions.
For the upper bound algorithm, we letQX̂ andQX̂|X be factorized Gaussians with learned parameters, predicting the parameters ofQX̂|X by a 1-layer MLP encoder. As shown in Fig. 2-Left, on a n = 1000 dimensional Gaussian (the results are similar across all the n we tested), our upper bound (yellow curve) accurately recovers the analytical R-D function. We also considered optimizing the variational distributions in a latent space Z with varying dimensions, using an MLP decoder to map from Z to X̂ (see Sec. 3). The resulting bounds show similarly good agreement with the ground truth R-D function when the latent dimension matches or exceeds the data dimension n (green, brown), but become increasingly loose as the latent dimension drops below n (red and purple curves), demonstrating the importance of a rich enough latent space for a tight R-D bound, as suggested by our Theorem A.3.
For the lower bound algorithm, we parameterize log u by a 2-layer MLP, and study the effect of source dimension n and the number of samples k used in our estimator Ck (and objective ˜̀k). To simplify comparison of R-D results across different source dimensions, here we consider standard Gaussian sources, whose R-D curve does not vary with n if we scale the rate by 1n (i.e., rate per sample per dimension); the results on randomly generated Gaussians are similar. First, we fix k = 1024 and examine the resulting lower bounds; as shown in Fig. 2-Middle, the bounds quickly become loose with increasing source dimension. This is likely due to the over-estimation bias of our estimator Ck for the sup-partition function, which causes under-estimation in the objective ˜̀k. While Ck is defined similarly to an M-estimator (Van der Vaart, 2000), analyzing its convergence behavior is not straightforward, as it depends on the function u being learned. Empirically, we observe the bias of Ck is amplified by an increase in the source dimension n or λ, such that an increasingly large k is required to effectively train a u model. In additional experiments, we found that for Gaussian sources, an increasing large k (seemingly exponential in n) is needed to reduce the gap in the lower bound. See results in Fig. 4, as well as a detailed discussion of this issue, in Appendix Sec. A.5.2.

heading: 6.2 HIGH-DIMENSION SOURCES WITH LOW INTRINSIC DIMENSION
text: The quickly deteriorating lower bound on higher-dimension Gaussians makes our goal of sandwiching the R(D) of a general source seem hopeless. However, real-world data often contains much more structure than Gaussian noise. Growing research points to the low intrinsic dimension of real-world data such as natural images, and identifies it as a key to the success of learning algorithms (Pope et al., 2021). Echoing these findings, our experiments demonstrate that we can, in fact, obtain tightly matching lower and upper bounds on high-dimension data with a sufficiently low intrinsic dimension.
First, we consider the 2D banana-shaped source from Ballé et al. (2021). Our upper bound algorithm used the same MLP autoencoder as the compression model of Ballé et al. (2021), and parameterized the variaitonal distributions by normalizing flows. Our lower bound algorithm used a simple 3-layer MLP for log u. As shown in Fig. 2-Right, our R-D upper bound (red) lies below the operational R-D curve of Ballé et al. (2021) (orange), and on average of 0.18 bits above our lower bound (maroon), with the sandwiched region shaded in red. The BA algorithm (green) largely agrees with our R-D estimates, but likely underestimates the true R-D in the high distortion region due to the finite discretization region. Next, we mapped the banana source to increasingly high dimensions by a randomly sampled linear transform. The BA algorithm is no longer feasible, even for n = 4. As shown in Fig. 2-Right, we still obtain tight sandwich bounds in n = 16 dimension (blue); the results are similar for all other n we tested, up to 1000. Unlike in the Gaussian experiment, where increasing n required (seemingly exponentially) larger k for a good lower bound, here a constant k = 2048 worked well for all n. The difference is that despite the high dimension, here the data still lies on a low-dimension manifold of the ambient space, a key property that has been observed to facilitate learning from data (Pope et al., 2021).
Following this observation, we experimented on GAN-generated images with controlled intrinsic dimension, and obtain R-D sandwich bounds to assess the performance of actual image compression algorithms. We borrowed the setup from Pope et al. (2021), and experimented on 128 × 128 images of basenji generated by a BigGAN. The intrinsic dimension d of the images are controlled by zeroing out all except d dimensions of the noise input to the GAN. As shown in Fig. 3-Left, the images appear highly realistic, showing the dog with subtly different body features and in various poses. We implemented a 6-layer ResNet-VAE (Kingma et al., 2016) for our upper bound model, and a simple convolutional network for our lower bound model. Fig. 3-Middle plots our resulting R-D bounds and the sandwiched region (red), along with the operational R-D curves (blue, green) of two learned image compression methods (Minnen et al., 2018; Minnen & Singh, 2020) trained on the GAN images, for d = 2 and d = 4. We see that despite the high dimensionality (n = 128× 128× 3), the images require few nats to compress; e.g., for d = 4, we estimate R(D) to be between ∼ 4 and 8 nats per sample when D = 0.001 (30 dB in PSNR). Notably, the R-D curve of Minnen & Singh (2020) stays roughly within a factor of 2 above our estimated true R(D) region. Both neural compression methods show improved performance as d decreases from 4 to 2, following the same
trend as our R-D bounds. This demonstrates the ability of learned compression methods to exploit the low intrinsic dimension of data, contributing to its success over traditional methods such as JPEG (whose R-D curve on this data does not appear to vary with d, and lies orders of magnitude higher).

heading: 6.3 NATURAL IMAGES
text: To establish upper bounds on the R-D function of natural images, we define variational distributions (QZ , QZ|X) on a Euclidean latent space for simplicity, and parameterize them as well as a learned decoder ω via hierarchical VAEs. We borrow the convolutional autoencoder architecture of a state-ofthe-art image compression model (Minnen & Singh, 2020), and set the variational distributions to be factorized Gaussians with learned means and variances (we still use the deep factorized hyperprior, but no longer convolve it with a uniform prior). We also reuse our (β-)ResNet-VAE model with 6 layers of latent variables from the GAN experiments (Sec. 6.2). We trained models with mean-squared error (MSE) distortion and various λ on the COCO 2017 (Lin et al., 2014) images, and evaluated them on the Kodak (1993) and Tecnick (Asuni & Giachetti, 2014) datasets. Following image compression conventions, we report the rate in bits-per-pixel, and the quality (i.e., negative distortion) in PSNR averaged over the images for each (λ, model) pair 2. The resulting quality-rate (Q-R) curves can be interpreted as giving upper bounds on the R-D functions of the image-generating distributions. We plot them in Fig. 3, along with the Q-R performance (in actual bitrate) of various traditional and learned image compression methods (Ballé et al., 2017; Minnen et al., 2018; Minnen & Singh, 2020), for the Kodak dataset (see similar results on Tecnick in Appendix Fig. 9). Our β-VAE model based on (Minnen & Singh, 2020) (orange) lies on average 0.7 dB higher than the operational Q-R curves of the original model (red) and VTM (green). With a deeper latent hierarchy, our (β-)ResNet-VAE gives a higher Q-R curve (blue) that lies on average 1 dB above the state-of-the-art Q-R curves (gap shaded in blue). We leave it to future work to investigate which choice of autoencoder architecture and variational distributions are most effective, as well as how the theoretical R-D performance of such a β-VAE can be realized by a practical compression algorithm (see discussions in Sec. 5).

heading: 7 DISCUSSIONS
text: In this work, we proposed machine learning techniques to computationally bound the rate-distortion function of a data source, a key quantity that characterizes the fundamental performance limit of all lossy compression algorithms, but is largely unknown. Departing from prior work in the information theory community (Gibson, 2017; Riegler et al., 2018), our approach applies broadly to general data sources and requires only i.i.d. samples, making it more suitable for real-world application.
Our upper bound method is a gradient descent version of the classic Blahut-Arimoto algorithm, and closely relates to (and extends) variational autoencoders from neural lossy compression research. Our lower bound method optimizes a dual characterization of the R-D function, which has been known for some time but seen little application outside of theoretical work. Due to difficulties involving global optimization, our lower bound currently requires a squared error distortion for tractability in the continuous case, and is only tight on data sources with a low intrinsic dimension. We hope that a better understanding of the lower bound problem will lead to improved algorithms in the future.
To properly interpret bounds on the R-D function, we emphasize that the significance of the R-D function is two-fold: 1. for a given distortion tolerance D, no coding procedure can operate with a rate less than R(D), and that 2. this rate is asymptotically achievable by some (potentially expensive) procedure. Therefore, a lower bound makes a universal statement about what kind of rate-distortion performance is “too good to be true”. The story is more subtle for the upper bound, due to the asymptotic nature of R(D). The achievability proof relies on a random coding procedure that jointly compresses multiple data samples in increasingly long blocks (Shannon, 1959). When compressing at a finite block length b (e.g., b = 1 when compressing individual images), R(D) is generally no longer achievable, due to a rate overhead that scales like b− 1 2 (Kostina & Verdú, 2012). Extending our work to the setting of finite block lengths could be another useful future direction.
2Technically, to compute an R-D upper bound with MSE ρ, the distortion needs to be evaluated by averaging MSE (instead of PSNR) on samples.

heading: ETHICS AND REPRODUCIBILITY STATEMENTS
text: As our work deals with theoretical aspects of lossy data compression, we are not aware of any potential ethical implications. Our code and instructions for reproducing the datasets and experimental results can be found in this anonymous repo https://drive.google.com/drive/ folders/1uI3q7sVU6T8iF2PoAjNXSAwVHpYLeyHC?usp=sharing.

heading: A APPENDIX

heading: A.1 TECHNICAL DEFINITIONS AND PREREQUISITES
text: In this work we consider the source to be represented by a random variable X : Ω → X , i.e., a measurable function on an underlying probability space (Ω,F ,P), and PX is the image measure of P under X . We suppose the source and reproduction spaces are standard Borel spaces, (X ,AX ) and (X̂ ,AX̂ ), equipped with sigma-algebras AX and AX̂ , respectively. Below we use the definitions of standard quantities from Polyanskiy & Wu (2014).
Conditional distribution The notation QX̂|X denotes an arbitrary conditional distribution (also known as a Markov kernel), i.e., it satifies
1. For any x ∈ X , QX̂|X=x(·) is a probability measure on X̂ ; 2. For any measurable set B ∈ AX̂ , x→ QX̂|X=x(B) is a measurable function on X .
Induced joint and marginal measures Given a source distribution PX , each test channel distribution QX̂|X defines a joint distribution PXQX̂|X on the product space X × X̂ (equipped with the usual product sigma algebra, AX ×AX̂ ) as follows:
PXQX̂|X(E) := ∫ X PX(dx) ∫ X̂ 1{(x, x̂) ∈ E}QX̂|X=x(dx̂),
for all measurable sets E ∈ AX ×AX̂ . The induced x̂-marginal distribution PX̂ is then defined by
PX̂(B) = ∫ X QX̂|X=x(B)PX(dx),
for all measurable sets ∀B ∈ AX̂ .
KL Divergence We use the general definition of Kullback-Leibler (KL) divergence between two probability measures P,Q defined on a common measurable space:
D(P‖Q) := {∫
log dPdQdP, if P Q ∞, otherwise.
P Q denotes that P is absolutely continuous w.r.t. Q (i.e., for all measurable sets E, Q(E) = 0 =⇒ P (E) = 0). dPdQ denotes the Radon-Nikodym derivative of P w.r.t. Q; for discrete distributions, we can simply take it to be the ratio of probability mass functions; and for continuous distributions, we can simply take it to be the ratio of probability density functions.
Mutual Information Given PX and QX̂|X , the mutual information I(X;Y ) is defined as
I(X; X̂) := D(PXQX̂|X‖PX ⊗ PX̂) = Ex∼PX [D(QX̂|X=x‖PX̂)],
where PX̂ is the x̂-marginal of the joint PXQX̂|X , PX ⊗ PX̂ denotes the usual product measure, and D(·‖·) is the KL divergence. For the mutual information upper bound, it’s easy to show that
IU (QX̂|X , QX̂) := Ex∼PX [KL(QX̂|X=x‖QX̂)] = I(X; X̂) +KL(PX̂‖QX̂), (10)
so the bound is tight when PX̂ = QX̂ .
Obtaining R(D) through the Lagrangian. For each λ ≥ 0, we define the Lagrangian by incorporating the distortion constraint in the definition of R(D) through a linear penalty:
L(QX̂|X , λ) := I(X; X̂) + λEPXQX̂|X [ρ(X, X̂)], (11)
and define its infimum w.r.t. QX̂|X by the function
F (λ) := inf QX̂|X
I(X; X̂) + λE[ρ(X, X̂)]. (12)
Geometrically, F (λ) is the maximum of the R-axis intercepts of straight lines of slope −λ, such that they have no point above the R(D) curve (Csiszár, 1974).
Define Dmin := inf{D′ : R(D′) <∞}. Since R(D) is convex, for each D > Dmin, there exists a λ ≥ 0 such that the line of slope −λ through (D,R(D)) is tangent to the R(D) curve, i.e.,
R(D′) + λD′ ≥ R(D) + λ(D) = F (λ), ∀D′. When this occurs, we say that λ is associated to D.
Consequently, the R(D) curve is then the envelope of lines with slope −λ and R-axis intercept F (λ). Formally, this can be stated as:
Lemma A.1. (Lemma 1.2, Csiszár (1974); Lemma 9.7, Gray (2011)). For every distortion tolerance D > Dmin, where Dmin := inf{D′ : R(D′) <∞}, it holds that
R(D) = max λ≥0
F (λ)− λD (13)
We can draw the following conclusions:
1. For each D > Dmin, the maximum above is attained iff λ is associated to D.
2. For a fixed λ, if Q∗Y |X achieves the minimum of L(·, λ), then λ is associated to the point (I(Q∗Y |X), ρ(Q∗Y |X)); i.e., there exists a line with slope −λ that is tangent to the R(D) curve at (I(Q∗Y |X), ρ(Q∗Y |X)).

heading: A.2 FULL VERSION OF THEOREM 4.1
text: Theorem A.2. (Theorem 1, Kostina (2016).) Suppose that the following basic assumptions are satisfied.
1. R(D) is finite for some D, i.e., Dmin := inf{D : R(D) <∞} <∞;
2. The distortion metric ρ is such that there exists a finite set E ⊂ X̂ such that E[min x̂∈E ρ(X, x̂)] <∞
Then, for each D > Dmin, it holds that
R(D) = max g(x),λ
{E[− log g(X)]− λD} (14)
where the maximization is over g(x) ≥ 0 and λ ≥ 0 satisfying the constraint E [
exp(−λρ(X, x̂)) g(X)
] = ∫ exp(−λρ(x, x̂))
g(x) dPX(x) ≤ 1,∀x̂ ∈ X̂ (15)
Note: the basic assumption 2 is trivially satisfied when the distortion ρ is bounded from above; the maximization over g(x) ≥ 0 can be restricted to only 1 ≥ g(x) ≥ 0. Unless stated otherwise, we use log base e in this work, so the R(D) above is in terms of nats (per sample).
Theorem A.2 can be seen as a consequence of Theorem 4.1 in conjunction with Lemma A.1. We implemented an early version of our lower bound algorithm based on Theorem A.2, generating each R-D lower bound by fixing a target D value and optimizing over both λ and g as in Equation 14. However, the algorithm often diverged due to drastically changing λ. We therefore based our current algorithm on Theorem 4.1, producing R-D lower bounds by fixing λ and only optimizing over g (or u, in our formulation).

heading: A.3 THEORETICAL RESULTS
text: Theorem A.3. (A suitable β-VAE defines an upper bound on the source R-D function). Let X ∼ PX be a memoryless source under distortion ρ. Let Z be any measurable space (“latent space” in a VAE), and ω : Z → X̂ any measurable function (“decoder” in a VAE). This induces a new lossy compression problem with Z being the reproduction alphabet, under a new distortion function ρω : X × Z → [0,∞), ρω(x, z) = ρ(x, ω(Z)). Define the corresponding rate-distortion function
Rω(D) := inf QZ|X :E[ρω(X,Z)]≤D I(X;Z) = inf QZ|X :E[ρ(X,ω(Z))]≤D I(X;Z).
Then for any D ≥ 0, Rω(D) ≥ R(D). A sufficient condition for D ≥ 0, Rω(D) = R(D) is for ω to be bijective.
Proof. Fix D. Take any admissible conditional distribution QZ|X that satisfies E[ρ(X,ω(Z))] ≤ D in the definition of Rω(D). Define a new kernel QX̂|X between X and X̂ by QX̂|X=x := QZ|X=x ◦ ω−1,∀x ∈ X , i.e.,QX̂|X=x is the image measure ofQZ|X=x induced by ω. Applying data processing inequality to the Markov chain X
QZ|X−−−→ Z ω−→ X̂ , we have I(X;Z) ≥ I(X; X̂). Moreover, since QX̂|X is admissible in the definition of R(D), i.e.,
EPXQX̂|X [ρ(X, X̂)] = EPXQZ|X [ρ(X,ω(Z))] ≤ D
we therefore have
I(X;Z) ≥ I(X; X̂) ≥ R(D) = inf QX̂|X :E[ρ(X,X̂)]≤D I(X; X̂).
Finally, since I(X;Z) ≥ R(D) holds for any admissible QZ|X , taking infimum over such QZ|X gives Rω(D) ≥ R(D). To prove Rω(D) = R(D) when ω is bijective, it suffices to show that R(D) ≥ Rω(D). We use the same argument as before. Take any admissible QX̂|X in the definition of R(D). We can then construct a QZ|X by the process X QX̂|X−−−−→ X̂ ω −1
−−→ Z. Then by DPI we have I(X; X̂) ≥ I(X;Z). Morever, QZ|X is admissible: EPXQZ|X [ρ(X,ω(Z))] = EPXQX̂|X [ρ(X,ω(ω −1(X̂))] = EPXQX̂|X [ρ(X, X̂)] ≤ D. So I(X; X̂) ≥ I(X;Z) ≥ Rω(D). Taking infimum over such QX̂|X concludes the proof.
Remarks: Although ω−1(ω(X̂)) = X̂ for an injective ω, we needed ω(ω−1(X̂)) = X̂ in the proof of the other direction, which requires ω−1 to always be well-defined. Several learned image compression methods have advocated for the use of sub-pixel convolution, i.e, convolution followed by (invertible) reshaping of the results, over upsampled convolution in the decoder, in order to produce better reconstructions (Theis et al., 2017; Cheng et al., 2020). This can be seen as making the decoder more bijective, therefore reducing the gap of Rω(D) over R(D), in light of our above theorem. Theorem A.4. (Basic properties of the proposed estimator Ck of the sup-partition function.) Let X1, X2, ... ∼ PX be a sequence of i.i.d. random variables. Let ψ : X × X̂ → R be a measurable function. For each k, define the random variable Ck := supx̂ 1 k ∑ i ψ(Xi, x̂). Then
1. Ck is an overestimator of the sup-partition function c, i.e., E[Ck] = EX1,...,Xk [supx̂ 1k ∑ i ψ(Xi, x̂)] ≥ supx̂ E[ψ(X, x̂)] =: c;
2. The bias of Ck decreases with increasing k, i.e., E[C1] ≥ E[C2] ≥ ... ≥ E[Ck] ≥ E[Ck+1] ≥ ... supx̂ E[ψ(X, x̂)] = c;
3. Ifψ(x, x̂) is bounded and continuous in x̂, and if X̂ is compact, thenCk is strongly consistent, i.e., Ck converges to c almost surely (as well as in probability, i.e., limk→∞ P(|Ck − c| > ) = 0,∀ > 0), and limk→∞ E[Ck] = c.
Proof. We prove each in turn:
1. E[Ck] = E[supx̂ 1k ∑ i ψ(Xi, x̂)] ≥ supx̂ E[ 1k ∑ i ψ(Xi, x̂)] = supx̂ E[ψ(X, x̂)] = c
2. First, note that E[C1] ≥ E[Ck] since
E[C1] = E[sup x̂ ψ(X1, x̂)] = E[
1
k ∑ i sup x̂ ψ(Xi, x̂)] ≥ E[sup x̂ 1 k ∑ i ψ(Xi, x̂)] = E[Ck]
We therefore have
E[Ck+1] = E[sup x̂
1
k + 1 k+1∑ i=1 ψ(Xi, x̂)]
= E[sup x̂ { 1 k + 1 k∑ i=1 ψ(Xi, x̂) + 1 k + 1 ψ(Xk+1, x̂)}]
≤ E[sup x̂ { 1 k + 1 k∑ i=1 ψ(Xi, x̂)}+ sup x̂ { 1 k + 1 ψ(Xk+1, x̂)}]
= k
k + 1 E[Ck] +
1
k + 1 E[C1]
≤ E[Ck]
3. The proof for this resembles that of Theorem 1 of Burda et al. (2015). We use standard arguments from probability theory and real analysis. Fix x̂ ∈ X̂ , and consider the random variable Mk = 1k ∑k i=1 ψ(Xi, x̂). If ψ is bounded, then it follows from the Strong Law of
Large Numbers that Mk converges to E[M1] = E[ψ(X, x̂)] almost surely; in other words, for every ω outside a set of measure zero,
lim k→∞
1
k k∑ i=1 ψ(Xi(ω)), x̂) = E[ψ(X(ω), x̂)],
Then, for every such ω
lim k→∞ sup x̂
1
k k∑ i=1 ψ(Xi(ω)), x̂) = sup x̂ lim k→∞ 1 k k∑ i=1 ψ(Xi(ω)), x̂) = sup x̂ E[ψ(X(ω), x̂)],
where we used the fact that the sequence of continuous functions γk(x̂) := 1 k ∑k i=1 ψ(Xi(ω)), x̂) converges pointwise to γ(x̂) := E[ψ(X(ω), x̂)] on a compact set X̂ , so γk converges to γ also uniformly, so we are allowed to exchange limit and supremum, i.e., limk→∞ supx̂ γk(x̂) = supx̂ limk→∞ γk(x̂) = supx̂ γ(x̂). But the above equation precisely means that Ck converges to c = supx̂ E[ψ(X, x̂)] almost surely. Therefore Ck also converges to c in probability, and limk→∞ E[Ck] = c.

heading: A.4 PROPOSED LOWER BOUND ALGORITHM
text: We give a pseudo-code implementation of the algorithm in Algorithm 1. The code largely follows Python semantics, and assumes an auto-differentiation package is available, such as GradientTape as provided in Tensorflow. We use γk to denote the function in the supremem definition of Ck , i.e., γk(x̂) := 1 k ∑k i=1 exp{−λρ(xi, x̂)}/uθ(xi).
For a given λ > 0 and model uθ : x→ R+, the lower bound algorithm works by (stochastic) gradient ascent on the objective ˜̀k(θ) (Eq. 9) w.r.t. the model parameters θ. To compute the gradient, we need an estimate of E[Ck] (ultimately logE[Ck]) as part of the objective, which requires us to draw m samples of Ck. In the pseudo-code, a separate set of m samples of Ck are also drawn to set the linear
Algorithm 1: Example implementation of the proposed algorithm for estimating ratedistortion lower bound RL(D). 1 Requires: λ > 0, model uθ (e.g., a neural network) parameterized by θ, batch sizes k,m,
and gradient ascent step size . 2 while ˜̀k not converged do
// Draw 2m samples of Ck; use the last m samples to set α,
and the first m samples to estimate E[Ck] in ˜̀k. 3 for j := 1 to 2m do 4 Draw k data samples, {xj1, ..., xjk} 5 x̂j , Cjk := compute Ck(θ, λ, {x1, ..., xk}) 6 end 7 α := 1m ∑2m j=m+1 C j k // Compute objective ˜̀k(θ) and update θ by gradient ascent. 8 with GradientTape() as tape: 9 E := 1m ∑m j=1 γk(x̂
j , θ, {xj1, ..., xjk}) // Estimate of E[Ck] 10 ˜̀k := − 1m ∑m j=1 1 k ∑k i=1 log uθ(x j i )− 1αE − logα+ 1 11 gradient := tape.gradient(˜̀k, θ) 12 θ := θ + gradient 13 end 14 Subroutine γk(x̂, θ, λ, {x1, ..., xk}) :
// Evaluate the global maximization objective at x̂. 15 Return 1k ∑k i=1 exp{−λρ(xi, x̂)}/uθ(xi) 16 Subroutine compute Ck(θ, λ, {x1, ..., xk}) : // Compute the global optimum of γk(x̂), assuming squared distortion ρ(x, x̂) ∝ ‖x− x̂‖2. 17 opt loss := −∞ 18 for i := 1 to k do
// Run gradient ascent from the ith mixture component. 19 x̂ := xi 20 while x̂ not converged do 21 with GradientTape() as tape: 22 loss := γk(x̂, θ, λ, {x1, ..., xk}) 23 gradient := tape.gradient(loss, x̂) 24 x̂ := x̂+ gradient 25 end 26 if loss > opt loss then 27 opt loss := loss 28 x̂∗ := x̂ 29 end 30 end 31 Ck = opt loss 32 return (x̂∗, Ck)
expansion point α. In our actual implementation, we set α to an exponential moving average of E[Ck] estimated from previous iterations (e.g., replacing line 7 of the pseudo-code by α := 0.2α+ 0.8E), so only m samples of Ck need to be drawn for each iteration of the algorithm. We did not find this approximation to significantly affect the training or results.
Recall Ck(θ) is defined by the result of maximizing w.r.t. x̂. When computing the gradient with respect to θ, we differentiate through the maximization operation by evaluating Ck = γk(x̂∗) on the forward pass, using the argmax x̂∗ found by the subroutine compute Ck. This is justified by appealing to a standard envelope theorem (Milgrom & Segal, 2002).
By Lemma A.1, each u∗θ trained with a given value of λ yields a linear under-estimator of R(D) ,
RλL(D) = −λD + ˜̀k(θ∗). We obtain the final RL(D) lower bound by taking the upper envelope of all such lines, i.e.,
RL(D) := max λ∈Λ
RλL(D),
where Λ is the set of λ values we trained with.
Tips and tricks:
To avoid numerical issues, we always parameterize log u, and all the operations involving Ck and α are performed in the log domain; e.g., we optimize log γk instead of γk (which does not affect the result since log is monotonically increasing), and use logsumexp when taking the average of multiple Ck samples in the log domain.
During training, we only run an approximate version of the global optimization subroutine compute Ck for efficiency. Instead of hill-climbing from each of the k mixture component centroids, we only do so from the t highest-valued centroids under the objective γk, for a small number of t (say 10). This approximation is not used when reporting the final results.

heading: A.5 ADDITIONAL EXPERIMENTAL DETAILS AND RESULTS

heading: A.5.1 COMPUTATIONAL ASPECTS
text: Our methods are implemented using the Tensorflow library. Our experiments with learned image compression methods used the tensorflow-compression3 library. Our experiments on Gaussian and banana-shaped sources were run on Intel(R) Xeon(R) CPUs, and experiments on images were run on Titan RTX GPUs. We used the Adam optimizer for gradient based optimization in all our experiments, typically setting the learning rate to 1e− 4. Training the β-VAEs for the upper bounds required from a few thousand gradient steps on the lower-dimension problems (under an hour), to a few million gradient steps on the image compression problem (a couple of days; similar to reported in Minnen & Singh (2020)). With our approximate mode finding procedure (starting hill climbing only from a small number of centroids, see Sec. A.4), the lower bound models only required a few thousand steps to train, even on the image experiments (under a day).

heading: A.5.2 GAUSSIANS
text: Data Here the data source is an n-dimensional Gaussian distribution with a diagonal covariance matrix, whose parameters are generated randomly. We sample each dimension of the mean uniformly randomly from [−0.5, 0.5] and variance from [0, 2]. The ground-truth R-D function of the source is computed analytically by the reverse water-filling algorithm (Cover & Thomas, 2006).
Upper Bound Experiments For the Z = X̂ (no decoder) experiment, we chose QX̂ and QX̂|X to be factorized Gaussians; we let the mean and variance vectors of QX̂ be trainable parameters, and predict the mean and variance of QX̂|X=x by one fully connected layer with 2n output units, using softplus activation for the n variance components.
For our experiments involving a decoder, we parameterize the variational distributions QZ and QZ|X similarly to before, and use an MLP decoder with one hidden layer (with the number of units equal the data dimension n, and leaky ReLU activation) to map from the Z to X̂ space. We observe the best performance with a linear (or identity) decoder and a simple linear encoder; using deeper networks with nonlinear activation required more training iterations for SGD to converge, and often to a poorer upper bound. In fact, for a factorized Gaussian source, it can be shown analytically that the optional Q∗Y |X=x is a Gaussian whose mean depends linearly on x, and an identity (no) decoder is optimal.
Lower Bound Experiments In our lower bound experiments, we parameterize log u by an MLP with 2 hidden layers with 20n hidden units each and SeLU activation, where n is the dimension of
3https://github.com/tensorflow/compression
the Gaussian source. We fixed k = 1024 for the results with varying n in Figure 2-Middle. We vary both k and n in the additional experiment below.
To investigate the necessary k needed to achieve a tight lower bound, and its relation to λ or the Gaussian source dimension n, we ran a series of experiments with k ranging from 64 to 1024 on increasingly high dimensional standard Gaussian sources, each time setting λ = n2 . For an n-dimensional standard Gaussian, the true R-intercept, Fn(λ), has an analytical form; in particular, Fn( n 2 ) = n 2 . In Figure 4, we plot the final objective estimate, ˜̀ k, using the converged MLP model, one for each k and n. As we can see, the maximum achievable ˜̀k plateaus to the value log k as we increase n, and for sufficiently high dimension (e.g., n = 20 here), doubling k only brings out a constant (log 2) improvement in the objective. This phenomenon relates to the over-estimation bias of Ck when k is too low compared to λ or the “spread” of the data, and can be understood as follows.
For a given sample size k, there is a “bad” regime produced by increasingly large λ (corresponding to an exceedingly narrow Gaussian kernel), or increasingly high (intrinsic) dimension of the data, so that the k data samples appear very far apart. Numerically, this is exhibited by very quick termination of the k hill-climbing runs when computing Ck (subroutine compute Ck in Algorithm. 1), since the mixture components are well separated, and the mixture centroids are nearly stationary points of the mixture density. The value of the mixture density at each centroid xi can be approximated as γk(xi) = 1 k 1 u(xi) + 1k ∑ j 6=i exp{−λρ(xj , xi)}/uθ(xj) ≈ 1ku(xi) , since exp{−λρ(xj , xi)} ≈ 0 for all j 6= i. The maximization problem defining Ck then essentially reduces to checking which mixture component is the highest, and returning the corresponding centroid (or a point very close to it), so that Ck ≈ supi=1,...,k 1ku(xi) . This implies
E[Ck] ≈ E[ sup i=1,...,k
1
ku(Xi) ] ≤ E[sup x
1
ku(x) ] = sup
x
1
ku(x) ,
therefore
`k := E[− log u(X)]− logE[Ck] (16) ≤ E[− log u(X)] + log k + inf
x log u(x) (17)
= log k + E[− log u(X) + inf x′ log u(x′)] ≤ log k. (18)
So, when these approximations hold, the maximum achievable lower bound objective ˜̀k can never exceed log k. On sources such as high-dimension (e.g., n = 10000) Gaussians, or 256× 256 patches of natural images, λ needs to be on the order of 106 to target low-distortion region of the R-D curve, and the above analysis well describes the behavior of the lower bound algorithm.

heading: A.5.3 BANANA-SHAPED SOURCE
text: Data The 2D banana-shaped source (Ballé et al., 2021) is a RealNVP transform of a 2D Gaussian 4. See Figure 5 for a plot of its density function. We also consider mapping it to a higher dimension n by a n× 2 matrix. We implemented this by simply passing samples of the banana source through a linear/dense MLP layer with n output units and no activation, with its weight matrix randomly drawn from a Gaussian by the Glorot initialization procedure.
Upper Bound Experiments For the 2D banana-shaped source, we base our β-VAE model architecture on the compressive autoencoder in Ballé et al. (2021), using a 2-dimensional latent space and two-layer MLPs for both the encoder and decoder, with 100 hidden units and softplus activation in each hidden layer. We parameterize QZ by a MAF (Papamakarios et al., 2017), and QZ|X by an IAF (Kingma et al., 2016) (we found that a factorized Gaussian QZ|X works just as well, as long as the MAF prior QZ is sufficiently powerful).
For the higher-dimension projections of the banana-shaped source, we use the same β-VAE as on the 2D source, but set the number of hidden units in each hidden layer as 100n, where n is the dimension of the source. We cap the number of hidden units per layer at 2000 for n > 20, and do not find this to adversely affect the results.
Lower Bound Experiments We use an MLP with three hidden layers and SeLU activations for the log u model; as in the upper bound models, here we set the number of hidden units in each layer to be 100n, and cap it at 1000. In all the experiments we set k = 2048.
As illustrated by Figure 6, the tightness of our sandwich bounds do not appear to be affected by the dimension n of the data (we verified this up to n = 1000), unlike on the Gaussian experiment where for a fixed k the lower bound became increasingly loose with higher n.

heading: A.5.4 GAN-GENERATED IMAGES
text: Data We adopt the same setup as in the GAN experiment by Pope et al. (2021). We use a BigGAN (Brock et al., 2019) pretrained on ImageNet at 128× 128 resolution (so that n = 128× 128× 3 = 49152), downloaded from https://tfhub.dev/deepmind/biggan-deep-128/1. To generate an image of an ImageNet category, we provide the corresponding one-hot class vector as well as a noise vector to the GAN. Following Poole et al. (2019), we control the intrinsic dimension d of the generated images by setting all except the first d dimensions of the 128-dimension noise vector to zero, and use a truncation level of 0.8 for the sample diversity factor. Since the GAN generates values in [−1, 1], we rescale them linearly to [0, 1] to correspond to images, so that the alphabets
4Source code taken from the authors’ GitHub repo https://github.com/tensorflow/ compression/blob/66228f0faf9f500ffba9a99d5f3ad97689595ef8/models/toy_ sources/toy_sources.ipynb.
X = X̂ = [0, 1]n. As in (Pope et al., 2021), we experimented with images from the ImageNet class basenji. In Figure 7 and 8, we plot additional random samples, for d = 2 and d = 4.
Upper Bound Experiments We implemented a version of ResNet-VAE following the appendix of Kingma et al. (2016), using the 3 × 3 convolutions of Cheng et al. (2020) in the encoder and decoder of our model. Our model consists of 6 layers of latent variables, and implements bidirectional inference using factorized Gaussian variational distributions at each stochastic layer. During an inference pass, an input image goes through 6 stages of convolution followed by downsampling (each time reducing the height and width by 2), and results in a stochastic tensor at each stage. In encoding order, the stochastic tensors have decreasing spatial extent and an increasing number of channels equal to 4, 8, 16, 32, 64, and 128. The latent tensor Z0 at the topmost generative hierarchy is flattened and modeled by a MAF prior QZ0 (aided by the fact that all the images have a fixed shape).
Lower Bound Experiments We parameterize the log u model by a simple feedforward convolutional neural network. It contains three convolutional layers with 4, 8, and 16 filters, each time downsampling by 2, followed by a fully connected layer with 1024 hidden units. Again we use SeLU activation in the hidden units, and set k = 2048 in all the experiments. When evaluating the global optimization objective γk (see Algorithm 1), we loop through a small batch of 32 xi sample at a time to avoid running out of GPU memory.

heading: A.5.5 NATURAL IMAGES
text: Data For signals such as images or video, which can have arbitrarily high spatial dimension, it is not immediately clear how to define i.i.d. samples. But since our focus is on image compression, we follow the common practice of training convolutional autoencoders on random 256×256-pixel crops of high resolution images as in learned image compression research (Ballé et al., 2017), noting that current methods cannot effectively exploit spatial redundancies larger than this scale (Minnen & Singh, 2020). As a representative dataset of natural images, we used images from the COCO 2017 (Lin et al., 2014) training set that are larger than 512×512, and downsampled each by a random factor in [0.6, 1] to remove potential compression artifacts from JPEG. We trained image compression models from (Minnen et al., 2018; Minnen & Singh, 2020) on our dataset and were able to closely reproduce their performance.
Upper Bound Experiments Instead of working with variational distributions over the set of pixel values X̂ = X = {0, 1, ..., 255}n, which would require specialized tools for representing and optimizing high-dimensional discrete distributions (Salimans et al., 2017; Hoogeboom et al., 2019; Maddison et al., 2016; Jang et al., 2016), here we parameterize the variational distributions in Euclidean latent spaces for simplicity (see Sec. 3).
We applied our upper bound algorithm based on two VAE architectures:
1. The learned image compression model from Minnen & Singh (2020) 5. This corresponds to a hierarchical VAE with two layers of latent variables, with the lower level latents modeled by a channel-wise autoregressive prior conditioned on the top-level latents. We adapt it for our purpose, by using factorized Gaussians for the variational posterior distributions QZ|X (instead of factorized uniform distributions used to simulate rounding in the original model), and no longer convolving the variational prior QZ with a uniform distribution as in the original model.
2. The ResNet-VAE from our experiment on GAN images. Following learned image compression models, we maintain the convolutional structure of the topmost latent tensor (rather than flattening it in the GAN experiments) in order to make the model applicable to images of variable sizes. Instead of a MAF, we model the topmost latent tensor by a deep factorized (hyper-)prior QZ0 as proposed by Ballé et al. (2018).
For the β-VAE based on Minnen & Singh (2020), dim(Z) ≈ 0.28 dim(X ); and for our β-ResNetVAE, dim(Z) ≈ 0.66 dim(X ). We did not observe improved results by simply increasing the number of latent channels/dimensions in the former model. We follow the same model training practice as reported by Minnen & Singh (2020), training both models to around 5 million gradient steps.
In our problem setup, we consider the reproduction alphabet to be the discrete set of pixel values X = {0, 1, ..., 255}n, so the decoder network ω needs to discretize the output of the final convolutional layer (in Tensorflow, this can be implemented by a saturate cast to uint8). However, since the discretization operation has no useful gradient, we skip it during training, and only apply it for evaluation.
The operational quality-rate curves of various compression methods are taken from the tensorflow-compression6 and compressAI7 libraries. The results of Minnen & Singh (2020) represent the current state-of-the-art of neural image compression, to the best of our knowledge.
Lower Bound Experiments When running our lower bound algorithm on 256 × 256 crops of natural images, we observed similar behavior as on high-dimension Gaussians, with the loss ˜̀k
5https://github.com/tensorflow/compression/blob/f9edc949fa58381ffafa5aa8cb37dc8c3ce50e7f/ models/ms2020.py
6https://github.com/tensorflow/compression/tree/8692f3c1938b18f123dbd6e302503a23ce75330c/ results/image_compression
7https://github.com/InterDigitalInc/CompressAI/tree/ baca8e9ff070c9f712bf4206b8f2da942a0e3dfe/results
quickly plateauing to a value around log k. This issue persisted when we tried different u model architectures. See a discussion on this phenomenon in Section A.5.2.

heading: A.6 MEASURES OF VARIABILITY
text: As alluded to in Sec. 3, the proposed R-D bounds are estimated as empirical averages from samples, which only equal the true quantities in expectation. Before we characterize the variability of the bound estimates, we first define the exact form of the estimators used.
Upper Bound Estimator For the proposed upper bound, consider the variational problem solved by the Lagrangian Eq. 3, and define the random variables from the rate and distortion terms of the Lagrangian: R(Z,X) := log q(Z|X)− log q(Z), and D(Z,X) := ρ(X,ω(Z)). Above,X and Z follow the joint distribution PXQZ|X , ω is the decoder function, and q(z|x) and q(z) are the Lebesgue density functions of absolutely continuousQZ|X=x andQZ variational distributions (we only worked with absolutely continuous variational distributions for simplicity). Then it is clear that the expectations of the two random variables equal the rate and distortion terms, i.e.,
E[R(Z,X)] = Ex∼PX [KL(QZ|X=x‖QZ)], and E[D(Z,X)] = EPXQZ|X [ρ(X,ω(Z))]. Recall for any given decoder and variational distributions, the above rate and distortion terms define a point (E[D],E[R]) that lies on an upper bound RU (D) of the source R(D) curve. Given m i.i.d. pairs of (Z1, X1), (Z2, X2), ..., (Zm, Xm) ∼ PXQZ|X , we estimate such a point by (µD, µR), using the usual sample mean estimators µR := 1m ∑m j=1R(Zj , Xj) and µD := 1m ∑m j=1D(Zj , Xj).
Lower Bound Estimator For the proposed lower bound, given i.i.d. X1, ..., Xk ∼ PX , define the random variable
ξ := −1 k k∑ i log u(Xi)− Ck α − logα+ 1,
where the various quantities are defined as in Sec. 4. Then it is clear that its expectation equals the lower bound objective Eq. 9, i.e., E[ξ] = ˜̀k := E[− log u(X)] − E[Ck]/α − logα + 1. Given m i.i.d. samples of ξ, we form the usual sample mean estimator µξ := 1m ∑ j ξj , and form an R-D lower bound estimate by R̂L(D) := −λD + µξ.
In computing µξ, we fix the constant α to a sample mean estimate of E[Ck] computed separately (as in line 7 of Algorithm 1). As explained in Sec. A.4, we repeat this procedure with different λ, and our overall R-D lower bound is obtained by taking the point-wise maximum of such linear under-estimators over λ.
Results Below we give measures of variability of for the sandwich bound estimates obtained on GAN generated images (Section 6.2). We report detailed statistics in the following tables, and plot them in Figure 10. Our results on other experiments exhibit similar or comparatively less variability.
For the upper bound, given each trained β-VAE, we compute m samples of (R,D), essentially passing m samples of X through the autoencoder and collect the rate and distortion values; then we report the sample mean (µ), sample standard deviation (s), and 95% large-sample confidence interval for the population mean (the true expected value) forR and D, respectively, in Tables 1 and 2. For the lower bound, given each trained log u model, we compute m samples of ξ, then report the sample mean (µξ), sample standard deviation (sξ), and 90% large-sample confidence lower bound for the population mean (the true expected value E[ξ] = ˜̀k) in Tables 3 and 4.

heading: 300 0.265 0.689 (0.15, 0.38) 0.00978 0.00376 (0.00916, 0.0104)
text: λ µξ sξ E[ξ] LCB 49.80 0.37 0.01 0.36 99.30 0.74 0.02 0.73 199.82 1.46 0.03 1.45 299.58 2.01 0.03 2.00 499.87 2.76 0.05 2.74 752.70 3.29 0.05 3.27 799.92 3.41 0.05 3.39 999.44 3.63 0.07 3.60
1514.20 4.04 0.06 4.01 1999.30 4.21 0.08 4.18 2275.80 4.33 0.07 4.30 3000.00 4.50 0.07 4.47

Review summaries:
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.

The paper proposes to use ML to establish lower and upper bounds on rate distortion for general sources, thus going beyond the Blahut-Arimoto algorithm (which assumes discrete sources with known PMF, and has complexity exponential in the dimension of the data). Some theoretical results are provided, extending prior results. Some numerical results are provided 

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.