Practical Full Resolution Learned Lossless Image Compression

Introduction

Since likelihood-based discrete generative models learn a probability distribution over pixels, they can in theory be used for lossless image compression . However, recent work on learned compression using deep neural networks has solely focused on lossy compression . Indeed, the literature on discrete generative models has largely ignored the application as a lossless compression system, with neither bitrates nor runtimes being compared with classical codecs such as PNG , WebP , JPEG2000 , and FLIF . This is not surprising as (lossless) entropy coding using likelihood-based discrete generative models amounts to a decoding complexity essentially identical to the sampling complexity of the model, which renders many of the recent state-of-the-art autoregressive models such as PixelCNN , PixelCNN++ , and Multiscale-PixelCNN impractical, requiring minutes or hours on a GPU to generate moderately large images, typically ${<}256\times 256$px (see Table 2). The computational complexity of these models is mainly caused by the sequential nature of the sampling (and thereby decoding) operation, where a forward pass needs to be computed for every single (sub) pixel of the image in a raster scan order.

In this paper, we address these challenges and develop a fully parallelizeable learned lossless compression system, outperforming the popular classical systems PNG, WebP and JPEG2000.

Our system (see Fig. 1 for an overview) is based on a hierarchy of fully parallel learned feature extractors and predictors which are trained jointly for the compression task. Our code is available onlinegithub.com/fab-jul/L3C-PyTorch. The role of the feature extractors is to build an auxiliary hierarchical feature representation which helps the predictors to model both the image and the auxiliary features themselves. Our experiments show that learning the feature representations is crucial, and heuristic (predefined) choices such as a multiscale RGB pyramid lead to suboptimal performance.

In more detail, to encode an image $x$, we feed it through the $S$ feature extractors $E^{(s)}$ and predictors $D^{(s)}$. Then, we obtain the predictions of the probability distributions $p$, for both $x$ and the auxiliary features $z^{(s)}$, in parallel in a single forward pass. These predictions are then used with an adaptive arithmetic encoder to obtain a compressed bitstream of both $x$ and the auxiliary features (Sec. 3.1 provides an introduction to arithmetic coding). However, the arithmetic decoder now needs $p$ to be able to decode the bitstream. Starting from the lowest scale of auxiliary features $z^{(S)}$, for which we assume a uniform prior, $D^{(S)}$ obtains a prediction of the distribution of the auxiliary features of the next scale, $z^{(S-1)}$, and can thus decode them from the bitstream. Prediction and decoding is alternated until the arithmetic decoder obtains the image $x$. The steps are visualized in Fig. A4.

In practice, we only need to use $S=3$ feature extractors and predictors for our model, so when decoding we only need to perform three parallel (over pixels) forward passes in combination with the adaptive arithmetic coding.

The parallel nature of our model enables it to be orders of magnitude faster for decoding than autoregressive models, while learning enables us to obtain compression rates competitive with state-of-the-art engineered lossless codecs.

In summary, our contributions are the following:

We propose a fully parallel hierarchical probabilistic model, learning both the feature extractors that produce an auxiliary feature representation to help the prediction task, as well as the predictors which model the joint distribution of all variables (Sec. 3).

We show that entropy coding based on our non-autoregressive probabilistic model optimized for discrete log-likelihood can obtain compression rates outperforming WebP, JPEG2000 and PNG, the latter by a large margin. We are only marginally outperformed by the state-of-the-art, FLIF, while being conceptually much simpler (Sec. 5.1).

At the same time, our model is practical in terms of runtime complexity and orders of magnitude faster than PixelCNN-based approaches. In particular, our model is $5.31\cdot 10^{4}\times$ faster than PixelCNN++ and $5.06\cdot 10^{2}\times$ faster than the highly speed-optimized MS-PixelCNN (Sec. 5.2).

Related Work

