Channel-wise Autoregressive Entropy Models for Learned Image Compression

Introduction

Most recent research in learned image compression uses deep neural networks, and a wide range of model architectures have been explored including recurrent networks and autoencoders with an entropy-constrained bottleneck . In models that use an autoencoder, an analysis network transforms pixels into a quantized latent representation suitable for compression by standard entropy coding algorithms, while a synthesis network is jointly optimized to transform the latent representation back into pixels.

To date, the most effective models make use of both forward and backward-adaptive components to improve the predictive power of the entropy model, which leads to higher compression rates without increasing distortion. Forward-adaption typically makes use of side information, for example in the form of local histograms over the quantized latent representation or a learned hyperprior . The hyperprior approach is particularly popular since it can easily be integrated into an end-to-end optimized network and allows for efficient encoding and decoding.

Backward-adaptation, on the other hand, typically incorporates predictions from the causal context of each symbol, i.e. neighboring symbols above and to the left of the current symbol as well as symbols in previously decoded channels . In such context-adaptive models, encoding can still be performed efficiently using masked convolution, which will run in parallel across the entire latent tensor on a GPU or TPU . Decoding, however, is inherently serial, and thus does not effectively utilize massively parallel hardware.

Our goal is to develop an image compression architecture capable of matching the rate-distortion (RD) performance of a context-adaptive model while minimizing serial processing that can lead to slow decoding times. Toward this goal, we explore two architectural enhancements: channel-conditioning (CC) and latent residual prediction (LRP). In addition, we show how training synthesis transforms with rounded latent values interacts positively with CC and LRP to further boost RD performance.

The combined effect of these improvements is a highly parallelizable architecture that outperforms recently proposed context-adaptive models by 6.7% on Kodak and 11.4% on the Tecnick image set . We see even larger gains compared to standard codecs and learning-based models that do not use context (see Figures 2 and 3). The coding improvements provided by CC and LRP are most effective at low bit rates where our model saves more than 16% compared to the context-adaptive baseline and as much as 25% relative to BPG . The following three sections describe channel-conditioning, latent residual prediction, and round-based training. A detailed analysis of the empirical results is presented in Section 5 and discussed in Section 6.

Channel-Conditional Entropy Models

Our model builds on the hyperprior architecture introduced in . This model learns to generate an image-dependent hyper-latent tensor that is compressed and transmitted as side information. It jointly learns to transform this tensor into the entropy parameters used to compress the symbols that represents the input image (see the Hyperprior block at the right of Figure 1). Hyperprior models typically use a conditional Gaussian model parameterized by scale or both scale and mean, and the most effective models combine information from the hyperprior (forward-adaptation) with a spatially autoregressive model (backward-adaptation) before predicting the entropy parameters $\mu$ and $\sigma$ .

Conditioning on the causal context allows for better modeling of spatial correlation and is commonly used in standard image codecs and for intra-frame prediction in video codecs [HEVC, 23, 24]. In a learning-based codec, the model must estimate the parameters of a spatially autoregressive (AR) model. This approach is effective but requires running the AR model sequentially to decode each symbol, which can slow down decoding times on GPUs and TPUs compared to architectures that better utilize the massively parallel processing abilities of such hardware. For this reason, we explore channel-conditional (CC) models, which split the latent tensor along the channel dimension into $N$ roughly equal-size slices, and conditions the entropy parameters for each slice on previously decoded slices.

Figure 1 provides a high-level overview of this architecture where the blue arrows show how $y_{2}$ (the second slice) is conditioned on $\hat{y}_{1}$ (the first slice). In a model with more splits, the third slice ($y_{3}$) would be conditioned on the hyperprior along with both $\hat{y}_{1}$ and $\hat{y}_{2}$, etc.

We can interpret CC models as autoregressive along the channel dimension rather than the spatial dimensions. Although this structure also introduces some serial processing (slice $y_{i}$ can only be decoded after slices $[y_{1}\ldots y_{i-1}]$), we typically use relatively few slices due to diminishing benefits to RD performance (see Figure 4). Note that in a model with $N$ slices, each slice contains $W\times H\times\frac{C}{N}$ values that can be processed in parallel (where $W$, $H$, and $C$ correspond to the width, height and number of channels, respectively). Contrast with a spatially autoregressive model where a naive implementation requires $W\times H$ sequential steps with only $C$ values computed during each run. A more careful implementation using wavefront processing adds some parallelization but still far less than channel-conditioning.

Latent Residual Prediction

