LEARNING ENERGY-BASED MODELS BY DIFFUSION RECOVERY LIKELIHOOD
Year: 2021

Abstract:
While energy-based models (EBMs) exhibit a number of desirable properties, training and sampling on high-dimensional datasets remains challenging. Inspired by recent progress on diffusion probabilistic models, we present a diffusion recovery likelihood method to tractably learn and sample from a sequence of EBMs trained on increasingly noisy versions of a dataset. Each EBM is trained with recovery likelihood, which maximizes the conditional probability of the data at a certain noise level given their noisy versions at a higher noise level. Optimizing recovery likelihood is more tractable than marginal likelihood, as sampling from the conditional distributions is much easier than sampling from the marginal distributions. After training, synthesized images can be generated by the sampling process that initializes from Gaussian white noise distribution and progressively samples the conditional distributions at decreasingly lower noise levels. Our method generates high fidelity samples on various image datasets. On unconditional CIFAR-10 our method achieves FID 9.58 and inception score 8.30, superior to the majority of GANs. Moreover, we demonstrate that unlike previous work on EBMs, our long-run MCMC samples from the conditional distributions do not diverge and still represent realistic images, allowing us to accurately estimate the normalized density of data even for high-dimensional datasets. Our implementation is available at https://github.com/ruiqigao/recovery_likelihood.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND
text: 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.

heading: 3 RECOVERY LIKELIHOOD

heading: 3.1 FROM MARGINAL TO CONDITIONAL
text: 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).

heading: 3.2 MAXIMIZING RECOVERY LIKELIHOOD
text: 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.

heading: 3.3 NORMAL APPROXIMATION TO CONDITIONAL
text: 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.

heading: 3.4 CONNECTION TO VARIATIONAL INFERENCE AND SCORE MATCHING
text: The normal approximation to the conditional distribution leads to a natural connection to diffusion probabilistic models (Sohl-Dickstein et al., 2015; Ho et al., 2020) and denoising score matching (Vincent, 2011; Song & Ermon, 2019; 2020; Saremi et al., 2018; Saremi & Hyvarinen, 2019). Specifically, in diffusion probabilistic models, instead of modeling pθ(x) as an energy-based model, it recruits variational inference and directly models the conditional density as
pθ(x|x̃) ∼ N ( x̃+ σ2sθ(x̃), σ 2 ) , (13)
which is in agreement with the normal approximation (equation 11), with sθ(x) = ∇xfθ(x). On the other hand, the training objective of denoising score matching is to minimize
1
2σ2 Ep(x,x̃)[‖x− (x̃+ σ2sθ(x̃))‖2], (14)
where sθ(x) is the score of the density of x̃. This objective is in agreement with the objective of maximizing log-likelihood of the normal approximation (equation 10), except that for normal approximation,∇xfθ(·) is the score of density of x, instead of x̃. However, the difference between the scores of density of x and x̃ is of O(σ2), which is negligible when σ is sufficiently small (see Appendix A.4 for details). We can further show that the learning gradient of maximizing log-likelihood of the normal approximation is approximately the same as the learning gradient of maximizing the original recovery log-likelihood with one step of Langevin dynamics (see Appendix A.5). It indicates that the training process of maximizing recovery likelihood agrees with the one of diffusion probabilistic models and denoising score matching when σ is small.
As the normal approximation is accurate only when σ is small, it requires many time steps in the diffusion process for this approximation to work well, which is also reported in Ho et al. (2020) and Song & Ermon (2020). In contrast, the diffusion recovery likelihood framework can be more flexible in choosing the number of time steps and the magnitude of σ.

