Contextformer: A Transformer with Spatio-Channel Attention for Context Modeling in Learned Image Compression

TL;DR

This work introduces Contextformer, a transformer-based context model that extends attention to spatio-channel dependencies for learned image compression. By reorganizing latent representations into sequences and employing both spatial and cross-channel attention, the method achieves state-of-the-art rate-distortion performance, outperforming VTM 16.2 and several learning-based baselines on standard benchmarks. The paper also contributes practical runtime optimizations (BDS, SCS, and wavefront decoding) to mitigate the quadratic cost of attention, and provides extensive analyses on model size, generalization, and qualitative quality. Overall, Contextformer demonstrates the potential of adaptive, multi-dimensional attention for entropy modeling in learned compression, with meaningful gains in both objective metrics and visual fidelity.

Abstract

Entropy modeling is a key component for high-performance image compression algorithms. Recent developments in autoregressive context modeling helped learning-based methods to surpass their classical counterparts. However, the performance of those models can be further improved due to the underexploited spatio-channel dependencies in latent space, and the suboptimal implementation of context adaptivity. Inspired by the adaptive characteristics of the transformers, we propose a transformer-based context model, named Contextformer, which generalizes the de facto standard attention mechanism to spatio-channel attention. We replace the context model of a modern compression framework with the Contextformer and test it on the widely used Kodak, CLIC2020, and Tecnick image datasets. Our experimental results show that the proposed model provides up to 11% rate savings compared to the standard Versatile Video Coding (VVC) Test Model (VTM) 16.2, and outperforms various learning-based models in terms of PSNR and MS-SSIM.

Figures (16)

• Figure 1: Illustration of the latent elements used by the context model $\left(\right)$ to estimate the entropy of the current latent $\left(\right)$ in (a– e) for the prior-arts and (f– k) our proposed context model. Previously coded and yet to be coded elements are displayed as $\left(\right)$ and $\left(\right)$, respectively. The displayed prior-art models are (a) multi-scale 2D context cui2020gcui2021asymmetriczhou2019multi, (b) 3D context liu2019nonliu2019practicalmentzer2018conditional, (c) channel-wise autoregressive context minnen2020channel, (d) 2D context with global reference qian2020learning, and (e) context with advanced global reference guo2021causal. Note that in (c) each $\left(\right)$ is coded simultaneously by using only a part of the elements presented as $\left(\right)$, and in (e), the primary channel segment is shown at the bottom for better visibility. Our models with different configurations are shown in (f) Contextformer($N_{cs}{=}1$), (g) Contextformer($N_{cs}{>}1, sfo$), (h) Contextformer($N_{cs}{>}1, cfo$). Note that the (serial) transformer-based context model of qian2021entroformer employs similar mechanism as (f). (i– k) show the versions of our models (f– h) using the sliding window attention.
• Figure 2: Illustration of (a) our proposed model with the Contextformer, (b) sequence generator for the Contextformer with spatial attention, Contextformer$(N_{cs}{=}1)$, and (c) sequence generator for the Contextformer with spatio-channel attention, Contextformer$(N_{cs}{>}1)$. The prepended start token is shown in dark gray in (b-c). Inspired by lee2018context, we use channel-wise local hyperprior neighbors to increase performance; thus, regardless of the selected $N_{cs}$, we apply the sequence generator depicted in (b) to the output of the hyperdecoder.
• Figure 3: Illustration of (a) the rate-distortion performance and (b) the rate savings relative to VTM 16.2 as a function of PSNR on the Kodak dataset showing the performance of our model compared to various learning-based and classical codecs.
• Figure 4: Illustration of the rate-distortion performance in terms of MS-SSIM on Kodak dataset showing the performance of our model compared to various learning-based and classical codecs. All learned methods were optimized for MS-SSIM.
• Figure 5: Comparison of the rate-distortion performance (a) on CLIC2020-Professional (solid line, left vertical axis) and CLIC2020-Mobile (dashed line, right vertical axis) datasets, and (b) on Tecnick dataset.