Learned Image Compression with Gaussian-Laplacian-Logistic Mixture Model and Concatenated Residual Modules

Table I summarizes some representative encoder/decoder architectures for learned image coding.

One of the first learned image compression schemes was proposed in [5], which was based on the long short-term memory (LSTM) recurrent neural network (RNN) and was used to compress thumbnail images. In [6], the scheme in [5] is generalized to full-resolution images.

In [7, 26], an encoder network that includes convolution, downsampling, and the generalized divisive normalization (GDN) is proposed. It was the first learned image compression scheme that achieved better performance than JPEG2000 in terms of both PSNR and MS-SSIM.

Most recent schemes are based on the autoencoder framework, whose latent representations have much lower dimension than the input, and is very suitable for data compression.

A powerful building block of many cutting-edge neural networks is the residual block first proposed in the ResNet [27], which uses shortcut connections to facilitate the design and training of deep networks, and can effectively improve the performances of many computer vision tasks. As a result, the residual block has also been used in many learning-based image compression schemes.

In [8], the residual block concept in [27] was used in the autoencoder architecture, which also achieved comparable performance to JPEG2000.

Generative adversarial network (GAN) is another powerful framework for many applications, and has also been used in several learned image coding schemes [9, 28, 29].

In [10], the importance map is introduced to achieve content-adaptive bit allocation in different regions, but the importance mask needs to be sent to the decoder. In [11], non-local attention module is used to capture long-range correlation, and generate attention mask without sending side information, and it is used in both core network and hyper network. In [12], the attention module in [11] is simplified.

In [30, 13], a multi-resolution network architecture based on the octave convolution in [31] is developed, similar to the wavelet transform. In [14], another wavelet-like scheme is proposed based on the lifting scheme, which supports both lossy and lossless compression.

In [15], the vision transformer framework is introduced and combined with convolutional layers to boost the image compression performance.

Table II is an overview of typical entropy coding models used for learned image coding.

In earlier schemes [32, 8, 26], the quantized latent representations or latents are assumed to be independent and identically distributed and followed a simple marginal distribution. Once the parameters of the distribution are trained, the probability of all latents is fixed for all images, which is used by the entropy encoding and decoding of the quantized latents. Since a fixed entropy model is used, the performances of these methods are compromised.

To improve the entropy coding performance, recent methods consider the correlation of neighboring pixels or latents (contexts), and achieve image-adaptive entropy coding. To estimate the joint probability with manageable complexity, the chain rule is usually used, which factorizes the joint probability into products of conditional probabilities.

In PixelCNN [16], masked convolutional networks are used to predict the distribution of each pixel value, conditioned on the causal neighbors in an autoregressive manner, which is used for image generation and image decoder in an autoencoder. In [17], the PixelCNN++ is proposed, which uses the discretized logistic mixture distribution to reduce the complexity of [16]. In [23], PixelCNN++ is applied to lossless image compression.

To help the estimation of the conditional probabilities of the latents, the hyperprior network is introduced in [19], where a hyper encoder network is used to extract some hyperpriors from the latents, which are coded as side information and sent to a hyper decoder network. The latter uses the reconstructed hyperpriors to estimate the conditional probability of the latents, thereby making the entropy model image-dependent and spatially adaptive.

In [19], the conditional probability of each latent is assumed to follow a zero-mean Gaussian scale mixture (GSM) model, which achieves better performance than BPG (4:4:4). In [20], the non-zero-mean Gaussian mixture model (GMM) is further proposed. In addition, the autoregressive context model in [16] is introduced to estimate the conditional probability of each latent from both the hyperpriors and its spatial context.

Other models have also been proposed to encode the latents. For example, in [22], the Laplacian distribution was used. In [12] and [24], the GMM model is adopted. A quality enhancement subnetwork is also introduced in [24] to improve the performance. It is the first learned image compression approach that achieves better performance than H.266/VVC (4:4:4) in both PSNR and MS-SSIM. As far as we know, the results in [24] are currently the best in the literature.

The entropy coding approach above predicts the conditional probability of a latent given some causal neighbors and the hyperpriors, according to the chain rule, which approximates the joint probability implicitly. Another possible approach is to use the more general multivariate models to represent the joint probability explicitly.

