Efficient and Effective Context-Based Convolutional Entropy Modeling for Image Compression

TL;DR

This paper introduces context-based convolutional networks (CCNs) for improved entropy modeling in image compression, achieving comparable performance to state-of-the-art methods with increased speed by capturing dependencies in image representations.

Contribution

The paper proposes a novel CCN architecture with 3D scanning and code dividing techniques for efficient, context-aware entropy modeling in both lossless and lossy image compression.

Findings

CCNs achieve comparable compression performance to state-of-the-art methods.

CCNs significantly increase the speed of entropy decoding.

The approach effectively models dependencies in image representations.

Abstract

Precise estimation of the probabilistic structure of natural images plays an essential role in image compression. Despite the recent remarkable success of end-to-end optimized image compression, the latent codes are usually assumed to be fully statistically factorized in order to simplify entropy modeling. However, this assumption generally does not hold true and may hinder compression performance. Here we present context-based convolutional networks (CCNs) for efficient and effective entropy modeling. In particular, a 3D zigzag scanning order and a 3D code dividing technique are introduced to define proper coding contexts for parallel entropy decoding, both of which boil down to place translation-invariant binary masks on convolution filters of CCNs. We demonstrate the promise of CCNs for entropy modeling in both lossless and lossy image compression. For the former, we directly apply a CCN to the binarized representation of an image to compute the Bernoulli distribution of each code for entropy estimation. For the latter, the categorical distribution of each code is represented by a discretized mixture of Gaussian distributions, whose parameters are estimated by three CCNs. We then jointly optimize the CCN-based entropy model along with analysis and synthesis transforms for rate-distortion performance. Experiments on the Kodak and Tecnick datasets show that our methods powered by the proposed CCNs generally achieve comparable compression performance to the state-of-the-art while being much faster.

Tables

Table 1. TABLE I: Running time in seconds and model complexity of different DNN-based entropy models on the Kodak dataset

	SIN	MCN	Multiscale PixelCNN	CCN	${\text{CCN}}_{\mathrm{light}}$
Encoding time	0.155	0.323	0.400	0.323	0.074
Decoding time	0.155	3079.68	0.400	35.28	0.984
# of parameters ($\times {10}^6)$	$6.70$	$3.89$	$3.67$	$3.89$	$0.10$
# of FLOPs ($\times {10}^{12})$	$0.58$	$2.78$	$2.87$	$2.78$	$0.07$

Table 2. TABLE II: Running time in seconds of our CCN-based entropy model at six bit rates on the Kodak dataset

Average bpp	0.100	0.209	0.362	0.512	0.671	0.794
Encoding	0.013	0.025	0.044	0.066	0.085	0.103
Decoding	0.116	0.227	0.457	0.735	1.150	1.232

Equations

$$P(\mathbf{y})\approx \prod _{i=0}^{M}P(y_i|\mathrm{PTX}(y_i,\mathbf{y})),$$

$$v_i^{(t)}(p,q)=\sum _{j=1}^{N_t}\left(u_j^{(t)}\ast \left(m^{(t)}\odot w_{i,j}^{(t)}\right)\right)(p,q)+b_i^{(t)},$$

$$\mathrm{SS}(v_i^{(0)}(p,q))=\mathrm{CTX}(y(p,q),{\Omega }_{p,q})\subset \mathrm{CTX}(y(p,q),\mathbf{y}),$$

