Improved Contrastive Divergence Training of Energy Based Models

1 INTRODUCTION . Energy-Based models ( EBMs ) have received an influx of interest recently and have been applied to realistic image generation ( Han et al. , 2019 ; Du & Mordatch , 2019 ) , 3D shapes synthesis ( Xie et al. , 2018b ) , out of distribution and adversarial robustness ( Lee et al. , 2018 ; Du & Mordatch , 2019 ; Grathwohl et al. , 2019 ) , compositional generation ( Hinton , 1999 ; Du et al. , 2020a ) , memory modeling ( Bartunov et al. , 2019 ) , text generation ( Deng et al. , 2020 ) , video generation ( Xie et al. , 2017 ) , reinforcement learning ( Haarnoja et al. , 2017 ; Du et al. , 2019 ) , protein design and folding ( Ingraham et al . ; Du et al. , 2020b ) and biologically-plausible training ( Scellier & Bengio , 2017 ) . Contrastive divergence is a popular and elegant procedure for training EBMs proposed by ( Hinton , 2002 ) which lowers the energy of the training data and raises the energy of the sampled confabulations generated by the model . The model confabulations are generated via an MCMC process ( commonly Gibbs sampling or Langevin dynamics ) , leveraging the extensive body of research on sampling and stochastic optimization . The appeal of contrastive divergence is its simplicity and extensibility . It does not require training additional auxiliary networks ( Kim & Bengio , 2016 ; Dai et al. , 2019 ) ( which introduce additional tuning and balancing demands ) , and can be used to compose models zero-shot . Despite these advantages , training EBMs with contrastive divergence has been challenging due to training instabilities . Ensuring training stability required either combinations of spectral normalization and Langevin dynamics gradient clipping ( Du & Mordatch , 2019 ) , parameter tuning ( Grathwohl et al. , 2019 ) , early stopping of MCMC chains ( Nijkamp et al. , 2019b ) , or avoiding the use of modern deep learning components , such as self-attention or layer normalization ( Du & Mordatch , 2019 ) . These requirements limit modeling power , prevent the compatibility with modern deep learning architectures , and prevent long-running training procedures required for scaling to larger datasets . With this work , we aim to maintain the simplicity and advantages of contrastive divergence training , while resolving stability issues and incorporating complementary deep learning advances . An often overlooked detail of contrastive divergence formulation is that changes to the energy function change the MCMC samples , which introduces an additional gradient term in the objective function ( see Section 2.1 for details ) . This term was claimed to be empirically negligible in the original formulation and is typically ignored ( Hinton , 2002 ; Liu & Wang , 2017 ) or estimated via highvariance likelihood ratio approaches ( Ruiz & Titsias , 2019 ) . We show that this term can be efficiently estimated for continuous data via a combination of auto-differentiation and nearest-neighbor entropy estimators . We also empirically show that this term contributes significantly to the overall training gradient and has the effect of stabilizing training . It enables inclusion of self-attention blocks into network architectures , removes the need for capacity-limiting spectral normalization , and allows us to train the networks for longer periods . We do not introduce any new objectives or complexity - our procedure is simply a more complete form of the original formulation . We further present techniques to improve mixing and mode exploration of MCMC transitions in contrastive divergence . We propose data augmentation as a useful tool to encourage mixing in MCMC by directly perturbing input images to related images . By incorporating data augmentation as semantically meaningful perturbations , we are able to greatly improve mixing and diversity of MCMC chains . We further propose to maintain a reservoir sample of past samples , improving the diversity of MCMC chain initialization in contrastive divergence . We also leverage compositionality of EBMs to evaluate an image sample at multiple image resolutions when computing energies . Such evaluation and coarse and fine scales leads to samples with greater spatial coherence , but leaves MCMC generation process unchanged . We note that such hierarchy does not require specialized mechanisms such as progressive refinement ( Karras et al. , 2017 ) Our contributions are as follows : firstly , we show that a gradient term neglected in the popular contrastive divergence formulation is both tractable to estimate and is important in avoiding training instabilities that previously limited applicability and scalability of energy-based models . Secondly , we highlight how data augmentation and multi-scale processing can be used to improve model robustness and generation quality . Thirdly , we empirically evaluate stability of model architectures and show improved performance on a host of benchmarks and use cases , such as image generation , OOD detection , and compositional generation . 2 AN IMPROVED CONTRASTIVE DIVERGENCE FRAMEWORK FOR ENERGY BASED MODELS . Energy based models ( EBMs ) represent the likelihood of a probability distribution for x ∈ RD as pθ ( x ) = exp ( −Eθ ( x ) ) Z ( θ ) where the function Eθ ( x ) : R D → R , is known as the energy function , and Z ( θ ) = ∫ x exp−Eθ ( x ) is known as the partition function . Thus an EBM can be represented by an neural network that takes x as input and outputs a scalar . Training an EBM through maximum likelihood ( ML ) is not straightforward , as Z ( θ ) can not be reliably computed , since this involves integration over the entire input domain of x . However , the gradient of log-likelihood with respect to a data sample x can be represented as ∂ log pθ ( x ) ∂θ = − ( ∂Eθ ( x ) ∂θ − Epθ ( x′ ) [ ∂Eθ ( x ′ ) ∂θ ] ) . ( 1 ) Note that Equation 1 is still not tractable , as it requires using Markov Chain Monte Carlo ( MCMC ) to draw samples from the model distribution pθ ( x ) , which often takes exponentially long to mix . As a practical approximation to the above objective , ( Hinton , 2002 ) proposes the contrastive divergence objective KL ( p ( x ) || pθ ( x ) ) − KL ( Πtθ ( p ( x ) ) || pθ ( x ) ) , ( 2 ) where Πθ represents a MCMC transition kernel for pθ , and Πtθ ( p ( x ) ) represents t sequential MCMC transitions starting from p ( x ) . The above objective can be seen as an improvement operator , where KL ( p ( x ) || pθ ( x ) ) ≥ KL ( Πtθ ( p ( x ) ) || pθ ( x ) ) , because Πθ is converging to equilibrium distribution pθ ( x ) ( Lyu , 2011 ) . Furthermore , the above objective is only zero ( at its fixed point ) , when Πθ does not change the distribution of p ( x ) , which corresponds to pθ ( x ) = p ( x ) . 2.1 A MISSING TERM IN CONTRASTIVE DIVERGENCE . When taking the negative gradient of the contrastive divergence objective ( Equation 2 ) , we obtain the expression − ( Ep ( x ) [ ∂Eθ ( x ) ∂θ ] − Eqθ ( x′ ) [ ∂Eθ ( x ′ ) ∂θ ] + ∂q ( x′ ) ∂θ ∂KL ( qθ ( x ) || pθ ( x ) ) ∂qθ ( x ) ) , ( 3 ) where for brevity , we summarize Πtθ ( p ( x ) ) = qθ ( x ) . The first two terms are identical to those of Equation 1 and the third gradient term ( which we refer to as the KL divergence term ) corresponds to minimizing the divergence between qθ ( x ) and pθ ( x ) . In practice , past contrastive divergence approaches have ignored the third gradient term , which was difficult to estimate and claimed to be empirically negligible ( Hinton , 1999 ) . These gradients correspond to a joint loss expression LFull , consisting of traditional contrastive loss LCD and a new loss expression LKL . Specifically , we have LFull = LCD + LKL where LCD is LCD = Ep ( x ) [ Eθ ( x ) ] − Estop gradient ( qθ ( x′ ) ) [ Eθ ( x ′ ) ] , ( 4 ) and the ignored KL divergence term corresponding to the loss LKL = Eqθ ( x ) [ Estop gradient ( θ ) ( x ) ] + Eqθ ( x ) [ log ( qθ ( x ) ) ] . ( 5 ) Despite being difficult to estimate , we show that LKL is a useful tool for both speeding up and stabilizing training of EBMs . Figure 2 illustrates the overall effects of both losses . Equation 4 encourage the energy function to assign low energy to real samples and high energy for generated samples . However , only optimizing Equation 4 often leads to an adversarial mode where the energy function learns to simply generate an energy landscape that makes sampling difficult . The KL divergence term counteracts this effect , and encourages sampling to closely approximate the underlying distribution pθ ( x ) , by encouraging samples to be both low energy under the energy function as well as diverse . Next , we discuss our approach towards estimating this KL divergence , and show that it significantly improves the stability when training EBMs . 2.2 ESTIMATING THE MISSING GRADIENT TERM . Estimating LKL can further be decomposed into two separate objectives , minimizing the energy of samples from qθ ( x ) , which we refer to as Lopt ( Equation 6 ) and maximizing the entropy of samples from qθ ( x ) which we refer to as Lent ( Equation 7 ) . Minimizing Sampler Energy . To minimize the energy of samples from qθ ( x ) we can directly differentiate through both the energy function and MCMC sampling . We follow recent work in EBMs and utilize Langevin dynamics ( Du & Mordatch , 2019 ; Nijkamp et al. , 2019b ; Grathwohl et al. , 2019 ) for our MCMC transition kernel , and note that each step of Langevin sampling is fully differentiable with respect to underlying energy function parameters . Precisely , gradient of Lopt becomes ∂Lopt ∂θ = Eqθ ( x′0 , x′1 , ... , x′t ) [ ∂Estop gradient ( θ ) ( x ′ t−1 −∇x′t−1Eθ ( x ′ t−1 ) + ω ) ∂θ ] , ω ∼ N ( 0 , λ ) ( 6 ) where x′i represents the i th step of Langevin sampling . To reduce to memory overhead of this differentiation procedure , we only differentiate through the last step of Langevin sampling ( though we show it the appendix that leads to the same effect as differentiation through Langevin sampling ) . Entropy Estimation . To maximize the entropy of samples from qθ ( x ) , we use a non-parametric nearest neighbor entropy estimator ( Beirlant et al. , 1997 ) , which is shown to be mean square consistent ( Kozachenko & Leonenko , 1987 ) with root-n convergence rate ( Tsybakov & Van der Meulen , 1996 ) . The entropy H of a distribution p ( x ) can be estimated through a set X = x1 , x2 , . . . , xn of n different points sampled from p ( x ) as H ( pθ ( x ) ) = 1n ∑n i=1 ln ( n · NN ( xi , X ) ) + O ( 1 ) where the function NN ( xi , X ) denotes the nearest neighbor distance of xi to any other data point in X . Based off the above entropy estimator , we write Lent as Lent = Eq ( x ) [ log ( NN ( x , B ) ) ] ( 7 ) where we measure the nearest neighbor with respect to a set B of 1000 past samples from MCMC chains ( see Section 2.5 for more details ) . We utilize L2 distance as the metric for computing nearest neighbors . Alternatively , Stein ’ s identity may also be used to estimate entropy , but this requires considering all samples , as opposed to the nearest , becoming computationally intractable . Our entropy estimator serves a simple , quick to compute estimator of entropy , that prevents sampling from collapsing . Empirically , we find that the combination of the above terms in LKL significantly improves both the stability and generation quality of EBMs , improving robustness across different model architectures .

Summary: Review: This paper studies how to improve contrastive divergence (CD) training of energy-based models (EBMs) by revisiting the gradient term neglected in the traditional CD learning. This paper also introduces some useful techniques, such as data augmentation, multi-scale energy design, and reservoir sampling to improve the training of energy-based model. Empirical studies are performed to validate the proposed learning strategy on the task of image generation, OOD detection, and compositional generation.