SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In order to efficiently model both  $q_{\phi_k}(z_i|\pmb{x}_i)$  and  $q_{\phi_k}(z_i)$ , we follow the method of [40] to train a single network to parameterize both models, where  $q_{\phi_k}(z_i|\pmb{x}_i)$  can be viewed as a conditional DDPM with the embedding of  $\pmb{x}_i$  produced by an encoder network as its condition, and  $q_{\phi_k}(z_i)$  an unconditional one. For  $q_{\phi_k}(z_i)$ , we can input a null token  $\varnothing$  as its condition when predicting the noise  $\epsilon$ . We jointly train both models by randomly nullifying the inputs with the probability  $p_{\mathrm{uncond}} = 0.2$ . During training, we use samples from  $q_{\phi_k}(z_i|\pmb{x}_i)$  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
Table 1: MSE(↓) and FID(↓) obtained from models trained on different datasets. The FID scores are computed based on 50k generated images and training images for the first three datasets and 5k images for the CelebA-HQ dataset. The MSEs are computed based on unseen testing images. We highlight our model results in gray color. The best and second-best performances are marked in bold numbers and underlines, respectively; tables henceforth follow this format. *[41] uses a prior model with 4x parameters compared with [22] and ours.
Model	SVHN	CelebA	CIFAR-10	CelebA-HQ
MSE	FID	MSE	FID	MSE	FID	MSE	FID
VAE [1]	0.019	46.78	0.021	65.75	0.057	106.37	0.031	180.49
2s-VAE [48]	0.019	42.81	0.021	44.40	0.056	72.90	-	-
RAE [49]	0.014	40.02	0.018	40.95	0.027	74.16	-	-
NCP-VAE [50]	0.020	33.23	0.021	42.07	0.054	78.06	-	-
Adaptive CE* [41]	0.004	26.19	0.009	35.38	0.008	65.01	-	-
ABP [51]	-	49.71	-	51.50	0.018	90.30	0.025	160.21
SRI [24]	0.018	44.86	0.020	61.03	-	-	-	-
SRI (L=5) [24]	0.011	35.32	0.015	47.95	-	-	-	-
LEBM [22]	0.008	29.44	0.013	37.87	0.020	70.15	0.025	133.07
Ours-LEBM	0.002	21.17	0.005	35.67	0.015	60.89	0.023	89.54
Ours-DAMC	18.76	30.83	57.72	85.88
size of  $s = 0.1$ . For test time sampling from  $\mathcal{K}_{T,\boldsymbol{z}_i|\boldsymbol{x}_i}q_{\phi_k}(\boldsymbol{z}_i|\boldsymbol{x}_i)$ ,  $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.
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], CIFAR-10 [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 VAE [1], as well as recent two-stage methods such as 2-stage 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
Figure 2: Samples generated from the DAMC sampler and LEBM trained on SVHN, CelebA, CIFAR-10 and CelebA-HQ datasets. In each sub-figure, the first four rows are generated by the DAMC sampler. The last four rows are generated by LEBM trained with the DAMC sampler.
(d) CelebA-HQ
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 \times 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 high-dimensional 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 GAN inversion 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
Model	FFHQ	LSUN-T
MSE	FID	MSE	FID
Opt. [53]	0.055	149.39	0.080	240.11
Enc. [16]	0.028	62.32	0.079	132.41
[22] w/ 1x	0.054	149.21	0.072	239.51
[22] w/ 2x	0.039	101.59	0.066	163.20
[22] w/ 4x	0.032	84.64	0.059	111.53
Ours	0.025	52.85	0.059	80.42
Table 2:  $\mathbf{{MSE}}\left( \downarrow \right)$  and  $\mathbf{{FID}}\left( \downarrow \right)$  for GAN inversion on different datasets. Opt. and Enc. denotes the optimization-based and encoder-based methods.
(a) Observation
(b) Inversion
Figure 3: Qualitative results of StyleGAN inversion using the DAMC sampler. In each sub-figure, the left panel contains samples from the FFHQ dataset, and the right panel contains samples from the LSUN-T dataset.
(a) 200 steps
(b) 2500 steps
Figure 4: Transition of Markov chains initialized from  $\mathcal{N}(0,\mathbf{I}_d)$  towards  $p_{\alpha}(z)$ . We present results by running LD for 200 and 2500 steps. In each sub-figure, the top panel displays the trajectory in the data space uniformly sampled along the chain. The bottom panel shows the energy score  $f_{\alpha}(z)$  over the iterations.
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.3s$ , while 100 steps of short-run LD with LEBM takes  $0.2s$ . The DAMC posterior sampling takes  $1.0s$ , while LEBM takes  $8.0s$ . 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_{\alpha}(z)$  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  $\mathcal{N}(0,\mathbf{I}_d)$  towards  $p_{\alpha}(z)$  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 oversaturation 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_{\theta,\phi}(z|x)$  could form a discriminative latent space that has separated probability densities for in-distribution (normal) and out-of-distribution (anomalous) data. Given the testing sample  $x$ , we use un-normalized log joint density  $p_{\theta,\phi}(z|x) \propto p_{\theta,\phi}(x,z) \approx p_{\beta}(x|z)p_{\alpha}(z)|_{z \sim \mathcal{K}_{T,z|x}q_{\phi}(z|x)}$  as our decision function. This means that we draw samples from  $\mathcal{K}_{T,z|x}q_{\phi}(z|x)$  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
Table 3: AUPRC(↑) scores for unsupervised anomaly detection on MNIST. Numbers are taken from [41]. Results of our model are averaged over the last 10 trials to account for variance.
Heldout Digit	1	4	5	7	9
VAE [1]	0.063	0.337	0.325	0.148	0.104
ABP [51]	0.095 ± 0.03	0.138 ± 0.04	0.147 ± 0.03	0.138 ± 0.02	0.102 ± 0.03
MEG [55]	0.281 ± 0.04	0.401 ± 0.06	0.402 ± 0.06	0.290 ± 0.04	0.342 ± 0.03
BiGAN-σ [56]	0.287 ± 0.02	0.443 ± 0.03	0.514 ± 0.03	0.347 ± 0.02	0.307 ± 0.03
LEBM [22]	0.336 ± 0.01	0.630 ± 0.02	0.619 ± 0.01	0.463 ± 0.01	0.413 ± 0.01
Adaptive CE [41]	0.531 ± 0.02	0.729 ± 0.02	0.742 ± 0.01	0.620 ± 0.02	0.499 ± 0.01
Ours	0.684 ± 0.02	0.911 ± 0.01	0.939 ± 0.02	0.801 ± 0.01	0.705 ± 0.01
Table 4: Ablation study on CIFAR-10 dataset. VI denotes learning LEBM using variational methods. SR denotes learning LEBM with short-run LD. DAMC-G replaces the LEBM in DAMC-LEBM with a standard Gaussian distribution. NALR denotes the non-amortized DDPM setting. For each set-up, we provide results using the vanilla sampling method, denoted as V, and the ones using the DAMC sampler, denoted as D.
Model	VI-LEBM	SR-LEBM	DAMC-G	NALR-LEBM	DAMC-LEBM
V.	D.	V.	D.	V.	D.	V.	D.	V.	D.
MSE	0.054	-	0.020	-	0.018	0.015	0.028	0.016	0.021	0.015
FID	78.06	-	70.15	-	90.30	66.93	68.52	64.38	60.89	57.72
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.
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  $\mathcal{N}(0, \mathbf{I}_d)$ . 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 DAMC sampler 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.
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 MCMC for 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  $\epsilon$ -prediction parametrization resembles LD of an EBM;  $\epsilon$  (predicted noise) plays a similar role to the gradient of the log density [8]. To be specific, learning a DDPM with  $\epsilon$ -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_{\phi}$ . 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.,  $\mathcal{D}[\mathcal{K}_T q_{\phi_{k-1}} || q_{\phi}]$ , where  $\mathcal{K}_T$  represents  $T$ -step MCMC transition kernel and  $q_{\phi_{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 VAE [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.
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.
Y. N. Wu was supported by NSF DMS-2015577. We would like to thank the anonymous reviewers for their constructive comments.