Image Compression using only Attention based Neural Networks
In recent research, Learned Image Compression has gained prominence for its
capacity to outperform traditional handcrafted pipelines, especially at low
bit-rates. While existing methods incorporate convolutional priors with
occasional attention blocks to address long-range dependencies, recent advances
in computer vision advocate for a transformative shift towards fully
transformer-based architectures grounded in the attention mechanism. This paper
investigates the feasibility of image compression exclusively using attention
layers within our novel model, QPressFormer. We introduce the concept of
learned image queries to aggregate patch information via cross-attention,
followed by quantization and coding techniques. Through extensive evaluations,
our work demonstrates competitive performance achieved by convolution-free
architectures across the popular Kodak, DIV2K, and CLIC datasets.
Natacha Luka, Romain Negrel, David Picard
# Image Compression using only Attention based Neural Networks

###### Abstract

In recent research, Learned Image Compression has gained prominence for its capacity to outperform traditional handcrafted pipelines, especially at low bit-rates. While existing methods incorporate convolutional priors with occasional attention blocks to address long-range dependencies, recent advances in computer vision advocate for a transformative shift towards fully transformer-based architectures grounded in the attention mechanism. This paper investigates the feasibility of image compression exclusively using attention layers within our novel model, _QPressFormer_. We introduce the concept of learned _image queries_ to aggregate patch information via cross-attention, followed by quantization and coding techniques. Through extensive evaluations, our work demonstrates competitive performance achieved by convolution-free architectures across the popular Kodak, DIV2K, and CLIC datasets.

_Keywords --_ Learned Image Compression, Vision Transformers, Visual Attention

## 1 Introduction

Image compression stands as a pivotal research domain, driven by the ubiquity of images in today's digital landscape. The persistent pursuit of higher image quality, while simultaneously optimizing bit size for storage and transmission, underscores its paramount importance. For many years, image compression codecs were meticulously handcrafted, yielding iconic solutions like the widely-used JPEG algorithm. In the post-JPEG era, a lineage of handcrafted codecs emerged, including JPEG2000 [2], BPG format [3] (leveraging the HEVC video compression standard [4]), and Google's WebP [5]. Nevertheless, recent years have witnessed a transformative shift towards learned image compression, propelled by deep learning approaches. In the contemporary landscape, these learned codecs have not only achieved parity with their handcrafted counterparts but have also demonstrated the potential to surpass them.

In both hand-designed and learned image compression paradigms, common components are shared within the encoding and decoding processes. During encoding, an analysis transform maps the image from its pixel space representation to a more compression-friendly domain, followed by quantization, a pivotal step in lossy codecs. Subsequently, a lossless coding scheme translates the image data into a compact bitstream. In the decoding phase, inverse operations unfold: the received bits are decoded, followed by synthesis through a suitable transform, ultimately restoring the image to its pixel space. A distinctive advantage inherent in learned compression methods arises from their holistic optimization approach. In this paradigm, all constituent blocks are collectively learned and end-to-end optimized. This optimization strategy diligently works to minimize distortion while concurrently minimizing bit-rate consumption. Conversely, traditional hand-designed methods require the separate optimization of individual blocks, which may not inherently synergize when used together. Furthermore, hand-designed approaches involve human-centric decisions in selecting the transform, introducing inherent bias. In stark contrast, learned methods offer greater flexibility, as the network autonomously adapts both the analysis and synthesis transforms, determining the optimal compromise between rate and quality within specified constraints.

In learned image compression, numerous approaches have emerged for the analysis stage. Remarkably, transformer-based architectures, introduced by Vaswani et al. [6], have remained relatively unexplored, despite their pervasive adoption in Computer Vision where they have demonstrated their supremacy over CNNs in many computer vision tasks [7]. It is interesting to note that while attention or transformer blocks find occasional integration, they are almost invariably utilized _in conjunction with_ other usual convolution blocks, in both the
analysis and synthesis processes [8, 9]. Furthermore, a critical prerequisite for effective image compression is the reduction of spatial redundancy, and the transformer's attention mechanism, which aggregates related elements within a sequence, appears especially well-suited for this purpose. With this context in mind, we propose in this paper a comprehensive investigation of the potential benefits of relying _exclusively_ on the transformer's attention mechanism for image compression, thereby eliminating convolutional priors from the equation.

