Online Learned Continual Compression with Stacked Quantization Modules

1 INTRODUCTION . Interest in machine learning in recent years has been fueled by the plethora of data being generated on a regular basis . Effectively storing and using this data is critical for many applications , especially those involving continual learning . In general , compression techniques can greatly improve data storage capacity , and , if done well , reduce the memory and compute usage in downstream machine learning tasks ( Gueguen et al. , 2018 ; Oyallon et al. , 2018 ) . Thus , learned compression has become a topic of great interest ( Theis et al. , 2017 ; Ballé et al. , 2016 ; Johnston et al. , 2018 ) . Yet its application in reducing the size of datasets bound for machine learning applications has been limited . This work focuses on the following familiar setting : new training data arrives continuously for a learning algorithm to exploit , however this data might not be iid , and furthermore there is insufficient storage capacity to preserve all the data uncompressed . We may want to train classifiers , reinforcement learning policies , or other models continuously from this data , or use samples randomly drawn from it at a later point for a downstream task . For example , an autonomous vehicle ( with bounded memory ) collects large amounts of high-dimensional training data ( video , 3D lidar ) in a non-stationary environment ( e.g . changing traffic patterns ) , and overtime applies an ML algorithm to improve its behavior using this data . This data might be transferred at a later point for use in downstream supervised learning . Current standard learned compression algorithms , e.g . Torfason et al . ( 2018 ) , are not well designed to deal with this case , as their convergence speed is too slow to be usable in an online setting . In the field of continual/lifelong learning ( Thrun & Mitchell , 1995 ) , which has for now largely focused on classification , approaches based on storing memories for later use have emerged as some of the most effective in online settings ( Lopez-Paz et al. , 2017 ; Aljundi et al. , 2018 ; Chaudhry et al . ; 2019 ; Aljundi et al. , 2019 ) . These memories can be stored as is , or via a generative model ( Shin et al. , 2017 ) . Then , they can either be used for rehearsal ( Chaudhry et al. , 2019 ; Aljundi et al. , 2019 ) or for constrained optimization ( Lopez-Paz et al. , 2017 ; Chaudhry et al . ; Aljundi et al. , 2018 ) . Indeed many continual learning applications would be nearly solved with replay approaches if one could afford to store all samples . These approaches are however inherently limited by the amount of data that can be stored1 . Learning a generative model to compress the previous data stream thus seems like an appealing idea . However , learning generative models , particularly in the online ( possibly non-stationary ) setting , continues to be challenging , and can greatly increase the complexity of the continual learning task . Furthermore , such models are susceptible to catastrophic forgetting Aljundi et al . ( 2019 ) . An alternate approach is to simply learn a compressed representation of the data ; this is typically faster and more stable than learning to generate the whole data distribution . While the learned compression may itself exhibit forgetting and representation drift , causing challenges for continual and online cases , a learned compression method that can learn continuously and online would allow the storing of far larger amount of samples for replay . In this work we investigate the use of quantized autoencoders , specifically the VQ-VAE framework ( van den Oord et al. , 2017 ) , observing that these can learn continuously and online with minimal forgetting , particularly when augmented with their own internal rehearsal mechanisms . We propose a multi-level stacked model that allows the compressor to adaptively store samples at different compression scales , based on the amount of data , storage capacity , and effectiveness of the model in compressing samples . Furthermore , the learned compressed representation allows multiple continual learning models to be trained from the same data . The main contributions in this work are as follows : ( a ) we introduce and highlight the importance of the online continual learned compression problem ; ( b ) we demonstrate how Multi-level VQVAE , combined with internal replay , can effectively learn compressed representations of online data , ( c ) we show online learned compression can yield state-of-the-art performance in standard online continual image classification benchmarks . 2 RELATED WORK . Learned compression has been recently studied for the specific case of image compression . Work by Theis et al . ( 2017 ) ; Ballé et al . ( 2016 ) ; Johnston et al . ( 2018 ) have shown learned compressions can outperform standard compression algorithms like JPEG . Some of these methods however are challenging to train , thus in our work we focus on the VQ-VAE approach van den Oord et al . ( 2017 ) ; Razavi et al . ( 2019 ) which allows us to address online continual learning settings and permit a multilevel storage . Continual Learning research currently focuses on overcoming catastrophic forgetting ( CF ) in the supervised learning setting , with some limited work in the generative modeling and reinforcement learning settings . Most continual learning methods can be grouped into three major families . Some algorithms dynamically change the model ’ s architecture to incorporate learning from each task separately . Popular methods are Rusu et al . ( 2016 ) , Li & Hoiem ( 2018 ) and Fernando et al . ( 2017 ) . Although these methods can perform well in practice , their introduction of task-specific weights requires growing compute and memory costs which are problematic for the online setting . Another set of techniques employ regularization to constrain weights updates in the hope of maintaining knowledge from previous tasks . Notable methods in this class include ( Kirkpatrick et al. , 2017 ; Huszár , 2017 ; Zenke et al. , 2017 ; Nguyen et al. , 2017 ; Chaudhry et al. , 2018 ) . This set of approaches is inefficient in the online setting Chaudhry et al . ( 2019 ) . The last family of methods encapsulates all that have a mechanism to store information about the previous data distributions . This memory then serves as a tool for the continual learner to rehearse previous tasks . The simplest instantiation of this method is to keep and sample from a buffer of old data to retrain the model after every update ( Chaudhry et al. , 2019 ) . This approach is widely used in RL where it is known as Experience Replay ( ER ) ( Lin , 1993 ) . Another method , known as Generative Replay ( GR ) ( Shin et al. , 2017 ) , uses generative modeling to store past tasks distributions . The continual learner then trains on generated samples to alleviate CF . Other notable examples are Gradient Episodic Memory ( GEM ) ( Lopez-Paz et al. , 2017 ) , iCarl ( Rebuffi et al. , 2017 ) , and Maximally Interfered Retrieval ( MIR ) Aljundi et al . ( 2019 ) , as well as ( Aljundi et al. , 2018 ; Hu et al. , 1This is typically not feasible for on-device machine learning . And while many ML services currently run on the cloud , the move towards increasingly higher data privacy standards is likely to push many ML algorithms to run locally on device . 2018 ) . Most closely related to our work Riemer et al . ( 2017 ) consider compressing memories for use in the continual classification task . They also employ a discrete latent variable model but with the Gumbel approximation which shows to be far less effective than our approach . Most notably a separate offline iid pre-training step for the learned compression is required in order to surpass the ER baseline , distinctly different from the online continual compression we consider . Lidar compression is considered in Tu et al . ( 2019 ) and Caccia et al . ( 2018 ) . Both approaches use a similar projection from 3D ( x , y , z ) coordinates to 2D cylindrical coordinates , and leverage deep generative models to compress the data . However , neither method was designed to account for potential distribution shift , nor for online learning . In this work we show that reusing this 2D projection in conjunction with out model allows us to mitigate the two issues above for lidar data . Algorithm 1 : SQM Replay Input : Learning rates αcls , αae 1 Initialize : MemoryM ; θae 2 for t ∈ 1 .. T do 3 % Fetch data from current task 4 for Binc ∼ Dt do 5 for n ∈ 1 .. N do 6 B ← Binc 7 if t > 1 then 8 % Fetch data from buffer 9 Bre ∼ SAMPLE ( M ) 10 B ← ( Binc , Bre ) 11 end 12 % Train Compressor Network 13 θgen ← ADAM ( Lvq , B , αae ) 14 if t > 1 then UPDATEBUFFERREP ( M , θae ) 15 % Save current indices 16 ADDTOMEMORY ( M , Binc , θae ) 17 end 18 end 19 end Algorithm 2 : AdaptiveCompress Input : datapoint x , SQM with L modules , distortion threshold dth 1 Initialize : MemoryM ; θae 2 % Compute hidden rep. for all blocks 3 ENCODE ( SQM , x ) 4 % Iterate over blocks , from most compressed to least 5 for i ∈ L .. 1 do 6 % Fetch latent indices and hidden rep. computed in forward pass 7 argmini , zq = ARGMIN ( modulei ) 8 % Decode to output space 9 for j ∈ i ... 1 do zq = DECODE ( modulej , zq ) 10 % Test reconstruction error 11 if MSE ( x , zq ) < dth then return argmini , i 12 end 13 % If no compression is good enough , return original image 14 return x , 0 3 METHODOLOGY . In this section we outline our approach to the online continual compression problem . Our learned compression network consists of a set of Multi-resolution VQ-VAE blocks . These blocks only communicate information forward and are not learned jointly unlike the architecture presented in Razavi et al . ( 2019 ) . Memories are compressed using an adaptive scheme that controls what resolution the sample is stored , and therefore how compressed the sample should be . Furthermore memories can be revisited and further compressed as the learned compression module improves . Finally a rehearsal phase that utilizes the stored memory is used to minimize forgetting and update representations stored in the memory . 3.1 PROBLEM SETTING . We consider the problem setting where a stream of samples x ∼ Dt arrives from different distributions Dt changing over time t = 1 . . . T . We have a fixed storage capacity of C bytes where we would like to store the most representative information from all data distributions D1 , ... DT . There is notably a trade-off in quality of information versus samples stored . We propose to use a learned compression model , and most crucially , this model must also be stored within the C bytes , to encode and decode the data samples . Another critical requirement is that at anytime t the content of the storage ( data and/or compression model ) be usable for downstream applications . A key challenge is that the learned compression module will change over time , while we still need to be able to decode the memories in storage . The high level training of the online learned compression is described in Algo 1 . Random memories are decoded and used to train the current compression module , at the same time this also allows us to re-encode those memories using the update weights . The approach also incorporates a form of replay to reduce forgetting and address drift , whereby the memories ’ quality degrade away from their original values . In the rest of this section we discuss the architecture , objective and storage we propose . 3.2 VECTOR QUANTIZED VAE Variational Autoencoders ( VAE ) ( Kingma & Welling , 2013 ) consist of two parts : the encoder network parameterizes the posterior distribution q ( z|x ) and the decoder network p ( x|z ) aims to reconstruct the original input x from the inferred latent variables z . In a standard VAE the prior and posterior distributions are usually Gaussian with diagonal covariance . On the other hand Vector Quantized Autoencoders ( VQ-VAE ) use a discrete latent representation instead ( van den Oord et al. , 2017 ) . This model additionally keeps an embedding table E ∈ RK×D , consisting of K vectors of size D. Given an input ( e.g . an RGB image ) , the encoder first encodes it as a Hh×Wh×D tensor , where Hh and Wh denote the height and width of the latent representation . Then , every D dimensional vector goes through a discretization bottleneck using a nearest-neighbor lookup on the embedding table . Specifically , zij = argmine∈E ||enc ( x ) ij − e||2 . The output of the discretization step is then fed through the decoder . The gradient of this non-differentiable step is approximated using the straight-through estimator . A key property to notice is that to reconstruct the input , only the Hh × Wh indices are required , thus yielding very powerful compression ( van den Oord et al. , 2017 ) . The full VQ-VAE objective , Lvq is given in van den Oord et al . ( 2017 ) .

Summary: This work contributes to introducing a problem called Online Continual Compression. This problem requires to avoid catastrophic forgetting and learn in an online way. Generative methods should be one of the popular ways to do continual learning. This work’s model can be categorized into this clue since it also aims to save samples from old tasks by learning a generative model. In this way, the generator plays a similar role Experience Replay (ER) (here is called Generative Replay). The main core of this work should be the stacked quantization modules (SQM) which can be regarded as a hierarchical variant of the VQ-VAE model. In their SQM, hidden encodings z_q^i will be encoded and its input is z_q^{i-1} which is from previous layer.