Checkerboard Context Model for Efficient Learned Image Compression

Dailan He; Yaoyan Zheng; Baocheng Sun; Yan Wang; Hongwei Qin

Source-linked AI summary

Checkerboard Context Model for Efficient Learned Image Compression

Dailan He, Yaoyan Zheng, Baocheng Sun, Yan Wang, Hongwei Qin

arXiv:2103.15306v2Mar 29, 2021eess.IVcs.CV

TL;DR

Serial context models improve learned image compression but constrain decoding parallelism and computational efficiency. The paper replaces them with a checkerboard context model and two-pass decoding scheme, achieving highly parallel decoding while retaining competitive compression performance. Experiments report more than 40× decoding acceleration with almost the same RD performance.

Problem
Serial context models improve rate-distortion performance but lack parallelization, making their decoding computationally inefficient for practical learned image compression.
Method
The paper uses a checkerboard-shaped spatial context model with two-pass decoding, encoding anchors with hyperprior information before decoding non-anchors using parallel context calculation.
Results
The proposed models greatly speed up decoding on parallel devices while maintaining competitive compression performance, with BDBR reductions of −17.0% and −27.4% versus Ballé2018 for Minnen2018 and Cheng2020.
Takeaways & Limitations
The checkerboard approach provides an almost drop-in replacement that improves the practicality of widely used serial context models without changing model structure or capacity.
Takeaways & Limitations
Combining spatial context with channel-wise adaptation is promising but beyond this paper’s scope.

Abstract

from arXiv · show

For learned image compression, the autoregressive context model is proved effective in improving the rate-distortion (RD) performance. Because it helps remove spatial redundancies among latent representations. However, the decoding process must be done in a strict scan order, which breaks the parallelization. We propose a parallelizable checkerboard context model (CCM) to solve the problem. Our two-pass checkerboard context calculation eliminates such limitations on spatial locations by re-organizing the decoding order. Speeding up the decoding process more than 40 times in our experiments, it achieves significantly improved computational efficiency with almost the same rate-distortion performance. To the best of our knowledge, this is the first exploration on parallelization-friendly spatial context model for learned image compression.

1. Introduction

Learned compression improves entropy modeling and RD performance, but widely used spatial context models impose serial decoding that limits efficiency. The paper proposes a checkerboard context model with two-pass decoding to balance compression performance and running efficiency.

1. Introduction: Recent learned codecs outperform JPEG2000 and BPG on both PSNR and MS-SSIM.A learning-based method has also approached VVC intra coding performance.
1. Introduction: Learned image compression methods use nonlinear autoencoder transforms and entropy modeling to reduce latent code length and optimize rate-distortion performance.Hyperprior modeling further improves latent-distribution estimation.
1. Introduction: Context models predict unknown codes from already decoded neighboring latents and, with hyperpriors, improve entropy-parameter estimation.Mean-scale hyperprior models jointly use hyper latent and context to predict location and scale parameters.
1. Introduction: Most recent context models have very low computational efficiency because decoding lacks parallelization, leading practical deployments to omit them.CCN offers only moderate parallelizability and remains limited by image size.
1. Introduction: The proposed checkerboard context model and two-pass decoding method aim to improve the balance between RD performance and running efficiency.The model reorganizes decoding to enable parallel context calculation.

2. Related Work

Related work develops learned compression architectures, entropy models, and deployment-oriented alternatives that address rate-distortion quality or serial processing. The checkerboard approach targets serial processing in the spatial dimension.

2. Related Work: Prior work establishes convolutional autoencoder frameworks for learned image compression.These works form the main architectural foundation discussed in the paper.
2. Related Work: A channel-wise autoregressive entropy model reduces element-level serial processing and achieves SOTA RD performance with additional techniques.Its adaptation operates across channels, whereas the proposed checkerboard context operates spatially.
2. Related Work: Related studies improve entropy models, reconstruction, variable-rate operation, content awareness, architectures, and quantization.The cited directions cover multiple components of learned compression systems.
2. Related Work: Operational autoencoder diagrams distinguish hyperprior-based compression from architectures that also incorporate context modeling.The diagrams show data flow and the training-versus-inference role of quantization.

3. Preliminary

