Entroformer: A Transformer-based Entropy Model for Learned Image Compression
Authors: Yichen Qian
Xiuyu Sun
Ming Lin
Zhiyu Tan
Rong Jin
Venue: ICLR 2022 Poster (2022)
Keywords: Image compression
Entropy Model
Global Dependencies

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.

Published as a conference paper at ICLR 2022
ENTROFORMER: A TRANSFORMER-BASHR ENTROPY MODEL FOR LEARNED IMAGE COMPRESSION
Yichen Qian
Ming Lin
Xiuyu Sun'
Alibaba Group
Alibaba Group
Alibaba Group Hangzhou, China
Hangzhou, China
Bellevue, WA, 98004, USA
yichen. www.aaybabbbbin........om com
ming. 18alibaba-inn com
xwuyu. www.gallbbbb---nn..c... con
Zhiyu Tan
Rong Jin
Alibaba Group
Alibaba Group
Hangzhou, China
Bellevue, WA, 98004, USA
shiyu, tww.tttyyyabbaaaaaa.aa -inaaaaaa-nnnn.eeeeeeeee com
3 jinrong. jr@alibaba--an com
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 de- pendencies. In this work, we propose a novel transformer-basee entropy model, termed Entroformer, to capture long-range dependencies in probability distribu- tion estimation effectively and efficiently. Different from vision transformers image classification, the Entroformer is highly optimized for image compression, including a top-k self-attention and a diamond relative position encoding. Mean- while, further expand this architecture with a parallel bidirectional context model to speed up the decoding process. The experiments show that the Entro- former achieves state-of-ther ahe performance on image compression while being time-efficient. Code is available at
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-of-
the-art (SOTA) deep image compression models are built on the auto-encoder framework (Hinton & Salakbutdinov (2006) with an entropy-conse conntrycd bottleneck (Theis et al.
2017 Ballé et al. 2017:
Ballé al.] 2018 Mentzer et al., 2018; Minnen et al.j 201 Sa Lee et al. 2019 Guo et 2021
Sun et al.) 2021). An entropy model estimates the conditional probability distribution of latents for compression by standard entropy coding algorithms. 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 learn representation for various depen- dencies. For example, as shown in Figure [ 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. 2018 Sa Lee et al. 2019 Mentzer et al., 2018) or additional side information (Ballé et 2018 Hu al. 2020 Minnen et al. 2018b) for short-range spatial dependencies, while
others use local mechanism (Li et al.] 2020 Qian et al. [202]] 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 method uses unidirectional context model to improve the predictive ability, it suffers
"Corresponding author

---

Published as a conference paper at ICLR 2022
(a)
(b)
(c)
Figure 1: Dependencies exist in (a) texture, (b) semantic and (c) color space.
from slow decoding speed. It decodes symbols in raster-scan order with Onn) serial process that can not be accelerated by modern GPUs. A two-pass parallel context model (He et al. (021) is introduced for acceleration, which decodes symbols in particular order to minimize serial pro- cessing. 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 ad- dress the above two challenges in one shot. Following the Transformer (Vaswani al., 2017). self-attention is used to relate different positions of a single latent in order to compute a represen- tation 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 com- pression, the following designs are innovated: (1) Entroformer uses a multi-heas head attention with top-k selection to distill information in representation subspaces: The multi-head attention provides a learning-basd based partition to learn dependencies
different views; and the top-k scheme filters noisy signals by select the most similar & latents in the self-attention module, which is crucial for the convergence of Entroformer. (2) To inherit the local bias CNNs, novel position encoding unit is designed 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), (3) The two-pass decoding framework (He et al. 202I) 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:
which embeds { -norm distance prior knowledge in image compression.
In this paper, Entroformer is proposed to improve learned image compression. To our knowledge, it the first successful attempt to introduce Transformer based method to im- age 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% low bit rates.
To further optimize Entroformer for image compression, the multi-head scheme with top-k selection proposed to capture precise dependencies for probability distribution estima- tion, 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 En- troformer 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 hyperprior (Minnen et al.) 2018b Ballé et al., 2018) context model (Li et al. 2018 Mentzer et al. 2018), and the combined method (Lee et al. 2019 Minnen et al.] 201 Sa. Minnen & Singh 2020 Hyperprior methods typically make use of side information of the quantized latent representation with additional bits. Context models learn auto- regressive priors incorporating prediction from the causal context of the latents, which is bit-free. Our Entroformer combines a context model with hyperprior based on previous works (Minnen et al.
2018h Ballé et al., 2018).
2