As previously mentioned, essentially all likelihood-based discrete generative models can be used with an arithmetic coder for lossless compression. A prominent group of models that obtain state-of-the-art performance are variants of the auto-regressive PixelRNN/PixelCNN . PixelRNN and PixelCNN organize the pixels of the image distribution as a sequence and predict the distribution of each pixel conditionally on (all) previous pixels using an RNN and a CNN with masked convolutions, respectively. These models hence require a number of network evaluations equal to the number of predicted sub-pixelsA RGB “pixel” has 3 “sub-pixels”, one in each channel. ($3\cdot\!W\!\cdot\!H$). PixelCNN++ improves on this in various ways, including modeling the joint distribution of each pixel, thereby eliminating conditioning on previous channels and reducing to $W\!\cdot\!H$ forward passes. MS-PixelCNN parallelizes PixelCNN by reducing dependencies between blocks of pixels and processing them in parallel with shallow PixelCNNs, requiring $O(\log WH)$ forward passes. equips PixelCNN with auxiliary variables (grayscale version of the image or RGB pyramid) to encourage modeling of high-level features, thereby improving the overall modeling performance. propose autoregressive models similar to PixelCNN/PixelRNN, but they additionally rely on attention mechanisms to increase the receptive field.

The well-known PNG operates in two stages: first the image is reversibly transformed to a more compressible representation with a simple autoregressive filter that updates pixels based on surrounding pixels, then it is compressed with the deflate algorithm . WebP uses more involved transformations, including the use of entire image fragments to encode new pixels and a custom entropy coding scheme. JPEG2000 includes a lossless mode where tiles are reversibly transformed before the coding step, instead of irreversibly removing frequencies. The current state-of-the-art (non-learned) algorithm is FLIF . It relies on powerful preprocessing and a sophisticated entropy coding method based on CABAC called MANIAC, which grows a dynamic decision tree per channel as an adaptive context model during encoding.

In lossy compression, context models have been studied as a way to efficiently losslessly encode the obtained image representations. Classical approaches are discussed in . Recent learned approaches include , where shallow autoregressive models over latents are learned. presents a model somewhat similar to L3C: Their autoencoder is similar to our fist scale, and the hyper encoder/decoder is similar to our second scale. However, since they train for lossy image compression, their autoencoder predicts RGB pixels directly. Also, they predict uncertainties $\sigma$ for $z^{(1)}$ instead of a mixture of logistics. Finally, instead of learning a probability distribution for $z^{(2)}$, they assume the entries to be i.i.d. and fit a univariate non-parametric density model, whereas in our model, many more stages can be trained and applied recursively.

Method

is the resulting expected (sub-optimal) bits per symbol, and is called cross-entropy .

The following sections describe how a hierarchical causal factorization of $p_{\text{img}}$ for natural images can be used to efficiently do learned lossless image compression (L3C).

2 Architecture

A high-level overview of the architecture is given in Fig. 1, while Fig. 2 shows a detailed description for one scale $s$. Unlike autoregressive models such as PixelCNN and PixelRNN, which factorize the image distribution autoregressively over sub-pixels $x_{t}$ as $p(x)=\prod_{t=1}^{T}p(x_{t}|x_{t-1},\ldots,x_{1})$, we jointy model all the sub-pixels and introduce a learned hierarchy of auxiliary feature representations $z^{(1)},\ldots,z^{(S)}$ to simplify the modeling task.We fix the dimensions of $z^{(s)}$ to be $C{\times}H^{\prime}{\times}W^{\prime}$, where the number of channels $C$ is a hyperparameter ($C=5$ in our reported models), and $H^{\prime}=H/2^{s},W^{\prime}=W/2^{s}$ given a $H{\times}W$-dimensional image.Considering that $z^{(s)}$ is quantized, this conveniently upper bounds the information that can be contained within each $z^{(s)}$, however, other dimensions could be explored. Specifically, we model the joint distribution of the image $x$ and the feature representations $z^{(s)}$ as

where $p(z^{(S)})$ is a uniform distribution. The feature representations can be hand designed or learned. Specifically, on one side, we consider an RGB pyramid with $z^{(s)}=\mathcal{B}_{2^{s}}(x)$, where $\mathcal{B}_{2^{s}}$ is the bicubic (spatial) subsampling operator with subsampling factor $2^{s}$. On the other side, we consider a learned representation $z^{(s)}=F^{(s)}(x)$ using a feature extractor $F^{(s)}$. We use the hierarchical model shown in Fig. 1 using the composition $F^{(s)}=Q\circ E^{(s)}\circ\dots\circ E^{(1)}$, where the $E^{(s)}$ are feature extractor blocks and $Q$ is a scalar differentiable quantization function (see Sec. 3.3). The $D^{(s)}$ in Fig. 1 are predictor blocks, and we parametrize $E^{(s)}$ and $D^{(s)}$ as convolutional neural networks.

