SlowGuess/NeurIPS_2023_OCR huggingface dataset:
def sample_langevin_prior_z(z, netE):
    s = step_size
    for i in range(n_steps):
        en = netE(z).sum()
        z_norm = 1.0 / 2.0 * torch.sum(z**2)
        z_grad = torch.autograd.grad(en + z_norm, z)[0]
        w = torch.random_like(z)
        # Prior LD Update
        z.data = z.data - 0.5 * (s ** 2) * z_grad + s * w
    return zdetach()
def sample_langevin_posterior_z(z, x, netG, netE):
    s = step_size
    sigma_inv = 1.0 / (2.0 * sigma ** 2)
    for i in range(n_steps):
        x_hat = netG(z)
        g_log_lkhd = sigma_inv * torch.sum((x_hat - x) ** z_n = 1.0 / 2.0 * torch.sum(z**2)
        en = netE(z).sum()
        total_en = g_log_lkhd + en + z_n
        z_grad = torch.autograd.grad(total_en, z)[0]
        w = torch.random_like(z)
        # Posterior LD Update
        z.data = z.data - 0.5 * (s ** 2) * z_grad + s * w
    return zdetach()
Listing 2: Learning LEBM with DAMC.
for x in dataset: # mask for unconditional learning of DAMC z_maskprob  $=$  torch.rand((len(x)),device  $\equiv$  x_device) z_mask  $=$  torch. ones(len(x), device=x_device) z_mask[z_maskprob  $<  0.2] = 0.0$  z_mask  $=$  z_mask unsqueeze(-1) # draw DAMC samples  $\mathrm{z}0 = \mathrm{Q}(\mathrm{x})$  zk_pos,zk_neg  $= \mathrm{z0}$  .detach().clone(),z0detach().clone() # prior and posterior updates zk_pos  $=$  sample_langevin_posterior_z(  $\mathrm{z = zk\_pos}$  ,x=x,netG=G，netE=E) zk_neg  $=$  sample_langevin_prior_z( z=torch.cat( [zk_neg,torch.randin_like(zk_neg)],dim=0), netE=E) #update Q for in range(6): Q_OPTizer.zero_grad () Q_loss  $=$  Qcalculate_loss (  $\mathrm{x = x}$  ，z=zk_pos，mask  $\equiv$  z_mask).mean() Q_loss.Backward () Q_OPTizer step () # update G G_OPTizer.zero_grad () x_hat  $=$  G(zk_pos) g_loss  $=$  torch-sum((x_hat-x)\*\*2，dim=[1,2,3]).mean() g_loss.backup () G_OPTizer step () # update E E_OPTizer.zero_grad () e_pos,e_neg  $=$  E(zk_pos)，E(zk_neg) E_loss  $=$  e_pos.mean() -e_neg.mean() E_loss.backup () E_OPTizer step ()
Datasets We include the following datasets to study our method: SVHN  $(32 \times 32 \times 3)$ , CIFAR-10  $(32 \times 32 \times 3)$ , CelebA  $(64 \times 64 \times 3)$ , CeleAMask-HQ  $(256 \times 256 \times 3)$  and MNIST  $(28 \times 28 \times 1)$ . Following Pang et al. [22], we use the full training set of SVHN (73,257) and CIFAR-10 (50,000), and take 40,000 samples of CelebA as the training data. We take 29,500 samples from the CelebAMask-HQ dataset as the training data, and test the model on 500 held-out samples. For anomaly detection on MNIST dataset, we follow the experimental settings in [22, 41, 55, 56] and use  $80\%$  of the in-domain data to train the model. The images are scaled to  $[-1,1]$  and randomly horizontally flipped with a prob. of .5 for training.
GAN inversion settings We attempt to use the DAMC sampler for GAN version on the FFHQ (256 x 256 x 3) and LSUN-Tower (256 x 256 x 3) datasets. We take 69,500 samples from the CelebAMask-HQ dataset as the training data, and use the held-out 500 samples for testing. We follow the default data splits of the LSUN dataset.
For the LEBM-based inversion method, we train a LEBM in the  $14 \times 512 = 7168$  dim. latent space. During training, we add  $l_{2}$ -regularization on the energy score of LEBM to stabilize training, as suggested in [59]. For the DAMC sampler and the encoder-based inversion method [16], we generate initial posterior samples of the training data using [18] and train the DAMC sampler and the encoder-based method for 5K iterations using these samples as a warm-up step. The encoder-based method is trained by minimizing the  $l_{2}$  distance between the encoder output and the target samples. After that, these two methods are trained with the default learning algorithms.
E.1 Toy Examples
Figure S1: The 2-arm pinwheel-shaped prior distribution used in the toy example.
(a) Evolution of the learned posterior distributions
(b) GT
Figure S2: Evolution of the posterior distributions learned by the DAMC sampler. In each row, we display from left to right the evolution of the learned posterior distributions from the DAMC sampler through training iterations. The last column shows the corresponding ground-truth posterior distribution obtained by running 1K-3K steps of Langevin Dynamics until convergence for posterior sampling.
As the proof-of-concept toy examples, we implement the neural likelihood experiments following the same set-up mentioned in Section 5.1 and B.1 in [80]. We choose to use a more complex prior distribution, i.e., 2-arm pinwheel-shaped prior distribution (shown in Fig. S1) instead of a standard normal one to make sure that the true posterior distributions are multimodal. The ground-truth posterior distributions are obtained by performing long-run LD sampling until convergence. We visualize the convergence trajectory of the posterior distributions learned by our model to the ground-truth ones. In Fig. S2, we can see that our model can faithfully reproduce the ground-truth distributions with sufficient training iterations, which indicates that the learned sampler successfully amortizes the long-run sampling chain with the length of 1,000-3,000 iterations.
E.2 Generation
We provide additional generated samples from our models trained on SVHN (Fig. S3), CelebA (Fig. S4), CIFAR-10 (Fig. S5) and CelebA-HQ (Fig. S6).
(a) Ours-DAMC
(b) Ours-LEBM
Figure S3: Samples generated from the DAMC sampler and LEBM trained on the SVHN dataset.
(a) Ours-DAMC
Figure S4: Samples generated from the DAMC sampler and LEBM trained on the CelebA dataset.
E.3 Reconstruction
We provide qualitative examples about the reconstruction results from our models trained on SVHN (Fig. S7), CelebA (Fig. S8), CIFAR-10 (Fig. S9) and CelebA-HQ (Fig. S10). Observed images are sampled from the testing set unseen during training.
Figure S5: Samples generated from the DAMC sampler and LEBM trained on the CIFAR-10 dataset.
(a) Ours-DAMC
Figure S6: Samples generated from the DAMC sampler and LEBM trained on the CelebA-HQ dataset.
E.4 Visualization of Transitions
We provide additional visualization results of LD transitions initialized from  $\mathcal{N}(0,\mathbf{I}_d)$  on SVHN (Fig. S11) and CIFAR-10 datasets (Fig. S12). For the 200-step set-up, we can see that the generation quality quickly improves by exploring the local modes with LD. For the 2500-step long-run set-up, we can see that the LD produces consistently valid results without the oversaturation issue of the long-run chain samples.
(a) Observation
Figure S7: Reconstructed samples from the posterior DAMC sampler trained on the SVHN dataset.
(b) Reconstruction
(a) Observation
Figure S8: Reconstructed samples from the posterior DAMC sampler trained on the CelebA dataset.
(b) Reconstruction
Figure S9: Reconstructed samples from the posterior DAMC sampler trained on the CIFAR-10 dataset.
Figure S10: Reconstructed samples from the posterior DAMC sampler trained on the CelebA-HQ dataset.
(a) 200 steps
Figure S11: Transition of Markov chains initialized from  $\mathcal{N}(0,\mathbf{I}_d)$  towards  $p_{\alpha}(z)$  on SVHN. 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.
(b) 2500 steps
(b) 2500 steps
Figure S12: Transition of Markov chains initialized from  $\mathcal{N}(0,\mathbf{I}_d)$  towards  $p_{\alpha}(z)$  on CIFAR-10. 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.
F Additional Quantitative Results
F.1 Learning DAMC in the Latent Space of Other DLVMs
We have compared our method with directly training diffusion models in the latent space. In the NALR-LEBM column in Table 4, we compare our method with learning a diffusion model in the pre-trained energy-based latent space on CIFAR-10 dataset. We further add experiments of learning this model in the pre-trained VAE [1] and ABP [51] model latent spaces summarized in Table S5. We can see that our method greatly outperforms these baselines using the same network architectures.
Table S5: Further baseline results for learning DAMC on CIFAR-10 dataset. These models are implemented using the same encoder and decoder network architectures for fair comparison
VAE	ABP	NALR	Ours
FID	102.54	69.93	64.38	57.52
MSE	0.036	0.017	0.016	0.015
F.2 Further Ablation Studies
We conduct more ablation studies to inspect the impact of some hyper-parameters. For prior LD sampling, we follow [22] and set the sampling step as 60 throughout the experiments for fair comparison. We further add experiments on CIFAR-10 for training our model with different posterior LD sampling steps T and with different capacities of the amortizer  $q_{\phi}$ , summarized in Table S6. We observe that larger posterior sampling steps and larger model capacities for training brings marginal improvement compared with our default set-up. Less sampling steps and lower model capacities, however, may have negative impact on the performances.
(a) Posterior sampling steps T
(b) Model capacity of  $q_{\phi}$
T=10	T=30	T=50	f=1/4	f=1/2	f=1	f=2	f=4
FID	74.20	57.72	57.03	FID	116.28	80.28	57.52	57.82	57.56
MSE	0.016	0.015	0.015	MSE	0.017	0.016	0.015	0.015	0.015
Table S6: Ablation Studies for the choice of Langevin steps and model capacity. We highlight the results of our set-up reported in the main text. T is the posterior sampling step. f stands for the factor for the model capacity, e.g.,  $\mathrm{f} = 2$  means  $2\mathrm{x}$  the size of the original model.
G Further Discussion
We mentioned in the main text that one potential disadvantage of our method is its parameter inefficiency for introducing an extra DDPM. Although fortunately, our models are in the latent space so the network is lightweight. To be specific, on SVHN, CelebA, CIFAR-10 and CelebA-HQ datasets the number of parameters in the diffusion network is around  $10\%$  of those in the generator.
Another issue is the time efficiency. We mentioned in the main text that the time efficiency for sampling is competitive. With the batch size of 64, on these datasets the DAMC prior sampling takes  $0.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$ . However, during training we need to run 30 steps of posterior LD sampling and 60 steps of prior LD sampling in each training iteration. We observe that the proposed learning method takes 15.2 minutes per training epoch, while the short-run LD-based learning method takes 14.8 minutes per epoch. These methods are slower than the VAE-base method, which takes 5.5 minutes for an training epoch. We can see that the time efficiency for training is generally bottlenecked by the LD sampling process, and could be improved in future works.
Generative models could be misused for disinformation or faking profiles. Our work focuses on the learning algorithm of energy-based prior model. Though we consider our work to be foundational and not tied to particular applications or deployments, it is possible that more powerful energy-based generative models augmented with this method may be used maliciously. Work on the reliable detection of synthetic content could be important to address such harms from generative models.