PaperlayerView exact revision on arXiv

Source-linked AI summary

Learned Image Compression with Mixed Transformer-CNN Architectures

Jinming Liu, Heming Sun, Jiro Katto

arXiv:2303.14978v1eess.IVcs.CV

TL;DR

LIC methods must balance the complementary local modeling of CNNs and non-local modeling of transformers with suitable complexity. The paper proposes parallel TCM blocks and a channel-wise entropy model with channel-squeezed SWAtten modules; it reports state-of-the-art rate-distortion performance on Kodak, Tecnick, and CLIC.

  • Problem

    Existing LIC methods are mainly CNN- or transformer-based, creating a need to exploit their complementary abilities while effectively fusing them at suitable complexity.

  • Method

    The paper combines parallel Transformer-CNN Mixture blocks with a channel-wise entropy model using parameter-efficient swin-transformer attention and channel squeezing.

  • Results

    12.30%, 13.71%, and 11.85% BD-rate improvements over VVC (VTM-12.1) are reported on Kodak, Tecnick, and CLIC, respectively.

  • Takeaways & Limitations

    The experiments report state-of-the-art rate-distortion performance across Kodak, Tecnick, and CLIC Professional Validation datasets.

Abstract

from arXiv · show

Learned image compression (LIC) methods have exhibited promising progress and superior rate-distortion performance compared with classical image compression standards. Most existing LIC methods are Convolutional Neural Networks-based (CNN-based) or Transformer-based, which have different advantages. Exploiting both advantages is a point worth exploring, which has two challenges: 1) how to effectively fuse the two methods? 2) how to achieve higher performance with a suitable complexity? In this paper, we propose an efficient parallel Transformer-CNN Mixture (TCM) block with a controllable complexity to incorporate the local modeling ability of CNN and the non-local modeling ability of transformers to improve the overall architecture of image compression models. Besides, inspired by the recent progress of entropy estimation models and attention modules, we propose a channel-wise entropy model with parameter-efficient swin-transformer-based attention (SWAtten) modules by using channel squeezing. Experimental results demonstrate our proposed method achieves state-of-the-art rate-distortion performances on three different resolution datasets (i.e., Kodak, Tecnick, CLIC Professional Validation) compared to existing LIC methods. The code is at https://github.com/jmliu206/LIC_TCM.

1. Introduction

Learned image compression offers an end-to-end alternative to classical standards, while CNN- and transformer-based methods provide complementary local and non-local modeling abilities. The paper combines these strengths with controllable complexity and adds a parameter-efficient entropy-model attention design, reporting state-of-the-art results across three datasets.

  • Learned image compression is optimized end-to-end and has outperformed VVC on PSNR and MS-SSIM in recent work.
  • CNNs model local patterns, whereas transformers capture non-local information, motivating their efficient combination in LIC.
  • Attention modules are moved to the smaller channel-wise entropy-model path, and the number of latent slices is reduced from 10 to 5 to balance speed and rate-distortion performance.
  • The proposed TCM blocks combine CNN and transformer modeling abilities while maintaining controllable complexity.
  • The paper introduces a channel-wise autoregressive entropy model using parameter-efficient SWAtten modules with channel squeezing.
  • 12.30%, 13.71%, and 11.85% BD-rate reductions versus VVC (VTM-12.1) are reported on Kodak, Tecnick, and CLIC, respectively.

2. Related Work

Prior learned image-compression research developed CNN, transformer, context, and attention-based approaches. These works established complementary modeling strategies and exposed computation-latency trade-offs that motivate the paper’s architecture.

  • Early LIC work introduced end-to-end CNN models, VAE architectures with hyper-priors, and local or global context models for entropy estimation.
  • Recent transformer-based LIC models use Swin Transformers, ViT-assisted entropy modeling, or sliding windows to capture global information with reduced complexity.
  • Image-compression attention modules improve rate-distortion performance, but non-local attention can be time-consuming and local attention captures only local information.
  • The paper’s framework and TCM block are illustrated with encoder and decoder coding operations, while arrows indicate twofold feature-size changes between layers.

3. Proposed Method