Letting $z^{(0)}=x$, we parametrize the conditional distributions for all $s\in\{0,\dots,S\}$ as

using the predictor features $f^{(s)}=D^{(s)}(f^{(s+1)},z^{(s)})$. The final predictor only sees $z^{(S)}$, i.e., we let $f^{(S+1)}=0$. Note that $f^{(s+1)}$ summarizes the information of $z^{(S)},\dots,z^{(s+1)}$.

The predictor is based on the super-resolution architecture from EDSR , motivated by the fact that our prediction task is somewhat related to super-resolution in that both are dense prediction tasks involving spatial upsampling. We mirror the predictor to obtain the feature extractor, and follow in not using BatchNorm . Inspired by the “atrous spatial pyramid pooling” from , we insert a similar layer at the end of $D^{(s)}$: In $A*$, we use three atrous convolutions in parallel, with rates 1, 2, and 4, then concatenate the resulting feature maps to a $3C_{f}$-dimensional feature map.

3 Quantization

but use differentiable “soft quantization”

4 Mixture Model

For ease of notation, let $z^{(0)}=x$ again. We model the conditional distributions $p(z^{(s)}|z^{(s+1)},\ldots,z^{(S)})$ using a generalization of the discretized logistic mixture model with $K$ components proposed in , as it allows for efficient training: The alternative of predicting logits per (sub-)pixel has the downsides of requiring more memory, causing sparse gradients (we only get gradients for the logit corresponding to the ground-truth value), and does not model that neighbouring values in the domain of $p$ should have similar probability.

Let $c$ denote the channel and $u,v$ the spatial location. For all scales, we assume the entries of $z^{(s)}_{cuv}$ to be independent across $u,v$, given $f^{(s+1)}$. For RGB ($s=0$), we define

where we use a weak autoregression over RGB channels to define the joint probability distribution via a mixture $p_{m}$ (dropping the indices $uv$ for shorter notation):

We define $p_{m}$ as a mixture of logistic distributions $p_{l}$ (defined in Eq. (10) below). To this end, we obtain mixture weightsNote that in contrast to we do not share mixture weights $\pi^{k}$ across channels. This allows for easier marginalization of Eq. (5). $\pi^{k}_{cuv}$, means $\mu^{k}_{cuv}$, variances $\sigma^{k}_{cuv}$, as well as coefficients $\lambda^{k}_{cuv}$ from $f^{(1)}$ (see further below), and get

For all scales, the individual logistics $p_{l}$ are given as

Here, $b$ is the bin width of the quantization grid ($b=1$ for $s=0$ and $b=1/12$ otherwise). The edge-cases $z=0$ and $z=255$ occurring for $s=0$ are handled as described in [35, Sec. 2.1].

For all scales, we obtain the parameters of $p(z^{(s-1)}|f^{(s)})$ from $f^{(s)}$ with a $1{\times}1$ convolution that has $C_{p}^{(s-1)}$ output channels (see Fig. 2). For RGB, this final feature map must contain the three parameters $\pi,\mu,\sigma$ for each of the 3 RGB channels and $K$ mixtures, as well as $\lambda_{\alpha},\lambda_{\beta},\lambda_{\gamma}$ for every mixture, thus requiring $C_{p}^{(0)}=3\cdot 3\cdot K+3\cdot K$ channels. For $s>0$, $C_{p}^{(s)}=3\cdot C\cdot K$, since no $\lambda$ are needed. With the parameters, we can obtain $p(z^{(s)}|f^{(s+1)})$, which has dimensions $3{\times}H{\times}W{\times}256$ for RGB and $C{\times}H^{\prime}{\times}W^{\prime}{\times}L$ otherwise (visualized with cubes in Fig. 1).