The remaining of this paper is structured as follows: First, we present recent achievements in learned image compression, as well as basics on attention mechanisms. Then, we present and explain our transformer-based analysis/synthesis model based on our _image queries_ concept. These _image queries_ are learned prototypes which aggregate information from image step by step thanks to transformer cross-attention blocks. Finally, we show and analyze results of the proposed architecture on several datatsets to give some insights about the inner working of our proposed model _QPressFormer_, before we conclude.

## 2 Related Work

### Learned Image Compression

In the realm of image compression, conventional lossless and lossy codecs, typified by ubiquitous standards like JPEG and its successors such as JPEG2000 [2], WebP [5], and BPG [3], have long prevailed as the industry norm. These well-established codecs continue to be indispensable in contemporary applications, exemplified by the ongoing development of the VVC codec [10] for video, with adaptions tailored for image compression. Nonetheless, the recent years have witnessed a burgeoning interest in deep learning-driven approaches for both lossless and near-lossless image compression [11], marking a noteworthy departure from conventional methodologies.

Todericci et al. introduced an innovative approach using recurrent neural networks (RNNs) [12, 13] for variable-rate image compression. Their method involves multiple passes of the original image through the network, generating an \(m\)-bit representation at each step to control the final bit-rate. While they explored various scenarios, they did not incorporate attention mechanisms.

Another strategy entails training multiple networks for distinct bit-rates. Balle et al. introduced an innovative approach [14, 1] centered on learning the probability distribution of images within a latent space, with a concurrent focus on minimizing both the distortion and the entropy of this learned distribution. As per Shannon's theory, entropy serves as a lower bound for code length, rendering it synonymous with bit-rate minimization. The delicate equilibrium between these two objectives delineates the distortion/bit-rate trade-off. Their methodology incorporates convolutional transforms, paired with normalization functions [15] for image compression. To enhance the estimation of the latent probability distribution, Balle and colleagues introduced side information through latent variable models, where the variables obtained from the analysis transform are

Figure 1: Image reconstruction from DIV2K. Zoom comparisons with original, convolutional equivalent model from [1] and BPG are shown on the sides. Even at lower PSNR, our method significantly improves on the details of the penguin’s fur, as captured by the lpips metrics. Examples on Kodak and CLIC in appendix A.
treated as observed variables. A prior, referred to as the hyper-prior, is learned on these variables, effectively serving as a 'prior of the prior' to better account for spatial dependencies within the latent space. The ensuing hyper-latent variables are subsequently learned, contributing to the distribution in the hyper-latent space whose entropy is minimized to minimize bit-rate of this side-information. These hyper-variables allowing the variances inference of the original latent variables modeled as zero-mean Gaussian variables.

In recent works [16], significant advancements were made in optimizing the image compression pipeline, with two notable steps. Causal consideration of context began to be incorporated to enhance the inference of parameters within the distribution of subsequent latent variables. Building upon this foundation, attention mechanisms were integrated to encompass contextual information, wherein latent variables were modeled as Gaussian mixture models (GMMs), with parameter inferences facilitated by hyper latent variables [17]. Concurrently, several works embraced attention mechanisms to more effectively account for contextual nuances, particularly concerning bit allocation. These mechanisms dynamically allocate more bits to regions of higher attention, signifying their increased complexity and importance [18, 19]. Notably, while transformers have occasionally served as attention modules within these architectures [9, 8], they have been employed alongside convolutional components, distinguishing our approach.

Finally, generative models have also been investigated [20, 21, 22]. This results in better perception quality image especially in low bit-rate setup even if it is sometimes not reflected by distortion metrics like PSNR or MS-SSIM due to the Perception-Distortion trade-off [23]. Similarly to other approaches, even when attention is used it is in addition to convolutions.

### Transformers

Transformers, initially introduced by Vaswani et al. in the domain of Natural Language Processing (NLP) for translation tasks [6], have since permeated diverse NLP applications, from BERT [24] to LlaMA [25] models. Within Computer Vision (CV), Dosovitskiy et al. inaugurated the integration of transformers [26], which are now prevalent in various vision tasks encompassing object detection [27], image synthesis [28], or cross-modal endeavors such as Text-Image Retrieval and Visual Question Answering [29]. Notably, transformers have supplanted traditional CNNs [7].