Autoencoder models learn to transform pixel values ($x$) into real-valued latents ($y$) that are quantized before they are losslessly compressed. This process inevitably leads to a residual error in the latent space ($r=y-Q[y]$) that manifests as extra distortion when $Q[y]$ is transformed back into the pixel domain ($\hat{x}$).

Latent residual prediction attempts to reduce this quantization error by predicting the residual based on the hyperprior and any previously decoded slices. The predicted residual is added to the quantized latents slice-by-slice, which allows LRP to improve results both by decreasing distortion and by decreasing entropy since the entropy parameters used to code later slices are conditioned on previous ones that include LRP.

Previous approaches for augmenting the input to the synthesis transform either re-used the mean prediction directly or used dilated convolution to provide additional features based on a larger receptive field . In both cases, however, the extra features were concatenated with the latent tensor, which increases computation, and neither used channel-conditioning, which means that potential improvements could only affect distortion.

Training with Rounded Latent Values

All compression models trained using gradient-based optimization are hindered by quantization, which yields gradients that are either zero or infinite at all values. Typically, researchers avoid this problem by either training with uniform noise, which simulates “noisy quantization” without destroying the gradient , or they use straight-through gradients where rounding is applied but the true gradient function is replaced with the identity function .

Although space constraints preclude a full report on the effects of different training methods, we empirically found that a mixed approach improves RD performance. Our baseline models replace quantization with uniform noise during training: $Q[y]\doteq y+\mathcal{U}(-\frac{1}{2},\frac{1}{2})$. The mixed approach uses the same uniform noise for learning entropy models but replaces the noisy tensor with a rounded one whenever the quantized tensor is passed to a synthesis transform. Looking at Figure 1, the difference is essentially whether the quantized tensor is flowing to the right (add noise) or left (round with straight-through gradients). We experimented with using the rounding-based method everywhere, but this approach performed worse than the noise-based baseline.

Experimental Results

In this section, we evaluate the effects of using CC, LRP, and round-based training in a learned image codec. Figure 2 compares RD curves averaged over the Kodak image set . The graph shows that our full model (10 CC slices + LRP + round-based training) outperforms all of the standard codecs (BPG, JPEG2000, WebP, and JPEG) as well as learning-based codecs that combine spatial context with a hyperprior . To improve clarity, earlier learning-based methods, including , are not shown in Figure 2, but all of these methods have worse RD performance than both BPG and our CC + LRP model.

Additional results are shown in Figure 3, which plots the relative rate savings compared to BPG at different quality levels. Larger values correspond to larger relative rate savings and thus better compression. This graph generalizes a Bjøntegaard Delta (BD) chart by plotting rate savings as a function of quality, rather than only presenting the average savings. Our largest model, which uses 10 CC slices, provides a significant rate savings over BPG, ranging from 10% at higher quality levels up to 25% at the lowest. This corresponds to an average BD rate savings of 13.9% over BPG and 6.7% over the context-adaptive baseline . The following sections analyze how each proposed improvements contributes to the final result.

Figure 4 shows the average rate savings as the number of channel-conditioning slices increases. When we split the latent tensor into more slices, there are more opportunities to model the dependencies between channels, which reduces entropy. This benefit, however, comes at the cost of extra computation, and we also see diminishing returns as the number of slices increases.

2 Latent Residual Prediction

Figure 5 shows the effect of LRP for different numbers of channel-conditioning splits. Each curve compares a model trained with LRP to an identical model without LRP by plotting the relative rate savings when LRP is used.

The figure shows several effects. First, LRP has almost no benefit for models that do not use channel-conditioning, which we can see because the blue “CC(0 splits)” curve is always close to zero. Second, regardless of the number of CC splits, LRP slightly reduces RD performance at high bit rates. At low bit rates, however, the benefit of LRP increases with the number of CC slices and improves compression by more than 6% for the model with 10 splits.

3 Rounding-based Optimization

Figure 6 shows the impact of mixed training with noise and round-based handling of quantized tensors as described in Section 4. The figure shows results for two CC models (zero and five splits) and plots both variants with and without LRP. Each curve shows the rate savings relative to an identical model optimized using uniform-noise everywhere, which means that the rate savings are due entirely to the change in how quantization is handled. We see the same trend in all cases: the benefit is minimal at higher quality levels but becomes significant at lower bit rates. For the “CC (5 splits) + LRP” model, the savings exceed 15% at the lowest bit rates.

Discussion

