Real-Time Adaptive Image Compression

Introduction

Streaming of digital media makes 70% of internet traffic, and is projected to reach 80% by 2020 (CIS, 2015). However, it has been challenging for existing commercial compression algorithms to adapt to the growing demand and the changing landscape of requirements and applications. While digital media are transmitted in a wide variety of settings, the available codecs are “one-size-fits-all”: they are hard-coded, and cannot be customized to particular use cases beyond high-level hyperparameter tuning.

In the last few years, deep learning has revolutionized many tasks such as machine translation, speech recognition, face recognition, and photo-realistic image generation. Even though the world of compression seems a natural domain for machine learning approaches, it has not yet benefited from these advancements, for two main reasons. First, our deep learning primitives, in their raw forms, are not well-suited to construct representations sufficiently compact. Recently, there have been a number of important efforts by Toderici et al. (2015, 2016), Theis et al. (2016), Ballé et al. (2016), and Johnston et al. (2017) towards alleviating this: see Section 2.2. Second, it is difficult to develop a deep learning compression approach sufficiently efficient for deployment in environments constrained by computation power, memory footprint and battery life.

In this work, we present progress on both performance and computational feasibility of ML-based image compression.

Our algorithm outperforms all existing image compression approaches, both traditional and ML-based: it typically produces files 2.5 times smaller than JPEG and JPEG 2000 (JP2), 2 times smaller than WebP, and 1.7 times smaller than BPG on the Kodak PhotoCD and RAISE-1k $512\times 768$ datasets across all of quality levels. At the same time, we designed our approach to be lightweight and efficiently deployable. On a GTX 980 Ti GPU, it takes around 9ms to encode and 10ms to decode an image from these datasets: for JPEG, encode/decode times are 18ms/12ms, for JP2 350ms/80ms and for WebP 70ms/80ms. Results for a representative quality level are presented in Table 1.

To our knowledge, this is the first ML-based approach to surpass all commercial image compression techniques, and moreover run in real-time.

We additionally supplement our algorithm with adversarial training specialized towards use in a compression setting. This enables us to produce convincing reconstructions for very low bitrates.

Background & Related Work

Compression, in general, is very closely related to pattern recognition. If we are able to discover structure in our input, we can eliminate this redundancy to represent it more succinctly. In traditional codecs such as JPEG and JP2, this is achieved via a pipeline which roughly breaks down into 3 modules: transformation, quantization, and encoding (Wallace (1992) and Rabbani & Joshi (2002) provide great overviews of the JPEG standards).

In traditional codecs, since all components are hard-coded, they are heavily engineered to fit together. For example, the coding scheme is custom-tailored to match the distribution of the outputs of the preceding transformation. JPEG, for instance, employs $8\times 8$ block DCT transforms, followed by run-length encoding which exploits the sparsity pattern of the resultant frequency coefficients. JP2 employs an adaptive arithmetic coder to capture the distribution of coefficient magnitudes produced by the preceding multi-resolution wavelet transform.

However, despite the careful construction and assembly of these pipelines, there still remains significant room for improvement of compression efficiency. For example, the transformation is fixed in place irrespective of the distribution of the inputs, and is not adapted to their statistics in any way. In addition, hard-coded approaches often compartmentalize the loss of information within the quantization step. As such, the transformation module is chosen to be bijective: however, this limits the ability to reduce redundancy prior to coding. Moreover, the encode-decode pipeline cannot be optimized for a particular metric beyond manual tweaking: even if we had the perfect metric for image quality assessment, traditional approaches cannot directly optimize their reconstructions for it.

2 ML-based lossy image compression

In approaches based on machine learning, structure is automatically discovered, rather than manually engineered.

One of the first such efforts by Bottou et al. (1998), for example, introduced the DjVu format for document image compression, which employs techniques such as segmentation and K-means clustering separate foreground from background, and analyze the document’s contents.

