Title: Learning Energy-Based Prior Model with Diffusion-Amortized MCMC
Venue: NeurIPS 2023
Accepted: yes
Keywords: Energy-Based Model

Denoising Diffusion Probabilistic Model

MCMC
Abstract: Latent space EBMs, also known as energy-based priors, have drawn growing interests in the field of generative modeling due to its flexibility in the formulation and strong modeling power of the latent space. However, the common practice of learning latent space EBMs with non-convergent short-run MCMC for prior and posterior sampling is hindering the model from further progress; the degenerate MCMC sampling quality in practice often leads to degraded generation quality and instability in training, especially with highly multi-modal and/or high-dimensional target distributions. To remedy this sampling issue, in this paper we introduce a simple but effective diffusion-based amortization method for long-run MCMC sampling and develop a novel learning algorithm for the latent space EBM based on it. We provide theoretical evidence that the learned amortization of MCMC is a valid long-run MCMC sampler. Experiments on several image modeling benchmark datasets demonstrate the superior performance of our method compared with strong counterparts.

Learning Energy-Based Prior Model with Diffusion-Amortized MCMC

Abstract

Latent space Energy-Based Models (EBMs), also known as energy-based priors, have drawn growing interests in the field of generative modeling due to its flexibility in the formulation and strong modeling power of the latent space. However, the common practice of learning latent space EBMs with non-convergent shortrun MCMC for prior and posterior sampling is hindering the model from further progress; the degenerate MCMC sampling quality in practice often leads to degraded generation quality and instability in training, especially with highly multimodal and/or high-dimensional target distributions. To remedy this sampling issue, in this paper we introduce a simple but effective diffusion-based amortization method for long-run MCMC sampling and develop a novel learning algorithm for the latent space EBM based on it. We provide theoretical evidence that the learned amortization of MCMC is a valid long-run MCMC sampler. Experiments on several image modeling benchmark datasets demonstrate the superior performance of our method compared with strong counterparts 1 .

1 Introduction

Generative modeling of data distributions has achieved impressive progress with the fast development of deep generative models in recent years [1-9]. It provides a powerful framework that allows successful applications in synthesizing data of different modalities [10-15], extracting semantically meaningful data representation [16-18] as well as other important domains of unsupervised or semisupervised learning [19-21]. A fundamental and powerful branch of generative modeling is the Deep Latent Variable Model (DLVM). Typically, DLVM assumes that the observation ( e.g ., a piece of text or images) is generated by its corresponding low-dimensional latent variables via a top-down generator network [1-3]. The latent variables are often assumed to follow a non-informative prior distribution, such as a uniform or isotropic Gaussian distribution. While one can directly learn a deep top-down generator network to faithfully map the non-informative prior distribution to the data distribution, learning an informative prior model in the latent space could further improve the ex-

2

Song-Chun Zhu

4

pressive power of the DLVM with significantly less parameters [22]. In this paper, we specifically consider learning an EBM in the latent space as an informative prior for the model.

Learning energy-based prior can be challenging, as it typically requires computationally expensive Markov Chain Monte Carlo (MCMC) sampling to estimate learning gradients. The difficulty of MCMC-based sampling is non-negligible when the target distribution is highly multi-modal or high-dimensional. In these situations, MCMC sampling can take a long time to converge and perform poorly on traversing modes with limited iterations [23]. Consequently, training models with samples from non-convergent short-run MCMC [24], which is a common choice for learning latent space EBMs [22], often results in malformed energy landscapes [15, 24, 25] and biased estimation of the model parameter. One possible solution is to follow the variational learning scheme [1], which however requires non-trivial extra efforts on model design to deal with problems like posterior collapse [26-28] and limited expressivity induced by model assumptions [1, 29, 30].

To remedy this sampling issue and further unleash the expressive power of the prior model, we propose to shift attention to an economical compromise between unrealistically expensive long-run MCMC and biased short-run MCMC: we consider learning valid amortization of the potentially long-run MCMC for learning energy-based priors . Specifically, inspired by the connection between MCMC sampling and denoising diffusion process [7, 8, 31], in this paper we propose a diffusionbased amortization method suitable for long-run MCMC sampling in learning latent space EBMs. The learning algorithm derived from it breaks the long-run chain into consecutive affordable shortrun segments that can be iteratively distilled by a diffusion-based sampler. The core idea is simple and can be summarized by a one-liner (Fig. 1). We provide theoretical and empirical evidence that the resulting sampler approximates the long-run chain (see the proof-of-concept toy examples in Appendix E.1), and brings significant performance improvement for learning latent space EBMs on several tasks. We believe that this proposal is a notable attempt to address the learning issues of energy-based priors and is new to the best of our knowledge. We kindly refer to Section 5 for a comprehensive discussion of the related work.

Contributions i) We propose a diffusion-based amortization method for MCMC sampling and develop a novel learning algorithm for the latent space EBM. ii) We provide a theoretical understanding that the learned amortization of MCMC is a valid long-run MCMC sampler. iii) Our experiments demonstrate empirically that the proposed method brings higher sampling quality, a better-learned model and stronger performance on several image modeling benchmark datasets.

2 Background

2.1 Energy-Based Prior Model

We assume that for the observed sample x P R D , there exists z P R d as its unobserved latent variable vector. The complete-data distribution is

where p α p z q is the prior model with parameters α , p β p x | z q is the top-down generation model with parameters β , and θ ' p α , β q . The prior model p α p z q can be formulated as an energy-based model, which we refer to as the Latent-space Energy-Based Model (LEBM) [22] throughout the paper. In this formulation, f α p z q is parameterized by a neural network with scalar output, Z α is the partition function, and p 0 p z q is standard normal as a reference distribution. The prior model in Eq. (1) can be interpreted as an energy-based correction or exponential tilting of the original prior distribution p 0 . The generation model follows p β p x | z q ' N p g β p z q , σ 2 I D q , where g β is the generator network and σ 2 takes a pre-specified value as in V AE [1]. This is equivalent to using l 2 error for reconstruction.

The parameters of LEBM and the generation model can be learned by Maximum Likelihood Estimation (MLE) [22]. To be specific, given the training data x , the gradients for updating α , β are,

In practice, one may use the Monte-Carlo average to estimate the expectations in Eq. (2). This involves sampling from the prior p α p z q and the posterior p θ p z | x q distribution using MCMC, specifically Langevin Dynamics (LD) [32], to estimate the expectations and hence the gradient. For a target

z

t

z

t

`

1

z

S

z

s

z

0

¨ ¨ ¨

¨ ¨ ¨

x

q

ϕ

k

p

z

s

|

z

S

,

x

q

q

ϕ

k

p

z

0

|

z

s

,

x

q

ϕ

k

π

t

`

1

'

K

π

t

t

'

0

, ..., T

´

1

ϕ

k