The transformer architecture relies on an attention mechanism, facilitating interactions between two distinct sequences of tokens. Its fundamental objective is to iteratively aggregate tokens from a sequence denoted as _values_, organized in a matrix \(V\), into another sequence known as _queries_, stored in a matrix \(Q\). This aggregation is guided by the similarity between elements in \(Q\) and \(V\). To achieve this, the _values_ are subject to \(h\) projections, distributed across different sub-units termed _heads_, yielding \(h\) sets of _keys_, amalgamated within a matrix \(K\). For each set of _keys_, a similarity score \(s_{i,j}\) is computed using the dot product between \(q_{i}\) and \(k_{j}\), where \(k_{j}\) corresponds to specific features of \(v_{j}\). These similarity scores are gathered within a matrix, normalized through the softmax function, and designated as the attention matrix \(A\). Thus, at each iteration \(n\), \(h\) attention matrices are computed using the following equation:

\[A^{n}=\sigma\left(Q^{n}\left(K^{n}\right)^{T}\right) \tag{1}\]

Once the attention matrices are obtained, they are employed to compute a weighted sum of the _values_, subsequently added to the _queries_. Consequently, the _queries_ accumulate the most akin _values_ corresponding to \(h\) specific features, as determined through the projection of _values_ into \(h\) sets of _keys_. This process is concisely captured in the following equation, where \([\dots]_{h}\) denotes the aggregation of the \(h\) sub-units:

\[Q^{n+1}=Q^{n}+\left[A^{n}V^{n}\right]_{h}=Q^{n}+\left[\left(\sum_{k}a_{ik}^{n }v_{kj}^{n}\right)_{i,j}\right]_{h} \tag{2}\]

The equation 2 correspond to the main module of transformer called Multi-Head-Attention denoted MHA in the following. In order to preserve the dimension \(d\) of the whole transformer, all elements in the \(h\) sub-units are projected in a space of dimension \(d_{h}=\frac{d}{h}\), which implies \(d\equiv 0\left(mod\ h\right)\).

There are 2 different MHA modules. The Self-Attention module, \(\text{SA}\left[Q\right]\), correspond to the case where each _query_ aggregates itself with its most similar _queries_ according to particular features. In this case _queries_ and _values_ are identical. The other case correspond to the Cross-Attention, \(\text{CA}\left[Q,V\right]\). Here, the set of _queries_ and _values_ are different. The sequence of _queries_ aggregates information from another sequence. More details are available in [6].

## 3 Method

In this section, we introduce our convolution-free architecture, denoted QPressFormer, which exclusively relies on attention blocks. For the compression component, we adopt a factorized prior model similar to [1]. To ensure
comparability, our aim is to develop a model with similar target bitrate, while replacing convolutional analysis and synthesis transforms with attention-based counterparts. The whole pipeline of our QPressFormer model is illustrated in Figure 2.

An image \(I\) is first divided into patches, which are unrolled and augmented with positional encodings corresponding to their patch index using the transformation denoted as \(\varphi\). The resulting set of vectors serves as the _values_/_keys_ (\(V_{E}/K_{E}\)) for the encoder \(E\), which employs attention blocks exclusively. Additionally, a set of \(N\) learned vectors, referred to as 'Image queries' (\(Q_{E}^{in}\)), serves as _queries_. The set \(Q_{E}\) aggregates information from the image at each step \(n\) through the attention process outlined in Section 2.2. In-depth, the blocks of \(E\) consist of transformer decoder blocks [6], governed by the following equations for each layer \(n\):

\[\begin{split} Q^{n+1/2}=\operatorname{CA}\left[\operatorname{ SA}\left[Q^{n}\right],\varphi(I)\right]\\ Q^{n+1}=Q^{n+1/2}+\operatorname{FFW}^{n}(Q^{n+1/2})\end{split} \tag{3}\]

At the final stage, the output _queries_ (\(Q_{E}^{out}\)) containing the compressed image. In this study, we work with 256x256 RGB images divided into 16x16 patches. This results in 256 patches, flattened as 768-dimensional vectors, which sets the embedded dimension of our architecture. We opt for \(N=64\)_Image Queries_. \(Q_{E}^{out}\) corresponds to the latent variables so the probability distribution of \(Q_{E}^{out}\) is estimated using a network and its entropy minimized to reduce the bit-rate, as discussed in Section 2.1. Importantly, our reconstruction perception metric is not the Mean-Square-Error (MSE) but lips [30], a perceptual metric that aligns better with human perception.

