Latent-variable models have become a mainstay of modern machine learning. Scalable approximate Bayesian inference methods, in particular Black Box Variational Inference (Ranganath et al., 2014;Rezende et al., 2014), have spurred the development of increasingly large and expressive probabilistic models, including deep generative probabilistic models such as variational autoencoders (Kingma & Welling, 2014b) and Bayesian neural networks (MacKay, 1992;Blundell et al., 2015). One natural application of deep latent variable modeling is data compression, and recent work has focused on end-to-end procedures that optimize a model for a particular compression objective. Here, we study a related but different problem: given a trained model, what is the best way to encode the information contained in its continuous latent variables?
As we show, our proposed solution provides a new "plugand-play" approach to lossy compression of both data instances (represented by local latent variables, e.g., in a VAE) as well as model parameters (represented by global latent variables that serve as parameters of a Bayesian statistical model). Our approach separates the compression task from model design and training, thus implementing variablebitrate compression as an independent post-processing step in a wide class of existing latent variable models.
At the heart of our proposed method lies a novel quantization scheme that optimizes a rate-distortion trade-off by exploiting posterior uncertainty estimates. Quantization is central to lossy compression, as continuous-valued data like natural images, videos, and distributed representations ultimately need to be discretized to a finite number of bits for digital storage or transmission. Lossy compression algorithms therefore typically find a discrete approximation of some semantic representation of the data, which is then encoded with a lossless compression method.
In classical lossy compression methods such as JPEG or MP3, the semantic representation is carefully designed to support compression at variable bitrates. By contrast, stateof-the-art deep learning based approaches to lossy data compression (Ballé et al., 2017;2018;Rippel & Bourdev, 2017;Mentzer et al., 2018;Lombardo et al., 2019) are trained to minimize a distortion metric at a fixed bitrate. To support variable-bitrate compression in this approach, one has to train several models for different bitrates. While training several models may be viable in many cases, a bigger issue is the increase in decoder size as the decoder has to store the parameters of not one but several deep neural networks for each bitrate setting. In applications like video streaming under fluctuating connectivity, the decoder further has to load a new deep learning model into memory every time a change in bandwidth requires adjusting the bitrate.
By contrast, we propose a a quantization method for latent variable models that decouples training from compression, and that enables variable-bitrate compression with a single model. We generalize a classical entropy coding algorithm, Arithmetic Coding (Witten et al., 1987;MacKay, 2003), from the discrete to continuous domain. Our proposed algorithm, Variational Bayesian Quantization, exploits posterior uncertainty estimates to automatically reduce the quantization accuracy of latent variables for which the model is uncertain anyway. This strategy is analogous to the way humans communicate quantitative information. For example, Wikipedia lists the population of Rome in 2017 with the specific number 2,879,728. By contrast, its population in the year 500 AD is estimated by the round number 100,000 because the high uncertainty would make a more precise number meaningless. Our ablation studies show that this posterior-informed quantization scheme is crucial to obtaining competitive performance.
In detail, our contributions are as follows:
• A new discretization scheme. We present a novel approach to discretizing latent variables in a variational inference framework. Our approach generalizes arithmetic coding from discrete to continuous distributions and takes posterior uncertainty into account.
• Single-model compression at variable bitrates. The decoupling of modeling and compression allows us to adjust the trade-off between bitrate and distortion in post-processing. This is in contrast to existing approaches to both data and model compression, which often require specially trained models for each bitrate.
• Automatic self-pruning. Deep latent variable models often exhibit posterior collapse, i.e., the variational posterior collapses to the model prior. In our approach, latent dimensions with collapsed posteriors require close to zero bits, thus don't require manual pruning.
• Competitive experimental performance. We show that our method outperforms JPEG over a wide range of bitrates using only a single model. We also show that we can successfully compress word embeddings with minimal loss, as evaluated on semantic reasoning task.
The paper is structured as follows: Section 2 reviews related work in neural compression; Section 3 proposes our Variational Bayesian Quantization algorithm. We give empirical results in Section 4, and conclude in Section 5. Section 6 provides additional theoretical insight about our method.
Compressing continuous-valued data is a classical problem in the signal processing community. Typically, a distortion measure (often the squared error) and a source distribution are assumed, and the goal is to design a quantizer that optimizes the rate-distortion (R-D) performance (Lloyd, 1982;Berger, 1972;Chou et al., 1989). Optimal vector quantization, although theoretically well-motivated (Gallager, 1968), is not tractable in high-dimensional spaces (Gersho & Gray, 2012) and not scalable in practice. Therefore most classical lossy compression algorithms map data to a suitably designed semantic representation, in such a way that coordinate-wise scalar quantization can be fruitfully applied.
Recent machine-learning-based data compression methods learn such hand-designed representation from data, but similar to classical methods, most such ML methods directly take quantization into account in the generative model design or training. Various approaches have been proposed to approximate the non-differentiable quantization operation during training, such as stochastic binarization (Toderici et al., 2016;2017), additive uniform noise (Ballé et al., 2017;2018;Habibian et al., 2019), or other differentiable approximation (Agustsson et al., 2017;Theis et al., 2017;Mentzer et al., 2018;Rippel & Bourdev, 2017); many such schemes result in quantization with a uniformly-spaced grid, with the exception of (Agustsson et al., 2017), which optimizes for quantization grid points. Yang et al. (2020) considers optimal quantization at compression time, but assumes a fixed quantization scheme of (Ballé et al., 2017) during training.
We depart from such approaches by treating quantization as a post-processing step decoupled from model design and training. Crucial to our approach is a new quantization scheme that automatically adapts to different length scales in the representation space based on posterior uncertainty estimates. To our best knowledge, the only prior work that uses posterior uncertainty for compression is in the context of bits-back coding (Honkela & Valpola, 2004;Townsend et al., 2019), but these works focus on lossless compression, with the recent exception of (Yang et al., 2020).
Most existing neural image compression methods require training a separate machine learning model for each desired bitrate setting (Ballé et al., 2017;2018;Mentzer et al., 2018;Theis et al., 2017;Lombardo et al., 2019). In fact, Alemi et al. (2018) showed that any particular fitted VAE model only targets one specific point on the rate-distortion curve.
Our approach has the same benefit of variable-bitrate singlemodel compression as methods based on recurrent VAEs (Gregor et al., 2016;Toderici et al., 2016;2017;Johnston et al., 2018); but unlike these methods, which use dedicated model architecture for progressive image reconstruction, we instead focus more broadly on quantizing latent representations in a given generative model, designed and trained for specific application purposes (possibly other than compression, e.g., modeling complex scientific observations).
We now propose an algorithm for quantizing latent variables in trained models. After describing the problem setup and assumptions (Subsection 3.1), we briefly review Arithmetic Coding (Subection 3.2). Subsection 3.3 describes our proposed lossy compression algorithm, which generalizes Arithmetic Coding to the continuous domain.
Generative Model and Variational Inference. We consider a wide class of generative probabilistic models with data x and unknown (or "latent") variables z ∈ R K from some continuous latent space with dimension K. The generative model is defined by a joint probability distribution,
with a prior p(z) and a likelihood p(x|z). Although our presentation focuses on unsupervised representation learning, our framework also captures the supervised setup.foot_0
Our proposed compression method uses z as a proxy to describe the data x. This requires "solving" Eq. 1 for z given x, i.e., inferring the posterior p(z|x) = p(x, z)/ p(x, z) dz.
Since exact Bayesian inference is often intractable, we resort to Variational Inference (VI) (Jordan et al., 1999;Blei et al., 2017;Zhang et al., 2019), which approximates the posterior by a so-called variational distribution q φ (z|x) by minimizing the Kullback-Leibler divergence D KL (q φ (z|x) || p(z|x)) over a set of variational parameters φ.
Factorization Assumptions. We assume that both the prior p(z) and the variational distribution q φ (z|x) are fully factorized (mean-field assumption). For concreteness, our examples use a Gaussian variational distribution. Thus,
(2)
where p(z i ) is a prior for the i th component of z, and the means µ i and standard deviations σ i together comprise the variational parameters φ over which VI optimizes.foot_1
Prominently, the model class defined by Eqs. 1-3 includes variational autoencoders (VAEs) (Kingma & Welling, 2014a) for data compression, but we stress that the class is much wider, capturing also Bayesian neural nets (MacKay, 2003), probabilistic word embeddings (Barkan, 2017;Bamler & Mandt, 2017), matrix factorization (Mnih & Salakhutdinov, 2008), and topic models (Blei et al., 2003).
Protocol Overview. We consider two parties in communication, a sender and a receiver. In the case of data compression, both parties have access to the model, but only the sender has access to the data point x, which it uses to fit a variational distribution q φ (z|x). It then uses the algorithm proposed below to select a latent variable vector ẑ that has high probability under q φ , and that can be encoded into a compressed bitstring, which gets transmitted to the receiver. The receiver losslessly decodes the compressed bitstring back into ẑ and uses the likelihood p(x|ẑ) to generate a reconstructed data point x, typically setting x = arg max x p(x|ẑ). In the case of model compression, the sender infers a distribution q φ (z|x) over model parameters z given training data x, and uses our algorithm to select a suitable vector ẑ of quantized model parameters. The receiver receives ẑ and uses it to reconstruct the model.
The rest of this section describes how the proposed algorithm selects ẑ and encodes it into a compressed bitstring.
Our quantization algorithm, introduced in Section 3.3 below, is inspired by a lossless compression algorithm, arithmetic coding (AC) (Witten et al., 1987;MacKay, 2003), which we generalize from discrete data to the continuous space of latent variables z ∈ R K . To get there, we first review the main idea of AC that our proposed algorithm borrows.
AC is an instance of so-called entropy coding. It uniquely maps messages m ∈ M from a discrete set M to a compressed bitstring of some length R m (the "bitrate"). Entropy coding exploits prior knowledge of the distribution p(m) of messages to map probable messages to short bitstrings while spending more bits on improbable messages. This way, entropy coding algorithms aim to minimize the expected rate
For lossless compression, the expected rate has a fundamental lower bound, the entroy
where h(m) =log 2 p(m) is the Shannon information content of m. AC provides near optimal lossless compression as it maps each message m ∈ M to a bitstring of length R m = ⌈h(m)⌉, where ⌈•⌉ denotes the ceiling function.
AC is usually discussed in the context of streaming compression where m is a sequence of symbols from a finite alphabet, as AC improves on this task over the more widely known Huffman coding (Huffman, 1952). In our work, we focus on a different aspect of AC: its use of a cumulative probability distribution function to map a nonuniformly distributed random variable m ∼ p(m) to a number ξ that is nearly uniformly distributed over the interval [0, 1).
decreasing bitrate, increasing distortion MNIST Frey Faces
While Section 4.1 demonstrated the proposed VBQ method for model compression, we now apply the same method to data compression using a variational autoencoder (VAE).
We first provide qualitative insight on small-scale images, and then quantitative results on full resolution color images.
Model. For simplicity, we consider regular VAEs with a standard normal prior and Gaussian variational posterior.
The generative network parameterizes a factorized categorical or Gaussian likelihood model in experiments in Sec. 4.2.1 or 4.2.2, respectively. Network architectures are described below and in more detail in Supplementary Material.
Baselines. We consider the following baselines:
• Uniform quantization: for a given image x, we quantize each dimension of the posterior mean vector µ(x) to a uniform grid. We report the bitrate for encoding the resulting quantized latent representation via standard entropy coding (e.g., arithmetic coding). Entropy coding requires prior knowledge of the probabilities of each grid point. Here, we use the empirical frequencies of grid points over a subset of the training set;
• k-means quantization: similar to "uniform quantization", but with the placement of grid points optimized via k-means clustering on a subset of the training data;
• Quantization with generalized Lloyd algorithm: similar to above, but the grid points are optimized using generalized Lloyd algorithm (Chou et al., 1989), a widelyused state-of-the-art classical quantization method;
• JPEG: we used the libjpeg implementation packaged with the Python Pillow library, using default configurations (e.g., 4:2:0 subsampling), and we adjust the quality parameter to vary the rate-distortion trade-off; • Deep learning baseline: we compare to Ballé et al. (2017), who directly optimized for the rate and distortion, training a separate model for each point on the R-D curve. In our large-scale experiment, we adopte their model architecture, so their performance essentially represents the end-to-end optimized performance upper bound for our method (which uses a single model).
We trained a VAE on the MNIST dataset and the Frey Faces dataset, using 5 and 4-dimensional latent spaces, respectively. See Supplemental Material for experimental details.
Figure 4 shows example image reconstructions from our VBQ algorithm with increasing λ, and thus decreasing bitrate. The right-most column is the extreme case λ → ∞, resulting in the shortest possible bistring encoding ξi = (0.1) 2 = 1 2 (i.e., ẑi being the median of the prior p(z i )) for every dimension i. As the bitrate decreases (as R( ξ) → 0), our method gradually "confuses" the original image with a generic image (roughly in the center of the embedding space), while preserving approximately the same level of sharpness. This is in contrast to JPEG which typically introduces blocky and/or pixel-level artifacts at lower bitrates.
We apply our VBQ method to a VAE trained on color images, and obtain practical image compression performance rivaling JPEG, while outperforming baselines that ignore posterior uncertainty and directly quantize latent variables.
Model and Dataset. The inference and generative networks of our VAE are identical to the analysis and synthesis networks of Ballé et al. (2017), using 3 layers of 256 filters each in a convolutional architecture. We used a diagonal Gaussian likelihood model, whose mean is computed by the generative net and the variance σ 2 is fixed as a hyper-parameter, similar to a β-VAE (Higgins et al., 2017) approach (σ 2 was tuned to 0.001 to ensure the VAE achieved overall good R-D trade-off; see (Alemi et al., 2018)). We trained the model on the same subset of the ImageNet dataset as used in (Ballé et al., 2017). We evaluated performance on the standard Kodak (Kodak) dataset, a separate set of 24 uncompressed color images. As in the word embedding experiment, we also observed that using an empirical prior for our method improved the bitrate; for this, we used the flexible density model of Ballé et al. (2018), fitting a different distribution per latent channel, on samples of posterior means µ (treating spatial dimensions as i.i.d.).
For supervised learning with labels y, we would consider a conditional generative model p(y, z|x) = p(y|z, x) p(z) with conditional likelihood p(y|z, x), where z are the model parameters, treated as a Bayesian latent variable with associated prior p(z).
These parameters are often amortized by a neural network (in which case µi and σi depend on x), but don't have to (in which case µi and σi do not depend on x and are directly optimized).