Year: 2021

Abstract:
One critical component in lossy deep image compression is the entropy model, which predicts the probability distribution of the quantized latent representation in the encoding and decoding modules. Previous works build entropy models upon convolutional neural networks which are inefficient in capturing global dependencies. In this work, we propose a novel transformer-based entropy model, termed Entroformer, to capture long-range dependencies in probability distribution estimation effectively and efficiently. Different from vision transformers in image classification, the Entroformer is highly optimized for image compression, including a top-k self-attention and a diamond relative position encoding. Meanwhile, we further expand this architecture with a parallel bidirectional context model to speed up the decoding process. The experiments show that the Entroformer achieves state-of-the-art performance on image compression while being time-efficient.

heading: 1 INTRODUCTION
text: 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.

heading: 2 COMPRESSION WITH HYPERPRIOR AND CONTEXT
text: 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).

heading: 3 TRANSFORMER-BASED ENTROPY MODEL
text: 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.

heading: 3.1 TRANSFORMER ARCHITECTURE
text: 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.

heading: 3.2 POSITION ENCODING
text: 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) .

heading: 3.3 TOP-K SCHEME IN SELF-ATTENTION
text: The vanilla attention in the transformer is the scaled dot-product attention. It computes the dot products of the query with all keys, divide each by √ dk, and apply a softmax function to obtain the weights on the values.
Instead of computing the attention matrix for all the query-key pairs as in vanilla attention, we select the top-k most similar keys for each query. The queries, keys and values are packed together into matrices Q ∈ Rn×dk , K ∈ Rn×dk and V ∈ Rn×dv . P ∈ Rn×dk denotes relative position encodings described in Sec 3.2. The row-wise top-k elements in attention matrix are selected for softmax function:
V ′ = softmax ( fk ( Q(K + P )T√
dk
)) V (6)
where fk(·) denotes the row-wise top-k selection operator:
fk (eij) = { eij if eij is within the top-k largest elements in row j −∞ otherwise. (7)
Albeit easing sequence length mismatching problem by relative position encoding, transformerbased entropy model also suffers from token quantity mismatching in self-attention. For extremely larger sequence length during inference than training, a large quantity of irrelevant context latents would push the relevant ones to a small weight. To counteract this effect, we modify the selfattention with top-k scheme in the Entroformer. It helps to distill irrelevant context latents and stable training procedure.

heading: 3.4 PARALLEL BIDIRECTIONAL CONTEXT MODEL
text: The transformer architecture can be used to efficiently train autoregressive model by masking the attention computation such that the i-th position can only be influenced by a position j if and only if j < i. This causal masking uses a limited unidirectional context, and makes autoregressive model impossible to be parallel.
Following the work of He et al. (2021), the latents are split into two slices along the spatial dimension. Figure 4 provides a high-level overview of this architecture. The decoding of the first slice (ŷ1) is run in parallel by using a Gaussian entropy model conditioned solely on the hyperprior. Then the decoded latents become visible for decoding the second slice (ŷ2) in the next pass. Since the context are accessible for each latent in second slice, the second slice can be processed in parallel. The decoding of the second slice is also run in parallel conditioned on both the hyperprior and the decoded context latents (i.e. first slice). Therefore, the decoding of this structure can be performed in two serial processing.
When decoding the second slice, we adjust the mask in the self-attention module to build a bidirectional context model, which introduces the future context to benefit accurate prediction. For this reason, though the performance of the first slice degrades (due to hyperprior only), Entroformer counteracts this effect by providing a more powerful context model to promote the performance of the second slice. In the other hand, compared to the CNN-based accelerated method (about 4% performance degradation) (He et al., 2021), our transformer-based method can utilize rich context and achieve a better balance between speed and performance (about 1% performance degradation).

heading: 4 EXPERIMENTAL RESULTS
text: We evaluate the effects of our transformer-baed entropy model by calculating the rate–distortion (RD) performance. Figure 5 shows the RD curves over the publicly available Kodak dataset (Kodak, 1993) by using peak signal-to-noise ratio (PSNR) as the image quality metric. As shown in the left part, our Entroformer with joint the hyperprior module and the context module outperforms the state-of-the-art CNNs methods by 5.2% and the BPG by 20.5% at low bit rates. As shown in the right part, two variant entropy models, the hyperprior-only one and the context-only one, are evaluated to isolate the effect of Entroformer architecture. Additionally, the parallel Entroformer with bi-directional context is also competitive and is better than the previous baselines.
To evaluate the importance of different components of the Entroformer, we varied our base model in different ways, measuring the change in performance on Kodak test set. When evaluating the effects of using the diamond relative position encoding and the top-k scheme, we adapt the Entroformer with only the context module. When evaluating the effects of bidirectional context model, we employ the Entroformer with the context module and the hyperprior module.
Architecture The architecture of our approach extends on the work of Minnen et al. (2018b). As our focus is on the entropy model, we only replace the CNNs entropy model by proposed Entroformer (i.e. the main autoencoder is not changed). We used 6 transformer encoder blocks for hyperprior and 6 transformer decoder blocks for context. The outputs of these two modules are combined by a linear layer. The dimension of the latent embedding is 384. Each transformer layer
has a multi-head self-attention (heads=6), and a feed-forward network (ratio=4). If not specified, we use serial Entroformer with the default parameters: k=32 (top-k scheme), h=3 (diamond RPE).
Training Details We choose 14886 images from OpenImage (Krasin et al., 2017) as our training data. During training, all images are randomly cropped to 384 × 384 patches. Note that large patch size is necessary for the training of the transformer-based model. We optimize the networks using distortion terms of MSE. For each distortion type, the average bits per pixel (bpp) and the distortion PSNR, over the test set are measured for each model configurations. The detailed structure and the experimental settings are described in appendix.

heading: 4.1 ANALYSIS ON POSITION ENCODING
text: Position Encoding Methods We perform a component-wise analysis to study the effects of different position encodings in our Entroformer. We construct transformer-based context models with different position encoding. A CNN-based context model is implemented for comparison. As shown in Table 1, When applying absolute position encoding or non-position encoding in Entroformer, we can achieve a limited bpp saving than the baseline. If we apply the relative position encoding, it can perform better than the absolute position one ( 4.0% bpp saving vs. 0.9% bpp saving ). Additionally, it is crucial to extend the position encoding from 1D to 2D ( 4.9% bpp saving vs. 4.0% bpp saving ) . Finally, combining with the diamond boundary, we could get more bit-rate saving than the others, in particular, 5.8% bpp saved than the CNNs one.
Clipping Distance We evaluate the effect of varying the clipping Hamming distance, h, of diamond RPE. Figure 6 compares the number of h by showing the relative reduction in bit rate at a single rate-point. The baseline is implemented by a transformer-based context model without position encoding. It achieves the best result when h = 3. The conclusion is similar as the analysis shown in Figure 3.
4.2 NUMBER k OF TOP-k SELF-ATTENTION
The impact of top-k scheme is analyzed in Figure 7. The parameter k specifies the involved token number of self-attention. We report the full learning bpp curves on Kodak test set for different number k and not just single RD point. Additional curve of original self-attention mechanism is presented (dash) for comparison. The spatial size of latents is 24 × 24 during training and 32 × 48 during testing. Notably, the compression performance increase for k ≤ 64, which is much smaller than the sequance length of 576 and 1536 for training and testing respectively. Moreover, the top-k scheme also influences convergence speed of Entroformer. Interestingly, when increase k larger than 64, we see quite different results that top-k scheme has no difference with original self-attention. One hypothesis is that there is a vast amount of irrelevant information in dense selfattention for image compression. These observations reflect that removing the irrelevant tokens benefits the convergence of Entroformer training.

heading: 4.3 PARALLEL BIDIRECTIONAL CONTEXT MODEL
text: In Table 2, we report the inference latency of Entroformer (serial architecture and parallel architecture) comparing to other learned entropy models. We benchmark models’ inference latency on 16GB Tesla V100 GPU card. From Table 2, we see that parallel bidirectional context model is much more time-efficient than serial unidirectional context model.
As shown in Figure 5, by employing Entroformer, performance of the parallel bidirectional context is on par with the serial unidirectional context. For checkboard parallel context model (He et al., 2021), it speeds up the decoding process with a loss of compression quality. We also study the impact of bidirectional context for compression performance on the second slice (i.e. ŷ2) solely. Table 3 provides an empirical evaluation of three model variants (hyperprior-only, hyperprior with unidirectional context and hyperprior with bidirectional context). Table 3 shows that bidirectional context has more accurate prediction than unidirectional context in our Entroformer. This experiment explains the ability of the parallel Entroformer with bidirectional context.

heading: 4.4 VISUALIZATION
text: Attention mechanism in transformer is the key component which models relations between feature representations. In this section, we visualize the self-attentions mechanism of our Entroformer, focusing on a few points in the image. The attention map is rescaled for a better view. In Figure 8, we visualize self-attention heads from different points These attention heads exhibit different behaviours that seem related to structure, semantic and color respectively. These visualizations show how the Entroformer finds related context to support its distribution prediction for the current latent. This allows the Entroformer to capture richer dependencies of the latents and achieve better compression performance. In Figure 9, we visualize self-attentions heads separately from one point. We can see that different heads attend to different region of an image. This means that separate
sections of the representation can learn different aspects of the meanings of each latent. In Figure 10, we visualize self-attentions heads with different k. Compared with dense attention, the topk attention filters out irrelevant information from background. It successfully makes the attention concentrating on the most informative context.

heading: 5 DISCUSSION
text: In this paper, Entroformer is proposed as the entropy model to improve learned image compression by accuracy probability distribution estimation. To our knowledge, this is the first successful attempt to introduce transformers in image compression task. The Entroformer follows the standard transformer framework with a novel positional encoding and top-k self-attention module. The proposed diamond RPE is a 2D relative position encoding with a diamond shape boundary. The top-k self-attention is helpful in dependencies distillation. In particular, the top-k module benefits the convergence of neural networks training. Last but not least, the Entroformer with bidirectional context can work well in the two-pass decoding framework for fast speed without performance degradation. Experiments show that the Entroformer achieves the SOTA results in standard benchmarks.
There are still some issues to be discussed in the future. For example, (1) Whether we could build a hierarchical decoding framework to achieve a better balance between speed and RD performance, even beyond the raster-scan mode; (2) Whether the fine-grained assignment of k to different images could increase performance; (3) Whether the rate control module plays a more important role in the transformer-based entropy model than in the convolution-based model. We would like to explore these issues in future work.

heading: A APPENDIX
text: A.1 PSNR (RGB) ON KODAK DATASET
A.2 MS-SSIM (RGB) ON KODAK DATASET
A.3 POSITION ENCODING FOR VARIOUS IMAGE SIZE
A.4 ARCHITECTURE DETAILS
Details about the individual network layers in each component of our models are outlined in Table 5. The output of the last layer in encoder corresponds to the latents ŷ, and the output of the last layer in hyper-encoder corresponds to the hyper-latents ẑ. The output of the last layer in decoder corresponds to the generated RGB image x̂. The hyperprior feature and context feature are concatenated and fed into two “Linear” layers to yield the parameter, mean and deviation of a Gaussian distribution N (µ, σ2), for each latent. The output of the linear network must have exactly twice as many channels as the latents.
A.5 TRAINING DETAILS
We implement all our models with PyTorch(Paszke et al., 2019). We use the Adam optimizer (Kingma & Ba, 2014) with β1 = 0.9, β2 = 0.999, = 1× 10−8, and base learning rate= 1× 10−4. When training transformers, it is standard practice to use a warmup phase at the beginning of learning, during which the learning rate increases from zero to its peak value (Vaswani et al., 2017). We use a warmup with 0.05 proportion of the total epochs. And then the learning rate decays stepwise for every 1/5 proportion epochs by a factor of 0.3. Gradient clipping is also helpful in the compression setup, which is set to 1.0. No weight decay is used in our models because rate constrain is inherently a regularization. We initialize the weights of our Entroformer with a truncated normal distribution.
Inspired by the self-supervised pre-training of Transformer, we also perform a random mask for pretrain. Different from masked patch prediction, masked latents are predicted solely by hyperprior while other latents are predicted by hyperprior and context. To do so we corrupt 20% of latents by replacing their value with 0. After pretrain, we finetune the models with regular manner.
All models are trained for 300 epochs with a batchsize of 16 and a patch size of 384× 384 on 16GB Tesla V100 GPU card. To obtain models for different bit rate, we adapt the hyper-parameter λ to cover a range of rate-distortion tradeoffs.

Review summaries:
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. 

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.

The authors propose a transformer-based entropy modeling order to capture long-range dependencies in probability distribution estimation. This model is optimized for image compression. The authors extend this architecture with a parallel bidirectional context model to speed up the decoding process.