The proposed LIC framework combines local CNN modeling with non-local transformer modeling through parallel TCM blocks and uses a channel-wise entropy model with SWAtten and channel squeezing. The design targets rate-distortion performance while controlling complexity through parallel feature processing, sliced entropy modeling, and reduced channel dimensions.

  • Channel-wise Entropy Model: The channel-wise entropy model sequentially processes latent slices, using previously encoded slices to estimate distribution parameters for each current slice.The hyper-prior provides mean and scale features, while each slice network estimates Φ_i = (µ_i, σ_i) for entropy coding.
  • Transformer-CNN Mixture Blocks: TCM blocks split features into parallel CNN and Swin-Transformer branches, then concatenate and fuse their local and non-local representations with a skip connection.The split reduces channels entering subsequent branches, while two stages use W-MSA and SW-MSA for feature fusion.
  • Transformer-CNN Mixture Blocks: The main and hyper-prior paths attach TCM blocks after residual processing modules to obtain both local and non-local information.TCM blocks are inserted throughout the main path except its final layer and are also incorporated into redesigned hyper-prior encoder and decoder paths.
  • Transformer-CNN Mixture Blocks: ERF comparisons evaluate whether TCM combines local and long-distance information relative to transformer-based and CNN-based models.The analysis uses clipped gradients at thresholds t = 0.01 and t = 0.0001 to visualize information sensitivity at different distances.
  • SWAtten Module: Reducing entropy-model slices from 10 to 5 and squeezing every slice input to 128 channels balances model complexity and rate-distortion performance.An unsqueeze operation restores the output channels to the original slice-dependent dimensionality.

4. Experiments

Experiments evaluate training settings, datasets, model complexity, rate-distortion performance, component ablations, and qualitative reconstruction quality. The proposed models achieve strong performance across Kodak, Tecnick, and CLIC Professional Validation, with TCM, SWAtten, and channel squeezing contributing to the reported results.

  • Experimental Settings: Experiments use Kodak, Tecnick, and CLIC Professional Validation datasets, measuring distortion with PSNR and MS-SSIM and bitrate with bpp.Kodak images are 768 × 512, Tecnick images are 1200 × 1200, and CLIC images have 2k resolution.
  • Experimental Settings: Three model sizes vary the middle-layer channel number C as 128, 192, or 256, while attention models use five entropy slices instead of ten.The small, medium, and large models are designed to explore performance under different complexities.
  • Rate-Distortion Performance: At the same bitrate on Kodak, the method improves up to about 0.4dB PSNR and 0.5dB MS-SSIM compared with state-of-the-art methods.Both PSNR and MS-SSIM are evaluated, with MS-SSIM converted to −10log10(1 − MS-SSIM) for comparison.
  • Rate-Distortion Performance: 12.30%, 13.71%, and 11.85% BD-rate improvements over VVC (VTM-12.1) are reported on Kodak, Tecnick, and CLIC, respectively.The VVC anchor is assigned BD-rate=0% for the comparisons.
  • Ablation and Component Comparisons: TCM outperforms Transformer-only and CNN-only models, while SWAtten improves rate-distortion performance and channel squeezing saves parameters with comparable performance.SWAtten also achieves the best rate-distortion performance among the compared attention modules.
  • Complexity and Qualitative Results: All small, medium, and large models achieve state-of-the-art rate-distortion performance, with performance improving as complexity increases.On Kodak, the proposed method also outperforms two other state-of-the-art works in efficiency and rate-distortion performance.
  • Complexity and Qualitative Results: The Kodak reconstruction example retains more detail in complex textures than WebP and VVC (VTM 12.1), including a clearer feather outline.The comparison uses decompressed images from kodim23.

5. Conclusion

The paper combines Transformer and CNN modeling in an efficient TCM block and adds SWAtten to channel-wise entropy models. Experiments report improved comparisons with single-architecture models and state-of-the-art performance across three datasets.

  • Conclusion: The proposed TCM block combines CNN local modeling and Transformer non-local modeling in an efficient parallel architecture.The architecture is designed with controllable complexity for learned image compression.
  • Conclusion: A swin-transformer-based attention module improves the channel-wise entropy model.The module uses channel squeezing for parameter efficiency.
  • Conclusion: The TCM-based model outperforms CNN-only and Transformer-only models under a suitable complexity, while SWAtten surpasses previous image-compression attention modules.The conclusion reports state-of-the-art performance on Kodak, Tecnick, and CLIC Professional Validation.
  • Evaluation Details: The reported evaluation uses VTM-12.1 for VVC and Pillow’s API for WebP comparisons.The WebP quality values are set to 5, 10, 15, 20, 25, 30, 35, 40, 45, and 50.