In [25], the multivariate GMM is used in learned image compression to capture the inter-channel correlation of co-located latents across feature maps (this is captured implicitly by the approach above using 1x1 filters in the hyper decoder). The means and covariance matrices are estimated in a cascaded approach. Vector quantization is used to encode the latents, as in [32]. The hyper coding network is not used. However, since no spatial context is used, the R-D performance of [25] is only comparable to VVC on the Kodak test set, and not as good as other leading methods such as [24]. This indicates that if the learned image compression framework can consider the joint distribution via context model, non-local operator, and attention module, it is sufficient to use univariate models to model the remaining redundancy in the latents.

On the other hand, the overall complexity of multivariate model-based methods is usually quite high, as noted in [25], because the number of parameters of multivariate models increases quadratically with the size of the covariance matrix. Therefore, the complexity of the network model and the complexity of calculating multivariate probability increase dramatically. The training of the network is also more difficult, because the reverse gradients of different variables affect each other. Although the method in [25] can be parallelized because no spatial context is used, it still has very high demands for GPU and memory. In fact, other methods using spatial context can also be parallelized [33, 34, 35].

Multivariate models have also been used in other applications. Most of them are based on the Gaussian distribution. In [36], a deep GMM is proposed for classification by using the factor-analytic representation of the GMM in each layer. The Expectation-Maximization (EM) algorithm is used to estimate the parameters. However, it is found in [37] that sometimes it is quite challenging for this method to infer its parameters, even for a small-scale problem with 3 layers and 76000 latents.

In [38], a novel face recognition framework based on the multivariate GMM is proposed to make feature embeddings extract more identity-relevant information. In [39], a learned multiple graph Gaussian embedding model is developed to learn highly informative network features by mapping high-dimensional networks into a low-dimensional latent space. In [40], the multivariate skewed t-distribution is proposed for hyperspectral anomaly detection. In [41], the last fully-connected (FC) layer of a deep network is combined with a mixture of GMM (MoGMM) for image recognition. In [42], the multivariate GMM is considered as a prior to recover degraded images. In [43], a framework for deep unconstrained face verification is proposed to map learned discriminative facial features to a regularized metric space, in which matching and non-matching pairs follow multivariate Gaussian distributions.

Although the learned methods have made significant progresses, the compression performance can be further improved. First, the previous entropy models only use a single distribution, which is not optimal for every latent. Second, the latents still exist some spatial redundancy. In order to address these issues, we make the following contributions in this paper:

$\bullet\,$ Instead of using a single probability model, we propose a discretized Gaussian-Laplacian-Logistic mixture model (GLLMM), which is more flexible and efficient in estimating the conditional probabilities of the latents, given the same complexity. In fact, the previous models are special cases of the proposed model. Ablation studies demonstrate that the proposed joint model outperforms all previously proposed single models. For example, it achieves 0.1-0.15 dB gain in PSNR compared to the GMM model for the Kodak dataset.

$\bullet\,$ We improve the basic building block of the encoder network by developing a concatenated residual module (CRM), with additional shortcut connections. The CRM improves the information flow, reduces the correlation of the output, and improves the training of the network. Ablation studies using the Kodak dataset show that the proposed CRM can achieve 0.2-0.3 dB gain in terms of PSNR compared to the original residual blocks. We also conduct ablation studies to compared the performances of CRMS with two and three stages of residual blocks.

$\bullet\,$ One of the key contributions of [24] is the post-processing component, which achieves 0.5 dB gain in [24] compared to its baseline. We also investigate the performance of this component in our scheme. However, our experiments show that there is no need to apply the post-processing in our method, because the proposed CRM and entropy model have done a good job. This also reduces the complexity of the scheme.

Experimental results using the Kodak[44], Tecnick-100, and Tecnick-40 datasets in [45] show that the proposed scheme outperforms all the previous learning-based methods including [24, 12] and traditional codecs including VVC in both PSNR and MS-SSIM. For example, in terms of PSNR, for the Kodak dataset, when the bit rate is higher than 0.4 bpp, our method is 0.2-0.3 dB higher than [24, 12], and 0.3-0.4 dB higher than VVC (4:4:4). For the Tecnick dataset, when the bit rate is higher than 0.2 bpp, our method is 0.3-0.4 dB better than VVC (4:4:4). This represents the new state of the art in learned image compression.