We emphasize that in contrast to , our model is not autoregressive over pixels, i.e., $z_{cuv}^{(s)}$ are modelled as independent across $u,v$ given $f^{(s+1)}$ (also for $z^{(0)}=x$).

5 Loss

Note that the loss decomposes into the sum of the cross-entropies of the different representations. Also note that this loss corresponds to the negative log-likelihood of the data w.r.t. our model which is typically the perspective taken in the generative modeling literature (see, e.g., ).

We emphasize that in contrast to the generative model literature, we learn the representations, propagating gradients to both $E^{(s)}$ and $D^{(s)}$, since each component of our loss depends on $D^{(s+1)},\dots,D^{(S)}$ via the parametrization of the logistic distribution and on $E^{(s)},\dots,E^{(1)}$ because of the differentiable $Q$. Thereby, our network can autonomously learn to navigate the trade-off between a) making the output $z^{(s)}$ of feature extractor $E^{(s)}$ more easily estimable for the predictor $D^{(s+1)}$ and b) putting enough information into $z^{(s)}$ for the predictor $D^{(s)}$ to predict $z^{(s-1)}$.

6 Relationship to MS-PixelCNN

When the auxiliary features $z^{(s)}$ in our approach are restricted to a non-learned RGB pyramid (see baselines in Sec. 4), this is somewhat similar to MS-PixelCNN . In particular, combines such a pyramid with upscaling networks which play the same role as the predictors in our architecture. Crucially however, they rely on combining such predictors with a shallow PixelCNN and upscaling one dimension at a time ($W{\times}H{\rightarrow}2W{\times}H{\rightarrow}2W{\times}2H$). While their complexity is reduced from $O(WH)$ forward passes needed for PixelCNN to $O(\log WH)$, their approach is in practice still two orders of magnitude slower than ours (see Sec. 5.2). Further, we stress that these similarities only apply for our RGB baseline model, whereas our best models are obtained using learned feature extractors trained jointly with the predictors.

Experiments

We compare our main model (L3C) to two learned baselines: For the RGB Shared baseline (see Fig. A2) we use bicubic subsampling as feature extractors, i.e., $z^{(s)}=\mathcal{B}_{2^{s}}(x)$, and only train one predictor $D^{(1)}$. During testing, we obtain multiple $z^{(s)}$ using $\mathcal{B}$ and apply the single predictor $D^{(1)}$ to each. The RGB baseline (see Fig. A3) also uses bicubic subsampling, however, we train $S=3$ predictors $D^{(s)}$, one for each scale, to capture the different distributions of different RGB scales. For our main model, L3C, we additionally learn $S=3$ feature extractors $E^{(s)}$.We chose $S=3$ because increasing $S$ comes at the cost of slower training, while yielding negligible improvements in bitrate. For an image of size $H{\times}W$, the last bottleneck has $5{\times}H/8{\times}W/8$ dimensions, each quantized to $L=25$ values. Encoding this with a uniform prior amounts to ${\approx}4\%$ of the total bitrate. For the RGB Shared baseline, we apply $D^{(1)}$ 4 times, as only one encoder is trained. Note that the only difference to the RGB baseline is that the representations $z^{(s)}$ are learned. We train all these models until they converge at 700k iterations.

We train our models on 362 551 images randomly selected from the Open Images training dataset . We randomly downscale the images to at least 512 pixels on the longer side to remove potential artifacts from previous compression, discarding images where rescaling does not result in at least $1.25\times$ downscaling. Further, following we discard high saturation/ non-photographic images, i.e., images with mean $S{>}0.9$ or $V{>}0.8$ in the HSV color space. We evaluate on 500 images randomly selected from the Open Images validation set, preprocessed like the training set (available on our github1), the 100 images from the commonly used super-resolution dataset DIV2K , as well as on RAISE-1k , a “real-world image dataset” with 1000 images. We automatically split images into equally-sized crops if they do not fit into GPU memory, and process crops sequentially. Note that this is a bias against our method.

In order to compare to the PixelCNN literature, we additionally train L3C on the ImageNet32 and ImageNet64 datasets , each containing 1 281 151 training images and 50 000 validation images, of $32\times 32$ resp. $64\times 64$ pixels.