From a theoretical perspective, the positive results from both CC and LRP are somewhat surprising. Ideally, the optimization process should expand the range of each channel to balance the rate-distortion trade-off, which means that using additional bits in the hyperprior to drive LRP would not be helpful. Essentially, channels that significantly reduce distortion would use more symbols, which can be interpreted as finer precision, e.g. consider a channel that uses values $$vs. one that uses$[-100,-99,\ldots,99,100]$and is scaled by$\frac{1}{100}$ in the next convolutional layer. Since the most useful channels should already have higher effective granularity, there is less opportunity for LRP to provide a benefit.

Similarly, the analysis transform would ideally learn to map pixels into a latent space such that each channel is conditionally independent given the hyperprior. If this is not the case, it means there is redundant information, which will increase entropy without reducing distortion.

Empirically, we see significant improvements using both CC and LRP, which implies that existing models are far from ideal. Further research is needed to understand why the models are failing to reach an optimal state, but we can theorize that the relatively simple 4-layer convolutional networks that make up the analysis and synthesis transforms lack the capacity to generate/decode a latent representation with conditionally independent channels. Alternatively, the networks may have the necessary capacity, but our learning procedure, which uses the Adam optimizer , is unable to find a suitable minimum despite training for five million steps.

By combining channel-conditioning, latent residual prediction, and round-based training, we have developed a neural image compression architecture that outperforms a corresponding context-adaptive model while minimizing serial processing. In future research, we plan to investigate combining channel-conditioning with spatial context modeling to see if the two approaches are complementary.

Appendix A Architecture Details

Figure 1 provides a high-level overview of the network architecture for the channel-conditional model. Details about the individual transforms and layer configurations are missing due to space constraints. To facilitate reproducibility, Table 1 provides detailed layer specifications for all of the transforms in a model with 10 slices.

The channel-conditional (CC) and latent residual prediction (LRP) transforms are trained separately for each slice. In each case, the output depth will be the same: latent_depth / num_slices Typically, we use architectures where num_slices evenly divides latent_depth. If it doesn’t, all slices have floor(latent_depth / num_slices) except for the final slice, which is set to the remaining number of channels: slice_depthn = latent_depth - $\sum_{i=1}^{n-1}$ slice_depthi for a model with $n$ slices.. The input depth, however, will vary since the input to later slices include the concatenation of all previous slices. In our example, the input to slice1 is 320, the input to slice2 is 352, and the input to slice10 is 608 ($320+32\times 9$). To account for the different input depths, each CC and LRP transform is programmatically defined to linearly interpolate between the input and the output depth. For example, the tenth slice will have depths: 416, 224 and 32. Finally, note that the LRP transform includes the decoded values from the current block, whereas the CC transforms that predict $\mu_{i}$ and $\sigma_{i}$ values do not since it’s not yet available. The input depth for the LRP transforms are thus larger than the input depths for the CC transforms by slicei channels, i.e. 32 extra channels in the example model.

At low and moderate bit rates, we found that a channel depth of 320 in the latent tensor (the output of the analysis transform) yielded good rate-distortion performance. For high bit rates, typically above 2.0 bpp on the Kodak image set , a larger bottleneck boosts RD performance. For all reported results, we used 512 channels for such high bit rate models.

Finally, we use a simplified version of generalized divisive normalization (GDN) where $\alpha_{ij}$ and $\varepsilon_{i}$ are both set to 1.0. The full formula for GDN is:

where $x_{i}$ and $z_{i}$ denote the input and output vectors, respectively, $\alpha_{ij}$, $\beta_{i}$, $\gamma_{ij}$, and $\varepsilon_{i}$ represent trainable parameters, and $i$, $j$ represent channel indices. By fixing $\alpha_{ij}$ and $\varepsilon_{i}$ to 1.0, the simplified formula becomes:

This change leads to slightly faster and more stable training without reducing RD performance .

Appendix B Training Details

For the experimental results in this paper, all models were trained for 5,000,000 steps using the Adam optimizer with $\beta_{1}$ = 0.9, $\beta_{2}$ = 0.999, and $\epsilon$ = 1e-8. The learning rate started at 1e-4 and dropped to 3e-5 at 3M steps, 1e-5 at 3.6M steps, 3e-6 at 4.2M steps, and 1e-6 at 4.8M steps. Later experimentation found that a better learning rate schedule would improve RD performance by 1-2%, even when the training duration was reduced to 4M total steps.

In addition to adjusting the learning rate, the rate-distortion trade-off parameter, $\lambda$, is also adjusted. For all experiments presented here, a model targeting $R+\lambda\cdot D$ is trained using $2\cdot\lambda$ for the first 2.5M steps (half of the total training time). The loss function is then adjusted to use the target $\lambda$. Training with a higher $\lambda$ encourages lower distortion and thus a higher bit rate. This appears to help low bit rate models avoid a sub-optimal entropy model, but more analysis and experimentation is needed to understand exactly why this occurs and how to optimally adjust $\lambda$ during training.