heading: 3.5 DIFFUSION RECOVERY LIKELIHOOD
text: As we discuss, sampling from pθ(x|x̃) becomes simple only when σ is small. In the extreme case when σ → ∞, pθ(x|x̃) converges to the marginal distribution pθ(x), which is again highly multimodal and difficult to sample from. To keep σ small and meanwhile equip the model with the ability to generate new samples initialized from white noise, inspired by Sohl-Dickstein et al. (2015) and Ho et al. (2020), we propose to learn a sequence of recovery likelihoods, on gradually perturbed observed data based on a diffusion process. Specifically, assume a sequence of perturbed observations x0,x1, ...,xT such that
x0 ∼ pdata(x); xt+1 = √ 1− σ2t+1xt + σt+1 t+1, t = 0, 1, ...T − 1. (15)
The scaling factor √
1− σ2t+1 ensures that the sequence is a spherical interpolation between the observed sample and Gaussian white noise. Let yt = √ 1− σ2t+1xt, and we assume a sequence of conditional EBMs
pθ(yt|xt+1) = 1
Z̃θ,t(xt+1) exp
( fθ(yt, t)− 1
2σ2t+1 ‖xt+1 − yt‖2
) , t = 0, 1, ..., T − 1, (16)
where fθ(yt, t) is defined by a neural network conditioned on t.
We follow the learning algorithm in section 3.2. A question is how to determine the step size schedule δt of Langevin dynamics. Inspired by the sampling procedure of the normal approximation (equation 12), we set the step size δt = bσt, where b < 1 is a tuned hyperparameter. This schedule turns out to work well in practice. Thus the K steps of Langevin dynamics iterates
yτ+1t = y τ t + b2σ2t 2 (∇yfθ(yτt , t) + 1 σ2t (xt+1 − yτt )) + bσt τ . (17)
Algorithm 1 summarizes the training procedure. After training, we initialize the MCMC sampling from Gaussian white noise, and the synthesized sample at each time step serves to initialize the MCMC that samples from the model of the previous time step. See algorithm 2. To show the efficacy of our method, Figures 3 and 2 display several 2D toy examples learned by diffusion recovery likelihood.
Algorithm 1 Training repeat
Sample t ∼ Unif({0, ..., T − 1}). Sample pairs (yt,xt+1). Set synthesized sample y−t = xt+1. for τ ← 1 to K do
Update y−t according to equation 17. end for Update θ following the gradients ∂ ∂θfθ(yt, t)− ∂ ∂θfθ(y − t , t).
until converged.
Algorithm 2 Progressive sampling Sample xT ∼ N (0, I). for t← T − 1 to 0 do
yt = xt+1. for τ ← 1 to K do
Update yt according to equation 17. end for xt = yt/ √ 1− σ2t+1.
end for return x0.

heading: 4 EXPERIMENTS
text: To show that diffusion recovery likelihood is flexible for diffusion process of various magnitudes of noise, we test the method under two settings: (1) T = 6, with K = 30 steps of Langevin dynamic per time step; (2) T = 1000, with sampling from normal approximation. (2) resembles the noise schedule of Ho et al. (2020) and the magnitude of noise added at each time step is much smaller compared to (1). For both settings, we set σ2t to increase linearly. The network structure of fθ(x, t) is based on Wide ResNet (Zagoruyko & Komodakis, 2016) and we remove weight normalization. t is encoded by Transformer sinusoidal position embedding as in (Ho et al., 2020). For (1), we find that adding another scaling factor ct to the step size δt helps. Architecture and training details are in Appendix B. Henceforth we simply refer the two settings as T6 and T1k.

