Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors

TL;DR

A Markov chain Monte Carlo algorithm is proposed that performs posterior sampling for general inverse problems by reducing it to sampling the posterior of a Gaussian denoising problem.

Abstract

Diffusion models (DMs) have recently shown outstanding capabilities in modeling complex image distributions, making them expressive image priors for solving Bayesian inverse problems. However, most existing DM-based methods rely on approximations in the generative process to be generic to different inverse problems, leading to inaccurate sample distributions that deviate from the target posterior defined within the Bayesian framework. To harness the generative power of DMs while avoiding such approximations, we propose a Markov chain Monte Carlo algorithm that performs posterior sampling for general inverse problems by reducing it to sampling the posterior of a Gaussian denoising problem. Crucially, we leverage a general DM formulation as a unified interface that allows for rigorously solving the denoising problem with a range of state-of-the-art DMs. We demonstrate the effectiveness of the proposed method on six inverse problems (three linear and three nonlinear), including a real-world black hole imaging problem. Experimental results indicate that our proposed method offers more accurate reconstructions and posterior estimation compared to existing DM-based imaging inverse methods.

Figures (21)

• Figure 1: Demonstration of the proposed method, PnP-DM, for posterior sampling using the real data for the M87 black hole from April 6$^\text{th}$, 2017 event2019first_paper4. The black hole imaging problem is non-convex and highly ill-posed due to severe noise corruption and measurement sparsity. Our method rigorously integrates measurements from a real-world imaging system with an expressive image prior in the form of a diffusion model, which was trained with images from the GRMHD black hole simulation event2019first_paper5 in this case. Besides having high visual quality, our posterior samples accurately capture key features of the M87 black hole such as the bright spot location and ring diameter.
• Figure 2: A schematic diagram of our method. Our method alternates between a likelihood step that enforces data consistency and a prior step that solves a denoising posterior sampling problem by leveraging the Split Gibbs Sampler vono2019split. An annealing schedule controls the strength of the two steps at each iteration to facilitate efficient and accurate sampling. A crucial part of our design is the prior step, where we identify a key connection to a general diffusion model framework called the EDM karras2022elucidating. This connection allows us to easily incorporate a family of state-of-the-art diffusion models as priors to conduct posterior sampling in a principled way without additional training. Our method demonstrates strong performance on a variety of linear and nonlinear inverse problems.
• Figure 3: A conceptual illustration of the non-stationary and stationary time-continuous processes as interpolations of $K$ discretize iterations of PnP-DM.
• Figure 4: Results on a synthetic problem with the ground truth posterior available. PnP-DM can sample it more accurately that DPS chung2023diffusion.
• Figure 5: Visual examples for the motion deblur problem ($\sigma_{\bm{y}}=0.05$). We visualize one sample generated by each sampling algorithm.

Theorems & Definitions (10)

• Theorem 3.1
• Proposition A.2: Brownian bridge for the likelihood step
• proof
• Proposition A.3: EDM reverse diffusion for the prior step
• proof
• Lemma A.4
• proof : Proof of Lemma \ref{['lem:key_lemma']}
• proof : Proof of Theorem \ref{['thm:main_result']}
• Proposition A.5
• proof