The decoding process mirrors the encoding procedure. The vectors \(\{Q_{E}^{out}\}\), containing the compressed image, serve as _values_ (\(V_{D}\))/_keys_ (\(K_{D}\)) for the image decoder \(D\). The _queries_ in \(D\) are learned vectors corresponding to patch prototypes. They aggregate the compressed image progressively, and in the final stage, are reshaped and reassembled to yield the decompressed image \(\hat{I}\).

Figure 2: QPressFormer : Our Proposed Architecture.
The architecture is trained on the ImageNet using Adam with a learning rate of 1e-4. Beyond 100,000 steps, optimizing lpips becomes more challenging. We observed that upscaling images to 512x512 before using lpips enhances results without requiring other changes. Both \(E\) and \(D\) have 12 heads and a depth of 12.

## 4 Results

### Quantitative Results

In Table 1, we present a comparative analysis between our model, BPG and equivalent convolutional factorized prior models from [1] with pretrained weights obtained from CompressAI [31] on three datasets (Kodak [32], CLIC [33] and DIV2K [34]). For the MSE and MS-SSIM optimized models, we select weights corresponding to the closest bit per pixel (bpp) to our model. Our evaluation employs a range of metrics, including PSNR, MS-SSIM [35], lpips [30], and FID [36], which assesses the realism of generated images. Given that our architecture was trained on 256x256 images, we decompose all images into 256x256 segments and input each 256x256 part to the network before reassembling them. In cases where the image dimensions are not divisible by 256x256, we employ center cropping to obtain the largest segment divisible by 256x256. To ensure a fair comparison with convolutional counterparts or hand-designed BPG format, we apply the same image decomposition for reconstruction.

As shown in Table 1, the convolutional model optimized on MSE (respectively on MS-SSIM) performs better in terms of PSNR (respectively on MS-SSIM). In the same manner, our model performs significantly better on lpips. Moreover, our model performs better on FID even if it was not trained specifically for it. Besides, even if the PSNR and MS-SSIM of our model is lower than convolutionnal counterpart, it captures better perceptual details - as expected with the lpips optimization - as shown in figure 1. Convolution bias is thus not a prerequisite for Image compression and a transformer architecture seems at least equivalent or better - in terms of perceptual metrics FID and lpips - than convolutional networks and BPG.

### Visual analysis

To understand how the attention blocks process the image during encoding, we perform several visualizations. During the attention process, the attention weights are computed for a given pair of _query_ and _key_. As each _key_ corresponds to a patch, we can see which part of the image one _query_ pays attention to by plotting its attention matrix. By combining the attention matrices of all queries, layers and heads (with **maximum**), we obtain a heatmap showing which parts of the image have been processed (Figure 3). As expected for good reconstruction, the whole image is attended, almost uniformly. In other words, at least one query pays attention only once in each image area.

To better understand what information is captured by a _query_, we visualize the layers and heads **mean** attention map for each _query_ individually. An example is given in Figure 3(c). It is worth noting the **maximum** of attention maps and the **mean** of attention maps able us to make two distinct observations. The mean

\begin{table}
\begin{tabular}{l|c||c|c|c|c} Method & bpp & PSNR\(\uparrow\) & MS-SSIM\(\uparrow\) & lpips\(\downarrow\) & FID\(\downarrow\) \\ \hline  & & \multicolumn{4}{c}{Kodak} \\ QpressFormer & 0.296 & 27.19 & 0.9203 & **0.2126** & **68.60** \\ Conv+mse [1] & 0.258 & 31.59 & 0.9608 & 0.3386 & 154.06 \\ Conv+mssim [1] & 0.320 & 30.66 & **0.9745** & 0.2913 & 117.94 \\ BPG [3] & 0.301 & **32.78** & 0.9626 & 0.2872 & 110.90 \\  & & \multicolumn{4}{c}{DIV2K (val-set)} \\ QpressFormer & 0.300 & 25.50 & 0.9159 & **0.2267** & **19.19** \\ Conv+mse [1] & 0.290 & 30.10 & 0.9642 & 0.3429 & 63.12 \\ Conv+mssim [1] & 0.322 & 28.94 & **0.9716** & 0.3237 & 49.78 \\ BPG [3] & 0.307 & **30.84** & 0.9604 & 0.3155 & 54.73 \\  & & \multicolumn{4}{c}{CLIC (test-set)} \\ QpressFormer & 0.286 & 27.96 & 0.9339 & **0.2244** & **9.63** \\ Conv+mse [1] & 0.230 & 32.33 & 0.9657 & 0.3820 & 51.85 \\ Conv+mssim [1] & 0.277 & 31.32 & **0.9758** & 0.3536 & 36.83 \\ BPG [3] & 0.297 & **34.31** & 0.9712 & 0.3216 & 40.78 \\ \end{tabular}
\end{table}
Table 1: Quantitative results on Kodak, CLIC and DIV2K datasets.
correspond to the area which is the most watched by the _query_ whereas the **maximum** allows us to ensure that all areas are watched by _queries_ at least one time. An original 256x256 crop from Kodak and the decoded patch are given in Figure 3(a) and 3(b). The figure 3(c) shows the **mean** attention map - across heads and layers - for _query_ 11. The attention is more focused on a part of the roof. To verify that this specific query encodes this particular region, we reconstruct the image without it. We can then see on the figure 3(d) that the area where the removed _query_ paid attention to is not reconstructed properly. Interestingly, the area is inferred and in-painted with the roof texture, indicating that information from others patches able the encoder to infer information in that area.

