Deep Lossy Plus Residual Coding for
Lossless and Near-lossless Image Compression

In LIC, we employ sophisticated image encoder and decoder while efficient hyper-prior model [25], as shown in Fig. 2. The image encoder $g_{e}(\cdot)$ and decoder $g_{d}(\cdot)$ are composed of analysis, synthesis and Swin-Attention blocks following the philosophy of residual and self-attention learning [56, 57, 33], as detailed in Fig. 3. In Swin-Attention blocks, we adopt the window and shifted window based multi-head self-attention (W-MSA/SW-MSA) in [58] to aggregate information within and across local windows adaptively, which improve the representation ability of $g_{e}(\cdot)$ and $g_{d}(\cdot)$ with moderate computational complexity. With $g_{e}(\cdot)$ and $g_{d}(\cdot)$, we transform an input raw image $\mathbf{x}$ to its latent representation $\mathbf{y}=g_{e}(\mathbf{x})$, quantize $\mathbf{y}$ to ${\hat{\mathbf{y}}}=Q(\mathbf{y})$, and reversely transform ${\hat{\mathbf{y}}}$ to the lossy reconstruction ${\tilde{\mathbf{x}}}=g_{d}({\hat{\mathbf{y}}})$.

Because the sophisticated image encoder and decoder can largely reduce the spatial redundancies in $\mathbf{x}$, the burden of entropy coding ${\hat{\mathbf{y}}}$ is relieved and we decide to employ the efficient hyper-prior model [25] without any context models to ensure the coding parallelization on GPUs. The hyper-prior model extracts side information ${\hat{\mathbf{z}}}=Q(h_{e}(\mathbf{y}))$ to model the probability distribution of ${\hat{\mathbf{y}}}$, where $h_{e}(\cdot)$ is the hyper-encoder. We assume a factorized Gaussian distribution model $\mathcal{N}({\boldsymbol{\mu}},{\boldsymbol{\sigma}})=h_{d}({\hat{\mathbf{z}}})$ for $p_{\theta}({\hat{\mathbf{y}}}|{\hat{\mathbf{z}}})$ where $h_{d}(\cdot)$ is the hyper-decoder, and a non-parametric factorized density model for $p_{\theta}({\hat{\mathbf{z}}})$. The $R_{\hat{\mathbf{y}}}$ in (7) and (10) is thus extended by

where $R_{{\hat{\mathbf{y}}},{\hat{\mathbf{z}}}}$ is the cost of encoding both ${\hat{\mathbf{y}}}$ and ${\hat{\mathbf{z}}}$.

Given the raw image $\mathbf{x}$ and its lossy reconstruction ${\tilde{\mathbf{x}}}$ from LIC, we have the residual $\mathbf{r}=\mathbf{x}-{\tilde{\mathbf{x}}}$. We next introduce RC to estimate the probability mass function (PMF) of $\mathbf{r}$ and compress $\mathbf{r}$ with arithmetic coding [46] accordingly.

Denote by $\mathbf{u}=g_{\mathbf{u}}({\hat{\mathbf{y}}})$, where the feature $\mathbf{u}$ is generated from ${\hat{\mathbf{y}}}$ by $g_{\mathbf{u}}(\cdot)$. The $g_{\mathbf{u}}(\cdot)$ and the image decoder $g_{d}(\cdot)$ share the network except the last convolutional layer, as shown in Fig. 1. We interpret $\mathbf{u}$ as the feature of the residual $\mathbf{r}$ given ${\tilde{\mathbf{x}}}$ and ${\hat{\mathbf{y}}}$. The feature $\mathbf{u}$ shares the same height and width with $\mathbf{r}$ and has 256 channels. Unlike the latent representation ${\hat{\mathbf{y}}}$ of which the spatial redundancies are largely reduced by the image encoder, the residual $\mathbf{r}$ in the pixel domain has spatial redundancies that cannot be fully exploited by only the feature $\mathbf{u}$. Therefore, we further introduce the autoregressive model into the statistical modeling of $\mathbf{r}$, leading to

where $r_{<(i,c)}$ denotes the elements of $\mathbf{r}$ encoded or decoded before $r_{i,c}$ in a pre-defined scan order. In practice, we implement spatial autoregressive model using a mask convolutional layer with a specific receptive field, rather than depending on all elements in $r_{<(i,c)}$. We regard the receptive field of the mask convolutional layer as the context $C_{\mathbf{r}}$. Based on (12) and $C_{\mathbf{r}}$, we reformulate the $R_{\mathbf{r}}$ in (7) as

