VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Models

1 INTRODUCTION . Deep generative learning is a central problem in machine learning . It has found diverse applications , ranging from image ( Brock et al. , 2018 ; Karras et al. , 2019 ; Razavi et al. , 2019 ) , music ( Dhariwal et al. , 2020 ) and speech ( Ping et al. , 2020 ; Oord et al. , 2016a ) generation , distribution alignment across domains ( Zhu et al. , 2017 ; Liu et al. , 2017 ; Tzeng et al. , 2017 ) and semi-supervised learning ( Kingma et al. , 2014 ; Izmailov et al. , 2020 ) to 3D point cloud generation ( Yang et al. , 2019 ) , light-transport simulation ( Müller et al. , 2019 ) , molecular modeling ( Sanchez-Lengeling & AspuruGuzik , 2018 ; Noé et al. , 2019 ) and equivariant sampling in theoretical physics ( Kanwar et al. , 2020 ) . Among competing frameworks , likelihood-based models include variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) , normalizing flows ( Rezende & Mohamed , 2015 ; Dinh et al. , 2016 ) , autoregressive models ( Oord et al. , 2016b ) , and energy-based models ( EBMs ) ( Lecun et al. , 2006 ; Salakhutdinov et al. , 2007 ) . These models are trained by maximizing the data likelihood under the model , and unlike generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) , their training is usually stable and they cover modes in data more faithfully by construction . Among likelihood-based models , EBMs model the unnormalized data density by assigning low energy to high-probability regions in the data space ( Xie et al. , 2016 ; Du & Mordatch , 2019 ) . EBMs are appealing because they require almost no restrictions on network architectures ( unlike normalizing flows ) and are therefore potentially very expressive . They also exhibit better robustness and out-of-distribution generalization ( Du & Mordatch , 2019 ) because , during training , areas with high probability under the model but low probability under the data distribution are penalized explicitly . However , training and sampling EBMs usually requires MCMC , which can suffer from slow mode mixing and is computationally expensive when neural networks represent the energy function . ∗Work done during an internship at NVIDIA On the other hand , VAEs are computationally more efficient for sampling than EBMs , as they do not require running expensive MCMC steps . VAEs also do not suffer from expressivity limitations that normalizing flows face ( Dupont et al. , 2019 ; Kong & Chaudhuri , 2020 ) , and in fact , they have recently shown state-of-the-art generative results among non-autoregressive likelihood-based models ( Vahdat & Kautz , 2020 ) . Moreover , VAEs naturally come with a latent embedding of data that allows fast traverse of the data manifold by moving in the latent space and mapping the movements to the data space . However , VAEs tend to assign high probability to regions with low density under the data distribution . This often results in blurry or corrupted samples generated by VAEs . This also explains why VAEs often fail at out-of-distribution detection ( Nalisnick et al. , 2019 ) . In this paper , we propose a novel generative model as a symbiotic composition of a VAE and an EBM ( VAEBM ) that combines the best of both . VAEBM defines the generative distribution as the product of a VAE generator and an EBM component defined in pixel space . Intuitively , the VAE captures the majority of the mode structure in the data distribution . However , it may still generate samples from low-probability regions in the data space . Thus , the energy function focuses on refining the details and reducing the likelihood of non-data-like regions , which leads to significantly improved samples . Moreover , we show that training VAEBM by maximizing the data likelihood easily decomposes into training the VAE and the EBM component separately . The VAE is trained using the reparameterization trick , while the EBM component requires sampling from the joint energy-based model during training . We show that we can sidestep the difficulties of sampling from VAEBM , by reparametrizing the MCMC updates using VAE ’ s latent variables . This allows MCMC chains to quickly traverse the model distribution and it speeds up mixing . As a result , we only need to run short chains to obtain approximate samples from the model , accelerating both training and sampling at test time . Experimental results show that our model outperforms previous EBMs and state-of-the-art VAEs on image generation benchmarks including CIFAR-10 , CelebA 64 , LSUN Church 64 , and CelebA HQ 256 by a large margin , reducing the gap with GANs . We also show that our model covers the modes in the data distribution faithfully , while having less spurious modes for out-of-distribution data . To the best of knowledge , VAEBM is the first successful EBM applied to large images . In summary , this paper makes the following contributions : i ) We propose a new generative model using the product of a VAE generator and an EBM defined in the data space . ii ) We show how training this model can be decomposed into training the VAE first , and then training the EBM component . iii ) We show how MCMC sampling from VAEBM can be pushed to the VAE ’ s latent space , accelerating sampling . iv ) We demonstrate state-of-the-art image synthesis quality among likelihood-based models , confirm complete mode coverage , and show strong out-of-distribution detection performance . 2 BACKGROUND . Energy-based Models : An EBM assumes pψ ( x ) to be a Gibbs distribution of the form pψ ( x ) = exp ( −Eψ ( x ) ) /Zψ , where Eψ ( x ) is the energy function with parameters ψ and Zψ =∫ x exp ( −Eψ ( x ) ) dx is the normalization constant . There is no restriction on the particular form of Eψ ( x ) . Given a set of samples drawn from the data distribution pd ( x ) , the goal of maximum likelihood learning is to maximize the log-likelihood L ( ψ ) = Ex∼pd ( x ) [ log pψ ( x ) ] , which has the derivative ( Woodford , 2006 ) : ∂ψL ( ψ ) = Ex∼pd ( x ) [ −∂ψEψ ( x ) ] + Ex∼pψ ( x ) [ ∂ψEψ ( x ) ] ( 1 ) For the first expectation , the positive phase , samples are drawn from the data distribution pd ( x ) , and for the second expectation , the negative phase , samples are drawn from the model pψ ( x ) itself . However , sampling from pψ ( x ) in the negative phase is itself intractable and approximate samples are usually drawn using MCMC . A commonly used MCMC algorithm is Langevin dynamics ( LD ) ( Neal , 1993 ) . Given an initial sample x0 , Langevin dynamics iteratively updates it as : xt+1 = xt − η 2 ∇xEψ ( xt ) + √ ηωt , ωt ∼ N ( 0 , I ) , ( 2 ) where η is the step-size.1 In practice , Eq . 2 is run for finite iterations , which yields a Markov chain with an invariant distribution approximately close to the original target distribution . 1In principle one would require an accept/reject step to make it a rigorous MCMC algorithm , but for sufficiently small stepsizes this is not necessary in practice ( Neal , 1993 ) . Variational Autoencoders : VAEs define a generative model of the form pθ ( x , z ) = pθ ( z ) pθ ( x|z ) , where z is the latent variable with prior pθ ( z ) , and pθ ( x|z ) is a conditional distribution that models the likelihood of data x given z . The goal of training is to maximize the marginal log-likelihood log pθ ( x ) given a set of training examples . However since the marginalization is intractable , instead , the variational lower bound on log pθ ( x ) is maximized with qφ ( z|x ) as the approximate posterior : log pθ ( x ) ≥ Ez∼qφ ( z|x ) [ log pθ ( x|z ) ] −DKL [ qφ ( z|x ) ‖pθ ( z ) ] : = Lvae ( x , θ , φ ) . ( 3 ) The state-of-the-art VAE , NVAE ( Vahdat & Kautz , 2020 ) , increases the expressivity of both prior and approximate posterior using hierarchical latent variables ( Kingma et al. , 2016 ) where z is decomposed into a set of disjoint groups , z = { z1 , z1 , . . . , zL } , and the prior pθ ( z ) = ∏ l pθ ( zl|z < l ) and the approximate posterior qφ ( z|x ) = ∏ l qφ ( zl|z < l , x ) are defined using autoregressive distributions over the groups . We refer readers to Vahdat & Kautz ( 2020 ) for more details . 3 ENERGY-BASED VARIATIONAL AUTOENCODERS . One of the main problems of VAEs is that they tend to assign high probability to regions in data space that have low probability under the data distribution . To tackle this issue , we propose VAEBM , a generative model constructed by the product of a VAE generator and an EBM component defined in the data space . This formulation allows our model to capture the main mode structure of the data distribution using the VAE . But when training the joint VAEBM , in the negative training phase we sample from the model itself and can discover non-data-like samples , whose likelihood is then reduced by the energy function explicitly . The energy function defined in the pixel space also shares similarities with discriminator in GANs , which can generate crisp and detailed images . Formally , we define the generative model in VAEBM as hψ , θ ( x , z ) = 1Zψ , θ pθ ( x , z ) e −Eψ ( x ) where pθ ( x , z ) = pθ ( z ) pθ ( x|z ) is a VAE generator and Eψ ( x ) is a neural network-based energy function , operating only in the x space , and Zψ , θ = ∫ pθ ( x ) e −Eψ ( x ) dx is the normalization constant . VAEBM is visualized in Fig . 1 . Marginalizing out the latent variable z gives hψ , θ ( x ) = 1 Zψ , θ ∫ pθ ( x , z ) e −Eψ ( x ) dz = 1 Zψ , θ pθ ( x ) e −Eψ ( x ) . ( 4 ) Given a training dataset , the parameters of VAEBM , ψ , θ , are trained by maximizing the marginal log-likelihood on the training data : log hψ , θ ( x ) = log pθ ( x ) − Eψ ( x ) − logZψ , θ ( 5 ) ≥ Ez∼qφ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qφ ( z|x ) ||p ( z ) ) ︸ ︷︷ ︸ Lvae ( x , θ , φ ) −Eψ ( x ) − logZψ , θ︸ ︷︷ ︸ LEBM ( x , ψ , θ ) , ( 6 ) where we replace log pθ ( x ) with the variational lower bound from Eq . 3 . Eq . 6 forms the objective function for training VAEBM . The first term corresponds to the VAE objective and the second term corresponds to training the EBM component . Next , we discuss how we can optimize this objective .

Summary: This paper proposes a model that corrects VAE by an energy-based model defined on image space. The model is learned in two phase. The first phase learns the VAE model, while the second phase learns the EBM correction term by MLE. Experimental results show that the proposed method outperforms pure EBM defined on image space and also pure VAE models by large margins.