`

1

k

'

0

, ..., K

´

1

distribution π p z q , the dynamics iterates

where s is a small step size. One can draw z 0 ' N p 0 , I d q to initialize the chain. For sufficiently small step size s , the distribution of z t will converge to π as t Ñ8 [32]. However, it is prohibitively expensive to run LD until convergence in most cases, for which we may resort to limited iterations of LDfor sampling in practice. This non-convergent short chain yields a moment-matching distribution close to the true π p z q but is often biased, which was dubbed as short-run LD [15, 23-25].

2.2 Denoising Diffusion Probabilistic Model

Closely related to EBMs are the Denoising Diffusion Probabilistic Models (DDPMs) [5, 7, 8]. As pointed out in [5, 8], the sampling procedure of DDPM with ϵ -prediction parametrization resembles LD; ϵ (predicted noise) plays a similar role to the gradient of the log density [8].

In the formulation proposed by Kingma et al. [33], the DDPM parameterized by ϕ is specified by a noise schedule built upon λ s ' log r β 2 s { σ 2 s s , i.e ., the log signal-to-noise-ratio, that decreases monotonically with s . β s and σ 2 s are strictly positive scalar-valued functions of s . We use z 0 to denote training data in R d . The forward-time diffusion process q p z | z 0 q is defined as:

where 0 ď s ă s 1 ď S and σ 2 s 1 | s ' p 1 ´ e λ s 1 ´ λ s q σ 2 s . Noticing that the forward process can be reverted as q p z s | z s 1 , z 0 q ' N p z s ; ˜ µ s | s 1 p z s 1 , z 0 q , ˜ σ 2 s | s 1 I d q , an ancestral sampler q ϕ p z s | z s 1 q [8] that starts at z S ' N p 0 , I d q can be derived accordingly [33]:

where ϵ is standard Gaussian noise, ˆ z 0 is the prediction of z 0 by the DDPM ϕ , and γ is a hyperparameter that controls the noise magnitude, following [34]. The goal of DDPM is to recover the distribution of z 0 from the given Gaussian noise distribution. It can be trained by optimizing E ϵ ,λ ' } ϵ p z λ q ´ ϵ } 2 2 ‰ , where ϵ ' N p 0 , I d q and λ is drawn from a distribution of log noise-to-signal ratio p p λ q over uniformly sampled times s P r 0 , S s . This loss can be justified as a lower bound on the data log-likelihood [8, 33] or as a variant of denoising score matching [7, 35]. We will exploit in this paper the connection between DDPMs and LD sampling of EBMs, based upon which we achieve better sampling performance for LEBM compared with short-run LD.

3 Method

In this section, we introduce the diffusion-based amortization method for long-run MCMC sampling in learning LEBM in Section 3.1. The learning algorithm of LEBM based on the proposed method and details about implementation are then presented in Section 3.2 and Section 3.3, respectively.

3.1 Amortizing MCMC with DDPM

Amortized MCMC Using the same notation as in Section 2, we denote the starting distribution of LD as π 0 p z q , the distribution after t -th iteration as π t p z q and the target distribution as π p z q . The trajectory of LD in Eq. (3) is typically specified by its transition kernel K p z | z 1 q . The process starts with drawing z 0 from π 0 p z q and iteratively sample z t at the t -th iteration from the transition kernel conditioned on z t ´ 1 , i.e ., π t p z q ' K π t ´ 1 p z q , where K π t ´ 1 p z q : ' ş K p z | z 1 q π t ´ 1 p z 1 q d z 1 . Recursively, π t ' K t π 0 , where K t denotes the t -step transition kernel. LD can therefore be viewed as approximating a fixed point update in a non-parametric fashion since the target distribution π is a stationary distribution π p z q : ' ş K p z | z 1 q π p z 1 q d z 1 , @ z . This motivates several works for more general approximations of this update [36-38] with the help of neural samplers.

Inspired by [36], we propose to use the following framework for amortizing the LD in learning the LEBM. Formally, let Q ' t q ϕ u be the set of amortized samplers parameterized by ϕ . Given the transition kernel K , the goal is to find a sampler q ϕ ˚ to closely approximate the target distribution π . This can be achieved by iteratively distill T -step transitions of LD into q ϕ :

where D r¨||¨s is the Kullbeck-Leibler Divergence (KLD) measure between distributions. Concretely, Eq. (6) means that to recover the target distribution π , instead of using long-run LD, we can repeat the following steps: i) employ a T -step short-run LD initialized with the current sampler q ϕ k ´ 1 to approximate K T q ϕ k ´ 1 as the target distribution of the current sampler, and ii) update the current sampler q ϕ k ´ 1 to q ϕ k . The correct convergence of q ϕ to π with Eq. (6) is supported by the standard theory of Markov chains [39], which suggests that the update in Eq. (6) is monotonically decreasing in terms of KLD, D r q ϕ k || π s ď D r q ϕ k ´ 1 || π s . We refer to Appendix A.1 for a detailed explanation and discussion of this statement. In practice, one can apply gradient-based methods to minimize D r q ϕ k ´ 1 ,T || q ϕ s and approximate Eq. (6) for the update from q ϕ k ´ 1 to q ϕ k . The above formulation provides a generic and flexible framework for amortizing the potentially long MCMC.

Diffusion-based amortization To avoid clutter, we simply write q ϕ k ´ 1 ,T as q T . We can see that

where H represents the entropy of distributions. The selection of the sampler q ϕ is a matter of design. According to Eq. (7), we may expect the following properties from q ϕ : i) having analytically tractable expression of the exact value or lower bound of log-likelihood, ii) easy to draw samples from and iii) capable of close approximation to the given distribution t q T u . In practice, iii) is important for the convergence of Eq. (6). If q ϕ is far away from q T in each iteration, then non-increasing KLDproperty D r q ϕ k || π s ď D r q ϕ k ´ 1 || π s might not hold, and the resulting amortized sampler would not converge to the true target distribution π p z q .

For the choice of q ϕ , let us consider distilling the gradient field of q T in each iteration, so that the resulting sampler is close to the q T distribution. This naturally points to the DDPMs [8]. To be specific, learning a DDPM with ϵ -prediction parameterization is equivalent to fitting the finite-time marginal of a sampling chain resembling annealed Langevin dynamics [7, 8, 31]. Moreover, it also fulfills i) and ii) of the desired properties mentioned above. We can plug in the objective of DDPM (Section 2.2), which is a lower bound of log q ϕ , to obtain the gradient-based update rule for q ϕ :

where ϵ ' N p 0 , I d q . λ is drawn from a distribution of log noise-to-signal ratio p p λ q . η is the step size for the update, and M is the number of iterations needed in Eq. (8). In practice, we find that when amortizing the LD sampling chain, a light-weight DDPM q ϕ updated with very small M , i.e ., few steps of Eq. (8) iteration, approximates Eq. (6) well. We provide a possible explanation using the Fisher information by scoping the asymptotic behavior of this update rule in the Appendix A.2. We term the resulting sampler as Diffusion-Amortized MCMC (DAMC).

3.2 Approximate MLE with DAMC

In this section, we show how to integrate the DAMC sampler into the learning framework of LEBM and form a symbiosis between these models. Given a set of N training samples t x i u N i ' 1 independently drawn from the unknown data distribution p data p x q , the model p θ (Section 2.1) can be trained

by maximizing the log-likelihood over training samples L p θ q ' 1 N ř N i ' 1 log p θ p x i q . Doing so typically requires computing the gradients of L p θ q , where for each x i the learning gradient satisfies:

Intuitively, based on the discussion in Section 2.1 and Section 3.1, we can approximate the distributions in Eq. (9) by drawing samples from r z i | x i s ' K T, z i | x i q ϕ k p z i | x i q , z i ' K T, z i q ϕ k p z i q , to estimate the expectations and hence the learning gradient. Here we learn the DAMC samplers q ϕ k p z i | x i q and q ϕ k p z i q for the posterior and prior sampling chain, respectively. ϕ k represents the current sampler as in Section 3.1; K T, z i | x i and K T, z i are the transition kernels for posterior and prior sampling chain. Equivalently, this means to better estimate the learning gradients we can i) first draw approximate posterior and prior MCMC samples from the current q ϕ k model, and ii) update the approximation of the prior p α p z q and posterior p θ p z | x q distributions with additional T -step LD initialized with q ϕ k samples. These updated samples are closer to p θ p z i | x i q and p α p z i q compared with short-run LD samples based on our discussion in Section 2.2. Consequently, the diffusionamortized LD samples provide a generally better estimation of the learning gradients and lead to better performance, as we will show empirically in Section 4. After updating θ ' p α , β q based on these approximate samples with Eq. (9), we can update q ϕ k with Eq. (8) to distill the sampling chain into q ϕ k ` 1 . As shown in Fig. 1, we can see that the whole learning procedure iterates between the approximate MLE of p θ and the amortization of MCMC with q ϕ . We refer to Appendix A.3 for an extended discussion of this procedure.