The preliminary framework uses hyperpriors and entropy models to represent and reconstruct images, then adds masked-convolution context modeling to predict latent entropy parameters. The serial mask exposes only causal neighbors, requiring Z-ordered decoding.

3. Preliminary: The scale hyperprior framework transforms an image into quantized latents and hyper latents, using the quantized hyper latent to estimate latent scale parameters.Neural transforms implement analysis, synthesis, and hyperprior encoding and decoding.
3. Preliminary: The rate-distortion objective combines estimated entropy terms for latents and hyper latents with a distortion term weighted by λ.Distortion commonly uses MSE or MS-SSIM.
3. Preliminary: Scale hyperprior entropy modeling uses a conditional Gaussian scale mixture with zero location parameter and scale predicted from the hyperprior.The hyper latent probability is modeled with a non-parametric fully factorized density model.
3. Preliminary: The mean-scale hyperprior adds a context model and parameter-inference network to predict latent entropy location and scale parameters.The context feature is computed from causal neighboring latents that are already visible.
3. Preliminary: Masked convolutions implement context modeling, while the conventional left-top mask permits only already decoded latents and therefore requires strict Z-ordered serial decoding.The mask pattern determines which spatial locations contribute to each prediction.
3. Preliminary: The paper also notes that GMM entropy models can use K groups of predicted mixture parameters, while context is not applied to the hyper latent.These are settings inherited from prior frameworks.

4. Parallel Context Modeling

The paper analyzes context-modeling patterns with a random-mask model and proposes a checkerboard context model with two-pass decoding to improve parallel efficiency while preserving rate-distortion performance.

Motivation: Existing serial context models require strict Z-ordered decoding, limiting computational efficiency despite their strong rate-distortion performance.The paper motivates replacing this serial design with a parallelizable alternative.
Random-Mask Analysis: The random-mask model compares arbitrary context patterns by sharing the latent representation and reconstruction, thereby isolating their bit-rate effects.A full-zero mask provides the context-free baseline, while η(M) quantifies rate saving relative to it.
Random-Mask Analysis: Single-reference-mask experiments on Kodak show that mutual information decays quickly with spatial distance, supporting emphasis on close neighbors.The experiments evaluate 24 single-reference masks with one active location each.
Random-Mask Analysis: Mutual information between visible and current latents explains why nearer spatial neighbors generally provide more useful context for rate saving.The analysis connects entropy reduction to the mutual information term I(Ŷ_i; Ŷ_j<i).
Checkerboard Context Model: The checkerboard context model uses four nearest neighbors and is reported to outperform 3 × 3 and 5 × 5 serial models in random-mask experiments.Its design avoids applying mutually dependent adjacent contexts across the whole latent map.
Parallel Decoding: Two-pass decoding divides latents into anchors and non-anchors: anchors use hyperprior information, while non-anchors use checkerboard context and hyperprior features.The entropy parameters Φ = (µ, σ) are estimated for the latent elements under this arrangement.
Parallel Decoding: The proposed parallel model reduces decoding from H×W sequential steps to a constant 2 steps for an H×W×M latent feature map.The encoding process can likewise obtain all latent entropy parameters in parallel within one pass.

5. Experiments

Experiments compare checkerboard and serial context models under matched architectures, showing much faster parallel decoding with only slight RD degradation and substantial gains over context-free baselines.

Experimental setup: The study trains and tests models on ImageNet-derived data, Kodak, and Tecnick, using matched settings for serial and parallel context comparisons.Minnen2018 and Cheng2020 are adapted by replacing only the serial context model, supporting fair comparisons.
Toy mask experiments: Closer spatial neighbours contribute more to context modeling than simply increasing the number of referenced neighbours.Checkerboard 3 × 3 outperforms serial 5 × 5 and all-neighbour 3 × 3 masks, while masks with more adjacent neighbours perform better.
Model comparison: The proposed checkerboard context replaces serial context calculation with a two-pass parallel design while preserving the surrounding autoregressive architecture.Figure 6 illustrates the replacement, and the experiments retain other architecture components for comparison.
Rate-distortion performance: BDBR −17.0%/−27.4% is achieved versus Ballé2018 for Minnen2018/Cheng2020 with checkerboard context, despite only slight RD reductions versus the original models.The authors describe the quality loss as acceptable because the parallel model removes the computational-efficiency limitation.
Additional evaluations: The checkerboard model reaches the same overall conclusion on Tecnick and remains effective when models are optimized for MS-SSIM.On Tecnick, Minnen2018 with checkerboard context slightly exceeds its original baseline at higher bit rates.