At a high level, one natural approach to implement the encoder-decoder image compression pipeline is to use an autoencoder to map the target through a bitrate bottleneck, and train the model to minimize a loss function penalizing it from its reconstruction. This requires carefully constructing a feature extractor and synthesizer for the encoder and decoder, selecting an appropriate objective, and possibly introducing a coding scheme to further compress the fixed-size representation to attain variable-length codes.

Many of the existing ML-based image compression approaches (including ours) follow this general strategy. Toderici et al. (2015, 2016) explored various transformations for binary feature extraction based on different types of recurrent neural networks; the binary representations were then entropy-coded. Johnston et al. (2017) enabled another considerable leap in performance by introducing a loss weighted with SSIM (Wang et al., 2004), and spatially-adaptive bit allocation. Theis et al. (2016) and Ballé et al. (2016) quantize rather than binarize, and propose strategies to approximate the entropy of the quantized representation: this provides them with a proxy to penalize it. Finally, Pied Piper has recently claimed to employ ML techniques in its Middle-Out algorithm (Judge et al., 2016), although their nature is shrouded in mystery.

3 Generative Adversarial Networks

One of the most exciting innovations in machine learning in the last few years is the idea of Generative Adversarial Networks (GANs) (Goodfellow et al., 2014). The idea is to construct a generator network $\mathscr{G}_{\boldsymbol{\Phi}}(\cdot)$ whose goal is to synthesize outputs according to a target distribution $p_{\textrm{true}}$, and a discriminator network $\mathscr{D}_{\boldsymbol{\Theta}}(\cdot)$ whose goal is to distinguish between examples sampled from the ground truth distribution, and ones produced by the generator. This can be expressed concretely in terms of the minimax problem:

This idea has enabled significant progress in photo-realistic image generation (Denton et al., 2015; Radford et al., 2015; Salimans et al., 2016), single-image super-resolution (Ledig et al., 2016), image-to-image conditional translation (Isola et al., 2016), and various other important problems.

The adversarial training framework is particularly relevant to the compression world. In traditional codecs, distortions often take the form of blurriness, pixelation, and so on. These artifacts are unappealing, but are increasingly noticeable as the bitrate is lowered. We propose a multiscale adversarial training model to encourage reconstructions to match the statistics of their ground truth counterparts, resulting in sharp and visually pleasing results even for very low bitrates. As far as we know, we are the first to propose using GANs for image compression.

Model

Our model architecture is shown in Figure 2, and comprises a number of components which we briefly outline below. In this section, we limit our focus to operations performed by the encoder: since the decoder simply performs the counterpart inverse operations, we only address exceptions which require particular attention.

Images feature a number of different types of structure: across input channels, within individual scales, and across scales. We design our feature extraction architecture to recognize these. It consists of a pyramidal decomposition which analyzes individual scales, followed by an interscale alignment procedure which exploits structure shared across scales.

This module is responsible for further compressing the extracted features. It quantizes the features, and encodes them via an adaptive arithmetic coding scheme applied on their binary expansions. An adaptive codelength regularization is introduced to penalize the entropy of the features, which the coding scheme exploits to achieve better compression.

We employ adversarial training to pursue realistic reconstructions. We dedicate Section 4 to describing our GAN formulation.

1 Feature extraction

Our pyramidal decomposition encoder is loosely inspired by the use of wavelets for multiresolution analysis, in which an input is analyzed recursively via feature extraction and downsampling operators (Mallat, 1989). The JPEG 2000 standard, for example, employs discrete wavelet transforms with the Daubechies 9/7 kernels (Antonini et al., 1992; Rabbani & Joshi, 2002). This transform is in fact a linear operator, which can be entirely expressed via compositions of convolutions with only two hard-coded and separable $9\times 9$ filters applied irrespective of scale, and independently for each channel.

The idea of a pyramidal decomposition has been employed in machine learning: for instance, Mathieu et al. (2015) uses a pyramidal composition for next frame prediction, and Denton et al. (2015) uses it for image generation. The spectral representations of CNN activations have also been investigated by Rippel et al. (2015) to enable processing across a spectrum of scales, but this approach does not enable FIR processing as does wavelet analysis.

