SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Peiyu Yu
yupeiyu98@g.ucla.edu
Yaxuan Zhu
yaxuanzhu@g.ucla.edu
Sirui Xie
srxie@ucla.edu
Xiaojian Ma2,4
xiaojian.ma@ucla.edu
Ruiqi Gao<sup>3</sup>
ruiqig@google.com
Song-Chun Zhu
sczhu@stat.ucla.edu
Ying Nian Wu
$^{1}$ UCLA Department of Statistics  $^{2}$ UCLA Department of Computer Science  $^{3}$ Google DeepMind  $^{4}$ Beijing Institute for General Artificial Intelligence (BIGAI)
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 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 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<sup>1</sup>.
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 semi-supervised 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
$^{1}$ Code and data available at https://github.com/yuPeiyu98/Diffusion-Amortized-MCMC.
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 diffusion-based 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 short-run 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.1 Energy-Based Prior Model
We assume that for the observed sample  $\pmb{x} \in \mathbb{R}^D$ , there exists  $\pmb{z} \in \mathbb{R}^d$  as its unobserved latent variable vector. The complete-data distribution is
$$
p _ {\boldsymbol {\theta}} (\boldsymbol {z}, \boldsymbol {x}) := p _ {\boldsymbol {\alpha}} (\boldsymbol {z}) p _ {\boldsymbol {\beta}} (\boldsymbol {x} | \boldsymbol {z}), p _ {\boldsymbol {\alpha}} (\boldsymbol {z}) := \frac {1}{Z _ {\boldsymbol {\alpha}}} \exp (f _ {\boldsymbol {\alpha}} (\boldsymbol {z})) p _ {0} (\boldsymbol {z}), \tag {1}
$$
where  $p_{\alpha}(z)$  is the prior model with parameters  $\alpha$ ,  $p_{\beta}(\boldsymbol{x}|\boldsymbol{z})$  is the top-down generation model with parameters  $\beta$ , and  $\theta = (\alpha, \beta)$ . The prior model  $p_{\alpha}(z)$  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_{\alpha}(z)$  is parameterized by a neural network with scalar output,  $Z_{\alpha}$  is the partition function, and  $p_0(z)$  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_{\beta}(\boldsymbol{x}|\boldsymbol{z}) = \mathcal{N}(g_{\beta}(\boldsymbol{z}), \sigma^2 \mathbf{I}_D)$ , where  $g_{\beta}$  is the generator network and  $\sigma^2$  takes a pre-specified value as in VAE [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  $\pmb{x}$ , the gradients for updating  $\alpha, \beta$  are,
$$
\delta_ {\boldsymbol {\alpha}} (\boldsymbol {x}) := \mathbb {E} _ {p _ {\boldsymbol {\theta}} (\boldsymbol {z} | \boldsymbol {x})} \left[ \nabla_ {\boldsymbol {\alpha}} f _ {\boldsymbol {\alpha}} (\boldsymbol {z}) \right] - \mathbb {E} _ {p _ {\boldsymbol {\alpha}} (\boldsymbol {z})} \left[ \nabla_ {\boldsymbol {\alpha}} f _ {\boldsymbol {\alpha}} (\boldsymbol {z}) \right], \delta_ {\boldsymbol {\beta}} (\boldsymbol {x}) := \mathbb {E} _ {p _ {\boldsymbol {\theta}} (\boldsymbol {z} | \boldsymbol {x})} \left[ \nabla_ {\boldsymbol {\beta}} \log p _ {\boldsymbol {\beta}} (\boldsymbol {x} | \boldsymbol {z}) \right]. \tag {2}
$$
In practice, one may use the Monte-Carlo average to estimate the expectations in Eq. (2). This involves sampling from the prior  $p_{\alpha}(z)$  and the posterior  $p_{\theta}(z|x)$  distribution using MCMC, specifically Langevin Dynamics (LD) [32], to estimate the expectations and hence the gradient. For a target
Figure 1: Learning the DAMC sampler. The training samples for updating the sampler to  $\phi_{k + 1}$  is obtained by  $T$ -step short-run LD, initialized with the samples from the current learned sampler  $\phi_k$ . Best viewed in color.
distribution  $\pi (z)$ , the dynamics iterates
$$
\boldsymbol {z} _ {t + 1} = \boldsymbol {z} _ {t} + \frac {s ^ {2}}{2} \nabla_ {\boldsymbol {z} _ {t}} \log \pi (\boldsymbol {z} _ {t}) + s \boldsymbol {w} _ {t}, t = 0, 1, \dots , T - 1, \boldsymbol {w} _ {t} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I} _ {d}), \tag {3}
$$
where  $s$  is a small step size. One can draw  $z_0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I}_d)$  to initialize the chain. For sufficiently small step size  $s$ , the distribution of  $z_t$  will converge to  $\pi$  as  $t \to \infty$  [32]. However, it is prohibitively expensive to run LD until convergence in most cases, for which we may resort to limited iterations of LD for sampling in practice. This non-convergent short chain yields a moment-matching distribution close to the true  $\pi(z)$  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  $\epsilon$ -prediction parametrization resembles LD;  $\epsilon$  (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  $\phi$  is specified by a noise schedule built upon  $\lambda_{s} = \log [\beta_{s}^{2} / \sigma_{s}^{2}]$ , i.e., the log signal-to-noise-ratio, that decreases monotonically with  $s$ .  $\beta_{s}$  and  $\sigma_s^2$  are strictly positive scalar-valued functions of  $s$ . We use  $z_0$  to denote training data in  $\mathbb{R}^d$ . The forward-time diffusion process  $q(z|z_0)$  is defined as:
$$
q (\boldsymbol {z} _ {s} | \boldsymbol {z} _ {0}) = \mathcal {N} (\boldsymbol {z} _ {s}; \beta_ {s} \boldsymbol {z} _ {0}, \sigma_ {s} ^ {2} \mathbf {I} _ {d}), \quad q (\boldsymbol {z} _ {s} ^ {\prime} | \boldsymbol {z} _ {s}) = \mathcal {N} (\boldsymbol {z} _ {s} ^ {\prime}; (\beta_ {s ^ {\prime}} / \beta_ {s}) \boldsymbol {z} _ {s}, \sigma_ {s ^ {\prime} | s} ^ {2} \mathbf {I} _ {d}), \tag {4}
$$
where  $0 \leqslant s < s' \leqslant S$  and  $\sigma_{s'|s}^2 = (1 - e^{\lambda_{s'} - \lambda_s})\sigma_s^2$ . Noticing that the forward process can be reverted as  $q(z_s|z_{s'}, z_0) = \mathcal{N}(z_s; \tilde{\mu}_{s|s'}(z_{s'}, z_0), \tilde{\sigma}_{s|s'}^2 \mathbf{I}_d)$ , an ancestral sampler  $q_\phi(z_s|z_{s'})$  [8] that starts at  $z_S \sim \mathcal{N}(\mathbf{0}, \mathbf{I}_d)$  can be derived accordingly [33]:
$$
\begin{array}{l} \tilde {\pmb {\mu}} _ {s | s ^ {\prime}} (\pmb {z} _ {s ^ {\prime}}, \pmb {z} _ {0}) = e ^ {\lambda_ {s ^ {\prime}} - \lambda_ {s}} (\alpha_ {s} / \alpha_ {s ^ {\prime}}) \pmb {z} _ {s ^ {\prime}} + (1 - e ^ {\lambda_ {s ^ {\prime}} - \lambda_ {s}}) \alpha_ {s} \pmb {z} _ {0}, \quad \tilde {\sigma} _ {s | s ^ {\prime}} ^ {2} = (1 - e ^ {\lambda_ {s ^ {\prime}} - \lambda_ {s}}) \sigma_ {s} ^ {2}, \\ \boldsymbol {z} _ {s} = \tilde {\boldsymbol {\mu}} _ {s | s ^ {\prime}} \left(\boldsymbol {z} _ {s ^ {\prime}}, \hat {\boldsymbol {z}} _ {0}\right) + \sqrt {\left(\tilde {\sigma} _ {s | s ^ {\prime}} ^ {2}\right) ^ {1 - \gamma} \left(\sigma_ {s ^ {\prime} | s} ^ {2}\right) ^ {\gamma}} \boldsymbol {\epsilon}, \tag {5} \\ \end{array}
$$
where  $\epsilon$  is standard Gaussian noise,  $\hat{z}_0$  is the prediction of  $z_0$  by the DDPM  $\phi$ , and  $\gamma$  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  $\mathbb{E}_{\epsilon,\lambda}\left[\|\epsilon(z_{\lambda}) - \epsilon\|_2^2\right]$ , where  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I}_d)$  and  $\lambda$  is drawn from a distribution of log noise-to-signal ratio  $p(\lambda)$  over uniformly sampled times  $s \in [0,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.
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  $\pi_0(z)$ , the distribution after  $t$ -th iteration as  $\pi_t(z)$  and the target distribution as  $\pi(z)$ . The trajectory of LD in Eq. (3) is typically specified by its transition kernel  $\mathcal{K}(z|z')$ . The process starts with drawing  $z_0$  from  $\pi_0(z)$  and iteratively sample  $z_t$  at the  $t$ -th iteration from the transition kernel conditioned on  $z_{t-1}$ , i.e.,  $\pi_t(z) = \mathcal{K}\pi_{t-1}(z)$ , where  $\mathcal{K}\pi_{t-1}(z) \coloneqq \int \mathcal{K}(z|z')\pi_{t-1}(z')dz'$ . Recursively,  $\pi_t = \mathcal{K}_t\pi_0$ , where  $\mathcal{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  $\pi$  is a stationary distribution  $\pi(z) \coloneqq \int \mathcal{K}(z|z')\pi(z')dz', \forall 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  $\mathcal{Q} = \{q_{\phi}\}$  be the set of amortized samplers parameterized by  $\phi$ . Given the transition kernel  $\mathcal{K}$ , the goal is to find a sampler  $q_{\phi*}$  to closely approximate the target distribution  $\pi$ . This can be achieved by iteratively distill  $T$ -step transitions of LD into  $q_{\phi}$ :
$$
q _ {\phi_ {k}} \leftarrow \underset {q _ {\phi} \in \mathcal {Q}} {\arg \min } \mathcal {D} \left[ q _ {\phi_ {k - 1}, T} | | q _ {\phi} \right], q _ {\phi_ {k - 1}, T} := \mathcal {K} _ {T} q _ {\phi_ {k - 1}}, q _ {\phi_ {0}} \approx \pi_ {0}, k = 0, \dots , K - 1. \tag {6}
$$
where  $\mathcal{D}[\cdot ||\cdot ]$  is the Kullbeck-Leibler Divergence (KLD) measure between distributions. Concretely, Eq. (6) means that to recover the target distribution  $\pi$ , 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_{\phi_{k - 1}}$  to approximate  $\kappa_T q_{\phi_{k - 1}}$  as the target distribution of the current sampler, and ii) update the current sampler  $q_{\phi_{k - 1}}$  to  $q_{\phi_k}$ . The correct convergence of  $q_{\phi}$  to  $\pi$  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,  $\mathcal{D}[q_{\phi_k}||\pi ]\leqslant \mathcal{D}[q_{\phi_{k - 1}}||\pi ]$ . 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  $\mathcal{D}[q_{\phi_{k - 1}},T||q_{\phi}]$  and approximate Eq. (6) for the update from  $q_{\phi_{k - 1}}$  to  $q_{\phi_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_{\phi_{k-1},T}$  as  $q_T$ . We can see that
$$
\underset {q _ {\phi}} {\arg \min } \mathcal {D} [ q _ {T} | | q _ {\phi} ] = \underset {q _ {\phi}} {\arg \min } - \mathcal {H} (q _ {T}) + \mathcal {H} (q _ {T}, q _ {\phi}) = \underset {q _ {\phi}} {\arg \min } - \mathbb {E} _ {q _ {T}} \left[ \log q _ {\phi} \right], \tag {7}
$$
where  $\mathcal{H}$  represents the entropy of distributions. The selection of the sampler  $q_{\phi}$  is a matter of design. According to Eq. (7), we may expect the following properties from  $q_{\phi}$ : 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  $\{q_T\}$ . In practice, iii) is important for the convergence of Eq. (6). If  $q_{\phi}$  is far away from  $q_{T}$  in each iteration, then non-increasing KLD property  $\mathcal{D}[q_{\phi_k}||\pi ]\leqslant \mathcal{D}[q_{\phi_{k - 1}}||\pi ]$  might not hold, and the resulting amortized sampler would not converge to the true target distribution  $\pi (z)$ .
For the choice of  $q_{\phi}$ , 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  $\epsilon$ -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_{\phi}$ , to obtain the gradient-based update rule for  $q_{\phi}$ :
$$
\phi_ {k - 1} ^ {(i + 1)} \leftarrow \phi_ {k - 1} ^ {(i)} - \eta \nabla_ {\phi} \mathbb {E} _ {\epsilon , \lambda} \left[ \| \epsilon (z _ {\lambda}) - \epsilon \| _ {2} ^ {2} \right], \phi_ {k} ^ {(0)} \leftarrow \phi_ {k - 1} ^ {(M)}, i = 0, 1, \dots , M - 1 \tag {8}
$$
where  $\epsilon \sim \mathcal{N}(0,\mathbf{I}_d)$ .  $\lambda$  is drawn from a distribution of log noise-to-signal ratio  $p(\lambda)$ .  $\eta$  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_{\phi}$  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  $\{\pmb{x}_i\}_{i=1}^N$  independently drawn from the unknown data distribution  $p_{\mathrm{data}}(\pmb{x})$ , the model  $p_\theta$  (Section 2.1) can be trained
by maximizing the log-likelihood over training samples  $\mathcal{L}(\pmb{\theta}) = \frac{1}{N}\sum_{i=1}^{N}\log p_{\pmb{\theta}}(\pmb{x}_i)$ . Doing so typically requires computing the gradients of  $\mathcal{L}(\pmb{\theta})$ , where for each  $\pmb{x}_i$  the learning gradient satisfies:
$$
\begin{array}{l} \nabla_ {\boldsymbol {\theta}} \log p _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {i}) = \mathbb {E} _ {p _ {\boldsymbol {\theta}} (\boldsymbol {z} _ {i} | \boldsymbol {x} _ {i})} \left[ \nabla_ {\boldsymbol {\theta}} \log p _ {\boldsymbol {\theta}} (\boldsymbol {z} _ {i}, \boldsymbol {x} _ {i}) \right] \\ = \left(\underbrace {\mathbb {E} _ {p _ {\boldsymbol {\theta}} \left(\boldsymbol {z} _ {i} \mid \boldsymbol {x} _ {i}\right)} \left[ \nabla_ {\boldsymbol {\alpha}} f _ {\boldsymbol {\alpha}} \left(\boldsymbol {z} _ {i}\right) \right] - \mathbb {E} _ {p _ {\boldsymbol {\alpha}} \left(\boldsymbol {z} _ {i}\right)} \left[ \nabla_ {\boldsymbol {\alpha}} f _ {\boldsymbol {\alpha}} \left(\boldsymbol {z} _ {i}\right) \right]} _ {\delta_ {\boldsymbol {\alpha}} \left(\boldsymbol {x} _ {i}\right)}, \underbrace {\mathbb {E} _ {p _ {\boldsymbol {\theta}} \left(\boldsymbol {z} _ {i} \mid \boldsymbol {x} _ {i}\right)} \left[ \nabla_ {\boldsymbol {\beta}} \log p _ {\boldsymbol {\beta}} \left(\boldsymbol {x} _ {i} \mid \boldsymbol {z} _ {i}\right) \right]}\right). \tag {9} \\ \end{array}
$$
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  $[z_i|x_i] \sim \mathcal{K}_{T,z_i|x_i} q_{\phi_k}(z_i|x_i)$ ,  $z_i \sim \mathcal{K}_{T,z_i} q_{\phi_k}(z_i)$ , to estimate the expectations and hence the learning gradient. Here we learn the DAMC samplers  $q_{\phi_k}(z_i|x_i)$  and  $q_{\phi_k}(z_i)$  for the posterior and prior sampling chain, respectively.  $\phi_k$  represents the current sampler as in Section 3.1;  $\mathcal{K}_{T,z_i|x_i}$  and  $\mathcal{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_{\phi_k}$  model, and ii) update the approximation of the prior  $p_{\alpha}(z)$  and posterior  $p_{\theta}(z|x)$  distributions with additional  $T$ -step LD initialized with  $q_{\phi_k}$  samples. These updated samples are closer to  $p_{\theta}(z|x_i)$  and  $p_{\alpha}(z_i)$  compared with short-run LD samples based on our discussion in Section 2.2. Consequently, the diffusion-amortized 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  $\theta = (\alpha, \beta)$  based on these approximate samples with Eq. (9), we can update  $q_{\phi_k}$  with Eq. (8) to distill the sampling chain into  $q_{\phi_{k+1}}$ . As shown in Fig. 1, we can see that the whole learning procedure iterates between the approximate MLE of  $p_{\theta}$  and the amortization of MCMC with  $q_{\phi}$ . 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_{\phi}(z_i)$  with Eq. (5). Prior sampling with LEBM still requires short-run LD initialized from  $\mathcal{N}(0,\mathbf{I}_d)$ . For posterior sampling, we may sample from  $\mathcal{K}_{T,z_i|\boldsymbol{x}_i}q_{\phi_k}(\boldsymbol{z}_i|\boldsymbol{x}_i)$ , i.e., first draw samples from DAMC and then run few steps of LD to obtain posterior samples.
Algorithm 1: Learning algorithm of DAMC.
Input: initial parameters  $(\alpha, \beta, \phi)$ ; learning rate  $\eta = (\eta_{\alpha}, \eta_{\beta}, \eta_{\phi})$ ; observed examples  $\{\pmb{x}^{(i)}\}_{i=1}^{N}$ ; prob. of uncond. training  $p_{\text{uncond}}$  for the DAMC sampler.
Output:  $\left(\pmb{\theta}^{(K)} = \{\pmb{\alpha}^{(K)},\pmb{\beta}^{(K)}\},\pmb{\phi}^{(K)}\right)$ .
1 for  $k = 0:K - 1$  do
2 Sample a minibatch of data  $\{\pmb{x}^{(i)}\}_{i = 1}^{B}$
3 Draw DAMC samples: For each  $x^{(i)}$ , draw  $z_{+}^{(i)}$  and  $z_{-}^{(i)}$  from  $q_{\phi_k}(z_i|x_i)$ .
4 | Prior LD update: For each  $\pmb{x}^{(i)}$ , update  $z_{-}^{(i)}$  using Eq. (3), with  $\pi(\pmb{z}_i) = p_{\pmb{\alpha}^{(k)}}(\pmb{z}_i)$ ;
5 Posterior LD update: For each  $\pmb{x}^{(i)}$ , update  $\pmb{z}_{+}^{(i)}$  using Eq. (3), with  $\pi(\pmb{z}_i) = p_{\beta^{(k)}}(\pmb{z}_i|\pmb{x}_i)$ ;
6 Update  $\pmb{\theta}^{(k)}$ : Update  $\alpha^{(k)}$  and  $\beta^{(k)}$  using Monte-Carlo estimates (i.e., Monte-Carlo average) of Eq. (9) with  $\{z_{+}^{(i)}\}_{i=1}^{B}$  and  $\{z_{-}^{(i)}\}_{i=1}^{B}$ .
7 Update  $\phi^{(k)}$  : Update  $\phi^{(k)}$  using Eq. (8) with  $p_{\mathrm{uncond}}$  and  $\{z_{+}^{(i)}\}_{i = 1}^{B}$  as the target.