We use the RMSProp optimizer , with a batch size of 30, minimizing Eq. (11) directly (no regularization). We train on $128\times 128$ random crops, and apply random horizontal flips. We start with a learning rate $\lambda=1\cdot 10^{-4}$ and decay it by a factor of 0.75 every 5 epochs. On ImageNet32/64, we decay $\lambda$ every epoch, due to the smaller images.

We find that adding BatchNorm slightly degrades performance. Furthermore, replacing the stacked atrous convolutions $A*$ with a single convolution, slightly degrades performance as well. By stopping gradients from propagating through the targets of our loss, we get significantly worse performance – in fact, the optimizer does not manage to pull down the cross-entropy of any of the learned representations $z^{(s)}$ significantly.

Additionally, we explored the impact of varying $C$ (number of channels of $z^{(s)}$) and the number of levels $L$ and found it more beneficial to increase $L$ instead of increasing $C$, i.e., it is beneficial for training to have a finer quantization grid.

We compare to FLIF and the lossless mode of WebP using the respective official implementations , for PNG we use the implementation of Pillow , and for the lossless mode of JPEG2000 we use the Kakadu implementation . See Sec. 2 for a description of these codecs.

Results

Table 1 shows a comparison of our approach (L3C) and the learned baselines to the other codecs, on our testsets, in terms of bits per sub-pixel (bpsp)We follow the likelihood-based generative modelling literature in measuring bpsp; $X$ bits per pixel (bpp) $=X/3$ bpsp, see also footnote 2. All of our methods outperform the widely-used PNG, which is at least $43\%$ larger on all datasets. We also outperform WebP and JPEG2000 everywhere by a smaller margin of up to $3.3\%$. We note that FLIF still marginally outperforms our model but remind the reader of the many hand-engineered highly specialized techniques involved in FLIF (see Section 2). In contrast, we use a simple convolutional feed-forward neural network architecture. The RGB baseline with $S=3$ learned predictors outperforms the RGB Shared baseline on all datasets, showing the importance of learning a predictor for each scale. Using our main model (L3C), where we additionally learn the feature extractors, we outperform both baselines: The outputs are at least $7.8\%$ larger everywhere, showing the benefits of learning the representation.

2 Comparison with PixelCNN

While PixelCNN-based approaches are not designed for lossless image compression, they learn a probability distribution over pixels and optimize for the same log-likelihood objective. Since they thus can in principle be used inside a compression algorithm, we show a comparison here.

Table 2 shows a speed comparison to three PixelCNN-based approaches (see Sec. 2 for details on these approaches). We compare time spent when sampling from the model, to be able to compare to the PixelCNN literature. Actual decoding times for L3C are given in Sec. 5.3.

While the runtime for PixelCNN and MS-PixelCNN is taken from the table in , we can compare with L3C by assuming that PixelCNN++ is not slower than PixelCNN to get a conservative estimatePixelCNN++ is in fact around $3\times$ faster than PixelCNN due to modelling the joint directly, see Sec. 2., and by considering that MS-PixelCNN reports a $105{\times}$ speedup over PixelCNN. When comparing on $320\times 320$ crops, we thus observe massive speedups compared to the original PixelCNN: ${>}1.63\cdot 10^{5}\times$ for batch size (BS) 1 and ${>}5.31\cdot 10^{4}\times$ for BS 30. We see that on $320\times 320$ crops, L3C is at least $5.06\cdot 10^{2}\times$ faster than MS-PixelCNN, the fastest PixelCNN-type approach. Furthermore, Table 2 makes it obvious that the PixelCNN based approaches are not practical for lossless compression of high-resolution images.

We emphasize that it is impossible to do a completely fair comparison with PixelCNN and MS-PixelCNN due to the unavailability of their code and the different hardware. Even if the same hardware was available to us, differences in frameworks/framework versions (PyTorch vs. Tensorflow) can not be accounted for. See also Sec. A.4 for notes on the influence of the batch size.

