Learning to Replace MCMC in Split-Gibbs Diffusion Posterior Sampling via Deep Unfolding
Authors: Yi Zhang, Rui Guo, Mengchu Xu, Zhaofeng Liu, Yonina C. Eldar
Organizations: Faculty of Mathematics and Computer Science, Weizmann Institute of Science, Rehovot, Israel · Department of Electrical and Computer Engineering, Northeastern University, Boston, MA, USA
Split Gibbs sampling enables diffusion posterior inference for general nonlinear inverse problems by decoupling prior and likelihood computations, allowing a pretrained diffusion prior to be reused across measurement models. However, its likelihood update often relies on iterative MCMC, which can hinder parallelization, require algorithm-specific tuning, and incur substantial computational cost. In this work, we propose a learning-based framework to replace this MCMC step by reformulating both Gibbs updates as Gaussian denoising problems and implementing them through ODE diffusion. The prior step reuses a pretrained denoiser, while the likelihood denoiser exploits known likelihood structure through a lightweight deep-unfolded network. Experiments on nonlinear phase retrieval demonstrate the effectiveness of the proposed method as an alternative to MCMC-based split Gibbs at lower likelihood-update cost.
Figures & tables
Figure 1: Deep-unfolded ODE-diffusion implementation of the alternating-denoising split Gibbs sampler.
Figure 2: Reconstruction quality versus computation time for different implementations of the split Gibbs sampler. Time includes both prior and likelihood updates; markers indicate the endpoints after 22 outer iterations.
MCMC
Proposed
Likelihood update
600 iter.
800 iter.
1000 iter.
1200 iter.
1400 iter.
5 NFEs
10 NFEs
15 NFEs
Likelihood time (ms)
289.964
386.634
483.589
579.937
676.557
48.624
97.061
145.566
PSNR (dB) ↑
20.147
20.325
20.412
20.463
20.479
20.092
20.412
20.518
SSIM ↑
0.5990
0.6032
0.6056
0.6068
0.6072
0.5997
0.6068
0.6091
Table 1: Average likelihood-update time per image per outer iteration and reconstruction quality after 22 outer iterations.
Diffusion posterior sampling combines generative diffusion priors with measurement models to infer plausible solutions and uncertainty in inverse problems.
Split Gibbs sampling (SGS) is a popular framework for posterior sampling in Bayesian imaging inverse problems. It decouples a Gaussian data-fidelity term from a complex prior through an auxiliary variable, so the data variable is updated exactly and only the prior-side conditional is hard to sample. Existing samplers treat this conditional in one of two ways. Plug-and-play SGS runs a multi-step diffusion denoiser at every iteration, which is expensive and lacks non-asymptotic guarantees. Langevin-within-SGS takes cheap overdamped Langevin steps but needs many iterations. We propose RED-KLwSGS, which keeps the exact Gaussian update for the data variable and updates the auxiliary variable with underdamped (kinetic) Langevin diffusions driven by a one-shot denoising score, at the same per-iteration cost as Langevin-within-SGS. We prove non-asymptotic Wasserstein-2 convergence in continuous and discrete time for strongly log-concave priors. We also introduce Joint-RED-KLwSGS, which applies kinetic Langevin diffusions to both variables. Experiments with Denoising diffusion probabilistic models as diffusion priors on FFHQ and ImageNet datasets show faster convergence and high-quality image reconstruction.
Dai Hai Nguyen, Duc Dung Nguyen
Hokkaido University · Institute of Information Technology, VAST
Pretrained diffusion models are powerful priors for inverse problems, but posterior sampling under nonlinear, non-differentiable forward models remain hard. We introduce diffusion waltz, an MCMC method using SDEdit-style noising-denoising as a proposal, corrected via Metropolis-Hastings for exact posterior sampling without prior evaluation. We further propose injecting observations into the proposal while preserving exactness, using a gradient-free ensemble Kalman update. On a non-differentiable Navier-Stokes initial condition recovery task, diffusion waltz outperforms existing baselines across different noise and nonlinearity regimes.
This paper addresses the issue of inversion in cases where (1) the observation system is modeled by a linear transformation and additive error, (2) the problem is ill-posed and regularization relies on a Bayesian strategy, (3)~the prior is modeled by a diffusion process adjusted on an available large set of examples. In this context, it is known that the issue of posterior sampling is a thorny one and the paper introduces a Gibbs algorithm. It appears that this avenue has not been explored, and we show that it is particularly effective and remarkably simple. In addition, it provides clear elements regarding convergence guarantees in a specific case and arguments supporting such guarantees in practical cases. The results are clearly confirmed by numerical simulations based on a toy example.
Jean-François Giovannelli
Groupe Signal-Image, IMS (Univ. Bordeaux, CNRS, BINP), Talence, France