SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We then use the estimated posterior to conduct (pseudo) Gibbs sampling to generate samples<sup>5</sup>. Specifically, we initialize the first sample  $x_0 \sim \mathcal{N}(0, 0.1)$  and run one Markov Chain with 10,000 time steps to generate 10,000 samples. In addition to the mixture of Gaussian datasets, we also train and generate samples from the 2D Swiss roll and two-ring datasets. For numerical evaluation, we calculate the Maximum Mean Discrepancy (MMD) [12] between 10k samples generated by a single-chain Gibbs sampler and 10k samples from the training dataset respectively. The kernel insides MMD is a sum over 5 Gaussian kernels with bandwidth ranging over  $[2^{-2}, 2^{-1}, 2^0, 2^1, 2^2]$ . The MMD results (including both mean and std) are calculated using 5 random seeds. We find that Gibbs sampling with the
Figure 1 visualizes the approximations to the denoising posterior  $p(x|\tilde{x})$  estimated by each of the three methods described in the previous sections. We surprisingly find that although the KL objective in Equation 12 encourages  $q_{\theta}(x|\tilde{x})$  to match the moments of
$p(\tilde{x}|\tilde{x})$ , the learned covariance in Figure 1e still underestimates the variance of the posterior. This shows the redundancy of covariance learning can degrade the variational approximation performance. Additionally, the  $\tilde{x}$ -independent covariance fails to account for the relative positions of  $\tilde{x}'$  and lacks the ability to predict the posterior's elliptical shape due to its isotropic nature. In contrast, our  $\tilde{x}$ -dependent full covariance approximation overcomes these limitations, enabling more accurate predictions that capture the intricate geometry of the posterior distribution.
(c) Analytic iso.
(d) Analytic full
(a) True data
(b) Learn. diag.
(g) Analytic iso.
(e) True data
(i) True data
Figure 2: Samples from a single chain Gibbs sampling
(f) Learn. diag.
(j) Learn. diag.
(k) Analytic iso.
(h) Analytic full
(1) Analytic full
proposed analytical full covariance achieves the best results; numerical results are in Table 1, with a visual comparison in Figure 2.
3 Scalable Implementations for Image Data
Scalable Diagonal Hessian Approximation As we discussed in Section 2.3, the proposed full covariance Gaussian approximation in Equation 10 requires calculating an  $D \times D$  Hessian  $\nabla_{\tilde{x}}^{2} \log \tilde{q}(\tilde{x})$  for each  $\tilde{x}$  with size  $D$ , which brings both memory and computation difficulties for high-dimensional data. A naive diagonal Hessian method (only using the diagonal entries in the Hessian) will address the memory bottleneck but still needs  $D$  times backward passes for the exact computation of the diagonal term [22]. In this paper, we use the following diagonal Hessian approximation [3],
$$
\operatorname {D i a g} (H) \approx 1 / S \sum_ {s = 1} ^ {S} v _ {s} \odot H v _ {s}, \tag {18}
$$
where  $v_{s} \sim p(v)$  is a Rademacher random variable with entries  $\pm 1$  and  $\odot$  denotes the element-wise product<sup>6</sup>. This estimator will converge to the exact Hessian diagonals when  $S \to \infty$  [3]. The computation for each  $v_{s}$  can be computed by two forward-backward passes. It is worth emphasizing that our  $\tilde{x}$ -dependent diagonal moment matching approach provides a comparable level of flexibility to the variational method proposed in [4] while eliminating the need for additional training of the
5The code of the experiments can be found in https://github.com/zmtomorrow/MMDGS_NeurIPS. 6This estimation should be distinguished from the Hutchinson's Trace estimation [15]:  $\operatorname {Tr}(H)\approx$ $\frac{1}{S}\sum_{s = 1}^{S}v_s^T Hv_s$  where a dot-product is used between  $v_{s}$  and  $Hv_{s}$
(a) Diagonal  $\tilde{x}$ -dependent covariance moment matching (Ours)
(b) Isotropic  $\tilde{x}$ -independent covariance moment matching [2]
(c) Diagonal covariance learned by KL minimization [4]
Figure 3: Figures (a,b,c) show the MNIST experiment comparisons, where we compare samples generated by pseudo-Gibbs sampling with three different  $q(x|\tilde{x})$ . We plot samples from 25 independent Markov Chains with  $t \in \{0,1,5,10,20\}$  time steps. We can find the samples generated by the proposed analytical covariance moment matching with diagonal approximation achieved the best sample quality.
diagonal covariance. Furthermore, our method remains more flexible than the isotropic  $\tilde{x}$ -independent moment matching method proposed by [2].
Energy or Score Parameterization For the full-covariance moment matching in Equation 10, we require  $\nabla_{\tilde{x}}^{2}\log \tilde{q}_{\theta}(\tilde{x})$  to be symmetric to obtain a valid Gaussian approximation. However, if we learn the score function  $\nabla_{\tilde{x}}\log p(\tilde{x}) = s_{\theta}(\tilde{x})$  using a network  $s_{\theta}(\cdot):\mathbb{R}^D\to \mathbb{R}^D$ , its Jacobian is not guaranteed to be symmetric. In this case, we follow [34] and directly parameterize the density function  $\tilde{q}_{\theta}(\cdot)$  with a neural network  $f_{\theta}(\cdot):\mathbb{R}^{D}\rightarrow \mathbb{R}$  and let the score function  $\nabla_{\tilde{x}}\log \tilde{q}_{\theta}(\tilde{x}) = -\nabla_{x}f_{\theta}(\tilde{x})$ . This can be obtained by AutoDiff packages like PyTorch [27], and this parameterization guarantees  $\nabla_{\tilde{x}}^{2}f_{\theta}(\tilde{x})$  to be symmetric. We also notice that when using the diagonal Hessian approximation (Equation 18), we only need entries in  $\mathrm{Diag}(H)$  to be positive in order to obtain a valid Gaussian approximation. In this case, the score parameterization remains applicable and offers more efficient training compared to the energy parameterization. Therefore, the combination of full/diagonal covariance and energy/score parameterization provides a tradeoff between flexibility and inference speed, allowing for a flexible approach while maintaining computational efficiency during training.
4 Image Generation with a Single Noise Level
We then apply the proposed method to model the grey-scale MNIST [21] dataset. We use the standard U-Net architecture [35, 31] with a single fixed noise level  $\sigma = 0.5$ ; the effect of varying  $\sigma$  is explored in Appendix C.1. For the KL training objective, the output channel size is 2 to generate both mean and log-standard at the same time. For DSM training, we take the sum of the U-Net output to obtain the scalar energy evaluation which also relates to the product-of-experts model described in [34]. We train both networks for 300 epochs with learning rate  $1\times 10^{-4}$  and batch-size 100.
As discussed in Section 1.1, the KL divergence is not well-defined for manifold data distributions. This limitation becomes evident when working with MNIST, where the presence of constant black pixels in the boundary areas leads to a rapid decrease towards 0 in the variance of  $q_{\theta}(x|\tilde{x})$  during training. Consequently, the likelihood value  $\log q_{\theta}(x|\tilde{x})$  tends to approach infinity, resulting in unstable training. In contrast, the DSM objective is well-defined for manifold data, providing
a stable training process even in the presence of such boundary effects. Figure 5 provides a visual comparison of the two training procedures, demonstrating the improved stability and effectiveness of the DSM objective in handling manifold data distributions.
For the sample generation process, calculating the full-covariance Gaussian posterior becomes challenging. We therefore apply the scalable diagonal Hessian approximation described in Section 3 to approximate the diagonal Gaussian covariance of  $p(x|\tilde{x})$ . We find that the estimated diagonal Hessian occasionally contains small negative values due to approximation error; we, therefore, use the  $\max(\cdot, \epsilon)$  function with  $\epsilon > 0$  to ensure the pos
itivity of the diagonal covariance. The  $\tilde{x}$ -independent isotropic covariance and the proposed  $\tilde{x}$ -dependent diagonal covariance share the same mean function.
(a) KL loss
(b) DSM loss
Figure 5: Training loss comparison of two objectives. We plot the training loss every iteration during a total 300 epochs.
We first visualize the covariance estimated by three different methods in Figure 4. We use 100 Rademacher samples in estimating the diagonal Hessian (Equation 18) and 50,000 samples in estimating the isotropic variance (Equation 17). We find that both  $\tilde{x}$ -dependent diagonal covariance approximations can capture the posterior structure whereas the isotropic
$\tilde{x}$ -independent covariance is just Gaussian noise since the variance is shared between different digit and pixel locations. In Figure 3, we plot the sample comparison for three methods.
Since the isotropic covariance has the same variance in each dimension, the generated samples in Figure 3b contain white noise in the black background, whereas the proposed full-covariance sampler
(a) Data  $x$
Figure 6: Visualizations of the diagonal covariance  $q(x|\tilde{x} = x + \sigma \epsilon)$  with different number of Rademacher samples.
(b) 1 sample
(c) 10 samples
(d) 100 samples
Figure 7: Samples from three Markov chains. We plot the samples every 10 Gibbs steps.
Algorithm 1 Sampling with Langevin Dynamics
Require:  $\{\sigma_t\}_{t = 0}^T,\delta ,K$  Initialize  $x_{T}^{0}$    
for  $t\gets T$  to 1 do  $\alpha_{t}\leftarrow \delta \sigma_{t}^{2} / \sigma_{0}^{2}$    
for  $k\gets 1$  to  $K$  do Draw  $z^{k}\sim \mathcal{N}(0,I)$ $x_{t}^{k}\gets x_{t}^{k - 1} + \alpha_{t}s_{\theta}(x_{t}^{k - 1},\sigma_{t}) + \sqrt{2\alpha_{t}} z^{k}$  end for  $x_{t - 1}^{0}\gets x_{t}^{K}$    
end for   
return  $x_0^0 +\sigma_0^2 s_\theta (x_0^0,\sigma_0)$
Algorithm 2 Sampling with the proposed pseudo Gibbs Sampling
Require:  $\{\sigma_t\}_{t = 0}^T,K$  Initialize  $x_{T}^{0}$    
for  $t\gets T$  to 1 do for  $k\gets 1$  to  $K$  do Draw  $x_{t + 1}^k\sim p(x_{t + 1}|x_t = x_t^k)$  Draw  $x_{t}^{k}\sim q_{\theta}(x_{t}|x_{t + 1} = x_{t + 1}^{k})$  end for  $x_{t - 1}^{0}\gets x_{t}^{K}$    
end for   
return  $x_0^0 +\sigma_0^2 s_\theta (x_0^0,\sigma_0)$
can generate a clean black background in Figure 3a. On the other hand, the samples generated by the KL-trained Gibbs sampler (Figure 3c) have worse sample quality due to the unstable training.
We then apply the same method to model the more complicated CIFAR 10 [20] dataset. We use the same U-Net structure as used in [36] and directly parameterize the score function rather than the energy function to speed up the training. The noise level is fixed at 0.3. We train the model using Adam optimizer with learning rate  $1 \times 10^{-4}$  and batch size 100 for 1000 epochs. We visualize the denoising posterior diagonal covariance in Figure 6 when using different numbers of Rademacher samples (Equation 18). We observe that better covariance estimation can be obtained by increasing the number of samples. To balance efficiency and accuracy, we use a sample number of 10 in the subsequent Gibbs sampling stage. Figure 7 shows three independent Markov chains with the samples plotted every 10 Gibbs steps, which demonstrates that sharp images can be generated with even one fixed level of noise.
Limitation: In the CIFAR experiment, we observe a mode collapse phenomenon when running multiple independent Markov chains for a longer time. This phenomenon is likely due to the small noise level  $\sigma = 0.3$ , which prevents the sampler from exploring the full space, as commonly found with MCMC methods [30]. This effect is visually represented in Figure 8, where we assess the Fréchet Inception Distance (FID) values for 50,000 images sampled with varying numbers of Gibbs steps. Notably, the FID increases beyond 40 Gibbs steps, and visual evidence of mode collapse is observed (Figure 12). In the ensuing section, we will demonstrate the application of our method to settings with multiple noise levels, an approach that may help mitigate the issue of mode collapse.
Figure 8: FID evaluation with Increased Gibbs Steps. We can find the FID increases after 40 Gibbs steps.
5 Image Generation Using Multiple Noise Levels
The success of diffusion models and lessons from prior work on score-based generative models point to the importance of using multiple noise levels [14, 35] when modelling data with complex multi-modal distributions. Intuitively, by learning to denoise data at a range of noise levels, a single network can learn both the fine and global structure of the distribution, which in turn allows for more effective sampling algorithms capable of efficiently exploring diverse modes [35]. We therefore propose to adapt the denoising Gibbs sampling procedure to sample from distributions corrupted with multiple noise levels. For this purpose, we use a noise-conditioned score network trained by Song and Ermon [36], who generated high-quality samples using a procedure inspired by annealed Langevin dynamics [43]. This procedure involves generating samples from a sequence of distributions  $p_{T}(x_{T}), \ldots, p_{0}(x_{0})$ , corrupted by progressively decreasing levels of Gaussian noise (parameterized via standard deviations  $\sigma_{t}$ , with  $\sigma_{T} > \dots > \sigma_{0}$ ). At a given step  $t$  in the sequence, Langevin dynamics is used to sample from the corresponding noised distribution  $p_{t}(x_{t})$ , using the
score network  $s_{\theta}(x_t, \sigma_t)$  to approximate the gradient of the noised distribution. The outputs of this Langevin dynamics run are then used to initialize the same procedure at the next noise level, leading the sampling procedure to converge gradually towards the data distribution as the noise level tends to zero (i.e.  $p_0(x_0) \approx p_d(x), \sigma_0 \approx 0$ ). The algorithm of the annealed Langevin dynamics with multi-level noise used in [35, 36] is summarized in Algorithm 1.
We show that the proposed Gibbs sampling scheme can be directly applied to a pre-trained score-based generative model as a drop-in replacement for Langevin dynamics MCMC in the generation stage. At each noise level, we use samples  $x_{t + 1}$  from the previous noise level to initialise a Gibbs sampling chain targeting the marginal distribution at the current noise level  $p_t(x_t)$ , which now plays the role of the 'clean' distribution in Equation 5. Therefore, the noisy distribution at time step  $t$  is a Gaussian  $p(x_{t + 1}|x_t) = \mathcal{N}(x_t,\sigma_{t + 1}^2 -\sigma_t^2)$ . The optimal denoising distribution  $q(x_{t}|x_{t + 1})$  is thus a function of the level of noise at step  $t$  relative to the level of noise at the previous step  $t + 1$ . For generation efficiency, we employ 3 Gibbs steps at each noise level, using 3 Rademacher samples to approximate the diagonal Hessian. The sampling procedure is summarized in Algorithm 2. In Figure 9 and 10, we visualize the samples from models that are trained on CIFAR10 and CelebA separately. Further experimental details can be found in Appendix C.2.
For direct comparison with the results of [36] on CIFAR10, we retain the same schedule of noise levels used to generate samples with Langevin dynamics. We generate 50000 samples using this approach and report FID and Inception scores in Table 2. Our multi-level Gibbs sampling scheme produces samples of equivalent quality to the multi-level Langevin
dynamics of [36], confirming its applicability to complex natural image data. The FID is also notably superior to that of the single-noise level Gibbs sampling, and the samples exhibit significant visual diversity (Figure 9). This underlines the importance of employing multi-level noise in our approach. Recent advances in sampling strategies for score-based models leveraging the framework of stochastic differential equations [39] have led to significant further improvements in generation quality as shown in Table 2; we leave the exploration of possible applications of our method to this framework to future work.
Table 2: CIFAR10 Inception and FID Scores
Model	Inception	FID
NCSNv2 (Langevin [36])	8.40 ± 0.07	10.87
NCSNv2 (Gibbs, Ours)	8.28 ± 0.07	10.75
DDPM [14]	9.46 ± 0.11	3.17
NCSN++ [39]	9.89	2.20
Figure 9: CIFAR 10 Samples
Figure 10: CelebA Samples
This paper focuses on addressing the inconsistency problem in training energy-based models (EBMs) using denoising score matching. Specifically, we identify the presence of an underlying clean model within a 'noisy' EBM and propose an efficient sampling scheme for the clean model. We demonstrate how this method can be effectively applied to high-dimensional data and showcase image generation results in both single and multi-level noise settings. More broadly, we hope our more accurate denoising posterior opens new avenues for future work on score-based methods in machine learning.
[1] M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
[2] F. Bao, C. Li, J. Zhu, and B. Zhang. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint arXiv:2201.06503, 2022.
[3] C. Bekas, E. Kokiopoulou, and Y. Saad. An estimator for the diagonal of a matrix. Applied numerical mathematics, 57(11-12):1214-1229, 2007.
[4] Y. Bengio, L. Yao, G. Alain, and P. Vincent. Generalized denoising auto-encoders as generative models. Advances in neural information processing systems, 26, 2013.
[5] C. M. Bishop and N. M. Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.
[6] A. J. Bose, H. Ling, and Y. Cao. Adversarial contrastive estimation. arXiv preprint arXiv:1805.03642, 2018.
[7] Y. Du and I. Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
[8] Y. Du, S. Li, J. Tenenbaum, and I. Mordatch. Improved contrastive divergence training of energy based models. arXiv preprint arXiv:2012.01316, 2020.
[9] B. Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106(496):1602-1614, 2011.
[10] R. A. Fisher. Theory of statistical estimation. In Mathematical proceedings of the Cambridge philosophical society, volume 22, pages 700-725. Cambridge University Press, 1925.
[11] R. Gao, Y. Lu, J. Zhou, S.-C. Zhu, and Y. N. Wu. Learning generative convnets via multi-grid modeling and sampling. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 9155–9164, 2018.
[12] A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Scholkopf, and A. Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
[13] G. E. Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
[14] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[15] M. F. Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simulation and Computation, 19(2):433-450, 1990.
[16] A. Hyvarinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[17] A. Jolicoeur-Martineau, R. Piche-Taillefer, I. Mitliagkas, and R. T. des Combes. Adversarial score matching and improved sampling for image generation. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id= eLfqM13z31q.
[18] T. Kim and Y. Bengio. Deep directed generative models with energy-based probability estimation. arXiv preprint arXiv:1606.03439, 2016.
[19] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[20] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[41] P. Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674, 2011.
[42] Z. Wang, S. Cheng, L. Yueru, J. Zhu, and B. Zhang. A wasserstein minimum velocity approach to learning unnormalized models. In International Conference on Artificial Intelligence and Statistics, pages 3728-3738. PMLR, 2020.
[43] M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688, 2011.
[44] H. Wendland. Scattered data approximation, volume 17. Cambridge university press, 2004.
[45] L. Wenliang and H. Kanagawa. Blindness of score-based methods to isolated components and mixing proportions. arXiv preprint arXiv:2008.10087, 2020.
[46] J. Xie, Y. Lu, S.-C. Zhu, and Y. Wu. A theory of generative convnet. In International Conference on Machine Learning, pages 2635-2644. PMLR, 2016.
[47] S. Zhai, Y. Cheng, R. Feris, and Z. Zhang. Generative adversarial networks as variational training of energy based models. arXiv preprint arXiv:1611.01799, 2016.
[48] M. Zhang, T. Bird, R. Habib, T. Xu, and D. Barber. Variational f-divergence minimization. arXiv preprint arXiv:1907.11891, 2019.
[49] M. Zhang, P. Hayes, T. Bird, R. Habib, and D. Barber. Spread divergence. In International Conference on Machine Learning, pages 11106-11116. PMLR, 2020.
[50] M. Zhang, P. Hayes, and D. Barber. Generalization gap in amortized inference. Advances in neural information processing systems, 2022.
[51] M. Zhang, O. Key, P. Hayes, D. Barber, B. Paige, and F.-X. Briol. Towards healing the blindness of score matching. arXiv preprint arXiv:2209.07396, 2022.
A Proof and Derivations
The existence is straightforward, since  $\mathrm{FD}(\tilde{p}_d||\tilde{q}_{\theta^*}) = 0 \to \tilde{p}_d = \tilde{q}_{\theta^*}$ , we can simply let  $q(x) = p_d(x)$ , which makes  $\int q(x)p(\tilde{x}|x)\mathrm{d}x = \int p_d(x)p(\tilde{x}|x)\mathrm{d}x = \tilde{p}_d$ . To show the uniqueness, we denote density  $k(\epsilon) = \mathcal{N}(0,\sigma^2 I)$ , so  $\tilde{q}_{\theta}(\tilde{x})$  and  $\tilde{p}_d(x)$  can be written as convolutions
$$
\tilde {q} _ {\theta} (\tilde {x}) = q * k, \quad \tilde {p} _ {d} (\tilde {x}) = p _ {d} * k, \tag {19}
$$
$$
\tilde {p} _ {d} = \tilde {q} _ {\theta} \Leftrightarrow q * k = p _ {d} * k \Leftrightarrow \mathcal {F} (q) \mathcal {F} (k) = \mathcal {F} (p _ {d}) \mathcal {F} (k), \tag {20}
$$
where  $\mathcal{F}$  denotes the Fourier transform. Since the Fourier transform of a Gaussian is also a Gaussian, so  $\mathcal{F}(k) > 0$  everywhere, we have
$$
\tilde {p} _ {d} = \tilde {q} _ {\theta^ {*}} \Leftrightarrow \mathcal {F} (q) \mathcal {F} (\bar {k}) = \mathcal {F} (p _ {d}) \mathcal {F} (\bar {k}) \Leftrightarrow \mathcal {F} (q) = \mathcal {F} (p _ {d}) \Leftrightarrow q = p _ {d}. \tag {21}
$$
Therefore,  $q = p_d$  is the unique distribution that makes  $\tilde{p}_d = \tilde{q}_{\theta}$ . This technique has also been used to construct spread KL divergence (we denote as  $\widetilde{\mathrm{KL}}$ ) [49], which is defined as  $\widetilde{\mathrm{KL}}(p_d||q_{\theta}) \equiv \mathrm{KL}(p_d*k||q_{\theta}*k)$  where  $k(\epsilon) = N(0,\sigma^2 I)$ , to train implicit model  $q_{\theta}$ . Different from the DSM situation, when  $\widetilde{\mathrm{KL}}(p_d||q_{\theta}) = 0$ , the underlying model  $q_{\theta} = p_d$  is directly available, whereas the EBM  $\tilde{q}_{\theta}$  trained by DSM learns to be the noisy distribution  $\tilde{q}_{\theta} = p_d * k$ .
A.2 General Conditions Characterising the Existence of the Clean Model