After learning the models, we can use either DAMC or LEBM for prior sampling. For DAMC, we may draw samples from q ϕ p z i q with Eq. (5). Prior sampling with LEBM still requires short-run LD initialized from N p 0 , I d q . For posterior sampling, we may sample from K T, z i | x i q ϕ k p z i | x i q , i.e ., first draw samples from DAMC and then run few steps of LD to obtain posterior samples.

Algorithm 1: Learning algorithm of DAMC.

3.3 Implementation

In order to efficiently model both q ϕ k p z i | x i q and q ϕ k p z i q , we follow the method of [40] to train a single network to parameterize both models, where q ϕ k p z i | x i q can be viewed as a conditional DDPM with the embedding of x i produced by an encoder network as its condition, and q ϕ k p z i q an unconditional one. For q ϕ k p z i q , we can input a null token H as its condition when predicting the noise ϵ . We jointly train both models by randomly nullifying the inputs with the probability p uncond ' 0 . 2 . During training, we use samples from q ϕ k p z i | x i q to initialize both prior and posterior updates for training stability. For the posterior and prior DAMC samplers, we set the number of diffusion steps to 100 . The number of iterations in Eq. (8) is set to M ' 6 throughout the experiments. The LD runs 30 and 60 iterations for posterior and prior updates during training with a step

size of s ' 0 . 1 . For test time sampling from K T, z i | x i q ϕ k p z i | x i q , T ' 10 for the additional LD. For a fair comparison, we use the same LEBM and generator as in [22, 41] for all the experiments. We summarize the learning algorithm in Algorithm 1. Please see Appendices B and C for network architecture and further training details, as well as the pytorch-style pseudocode of the algorithm.

4 Experiments

In this section, we are interested in the following questions: (i) How does the proposed method compare with its previous counterparts ( e.g ., purely MCMC-based or variational methods)? (ii) How is the scalability of this method? (iii) How are the time and parameter efficiencies? (iv) Does the proposed method provide a desirable latent space? To answer these questions, we present a series of experiments on benchmark datasets including MNIST [42], SVHN [43], CelebA64 [44], CIFAR10 [45], CelebAMask-HQ [46], FFHQ [10] and LSUN-Tower [47]. As to be shown, the proposed method demonstrates consistently better performance in various experimental settings compared with previous methods. We refer to Appendix D for details about the experiments.

4.1 Generation and Inference: Prior and Posterior Sampling

Generation and reconstruction We evaluate the quality of the generated and reconstructed images to examine the sampling quality of DAMC. Specifically, we would like to check i) how well does DAMC fit the seen data, ii) does DAMC provide better MCMC samples for learning LEBM and iii) the generalizability of DAMC on unseen data. We check the goodness of fit of DAMC by evaluating the quality of the images generated with DAMC prior samples. If DAMC does provide better MCMC samples for learning LEBM, we would expect better fitting of data and hence an improved generation quality of LEBM. We evaluate the performance of posterior sampling given unseen testing images by examining the reconstruction error on testing data. We benchmark our model against a variety of previous methods in two groups. The first group covers competing methods that adopt the variational learning scheme, including V AE [1], as well as recent two-stage methods such as 2stage VAE [48], RAE [49] and NCP-VAE [50], whose prior distributions are learned with posterior samples in a second stage after the generator is trained. The second group includes methods that adopt MCMC-based sampling. It includes Alternating Back-Propogation (ABP) [51], Short-Run Inference (SRI) from [24] and the vanilla learning method of LEBM [22], which relies on short-run LD for both posterior and prior sampling. We also compare our method with the recently proposed Adaptive CE [41]. It learns a series of LEBMs adaptively during training, while these LEBMs are sequentially updated by density ratio estimation instead of MLE. To make fair comparisons, we follow the same evaluation protocol as in [22, 41].

(a) SVHN

(b) CelebA

(c) CIFAR-10

(d) CelebA-HQ

53100

8821456

1678300436

2

83

03

JES

For generation, we report the FID scores [52] in Table 1. We observe that i) the DAMC sampler, denoted as Ours-DAMC, provides superior generation performance compared to baseline models, and ii) the LEBM learned with samples from DAMC, denoted as Ours-LEBM, demonstrates significant performance improvement compared with the LEBM trained with short-run LD, denoted as LEBM. These results confirm that DAMC is a reliable sampler and indeed partly addresses the learning issue of LEBM caused by short-run LD. We would like to point out that the improvement is clearer on the CelebAMask-HQ dataset, where the input data is of higher dimension ( 256 ˆ 256 ) and contains richer details compared with other datasets. This illustrates the superiority of DAMC sampler over short-run LD when the target distribution is potentially highly multi-modal. We show qualitative results of generated samples in Fig. 2, where we observe that our method can generate diverse, sharp and high-quality samples. For reconstruction, we compare our method with baseline methods in terms of MSE in Table 1. We observe that our method demonstrates competitive reconstruction error, if not better, than competing methods do. Additional qualitative results of generation and reconstruction are presented in Appendices E.2 and E.3.

GAN inversion We have examined the scalability of our method on the CelebAMask-HQ dataset. Next, we provide more results on highdimensional and highly multi-modal data by performing GAN inversion [54] using the proposed method. Indeed, we may regard GAN inversion as an inference problem and a special case of posterior sampling. As a suitable testbed, the StyleGAN structure [10] is specifically considered as our generator in the experiments: [53] points out that to effectively infer the latent representation of a given image, the GANinversion method needs to consider an extended latent space of StyleGAN, consisting of

14 different 512-dimensional latent vectors. We attempt to use the DAMC sampler for GAN inversion. We benchmark our method against i) learning an encoder that maps a given image to the latent space [16], which relates to the variational methods for posterior inference, ii) optimizing a random initial latent code by minimizing the reconstruction error and perceptual loss [53], which can be viewed as a variant of LD sampling, and iii) optimizing the latent code by minimizing both the objectives used in ii) and the energy score provided by LEBM. We use the pretrained weights provided by [18] for the experiments. Both the DAMC sampler and the encoder-based method are augmented with 100 post-processing optimization iterations. We refer to Appendix D for more experiment details. We test LEBM-based inversion with different optimization iterations. To be specific, 1x, 2x, and 4x represent 100, 200, and 400 iterations respectively. We can see in Table 2 that DAMC performs better than all the baseline methods on the unseen testing data, which supports the efficacy of our method in high-dimensional settings. We provide qualitative results in Fig. 3.

Parameter efficiency and sampling time One potential disadvantage of our method is its parameter inefficiency for introducing an extra DDPM. Fortunately, our models are in the latent space

(a) Observation

911991

199971112

(b) Inversion

(a) 200 steps

(b) 2500 steps

500

500

1000

1000

1500

1500

2000

2000

2500

2500

0

25

50

