Learning Accurate Entropy Model with Global Reference for Image Compression

1 INTRODUCTION . Image compression is a fundamental research topic in computer vision . The goal of image compression is to preserve the critical visual information of the image while reducing the bit-rate for storage or transmission . The state-of-the-art image compression standards , such as JPEG ( Wallace , 1992 ) , JPEG2000 ( Rabbani & Joshi , 2002 ) , HEVC/H.265 ( Sullivan et al. , 2012 ) and Versatile Video Coding ( VVC ) ( Ohm & Sullivan , 2018 ) , are carefully engineered and highly tuned to achieve better performance . Albeit widely deployed , the conventional human-designed codecs take decades of development to achieve impressive compression rate today . Any further improvement is expected to be even more difficult . Inspired by the successful stories of deep learning in many vision tasks , several pioneer works ( Toderici et al. , 2016 ; Agustsson et al. , 2017 ; Theis et al. , 2017 ; Ballé et al. , 2017 ; Ballé et al. , 2018 ; Mentzer et al. , 2018 ; Lee et al. , 2019 ; Minnen et al. , 2018a ) demonstrate that the image compression task can be effectively solved by deep learning too . This breakthrough allows us to use data-driven learning system to design novel compression algorithms automatically . As a result , a majority of deep image compression ( DIC ) models are based on autoencoder framework . In this framework , an encoder transforms pixels into a quantized latent representation suitable for compression , while a decoder is jointly optimized to transform the latent representation back into pixels . ∗Corresponding author . The latent representation can be losslessly compressed to create a bitstream by using entropy coding method ( Rissanen & Langdon , 1981 ) . In the entropy coding , the compression quality is controlled by the entropy estimation of latent features generated by the encoder . It is therefore important to learn an accurate entropy model . To this end , several solutions have been considered . With additional bits , some methods propose entropy model conditioned on a hyperprior , using side information of local histograms over the latent representation ( Minnen et al. , 2018b ) or a hierarchical learned prior ( Ballé et al. , 2018 ) . Contextadaptive models ( Minnen et al. , 2018a ; Lee et al. , 2019 ) incorporate predictions from neighboring symbols to avoid storing the additional bits . While these methods improve the accuracy of the entropy models , they are unable to use global context information during the compression , leading to suboptimal performance . In this work , we observe that global spatial redundancy remains in the latents , as shown in Figure 1 . Motivated by this , we propose to build up a global relevance throughout the latents . Inspired by the recent reference-based Super-Resolution ( SR ) methods ( Zheng et al. , 2018 ; Yang et al. , 2020 ) , we empower the entropy model with global vision by incorporating a reference component . Unlike the super-resolution scenario , incorporating global reference information is non-trivial in deep image compression . The image during decoding is often incomplete which means that the information is badly missing . Besides , our target is to reduce the bit rates and recover the image from the bitstream faithfully , rather than inpainting a low-resolution image with vivid generated details . To address the above challenges , in our proposed method , a global reference module searches over the decoded latents to find the relevant latents to the target latent . The feature map of the relevant latent is then combined with local context and hyperprior to generate a more accurate entropy estimation . A key ingredient in the global reference ensemble step is that we consider not only the similarity between the relevant and the target but also a confidence score to measure the high-order statictics in the latent feature distribution . The introduction of the confidence score enhances the robustness of the entropy model , especially for images with noisy backgrounds . Also , we found that the widely used Generalized Divisive Normalization ( GDN ) in image compression suffers from a mean-shifting problem . Since the GDN densities are zero-mean by definition , mean removal is necessary to fit the density ( Ballé et al. , 2016b ) . Therefore we propose an improved version of GDN , named GSDN ( Generalized Subtractive and Divisive Normalization ) to overcome this difficulty . We summarize our main contributions as follows : • To the best of our knowledge , we are the first to introduce global reference into the entropy model for deep image compression . We develop a robust reference algorithm to ensemble local context , global reference and hyperprior in a novel architecture . When estimating the latent feature entropy , both similarity score and confidence score of the reference area are considered to battle with the noisy background signals . • We propose a novel GSDN module that corrects the mean-shifting problem . • Experiments show that our method outperforms the most advanced codes available today on both PSNR and MS-SSIM quality metrics . Our method saves by 6.1 % compared to the context-adaptive deep models ( Minnen et al. , 2018b ; Lee et al. , 2019 ) and as much as 21.0 % relative to BPG ( Bellard. , 2014 ) . The remainder of this work is organized as follows . In Section 2 , we introduce the backbone of the end-to-end deep image compression network as well as the reference-based component for the entropy model . Section 3 demonstrates the structure of our combined entropy model . The GSDN with mean-shifting correction is given in Section 4 . We present experimental comparison and visualization in Section 5 . Finally , we enclose this work with an open discussion in Section 6 . 2 LEARNED IMAGE COMPRESSION . Learned image compression using deep neural networks has attracted considerable attention recently . The work of Toderici et al . ( 2016 ) first explored a recurrent architecture using an LSTMbased entropy model . A wide range of models ( Ballé et al. , 2017 ; Ballé et al. , 2018 ; Mentzer et al. , 2018 ; Minnen et al. , 2018a ; Lee et al. , 2019 ; Hu et al. , 2020 ; Cheng et al. , 2020 ) used a CNN-based autoencoder with constrained entropy . General learned image compression consists of an encoder , a quantizer , a decoder , and an entropy model . An image x is transformed into a latent representation y via the encoder ga ( x ) , which is discretized by the quantizer Q ( y ) to form ŷ . Given the entropy model pŷ , the discretized value ŷ can be compressed into a bitstream using entropy coding techniques such as arithmetic coding ( Rissanen & Langdon , 1981 ) . The decoder gs ( ŷ ) then forms the reconstructed image x̂ from the quantized latent representation ŷ , which is decompressed from the bitstream . The training goal for learned image compression is to optimize the trade-off between the estimated coding length of the bitstream and the quality of the reconstruction , which is a rate-distortion optimization problem : L = R+ λD = Ex∼px [ − log2 pŷ ( Q ( ga ( x ) ) ) ) ] + λEx∼px [ d ( x , gs ( ŷ ) ] , ( 1 ) where λ is the coefficient which controls the rate-distortion trade-off , px is the unknown distribution of natural images . The first term represents the estimated compression rate of the latent representation . The second term d ( x , x̂ ) represents the distortion value under given metric , such as mean squared error ( MSE ) or MS-SSIM ( Wang et al . ( 2003 ) ) . Entropy coding relies on an entropy model to estimate the prior probability of the latent representation . Ballé et al . ( 2017 ) propose a fully factorized prior for entropy estimation as shown in Figure 2 ( a ) , while the prior probability of discrete latent representations is not adaptive for different images . As shown in Figure 2 ( b ) , Ballé et al . ( 2018 ) model the latent representation as a zero-mean Gaussian distribution based on a spatial dependency with additional bits . In Lee et al . ( 2019 ) and Minnen et al . ( 2018a ) , they introduce an autoregressive component into the entropy model . Taking advantage of high correlation of local dependency , context-adaptive models contribute to more accurate entropy estimation . However , since their context-adaptive entropy models only capture the spatial information of neighboring latents , there is redundant spatial information across the whole image . To further remove such redundancy , our method incorporates a reference-based model to capture global spatial dependency . Specially for learned image compression , a generalized divisive normalization ( GDN ) ( Ballé et al. , 2016a ) transform with optimized parameters has proven effective in Gaussianizing the local joint statistics of natural images . Unlike many other normalization methods whose parameters are typically fixed after training , GDN is spatially adaptive therefore is highly nonlinear . As the referencebased model calculates relevance over the latents , it is crucial to align the distribution of the latents . To better align the latents , the proposed GSDN incorporates a subtracting factor with GDN . We also present an effective method of inverting it when decompressing the latent representation back to image . 3 COMBINED LOCAL , GLOBAL AND HYPERPRIOR ENTROPY MODEL . The models we analyze in this paper build on the architecture introduced in Minnen et al . ( 2018a ) , which combined an autoregressive model with the hyperprior . Figure 3 provides a high-level overview of our approach . The compression model contains two main sub-networks . The first is the core autoencoder , which learns the transform and the inverse transform between image and latent representation . Q represents the quantization function . The gradient-based optimization in learned methods is hindered by quantization . Here , we make use of a mixed approach that has proven efficient in Minnen & Singh ( 2020 ) . The second sub-network is the combined entropy model , which is responsible for estimating a probabilistic model over the latents for entropy coding . The combined entropy model consists of a context model , a reference model , and a hyper-network ( hyper encoder and hyper decoder ) . The three components are combined progressively . Then three parameter networks generate the mean and scale parameters for a conditional Gaussian entropy model respectively . Following the work of Minnen et al . ( 2018a ) , we model each latent , ŷi , as a Gaussian with mean µi and deviation σi convolved with a unit uniform distribution : pŷ ( ŷ|ẑ , θ ) = ∏ i=1 ( N ( µi , σ2i ) ∗ U ( −0.5 , 0.5 ) ) ) ( ŷi ) ( 2 ) where µ and σ are the predicted parameters of entropy model , ẑ is the quantized hyper-latents , θ is the entropy model parameters . The entropy model for the hyperprior is the same as in Ballé et al . ( 2018 ) , which is a non-parametric , fully factorized density model . As the hyperprior is part of the compressed bitstream , we extend the rate of Equation 1 as follows : R = Ex∼px [ − log2 pŷ ( ŷ ) ] + Ex∼px [ − log2 pẑ ( ẑ ) ] ( 3 ) The compressed latents and the compressed hyper-latents are part of the bitstream . The reference-based SR methods ( Zheng et al. , 2018 ; Yang et al. , 2020 ) adopt “ patch match ” to search for proper reference information . However , in the serial processing of image compression , the latent representation during decoding is often incomplete . We extend this search method by using a masked patch . Figure 4 illustrates how the relevance embedding module estimates similarity and fetches the relevant latents . When decoding the target latent , we use neighboring latents ( left and top ) as a basis to compute the similarities between the target latent and its previous latents . Particularly , the latents are unfolded into patches and then masked , denoted as q ∈ [ H×W , k×k×C ] ( whereH , W , k , C correspond to height , width , unfold kernel size and channels , respectively ) . We calculate the similarity matrix r ∈ [ H×W , H×W ] throughout the masked patches by using cosine similarity , ri , j = 〈 qi ‖qi‖ , qj ‖qj‖ 〉 ( 4 ) Similarity 𝑦 '' 𝑦 '' Mask Conv 𝑡 ! ̂ Relevant Index S Figure 4 : A mask slide patch searches on all the decoded latents ( tan area ) . The relevant latents are fetched and learned with a masked convolution . Note that we can only see the decoded latents , so the lower triangular of the similarity matrix is set to zero . We get the most relevant position for each latent as well as the similarity score . According to the position , we fetch the neighboring latents ( left and top ) as well as the center latent , which is named as “ relevant latents ” . We use a masked convolution as in Van den Oord et al . ( 2016 ) to transfer the relevant latents . To measure how likely a reference patch perfectly matches the target patch , Yang et al . ( 2020 ) propose a soft-attention module to transfer features by using the similarity map S. However , we found that a similarity score is not sufficient to reflect the quality of reference latent in image compression . For this reason , a confidence score is introduced to measure the texture complexity of the relevant latent . We use the context model to predict the Gaussian parameters ( i.e. , µ1 , σ1 ) of latents solely . The latents ŷ are now modeled as Gaussian with mean µ1 and standard deviation σ1 . The probabilities of the latents are then calculated according to ( µ1 , σ1 ) as in Equation 2 . As reference model is designed in spatial dimension , the confidence map U is obtained by averaging the probabilities across channel . With the above two parameters , the more relevant latent combination would be enhanced while the less relevant one would be relived . The similarity S and the confidence U are both 2D feature maps . Figure 5 provides the structure of our combined entropy model . For the context model , we transfer the latents ( i.e. , ŷ ) with a masked convolution . For the reference model , we transfer the unfolded relevant latents with a masked convolution . We use 1 × 1 convolution in the parameter networks . Local , global and hyperprior features are ensembled stage by stage , as well as the predicted Gaussian parameters . The mean parameters are estimated by the context model first , and then updated by the global model and the hyperprior model . We use the Log-Sum-Exp trick for resolving the under or overflow issue of deviation parameters . The output of global reference is further multiplied by the similarity S and the confidence U . The context model is based on the neighboring latents of the target latent to reduce the local redundancy . From the perspective of the global context , the reference model makes further efforts to capture spatial dependency . As the first two models predict by the decoded latents , there exists uncertainty that can not be eliminated solely . The hyperprior model learns to store information needed to reduce the uncertainty . This progressive mechanism allows an incremental accuracy for distribution estimation .

Summary: This paper propose two methods for improve deep image compression performance: (i) Global Reference Module and (ii) Mean-shifting GDN Module (GSDN). (i) Global Reference Module searches over the decoded latents to find the relevant latents to the target latent for improve accuracy of entropy estimate. Authors extended Yang et al. 2020 method to using masked patch. (ii) GSDN extends GDN to use subtractive operation.