Source-linked AI summary

Causal Contextual Prediction for Learned Image Compression

Zongyu Guo, Zhizheng Zhang, Runsen Feng, Zhibo Chen

arXiv:2011.09704v5cs.CVeess.IV

TL;DR

Learned image compression entropy models do not fully capture global spatial or cross-channel latent correlations. This paper introduces separate entropy coding with causal context and global prediction models, plus group-separated attention, and reports state-of-the-art rate-distortion performance, including 5.1% BD rate savings against VTM 8.01 on Kodak.

  • Problem

    Existing entropy models capture local and hierarchical latent dependencies but remain suboptimal because global correlations are missed and cross-channel relationships are underexplored.

  • Method

    The paper separates latents into channel groups for causal cross-channel context prediction and global prediction, while adding group-separated attention to the transform networks.

  • Results

    5.1% BD rate savings on the Kodak dataset are achieved against VTM 8.01, with performance improvements over previous learned schemes in PSNR and MS-SSIM.

  • Takeaways & Limitations

    Causal contextual prediction improves entropy estimation by exploiting cross-channel and global latent redundancies without transmitting overhead.

Abstract

from arXiv · show

Over the past several years, we have witnessed impressive progress in the field of learned image compression. Recent learned image codecs are commonly based on autoencoders, that first encode an image into low-dimensional latent representations and then decode them for reconstruction purposes. To capture spatial dependencies in the latent space, prior works exploit hyperprior and spatial context model to build an entropy model, which estimates the bit-rate for end-to-end rate-distortion optimization. However, such an entropy model is suboptimal from two aspects: (1) It fails to capture spatially global correlations among the latents. (2) Cross-channel relationships of the latents are still underexplored. In this paper, we propose the concept of separate entropy coding to leverage a serial decoding process for causal contextual entropy prediction in the latent space. A causal context model is proposed that separates the latents across channels and makes use of cross-channel relationships to generate highly informative contexts. Furthermore, we propose a causal global prediction model, which is able to find global reference points for accurate predictions of unknown points. Both these two models facilitate entropy estimation without the transmission of overhead. In addition, we further adopt a new separate attention module to build more powerful transform networks. Experimental results demonstrate that our full image compression model outperforms standard VVC/H.266 codec on Kodak dataset in terms of both PSNR and MS-SSIM, yielding the state-of-the-art rate-distortion performance.

I. INTRODUCTION

Learned image compression entropy models remain limited by incomplete modeling of global spatial and cross-channel latent correlations. The paper addresses these gaps with separate entropy coding, causal contextual and global prediction models, and group-separated attention, reporting improved rate-distortion performance.

  • Motivation: Learned codecs transform images into quantized latent representations and use entropy models to estimate their discrete entropy for lossless compression.Hyperpriors and autoregressive context models improve this estimation by modeling hierarchical and local latent dependencies.
  • Motivation: Adjacent-context and hierarchical-prior entropy models miss global latent correlations, while cross-channel relationships remain underexplored.The paper identifies both limitations as targets for more accurate entropy estimation.
  • Proposed approach: Separate entropy coding divides latents into two channel groups so decoded channels provide cross-channel information for causal context prediction.The causal context model aggregates spatially adjacent latents and first-half channel elements after the first group is decoded.
  • Proposed approach: The causal global prediction model computes causal global correlations from decoded channels to guide predictions for undecoded channels without transmitting overhead.A correlation matrix from the first channel group describes global dependencies for the second group.
  • Proposed approach: A group-separated attention module independently processes feature-map attention across separated groups to strengthen the transform networks.The paper presents this as a side contribution alongside the causal entropy-modeling components.
  • Results: 5.1% BD rate savings on the Kodak dataset are achieved against traditional codec VTM 8.01, while the method outperforms previous learned schemes in PSNR and MS-SSIM.The full network integrates the causal context model, causal global prediction model, and group-separated attention module; ablation studies assess each component.

II. RELATED WORK

Traditional image compression standards rely on manually designed transform-coding modules and increasingly complex intra-prediction methods. Intra block compensation extends this paradigm with global reference-block searching but requires displacement vectors.

  • Traditional codecs: JPEG, JPEG2000, and BPG are traditional standards that rely heavily on manually designed modules while following transform coding.Intra prediction became an important component of image compression after H.264/AVC standardization.
  • Traditional codecs: H.264 uses 9 intra-prediction modes, whereas H.265/HEVC adopts 35 optional intra-prediction modes.These standards illustrate the increasing complexity of manually designed intra-prediction methods.
  • Traditional codecs: Intra block compensation uses a bit-consuming displacement vector to represent the relative displacement between the current prediction unit and a reference block.The technique was standardized in HEVC screen-content extensions and supports global searching for efficient prediction.

