E NTROFORMER : A T RANSFORMER - BASED E NTROPY M ODEL FOR L EARNED I MAGE C OMPRESSION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 inefﬁcient 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 efﬁciently.
Different from vision transformers in image classiﬁcation, 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-efﬁcient.
I NTRODUCTION
Image compression is a fundamental research ﬁeld 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; Ball´e et al., 2017;
Ball ´ e 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 ﬁve 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 (Ball ´ e 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
[Figure]
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 ﬁlters 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.
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 ﬁrst 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-efﬁcient than the serialized one on modern GPU devices without performance degradation.
C OMPRESSION WITH H YPERPRIOR AND C ONTEXT
There are types of entropy models such as hyperprior (Minnen et al., 2018b; Ball ´ e 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; Ball´e 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
Ball´e 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:
[Equation]
[Figure] [Figure]
Q 𝑧 𝑁(𝜇,σ) 𝑦( 𝑧̂ (a) Q
Mask Self-Attention
FeedForward x 2N
Hyperprior 𝑦0 !"
Linear 𝜇, σ
Context
Self-Attention
Downscale
FeedForward N x
Q 𝑦 𝑧 𝑧̂ (b)
Self-Attention
Upscale
FeedForward N x cat HyperEncoder
Encoder
Decoder HyperDecoder ContextModel AE
AD
AE AD 𝑦( 𝑧̂ 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
[Equation]
where λ is the coefﬁcient which controls the rate-distortion trade-off, p x is the unknown distribution of natural images.
p ˆ z use a fully factorized density model.
The ﬁrst 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).
T RANSFORMER - BASED E NTROPY M ODEL
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 ﬁrst 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.
[Figure]
Closer context contributes to more bitrate inﬂuence, whose trend follows a particular pattern of diamond shape.
[Figure]
The ﬁrst slice ( ˆ y 1 ) is conditioned solely on the hyperprior.
The second slice ( ˆ y 2 ) is conditioned on the hyperprior and the decoded context ( i.e . ﬁrst slice).
3. T RANSFORMER A RCHITECTURE
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 ∈
R H
× 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 2 N 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 ∈ R n × d m :
[Equation]
W Q , W K ∈ R d m × d k and W V ∈ R d m × d v are learned parameter matrices.
3. P OSITION E NCODING
To explore the impact of position in image compression, we ﬁrst 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 inﬂuence on image compression performance.
Each attention operates on an input sequence, x = ( x 1 , ..., x n ) of n elements where x i ∈ R d m .
The edge between input elements x i and x j is represented by vectors p Kij ∈ R d k .
We modify the dot-product attention matrix in eq. 3 to a compatibility function that compares two input elements in relation-aware self-attention:
[Equation]
The weight coefﬁcient, A , is computed using a softmax on dot-product attention matrix.
On image plane, the relative position a i − 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 inﬂuences the bit saving of context modeling in learned image compression.
Therefore, we consider a 2D boundary with a maximum Hamming distance value of h .
[Equation]
We then learn the relative position encoding w K
[Equation]
[Equation]
[Equation]
3. T OP - K S CHEME IN S ELF -A TTENTION
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 √ d k , 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 ∈ R n × d k , K ∈ R n × d k and V ∈ R n × d v .
P ∈ R n × d k denotes relative position encodings described in Sec 3.2.
The row-wise top-k elements in attention matrix are selected for softmax function:
[Equation]
where f k ( · ) denotes the row-wise top-k selection operator:
[Equation]
is within the top- [Equation] largest elements in row
[Equation]
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.
3. P ARALLEL B IDIRECTIONAL C ONTEXT M ODEL
The transformer architecture can be used to efﬁciently train autoregressive model by masking the attention computation such that the i -th position can only be inﬂuenced 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 ﬁrst 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 . ﬁrst slice).
Therefore, the decoding of this structure can be performed in two serial processing. [Figure]
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 beneﬁt accurate prediction.
For this reason, though the performance of the ﬁrst 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).
E XPERIMENTAL R ESULTS
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 Model [Table] 2D 0.
[Table]
Transformer
Relative
[Table]
Encoding 1D+1D 0.
Relative Position Encoding 2D 0.
[Table]
[Table]
Relative Position Encoding 2D+Diamond 0.
[Figure]
c r
[Figure]
[Figure]
i
[Figure]
i
[Figure]
[Figure]
has a multi-head self-attention (heads=6), and a feed-forward network (ratio=4).
If not speciﬁed, 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 conﬁgurations.
The detailed structure and the experimental settings are described in appendix.
4. A NALYSIS ON P OSITION E NCODING 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. [Table] [Table]
4. N UMBER k OF TOP - k SELF - ATTENTION
The impact of top- k scheme is analyzed in Figure 7.
The parameter k speciﬁes 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 inﬂuences 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 reﬂect that removing the irrelevant tokens beneﬁts the convergence of Entroformer training.
4. P ARALLEL B IDIRECTIONAL C ONTEXT M ODEL
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-efﬁcient 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.
4. V ISUALIZATION
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 ﬁnds 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 [Figure] [Figure] 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 ﬁlters out irrelevant information from background.
It successfully makes the attention concentrating on the most informative context.
D ISCUSSION
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 ﬁrst 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 beneﬁts 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 ﬁne-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.
A A PPENDIX A. PSNR (RGB) ON K ODAK DATASET
[Figure]
A. MS-SSIM (RGB) ON K ODAK DATASET
[Figure]
A. P OSITION E NCODING FOR VARIOUS IMAGE SIZE
[Table]
A. A RCHITECTURE DETAILS
[Table]
layers are speciﬁed with the “Conv” preﬁx followed by the kernel size, number of output channels and downsampling stride ( e.g . the ﬁrst layer of the encoder uses 5 × 5 kernels with 192 output channels and a stride of 2).
The “Deconv” preﬁx corresponds to upsampled convolutions. The “Trans” corresponds to transformer block as in Vaswani et al. (2017), followed by the inner dimension and head number. A transformer block consists of a self-attention module and a feed-forward module (the dimension of the feed-forward is 4 times of inner embedding’s).
GDN stands for generalized and divisive normalization, and IGDN is inverse GDN (Ball ´ e 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. [Table]
[Table]
[Table] Reference [Table] 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. T RAINING 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 , (cid:15) = 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 ﬁnetune 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.