SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mingtian Zhang*
m.zhang@cs.ucl.ac.uk
Brooks Paige
b.paige@ucl.ac.uk
Alex Hawkins-Hooker
a.hawkins-hooker@cs.ucl.ac.uk
David Barber
david.barber@ucl.ac.uk
Energy-Based Models (EBMs) offer a versatile framework for modeling complex data distributions. However, training and sampling from EBMs continue to pose significant challenges. The widely-used Denoising Score Matching (DSM) method [41] for scalable EBM training suffers from inconsistency issues, causing the energy model to learn a 'noisy' data distribution. In this work, we propose an efficient sampling framework, (pseudo)-Gibbs sampling with moment matching, which enables effective sampling from the underlying clean model when given a 'noisy' model that has been well-trained via DSM. We explore the benefits of our approach compared to related methods and demonstrate how to scale the method to high-dimensional datasets.
1 Energy-Based Models
Energy-Based Models (EBMs) have attracted a lot of attention in the generative model literature [25, 46, 7, 35]. EBMs are a type of non-normalized probabilistic model that determines the probability density function without a known normalizing constant. For continuous data  $x$ , the density function of an EBM is specified as  $q_{\theta}(x) = \exp (-f_{\theta}(x)) / Z(\theta)$  where the  $f_{\theta}(x)$  is a nonlinear function with parameter  $\theta$  and  $Z(\theta) = \int \exp (-f_{\theta}(x)) \, \mathrm{d}x$  is the normalization constant that is independent of  $x$ . The energy parameterization allows for greater flexibility in model parameterization and the ability to model a wider range of probability distributions. However, the lack of a known normalizing constant makes training these models challenging. We start by giving a brief introduction of how to estimate  $\theta$  in EBMs and refer the reader to [37] for a detailed overview of different training techniques for continuous EBMs.
Likelihood-based training: A classic method to learn  $\theta$  is to minimize the KL divergence between the data distribution  $p_d$  and the model density  $q_{\theta}$ , which is defined as
$$
\mathrm {K L} \left(p _ {d} \| q _ {\theta}\right) \dot {=} - \int p _ {d} (x) \log q _ {\theta} (x) \mathrm {d} x \dot {=} - \int p _ {d} (x) f _ {\theta} (x) \mathrm {d} x - \log Z (\theta), \tag {1}
$$
where we use  $\doteq$  to denote the equivalence up to a constant that is independent of  $\theta$ . The integration of  $p_d(x)$  can be approximated by Monte Carlo with the training dataset  $\mathcal{X}_{train} = \{x_1, \dots, x_N\} \sim p_d(x)$ ; in this case, it is equivalent to the maximum likelihood estimate (MLE) [5]. However, for EBMs, minimizing the KL divergence requires the estimation of  $Z(\theta)$ , which is intractable for nonlinear  $f_{\theta}(x)$  defined by a neural network. Various methods have been proposed to alleviate the
*This work was partially done during an internship in Huawei Noah's Ark Lab.
intractability by introducing techniques like Markov chain Monte Carlo (MCMC) [13, 26, 8, 11] or adversarial training [18, 47, 6].
Score-based training: Alternatively, [16] proposes to minimize the Fisher divergence to learn  $\theta$ , which is defined as
$$
\mathrm {F D} \left(p _ {d} \| q _ {\theta}\right) = \frac {1}{2} \int p _ {d} (x) \| s _ {p _ {d}} (x) - s _ {q _ {\theta}} (x) \| _ {2} ^ {2} \mathrm {d} x, \tag {2}
$$
where we use  $s_p(x)$  to denote the score function of distribution  $p$ :  $s_p(x) \equiv \nabla_x \log p(x)$ . Under certain regularity conditions, the Fisher divergence is equivalent to the score-matching (SM) objective [16],
$$
\mathrm {F D} \left(p _ {d} \| q _ {\theta}\right) \doteq \frac {1}{2} \int p _ {d} (x) \left(\left| \left| s _ {q _ {\theta}} (x) \right| \right| _ {2} ^ {2} + 2 \operatorname {T r} \left(\nabla_ {x} s _ {q _ {\theta}} (x)\right)\right) \mathrm {d} x, \tag {3}
$$
which does not require estimation of the intractable  $Z(\theta)$ . However, this objective needs to calculate the Hessian trace  $\nabla_{x} s_{q_{\theta}}(x) = -\nabla_{x}^{2} f_{\theta}(x)$  in every gradient step during training, which is computationally expensive and does not scale to high dimensional data or requires approximation [38]. In this paper, we will focus on another training method, denoising score matching [41], which overcomes the tractability and scalability issues mentioned above, and is introduced in the next section.
1.1 Denoising Score Matching
For the target data density  $p_d(x)$ , a noise distribution  $p(\tilde{x} | x) = \mathcal{N}(x, \sigma^2 I)$  is introduced to construct a noised data distribution  $\tilde{p}_d(\tilde{x}) = \int p_d(x)p(\tilde{x} | x) \, \mathrm{d}x$ . Denoising score matching (DSM) [41] minimizes the Fisher divergence between the noised data distribution  $\tilde{p}_d(\tilde{x})$  and an energy-based model  $\tilde{q}_{\theta}(\tilde{x}) = \exp(-f_{\theta}(\tilde{x})) / Z(\theta)$ , with
$$
\begin{array}{l} \mathrm {F D} (\tilde {p} _ {d} | | \tilde {q} _ {\theta}) = \frac {1}{2} \int \tilde {p} (\tilde {x}) | | s _ {\tilde {p} _ {d}} (\tilde {x}) - s _ {\tilde {q} _ {\theta}} (\tilde {x}) | | _ {2} ^ {2} \mathrm {d} \tilde {x} \\ \dot {=} \frac {1}{2} \iint p (\tilde {x} | x) p _ {d} (x) | | \nabla_ {\tilde {x}} \log p (\tilde {x} | x) - s _ {\tilde {q} _ {\theta}} (\tilde {x}) | | _ {2} ^ {2} \mathrm {d} \tilde {x} \mathrm {d} x \\ \stackrel {=} {\frac {1}{2}} \iint p (\tilde {x} | x) p _ {d} (x) \left\| \frac {\tilde {x} - x}{\sigma^ {2}} + s _ {\tilde {q} _ {\theta}} (\tilde {x}) \right\| _ {2} ^ {2} \mathrm {d} \tilde {x} \mathrm {d} x, \tag {4} \\ \end{array}
$$
where the last equation is due to  $\nabla_{\tilde{x}}\log p(\tilde{x} |x)$  being tractable for the Gaussian distribution  $p(\tilde{x} |x)$ .
Compared to the KL or SM objectives, the DSM objective is scalable and well-defined when the data distribution is singular[49] and can alleviate the blindness problem of score matching [35, 45, 51]. On the other hand, there is a notable disadvantage associated with the DSM objective: for a fixed  $\sigma >0$ , the DSM objective is not a consistent objective for learning the underlying data distribution  $p_d$  since  $\mathrm{FD}(\tilde{p}_d||\tilde{q}_\theta) = 0\Rightarrow \tilde{p}_d = \tilde{q}_\theta \neq p_d$ . A common solution is to anneal  $\sigma \rightarrow 0$  during training. However, Equation 4 is not defined when  $\sigma = 0$  since the division in Equation 4 will make  $(\tilde{x} -x) / \sigma^2$  unbounded, which results in an inconsistent objective. Annealing  $\sigma$  increases the variance of the training gradients [37, 42], which makes the optimization challenging in practice.
To overcome the challenges, we propose an alternative data generation scheme: we use DSM with a fixed  $\sigma > 0$  to train a 'noisy' energy model and then construct a sampler which targets the underlying 'clean' model. Specifically, our contributions are summarized as follows:
- We demonstrate that for an EBM that learns a noisy data distribution, there exists a unique underlying clean model which recovers the true data distribution.
- We introduce a pseudo-Gibbs sampling scheme incorporating an analytical moment-matching approximation of the denoising distribution. This allows us to sample from the underlying clean model without requiring additional training.
- We illustrate how to scale our method for high-dimensional data and demonstrate the generation of high-quality images using only a single level of fixed noise. Furthermore, we showcase the application of our proposed method in multi-level noise scenarios, closely resembling a diffusion model.
2The singular distribution is not absolutely continuous  $(a.c.)$  with respect to the Lebesgue measure, thus doesn't allow a density function [40, p.172]. A typical example is a data distribution supported on a lower-dimensional manifold. In this case, the KL divergence is ill-defined and cannot be used to train the models. When using DSM, the distribution after Gaussian convolution would always be  $a.c.$ , thus can be a valid training objective, see Zhang et al. [49] or Arjovsky et al. [1] for a detailed introduction.
2 Clean Model Identification
For a fixed  $\sigma > 0$ , DSM can only learn a 'noisy' data distribution  $\tilde{q}_{\theta}(\tilde{x})$  even in the ideal case where the Fisher divergence is exactly minimized, since  $\mathrm{FD}(\tilde{p}_d||\tilde{q}_{\theta^*}) = 0 \to \tilde{p}_d = \tilde{q}_{\theta} \neq p_d$ . In this case, the following theorem shows that there exists a 'clean' model that is implicitly defined that learns the true data distribution.
Theorem 2.1 (Existence of the underlying clean model for optimal  $\tilde{q}_{\theta}(\tilde{x})$ ). When the Fisher divergence goes to 0,  $\mathrm{FD}(\tilde{p}_d||\tilde{q}_\theta) = 0 \to \tilde{p}_d = \tilde{q}_\theta$ , there exists an unique underlying clean model  $q(x)$  such that  $\tilde{q}_{\theta}(\tilde{x}) = \int q(x)p(\tilde{x} |x)\mathrm{d}x$  and  $q(x) = p_d(x)$ .
See Appendix A.1 for proof. This theorem shows that despite training an EBM on noisy data, there is an implicit model within it that can recover the true data distribution. Therefore, instead of annealing the noise  $\sigma \rightarrow 0$  to recover the true data distribution, we will demonstrate how to directly sample from the implicitly-defined clean model given the noisy energy-based model in the next section.
We want to highlight that the 'perfect fit' assumption, i.e. achieving  $\mathrm{FD}(\tilde{p}_d||\tilde{q}_\theta) = 0$ , may not hold for a complex data distribution  $p_d$  or underpowered EBMs. Therefore, we provide general sufficient conditions for the existence of the clean model for an imperfect EBM in Appendix A.2.
2.1 Gibbs Sampling with Gaussian Moment Matching
Given a well-trained noisy energy-based model  $\tilde{q}_{\theta}(\tilde{x}) = \tilde{p}_d(\tilde{x})$ , the clean model has the form
$$
q (x) = \int p (x | \tilde {x}) \tilde {q} _ {\theta} (\tilde {x}) d \tilde {x}, \tag {5}
$$
where the denoising distribution can be written as  $p(x|\tilde{x}) \propto p(\tilde{x} |x)q(x)$ . We notice that, since the noise distribution  $p(\tilde{x} |x) = \mathcal{N}(x,\sigma^2 I)$  is known, a Gibbs sampling scheme can be constructed to sample from the underlying clean model if we know the denoising distribution  $p(x|\tilde{x})$ , with
$$
\tilde {x} _ {k - 1} \sim p (\tilde {x} | x = x _ {k - 1}), \quad x _ {k} \sim p (x | \tilde {x} = \tilde {x} _ {k - 1}), \tag {6}
$$
where the initial sample  $x_0 \sim p_0(x)$  can be drawn from a standard Gaussian  $p_0(x) = \mathcal{N}(0, I)$ . However, as the denoising distribution  $p(x|\tilde{x})$  is usually intractable for complex  $p_d$ , we propose an analytical Gaussian moment matching approximation of  $p(x|\tilde{x})$ .
Denote the mean and covariance of  $p(x|\tilde{x})$  as
$$
\mu (\tilde {x}) = \langle x \rangle_ {p (x | \tilde {x})}, \quad \Sigma (\tilde {x}) = \langle x ^ {2} \rangle_ {p (x | \tilde {x})} - \langle x \rangle_ {p (x | \tilde {x})} ^ {2}. \tag {7}
$$
The classic Gaussian moment matching method [24] specifies a Gaussian approximation  $p(x|\tilde{x}) \approx \mathcal{N}(\mu(\tilde{x}), \Sigma(\tilde{x}))$ , which matches the first and second moment of  $p(x|\tilde{x})$ . When  $\tilde{q}_{\theta} = \tilde{p}_{d}$ , the first mean of the denoised distribution has a well-known analytical form [39, 2, 9, 29]
$$
\mu (\tilde {x}) = \tilde {x} + \sigma^ {2} s _ {\tilde {q} _ {\theta}} (\tilde {x}); \tag {8}
$$
we include the derivation in Appendix A.3. Using this identity, we can rewrite Equation 4 as
$$
\mathrm {F D} \left(\tilde {p} _ {d} \| \tilde {q} _ {\theta}\right) \doteq \frac {1}{2 \sigma^ {4}} \iint p (\tilde {x} | x) p _ {d} (x) \| x - \mu (\tilde {x}) \| _ {2} ^ {2} \mathrm {d} \tilde {x} \mathrm {d} x, \tag {9}
$$
where we can see that the Fisher divergence only depends on  $\mu (\tilde{x})$ . Since  $\mathrm{FD}(\tilde{p}_d||\tilde{q}_\theta) = 0\Leftrightarrow q = p_d$  (Theorem 2.1), the function  $\mu (\tilde{x})$  fully characterizes the distribution  $q$ . Therefore,  $\mu (\tilde{x})$  and  $p(\tilde{x}|x) = \mathcal{N}(x,\sigma^2 I)$  can provide sufficient information to determine  $p(x|\tilde{x})\propto q(x)p(\tilde{x} |x)$ . As a consequence, the following theorem shows that the covariance function can also be analytically derived.
Theorem 2.2 (Analytical Covariance Identity). Given a clean model  $q(x)$  such that  $\int q(x)p(\tilde{x} |x)\mathrm{d}x = \tilde{q}_{\theta}(\tilde{x}) = \tilde{p}_d(\tilde{x})$  with  $p(\tilde{x} |x) = \mathcal{N}(x,\sigma^2 I)$ , the  $\mu (\tilde{x})$  and  $\Sigma (\tilde{x})$  of the  $p(x|\tilde{x})\propto q(x)p(\tilde{x} |x)$  has the following relations
$$
\Sigma (\tilde {x}) = \sigma^ {2} \nabla_ {\tilde {x}} \mu (\tilde {x}) = \sigma^ {4} \nabla_ {\tilde {x}} ^ {2} \log \tilde {q} _ {\theta} (\tilde {x}) + \sigma^ {2} I. \tag {10}
$$
See Appendix A.3 for proof. This analytical covariance identity can be seen as a high-dimensional generalization of the 2nd-order Tweedie's Formula [9, 29]. Therefore, the analytical full-covariance moment matching approximation can be written as
$$
p (x | \tilde {x}) \approx \mathcal {N} (\tilde {x} + \sigma^ {2} \nabla_ {\tilde {x}} \log \tilde {q} _ {\theta} (\tilde {x}), \sigma^ {4} \nabla_ {\tilde {x}} ^ {2} \log \tilde {q} _ {\theta} (\tilde {x}) + \sigma^ {2} I). \tag {11}
$$
We want to highlight that since the Gaussian moment matching is only an approximation of  $p(x|\tilde{x})$ , the sampling scheme in Equation 6 is a 'pseudo' Gibbs sampler unless the true  $p(x|\tilde{x})$  is also a Gaussian distribution<sup>3</sup>, which is not true for general non-Gaussian  $p_d$ . However, since  $\mu(\tilde{x})$  and  $p(\tilde{x}|x) = \mathcal{N}(x,\sigma^2 I)$  are already sufficient to specify  $p(x|\tilde{x})$ , it should be possible to derive expressions for higher-order moments which themselves involve only  $\mu(\tilde{x})$  and  $\sigma$ ; we leave this to future work. To our knowledge, the  $\tilde{x}$ -conditioned full covariance Gaussian moment matching approximation to  $p(x|\tilde{x})$  has not been derived previously. In the next section, we briefly discuss the connections between our method and other related approaches.
2.2 Connection to Covariance Learning Approaches
Bengio et al. [4] proposes to approximate the true posterior  $p(x|\tilde{x}) \propto p_d(x)p(\tilde{x} |x)$  with a variational distribution  $q_{\theta}(x|\tilde{x})$ . The parameter  $\theta$  is then learned by minimizing the joint KL divergence
$$
\operatorname {K L} (p (\tilde {x} | x) p _ {d} (x) \| q _ {\theta} (x | \tilde {x})) \tilde {p} _ {d} (\tilde {x})) \doteq - \iint p _ {d} (x) p (\tilde {x} | x) \log q _ {\theta} (x | \tilde {x}) \mathrm {d} \tilde {x} \mathrm {d} x, \tag {12}
$$
where  $\tilde{p}_d(\tilde{x}) = \int p_d(x)p(\tilde{x} |x)\mathrm{d}x$ . The joint KL divergence in Equation 12 encourages  $q_{\theta}(x|\tilde{x})$  to match the moments of the true posterior  $p(x|\tilde{x})$ , and defines an upper bound of the marginal KL [48]
$$
\operatorname {K L} \left(p (\tilde {x} | x) p _ {d} (x) \| q _ {\theta} (x | \tilde {x})) \tilde {p} _ {d} (\tilde {x})\right) \geq \operatorname {K L} \left(p _ {d} (x) \| q _ {\theta} (x)\right), \tag {13}
$$
where the model is implicitly defined as the marginal of the joint  $q_{\theta}(x) = \int q_{\theta}(x|\tilde{x})\tilde{p}_d(\tilde{x})$ . When  $q_{\theta}(x|\tilde{x})$  is a consistent estimator of  $p(x|\tilde{x})$ , this asymptotic distribution of the Gibbs sampling will converge to the true data distribution  $p_d(x)$  [4]. For continuous data, the variational distribution is chosen as a Gaussian distribution  $q_{\theta}(x|\tilde{x}) = \mathcal{N}(\mu_{\theta}(\tilde{x}), \Sigma_{\theta}(\tilde{x}))$ , where the mean  $\mu_{\theta}(\cdot)$  and the covariance  $\Sigma_{\theta}(\cdot)$  are parameterized by neural networks. We note that the only difference between the KL and DSM objective (Equation 9) is that the KL objective additionally learns the covariance. We thus show that the optimal covariance under KL minimization is the proposed analytical covariance.
Theorem 2.3 (Optimal Gaussian Approximation). Let  $p(\tilde{x} |x) = \mathcal{N}(0,\sigma^2 I)$  and assume Gaussian distribution  $q_{\theta}(x|\tilde{x}) = \mathcal{N}(\mu_q(\tilde{x}),\Sigma_q(\tilde{x}))$ , then the optimal  $q^{*}$  such that
$$
q ^ {*} = \arg \min  _ {q} \mathrm {K L} (p (\tilde {x} | x) p _ {d} (x) \| q (x | \tilde {x}) \tilde {p} _ {d} (\tilde {x})) \tag {14}
$$
has the mean and covariance with the form
$$
\mu_ {q} ^ {*} (\tilde {x}) = \left\langle x \right\rangle_ {p (x | \tilde {x})}, \quad \Sigma_ {q} ^ {*} (\tilde {x}) = \sigma^ {2} \nabla_ {\tilde {x}} \mu_ {q} ^ {*} (\tilde {x}), \tag {15}
$$
see Appendix A.4 for proof. Therefore, when the optimal mean function is learned  $\mu_{\theta}(\tilde{x}) = \mu_q^* (\tilde{x})$  the optimal  $\Sigma (\tilde{x})$  can be analytically derived, making the learning of  $\Sigma_{\theta}(\tilde{x})$  redundant. In addition to the training inefficiency caused by more parameters, the amortized covariance network may suffer from poor generalization [50]. Moreover, the KL objective is also not well-defined for learning data distributions which lie on a low-dimensional manifold, e.g. MNIST, see Section 4 for a detailed discussion. In this case, the learned  $\Sigma_{\theta}(\tilde{x})$  may be a degenerate matrix, making the Gaussian density function  $q(x|\tilde{x})$  ill-defined [49] which impedes the training, see Figure 5 for an example.
Paper[23] proposes a higher-order score-matching loss to simultaneously learn both the first order score  $\nabla_{\tilde{x}}\log \tilde{q}_{\theta}(\tilde{x})$  and the second order score  $\nabla_{\tilde{x}}^2\log \tilde{q}_{\theta}(\tilde{x})$ . However, our findings indicate that the mean function  $\mu (\tilde{x})$  (or the first order score  $\nabla_{\tilde{x}}\log p(\tilde{x}))$  already contains all the moment information of the underlying true distribution  $p_d$ , and the optimal moment can be derived using the mean function. Therefore, learning the second-order score is redundant and may lead to sub-optimal inference.
2.3 Connection to Analytic DDPM
The recent paper Bao et al. [2] considers a constrained variational family  $q_{\theta}(x|\tilde{x}) = \mathcal{N}(\mu_{\theta}(\tilde{x}),\sigma_q^2 I)$  in the context of diffusion model and derive the optimal  $\sigma_q^*$  as
$$
\sigma_ {q} ^ {* 2} = \arg \min  _ {\sigma_ {q}} \mathrm {K L} (p (\tilde {x} | x) p _ {d} (x) \| q _ {\theta} (x | \tilde {x})) \tilde {p} _ {d} (\tilde {x})) = \frac {1}{d} \left\langle \operatorname {T r} \left(\operatorname {C o v} _ {q (x | \tilde {x})} [ x ]\right) \right\rangle_ {\tilde {p} _ {d} (\tilde {x})}, \tag {16}
$$
For example, when  $p_d(x) = \mathcal{N}(\mu_d, \sigma_d)$ , the true posterior  $p(x|\tilde{x}) \propto p_d(x)p(\tilde{x}|x)$  will be a Gaussian with mean  $\mu(\tilde{x}) = (\sigma^2\mu_d + \sigma_d^2\tilde{x}) / (\sigma^2 + \sigma_d^2)$  and variance  $\Sigma(\tilde{x}) = \sigma_d^2\sigma^2 / (\sigma^2 + \sigma_d^2)$ , we can verify that  $\Sigma(\tilde{x}) = \sigma^2\nabla_{\tilde{x}}\mu(\tilde{x})$ .
which can also be rewritten using the score function
$$
\sigma_ {q} ^ {* 2} = \sigma^ {2} - \sigma^ {4} / d \left\langle \| s _ {q _ {\theta}} (\tilde {x}) \| _ {2} ^ {2} \right\rangle_ {\tilde {p} _ {d} (\tilde {x})}. \tag {17}
$$
In Appendix B, we provide a detailed derivation to show how this approximation can be linked to our method using the Fisher information identity [10]. This approximation has two potential limitations: first, compared to full covariance moment matching, the assumed isotropic covariance structure may be insufficiently flexible to capture the true posterior; second, the covariance is independent of  $\tilde{x}$ .
The second assumption only holds when  $\mu (\tilde{x})$  is a linear function of  $\tilde{x}^4$  (e.g. when  $p_d(x)$  is Gaussian) and does not hold for other non-Gaussian  $p_d(x)$ . Therefore, our  $\tilde{x}$ -dependent full-covariance approximation offers a more versatile approximation family, which ultimately results in a more precise estimation. However, in certain applications such as accelerating the sampling procedure of a diffusion model [2], it is advantageous to use a  $\tilde{x}$ -independent isotropic covariance due to its inexpensive estimation. On the other hand, our  $\tilde{x}$ -dependent covariance necessitates the computation of the Hessian for each  $\tilde{x}$ , making it inefficient for high-dimensional data. In Section 3, we will explore approaches to mitigate this limitation.
2.4 Posterior Approximation Comparison
We now consider a toy example to compare the three denoising posterior approximations discussed above. Let  $p_d(x)$  be a Mixture of Gaussians (MoG)  $p_d(x) = \frac{1}{4} \sum_{k=1}^{k=4} g_k(x)$  whose components  $g_{[1:4]}$  are 2D Gaussians with means  $[-1, -1]$ ,  $[-1, 1]$ ,  $[1, 1]$ ,  $[1, -1]$  and isotropic covariance  $\sigma_g^2 I$  with  $\sigma_g = 0.2$ . The noise distribution is  $p(\tilde{x} | x) = \mathcal{N}(x, \sigma^2 I)$  with  $\sigma = 0.2$ , so  $\tilde{p}(\tilde{x}) = \int p(x) p(\tilde{x} | x) \, \mathrm{d}x$  is an MoG with the same component means and diagonal covariance  $(\sigma_g^2 + \sigma^2) I$ ; see Figure 1a for a visualization. In this case the true posterior  $p(x | \tilde{x})$  does not allow a tractable form. Fortunately, given a noisy sample  $\tilde{x}'$  and an evaluation point  $x'$ , we can evaluate the true density  $p(x | \tilde{x}')$  using Bayes rule:  $p(x'|\tilde{x} = \tilde{x}') = p(\tilde{x}' | x = x') p_d(x') / \tilde{p}_d(\tilde{x}')$ . Figure 1 shows the true posteriors given four different  $\tilde{x}'$  where we use grid data in  $x$ -space to visualize the density.
To train the model, we sample 10,000 data points from  $p_d$  as our training data. For the KL-trained
Gibbs sampler described in Section 2.2, we use a network with 3 hidden layers with 400 hidden units, Swish activation [28] and output size 4 to generate both mean and log standard deviation of the Gaussian approximation. For the moment-matching Gibbs sampler (including both full and isotropic covariance), we use the same network architecture but with output size 1 to get the scalar energy and DSM as the training objective. Both networks are trained with batch size 100 and Adam [19] optimizer with learning rate  $1 \times 10^{-4}$  for 100 epochs. For the  $\tilde{x}$ -independent isotropic covariance, we use the Monte Carlo approximation to estimate the variance [2] with 10000 samples from  $\tilde{p}_d(\tilde{x})$ .
(a) Gaussian mixtures visualizations
(b) Density of  $p(x|\tilde{x}^{\prime})$  (blue) and four  $\tilde{x}^{\prime}$  points (green)
(c) Analytical full covariance moment matching
(d) Analytical isotropic covariance moment matching
(e) Learned diagonal covariance
Figure 1: Figure (a) shows the clean data distribution  $p_d(x)$  and the corresponding noisy distribution  $\tilde{p}_d(\tilde{x})$ . Figure (b) shows 4 conditioned samples in the noisy space. Figures (c, d, e) visualize the true posterior  $p(x|\tilde{x})$  (green) and three posterior approximations (orange). We find that only the proposed  $\tilde{x}$ -dependent analytical full-covariance moment matching can capture the variance of the true posterior, whereas the other two methods underestimate the variance.
<sup>4</sup>Since when  $\mu (\tilde{x})$  is a linear function of  $\tilde{x}$ , using Theorem 2.2, we have  $\Sigma (\tilde{x}) = \sigma^2\nabla_{\tilde{x}}\mu (\tilde{x})$  will not depend on  $\tilde{x}$ , see also footnote 1 for an example.