As we saw in part 1 of this post, given a data (source) distribution \(P_X\) and a distortion function \(\rho\), the ratedistortion (RD) function of the data quantifies its compressibility. It gives us an idea about the best possible compression performance we can expect, and can help us design better compression algorithms and data communication systems in general.
In this part 2 of the post, I'll introduce a simple method to estimate the RD function, i.e., to train a family of \(\beta\)VAElike models. This method is much more scalable and widelyapplicable than the classic BlahutArimoto algorithm, and works in an online fashion as it is based on stochastic gradient descent (SGD). Perhaps not coincidentally, a subset of these \(\beta\)VAEs are also the basis of many learned lossy compression models today. You can find more details in the ICLR2022 paper, and a quick demo for how this all works in this ipython notebook.
Estimating the RD function.
Recall the definintion of the ratedistortion function from last time:
To be able to carry out the integrals, we need knowledge of \(P_X\), i.e., the probability that the distribution assigns to any event. For a Bernoulli or Gaussian source, the RD function can be worked out mathematically. Unfortunately, outside of these special cases, the RD function has no analytical form, and we have little knowledge of the RD function of most data sources.
If the source and reproduction alphabets are finite (thus necessarily discrete), and if we know the source probability mass function \(P_X\), then the BlahutArimoto (BA) algorithm can numerically compute the RD function by solving the mathematical optimization problem defining it. However, there are two major difficulties with establishing the RD function of realworld sources, which make the BA algorithm inapplicable:
1. the source distribution (e.g., the distribution of natural images) is almost always unknown; ^{1}
2. the data and the lossy reproductions may be continuously valued, e.g., in scientific applications.
The availability of massive amounts of data prompts us to consider estimating the RD function from data samples. Indeed, Harrison & Kontoyiannis (2008) proved that that in many general settings, the RD estimator from replacing the true data distribution by an empirical distribution of samples will asymptotically recover the true RD function. Thus, one might try running the BA algorithm on an empirical distribution of data samples, after discretizing the source and reproduction alphabets if they are continuous. However, besides the inaccuracy from discretization, this approach also quickly runs into computational difficulties. In higher dimensions, an exceedingly large number of samples is required to faithfully represent the source distribution. This in turn creates a blowup in the number of discretized states in the source and reproduction alphabets. As the BA algorithm operates on tabular representations of the various quantities over all the possible discrete states (e.g., the source PMF \(p(x_i)\), and the matrix of pairwise distortion values \(\rho(x_i, y_j)\), where \(x_i\) and \(y_j\) range over all values of source and reproduction alphabets), the memory requirement of BA alone can quickly become infeasible in just a handful of dimensions.
Alternative approach based on an RDVAE.
Now I'll introduce an alternative to the BA algorithm, which works with general alphabets, and can be viewed as an SGD version of the BA algorithm. Besides being much more scalable, this method can also free us from working with a fixed (possibly discretized) empirical data distribution: if a stream of i.i.d. samples from the source is available, we will be able to continually improve our estimate of the RD function in an online fashion, overcoming the finitesample bias of the BA algorithm.
The basic idea is straightforward, and may appear unrelated at first sight: fitting a VAE to the data source. For now, let's use a common Gaussian likelihood model in the VAE, so that \(p(xz) \sim \mathcal{N}(x\mu(z), \sigma^2(z))\), with the mean and variance parameters computed by a decoder neural network. Denoting the prior and posterior densities by \(q(z)\) and \(q(zx)\), the standard maximum likelihood training objective is the aggregate ELBO (evidence lower bound):
As I wrote in a previous note, fitting such a density may be problematic if the true data distribution does not have full support over its ambient space. For instance, the true distribution of MNIST images has no probability mass outside of a collection of integervalued points in \(\mathbb{R}^{784}\). In this case, the maximum likelihood training objective is illdefined (indeed, it seems silly to fit a density model to discrete data), and may reach arbitrarily high values.
One way to avoid this nonsensical scenario is to restrict, or regularize the density model in some way, so that the training objective has a finite maximum. For the Gaussian VAE, it is common to fix the variance of the likelihood model to a constant, like \(\sigma^2 I\), and only learn its mean. The authors of this paper did just that: they fit a series of Gaussian VAE models to MNIST digits, keeping \(\sigma^2\) fixed for each model, and varying it across models; they observe that models with smaller \(\sigma^2\) achieve better/higher aggregate ELBO and exhibit less posterior collapse.
How should we interpret the aggregate ELBO in this model, which could be unbounded had we not fixed \(\sigma^2\)? And how should we set \(\sigma^2\)? Similar questions can framed in terms of \(\beta\)VAEs; indeed, since the loglikelihood has the form \(\log p(xz) = \frac{1}{2\sigma^2} \x  \mu(z) \^2 + const\), the VAE with a fixed variance is essentially equivalent to a \(\beta\)VAE with a unit variance in the likelihood model, with \(2\sigma^2\) playing the same role as \(\beta\).
The answer: this \(\beta\)VAE is optimizing an upper bound on the RD function of the data distribution, under a suitable distortion metric. The key is to look at the aggregate ELBO objective just the right way, and identify a correspondence between the likelihood model and the distortion function. A Gaussian \(\beta\)VAE, for example, gives us an estimate of the compressibility of the data samples under a squared error distortion. A \(\beta\)VAE with a Laplace likelihood model then corresponds to an \(L_1\) distortion, \(\rho(x,y) = \x  y \_1\).
The general statement is something like this (see Theorem A.3 in the paper for details): Let \(q(z), q(zx)\) be arbitrary choices of prior and posterior distributions in a VAE. Suppose the likelihood density \(p(xz)\) has the form of an exponentiated distortion,
where \(\omega\) is an arbitrary decoder function, and the normalizing constant \(const = \int \exp\{\rho(x, \omega(z))\} d x\) does not depend on \(z\) and is allowed to be infinite. Then, the point \((\mathcal{D}, \mathcal{R})\) formed by the two terms of the negative ELBO,
and
lies on an upper bound of the ratedistortion function of the source under distortion function \(\rho\). Moreover, by optimizing the objective \(\mathcal{R} + \lambda \mathcal{D}\) (or, equivalently in the \(\beta\)VAE formulation, \(\beta \mathcal{R} + \mathcal{D}\)) over all possible choices of the prior, posterior, and decoder, the resulting point \((\mathcal{D}, \mathcal{R})\) will lie exactly on the source RD function. All of this pertains to a fixed \(\lambda\) value; the optimal \((\mathcal{D}, \mathcal{R})\) pairs resulting from different settings of \(\lambda\) then trace out the full \(R(D)\) curve.
I refer to a \(\beta\)VAElike model satisfying the above condition on the likelihoodmodel an RD (ratedistortion) VAE. The rate and distortion terms purposefully mirror the quantities in the definition of the RD function, and it's not hard to prove the above result by noting that \(\mathcal{R}\) is an upper bound on the mutual information \(I(X; Z)\), which in turn is an upper bound on \(I(X; Y), Y:= \omega(Z)\) by data processing inequality. The above result essentially allows us to train an RDVAE on a NELBOlike objective to optimize an upper bound on the source RD function. The lossless counterpart of this result is that when the data is discrete, we can train a standard VAE with a discrete likelihood by optimizing the standard negative ELBO, which upper bounds the entropy of the source (this is explained towards the end of my note on ELBO decompositions).
Thus, we can think of an RDVAE as a theoretical data compressor. Compared to the lossless case, where a single algorithm, bitsback coding, has made it possible (and in many cases practical) to do lossless compression with generic VAEs and achieve compression rates close to the negative ELBO, the lossy case appears more challenging, and research is still ongoing to develop practical compression algorithms that realize the theoretical performance of RDVAEs, often with restrictions on the "prior" and "posterior" distributions to make things practical. You can learn more about these topics in Section 3.4 ("Compression without quantization") of my intro article to neural data compression with Stephan Mandt and Lucas Theis.
The need for a lower bound.
Although training an RDVAE is straightforward with deep learning libraries like Tensorflow or PyTorch, at the end of the day this only gives us an upper bound on the RD function. We know that in theory, this upper bound becomes arbitrarily tight as we increase the expressive power of the RDVAEs, but deeper and bigger neural networks may also mean more local minima in the SGD optimization landscape, and not necessarily tighter bounds in practice.
To verify the tightness of an upper bound (and ultimately, to decide if an algorithm's RD performance is close to theoretically optimal), we also need to estimate a lower bound on the RD function. This the second half of the "sandwich bounds" proposed in the paper. I won't go into the details as they are a bit involved, but the basic idea is similar  optimizing over functions parameterized by neural networks with SGD. It's really amazing how well SGD scales to largescale datasets in the RDVAE upper bound method (afterall, it's just training a VAE), and by comparison, how much more challenging it appears to solve a largescale optimization problem when SGD cannot be directly applied. Indeed, the bulk of my technical work on the lower bound went into reformulating the problem to make it suitable for stochastic optimization.

In some cases where we are in full control of the data generation process, we might be able to assert knowledge of the true data distribution, e.g., the toss of a fair coin is, for all intents and purposes, Bernoulli(0.5); but these cases tend to be rare or toyish. ↩