B. Detailed Network Architecture

The detailed architecture specifies TCM placement, attention-head dimensions, and channel configurations for the main and hyper-prior paths. Small, medium, and large models use progressively larger channel widths.

  • TCM Placement: TCM blocks follow each RBS or RBU in the main path, except after the final layers of ga and gs.TCM blocks are also added to the hyper-prior path by redesigning ha and hs.
  • TCM Configuration: TCM head dimensions in ga and gs are {8, 16, 32, 32, 16, 8}, while those in ha and hs are 32.The architecture uses different head configurations for the main and hyper-prior paths.
  • Model Sizes: TCM channel numbers C are 128, 192, and 256 for the Small, Medium, and Large models, respectively.The corresponding middle convolutional-layer channel numbers in RBS and RBU are 64/96/128.

C. Comparison with Recent LIC Works

The comparison evaluates BD-rate relative to a VVC anchor across datasets with different resolutions. The reported results indicate that the proposed method outperforms prior works and achieves state-of-the-art performance.

  • Evaluation Protocol: BD-rate is computed from PSNR-BPP curves, with VVC results on each dataset serving as the BD-rate=0% anchor.Lower BD-rate represents higher performance.
  • Evaluation Protocol: The VVC implementation uses VTM-12.1, while CompressAI scripts are used to evaluate the learned compression models.These tools support the reported comparative evaluation.
  • Comparison Results: The proposed method outperforms previous works and achieves state-of-the-art performance across Kodak, Tecnick, and CLIC datasets with different resolutions.The paper reports this conclusion from the BD-rate comparisons.

D. Ablation Studies on Various Entropy Estimation Models

The TCMmain model is evaluated with a spatial-wise entropy model against CNN- and swin-transformer-based alternatives. It achieves the best rate-distortion performance and remains effective across entropy-model designs.

  • TCMmain uses a TCM block in its main path and a spatial-wise entropy model for comparison.It is compared with the models in [6] and SwinT-Hyperprior, which use CNN and swin-transformer main paths, respectively.
  • TCMmain achieves the best rate-distortion performance among the compared spatial-wise entropy models.The comparison includes TCMmain, the CNN-based model in [6], and the swin-transformer-based SwinT-Hyperprior model.
  • The TCM blocks improve image compression and are robust to different entropy models.

E. Ablation Studies on the Numbers of Slices

The number of slices controls the trade-off between rate-distortion performance and efficiency in the channel-wise entropy model. Experiments select five slices because larger settings provide little or negative additional improvement.

  • A larger slice count lowers efficiency, whereas a smaller count worsens rate-distortion performance.
  • Rate-distortion performance improves as the slice count increases from low settings, but gains become unclear when s > 5.The tested settings are s = {2, 4, 5, 8, 10}.
  • The model uses s = 5 to balance running speed and rate-distortion performance.The authors report that five slices learn enough information for their model.

F. Abaltion Studies on the Design of SWAtten

The SWAtten ablation evaluates the effects of its CNN-based attention-map component and swin transformer. Either component improves PSNR while reducing bitrate.

  • Removing either the CNN or swin-transformer component from SWAtten still allows about 0.1dB PSNR improvement with fewer bitrates.The ablation compares SWAtten without CNN for the attention map and without the swin transformer.
  • The ablation evaluates SWAtten without CNN for the attention map and without the swin transformer.

G. Visualization

The visualization compares TCM-based, CNN-based, and transformer-based models on Kodak images kodim19 and kodim20. It examines differences in two local regions using PSNR and bitrate annotations.

  • The comparison uses Kodak images kodim19 and kodim20 with TCM-based, CNN-based, and transformer-based models.
  • The visualization focuses on two local regions where reconstruction differences between the three models are noticeable.
  • In the upper local area of kodim19, the TCM-based method reconstructs some road-sign details and makes back fences easier to differentiate.
  • Figure 15 lists PSNR|Bit-rate in the last column, while Figure 16 places PSNR|Bit-rate on the subfigures’ right.
Loading 2303.14978v1…