Specifically, we utilize a $7\times 7$ mask convolutional layer with 256 channels to extract the context $C_{r_{i}}\in C_{\mathbf{r}}$ from $r_{<(i,c)}$. The $C_{r_{i}}$ is shared by $r_{i,c}$ of all channels. For RGB images with three channels, we have

We further adopt a channel autoregressive scheme over $r_{i,1}$, $r_{i,2}$, $r_{i,3}$ [20] and reformulate $p_{\boldsymbol{\theta}}(r_{i,1},r_{i,2},r_{i,3}|u_{i},C_{r_{i}})$ as

We model the PMF of $r_{i,c}$ with discrete logistic mixture likelihood [20] and propose a sub-network to estimate the corresponding entropy parameters, including mixture weights $\pi_{i}^{k}$, means $\mu_{i,c}^{k}$, variances $\sigma_{i,c}^{k}$ and mixture coefficients $\beta_{i,t}$. $k$ denotes the index of the $k$-th logistic distribution. $t$ denotes the channel index of $\beta$. The network architecture of the entropy model is shown in Fig. 4. We utilize a mixture of $K=5$ logistic distributions. The channel autoregressive scheme over $r_{i,1},r_{i,2},r_{i,3}$ is implemented by updating the means using:

With $\pi_{i}^{k}$, $\tilde{\mu}_{i,c}^{k}$ and $\sigma_{i,c}^{k}$, we have

where $\operatorname*{logistic}(\cdot)$ denotes the logistic distribution. For discrete $r_{i,c}$, we evaluate $p_{\boldsymbol{\theta}}(r_{i,c}|r_{i,<c},u_{i},C_{r_{i}})$ as [20]:

where $S(\cdot)$ denotes the sigmoid function. $r_{i,c}^{+}=r_{i,c}+0.5$ and $r_{i,c}^{-}=r_{i,c}-0.5$. The probability inference scheme of $\mathbf{r}$ is sketched in Fig. 5a.

We finally introduce SQRC to realize scalable near-lossless image compression with variable $\ell_{\infty}$ bound $\tau\in\{1,2,\ldots\}$. Though near-lossless image compression given a specific $\tau$ can be realized by optimizing (10), this $\tau$-specific scheme in Fig. 5b leads to two problems:

• •

  Relaxation problem of residual quantization. Unlike rounding quantization, the bin size of the residual quantization (9) is much larger. Moreover, the original residuals are not uniformly distributed in each bin, and thus cannot be relaxed by adding uniform noise.

• •

  Storage problem of multiple networks. To deploy the near-lossless codec, we have to transmit and store multiple networks for different $\tau$’s, which is storage-inefficient.

Instead, we propose a scalable near-lossless image compression scheme, which can circumvent the relaxation of residual quantization and utilize a single network to satisfy variable $\ell_{\infty}$ error bound $\tau\in\{1,2,\ldots\}$. Specifically, the scalable compression scheme is based on the learned lossless compression with the DLPR coding framework. We keep the lossy reconstruction ${\tilde{\mathbf{x}}}$ fixed and quantize the original residual $\mathbf{r}$ to ${\hat{\mathbf{r}}}$ with variable $\tau$’s by (9). To encode the quantized ${\hat{\mathbf{r}}}$, we can derive the PMF of ${\hat{\mathbf{r}}}$ from the learned PMF of the original $\mathbf{r}$. Given $\tau$ and the learned PMF $p_{\boldsymbol{\theta}}(r_{i,c}|r_{i,<c},u_{i},C_{r_{i}})$ of original $r_{i,c}$, the PMF $\hat{p}_{\boldsymbol{\theta}}(\hat{r}_{i,c}|r_{i,<c},u_{i},C_{r_{i}})$ of quantized $\hat{r}_{i,c}$ can be computed by the following PMF quantization:

We show an illustrative example in Fig. 6. Together with (14) and (15), we can derive the probability model $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\mathbf{r}})$ of ${\hat{\mathbf{r}}}$, which is optimal given the learned $p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ of $\mathbf{r}$. The resulting cost of encoding ${\hat{\mathbf{r}}}$, denoted by $R^{\tau}_{{\hat{\mathbf{r}}}}$, is reduced significantly with the increase of $\tau$.