heading: 4.1 IMAGE GENERATION
text: Figures 1 and 4 display uncurated samples generated from learned models on CIFAR-10, CelebA 64 × 64, LSUN 64 × 64 and 128 × 128 datasets under T6 setting. The samples are of high fidelity and comparable to GAN-based methods. Appendix C.5 provides more generated samples. Tables 1 and 3 summarize the quantitative evaluations on CIFAR-10 and CelebA datasets, in terms of Frechet Inception Distance (FID) (Heusel et al., 2017) and inception scores (Salimans et al., 2016). On CIFAR-10, our model achieves FID 9.58 and inception score 8.30, which outperforms existing methods of learning explicit energy-based models to a large extent, and is superior to a majority of GAN-based methods. On CelebA, our model obtains results comparable with the state-of-the-art GAN-based methods, and outperforms score-based methods (Song & Ermon, 2019; 2020). Note that the score-based methods (Song & Ermon, 2019; 2020) and diffusion probabilistic models (Ho et al., 2020) directly parametrize and learn the score of data distribution, whereas our goal is to learn explicit energy-based models.
Table 1: FID and inception scores on CIFAR-10.
Model FID↓ Inception↑ GAN-based
WGAN-GP (Gulrajani et al., 2017) 36.4 7.86 ± .07 SNGAN (Miyato et al., 2018) 21.7 8.22 ± .05 SNGAN-DDLS (Che et al., 2020) 15.42 9.09 ± .10 StyleGAN2-ADA (Karras et al., 2020) 3.26 9.74 ± .05 Score-based
NCSN (Song & Ermon, 2019) 25.32 8.87 ± .12 NCSN-v2 (Song & Ermon, 2020) 10.87 8.40 ± .07 DDPM (Ho et al., 2020) 3.17 9.46 ± .11 Explicit EBM-conditional
Table 2: Ablation of training objectives, time steps T and sampling steps K on CIFAR-10. K = 0 indicates that we sample from the normal approximation.
Setting / Objective FID↓ Inception↑ T = 1, K = 180 32.12 6.72 ± 0.12 T = 1000, K = 0 22.58 7.71 ± 0.08 T = 1000, K = 0 (DSM) 21.76 7.76 ± 0.11 T = 6, K = 10 - - T = 6, K = 30 9.58 8.30 ± 0.11 T = 6, K = 50 9.36 8.46 ± 0.13
Table 3: FID scores on CelebA 642.
Interpolation. As shown in Figure 5, our model is capable of smooth interpolation between two generated samples. Specifically, for two samples x(0)0 and x (1) 0 , we do a sphere interpolation between
the initial white noise images x(0)T and x (1) T and the noise terms of Langevin dynamics (0) t,τ and (1) t,τ for every sampling step at every time step. More interpolation results can be found in Appendix C.3.
Image inpainting. A promising application of energy-based models is to use the learned model as a prior model for image processing, such as image inpainting, denoising and super-resolution (Gao et al., 2018; Du & Mordatch, 2019; Song & Ermon, 2019). In Figure 6, we demonstrate that the learned models by maximizing recovery likelihoods are capable of realistic and semantically meaningful image inpainting. Specifically, given a masked image and the corresponding mask, we first obtain a sequence of perturbed masked images at different noise levels. The inpainting can be easily achieved by running Langevin dynamics progressively on the masked pixels while keeping the observed pixels fixed at decreasingly lower noise levels. Additional image inpainting results can be found in Appendix C.4.
Ablation study. Table 2 summarizes the results of ablation study on CIFAR-10. We investigate the effect of changing the numbers of time steps T and sampling steps K. First, to show that it is beneficial to learn by diffusion recovery likelihood, we compare against a baseline approach (T = 1,K = 180) where we use only one time step, so that the recovery likelihood becomes marginal likelihood. The approach is adopted by Nijkamp et al. (2019b) and Du & Mordatch (2019). For fair comparison, we equip the baseline method the same budget of MCMC sampling as our T6 setting (i.e., 180 sampling steps). Our method outperforms this baseline method by a large margin. Also the models are trained more efficiently as the number of sampling steps per iteration is reduced and amortized by time steps.
Next, we report the sample quality of setting T1k. We test two training objectives for this setting: (1) maximizing recovery likelihoods (T = 1000, K = 0) and (2) maximizing the approximated normal distributions (T=1000, K=0 (DSM)). As mentioned in section 3.4, (2) is equivalent to the training objectives of denoising score matching (Song & Ermon, 2019; 2020) and diffusion probabilistic model (Ho et al., 2020), except that the score functions are taken as the gradients of explicit energy functions. In practice, for a direct comparison, (2) follows the same implementation as in Ho et al. (2020), except that the score function is parametrized as the gradients of the explicit energy function used in our method. (1) and (2) achieve similar sample quality in terms of quantitative metrics, where (2) results in a slightly better FID score yet a slightly worse inception score. This verifies the fact that the training objectives of (1) and (2) are consistent. Both (1) and (2) performs worse than setting T6. A possible explanation is that the sampling error may accumulate with many time steps, so that a more flexible schedule of time steps accompanied with certain amount of sampling steps is preferred.
Last, we examine the influence of varying the number of sampling steps while fixing the number of time steps. The training becomes unstable when the number of sampling steps are not enough (T = 6,K = 10), and more sampling steps lead to better sample quality. However, since K = 50 does not gain significant improvement versus K = 30, yet of much higher computational cost, we keep K = 30 for image generation on all datasets. See Appendix C.1 for a plot of FID scores over iterations.

