Learning Energy-Based Prior Model with Diffusion-Amortized MCMC

Abstract: 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 multi-modal 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

• The paper introduces Diffusion-Amortized MCMC to efficiently sample latent space EBMs, enhancing training stability and generation quality.
• It leverages denoising diffusion processes to replace long-run MCMC with a learnable neural sampler, reducing computational costs.
• Experiments on SVHN, CelebA, and CIFAR-10 show improved FID and MSE performance, demonstrating robust handling of complex generative tasks.

Learning Energy-Based Prior Model with Diffusion-Amortized MCMC

Introduction

The paper introduces a novel approach for improving latent space Energy-Based Models (EBMs), termed Energy-Based Prior Models, used in generative modeling. The focus is on ameliorating issues related to non-convergent short-run Markov Chain Monte Carlo (MCMC) sampling, which typically leads to suboptimal generation quality and training instability. The authors propose a Diffusion-Amortized MCMC (DAMC) method, designed to offer a more effective sampling strategy for learning latent Space EBMs.

Methodology

Diffusion-Amortized MCMC

The core contribution is the introduction of the DAMC method, which leverages a diffusion process for amortizing long-run MCMC sampling. This is inspired by the connection between MCMC sampling and denoising diffusion probabilistic models (DDPMs), where the noise prediction plays a similar role to the gradient of log density.

The DAMC approach involves iteratively distilling the MCMC process into a learnable neural sampler, effectively replacing expensive long-run simulations with computationally efficient amortized inference.

Figure 1: Transition of Markov chains initialized from $N(0, I_d)$ towards $p_{(z)}$.

Learning Algorithm

The learning algorithm integrates DAMC with Maximum Likelihood Estimation (MLE) for training latent space EBMs. The key steps include:

1. Sampling: Using the current DAMC sampler to initialize short-run MCMC chains.
2. Distillation: Iteratively refine the DAMC sampler to approximate the long-run MCMC distribution.
3. Optimization: Utilizing gradient-based methods to update the parameters of the EBM and the DAMC sampler.

This iterative scheme ensures that both prior and posterior distributions are accurately represented, enhancing model expressivity and performance.

Experiments

Performance Evaluation

The authors conduct experiments on several benchmark datasets, including SVHN, CelebA, and CIFAR-10. The results demonstrate that models trained with DAMC exhibit superior performance in terms of FID scores and mean squared error (MSE), validating the efficacy of the approach.

Figure 2: Samples generated from the DAMC sampler and LEBM.

The study also highlights the robustness of DAMC in handling high-dimensional and multi-modal distributions, showcasing its applicability to complex generative tasks.

Implications and Future Work

The integration of diffusion-based amortization in training EBMs presents significant advancements in generative modeling, particularly for applications requiring rich latent representations. The proposed method offers a scalable and efficient framework, potentially extendable to other forms of unnormalized continuous densities.

Future work could explore the application of DAMC to other domains involving latent variable models and extend the approach to concurrent learning of EBMs with other generative paradigms.

Conclusion

The paper successfully addresses longstanding challenges in sampling for latent space EBMs, providing both empirical and theoretical validation for the proposed DAMC approach. By marrying diffusion processes with MCMC techniques, the method paves the way for more effective and efficient generative model training.