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 into the generative model 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 for image compression with artiﬁcial neural networks (ANNs).
Unlike existing autoencoder compression methods, our model trains a powerful entropy model 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)
I NTRODUCTION
Recent machine learning methods for lossy image compression have generated signiﬁcant interest in both the machine learning and image processing communities (Rippel and Bourdev, 2017; Johnston et al., 2017; Ball ´ e et al., 2016; Minnen et al., 2017).
The majority of these methods are based on autoencoders with a continuous latent representation that is quantized and then losslessly compressed using entropy coding to create the bitstream.
The expected code length of an optimal entropy coder is the cross entropy between the true distribution of the latent representation and the probability model used by the entropy coder (the entropy model ).
The minimal code length is achieved when the observed distribution and the model distribution are identical, and any divergence can lead to suboptimal compression performance.
In all of the existing machine learning methods for image compression, the entropy model used to encode the latent representation is assumed to be ﬁxed and fully known by both the encoder and decoder.
Empirically, however, the distribution of the latent representation of natural images varies signiﬁcantly from image to image as well as spatially within a single image.
This observation implies a potential mismatch between the entropy model and the local distributions within each image, which means that it may be possible to improve existing compression models by making the entropy model image-dependent and spatially adaptive.
Conventional compression methods achieve this by transmitting side information that allows for more effective entropy coding.
For instance, while JPEG always uses ﬁxed-size blocks of 8 ×pixels, more recent methods such as WebP and BPG (derived from HEVC) partition each image into variable-size blocks, convey the partition structure to the decoder, and then transmit the actual block representations.
Because the encoder can optimize the partitioning for each image, it can be used to achieve higher compression rates without degrading image quality.
In conventional codecs, the structure of this side information is hand-designed.
In contrast, the models presented in this paper learn to generate side information that adapts to the image content.
Essentially, they learn a representation of the entropy model in the same way that the underlying model learns a representation of the image.
Furthermore, because our models are optimized endto-end, they minimize the total expected code length by learning to balance the amount of side information with the expected improvement of the entropy model.
This optimization problem can be formalized by expressing it in terms of variational autoencoders (VAEs), probabilistic generative models augmented with approximate inference models (Diederik P.
Kingma and Welling, 2014).
Ball´e et al.
(2016) and Theis et al.
(2017) previously identiﬁed that
[Figure]
˜ y
[Figure]
g
[Figure]
x g a
U | Q g s ˜ x | ˆ x y ˜ y
| ˆ y
[Figure]
some autoencoder-based compression methods are formally equivalent to VAEs, where the entropy model 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´e et al.
(2016) with a hyperprior that captures the fact that spatially neighboring elements of the latent representation vary together in their scales.
We demonstrate that our model leads to state-of-the-art image compression when measured using the
MS-SSIM quality index, and it provides signiﬁcantly better rate-distortion performance compared to other machine learning methods when measured using PSNR.
Finally, we present novel results studying 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, 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 () 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.
The parameters θ g and φ g then encapsulate the weights of the neurons, etc.
(refer to Section 4 for details).
Compression can be cast as a rate–distortion optimization problem.
Rate is the expected code length (bit rate) of the compressed representation: assuming the entropy coding technique is operating at high efﬁciency, this can be approximated as the Shannon entropy of ˆ y .
Distortion is the expected difference between the reconstruction ˆ x and the original image x .
A trade-off exists between rate and distortion 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.
The minimization problem is parameterized by λ , the weight on the distortion term.
Different applications require a different trade-off between the rate and distortion, 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 (J. Ball´e et al., 2016).
Here, we follow the latter method, which switches back to actual quantization for 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 problem can be formally represented as a variational autoencoder (Diederik P. 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 true posterior p ˜ y | x ( ˜ y
| x ) is approximated with a parametric variational density q
( ˜ y
| x ) by minimizing the Kullback–Leibler divergence
[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] with y = g a ( x
[Equation]
φ g ) , where U denotes the uniform distribution.
Since the width of the uniform distribution is constant, 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
[Equation]
g [Equation] 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.
Minimizing squared error 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.
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
The third term in the KL divergence is easily seen to be identical to the differential cross-entropy between the variational density q and the prior.
It reﬂects the cost of encoding ˜
y , as produced by the inference model, assuming p ˜
y
as the entropy model.
Under certain provisions, this closely approximates the discrete (Shannon) cross-entropy of encoding ˆ y (for details, see Ball ´ e et al., 2016), which represents the expected bitrate of an optimal arithmetic coder.
Most previous work assumes that p ˜
y is a parametric, but factorized distribution, i.e., assumes independence between the elements of ˜ y
(Theis et al., 2017; Ball ´ e et al., 2016).
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: the 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 factored 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.
I NTRODUCTION OF A SCALE HYPERPRIOR
As evident from the center pane of ﬁgure 2, there are signiﬁcant spatial dependencies among ˆ y .
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, 1998).
We introduce an additional set of random variables ˜
z to capture the spatial dependencies and propose the following extension of the model (ﬁgure 3):
Each element ˜
y
i is 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 :
[Equation]
with ˜
[Equation]
=
[Equation]
Note that we convolve each Gaussian with a standard uniform density.
This is to take into account the uniform noise added by the inference model – if we modeled ˜
y
i simply as Gaussian, it would create a potential model mismatch between the variational encoder and decoder.
This technique is explained in more detail in appendix 7.2.
We extend the inference model simply by stacking another parametric transform h a on top of y , effectively creating a single joint factored variational posterior, as follows:
[Equation]
with
y =
g a ( x
φ
g
) , z
= h a ( y
[Equation]
φ h )
C o n v
N x 3 x 3 /
S o f t p l u s
C o n v
N x 5 x 5 / 2 ↓
S o f t p l u s
C o n v
N x 5 x 5 / 2 ↓
A b s
C o n v
N x 3 x 3 /
S o f t p l u s
C o n v
N x 5 x 5 / 2 ↑
S o f t p l u s
C o n v
N x 5 x 5 / 2 ↑
S o f t p l u s h h g
C o n v
N x 5 x 5 / 2 ↓
G D N
C o n v
N x 5 x 5 / 2 ↓
G D N
C o n v
N x 5 x 5 / 2 ↓
G D N
I n p u t
I m a g e
C o n v
N x 5 x 5 / 2 ↑
I G D N
C o n v
N x 5 x 5 / 2 ↑
I G D N
C o n v
N x 5 x 5 / 2 ↑
I G D N
C o n v
N x 5 x 5 / 2 ↑
R e c o n s t r u c t
i
o
n g
C o n v
N x 5 x 5 / 2 ↓
[Figure]
a s s y ŷ z ẑ σ̂
[Figure]
ẑ ẑ |
[Figure]
AE
[Figure]
|
[Figure]
AE
AD
Q
Convolutional layers are represented as (num channels × height × width / stride).
The left half is akin to an autoencoder with Q representing a quantization block and AE, AD representing arithmetic encoder and arithmetic decoder blocks, respectively.
The right half corresponds to the network learning hyperpriors with p ˆ y | ˆ σ representing priors dependent on hyperpriors and p ˆ z | ψ representing the distribution of hyperparameters.
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 on ˜ z , we model it using a non-parametric, fully factorized density model (refer to appendix 7.1 for details):
[Equation]
where the vectors ψ ( i ) encapsulate the parameters of each marginal distribution
p z i | ψ ( i ) (we denote all these parameters collectively as ψ ).
Again, we use the uniform noise distribution to reduce model mismatch (appendix 7.2).
The loss function of this model works out to be:
[Equation]
Again, the ﬁrst term is constant, since q is a product of uniform densities with constant width.
The second term (the likelihood) encapsulates the distortion, as before.
The third and fourth term measure the conditional cross-entropy/cross-entropy needed to encode ˜ y
and ˜ z , respectively.
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 in z .
z is then quantized, compressed, and sent over the channel.
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 responses ˆ 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 assess the compression performance of our proposed model, 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 and are composed of convolutions and softplus nonlinearities.
The fully factorized nonparametric model p z i | ψ ( i ) is described in appendices 7.1 and 7.2.
To maintain the translation invariance, all elements of z with the same channel index are assumed to follow the same marginal 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 (J. Ball´e et al., 2016).
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 (D. P. Kingma and Ba, 2014) 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, Eero P 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, n.d.).
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 such as JPEG 2000,
WebP, Johnston et al.
(2017), are provided in appendices 7.3 and 7.4.
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 ﬁxed prior outperforms existing conventional codecs such as JPEG, as well as other machine learning based methods which have been trained for squared error (Theis et al., 2017; Ball´e et al., 2016).
Note that other published methods not shown here underperform compared to the ones that are shown, or have not made their data available to us.
Our ﬁxed prior model does not outperform BPG, an encapsulation of HEVC (Bross et al., 2012) targeted
[Figure]
at still image compression.
When training our hyperprior model for squared error, we roughly match BPG performance, with better results at higher bitrates than lower ones, but still substantially outperforming all published machine learning 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 at low bitrates, and substantially surpassed at higher bitrates, by our ﬁxed prior model, even though the model presented in Rippel and Bourdev (2017) is conceptually much more complex (due to its multiscale architecture, GAN loss, and context-adaptive entropy coder).
The hyperprior model adds even further gains above 0.5 bpp, although it loses some at the lowest bitrates.
We believe that this may be caused by a suboptimal architecture choice for the low bitrate models (such as number of channels or layers).
What is the qualitative difference of image reconstructions when optimizing for different distortion losses?
When comparing images compressed to similar bit rates by models optimized with an MSSSIM distortion loss compared to a squared loss (ﬁgure 6), 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, is a metric designed to model contrast perception.
Compared to a simple 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 (Wang, Bovik, et al., 2004).
The result is that the compression model trained for MS-SSIM assigns more detail to lower contrast regions, such as the grass, and removes detail from high contrast regions, such as the text on the side of the airplane.
More frequently than expected, this is inconsistent with human expectations: 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 some cases.
It is important to note that neither distortion metric is powerful enough to capture semantics, which makes the choice of distortion loss a difﬁcult one.
Prior work on machine learning 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 ´ e et al., 2016; Theis et al., 2017; Johnston 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 even under severe rate limitations.
When training a model with or without the hyperprior for equal values of λ , 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 ﬁxed prior (ﬁgure 5; refer to appendix 7.4 for visual results).
R ELATED WORK
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) approaches the problem of generating gradient descent directions for quantization functions by replacing their (unhelpful) gradient with the identity function, and derives a differentiable upper bound for the discrete rate term.
J. Ball ´ e et al.
(2016) replace the quantizer with additive uniform noise during training, and the discrete rate term with differential entropy.
This method allows a formal connection to variational autoencoders; however, no bound for the rate approximation is given.
Instead, Ball ´ e et al.
(2016) verify empirically that the discrete entropy cost matches the differentiable one.
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.
Wainwright and Eero P. 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.
Our model can be seen as an extension where the ﬁlter responses are replaced with responses of a nonlinear transform, spatially local model is generalized to a global one, and an approximate inference model is added.
Theis et al.
(2017) use discrete Gaussian scale mixtures as a fully factored density model of the latent representation.
Crucially, they model all elements as independent, not fully exploiting their statistical dependencies.
Santurkar et al.
(2017) formulate their model 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 a encoder is trained to minimize the reconstruction from decoder.
However, the step-wise training breaks the end-to-end paradigm and may lead to sub-optimal models.
Baig and Torresani, 2017 propose a colorization based compression scheme where, given the luminance channel, color channels are predicted 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 rate and is limited by the compression method used for luminance for preserving details.
Among traditional approaches, JPEG (Wallace, 1991) is the most common lossy image compression method (Bull, 2014).
However, more advanced standards have since been developed including
JPEG2000, WebP (Google, 2010), and Better Portable Graphics (BPG) (Bellard, 2014).
D ISCUSSION
We implement a variational image compression model, conceptually identical to the model presented in (Ball´e et al., 2016), and augment it with a more powerful entropy model by introducing a hyperprior on its local scale parameters.
The hyperprior is trained end-to-end with the rest of the model.
This sets our work apart from other image compression models based on latent image representations:
In classical transform coding methods, compression researchers have exploited 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 is a much more difﬁcult problem than relying on an entropy coder that assumes independence between all the elements of ˆ y ; transitioning to nonlinear transforms whose parameters are determined through training (and thus may be different for each re-training) only aggravates this 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, it corresponds to a very inﬂexible prior (a uniform prior on a binary representation, with no trainable parameters).
The models compensate for this by introducing higher capacity transforms (e.g., based on recurrent networks).
Johnston et al.
(2017) improve this by designing an entropy coder in a way that adapts to the spatial probability structure of the latent representation.
However, no feedback (in terms of gradients) is returned from the entropy coder back to the prior during training.
This breaks the paradigm of end-to-end optimization, and may lead to suboptimal results.
Rippel and Bourdev (2017) impose a hand-designed energy function without trainable parameters as the prior for training the autoencoder.
They also train an adaptive entropy coder without feedback to the autoencoder.
The fact that our ﬁxed prior model outperforms their method when optimized on the same metric, despite their state-of-the-art results, may point towards this disconnect.
Like all recent image compression methods based on machine learning, our method can be directly optimized for distortion losses that are more complex than simpler 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.
Figuredemonstrates 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.
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 machine learning).
This discrepancy may indicate that methods based on machine learning 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 powerful priors, and a big step in the right direction.
A PPENDIX 7. F ULLY FACTORIZED , NON - PARAMETRIC DENSITY MODEL
J. Ball ´ e et al.
(2016) use a non-parametric piecewise linear model to capture the marginals of the prior.
By increasing the number of samples per unit interval, this can in principle be used to model any density with arbitrary precision.
However, it has two practical issues: The range of 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 normalization 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 n as f (cid:48) n .
We’ll allow the f n to be vector functions:
[Equation]
In general, the f (cid:48) n 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 N need to be one dimensional ( d 1 = r N = 1 ).
To guarantee that p ( x ) is a density, we just need f N to map to the range between 0 and 1, and ensure that p ( x ) ≥ 0 .
To do that, let us require all the Jacobian elements to be non-negative:
[Equation]
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 n is the following (as a shorthand, we deﬁne tanh , sigmoid , and softplus as elementwise functions, when applied to vectors or matrices):
[Equation]
[Equation]
where H ( n ) are matrices, b ( n ) are vectors, and g n are nonlinearities deﬁned as
[Equation]
16) where a ( n ) 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 ( n ) controls the rate of expansion (when positive) or contraction (when negative).
If a ( n ) were ﬁxed to a positive value, “peaks” in the density would become easier to model than “troughs”.
The derivatives work out as follows:
[Equation]
, with
[Equation] [Equation]
and
[Equation] [Equation]
For the derivatives to be non-negative, we need to constrain H ( n ) to have all non-negative elements, and the elements of a ( n ) to be lower bounded by − 1 .
This is easily done by reparameterization:
[Equation]
where the quantities with the hat are the actual parameters.
In the special case that N = 1 , we obtain an interesting special case:
[Equation]
[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 7.
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 N = 4 , with the dimensionalities r 1 = r 2 = r 3 = 3 .
Each marginal distribution model is associated with its own set of parameters a ( n ) , b ( n ) , H ( n ) (which, together, form ψ ( i ) ).
7. 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 make sure that the priors have enough ﬂexibility to match the variational posterior ( q ).
To see this, consider the following special case:
The inference model may effectively “disable” part of the representation by collapsing some of its dimensions.
When this happens, the marginal distribution of individual elements of ˜ y or ˜ z will collapse to a standard uniform density due to the added uniform noise.
To correctly capture the expected entropy cost of this distribution (zero – the quantized version of the element will have a deterministic value independent of the image), the prior or hyperprior needs to be able to model a uniform distribution, or it will assign a cross entropy to these elements that is non-zero, which may lead to suboptimal results.
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 marginal model described in appendix 7.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] [Equation]
where c is the cumulative of the underlying density model.
Since the non-parametric model is deﬁned via its cumulative, and the cumulative of a Gaussian is available in most computational frameworks, this is easy to implement in practice.
7. P ERFORMANCE COMPARISONS FOR THE K ODAK IMAGE SET
Bits per pixel (BPP)
P S N
R R
G B
RD curves averaged over Kodak (PSNR)
BPG (4:4:4)
BPG (4:2:0)
Hyperprior (Optimized for L2)
Fixed Prior (Optimized for L2)
Ballé (2017)
JPEG2000 (OpenJPEG)
WebP
JPEG2000 (Kakadu)
Johnston (2017)
Theis (2017)
Fixed Prior (Optimized for MS-SSIM)
Hyperprior (Optimized for MS-SSIM)
JPEG (4:2:0)
Toderici (2017)
Bits per pixel (BPP)
M SSS I M
R G B ( d B ) RD
MS Hyperprior (Optimized for MS-SSIM)
Fixed Prior (Optimized for MS-SSIM)
Rippel (2017)
Hyperprior (Optimized for L2)
Johnston (2017)
BPG (4:4:4)
Ballé (2017)
BPG (4:2:0)
Fixed Prior (Optimized for L2)
JPEG2000 (Kakadu)
WebP
JPEG2000 (OpenJPEG)
Theis (2017)
Toderici (2017)
JPEG (4:2:0)
[Figure]
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]