heading: 4.2 LONG-RUN CHAIN ANALYSIS
text: Besides achieving high quality generation, a perhaps equally important aspect of learning EBMs is to obtain a faithful energy potential. A principle way to check the validity of the learned potential is to perform long-run sampling chains and see if the samples still remain realistic. However, as pointed out in Nijkamp et al. (2019a), almost all existing methods of learning EBMs fail in getting realistic long-run chain samples. In this subsection, we demonstrate that by composing a thousand diffusion time steps (T1k setting), we can form steady long-run MCMC chains for the conditional distributions.
First we prepare a faithful sampler for conducting long-run sampling. Specifically, after training the model under T1k setting by maximizing diffusion recovery likelihood, for each time step, we first sample from the normal approximation and count it as one sampling step, and then use Hamiltonian Monte Carlo (HMC) (Neal et al., 2011) with 2 leapfrog steps to perform the consecutive sampling steps. To obtain a reasonable schedule of sampling step size, for each time step we adaptively adjust the step size of HMC to make the average acceptance rate range in [0.6, 0.9], which is computed
over 1000 chains for 100 steps. Figure 7 displays the adjusted step size (left) and acceptance rate (center) over time step. The adjusted step size increases logarithmically. With this step size schedule, we generate long-run chains from the learned sequence of conditional distributions. As shown in Figure 8, images remain realistic for even 100k sampling steps in total (i.e., 100 sampling steps per time step), resulting in FID 24.89. This score is close to the one computed on samples generated by 1k steps (i.e., sampled from normal approximation), which is 25.12. As a further check, we recruit a No-U-Turn Sampler (Hoffman & Gelman, 2014) with the same step size schedule as HMC to perform long-run sampling, where the samples also remain realistic. See Appendix C.2 for details.
More interestingly, given the faithful long-run MCMC samples from the conditional distributions, we can estimate the log ratio of the partition functions of the marginal distributions, and further estimate the partition function of pθ(y0). The strategy is based on annealed importance sampling (Neal, 2001). See Appendix A.6 for the implementation details. The right subfigure of Figure 7 depicts the estimated log partition function of pθ(y0) over the number of MCMC samples used. To verify the estimation strategy and again check the long-run chain samples, we conduct multiple runs using samples generated with different numbers of HMC steps and display the estimation curves. All the
curves saturate to values close to each other at the end, indicating the stability of long-run chain samples and the effectiveness of the estimation strategy. With the estimated partition function, by change of variable, we can estimate the normalized density of data as gθ(x0) = √ 1− σ21pθ( √ 1− σ21x0). We report test bits per dimension on CIFAR-10 in Table 4. Note that the result should be taken with a grain of salt, because the partition function is estimated by samples and as shown in Appendix A.6, it is a stochastic lower bound of the true value, that will converge to the true value when the number of samples grows large.

heading: 5 CONCLUSION
text: We propose to learn EBMs by diffusion recovery likelihood, a variant of MLE applied to diffusion processes. We achieve high quality image synthesis, and with a thousand noise levels, we obtain faithful long-run MCMC samples that indicate the validity of the learned energy potentials. Since this method can learn EBMs efficiently with small budget of MCMC, we are also interested in scaling it up to higher resolution images and investigating this method in other data modalities in the future.

heading: ACKNOWLEDGEMENT
text: The work was done while Ruiqi Gao and Yang Song were interns at Google Brain during the summer of 2020. The work of Ying Nian Wu is supported by NSF DMS-2015577. We thank Alexander A. Alemi, Jonathan Ho, Tim Salimans and Kevin Murphy for their insightful discussions during the course of this project.

