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: This paper proposes the use of pseudo-coresets to accelerate inference in Examplar VAE. The Examplar VAE ELBO (Eq.4/5 in [ExamplarVAE]) is updated to use pseudo-coreset (Eq.8/9 in this paper). Other quantities of interest are updated accordingly and a two-stage approach is proposed to update two types of parameters (VAE and pseudo-coresets).  [ExamplarVAE] Norouzi et al (2020) "Exemplar VAE: Linking Generative Models, Nearest Neighbor Retrieval, and Data Augmentation", NeurIPS 2020.