Deep unfolding of MCMC kernels: scalable, modular & explainable GANs for high-dimensional posterior sampling
Authors: Jonathan Spence, Tobías I. Liaudat, Konstantinos Zygalakis, Marcelo Pereyra
Organizations: Maxwell Institute for Mathematical Sciences, Bayes Centre, 47 Potterrow, Edinburgh, Scotland, UK · School of Mathematical and Computer Sciences, Heriot-Watt University, Edinburgh, Scotland, UK · IRFU, CEA, Université Paris-Saclay, F-91191 Gif-sur-Yvette, France · School of Mathematics, University of Edinburgh, Edinburgh, Scotland, UK
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.
Figures & tables
Figure 1: (a): Illustration of unfolded optimization architecture. The parameters θ,ϑ0,…,ϑL are trained end-to-end such that the output samples from layers x^L is a close approximation of E[x∣y] . (b): Illustration of the proposed unfolded MCMC architecture. Parameters θ,ϑ0,…,ϑL are trained end-to-end such that the output samples from layers ℓ≥L0 closely resemble p(x∣y) over a training dataset.
Method (NFEs)
PSNR (sample)
PSNR (mean)
LPIPS
SW ( ×105 )
W2Latent ( ×102 )
CMMD ( ×103 )
Time (ms)
Size (MB)
U-SGS (4)
23.77
24.96
0.007
9.53
3.70
12.70
0.13
4.0
U-SGS (8)
24.95
26.29
0.006
8.49
2.58
3.91
0.24
4.0
U-SGS (16)
25.78
27.34
0.005
7.08
0.96
5.86
0.48
4.0
U-SGS (32)
26.93
28.12
0.003
6.72
0.54
1.95
0.96
4.0
U-SGS (64)
27.59
28.46
0.003
6.95
0.35
0.98
1.89
4.0
RCGAN (1)
23.28
26.69
0.009
9.21
9.08
21.48
0.10
11.83
Table 1: Reconstruction metrics for the Random Motion deblurring problem. Results are averaged over the MNIST test dataset, using our unrolled Langevin model with L=2i iterations, 2≤i≤6 . Comparisons are made with the rcGAN architecture [ 4 ] and a zero-shot split Gibbs sampler with a data-driven VAE prior (VAE-SGS). LPIPS is computed using a single approximate posterior sample from each model. For each observation y , PSNR results are computed for a single posterior sample and for the posterior mean, approximated using 512 Monte Carlo samples, with the exception of VAE-SGS for which the posterior mean is approximated by the ergodic average. The final columns show the time to generate a single posterior sample for each model, and the storage cost of model parameters. Metrics are scaled by the stated factors to assist readability.
Figure 2: Convergence comparison of U-SGS and conventional SGS on small truncated chains.
Figure 3: Qualitative comparison for motion deblurring applied to the MNIST dataset. For each method, we display trained posterior samples from pθ(x∣y) for a range of three observations from the test dataset. In addition, the residual error and an estimate of the principal three eigen-directions from a PCA expansion are shown. For each eigenvector, the Pearson correlation to the residual is shown in the upper-left corner.
Model (NFEs)
PSNR (Sample)
PSNR (Mean)
LPIPS
SSIM
CFID
Time (s)
RAM (MB)
Weights (M)
RIGAN (1)
41.51
43.77
0.07
0.91
0.52
0.04
1078
195.8
U-LATINO (8)
43.40
45.61
0.01
0.88
0.06
0.08
683
88.3
LATINO (8)
37.33
37.90
0.08
0.64
0.19
0.48
683
88.0
IRIS (1000)
48.08
48.99
0.07
0.86
2.24
55.93
718
88.0
IRIS (64)
41.15
46.09
0.07
0.82
4.06
3.91
718
88.0
Table 2: Reconstruction metrics for the radio-interferometry problem averaged over samples from the PROBES dataset which were not used to fine-tune our model, but which were used in pre-training the SBM in [ 11 ] . For both datasets, we compare PSNR for a single sample and for the estimated posterior mean (using the ergodic mean of L−L0=6 samples for U-LATINO and 8 independent posterior samples for all other methods), LPIPS, SSIM and CFID (using a single posterior sample for each reconstruction). The average time per sample, GPU RAM and number of model parameter weights used to compute a single posterior sample are shown in the final column.
Figure 4: Qualitative comparison of example reconstructions on ground-truth images from the PROBES dataset. The left column displays the ground truth, the pseudo-inverse A†y , and the sampled mask m applied in the Fourier domain. We compare the posterior mean along, the residual and predicted residual (approximated via the standard dev. of 8 posterior samples) on a 32×32 scale. In the upper-left corner of each predicted residual, we report the Pearson correlation to the true residual.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Method (NFEs)
PSNR (sample)
PSNR (mean)
LPIPS
SW ( ×105 )
W2Latent ( ×102 )
CMMD ( ×103 )
U-SGS (4)
22.11
23.55
0.013
10.30
4.95
23.44
U-SGS (8)
23.14
24.76
0.010
9.14
3.95
11.72
U-SGS (16)
23.78
25.59
0.009
7.67
2.16
17.58
U-SGS (32)
25.19
26.73
0.006
7.01
0.81
7.81
U-SGS (64)
25.85
27.02
0.005
7.17
0.54
6.83
RCGAN (1)
21.94
24.81
0.015
7.71
11.46
57.63
Appendix
Table 3: Reconstruction, perceptual, and inception metrics on a single posterior sample for the Random Motion deblurring problem with out-of-distribution observation kernel k∼GP(19,λMatern=0.5,σMatern=0.4) .
Figure 5: Qualitative comparison for motion deblurring applied to the MNIST dataset with out-of-distribution kernels k∼GP(19,λMatern=0.5,σMatern=0.4) .
Figure 6: Qualitative comparison of an example reconstruction on ground-truth images from the PROBES dataset. The left column displays the ground truth data, the pseudo-inverse of A applied to y , and the sampled mask m applied in the Fourier domain. For each method, we compare the posterior mean, standard deviation and the absolute error computed between the posterior mean and ground truth.
Figure 7: Qualitative comparison of an example reconstruction on ground-truth images from the PROBES dataset. The left column displays the ground truth data, the pseudo-inverse of A applied to y , and the sampled mask m applied in the Fourier domain. For each method, we compare the posterior mean, standard deviation and the absolute error computed between the posterior mean and ground truth.
Figure 8: Output of final 6 MCMC samples comparing U-LATINO to LATINO. The ground truth x⋆ is taken from the PROBES hold-out validation data and observation kernel k∈K4h .
Figure 9: Qualitative comparison of an example reconstruction on synthetic ground-truth validation images sampled from the score-based model prior. The left column displays the ground truth data, the pseudo-inverse of A applied to y , and the sampled mask m applied in the Fourier domain. For each method, we compare the posterior mean, standard deviation and the absolute error computed between the posterior mean and ground truth. To ease visual comparison, the mask m , the standard deviation and the absolute error are visualised on a log-scale colour gradient.
Model (NFEs)
PSNR
LPIPS
SSIM
CFID
RIGAN (1)
43.22
0.06
0.91
0.48
U-LATINO (8)
44.97
0.01
0.79
0.02
LATINO (8)
39.00
0.03
0.92
0.08
IRIS (1000)
50.87
0.05
0.87
0.08
IRIS (64)
44.36
0.08
0.75
2.45
Appendix
Table 4: Reconstruction metrics for the RI task using the PROBES score-based prior images.
Figure 10: Qualitative comparison of an example reconstruction on a ground-truth image from the PROBES dataset using an out-of-distribution kernel k∼U(K2h) .
Model (NFEs)
PSNR (Sample)
PSNR (Mean)
LPIPS
SSIM
CFID
RIGAN (1)
40.35
42.19
0.08
0.91
0.55
U-LATINO (8)
43.57
45.50
0.02
0.82
0.06
LATINO (8)
37.53
37.89
0.07
0.62
0.20
IRIS (1000)
47.74
48.53
0.06
0.83
0.32
IRIS (64)
40.94
45.70
0.08
0.81
3.22
Appendix
Table 5: Reconstruction metrics for the RI problem using out-of-distribution kernels k∼U(K2h).
Coding and computation remain major bottlenecks in Markov chain Monte Carlo (MCMC) workflows, especially as modern sampling algorithms have become increasingly complex and existing probabilistic programming systems remain limited in model support, extensibility, and composability. We introduce AI4BayesCode, an extensible LLM-driven system that translates natural-language Bayesian model descriptions into runnable, validated MCMC samplers. To improve reliability, AI4BayesCode adopts a modular design that decomposes models into modular sampling blocks and maps each block to a built-in sampling component, reducing the need to implement complex sampling algorithms from scratch. Reliability is further improved through pre-generation validation of model specifications and post-generation validation of generated sampler code. AI4BayesCode also introduces a novel recursively stateful coding paradigm for MCMC, allowing modular sampling components, potentially developed by different contributors, to be composed coherently within larger MCMC procedures. We develop a benchmark suite to evaluate AI4BayesCode for sampler-generation. Experiments show that AI4BayesCode can implement a wide range of Bayesian models from natural-language descriptions alone. As an open-ended system, its capability can continue to expand with improvements in the underlying AI agent and the addition of new built-in blocks.
Efficient sampling from Bayesian posterior distributions in high dimensions remains a hard problem, despite advances in theory and algorithms. Multiple steps are usually required to draw unbiased samples, for example, iterations of an MCMC algorithm or integration steps in a diffusion-based amortised sampler. However, scalable one-step amortised inference remains a challenging open problem. In this work, we build upon recent progress in data-driven generative modelling to offer a possible solution for one-step amortised sampling. Our one-step implicit sampler (OSIS) uses short-time MCMC dynamics to define a local target for the evolution of an implicitly modelled distribution. In the asymptotic limit, the modelled distribution is a fixed point of the MCMC kernel and thus samples from the target. We analyse the convergence properties of OSIS and show connections with other generative modelling and sampling algorithms. We study various design choices and show that OSIS performs well compared to multi-step amortised samplers and non-amortised MCMC baselines while being far more efficient at sampling time.
Vincent Pauline, Kirill Tamogashev, Arran Carter +3
Technical University of Munich · Munich Center for Machine Learning · University of Edinburgh +1
Sampling from complex, unnormalized probability densities is a fundamental challenge in Bayesian inference and probabilistic modeling. While Markov chain Monte Carlo (MCMC) methods provide asymptotic guarantees, they often suffer from slow mixing and high computational costs due to fixed or manually tuned trajectory lengths. In this work, we propose a novel framework that treats trajectory termination as a learnable component of the sampling dynamics. By framing MCMC within the theory of non-acyclic generative flow networks (GFlowNets), we train state-dependent neural classifiers to decide when a trajectory has reached a high-density region and should terminate. We theoretically establish the connection between optimal classifiers and the target density via detailed balance conditions and introduce a multilevel training scheme to facilitate exploration in complex geometries. Experimental results across various benchmark densities demonstrate that our approach significantly reduces average trajectory lengths while improving mode coverage and mixing compared to standard MCMC baselines.