Our pyramidal decomposition architecture is illustrated in Figure 3. In practice, we extract across a total of $M=6$ scales. The feature extractors for the individual scales are composed of a sequence of convolutions with kernels $3\times 3$ or $1\times 1$ and ReLUs with a leak of $0.2$. We learn all downsamplers as $4\times 4$ convolutions with a stride of 2.

1.2 Interscale alignment

2 Code computation and regularization

2.1 Quantization

For the special case $B=1$, this reduces exactly to a binary quantization scheme. While some ML-based approaches to compression employ such thresholding, we found better performance with the smoother quantization described. We quantize with $B=6$ for all models in this paper.

2.2 Bitplane decomposition

2.3 Adaptive arithmetic coding

During decoding, we decompress the code by performing the inverse operation. Namely, we interleave between computing the context of a particular bit using the values of previously decoded bits, and using this context to retrieve the activation probability of the bit and decode it. We note that this constrains the context of each bit to only include features composed of bits already decoded.

2.4 Adaptive codelength regularization

One problem with classic autoencoder architectures is that their bottleneck has fixed capacity. The bottleneck may be too small to represent complex patterns well, which affects quality, and it may be too large for simple patterns, which results in inefficient compression. What we need is a model capable of generating long representations for complex patterns and short for simple ones, while maintaining an expected codelength target over large number of examples. To achieve this, the AAC is necessary, but not sufficient.

for iteration $t$ and difference index set $S=\{(0,1),(1,0),(1,1),(-1,1)\}$. The first term penalizes the magnitude of each tensor element, and the second penalizes deviations between spatial neighbors. These enable better prediction by the AAC.

As we train our model, we continuously modulate the scalar coefficient $\alpha_{t}$ to pursue our target codelength. We do this via a feedback loop. We use the AAC to monitor the mean number of effective bits. If it is too high, we increase $\alpha_{t}$; if too low, we decrease it. In practice, the model reaches an equilibrium in a few hundred iterations, and is able to maintain it throughout training.

Realistic Reconstructions via Multiscale Adversarial Training

In our compression approach, we take the generator as the encoder-decoder pipeline, to which we append a discriminator — albeit with a few key differences from existing GAN formulations.

In many GAN approaches featuring both a reconstruction and a discrimination loss, the target and the reconstruction are treated independently: each is separately assigned a label indicating whether it is real or fake. In our formulation, we consider the target and its reconstruction jointly as a single example: we compare the two by asking which of the two images is the real one.

To do this, we first swap between the target and reconstruction in each input pair to the discriminator with uniform probability. Following the random swap, we propagate each set of examples through the network. However, instead of producing an output for classification at the very last layer of the pipeline, we accumulate scalar outputs along branches constructed along it at different depths. We average these to attain the final value provided to the terminal sigmoid function. This multiscale architecture allows aggregating information across different scales, and is motivated by the observation that undesirable artifacts vary as function of the scale in which they are exhibited. For example, high-frequency artifacts such as noise and blurriness are discovered by earlier scales, whereas more abstract discrepancies are found in deeper scales.

We apply our discriminator $\mathscr{D}_{\boldsymbol{\Theta}}$ on the aggregate sum across scales, and proceed to formulate our objectives as described in Section 2.3. The complete discriminator architecture is illustrated in Figure 10.

2 Adversarial training

Training a GAN system can be tricky due to optimization instability. In our case, we were able to address this by designing a training scheme adaptive in two ways. First, the reconstructor is trained by both the confusion signal gradient as well as the reconstruction loss gradient: we balance the two as function of their gradient magnitudes. Second, at any point during training, we either train the discriminator or propagate confusion signal through the reconstructor, as function of the prediction accuracy of the discriminator.

More concretely, given lower and upper accuracy bounds $L,U\in$ and discriminator accuracy $a(\mathscr{D}_{\boldsymbol{\Theta}})$, we apply the following procedure:

If $a<L$: freeze propagation of confusion signal through the reconstructor, and train the discriminator.