To put the runtimes reported in Table 2 into perspective, we also evaluate the bitcost on ImageNet32, for which PixelCNN and MS-PixelCNN were trained, in Table 3. We observe our outputs to be $20.6\%$ larger than MS-PixelCNN and $24.4\%$ larger than the original PixelCNN, but smaller than all classical approaches. However, as shown above, this increase in bitcost is traded against orders of magnitude in speed. We obtain similar results for ImageNet64, see Sec. A.3.

3 Encoding / Decoding Time

To encode/decode images with L3C (and other methods outputting a probability distribution), a pass with an entropy coder is needed. We implemented a relatively simple pipeline to encode and decode images with L3C, which we describe in the supplementary material, in Section A.2. The results are shown in Tables 4 and A1. As noted in Section A.2, we did not optimize our code for speed, yet still obtain practical runtimes. We also note that to use other likelihood-based methods for lossless compression, similar steps are required. While our encoding time is in the same order as for classical approaches, our decoder is slower than that of the other approaches. This can be attributed to more optimized code and offloading complexity to the encoder – while in our approach, decoding essentially mirrors encoding. However, combining encoding and decoding time we are either faster (FLIF) or have better bitrate (PNG, WebP, JPEG2000).

4 Sampling Representations

Conclusion

We proposed and evaluated a fully parallel hierarchical probabilistic model with auxiliary feature representations. Our L3C model outperforms PNG, JPEG2000 and WebP on all datasets. Furthermore, it significantly outperforms the RGB Shared and RGB baselines which rely on predefined heuristic feature representations, showing that learning the representations is crucial. Additionally, we observed that using PixelCNN-based methods for losslessly compressing full resolution images takes two to five orders of magnitude longer than L3C.

To further improve L3C, future work could investigate weak forms of autoregression across pixels and/or dynamic adaptation of the model network to the current image. Moreover, it would be interesting to explore domain-specific applications, e.g., for medical image data.

Acknowledgments The authors would like to thank Sergi Caelles for the insightful discussions and feedback. This work was partly supported by ETH General Fund (OK) and Nvidia through the hardware grant.

References

Appendix A Practical Full Resolution Learned Lossless Image Compression – Supplementary Material

Previously, our preprocessing script saved all training images and validation images as JPGs with a high quality factor of $Q=95$, downscaled by a factor 0.75. It turns out that the resulting images have a specific enough distribution that the neural network picks up on it, and the images are also easier to compress for the non-learned codecs.

For correctness, we have thus re-created the training and validation sets. The new preprocessing script and more details is available on github1. The important differences are:

We do not rescale validation sets in any way, and instead divide the images into crops such that everything fits into memory.

For the training set, we use a random downscaling factor, instead of fixed 0.75x: this provides a wider variety of downscaling artefacts.

A.2 Encoding and Decoding Details

Table A1 shows the time required to decode each scale $s$. We first obtain the CDF as a matrix on the CPU to be able to use the arithmetic decoder (see below), and then do a pass with the arithmetic decoder. We did not optimize either part for speed, as noted in Sec. A.2.2.

The following shows detailed steps, using again $z^{(0)}=x$. The steps are also visualized in Fig. A4.

Forward pass through network to obtain $\forall s:\;z^{(s)}$, $f^{(s)}$.

Encode $z^{(S)}$ assuming a uniform prior, i.e., assuming each of the $L$ symbols is equally likely. This requires $\log_{2}(L)$ bits per symbol.

In practice, the division into intervals $[a,b)$ required for arithmetic coding described in Sec. 3.1 is most efficiently done by having access to the cumulative distribution function (CDF) of the symbol to encode. Thus, for the RGB scale ($s=0$), we obtain the CDF analogously to Eq. (6):

And, analogously to Eq. (9), the CDF for $s>0$ for each channel $c$ is

$C_{l}$ in Eqs. (12), (13) is the CDF of the logistic distribution,

For each $s,c$ the CDF $C(z^{(s)}_{c}|f^{(s+1)})$ is a $H^{\prime}\times W^{\prime}\times L$-dimensional matrix, where $L=257$ for RGB and $L=26$ otherwise, and $H^{\prime}=H/2^{s},W^{\prime}=W/2^{s}$.