---

Published as a conference paper at ICLR 2022
Imput Image
Output Image
x
Encoder
Decoder
Self-Attention
Nx
Forward
Downscalle
AE
AD
Self-Attention
Context
N N 6)
FeedForrart
Mask Self- Gell-Attention
N
Unsmle
Forward
x 2N
Hyper
Hyper Decnder
Hyperprior
Context
Linear
Q
AE
AD
14, 0
(a)
(b)
(a), we show the architecture of image compression model with with hyerprprior and context model. The main autoencoder learns latent representation of image, while hyper- autoencoder learns a hyperprior representation. Context model learns correlation of latents from their decoded context. Real-valued latent representation (y, 2) is quantized (Q) to create latents (ÿ) and hyper-latents (#), which are compressed into a bitstream using an arithmetic encoder (AF) and decompressed by arithmetic decoder (AD). The latents tioned on hyperprior and context. The region in dashed line corresponds to our Entroformer, which shown in (b). A transformer encoder learns the hyerprprior, while a transformer decoder learns the context using masked self-attention Linear component is used to joint these features to generate the Gaussian parameters (i.e. μ., O).
Figure 2:
Gaussian entropy model condi-
We describe the architecture of compression model with Entroformer in Figure 2 The main autoen- coder learns quantized latent representation I of image æ. The reconstructed image is denoted as t. The hyper-autoencooer learns a quantized hyper-latents representation é. Following the work of Ballé al. (2016), the quantization operation is approximated by additive uniform noise during training. This ensures good match between encoder and decoder distributions of both the quantized latents, and continuous-valued latents subjected additive uniform noise during training. Follow- ing the work of Minnen et al. (2018a), we model each latent, Di. as a Gaussian with mean μa and deviation a; convolved with a unit uniform distribution:
Po(0 0)
II (N(,,, ((550., (0:)
where μ and are predicted by the entropy model, 0 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, 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 the bitstream and the quality of the reconstruction, which is a rate-distortion opti- mization problem:
- R= = Encp, [- logg (g(
+ Excep. [- log, Pa (2()) + X En-p, | - EIl
(2)
(latents)
(hyper-latents)
distortion
where À is the coefficient which controls the rate-distortion trade-off, Pa is the unknown distribution of natural images. ps 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 com- pression rate of the hyper-latent representation. The third term represents the distortion value under given metric, such as mean squared error (MSE).
3

---

Published as a conference paper at ICLR 2022
Hyperprior
0.30
m. 00.09%%
0.67%
6.75%
0.90%
0.78%
0.73%
0.58%
0.255
Model
Model
-0.66%
0.84%
0.94%
1.22%
1.72%
2.05%
1.71%
1.23%
0.22
1.33%
1.74%
3.02%
3.04%
1.76%
1.33%
0..44%
1.19%
1.00%
2.55%
3.47%
1.24%
0.98%
82
91
0.91%
127%
3 00.7%% 6.79%
0.82%
1.00%
0.06% 0.77%
0.67%
0.000
Figure 3: Rate increase ratios of different posi- tion when the corresponding context masked. Closer context contributes to more bitrate influ- ence, whose trend follows a particular pattern diamond shape.
Figure 4: Parallel bidirectional context model. The first slice (/1) is conditioned solely on the hyperprior. The second slice (f22) is conditioned on the hyperprior and the decoded context (i.e. first slice).
TRANSFORMER-BASED ENTROPY MODEL
Our model contains two main components, CNN-based autoencoder to learn latent representa- tion, transformer-bart based entropy model to predict latents. We deseribe our architecture in Figure (a), and our Entroformer in detail in Figure (b). In this section, we first propose two ingredients, a diamond relative position encoding (diamond RPE) and top-k scheme, which are essential for image compression. Then, we extend the checkboard context model (He et al. 202T) 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 ý as a sequence. The latents sequence is then mapped to D dimensions with trainable linear projection. For hyperprior, we stack N transformer encoder layers to yield hyper-latent z, which quantized and fed N transformer- encoder layers to generate hyperprior. To produce hierarchical representation, the resolution of feature changed by downsample and upscale module. For autoregressive prior, we stack 2NV transformer decoder layers 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 Raxdm:
( XW
Attention X) softmax
(3)
vdz
Wº, W Rd xd, and V e pdd.. xd, 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, evaluate the effect of each position i by employing corresponding mask, which are set to except position i. Figure 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.
4

---

Published as a conference paper at ICLR 2022
In order for spatial information of the latents, common approach is biased attention weights based on relative relation (Shaw et al., 2018). Based on the result in Figure 3 we propose extension to this relation-aware self-attention to consider the element position influence image compression performance.
input sequence, x ..,, of elements where H; E Rd The edge between input elements ž; and 2j is represented by vectors p E Rd We modify the dot-product attention matrix in eq. 3 compatibility function that compares two input elements in relation-aware self-attention:
Each attention operates
(2,
D
Vd
The weight coefficient, A, is computed using softmax dot-product attention matrix.
On image plane, the relative position O,-, is 2D coordinate. And we modify the maximum absolute relative position to 2D boundary with diamond shape. As shown in Figure 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, consider a 2D boundary with maximum & distance value of h.
P.j = w clip((s--.h)
clip (at-s, h) = ai-j |aa-5llı <= h
(5)
(h, otherwise.
We then learn the relative position encoding w (w - w
w
hh hhh)
w
3.3 TOP-K K SCHEME IN SELF-AIE ALLE-TTTTNIOON
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 vdg. and apply a softmax function obtain the weights on the values.
2022) have demonstrated promising results on vision tasks. Inspired by these works, we select the top-k most similar keys for each query instead of computing the attention matrix for all the query-key pairs as vanilla attention.
Content-based sparse attention methods (Wang et 2021
The queries, keys and values packed together into matrices Q Raxds K E Raxd4 and
V E Igaxd. P Raxd denotes relative position encodings described See 3.2
The row-wise
top-k elements in attention matrix selected for softmax function:
(Q(K+ (((K+t)))))) )))
V' = softmax
(6)
Vile
where fa(-) denotes the row-wise top-k selection operator:
Cij if eig is within the top-k largest elements in row otherwise.
fk (c(g)
-00
Albeit easing sequence length mismatching problem by relative position encoding, transformer- based 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 self- attention with top-k scheme in the Entroformer. helps to distill irrelevant context latents and stable training procedure.
3.4 PARALLEL BIDIRECTIONAL CONTEXT MODEL
The transformer architecture can be used to efficiently train autoregressive model by masking the attention computation such that the i-th position only be influenced by position if and only j < i. This causal masking uses a limited unidirectional context, and makes autoregressive model impossible be parallel.
5

---

Published as a conference paper at ICLR 2022
PSNR (RS8) on Kodaik vs. baseline models
PSNR, IREBI Kodak vs. model variants
42
40
36
38
34-
30
PSWR (REB)
02-
Fntroforme: Contest-typarpriie
CNNN: Hypergeior
JPE:2000 (Open)PES)
28
CNNs: Context++iyperpree
BPG (4-444 12018) Coptext--yypeptor Lee ContextAdaptive
Entroformen Hyperprir
Entroformer Context
Entroformen Sorteate-Hperprria
26
Qlan (2021) Raference
26-
Paraliel Entrofemer: +iyperrt
0.0
1.0 Bits per plwel
per pèxel (BPP)
Figure 5: The Entroformer (context hyperprior) achieves better RD performance than previous CNNs-based methods and standard codecs on the Kodak image set as measured by PSNR (RGB) (left). The right graph compares the relative performance of different versions of our method. Our Entroformer architecture helps significantly at hyperprior-only and context-only model by 12.4% and 11.3% respectively low bit rates. Note that parallel Entroformer has almost no effect ( about 1% bpp increase) for RD performance.
Following the work of He et al. (202T), the latents are split into two slices along the spatial dimen- sion. Figure provides high-leve overview of this architecture. The decoding of the first slice (g1) 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 (ú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 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 bidi- rectional 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-bast method can utilize rich context and achieve a better balance between speed and performance (about 1% performance degradation).
EXPERIMENTAL RESULTS
We evaluate the effects of our munsformer-baed entropy model by calculating the rate- distoristo (RD) performance. Figure/JJshows 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 also competitive and is better than the previous baselines.
To evaluate the importance of different components of the Entroformer, 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, adapt the Entro-

---

Published as a conference paper at ICLR 2022
Model
Position Information
Mode
bpp
A bpp (%)
CNNs
0.6621
Transformer
0.6569
0.8
Transformer
Absolute Position Encoding
2D
0.6563
0.9
Relative Position Encoding
ID+1D
0.6358
4.0
Transformer
Relative Position Encoding
2D
0.6296
4.9
Relative Position Encoding
2D>+Dammond
0.6238
5.8
Table 1: Ablation of position encoding methods. All models optimized with À 0.02 (PNSR R 35.0 on Kodak).
former with only the context module. When evaluating the effects of bidirectional context model, 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. 2011h) 8b). As our focus is on the entropy model, we only replace the CNNs entropy model by proposed En- troformer (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 combined by a linear layer. The dimension of the latent embedding is 384, Each transformer layer has multi-head self-attention (heads=6), and 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).
We choose 14886 images from OpenImage (Krasin et al. 2017) as our training data. During training, all images are randomly cropped to 384 x 384 patches. Note that large patch size necessary for the training of the transformer-batt 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.
Training Details
4.1 ANALYSIS ON POSITION ENCODING
Position Encoding Methods We perform a component-wiss analysis study the effects of dif- ferent position encodings in our Entroformer. We construct transformer-based context models with different position encoding. A CNN-based context model implemented for comparison. As shown in Table When applying absolute position encoding or non non-position encoding in Entroformer, we can achieve a limited hpp saving than the baseline. we apply the relative position encoding, can perform better than the absolute position one ( (.00% bpp saving vs. 0.9% bpp saving ). Additionally, is crucial to extend the position encoding from 1D to 2D ( (.99% hpp 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 & distance, h, of diamond RPE. Figure comppare the number of by showing the relative reduction in bit rate at a single rate-point. The baseline is implemented by transformer-bast based context model without position encoding. It achieves the best result when = 3. The conclusion is similar as the analysis shown in Figure
4.2 NUMBER 1: OF T 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 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 x 24 during training and 32 x 48 during testing. Notably, the compression performance increase for k S 64, which is much smaller than the sequance length 576 and 1536 for training and testing respectively. Moreover, the top-k scheme also influences convergence speed of Entroformer. Interestingly, when increase 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 vast amount of irrelevant information in dense self- attention for image compression. These observations reflect that removing the irrelevant tokens benefits the convergence of Entroformer training.

---

Published as a conference paper at ICLR 2022
Rate Saveing Baseline (na position encoding)
Vercenee Size decreasse (higher bettei)
No topk
k-8
0.75
0.750
km64
0.723
0.700
0.675
2.650
0.690
Training apoch
Figure 6: Results for varying the clipping & dis- tance of diamond RPE. The base implementa- tion uses a transformer context model without position encoding. All models are optimized with x = 0.02 (PNSR x 35.0 on Kodak).
Figure Effect of restricting the self-attention to various top-k similar content. All models use transformer context model and are optimized with 0.02 (PNSR 35.0 Kodak).
Entroformer
Ballé2018 (Hyperprior)
Minne2018
Test set
(Hyperprior+Conteex,
(Hypprpprio+Conneex)
Serial parallel
Kodak (768 512)
21.6
4229.9
51678.5
170.6
Table Inference latency (ms) of entropy model on Kodak, [Ballé et al. (2018) uses is paralleled, while Minnen al. 28) is inherently serial. For direct comparison, we test our Entroformer with serial architecture and parallel architecture.
Hyperprior & unidirectional context
Hyperprior & bidirectional context
Half latents (j2) bpp
Hyperprior-only
0.3267
0.3105
0.2972
Table 3: Bidirectional context unidirectional context. Combined with bidirectional context, context model achieves bit rate saving significantly. The results are test by the same model.
4.3 PARALLEL BIDIRECTIONAL CONTEXT MODEL
In Table we report the inference latency Entroformer (serial architecture and parallel archi- tecture) comparing other learned entropy models. We benchmark models' inference latency 16GB Tesla V100 GPU card. From T we see that parallel bidirectional context model much more time-efficient than serial unidirectional context model.
As shown in Figure by employing Entroformer, performance of the parallel bidirectional context on par with the serial unidirectional context. For checkboard parallel context model (He al. 2021). it speeds up the decoding process at cost of compression quality. We also study the impact of bidirectional context for compression performance on the second slice (i.e. ÿ2) solely. Table provides an empirical evaluation of three model variants (hyperprior-only hyperprior with unidi- rectional context and hyperprior with bidirectional context). Table 3 shows that the bidirectional context is more accurate than the unidirectional context in our Entroformer.
4.4 VISUALIZATION
Attention mechanism in transformer the key component which models relations between feature representations. In this section, we visualize the self-attentions mechanism of our Entroformer, fo- cusing on a few points in the image. In Figure g the attention maps from different points exhibit different behaviours that seem related structure, semantic and color respectively. It show how the Entroformer finds related context to support its distribution prediction for the current latent. In Figure 9 visualize self-attentions heads separately from point. We that different
8

---

Published as a conference paper at ICLR 2022
Figure 8: Examples of the self-attention mechanism for three points.
head
head 2
head 3
head
head
head
Figure Self-attention heads following long-range dependencies with different patterns.
top-k
k=32
Figure 10: Self-attention from models with different k.
heads attend different region of an image. This means that separate sections of the representation can learn different aspects meanings of each latent. In Figure T0 we visualize self-attentions heads with different k. Compared with dense attention, the top-k attention filters out irrelevant infor- mation from background. It successfully makes the attention concentrating on the most informative context.
DISCUSSION
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 tempt to introduce transformers in image compression task. The Entroformer follows the standard transformer framework with novel positional encoding and top-k self-attention module. The pro- posed diamond RPE is 2D relative position encoding with diamond shape boundary. The top-k self-attention is helpful in dependencies distillation. In particular, the top-k module benefits the con- vergence of neural networks training. Last but not least, the Entroformer with bidirectional context can work well 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 the future. For example, (1) Whether we could build hierarchical decoding framework to achieve better balance between speed and RD performance, even beyond the raster-scan mode: (2) Whether the fine-graine are-grained assignment
to different images could increase performance; (3) Whether the rate control module plays 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.

---

Published as a conference paper at ICLR 2022
REFERENCES
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimization of nonlinear trans- form codes for perceptual quality. Picture Coding Symposium, pp. 1-5. IEEE, 2016.
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. International Conference on Learning Representations. 2017.
Johannes Ballé, David Minnen, Saurabh Singh. Sung Jin Hwang, and Nick Johnston. Variational image compression with a scale hyperprior. In International Conference Learning Represen- tations, 2018.
Fabrice Bellard. Bpg image format, 2014, URL
Tong Chen, Haojie Liu, Zhan Ma, Qiu Shen, Xun Cao, and Yao Wang. End-to-end learned image compression via non-local attention optimization and improved context modeling. IEEE Trans- actions on Image Processing, 30:3179-3191, 2021.
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto. Learned image compression with discretized gaussian mixture likelihoods and attention modules. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7939-7948, 2020.
Zongyu Guo, Zhizheng Zhang, Runsen Feng, and Zhibo Chen. Causal contextual prediction for learned image compression. IEEE Transactions on Circuits and Systems for Video Technology, 2021.
Dailan He, Yaoyan Zheng, Baocheng Sun, Yan Wang, and Hongwei Qin. Checkerboard context model for efficient learned image compression. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 14771-14780, 2021.
Geoffrey E Hinton and Ruslan R Salakhutdinov. Reducing the dimensionality of data with neural networks. science, 313(5786):504-50) 2006.
Yueyu Hu, Wenhan Yang, and Jiaying Liu. Coarse-to-fin hyper-prior modeling for learned image compression. In AAAL, pp. 1101--11220 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.69880 2014.
Eastman Kodak.
Kodak lossless true color image suite (photoed ped0992), 1993. URL http:
Ivan Krasin, Tom Duerig, Neil Alldrin, Vittorio Ferrari, Sami Mi--El-Haija, Alina Kuznetsova, Hassan Rom, Jasper Uijlings, Stefan Popov, Shahab Kamali, Matteo Malloci, Jordi Pont- Tuset, Andreas Veit, Serge Belongie, Victor Gomes, Abhinav Gupta, Chen Sun, Gal Chechik, David Cai, Zheyun Feng, Dhyanesh Narayanan, and Kevin Murphy. Openimages: A public dataset for turge-scale multi-label and multi-class image classification. Dataset available from https://strage.googleeeee wwtppsss.....ooeionnn/eeeeeeeeee
2017.
Jooyoung Lee, Seunghyun Cho, and Seung-K Beack. Context-adaptive entropy model for end- to-end optimized image compression. International Conference on Learning Representations, 2019.
Mu Li, Wangmeng Zuo, Shuhang Gu, Debin Zhao, and David Zhang. Learning convolutional networks for wontent-weighted image compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3214-3223, 2018
Mu Li, Kai Zhang, Wangmeng Zuo, Radu Timofte, and David Zhang. Learning context-based non- local entropy modeling for image compression. arXiv preprint arXiv:2005.04661 2020.
Fabian Mentzer, Eirikur Agustsson, Michael Tschannen, Radu Timofte, and Lue Van Gool. Condi- tional probability models for deep image compression. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 4394-4402, 2018
10

---

Published as a conference paper at ICLR 2022
David Minnen and Saurabh Singh. Channel-wist autoregressive entropy models for learned image compression. In 2020 IEEE International Conference on Image Processing, pp. 3339-3343.
IEEE, 2020.
David Minnen, Johannes Ballé, and George Toderici. Joint autoregressive and hierarchical priors for learned image compression. In Advances in Neural Information Processing Systems, pp. 10771-10780, 2018a.
David Minnen, George Toderici, Saurabh Singh, Sung Jin Hwang, and Michele Covell. Image- dependent local entropy models for image compression with deep networks. 2018b.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lever, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
Yichen Qian, Zhiyu Tan, Xiuyu Sun, Ming Lin, Dongyang Li, Zhenhong Sun, Hao Li, and Rong Jin. Learning accurate entropy model with global reference for image compression. International Conference on Learning Representations, 2021.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. arXiv preprint 1803 803.00255,, 2018.
Zhenhong Sun, Zhiyu Tan, Xiuyu Sun, Fangyi Zhang, Yichen Qian, Dongyang Li, and Hao Li. Interpolation variable rate image compression. In Proceedings of the 29th ACM International Conference on Multimedia, 5574-5582, 2021.
Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. International Conference on Learning Representations, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Pichao Wang, Xue Wang, Fan Wang, Ming Lin, Shuning Chang, Wen Xie, Hao Li, and Rong Jin. Kvt: k-nn attention for boosting vision transformers. arXiv preprint arXiv:2106. 00515, 2021.
Yaohua Wang, Yaobin Zhang, Fangyi Zhang, Senzhang Wang, Ming Lin, YuQi Zhang, and Xiuyu Sun. Ada-nets: Face clustering via adaptive neighbour discovery in the structure space. arXiv preprint 1rXiv:2109. 0666, 2022.

---

Published as a conference paper at ICLR 2022
A
APPENDIX
A.1 PSNR (RGB) ON KODAK DATASET
PSNR (RGB)
42
PSWR (RGB)
34
32
Entroformer: Context-Hyyerprror
Guo (2021)*
Cheng (2020)
(4:4:4)
Minnen (2020)
Qian (2021)
(2018)
(2019)
BPG (4:4:4)
Ballé (2018)
JPEG2000 (OpenjPEG)
JPEG (4:2:0)
0.0
0.5
1.0
1.5
2.0
Bits per pixel (BPP)
Rate-Distortion curves the Kodak image using PSNR. Each point on the RD curves is calculated by averaging over the MS-SSIM and bit rate. Note that the models with use deeper main auto-encoder (transform between image and latent).
Figure 11:
12

---

Published as a conference paper at ICLR 2022
A.2 MS-SSIM (RGB) ON KODAK DATASET
(RGB) baseline
28
24
20
MS
16
Entroformer: Context- Hoyperi-- for MS-SSIM Guo (2021) for MS SSIM
Cheng (2020) for MS-SSIM
Qian (2021) opt. for MS-SSIM
Lee (2019) opt. for MS-SSIM
12
Minnen (2018) opt. for MS-SSIM
Ballé (2016) opt. MS-55IM
(4:4:4)
JPEG2000 (Open)PEG)
JPEG (4:2:0)
10
0.0
1.2 Bits pixe! (BPP)
1.6
2.0
Figure 12: Rate-Distores Distoorttittn curves the Kodak image set using MS- Each point the RD curves calculated by averaging over the MS-SSIM and bit rate. To improve readability, this graph shows AS-SSIM scores in dB using the formula: MS-SSIMaB -0 1 MS-SSIM). Note that the models with use deeper main auto-encoder (transform between image and latent).
13

---

Published as a conference paper at ICLR 2022
A.3 POSITION ENCODING FOR VARIOUS IMAGE SIZE
Position Encoding
Image size
bpp
A bpp (%)
Absolute PE
300 300
0.7654
Relative PE
300 x 300
0.7180
6.2%
Absolute PE
600 600
0.5360
Relative PE
600 600
0.5057
5.7%
Absolute PE
900 x 900
0.4938
Relative PE
900 900
0.4154
15.9%
Absolute PE
1200 1200
0.4808
Relative PE
1200 1200
0.4099
14.7%
Table 4: Performance of position encoding for varying the image size on Tecnick dataset. shows that relative position encoding can directly generalize to larger image size. The model with absolute position encoding degrades when process larger image size than training.
A.4 ARCHITECTURE DETAILS
Encoder
Decoder
Hyper Encoder
Hyper Decoder
Context Model
Lincar
Conv: 19cc2
Deconv: k5c192s2
Trans: dimb38hh
Trans: dim384hheu66
Trans: dim384beadd
Linear: dim768
GDN
IGDN
Downscale
Upscalc
Trans: dim38 38-iadd
Leaky RcLU
Conv: k5c192cc
Decony: A5c192<2
Trans: dim18-hhead6
Trans:
dim 384hend6
Trans dim 184head6
Linear: dim768
GDN
IGDN
Downscale
Upscale
Trans: 38mhea6
Conv: k5c19242
Decor::
k5c192s2
Trans:
dim3384beed66
Trans: dim384heead6
Trans: dim384bbead6
GDN
IGDN
Trans: dim384bbead6
Conv; k5c384s2
Deconv:
k5c192<2
Table 5: Each row corresponds to a layer of our generalized model. Convolutional layers are spec- ified with the "Conv" prefix followed by the kernel size, number of output channels and downsam- pling stride (e.g. the first layer of the encoder uses 5x5 kernels with 192 output channels and a stride of 2). The "Deconv" prefix corresponds upsampled convolutions. The "Trans" corresponds to transformer block as in Vaswani et al. (2017), followed by the inner dimension and head number. transformer block consists of a self-attention module and feed-forward module (the dimension of the feed-forward is 4 times of inner embedding's s). GDN stands for generalized and divisive nor- malization, and IGDN is inverse GDN (Ballé et al., 2016). Group convolution with kernel=3 and stride=2 is used for downscale operation. Pixel shuffule (by group convolution with kernel=3 and stride=2) is used for upscale operation.
Minnen (2018) Context++yperprrior
Minnen (2020)
Qian (2021)
Method
Ours
Channel Auto-regressive
Reference
Weights of entropy model
90.0M
237.3M
145.7M
142.7M
Table 6: Comparison on weight with entropy model variants.
Details about the individual network layers in each component of our models are outlined in Table 5 The output the last layer in encoder corresponds to the latents ÿ. and the output the last layer in hyper-encoder corresponds to the hyper-latents ?. The output of the last layer in decoder corresponds to the generated RGB image â. The hyperprior feature and context feature are con- catenated and fed into two "Linear" layers to yield the parameter,
and deviation of a Gaussian distribution. N(u, o²), for each latent. The output of the linear network must have exactly twice as many channels the latents.
14

---

À to Published as a conference paper at ICLR 2022
A.5 TRAINING DETAILS
We implement all our models with PyTorch(Pszkke et al. 2019). We use the Adam optimizer (Knggma & Ba 2014) with 31 0.9, 0.999, e 1 10 and base learning rate x 10 When training transformers, is standard practice to use a warmup phase at the beginning of learn- ing, during which the learning rate increases from zero its peak value (Vaswani et al.] 2017). We use 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.75. Gradient clipping is also helpful in the com- pression setup. which is set to 1.0. No weight decay is used in our models because rate constrain inherently a regularization. distribution.
initialize the weights our Entroformer with a truncated normal
Inspired by the self-supervised pre-training of Transformer, we also perform a random mask for pretrain. Different from masked patch prediction, masked latents predicted solely by hyperprior while other latents are predicted by hyperprior and context. To do so we corrupt 50% of latents by replacing their value with O. After pretrain, we finetune the models with regular manner.
All models are trained for 300 epochs with a batchsize of 16 and patch size of 384 x 384 on 16GB Tesla V100 GPU card. To obtain models for different bit rate, we adapt the cover a range of rate-distortion tradeoffs.
15


Reviews:
[{"review_id": "zQ_nBMLpJj", "reviewer": "Reviewer_3XoD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strenghts:\n1) First approach using Transformer based method to image compression task.\n2)  The proposed method outperforms previous methods based on CNNs.\n3) An efficient parallel architecture is proposed and is more time-efficient than the serialized one on modern GPU device.\n\nWeaknesses\nExperiments could be more extensive.", "summary_of_the_paper": "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\nmodel to speed up the decoding process.", "summary_of_the_review": "The paper present a novel contribution by using first Transformer based method to image compression task. Experiments validate the performance of the proposed approach named Entroformer compared to CNNs based methods.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "tpr-FWIkaXk", "reviewer": "Reviewer_JhCm", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have not seen transformers used as entropy models for image compression before this paper (though there appears to be other submissions ICLR 2022 exploring the same idea). Transformers seem like a natural fit for this problem so an empirical evaluation is useful for the community. The diamond-shaped relative position encoding scheme introduced in this paper is also novel (as far as I know) and makes sense given the falloff in spatial correlation typical to natural images.\n\nThe use of a checkerboard decomposition for the latent tensor to create a \"parallel bidirectional context model\" is not novel (the authors appropriately cite He et al. 2021). The authors do provide runtime information (Table 2) showing that \"serial context\" models can be quite slow, while the parallel approach is much faster (170ms vs. thousands of milliseconds, though simpler models allow for decoding in under 22ms with the trade-off being worse compression rates).\n\nMy main concern with this paper is that the empirical evaluation does not show better results compared to the best published models. For example, Figure 5 compares again Qian (2021) but not He (2021), Minnen (2020), or Cheng (2020), all of which are cited elsewhere in the paper. It's hard to compare curves by eye across papers, but it appears that all of these methods have similar rate-distortion (RD) results. Furthermore, other models yield better RD results, for example:\n\nZongyu Guo, Zhizheng Zhang, Runsen Feng, Zhibo Chen. Causal Contextual Prediction for Learned Image Compression. IEEE Transactions on Circuits and Systems on Video Technology. 2021.\n\nChangyue Ma, Zhao Wang, Ruling Liao, Yan Ye. A Cross Channel Context Model for Latents in Deep Image Compression. \tarXiv:2103.02884. 2021.\n\nIt appears (again, I'm eyeballing but you can trace the graph or contact the authors for more accurate numbers) that the Entroformer model hits 33.1 dB at 0.4 bpp, while Guo (2021) is closer to 33.5 dB and Ma (2021) is around 33.3 dB.\n\nI do think that the entroformer model has an advantage since it allows for much faster decoding speeds (both Guo and Ma use a serial context model), but this speed benefit comes from the parallel decoding introduced by He (2021), not by the transformer introduced in this paper.\n\n", "summary_of_the_paper": "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.\n\nThe 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.", "summary_of_the_review": "I think this paper is \"marginally above the acceptance threshold\" since the use of transformers as entropy models for image compression is novel and interesting. The paper needs a more thorough empirical evaluation to warrant a higher score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "l2UuYpr5tLn", "reviewer": "Reviewer_FHR4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "** Strengths **\n\n(1) The improvements over the CNN-based hyperprior and context models are promising. \n\n(2) The ablation study justifies well the design choices. \n\n** Weaknesses **\n\n(1) Transformers are costly in terms of complexity and runtimes as shown in Table 2 (Transformer Serial). Some prior works introduce the non-local attention module (NLAM) to the hyperprior codec without using layers of transformer encoders. One example is \"End-to-End Learnt Image Compression via Non-Local Attention Optimization and Improved Context Modeling\", TIP2021. It is unclear how the proposed transformer-based hyperprior compares with such NLAM-based hyperprior design. \n\n(2)  It is expected that transformer-based design may not scale up easily from the complexity perspective when it comes to processing higher-resolution images. Essentially, its complexity may grow exponentially with the spatial resolution of the image latents. \n\n(3) In some ablation studies, the performance differences are expressed in BPP or delta BPP. It is unclear at what PSNR (or quality) level these numbers are reported. Why not report BD-rate savings and present a comparison of complete RD curves? The same argument applies to Figs. 6 & 7. \n\n(4)  The number of network weights spent on the entropy coder should be compared.  ", "summary_of_the_paper": "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. ", "summary_of_the_review": "Overall, the idea is novel and interesting. The results look promising. However, there are some critical details missing. Further clarifications are needed to understand better the pros and cons of the transformer-based scheme as compared to the prior work. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The three reviewers all felt the paper was above threshold for acceptance to ICLR. To improve the final version, they suggest some additional discussion and experiments may help the paper.
Decision: Accept (Poster)