75

100

125

150

175

Q

00S

1000

1500

2000

2500

so the network is lightweight. To be specific, on CIFAR-10 dataset the number of parameters in the DDPM is only around 10% (excluding the encoder) of those in the generator. The method has competitive time efficiency. With the batch size of 64 , the DAMC prior sampling takes 0 . 3 s , while 100 steps of short-run LD with LEBM takes 0 . 2 s . The DAMC posterior sampling takes 1 . 0 s , while LEBM takes 8 . 0 s . Further discussions about the limitations can be found in the Appendix G.1.

4.2 Analysis of Latent Space

Long-run langevin transition In this section, we examine the energy landscape induced by the learned LEBM. We expect that a well-trained p α p z q fueled by better prior and posterior samples from the DAMC sampler would lead to energy landscape with regular geometry. In Fig. 4, we visualize the transition of LD initialized from N p 0 , I d q towards p α p z q on the model trained on the CelebA dataset. Additional visualization of transitions on SVHN and CIFAR-10 datasets can be found in the Appendix E.4. The LD iterates for 200 and 2500 steps, which is longer than the LD within each training iteration ( 60 steps). For the 200 -step set-up, we can see that the generation quality quickly improves by exploring the local modes (demonstrating different facial features, e.g ., hairstyle, facial expression and lighting). For the 2500 -step long-run set-up, we can see that the LD produces consistently valid results without the oversaturating issue of the long-run chain samples [23]. These observations provide empirical evidence that the LEBM is well-trained.

Anomaly detection We further evaluate how the LEBM learned by our method could benefit the anomaly detection task. With properly learned models, the posterior p θ , ϕ p z | x q could form a discriminative latent space that has separated probability densities for in-distribution (normal) and outof-distribution (anomalous) data. Given the testing sample x , we use un-normalized log joint density p θ , ϕ p z | x q9 p θ , ϕ p x , z q « p β p x | z q p α p z q| z ' K T, z | x q ϕ p z | x q as our decision function. This means that we draw samples from K T, z | x q ϕ p z | x q and compare the corresponding reconstruction errors and energy scores. A higher value of log joint density indicates a higher probability of the test sample being a normal sample. To make fair comparisons, we follow the experimental settings in [22, 41, 55, 56] and train our model on MNIST with one class held out as an anomalous class. We consider the

baseline models that employ MCMC-based or variational inferential mechanisms. Table 3 shows the results of AUPRC scores averaged over the last 10 trials. We observe significant improvements in our method over the previous counterparts.

4.3 Ablation Study

In this section, we conduct ablation study on several variants of the proposed method. Specifically, we would like to know: i) what is the difference between the proposed method and directly training a DDPM in a fixed latent space? ii) What is the role of LEBM in this learning scheme? iii) Does DAMC effectively amortize the sampling chain? We use CIFAR-10 dataset for the ablative experiments to empirically answer these questions. More ablation studies can be found in Appendix F.

Non-Amortized DDPM vs. DAMC We term directly training a DDPM in a fixed latent space as the non-amortized DDPM. To analyze the difference between non-amortized DDPM and DAMC, we first train a LEBM model with persistent long-run chain sampling [57] and use the trained model to obtain persistent samples for learning the non-amortized DDPM. In short, the non-amortized DDPM can be viewed as directly distilling the long-run MCMC sampling process, instead of progressively amortizing the chain. We present the FID and MSE of the non-amortized model (NALR) in Table 4. We observe that directly amortizing the long-run chain leads to degraded performance compared with the proposed method. The results are consistently worse for both posterior and prior sampling and the learned LEBMs, which verify the effectiveness of the proposed iterative learning scheme.

The contribution of LEBM One may argue that since we have the DAMC as a powerful sampler, it might not be necessary to jointly learn LEBM in the latent space. To demonstrate the necessity of this joint learning scheme, we train a variant of DAMC by replacing the LEBM with a Gaussian prior. The results are presented in Table 4. We observe that models trained with non-informative Gaussian prior obtain significantly worse generation results. It suggests that LEBM involved in the learning iteration provides positive feedback to the DAMC sampler. Therefore, we believe that it is crucial to jointly learn the DAMC and LEBM.

Vanilla sampling vs. DAMC We compare the vanilla sampling process of each model with DAMC. The vanilla sampling typically refers to short-run or long-run LD initialized with N p 0 , I d q . We also provide results of learning LEBM using variational methods for comparison. We can see in Table 4 that sampling with DAMC shows significantly better scores of the listed models, compared with vanilla sampling. The result is even better than that of the persistent chain sampler (V. of NALR-LEBM). This indicates that DAMC effectively amortizes the sampling chain. Comparing DAMCsampler with the variational sampler also indicates that DAMC is different from general variational approximation: it benefits from its connection with LD and shows better expressive power.

5 Related Work

Energy-based prior model EBMs [23, 24, 58-60] play an important role in generative modeling. Pang et al. [22] propose to learn an EBM as a prior model in the latent space of DLVMs; it greatly improves the model expressivity over those with non-informative priors and brings strong performance on downstream tasks, e.g ., image segmentation, text modeling, molecule generation, and trajectory prediction [12, 15, 61, 62]. However, learning EBMs or latent space EBMs requires MCMC sampling to estimate the learning gradients, which needs numerous iterations to converge when the target distributions are high-dimensional or highly multi-modal. Typical choices of sampling with non-convergent short-run MCMC [24] in practice can lead to poor generation quality, malformed energy landscapes [15, 24, 25], biased estimation of the model parameter and instability in training [23, 25, 59, 60]. In this work, we consider learning valid amortization of the long-run MCMCfor energy-based priors; the proposed model shows reliable sampling quality in practice.

Denoising diffusion probabilistic model DDPMs [7, 8, 31], originating from [5], learn the generative process by recovering the observed data from a sequence of noise-perturbed versions of the data. The learning objective can be viewed as a variant of the denoising score matching objective [35]. As pointed out in [5, 8], the sampling procedure of DDPM with ϵ -prediction parametrization resembles LD of an EBM; ϵ (predicted noise) plays a similar role to the gradient of the log density [8]. To be specific, learning a DDPM with ϵ -prediction parameterization is equivalent to fitting the finite-time marginal of a sampling chain resembling annealed Langevin dynamics [7, 8, 31]. Inspired by this connection, we propose to amortize the long-run MCMC in learning energy-based prior by iteratively distilling the short-run chain segments with a DDPM-based sampler. We show empirically and theoretically that the learned sampler is valid for long-run chain sampling.

Amortized MCMC The amortized MCMC technique is formally brought up by Li et al. [36], which incorporates feedback from MCMC back to the parameters of the amortizer distribution q ϕ . It is concurrently and independently proposed by Xie et al. [37] as the MCMC teaching framework. Methods under this umbrella term [36-38, 63-66] generally learns the amortizer by minimizing the divergence (typically the KLD) between the improved distribution and its initialization, i.e., D r K T q ϕ k ´ 1 || q ϕ s , where K T represents T -step MCMC transition kernel and q ϕ k ´ 1 represents the current amortizer. The diffusion-based amortization proposed in this work can be viewed as an instantiation of this framework, while our focus is on learning the energy-based prior. Compared with previous methods, our method i) specifically exploits the connection between EBMs and DDPMs and is suitable for amortizing the prior and posterior sampling MCMC of energy-based prior, and ii) resides in the lower-dimensional latent space and enables faster sampling and better convergence.