$$m^{(0)}(\mu ,\nu )=\{\begin{array}{ll}1,& \text{mbox}if{\Omega }_{p,q}(\mu ,\nu )\in \mathrm{CTX}(y(p,q),\mathbf{y})\\ 0,& \text{mbox}otherwise.\end{array}$$

$$\mathrm{SS}(u_j^{(t)}(p+\mu ,q+\nu ))\subset$$

$$\subset$$

$${\mathbf{m}}^{(t)}(\mu ,\nu )=\{\begin{array}{ll}{\mathbf{m}}^{(0)}(\mu ,\nu ),& \text{mbox}if(\mu ,\nu )\ne (0,0)\\ 1,& \text{mbox}otherwise.\end{array}$$

$$m^{(t)}(\mu ,\nu )=\{\begin{array}{ll}{m}^{(0)}(\mu ,\nu ),& \text{mbox}if\mu +\nu \ne 0\\ 1,& \text{mbox}otherwise.\end{array}$$

$$v_{i,r}^{(t)}(p,q)=\sum _{j=1}^{N_t}\sum _{s=1}^M\left(u_{j,s}^{(t)}\ast \left(m_{r,s}^{(t)}\odot w_{i,j,r,s}^{(t)}\right)\right)(p,q)+b_{i,r}^{(t)},$$

$$m_{r,s}^{(0)}(\mu ,\nu )=\{\begin{array}{ll}1,& \text{mbox}if{\Omega }_{p,q}(\mu ,\nu )\in \mathrm{PTX}(y_r(p,q),\mathbf{y})\\ 0,& \text{mbox}otherwise.\end{array}$$

$$m_{r,s}^{(t)}(\mu ,\nu )=\{\begin{array}{ll}{m}_{r,s}^{(0)}(\mu ,\nu ),& \text{mbox}ifs+\mu +\nu \ne r\\ 1,& \text{mbox}otherwise,\end{array}$$

$$\mathrm{DOP}=\frac{1}{K}\sum _{k=0}^{K-1}{Z}_k=\frac{MHW}{M+H+W-2},$$

$$y_r(p,q)=\lfloor \frac{x(p,q)}{2^{7-r}}\rfloor \mathrm{mod}2,r=0,1,\dots ,7,$$

$$\ell (\mathbf{\theta })={\mathbb{E}}_{\mathbf{y}}$$

$$\displaystyle-\mathbbm{1}(y_{r}(p,q)=0)\log_{2}(1-v_{r}(p,q))\big{)}\Big{]},$$

$$y_r(p,q)=g_d(z_r(p,q))=\underset{\{{\omega }_{r,l}\}}{\mathrm{argmin}}\Vert z_r(p,q)-{\omega }_{r,l}{\Vert }_2^2.$$

$$\ell (\mathbf{\omega })=\frac{1}{MHW}\sum _{r,p,q}\Vert z_r(p,q)-y_r(p,q)){\Vert }_2^2,$$

$$y_r(p,q)\sim \sum _{i=0}^{C-1}{\pi }_i\mathcal{N}(y_r(p,q);{\mu }_i,{\sigma }_i^2),$$

$$P(y_r(p,q))={\int }_{\Delta }\sum _{i=0}^{C-1}{\pi }_i\mathcal{N}(\xi ;{\mu }_i,{\sigma }_i^2)d\xi .$$

$$\displaystyle\ell(\bm{\phi},\bm{\psi},\bm{\theta})=\mathbb{E}_{\bm{x}}\Bigg{[}$$

$$\displaystyle+\lambda\ell_{d}\bigg{(}g_{s}\Big{(}g_{d}\big{(}g_{a}(\bm{x};\bm{\phi})\big{)};\bm{\psi}\Big{)},\bm{x}\bigg{)}\Bigg{]}.$$

Full text

I Introduction

Data compression has played a significant role in engineering for centuries [1]. Compression can be either lossless or lossy. Lossless compression allows perfect data reconstruction from compressed bitstreams with the goal of assigning shorter codewords to more “probable” codes. Typical examples include Huffman coding [2], arithmetic coding [3], and range coding [4]. Lossy compression discards “unimportant” information of the input data, and the definition of importance is application-dependent. For example, if the data (such as images and videos) are meant to be consumed by the human visual system, importance should be measured in accordance with human perception, discarding features that are perceptually redundant, while keeping those that are most visually noticeable. In lossy compression, one must face the rate-distortion trade-off, where the rate is computed by the entropy of the discrete codes [5] and the distortion is measured by a signal fidelity metric. A prevailing scheme in the context of lossy image compression is transform coding, which consists of three operations - transformation, quantization, and entropy coding. Transforms map an image to a latent code representation, which is better-suited for exploiting aspects of human perception. Early transforms [6] are linear, invertible, and fixed for all bit rates; errors arise only from quantization. Recent transforms take the form of deep neural networks (DNNs) [7], aiming for nonlinear and more compressible representations. DNN-based transforms are mostly non-invertible, which may, however, encourage discarding perceptually unimportant image features during transformation. This gives us an opportunity to learn different transforms at different bit rates for optimal rate-distortion performance. Entropy coding is responsible for losslessly compressing the quantized codes into bitstreams for storage and transmission.

In either lossless or lossy image compression, a discrete probability distribution of the latent codes shared by the encoder and the decoder (i.e., the entropy model) is essential in determining the compression performance. According to the Shannon’s source coding theorem [5], given a vector of code intensities $\bm{y}=\{y_{0},\ldots,y_{M}\}$, the optimal code length of $\bm{y}$ should be $\lceil-\log_{2}P(\bm{y})\rceil$, where binary symbols are assumed to construct the codebook. Without further constraints, it is intractable to estimate $P(\bm{y})$ in high-dimensional spaces, a problem commonly known as the curse of dimensionality. For this reason, most entropy coding schemes assume $\bm{y}$ is fully statistically factorized with the same marginal distribution, leading to a code length of $\lceil-\sum_{i=0}^{M}\log_{2}P(y_{i})\rceil$. Alternatively, the chain rule in probability theory offers a more accurate approximation

[TABLE]

where $\mathrm{PTX}(y_{i},\bm{y})\subset\{y_{0},\ldots,y_{i-1}\}$ represents the partial context of $y_{i}$ coded before it in $\bm{y}$. A representative example is the context-based adaptive binary arithmetic coding (CABAC) [8] in H.264/AVC, which considers the two nearest codes as partial context, and obtains noticeable improvements over previous image/video compression standards. As the size of $\mathrm{PTX}(y_{i},\bm{y})$ becomes large, it is difficult to estimate this conditional probability by constructing histograms. Recent methods such as PixelRNN [9] and PixelCNN [10] take advantage of DNNs in modeling long range relations to increase the size of partial context, but are computationally intensive.

In this work, we present context-based convolutional networks (CCNs) for effective and efficient entropy modeling. Given $\bm{y}$, we specify a 3D zigzag coding order such that the most relevant codes of $y_{i}$ can be included in its context. Parallel computation during entropy encoding is straightforward as the context of each code is known and readily available. However, this is not always the case during entropy decoding. The partial context of $y_{i}$ should first be decoded sequentially for the estimation of $P(y_{i}|\mathrm{PTX}(y_{i},\bm{y}))$, which is prohibitively slow. To address this issue, we introduce a 3D code dividing technique, which partitions $\bm{y}$ into multiple groups in compliance with the proposed coding order. The codes within each group are assumed to be conditionally independent given their respective contexts, and therefore can be decoded in parallel. In the context of CCNs, this amounts to applying properly designed translation-invariant binary masks to convolutional filters.

To validate the proposed CCNs, we combine them with arithmetic coding [3] for entropy modeling. For lossless image compression, we convert the input grayscale image to eight binary planes and train a CCN to predict the Bernoulli distribution of $y_{i}$ by optimizing the entropy loss in information theory [11]. For lossy image compression, we parameterize the categorical distribution of $y_{i}$ with a discretized mixture of Gaussian (MoG) distributions, whose parameters (i.e., mixture weights, means, and variances) are estimated by three CCNs. The CCN-based entropy model is jointly optimized with analysis and synthesis transforms (i.e., mappings between raw pixel space and latent code space) over a database of training images, trading off the rate and the distortion. Experiments on the Kodak and Tecnick datasets show that our methods for lossless and lossy image compression perform favorably against image compression standards and DNN-based methods, especially at low bit rates.

II Related Work

In this section, we provide a brief overview of entropy models and lossy image compression methods based on DNNs. For traditional image compression techniques, we refer interested readers to [12, 13, 14].

II-A DNN-Based Entropy Modeling

The first and the most important step in entropy modeling is to estimate the probability $P(\bm{y})$. For most image compression techniques, $\bm{y}$ is assumed to be statistically independent, whose entropy can be easily computed through the marginal distributions [7, 15, 16, 17]. Arguably speaking, natural images undergoing a highly nonlinear analysis transform still exhibit strong statistical redundancies [18]. This suggests that incorporating context into probability estimation has great potentials in improving the performance of entropy coding.

DNN-based context modeling for natural languages and images has attracted considerable attention in the past decade. In natural language processing, recurrent neural networks (RNN) [19], and long short-term memory (LSTM) [20] are two popular tools to model long-range dependencies. In image processing, PixelRNN [9] and PixelCNN [10] are among the first attempts to exploit long-range pixel dependencies for image generation. The above-mentioned methods are computationally inefficient, requiring one forward propagation to generate (or estimate the probability of) a single pixel. To speed up PixelCNN, Reed et al. [21] proposed Multiscale PixelCNN, which is able to sample a twice larger intermediate image conditioning on the initial image. This process may be iterated to generate the final high-resolution result. When viewing Multiscale PixelCNN as an entropy model, we must losslessly compress and send the initial image as side information to the decoder for entropy decoding.

Only recently have DNNs for context-based entropy modeling become an active research topic. Ballé et al. [18] introduced a scale prior, which stores a variance parameter for each $y_{i}$ as side information. Richer side information generally leads to more accurate entropy modeling. However, this type of information should also be quantized, compressed and considered as part of the codes, and it is difficult to trade off the bits saved by the improved entropy model and the bits introduced by storing this side information. Li et al. [22] extracted a small code block for each $y_{i}$ as its context, and adopted a simple DNN for entropy modeling. The method suffers from heavy computational complexity similar to PixelRNN [9]. Li et al. [23] and Mentzer et al. [24] implemented parallel entropy encoding with masked DNNs. However, sequential entropy decoding has to be performed due to the context dependence, which remains painfully slow. In contrast, our CCN-based entropy model permits parallel entropy encoding and decoding, making it more attractive for practical applications.

II-B DNN-Based Lossy Image Compression

A major problem in end-to-end lossy image compression is that the gradients of the quantization function are zeros almost everywhere, making gradient descent-based optimization ineffective. Different strategies have been proposed to alleviate the zero-gradient problem resulting from quantization. From a signal processing perspective, the quantizer can be approximated by additive i.i.d. uniform noise, which has the same width as the quantization bin [25]. A desired property of this approximation is that the resulting density is a continuous relaxation of the probability mass function of $\bm{y}$ [7]. Another line of research introduced continuous functions (without the zero-gradient problem) to approximate the quantization function. The step quantizer is used in the forward pass, while its continuous proxy is used in the backward pass. Toderici et al. [26] learned an RNN to compress small-size images in a progressive manner. They later tested their models on large-size images [27]. Johnston et al. [28] exploited adaptive bit allocations and perceptual losses to boost the compression performance especially in terms of MS-SSIM [29].

The joint optimization of rate-distortion performance is another crucial issue in DNN-based image compression. The methods in [26, 27, 28] treat entropy coding as a post-processing step. Ballé et al. [7] explicitly formulated DNN-based image compression under the framework of rate-distortion optimization. Assuming $\bm{y}$ is statistically factorized, they learned piece-wise linear density functions to compute differential entropy as an approximation to discrete entropy. In a subsequent work [18], each $y_{i}$ is assumed to follow zero-mean Gaussian with its own variance separately predicted using side information. Minnen et al. [30] combined the autoregressive and hierarchical priors, leading to improved rate-distortion performance. Theis et al. [15] introduced a continuous upper bound of the discrete entropy with a Gaussian scale mixture. Rippel et al. [17] described pyramid-based analysis and synthesis transforms with adaptive code length regularization for real-time image compression. An adversarial loss [31] is incorporated to generate visually realistic results at low bit rates [17].

III CCNs for Entropy Modeling

In this section, we present in detail the construction of CCNs for entropy modeling. We work with a fully convolutional network, consisting of $T$ layers of convolutions followed by point-wise nonlinear activation functions, and assume the standard raster coding order (see Fig. 1). In order to perform efficient context-based entropy coding, two assumptions are made on the network architecture:

• •

For a code block $\bm{y}\in\mathbb{Q}^{M\times H\times W}$, where $M$, $H$, and $W$ denote the dimensions along channel, height, and width directions, respectively, the corresponding output of the $t$-th convolution layer $\bm{v}^{(t)}$ has a size of ${M\times H\times W\times N_{t}}$, where $N_{t}$ denotes the number of feature blocks to represent $\bm{y}$.

• •

Let $\mathrm{CTX}(y_{i}(p,q),\bm{y})$ be the set of codes encoded before $y_{i}(p,q)$ (i.e., full context), and $\mathrm{SS}(v^{(t)}_{i,j}(p,q))$ be the set of codes in the receptive field of $v^{(t)}_{i,j}(p,q)$ that contributes to its computation (i.e., support set), respectively. Then, $\mathrm{SS}(v^{(t)}_{i,j}(p,q))\subset\mathrm{CTX}(y_{i}(p,q),\bm{y})$.

Assumption I establishes a one-to-many correspondence between the input code block $\bm{y}$ and the output feature representation $\bm{v}^{(T)}$. In other words, the feature $v^{(t)}_{i,j}(p,q)$ in $i$-th channel and $j$-th feature block at spatial location $(p,q)$ is uniquely associated with $y_{i}(p,q)$. Assumption II ensures that the computation of $v^{(t)}_{i}(p,q)$ depends only on a subset of $\mathrm{CTX}(y_{i}(p,q),\bm{y})$. Together, the two assumptions guarantee the legitimacy of context-based entropy modeling in fully convolutional networks, which can be achieved by placing translation-invariant binary masks to convolution filters.

We start with the case of a 2D code block, where $\bm{y}\in\mathbb{Q}^{H\times W}$, and define masked convolution at the $t$-th layer as

[TABLE]

where $\ast$ and $\odot$ denote 2D convolution and Hadamard product, respectively. $\bm{w}^{(t)}_{ij}$ is a 2D convolution filter, $\bm{m}^{(t)}$ is the corresponding 2D binary mask, and $b_{i}^{(t)}$ is the bias.

According to Assumption I, the input $\bm{u}^{(t)}_{i}$ and the output $\bm{v}^{(t)}_{i}$ are of the same size as $\bm{y}$. The input code block $\bm{y}$ corresponds to $\bm{u}^{(0)}_{0}$.

For the input layer of a fully convolutional network, the codes to produce $v^{(0)}_{i}(p,q)$ is $\Omega_{p,q}=\{y(p+\mu,q+\nu)\}$, where $(\mu,\nu)\in\Psi$ is the set of local indices centered at $(0,0)$. We choose

[TABLE]

which can be achieved by setting

[TABLE]

Fig. 1 illustrates the concepts of full context $\mathrm{CTX}(y(p,q),\bm{y})$, support set $\mathrm{SS}(v^{(0)}(p,q))$, and translation-invariant mask $\mathbf{m}^{(0)}$, respectively. At the $t$-th layer, if we let $\bm{m}^{(t)}=\bm{m}^{(0)}$, for a code $y(p+\mu,q+\nu)\in\mathrm{CTX}(y(p,q),\bm{y})$, we have

[TABLE]

where the first line follows by induction and the second line follows from the definition of context. That is, as long as $y(p+\mu,q+\nu)$ is in the context of $y(p,q)$, we are able to compute $v^{(t)}_{i}{(p,q)}$ from $u^{(t)}_{j}(p+\mu,q+\nu)$ without violating Assumption II. In addition, $u^{(t)}_{j}(p,q)$ for $t>0$ is also generated from $\mathrm{CTX}(y{(p,q)},\bm{y})$, and can be used to compute $v^{(t)}_{i}{(p,q)}$. Therefore, we may modify the mask at the $t$-th layer

[TABLE]

III-A Proposed Strategies for Parallel Entropy Decoding

With the translation-invariant masks designed in Eqn. (4) and Eqn. (6), the proposed CCN can efficiently encode $\bm{y}$ in parallel. However, it remains difficult to parallelize the computation in entropy decoding. As shown in Fig. 2 (a) and (b), the two nearby codes in the same row (highlighted in orange and blue, respectively) cannot be decoded simultaneously because the orange code is in the support set (or context) of the blue code given the raster coding order. To speed up entropy decoding, we may further remove dependencies between codes at the risk of model accuracy. Specifically, we partition $\bm{y}$ into $K$ groups, namely, $\mathrm{GP}_{0}(\bm{y}),\ldots,\mathrm{GP}_{K-1}(\bm{y})$, and assume the codes within the same group are statistically independent. This results in a partial context $\mathrm{PTX}(y(p,q),\bm{y})=\{\mathrm{GP}_{0}(\bm{y}),\ldots,\mathrm{GP}_{k-1}(\bm{y})\}$ for $y(p,q)\in\mathrm{GP}_{k}(\bm{y})$. In other words, all codes in the $k$-th group share the same partial context, and can be decoded in parallel. Note that code dividing schemes are largely constrained by pre-specified coding orders. For example, if we use a raster coding order, it is straightforward to divide $\bm{y}$ by row. In this case, $y(p,q-1)$ ($p$ and $q$ index vertical and horizontal directions, respectively), which is extremely important in predicting the probability of $y(p,q)$ according to CABAC [8], has been excluded from its partial context. To make a good trade-off between modeling efficiency and accuracy, we switch to a zigzag coding order as shown in Fig. 2 (d), where $\mathrm{GP}_{k}(\bm{y})=\{y(p,q)|p+q=k\}$ and $\mathrm{PTX}(y(p,q),\bm{y})=\{y(p^{\prime},q^{\prime})|p^{\prime}+q^{\prime}<k\}$. As such, we retain the most relevant codes in the partial context for better entropy modeling (see Fig. 9 for quantitative results). Accordingly, the mask at the $t$-th layer becomes

[TABLE]

Now, we extend our discussion to a 3D code block, where $\bm{y}\in\mathbb{Q}^{M\times H\times W}$. Fig. 3 (a) shows the proposed 3D zigzag coding order and 3D code dividing technique. Specifically, $\bm{y}$ is divided into $K=M+H+W-2$ groups in the shape of diagonal planes, where the $k$-th group is specified by $\mathrm{GP}_{k}(\bm{y})=\{y_{r}(p,q)|r+p+q=k\}$. The partial context of $y_{r}(p,q)\in\mathrm{GP}_{k}(\bm{y})$ is defined as $\mathrm{PTX}(y_{r}(p,q),\bm{y})=\{y_{r^{\prime}}(p^{\prime},q^{\prime})|r^{\prime}+p^{\prime}+q^{\prime}<k\}$. We then write masked convolution in the 3D case as

[TABLE]

where $\{i,j\}$ and $\{r,s\}$ are indices for the feature block and the channel, respectively. For the 2D case, each layer shares the same mask ($M=1$). When extending to 3D code blocks, each channel in a layer shares a mask, and there are a total of $M$ 3D masks. For the input layer, the codes to produce $v^{(0)}_{i,r}(p,q)$ is $\Omega_{p,q}=\{y_{s}{(p+\mu,q+\nu)}\}_{(\mu,\nu)\in\Psi,0\leq s<M}$, based on which we define the mask as

[TABLE]

For the $t$-th layer, we modify the mask to include the current diagonal plane

[TABLE]

as shown in Fig. 4.

III-B Analysis of the Degree of Parallelism

We analyze the degree of parallelism (DOP) of the proposed CCN, which is defined as the number of operations that can be or are being simultaneously executed by a computer. For entropy encoding, all codes can be estimated in parallel, leading to a DOP of $MHW$. The proposed CCN divides the code block into $K=M+H+W-2$ groups and allows the codes within each group to be processed in parallel. Ideally, the DOP should be the group size (i.e., the number of codes in the group). However, different groups have different sizes in the shape of irregular diagonal planes. Here, the average group size is adopted as a rough indication of DOP,

[TABLE]

where $Z_{k}$ denotes the number of codes in the $k$-th group. It is clear that DOP is a function of height, width, and channel of the code block.

IV CCN-Based Entropy Models for Lossless Image Compression

In this section, we combine our CCN-based entropy model with the arithmetic coding algorithm for lossless image compression.

As a starting point, we binarize the grayscale image $\bm{x}\in\mathbb{R}^{H\times W}$ to obtain a 3D code block

[TABLE]

where we index the most significant bit-plane with $r=0$. Our CCN takes $\bm{y}$ as input and produces a feature block $\bm{v}$ (the superscript $(T)$ is omitted for notation convenience) of the same size to compute the mean estimates of Bernoulli distributions $P(y_{r}(p,q)|\mathrm{SS}(v_{r}(p,q)))$. Fig. 5 shows the network architecture, which has eleven masked convolution layers with parametric ReLU nonlinearities in between. The last convolution responses undergo a sigmoid nonlinearity to constrain their range between $[0,1]$. We make four residual connections as suggested in [32] to accelerate training. We will experiment with two hyper-parameters in CCN: the convolution filter size for all layers $S$ and the number of feature blocks in hidden layers $N$.

To optimize the network parameters collectively denoted by $\bm{\theta}$, we adopt the expected code length as the empirical loss

[TABLE]

where $\mathbbm{1}(\cdot)$ is an indicator function and the expectation may be approximated by averaging over a mini-batch of training images. Finally, we implement our own arithmetic coding method with the learned CCN-based entropy model to compress $\bm{y}$ to bitstreams, and report performance using actual bit rates.

V CCN-Based Entropy Models for Lossy Image Compression

In lossy image compression, our objective is to minimize a weighted sum of rate and distortion, $\ell_{r}+\lambda\ell_{d}$, where $\lambda$ governs the trade-off between the two terms. As illustrated in Fig. 6, the proposed compression method has four components: an analysis transform $g_{a}$, a quantizer $g_{d}$, a CCN-based entropy model, and a synthesis transform $g_{s}$. The analysis transform $g_{a}$ takes a color image $\bm{x}$ as input and produces the latent code representation $\bm{z}$, which is further quantized to generate the discrete code block $\bm{y}$. $g_{a}$ consists of three convolutions, each of which is followed by down-sampling with a factor of two. A dense block [33] comprised of seven convolutions is employed after each down-sampling. After the last dense block, we add another convolution layer with $M$ filters to produce $\bm{z}$. Empirically, the parameter $M$ sets the upper bound of the bit rate that a general DNN-based compression method can achieve. The parameters of $g_{a}$ constitute the parameter vector $\bm{\phi}$ to be optimized.

The synthesis transform $g_{s}$ has a mirror structure of the analysis transform. Particularly, the depth-to-space reshaping [27, 34] is adopted to up-sample the feature maps. The last convolution layer with three filters is responsible for producing the compressed image in RGB space. The parameters of $g_{s}$ constitute the parameter vector $\bm{\psi}$ to be optimized.

For the quantizer $g_{d}$, we parameterize its quantization centers for the $r$-th channel by $\{\omega_{r,0},\ldots,\omega_{r,L-1}\}$, where $L$ is the number of quantization centers and $\omega_{r,0}\leq\ldots\leq\omega_{r,L-1}$. The monotonicity of $\bm{\omega}$ can be enforced by a simple re-parameterization based on cumulative functions. Given a fixed set of $\bm{\omega}$, we perform quantization by mapping $z_{r}(p,q)$ to its nearest center that minimizes the quantization error

[TABLE]

$g_{d}$ has zero gradients almost everywhere, which hinders training via back-propagation. Taking inspirations from binarized neural networks [35, 36, 37], we make use of an identify mapping $\hat{g}_{d}(y_{r}(p,q))=y_{r}(p,q)$ as a continuous proxy to the step quantization function. During training, we use $g_{d}$ and $\hat{g}_{d}$ in the forward and backward passes, respectively.

The quantization centers $\bm{\omega}$ should be optimized by minimizing the mean squared error (MSE),

[TABLE]

which is essentially a $k$-means clustering problem, and can be solved efficiently by the Lloyd’s algorithm [38]. Specifically, we initialize $\bm{\omega}$ using uniform quantization, which appears to work well in all experiments. To make parameter learning of the entire model smoother, we adjust $\bm{\omega}$ using stochastic gradient descent instead of a closed-form update.

Without prior knowledge of the categorical distributions of the quantized codes ${\bm{y}}$, we choose to work with discretized MoG distributions, whose parameters are predicted by the proposed CCNs. We write the differentiable MoG distribution with $C$ components as

[TABLE]

where $\pi_{i}$, $\mu_{i}$ and $\delta^{2}_{i}$ are the mixture weight, mean, and variance of the $i$-th component, respectively. Then, the probability of $y_{r}(p,q)$ is calculated as the integral over the quantization bin $\Delta$ that the code lies in

[TABLE]

For example, suppose $y_{r}(p,q)=\omega_{r,l}$, then $\Delta=[(\omega_{r,l-1}+\omega_{r,l})/2,(\omega_{r,l}+\omega_{r,l+1})/2]$, where we define $\omega_{r,-1}=-\infty$ and $\omega_{r,L}=\infty$, respectively.

Next, we describe the proposed entropy model in lossy image compression, which is comprised of three CCNs with the same structure, as shown in Fig. 6. Each CCN consists of nine masked convolutions with three residual connections to produce $C$ feature blocks, matching the number of components in MoG. They separately output mixture weights, means and variances to build the discretized MoG distributions. The network parameters of our CCNs constitute the parameter vector $\bm{\theta}$ to be optimized.

Finally, we are able to write the empirical rate-distortion objective for the parameters $\{\bm{\phi},\bm{\psi},\bm{\theta}\}$ as

[TABLE]

$\ell_{d}$ is the distortion term, which is more preferable to be assessed in perceptual space. In this paper, we optimize and evaluate our lossy image compression methods using the standard MSE and a perceptual metric MS-SSIM [29]. Similar to lossless image compression, we combine the optimized entropy model with arithmetic coding and report the actual bit rates.

VI Experiments

In this section, we test the proposed CCN-based entropy models in lossless and lossy image compression by comparing to state-of-the-art image coding standards and recent deep image compression algorithms. We first collect $10,000$ high-quality and high-definition images from Flickr, and down-sample them to further reduce possibly visible artifacts. We crop $1,280,000$ grayscale patches of size $128\times 128$ and $640,000$ color patches of size $3\times 256\times 256$ as the training sets for lossless and lossy image compression, respectively. We test our models on two independent datasets - Kodak and Tecnick [39], which are widely used to benchmark image compression performance. The Caffe implementations along with the pre-trained models are made available at https://github.com/limuhit/CCN.

VI-A Lossless Image Compression

We train our CCN-based entropy models using the Adam stochastic optimization package [40] by minimizing the objective in Eqn. (13). We start with a learning rate of $10^{-4}$, and subsequently lower it by a factor of $10$ when the loss plateaus, until $10^{-6}$. The (actual) bit rate in terms of bits per pixel (bpp) is used to quantify the compression performance, which is defined as the ratio between the total amount of bits used to code the image and the number of pixels in the image. A smaller bpp indicates better performance. For example, an uncompressed grayscale image has eight bpp.

We first compare the proposed CCNs with masked convolutional networks (MCNs) [24, 23], Multiscale PixelCNN [21], and side information networks (SINs)  [18] for entropy modeling. As a special case of CCNs, MCNs specify the raster coding order without using any code dividing technique (see Fig. 2). We implement our own version of MCN that inherits the network architecture from CCN with $N=16$ (number of feature blocks) and $S=5$ (filter size).

Multiscale PixelCNN [21] is originally designed for image generation. Here we slightly modify it for entropy modeling. Specifically. we approximate the iterative process of Multiscale PixelCNN with a single network by mapping the initial grayscale image directly to the probability distributions of the original image. The network includes an up-sampling convolution layer followed by five residual blocks. We test various configurations of Multiscale PixelCNN, and choose to down-sample the original image by a factor of 4 as initialization, which delivers the best overall compression performance.

SINs summarize the side information of ${\bm{y}}$ with a separate DNN, which is beneficial in probability estimation. We adopt a DNN-based autoencoder including three stride convolutions and two residual connections to generate the side information, which is further quantized and compressed with arithmetic coding for performance evaluation. We test five variants of SINs with different amount of side information by changing the number of output channels $M$. All competing models are trained on the same dataset described at the beginning of Section VI.

We also introduce a light-weight version of our method, which we name CCNlight, by making the network architecture lighter (with $N=4$ and $S=3$) and by dividing the input image into non-overlapping patches for parallel processing.

We split the image into a total of $R\times R$ smaller patches with the size of $\frac{H}{R}\times\frac{W}{R}$. The DOP now becomes $\frac{MHW}{M+H/R+W/R-2}$, which further speeds up the original CCN by about $R$ times.

Fig. 7 shows the bit rates of the competing methods on the Kodak dataset. The proposed CCN matches the best performing model MCN, which suggests that with the proposed zigzag coding order and code dividing technique, CCN includes the most important codes as the partial context of the current code being processed. The bit rates of SINs come from two parts - the image itself and the side information. It is clear from the figure that increasing the amount of side information leads to bit savings of the image, at the cost of additional bits introduced to code the side information. In general, it is difficult to determine the right amount of side information for optimal compression performance. Multiscale PixelCNN can also be regarded as a special case of SINs, whose side information is a small image without further processing, leading to the worst performance.

We also compare CCN with the widely used lossless image compression standards, including TIFF, GIF, PNG, JPEG-LS, JPEG2000-LS, and BPG. All test results are generated by MATLAB2017. From Fig. 8, we find that CCN (along with its light-weight version CCNlight) overwhelms all competing methods on the Kodak/Tecnick dataset, achieving more than $5.9\%/6.2\%$ bit savings compared to the best lossless image compression standard, JPEG2000-LS.

The running time of the competing DNN-based entropy models is tested on an NVIDIA TITAN Xp machine, whose results on the Kodak dataset are listed in Table I. For encoding, CCNlight enjoys the fastest speed, followed by Multiscale PixelCNN and SIN (the best performing variant). Despite similar encoding time, they have substantially different decoding complexities. Multiscale PixelCNN has the fastest decoder, followed by SIN. Due to the sequential decoding nature, MCN is the slowest, taking nearly one hour to decode a grayscale image of size $752\times 496$. Our CCN achieves a significant improvement upon MCN with the proposed code dividing technique, while maintaining nearly the same bit rates. Moreover, CCNlight speeds up CCN more than $30$ times, striking a good balance between model efficiency and model accuracy. We also report the model complexity of the competing models in terms of the number of parameters and floating point operations (FLOPs) in Table I. Except for the CCNlight, all the competing models have comparable complexity.

We conduct thorough ablation experiments to analyze the impact of individual components of CCN to final compression performance. Fig. 9 shows the bit rates of CCNs with three different numbers of feature blocks ($N\in\{4,8,16\}$) and two filter sizes ($S\in\{3,5\}$). When the filter size and the network depth are fixed, adding more feature blocks effectively increases model capability and thus boosts the compression performance. Similarly, using a larger filter size with fixed network depth and feature block number increases the partial context, leading to better entropy modeling. Moreover, we replace the proposed zigzag coding order in CCN with the raster coding order, whose model is denoted by CCNraster. From Fig. 9, we observe that the performance of CCNraster drops significantly, only comparable to the CCN with four feature blocks. This verifies the advantages of the proposed coding order.

VI-B Lossy Image Compression

In lossy image compression, the analysis transform $g_{a}$, the non-uniform quantizer $g_{d}$, the CCN-based entropy model, and the synthesis transform $g_{s}$ are jointly optimized for rate-distortion performance. In the early stages of training, the probability $P({\bm{y}})$ often changes rapidly, which makes it difficult to keep track of, and may cause instability in learning the entropy model. We find that this issue can be alleviated by a simple warmup strategy. Specifically, $g_{a}$ and $g_{s}$ are trained using the distortion term $\ell_{d}$ only for the first epoch. We then fix $g_{a}$ and train the CCN-based entropy model until it reasonably fits the current distribution of the codes. After that, we end-to-end optimize the entire method for the remaining epochs. We use Adam with a learning rate of $10^{-5}$ and gradually lower it by a factor of $10$, until $10^{-7}$. The number of quantization centers is $L=8$, and the number of Gaussian components in MoG is $C=3$.

Fig. 10 shows the learned continuous distributions of representative codes, which are typically multimodal, and therefore cannot be well fitted by unimodal distributions (e.g., Gaussian). The optimization is performed separately for each $\lambda$ and for each distortion measure. We optimize twelve models for six bit rates and two distortion metrics (MSE and MS-SSIM). MSE is converted to peak signal-to-noise ratio (PSNR) for quantitative analysis.

We also evaluate the effect of several components, i.e., the transforms, entropy model, and joint optimization, of the proposed lossy image compression method. All the variants of our method are optimized for MSE distortion and tested on the Kodak dataset. The compression performance of these methods is shown in Fig. 11. In terms of the transforms, we replace each dense block of our method with three residual blocks. To keep the same network depth, the first residual block consists of three convolution layers, while the last two residual blocks have two convolution layers. Our method performs better than this variant, especially at higher bit rates. As for the entropy model, we replace the proposed CCN with CCNraster. Similar to the results in lossless image compression, our method with the zigzag coding order outperforms this variant with the raster coding order. Moreover, we investigate the effect of joint optimization of the transforms and the entropy model. A baseline method is introduced, which first optimizes the transforms for MSE ($\lambda=\infty$), and then trains the CCN-based entropy model by fixing the learned transforms. The baseline method performs worse than the jointly optimized one by 1 dB. Similar trends can be observed for the variants with residual blocks.

We compare our methods with existing image coding standards and recent DNN-based compression models. These include JPEG [13], JPEG2000 [14], BPG [41], Agustsson17 [16], Theis17 [15], Toderici17 [27], Rippel17 [17], Mentzer18 [24], Johnston17 [28], Ballé17 [7], and Li18 [22]. Both JPEG (with 4:2:0 chroma subsampling) and JPEG2000 are based on the optimized implementations in MATLAB2017. For BPG, we adopt the latest version from its official website with the default setting. When it comes to DNN-based models for lossy image compression, the implementations are generally not available. Therefore, we copy the results from their respective papers.

Fig. 12 shows the rate-distortion curves on the Kodak dataset. We find that our method optimized for MSE outperforms all competing methods at low bit rates ($<0.5$ bpp), except for BPG. When optimized for MS-SSIM, our method performs on par with Rippel17 and is much better than the rest. Fig. 13 shows the rate-distortion curves on the Tecnick dataset, where we observe similar trends for both PSNR and MS-SSIM. An interesting observation is that when we continue increasing the bit rate, PSNR/MS-SSIM starts to plateau, which may be due to the limited model capability. Without any constraint on rate ($\lambda=\infty$) and quantization ($L=\infty$), our method optimized for MSE only reaches $38.2$ dB on the Kodak dataset, which can be treated as an empirical upper-bound for our network structure. Preliminary results from [18] indicate that increasing the depth and width of the network leads to performance improvements at high bit rates.

We visually compare the compressed images by our method against Ballé17, Li18, JPEG2K, and BPG. Fig. 14 and Fig. 15 show sample compressed results on the Kodak and Tecnick datasets, respectively. JPEG2K and BPG exhibit artifacts (such as blocking, ringing, blurring, and aliasing) that are common to all handcrafted transform coding methods, reflecting the underlying linear basis functions. Ballé17 is effective at suppressing ringing artifacts at the cost of over-smoothing fine structures. Li18 allocates more bits to preserve large-scale strong edges, while tends to eliminate small-scale localized features (e.g., edges, contours, and textures). In contrast, our method generates compressed images with more faithful details and less visible distortions.

We report the running time of our method at six bit rates on the Kodak dataset using the same machine. It takes $0.024$ second to generate ${\bm{y}}$ and $0.032$ second to reconstruct the image. The entropy coding time is listed in Table II, where we see that more time is needed to encode and decode images at higher bit rates. This is because,

in our experiments, we use a larger $M$ at higher bit rates as a way of preserving more perceptually meaningful information. This corresponds to more convolution operations and longer processing time.

With the help of the proposed code dividing technique, our method performs entropy decoding in around one second for images of size $752\times 496$.

We also note that the entropy coding time for lossy image compression is much faster than that of lossless image compression for two main reasons. First, in lossless compression, the size of the code block to CCN for entropy coding is $8\times H\times W$. In lossy compression, the input image is first transformed to the code block with the size of $M\times\left(H/8\right)\times\left(W/8\right)$ ($M\leq 32$). With a smaller code block size, the lossy CCN is expected to be faster than the lossless CCN. Moreover, $N$ and $S$ also highly affect the speed. For the CCN in lossless compression, we set $N=16$ and $S=5$, whereas for the CCN in lossy compression, we set $N=3$ and $S=5$, leading to the faster speed.

VII Conclusion and Discussion

We have introduced CCNs for context-based entropy modeling. Parallel entropy encoding and decoding are achieved with the proposed coding order and code dividing technique, which can be efficiently implemented using masked convolutions. We tested the CCN-based entropy model (combined with arithmetic coding) in both lossless and lossy image compression. For the lossless case, our method achieves the best compression performance, which we believe arises from the more accurate estimation of the Bernoulli distributions of the binary codes. For the lossy case, our method offers improvements both visually and in terms of rate-distortion performance over image compression standards and recent DNN-based models.

The application scope of the proposed CCN is far beyond building the entropy model in image compression. As a general probability model, CCN appears promising for a number of image processing applications. For example, we may use CCN to learn a probability model $P(\bm{x})$ for natural images, and use it as a prior in Bayesian inference to solve various vision tasks such as image restoration [42, 43], image quality assessment [44, 45], and image generation[9, 46].

Acknowledgment

This work was supported in part by Shenzhen Research Institute of Big Data and Shenzhen Institute of Artificial Intelligence and Robotics for Society. The authors would like to thank the NVIDIA Corporation for donating a TITAN Xp GPU used in this research.