However, encoding ${\hat{\mathbf{r}}}$ with $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\mathbf{r}})$ results in undecodable bitstreams, since the original residual $\mathbf{r}$ is unknown to the decoder. $\hat{p}_{\boldsymbol{\theta}}(\hat{r}_{i,c}|r_{i,<c},u_{i},C_{r_{i}})$ cannot be evaluated without $r_{i,<c}$ and causal context $C_{r_{i}}$. Instead, we can evaluate PMF using the quantized residual ${\hat{\mathbf{r}}}$, i.e., we evaluate $p_{\boldsymbol{\theta}}(r_{i,c}|\hat{r}_{i,<c},u_{i},C_{\hat{r}_{i}})$ and derive $\hat{p}_{\boldsymbol{\theta}}(\hat{r}_{i,c}|\hat{r}_{i,<c},u_{i},C_{\hat{r}_{i}})$ with (19), leading to $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\hat{\mathbf{r}}})$ for the encoding of ${\hat{\mathbf{r}}}$. Because of the mismatch between training (with $\mathbf{r}$) and inference (with ${\hat{\mathbf{r}}}$) phases, it leads to biased PMF $p_{\boldsymbol{\theta}}(r_{i,c}|\hat{r}_{i,<c},u_{i},C_{\hat{r}_{i}})$, $\hat{p}_{\boldsymbol{\theta}}(\hat{r}_{i,c}|\hat{r}_{i,<c},u_{i},C_{\hat{r}_{i}})$ and $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\hat{\mathbf{r}}})$. The above probability inference scheme is sketched in Fig. 5c.

SQRC for Bias Correction: Because of the discrepancy between the oracle $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\mathbf{r}})$ and the biased $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\hat{\mathbf{r}}})$, encoding ${\hat{\mathbf{r}}}$ with $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\hat{\mathbf{r}}})$ degrades the compression performance. In order to tackle this problem, we propose SQRC for bias correction to close the gap between the oracle $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\mathbf{r}})$ and the biased $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\hat{\mathbf{r}}})$, while the resulting bitstreams are still decodable. The components of SQRC are illustrated in Fig. 1. The masked convolutional layer in SQRC is shared with that in RC. The conditional entropy model has the same network architecture as the entropy model illustrated in Fig. 4, but replaces the convolutional layers with the conditional convolutional layers [19, 60] illustrated in Fig. 7.