For each $s\in\{S+1,\dots,0\}$, encode each channel $c$ of $z^{(s)}$ with the predicted $C(z^{(s)}_{c}|f^{(s+1)})$, using adaptive arithmetic coding (see Sec. 3.1). To be able to uniquely decode, the sub-bitstream for $z^{(s)}$ always starts with a triplet encoding its dimensions $C,H^{\prime},W^{\prime}$ as uint16. The final bitstream is the concatenation of all sub-bitstreams.

Decoding

Obtain the final $z^{(S)}$ from the bitstream, which was encoded with a uniform prior.

Feed $z^{(S)}$ to $D^{(S)}$ to obtain $f^{(S)}$, and thereby also $C(z^{(S-1)}_{cuv}|f^{(S)})$ for all $c$. Since the decoder now has access to the same CDF as the encoder, we can decode $z^{(S-1)}$ from the bitstream with our adaptive arithmetic decoder.

Analogously, we repeat the previous step to obtain $z^{(S)},\dots,z^{(1)}$, as well as $f^{(S)},\dots,f^{(1)}$ using the accompanying CDFs.

Concatenating the channels $x_{1},x_{2},x_{3}$, we finally obtain the decoded image $x$.

A.2.1 Hardware Used

Our timings were obtained on a machine with a Titan X (Pascal) GPU and Intel Xeon E5-2680 v3 CPU.

A.2.2 Notes on Code Optimization

The encoder can be run in parallel over all scales, as all CDFs are known after one forward pass. Further, we do not need to know the CDF for all symbols, but only for the symbols $z$ we encode and $z+1$, since this specifies the interval $[a,b)$. The decoder is sequential in the scales since $z^{(s)}$ is required to predict the distribution of $z^{(s-1)}$. Still, for $s>0$, the decoding of the channels of the $z^{(s)}$ could be parallelized, as the channels are modelled fully independently. However, we did not implement either of these improvements, keeping the code simple.

For both encoder and decoder, the CDFs must be available to the CPU, as the arithmetic coder runs there. However, the CDFs are huge tensors for real-world images ($H\times W\times 257$ for RGB, which amounts to $257$MB for each channel of a $512\times 512$ image). To save the expensive copying from GPU to CPU, we implemented our own CUDA kernel to store the claculated $C$ directly into “managed memory”, which can be accessed from both CPU and GPU. However, we did not optimize this CUDA kernel for speed.

Finally, while state-of-the-art adaptive entropy coders typically require on the order of milliseconds per MB (see and in particular for benchmarks on adaptive entropy coding), we implemented a simple arithmetic coding module to obtain the times in our tables. Please see the code1 for details.

A.3 Comparison on ImageNet64

We show a bpsp comparison on ImageNet64 in Table A2. Similar to what we observed on ImageNet32 (see Section 5.2), our outputs are $23.8\%$ larger than MS-PixelCNN and $19.4\%$ larger than the original PixelCNN, but smaller than all classical approaches. We note again that increase in bitcost is traded against orders of magnitude in speed.

We also note that the gap between classical approaches and PixelCNN becomes smaller compared to ImageNet32.

A.4 Note on Comparing Times for 32×32323232\times 32 Images

In Table 2, we report run times for batch size 30 to be able to compare with the run times reported in . However, this comparison is biased against us, as can be seen in Table A3: Since our network is fairly small, we can process up to 480 images of size $32{\times}32$ in parallel. We observe that the time to sample one image drops as the batch size increases, indicating that for BS=30, some overhead dominates.

A.5 Visualizing Representations

We visualize the representations $z^{(1)},z^{(2)},z^{(3)}$ in Fig. A1. It can be seen that the global image structure is preserved over scales, with representations corresponding to smaller $s$ modeling more detail. This shows potential for efficiently performing image understanding tasks on partially decoded images similarly as described in for lossy learned compression: instead of training a feature extractor for a given task on $x$, one could directly use the features $z^{(s)}$ from our network.

A.6 Architectures of Baselines

Figs. A2, A3 show the architectures for the RGB Shared and RGB baselines. The dots in Fig. A2 indicate that the model could in theory be applied more since $D^{(1)}$ is used for every scale.

A.7 Encoding and Decoding Visualized

We visualize the steps needed to encode the different $z^{(s)}$ in Fig. A4 on the next page.