Deep unfolding of MCMC kernels: scalable, modular & explainable GANs for high-dimensional posterior sampling

TL;DR

This paper introduces a novel deep unfolding approach to MCMC algorithms, creating modular, interpretable GAN-based samplers that are efficient and adaptable for high-dimensional Bayesian posterior sampling.

Contribution

It presents a new architecture that combines deep unfolding with GANs for scalable, explainable, and flexible posterior sampling in high-dimensional Bayesian inference.

Findings

Achieves high sampling accuracy in Bayesian imaging tasks.

Offers computational efficiency comparable to traditional MCMC methods.

Maintains interpretability and adaptability of classical MCMC strategies.

Abstract

Markov chain Monte Carlo (MCMC) methods are fundamental to Bayesian computation, but can be computationally intensive, especially in high-dimensional settings. Push-forward generative models, such as generative adversarial networks (GANs), variational auto-encoders and normalising flows offer a computationally efficient alternative for posterior sampling. However, push-forward models are opaque as they lack the modularity of Bayes Theorem, leading to poor generalisation with respect to changes in the likelihood function. In this work, we introduce a novel approach to GAN architecture design by applying deep unfolding to Langevin MCMC algorithms. This paradigm maps fixed-step iterative algorithms onto modular neural networks, yielding architectures that are both flexible and amenable to interpretation. Crucially, our design allows key model parameters to be specified at inference time, offering robustness to changes in the likelihood parameters. We train these unfolded samplers end-to-end using a supervised regularized Wasserstein GAN framework for posterior sampling. Through extensive Bayesian imaging experiments, we demonstrate that our proposed approach achieves high sampling accuracy and excellent computational efficiency, while retaining the physics consistency, adaptability and interpretability of classical MCMC strategies.