B. Learned Image Compression

Learned image compression uses transforms, quantization, entropy coding, and probabilistic priors to optimize rate-distortion performance. Existing context models capture local structure but leave global and cross-channel latent relationships insufficiently modeled.

  • Limitations: Prior context models cover local regions but do not attend to global latent correlations, while cross-channel relationships remain an additional modeling concern.
  • Framework: Learned codecs transform images into quantized latent representations, then losslessly entropy-code them and reconstruct images through a synthesis transform.
  • Framework: The rate-distortion objective balances estimated coding rate against reconstruction distortion through a weighted loss.The rate term is the cross entropy under the learned entropy model, while distortion may use PSNR or MS-SSIM.
  • Entropy modeling: Hyperprior models transmit side information describing latent spatial structure and improve estimation of the latent probability distribution.
  • Entropy modeling: Serial mask-convolution context models use previously decoded neighboring latents, primarily the points to the left and above the current point.

B. Global Prediction

Global reference information can improve entropy estimation, but serial decoding requires a causal way to establish those references without impractical side-information costs. The paper therefore motivates a global prediction model based on already decoded latents.

  • Global prediction: Global prediction uses all previously decoded latents to predict the current decoding point.
  • Global prediction: Causal correlation matrices compare previously decoded points with the current point to search for global reference points available during serial decoding.The paper measures similarity with cosine distance and notes that the full correlation matrix is unavailable directly at the decoder.
  • Global prediction: Table I contrasts an overhead-free evaluation with a practical case that transmits reference-point coordinates in advance.The caption states that the overhead-free case is not a practical codec, whereas the practical case assigns additional bits to coordinate transmission.
  • Global prediction: Experiments demonstrate that accurate global context benefits entropy estimation, while explicitly transmitting global reference information is bit-consuming.

IV. PROPOSED METHOD

The proposed separate entropy coding model exploits serial decoding by splitting each latent location across channels. Its improved mask convolution uses decoded channel groups to provide cross-channel context for entropy estimation.

  • A. Causal Context Model: Separate entropy coding divides each latent location into two channel groups decoded in sequence.
  • A. Causal Context Model: The first channel group is predicted with a conventional mask convolution and hyperpriors before the second group is decoded.
  • A. Causal Context Model: Mask Conv + aggregates adjacent decoded latents and the first channel group at the current location to generate richer context for the second group.
  • A. Causal Context Model: The causal context model captures cross-channel relationships, while separate decoding stages use independent parameters and non-shared entropy-parameter modules.

B. Causal Global Prediction Model

The causal global prediction model extends local context modeling with long-range latent correlations. It separates latent channels so the first group can provide causal global references for predicting the second group.

  • B. Causal Global Prediction Model: The causal context model captures local correlations but ignores global correlations, motivating the causal global prediction model.
  • B. Causal Global Prediction Model: The first latent group remains encoded using only hyperpriors and local context.
  • B. Causal Global Prediction Model: The second latent group is estimated using both improved adjacent context and global context.The adjacent context comes from the causal context model, while the global context comes from the causal global prediction model.
  • B. Causal Global Prediction Model: The model uses latent separation to establish global reference information at the decoder without transmitting reference coordinates.

C. Global Prediction via Latent Separation

Latent separation enables causal global prediction without transmitting explicit global references. The decoder uses already decoded latent groups to identify correlated points and generate predictions.

  • Global reference construction: The first latent group provides bit-free global correlation information for predicting the second group.Similarities among half latent vectors approximate similarities among entire latent vectors, enabling decoder-side reference construction.
  • Causal global prediction: After decoding the first group, the model computes masked pairwise distances and retains the top-k correlated decoded points.The selected indices gather reference features for each undecoded point.
  • Separate entropy coding: Separate entropy coding divides latents into two channel groups and supports serial decoding for causal prediction.The first group is inferred directly, while the second is inferred conditionally on the prior group.

D. Analysis