All models are trained on the same images modulo random shuffling and patch extraction. The image set is made up of nearly 2M web images filtered for resolution and compression quality. The models are trained using a batch size of eight with $256\times 256$ patches randomly cropped from the input images after random downscaling. The downscaling is useful both to reduce pre-existing compression artifacts in the training data and to discourage overfitting to a particular scale. We found that smaller patches reduced RD performance, while larger patches provided little benefit relative to slower training time.

Appendix C Additional Experiments

In all of the channel-conditional models explored in this paper, each slice is conditioned on all previous slices. Thus, for example, slice10 is conditioned on nine previous slices (slice1– slice9), which leads to a relatively large input depth for the later slices and thus slower models. Figure 7 shows how RD performance degrades as slices are conditioned on fewer previous slices. For example, in a 10 slice model, if the slice support is five, the final slice will only be conditioned on the first five slices and will be conditionally independent from slice6– slice9. Reducing the slice support decreases the size of the channel-conditional transform for later slices and creates more opportunity for parallel calculations, both of which improve runtime.

In addition to exploring the effect of conditioning on a limited number of early slices (slice1– sliceN), we also explored conditioning on the previous N slices. For example, if conditioning on five slices, slice7 would be conditioned on slice2–slice6. This approach led to slightly worse RD performance and provides fewer opportunities for parallelization compared to always conditioning on the first N slices.

Figure 7 shows that reducing the slice support does reduce RD performance, but the reduction is relatively small. Further research is needed to fully understand how the latent representation in the early slices changes, but our theory is that most of the benefit of channel-conditioning comes from a relatively small amount of high-level information. Models that condition on fewer slices learn to represent this information in the early slices, thus preserving the overall effectiveness of the model.

C.2 Smaller Hyperprior and μ′,σ′\mu^{\prime},\sigma^{\prime} Tensors

In channel-conditional models, all slices are conditioned on previous slices as well as a latent tensor predicted from the hyperprior ($\mu^{\prime}$ and $\sigma^{\prime}$ in Figure 1). Each slice has its own transform for predicting the mean and scale values in that slice ($\mu_{i}$ and $\sigma_{i}$), and the size and speed of the transform depends on the size of $\mu^{\prime}$ and $\sigma^{\prime}$. All of the models in the main paper use 320 channels for these tensors, which is relatively large. We therefore explored the effect on RD performance as the depth of these tensors is reduced (see Figure 8). The experiment confirms that larger tensors help, at least up to 320 channels, but the RD penalty for shrinking the tensor is minimal (~2% in the worst case) down to a depth of 160 channels.

Further experiments explored the effect of shrinking the hyper-analysis transform. Our typical model transforms a 320 channel latent into a 192 channel hyperprior (see the Hyper-Analysis column in Table 1). Empirically, we found that shrinking the hyper-analysis transform from layers with depths of to depths of had only a small impact on RD performance.

Appendix D Rate-Distortion Comparisons

Figure 2 in the main paper provides a rate-distortion comparison between our method and a small set of recent learning-based methods as well as several standard methods. Figure 14 extends this comparison by including many more compression methods and providing a larger graph to aid readability.

Figure 9 shows the same data in a format that highlights the rate savings of different compression methods relative to JPEG (4:2:0). In this case, the curves represent the file size reduction as a percentage of the size of a JPEG encoding with equivalent PSNR. Larger values imply higher savings, and the graph shows that our method can shrink images by more than 70% compared to JPEG at lower quality levels and by more than 55% at all quality levels.

Appendix E Reconstructed Images

Reconstructed images are shown in Figures 15, 16, and 17 comparing BPG with our channel-conditional model optimized for three different metrics (MSE, L1, and MS-SSIM). These example images were compressed at very high rates (ranging from roughly 250x to nearly 360x compression) since high compression rates help highlight the kinds of distortions typical for each method and quality metric. For example, optimizing for MS-SSIM typically preserves texture better than MSE or L1, e.g. in the grass in Figures 15 and 17 and the red sweater in Figure 16. MS-SSIM performs the worst, however, on high-contrast and high-frequency content like text, as shown in Figure 15, and in the receding fence in Figure 17.

