OSOA: One-Shot Online Adaptation of Deep Generative Models for Lossless Compression

1 Introduction . Lossless compression has always been an important part in the information and communications technology industry and has become more and more essential with the data explosion of the Big Data era [ 10 , 39 , 40 , 8 ] . Recently , a large number of machine learning approaches , predominantly deep generative models ( DGMs ) , have been proposed for lossless compression [ 34 , 48 , 19 , 17 , 51 ] . As indicated by Shannon ’ s source coding theorem , the compression limit is determined by the entropy of the data ’ s ground truth distribution [ 30 ] ; any mismatch between an approximate distribution for coding and the unknown ground truth distribution will cause extra coding cost . It is the powerful approximation ability of DGMs that enables more precise estimation of the ground truth distribution [ 23 , 41 , 24 , 29 ] and enables DGMs based compression algorithms to outperform the conventional counter-parties in terms of the compression ratio . However , in real production scenarios , both space and time efficiencies are desirable . While more space efficiency is provided by higher compression ratio with DGMs based lossless compression algorithms , training of the probabilistic models is time consuming . Training a separate model for each different dataset to compress would scale the expensive time cost linearly . Moreover , to decode the dataset , one has to store the probabilistic model together with the codeword of the data , which requires additional space for local decoding and additional transmission cost for remote decoding . Apart from one model per dataset , another extreme on the spectrum is using one probabilistic model ∗Work done during employment at Huawei Technologies R & D UK 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . for many different datasets . Since each dataset can follow a very different distribution , using one model for all the datasets could lead to high inefficiency in terms of compression ratio . To solve the above dilemma , we propose the setting of One-Shot Online Adaptation ( OSOA ) of deep generative models for lossless compression . A DGM based lossless compression algorithm under the OSOA setting ( referred to as OSOA algorithm ) includes one pretrained deep generative model ( referred to as base model ) and one deterministic dynamical system . The base model should be able to capture essential structures of a class of datasets to serve as a good starting point of the adaptation process . The dynamical system is the mechanism that can update the current distribution to a new one based on the current batch at each step in the OSOA process . At each OSOA step , one first uses the current deep generative model to evaluate the probabilistic mass function values to compress the current data batch before using the same batch to update the model with the dynamical system , which gives the name of online adaptation . The online adaptation is only conducted for one epoch and a distribution sequence will be unrolled during this process . Moreover , the dynamical system should be designed to optimise the objective functional associated with the deep generative model so that the sequence could provide better approximations on the unrolling procedure . Under this new setting , we can start from a proper base model and benefit from improved compression ratio as the sequence being unrolled . As the sequence is fully determined by the base model and the dynamical system , this approach allows to not only eliminate the need of extensively training a new model on each new dataset , but also to bypass the associated model storage cost that would normally be required . Apart from the novel setting formulation , we showcase with experiments on large datasets that a vanilla OSOA algorithm can significantly save time versus training bespoke models and save space versus using one model for all targets . Moreover , with the same adaptation step number or adaptation time , vanilla OSOA can exhibit better space efficiency , e.g. , 47 % less space , than fine-tuning the pretrained model and saving the fine-tuned model . We also show that further trade-off between time and space cost is available with varies strategies , e.g. , multiple updates per batch and early stopping , and motivate more sophisticated algorithms with preliminary explorations . Further , we use sample images from OSOA and baselines to justify adaptation based compression algorithms and to discuss societal considerations for DGMs based lossless compression . The paper is organised as follows . In Sec . 2 , we review background knowledge regarding deep generative models for lossless compression . In Sec . 3 , we formulate the setting of OSOA , discuss associated codec designing principles and propose a vanilla OSOA algorithm . In Sec . 4 , we validate the setting by comparison with baselines , explore the strategies for further trade-off and discuss societal considerations . In Sec . 5 , we discuss some related work . In Sec . 6 , we summarise the paper , offer insights into our findings and point out promising future work directions . 2 Preliminary . Entropy Coding The core component underlying any lossless compression algorithm is the entropy coder , which converts between the data symbols and codewords . Popular entropy coders include Huffman coding [ 21 ] , arithmetic coding ( AC ) [ 54 ] and asymmetric numeral systems ( ANS ) [ 9 ] . Huffman coding mainly involves binary trees operations , which is usually more time efficient than AC and ANS but can admit suboptimal codelength for probabilities not being power of 1/2 [ 30 ] . As a result , popular codecs adopted by DGM based algorithms are normally based on AC or ANS , e.g. , [ 34 , 19 , 35 ] . The main difference between AC and ANS is that the former encodes and decodes in a first-in-first-out ( FIFO ) style while the latter performs in a first-in-last-out ( FILO ) style [ 48 ] . All entropy coders require the knowledge of the data symbol alphabet and the associated distribution , in order to conduct compression and decompression . It is worth noting that the distribution used for entropy coding should be discrete and the average codelength is expected to be the entropy of the distribution used for coding . Recall that for a discrete distribution with the probabilistic mass function p ( x ) , the ( Shannon ) entropy is defined as Ex∼p [ − log p ( x ) ] . Moreover , the optimal average codelength , given by the entropy of the ground truth distribution , is achieved only if the ground truth distribution is used for coding as per Shannon ’ s source coding theorem [ 30 ] . Specifically , compressing data from distribution p ( x ) with distribution q ( x ) will lead to the expected extra codelength at KL ( p ( x ) ||q ( x ) ) . Coding with Explicit Deep Generative Models Since the entropy coders require probability values , implicit DGMs , e.g. , generative adversarial networks ( GANs ) [ 11 ] , can not be readily used for lossless compression . For explicit DGMs , the learning principle is either maximum likelihood estimation ( MLE ) [ 41 , 24 , 16 ] : θ̂ = arg min θ∈Θ Ex∼p∗ ( x ) [ − log pθ ( x ) ] ( 1 ) or maximising the evidence lower bound ( ELBO ) of the logarithm evidence log pθ ( x ) [ 23 , 49 ] : θ̂ , φ̂ = arg minθ∈Θ , φ∈ΦEx∼p∗ ( x ) [ −L ( θ , φ , x ) ] , ( 2 ) where L ( θ , φ , x ) = Eqφ ( z|x ) [ − log qφ ( z|x ) + log pθ ( x|z ) pθ ( z ) ] . Note that x is a discrete random variable , so p∗ ( x ) refers to the probabilistic mass function of the ground truth data distribution . When learning with MLE , one can either directly stipulate a discrete model or stipulate a continuous model with the dequantization technique [ 50 , 45 ] . If one stipulates a discrete model , where pθ ( x ) denotes the probabilistic mass function of the model , the objective function in Eq . 1 delegates the expected codelength when using pθ ( x ) for entropy coding . If one stipulates a continuous model , one normally dequantizes the discrete data x with a continuous noise u uniformly distributed on the unit hypercube , and gets the dequantized variable y = x + u with distribution p∗ ( y ) . After fitting the probabilistic density function pθ ( y ) of the model to the the probability density p∗ ( y ) of the dequantized distribution , one needs to discretise pθ ( y ) and use the discretised distribution for entropy coding . In this case , the MLE objective Ey∼p∗ ( y ) [ − log pθ ( y ) ] delegates an upper bound of the expected codelength [ 45 ] . For how to design sophisticated dequantization schemes or how to design discrete deep generative models , we refer the readers to [ 16 , 18 , 19 , 51 , 17 , 55 ] . When learning with ELBO , one usually stipulates discrete models for x , i.e. , pθ ( x|z ) is a probabilistic mass function , and continuous models for z , i.e. , qφ ( z|x ) and pθ ( z ) are probabilistic density functions . Furthermore , one usually does not have direct access to the close form of pθ ( x ) but the generative components pθ ( x|z ) and pθ ( z ) . To use pθ ( x|z ) and pθ ( z ) for efficient coding , a technique called bitsback [ 53 , 15 ] is employed , which additionally involves the approximate posterior distribution qφ ( z|x ) learned together with the generative components [ 48 , 49 , 26 ] . In these cases , the discretisation precision of the latent variable will not effect the average codelength , delegated by the objective in Eq 2 , with consistent discretisation adopted for the latent distributions [ 48 ] . It is worth noting that , since bits-back works more compatibly with ANS , bits-back ANS [ 47 ] is usually the first choice for models learned with the ELBO functional . 3 One-shot online adaptation . In this section , we formulate the setting into three stages , i.e. , 1 ) Pretraining , 2 ) OSOA Encoding and 3 ) OSOA Decoding , and summarise the OSOA Encoding / Decoding stages in Alg . 1 . To unify the narrative of FIFO style and FILO style entropy coders , we aggregate the operations in Stage 2 & 3 into two modules , namely the dynamical system module D and the codec module C. The deterministic dynamical system module D contains an optimiser , which can be a gradient based optimiser or a learning based meta optimiser , and associated parameters , to update a model pθt−1 ( x ) to pθt ( x ) with a data batch Bt . The codec module C contains an entropy coder , a codeword variable and methods to modify the codeword with the entropy coder , cf . following discussions and Alg . 2 . Algorithm 1 One Shot Online Adaptation : Encoding and Decoding OSOA Encoding 1 : Inbuilt attributes : pθ0 ( x ) , D , C 2 : Input : data { xi } Ni=1 3 : Form data batches { Bt } Tt=1 from data { xi } Ni=1 4 : for t = 1 to T do 5 : C.encode_or_cache ( pθt−1 ( x ) , Bt ) ; 6 : pθt ( x ) = D ( pθt−1 ( x ) , Bt ) ; 7 : end for 8 : Output : cf from C. OSOA Decoding 1 : Inbuilt attributes : pθ0 ( x ) , D , C 2 : Input : code cf 3 : Initialise C with cf 4 : for t = 1 to T do 5 : Bt ← C.decode ( pθt−1 ( x ) ) ; 6 : pθt ( x ) = D ( pθt−1 ( x ) , Bt ) ; 7 : end for 8 : Output : data { xi } Ni=1 from { Bt } Tt=1 Stage 1 : Pretraining A base model pθ0 ( x ) is obtained by solving the problem defined in Eq . 1 or Eq . 2 on a pretraining dataset . Stage 2 : OSOA Encoding A sequence of distributions { pθt ( x ) } Tt=1 will be unrolled through the adaptation procedure by the deterministic dynamical system D with batches { Bt } Tt=1 . Specifically , at each step t , one 1 ) inputs the model pθt−1 ( x ) and the data batch Bt to the codec module and 2 ) performs the dynamical system D on pθt−1 ( x ) with Bt and obtains the updated model pθt ( x ) . In 1 ) , the method C.encode_or_cache determines whether to directly encode the current batch ( FIFO style ) or save the needed information , i.e. , the variable values and associated probabilistic mass function ( pmf ) values , for asynchronous encoding ( FILO style ) , cf . Alg . 2 . At the end of OSOA Encoding , the final state of the codeword variable cf will be output from the codec module . Stage 3 : OSOA Decoding One uses the codeword cf from the OSOA Encoding stage to initialise the codeword variable in C and starts the iterative decoding process as follows . At each step t , one 1 ) uses the method C.decode to decode Bt from the codeword variable with pθt−1 ( x ) and 2 ) updates the model with the decoded batch Bt by the dynamical system D. As a result , the same sequence of distributions { pθt ( x ) } Tt=1 will be unrolled and all the data batches { Bt } Tt=1 will also be decoded . Methods in the codec module To recover the exact distribution sequence { pθt ( x ) } Tt=1 in the decoding stage , one has to decode the data batches in the order of { B1 , B2 , ... , BT } . For DGMs based algorithms with an FIFO style coder , e.g. , AC , one can consecutively encode the batch sequence { Bt } Tt=1 in the original order and recover the correct order in the decoding stage . However , for FILO style coders , e.g. , ANS , one has to temporarily save associated pmf values for encoding purpose in cache and encode them with the entropy coder in the reverse order . To reduce the size of cache memory , one can collect m consecutive batches into one chunk and reversely compress the batches within the chunk , cf . Fig . 1 . In this way , one effectively compresses different chunks into different code files . We refer readers to Appendix A for complete demos of OSOA with FIFO and FILO style coders . To distinguish the difference between the implementations of the encode_or_cache method in Alg . 1 for coders of these two styles , we summarise the processes in Alg . 2 . Note that coder.encode in Alg . 2 and C.decode in Alg . 1 are respectively encoding and decoding functions already implemented with DGMs for lossless compression [ 34 , 48 , 49 , 26 ] . Due to this flexibility , one can see that OSOA is generally applicable to DGMs based lossless compression algorithms . Algorithm 2 The encode_or_cache method in OSOA Encoding FIFO coders 1 : Inbuilt attributes : coder 2 : Input : pθt−1 ( x ) , Bt 3 : coder.encode ( pθt−1 ( x ) , Bt ) ; FILO coders 1 : Inbuilt attributes : coder , cache , m 2 : Input : pθt−1 ( x ) , Bt 3 : cache.append ( ( pθt−1 ( x ) , Bt ) ) ; 4 : if len ( cache ) = m or last_batch then 5 : for p ( x ) , B in reversed ( cache ) do 6 : coder.encode ( p ( x ) , B ) ; 7 : end for 8 : cache.empty ( ) 9 : end if Example : Vanilla OSOA Following formulations of the proposed setting , we present the most straightforward strategy to conduct OSOA , referred to as vanilla OSOA . Specifically , the base model pθ0 ( x ) is pretrained on the same data modality , e.g. , HiLLoC [ 49 ] pretrained on CIFAR10 for image compression ; and the deterministic dynamical system D is a gradient based optimiser , e.g. , AdaMax [ 22 ] . Stage 1 . Find or train base model pθ0 ( x ) for certain data modality . Stage 2 . Perform OSOA with base model from the last stage . The dataset to compress S should be first split to T batches , i.e. , S = { Bt } Tt=1 . For 1 ≤ t ≤ T , we first use pθt−1 ( x ) to compute pmf values to encode the t-th batch . Then we use the t-th batch with the specific optimiser and certain learning rate to update and get the new model pθt ( x ) . Stage 3 . Perform OSOA with base model from Stage 1 . The data splitting strategy is the same as Stage 2 . For 1 ≤ t ≤ T , we first use pθt−1 ( x ) to decode the t-th batch from the code . Then we use the decoded batch with the same D ( e.g . optimiser , learning rate , etc . ) as Stage 2 to update and get the new model pθt ( x ) . Note the entropy coding part is omitted for narrative simplicity , and we refer readers to Appendix A for more details .

Summary: This paper presents a method for adapting a learned compression model to improve the lossless compression rate over a set of data to be transmitted, e.g. a set of images. The method assumes a generative model has been optimized over a large dataset, and several models are explored empirically including HiLLoC and IDF++. The model represents a distribution over data items that can be used as the entropy model for a lossless compression algorithm (Huffman coding, arithmetic coding, ANS, etc.). The core idea in OSOA is to divide the data to be compressed into separate batches and to adapt the pre-trained model on each batch before coding the next batch.  The authors show that this approach reduces total bit rate, outperforms training and transferring a separate model for each batch (at least in several cases), that the method can be adapted to FIFO and FILO algorithms (e.g. range coding vs ANS), and that there is a trade-off between runtime and compression rate based on early stopping of the update algorithm.