Learning Energy-Based Models by Diffusion Recovery Likelihood

1 INTRODUCTION . EBMs ( LeCun et al. , 2006 ; Ngiam et al. , 2011 ; Kim & Bengio , 2016 ; Zhao et al. , 2016 ; Goyal et al. , 2017 ; Xie et al. , 2016b ; Finn et al. , 2016 ; Gao et al. , 2018 ; Kumar et al. , 2019 ; Nijkamp et al. , 2019b ; Du & Mordatch , 2019 ; Grathwohl et al. , 2019 ; Desjardins et al. , 2011 ; Gao et al. , 2020 ; Che et al. , 2020 ; Grathwohl et al. , 2020 ; Qiu et al. , 2019 ; Rhodes et al. , 2020 ) are an appealing class of probabilistic models , which can be viewed as generative versions of discriminators ( Jin et al. , 2017 ; Lazarow et al. , 2017 ; Lee et al. , 2018 ; Grathwohl et al. , 2020 ) , yet can be learned from unlabeled data . Despite a number of desirable properties , two challenges remain for training EBMs on highdimensional datasets . First , learning EBMs by maximum likelihood requires Markov Chain Monte Carlo ( MCMC ) to generate samples from the model , which can be extremely expensive . Second , as pointed out in Nijkamp et al . ( 2019a ) , the energy potentials learned with non-convergent MCMC do not have a valid steady-state , in the sense that samples from long-run Markov chains can differ greatly from observed samples , making it difficult to evaluate the learned energy potentials . Another line of work , originating from Sohl-Dickstein et al . ( 2015 ) , is to learn from a diffused version of the data , which are obtained from the original data via a diffusion process that sequentially adds Gaussian white noise . From such diffusion data , one can learn the conditional model of the data at a certain noise level given their noisy versions at the higher noise level of the diffusion process . After learning the sequence of conditional models that invert the diffusion process , one can then generate synthesized images from Gaussian white noise images by ancestral sampling . Building on Sohl-Dickstein et al . ( 2015 ) , Ho et al . ( 2020 ) further developed the method , obtaining strong image synthesis results . Inspired by Sohl-Dickstein et al . ( 2015 ) and Ho et al . ( 2020 ) , we propose a diffusion recovery likelihood method to tackle the challenge of training EBMs directly on a dataset by instead learning a sequence of EBMs for the marginal distributions of the diffusion process . The sequence of marginal EBMs are learned with recovery likelihoods that are defined as the conditional distributions that invert the diffusion process . Compared to standard maximum likelihood estimation ( MLE ) of EBMs , learning marginal EBMs by diffusion recovery likelihood only requires sampling from the conditional distributions , which is much easier than sampling from the marginal distributions . After learning the marginal EBMs , we can generate synthesized images by a sequence of conditional samples initialized from the Gaussian white noise distribution . Unlike Ho et al . ( 2020 ) that approximates the reverse process by normal distributions , in our case the conditional distributions are derived from the marginal EBMs , which are more flexible . The framework of recovery likelihood was originally proposed in Bengio et al . ( 2013 ) . In our work , we adapt it to learning the sequence of marginal EBMs from the diffusion data . Our work is also related to the denoising score matching method of Vincent ( 2011 ) , which was further developed by Song & Ermon ( 2019 ; 2020 ) for learning from diffusion data . The training objective used for diffusion probabilisitic models is a weighted version of the denoising score matching objective , as revealed by Ho et al . ( 2020 ) . These methods learn the score functions ( the gradients of the energy functions ) directly , instead of using the gradients of learned energy functions as in EBMs . On the other hand , Saremi et al . ( 2018 ) parametrizes the score function as the gradient of a MLP energy function , and Saremi & Hyvarinen ( 2019 ) further unifies denoising score matching and neural empirical Bayes . We demonstrate the efficacy of diffusion recovery likelihood on CIFAR-10 , CelebA and LSUN datasets . The generated samples are of high fidelity and comparable to GAN-based methods . On CIFAR-10 , we achieve FID 9.58 and inception score 8.30 , exceeding existing methods of learning explicit EBMs to a large extent . We also demonstrate that diffusion recovery likelihood outperforms denoising score matching from diffusion data if we naively take the gradients of explicit energy functions as the score functions . More interestingly , by using a thousand diffusion time steps , we demonstrate that even very long MCMC chains from the sequence of conditional distributions produce samples that represent realistic images . With the faithful long-run MCMC samples from the conditional distributions , we can accurately estimate the marginal partition function at zero noise level by importance sampling , and thus evaluate the normalized density of data under the EBM . 2 BACKGROUND . Let x ∼ pdata ( x ) denote a training example , and pθ ( x ) denote a model ’ s probability density function that aims to approximates pdata ( x ) . An energy-based model ( EBM ) is defined as : pθ ( x ) = 1 Zθ exp ( fθ ( x ) ) , ( 1 ) where Zθ = ∫ exp ( fθ ( x ) ) dx is the partition function , which is analytically intractable for highdimensional x . For images , we parameterize fθ ( x ) with a convolutional neural network with a scalar output . The energy-based model in equation 1 can , in principle , be learned through MLE . Specifically , suppose we observe samples xi ∼ pdata ( x ) for i = 1 , 2 , ... , n. The log-likelihood function is L ( θ ) = 1 n n∑ i=1 log pθ ( xi ) . = Ex∼pdata [ log pθ ( x ) ] . ( 2 ) In MLE , we seek to maximize the log-likelihood function , where the gradient approximately follows ( Xie et al. , 2016b ) − ∂ ∂θ DKL ( pdata‖pθ ) = Ex∼pdata [ ∂ ∂θ fθ ( x ) ] − Ex∼pθ [ ∂ ∂θ fθ ( x ) ] . ( 3 ) The expectations can be approximated by averaging over the observed samples and the synthesized samples drawn from the model distribution pθ ( x ) respectively . Generating synthesized samples from pθ ( x ) can be done with Markov Chain Monte Carlo ( MCMC ) such as Langevin dynamics ( or Hamiltonian Monte Carlo ( Girolami & Calderhead , 2011 ) ) , which iterates xτ+1 = xτ + δ2 2 ∇xfθ ( xτ ) + δ τ , ( 4 ) where τ indexes the time , δ is the step size , and τ ∼ N ( 0 , I ) . The difficulty lies in the fact that for highdimensional and multi-modal distributions , MCMC sampling can take a long time to converge , and the sampling chains may have difficulty traversing modes . As demonstrated in Figure 2 , training EBMs with synthesized samples from non-convergent MCMC results in malformed energy landscapes ( Nijkamp et al. , 2019b ) , even if the samples from the model look reasonable . 3 RECOVERY LIKELIHOOD . 3.1 FROM MARGINAL TO CONDITIONAL . Given the difficulty of sampling from the marginal density pθ ( x ) , following Bengio et al . ( 2013 ) , we use the recovery likelihood defined by the density of the observed sample conditional on a noisy sample perturbed by isotropic Gaussian noise . Specifically , let x̃ = x+ σ be the noisy observation of x , where ∼ N ( 0 , I ) . Suppose pθ ( x ) is defined by the EBM in equation 1 , then the conditional EBM can be derived as pθ ( x|x̃ ) = 1 Z̃θ ( x̃ ) exp ( fθ ( x ) − 1 2σ2 ‖x̃− x‖2 ) , ( 5 ) where Z̃θ ( x̃ ) = ∫ exp ( fθ ( x ) − 12σ2 ‖x̃− x‖ 2 ) dx is the partition function of this conditional EBM . See Appendix A.1 for the derivation . Compared to pθ ( x ) ( equation 1 ) , the extra quadratic term 1 2σ2 ‖x̃−x‖ 2 in pθ ( x|x̃ ) constrains the energy landscape to be localized around x̃ , making the latter less multi-modal and easier to sample from . As we will show later , when σ is small , pθ ( x|x̃ ) is approximately a single mode Gaussian distribution , which greatly reduces the burden of MCMC . A more general formulation is x̃ = ax+ σ , where a is a positive constant . In that case , we can let y = ax , and treat y as the observed sample . Assume pθ ( y ) = 1Zθ exp ( fθ ( y ) ) , then by change of variable , the density function of x can be derived as gθ ( x ) = apθ ( ax ) . 3.2 MAXIMIZING RECOVERY LIKELIHOOD . With the conditional EBM , assume we have observed samples xi ∼ pdata ( x ) and the corresponding perturbed samples x̃i = xi + σ i for i = 1 , ... , n. We define the recovery log-likelihood function as J ( θ ) = 1 n n∑ i=1 log pθ ( xi|x̃i ) . ( 6 ) The term recovery indicates that we attempt to recover the clean sample xi from the noisy sample x̃i . Thus , instead of maximizing L ( θ ) in equation 2 , we can maximize J ( θ ) , whose distributions are easier to sample from . Specifically , we generate synthesized samples by K steps of Langevin dynamics that iterates xτ+1 = xτ + δ2 2 ( ∇xfθ ( xτ ) + 1 σ2 ( x̃− xτ ) ) + δ τ . ( 7 ) The model is then updated following the same learning gradients as MLE ( equation 3 ) , because the quadratic term − 12σ2 ‖x̃ − x‖ 2 is not related to θ . Following the classical analysis of MLE , we can show that the point estimate given by maximizing recovery likelihood is an unbiased estimator of the true parameters , which means that given enough data , a rich enough model and exact synthesis , maximizing the recovery likelihood learns θ such that pdata ( x ) = pθ ( x ) . See Appendix A.2 for a theoretical explanation . 3.3 NORMAL APPROXIMATION TO CONDITIONAL . When the variance of perturbed noise σ2 is small , pθ ( x|x̃ ) can be approximated by a normal distribution via a first order Taylor expansion at x̃ . Specifically , the negative conditional energy is −Eθ ( x|x̃ ) = fθ ( x ) − 1 2σ2 ‖x̃− x‖2 ( 8 ) . = fθ ( x̃ ) + 〈∇xfθ ( x̃ ) , x− x̃〉 − 1 2σ2 ‖x̃− x‖2 ( 9 ) = − 1 2σ2 [ ‖x− ( x̃+ σ2∇xfθ ( x̃ ) ) ‖2 ] + c , ( 10 ) where c include terms irrelevant of x ( see Appendix A.3 for a detailed derivation ) . In the above approximation , we do not perform second order Taylor expansion because σ2 is small , and ‖x̃ − x‖2/2σ2 will dominate all the second order terms from Taylor expansion . Thus we can approximate pθ ( x|x̃ ) by a Gaussian approximation p̃θ ( x|x̃ ) : p̃θ ( x|x̃ ) = N ( x ; x̃+ σ2∇xfθ ( x̃ ) , σ2 ) . ( 11 ) We can sample from this distribution using : xgen = x̃+ σ 2∇xfθ ( x̃ ) + σ , ( 12 ) where ∼ N ( 0 , I ) . This resembles a single step of Langevin dynamics , except that σ is replaced by √ 2σ in Langevin dynamics . This normal approximation has two traits : ( 1 ) it verifies the fact that the conditional density pθ ( x|x̃ ) can be generally easier to sample from when σ is small ; ( 2 ) it provides hints of choosing the step size of Langevin dynamics , as discussed in section 3.5 .

Summary: This paper describes training a sequence of conditional EBMs (inspired by  Ho et al. (2020)) instead of training unconditional EBMs.  Each conditional energy describes the probability of recovering x, given its noisy version \hat{x}. The noisy version of x can be described as a normal distribution centered at x, so the condition EBM has an additional term ||x - \hat{x}||^2, which constrains the Langevin dynamics to remain in the vicinity of \hat{x}, so it converges faster!