Kinetic langevin speeds up posterior sampling for image reconstruction
Kinetic Langevin Meets Split Gibbs: Accelerated Posterior Sampling for Imaging Inverse Problems with Diffusion Priors
Machine Learning
Summary
Imaging inverse problems involve reconstructing images from incomplete or noisy data, which can be challenging due to complex underlying patterns. The authors focus on improving sampling methods that combine data information with learned image models to produce better reconstructions. They introduce a new approach that uses kinetic Langevin dynamics, a way of exploring image possibilities more efficiently, alongside a splitting technique that simplifies computations. Their method converges faster and produces high-quality images in tests on popular datasets.
What this means in practice
- •For medical imaging teams: Improve reconstruction speed and quality for 3D scans by using kinetic Langevin sampling with diffusion-based priors.
- •For computational photography developers: Enhance image restoration and super-resolution results with faster posterior sampling methods incorporating kinetic Langevin dynamics.
Authors
Dai Hai Nguyen, Duc Dung Nguyen
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.