More methods for learning EBM Several techniques other than short-run MCMC have been proposed to learn the EBM. In the seminal work, Hinton [67] proposes to initialize Markov chains using real data and run several steps of MCMC to obtain samples from the model distribution. Tieleman [57] proposes to start Markov chains from past samples in the previous sampling iteration, known as Persistent Contrastive Divergence (PCD) or persistent chain sampling, to mimic the long-run sampling chain. Nijkamp et al. [23] provide comprehensive discussions about tuning choices for LD such as the step size s and sampling steps T to obtain stable long-run samples for persistent training. [59, 60] employ a hybrid of persistent chain sampling and short-run sampling by maintaining a buffer of previous samples. The methods draw from the buffer or initialize the short-run chain with noise distribution with some pre-specified probability. Another branch of work, stemmed from [68], considers discriminative contrastive estimation to avoid MCMC sampling. Gao et al. [69] use a normalizing flow [4] as the base distribution for contrastive estimation. Aneja et al. [50] propose to estimate the energy-based prior model based on the prior of a pre-trained V AE [70] by noise contrastive estimation. More recently, Xiao and Han [41] learn a sequence of EBMs in the latent space with adaptive multi-stage NCE to further improve the expressive power of the model.

6 Conclusion

In this paper, we propose the DAMC sampler and develop a novel learning algorithm for LEBM based on it. We provide theoretical and empirical evidence for the effectiveness of our method. We notice that our method can be applied to amortizing MCMC sampling of unnormalized continuous densities in general. It can also be applied to sampling posterior distributions of continuous latent variables in general latent variable models. We would like to explore these directions in future work.

References

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of International Conference on Machine Learning (ICML) , 2014.

Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems (NeurIPS) , 2014.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning , pages 1530-1538. PMLR, 2015.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of International Conference on Machine Learning (ICML) , 2015.

Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems , 2019.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR) , 2020.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4401-4410, 2019.

Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761 , 2020.

Peiyu Yu, Sirui Xie, Xiaojian Ma, Yixin Zhu, Ying Nian Wu, and Song-Chun Zhu. Unsupervised foreground extraction via deep region competition. Advances in Neural Information Processing Systems , 2021.

Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2021.

Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Peiyu Yu, Sirui Xie, Xiaojian Ma, Baoxiong Jia, Bo Pang, Ruigi Gao, Yixin Zhu, Song-Chun Zhu, and Ying Nian Wu. Latent diffusion energy-based model for interpretable text modeling. arXiv preprint arXiv:2206.05895 , 2022.

Jun-Yan Zhu, Philipp Krähenbühl, Eli Shechtman, and Alexei A Efros. Generative visual manipulation on the natural image manifold. In Proceedings of European Conference on Computer Vision (ECCV) , 2016.

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 2017.

Jiapeng Zhu, Yujun Shen, Deli Zhao, and Bolei Zhou. In-domain gan inversion for real image editing. In Proceedings of European Conference on Computer Vision (ECCV) , 2020.

Yilun Du, Shuang Li, and Igor Mordatch. Compositional visual generation with energy based models. Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. In International Conference on Learning Representations (ICLR) , 2020.

Weili Nie, Arash Vahdat, and Anima Anandkumar. Controllable and compositional generation with latent-space energy-based models. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Erik Nijkamp, Mitch Hill, Tian Han, Song-Chun Zhu, and Ying Nian Wu. On the anatomy of mcmc-based maximum likelihood learning of energy-based models. In Proceedings of AAAI Conference on Artificial Intelligence (AAAI) , 2020.

Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu. Learning non-convergent nonpersistent short-run mcmc toward energy-based model. In Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energybased models by diffusion recovery likelihood. In International Conference on Learning Representations (ICLR) , 2020.

Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. In Proceedings of the Conference on Computational Natural Language Learning (CoNLL) , 2016.

Ali Razavi, Aäron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with delta-vaes. arXiv preprint arXiv:1901.03416 , 2019.

James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! a linear vae perspective on posterior collapse. Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In Proceedings of International Conference on Machine Learning (ICML) , 2016.

Jakub Tomczak and Max Welling. Vae with a vampprior. In Proceedings of the International Conference on Artificial Intelligence and Statistics (AISTATS) , 2018.

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of International Conference on Machine Learning (ICML) , 2011.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Proceedings of International Conference on Machine Learning (ICML) , 2021.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Yingzhen Li, Richard E Turner, and Qiang Liu. Approximate inference with amortised mcmc. arXiv preprint arXiv:1702.08343 , 2017.

Jianwen Xie, Yang Lu, Ruiqi Gao, and Ying Nian Wu. Cooperative learning of energy-based model and latent variable model via mcmc teaching. In Proceedings of AAAI Conference on Artificial Intelligence (AAAI) , 2018.

Francisco Ruiz and Michalis Titsias. A contrastive divergence for combining variational inference and mcmc. In Proceedings of International Conference on Machine Learning (ICML) , 2019.

Thomas M Cover. Elements of information theory . John Wiley & Sons, 1999.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications .

Zhisheng Xiao and Tian Han. Adaptive multi-stage density ratio estimation for learning latent space energy-based model. Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mnist/ , 1998.

Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision , pages 37303738, 2015.

Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.

Cheng-Han Lee, Ziwei Liu, Lingyun Wu, and Ping Luo. Maskgan: Towards diverse and interactive facial image manipulation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2020.

Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:1506.03365 , 2015.

Bin Dai and David Wipf. Diagnosing and enhancing vae models. International Conference on Learning Representations (ICLR) , 2019.

Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. International Conference on Learning Representations (ICLR) , 2020.

Jyoti Aneja, Alex Schwing, Jan Kautz, and Arash Vahdat. A contrastive learning approach for training variational autoencoder priors. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Tian Han, Yang Lu, Song-Chun Zhu, and Ying Nian Wu. Alternating back-propagation for generator network. In Proceedings of AAAI Conference on Artificial Intelligence (AAAI) , 2017.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Rameen Abdal, Yipeng Qin, and Peter Wonka. Image2stylegan: How to embed images into the stylegan latent space? In Proceedings of International Conference on Computer Vision (ICCV) , 2019.

Weihao Xia, Yulun Zhang, Yujiu Yang, Jing-Hao Xue, Bolei Zhou, and Ming-Hsuan Yang. Gan inversion: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2022.

Rithesh Kumar, Sherjil Ozair, Anirudh Goyal, Aaron Courville, and Yoshua Bengio. Maximum entropy generators for energy-based models. arXiv preprint arXiv:1901.08508 , 2019.

Houssam Zenati, Chuan Sheng Foo, Bruno Lecouat, Gaurav Manek, and Vijay Ramaseshan Chandrasekhar. Efficient gan-based anomaly detection. Proceedings of the 20th IEEE International Conference on Data Mining (ICDM), 2018 , 2018.

Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of International Conference on Machine Learning (ICML) , 2008.

Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu. A theory of generative convnet. In Proceedings of International Conference on Machine Learning (ICML) , 2016.

Yilun Du and Igor Mordatch. Implicit generation and generalization in energy-based models. arXiv preprint arXiv:1903.08689 , 2019.

Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. In Proceedings of International Conference on Machine Learning (ICML) , 2021.

Bo Pang, Tian Han, and Ying Nian Wu. Learning latent space energy-based prior model for molecule generation. arXiv preprint arXiv:2010.09351 , 2020.

Bo Pang, Tianyang Zhao, Xu Xie, and Ying Nian Wu. Trajectory prediction with latent belief energy-based model. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2021.

