Kinetic Langevin Meets Split Gibbs: Accelerated Posterior Sampling for Imaging Inverse Problems with Diffusion Priors
Organizations: Hokkaido University · Institute of Information Technology, VAST
Abstract
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.
Figures & tables
| Method | Time/iter (s) | Iterations | Total time (s) |
|---|---|---|---|
| PnP-SGS | 2.121 | 22 | 46.64 |
| RED-LwSGS | 0.087 | 600 | 52.2 |
| RED-KLwSGS (ours) | 0.091 | 350 | 31.5 |
| SPA | TV-ADMM | PnP-ADMM | DDRM | PnP-SGS | RED-LwSGS | RED-KLwSGS (ours) | Joint-RED-KLwSGS (ours) | ||
|---|---|---|---|---|---|---|---|---|---|
| Deblurring (Gaussian) | PSNR | 23.17 | 22.37 | 24.93 | 23.36 | 27.96 | 28.01 | 29.51 | 29.32 |
| SSIM | 0.499 | 0.801 | 0.812 | 0.767 | 0.837 | 0.842 | 0.865 | 0.853 | |
| LPIPS | 0.452 | 0.507 | 0.441 | 0.332 | 0.331 | 0.352 | 0.336 | 0.341 | |
| Deblurring (motion) | PSNR | 17.73 | 21.36 | 24.65 | N/A | 28.46 | 28.22 | 29.74 | 29.43 |
| SSIM | 0.211 | 0.751 | 0.825 | N/A | 0.828 | 0.823 | 0.858 | 0.851 | |
| LPIPS | 0.446 | 0.508 | 0.405 | N/A | 0.294 | 0.292 | 0.302 | 0.315 |
| SPA | TV-ADMM | PnP-ADMM | DDRM | PnP-SGS | RED-LwSGS | RED-KLwSGS (ours) | Joint-RED-KLwSGS (ours) | ||
|---|---|---|---|---|---|---|---|---|---|
| Deblurring (Gaussian) | PSNR | 21.08 | 19.99 | 21.81 | 22.73 | 21.76 | 23.12 | 23.92 | 23.85 |
| SSIM | 0.577 | 0.634 | 0.669 | 0.705 | 0.701 | 0.712 | 0.735 | 0.729 | |
| LPIPS | 0.537 | 0.588 | 0.519 | 0.427 | 0.399 | 0.359 | 0.349 | 0.365 | |
| Deblurring (motion) | PSNR | 20.49 | 20.79 | 21.98 | N/A | 21.47 | 22.73 | 24.52 | 24.21 |
| SSIM | 0.681 | 0.677 | 0.702 | N/A | 0.695 | 0.709 | 0.743 | 0.739 | |
| LPIPS | 0.538 | 0.525 | 0.483 | N/A | 0.372 | 0.339 | 0.342 | 0.325 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.