Learned Image Compression with Dual-Branch Encoder and Conditional Information Coding

Image compression is a fundamental and crucial topic in the field of signal processing. Over the past few decades, several classical standards have emerged, including JPEG [1], JPEG2000 [2], BPG [3], and VVC, which generally follow the same transform coding paradigm: linear transformation, quantization, and entropy encoding.

Recent learned image compression methods [4, 5, 6] have outperformed the current best classical image and video encoding standard VVC in terms of both peak signal-to-noise ratio (PSNR) and multi-scale structural similarity (MS-SSIM). This indicates that learned image compression methods hold tremendous potential for the next generation of image compression technologies.

Most learning-based image compression methods are Convolutional Neural Networks-based (CNN-based) approaches [7, 5, 8] that use the variational autoencoder (VAE) [9]. However, with the recent development of vision Transformers [10], several transformer-based learning methods [11, 12, 13] have been introduced. For instance, in CNN-based approaches, a residual block-based image compression model is proposed to achieve performance comparable to VVC in terms of PSNR. Meanwhile, in the realm of transformer-based methods, a swin-transformer-based image compression model has been proposed to enhance rate-distortion performance. Both CNN-based and transformer-based approaches offer distinct advantages: CNNs excel in local modeling with lower complexities, whereas transformers are adept at capturing non-local information. Nevertheless, the complexity of swin-transformer-based methods outperforms that of the CNN schemes.

The design of entropy models is also a critical aspect of learned image compression. A common approach introduces additional latent variables as hyper-priors, thereby transforming the compact encoded symbol probability model into a joint model [9]. Based on this, several methods [14, 7, 4] have been developed. For instance, masked convolution [14] is proposed to capture context information. More accurate probabilistic models, such as GMM [7] and GLLMM [4], have been introduced to enhance compression performance. Furthermore, a parallel channel autoregressive entropy model has been proposed [15], wherein the latent representation is divided into 10 slices. The encoded slices can aid the encoding of subsequent slices by providing side information in a pipelined manner.

Recently, some attention modules [16, 7] have been proposed to enhance image compression. Attention modules can be incorporated into the image compression framework to assist the model in focusing more on detailed information. However, many schemes are time-consuming or can only capture local information [16]. To reduce the complexity of the attention module, a simplified attention model is placed in the main encoder and decoder to enhance image compression. We also introduce an attention module [7] to improve the rate-distortion performance.

The contributions of this paper can be summarized as follows:

First, we employ two branches of encoders to learn latent representations at different resolutions of the input. This strategy allows us to better capture both global and local information from the input image. Notably, each branch operates independently without any information exchange.

Second, we utilize the high-resolution latent representation as conditional information to assist in the encoding and decoding of the low-resolution latent representation. This strategy helps to remove spatial redundancy in the low-resolution latent representation, enhancing the overall performance of the framework. We do not utilize any serial entropy models. Instead, we employ a parallel channel-wise auto-regressive entropy model [15, 13] for encoding and decoding low-resolution and high-resolution latent representations.

Third, extensive experiments demonstrate that our method achieves state-of-the-art performance when compared to recent learning-based methods and traditional image codecs on the Kodak dataset. Our method is approximately twice as fast in both encoding and decoding compared to the parallelizable checkerboard context model [6], and it also achieves a 1.2% improvement in R-D performance compared to state-of-the-art learned image compression schemes. Our method also outperforms the latest classical image codec in H.266/VVC-Intra (4:4:4) and other leading learned schemes such as [6] in both PSNR and MS-SSIM metrics. The decoding time and BD-Rate comparisons with VVC for various methods are presented in Fig. 1.

In this section, we describe the whole framework of the proposed method. Subsequently, we will detail its major components, including the dual-branch encoding network, conditional coding of low-resolution latent representations, and the associated training methodology.

The proposed framework is illustrated in Fig. 2. The input image, represented by $x$, has dimensions $W\times H\times 3$, where $W$ and $H$ are its width and height, respectively. The framework primarily consists of the core networks ($g_{a1}$, $g_{a2}$ and $g_{s}$) and the hyper networks ($h_{a}$ and $h_{s}$).

