Source-linked AI summary
The Devil Is in the Details: Window-based Attention for Image Compression
Renjie Zou, Chunfeng Song, Zhaoxiang Zhang
TL;DR
Learned image compression models achieve strong rate-distortion performance, but CNN structures do not adequately capture local, non-repetitive textures. The paper studies local-aware attention, introduces window-attention and the STF architecture, and reports effective methods that exceed state-of-the-art rate-distortion performance while identifying Transformer design constraints.
Problem
CNN-based learned compression models are not designed to capture local redundancy, especially non-repetitive textures, motivating joint use of global structure and local texture.
Method
The paper studies local-aware attention, proposes window-attention as a plug-and-play module, and designs STF with Transformer blocks in down-sampling and up-sampling stages.
Results
The proposed methods are reported to be effective and to exceed state-of-the-art rate-distortion performance.
Takeaways & Limitations
Window-attention combines global structure learning with local texture modeling and can enhance both CNN and Transformer compression models.
Takeaways & Limitations
GDN is unstable in deep Transformer architectures and incompatible with Transformer-block attention.
Abstract
from arXiv · showhide
Learned image compression methods have exhibited superior rate-distortion performance than classical image compression standards. Most existing learned image compression models are based on Convolutional Neural Networks (CNNs). Despite great contributions, a main drawback of CNN based model is that its structure is not designed for capturing local redundancy, especially the non-repetitive textures, which severely affects the reconstruction quality. Therefore, how to make full use of both global structure and local texture becomes the core problem for learning-based image compression. Inspired by recent progresses of Vision Transformer (ViT) and Swin Transformer, we found that combining the local-aware attention mechanism with the global-related feature learning could meet the expectation in image compression. In this paper, we first extensively study the effects of multiple kinds of attention mechanisms for local features learning, then introduce a more straightforward yet effective window-based local attention block. The proposed window-based attention is very flexible which could work as a plug-and-play component to enhance CNN and Transformer models. Moreover, we propose a novel Symmetrical TransFormer (STF) framework with absolute transformer blocks in the down-sampling encoder and up-sampling decoder. Extensive experimental evaluations have shown that the proposed method is effective and outperforms the state-of-the-art methods. The code is publicly available at https://github.com/Googolxx/STF.
1. Introduction
Learned compression improves rate-distortion performance over classical standards, but CNN-based models struggle with local, non-repetitive textures. The paper addresses this by combining global feature learning with local-aware attention through window-attention and the STF framework.
- Motivation: Classical standards rely on handcrafted transforming, quantization, and entropy-coding rules for lossy image compression.
- Motivation: Learned VAE-based compression achieves better rate-distortion performance than conventional methods on PSNR and MS-SSIM.
- Motivation: Non-local attention can allocate more bits to challenging edges and textures, but CNN structures retain an intrinsically global-aware character.
- Contributions: Comparative experiments find local attention more suitable for reconstructing local textures than global attention.
- Contributions: Window-attention captures correlations among neighboring spatial elements and functions as a plug-and-play component for CNN and Transformer models.
- Contributions: The STF framework places absolute Transformer blocks in both the down-sampling encoder and up-sampling decoder.
2. Related Works
Prior learned compression work develops CNN, probabilistic, GAN, and attention-based architectures, while this paper identifies compatibility and efficiency constraints relevant to deeper Transformer designs.
- Learned Image Compression: CNN-based VAE architectures improved learned image compression through end-to-end optimization, hyper-priors, autoregressive components, and residual attention modules.
- Learned Image Compression: GMM entropy models can improve rate-distortion performance but require dynamically generated distributions for every element, increasing redundancy and encoding-decoding time.
- Learned Image Compression: SGM entropy models use fixed probability and cumulative distribution tables, making entropy coding less computationally expensive than GMM-based models.
- Learned Image Compression: GAN-based compression methods learn image distributions directly and frame compression as a rate-distortion-perception trade-off.
- Attention Mechanism: The paper reports that GDN is unstable in deep Transformer architectures and incompatible with Transformer-block attention.
3. Method
The method combines local window-based attention with global feature learning in CNN and Transformer architectures for learned image compression. It introduces WAM and a symmetrical Transformer framework, alongside a hyper-prior and channel-wise autoregressive entropy model.
- Formulation: The compression pipeline maps an image x to latent y, quantizes it into discrete ŷ, and decodes ŷ into reconstructed image x̂.Quantization introduces clipping error, which contributes to reconstruction distortion.
- Formulation: The training objective combines latent bit rate with distortion, with λ controlling the rate-distortion trade-off.The rate includes the latents ŷ and ẑ, while distortion compares x with x̂.
- Window-based Attention: Window-based attention divides feature maps into non-overlapping M × M windows and computes attention maps separately within each window.The design targets spatially neighboring elements while reducing the cost associated with large-field attention.
- Window-based Attention: The Window Attention Module replaces the non-local block and allocates more bits to high-contrast complex regions while producing sharper reconstructed texture details.Compared with NLAM and removing attention, WAM allocates fewer bits to low-contrast simple regions such as sky and clouds.
- Transformer-based Architecture: The CNN-based model inserts WAMs into both encoder and decoder, while STF removes convolution layers and uses Transformer blocks in a symmetrical encoder-decoder.STF uses local-window attention, small patches, layer normalization, GELU, patch merging, patch splitting, and de-embedding.
- Transformer-based Architecture: The Transformer architecture addresses patch redundancy and GDN incompatibility by using small patches, layer normalization, GELU, and local-window attention.The receptive field expands gradually through the Transformer architecture while attention remains focused on spatially neighboring patches.
4. Experiments
Experiments evaluate the proposed CNN-based and STF models on Kodak and CLIC datasets, comparing rate-distortion performance, visual quality, efficiency, and ablations. Results show strong rate-distortion performance and improved visual detail, while discussion identifies limitations involving structural modeling, normalization, and perception metrics.
- Experimental Setup: Models are trained on 300k OpenImages crops for 1.8M steps and evaluated using PSNR and MS-SSIM on Kodak and CLIC datasets.Kodak contains 24 images, while the CLIC professional validation dataset contains 41 high-resolution, high-quality images.
- RD Performance: On Kodak, CNN-based and STF models closely match each other under MSE/PSNR and outperform other learned compression methods, while MS-SSIM gains are slight.The paper attributes the limited MS-SSIM improvement to a mismatch between attention-focused bit allocation and MS-SSIM’s treatment of contrast regions.
- RD Performance: On CLIC, comparisons indicate the same conclusion and demonstrate the robustness of the CNN-based and STF models.The evaluation uses average PSNR and MS-SSIM rate-distortion performance.
- Visual Quality: Reconstructed images retain more details at approximately comparable bpp, and MS-SSIM optimization substantially improves visual quality for both CNN-based and STF models.The visual comparison includes JPEG, BPG, and VVC alongside the proposed methods.
- Ablation Study: Window-based attention improves a Minnen2020-based model in ablations and enhances current SOTA CNN compression with negligible computation overhead.The contrast experiments compare removing WAMs or using non-local attention modules.
- Discussion: The CNN architecture remains limited because window attention does not capture structural information sufficiently, while GDN is reported as incompatible with Transformer attention.The authors also report that convolution-generated attention masks achieve comparable rate-distortion performance.
5. Conclusion
The paper concludes that local-aware attention should be combined with globally learned structure for image compression. It presents window-based attention and the STF framework as effective approaches, while identifying convolution-kernel shaping and normalization as future directions.
- 5. Conclusion: Local-aware attention is presented as crucial for combining global structure with local texture in image compression.The proposed window-based module captures correlations among spatially neighboring elements and can enhance CNN or Transformer models.
- 5. Conclusion: The STF framework uses absolute Transformer blocks in both the down-sampling encoder and up-sampling decoder.The paper reports extensive experimental results showing that the proposed methods exceed state-of-the-art rate-distortion performance.
- 5. Conclusion: Future work will investigate convolution-kernel shaping and normalization modes as additional factors affecting local-detail reconstruction.