The probability inference scheme with SQRC is sketched in Fig. 5d. For $\tau=0$, we still select the entropy model in RC to estimate $p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ to encode $\mathbf{r}$. For $\tau\in\{1,2,\ldots,N\}$, we select the conditional entropy model in SQRC to estimate $p_{\boldsymbol{\varphi}}(\mathbf{r}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$ conditioned on $\tau$. We then derive $\hat{p}_{\boldsymbol{\varphi}}({\hat{\mathbf{r}}}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$ with (19) to encode ${\hat{\mathbf{r}}}$, where ${\boldsymbol{\varphi}}$ denote the parameters of SQRC. As $\hat{p}_{\boldsymbol{\varphi}}({\hat{\mathbf{r}}}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$ approximates the oracle $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\mathbf{r}})$ better than the biased $\hat{p}_{\boldsymbol{\theta}}({\hat{\mathbf{r}}}|\mathbf{u},C_{\hat{\mathbf{r}}})$, the compression performance can be improved. Since evaluating $\hat{p}_{\boldsymbol{\varphi}}({\hat{\mathbf{r}}}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$ is independent of $\mathbf{r}$, the resulting bitstreams are decodable. In experiments, we demonstrate that the proposed scalable near-lossless compression scheme with SQRC in Fig. 5d can outperform both the $\tau$-specific near-lossless scheme in Fig. 5b and the scalable near-lossless compression scheme without SQRC in Fig. 5c.

The full loss function for jointly optimizing LIC and RC, i.e., DLPR coding for lossless image compression, is

where ${\boldsymbol{\theta}}$ and ${\boldsymbol{\phi}}$ are the learned parameters of LIC and RC. Besides rate terms $R_{{\hat{\mathbf{y}}},{\hat{\mathbf{z}}}}$ in (11) and $R_{\mathbf{r}}$ in (13), we further introduce a distortion term $D_{ls}(\mathbf{x},{\tilde{\mathbf{x}}})$ to minimize the mean square error (MSE) between the raw image $\mathbf{x}$ and its lossy reconstruction ${\tilde{\mathbf{x}}}$:

As discussed in [25], minimizing MSE loss is equivalent to learn a LIC that fits residual $\mathbf{r}$ to a factorized Gaussian distribution. However, the discrepancy between the real distribution of $\mathbf{r}$ and the factorized Gaussian distribution is usually large. Therefore, we utilize a sophisticated discrete logistic mixture likelihood model to encode $\mathbf{r}$ in our DLPR coding framework.

The $\lambda$ in (20) is a “rate-distortion” trade-off parameter between the lossless compression rate and the MSE distortion. When $\lambda=0$, the loss function (20) is consistent with the theoretical DLPR coding formulation (7), and ${\tilde{\mathbf{x}}}$ becomes a latent variable without any constraints. In experiments, we study the effects of $\lambda$’s on the lossless and near-lossless image compression performance. We set $\lambda=0$ leading to the best lossless image compression, while set $\lambda=0.03$ leading to robust near-lossless image compression with variable $\tau$’s.

For training SQRC, we generate random $\tau\in\{1,2,\ldots,N\}$ and quantize $\mathbf{r}$ to ${\hat{\mathbf{r}}}$ with (9). Given $\mathbf{u}$ and the extracted context $C_{\hat{\mathbf{r}}}$ from quantized ${\hat{\mathbf{r}}}$, we use the conditional entropy model to estimate $-\log p_{\boldsymbol{\varphi}}(\mathbf{r}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$ conditioned on different $\tau$’s, and minimize

where ${\boldsymbol{\varphi}}$ denote the learned parameters of the conditional entropy model. $-\log p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ is estimated by the entropy model in RC. $\mathcal{L}({\boldsymbol{\varphi}})$ can be considered as an approximate KL-divergence or relative entropy [55] between $p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ and $p_{\boldsymbol{\varphi}}(\mathbf{r}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$.

SQRC is trained together with LIC and RC, but minimizing (22) only updates the parameters of the conditional entropy model as shown in Fig. 1. The masked convolutional layer is shared with that in RC, and thus can be updated by minimizing (20). This leads to three advantages: 1) We can achieve the target conditional entropy model to close the gap between training with $\mathbf{r}$ and inference with ${\hat{\mathbf{r}}}$; 2) We can circumvent the aforementioned relaxation problem of residual quantization; 3) We can avoid degrading the estimation of $p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ in RC caused by training with randomly generated $\tau$. Because the entropy model in RC receives the context $C_{\mathbf{r}}$ extracted from the original residual $\mathbf{r}$, $p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ approximates the true distribution $p(\mathbf{r}|{\tilde{\mathbf{x}}},{\hat{\mathbf{y}}})$ better than $p_{\boldsymbol{\varphi}}(\mathbf{r}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$. Thus, $-\log p_{\boldsymbol{\theta}}(\mathbf{r}|\mathbf{u},C_{\mathbf{r}})$ is the lower bound of $-\log p_{\boldsymbol{\varphi}}(\mathbf{r}|\mathbf{u},C_{{\hat{\mathbf{r}}}},\tau)$ on average.

Generally, autoregressive models suffer from serialized decoding and cannot be efficiently implemented on GPUs. Given an $H\times W$ image, we need to compute $HW$ times context model to decode all pixels sequentially. In lossy image compression, checkerboard context model [45] and channel-wise context model [44] were introduced to accelerate the probability inference of latent variables. However, these two context models are too weak for our residual coding and result in significant performance degradation, without the help of transform coding.

To improve the parallelization of residual coding, we first adopt a common operation to split an $H\times W$ image into multiple non-overlapping $P\times P$ patches and code all $P\times P$ patches in parallel, reducing $HW$ times sequential context computations to $P^{2}$ times. We next propose a novel design of context coding to improve the algorithm parallelization given the $P\times P$ patches and $k\times k$ mask convolution, as illustrated in Fig. 8. Assuming that $P=9$ and $k=7$, we need $P^{2}=81$ sequential decoding steps in raster scan order for the commonly used context model shown in Fig. 8a. The number in each pixel denotes the time step $t$ at which the pixel is decoded. Since $P>\lceil\frac{k}{2}\rceil$ is usually satisfied, the currently decoded pixel only depends on some of the previously decoded pixels. For example, the red pixel is decoded currently and the yellow pixels are its context. The blue pixels are previously decoded pixels but are not included in the context of the red pixel. Hence, there are pixels that can be potentially decoded in parallel by revising the scan order. By using $\frac{180}{\pi}\cdot\arctan(\frac{2}{k+1})$ degree parallel scan, the pixels with the same number $t$ can be decoded simultaneously, as shown in Fig. 8b. The number of decoding steps is reduced from $P^{2}$ to $\frac{k+3}{2}\cdot P-\frac{k+1}{2}$. In this case, we use $14.04^{\circ}$ parallel scan, leading to $5P-4=41$ sequential decoding steps. The similar scan order was also used in [61]. Moreover, we can remove one context pixel in the upper right of the currently decoded pixel. The newly designed context model leads to $\frac{180}{\pi}\cdot\arctan(\frac{2}{k-1})$ degree parallel scan and reduces decoding steps to $\frac{k+1}{2}\cdot P-\frac{k-1}{2}$. As shown in Fig. 8c, we use $18.43^{\circ}$ parallel scan and $4P-3=33$ decoding steps with this context model. When more upper-right context pixels are removed, the coding parallelization can be further improved while the compression performance is gradually compromised on. As shown in Fig. 8d, the context model leads to $\frac{180}{\pi}\cdot\arctan(\frac{2}{k-3})$ degree parallel scan and $\frac{k-1}{2}\cdot P-\frac{k-3}{2}$ decoding steps, i.e., $26.57^{\circ}$ parallel scan and $3P-2=25$ decoding steps in our example. When $45^{\circ}$ parallel scan is reached, this special case is the zig-zag scan [43] and we need $2P-1=17$ decoding steps, as shown in Fig. 8e. Finally, the fastest case is shown in Fig. 8f. We can use $90^{\circ}$ parallel scan and only $P=9$ decoding steps.

In summary, the proposed design of context coding demonstrates that: given $P\times P$ image patches and $k\times k$ mask convolution with $P>\lceil\frac{k}{2}\rceil$, we can design a series of context models $\{M^{(k+3)/2}_{k},M^{(k+1)/2}_{k},\ldots,M^{1}_{k}\}$ leading to $\{\frac{k+3}{2}\cdot P-\frac{k+1}{2},\frac{k+1}{2}\cdot P-\frac{k-1}{2},\ldots,P\}$ parallel decoding steps, by gradually adjusting the context pixels. The corresponding scan angles are $\{\frac{180}{\pi}\cdot\arctan(\frac{2}{k+1}),\frac{180}{\pi}\cdot\arctan(\frac{2}{k-1}),\ldots,90^{\circ}\}$, respectively. In experiments, we set $P=64$, $k=7$ and select the context model $M^{3}_{7}$ shown in Fig. 8d. The $M^{3}_{7}$ enjoys almost the same compression performance as $M_{7}^{5}$ in Fig. 8a and Fig. 8b, but needs much fewer coding steps.

Since the pixels of both a raw image $\mathbf{x}$ and its lossy reconstruction ${\tilde{\mathbf{x}}}$ are in the interval $[0,255]$, the element $r_{i,c}$ of the corresponding residual $\mathbf{r}$ is in the interval $[-255,255]$. To entropy coding each $r_{i,c}$, we need to compute and utilize PMF with $511$ elements, which is relatively large and slows down the entropy coding process.

In practice, the theoretical interval $[-255,255]$ of $r_{i,c}$ can hardly be filled up. Hence, we can compute and record $r_{\min}=\min_{i,c}r_{i,c}$ and $r_{\max}=\max_{i,c}r_{i,c}$ of each image, and reduce the domain of PMF to the adaptive interval $[r_{\min},r_{\max}]$ with $r_{\max}-r_{\min}+1$ elements. The overheads of recording the $r_{\min}$ and $r_{\max}$ can be amortized and ignored. For near-lossless image compression with $\tau>0$, we can similarly compute and record the quantized $\hat{r}_{\min}=\min_{i,c}\hat{r}_{i,c}$ and $\hat{r}_{\max}=\max_{i,c}\hat{r}_{i,c}$ of each image, and the domain of the quantized PMF can be further reduced to the adaptive interval $\{\hat{r}_{\min},\hat{r}_{\min}+2\tau+1,\hat{r}_{\min}+2\cdot(2\tau+1),\ldots,\hat{r}_{\max}\}$, i.e., $\frac{\hat{r}_{\max}-\hat{r}_{\min}}{2\tau+1}+1$ elements in total. The reduction of residual intervals can significantly accelerate the entropy coding on average.