Causal Context Adjustment Loss for Learned Image Compression

TL;DR

This paper makes the first attempt in investigating the way to explicitly adjust the causal context with the proposed Causal Context Adjustment loss (CCA-loss), and enables the neural network to spontaneously adjust important information into the early stage of the autoregressive entropy model.

Abstract

In recent years, learned image compression (LIC) technologies have surpassed conventional methods notably in terms of rate-distortion (RD) performance. Most present learned techniques are VAE-based with an autoregressive entropy model, which obviously promotes the RD performance by utilizing the decoded causal context. However, extant methods are highly dependent on the fixed hand-crafted causal context. The question of how to guide the auto-encoder to generate a more effective causal context benefit for the autoregressive entropy models is worth exploring. In this paper, we make the first attempt in investigating the way to explicitly adjust the causal context with our proposed Causal Context Adjustment loss (CCA-loss). By imposing the CCA-loss, we enable the neural network to spontaneously adjust important information into the early stage of the autoregressive entropy model. Furthermore, as transformer technology develops remarkably, variants of which have been adopted by many state-of-the-art (SOTA) LIC techniques. The existing computing devices have not adapted the calculation of the attention mechanism well, which leads to a burden on computation quantity and inference latency. To overcome it, we establish a convolutional neural network (CNN) image compression model and adopt the unevenly channel-wise grouped strategy for high efficiency. Ultimately, the proposed CNN-based LIC network trained with our Causal Context Adjustment loss attains a great trade-off between inference latency and rate-distortion performance.

Figures (9)

• Figure 1: Left: A systematic overview of our method. We adopt the VAE-based framework balle2017endtoend with hyperprior balle2018variational and channel-wise autoregressive entropy model minnen2018joint; besides the original Rate-Distortion loss ($\mathcal{L}_R$, $\mathcal{L}_D$), we introduce an auxiliary entropy model and propose the causal context adjustment loss ($\mathcal{L}_{CCA}$) for better training the entropy model. Right: An illustration of the entropy model and the auxiliary entropy model. The auxiliary entropy model does not use the information to be encoded to predict the following representations, our $\mathcal{L}_{CCA}$ encourage the predicting gap between the two models, so as to enhance the importance of causal context in early stages.
• Figure 2: The comparison of averaged information distributed ratios of various models in Table \ref{['ablation1']}.
• Figure 3: Rate-Distortion performance evaluation of PSNR on Kodak dataset (left), CLIC Professional Validation dataset (middle), Tecnick dataset (right), respectively.
• Figure 4: Visualization of the reconstructed images (top: kodim19, middle: kodim10, bottom: kodim4) from Kodak dataset. The titles under the sub-figures are represented as [bpp | PSNR(dB)].
• Figure 5: Architecture of NAF-block, Entropy Model and Latent Residual Prediction (LRP).