The analysis examines how channel separation supports global prediction and how the number of groups affects complexity. It also distinguishes the proposed global context from local context and hyperpriors.

  • Channel separation: When channel groups are equal in size, the first group accounts for around 62% of the bitstream across bitrates.This suggests the first group contains representative, bit-consuming channels for modeling global correlations.
  • Channel separation: Using more than two latent groups would further increase both space and time complexity during decoding.The existing two-group design already divides decoding into two stages and uses separate parameter-estimation weights.
  • Context comparison: Unlike local context, causal global prediction searches all decoded points and selects top-k correlated references for the current point.Local context aggregates available neighboring information through masked convolution without explicit point prediction.
  • Context comparison: Unlike hyperpriors, the causal global prediction model uses accessible decoded points without transmitting side-information overhead.Hyperpriors transmit additional information to represent rough global structures.

4) The influence of the value of k:

The study evaluates the number of reference points used by causal global prediction and relates this choice to residual cross-channel redundancy and the model’s attention design.

  • The influence of the value of k: k = 4 achieves the best performance, slightly surpassing dense prediction in the comparison of reference-point counts.Dense prediction uses all decoded points, whereas top-2, top-4, and top-6 use selected relevant points.
  • The influence of the value of k: Some global references are inaccurate because the correlation matrix is approximated using half the latent channels.This helps explain why using every decoded point does not outperform the selected-reference setting.
  • The influence of the value of k: Cross-channel redundancy remains after the nonlinear transform and is jointly exploited with spatial correlation.The proposed entropy-model components target both cross-channel and global redundancies.
  • The influence of the value of k: The group-separated attention module enables separate feature-map attention across channel groups.It is empirically presented as more powerful than the attention mechanism used in a previous compression method.

F. Network Architecture

The compression network combines the proposed entropy-model components with a separate attention module and residual parameter-estimation blocks. Its architecture and training use bitrate-dependent capacity settings and staged optimization.

  • Network architecture: The main network replaces the original attention module with separate attention and integrates causal context and causal global prediction models.Residual blocks also enhance the parameter-estimation module, while the architecture otherwise follows a prior compression design.
  • Entropy modeling: The entropy model estimates mixture parameters for a continuous density, which is convolved with unit uniform noise to model discrete quantized latents.The model uses I = 3 estimated components following prior work.
  • Training: Hard-quantized latents are sent into the synthesis decoder during training, alongside additive uniform noise for nondifferentiable quantization.The paper reports that this training choice achieves better performance.
  • Training: The model is trained on 256×256 ImageNet crops with minibatches of 8 through three optimization stages.The first stage trains the main compression network for 1.5 million iterations before later optimization steps.
  • Network architecture: For bitrates below 0.50 bpp, the network uses N=192 and M=192; at higher bitrates, M increases to 320.N remains 192 in both bitrate regimes.

B. Performance

On Kodak, the proposed method achieves state-of-the-art rate-distortion performance, while screen-content results favor it at low bitrate but VVC at higher reconstruction quality. Ablation studies attribute improvements to the causal context model, causal global prediction model, and separate attention module.

  • 5.1% BD rate savings against VTM 8.0 are achieved on Kodak across 0.15, 0.4, 0.7 and 1.0 bpp.
  • The method outperforms VTM 8.0 at all bitrates in PSNR and exceeds the previous state-of-the-art method at most bitrates in MS-SSIM.
  • Screen-captured images show better performance than VVC at low bitrate, but VVC performs better as bitrate and PSNR increase.The paper attributes the high-quality regime gap to the autoencoder limit of the non-invertible neural transform.
  • Ablation Study: Causal adjacent context and causal global prediction each improve rate-distortion performance relative to the main compression network.
  • Ablation Study: Replacing the proposed separate attention module with the attention layer used in causes a performance drop.

D. Coding time

Separate entropy coding improves rate-distortion performance but increases decoding time because serial channel decoding and global reference searching are computationally costly.

  • The three evaluated methods have similar encoding time because encoding can be computed in parallel.
  • Decoding time increases from average 6.7 seconds to average 7.9 seconds when serial decoding separates the process across channels.The measurements cover all 24 Kodak images at 768×512 resolution.
  • 37.7 seconds is the average decoding time for Causal Context + Causal Global Prediction on the Kodak dataset.Global searching of reference points consumes much more time, and the paper notes room for software and hardware optimization.
  • The separate entropy model suffers from increased decoding time while significantly improving rate-distortion performance.
Loading 2011.09704v5…