BooVAE: Boosting Approach for Continual Learning of VAE

1 Introduction . Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) are deep generative models used in various domains ( Lee et al. , 2017 ; Zhou et al. , 2020 ) . The VAE model consists of deep neural networks ( DNNs ) : an encoder ( inference network ) and decoder ( generative network ) . DNNs are known to reduce their quality on previously learned tasks when trained on data from a new task . Several directions to address this problem of catastrophic forgetting were suggested . But this phenomenon is mainly considered without attention to the specific properties of the VAE . We want to discuss current approaches to continual learning and formulate requirements for the ideal solution . Dynamic architecture approach adds task-specific last layers ( multi-heads ) to encoder and decoder for each task ( Rusu et al. , 2016 ; Nguyen et al. , 2018 ; Li & Hoiem , 2018 ) . We suppose that practical applications of VAE require the common latent space , which is violated in multi-heads . Using multiheads requires deciding which head to apply to the new data and when to expand the architecture . They reduce reuse of similarities between tasks . Hence , we suppose that the approach should keep the static architecture for both encoder and decoder . The generative replay ( Shin et al. , 2017 ; Rao et al. , 2019 ) uses a `` teacher '' generative model for generating `` fake '' data that mimics former training examples . Then the `` student '' model is trained on joint `` fake '' and new data . This approach is conceptually simple , model-agnostic and overcomes forgetting . However , these benefits come with a computational price . We need to retrain model while generating the dataset from all the past tasks , asses samples quality and the task-balance . Thus the approach should avoid generative replay ) . ⇤These authors contributed equally †The work was done while the author was at the Skoltech , Moscow , 2019 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Weight penalty approach ( Liu et al. , 2018 ; Kirkpatrick et al. , 2017 ) forms the trust-region around the optimum of the previous task for protecting parameters . This approach preserves the architecture and avoids generative replay . However , for DNNs a change in the weights is a poor proxy for the difference in the outputs ( Benjamin et al. , 2018 ) . It is an even more critical issue for VAE model as it consists of a pair DNNs : encoder and decoder . Thus the approach should link the data-space and the latent-space . We propose a novel continual learning approach for the VAEs . For each task , we expand current prior to get the approximation of an aggregated posterior over the whole data . We parametrise approximation as an additive mixture of distributions induced by encoder evaluated at trainable pseudo-inputs . These pseudo-inputs link the data-space and latent-space and help to memorise knowledge about past tasks . The problem of matching the aggregated posterior is ill-posed since we observe only its empirical version . As a solution , we use a greedy boosting-like approach with entropy regularisation . This method encourages components in the learned approximation to be diverse , which is essential as we aim at memorising the current task with the fewest components possible . The proposed approach is orthogonal to other methods mentioned above and can be applied in combination with them . Our main contributions are the following : • We relate the approximation of optimal prior , the aggregated posterior and the continual learning task for VAE model . We find optimal additive perturbation in order to approximate optimal prior distribution . We derive the algorithm of effective approximation of the optimal prior for the continual learning framework . • We use this result and present Boosting Approach for Continual Learning of VAE ( BooVAE ) , a framework for training VAE models in the continual framework with static architecture . • We empirically validate the proposed algorithm on commonly used benchmarks ( MNIST , Fashion-MNIST , NotMNIST ) and CelebA for disjoint sequential image generation tasks . The proposed generative model could be efficiently used in a generative replay for discriminative models . We train both generative and discriminative models incrementally , avoiding retraining the generative model for each task from scratch or storing several generative models . We provide code at https : //github.com/AKuzina/BooVAE . 2 Background . Variational Autoencoders ( VAEs ) Let p✓ ( x , z ) be the joint distribution of observed variable x2RD and hidden latent variable z , with the distribution ⇡ ( z ) . Given the distribution x⇠pe ( x ) we aim to find ✓ which maximizes the marginal log-likelihood Epe ( x ) ⇥ log R p✓ ( x , z ) dz ⇤ . Since the marginal likelihood is often intractable , we solve this optimization problem by the variational inference ( Jordan et al. , 1999 ) . Variational autoencoders ( VAE ) ( Kingma & Welling , 2013 ) amortize inference with q ( z|x ) as variational posterior ( encoder ) and p✓ ( x|z ) as likelihood ( decoder ) . The encoder and decoder are parameterized by neural networks with parameters , ✓ . Nets are optimized simultaneously via maximisation of the evidence lower bound ( ELBO ) objective : L ( ✓ , , ) , E pe ( x ) q ( z|x ) ( log p✓ ( x|z ) ⇡ ( z ) log q ( z|x ) ) Epe ( x ) log Z p✓ ( x , z ) dz . ( 1 ) Typically , instead of the density pe ( x ) we are given a dataset { xn } Nn=1 and consider an empirical distribution 1N PN n=1 xn ( x ) . Optimal Prior and Aggregated Posterior Hoffman & Johnson ( 2016 ) ; Goyal et al . ( 2017 ) discuss the choice of the prior distribution over latent space ⇡ ( z ) and observe that default choice of the Gaussian prior significantly restricts the expressiveness of the model . Tomczak & Welling ( 2018 ) proposes to use a prior that optimizes the ELBO ( 1 ) . The solution of this problem is the aggregated variational posterior : q̂ ( z ) = Epe ( x ) q ( z|x ) . ( 2 ) For an empirical pe ( x ) , q̂ ( z ) = 1N PN n=1 q ( z|xn ) and Tomczak & Welling ( 2018 ) proposes a parametric approximation to avoid over-fitting and reduce computational costs : ⇡ ( z ) = 1K PK k=1 q ( z|uk ) . Parameter { uk } Kk=12RK⇥D is optimized simultaneously with ( ✓ , ) by the ELBO ( 1 ) maximization , and K is a hyper-parameter of the algorithm . The additive structure of the optimal prior ⇡⇤ ( z ) over the dataset points motivates us to consider such a prior distribution for continual learning . Continual learning framework In the continual learning framework , we do not have access to the whole dataset . We define the sequence of tasks to be solved t=1 , . . . , T . Subsets of the data for each task D1 , . . . , DT arrive sequentially and we have access only to the data of the current task . For each task t , we consider the empirical distribution pte ( x ) = 1 |Dt| P xn2Dt xn ( x ) . 3 BooVAE Algorithm ( Proposed Method ) . We start from defining an optimal prior for the VAE model in the continual learning framework . Next , we find the optimal additive expansion of the current prior to match the innovation coming from the new task . We use obtained result and provide a general algorithm for training the VAE model in the continual learning framework . It works as an iterative Minorization-Maximization algorithm . In the minorization step , we expand the current approximation of the prior and learn pseudo-inputs . At the Maximization step , we update parameters of the encoder and the decoder with the prior being fixed . 3.1 Optimal prior in continual learning . In this section , we derive an optimal prior for the VAE model in the continual framework and provide the algorithm to it approximation . We provide skipped technical details in Sup. ( A.1 ) . We start from the following decomposition of the ELBO ( 1 ) : L ( ✓ , , ⇡ ) =Epe ( x ) [ Eq ( z|x ) log p✓ ( x|z ) ( KL [ q ( z|x ) |q̂ ( z ) ] +KL [ q̂ ( z ) |⇡ ( z ) ] ) ] , ( 3 ) where q̂ ( z ) =Epe ( x ) q ( z|x ) is the aggregated variational posterior . As KL-divergence is non-negative , the global maximum of Eq . ( 3 ) over ⇡ is reached when : ⇡ ⇤ ( z ) = q̂ ( z ) =Epe ( x ) q ( z|x ) . ( 4 ) We assume that for the sequence of the two tasks , we can write the data distribution as the discrete mixture of two distributions : pe ( x ) =↵ p1e ( x ) + ( 1 ↵ ) p2e ( x ) , ↵2 ( 0 ; 1 ) . This assumption holds for empirical data distribution , which is of the most interest for us : pe ( x ) = 1N PN n=1 xn ( x ) = |D1| |D1|+|D2|p 1 e ( x ) + |D2| |D1|+|D2|p 2 e ( x ) . Then we can decompose the ELBO ( 1 ) as following : L ( ✓ , , ⇡ ) =E pe ( x ) q ( z|x ) log p✓ ( x|z ) X i=1,2 ↵i ( Epie ( x ) KL [ q ( z|x ) |q̂i ( z ) ] +KL [ q̂i ( z ) |⇡ ( z ) ] ) , ( 5 ) where ↵1=↵ , ↵2=1 ↵ and q̂i ( z ) =Epie ( x ) q ( z|x ) . Keeping only terms with dependence over the prior distribution , we conclude to the optimization problem over the probability density space P : min ⇡ ( z ) 2P ↵ KL [ q̂1 ( z ) |⇡ ( z ) ] + ( 1 ↵ ) KL [ q̂2 ( z ) |⇡ ( z ) ] . ( 6 ) Hence , if we can store the data for both tasks we can use the same optimal prior Ex q ( z|x ) , where expectation is taken with respect to ↵ p1e ( x ) + ( 1 ↵ ) p2e ( x ) . However , in continual setting we don ’ t have access to the p1e ( x ) , neither it is reasonable to store q̂1 ( z ) . Thus after Task 1 , we consider using the approximation of the aggregated variational posterior of the first task q̂a1 ( z ) ⇡ q̂1 ( z ) . This modification of the problem ( 6 ) leads to the following prior : ⇡ 1,2 ( z ) = ↵q̂a1 ( z ) + ( 1 ↵ ) q̂2 ( z ) = arg min ⇡ ( z ) 2P ↵ KL [ q̂a ( z ) |⇡ ( z ) ] + ( 1 ↵ ) KL [ q̂2 ( z ) |⇡ ( z ) ] . ( 7 ) Next , we consider that the optimal prior for the first task ⇡1 ( z ) is a good start for a new prior after observing the new task . Hence , we propose using an additive expansions of the current prior to match innovation induced by the new task . Then the objective of the problem ( 7 ) transforms to : min h2P ↵ KL [ q̂a1 ( z ) | ( 1 ) ⇡1 ( z ) + h ( z ) ] + ( 1 ↵ ) KL [ q̂2 ( z ) | ( 1 ) ⇡1 ( z ) + h ( z ) ] ) . ( 8 ) The optimization problem ( 8 ) is bi-convex over h and . We consider the Functional Frank-Wolfe algorithm ( informally , `` boosting '' ) ( Wang et al. , 2015 ) for the corresponded objective : F↵ , [ q̂a1 ( z ) , q̂2 ( z ) ; ⇡1 ( z ) , h ( z ) ] = =↵ KL [ q̂a1 ( z ) | ( 1 ) ⇡1 ( z ) + h ( z ) ] + ( 1 ↵ ) KL [ q̂2 ( z ) | ( 1 ) ⇡1 ( z ) + h ( z ) ] ) . ( 9 ) We aim to find h as the projection of the functional gradient to the probability density space P and then optimize over with the stochastic gradient method , while h fixed . The functional gradient of the objective ( 8 ) with respect to perturbation h is ↵ q̂ a 1 ( z ) ⇡1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ⇡1 ( z ) in the following sense : F↵ , [ q̂a1 ( z ) , q̂2 ( z ) ; ⇡1 ( z ) , h ( z ) ] = =F↵ [ q̂a1 ( z ) , q̂2 ( z ) ; ⇡1 ( z ) ] ✓Z dz h ( z ) ↵ q̂ a 1 ( z ) ⇡1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ⇡1 ( z ) 1 ◆ +o ( ) . ( 10 ) We consider projection to the probability space as the solution of the optimization problem : max h2P log Z h ( z ) ↵ q̂ a 1 ( z ) ⇡1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ⇡1 ( z ) dz max h2P Z h ( z ) log ↵ q̂ a 1 ( z ) ⇡1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ⇡1 ( z ) dz . ( 11 ) We use the lower bound in the ( 11 ) in order to use Monte-Carlo estimates of the gradients . The problem is linear over h ( z ) and the solution is degenerate distribution . To this end , we add the entropy regularization H [ h ] = R h ( z ) log h ( z ) dz and obtain the final objective : min h2P KL h ↵q̂ a 1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ⇡1 ( z ) . ( 12 ) As far as we update the initial prior r0 ( z ) =⇡1 ( z ) with the mixture component h : r1 ( z ) = ( 1 ) ⇡1 ( z ) + h ( z ) , we can improve current approximation by finding the new component , solving the same optimization problem , but using r1 instead of ⇡1 . The objective ( 12 ) encourages a new component to be different from the already constructed approximation : KL h ↵q̂ a 1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ⇡1 ( z ) =KL [ h|↵q̂a1 ( z ) + ( 1 ↵ ) q̂2 ( z ) ] | { z } match the target prior Z h ( z ) log 1 ⇡1 ( z ) dz | { z } be different from the current approximation . ( 13 ) To this end , we have defined the algorithm of learning optimal prior for VAE in continual setting . Now we will use it to formulate the algorithm of VAE training in continual learning framework .

Summary: The authors propose a novel approach to tackle catastrophic forgetting. In more detail, they focus on variational autoencoders (VAEs) and investigate the role of the prior in the aforementioned setting. This work -- which is built on top of (Tomczak & Welling, 2018) that introduces a learnable prior through pseudo-inputs -- provides a definition of “optimal prior” in continual learning as well as an algorithm to update it when new tasks arrive. The authors evaluate the proposed approach on image generation (MNIST, its variants, and CelebA) and, moreover, as a generative memory that mitigates the forgetting of a discriminative model.