Jianwen Xie, Zilong Zheng, Ruiqi Gao, Wenguan Wang, Song-Chun Zhu, and Ying Nian Wu. Generative voxelnet: learning energy-based models for 3d shape synthesis and analysis. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(5):2468-2484, 2020.

Jianwen Xie, Zilong Zheng, and Ping Li. Learning energy-based model with variational autoencoder as amortized sampler. In Proceedings of AAAI Conference on Artificial Intelligence (AAAI) , 2021.

Jianwen Xie, Yaxuan Zhu, Jun Li, and Ping Li. A tale of two flows: cooperative learning of langevin flow and normalizing flow toward energy-based model. arXiv preprint arXiv:2205.06924 , 2022.

Yaxuan Zhu, Jianwen Xie, Yingnian Wu, and Ruiqi Gao. Learning energy-based models by cooperative diffusion recovery likelihood. arXiv preprint arXiv:2309.05153 , 2023.

Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation , 14(8):1771-1800, 2002.

Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , 2010.

Ruiqi Gao, Erik Nijkamp, Diederik P Kingma, Zhen Xu, Andrew M Dai, and Ying Nian Wu. Flow contrastive estimation of energy-based models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2020.

Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems , 2020.

JM Bernardo, MJ Bayarri, JO Berger, AP Dawid, D Heckerman, AFM Smith, M West, et al. The variational bayesian em algorithm for incomplete data: with application to scoring graphical model structures. Bayesian statistics , 7(453-464):210, 2003.

Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434 , 2015.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2015.

Alexander Li and Deepak Pathak. Functional regularization for reinforcement learning via learned fourier features. Advances in Neural Information Processing Systems , 2021.

Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:1607.08022 , 2016.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Shohei Taniguchi, Yusuke Iwasawa, Wataru Kumagai, and Yutaka Matsuo. Langevin autoencoders for learning deep latent variable models. Advances in Neural Information Processing Systems , 35:13277-13289, 2022.

A Theoretical Discussion

A.1 Monotonically Decreasing KLD

We state in the main text that D r q ϕ k || π s ď D r q ϕ k ´ 1 || π s , where π is the stationary distribution. To show this, we first provide a proof of D r π t ` T || π s ď D r π t || π s , where π t and π t ` T are the distributions of z at t -th and p t ` T q -th iteration, respectively. This is a known result from [39], and we include it here for completeness.

where we denote the forward-time transition kernel as K and the reverse-time kernel as K 1 . (i) holds because we are just re-factorizing the joint density of r z t , z t ` 1 s : π t p z t q K p z t ` 1 | z t q ' K 1 π t ` 1 p z t | z t ` 1 q π t ` 1 p z t ` 1 q and π p z t q K p z t ` 1 | z t q ' K 1 π p z t | z t ` 1 q π p z t ` 1 q . (ii) holds because the KLD is non-negative. We can see that D r π t ` T || π s ď D r π t || π s is a direct result from Eq. (10), and that π t Ñ π as t Ñ8 under proper conditions [32].

In the main text, we describe the update rule of the sampler q ϕ as follows:

In the ideal case, we can assume that the objective in Eq. (6) is properly optimized and that t q ϕ u is expressive enough to parameterize each q ϕ k ´ 1 ,T . With q ϕ 0 « π 0 and q ϕ k « K T q ϕ k ´ 1 , we can conclude that D r q ϕ k || π s ď D r q ϕ k ´ 1 || π s for each k ' 1 , ..., K according to Eq. (10). We will discuss in the following section the scenario where we apply gradient-based methods to minimize D r q ϕ k ´ 1 ,T || q ϕ s and approximate Eq. (6) for the update from q ϕ k ´ 1 to q ϕ k .

A.2 Discussion about Diffusion-Based Amortization

We can see that the statement in Appendix A.1 holds when q ϕ k is a close approximation of q T : ' K T q ϕ k ´ 1 . This motivates our choice of employing DDPM to amortize the LD transition, considering its capability of close approximation to the given distribution t q T u . Based on the derivation of DDPM learning objective in [1], we know that

where z λ is draw from q p z λ | z 0 q ' N p z λ ; β λ z 0 , σ 2 λ I d q . z 0 ' q T . ϵ ' N p 0 , I d q and λ is drawn from a distribution of log noise-to-signal ratio p p λ q . In practice, we use Monte-Carlo average to approximate the objective and employ a gradient-based update rule for q ϕ :

This can be viewed as a M-estimation of ϕ . Recall that q T ' K T q ϕ k ´ 1 , we can construct ϕ k ' ´ ϕ k ´ 1 , ˜ ϕ ¯ to minimize the KLD, where ˜ ϕ models the transition kernel K T . Therefore, initializing q ϕ to be optimized with q ϕ k ´ 1 , we are effectively maximizing L p ˜ ϕ q ' 1 N ř N j ' 1 log p ˜ ϕ p ˆ z j | z j q ,

where t z j u are from q ϕ k ´ 1 and t ˆ z j u are from K T q ϕ k ´ 1 . Let ˜ ϕ be the result of M-estimation and ˜ ϕ ˚ be the true target parameter. Then based on the derivation in [25], asymptotically we have

where I ´ ˜ ϕ ˚ ¯ ' E ˆ z , z ' ´ ∇ 2 log p ˜ ϕ p ˆ z | z q ı is the Fisher information matrix. This interpretation tells us that i) when the sample size N is large, the estimation ˜ ϕ is asymptotically unbiased, and ii) if we want to obtain the estimation ˜ ϕ with a few gradient-based updates, then the eigenvalues of the Fisher information matrix would be relatively small but non-zero. ii) suggests that K T q ϕ k ´ 1 should be significantly different from q ϕ k ´ 1 , which is confirmed by [38] and our preliminary experiments, but it should not be too far away because that would require more gradient-based updates. We find that setting T ' 30 and M ' 6 works well in the experiments.

A.3 Further Discussion about the Learning Algorithm

For completeness, we first derive the learning gradients for updating θ .

We can see that θ is estimated in a MLE-by-EM style. The learning gradient is the same as that of directly maximizing the observed data likelihood, while we need to approximate the expectations in Eq. (15). Estimating the expectations is like the E-step, and update θ with Eq. (15) is like the M-step in the EM algorithm. The proposed diffusion-based amortization brings better estimation of the expectations in the E-step, and incorporate the feedback from the M-step by running prior and posterior sampling LD as follows

to obtain training data. Here t ' 0 , 1 , ..., T . z 0 ' q ϕ p z | x q for posterior sampling and z 0 ' q ϕ p z q for prior sampling. w t ' N p 0 , I d q . Note that we plug-in p θ p x , z q for the target distribution of posterior sampling LD. This is because given the observed data x , by Bayes' rule we know that p θ p z | x q9 p θ p x , z q ' p β p x | z q p α p z q . The whole learning iteration can be viewed as a variant of the variational EM algorithm [71].

B Network Architecture and Training Details

Architecture The energy score network f α uses a simple fully-connected structure throughout the experiments. We describe the architecture in details in Table S1. The generator network has a simple deconvolution structure similar to DCGAN [72] shown in Table S2. The denoising diffusion model is implemented by a light-weight MLP-based U-Net [73] structure (Table S3). The encoder network to embed the observed images has a fully convolutional structure [74], as shown in Table S4.

