HiLLoC: lossless image compression with hierarchical latent variable models

1 INTRODUCTION . Bits back coding ( Wallace , 1990 ; Hinton & van Camp , 1993 ) is a method for performing lossless compression using a latent variable model . In an ideal implementation , the method can achieve an expected message length equal to the variational free energy , often referred to as the negative evidence lower bound ( ELBO ) of the model . Bits back was first introduced to form a theoretical argument for using the ELBO as an objective function for machine learning ( Hinton & van Camp , 1993 ) . The first implementation of bits back coding ( Frey , 1997 ; Frey & Hinton , 1996 ) made use of first-infirst-out ( FIFO ) arithmetic coding ( AC ) ( Witten et al. , 1987 ) . However , the implementation did not achieve optimal compression , due to an incompatibility between a FIFO coder and bits back coding , and its use was only demonstrated on a small dataset of 8×8 binary images . Recently , zero-overhead bits back compression with a significantly simpler implementation has been developed by Townsend et al . ( 2019 ) . This implementation makes use of asymmetric numeral systems ( ANS ) , a last-in-first-out ( LIFO ) entropy coding scheme ( Duda , 2009 ) . The method , known as ‘ Bits Back with Asymmetric Numeral Systems ’ ( BB-ANS ) was demonstrated by compressing the MNIST test set using a variational auto-encoder ( VAE ) model ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) , achieving a compression rate within 1 % of the model ELBO . More recently , Hoogeboom et al . ( 2019 ) and Ho et al . ( 2019 ) have proposed flow-based methods for lossless compression , and Kingma et al . ( 2019 ) have presented ‘ Bit-Swap ’ , extending BB-ANS to hierarchical models . In this work we present an alternative method for extending to hierarchical VAEs . This entails the following novel techniques : 1 . Direct coding of arbitrary sized images using a fully convolutional model . 2 . A vectorized ANS implementation supporting dynamic shape . 3 . Dynamic discretization to avoid having to calibrate a static discretization . 4 . Initializing the bits back chain using a different codec . We discuss each of these contributions in detail in Section 3 . We call the combination of BB-ANS using a hierarchical latent variable model and the above techniques : ‘ Hierarchical Latent Lossless ∗Equal contribution . 1Available at https : //github.com/hilloc-submission/hilloc . Compression ’ ( HiLLoC ) . In our experiments ( Section 4 ) , we demonstrate that HiLLoC can be used to compress color images from the ImageNet test set at rates close to the ELBO , outperforming all of the other codecs which we benchmark . We also demonstrate the speedup , of nearly three orders of magnitude , resulting from vectorization . We release an open source implementation based on ‘ Craystack ’ , a Python package which we have written for general prototyping of lossless compression with ANS . 2 BACKGROUND . In this section we briefly describe the BB-ANS algorithm first introduced by Townsend et al . ( 2019 ) . We begin by giving a high-level description of the ANS LIFO entropy coder ( Duda , 2009 ) , along with a new notation for describing the basic ANS operations . Throughout the rest of the paper we use log to mean the base two logarithm , usually denoted log2 , and we measure message lengths in bits . 2.1 ASYMMETRIC NUMERAL SYSTEMS . As an entropy coder , ANS was designed for compressing sequences of discretely distributed symbols . It achieves a compressed message length equal to the negative log-probability ( information content ) of the sequence plus an implementation dependent constant , which is usually less than 32 bits . For long sequences , the constant overhead has a negligible contribution to the overall compression rate . Thus , by Shannon ’ s source coding theorem ( Shannon , 1948 ) , ANS coding is guaranteed to be near-optimal for long sequences . There are two basic operations defined by ANS , which we will refer to as ‘ push ’ and ‘ pop ’ . Push encodes a symbol by adding it to an existing message . It has the signature push : ( message , symbol ) 7→ message′ . ( 1 ) Pop is the inverse of push , and may be used to decode a symbol and recover a message identical to that before pushing . pop : message′ 7→ ( message , symbol ) . ( 2 ) When multiple symbols are pushed in sequence , they must be popped using the precise inverse procedure , which means popping the symbols in the opposite order . Hence why ANS is referred to as a last-in-first-out coder , or a stack . The push and pop operations require access to a probabilistic model of symbols , summarized by a probability mass function p over the alphabet of possible symbols . The way that symbols are encoded depends on the model , and pushing a symbol s according to p results in an increase in message length of log 1p ( s ) . Popping s results in an equal reduction in message length . For details on how the ANS operations are implemented , see Duda ( 2009 ) . Note that any model/mass function can be used for the pop operation , i.e . there ’ s no hard restriction to use the distribution that was used to encode the message . In this way , rather than decoding the same data that was encoded , pop can actually be used to sample a symbol from a different distribution . The pop method itself is deterministic , so the source of randomness for the sample comes from the data contained within the message . This sampling operation , which can be inverted by pushing the sample back onto the stack , is essential for bits back coding . For convenience , we introduce the shorthand notation s→ p ( · ) for encoding ( pushing ) a symbol s according to p , and s← p ( · ) for decoding ( popping ) . 2.2 BITS BACK WITH ANS . Suppose we have a model for data x which involves a latent variable z . A sender and receiver wish to communicate a sample x . They have access to a prior on z , denoted p ( z ) , a likelihood p ( x | z ) and a ( possibly approximate ) posterior q ( z |x ) , but not the marginal distribution p ( x ) . Without access to p ( x ) , sender and receiver can not directly code x using ANS . However , BB-ANS specifies an indirect way to push and pop x . It does not require access to the marginal p ( x ) , but rather uses the prior , conditional , and posterior from the latent variable model . Table 1 ( a ) shows , in order from the top , the three steps of the BB-ANS pushing procedure which the sender can perform to encode x . The ‘ Variables ’ column shows the variables known to the sender before each step . 1 ( b ) shows the inverse steps which the receiver can use to pop x , with the ‘ Variables ’ column showing what is known to the receiver after each step . After decoding x , the third step of popping , z → q ( · |x ) , is necessary to ensure that BB-ANS pop is a precise inverse of push . The change in message length from BB-ANS can easily be derived by adding up the quantities in the ∆L column of Table 1 . For encoding we get ∆LBB−ANS = − log 1 q ( z |x ) + log 1 p ( x | z ) + log 1 p ( z ) ( 3 ) = − log p ( x , z ) q ( z |x ) . ( 4 ) Taking the expectation over z gives the expected message length for a datum x L ( x ) = −Eq ( z | x ) [ log p ( x , z ) q ( z |x ) ] ( 5 ) which is the negative evidence lower bound ( ELBO ) , also known as the free energy . This is a commonly used training objective for latent variable models . The above equation implies that latent variable models trained using the ELBO are implicitly being trained to minimize the expected message length of lossless compression using BB-ANS . Note that , as Table 1 shows , the first step of encoding a data point , x , using BB-ANS is to , counterintuitively , decode ( and thereby sample ) a latent z ← q ( · |x ) . This requires that there is already a buffer of random data pushed to the ANS coder , which can be popped . This data used to start the encoding process is recovered after the final stage of decoding , hence the name ‘ bits back ’ . If we have multiple samples to compress , then we can use ‘ chaining ’ , which is essentially repeated application of the procedure in Table 1 ( Townsend et al. , 2019 ) . In Section 3.4 we describe how we build up an initial buffer of compressed data by using a different codec to code the first images in a sequence . 3 SCALING UP BITS BACK WITH ANS . We now discuss the techniques we introduce to scale up BB-ANS . 3.1 FULLY CONVOLUTIONAL MODELS . When all of the layers in the generative and recognition networks of a VAE are either convolutional or elementwise functions ( i.e . the VAE has no densely connected layers ) , then it is possible to evaluate the recognition network on images of any height and width , and similarly to pass latents of any height and width through the generative network to generate an image . Thus , such a VAE can be used as a ( probabilistic ) model for images of any size . We exploit this fact , and show empirically in Section 4 that , surprisingly , a fully convolutional VAE trained on 32 × 32 images can perform well ( in the sense of having a high ELBO ) as a model for 64 × 64 images as well as far larger images . This in turn corresponds to a good compression rate , and we implement lossless compression of arbitrary sized images by using a VAE in this way . 3.2 VECTORIZED LOSSLESS COMPRESSION . The primary computational bottlenecks in the original BB-ANS implementation ( Townsend et al. , 2019 ) were loops over data and latent variables occurring in the Python interpreter . We have been able to vectorize these , achieving an implementation which can scale to large ImageNet images . The effect of vectorization on runtime is shown in Figure 4 . A vectorized implementation of ANS was described in Giesen ( 2014 ) using SIMD instructions . This works by expanding the size of the ANS stack head , from a scalar to a vector , and interleaving the output/input bit stream . We implement this in our lossless compression library , Craystack , using Numpy . Please refer to the Craystack code and to Giesen ( 2014 ) for more detail . We ensure that the compression rate overhead to vectorization is low by using the BitKnit technique described in Giesen ( 2015 ) , see Appendix D for more detail . Having vectorized , we found that most of the compute time for our compression was spent in neural net inference , whether running on CPU or GPU , which we know to already be reasonably well optimized . In Craystack , we further generalize the ANS coder using Numpy ’ s n-dimensional array view interface , allowing the stack head to be ‘ shaped ’ like an n-dimensional array , or a nested Python data-structure containing arrays . We can then use a shape which fits that of the data that we wish to encode or decode . When coding data according to a VAE we use an ANS stack head shaped into a pair of arrays , matching the shapes of the observation x and the latent z . This allows for a straightforward implementation and clarifies the lack of data dependence between certain operations , such as the x→ p ( · | z ) and z → p ( · ) during encoding , which can theoretically be performed concurrently . This vectorized encoding process is visualized in Figure 2 .

Summary: This paper focuses on lossless source compression with bits back coding for hierarchical fully convolutional VAEs. The focus/contribution is three-fold: 1. Improve the compression rate performance by adapting the discretization of latent space required for the entropy coder ANS. The newly proposed discretization scheme allows for a dependency structure that is not restricted to a Markov chain structure in the encoder model q(z|x) and in the generative part of the model p(x,z). This is in contrast with bit-swap[1], which requires a markov chain structure. The dependency structure that is allowed in the proposed method is widely known to perform better than a markov chain structure, which can explain why it improves significantly over Bit-swap [1] (another hierarchical VAE compression algorithm that uses bits back coding.) 2. Increasing compression speed by implementing a vectorized version of ANS, and heaving an ANS head in the shape of a pair of arrays matching that of the latent variable and the observed variable. The latter allows for simultaneous encoding of the latent with the prior distribution and the image with the decoder distribution. 3. Showing that a model trained on a low-resolution imagenet 32 dataset can generalize its compression capabilities to higher resolution datasets with convincing results.