T OWARDS E MPIRICAL S ANDWICH B OUNDS ON THE R ATE -D ISTORTION F UNCTION
A BSTRACT
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 ﬁdelity 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 scientiﬁc interest, but also reveals the possible room for improvement in 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 ﬁrst 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 for a variety of artiﬁcial and real-world data sources, in settings far beyond the feasibility of any known method, and shed light on the optimality of neural data compression (Ball ´ e et al., 2021; Yang et al., 2022).
Our R-D upper bound on natural images indicates theoretical room for improving state-of-the-art image compression methods by at least one dB in PSNR at various bitrates.
Our data and code can be found here.
I NTRODUCTION
From storing astronomical images captured by the Hubble telescope, to delivering familiar faces and voices over video calls data compression, i e., communication of 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 signiﬁcant 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 in 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 difﬁcult 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 long-standing 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 speciﬁcally,
Similarly to how a VAE with a discrete likelihood model minimizes an upper bound on the data entropy, we establish that any β -VAE with a likelihood model induced by a distortion metric minimizes an upper bound on the data rate-distortion function .
We thus open the deep generative modeling toolbox to the estimation of an upper bound on the R-D function.
We derive a lower bound estimator of the R-D function that can be made asymptotically exact and optimized by stochastic gradient ascent.
Facing the difﬁculty of the problem involving global optimization, we restrict to a squared error distortion for a practical implementation.
We perform extensive experiments and obtain non-trivial sandwich bounds on various data sources, including GAN-generated artiﬁcial sources and real-world data from speech and physics.
Our results shed light on the effectiveness of neural compression approaches (Ball´e et al., 2021; Minnen et al., ; Minnen & Singh and identify the intrinsic (rather than nominal) dimension of data as a key factor affecting the tightness of our lower bound.
Our estimated R-D upper bounds on high-resolution natural images (evaluated on the standard Kodak and Tecnick datasets) indicate theoretical room for improvement of state-ofthe-art image compression methods by at least 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.
B ACKGROUND
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 Y .
It asks the following question about the limit of lossy compression: for a given data source and a distortion metric (a.k.a., a ﬁdelity 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 Y , conventionally called the source and reproduction alphabets , respectively.
We deﬁne the data source formally by a random variable X ∈ X following a (usually unknown) distribution P X , and assume a distortion metric ρ : X × Y → [0 , ∞ ) has been given, such as the squared error ρ ( x, y )
= (cid:107) x − y (cid:107) 2 .
The rate-distortion function is then deﬁned by the following constrained optimization problem,
[Equation]
where we consider all random transforms Q Y | X whose expected distortion is within the given threshold D ≥ 0 , and minimize the mutual information between the source X and its reproduced representation Y 1 .
Shannon’s lossy source coding theorems (Shannon, 1948; 1959) gave operational signiﬁcance to the above mathematical deﬁnition 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 P X exist and can be approximated by sample averages.
When the source alphabet is ﬁnite, this assumption automatically holds, and R ( D ) also provides a lower bound on the Shannon entropy of discrete data.
U PPER B OUND A LGORITHM
R-D theory (Cover & Thomas, 2006) tells us that every (distortion, rate) pair lying above the R ( D ) - curve is in theory realizable by a (possibly expensive) compression algorithm.
An upper bound on
R ( D ) thus reveals what kind of compression performance is theoretically achievable.
Towards this goal, we borrow the variational principle of the Blahut-Arimoto (BA) algorithm, but extend it to a general (e.g., non-discrete) source requiring only its samples.
Our resulting algorithm optimizes a β -VAE whose likelihood model is speciﬁed by the distortion metric, a common case being a Gaussian likelihood with a ﬁxed variance.
For the ﬁrst time, we establish this class of models as computing a model-agnostic upper bound on the source R-D function, as deﬁned by a data compression task.
Variational Formulation.
Following the BA algorithm (Blahut, 1972; Arimoto, 1972), we consider a Lagrangian relaxation of the constrained problem deﬁning R ( D ) , which has the variational objective
[Equation]
where Q Y is a new, arbitrary probability measure on Y .
The ﬁrst ( rate ) term is a variational upper bound on the mutual information
I ( X ; Y ) , and the second ( distortion ) term enforces the distortion tolerance constraint in Eq.
For each ﬁxed λ > 0 , the BA algorithm globally minimizes L w.r.t.
the variational distributions Q Y | X and Q Y by coordinate descent; at convergence, the (distortion, rate) pair yields a point on the R ( D ) curve (Csisz´ar, 1974a).
Unfortunately, the BA algorithm only applies when X and Y are ﬁnite (hence discrete), and the source distribution known.
Otherwise, a preprocessing step is required to discretize a continuous source and/or estimate source probabilities by a histogram, which introduces a non-negligible bias.
This bias, along with its exponential complexity in the data dimension, also makes BA infeasible beyond a couple of (usually 2 or 3) dimensions.
Proposed Method.
To avoid these difﬁculties, we propose to apply (stochastic) gradient descent on
L w.r.t.
ﬂexibly parameterized variational distributions Q Y | X and Q Y .
In this work we parameterize the distributions by neural networks, and predict the parameters of each Q Y | X = x by an encoder network φ ( x ) as in amortized inference (Kingma & Welling, 2014).
Given data samples, the estimates of rate and distortion terms of L yield a point that in expectation lies on an R-D upper bound R U ( D ) , and we tighten this bound by optimizing L ; repeating this procedure for various λ traces out R U ( D ) .
The objective L closely resembles the negative ELBO (NELBO) objective of a β -VAE (Higgins et al., 2017) if we view the reproduction space Y as the “latent space”.
The connection is immediate when X is continuous and a squared error ρ speciﬁes the density of a Gaussian likelihood p ( x | y )
∝ exp( −(cid:107) x − y | 2 ) .
However, unlike in data compression, where Y 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 , and a decoder network is 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 deﬁne variational distributions Q Z | X , Q Z , and a (possibly stochastic) decoder function ω : Z → Y .
This results in an extended objective, resembling a β -VAE with a likelihood density p ( x | z ) ∝ exp {− ρ ( x, ω ( z ) } ,
[Equation]
This objective is closely related to the original data compression task and provides an upper bound on the source R ( D ) , as follows.
Treating Z as the reproduction alphabet, we can deﬁne a new distortion ρ ω ( x, z ) :
= ρ ( x, ω ( z )) , and a ω -induced R-D function, R ω ( D ) :=
inf Q Z | X : E [ ρ ω ( X,Z )] ≤ D I
( X ; Z ) , for each choice of ω .
Our Theorem A.3 then guarantees that R ω ( D ) ≥ R ( D ) , for any ω , and consequently the (distortion, rate) of J always lies above R ( D ) .
Moreover, R ω ( D ) =
R ( D ) for a bijective ω , which offers some theoretical support for the use of invertible pixel-shufﬂe operations 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-like objective J w.r.t.
the parameters of ( Q Z | X , Q Z , ω ) similar to training a β -VAE, knowing that we are optimizing an upper bound on the rate-distortion function of the source.
This can be seen as the lossy counterpart to the lossless setting, where it is well-established that minimizing the NELBO minimizes an upper bound on the Shannon entropy of the source (Frey & Hinton, 1997; MacKay, 2003), the limit of lossless data compression.
The extended objective offers the freedom to deﬁne variational distributions on any suitable latent space Z , rather than Y , which we found to simplify the modeling task and yield tighter bounds.
E.g., even if Y is high-dimensional and discrete, we can still work with densities on a continuous and lower-dimensional Z and draw upon tools such as normalizing ﬂows (Kobyzev et al., 2021).
We can also treat Z as the concatenation of sub-vectors [ Z 1 , Z 2 , ..., Z L ] , and parameterize Q Z in terms of simpler component distributions Q Z = (cid:81)
Ll =1 Q Z l | Z <l (similarly for Q Z | X ) as in a hierarchical VAE.
L OWER B OUND A LGORITHM
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 ﬁnd a matching lower bound on R ( D ) .
Unfortunately, the problem turns out to be much more difﬁcult computationally.
Indeed, every compression algorithm, or every pair of variational distributions ( Q Y , Q Y | 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 difﬁculty, 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 obtain useful lower bounds especially on data with low intrinsic dimension (see Sec. 6).
[Figure]
we seek to maximize the rate -axis intercept, E [ − log g ( X )] , over all g ≥ functions admissible according to
Dual characterization of R ( D ) .
While upper bounds on R ( D ) arise naturally out of its deﬁnition as a minimization problem, a variational lower bound requires expressing R ( D ) through a maximization problem.
For this, we introduce a “dual” function as the optimum of the
Lagrangian Eq.
2 ( Q Y is eliminated by replacing the rate upper bound with the exact mutual information
I ( X ; Y ) ):
[Equation]
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:
Theorem 4.1.
(Csisz´ar, 1974b)
(As follows, all the expectations are with respect to the data source r.v.
X ∼ P X .)
Under basic conditions (e.g., satisﬁed by a bounded ρ ; see Appendix A.2), it holds that
[Equation]
where the maximization is over all non-negative functions g : X → [0 , ∞ ) satisfying the constraint
[Equation]
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 of this result in related work in Section 5.
Proposed Unconstrained Formulation.
The constraint in Eq.
6 is concerning—it is a family of possibly inﬁnitely many constraints, one for each y .
To make the problem easier to work with, we propose to eliminate the constraints by the following transformation.
Let g be deﬁned in terms of another function u ( x ) ≥ 0 and a scalar c depending on u , such that
[Equation]
This reparameterization of g is without loss of generality, and can be shown to always satisfy the constraint in Eq.
While this form of g bears a superﬁcial resemblance to an energy-based model
(LeCun et al., 2006), with 1 c resembling a normalizing constant, there is an important difference: c = sup y Ψ u ( y ) is in fact the supremum of a family of “partition functions” Ψ u ( y ) indexed by y ; 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 ﬁxed 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),
[Equation]
Deﬁne the θ -parameterized objective (cid:96) ( θ ) :=
E [ − log u θ ( X )] − log c ( θ ) , with c ( θ ) =
sup y ∈Y Ψ θ ( y ) .
Given samples of X , we can in principle maximize (cid:96) ( θ ) by (stochastic) gradient ascent.
However, computing the sup-partition function c ( θ ) poses serious computation challenges: even evaluating Ψ θ ( y ) for a single y value involves a potentially high-dimensional integral w.r.t.
P X ; this is only exacerbated by the need to globally optimize w.r.t.
y , an NP-hard problem even in one-dimension.
Proposed Method.
To tackle this problem, we propose an over-estimator of the sup-partition function inspired by IWAE (Burda et al., 2015).
Fix θ for now; we denote the integrand in
Eq. 7 by ψ ( x, y ) :
= exp − λρ ( x,y ) u ( x ) (so c = sup y ∈Y E [ ψ ( X, y )] ), and omit the dependence on θ to simplify notation.
Given k ≥ 1 i.i.d.
random variables X 1 , ..., X k ∼ P X , deﬁne the estimator
C k :=
sup y 1 k
(cid:80)
i ψ ( X i , y ) .
We prove in Theorem A.4 that E [ C 1 ] ≥ E [ C 2 ] ≥ ... ≥ c , i.e., C k is in expectation an over-estimator of the sup-partition function c .
Similarly to the Importance-Weighted
ELBO (Burda et al., 2015), the bias of this estimator decreases monotonically as k → ∞ , and asymptotically vanishes under regularity assumptions.
In light of this, we replace c by E [ C k ] and obtain a k -sample under-estimator of the objective (cid:96) ( θ ) (which in turn underestimates F ( λ ) ):
[Equation]
In order to apply stochastic gradient ascent, we overcome two more technical hurdles.
First, each draw of C k requires solving a global maximization problem.
We note that by restricting to a squared-error ρ and Y =
X , C k can be computed by ﬁnding the mode of a Gaussian mixture density; for this we use the method of Carreira-Perpinan (2000), essentially by hill-climbing from each of the k centroids.
Second, to turn − log E [ C k ] into an expectation, we follow Poole et al.
and underestimate it by linearizing − log around a scalar parameter α > 0 , resulting in the following lower bound objective:
[Equation]
˜
[Equation]
k ( θ ) can ﬁnally be estimated by sample averages, and yields a lower bound on the optimal intercept
F ( λ ) , via ˜ (cid:96)
k ( θ ) ≤ (cid:96) k
( θ ) ≤ (cid:96) ( θ ) ≤ F ( λ ) .
A trained model u θ ∗ then yields an R-D lower bound,
R L ( D ) = − λD + ˜ (cid:96) k
( θ ∗ ) .
We give a more detailed derivation and pseudocode in Appendix A.4.
R ELATED W ORK
Machine Learning:
The past few years have seen signiﬁcant progress in applying machine learning to lossy data compression.
Theis et al.
(2017); Ball´e et al.
ﬁrst showed that a particular type of β -VAE can be trained to perform data compression using the same objective as Eq.
The variational distributions in such a model have shape restrictions to simulate quantization and entropy coding (Ball´e 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 sufﬁciently expressive variational distributions and a “rich enough” decoder (see 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) use Gaussian β -VAEs for data compression and exploit the ﬂexibility of variable-width Gaussian posteriors.
Flamich et al.
(2020)’s method, and
more generally, reverse channel coding (Theis & Yosri, 2021), can transmit a sample of Q Z | X with a rate close to that optimized by our upper bound model in Eq.
3, more precisely, I ( X ; Z ) + log( I ( X ; Z ) + 1) + O (1) .
i.e., in this one-shot setting (which is standard for neural image compression), R ( D ) is no longer achievable; rather, the achievable R-D performance is characterized by R ( D ) + log( R ( D ) + 1) + O (1) .
Therefore, our R-D bounds can be shifted upwards by this logarithmic factor to give an estimate of the achievable R-D performance in this setting.
Information theory has also broadly inﬂuenced unsupervised learning and representation learning.
The Information Bottleneck method (Tishby et al., 2000) was directly motivated by, and extends
R-D theory and 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 ρ ( y, x ) :
= − log p ( x | y ) with an arbitrary Y 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 Q Y ﬁxed); 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) computes the R ( D ) of a discrete source with a known distribution, no tool currently exists 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 by Gibson (2017), who estimates lower bounds on R ( D ) of speech and video using Gaussian autoregressive models of the source.
However, the correctness of the resulting 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 ﬁnite-alphabet case, the maximization characterization of R ( D ) follows from taking the Lagrange dual of its standard deﬁnition in Eq.
1; the dual problem can then be solved by convex optimization (Chiang & Boyd, 2004), but faces the same computational difﬁculties as the BA algorithm.
Csisz ´ ar (1974b) proved the general result in Theorem 4.1, applicable to abstract alphabets (in particular, with source X taking values in an arbitrary probability space), by analyzing the ﬁxed-point conditions of the BA algorithm.
E XPERIMENTS
We estimate the R-D functions of a variety of artiﬁcial and real-world data sources, in settings where the BA algorithm is infeasible and no prior known method has been applied.
On Gaussian sources , our upper bound algorithm is shown to converge to the exact R-D function, while our lower bounds become increasingly loose in higher dimensions, an issue we investigate subsequently.
We obtain tighter sandwich bounds on particle physics and speech data than on similar dimensional Gaussians, and compare with the performance of neural compression.
We further investigate looseness in the lower bound, experimentally establishing the intrinsic dimension of the data as a much more critical contributing factor than the nominal/ambient dimension.
Indeed, we obtain tight sandwich bounds on high-dimension GAN-generated images with a low intrinsic dimension, and compare with popular neural image compression methods.
Finally, we estimate bounds on the R-D function of natural images .
The intrinsic dimension is likely too high for our lower bound to be useful, while our upper bounds on the Kodak and Tecnick datasets imply at least one dB (in PSNR) of theoretical room for improving state-of-the-art image compression methods, at various bitrates.
We also validate our bounds against the BA algorithm over the various data sources, using a 2D marginal of the source to make BA feasible.
We provide experimental details and additional results in Appendix A.5 and A.6.
6. G AUSSIAN SOURCES
We start by applying our algorithms to 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 let Q Y and Q Y | X be factorized Gaussians with learned parameters, predicting the parameters of Q Y | X by a 1-layer MLP encoder.
As shown in Fig.
2a- top , on a n =dimensional Gaussian (the results are similar across all the n we tested), our upper bound ( yellow
[Figure]
curve) accurately recovers the analytical R ( D ) .
We also optimized the variational distributions in a latent space Z with varying dimensions, using an MLP decoder to map from Z to Y (see Sec. 3).
The resulting bounds are similarly tight when the latent dimension matches or exceeds the data dimension n ( green , brown ), but become loose otherwise ( 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 C k (and objective ˜ (cid:96) k ).
To simplify comparison of 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 1 n (i.e., rate per sample per dimension); the results on randomly generated Gaussians are similar.
First, we ﬁx k = 1024 ; Fig.
2a- bottom shows that the resulting bounds quickly become loose with increasing source dimension.
This is due to the bias of our estimator C k for the sup-partition function, which causes under-estimation in the objective ˜ (cid:96) k .
While C k is deﬁned 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.
In this experiment, we observe the bias of C k is ampliﬁed by an increase in n or λ , such that an increasingly large k is required for effective training.
In another experiment, we estimate that the k needed to close the gap in the lower bound increases exponentially in n ; see results in Fig. 4, and a detailed discussion on this, in Appendix A.5.2.
Fortunately, as we see in Sec. 6.3, the bias in our lower bound appears to depend on the intrinsic rather than (often much higher) nominal dimension of data, giving us a more favorable trade-off between computation and a tighter lower bound as controlled by k .
6. D ATA FROM PARTICLE PHYSICS AND SPEECH
The quickly deteriorating lower bound on higher (even 16) dimensional Gaussians may seem disheartening.
However, the Gaussian is also the hardest continuous source to compress under squared error (Gibson, 2017), and real-world data often exhibits considerably more structure than Gaussian noise.
In this subsection, we experiment on data from particle physics (Howard et al., 2021) and speech (Jackson et al., 2018), and indeed obtain improved sandwich bounds compared to Gaussians.
First, we consider the Z -boson decay dataset from Howard et al.
(2021), containing n =16-dimensional vectors of four-momenta information from independent particle decay events.
We ran the neural compression method from (Ball´e et al., 2021), as well as our bounding algorithms with similar conﬁgurations to before, except we ﬁx k = 2048 for the lower bound, and use a normalizing ﬂow for
Q Z in the upper bound model for better expressiveness.
Fig.
2b top shows the resulting estimated
R-D bounds (sandwiched region colored in red ) and the operational R-D curve for Ball ´ e et al.
[Figure]
( blue ).
The resulting sandwich bounds appear signiﬁcantly tighter here than on the Gaussian source with equal dimension ( n = 16 , bottom curve in Fig.
2a bottom ), and the neural compression method operates with a relatively small average overhead of 0.5 nat/sample relative to our upper bound.
To also compare to the ground-truth R ( D ) as estimated by the BA algorithm, we created a 2-dimensional marginal data distribution (plotted in Fig.
2d top ), so that the BA algorithm can be feasibly run with a ﬁne discretization grid.
As shown in Fig.
2b, the BA estimate of R ( D ) ( green ) almost overlaps with our upper bound on the 2D marginal, and is tightly sandwiched from below by our lower bound.
We then repeat the same experiments on speech data from the Free Spoken Digit Dataset (Jackson et al., 2018).
We constructed our dataset by pre-processing the audio recordings into spectrograms (see, e.g., Fig. 2d bottom ), then treating the resulting n = 33 -dimensional frequency feature vectors as independent across time.
As shown in Fig.
2c, the gap in our R-D bounds appears wider than on the physics dataset ( top ), but the results on the corresponding 2D marginal appear similar ( bottom ).
6. T HE E FFECT OF I NTRINSIC V . S . N OMINAL D IMENSION OF D ATA
It is known that learning a manifold has a complexity that depends on the intrinsic dimension of the manifold, but not on its ambient dimension (Narayanan & Mitter, 2010).
Our experiments suggest a similar phenomenon for our lower bound, and show that we can still obtain tight sandwich bounds on data with a sufﬁciently low intrinsic dimension despite the high ambient dimension.
First, we explore the effect of increasing the ambient dimension, while keeping the intrinsic dimension of the data ﬁxed.
We borrow the 2D banana-source from Ball´e et al.
(2021), and randomly embed it in R n .
As shown in Fig.
7 and Fig.
8 (in Appendix due to space constraint), our sandwich bounds on the 2D source appear tight, and closely agree with BA (similar to the results in Fig. 2); moreover, the tightness appears unaffected by the increase in ambient dimension n to 4, 16, and 100 (we veriﬁed this for n 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 = 1024 worked well for all n .
Next, we experiment on high-dimension GAN-generated images with varying intrinsic dimension, and obtain R-D sandwich bounds that help assess neural image compression methods.
Following
Pope et al.
(2021), we generate 128 × 128 images of basenji from a pre-trained GAN, and control the intrinsic dimension 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 dogs 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 neural 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 at 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.
The 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 effectiveness of learned compression methods at adapting to low-dimensional structure in data, in contrast to 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.
N ATURAL I MAGES
To establish upper bounds on the R-D function of natural images, we deﬁne variational distributions ( Q Z , Q Z | X ) on a Euclidean latent space for simplicity, and parameterize them as well as a learned decoder ω via hierarchical VAEs.
We consider two VAE architectures: 1. we borrow the convolutional autoencoder architecture of a state-of-the-art image compression model (Minnen & Singh, 2020), but use factorized Gaussians for the variational distributions (we still keep the deep factorized hyperprior, but no longer convolve it with a uniform prior); 2.
we also reuse our ResNet-VAE architecture withstochastic layers from the GAN experiments (Sec. 6.3).
We trained the models with mean-squared error (MSE) distortion and various λ on the COCO
(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 (Ball ´ e et al., 2017; Minnen et al., 2018; Minnen & Singh, 2020), for the Kodak dataset (see similar results on Tecnick in Appendix Fig. 11).
Our β -VAE version of (Minnen & Singh, 2020) ( orange ) lies on average 0.7 dB higher than the Q-R curves of the original compression model ( red ) and VTM ( green ).
With a deeper latent hierarchy, our ResNet-( β -)VAE gives a higher Q-R curve ( blue ) that is on average 1.1 dB above the state-of-the-art Q-R curves (gap shaded in cyan ).
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).
D ISCUSSIONS
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 difﬁculties 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 signiﬁcance 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 a potentially expensive block code.
Thus, while a lower bound makes a universal statement about what performance is “too good to be true”, the story is more subtle for the upper bound.
The achievability proof relies on a random coding procedure that jointly compresses multiple data samples in arbitrarily long blocks (Shannon, 1959).
When compressing at a ﬁnite block length b , R ( D ) is generally no longer achievable due to a O ( 1 √ b ) rate overhead (Kontoyiannis, 2000; Kostina & Verd´u, 2012).
Extending our work to the settings of ﬁnite block lengths or non-memoryless sources may be additional future directions.
A CKNOWLEDGEMENTS
Yibo Yang acknowledges support from the Hasso Plattner Institute at UCI.
This material is in part based upon work supported by DARPA under Contract No. HR001120C0021.
Stephan Mandt acknowledges support by the National Science Foundation (NSF) under the NSF CAREER Award 2047418; NSF Grants 1928718, 2003237 and 2007719; the Department of Energy under grant
DESC0022331, as well as Intel, Disney and Qualcomm.
Any opinions, ﬁndings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reﬂect the views of DARPA or NSF.
A A PPENDIX A. T ECHNICAL D EFINITIONS AND P REREQUISITES
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 P X is the image measure of P under X .
We suppose the source and reproduction spaces are standard Borel spaces, ( X , A X ) and ( Y , A Y ) , equipped with sigma-algebras A X and A Y , respectively.
Below we use the deﬁnitions of standard quantities from Polyanskiy & Wu (2014).
Conditional distribution The notation Q Y | X denotes an arbitrary conditional distribution (also known as a Markov kernel), i.e., it satisﬁes
For any x ∈ X , Q Y | X = x ( · ) is a probability measure on Y ; 2.
For any measurable set B ∈ A Y , x → Q Y | X = x ( B ) is a measurable function on X .
Induced joint and marginal measures Given a source distribution P X , each test channel distribution Q Y | X deﬁnes a joint distribution P X Q Y | X on the product space
X × Y (equipped with the usual product sigma algebra,
A X ×
[Equation]
) as follows:
[Equation]
for all measurable sets E ∈
A X × A Y .
The induced y -marginal distribution P Y is then deﬁned by
[Equation]
for all measurable sets ∀ B ∈
A Y .
KL Divergence We use the general deﬁnition of Kullback-Leibler (KL) divergence between two probability measures P, Q deﬁned on a common measurable space:
[Equation]
P (cid:28) Q denotes that P is absolutely continuous w.r.t.
Q (i.e., for all measurable sets E , Q ( E ) =
=
⇒
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 P X and Q Y | X , the mutual information I ( X ; Y ) is deﬁned as
[Equation]
where P Y is the Y -marginal of the joint P X Q Y | X , P X ⊗ P Y denotes the usual product measure, and
KL ( ·(cid:107)· ) is the KL divergence.
For the mutual information upper bound, it’s easy to show that
[Equation]
so the bound is tight when P Y = Q Y .
Obtaining R ( D ) through the Lagrangian.
For each λ ≥ 0 , we deﬁne the Lagrangian by incorporating the distortion constraint in the deﬁnition of R ( D ) through a linear penalty:
[Equation] [Equation]
and deﬁne its inﬁmum w.r.t.
Q Y | X by the function
[Equation]
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 (Csisz´ar, 1974b).
Deﬁne D min := inf { D (cid:48) :
R ( D (cid:48) ) < ∞} .
Since R ( D ) is convex, for each D > D min , there exists a λ ≥ 0 such that the line of slope − λ through ( D, R ( D )) is tangent to the R ( D ) curve, i.e.,
[Equation]
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, Csisz´ar (1974b); Lemma 9.7, Gray (2011)).
For every distortion tolerance
D > D min , where D min := inf { D (cid:48) :
R ( D (cid:48) ) < ∞} , it holds that
[Equation]
We can draw the following conclusions:
For each D > D min , the maximum above is attained iff λ is associated to D .
For a ﬁxed λ , if Q ∗ Y | X achieves the minimum of L ( · , λ ) , then λ is associated to the point ( ρ ( Q ∗ Y | X ) , I ( Q ∗ Y | X ) , ) ;
i
e., there exists a line with slope − λ that is tangent to the R ( D ) curve at ( ρ ( Q ∗ Y | X ) , I ( Q ∗ Y | X )) , where we deﬁned the shorthand ρ ( Q Y | X ) :
=
E [ ρ ( X, Y )] and I ( Q Y | X ) :
= I ( X ; Y ) as induced by P X Q Y
|
X
A. F ULL V ERSION OF T HEOREM 4.
Theorem A.2.
(Theorem 2.3, Csisz ´ ar (1974b); Theorem 1, Kostina (2016).)
Suppose that the following basic assumptions are satisﬁed.
R ( D ) is ﬁnite for some D , i.e., D min :
= inf { D : R ( D ) < ∞} < ∞ ; 2.
The distortion metric ρ is such that there exists a ﬁnite set E ⊂ Y such that
[Equation]
Then, for each D > D min , it holds that
[Equation]
13) where the maximization is over g ( x ) ≥ 0 and λ ≥ 0 satisfying the constraint
[Equation] [Equation]
Note: the basic assumption 2 is trivially satisﬁed 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 ﬁxing a target D value and optimizing over both λ and g as in Equation 13.
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 ﬁxing λ and only optimizing over g (or u , in our formulation).
A. T HEORETICAL R ESULTS
Theorem A.3.
(A decoder-induced R-D function bounds the source R-D function from above).
Let
X ∼ P X be a memoryless source subject to lossy compression under distortion ρ .
Let Z be any measurable space (“latent space” in a VAE), and ω : Z → Y 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 ρ ω :
[Equation]
[Equation]
Deﬁne the corresponding ω -dependent rate-distortion function
[Equation]
Then for any D ≥ 0 , R ω ( D ) ≥ R ( D ) .
Moreover, the inequality is tight if ω is bijective (so that there is “no information loss”).
Proof.
Fix D .
Take any admissible conditional distribution Q Z | X that satisﬁes E [ ρ ( X, ω ( Z ))] ≤ D in the deﬁnition of R ω ( D ) .
Deﬁne a new kernel Q Y | X between X and Y by Q Y | X = x :=
Q Z | X = x ◦ ω − 1 , ∀ x ∈ X , i.e., Q Y | X = x is the image measure of Q Z | X = x induced by ω .
Applying data processing inequality to the Markov chain X Q Z | X −−−→ Z ω −→ Y , we have I ( X ; Z ) ≥ I ( X ; Y ) .
Moreover, since Q Y | X is admissible in the deﬁnition of R ( D ) , i.e.,
[Equation]
we therefore have
[Equation]
Finally, since I ( X ; Z ) ≥ R ( D ) holds for any admissible Q Z | X , taking inﬁmum over such Q Z | X gives R ω ( D ) ≥ R ( D ) .
To prove R ω ( D )
=
R ( D ) if ω is bijective, it sufﬁces to show that R ( D ) ≥ R ω ( D ) .
We use the same argument as before.
Take any admissible Q Y | X in the deﬁnition of R ( D ) .
We can then construct a Q Z | X by the process X Q Y | X −−−−→ Y ω − 1 −−→ Z .
Then by DPI we have I ( X ; Y ) ≥
I ( X ; Z ) .
Morever, Q Z | X is admissible: E P X Q Z | X
[ ρ ( X, ω ( Z ))] =
E P X Q Y | X [ ρ ( X, ω ( ω − 1 ( Y ))]
=
E
P X Q Y | X
[ ρ ( X, Y )] ≤ D .
So I ( X ; Y ) ≥ I ( X ; Z ) ≥ R ω ( D ) .
Taking inﬁmum over such Q Y | X concludes the proof.
Remark.
Although ω − 1 ( ω ( ◦ ))
=
Identity( ◦ ) for any injective ω , our construction of Q Z | X in the last argument requires the inverse of ω to exist, so that additionally ω ( ω − 1 ( Y ))
=
Y .
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.
Corollary A.3.1.
(A suitable β -VAE deﬁnes an upper bound on the source R-D function).
Let
X ∼ P X be a data source, ρ : X × Y → [0 , ∞ ) a distortion function, and Z be any latent space.
Consider any β -VAE consisting of a prior distribution Q Z on
Z , and an encoder which speciﬁes a conditional distribution Q Z | X = x for each x .
Suppose further that the likelihood (observation) model is chosen to have density p ( x | z ) ∝ exp {− ρ ( x, ω
( z ) } for some decoder function ω : Z → Y (a common example being an isotropic Gaussian p ( x | z ) with mean ω ( z ) and ﬁxed variance, speciﬁed by a squared error distortion).
Then the two terms of the negative ELBO — speciﬁcally, the “KL term”
[Equation]
] , and the “log-likelihood term” up to a constant,
[Equation]
deﬁne a point ( D , R ) that lies above the source R ( D ) -curve; i.e., R ≥ R ( D ) .
Proof.
[Equation]
The ﬁrst inequality is the variational upper bound on mutual information (Eq. 9), the second inequality follows from the deﬁnition of the ω -induced R-D function (and the fact that E P X Q Z | X [ ρ ω ( X, Z )] =:
D ), and the last inequality is Theorem A.3.
Remark.
As also remarked by Theis et al.
(2017), not all distortion functions lead to the β -VAE interpretation above.
For this to work, the function exp {− ρ ( x, ω ( z ) } must be normalizable in x (w.r.t.
a suitable base measure on X ) for every z , and the normalizing constant may not depend on z .
Theorem A.4.
(Basic properties of the proposed estimator C k of the sup-partition function.)
Let
X 1 , X 2 , ... ∼ P X be a sequence of i.i.d.
random variables.
Let ψ : X × Y → R + be a measurable function.
For each k , deﬁne the random variable C k := sup y 1 k
(cid:80) i ψ ( X i , y ) .
Then
C k is an overestimator of the sup-partition function c , i.e., E
[Equation]
...,X k
[Equation]
sup
[Equation]
[Equation]
(cid:80)
[Equation]
, y
[Equation]
The bias of C k decreases with increasing k , i.e., E
[Equation]
If ψ ( x, y ) is bounded and continuous in y , and if Y is compact, then C k is strongly consistent, i.e., C k converges to c almost surely (as well as in probability, i.e., lim k →∞ P ( | C k − c | > (cid:15) )
= 0 , ∀ (cid:15) > 0 ), and lim k →∞ E
[ C k ] = c .
Proof.
We prove each in turn:
E [ C k ] =
E [Equation] sup
[Equation] X [Equation] , y )
[Equation]
X [Equation] , y )
[Equation] sup [Equation] [Equation]
X, y [Equation] ] = [Equation]
First, note that E [ C 1 ] ≥ E [ C k ] since
[Equation]
We therefore have
[Equation]
Deﬁne the shorthand M k ( y ) :
= 1 k (cid:80) ki =1 ψ ( X i , y ) , Ψ( y )
=
E [ ψ ( X, y )] .
Denote the sup norm on the set of continuous bounded functions on Y by (cid:107) f ( · ) (cid:107) ∞ :
= sup y ∈Y | f ( y ) | .
Then it holds that
[Equation]
where we made use of the fact that ψ ≥ 0 in the ﬁrst equality, and the reverse triangle inequality in the second-to-last step.
By the uniform strong law of large numbers (Ferguson, 2017),
[Equation]
almost surely.
Therefore, by the inequality we just showed, we also have
[Equation]
almost surely, i.e.,
[Equation]
almost surely.
Then it trivially follows that C k also converges to c in probability, and
that lim k →∞ E [ C k ] =
c by the bounded convergence theorem.
A. D ETAILED L OWER B OUND M ETHOD
Approximations for stochastic gradient ascent.
Recall the lower bound objective obtained by replacing the sup-partition function c by the mean of the k -sample underestimator C k :
[Equation]
Unfortunately, we still cannot apply stochastic gradient ascent to (cid:96) k , as two more difﬁculties remain, which we solve by further approximations.
First, C k is still hard to compute, as it is deﬁned through a global maximization problem.
We note that by restricting to a symmetric ρ and Y
= X , the maximization objective of C k has the form of a reweighted kernel density estimate,
[Equation]
where each data point is reweighted by 1 u ( x ) .
For a squared error distortion, this becomes a Gaussian mixture density, 1 k (cid:80) i ψ ( x i , y )
∝ (cid:80)
i π
i exp( − λ (cid:107) x
i − y (cid:107) 2 ) , with centroids deﬁned by the samples x 1 , ..., x k , and mixture weights π i = ( ku ( x i )) − 1 .
The global mode of a Gaussian mixture can generally be found by hill-climbing from each of the k centroids, except in rare artiﬁcial examples (Carreira-Perpinan, 2000; 2020); we therefore use this procedure to compute C k , but note that other methods exist (Lee et al., 2019; Carreira-Perpinan, 2007).
The second difﬁculty is that even if we could estimate E [ C k ] (with samples of C k computed by global optimization), the objective (cid:96) k requires an estimate of its logarithm; a naive application of Jensen’s inequality − log E [ C k ] ≥ E [ − log C k ] results in an over -estimator (as does the IWAE estimator), whereas we require a lower bound.
Following Poole et al.
(2019), we use the following basic lower bound of − log( x ) by its linearization around some point α , i.e.,
[Equation]
with equality achieved by α = x .
This results in the ﬁnal lower bound objective we optimize:
[Equation]
Since the linear underestimator of − log E [ C k ] is tightest near α = E
[ C k ] , in our algorithm we set α adaptively to an estimate of E [ C k ] (separately from the E [ C k ] term we estimate in the objective function). The algorithm.
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 deﬁnition of C γ λρ ( i y ) } u i ) .
Draw 2 m samples of C k use the last m samples to set α , and the first m samples to estimate E [ C k ] in ˜ (cid:96) k squared distortion ρ ( x, y ) ∝ (cid:107) x − y (cid:107) 2 .
opt loss :=
−∞ for i := 1 to k do //
Run gradient ascent from the i th mixture component.
y :
= x i while y not converged do with GradientTape() as tape: loss := γ k ( y, θ, λ, { x 1 , ..., x k } ) gradient :=
tape.gradient(loss, y )
y :
= y + (cid:15) gradient end if loss > opt loss then opt loss := loss y ∗ :
= y end end
C k = opt loss return ( y ∗ , C k )
For a given λ > 0 and model u θ : x → R + , the lower bound algorithm works by (stochastic) gradient ascent on the objective ˜ (cid:96) k
( θ ) (Eq.
8) w.r.t.
the model parameters θ .
To compute the gradient, we need an estimate of E [ C k ] (ultimately log E [ C k ] ) as part of the objective, which requires us to draw m samples of C k .
In the pseudo-code, a separate set of m samples of C k are also drawn to set the linear expansion point α .
In our actual implementation, we set α to an exponential moving average of E [ C k ] estimated from previous iterations (e.g., replacing line 7 of the pseudo-code by α := 0 .
2 α + 0 .
8 E ), so only m samples of C k need to be drawn for each iteration of the algorithm.
We did not ﬁnd this approximation to signiﬁcantly affect the training or results.
Recall C k ( θ ) is deﬁned by the result of maximizing w.r.t.
y .
When computing the gradient with respect to θ , we differentiate through the maximization operation by evaluating C k = γ k ( y ∗ ) on the forward pass, using the argmax y ∗ found by the subroutine compute C k .
This is justiﬁed 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 ) ,
[Equation]
We obtain the ﬁnal R L ( D ) lower bound by taking the upper envelope of all such lines, i.e.,
[Equation]
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 C k 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 C k samples in the log domain.
During training, we only run an approximate version of the global optimization subroutine compute C k for efﬁciency, as follows.
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 ﬁnal results.
A. A DDITIONAL E XPERIMENTAL D ETAILS AND R ESULTS A.5. C OMPUTATIONAL A SPECTS
Our methods are implemented using the Tensorflow library.
Our experiments with learned image compression methods additionally used the tensorflow-compression 3 library.
Our experiments on images were run on Titan RTX GPUs, while the rest of the experiments were run on Intel(R)
Xeon(R) CPUs.
We used the Adam optimizer for gradient based optimization in all our experiments, typically setting the learning rate between 1 e − 4 and 5 e − 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 ﬁnding procedure (starting hill climbing only from a small number of centroids, see Sec.
A.4), the lower bound models required less than 10000 gradient steps to converge in all our experiments.
A.5. G AUSSIANS
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-ﬁlling algorithm (Cover & Thomas, 2006).
Upper Bound Experiments For the Z =
Y (no decoder) experiment, we chose Q Y and Q Y | X to be factorized Gaussians; we let the mean and variance vectors of Q Y be trainable parameters, and predict the mean and variance of Q Y | X
= x by one fully connected layer with 2 n output units, using softplus activation for the n variance components.
For our experiments involving a decoder, we parameterize the variational distributions Q Z and Q Z | 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 Y 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.
28  32 Dimension ( n )
R ( D ) i n t e r ce p t a t s l o p e − ntrue F n ( n 2 ) ˜ l˜ l˜ l˜ l˜ l
Lower Bound Experiments In our lower bound experiments, we parameterize log u by an MLP with 2 hidden layers with 20 n hidden units each and SeLU activation, where n is the dimension of the Gaussian source.
We ﬁxed k = 1024 for the results with varying n in Figure 2a- bottom .
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 toon increasingly high dimensional standard Gaussian sources, each time setting λ = n 2 .
For an n -dimensional standard Gaussian, the true R -intercept, F n ( λ ) , has an analytical form; in particular,
F n ( n 2 ) = n 2 .
In Figure 4, we plot the ﬁnal objective estimate, ˜ (cid:96) k , using the converged MLP model, one for each k and n .
As we can see, the maximum achievable ˜ (cid:96) k plateaus to the value log k as we increase n , and for sufﬁciently 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 C k 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 “catestrophic” 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 C k (subroutine compute C k 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 x i can be approximated as γ k ( x i )
= 1 k 1 u ( x i ) + 1 k (cid:80) j (cid:54) = i exp {− λρ ( x j , x i ) } /u θ ( x j ) ≈
1 ku ( x i ) , since exp {− λρ ( x j , x i ) } ≈ 0 for all j (cid:54) = i .
The maximization problem deﬁning C k 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 C k
≈ sup i =1 ,...,k 1 ku ( x i ) .
This implies
[Equation]
therefore
[Equation]
[Equation] [Equation] Thus, when these approximations hold, the maximum achievable lower bound objective ˜ (cid:96) k cannot 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 10 6 or higher to target even the low-distortion region of the R-D curve, and the above analysis well describes the behavior of the lower bound algorithm for any value of k we feasibly experimented with.
A.5. P ARTICLE P HYSICS
Data We borrowed the Z → e + e − dataset from Howard et al.
(2021), containing 331699 data observations.
Each n = 16 -dimensional data vector corresponds to an independent Z -boson decay event, and is the concatenation of the four-momenta of an electron ( e − ) and positron ( e + ) both in the theoretical prior space and observed space.
The 2D marginal distribution (where we additionally compare with the BA algorithm) is formed by taking the two data coordinates with the least absolute value of correlation.
A histogram is given in Figure 5a.
Unlike in the Gaussian experiments, where we generate new random samples from the source on the ﬂy for training and evaluation, here the true underlying source is only approximated by a ﬁnite number of samples in the given dataset.
Since our R-D upper/lower bounds have the form of expectations of learned functions w.r.t.
the true underlying source (explained in more detail in Section.
A.6), we create a separate held-out test set of 10000 samples for our ﬁnal estimates of R-D bounds using the trained models.
We also use a small subset of the training data as a validation set for model selection, typically using the most expressive model we can afford without overﬁtting.
We use the same train/test split and the same model architectures on the 2D marginal data as on the n = 16 data.
Upper Bound Experiments For our R-D upper bound model, we use a VAE with a MLP decoder, a normalizing ﬂow “prior” Q Z , and an MLP encoder computing a factorized Gaussian Q Z | X .
We set the latent space to have the same dimensionality n as the data.
The encoder and decoder MLPs have three layers with [500 , 500 , l ] units in each, where l = n for the deocder MLP and l = 2 n for the encoder MLP (as it computes both a mean and a variance vector for Q Z | X ); we use softplus activation in each hidden layer.
Q Z is parameterized as a base Gaussian distribution transformed by three stacks of MAF (Papamakarios et al., 2017).
For the Nonlinear Transform Coding (NTC) (Ball´e et al., 2021) model, we use a VAE with a similar 500-500- n MLP architecture for the encoder and decoder (thus the latent space has the same dimensionality n as the data), but modify the variational distributions to simulate quantization and end-to-end compression.
As in (Ball ´ e et al., 2018; Ball ´ e et al., 2021), we let Q Z | X be a factorized uniform distribution with unit width, and Q Z be a factorized neural density model convolved with a uniform noise.
For the BA algorithm on the 2D marginal, we obtain discretized alphabets and source probabilities as follows.
First, we determine the smallest box set S in R 2 containing all the data samples (simply the Cartesian product of sample ranges along the x and y axes); next, we ﬁnely tile up S using small rectangle bins, and compute a histogram of data samples over them.
We then let the source and reproduction alphabets to be the bin centers, and set the source distribution to be the frequency of samples in each bin.
Finally, for computation efﬁciency, we remove the bin centers from the source alphabet associated with 0 samples; this does not affect the results.
We ran a maximum of 2000 steps of the BA algorithm on the test set, terminating early if the objective has not improved by more than 0.0001% in consecutive steps.
Lower Bound Experiments We parameterize our lower bound model log u by an MLP withhidden layers with 200 hidden units each and SeLU activation.
We ﬁxed k = 2048 .
A.5. S PEECH
Data We use the Free Spoken Digit Dataset (Jackson et al., 2018) consisting of recordings of spoken digits from ‘zero’ to ‘nine’ at a sampling rate of 8kHz.
We use a subset of the recordings with a randomly chosen speaker id (‘theo’), and keep the original train-test split 4 consisting ofand 50 audio ﬁles for the train and test sets, respectively.
From each audio ﬁle, we extracted Shorttime Fourier Transform (STFT) coefﬁcients using Tensorﬂow’s tf.signal.stft routine with a 63-sample FFT window ( frame length=63 ) in 1-sample steps ( frame step=1 ), took the magnitude and converted to log domain, a common pre-processing step for computing spectrograms [Figure] and Mel-frequency cepstral coefﬁcients.
The resulting n = 33 -dimensional Fourier feature vectors are then shufﬂed across time steps and treated as i.i.d.
samples in our experiments, with over 1 million samples in the training set and 120000 samples in the test set.
We follow the same procedure as in the physics experiment, creating a 2D marginal from the two coordinates with the least absolute correlation; a histogram is given in Figure 5b.
Experiments We use the same model architectures and experimental procedures/hypermarameters as in the physics experiment.
A.5. B ANANA - SHAPED S OURCE
Data
The 2D banana-shaped source (Ball´e et al., 2021) is a RealNVP transform of a 2D Gaussian 5 .
See Figure 6 for a plot of its density function.
We embed the banana source in a higher dimension n by multiplying with a randomly sampled n ×matrix.
This is done by simply passing samples of the 2D banana source through a linear 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 We reproduced the NTC experiment using the same compressive autoencoder as in Ball ´ e et al.
(2021), with 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.
For our upper bound model, we use the same MLP architecture as in Ball´e et al.
but parameterize Q Z by a MAF (Papamakarios et al., 2017) and Q Z | X by a factorized Gaussian (doubling the number of output of the encoder), similar to on previous experiments.
For the higher-dimension embeddings 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 to 100 n , 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 ﬁnd 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 100 n , and cap it at 1000.
In all the experiments we set k = 1024 .
[Figure] [Figure] As illustrated by Figure the tightness of our sandwich bounds do not appear to be affected by the dimension n of the data (we veriﬁed this up to n = 1000 ), unlike on the Gaussian experiment where for a ﬁxed k the lower bound became increasingly loose with higher n .
A
GAN- GENERATED I MAGES
Data We adopt the same setup as in the GAN experiment by Pope et al.
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 ﬁrst 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
X =
Y =
[0 , 1] n .
As in (Pope et al., 2021), we experimented with images from the ImageNet class basenji .
In Figure 9 and 10, 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 [Figure]
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 Z 0 at the topmost generative hierarchy is ﬂattened and modeled by a MAF prior Q Z 0 (aided by the fact that all the images have a ﬁxed 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 ﬁlters, 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.
A.5. N ATURAL I MAGES
Data For signals such as images or video, which can have arbitrarily high spatial dimension, it is not immediately clear how to deﬁne 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 (Ball ´ e et al., 2017), noting that current methods cannot effectively exploit spatial redundancies larger than this scale (Minnen & Singh, 2020).
We can then use the trained models to estimate R-D upper bounds on the imagegenerating distributions underlying standard benchmark datasets such as Kodak Kodak (1993) and
Tecnick Asuni & Giachetti (2014).
As a representative dataset of natural images, we used images from the COCO
(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 Y =
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:
The learned image compression model from Minnen & Singh (2020) 6 .
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 Q Z | X (instead of factorized uniform distributions used to simulate rounding in the original model), and no longer convolving the variational prior Q Z with a uniform distribution as in the original model.
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 ﬂattening 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 Q Z 0 as proposed by Ball´e et al.
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 ﬁnal 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-compression 7 and compressAI 8 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 ˜ (cid:96) k 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.
A. M EASURES OF V ARIABILITY
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 ﬁrst deﬁne 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 deﬁne the random variables from the rate and distortion terms of the
Lagrangian:
[Equation]
and
[Equation]
Above, X and Z follow the joint distribution P X Q Z | X , ω is the decoder function, and q ( z | x ) and q ( z ) are the Lebesgue density functions of absolutely continuous Q Z | X = x and Q Z 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.,
[Equation]
and
[Equation]
Recall for any given decoder and variational distributions, the above rate and distortion terms deﬁne a point ( E [ D ] , E
[ R ]) that lies on an upper bound R U ( D ) of the source R ( D ) curve.
Given m i.i.d.
https results [Table] [Table] pairs of ( Z 1 , X 1 ) , ( Z 2 , X 2 ) , ..., ( Z m , X m ) ∼ P X Q Z | X , we estimate such a point by ( µ D , µ R ) , using the usual sample mean estimators µ R := 1 m (cid:80) mj
=1 R ( Z j , X j ) and µ D := 1 m (cid:80) mj
=1 D ( Z j , X j ) .
Lower Bound Estimator For the proposed lower bound, given i.i.d.
X 1 , ..., X k ∼ P X , deﬁne the random variable
[Equation]
where the various quantities are deﬁned as in Sec. 4.
Then it is clear that its expectation equals the lower bound objective Eq. 8, i.e., E [ ξ ] =
˜ (cid:96) k :=
E [ − log u
( X )] − E
[ C k ] /α − log α + 1 .
Given m i.i.d.
samples of ξ , we form the usual sample mean estimator µ ξ := 1 m (cid:80) j ξ j , and form an R-D lower bound estimate by
[Equation]
In computing µ ξ , we ﬁx the constant α to a sample mean estimate of E [ C k ] 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.3).
We report detailed statistics in the following tables, and plot them in Figure 12.
Our results on other experiments have comparable 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 conﬁdence interval for the population mean (the true expected value) for R 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 conﬁdence lower bound for the population mean (the true expected value E [ ξ ] =
˜ (cid:96) k ) in Tables 3 and 4.
[Table]
[Table]
[Figure]
[Table]
[Figure]
[Table]
R a t e ( n a t s p e r s a m p l e ) d =
Minnen 2018 (hyperprior) Minnen95 % conf.
interval of proposed ˆ R U ( D ) proposed ˆ R L ( D ) 90% conf.
lower bound of proposed ˆ R L ( D )