Hyperparameters and training details As mentioned in the main text, for the posterior and prior DAMCsamplers, we set the number of diffusion steps to 100 . The number of iterations in Eq. (8) is set to M ' 6 for the experiments. The LD runs T ' 30 and T ' 60 iterations for posterior and prior updates during training with a step size of s ' 0 . 1 . For test time sampling from K T, z i | x i q ϕ k p z i | x i q , we set T ' 10 for the additional LD. For test time prior sampling of LEBM with LD, we follow [22, 41] and set T ' 100 . To further stabilize the training procedure, we i) perform gradient clipping by setting the maximal gradient norm as 100 , ii) use a separate target diffusion network which is the EMAof the current diffusion network to initialize the prior and posterior updates and iii) add noiseinitialized prior samples for the prior updates. These set-ups are identical across different datasets.

The parameters of all the networks are initialized with the default pytorch methods [77]. We use the Adam optimizer [78] with β 1 ' 0 . 5 and β 2 ' 0 . 999 to train the generator network and the energy score network. We use the AdamW optimizer [79] with β 1 ' 0 . 5 , β 2 ' 0 . 999 and weight_decay=1e-4 to train the diffusion network. The initial learning rates of the generator and diffusion networks are 2e-4 , and 1e-4 for the energy score network. The learning rates are decayed with a factor of 0 . 99 every 1K training iterations, with a minimum learning rate of 1e-5 . We run the experiments on a A6000 GPU with the batch size of 128 . For GAN inversion, we reduce the batch size to 64. Training typically converges within 200K iterations on all the datasets.

C Pytorch-style Pseudocode

We provide pytorch-style pseudocode to help understand the proposed method. We denote the generator network as G , the energy score network as E and the diffusion network as Q . The first page sketches the prior and posterior sampling process. The second page outlines the learning procedure.

Listing 2: Learning LEBM with DAMC.

D Dataset and Experiment Settings

Datasets We include the following datasets to study our method: SVHN (32 × 32 × 3), CIFAR10 (32 × 32 × 3), CelebA (64 × 64 × 3), CeleAMask-HQ (256 x 256 x 3) and MNIST (28 x 28 x 1). Following Pang et al. [22], we use the full training set of SVHN (73,257) and CIFAR-10 (50,000), and take 40,000 samples of CelebA as the training data. We take 29,500 samples from the CelebAMask-HQ dataset as the training data, and test the model on 500 held-out samples. For anomaly detection on MNIST dataset, we follow the experimental settings in [22, 41, 55, 56] and use 80% of the in-domain data to train the model. The images are scaled to r´ 1 , 1 s and randomly horizontally flipped with a prob. of .5 for training.

GAN inversion settings We attempt to use the DAMC sampler for GAN version on the FFHQ (256 x 256 x 3) and LSUN-Tower (256 x 256 x 3) datasets. We take 69,500 samples from the CelebAMask-HQ dataset as the training data, and use the held-out 500 samples for testing. We follow the default data splits of the LSUN dataset.

For the LEBM-based inversion method, we train a LEBM in the 14 x 512 = 7168 dim. latent space. During training, we add l 2 -regularization on the energy score of LEBM to stabilize training, as suggested in [59]. For the DAMC sampler and the encoder-based inversion method [16], we generate initial posterior samples of the training data using [18] and train the DAMC sampler and the encoderbased method for 5K iterations using these samples as a warm-up step. The encoder-based method is trained by minimizing the l 2 distance between the encoder output and the target samples. After that, these two methods are trained with the default learning algorithms.

E Additional Qualitative Results

E.1 Toy Examples

(a) Evolution of the learned posterior distributions

(b) GT

As the proof-of-concept toy examples, we implement the neural likelihood experiments following the same set-up mentioned in Section 5.1 and B.1 in [80]. We choose to use a more complex prior distribution, i.e., 2-arm pinwheel-shaped prior distribution (shown in Fig. S1) instead of a standard normal one to make sure that the true posterior distributions are multimodal. The ground-truth posterior distributions are obtained by performing long-run LD sampling until convergence. We visualize the convergence trajectory of the posterior distributions learned by our model to the ground-truth ones. In Fig. S2, we can see that our model can faithfully reproduce the ground-truth distributions with sufficient training iterations, which indicates that the learned sampler successfully amortizes the long-run sampling chain with the length of 1 , 000 -3 , 000 iterations.

E.2 Generation

We provide additional generated samples from our models trained on SVHN (Fig. S3), CelebA (Fig. S4), CIFAR-10 (Fig. S5) and CelebA-HQ (Fig. S6).

(a) Ours-DAMC

8136518983

37916

16166

78

513

070114

224

260S

8

851218213

22

570898

1

81

261

6043612

5F

2553

8191

873343S18L

U

(a) Ours-DAMC

E.3 Reconstruction

We provide qualitative examples about the reconstruction results from our models trained on SVHN (Fig. S7), CelebA (Fig. S8), CIFAR-10 (Fig. S9) and CelebA-HQ (Fig. S10). Observed images are sampled from the testing set unseen during training.

(a) Ours-DAMC

(a) Ours-DAMC

E.4 Visualization of Transitions

We provide additional visualization results of LD transitions initialized from N p 0 , I d q on SVHN (Fig. S11) and CIFAR-10 datasets (Fig. S12). For the 200-step set-up, we can see that the generation quality quickly improves by exploring the local modes with LD. For the 2500-step long-run set-up, we can see that the LD produces consistently valid results without the oversaturating issue of the long-run chain samples.

(a) Observation

(b) Reconstruction

2

10

9

2

6

181893

4

4

181833

4

3

4

1919

344

0

1519

2

5

4

25

4

1631

60

16531

9

20

9

20

0

659

1515

20

659

1515

8

日

5157

6656

57157

29

99

1

1

(a) Observation

(b) Reconstruction

(a) Observation

(b) Reconstruction

(a) Observation

(b) Reconstruction

(a) 200 steps

(b) 2500 steps

1736822257319

1508989042961815.40539

58

0238131018

200

200

400

400

600

600

800

008-

1000

1000

1200

1200

0

50

75

100

125

150

175

0

500

1000

1500

2000

2500

(a) 200 steps

(b) 2500 steps

-200

-200

-400

400

600

600

-800

800

1000

50

95

100

125

150

175

500

1000

1500

2000

2500

F Additional Quantitative Results

F.1 Learning DAMC in the Latent Space of Other DLVMs

We have compared our method with directly training diffusion models in the latent space. In the NALR-LEBM column in Table 4, we compare our method with learning a diffusion model in the pre-trained energy-based latent space on CIFAR-10 dataset. We further add experiments of learning this model in the pre-trained V AE [1] and ABP [51] model latent spaces summarized in Table S5. We can see that our method greatly outperforms these baselines using the same network architectures.

F.2 Further Ablation Studies

We conduct more ablation studies to inspect the impact of some hyper-parameters. For prior LD sampling, we follow [22] and set the sampling step as 60 throughout the experiments for fair comparison. We further add experiments on CIFAR-10 for training our model with different posterior LD sampling steps T and with different capacities of the amortizer q ϕ , summarized in Table S6. We observe that larger posterior sampling steps and larger model capacities for training brings marginal improvement compared with our default set-up. Less sampling steps and lower model capacities, however, may have negative impact on the performances.

Table S6: Ablation Studies for the choice of Langevin steps and model capacity . We highlight the results of our set-up reported in the main text. T is the posterior sampling step. f stands for the factor for the model capacity, e.g ., f=2 means 2x the size of the original model.

G Further Discussion

G.1 Limitations

We mentioned in the main text that one potential disadvantage of our method is its parameter inefficiency for introducing an extra DDPM. Although fortunately, our models are in the latent space so the network is lightweight. To be specific, on SVHN, CelebA, CIFAR-10 and CelebA-HQ datasets the number of parameters in the diffusion network is around 10% of those in the generator.