The remainder of the paper is organized as follows. In Section II, we propose our image compression framework with discretized Gaussian-Laplacian-Logistic mixture model and the concatenated residual modules. In Sec. III, we compare our method with some state-of-the-art learning-based methods and classical image compression methods. Ablation experiments are carried out to investigate the performance gain of the the proposed scheme. The conclusions are reported in Sec. IV.

The learned image coding framework adopted in this paper is illustrated in Fig. 2, which is mainly based on [12]. It is comprised of two main parts: the core autoencoder and the hyperprior coding.

The size of the input color image $\bm{x}$ is $W\times H\times 3$, where $W$ and $H$ represents the width and height of the image, respectively. In the training of the network, the pixel values are normalized to $[-1,1]$ by ($\frac{\bm{x}}{127.5}-1.0$).

The core autoencoder includes an encoder ($g_{a}$) and a decoder ($g_{s}$). The input image $\bm{x}$ is first sent to the encoder network $g_{a}$, which aims to reduce the redundancy and learn a compact latent representation $\bm{y}$ of the input image. The latent is then quantized and entropy coded. The quantized latent is denoted by $\hat{\bm{y}}$.

The core encoder network includes various convolution layers and four stages of pooling operators to get the latents. The residual blocks with shortcut connections are used extensively. The GDN operator is used when the size is changed.

To improve the R-D performance, the simplified attention module in [12] is applied at two resolutions to capture long-range correlation, and strengthen more important areas, so that they will get more bit allocation.

To estimate the distribution of the latents and improve entropy coding efficiency, the hyperprior coding part is introduced in [19], which consists of a hyper encoder ($h_{a}$) and a hyper decoder ($h_{s}$). The hyper encoder extracts the hyperprior $\bm{z}$ from the latents, which is quantized into $\hat{\bm{z}}$ and entropy coded as a side information to the hyper decoder. The hyper decoder first recovers $\hat{\bm{z}}$ via entropy decoding, and then uses hyper decoder network $h_{s}$ to estimate the parameters of the conditional distribution of $\hat{\bm{y}}$, which is used by the entropy encoding and decoding of $\hat{\bm{y}}$.

Since the input dimension of the hyper network is much smaller than the original image, the hyper network is much simpler than the core network. Leaky ReLU is utilized in most convolution layers, except for the last layer in hyper encoder and decoder, which do not have any activation function.

To further improve the entropy coding efficiency, the context model network $c_{m}$ in [20] and [12] is also used in our system, which uses masked convolutions to capture the correlation of neighboring latents. The output layer of the context model is concatenated with the output of the first part of hyper decoder, and then further processed by some additional convolutional layers to estimate the parameters of the conditional distributions of the latents, which are then used to entropy encode and decode $\hat{\bm{y}}$.

Since autoregressive context model is used, the symbols have to be decoded in serial manner. After all the symbols of $\hat{\bm{y}}$ are decoded, they are sent to the core decoder ($g_{s}$) to generate the reconstructed image. Note that the context model can be sped up using the methods in [33, 34, 35].

In Fig. 2, the encoders and decoders in the core and hyper networks are symmetric, except that they use convolution and deconvolution operators, respectively.

When the size does not change, to further remove the spatial correlation in the latent representation, we develop two deeper residual blocks in this paper, which is illustrated in Fig. 3. The basic building block is the standard residual block developed in the ResNet [27], as described in Fig. 3. As in [12], we first employ the leaky ReLU activation function to replace the ReLU function and remove the batch normalization layer from the residual block. The leaky ReLU activation function can speed up the convergence of the network. The detailed structure is shown in Fig. 3, which is used in [12]. Based on this, we develop two concatenated residual modules, as shown in Fig. 3 and Fig. 3.

In Fig. 3, two residual blocks in Fig. 3 are concatenated, and another short connection is added between the input and the output. In Fig. 3, three residual blocks are concatenated, with an additional shortcut connection as well.

Compared to the standard residual block in Fig. 3, the concatenated modules have larger receptive fields. They can remove more spatial correlation, which can also help the attention module in the network. In Sec. III, we will compare the performances of the two types of concatenated modules.

In this paper, we focus on optimizing the learned image compression to achieve the best rate-distortion (R-D) performance, based on the information theory. Let $R$ be the expected length of the bitstream, and $D$ be the reconstruction error between the source image and reconstructed image. The tradeoff between the rate and distortion is adjusted by a Lagrange multiplier denoted by $\lambda$. The objective cost function is then defined as follows:

where the distortion $D(\bm{x},\hat{\bm{x}})$ is the reconstruction error between origin image $\bm{x}$ and the decompressed image $\hat{\bm{x}}$. $H(\hat{\bm{y}})$ and $H(\hat{\bm{z}})$ are the entropies of the latents and hyperpriors, based on the estimated conditional probabilities, as a measure of the bits needed to encode them.

The mean square error (MSE) and MS-SSIM are the most widely used distortion metrics, which are chosen in this paper. Other terms such as the GAN cost can be added to the loss function to achieve other goals. However, when more terms are introduced, the optimized result has to trade off these different constraints.

The CLIC dataset [52] and LIU4K dataset [53] are employed to train our models. The images of the training dataset are rescaled to a resolution of $2000\times 2000$, which is better for training. Data augmentation algorithms (i.e., rotation and scaling) are used to randomly crop 81,650 patches with a resolution of $384\times 384$. The patches are stored as lossless PNG images.

We optimize the proposed models using mean square error (MSE) and MS-SSIM respectively. When optimized for MSE metric, the parameter $\lambda$ is selected from the set $\{0.0016,0.0032,0.0075,0.015,0.023,0.03,0.045\}$. Each value trains a network for a particular bit rate. The number of filters $N$ is set to 128 for the three lower bit rates, and is set to 256 for the four higher bit rates. When the MS-SSIM metric is targeted, the parameter $\lambda$ is in the set $\{12,40,80,120\}$. The number of filters $N$ is set to 128 for the two lower bit rate, and is set to 256 for the two higher bit rates. Each model was trained up to $1.5\times 10^{6}$ iterations to obtain stable performance. The Adam solver with a batch size of 8 is adopted. The learning rate is set to $1\times 10^{-4}$ in the first 750,000 iterations, and we gradually reduce the learning rate by 0.5 after every 100,000 iterations in the last 750,000 iterations.

The average MS-SSIM and PSNR performances over the 24 Kodak images are illustrated in Fig. 5. Note that in Fig. 5 to Fig. 7, the notations Scheme (PSNR) and Scheme (MS-SSIM) in the legends mean that the model in the scheme is optimized for PSNR and MS-SSIM respectively.

In Fig. 5, when optimized for PSNR, Lee2020 (PSNR) [24] is the best among previous methods, which obtains even better performance than VVC (4:4:4) at high rates. The next closest method to VVC (4:4:4) is Cheng2020 [12]. When the bit rate is less than 0.3 bpp, our method has similar performance to [24] and VVC (4:4:4). When the bit rate is higher than 0.4 bpp, our method achieves the best performance, which is 0.2-0.3 dB over [24] and 0.3-0.4 dB over VVC (4:4:4).

In the MS-SSIM results in Fig. 5, Lee2020 (MS-SSIM) [24] also achieves better performance than previous learning-based methods and all the traditional image codecs including VVC (4:4:4). Our proposed method optimized for MS-SSIM achieves slightly better results than Lee2020 (MS-SSIM).

Fig. 6 compares the performances of different methods on the Tecnick-100 dataset. Our scheme also outperforms all available learning-based methods and all the traditional image codecs including VVC (4:4:4) in term of both PSNR and MS-SSIM. Our method is the only method better than VVC. When the bit rate is above 0.2bpp, our method is about 0.2-0.3 dB higher than VVC (4:4:4).

Fig. 7 compares the performances of some methods on the Tecnick-40 dataset. The results of other methods are not available. Our scheme also outperforms other learned methods including Lee2020 [24] and VVC (4:4:4) in both PSNR and MS-SSIM.

One example are shown in Fig. 8 to compare the visual quality of different methods. Our method produces the most visually pleasing results.

Table V compares the complexities of different approaches. Since VVC, Hu2020 [54] and Cheng2020 [12] only run on CPU, we evaluate the encoding and decoding time of different methods at the similar bit rate on an 2.9GHz Intel Xeon Gold 6226R CPU. The average time over all Kodak images is used. The average model sizes at low bit rates and high bit rates are also reported.

It can be seen from Table V that compared to Chen2021 [11], the proposed scheme is faster in both encoding and decoding, but is slower than other methods. Our model size is also smaller than [11]. Compared to VVC, our encoding is faster, but the decoding is much slower.