A DAPTIVE S AMPLE - SPACE & A DAPTIVE P ROBABILITY CODING : A NEURAL - NETWORK BASED APPROACH FOR COMPRESSION
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose Adaptive Sample-space & Adaptive Probability (ASAP) coding, an efﬁcient neural-network based method for lossy data compression.
Our ASAP coding distinguishes itself from the conventional methods based on adaptive arithmetic coding in that it models the probability distribution for the quantization process in such a way that one can conduct back-propagation for the quantization width that determines the support of the distribution.
Our ASAP also trains the model with a novel, hyper-parameter free multiplicative loss for the rate-distortion tradeoff.
With our ASAP encoder, we are able to compress the image ﬁles in the
Kodak dataset to as low as one ﬁfth the size of the JPEG-compressed image without compromising their visual quality, and achieved the state-of-the-art result in terms of MS-SSIM based rate-distortion tradeoff
I NTRODUCTION
In terse terms, lossy data compression is a task in which one seeks to encode a data ﬁle into as short a code as possible without losing the essential information.
Extensive research has been conducted in the ﬁeld of lossy data compression; JPEG, WebP and BPG are well known lossy image compression codec.
The recent advances in machine learning methods are accelerating the pace of the research in this ﬁeld.
For example, studies like Toderici et al.
; Ball ´ e et al.
; Theis et al.
(2017); Johnston et al.
(2017); Rippel & Bourdev (2017)
; Mentzer et al.
; Agustsson et al.
; Nakanishi et al.
; Ball ´ e et al. (2018) have incorporated the methods of deep learning into lossy compression, and some of them succeeded in producing a result that far surpasses the classical, neural-network-free methods.
Almost all lossy codecs to date—including the state-of-the-art codec—are built on autoencoder.
After transforming the input data by the autoencoder, the lossy codecs discretize the latent space using a quantizer function and further converts the transformed data so that it can be stored as a binary code.
In order to quantize the latent feature into as short a code as possible, one usually decompose the feature into pieces and use the information of the pieces that it has quantized so far (i.e., context) in order to adaptively choose a discrete probability space for the quantization and the entropy coding of the next piece of information.
By construction, the choice of the probability space for the quantization and the entropy coding of the feature vector is a key factor that determines the performance of the compression algorithm.
In this study, we present a novel compression architecture that allows the user to optimize the discrete support of the probability space to be used for quantization and entropy coding (PSQE).
Optimization of the support of the PSQE in a NN-based architecture is a daunting task.
To the authors’ best knowledge, there has been no study to date that enabled this optimization for NN-based architecture.
Our Adaptive Sample-space & Adaptive Probability (ASAP) coding partially achieves this task by assigning to each latent feature a discretized normal distribution with different mean, variance and quantization width. Our method automatically chooses appropriate quantization width for each latent feature. For the compression with small models in particular, we were able to conﬁrm that this adaptive quantization width beneﬁts the performance in terms of the MS-SSIM score of the reconstructed images.
We also present a novel objective function for the training of the quantizer.
In general, the problem of minimizing the reconstruction error ( distortion ) under a constraint of code-length (or bpp ) can be
JPEG
BPG
Proposal
[Figure]
Original
[Figure] [Figure]
A
B authors formulated in an equation of the form bpp + λ × distortion with a Lagrange multiplier λ .
Needless to say, however, the appropriate value of λ depends on the user-chosen weight of the importance for the code-length, which in turn heavily depends on the situation.
Training a separate model and conducting a separate hyper-parameter search at every different situation is cumbersome.
In order to circumvent this inconvenience, we propose a novel, hyper-parameter free multiplicative loss function that measures the rate-distortion tradeoff.
Equipped with the greater freedom for the choice of PSQE, our ASAP model trained with our multiplicative-loss is not only able to attain compression quality that is on par with the state-of-theart methods trained with an extensive hyper-parameter search about λ , but also is able to outperform all the predecessors in terms of the tradeoff ability for almost all region of bpp constraint.
Vice versa, under the same distortion constraint, our codec can compress the image ﬁle to as low as 2 ∼
5 times smaller code in comparison to JPEG.
This is 2 ∼ 2.5 times smaller a size in comparison to
BPG, which is one of the latest compression algorithms used in industry (see Figs.
2 and 15).
[Figure]
[Figure]
M ETHODS
By all means, we intend our algorithm to be applied to various tasks including audio and video compression.
For ease of explanation, however, we will assume for now that our problem is the compression of image, and describe the overall ﬂow of our algorithm in the context of image compression.
In a nutshell, our method follows the usual routine of transform coding.
Throughout, let x ∈ R C
× H 0 × W 0 be the original image, where C 0 , H 0 , W 0 respectively represent the number of the channels, the height and the width of the image.
The algorithm begins by ﬁrst applying the transformation function F to x in order to obtain a latent feature representation, z ∈ R C × H × W .
The algorithm then quantizes the signal z into a ﬁnite space, and then covert the quantized signal ˆ z to a binary signal s by entropy coding.
For the reconstruction, we apply the synthesizer map G to the code ˆ z .
That is, ˆ x = G (ˆ z ) .
Fig.
3 illustrates the overall ﬂow.
In general, we would like to obtain a shorter s with smaller distortion d ( x , ˆ x ) , where d is the measure of the corruption in the reconstruction (e.g., 1 − MS-SSIM)
(Wang et al., 2004).
Our goal is therefore to train the networks that optimize the rate-distortion tradeoff.
By construction, the loss of information takes place during the transformation from z to ˆ z .
Our ASAP is the invention for this part of the compression.
In Sec. 2.1, we describe the distinguishing features of ASAP coding relative to the conventional coding methods.
We then describe in Sec.
2.2 the technicalities for the training of the ASAP models in the context of image compression.
2. T HE FEATURES OF ASAP CODING
In this section, we review the basics of entropy coding.
Let us assume that z ∈ [ − 1 , 1] C × H × W .
This z is to be quantized into ˆ z ∈ R C × H × W by the discretization of the continuous space [ − 1 , 1] C × H × W , so that it can be converted into a binary code via invertible compression algorithm.
For the entropy coding of ˆ z , we would have to prepare a probability distribution for ˆ z that represents its occurring probability.
Oftentimes, this occurring probability is modeled by some form of auto regressive model.
In our task of compressing the images, we adopted a model used in Nakanishi et al.
We partitioned ˆ z into K groups { ˆ z ( k ) ∈ R C ( k ) × H ( k ) × W ( k ) ; k = 1 , ..., K } based on a geometric pattern we illustrate in Fig.
Let C ( k ) , H ( k ) , and W ( k ) be the dimension of the k -th group.
We then estimated the true occurrence probability as p ∗ ( ˆ z ) =
p ∗ ( ˆ z (1) ) (cid:81)
Kk =2 p ∗ ( ˆ z ( k )
| ˆ z (1: k − 1) ) by recursively approximating p ∗ ( ˆ z ( k )
| ˆ z (1: k − 1) )
with some probability distribution π ( ˆ z ( k ) | ˆ z (1: k − 1) ) with discrete support.
For ease of notation, we would occasionally use π ( k ) (ˆ z ( k ) ) to denote the same distribution.
In this model, we assume that each coordinate of ˆ z ( k ) is approximately independent conditional to ˆ z (1: k − 1) .
This way, we can resort to the power of parallel computation and greatly speed up the algorithm.
For more details of this procedure, please see Nakanishi et al.
For the task of compressing other type of data, one can choose different type of partitioning that is natural to the situation.
Our method is novel in that we allow the optimization of the support of π (ˆ z ( k )
| ˆ z (1: k − 1) ) .
We begin by modeling p ∗ (ˆ z ( k ) i | ˆ z (1: k − 1) ) with a truncated Gaussian distribution p ( z ( k ) i
| ˆ z (1: k − 1) ) on [ − 1 , 1] with parameters µ ( k ) i , σ ( k ) i for each i -th coordinate of z ( k ) using the information of ˆ z (1: k − 1) .
We then quantize [ − 1 , 1] into the intervals { I ( k ) i
[ j ] } M ki j =1 with the centers { A ( k ) i
[ j ] } M ki j =1 using the quantization width { q ( k ) i } M ki j =1 , where M ki is the number of the intervals.
The larger the value of
M ki , ﬁner the quantization.
By construction, we construct π (ˆ z ( k ) i
| ˆ z (1: k − 1) ) by deﬁning π (ˆ z ( k ) i | ˆ z (1: k − 1) ) : = M ki (cid:88) j = p
[Equation]
With our method, we can train the objective function about µ ( k ) i , σ ( k ) i , and q ( k ) i .
The ability to optimize q ( k ) i is a distinctive feature of our ASAP coding.
With this distinctive feature, our ASAP can compress z into a code of much smaller size under the same distortion constraint.
By the ratedistortion argument we mentioned in the introduction, this also means that our ASAP can store the data with much smaller distortion under the the same constraint on the code-length.
2. ASAP CODING FOR IMAGE COMPRESSION
Next, we will explain the technicalities for the optimization of the parameters µ ( k ) i , σ ( k ) i , and q ( k ) i in the context of image compression.
We will also explain the multiplicative loss, a novel objective function for the training of R .
2.2. A DAPTIVE CONSTRUCTION OF DISCRETIZED SAMPLE SPACE AND PROBABILITY DISTRIBUTION
Fig.
6 illustrates the construction of ˆ z ( k ) and ( A ( k ) , π ( k ) ) .
From now on, let us write π ( k ) i ( a ) :
=
π (ˆ z ( k ) i = a | ˆ z (1: k − 1) ) for a ∈ A ( k ) i .
In order to recursively construct these probability measures, we trained a set of neural networks that map ˆ z (1: k − 1) to the triplets θ ( k ) i :=
( µ ( k ) i , σ ( k ) i , q ( k ) i ) .
For the base level of k = 1 , we prepared a trainable triplets θ init m :=
( σ init m , µ init m , q init m ) for each m -th channel
[Figure]
latent feature quantized feature probability space for when training: when test: entropycoding probability density function support of and set θ (1) i
= θ init m for all ˆ z (1) i belonging to the channel m (please see Fig. 14 in the appendix for the detail).
We deﬁne A ( k ) i as
[Equation]
so that I ( k ) i [ j ] = (cid:104) A ( k ) i
[ j ] − 0 .
5 q ( k ) i , A ( k ) i
[ j ] + 0 .
5 q ( k ) i (cid:17) .
This indeed amounts to discretizing the interval [ − 1 , 1] using the subintervals of width q ( k ) i ; M ki is determined naturally from this constraint.
To each subinterval, we would assign the probability mass computed with the formula
[Equation]
[Equation] where p ( k ) i is a truncated normal distribution with the parameters ( µ ( k ) i , σ ( k ) i ) that is deﬁned over the domain (cid:104) a ( k ) i − 0 .
5 q ( k ) i , a ( k ) i + 0 .
5 q ( k ) i (cid:17) .
We are aware that this is not technically the faithful discretization of the interval [ − 1 , 1] ; this compromise is for ease of implementation.
If I ( k ) i perfectly partitions [ − 1 , 1] , the support of p ( k ) i will be exactly [ − 1 , 1] 1 .
By construction, our π ( k ) i satisﬁes the required properties of the probability mass function.
By appealing to the basic result of information theory, any signal sampled from the distribution π ( k ) i can theoretically be encoded in a code-length:
[Equation]
In order to map z ( k ) i to A ( k ) i , we simply pick the member of A ( k ) i that is closest to z ( k ) i ; this member would be the value we assign to ˆ z ( k ) i .
By construction, the quantization error ˆ z ( k ) i
[Equation]
is bounded from above by q ( k ) i / 2 .
Training of the neural models
In order to make the training of the neural networks possible, we would take advantage of the fact that we can always write
[Equation]
for some δ ( k ) i ∈ [ − 0 .
We can therefore interpret δ ( k ) i as a noise and regard the equation above as a probabilistic model for ˆ z ( k ) i with a trainable noise parameter q ( k ) i .
Indeed, the demand [Figure] for smaller distortion (loss of information) would prefer smaller q ( k ) i and hence smaller variance for noise.
On the contrary, the demand for shorter code-length would prefer coarser discretization of the sample space and hence larger q ( k ) i .
This is indeed the very manifestation of the rate-distortion tradeoff for the training of q ( k ) i .
2.2. M ULTIPLICATIVE RATE - DISTORTION TRADEOFF FUNCTION
In this section, we will ﬁnally describe the objective function by which we would train the ASAP model we introduced above.
In order to make comparisons, we would introduce our objective function along with its variation and the classic counterpart: [Equation] [Equation] [Equation]
method
[Equation]
(8) where bpp :=
[Equation]
(ˆ z ( k ) ; π ( k ) ) .
An apparent disadvantage of Eq.
(6) is that one must search for different appropriate λ for different constraint on the compression rate (code-length).
In order to do away with the computationally heavy parameter search, we took advantage of our observation that the relation bpp × (1 − MS-SSIM ) = const approximately holds in experiments, and ﬁrst adopted this function as a candidate for the objective function (Eq.
As we show in Fig.
7, our ASAP trained with this objective function performs almost as well as the ASAP trained with the Eq.
(6) with the optimal λ for all choice of the compression rate.
However, as we can see in Fig. 7, the MS-SSIM value for both Eqs.
(6) and (7) plateaus at some value less than 1 .
We resolved this issue by modifying Eq.
(7) to Eq.
This solution works because the inclusion of the MSE cost would create still another pressure to the training process.
By this modiﬁcation, we were able to produce good results in the region of large bpp as well.
R ELATED WORKS
Application of the neural networks to the compression of image dataset is popular these days.
Most studies choose (recurrent) autoencoder based on CNN as the main architecture.
By its construction, the coordinates of the output from such autoencoder will have strong spatial correlation.
This tendency should be particularly true for the network trained for the compression of the images.
Moreover, the local correlative relations among the coordinates can greatly vary across the locations and the channels, not to mention across the types of images.
In order to optimize the rate-distortion trade off, one must therefore strive to choose appropriate quantization width for different situations.
Indeed, this effort must be jointly planned and implemented with the optimization of the parameters of the autoencoder, because the the correlations amongst the coordinates of the feature vector is obviously dependent on the behavior of the autoencoder.
The work of Ball´e et al.
(2017) is one of the pioneer studies that trained PSQE together with the autoencoder.
Ball´e et al.
(2018) further improved their work by introducing an additional latent variable that governs the spatial relation among the quantization probability distributions of the latent spaces.
The optimization of the quantization process is a challenging problem in general, and numerous studies proposed different approaches to this problem.
Toderici et al.
; Johnston et al.
(2017); Ball´e et al.
(2017; 2018) tackled this problem by interpreting the quantization process probabilistically.
In these approaches, the use of probabilistic noise can de-stabilize the training process; this tends to be particularly true for variational inferences.
The works of Theis et al.
; Nakanishi et al.
; Mentzer et al.
(2018) fall in to this category.
Mentzer et al.
(2018) in particular deploys a set of trainable parameters
C
= { c 1 , · · · , c L } from which to sample ˆ z with the encoding rule ˆ z i = arg min j (cid:107) z i − c j (cid:107) .
The compression method of Mentzer et al.
(2018), however, trains C of ﬁxed length (i.e., the number of bins).
On the other hand, in our study, the number of bins is a function of quantization width.
Our method therefore optimizes the number of bins together with the quantization width.
Numerous types of objective functions are also used for the training of the network for compression task.
Most of the previous works like Ball ´ e et al.
(2017); Rippel & Bourdev (2017)
; Mentzer et al.
(2018) use the objective function of type (6).
There is also a variety of distortion measure, such as MS-SSIM, L 2 loss, L 1 loss, as well as their weighted mixtures.
Some studies also use GANs’ loss (Rippel & Bourdev, 2017; Agustsson et al., 2018).
Johnston et al.
(2017) uses a multiplicative loss of the form SSIM × L 1 .
As mentioned in the previous sections, our multiplicative loss is based on the observation that the empirical trade-off functions resemble a function of the form bpp × (1 − MS-SSIM )
= const .
It shall be emphasized that we were able to produce superior result using a single, hyper-parameter free objective function.
E XPERIMENT
We evaluated the performance of our method on the benchmark image datasets.
We trained our proposed model over 0.3M iterations with batchsize 25.
For the optimization, we used Adam (Kingma & Ba, 2015) with AMSGrad (Reddi et al., 2018).
As for the hyper-parameters of Adam, we set α = 0 .
001 , β 1 = 0 .
9 , β 2 = 0 .
999 , and linearly decayed α from the 0.225M-th iteration so that it will reach 0 at the end of the training.
For the update of the parameters, we applied L 2 regularization with weight decay of rate 0 . 00001 , and clipped the magnitude of each coordinate of the gradient at the value of 5 .
As a similarity measure between two images, we used MS-SSIM (Wang et al., 2004), which is considered by the community to have close correlation with the subjective evaluation of humans.
Because MS-SSIM is originally intended for the evaluation of gray scale measures, we computed the MS-SSIM separately for each of RGB channels and reported their average.
For the training, we used ILSVRC2012 dataset (ImageNet)
(Russakovsky et al., 2015).
In order to produce the rate distortion (RD) curve, we trained a model for each choice of C , the number of channels for the latent feature.
It shall be reminded that C affects the code-length via entropy coding.
For the performance evaluation, we used Kodak 2 dataset and RAISE-1k dataset (Dang-Nguyen et al., 2015).
4. P ERFORMANCE
Kodak dataset is a widely used dataset for the performance evaluation of compression methods.
Fig.
2 compares the RD curves for different compression methods.
In terms of MS-SSIM, our method was slightly better than the method trained with the best ﬁxed quantization size, which was found through an extensive grid search.
Again in terms of MS-SSIM, our method outperformed the model of Ball´e et al. (2018) trained to optimize MS-SSIM (“Balle for MS-SSIM”), as well as the same model trained to optimize mean squared error (MSE) (“Balle for MSE”). In terms of PSNR, our method performed better than “Balle for MS-SSIM”, not to mention JPEG. However, it was not able to perfrom better than “Balle for MSE” in terms of PSNR.
[Figure] For the visuals of the reconstructed images, please see Figs. 1 and 16–18. Figs.
19–24 illustrate regions of medium-high bpp. Our method outperforms BPG on the medium-high bpp region.
Fig.
8 illustrates the heat map of ( µ ( k ) i , σ ( k ) i , q ( k ) i ) produced for a Kodak image by a trained ASAP model.
The signals from each z ( k ) was re-aligned in the way that respects the original geometrical relations prescribed by the pattern in Fig.
A pair of stacked images is a pair of images from two different channels.
Fig.
8 features several interesting properties of our ASAP coding.
Firstly, as we can readily see in the ﬁgure, the ‘bottom’ channel is capturing more contextual information than the ‘top’ channel.
Note that our ASAP is assigning larger values of q in general to the bottom channel than to the top channel.
We can also see that the algorithm is assigning larger value of q to the region within the picture with relatively small information.
Lastly, the value of σ assigned by the algorithm to each coordinate is reﬂective of the algorithm’s conﬁdence for the compression of the information on the coordinate.
By construction, z (1) is a group for which the algorithm would have least conﬁdence because it is the base level of the recursive inference.
The bright points (representing large value) scattered across the heat map in grid-pattern correspond to z (1) .
Kodak dataset consists of only 24 images.
To rule out the possibility that our model is overﬁtted to the Kodak dataset by chance, we also evaluated the performance of our models using RAISE-1k dataset that consists of 1000 raw images (Fig. 15).
We can conﬁrm the superiority of our method over all the current contemporaries on RAISE-1k dataset as well.
It shall be noted that our method is outperforming other methods not only in the low bit rate region, but also in the high bit rate region as well.
4. T HE EFFECTIVENESS OF THE MULTIPLICATIVE LOSS
We conducted still another comparative experiment in order to evaluate the effectiveness of using the multiplicative loss.
Fig.
7 shows the RD curves of our model trained with the three different objective functions we introduced in Sec. 2.2.
(Eqs.
We used the same experimental setup as the one we used in Sec.
4.2 except that we stopped the iteration at 0.1M-th step and initiated the linear decay of α from 0.075M-th step.
To be more fair, we tested the performance of Eq.
(6) using multiple good values of λ .
We can conﬁrm in Fig.
7 that the multiplicative loss outperforms all other loss functions.
We shall also note that the model trained to optimize our multiplicative loss yielded the results with higher MS-SSIM than the model trained to directly optimize the MS-SSIM loss.
We shall emphasize that our multiplicative loss does not require the hyperparameter search about λ for each bit rate.
4. A BLATION STUDY WITH SMALL ASAP MODEL
We conducted an ablation study to assess the effect of making the quantization width adaptive for each latent feature.
We conducted the compression with a ﬁxed quantization width for all latent features and compared the result with our ASAP coding. For the size of the the common quantization width, we conducted a grid search and selected the width that achieved the best MS-SSIM score. Figure 10a compares the result of our algorithm against those of different choices of ﬁxed quantization width. We can see in the plot that our algorithm performs equally well with the result with the optimal ﬁxed quantization width.
Co-incidentally, there was a choice of ﬁxed quantization width that works equally well with the adaptive quantization width.
We also conducted a set of experiments with a smaller version of our model, in which we reduced the number of channels for each convolution layer and replaced each residual block with a single convolution layer.
The ﬁgure 10b shows the results of the experiments with the small model.
As we can see in the ﬁgure, the adaptive quantization width performs better than all choices of ﬁxed quantization width we tested.
All in all, we were able to conﬁrm through the ablation study that our adaptive quantization width performs equally well or better in comparison to the optimal ﬁxed width quantization that was found through a grid-search.
C ONCLUSION
In this study, we proposed ASAP, a novel lossy compressor based on a neural architecture that allows the training of quantization width.
We also proposed multiplicative loss, a novel hyper- [Figure] parameter free objective function that can be used to train the compressor with better rate-distortion tradeoff.
Our study has shown that there is still a room for improvement in the neural network based approaches to compression, which once seemed to have reached its limit.
The application of our ASAP is not be limited to image compression.
The framework of ASAP and the idea of the multiplicative loss shall be applicable to various industrial needs in which the rate distortion tradeoff is an important factor. B A PPENDIX RESULTS A B
A
B
{+C=2 46 
+}
JPEG
BPG
Proposal
[Figure]
Original
[Figure]
JPEG
BPG
[Figure]
Original
[Figure]
{+C=
+}
JPEG
[Figure]
Original
[Figure]
{+C=
+}
[Figure]
Original
{+C=
+}
[Figure]
Original
{+C=
+}
[Figure]
Original
{+C=
+}
JPEG
[Figure]
Proposal
[Figure]
Original
{+C=
+}
[Figure]
{+C=
+}
JPEG
BPG
Proposal
Original
[Figure]