Adaptive Sample-space & Adaptive Probability coding: a neural-network based approach for compression

1 INTRODUCTION . In terse terms , lossy data compression is a task in which one seeks to encode a data file into as short a code as possible without losing the essential information . Extensive research has been conducted in the field 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 field . For example , studies like Toderici et al . ( 2015 ; 2017 ) ; Ballé et al . ( 2017 ) ; Theis et al . ( 2017 ) ; Johnston et al . ( 2017 ) ; Rippel & Bourdev ( 2017 ) ; Mentzer et al . ( 2018 ) ; Agustsson et al . ( 2018 ) ; Nakanishi et al . ( 2018 ) ; Ballé 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 confirm that this adaptive quantization width benefits 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 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 file 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 ) . 2 METHODS . 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 flow of our algorithm in the context of image compression . In a nutshell , our method follows the usual routine of transform coding . Throughout , let x ∈ RC0×H0×W0 be the original image , where C0 , H0 , W0 respectively represent the number of the channels , the height and the width of the image . The algorithm begins by first applying the transformation function F to x in order to obtain a latent feature representation , z ∈ RC×H×W . The algorithm then quantizes the signal z into a finite 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 flow . 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.1 THE 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̂ ∈ RC×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 . ( 2018 ) . We partitioned ẑ into K groups { ẑ ( k ) ∈ RC ( k ) ×H ( k ) ×W ( k ) ; k = 1 , ... , K } based on a geometric pattern we illustrate in Fig . 5 . 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 ) ) ∏K k=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 . ( 2018 ) . 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 ] } Mki j=1 with the centers { A ( k ) i [ j ] } Mki j=1 using the quantization width { q ( k ) i } Mki j=1 , where Mki is the number of the intervals . The larger the value of Mki , finer the quantization . By construction , we construct π ( ẑ ( k ) i |ẑ ( 1 : k−1 ) ) by defining π ( ẑ ( k ) i |ẑ ( 1 : k−1 ) ) : = Mki∑ j=1 p ( z ( k ) i ∈ I ( k ) i [ j ] |ẑ ( 1 : k−1 ) ) 1 z ( k ) i =A ( k ) i [ j ] . ( 1 ) 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.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.1 ADAPTIVE 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 θinitm : = ( σ init m , µ init m , q init m ) for each m-th channel 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 define A ( k ) i as A ( k ) i : = { a ∣∣∣a = µ ( k ) i + ( u+ 0.5 ) q ( k ) i , u ∈ Z , a− 0.5q ( k ) i < 1 , a+ 0.5q ( k ) i ≥ −1 } ( 2 ) so that I ( k ) i [ j ] = [ A ( k ) i [ j ] − 0.5q ( k ) i , A ( k ) i [ j ] + 0.5q ( k ) i ) . This indeed amounts to discretizing the in- terval [ −1 , 1 ] using the subintervals of width q ( k ) i ; Mki is determined naturally from this constraint . To each subinterval , we would assign the probability mass computed with the formula π ( k ) i ( a ) : = ∫ a+0.5q ( k ) i a−0.5q ( k ) i p ( k ) i ( x ) dx , ( 3 ) where p ( k ) i is a truncated normal distribution with the parameters ( µ ( k ) i , σ ( k ) i ) that is defined over the domain [ a ( k ) i − 0.5q ( k ) i , a ( k ) i + 0.5q ( k ) i ) . 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 satisfies 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 : ` ( ẑ ( k ) ; π ( k ) ) : = − C ( k ) H ( k ) W ( k ) ∑ i=1 log2 π ( k ) ( ẑ ( k ) i |ẑ ( 1 : k−1 ) ) . ( 4 ) 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 − z ( k ) i 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 ẑ ( k ) i = z ( k ) i − δ ( k ) i q ( k ) i . ( 5 ) for some δ ( k ) i ∈ [ −0.5 , 0.5 ) . 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 1In fact , in order to avoid the creation of the interval with too small a mass , we have occasionally augmented p ( k ) i and renormalized the integrals . 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 .

Summary: The paper proposes Adaptive Sample-space & Adaptive Probability (ASAP) coding for image compression based on neural networks. In contrast to most prior methods, which adhere to a fixed quantization scheme (i.e. with fixed number of quantization levels, and fixing the level themselves), the proposed method jointly learns a probability model of the quantized representation (the bottleneck of an autoencoder model) for coding and a corresponding adaptive quantization scheme. The distribution of each entry in the bottleneck before quantization is modeled as a Gaussian, whose mean and variance are predicted by a neural network conditionally on bottleneck entries on a grid at different scales (similar as in Nakanishi et al. 2018). The same network also predicts quantization intervals to adaptively quantize the respective entry of the bottleneck. Together, the predicted means, variances, and quantization intervals are used to obtain an estimate of the code length. The proposed compression networks are trained with a novel multiplicative loss, showing clear improvements over prior methods Rippel & Bourdev 2017, Nakanishi et al. 2018 on the Kodak and Raise1k data sets in terms of MS-SSIM.