Another issue is the time efficiency. We mentioned in the main text that the time efficiency for sampling is competitive. With the batch size of 64 , on these datasets the DAMC prior sampling takes 0 . 3 s , while 100 steps of short-run LD with LEBM takes 0 . 2 s . The DAMC posterior sampling takes 1 . 0 s , while LEBM takes 8 . 0 s . However, during training we need to run 30 steps of posterior LD sampling and 60 steps of prior LD sampling in each training iteration. We observe that the proposed learning method takes 15.2 minutes per training epoch, while the short-run LD-based learning method takes 14.8 minutes per epoch. These methods are slower than the V AE-base method, which takes 5.5 minutes for an training epoch. We can see that the time efficiency for training is generally bottlenecked by the LD sampling process, and could be improved in future works.

G.2 Broader Impacts

Generative models could be misused for disinformation or faking profiles. Our work focuses on the learning algorithm of energy-based prior model. Though we consider our work to be foundational and not tied to particular applications or deployments, it is possible that more powerful energybased generative models augmented with this method may be used maliciously. Work on the reliable detection of synthetic content could be important to address such harms from generative models.

Peer reviews:
comments: SUMMARY:
The paper presents a new training and sampling procedures for learning energy based generative models.  The method is compared to earlier (few-step) MCMC based approaches and to diffusion models.  The procedure is evaluated on computer vision tasks.

STRENGTHS:
The paper presents an original algorithm, with a clear methodological description, and a large collection of benchmarking results.

WEAKNESSES:
The paper could make clearer for the general audience why additional energy based modeling methodology is needed.  As a researcher who has been working extensively on diffusion models for the past two years, but has not worked with (other) energy based models in approaches in recent years additional context in goals of this direction and the limitations of previous approaches would be helpful.

For example, the authors comment that EBMs are "closely related to" DDPMs.  Which do not suffer from the "sampling issue" associated with "non-convergent short-run MCMC" used for EBMs.   Is there something that can be accomplished with EBMs that is not solved by DDPMs that motivates their continued study?  

The paper is also unclear with respect to the claim for "theoretical evidence that the learned amortization of MCMC is a valid long-run MCMC sampler".  Could the authors make this theoretical claim more explicit (e.g. as a proposition or theorem)?

QUESTIONS:
No further questions.

LIMITATIONS:
Limitations are adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: SUMMARY:
This paper proposes the DAMC sampler (Diffusion-Amortized MCMC) and develops a new learning algorithm for LEBM (Latent-space Energy-Based Model) based on it. Theoretical and empirical evidences are provided for the effectiveness of our method.

STRENGTHS:
The paper is generally well-written. The idea of amortizing the LD with DDPM in learning the LEBM seems to be new.

WEAKNESSES:
The basic idea of using an auxiliary model to amortize the LD in learning energy-based models has been well known and explored in the literatures, e.g. in [a,b] to name a few. Methodology connection and experiment comparison with those methods are needed.

By looking at FID on CIFAR-10 (Table 1), the performance of the proposed method (FID 57.72) is far behind previous methods [a,b] (33.61, 20.9).

In the current literature, long-run MCMC analysis is typically conducted over a range of 10,000 to 100,000 Langevian Dynamics iterations. However, in this paper (Figure 4), the authors perform a much shorter 2,500 Langevian dynamics update, which is considerably less extensive compared to other studies in the field.

[a] Cooperative Training of Descriptor and Generator Networks, arXiv:1609.09408v3
[b] Learning Neural Random Fields with Inclusive Auxiliary Generators, arXiv:1806.00271v4

QUESTIONS:
see above.

LIMITATIONS:
see above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: SUMMARY:
The authors propose DAMC, an amortization of MCMC sampling, via a scheme based on diffusion models, as an alternative to pure MCMC sampling, which usually suffers from either long mixing time or from being short and biased,for  priors and posteriors in energy based models. The method is theoretically sound, and experimentally argued for.

STRENGTHS:
The theoretical analysis is sound. The algorithm is clear. The experiments are convincing.

WEAKNESSES:
N/A

QUESTIONS:
Are there some constraints that the latent space needs to satisfy for the proposed method to have an advantage over the vanilla MCMC counterpart ?

LIMITATIONS:
I see None for now.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-13

comments: SUMMARY:
This paper proposed a diffusion-based amortised method to address the short-run MCMC samplers issues in the latent-space energy-based models. One interesting part is that it interleaves the distill T-steps of Langevin dynamics and KL divergence minimisation to sample from the target distribution $\pi$. Regarding the choice of amortised sampler parameters $\phi$, it uses the gradient of DDPM objective function to optimise them. The experimental results has verified the effectiveness of the DAMC.

STRENGTHS:
The paper has the following strengths:

**1**, the usages of KL minimization and amortised MCMC transition is a well defined framework. Also, the usages of DDPM gradient for the long-run MCMC sampling in LEBMs is interesting. 

**2**, the experimental evaluation is sufficient and clearly verifies the advantages of DAMC.

WEAKNESSES:
Regarding the weakness, I think the paper may need another round of polish. Some notations are not well defined before use and some introductory statement is not consistent. For example:

**1**, p_{uncond} may need more explanation, rather than a short word in the input of Algorithm. 
**2**, $z_+^{(i)}, z_-^{(i)}$ are a bit confusing. 
**3**, it may be a better idea if the notations of section 2.1 and 2.2 are consistent

QUESTIONS:
Sorry I am not an expert in this topic. I do not have questions for the authors.

LIMITATIONS:
I believe the authors has addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-16

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a diffusion-based amortized MCMC method for sampling the prior and posterior in latent space energy-based models. The paper provides some theoretical evidence using directly the result from Li et al., 2017. The paper shows the effectiveness of the proposed method throughout an extensive campaign on a variety of tasks such as image generation and construction, anomaly detection.\n\nSTRENGTHS:\n- This paper aims at tackling an important task in deep latent variable models, which is to improve energy-based prior model.\n- The paper is well-motivated and is well-written.\n- The experiments on image datasets are extensive and diverse.\n\nWEAKNESSES:\n- The idea of amortizing the short-run MCMC sampling of the prior and posterior distributions is already proposed in the original work of Pang et al., 2020. This paper aims to improve this amortization using a diffusion model. In addition, the theoretical contribution is weak, as the convergence property is already shown by Li et al., 2017.\n- Although the authors claim that the theoretical evidence that the learned amortization of MCMC is a valid long-run MCMC sampler, this result is taken directly from Li et al., 2017. As such, the authors should show that the proposed amortization scheme can converge to the true distributions, at least via toy examples. For example, the authors could use examples from this work [1].\n- The improvement is expectable as the authors use a very powerful diffusion model to amortize the prior and posterior. However, this increases the computational costs, as the authors discussed at the end of Section 4.1.\n- As, again, the authors employ a power diffusion model in latent space. To be fair, the authors should consider the baselines of using diffusion models in the latent space of latent variable models; for example [2], at the very least.\n\n[1] Taniguchi et al., Langevin Autoencoders for Learning Deep Latent Variable Models. NeurIPS 2022.\n\n[2] Pandey et al., DiffuseVAE: Efficient, Controllable and High-Fidelity Generation from Low-Dimensional Latents. TMLR 2022.\n\nQUESTIONS:\n- There are a lot of hyper-parameters in the proposed methods. How do the authors choose $T$?\n- It would be great if the authors could ablate the amortization gap of the proposed method by considering different capacities of the $q_{\\phi}$\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-16"}