ByPE-VAE: Bayesian Pseudocoresets Exemplar VAE

1 Introduction . Deep generative models that learn implicit data distribution from the enormous amount of data have received widespread attention to generating highly realistic new samples in machine learning . In particular , due to the utilization of the reparameterization trick and variational inference for optimization , Variational Autoencoders ( VAEs ) [ 1 , 2 ] stand out and have demonstrated significant successes for dimension reduction [ 3 ] , learning representations [ 4 ] , and generating data [ 5 ] . In addition , various variants of VAE have been proposed conditioned on advanced variational posterior [ 6 , 7 , 8 ] or powerful decoders [ 9 , 10 ] . It is worth noting that the prior in the typical VAE is a simple standard normal distribution that is convenient to compute while ignores the nature of the data itself . Moreover , a large number of experiments have empirically demonstrated that simplistic priors could produce the phenomena of over-regularization and posterior collapse , and finally cause poor performance [ 5 , 11 ] . Hence , many researchers have worked to develop more complex priors to enhance the capacity of the variational posterior . In this line , Tomczak et al . [ 12 ] introduces a more flexible prior named VampPrior , which is a mixture of variational posteriors based on pseudo-inputs to alleviate the problems like overfitting ∗Corresponding Author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . and high computational cost . However , the way in which the pseudo-inputs are obtained is not interpretable . Recently , Norouzi et al . [ 13 ] develops an Exemplar VAE with a non-parametric prior based on an exemplar-based method , achieving excellent performance . To ensure the performance and the generation diversity , the exemplar set needs to be large enough , and usually the entire training data is utilized . Obviously , this leads to huge computational complexity . Even though Exemplar VAE further employs approximate nearest neighbor search to speed up the training process , the number of nearest neighbors should be as large as possible to ensure performance . In a nutshell , Exemplar VAE is computationally expensive . To address such issues , we develop a new prior for VAE that is inspired by the paradigm of coresets [ 14 ] . The coreset is a powerful tool aiming to find a small weighted subset for efficiently approximating the entire original dataset . Therefore , rather than using the large-scale training data directly , we seek to design a prior conditioned on a coreset , which greatly reduces the computational complexity and prevents overfitting . In practice , to better incorporate this idea with the framework of VAE , we further employ a specific form of coresets , namely Bayesian pseudocoresets [ 15 ] , which is known as a small weighted subset of the pseudodata points , resulting in a Bayesian pseudocoreset based prior . With this prior , we gain a new variant of VAE called Bayesian Pseudocoresets Exemplar VAE ( ByPE-VAE ) . To sample from the ByPE-VAE , we first take a pseudodata point from the pseudocoreset according to its weight , and then transform it into a latent representation using the learned prior . Then a decoder is used to transform the latent representation into a new sample . A crucial part of ByPE-VAE is to obtain the optimal pseudocoreset . We formulate this process as a variational inference problem where the pseudodata points and corresponding weights are the parameters of variational posterior approximation . More precisely , to seek the optimal pseudocoreset , we minimize the Kullback-Leibler ( KL ) divergence between the prior based on the pseudocoreset and that based on the entire dataset . This processing ensures that the learned prior is actually an approximation of the prior conditioned on whole training data . Thus , it is fundamentally different from general pseudodata based priors , like VampPrior . For optimization , we adopt a two-step alternative search strategy to learn two types of parameters of ByPE-VAE , which refer to the parameters in the VAE framework and the pseudodata points and corresponding weights in the pseudocoreset . In particular , we iteratively optimize one of the two types of parameters while keeping the other one fixed until convergence . Finally , we compare ByPE-VAE with several state-of-the-art VAEs in a number of tasks , including density estimation , representation learning and generative data augmentation . Experimental results demonstrate the effectiveness of ByPE-VAE on Dynamic MNIST , Fashion MNIST , CIFAR10 , and CelebA . Additionally , to validate the efficiency of our model , we measure the running time on a basic VAE architecture . Compared to the Exemplar VAE , ByPE-VAE is up to 3 times speed-up without losing performance on Dynamic MNIST , Fashion MNIST , and CIFAR10 . 2 Preliminaries . Before presenting the proposed model , we first introduce some preliminaries , namely Exemplar VAE and Bayesian pseudocoresets . Throughout this paper , vectors are denoted by bold lowercase letters , whose subscripts indicate their order , and matrices are denoted by upper-case letters . 2.1 Exemplar VAE . Exemplar VAE is regarded as a variant of VAE that integrates the exemplar-based method into VAE for the sake of seeking impressive image generation . Specifically , it first draws a random exemplar xn using uniform distribution from the training data X = { xn } Nn=1 ( where N denotes the sample amount ) , then transforms an exemplar xn into latent variable z using an example-based prior rφ ( z | xn ) , finally generates observable data x by a decoder pφ ( x | z ) . The parametric transition distribution Tφ , θ ( x | xn ) of exemplar-based generative models can be expressed as Tφ , θ ( x | xn ) = ∫ z rφ ( z | xn ) pθ ( x | z ) dz , ( 1 ) where φ and θ denote corresponding parameters . Assuming that x is independent to xn conditioned on the latent variable z and marginalizing over the latent variable z , the objective O ( θ , φ ; x , X ) of Exemplar VAE can be formulated as log p ( x ; X , θ , φ ) = log N∑ n=1 1 N Tφ , θ ( x|xn ) ≥ Eqφ ( z|x ) log pθ ( x|z ) − Eqφ ( z|x ) log qφ ( z|x ) ∑N n=1 rφ ( z|xn ) /N = O ( θ , φ ; x , X ) , ( 2 ) where O ( θ , φ ; x , X ) is known as the evidence lower bound ( ELBO ) . From the Eq . ( 2 ) , it can be derived that the difference between Exemplar VAE and typical VAE is the definition of the prior p ( z ) in the second term . The prior of Exemplar VAE is defined as a mixture form , that is p ( z | X ) = ∑N n=1 rφ ( z | xn ) /N . The variational posterior qφ ( z | x ) and the exemplar-based prior rφ ( z | xn ) are assumed to be Gaussian distributions whose parameters are fulfilled by neural networks . Note that , the computational cost of this training process is related to the number of exemplars which usually set to be the entire training data . This indicates that Exemplar VAE is computationally expensive since the amount of training data is generally huge . 2.2 Bayesian Pseudocoresets . Bayesian pseudocoresets is a method of coreset construction based on variational inference , which constructs a weighted set of synthetic “ pseudodata ” instead of the original dataset during inference . First , the goal of this method is to approximate expectations under the posterior π ( ψ ) with the parameter ψ , which is formulated as π ( ψ ) = 1 Z exp ( N∑ n=1 f ( xn , ψ ) ) π0 ( ψ ) , ( 3 ) where f ( · , ψ ) denotes a potential function , and usually is a log-likelihood function . π0 ( ψ ) is the prior , and Z is the normalization constant . Instead of directly inferring the posterior π ( ψ ) , the Bayesian pseudocoreset employs a weighted set of pseudodata points to approximate the true posterior π ( ψ ) , which is given by πU , w ( ψ ) = 1 Z̃U , w exp ( M∑ m=1 wmf ( um , ψ ) ) π0 ( ψ ) , ( 4 ) where U = { um } Mm=1 represents M pseudodata points um ∈ Rd , w = { wm } Mm=1 denotes nonnegative weights , and Z̃U , w is the corresponding normalization constant . Finally , this model obtains the optimal pseudodata points and their weights by minimizing the KL divergence , as follows , U ? , w ? = argmin U , w DKL ( πU , w‖π ) . ( 5 ) This formulation can reduce the computational cost by decreasing data redundancy . 3 Bayesian Pseudocoresets Exemplar VAE . To generate a new observation x , the Exemplar VAE requires a large collection of exemplars from X = { xn } Nn=1 to guide the whole process , as shown in Eq . ( 2 ) . One can see that the greater the number of exemplars set , the richer the prior information can be obtained . In practice , to ensure performance , the number of exemplars is relatively large which is generally set to the size of the entire training data . This leads to huge computational costs in the training process of the Exemplar VAE . To overcome such a issue , inspired by the paradigm of Bayesian pseudocoresets , we adopt M pseudodata points U = { um } Mm=1 with corresponding weights w = { wm } Mm=1 to denote exemplars , importantly M N . That is , the original exemplars X are approximated by a small weighted set of pseudodata points known as a Bayesian pseudocoreset . The framework can be expressed as log p ( x | X , θ ) = log N∑ n=1 1 N Tθ ( x | xn ) ≈ log M∑ m=1 wm N Tθ ( x | um ) = log p ( x | U , w , θ ) , ( 6 ) wherewm ≥ 0 ( m = 1 , · · · , M ) and ‖w‖1 = N . Further , we integrate this approximated framework with VAE by introducing a latent variable z . Then the parametric function is given by Tφ , θ ( x | um ) = ∫ z rφ ( z | um ) pθ ( x | z ) dz , ( 7 ) where rφ ( z | um ) with parameter φ denotes a pseudodata based prior for generating z from a pseudodata point um . pθ ( x | z ) with parameter θ represents the decoder for generating the observation x from z . Similarly , we assume that an observation x is independent from a pseudodata point um conditional on z to simplify the formulation and optimization . In general , we desire to maximize the marginal log-likelihood log p ( x ) for learning , however , this is intractable since we have no ability to integrate the complex posterior distributions out . Now , we focus on maximizing the evidence lower bound ( ELBO ) derived by Jensen ’ s inequality , as follows , log p ( x ; U , w , θ , φ ) = log M∑ m=1 wm N Tφ , θ ( x | um ) = log M∑ m=1 wm N ∫ z rφ ( z | um ) pθ ( x | z ) dz ( 8 ) ≥ Eqφ ( z|x ) log pθ ( x | z ) − Eqφ ( z|x ) log qφ ( z | x ) ∑M m=1 wmrφ ( z | um ) /N ≡ O ( θ , φ , U , w ; x ) , ( 9 ) where qφ ( z | x ) represents the approximate posterior distribution . And O ( θ , φ , U , w ; x ) is defined as the objective function of the ByPE-VAE to optimize parameters θ and φ . The specific derivation can be found in Supp.A . As we can see from Eq . ( 9 ) , the difference between the ByPE-VAE and other variants of VAE is the formulation of the prior p ( z ) in the second term . In detail , the prior of the ByPE-VAE is a weighted mixture model prior , in which each component is conditioned on a pseudodata point and the corresponding weight , i.e. , p ( z|U , w ) = ∑M m=1 wmrφ ( z|um ) /N . In contrast , the prior of the Exemplar VAE is denoted as p ( z|X ) = ∑N n=1 rφ ( z|xn ) /N . As shown in Eq . ( 9 ) , the ByPE-VAE includes two encoder networks , namely qφ ( z | x ) and rφ ( z | um ) . And the distributions of qφ ( z | x ) and rφ ( z | um ) are both designed as Gaussian distributions . According to the analysis in [ 16 ] , the optimal prior is the form of aggregate posterior . Inspired by this report , we make the prior be coupled with the variational posterior , as follows , qφ ( z | x ) = N ( z | µφ ( x ) , Λφ ( x ) ) , ( 10 ) rφ ( z | um ) = N ( z | µφ ( um ) , σ2I ) . ( 11 ) We employ the same parametric mean function µφ for two encoder networks for better incorporation . However , the covariance functions of the two encoder networks are different . Specifically , the variational posterior uses a diagonal covariance matrix function Λφ , while each component of the mixture model prior uses an isotropic Gaussian with a scalar parameter σ2 . Note that the scalar parameter σ2 is shared by each component of the mixture model prior for effective computation . Then , we can express the log of the weighted pseudocoreset based prior log pφ ( z | U , w ) as log pφ ( z | U , w ) = −dz log ( √ 2πσ ) − logN + log M∑ m=1 wm exp −‖z− µφ ( um ) ‖2 2σ2 , ( 12 ) where dz denotes the dimension of z . Based on the formulation of Eq . ( 12 ) , we can further obtain the objective function of the ByPE-VAE , as follows , O ( θ , φ , U , w ; X ) = Eqφ ( z|x ) [ log pθ ( x | z ) qφ ( z | x ) + log M∑ m=1 wm ( √ 2πσ ) dz exp −‖z− µφ ( um ) ‖2 2σ2 ] , ( 13 ) ∝ N∑ i=1 Eqφ ( z|xi ) [ log pθ ( xi | z ) qφ ( z | xi ) + log M∑ m=1 wm ( √ 2πσ ) dz exp −‖z− µφ ( um ) ‖2 2σ2 ] , ( 14 ) where the constant − logN is omitted for convenience . For Eqφ ( z|xi ) , we employ the reparametrization trick to generate samples . Note that , for the standard VAE with a Gaussian prior , the process of generating new observations involves only the decoder network after training . However , to generate a new observation from ByPE-VAE , we not only require the decoder network , but also the learned Bayesian pseudocoreset and a pseudodata point based prior rφ . We summarize the generative process of the ByPE-VAE in Algorithm 1 . Algorithm 1 : The Generative Process of ByPE-VAE Input : Pseudocoreset { U , w } , Decoder pθ , Prior rφ . Output : A generated observation x. Step.1 Sample um ∼ Multi ( · | U , wN ) for obtaining a pseudodata point um from the learned pseudocoreset . Step.2 Sample z ∼ rφ ( · | um ) using the pseudodata point based prior rφ for obtaining the latent representation z. Step.3 Sample x ∼ pθ ( · | z ) using the decoder pθ for generating a new observation x . However , more importantly , the whole process above holds only if the pseudocoreset is an approximation of all exemplars or training data . To ensure this , we first re-represent frameworks of pφ ( z | X ) and pφ ( z | U , w ) from the Bayesian perspective . Concretely , pφ ( z | X ) and pφ ( z | U , w ) are also viewed as posteriors conditioned on the likelihood function pθ ( · | z ) and a certain prior p0 ( z ) , as follows , pφ ( z | X ) = 1 Z exp ( N∑ n=1 log pθ ( xn | z ) ) p0 ( z ) , ( 15 ) pφ ( z | U , w ) = 1 Z̃U , w exp ( M∑ m=1 wm log pθ ( um | z ) ) p0 ( z ) , ( 16 ) where Z and Z̃U , w are their respective normalization constants , pθ ( · | z ) here specifically refers to the decoder . Then , we develop this problem into a variational inference problem , where the pseudodata points and the corresponding weights are the parameters of the variational posterior approximation followed [ 15 ] . Specifically , we construct the pseudocoreset by minimizing the KL divergence in terms of the pseudodata points and the weights , as follows , U ? , w ? = argmin U , w DKL ( pφ ( z | U , w ) ‖pφ ( z | X ) ) . ( 17 ) Further , the gradients of DKL in Eq . ( 17 ) with respect to the pseudodata point um and the weights w are given by ∇umDKL = −wm CovU , w [ ∇U log pθ ( um | z ) , log pθ ( X | z ) T1N − log pθ ( U | z ) Tw ] , ( 18 ) ∇wDKL = −CovU , w [ log pθ ( U | z ) , log pθ ( X | z ) T1N − log pθ ( U | z ) Tw ] , ( 19 ) where CovU , w denotes the covariance operator for the pφ ( z | U , w ) , and 1N ∈ RN represents the vector of all 1 entries . In practice , we adopt a black-box stochastic algorithm to obtain the optimal pseudodata points and weights . Details of these derivations are provided in Supp.B and C. Note that ByPE-VAE involves two types of parameters , namely the parameters θ and φ in the VAE framework and the parameters U and w in the pseudocoreset . For optimization , we use a two-step alternative optimization strategy . In detail , ( i ) update θ and φ with fixed pseudocoreset { U , w } , and ( ii ) update U and w with fixed θ and φ . Steps ( i ) and ( ii ) are iteratively implemented until convergence . The detailed optimization algorithm is shown in Algorithm 2 . One can see that the computation does not scale with N , but rather with the number of pseudocoreset points M , which greatly reduces the computational complexity and also prevents overfitting . Also note that , rather than be updated every epoch , the pseudocoreset { U , w } is updated by every k epochs . And k is set to 10 in the experiments .

Summary: The paper tackles the important problem of finding suitable learnable prior for VAE. The paper provides some improvements of Exemplar VAE that increase the efficiency of the model. The proposed model is simply aiming at reducing the examples for nearest neighbours used by Exemplar VAE by application of bayesian pseudocoresets  for efficiently approximating the entire original dataset. The proposed approach was evaluated using standard benchmark datasets in a number of tasks, including density estimation, representation learning and generative data augmentation.