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 presents a method for improving training of energy-based models. Rather than drawing samples using persistent contrastive divergence / MCMC, this approach parameterizes a separate model, which is trained to directly output samples. This effectively adds an additional KL divergence to the objective. The authors use a particular form of sampling model (a latent Gaussian model), borrowing a few tricks for getting entropy estimates out of the model. Results are demonstrated on a few qualitative setups, but most of the results are centered on improved JEM on sample quality, out-of-distribution detection, and semi-supervised learning. The main benefit of the approach seems to be speed and stability, however, the authors also claim that minimal tuning is needed.