Entroformer: A Transformer-based Entropy Model for Learned Image Compression

1 INTRODUCTION . Image compression is a fundamental research field in computer vision . With the development of deep learning , learned methods have led to several breakthroughs in this task . Currently , the state-ofthe-art ( SOTA ) deep image compression models are built on the auto-encoder framework ( Hinton & Salakhutdinov , 2006 ) with an entropy-constrained bottleneck ( Theis et al. , 2017 ; Ballé et al. , 2017 ; Ballé et al. , 2018 ; Mentzer et al. , 2018 ; Minnen et al. , 2018a ; Lee et al. , 2019 ; Guo et al. , 2021 ) . An entropy model estimates the conditional probability distribution of latents for compression by standard entropy coding algorithms . In this paper , we focus on improving the predictive ability of the entropy model , which leads to higher compression rates without increasing distortion . One of the main challenges of the entropy model is how to learn representation for various dependencies . For example , as shown in Figure 1 , the edges have long-range dependencies ; and the five hats are likely to be related in shape ; the region of the sky has identical color . Some works use local context ( Minnen et al. , 2018a ; Lee et al. , 2019 ; Mentzer et al. , 2018 ) or additional side information ( Ballé et al. , 2018 ; Hu et al. , 2020 ; Minnen et al. , 2018b ) as short-range spatial dependencies , while others use non-local mechanism ( Li et al. , 2020 ; Qian et al. , 2021 ; Cheng et al. , 2020 ; Chen et al. , 2021 ) to capture long-range spatial dependencies . However , the constraint of capturing long-range spatial dependencies still remains in these CNN-based methods . Another main challenge is the trade-off between performance and decoding speed . Though previous context based methods use unidirectional context model to improve the predictive ability , it suffers from slow decoding speed . It decodes symbols in a raster-scan order with O ( n ) serial process that can not be accelerated by modern GPUs . A two-pass parallel context model ( He et al. , 2021 ) is introduced for acceleration , which decodes symbols in a particular order to minimize serial processing . However , this parallel context model uses a weak context information , which degrades the compression performance . In this paper , we propose a novel transformer-based entropy model termed as Entroformer to address the above two challenges in one shot . Following the Transformer ( Vaswani et al. , 2017 ) , selfattention is used to relate different positions of a single latent in order to compute a representation of the latents . In Entroformer , spatial and content based dependencies are jointly taken into account in both hyperprior and context model . To further optimize Entroformer for image compression , the following designs are innovated : ( 1 ) Entroformer uses a multi-head attention with a top-k selection to distill information in representation subspaces : The multi-head attention provides a learning-based partition to learn dependencies in different views ; and the top-k scheme filters noisy signals by select the most similar k latents to the self-attention module , which is crucial for the convergence of Entroformer . ( 2 ) To inherit the local bias of CNNs , a novel position encoding unit is designed to provide better spatial representation for the image compression . This position encoding is based on a relative position encoding with a diamond-shape boundary ( Diamond RPE ) . Based on a empirical evaluation about the rate effect of relative position , Diamond RPE takes full advantage of prior information for image compression . ( 3 ) The two-pass decoding framework ( He et al. , 2021 ) is utilized to speed up the decoding of Entroformer . A bidirectional context with long-range context is introduced instead of the local checkeboard context , which helps counteract the performance degradation . In summary , our contributions include : • In this paper , Entroformer is proposed to improve learned image compression . To our knowledge , it is the first successful attempt to introduce Transformer based method to image compression task . Experiments show that our method outperforms the most advanced CNNs methods by 5.2 % and standard codec BPG ( Bellard. , 2014 ) by 20.5 % at low bit rates . • To further optimize Entroformer for image compression , the multi-head scheme with top-k selection is proposed to capture precise dependencies for probability distribution estimation , and the diamond relative position encoding ( diamond RPE ) is proposed for better positional encoding . Experiments show that these methods support image compression task to keep training stable and obtain superior results . • With the help of bidirectional context and two-pass decoding framework , our parallel Entroformer is more time-efficient than the serialized one on modern GPU devices without performance degradation . 2 COMPRESSION WITH HYPERPRIOR AND CONTEXT . There are types of entropy models such as hyperprior ( Minnen et al. , 2018b ; Ballé et al. , 2018 ) , context model ( Li et al. , 2018 ; Mentzer et al. , 2018 ) , and the combined method ( Lee et al. , 2019 ; Minnen et al. , 2018a ; Minnen & Singh , 2020 ) . Hyperprior methods typically make use of side information of the quantized latent representation with additional bits . Context models learn autoregressive priors incorporating prediction from the causal context of the latents , which is a bit-free . Our Entroformer combines a context model with hyperprior based on previous works ( Minnen et al. , 2018b ; Ballé et al. , 2018 ) . We describe the architecture of compression model with Entroformer in Figure 2 . The main autoencoder learns a quantized latent representation ŷ of image x . The reconstructed image is denoted as x̂ . The hyper-autoencoder learns a quantized hyper-latents representation ẑ . Following the work of Ballé et al . ( 2016 ) , the quantization operation is approximated by an additive uniform noise during training . This ensures a good match between encoder and decoder distributions of both the quantized latents , and continuous-valued latents subjected to additive uniform noise during training . 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 ) ( 1 ) where µ and σ are predicted by the entropy model , θ is the parameters of the entropy model . The entropy model consists of hyper-autoencoder , context model , and parameter networks . Since we do not make any assumptions about the distribution of the hyper-latents , a non-parametric , fully factorized density model is used . 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̂ ( ŷ ) ] ︸ ︷︷ ︸ rate ( latents ) +Ex∼px [ − log2 pẑ ( ẑ ) ] ︸ ︷︷ ︸ rate ( hyper-latents ) +λ · Ex∼px ||x− x̂||22︸ ︷︷ ︸ distortion , ( 2 ) where λ is the coefficient which controls the rate-distortion trade-off , px is the unknown distribution of natural images . pẑ use a fully factorized density model . The first term represents the estimated compression rate of the latent representation , while the second term represents the estimated compression rate of the hyper-latent representation . The third term represents the distortion value under given metric , such as mean squared error ( MSE ) . 3 TRANSFORMER-BASED ENTROPY MODEL . Our model contains two main components , a CNN-based autoencoder to learn a latent representation , an transformer-based entropy model to predict latents . We describe our architecture in Figure 2 ( a ) , and our Entroformer in detail in Figure 2 ( b ) . In this section , we first propose two ingredients , a diamond relative position encoding ( diamond RPE ) and a top-k scheme , which are essential for image compression . Then , we extend the checkboard context model ( He et al. , 2021 ) to a parallel bidirectional context model . 3.1 TRANSFORMER ARCHITECTURE . In entropy model design we follow the original Transformer ( Vaswani et al. , 2017 ) , which employs an encoder-decoder structure . We model the latent representation ŷ ∈ RH×W×C as a sequence ŷp ∈ R ( H·W ) ×C , where ( H , W ) is the resolution of the latents , C is the number of channels . The latents sequence is then mapped to D dimensions with a trainable linear projection . For hyperprior , we stack N transformer encoder layers to yield hyper-latent z , which is quantized and fed to N transformer-encoder layers to generate hyperprior . To produce a hierarchical representation , the resolution of feature is changed by downsample and upscale module . For autoregressive prior , we stack 2N transformer decoder layers to generate autoregressive features . To generate the Gaussian parameters , a linear layer is attached to the combination of hyperprior features and context features . For any sequence of length n , the vanilla attention in the transformer is the dot product attention ( Vaswani et al. , 2017 ) . Following the standard notation , we compute the matrix of outputs via the linear projection of the input matrix X ∈ Rn×dm : Attention ( X ) = softmax ( XWQ ( XWK ) T√ dk ) XWV ( 3 ) WQ , WK ∈ Rdm×dk and WV ∈ Rdm×dv are learned parameter matrices . 3.2 POSITION ENCODING . To explore the impact of position in image compression , we first design a transformer-based entropy model without any position encoding . We feed a random mask to self-attention during training . During testing , we evaluate the effect of each position i by employing a corresponding mask , which are set to 1 except position i . Figure 3 plots the impact in bit rate of each position , compared to the result with context of all positions . This result highlights how the rate effected by the position of context . This observation presents a comprehensive understanding and provides empirical guideline for new position encoding design . In order for spatial information of the latents , one common approach is to use biased attention weights based on relative relation ( Shaw et al. , 2018 ) . Based on the result in Figure 3 , we propose an extension to this relation-aware self-attention to consider the element position influence on image compression performance . Each attention operates on an input sequence , x = ( x1 , ... , xn ) of n elements where xi ∈ Rdm . The edge between input elements xi and xj is represented by vectors pKij ∈ Rdk . We modify the dot-product attention matrix in eq . 3 to a compatibility function that compares two input elements in relation-aware self-attention : eij = xiW Q ( xjW K + pKij ) T √ dk ( 4 ) The weight coefficient , A , is computed using a softmax on dot-product attention matrix . On image plane , the relative position ai−j is a 2D coordinate . And we modify the maximum absolute relative position to a 2D boundary with diamond shape . As shown in Figure 3 , we observe that the closer contexts latent save more bit rate in context model . That tells how the distance of context latents influences the bit saving of context modeling in learned image compression . Therefore , we consider a 2D boundary with a maximum Hamming distance value of h. pKij = w K clip ( ai−j , h ) clip ( ai−j , h ) = { ai−j ‖ai−j‖1 < = h ( h , h ) otherwise . ( 5 ) We then learn the relative position encoding wK = ( wK−h , −h , w K −h , −h+1 , ... , w K h , h−1 , w K h , h ) .

Summary: This work introduces a transformer-based entropy coding model for learned image compression. The backbones of the commonly used hyperprior encoder and decoder are replaced with transformer encoder layers, and the context model is replaced with the transformer decoder. The striking features of the proposed method include (1) a diamond-shaped relative position encoding with clipping, (2) a top-k selection scheme, and (3) a parallel bi-directional context model. The parallel bi-direction context model is a derivative work from He et al. 2021.