The two core encoder networks, labeled as $g_{a1}$ and $g_{a2}$, are tasked with learning two compact latent representations $y_{1}$ and $y_{2}$ from the input image. The architectures of $g_{a1}$ and $g_{a2}$ both integrate two simplified attention modules, three residual group blocks (highlighted in cyan in Fig. 2), and four stages of pooling operations. The residual group blocks are composed of four basic residual blocks [17] connected in series.

We employ two branches to capture different resolutions of the original image. The first branch learns the high-resolution latent representation $y_{1}$, utilizing a $3\times 3$ convolution as its downsampling module. In contrast, the second branch captures the low-resolution latent representation $y_{2}$ of the input image, leveraging a $1\times 1$ convolution for downsampling.

To enable parallel entropy decoding of the quantized latents $y_{1}$ and $y_{2}$, we do not use any serial context model. As in [15, 13], we use channel-wise entropy model to encode and decode $y_{1}$ and $y_{2}$ in parallel. However, it’s noted that we first encode and decode $y_{1}$ in parallel, and then use $y_{1}$ as conditional side information to encode and decode $y_{2}$ in parallel. We will provide a detailed explanation of this encoding and decoding process in Sec. 2.2 and illustrate it in Fig. 3.

Subsequently, arithmetic coding compresses $\hat{y_{1}}$ and $\hat{y_{2}}$ into a bitstream. The decoded values, $\hat{y_{1}}$ and $\hat{y_{2}}$, are concatenated and forwarded to the primary decoder network $g_{s}$. This decoder mirrors the core encoder network $g_{a}$, with convolutions replaced by deconvolutions. While most convolution layers employ the leaky ReLU activation function, the final layer in both the hyperprior encoder and decoder operates without any activation function

In this section, we compare some recent learned image compression methods and traditional image codecs with our proposed method in terms of Peak Signal-to-Noise Ratio (PSNR) and MS-SSIM metrics. The performance of different schemes are evaluated in two datasets with different resolutions. The Kodak PhotoCD dataset ²²2http://r0k.us/graphics/kodak/ is comprised of 24 images with a resolution of $768\times 512$. Some recent learning-based schemes includes GLLMM [4], He2021 [6], Hu2021 [19], and Cheng2020 [20]. The traditional image codecs includes the latest image codec VVC-Intra (4:4:4) ³³3https://vcgit.hhi.fraunhofer.de/jvet/VVCSoftware_VTM/tree/VTM-5.2, BPG-Intra (4:4:4), JPEG2000, and JPEG.

For a fair comparison, we have implemented the method described in Cheng2020 [20], increasing its number of filters $N$ from 192 to 256 at high rates. This modification leads to enhanced performance compared to the original results in [20]. The results from He2021 [6] are sourced from the code available at [21].

In this paper, we propose two techniques aimed at improving coding performance and speeding up the decoding process. These techniques consist of a dual-branch encoder network and a conditional information module. We employ two independent encoding networks to learn the latent representations of input images at different resolutions, which facilitates the reduction of spatial redundancy in the input images. The high-resolution latent representation primarily captures the global information of the input image, while the low-resolution latent representation predominantly represents its local details. We utilize the high-resolution latent information as side information for the low-resolution latent representation. This approach aids in reducing the spatial redundancy of the low-resolution latent representation, thereby enhancing encoding efficiency. We also employ a parallel channel-wise auto-regressive entropy model [15, 13] for encoding and decoding low-resolution and high-resolution latent representations.

This work was supported by the National Natural Science Foundation of China (No. 61474093), Industrial Field Project - Key Industrial Innovation Chain (Group) of Shaanxi Province (2022ZDLGY06-02), the Natural Sciences and Engineering Research Council of Canada (RGPIN-2020-04525), Google Chrome University Research Program, NSERC Discovery Grant RGPIN-2019-04613, DGECR-2019-00120, Alliance Grant ALLRP-552042-2020; CFI John R. Evans Leaders Fund.