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 paper addresses the problem of learned image compression using a transformer as the entropy model. The authors introduce a diamond-shaped relative position encoding scheme that makes sense for image modeling. They also adopt a two-step, bidirectional context model based on a checkerboard-style spatial decomposition of the latent tensor.  The modeling choices are well-supported by an ablation study. In particular, the authors show that the checkerboard model is almost as good as a much slower spatial context model that is often used for learned image compression. They also show that the transformer-based entropy model ("entroformer") leads to a model with better rate-distortion performance compared to earlier models.