If $L\leq a<U$: alternate between propagating confusion signal and training the disciminator.

If $U\leq a$: propagate confusion signal through the reconstructor, and freeze the discriminator.

In practice we used $L=0.8,U=0.95$. We compute the accuracy $a$ as a running average over mini-batches with a momentum of $0.8$.

Results

Since the human visual system is much more sensitive to variations in brightness than color, most codecs represent colors in the YCbCr color space to devote more bandwidth towards encoding luma rather than chroma. In quantifying image similarity, then, it is common to assign the Y, Cb, Cr components weights $6/8,1/8,1/8$. While many ML-based compression papers evaluate similarity in the RGB space with equal color weights, this does not allow fair comparison with standard codecs such as JPEG, JPEG 2000 and WebP, since they have not been designed to perform optimally in this domain. In this work, we provide comparisons with both traditional and ML-based codecs, and present results in both the RGB domain with equal color weights, as well as in YCbCr with weights as above.

We present both compression performance of our algorithm, but also its runtime. While the requirement of running the approach in real-time severely constrains the capacity of the model, it must be met to enable feasible deployment in real-life applications.

We trained and tested all models on a GeForce GTX 980 Ti GPU and a custom codebase. We trained all models on $128\times 128$ patches sampled at random from the Yahoo Flickr Creative Commons 100 Million dataset (Thomee et al., 2016).

2 Performance

Average MS-SSIM as function of the BPP fixed for each image, found in Figures 5 and 6, and Table 1.

Average compressed file sizes relative to ours as function of the MS-SSIM fixed for each image, found in Figures 5 and 6, and Table 1.

Encode and decode timings as function of MS-SSIM, found in Figure 7, in the appendix, and Table 1.

Visual examples of reconstructions of different compression approaches for the same BPP, found in Figure 1 and in the appendix.

To enable comparison with other approaches, we first present performance on the Kodak PhotoCD dataset5.1. While the Kodak dataset is very popular for testing compression performance, it contains only 24 images, and hence is susceptible to overfitting and does not necessarily fully capture broader statistics of natural images. As such, we additionally present performance on the RAISE-1k dataset (Dang-Nguyen et al., 2015) which contains 1,000 raw images. We resized each image to size $512\times 768$ (backwards if vertical): we intend to release our preparation code to enable reproduction of the dataset used.

We remark it is important to use a dataset of raw, rather than previously compressed, images for codec evaluation. Compressing an image introduces artifacts with a bias particular to the codec used, which results in a more favorable RD curve if it compressed again with the same codec. See Figure 9 for a plot demonstrating this effect.

We compare against commercial compression techniques JPEG, JPEG 2000, WebP, as well as recent ML-based compression work by Toderici et al. (2016)5.1, Theis et al. (2016)5.1, Ballé et al. (2016)5.1, and Johnston et al. (2017)5.1 in all settings in which results are available. We also compare to BPG5.1 (4:2:0 and 4:4:4) which, while not widely used, surpassed all other codecs in the past. We use the best-performing configuration we can find of JPEG, JPEG 2000, WebP, and BPG, and reduce their bitrates by their respective header lengths for fair comparison.

For each image in each test set, each compression approach, each color space, and for the selection of available compression rates, we recorded (1) the BPP, (2) the MS-SSIM (with components weighted appropriately for the color space), and (3) the computation times for encoding and decoding.

It is important to take great care in the design of the performance evaluation procedure. Each image has a separate RD curve computed from all available compression rates for a given codec: as Ballé et al. (2016) discusses in detail, different summaries of these RD curves lead to disparate results. In our evaluations, to compute a given curve, we sweep across values of the independent variable (such as bitrate). We interpolate each individual RD curve at this independent variable value, and average all the results. To ensure accurate interpolation, we sample densely across rates for each codec.

We are grateful to Trevor Darrell, Sven Strohband, Michael Gelbart, Robert Nishihara, Albert Azout, and Vinod Khosla for meaningful discussions and input.

References