This paper focuses on a more effective and more efficient entropy model based on channel-conditioning (CC) and latent residual prediction (LRP) compared to spatially autoregressive models or solely hierarchical priors. The benefits of the CC and LRP model are independent of the loss function used to quantify visual distortions. Interesting future research could look at combining our entropy model with more sophisticated image quality metrics such as perceptual metrics and adversarial loss .

Appendix F Sampling from the Compression Model

Since the compression model learns a factorized distribution over the hyperprior as well as a conditional distribution over the latents, we can treat the network as a generative model and sample random images. Doing so provides a visualization of what kind of images are typical according to the image distribution learned by the model. See Figure 11 for several examples for a CC(10) model.

The sampling process is straightforward. First, draw a sample from the factorized entropy model to generate a random hyperprior. Next, run the hyper-synthesis transform to get $\mu^{\prime}$ and $\sigma^{\prime}$. Then iterate over each slice of the latent tensor. For each slicei, run the $\mu_{i}$ and $\sigma_{i}$ transforms to get an entropy model for the slice and draw a random sample. After all slices have been sampled, concatenate the result and run the synthesis transform to generate an RGB image.

To better understand the sampled images in the context of learned image compression, Figure 12 shows typical samples drawn from six different architectures. All of the models are fully convolutional and use strided convolution to reduce the spatial extent of the data in deeper layers of the network. This downscaling creates block artifacts in the sampled images since all of the models use a factorized entropy model at the highest level (i.e. the samples are spatially i.i.d. at this level). For example, the Spatially Local Histograms approach uses 16x downscaling to form a latent representation and models each $13\times 13$ tile in latent space using its own histogram. This structure is evident in the sampled image since we can see the large blocks (corresponding to the local histograms) and the smaller blocks within each tile (corresponding to 16x downscaling). Only the Spatial AR model is capable of long-range spatial dependencies, as is visible in the sample.

The channel-wise AR model is more limited than the spatially AR model in terms of long-range coherence. Nonetheless, the channel-conditional structure allows for a much larger receptive field than any of the earlier models since the receptive field can grow with each success slice of the latent representation.

All of the models visualized here were optimized for mean squared error (MSE). More research is needed to determine if the lack of any obvious semantic information in the samples is due to this simple loss function or if it’s primarily due to insufficient capacity. Typically, generative models are optimized over constrained domains (faces, city views, flowers, bedrooms, etc.) and often use much larger networks than what is explored here. Interesting future work includes training our channel-wise AR model on a constrained domain and integrating an adversarial loss to see how each change affects the sampled images.

F.2 Samples with a Known Hyperprior

To better understand what information is represented in the hyperprior, we can partially sample from the compression model. In this case, a real image is encoded and the resulting hyperprior is saved. We then repeat the sampling process described above but use this known hyperprior instead of a random one. Figure 13 shows such partial samples based on the lighthouse image (kodim19) from the Kodak image set . In this example, the hyperprior requires 0.0143 bpp, which is an extremely compact representation (nearly 1680x compression) compared to typical rates used for image compression.

Two sampling approaches are explored. In the first (Figure 13b), random samples are drawn from the entropy model conditioned on the real hyperprior. In the second approach (Figure 13c), the mode of the conditional entropy model is used to form the latent tensor. Since the compression model uses a conditional Gaussian distribution and thus the mode is located at the mean, the difference between the methods is whether we sample from the Gaussian at each location or if we use the predicted mean.

From the images generated by this partial sampling procedure, we see that the hyperprior stores low frequency color data along with some texture and orientation information. For example, the strong vertical components in the fence are visible as well as the horizontal components in the slats on the building on the left side of the image. The hyperprior also stores a small amount of texture information as demonstrated by the different high-frequency patterns in the sky compared to the grass or the stone lighthouse.

Appendix G Progressive Decoding

Although the primary purpose of our channel-wise autoregressive model is to improve entropy coding, the structure naturally lends itself to progressive decoding. A rough image can already be displayed after the hyperprior is transferred as shown in Figure 13c. This reconstruction can then be improved after each slice in the latent space is decoded. Figure 18 shows the sequence of reconstructions recovered from a 10-slice model where the synthesis transform is executed after each slice is decoded, and missing latent values use the mode of the conditional distribution inferred from the hyperprior.

While progressive decoding is possible and easy to achieve using our channel-conditional model, two issues make it fairly impractical. First, the rate-distortion curve implied by the progressively decoded images is much worse than the result from separately optimized models (see Figure 10). This means that progressive decoding may be useful as a temporary preview, but it does not lead to an effective multi-rate model. Second, the computational cost is fairly high since the full synthesis transform must be run to generate each image. A more useful model would reduce the computational cost along with the bit rate to support progressive previews.

References