heading: A EXTENDED DERIVATIONS

heading: A.1 DERIVATION OF EQUATION 5
text: Let x̃ = x+ σ , where ∼ N (0, I). Given the marginal distribution of
pθ(x) = 1
Zθ exp(fθ(x)), (18)
We can derive the conditional distribution of x given x̃ as pθ(x|x̃) = pθ(x)p(x̃|x)/p(x̃) (19)
= 1
Zθ exp(fθ(x))
1 (2πσ2) n 2 exp(− 1 2σ2 ‖x̃− x‖2)/p(x̃) (20)
= 1
Z̃θ(x̃) exp
( fθ(x)− 1
2σ2 ‖x̃− x‖2
) , (21)
where we absorb all the terms that are irrelevant of x as Z̃θ(x̃).

heading: A.2 THEORETICAL UNDERSTANDING
text: In this subsection, we analyze the asymptotic behavior of maximizing the recovery log-likelihood.
For model class {pθ(x),∀θ}, suppose there exists θ∗ such that pdata = pθ∗ . According to the classical theory of MLE, let θ̂0 be the point estimate by MLE. Then we have θ̂ is an unbiased estimator of θ∗ with asymptotic normality:
√ n(θ̂0 − θ∗)→ N (0, I0(θ∗)−1), (22)
where I0(θ) = Ex∼pθ [−∇2θ log pθ(x)] is the Fisher information, and n is the number of observed samples.
Let θ̂ be the point estimate given by maximizing recovery log-likelihood, we can derive a result in parallel to that of MLE:
√ n(θ̂ − θ∗)→ N (0, I(θ∗)−1), (23)
where I(θ) = Epθ(x,x̃)[−∇2θ log pθ(x|x̃)]. The relationship between I0(θ) and I(θ) is that I0(θ) = I(θ) + Epθ(x,x̃)[−∇ 2 θ log pθ(x̃)]. (24) Thus there is loss of information, but θ̂ is still an unbiased estimator of θ∗ with asymptotic normality.
A.3 DETAILED DERIVATION OF NORMAL APPROXIMATION
−Eθ(x|x̃) = fθ(x)− 1
2σ2 ‖x̃− x‖2 (25)
. = fθ(x̃) + 〈∇xfθ(x̃),x− x̃〉 −
1
2σ2 ‖x̃− x‖2 (26)
= − 1 2σ2
[ ‖x‖2 − 2〈x̃,x〉+ ‖x̃‖2 ] + 〈∇xfθ(x̃),x〉 − 〈∇xfθ(x̃), x̃〉+ fθ(x̃) (27)
= − 1 2σ2
[ ‖x‖2 − 2〈x̃+ σ2∇xfθ(x̃),x〉 ] − 1
2σ2 ‖x̃‖2 − 〈∇xfθ(x̃), x̃〉+ fθ(x̃) (28)
= − 1 2σ2
[ ‖x− (x̃+ σ2∇xfθ(x̃))‖2 ] + c, (29)
A.4 DIFFERENCE BETWEEN THE SCORES OF p(x) AND p(x̃)
For notation clarity, with x̃ = x+ , we let p̃ be the distribution of x̃, and p be the distribution of x. Then for a smooth testing function with vanishing tails,
E[h(x̃)] = E[h(x+ )] (30) . = E[h(x) + h′(x) + h′′(x) 2/2] (31) = E[h(x)] + E[h′′(x)]σ2/2. (32)
Integral by parts, E[h′′(x)] = ∫ h′′(x)p(x)dx = − ∫ h′(x)p′(x)dx = ∫ p′′(x)h(x)dx. (33)
Thus we have the heat equation
p̃(x) = p(x) + p′′(x)σ2/2. (34)
The score
∇x log p̃(x) = ∇x log p(x) +∇x log(1 + p′′(x)/p(x)σ2/2) (35) . = ∇x log p(x) +∇x[p′′(x)/p(x)]σ2/2. (36)
Thus the difference between the score of p and p̃ is of the order σ2, which is negligible when σ2 is small.