In Figure 3(e), the reconstruction error without this particular query and averaged over the Kodak dataset shows that the spatial information encode by the _query_ is the same for all images. In fact, after looking at the 64 averaged reconstruction error, nearly all _queries_ encode spatial information. It is worth noting that this spatial encoding property in the queries is an emergent behavior of the model, as it was not designed in the architecture or in the optimization.

However, spatial information is not the whole information encoded by _queries_\(Q_{E}^{out}\). First, a look at **maximum** attention maps of queries allow us to discover that the same _query_ pays attention to areas in different locations for different 256x256 crops of the same kodak image. Figure 9 in appendix C gives more details. Second, if each _query_ was specialized spacially, in-painting wouldn't be possible.

To better visualize the elements encoded by _queries_, we perform a PCA on the whole queries for a Kodak image to reduce the dimensions of _queries_ from 768 to 3. In the following we refers to these reduce dimensional _queries_ as _meta-queries_. We can then visualize a RGB image corresponding to the main, or at least the most important, information decoded at each layer of the decoder. Remember that the output _queries_ of the encoder \(Q_{E}^{out}\) serve as _keys/values_ for the decoder wherein the _queries_ are learned patchs prototypes. After reducing the dimensionality of \(Q_{E}^{out}\), we can project on it the mean across heads of cross-attention matrices (instead of the decoder _values_ obtained from the full encoder _queries_) for a given layer of the decoder with the following

Figure 4: 256x256 Kodak Image Crop Reconstruction with all \(Q_{E}\) or with \(\left\{Q_{E}\right\}_{11}\) removed (denoted \({}^{\dagger}\)).

Figure 3: Attention map (max attn scores for all queries/layers/heads) for a Kodak image.
equation:

\[I_{visu}^{n}=\frac{1}{nb_{heads}}\sum_{h\in heads}\left(A_{h}^{n}\right)Q_{E;[:3]}^{ out} \tag{4}\]

In equation 4 \(Q_{E;[:3]}^{out}\) denotes the _meta-queries_ - _i.e._ output encoder _queries_ obtained after PCA for reducing dimension to 3. The figure 5 shows clearly that _queries_ contain spatial information extracted at some specific layers: figure 4(b) shows the same colors across different positions for different 256x256 blocks of an kodak image in addition to smooth transitioning between colors. In contrast figure 4(d) show a nearly segmentation of the original kodak image demonstrating some semantic information is extracted at others layers in the decoder, and so present in the queries. Thus, to perform a full reconstruction, _queries_ encodes at the same time semantic and spatial information. More visualizations are given in appendix B

## 5 Conclusion

In this study, we explore learned image compression exclusively through attention blocks, eliminating convolutional biases. Our model, _QPressFormer_, introduces "image queries" for cross-attention-based information aggregation. Our results show competitive results with convolutional counterparts with improved perceptual metrics. Future research should further investigate attention-based image compression. While we focused on analysis and synthesis components with a basic prior model, incorporating enhancements on prior model from convolutional models into attention-based architectures can assess transformer-based compression against state-of-the-art convolutional counterparts.

Figure 5: Visualization of information encoded by output encoder _queries_ and extracted at layer 1 and 11 of the decoder. 4(b) shows the layer 1 extract more spatial information from the _queries_ whereas 4(d) shows more semantics extraction with a nearly segmentation of image.