No MCMC for me: Amortized sampling for fast and stable training of energy-based models

1 INTRODUCTION . Energy-Based Models ( EBMs ) have recently regained popularity within machine learning , partly inspired by the impressive results of Du & Mordatch ( 2019 ) and Song & Ermon ( 2020 ) on largescale image generation . Beyond image generation , EBMs have also been successfully applied to a wide variety of applications including : out-of-distribution detection ( Grathwohl et al. , 2019 ; Du & Mordatch , 2019 ; Song & Ou , 2018 ) , adversarial robustness ( Grathwohl et al. , 2019 ; Hill et al. , 2020 ; Du & Mordatch , 2019 ) , reliable classification ( Grathwohl et al. , 2019 ; Liu & Abbeel , 2020 ) and semi-supervised learning ( Song & Ou , 2018 ; Zhao et al. ) . Strikingly , these EBM approaches outperform alternative classes of generative models and rival hand-tailored solutions on each task . Despite progress , training EBMs is still a challenging task . As shown in Table 1 , existing training methods are all deficient in at least one important practical aspect . Markov chain Monte Carlo ( MCMC ) methods are slow and unstable during training ( Nijkamp et al. , 2019a ; Grathwohl et al. , 2020 ) . Score matching mechanisms , which minimize alternative divergences are also unstable and most methods can not work with discontinuous nonlinearities ( such as ReLU ) ( Song & Ermon , 2019b ; Hyvärinen , 2005 ; Song et al. , 2020 ; Pang et al. , 2020b ; Grathwohl et al. , 2020 ; Vincent , 2011 ) . Noise contrastive approaches , which learn energy functions through density ratio estimation , typically don ’ t scale well to high-dimensional data ( Gao et al. , 2020 ; Rhodes et al. , 2020 ; Gutmann & Hyvärinen , 2010 ; Ceylan & Gutmann , 2018 ) . ∗Equal Contribution . Code available at github.com/wgrathwohl/VERA In this work , we present a simple method for training EBMs which performs as well as previous methods while being faster and substantially easier to tune . Our method is based on reinterpreting maximum likelihood as a bi-level variational optimization problem , which has been explored in the past for EBM training ( Dai et al. , 2019 ) . This perspective allows us to amortize away MCMC sampling into a GAN-style generator which is encouraged to have high entropy . We accomplish this with a novel approach to entropy regularization based on a fast variational approximation . This leads to the method we call Variational Entropy Regularized Approximate maximum likelihood ( VERA ) . Concretely , we make the following contributions : • We improve the MCMC-based entropy regularizer of Dieng et al . ( 2019 ) with a parallelizable variational approximation . • We show that an entropy-regularized generator can be used to produce a variational bound on the EBM likelihood which can be optimized more easily than MCMC-based estimators . • We demonstrate that models trained in this way achieve much higher likelihoods than methods trained with alternative EBM training procedures . • We show that our approach stabilizes and accelerates the training of recently proposed Joint Energy Models ( Grathwohl et al. , 2019 ) . • We show that the stabilization of our approach allows us to use JEM for semi-supervised learning , outperforming virtual adversarial training when little prior domain knowledge is available ( e.g. , for tabular data ) . 2 ENERGY BASED MODELS . An energy-based model ( EBM ) is any model which parameterizes a density as pθ ( x ) = efθ ( x ) Z ( θ ) ( 1 ) where fθ : RD → R and Z ( θ ) = ∫ efθ ( x ) dx is the normalizing constant which is not explicitly modeled . Any probability distribution can be represented in this way for some fθ . The energy-based parameterization has been used widely for its flexibility , ease of incorporating known structure , and relationship to physical systems common in chemistry , biology , and physics ( Ingraham et al. , 2019 ; Du et al. , 2020 ; Noé et al. , 2019 ) . The above properties make EBMs an appealing model class , but because they are unnormalized many tasks which are simple for alternative model classes become challenging for EBMs . For example , exact samples can not be drawn and likelihoods can not be exactly computed ( or even lowerbounded ) . This makes training EBMs challenging as we can not simply train them to maximize likelihood . The most popular approach to train EBMs is to approximate the gradient of the maximum likelihood objective . This gradient can be written as : ∇θ log pθ ( x ) = ∇θfθ ( x ) − Epθ ( x′ ) [ ∇θfθ ( x ′ ) ] . ( 2 ) MCMC techniques are used to approximately generate samples from pθ ( x ) ( Tieleman , 2008 ) . Practically , this approach suffers from poor stability and computational challenges from sequential sampling . Many tricks have been developed to overcome these issues ( Du & Mordatch , 2019 ) , but they largely still persist . Alternative estimators have been proposed to circumvent these challenges , including score matching ( Hyvärinen , 2005 ) , noise contrastive estimation ( Gutmann & Hyvärinen , 2010 ) , and variants thereof . These suffer from their own challenges in scaling to high dimensional data , and sacrifice the statistical efficiency of maximum likelihood . In Figure 1 we visualize densities learned with our approach and Persistent Contrastive Divergence ( Tieleman , 2008 ) ( PCD ) training . As we see , the sample quality of the PCD models is quite high but the learned density models do not match the true model . This is due to accrued bias in the gradient estimator from approximate MCMC sampling ( Grathwohl et al. , 2020 ) . Prior work ( Nijkamp et al. , 2019b ) has argued that this objective actually encourages the approximate MCMC samples to match the data rather than the density model . Conversely , we see that our approach ( with proper entropy regularization ) recovers a high quality model . 3 VARIATIONAL MAXIMUM LIKELIHOOD . We seek the energy function which maximizes likelihood given in Equation 1 . We examine the intractable component of the log-likelihood , the log partition-function logZ ( θ ) = log ∫ efθ ( x ) dx . We can re-write this quantity as the optimum of logZ ( θ ) = max q Eq ( x ) [ fθ ( x ) ] +H ( q ) ( 3 ) where q is a distribution and H ( q ) = −Eq ( x ) [ log q ( x ) ] denotes its entropy 1 ( see the Appendix A.1 for the derivation ) . Plugging this into our original maximum likelihood statement we obtain : θ̂ = argmax θ [ Epdata ( x ) [ fθ ( x ) ] −maxq [ Eq ( x ) [ fθ ( x ) ] +H ( q ) ] ] ( 4 ) which gives us an alternative method for training EBMs . We introduce an auxiliary sampler qφ which we train online to optimize the inner-loop of Equation 4 . This objective was used for EBM training in Kumar et al . ( 2019 ) ; Abbasnejad et al . ( 2019 ) ; Dai et al . ( 2017 ) , Dai et al . ( 2019 ) ( motivated by Fenchel Duality ( Wainwright & Jordan , 2008 ) ) . Abbasnejad et al . ( 2019 ) use an implicit generative model and Dai et al . ( 2019 ) propose to use a sampler which is inspired by MCMC sampling from pθ ( x ) and whose entropy can be computed exactly . Below we describe our approach which utilizes the same objective with a simpler sampler and a new approach to encourage high entropy . We note that when training pθ ( x ) and q ( x ) online together , the inner maximization will not be fully optimized . This leads our training objective for pθ ( x ) to be an upper bound on log pθ ( x ) . In Section 5.1 we explore the impact of this fact and find that the bound is tight enough to train models that achieve high likelihood on high-dimensional data . 1For continuous spaces , this would be the differential entropy , but we simply use entropy here for brevity . 4 METHOD . We now present a method for training an EBM pθ ( x ) = efθ ( x ) /Z ( θ ) to optimize Equation 4 . We introduce a generator distribution of the form qφ ( x ) = ∫ z qφ ( x|z ) q ( z ) dz such that : qφ ( x|z ) = N ( gψ ( z ) , σ2I ) , q ( z ) = N ( 0 , I ) ( 5 ) where gψ is a neural network with parameters ψ and thus , φ = { ψ , σ2 } . This is similar to the decoder of a variational autoencoder ( Kingma & Welling , 2013 ) . With this architecture it is easy to optimize the first and second terms of Equation 4 with reparameterization , but the entropy term requires more care . 4.1 ENTROPY REGULARIZATION . Estimating entropy or its gradients is a challenging task . Multiple , distinct approaches have been proposed in recent years based on Mutual Information estimation ( Kumar et al. , 2019 ) , variational upper bounds ( Ranganath et al. , 2016 ) , Denoising Autoencoders ( Lim et al. , 2020 ) , and nearest neighbors ( Singh & Póczos , 2016 ) . The above methods require the training of additional auxiliary models or do not scale well to high dimensions . Most relevant to this work are Dieng et al . ( 2019 ) ; Titsias & Ruiz ( 2019 ) which present a method for encouraging generators such as ours to have high entropy by estimating∇φH ( qφ ) .The estimator takes the following form : ∇φH ( qφ ) = ∇φEqφ ( x ) [ log qφ ( x ) ] = ∇φEp ( z ) p ( ) [ log qφ ( x ( z , ) ) ] ( Reparameterize sampling ) = Ep ( z ) p ( ) [ ∇φ log qφ ( x ( z , ) ) ] = Ep ( z ) p ( ) [ ∇x log qφ ( x ( z , ) ) T∇φx ( z , ) ] ( Chain rule ) ( 6 ) where we have written x ( z , ) = gψ ( z ) + σ . All quantities in Equation 6 can be easily computed except for the score-function ∇x log qφ ( x ) . The following estimator for this quantity can be easily derived ( see Appendix A.2 ) : ∇x log qφ ( x ) = Eqφ ( z|x ) [ ∇x log qφ ( x|z ) ] ( 7 ) which requires samples from the posterior qφ ( z|x ) to estimate . Dieng et al . ( 2019 ) ; Titsias & Ruiz ( 2019 ) generate these samples using Hamiltonian Monte Carlo ( HMC ) ( Neal et al. , 2011 ) , a gradientbased MCMC algorithm . As used in Dieng et al . ( 2019 ) , 28 sequential gradient computations must be made per training iteration . Since a key motivation of this work is to circumvent the costly sequential computation of MCMC sampling , this is not a favourable solution . In our work we propose a more efficient solution that we find works just as well empirically .

Summary: This paper proposes a new method on training energy-based models with maximum likelihood. Instead of using MCMC approaches to sample from the EBM, authors follow previous work on training neural generators for faster sample generation. In particular, authors consider a special generator where the output is convolved with Gaussian noise. The score function of this generator can be estimated with self-normalized importance sampling, which is then used to estimate the entropy term through the reparameterization trick. Authors demonstrate that their new method is able to train EBMs efficiently, and improves the stability and performance of JEMs compared to MCMC-based training approaches.