heading: A.5 LEARNING GRADIENTS OF NORMAL APPROXIMATION AND ORIGINAL RECOVERY LIKELIHOOD
text: In this subsection we demonstrate that the learning gradient of maximizing likelihood of the normal approximation is approximately the same as the gradient of maximizing the original recovery likelihood with one step of Langevin sampling. Specifically, the gradient of the normal approximation of recovery log-likelihood for an observed xobs is
∇θ ( 1 2σ2 [ ‖xobs − (x̃+ σ2f ′θ(x̃))‖2 ]) = ∇θf ′θ(x̃)(xobs − (x̃+ σ2f ′θ(x̃)). (37)
On the other hand, to maximize the original recovery likelihood, suppose we sample xsyn ∼ pθ(x|x̃), then the gradient ascent of the original recovery log-likelihood is
∇θfθ(xobs)− E[∇θfθ(xsyn)] = hθ(xobs)− E[hθ(xsyn)], (38)
where hθ(x) = ∇θfθ(x). Approximately, if we perform one step of Langevin dynamics from x̃ to obtain xsyn, i.e., xsyn = x̃+ σ2f ′θ(x̃) + √ 2σe, and assume fθ(x) is locally linear in x, then
∇θfθ(xobs)− E[∇θfθ(xinit)] (39) = hθ(xobs)− E[hθ(x̃+ σ2f ′θ(x̃) + σe)] (40) . = hθ(x̃) + h ′ θ(x̃)(xobs − x̃)− E[hθ(x̃) + h′θ(x̃)(σ2f ′θ(x̃) + σe)] (41)
= h′θ(x̃)(xobs − (x̃+ σ2f ′θ(x̃)) (42) = ∇θf ′θ(x̃)(xobs − (x̃+ σ2f ′θ(x̃)). (43)
Comparing equations 37 and 43, we see that the two gradients agree with each other.

heading: A.6 ESTIMATING THE PARTITION FUNCTION
text: We can utilize the sequence of learned distributions of yt (= √
1− σ2t+1xt) to estimate the partition function. Specifically, the marginal distribution of yt is
pθ(yt) = 1
Zθ,t exp (fθ(yt, t)) (44)
We can estimate the ratio of the partition functions at two consecutive time steps using importance sampling
Zθ,t Zθ,t+1 = Epθ(yt+1) [exp(fθ(y, t)− fθ(y, t+ 1))] (45)
. =
1
M M∑ i=1 [exp(fθ(yt+1,i, t)− fθ(yt+1,i, t+ 1))] , (46)
where yt+1,i are samples generated by progressive sampling. Starting from t = T , where pT (x) follows Gaussian distribution, we can compute logZθ,t along the reverse path of the diffusion process, until we reach t = 0:
Zθ,0 = Zθ,T T−1∏ t=0 Zθ,t Zθ,t+1 . (47)
In practice, since the ratio given by MCMC samples can vary across many orders of magnitude, it is more meaningful to estimate
logZθ,0 = logZθ,T + T−1∑ t=0 log Zθ,t Zθ,t+1 . (48)
Unfortunately, although equation 46 is an unbiased estimator of Zθ,t/Zθ,t+1, the logarithm of this estimator is generally a stochastic lower bound of log(Zθ,t/Zθ,t+1) (Grosse et al., 2016). However, as we show below, this bound will gradually converge to an unbiased estimator of log(Zθ,t/Zθ,t+1), as the number of samples becomes large. Specifically, let A be the estimator in equation 46, µ be the true value of Zθ,t/Zθ,t+1. We have E[A] = µ, then by second order Taylor expansion,
E[logA] .= E [ logµ+ 1
µ (A− µ)− 1 2µ2 (A− µ)2
] (49)
= logµ− 1 2µ2 Var(A). (50)
By law of large number, Var(A) → 0 as M → ∞, and thus E[logA] → logµ. This is also consistent with the estimation curves in the right subfigure of Figure 7: since Var(A) ≥ 0, the estimation curve increases from below as the number of samples becomes larger. When the curve becomes stable, it indicates the convergence.

heading: B EXPERIMENTAL DETAILS
text: Model architecture. Our network structure is based on Wide ResNet (Zagoruyko & Komodakis, 2016). Table 5 lists the detailed network structures of various resolutions. The number of ResBlocks at every level N is a hyperparameter that we sweep over. The values of N for various datasets are listed in Table 6. Each ResBlock consists of two Conv2D layers. For the second Conv2D layer, we use zero initialization for the weights, and add a trainable channel-wise scaling parameter to the output. We remove the weight normalization, and use leaky ReLU (slope = 0.2) as the activation function in ResBlocks. Spectral normalization (Miyato et al., 2018) is used to regularize parameters in Conv2D layer, ResBlocks and Dense layer. For encoding time step t, we follow the scheme in (Ho et al., 2020). Specifically, the time step t is first transformed into sinusoidal embedding, and then two Dense layers is added. The time embedding is added after the first Conv2D layer of each ResBlock.
Training. We use Adam (Kingma & Ba, 2014) optimizer for all the experiments. We find that for high resolution images, using a smaller β1 in Adam help stabilize training. We use learning rate 0.0001 for all the experiments. For the values of β1, batch sizes and the number of training iterations for various datasets, see Table 6.
Datasets. We use the following datasets in our experiments: CIFAR-10 (Krizhevsky et al., 2009), CelebA (Liu et al., 2018) and LSUN (Yu et al., 2015). CIFAR-10 is of resolution 32 × 32, and contains 50, 000 training images and 10, 000 test images. CelebA contains 202,599 face images, of which 162,770 are training images and 19,962 are test images. For processing, we first clip each image to 178× 178 and then resize it to 64× 64. For LSUN, we use church outdoor and bedroom categories, which contains 126,227 and 3,033,042 training images respectively. Both categories contain 300 test images. For processing, we first crop each image to a square image of the smaller size among the height and weight, and then we resize it to 64 × 64 or 128 × 128. For resizing, we set antialias to True. We apply horizontal random flip as data augmentation for all datasets during training.
Evaluation metrics. We use FID and inception scores as quantitative evaluation metrics of sample quality. On all the datasets, we calculate FID and inception scores on 50,000 samples using the original code from Salimans et al. (2016) and Heusel et al. (2017).

heading: C ADDITIONAL EXPERIMENTAL RESULTS

heading: C.1 FID SCORES OVER ITERATIONS
text: Figure 9 demonstrates FID scores computed on 2,500 samples every 15,000 iterations.

heading: C.2 LONG-RUN CHAIN SAMPLING WITH NUTS
text: As a further check, we use a No-U-Turn Sampler (Hoffman & Gelman, 2014) to perform the longrun chain sampling, with the same step size schedule obtained for HMC sampler. Figure 10 displays samples with different number of sampling steps. The samples remain realistic after 100k sampling steps in total and the FID score remains stable.

heading: C.3 ADDITIONAL INTERPOLATION RESULTS
text: Figures 11, 12 and 13 display more examples of interpolation between two generated samples on CelebA 642, LSUN church outdoor 1282 and LSUN bedroom 1282.

heading: C.4 ADDITIONAL IMAGE INPAINTING RESULTS
text: Figures 14 and 15 show additional examples of image inpainting on CelebA 642 and LSUN church outdoor 1282.

heading: C.5 ADDITIONAL UNCURATED SAMPLES
text: Figures 16, 17, 18, 19, 20 and 21 show uncurated samples from the learned models under T6 setting on CIFAR-10, CelebA 642, LSUN church outdoor 1282, LSUN bedroom 1282, LSUN church outdoor 642 and LSUN bedroom 642 datasets.

Review summaries:
The paper proposed a novel method to train EBMs based on diffusion recovery likelihood. It constructs a sequence of noisy version of data and learn a conditional between consecutive noisy pairs. Compare to working with the likelihood directly, doing so makes the training much easier. Besides, even using a potentially non-convergent MCMC for gradient estimation, it still leads to a well-behaved energy potential, unlike EBMs trained via maximising the likelihood directly.

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!