V ARIATIONAL IMAGE COMPRESSION WITH A SCALE HYPERPRIOR
Anonymous authors
Paper under double-blind review
A BSTRACT
We describe an end-to-end trainable model for image compression based on variational autoencoders.
The model incorporates a hyperprior to effectively capture spatial dependencies in the latent representation.
This hyperprior relates to side information, a concept universal to virtually all modern image codecs, but largely unexplored in image compression using artiﬁcial neural networks (ANNs).
Unlike existing autoencoder compression methods, our model trains a complex prior jointly with the underlying autoencoder.
We demonstrate that this model leads to state-of-the-art image compression when measuring visual quality using the popular MS-SSIM index, and yields rate–distortion performance surpassing published ANN-based methods when evaluated using a more traditional metric based on squared error (PSNR). Furthermore, we provide a qualitative comparison of models trained for different distortion metrics.
I NTRODUCTION
Recent machine learning methods for lossy image compression have generated signiﬁcant interest in both the machine learning and image processing communities (e.g., Ballé et al., 2017; Theis et al., 2017; Toderici et al., 2017; Rippel and Bourdev, 2017).
Like all lossy compression methods, they operate on a simple principle: an image, typically modeled as a vector of pixel intensities x , is quantized, reducing the amount of information required to store or transmit it, but introducing error at the same time. Typically, it is not the pixel intensitites that are quantized directly. Rather, an alternative (latent) representation of the image is found, a vector in some other space y , and quantization takes place in this representation, yielding a discrete-valued vector ˆ y . Because it is discrete, it can be losslessly compressed using entropy coding methods, such as arithmetic coding (Rissanen and Langdon, 1981), to create a bitstream which is sent over the channel.
Entropy coding relies on a prior probability model of the quantized representation, which is known to both encoder and decoder (the entropy model ). In the class of ANN-based methods for image compression mentioned above, the entropy model used to compress the latent representation is typically represented as a joint, or even fully factorized, distribution p ˆ y ( ˆ y ) . Note that we need to distinguish between the actual marginal distribution of the latent representation m ( ˆ y ) , and the entropy model p ˆ y ( ˆ y ) .
While the entropy model is typically assumed to have some parametric form, with parameters ﬁtted to the data, the marginal is an unknown distribution arising from both the distribution of images that are encoded, and the method which is used to infer the alternative representation y . The smallest average code length an encoder–decoder pair can achieve, using p ˆ y as their shared entropy model, is given by the Shannon cross entropy between the two distributions:
[Equation]
Note that this entropy is minimized if the model distribution is identical to the marginal. This implies that, for instance, using a fully factorized entropy model, when statistical dependencies exist in the actual distribution of the latent representation, will lead to suboptimal compression performance.
One way conventional compression methods increase their compression performance is by transmitting side information : additional bits of information sent from the encoder to the decoder, which signal modiﬁcations to the entropy model intended to reduce the mismatch. This is feasible because the marginal for a particular image typically varies signiﬁcantly from the marginal for the ensemble of images the compression model was designed for. In this scheme, the hope is that the amount of [Figure]
x [Figure] a U | Q g s ˜ x | ˆ x y ˜ y
| ˆ y
[Figure]
side information sent is smaller, on average, than the reduction of code length achieved in eq.
(1) by matching p ˆ y more closely to the marginal for a particular image.
For instance, JPEG (1992) models images as independent ﬁxed-size blocks of 8 × 8 pixels.
However, some image structure, such as large homogeneous regions, can be more efﬁciently represented by considering larger blocks at a time.
For this reason, more recent methods such as HEVC (2013) partition an image into variablesize blocks, convey the partition structure to the decoder as side information, and then compress the block representations using that partitioning.
That is, the entropy model for JPEG is always factorized into groups of 64 elements, whereas the factorization is variable for HEVC.
The HEVC decoder needs to decode the side information ﬁrst, so that it can use the correct entropy model to decode the block representations.
Since the encoder is free to select a partitioning that optimizes the entropy model for each image, this scheme can be used to achieve more efﬁcient compression.
In conventional compression methods, the structure of this side information is hand-designed.
In contrast, the model we present in this paper essentially learns a latent representation of the entropy model, in the same way that the underlying compression model learns a representation of the image.
Because our model is optimized end-to-end, it minimizes the total expected code length by learning to balance the amount of side information with the expected improvement of the entropy model.
This is done by expressing the problem formally in terms of variational autoencoders (VAEs), probabilistic generative models augmented with approximate inference models (Kingma and Welling, 2014).
Ballé et al.
(2017) and Theis et al.
(2017) previously noted that some autoencoder-based compression methods are formally equivalent to VAEs, where the entropy model, as described above, corresponds to the prior on the latent representation.
Here, we use this formalism to show that side information can be viewed as a prior on the parameters of the entropy model, making them hyperpriors of the latent representation.
Speciﬁcally, we extend the model presented in Ballé et al.
(2017), which has a fully factorized prior, with a hyperprior that captures the fact that spatially neighboring elements of the latent representation tend to vary together in their scales.
We demonstrate that the extended model leads to state-ofthe-art image compression performance when measured using the MS-SSIM quality index (Wang,
Simoncelli, et al., 2003).
Furthermore, it provides signiﬁcantly better rate–distortion performance compared to other ANN-based methods when measured using peak signal-to-noise ratio (PSNR), a metric based on mean squared error.
Finally, we present a qualitative comparison of the effects of training the same model class using different distortion losses.
C OMPRESSION WITH V ARIATIONAL M ODELS
In the transform coding approach to image compression (Goyal, 2001), the encoder transforms the image vector x using a parametric analysis transform g a ( x ; φ g ) into a latent representation y , which is then quantized to form ˆ y .
Because ˆ y is discrete-valued, it can be losslessly compressed using entropy coding techniques such as arithmetic coding (Rissanen and Langdon, 1981) and transmitted as a sequence of bits.
On the other side, the decoder recovers ˆ y from the compressed signal, and subjects it to a parametric synthesis transform g s ( ˆ y ; θ g ) to recover the reconstructed image ˆ x .
In the context of this paper, we think of the transforms g a
and g s as generic parameterized functions, such as artiﬁcial neural networks (ANNs), rather than linear transforms as in traditional compression methods.
The parameters θ g and φ g then encapsulate the weights of the neurons, etc.
(refer to section 4 for details).
The quantization introduces error, which is tolerated in the context of lossy compression, giving rise to a rate–distortion optimization problem.
Rate is the expected code length (bit rate) of the compressed representation: assuming the entropy coding technique is operating efﬁciently, this can again be written as a cross entropy: [Equation] [Equation] ) where Q represents the quantization function, and p ˆ
y is the entropy model, as described in the introduction.
In this context, the marginal distribution of the latent representation arises from the (unknown) image distribution p x and the properties of the analysis transform.
Distortion is the expected difference between the reconstruction ˆ x and the original image x , as measured by a norm or perceptual metric.
The coarseness of the quantization, or alternatively, the warping of the representation implied by the analysis and synthesis transforms, affects both rate and distortion, leading to a trade-off, where a higher rate allows for a lower distortion, and vice versa.
Various compression methods can be viewed as minimizing a weighted sum of these two quantities.
Formally, we can parameterize the problem by λ , a weight on the distortion term.
Different applications require different trade-offs, and hence different values of λ .
In order to be able to use gradient descent methods to optimize the performance of the model over the parameters of the transforms ( θ g and φ g ), the problem needs to be relaxed, because due to the quantization, gradients with respect to φ g are zero almost everywhere.
Approximations that have been investigated include substituting the gradient of the quantizer (Theis et al., 2017), and substituting additive uniform noise for the quantizer itself during training (Ballé et al., 2016b).
Here, we follow the latter method, which switches back to actual quantization when applying the model as a compression method.
We denote the quantities derived from this approximation with a tilde, as opposed to a hat; for instance, ˜ y represents the “noisy” representation, and ˆ y the quantized representation.
The optimization problem can be formally represented as a variational autoencoder (Kingma and
Welling, 2014); that is, a probabilistic generative model of the image combined with an approximate inference model (ﬁgure 1).
The synthesis transform is linked to the generative model (“generating” a reconstructed image from the latent representation), and the analysis transform to the inference model (“inferring” the latent representation from the source image).
In variational inference, the goal is to approximate the true posterior p ˜ y
| x ( ˜ y | x ) , which is assumed intractable, with a parametric variational density q ( ˜ y | x ) by minimizing the expectation of their Kullback–Leibler (KL) divergence over the data distribution p x :
[Equation]
By matching the parametric density functions to the transform coding framework, we can appreciate that the minimization of the KL divergence is equivalent to optimizing the compression model for rate–distortion performance.
We have indicated here that the ﬁrst term will evaluate to zero, and the second and third term correspond to the weighted distortion and the bit rate, respectively.
Let’s take a closer look at each of the terms.
First, the mechanism of “inference” is computing the the analysis transform of the image and adding uniform noise (as a stand-in for quantization), thus:
[Equation]
where U denotes a uniform distribution centered on y i .
Since the width of the uniform distribution is constant (equal to one), the ﬁrst term in the KL divergence technically evaluates to zero, and can be dropped from the loss function.
For the sake of argument, assume for a moment that the likelihood is given by:
[Equation]
with ˜ x
= g s ( ˜ y ; θ g )
The log likelihood then works out to be the squared difference between x and ˜ x , the output of the synthesis transform, weighted by λ .
Minimizing the second term in the KL divergence is thus equivalent to minimizing the expected distortion of the reconstructed image.
A squared error loss is equivalent to choosing a Gaussian distribution; other distortion metrics may have an equivalent distribution, but this is not guaranteed, as not all metrics necessarily correspond to a normalized density function.
The third term in the KL divergence is easily seen to be identical to the cross entropy between the marginal m ( ˜ y )
=
E x ∼ p x q ( ˜ y | x ) and the prior p ˜ y ( ˜ y ) .
It reﬂects the cost of encoding ˜
y , as produced by the inference model, assuming p ˜
y
as the entropy model.
Note that this term represents a differential cross entropy, as opposed to a Shannon (discrete) entropy as in eq. (2), due to the uniform noise approximation. Under the given assumptions, however, they are close approximations of each other (for an empirical evaluation of this approximation, see Ballé et al., 2017).
Similarly to Ballé et al. , we model the prior using a non-parametric, fully factorized density model (refer to appendix 6.1 for details): [Equation]
[Equation] where the vectors ψ ( i ) encapsulate the parameters of each univariate distribution p y i | ψ ( i ) (we denote all these parameters collectively as ψ ).
Note that we convolve each non-parametric density with a standard uniform density. This is to enable a better match of the prior to the marginal – for more details, see appendix 6.2. As a shorthand, we refer to this case as the factorized-prior model.
The center panel in ﬁgure 2 visualizes a subset of the quantized responses ( ˆ y ) of a compression model trained in this way.
Visually, it is clear that the choice of a factorized distribution is a stark simpliﬁcation: non-zero responses are highly clustered in areas of high contrast; i.e., around edges, or within textured regions.
This implies a probabilistic coupling between the responses, which is not represented in models with a fully factorized prior.
We would expect a better model ﬁt and, consequently, a better compression performance, if the model captured these dependencies.
Introducing a hyperprior is an elegant way of achieving this. x ˜ y ˜ z
θ g θ h ψ generative model x ˜ y ˜
z φ g φ h inference model x g a h a
U | Q h s
U | Q g s ˜ x | ˆ x y ˜ y | ˆ y z ˜ z | ˆ z ˜ σ | ˆ σ operational diagram  I NTRODUCTION OF A SCALE HYPERPRIOR
As evident from the center panel of ﬁgure 2, there are signiﬁcant spatial dependencies among the elements of ˆ y .
Notably, their scales appear coupled spatially. A standard way to model dependencies between a set of target variables is to introduce latent variables conditioned on which the target variables are assumed to be independent (Bishop, 1999).
We introduce an additional set of random variables ˜
z to capture the spatial dependencies and propose to extend the model as follows (ﬁgure 3).
Each element ˜
y i is now modeled as a zero-mean Gaussian with its own standard deviation σ i , where the standard deviations are predicted by applying a parametric transform h s to ˜ z (as above, we convolve each Gaussian density with a standard uniform; see appendix 6.2):
[Equation]
(7) with ˜ σ = h s ( ˜ z ; θ h )
[Equation]
We extend the inference model simply by stacking another parametric transform h a on top of y , effectively creating a single joint factorized variational posterior, as follows:
[Equation]
with y =
g a ( x ; φ g ) , z = h a ( y ; φ h ) .
This follows the intuition that the responses y should be sufﬁcient to estimate the spatial distribution of the standard deviations. As we have no prior beliefs about the hyperprior, we now model ˜ z using the non-parametric, fully factorized density model previously used for ˜ y
(appendix 6.1): [Equation] where the vectors ψ ( i ) encapsulate the parameters of each univariate distribution
p z
i | ψ ( i ) (collectively denoted as ψ ). The loss function of this model works out to be:
[Equation]
Again, the ﬁrst term is zero, since q is a product of uniform densities of unit width. The second term (the likelihood) encapsulates the distortion, as before. The third and fourth term represent the cross entropies encoding ˜ y and ˜ z , respectively.
In analogy to traditional transform coding, the fourth term can be seen as representing side information. The right-hand panel in ﬁgure 3 illustrates how the model is used as a compression method. The encoder subjects the input image x to g a , yielding the responses y with spatially varying standard deviations. The responses are fed into h a , summarizing the distribution of standard deviations ŷ x ŷ z c o n v
N x 3 x 3 /
R e L U c o n v
N x 5 x 5 / 2 ↓
R e L U c o n v
N x 5 x 5 / 2 ↓ a b s c o n v
M x 3 x 3 /
R e L U c o n v
N x 5 x 5 / 2 ↑
R e L U c o n v
N x 5 x 5 / 2 ↑
R e L U c o n v
N x 5 x 5 / [Figure]
G D N c o n v
N x 5 x 5 / [Figure]
G D N c o n v
N x 5 x 5 / [Figure]
G D N i
[Figure]
t i m a g e c o n v 3 x 5 x
[Figure] [Figure]
I G D
N c o n v
N x 5 x 5 / [Figure]
I G D
N c o n v
N x 5 x 5 / [Figure]
I G D
N c o n v
N x 5 x 5 / [Figure]
r e
[Figure] n
[Figure]
u c
[Figure]
i
[Figure]
n
[Figure] [Figure]
z
[Figure]
AE
AD
Q
[Figure]
AE
AD
Q ŷ z ŷ σ values, and N = 192 and M = 320 for the 3 higher values.
in z .
z is then quantized, compressed, and transmitted as side information.
The encoder then uses the quantized vector ˆ z to estimate ˆ σ , the spatial distribution of standard deviations, and uses it to compress and transmit the quantized image representation ˆ y .
The decoder ﬁrst recovers ˆ z from the compressed signal.
It then uses h s to obtain ˆ σ , which provides it with the correct probability estimates to successfully recover ˆ y as well.
It then feeds ˆ y into g s to obtain the reconstructed image.
E XPERIMENTS
To compare the compression performance of our proposed models, we conducted a number of experiments using the Tensorﬂow framework.
4. E XPERIMENTAL SETUP
We set up the transforms g a , g s , h a , and h s as alternating compositions of linear and nonlinear functions, as is common in artiﬁcial neural networks (ﬁgure 4).
Speciﬁcally, g a and g s are composed of convolutions and GDN/IGDN nonlinearities, which implement local divisive normalization, a type of transformation that has been shown to be particularly suitable for density modeling and compression of images (Ballé et al., 2016a; Ballé et al., 2017).
1 h a and h s are composed of convolutions and rectiﬁers (rectiﬁed linear units). To make the hyperprior model and the factorized-prior model comparable, we chose identical architectures for g a and g s , as shown in ﬁgure 4.
To maintain translation invariance across the model, all elements of z with the same channel index are assumed to follow the same univariate distribution.
This allows the model to be used with arbitrary image sizes.
Arithmetic coding is implemented using a simple non-adaptive binary arithmetic coder.
Each element of ˆ y and ˆ z is independently converted to its representation as a binary integer and arithmetically encoded from the most signiﬁcant to the least signiﬁcant bit.
Since the spatial distribution of standard deviations ( ˆ σ ) is known to the decoder by the time decoding of ˆ y is attempted, the arithmetic coder does not need to handle conditional dependencies.
It also does not need to be separately trained, since the binary probabilities needed for encoding are a direct function of the probability mass functions of ˆ y and ˆ z , and the probability mass functions in turn are direct functions of their “noisy” counterparts ˜ y , ˜
z by design (Ballé et al., 2017).
This is particulary important for ˆ y .
Since the prior is conditioned on ˆ σ , the probability mass functions p ˆ y i need to be constructed “on the ﬂy” during decoding of an image:
[Equation]
which can be evaluated in closed form.
The models were trained on a body of color JPEG images with heights/widths between 3000 and 5000 pixels, comprising approximately 1 million images scraped from the world wide web.
Images with excessive saturation were screened out to reduce the number of non-photographic images.
To reduce existing compression artifacts, the images were further downsampled by a randomized factor, such that the minimum of their height and width equaled between 640 and 1200 pixels.
Then, randomly placed 256 × 256 pixel crops of these downsampled images were extracted.
Minibatches of 8 of these crops at a time were used to perform stochastic gradient descent using the Adam algorithm (Kingma and Ba, 2015) with a learning rate of 10 − 4 .
Common machine learning techniques such as batch normalization or learning rate decay were found to have no beneﬁcial effect (this may be due to the local normalization properties of GDN, which contain global normalization as a special case).
With this setup, we trained a total of 32 separate models: half of the models with a hyperprior and half without; half of the models with mean squared error as the distortion metric (as described in the previous section), and half on the MS-SSIM distortion index (Wang, Simoncelli, et al., 2003); ﬁnally, each of these combinations with 8 different values of λ in order to cover a range of rate– distortion tradeoffs.
4. E XPERIMENTAL RESULTS
We evaluate the compression performance of all models on the publicly available Kodak dataset (Kodak, 1993).
Summarized rate–distortion curves are shown in ﬁgure 5.
Results for individual images, as well as summarized comparisons to a wider range of existing methods are provided in appendices 6.4 and 6.5.
We quantify image distortion using peak signal-to-noise ratio (PSNR) and MS-SSIM.
Each curve represents the rate–distortion tradeoffs for a given set of models, across different values of λ .
Since MS-SSIM yields values between 0 (worst) and 1 (best), and most of the compared methods achieve values well above 0.9, we converted the quantity to decibels in order to improve legibility.
Interestingly, but maybe not surprisingly, results differ substantially depending on which distortion metric is used in the loss function during training.
When measuring distortion in PSNR (ﬁgure 5, top), both our models perform poorly if they have been optimized for MS-SSIM.
However, when optimized for squared error, the model with the factorized prior outperforms existing conventional codecs such as JPEG, as well as other ANN-based methods which have been trained for squared error (Theis et al., 2017; Ballé et al., 2017).
Note that other published ANN-based methods not shown here underperform compared to the ones that are shown, or have not made their data available to us.
Our factorized prior model does not outperform BPG (Bellard, 2014), an encapsulation of HEVC (2013) targeted at still image compression.
When training our hyperprior model for squared error, we get close to BPG performance, with better results at higher bit rates than lower ones, but still substantially outperforming all published ANN-based methods.
When measuring distortion using MS-SSIM (ﬁgure 5, bottom), conventional codecs such as JPEG and BPG end up at the lower end of the performance ranking.
This is not surprising, since these methods have been optimized for squared error (with hand-selected constraints intended to ensure that squared error optimization doesn’t go against visual quality).
To the best of our knowledge, the state of the art for compression performance in terms of MS-SSIM is Rippel and Bourdev (2017).
Surprisingly, it is matched (with better performance at high bit rates, and slightly worse performance at low bit rates) by our factorized prior model, even though their model is conceptually much more complex (due to its multiscale architecture, GAN loss, and context-adaptive entropy model).
The hyperprior model adds further gains across all rate–distortion tradeoffs, consistently surpassing the state of the art.
With the results differing so heavily depending on which training loss is used, one has to wonder if there are any qualitative differences in the image reconstructions. When comparing images com- [Figure]
comparison to both stateof-the-art methods. Refer to ﬁgures 10 and 11 in the appendix for full-page RD curves that include a wider range of compression methods. [Figure]
s i d e i n f o r m a ti on [ [Figure] [Figure] pressed to similar bit rates by models optimized with an MS-SSIM distortion loss compared to a squared loss, we ﬁnd that the overall ﬁdelity in terms of how much detail is preserved appears similar.
However, the spatial distribution of detail changes substantially.
MS-SSIM, like its predecessor SSIM (Wang, Bovik, et al., 2004), is a metric designed to model human visual contrast perception.
Compared to squared loss, it has the effect of attenuating the error in image regions with high contrast, and boosting the error in regions with low contrast, because the human visibility threshold varies with local contrast. This behavior yields good results for images containing textures with different local contrast (refer to examples provided in appendix 6.5).
However, more frequently than expected, it can also produce results inconsistent with human expectations: for the image we show in ﬁgure 6, the compression model trained for MS-SSIM assigns more detail to the grass (low contrast), and removes detail from the text on the side of the airplane (high contrast).
Because semantic relevance is often assigned to high-contrast areas (such as text, or salient objects), the squared-error optimized models produce subjectively better reconstructions in these cases.
It is important to note that neither distortion metric is sophisticated enough to capture image semantics, which makes the choice of distortion loss a difﬁcult one.
Prior work on ANN-based image compression has shown that extending the transform coding concept from linear to nonlinear transforms fundamentally improves the qualitative nature of compression artifacts (Ballé et al., 2017).
It appears that nonlinear transforms with higher computational capacity adapt better to the statistics of natural images, imitating properties of the data distribution better than linear transforms.
When comparing image reconstructions visually between models with or without the hyperprior, we ﬁnd no changes to the qualitative nature of the artifacts.
Rather, the hyperprior model simply tends to produce image reconstructions with improved detail and a lower bit rate than the corresponding model with a factorized prior. Figure 7 shows how much of the total bit rate the hyperprior model uses as side information.
The amount of side information grows with the total bit rate, but stays far below 0.1 bpp, even for the highest total bit rates.
Still, the resulting improvement of the prior enables the performance gains over the factorized-prior model shown in ﬁgure 5.
Note that the architecture of the models does not explicitly constrain the bit rates in any way.
The illustrated trade-off in allocating bits for encoding ˆ z vs. ˆ y is simply the result of optimizing the loss function given in eq.  D ISCUSSION We implement a variational image compression model, conceptually identical to the model presented by Ballé et al. (2017), and augment it with a more powerful entropy model by introducing a hyperprior on the local scale parameters of the latent representation.
The hyperprior is trained end-to-end with the rest of the model.
Like all recent image compression methods based on ANNs, our method can be directly optimized for distortion losses that are more complex than pixel-wise losses such as mean squared error.
As one of the ﬁrst studies in this emerging ﬁeld, we examine the effect of optimizing for one of the most popular perceptual metrics, MS-SSIM, and compare it to optimizing for squared loss.
Note that Ballé et al.
(2016b) compare models trained for different metrics, but their results are limited by the choice of transforms.
Figure 6 demonstrates that the results can show signiﬁcant variation in terms of visual quality, depending on image content, which implies that unless human rating experiments are conducted to provide more reliable data, it is wise to compare methods based on more than a single type of metric.
Santurkar et al.
(2017) formulate their compression method in a hybrid VAE-GAN framework, adopting a stepwise training scheme where a decoder is ﬁrst trained using an adversarial loss.
It is then ﬁxed, and an encoder is trained to minimize the reconstruction error.
Rippel and Bourdev (2017) also employ an adversarial approach, but use a weighted combination of an MS-SSIM and an adversarial loss.
Baig and Torresani (2017) propose a compression scheme based on colorization, where color channels are predicted from the the luminance channel by making use of some model speciﬁc side information.
The luminance channel is compressed using a traditional method.
The proposed method exhibits signiﬁcant color distortions at low bit rates, and is limited by the compression method used for the luminance channel.
An early exploration of hierarchical generative models for compression of small images is found in
Gregor et al.
However, the aspect of quantization is not thoroughly considered, and hence, no actual compression method is designed.
Theis et al.
(2017) approach the problem of generating gradient descent directions for quantization functions by replacing their (unhelpful) gradient with the identity function, and derive a differentiable upper bound for the discrete rate term.
Ballé et al.
(2016b) instead replace the quantizer with additive uniform noise during training, and the discrete rate term with a differential entropy.
While this method doesn’t offer a bound for the approximation, it establishes a direct relationship between the discrete and continuous prior distributions p ˆ y and p ˜ y , which enables direct evaluation of the discrete prior as a function of the latents ˆ z as in eq.
(11), and hence makes use of a hyperprior feasible in practice.
The quality of the approximation is veriﬁed empirically by Ballé et al.
Wainwright and Simoncelli (2000) observe that linear ﬁlter responses (i.e., wavelet coefﬁcients obtained by ﬁltering an image) follow heavy-tailed marginal distributions, but can be represented as conditionally Gaussian when groups of neighboring coefﬁcients are linked by a common scale multiplier.
That is, the distributions of the ﬁlter responses can be modeled as Gaussian scale mixtures.
Lyu and Simoncelli (2009) extend this model from spatially localized groups of wavelet coefﬁcients to a global image model. Our model can be seen as a further extension of this, where the ﬁlter responses are replaced with responses of a nonlinear transform, and an approximate inference model is added.
Theis et al.
(2017) directly use Gaussian scale mixtures, but in the form of a fully factorized prior.
In the presented form, our variational model is perhaps most closely related to ladder
VAEs (Sønderby et al., 2016).
However, we choose different parametric forms to accommodate the approximation of the quantization and entropy coding process.
In classical transform coding methods, compression researchers have exploited statistical dependency in the latent variables (e.g., DCT or wavelet coefﬁcients) by carefully hand-engineering entropy codes modeling the dependencies in the quantized regime (Taubman and Marcellin, 2002).
This presents a much more difﬁcult engineering problem than relying on a fully factorized entropy model; transitioning to nonlinear transforms whose parameters are determined through training (and thus may be different for each re-training) only complicates the problem.
Toderici et al.
(2017) model images directly with a binarized latent representation, which technically removes the need for a separate entropy coding step.
However, this corresponds to a very inﬂexible entropy model (a uniform prior on a binary representation, with no trainable parameters).
The model apparently compensates for this by using higher capacity transforms (e.g., based on recurrent networks).
Johnston et al.
(2017) improve the method by designing an adaptive entropy model.
However, this entropy model is not included in the rate term while training the transforms, and hence no feedback (in terms of gradients) is returned from the entropy model back to the transforms during training.
This breaks the paradigm of end-to-end optimization, and may stand in the way of better compression performance.
Similarly, Rippel and Bourdev (2017) use a hand-designed energy function without trainable parameters as the prior for training the autoencoder, and design an adaptive entropy model post hoc.
The fact that our factorized prior model matches the performance of their method, when optimized on the same metric, may point towards this disconnect. Ágústsson et al.
(2017) extend the fully-factorized prior model by proposing to do vector quantization over small subtensors of the latent representation, which effectively relaxes the factorization. They train their method end-to-end. All of the models presented here make use of GDN, a type of nonlinearity implementing local normalization. As part of a Gaussianizing transformation, GDN has been shown to be more efﬁcient, in terms of number of parameters, at removing statistical dependencies in image data, than pointwise nonlinearities (Ballé et al., 2016a).
Furthermore, there has been a long history of generative models, starting with independent component analysis (Cardoso, 2003), which can successfully recover factorized representations just by maximizing likelihood assuming a fully factorized prior. Despite these facts, we observe that signiﬁcant dependencies between neighboring elements remain in the latent representation of our compression models (ﬁgure 2), even though we took care not to impose constraints on the transforms which might reduce their capacity to factorize the representation (refer to appendix 6.3 for details). We attribute this to the fact that the rate–distortion loss, unlike a maximum likelihood loss, trades off the rate term against expected distortion. It is easy to see that for increasing values of λ , the rate term containing the factorized prior becomes less and less important.
Hence, it is questionable whether rate–distortion optimality implies full independence of the representation, at least for arbitrary values of λ . Regardless of this, the fact that the hyperprior models consistently outperform models with a factorized prior illustrate that it is important for any compression method to reduce mismatch between the prior and the marginal, as in eq.
Our model, when trained on the appropriate loss, has the capacity to surpass the state of the art on MS-SSIM, but does not quite reach the performance of a heavily optimized traditional method such as BPG on PSNR (while outperforming all other methods based on ANNs).
This discrepancy may indicate that methods based on ANNs have not yet reached the expressive power of traditional methods.
As such, the introduction of a hyperprior – or, in traditional terms, side information – is an elegant way of introducing more ﬂexible priors, and a big step in the right direction.
A PPENDIX 6. U NIVARIATE NON - PARAMETRIC DENSITY MODEL
Ballé et al.
(2017) use a non-parametric piecewise linear density model to represent each factor of the fully factorized prior.
By increasing the number of samples per unit interval, it can in principle be used to model any univariate density with arbitrary precision.
However, it has two practical problems: The range of values with non-zero probability must be ﬁnite and known ahead of time, and its implementation is non-trivial with existing automatic differentiation frameworks, both due to numerical issues with normalizing the density and the fact that it typically relies on discrete operations such as array indexing.
For the compression models presented in this paper, we instead use the following model based on the cumulative.
We deﬁne a density p : R → R + using its cumulative c : R → [0 , 1] by satisfying the following constraints:
[Equation]
Note that the monotonicity constraint of the cumulative is established by requiring the density function p to be non-negative.
Suppose the cumulative is a composition of functions.
Then the density can be written using the chain rule of calculus:
[Equation]
where we write the derivative of f k as f (cid:48) k .
We’ll allow the f k to be vector functions:
[Equation]
In general, the f (cid:48) k are Jacobian matrices, and the dots are matrix multiplications.
To ensure p ( x ) is univariate, the domain of f 1 and the range of f K need to be one dimensional ( d 1 = r K = 1 ).
To guarantee that p ( x ) is a density, we just need f K to map to the range between 0 and 1, and ensure that p ( x ) ≥ 0 .
To do that, we require all the Jacobian elements to be non-negative.
Then the matrix product computing p ( x ) is non-negative as well, and we have deﬁned a valid density.
An effective choice of f k is the following (as a shorthand, we deﬁne tanh , sigmoid , and softplus as elementwise functions, when applied to vectors or matrices):
[Equation]
where H ( k ) are matrices, b ( k ) are vectors, and g k are nonlinearities deﬁned as
[Equation]
(18) where a ( k ) is a vector and (cid:12) denotes elementwise multiplication.
The rationale behind this particular nonlinearity is that it allows to expand or contract the space near x = 0 .
a ( k ) controls the rate of expansion (when positive) or contraction (when negative).
If a ( k ) were ﬁxed to a positive value, “peaks” in the density would become easier to model than “troughs”.
The derivatives work out as follows:
[Equation]
and
[Equation] [Equation]
For the derivatives to be non-negative, we need to constrain H ( k ) to have all non-negative elements, and the elements of a ( k ) to be lower bounded by − 1 .
This is easily done by reparameterization:
[Equation]
where the quantities with the hat are the actual parameters.
Setting K = 1 yields an interesting special case:
[Equation]
[Equation]
which is a density function similar to a “smoothed” Laplacian, with identical tail behavior.
A plot of this density function, as well as a plot of a ﬁt of this model to a “toy” mixture density is provided in ﬁgure 8.
It may seem odd to deﬁne a density function as an explicit derivative; however, in an automatic differentiation framework, this operation is very easy to implement, and the resulting density function is normalized by construction.
We have found the model to ﬁt well to arbitrary densities, and perform just as well as the piecewise linear model in the context of compression models.
For all experiments in this paper, we used K = 4 , with the dimensionalities r 1 = r 2 = r 3 = 3 .
Each univariate density model is associated with its own set of parameters a ( k ) , b ( k ) , H ( k ) (which, together, form ψ ( i ) ).
6. M ODELING PRIORS WITH ADDED UNIFORM NOISE
We model both the prior p ˜ y | ˜
z and the hyperprior p ˜ z
using densities that are convolved with a standard uniform density function.
This is to ensure that the priors have enough ﬂexibility to match the variational posterior q .
To see this, consider that in some cases, it is beneﬁcial in terms of rate– distortion performance for the model to “disable” part of the latent representation, leading to a lower effective dimensionality than the model architecture has been set up for.
For simplicity of notation, let’s assume that the variational posterior and the prior have just one dimension which has collapsed.
In this case, g a converges to always producing a constant value for the corresponding dimensions: [Equation] When this happens, the marginal distribution of that element during training is a uniform density centered on c , due to the added uniform noise, and the variational posterior matches it exactly: [Equation]
The cross entropy of this element is given by: [Equation] This entropy should evaluate to zero bits, as the quantized representation is deterministic (and hence, no information needs to be transmitted). For the cross entropy to evaluate to zero, however, the prior needs to be ﬂexible enough to assume the shape of the posterior – a unit-width uniform density.
[Figure] Due to its inﬁnitely steep edges, the uniform distribution is a corner case for not only the Gaussian density model, but also the non-parametric model described in appendix 6.1.
To ﬁx this, we incorporate the added noise directly into the prior/hyperprior by convolving the underlying density model p with a standard uniform:
[Equation]
where c is the cumulative of the underlying density model.
Now, whatever the underlying density p is, letting its scale go towards zero makes p ˜
y approach a unit-width uniform density. Since the non-parametric model is deﬁned via its cumulative, and the cumulative of a Gaussian is available in most computational frameworks, this solution is easy to implement in practice.
6. M ODEL CAPACITY
Our results seem to indicate that a certain degree of statistical dependency in the latent image representation y is preferred by the rate–distortion objective, and that the hyperprior model performs better by embracing this. However, it is possible that dependencies remain simply because the analysis and synthesis transforms g a and g s do not have enough capacity to factorize the image representation, or because the training algorithm did not succeed in ﬁnding the global optimum.
Although it is impossible to fully control for this, we attempted to minimize the chances that capacity limitations in the transforms lead to the wrong conclusions, by carefully selecting the number of ﬁlters across layers of the transforms (as given by N and M in ﬁgure 4).
We established in previous experiments that, for a given λ , there exist a certain number of ﬁlters per layer at which performance saturates, and no gains can be achieved by further increasing it (ﬁgure 9; note that for these experiments, we set N = M ).
The optimal number of ﬁlters increases with λ , indicating that models with higher bit rates require higher transform capacities.
Based on these previous experiments, we attempted to choose values close to the point of saturation, or a little higher, in order to control for capacity limitations while minimizing training time.
Additionally, we found that allowing a somewhat wider bottleneck M > N helps to achieve comparable performance with overall lower N , and we used this when choosing the model architectures.
6. P ERFORMANCE COMPARISONS FOR THE K ODAK IMAGE SET
The following plots show the same results as ﬁgure 5, but provide comparisons to a wider array of compression methods.
Note that the method to aggregate rate–distortion points across images differs between the PSNR and MS-SSIM plots: in the latter, we interpolate the RD curves for each image (as shown in appendix 6.5) using cubic splines at a predeﬁned set of bit rates, and then average across equal bit rates. In the former, no interpolation was used, averaging rate and distortion measurements across equal values of λ . As noted by Ballé et al. (2017), directly comparing RD curves with different methods of aggregation can give misleading results.
Because of this, we match our aggregation method to the data available for the current state of the art ( λ -aggregation for HEVC and
PSNR, and rate aggregation for Rippel and Bourdev (2017) and MS-SSIM).
Ultimately, a comparison based on individual images, as provided in the next section, should be considered more reliable; however, data on individual images for Rippel and Bourdev (2017) has not been available.
[Figure]
bits
[Figure]
pixel (bpp)
P S N
R R
G B
RD curves averaged over
BPG (4:4:4)
BPG (4:2:0) hyperprior (optimized for MSE) factorized prior (optimized for MSE)
Ballé (2017)
JPEG2000 (OpenJPEG)
WebP
JPEG2000 (Kakadu)
Johnston (2017)
Theis (2017) hyperprior (optimized for MS-SSIM) factorized prior (optimized for MS-SSIM)
JPEG (4:2:0)
Toderici (2017)
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure]
[Figure]
M SSS I M
R G B ( d
[Figure]
RD curves averaged over Kodak (MS-SSIM
[Figure]
6. P ERFORMANCE COMPARISONS FOR INDIVIDUAL K ODAK IMAGES
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]