6. Discussion

The checkerboard context model and two-pass decoding make spatial-context compression highly parallel while retaining competitive compression performance without changing model structure or capacity.

Discussion: The proposed approach is an almost drop-in replacement for serial context models that improves decoding efficiency without changing model structure or capacity.The paper presents it as a way to improve the practical potential of learned compression methods using spatial context.

1. Detailed Running Speed

Detailed timing identifies serial context computation as the main decoding bottleneck and reports large speedups from the parallel checkerboard model.

Timing breakdown: More than 95% of serial decoding time is spent calculating context features and entropy parameters.The bottleneck occurs across the evaluated serial models.
Minnen2018: 44.6× and 52.1× speedups are achieved for Minnen2018 on Kodak and Tecnick, respectively.The Kodak images are smaller and the Tecnick images larger in the reported evaluation.
Cheng2020: Cheng2020 decoding becomes 18.5× faster on Kodak and 20.4× faster on Tecnick with parallel context.The results come from the authors’ PyTorch implementation.
CPU reference: The parallel context model also performs well on CPUs when using matrix libraries such as Intel MKL.CPU measurements use an Intel Xeon E5-2620, although the primary target is parallel-device efficiency.

2. More RD Curves

Additional RD evaluations on Tecnick and under MS-SSIM optimization support the checkerboard model as an efficient replacement for serial spatial context.

Tecnick RD curves: On Tecnick, checkerboard Cheng2020 is slightly worse than Minnen2020 at low bit rates but better at high bit rates.The paper attributes the difference partly to GMM density estimation with fewer coding channels.
Tecnick RD curves: On Tecnick, Minnen2018 with checkerboard context slightly outperforms the original Minnen2018 baseline at higher bit rates.The result is presented as evidence that checkerboard convolution captures useful spatial causal relationships.
MS-SSIM evaluation: MS-SSIM optimization is evaluated on Kodak using separate λ settings for low- and high-rate models.The reported configurations use λ values {1, 3, 16} at low rates and {40, 120, 360} at high rates.

3. Cheap Operations for Parallel En/De-coding

The proposed checkerboard model uses lightweight tensor operations to divide, decode, and recombine anchor and non-anchor latents for parallel en/de-coding. Its implementation relies on slicing, reshaping, permuting, concatenating, and demultiplexing operations, with runtime evaluated on GPU and CPU.

Multiplexer: The multiplexer mixes two same-sized feature maps to generate the anchor and non-anchor representations needed for one-pass encoding.The inputs have size H × W × M, and the operation is arranged to avoid slowing encoding.
Multiplexer: Slicing and assignment with a step of 2 provides a direct way to interleave anchor and non-anchor features in supported frameworks.The implementation clones one input and assigns the other input to alternating checkerboard locations.
Multiplexer: Frameworks without slice-assignment can construct the same mixture by flattening 2 × 2 patches, splitting feature maps into chunks, and concatenating selected chunks.This shape-transforming implementation uses space-to-depth or equivalent permuting and reshaping operations.
Runtime and quality evaluation: The evaluation reports per-process decoding time on GPU in microseconds and CPU in seconds, with GPU speed expressed in million-pixels-per-second.The figures compare reconstructed Kodak images and PSNR or MS-SSIM evaluations, distinguishing serial-context models with dashed lines.
Demultiplexer: During encoding, latents are separated into anchor and non-anchor chunks before being encoded independently, enabling decoding to read anchors first.The demultiplexer is the inverse of the slice-and-concatenate multiplexer and recovers the two latent chunks from the bitstream.
Two-pass decoding: Decoded chunks are inserted into feature maps with filler values, reshaped into full-sized anchor or non-anchor tensors, and merged after both passes finish.The implementation uses zero tensors as fillers and combines the chunks with the multiplexer or, with zero fillers, an addition operation.