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

Introduction

Inverse problems arise in many computational imaging applications, where the goal is to recover an image x 2 R n from a set of sparse and noisy measurements y 2 R m . The relationship between x and y can be described by

where A(•) : R n ! R m is the forward operator (linear or nonlinear) and n is the random measurement noise in R m . Since the sparsity and noisiness of y often lead to significant uncertainty in x, it is preferable to sample the posterior distribution p(x|y) over all possible solutions based on some prior distribution p(x), rather than finding a single deterministic solution. Traditional posterior sampling methods often rely on simple image priors that do not reflect the sophistication of real-world image distributions. On the other hand, diffusion models (DMs) have recently emerged as a powerful tool for modeling highly complex image distributions [33,63]. Nevertheless, it remains a challenge to turn DMs into reliable imaging inverse solvers, which motivates us to develop a principled Bayesian method that leverages DMs as priors for posterior sampling.

Diffusion models generate samples from a distribution by reversing a diffusion process from the target distribution to a simple (usually Gaussian) distribution [33,63]. In particular, it estimates a clean image x 0 from a noise image x T by successively denoising noisy images, where x t ⇠ p t is the intermediate noisy image at time t 2 [0, T ]. Reversing diffusion requires one to estimate the time-varying gradient log density (score function) r log p t (x t ) along the diffusion process, or r log p t (x t |y) in the case of sampling the posterior p(x|y).

Mean image Posterior samples Official image by EHT Diffusion model image prior Real-world imaging system * Experiment is performed with real data for the M87 black hole Interferometry N-S frequency (v) E-W frequency (u) Fourier spectrum PnP-DM Proposed method 40μas 40μas

Figure 1: Demonstration of the proposed method, PnP-DM, for posterior sampling using the real data for the M87 black hole from April 6 th , 2017 [21]. 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 [22] 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.

To design generic DM-based inverse problem solvers, most existing methods attempt to approximate the time-varying gradient log density r log p t (x t |y) [17,73,84,62,39,58,60,44,15,77,18,55,8].

In particular they first apply Bayes' rule to separate the forward operator from an unconditional prior over the intermediate noisy image x t :

r log p t (x t |y) = r log p t (y|x t ) + r log p t (x t ).

(

By instead aiming to evaluate the right hand side, one can leverage the existing pre-trained DMs for the unconditional term r log p t (x t ). However, the main challenge in this case is that r log p t (y|x t ) is intractable to compute in general, as p t (y|x t ) involves an integral over all possible x 0 's that could give rise to x t [17]. Various methods have been proposed to circumvent the intractability and can mostly be categorized into two groups. One group of methods explicitly approximate r log p t (y|x t ) by making simplifying assumptions [62,17,60,8]. However, even for arguably the finest approximation to date proposed in the recent work [8], it is exact only when the prior distribution p(x) is Gaussian. For general prior distributions beyond Gaussian, these methods do not sample the true posterior p(x|y). The other group of methods do not make explicit approximations but instead substitute r log p t (y|x t ) with empirically designed updates where y is treated as a guidance signal [73,84,39,58,44,15,77,18,55]. Although these methods may have strong empirical performance, they have deviated from the Bayesian formulation and no longer aim to sample the target posterior. In summary, these existing DM-based inverse methods should be best viewed as guidance methods, where the generative process is guided towards the regions where the measurement y is more likely to be observed, not as posterior sampling methods [8]. We also note that some recent work considered combing DMs with Sequential Monte Carlo to ensure asymptotic consistency in posterior sampling [11,23], but the investigation has been limited to linear imaging inverse problems.

Our contributions In this work, we pursue a different path towards posterior sampling with DM priors by proposing a new Markov chain Monte Carlo (MCMC) algorithm, which we call Plug-and-Play Diffusion Models (PnP-DM). It incorporates DMs in a principled way and circumvents the approximation required when taking the approach in (2). The proposed algorithm is based on the Split Gibbs Sampler [71] that alternates between two sampling steps that separately involve the likelihood and prior. While the likelihood step can be tackled with traditional sampling techniques, the prior step involves a Bayesian denoising problem that requires careful design. Importantly, we identify a connection between the Bayesian denoising problem and the unconditional image generation problem under a general formulation of DMs presented in [37] (which is referred to as the EDM formulation hereafter). This connection allows us to perform rigorous posterior sampling for denoising using DMs without approximating the generative process and enables the use of a wide range of pretrained DMs through the unified EDM formulation. We present an analysis on the non-asymptotic behavior of PnP-DM by establishing a stationarity guarantee in terms of the average Fisher information. We further demonstrate the strong empirical performance of PnP-DM by investigating three linear and three nonlinear noisy inverse problems, including a black hole interferometric imaging problem involving real data that is both nonlinear and severely ill-posed (see Figure 1). Overall, PnP-DM outperforms existing baseline methods, achieving higher accuracy in posterior estimation.

Preliminaries

Split Gibbs Sampler (SGS) is an MCMC approach developed for Bayesian inference [71]. It is also related to the Proximal Sampler [42,14,25,79] and serves as the backbone for the Generative Plug-and-Play (GPnP) [6] and Diffusion Plug-and-Play (DPnP) [78] frameworks in computational imaging. The goal of SGS is to sample the posterior distribution p(x|y) / p(y|x)p(x) = exp( f (x; y) g(x))

where f (x; y) := log p(y|x) and g(x) := log p(x) are the potential functions of the likelihood and prior distribution, respectively. The dual dependence of (3) on both the likelihood and prior makes it nontrivial to directly sample from it in general. Instead, SGS leverages the composite structure of the posterior distribution by adopting a variable-splitting strategy and considers sampling an alternative distribution

where z 2 R n is an augmented variable and ⇢ > 0 is a hyperparameter that controls the strength of the coupling between x and z. We denote the x-and z-marginal distributions of (4) as ⇡ X (x) := R ⇡(x, z)dz and ⇡ Z (x) := R ⇡(x, z)dx, respectively. As ⇢ ! 0, ⇡ X converges to the target posterior p(x|y) in terms of total variation distance [71], so one can obtain approximate samples from the target posterior by sampling (4) instead.

SGS samples (4) via Gibbs sampling. Specifically, SGS starts from an initialization x (0) and, for iteration k = 0, • • • , K 1, alternates between

Note that the two conditional distributions separately involve f (•; y) and g(•). The likelihood and prior are decoupled so that these two steps can be designed in a modular way. A similar variablesplitting strategy is also adopted in optimization methods such as the Half-Quadratic Splitting (HQS) method [31] and the Alternating Direction Method of Multipliers (ADMM) [30,7]. In fact, SGS can be viewed as a sampling analogue of HQS. SGS is a principled approach to posterior sampling if the two sampling steps are rigorously implemented.

Existing works related to SGS Several works have designed algorithms for solving imaging inverse problems based on SGS [53,19,6,27,78]. The key distinction among these methods lies in their approaches to the prior step. For instance, the works [53,6,27] applied Langevin-based updates for sampling ⇡ X|Z=z such that the prior information is encoded by either traditional regularizers or off-the-shelf image denoisers. The work [19] tackled the prior step by heuristically customizing a diffusion model (i.e. DDPM [33]) for sampling ⇡ X|Z=z . A concurrent work [78] improved the implementation by devising two diffusion processes that rigorously solve the prior step. Our method differs from [78] by connecting the prior step to the EDM formulation [37]. This connection allows us to seamlessly integrate state-of-the-art DMs as expressive image priors for Bayesian inference through a unified interface, eliminating the need for additional customization for each model and leading to better empirical performance. We also note the recent work [43] that adopted the optimizationbased variable-splitting formulation of HQS and utilized general DMs as image priors. We instead considers the SGS formulation from a Bayesian posterior sampling standpoint. Additionally, while SGS-based methods theoretically accommodate general inverse problems, empirical evidence on real-world nonlinear inverse problems remains scarce in the literature. In this work, we demonstrate our method on three nonlinear inverse problems, including a black hole imaging problem. For a more comprehensive review of related works, see Appendix E.

Method

A schematic diagram for the proposed method is shown in Figure 2. Our method, dubbed PnP-DM, builds upon the SGS framework with rigorous implementations of the two sampling steps and an annealing schedule for the coupling parameter ⇢. We start with our implementations of the first step for solving both linear and nonlinear inverse problems.

{⇢k} Likelihood step: enforcing data-consistency Annealing schedule

General reverse diffusion under EDM < l a t e x i t s h a 1 _ b a s e 6 4 = " 7 q I l q 4 X 4 j H r 7 y x 7 C q 6 e c v t 4 p P o k = " > A A A C X n i c b Z H N S g M x F I X T U a t W q 1 U 3 g p v B I r i x z I h U V y K < l a t e x i t s h a 1 _ b a s e 6 4 = " 7 q I l q 4 X 4 j H r 7 y x 7 C q 6 e c v t 4 p P o k = " > A A A C X n i c b Z H N S g M x F I X T U a t W q 1 U 3 g p v B I r i x z I h U V y K Prior step: denoising with noise level < l a t e x i t s h a 1 _ b a s e 6 4 = " c F Y B E x c r t j + I i J N g M p X I J t 9 a e M A = " > A A A C W n i c b Z H L S g M x F I b T 8 T 7 e W n X n Z r A I r s q M S H U l R T c u F W w r 2 F L O p G k b m 8 u Y n F G H M u / g V t 9 M 8 G F M a x e 2 9 k D g z / e f A z l / 4 k R w i 2 H 4 V f C W l l d W 1 9 Y 3 / M 2 t 7 Z 3 d Y m m v Y X V q K K t T L b R 5 i M E y w R W r I 0 f B H h L D Q M a C N e P h 9 d h v v j B j u V b 3 m C W s L a G v e I 9 T Q I c a L T P Q n W G n W A 4 r 4 a S C / y K a i j K Z 1 m 2 n V L h s d T V N J V N I B V j 7 G I U J t k d g k F P B c r + V W p Y A H U K f P T q p Q D L b H k 2 e m w f H j n S D n j b u K A w m 9 O / E C K S 1 m Y x d p w Q c 2 H l v D B d 6 7 F m B M Z A t N G O Z z 9 7 n g U w F c q N f Z y l C n A o w b 7 l / P L M U 9 i 7 a I 6 6 S F J G r / u 9 S v V Q E q I N x z k G X G 0 Z R Z E 4 A N d z l E t A B G K D o f s N 3 i U f z + f 4 X j d N K V K 1 U 7 8 7 K t a t p 9 u v k k B y R E x K R c 1 I j N + S W 1 A k l T + S d f J D P w r f n e R v e 5 m + r V 5 j O 7 J O Z 8 g 5 + A D o 2 t 2 I = < / l a t e x i t > ⇢k start from at noise level < l a t e x i t s h a 1 _ b a s e 6 4 = " s B j + T S l L s N g E n D 3 7 E q w o K 1 l V s n g = " > A A A C W n i c b Z H L T g I x F I b L e A W 8 g L p z M 5 G Q u C I z x q A r Y 3 T j E h O 5 J E D I m V K g 0 s v Y n l E J 4 R 3 c 6 p u Z + D C W y 0 K Q k z T 5 + / 3 n J D 1 / o 1 h w i 0 H w n f I 2 N r e 2 d 3 b T m e z e / s F h L n 9 U s z o x l F W p F t o 0 I r B M c M W q y F G w R m w Y y E i w e j S 8 n / r 1 V 2 Y s 1 + o J R z F r S + g r 3 u M U 0 K F a y w x 0 J + j k C k E p m J X / X 4 Q L U S C L q n T y q Z t W V 9 N E M o V U g L X N M I i x P Q a D n A o 2 y b Q S y 2 K g Q + i z p p M K J L P t 8 e y 5 E 7 / o S N f v a e O O Q n 9 G / 0 6 M Q V o 7 k p H r l I A D u + p N 4 V q P v S g w B k Z r z U h O l u + r Q C Y C u d F v y x Q h S g S Y 9 0 m m u L Q U 9 q 7 b Y 6 7 i B J G r / n y p X i J 8 1 P 4 0 Z 7 / L D a M o R k 4 A N d z l 4 t M B G K D o f i P j E g 9 X 8 / 0 v a h e l s F w q P 1 4 W b u 8 W 2 e + S U 3 J G z k l I r s g t e S A V U i W U P J M P 8 k m + U j + e 5 6 W 9 7 L z V S y 1 m j s l S e S e / x 5 y 3 J w = = < / l a t e x i t > 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 [71]. 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 [37]. 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.

Prior step: denoising via the EDM framework

For the prior step at iteration k, we sample

A closer examination of (6) reveals that this prior step is essentially to draw posterior samples for a Gaussian denoising problem, where the "measurement" is z (k) , the noise level is ⇢, and the prior distribution is p(x) / exp( g(x)).

We tackle this denoising posterior sampling problem within SGS using DMs as image priors. In particular, we leverage the EDM framework [37], which was originally proposed to unify various formulations of DMs for unconditional image generation. To see the connection of the EDM framework to (6), consider a family of mollified distributions p(x; ) given by adding i.i.d Gaussian noise of standard deviation to the prior distribution p(x), i.e. x + ✏ ⇠ p(x; ). The core idea of the EDM framework is that a variety of state-of-the-art DMs can be unified into the following reverse SDE:

where wt is an n-dimensional Wiener process running backward in time, (t) > 0 is a pre-defined noise level schedule with (0) = 0, s(t) is a pre-defined scaling schedule, and ˙ (t), ṡ(t) are their time derivatives. As shown in [37], the defining property of ( 7) is that x t /s(t) ⇠ p(x; (t)) for any time t. Therefore, solving this SDE backward in time allows us to travel from any noise level (t) to the clean image distribution at t = 0. This means that we can use (7) to solve (6) with arbitrary noise level ⇢ as long as ⇢ is within the range of (t). Indeed, the distribution of x 0 conditioned on x t is

We highlight that the last line exactly matches (6) when x t = s(t)z (k) and (t) = ⇢. Therefore, we can naturally design a practical algorithm that samples (6) by following these three steps: (1) find k) , and (3) solve ( 7) backward from t ⇤ to 0 by choosing the discretization time steps and integration scheme. Through this unified interface, any DMs, once converted to the EDM formulation, can be directly turned into a rigorous solver for (6).

Leveraging the connection with EDM, our prior step implementation comes with a large design space that encompasses a variety of existing DMs, such as DDPM (or VP-SDE) [33], VE-SDE [63], and iDDPM [52]. In our experiments, we conduct posterior sampling with all these different models within our framework and all of them provide high-quality samples. The pseudocode of our implementation and more details on the EDM formulation for the prior step is given in Appendix C.2.

Putting it all together

The pseudocode of PnP-DM in complete form is presented in Algorithm 1. PnP-DM alternates between the two sampling steps with an annealing schedule {⇢ k } for the coupling parameter. We find that the annealing schedule on ⇢ accelerates the mixing time of the Markov chain and prevents the algorithm from getting stuck in bad local minima for solving highly ill-posed inverse problems. This is a common practice in both Langevin-based [40,34,65] and SGS-based [6,78] MCMC algorithms to improve the empirical performance in solving inverse problems.

Our work shares some similarities with PnP-SGS [19] but contains three main key differences. First, as demonstrated in our experiments, we investigate three nonlinear inverse problems, while nonlinear inverse problems are beyond the scope of [19]. Our experiments show that PnP-SGS struggles with challenging nonlinear inverse problems such as Fourier phase retrieval. Second, we adopt the EDM formulation to ensure that the prior step of PnP-DM is a rigorous mapping from the image manifold with the desired noise level to the clean image manifold, aligning with the theory of SGS. In contrast, the prior step of PnP-SGS [19] is heuristic (which is also pointed out by [78]) and not rigorously designed to sample (6). Third, unlike PnP-SGS [19] that uses a constant ⇢, we consider an annealing schedule {⇢ k } for the coupling parameter, which is important for highly ill-posed inverse problems.

Algorithm 1 Plug-and-Play Diffusion Models (PnP-DM)

. Section 3.2 4: end for 5: return x (k+1)

⌧

Prior step with error We provide some theoretical insights on the non-asymptotic behavior of PnP-DM. We start with the following definitions.

For two probability measures µ and e µ such that µ ⌧ e µ, the Kullback-Leibler (KL) divergence and Fisher information (or Fisher divergence) of µ with respect to e µ are defined, respectively, as

Both divergences are equal to zero if and only if µ = e µ. KL divergence is a common metric for quantifying the difference of one distribution with respect to another. Fisher information has been used for analyzing the stationarity of sampling algorithms [3,66].

We analyze PnP-DM via a continuous-time perspective, leveraging the interpolation techniques introduced for Langevin Monte Carlo [68,3,66]. We assume that the likelihood step ( 5) can be implemented exactly and the prior step (6) involves running the reverse diffusion process (7) with an approximated score function s t ⇡ r log p t := r log p( • ; (t)). Let ⌫ X 0 be the distribution of the initialization x (0) . Let ⌫ Z k and ⌫ X k+1 be the distributions of z (k) and x (k+1) at the k th iteration. Recall that the stationary distributions are ⇡ X and ⇡ Z . Our analysis is concerned with two continuous-time processes: (1) the non-stationary process from ⌫ X 0 , a non-stationary initialization, to ⌫ X K where ( 7) is run with the approximated score function s t and (2) the stationary process that alternates between stationary distributions ⇡ X and ⇡ Z . These two processes are the interpolation PnP-DM in non-stationary and stationary states and define continuous transitions over discrete iterations. A conceptual illustration of the two processes is provided in Figure 3 with the exact formulations in Appendix A. Now we present our main result:

Define ⌫ ⌧ and ⇡ ⌧ as the distributions at time ⌧ of the non-stationary and stationary process, respectively. Then, for over K iterations of PnP-DM, or equivalently over ⌧ 2 [0, T ] with T := K(t ⇤ + 1), we have 1

where we assume that the score estimation error ✏ score :=

The proof is provided in Appendix A. This theorem states that the average distance (measured by Fisher information) of the non-stationary process with respect to the stationary process over K iterations of PnP-DM goes to zero at a rate of O(1/K) under certain conditions up to the score approximation error. Note that our theory only requires L 2 -accurate score estimate under the measure ⇡ ⌧ , which is a relatively weaker condition than the common L 1 -accurate score estimate

GT prior

Std. deviation ( ) [17].

assumption in prior analysis of sampling methods involving score estimates [5,66]. This result resembles the first-order stationarity for Langevin Monte Carlo [3]. Unlike the non-asymptotic analysis in [78], we utilize the average Fisher information instead of the total variation distance, enabling us to obtain an explicit convergence rate. Here is the infimum of the diffusion coefficient along the reverse diffusion in (7); see further discussions on the role of in Appendix A.3. Our theory shows that the accurate implementations of the two sampling steps lead to a sampler that provably converges to the stationary process that alternates between the two target stationary distributions. 4 Experiments 4.1 Validation with ground truth posterior We first demonstrate the accuracy of PnP-DM for posterior sampling on a simulated compressed sensing problem with a Gaussian prior where the posterior distribution can be expressed in a closed form. The mean and per-pixel standard deviation of the prior are visualized on the bottom left of Figure 4. The linear forward model A 2 R m⇥n is a Gaussian matrix (m = n/2), i.e. A ij ⇠ N (0, 1).

A test image is randomly generated from the prior (see top left of Figure 4), and the measurement is calculated according to (1) with n ⇠ N (0, 0.01 2 I). We compare our method with the popular DM-based method DPS [17]. We draw 1,000 samples and visualize the empirical mean and per-pixel standard deviation for both algorithms. Compared with the true posterior (second column), we find that the both methods accurately estimate the mean. However, the standard deviation image estimated by DPS significantly deviates from the ground truth. In contrast, our standard deviation image matches the ground truth in terms of both absolute magnitude and spatial distribution. These results highlight the accuracy of our method over DPS by taking a more principled Bayesian approach.

Baselines and comparison protocols

We consider four variants of DMs as plug-in priors for our method, namely VP-SDE (VP) [33], VE-SDE (VE) [63], iDDPM [52], and EDM [37]. We compare our method with various baselines, including (1) optimization-based methods: PnP-ADMM [13], DPIR [80]; (2) conditional DMs: DDRM [39], DPS [18]; and (3) SGS-based method: PnP-SGS [19], DPnP [78]. For fair comparison, we use the same pre-trained score function checkpoint for all DM-based methods. Since the pre-trained score function was trained with the DDPM formu-lation (VP-SDE) [33], we convert it to the EDM formulation by applying the VP preconditioning [37]. We use the Peak Signal-to-Noise Ratio (PSNR), the Structural Similarity Index Measure (SSIM), and the Learned Perceptual Image Patch Similarity (LPIPS) distance for quantitative comparison. For each sampling method, we draw 20 randoms samples, calculate their mean, and report the metrics on the mean image. More experimental details are provided in Appendices B, C, D.

0 T + e k 1 q J J 6 T 0 I q + 3 y E q g A P x T X 6

b 2 2 9 2 d q t v 9 9 7 t H 9 T q 7 + 9 d m l s u 2 j x V q e 0 y c E J J I 9 o o U Y l u Z g V o p k S H P X 6 Z 6 p 3 v w j q Z m m 8 4 y U R f w 8 j I R H J A T w 1 q x 8 8 x A 1 v E T B f j s j y f 3 8 Results: linear problems A quantitative comparison is provided in Table 1. PnP-DM generally outperforms the baseline methods and that the VE and EDM variants consistently outperform the other two variants on these linear problems. Figure 5 contains visual examples for the motion deblur problem (see Appendix F.2 for the other two linear problems). PnP-DM provides high-quality reconstructions that are both sharp and consistent with the ground truth image. We also provide an uncertainty quantification analysis based on pixel-wise statistics in Figure 6 PnP-SGS, our approach captures a higher percentage (97.46%) of ground truth pixels than the baselines (96.20% and 88.77%). If the true posterior were truly Gaussian, 99% of the ground-truth pixels should lie within the 3-sigma CI; however, as the posterior is not Gaussian with a DM-based prior, we do not necessarily expect to reach 99% coverage.

Results: nonlinear problems We provide a quantitative comparison in Table 2. For the CDP reconstruction problem, PnP-DM performs on par with DPS but outperforms other SGS-based methods. We then consider the Fourier phase retrieval (FPR) problem, which is known to be a challenging nonlinear inverse problem. One challenge lies in its invariance to 180 rotation, so the posterior distribution have two modes, one with upright images and another with 180 -rotated images, that equally fit the measurement. To increase the chance of getting properly-oriented reconstructions, we run each algorithm with four different random initializations and report the metrics for the best run, following the practice in [18]. We find that PnP-DM significantly outperforms the baselines on this highly ill-posed inverse problem. As shown in Figure 7 (a), our method can provide high-quality reconstructions for both orientations, while the baseline methods fail to capture at least one of the two modes. We further run our method for a test image with 100 different random initialization and collect reconstructions in both orientations that are above 28dB in PSNR (90 out of 100 runs). The percentage of upright and rotated reconstructions are visualized by the pie chart in Figure 7 (b). With a prior on upright face images, our method generate mostly samples with the upright orientation. Nevertheless, it can also find the other mode that has an equal likelihood, demonstrating its ability to capture multi-modal posterior distributions.

Experiments on black hole imaging

Problem setup We finally validate PnP-DM on a real-world nonlinear imaging inverse problem: black hole imaging (BHI) (see Appendix B for more details). A visual illustration of BHI is provided in Figure 8 (a). This BHI inverse problem is severely ill-posed. Even with an Earth-sized telescope, only a small fraction of the Fourier frequencies of the target black hole can be measured (region within the red box); in reality, this region is further subsampled with a highly sparse pattern (black lines). Additionally, the atmospheric noise causes nonlinearity of this BHI problem that sometimes results in a multi-modal posterior distribution of the reconstructed image [64]. Here we demonstrate the effectiveness of PnP-DM in capturing a multi-modal posterior distribution. For brevity, we restrict our choice of diffusion models in PnP-DM to EDM and use DPS as the baseline.

Results on simulated data We use the simulated data from [64] where the measurements are generated assuming that the ground-truth black hole image were at the location of the Sagittarius A ⇤ black hole. Figure 8 (b) visually compares the results obtained by PnP-DM and DPS. We use the t-SNE method [67] to cluster the generated samples (100 for each method) and identify two modes in the samples generated by PnP-DM and three modes in those generated by DPS. We visualize the mean and three samples for each image mode. A metric for quantifying the degree of data mismatch is labeled on the top right corner of each image. As illustrated by both the mean and sample images, PnP-DM successfully captures the two modes previously identified for this dataset [64]. Note that PnP-DM generates high-fidelity samples from both modes with sharp details of the flux ring, and its samples from "Mode 1" align well with the ground truth image. In contrast, two out of the three modes sampled by DPS fail to exhibit a meaningful black hole structure and do not correspond with the observed measurements, as indicated by the significantly larger data mismatch values.

Results on real data Finally, we apply PnP-DM to the real M87 black hole data from April 6 th , 2017 [21], with the results shown in Figure 1. By leveraging an expressive DM-based image prior, PnP-DM generates high-quality samples that are both visually plausible and consistent with the ring diameters observed in the official EHT reconstruction. These results highlight the robustness and effectiveness of our method in tackling a highly ill-posed real-world inverse problem.

Conclusion

We have introduced PnP-DM, a posterior sampling method for solving imaging inverse problems. The backbone of our method is a split Gibbs sampler that iteratively alternates between two steps that separately involve the likelihood and prior. Crucially, we establish a link between the prior step and a general DM framework known as the EDM formulation. By leveraging this connection, we seamlessly integrate a diverse range of state-of-the-art DMs as priors through a unified interface.

Experimental results demonstrate that our method outperforms existing DM-based methods across n n 3 2 + / P i L 1 0 q k 0 m V n r g i F f O t Q x U I e s z P N d c j e J p L R y A n Z G + d i V N j f X D G p u I h P 9 T x h 7 y P q x 9 z j L t W A z p e / t d N 4 w q Q j q c s y + z K l k z t / 8 v P l 1 U F / s L W 9 / X S b D P q b m 8 O d j Q E M d n a 2 N w c 7 Z N g f l J 9 V q / 6 M z x 8 / + s W e C D e N W K z d k C r 1 2 3 C Q 6 P c Z l Z q 7 I c u X 7 F S x h L o X 1 G e Z I 8 S F p o 7 K S f V Z I 4 k U H l P F 0 m n 4 H a w g 5 H p O w C V k q h N J I + W J W L c i n R B M j q B y Q i K q A 6 L m k S P C h a g C m y C I u q I y F j r g s d 9 e x i t > | {z } DPS 3.24 41.17 788.44 Mode 3 19.79 45.79 15.67 5.58 1087.77 550.14 3.24 3.23 6.19 Mode 2 Mode 1 2.60 2.83 2.92 4.19 3.12 2.60 3.22 2.38 both linear and nonlinear inverse problems, including a nonlinear and severely ill-posed black hole interferometric imaging problem.

Limitations PnP-DM can be further improved in the following two aspects. First, PnP-DM currently requires evaluating the likelihood and prior steps for the entire image at a time. This potentially poses computational challenges in solving large-scale inverse problems (e.g. 3D imaging) or those with expensive likelihood evaluation (e.g. PDE inverse problems). Second, the current theoretical analysis does not consider the approximation error introduced in the likelihood step for general nonlinear inverse problems when running Langevin MCMC for finite iterations. Explicit incorporation of this error would offer further insights into the empirical performance of PnP-DM.

Broader impacts We expect this work to make a positive impact in computational imaging and related application domains. For many imaging problems, there is a need to facilitate image reconstruction with expressive image priors and quantify uncertainty, which could lead to better imaging systems that enables further understanding of the imaging target. Nonetheless, as we are introducing DMs as priors into the imaging process, it is inevitable to inherent the potential bias of these models.

A Theory A.1 Interpolation of PnP-DM

In this section, we formally introduce the interpolation of PnP-DM. We consider the case where the coupling strength ⇢ is constant, i.e. ⇢ k ⌘ ⇢ and make the following assumption. Assumption A.1. There exists a unique t ⇤ such that (t ⇤ ) = ⇢.

This assumption is satisfied for common diffusion models. Popular choices of the noise level schedule include (t) = t or (t) = p t, which are monotonically increasing functions of t. We first present two propositions showing that the two steps in SGS can be implemented by running two SDEs. Proposition A.2 (Brownian bridge for the likelihood step). For iteration k with iterate x (k) , the likelihood step of SGS is equivalent to solving the following SDE from t = 0 to t = 1:

where

Proof. This proposition is due to the Brownian bridge construction presented in Lemma 4 of [79]. This SDE satisfies that p(x

. Therefore, solving (9) from t = 0 to t = 1 is equivalent to taking a likelihood step.

Proposition A.3 (EDM reverse diffusion for the prior step). For iteration k with iterate z (k) , the prior step of SGS is equivalent to solving the following SDE from t = t ⇤ to t = 0:

where

, and p t is the distribution of s(t)x + s(t) (t)✏ with x following the prior distribution p(x) / exp( g(x)) and ✏ ⇠ N (0, I).

Proof. First note that ( 10) is exactly (7) written in terms of u(t) and v(t). We know that the (10) is the reverse SDE of the following SDE dx t = u(t)x t dt + v(t)dw t .

(11) where x 0 ⇠ p(x) and p t is the marginal distribution of x t . As we showed in the main text, it holds for (11) that

As (10) is the time-reversed process of (11), they share the same path distribution and thus the same conditional distribution p(x 0 |x t ). So, if we set x t ⇤ = s(t ⇤ )z (k) , we have that

, which is the desired conditional distribution of the prior step. Therefore, solving (10) from t = t ⇤ to t = 0 is equivalent to taking a prior step.

Due to Proposition A.2 and Proposition A.3, the SDEs ( 9) and ( 10) implement the two desired conditional distributions in SGS. In PnP-DM, the prior step involves a network that approximates the score function of the prior distribution, i.e. s t ⇡ r log p t , so the continuous-time process for the actual update is

We can then interpolate PnP-DM by considering a dynamic that alternates between running ( 9) and (12).

Since each likelihood step takes 1 unit of time and each prior step takes t ⇤ unit of time, the total time of the interpolating process for K iterations of PnP-DM is T := K(t ⇤ + 1). We use ⌧ to denote the time that has elapsed from initializing PnP-DM with x (0) . We define {⌫ ⌧ } and {⇡ ⌧ } as the distributions at time ⌧ of the non-stationary process initialized at x (0) ⇠ ⌫ X 0 (Figure 3 top) and the stationary process initialized at x (0) ⇠ ⇡ X (Figure 3 bottom), respectively. Therefore, we have

A.2 Proof of Theorem 3.1

Before proving our main result, we present a key lemma for our analysis, which quantifies the time-derivative of the KL divergence in terms of the Fisher information along a pair of general diffusion processes. Lemma A.4. Given the following pair of diffusion processes dx t = b(x t , t)dt + c(t)dw t (13)

where b(•,

, and c(•) : R ! R. Let µ t be the distribution of x t initialized with x 0 ⇠ µ 0 for (13), and let e µ t be the distribution of e x t initialized with e x 0 ⇠ e µ 0 for (14). Then we have

Proof of Lemma A.4. Writing b(•, t) as b t and e b(•, t) as e b t , by the Fokker-Planck equations of ( 13) and ( 14), we have that

Defining (x) := x log x and 0 (x) = d dx (x) = log x + 1, we can calculate

t µ t e µ t @ t e µ t ◆ + Z ✓ µ t e µ t ◆ @ t e µ t = Z 0 ✓ µ t e µ t ◆ ✓ div ✓ c(t) 2 2 r log µ t b t ◆ µ t µ t e µ t div ✓ c(t) 2 2 r log e µ t e b t ◆ e µ t ◆ + Z ✓ µ t e µ t ◆ div ✓ c(t) 2 2 r log e µ t e b t ◆ e µ t = Z ⌧ r 0 ✓ µ t e µ t ◆ , c(t) 2 2 r log µ t b t µ t + Z ⌧ r  0 ✓ µ t e µ t ◆ µ t e µ t , c(t) 2 2 r log e µ t e b t e µ t Z ⌧ r ✓ µ t e µ t ◆ , c(t) 2 2 r log e µ t e b t e µ t = Z ⌧ r 0 ✓ µ t e µ t ◆ , c(t) 2 2 r log ✓ µ t e µ t ◆ b t + e b t µ t + Z ⌧ r µ t e µ t , c(t) 2 2 r log e µ t e b t 0 ✓ µ t e µ t ◆ e µ t Z ⌧ r µ t e µ t , c(t) 2 2 r log e µ t e b t 0

where we used the fact that 1 2 a 2 ab  1 4 a 2 + b 2 , 8a, b 2 R for the inequality. Now we are ready to prove Theorem 3.1.

Proof of Theorem 3.1. We first consider the likelihood steps over K iterations of PnP-DM. Applying Lemma 2 of [79] to the likelihood steps ( 9) of the non-stationary and stationary processes, we have that

Then, applying Lemma A.4 to the prior steps (12) with

we have that

for ⌧ 2 [k(t ⇤ + 1) + 1, (k + 1)(t ⇤ + 1)] with k = 0, ..., K 1. Integrating both sides over ⌧ 2 [k(t ⇤ + 1) + 1, (k + 1)(t ⇤ + 1)], we get

where

Finally, combining ( 16) and ( 17) for k = 0, ..., K 1, we obtain

The proof is concluded by dividing T = K(t ⇤ + 1) on both sides.

A.3 Discussion

To facilitate the discussion, we first present the following proposition. Proposition A.5. Define a weighting function (⌧ ) over ⌧ 2 [0, T ] such that for k = 0, ..., K 1,

Then, under the same settings of Theorem 3.1, we have

where ✏ score :=

Proof. With the definition of (⌧ ), we can apply Lemma 2 of [79] to the likelihood steps and obtain

for k = 0, ..., K 1. Similarly, we can apply Lemma A.4 to the prior steps and obtain

where

Together, for ⌧ 2 [0, T ]. We can then get (18) by combining ( 19) and ( 20) for k = 0, ..., K 1 and dividing by T := K(t ⇤ + 1).

Unlike Theorem 3.1, this proposition calculates the weighted average of the Fisher information along the two processes with the weighting function (⌧ ). The bound in Theorem 3.1 on the unweighted average of Fisher information can be obtained by further lower-bounding the left hand side of ( 18) using the infimum of (⌧ ) over ⌧ 2 [0, T ]. Given this observation, we can see the role of in Theorem 3.1. With a strictly positive , the weighting function (⌧ ) is always strictly positive, so the (unweighted) average Fisher information must converge to 0. This is precisely the case for the VP-and VE-SDE [63]. On the other hand, if = 0, the Fisher information FI (⇡ ⌧ ||⌫ ⌧ ) may be increasingly large as (⌧ ) gets closer to 0. For iDDPM and EDM, this could happen near t = 0 in the reverse diffusion at v(0) = 0 for these diffusion processes. Nevertheless, we can instead consider a slightly adjusted diffusion coefficient ṽ(t) := v(t) + ✏ with ✏ > 0. Using the relation between scores and diffusions div(pr log p) = p, we get the following reverse SDE which has the same law as (7) at each t:

In this case, ṽ(t) = s(t) p 2 ˙ (t) (t) + ✏ is strictly positive, so the convergence on the unweighted average Fisher information is also guaranteed.

Super-resolution

The forward model is defined as

where P f 2 R n f ⇥n is a matrix that implements a block averaging filter to downscale the images by a factor of f . Specifically, we set f = 4 and used the SVD implementation from the codefoot_2 of [39].

Coded diffraction patterns (CDP) CDP is a measurement model originally proposed in [10]. The target x is illuminated by a coherent source and modulated by a phase mask D. The light field then undergoes the far-field Fraunhofer diffraction and is measured by a standard camera. Mathematically, the forward model of CDP is defined as

y I) where F denotes the 2D Fourier transform. We followed [75] to set D as a diagonal matrix with entries drawn randomly from the complex unit circle.

Fourier phase retrieval We adopted a similar setting as [18]. In particular, the forward model is defined as y ⇠ N (|F P x|, 2 y I), where P denotes the oversampling matrix that effectively pads x in 2D matrix form with zeros. We considered a 4⇥ oversampling ratio for grayscale images of size 256 ⇥ 256, so P x has a size of 512 ⇥ 512.

Black hole imaging We adopted the same BHI setup as in [64,66]. The relationship between the black hole image and each interferometric measurement, or so-called visibility, is given by

where a and b denote a pair of telescopes, t represents the time of measurement acquisition, i is the imaginary unit, and F t a,b (x) is the Fourier component of the image x corresponding to the baseline between telescopes a and b at time t. In practice, there are three main sources of noise in (21): gain error g a and g b at the telescopes, phase error t a and t b , and baseline-based additive white Gaussian noise ⌘ a,b . The gain and phase errors stem from atmospheric turbulence and instrument miscalibration and often cannot be ignored. To correct for these two errors, multiple noisy visibilities can be combined into data products that are invariant to these errors, which are called closure phase and log closure amplitude measurements [12]

where \ computes the angle of a complex number. Given a total of M telescopes, there are in total (M 1)(M 2) 2 closure phase and M (M 3) 2 log closure amplitude measurements at time t, after eliminating repetitive measurements. In our experiments, we used a 9-telescope array (M = 9) from the Event Horizon Telescope (EHT) and constructed the data likelihood term based on these nonlinear closure quantities. Additionally, because the closure quantities do not constrain the total flux (i.e. summation of the pixel values) of the underlying black hole image, we added a constraint on the total flux in the likelihood term. The overall potential function of the likelihood is given by

In this equation, y flux is the total flux of the underlying black hole, which can be accurately measured. We use y := (y cph , y logcamp , y flux ) to denote all the measurements and c, d as the indices for the closure phase and log closure amplitude measurements. Parameters cph , logcamp were given by the telescope system and flux was set to p 2 in our experiments to constrain the total flux. The data mismatch metric reported in Figure 8 is defined as the sum of the reduced 2 values for the closure phase and log closure amplitude measurements, which are calculated using the ehtim.obsdata.Obsdata.chisq function of the ehtim package 4 . Both 2 values should ideally be around 1 for data with high signal-to-noise ratio (SNR). Therefore, a data mismatch value around 2 to 3 is considered as fitting the measurements well.

C Technical details of PnP-DM C.1 Likelihood step

Linear forward model and Gaussian noise As we showed in the main paper, in case of linear forward models and Gaussian noise, the likelihood step is

where

The bottleneck here is that both the mean and the covariance involve the matrix inverse ⇤ 1 , which can be prohibitive to compute directly for high-dimensional problems. Nevertheless, the computational cost can be significantly alleviated when the noise is i.i.d. Gaussian, i.e. ⌃ = 2 y I, and A can be efficiently decomposed. For example, if one can efficiently calculate the SVD of the forward model A, i.e. A = U SV T , one can find the Cholesky decomposition of ⇤ 1 as

Since S is a diagonal matrix, the second term can be calculated with only O(n) complexity. Then, leveraging the property of multivariate Gaussian distribution, we can sample ⌘ ⇠ N (0, I) and calculate z = m(x)+L⌘ as a sample that exactly follows the target Gaussian distribution N (m(x), ⇤ 1 ).

An analogous derivation with Fourier transform can be done when A is a circulant convolution matrix.

Nonlinear forward model We provide the pseudocode of the LMC algorithm for sampling the likelihood step with general differentiable forward models in Algorithm 2.

Algorithm 2 Langevin Monte Carlo for the likelihood step under general A Input: state x, coupling strength ⇢ > 0, likelihood potential f ( • ; y) with measurements y Hyperparameter: step size > 0, number of iterations J > 0

where ✏ j ⇠ N (0, I) 4: end for 5: return u J Table 3 summarizes the hyperparameters we used for solving the nonlinear inverse problems considered in this work.

C.2 Prior step

The EDM framework We formally introduce the EDM formulation [37] using our notations. The forward diffusion process is defined as the following linear Itô SDE

where u(t) : R ! R, v(t) : R ! R are the drift and diffusion coefficients. The generative process is the time-reversed version of (23). According to [2], it is another Itô SDE of the form

where p t (x t ) is the marginal distribution of x t . There also exists a reverse probability flow ODE

which shares the same marginal distributions as (24). Based on (23), we have

where s(t)

s(⇠) 2 d⇠. We also have x t /s(t) ⇠ p(x; (t)) where p(x; (t)) is the distribution obtained by adding i.i.d. Gaussian noise of standard deviation (t)) to the prior data. The idea of the EDM formulation is to write the reverse diffuison directly in terms of the scaling and noise level of x t with respect to x 0 , which are more important than the drift and diffusion coefficients. With the relations between u(t), v(t), p t and s(t), (t), p(•; (t)), we can rewrite ( 24) and ( 25) as

and

Note that ( 26) is precisely the SDE (7) we considered for the prior step. Finally, due to the Tweedie's formula [24], we can approximate r xt log p ( • ; (t)) by a denoiser [D ✓ ( • ; (t)) • ]/ (t) 2 trained to minimize the `2 error of a denoising objective. Substituting the score function by the approximation with the denoiser and using the chain rule, we can further rewrite ( 26) and (27) as

and

Pseudocode We provide the pseudocode for our prior step in Algorithm 3. Note that the update rule is precisely the Euler discretization of ( 28) and (29). The discretization time steps {t i } N i=0 , scaling schedule s(•), and noise schedule (•), are kept the same as in Table 1 of [37]. For all experiments, we set the total number of time steps to 100, i.e. N = 100. We note that this does not imply that each prior step has a number of function evaluations (NFE) equal to 100. Since ⇢ is to a small value as the algorithm runs, the number of steps in later iterations of the algorithm are much fewer than 100. The prior step is similar to the image synthesis process in SDEdit [50] that starts from the middle of the reverse diffusion process. We used the pf-ODE solver for the CDP problem and the SDE solver for all other problems. Part of our code implementation is based on the repository 5 .

Algorithm 3 EDM for the prior step (Bayesian denoising with noise level ⇢) Input: noisy image z 2 R n , assumed noise level ⇢ > 0, pretrained model D ✓ ( • ; • ) that approximates r log p (x; ) with (D ✓ (x; ) x)/ 2 Hyperparameter: discretization time steps {t i } N i=0 (monotonically decreasing to t N = 0), scaling schedule s(•), noise schedule (•), solver (SDE or pf-ODE)

2 if solver is SDE else 1 5:

if i 6 = N 1 and solver is SDE then 8:

end if 10: end for 11: return v N Model checkpoint For experiments with FFHQ color images, we used the pretrained checkpoint from [16] available at the repository 6 . For experiments with synthetic data, FFHQ grayscale images, and black hole images, we trained our own models using the same repository. The model network is based on the U-Net architecture in [52] with BigGAN [9] residual blocks, multi-resolution attention, and multi-head attention with fixed channels per head. See the appendix of [16] for architecture details. Specifically, we changed the input and output channels to 1 and 2, respectively, to accommodate grayscale inputs, and reduced the number of down-pooling and up-pooling levels in the U-Net for smaller images. We trained all models until convergence using an exponential moving average (EMA) rate of 0.9999, 32-bit precision, and the AdamW optimizer [46]. Here is a list of training data we used for each model:

• For the Gaussian prior model, we randomly generated images from the constructed Gaussian prior distribution. • For the FFHQ grayscale model, we used the images with index 01000 to 69999 in the FFHQ dataset. • For the black hole model, we used 3068 simulated black hole images from the GRMHD simulation, which stands for general relativistic magnetohydrodynamic simulation [22]. See Figure 9 for some example training images. We applied data augmentation with random flipping and resizing, so that the flux spin rotation and the ring diameter vary from sample to sample. Preconditioning Since the checkpoints we used are all trained based on the DDPM (or VP-SDE) formulation [33], we converted them to the denoiser D ✓ under the EDM formulation via the VP preconditioning [37]. Specifically, if we denote the pretrained model as F ✓ ( • ; • ), the model we used for Algorithm 3 is

where c skip ( ) = 1, c out ( ) = , c in ( ) = 1/ p 2 + 1, and c noise ( ) = 999 1 VP ( ). Here Connection to DDS-DDPM in [78] A concurrent work [78] introduced a rigorous implementation of the prior step, called DDS-DDPM, by converting the DDPM [33] (or VP-SDE [63]) sampler into a reverse diffusion based on the VE-SDE [63]. The diffusion process after the conversion can be used to solve (6) rigorously by properly choosing the starting point. In fact, our formulation admits DDS-DDPM as a special case with the VP preconditioning and reverse diffusion based on the VE-SDE. Here we explicitly show this connection. For the VE-SDE, we have s VE (t) = 1, VE (t) = p t, u VE (t) = 0, and v VE (t) = 1. So (28) becomes

Applying the VP preconditioning (30) to (31), we obtain

We can then rescale the time range from [0, 1] to [0, 1000], discretize (32) backward in time over the time steps {⌧ t } from [78], and apply the exponential integrator [83] to the drift term, resulting in the following update rule:

Based on the definition

This is exactly the update rule of DDS-DDPM with F ✓ ( • ; t) denoting the noise estimate b ✏ t (•) of DDPM. One can also verify that the initialization in DDS-DDPM is equivalent to ours by checking that ↵t

where ⌘ ⌘ ⇢ is the assumed noise level in (6). As one can see, DDS-DDPM is equivalent to our prior step by choosing the VP-preconditioning, VE reverse diffusion, and a particular integration scheme. In fact, our prior step allows for more general definitions of diffusion processes and includes both the ODE and SDE solvers.

Sample collection

To collect multiple samples using our method, there are two main approaches:

(1) Run a single Markov chain and collect samples after a certain number of iterations, known as the burn-in period, to ensure the chain has converged. (2) Run several independent Markov chains and collect one sample from each chain after convergence. The first approach is more efficient, but the collected samples are not entirely independent and thus may have a small effective sample size. The second approach ensures all samples are fully independent but takes longer to run. In our experiments, we used the first approach for all tests involving 256 ⇥ 256 images to enhance efficiency. Specifically, we set the burn-in period to 40 iterations and collected 20 random samples from the remaining 60 iterations (one every 3 iterations). For other experiments, due to the smaller image sizes, we employed the second approach to obtain fully independent samples.

Compute All experiments were performed on NVIDIA RTX A6000 and A100 GPUs. The runtime per image depends on several factors, such as the choice of GPU, the total number of iterations and the coupling strength schedule {⇢ k } (as it takes more network evaluations for larger ⇢ for our EDM-based denoiser). In our actual experiments, we ran each image for at least 100 iterations to ensure convergence, which took around 1 minute for a single Markov chain. Here we present a comparison of computational efficiency with the major baselines on a linear super-resolution and a nonlinear coded diffraction patterns problem in Table 5. The clock time in seconds and number of function evaluations (NFE) are calculated for each method to measure its computational efficiency. All hyperparameters are kept the same for each method as those used for Table 1 and Table 2 in the manuscript. As expected, DM-based approaches (DDRM & DPS) generally yield shorter runtimes due to their lower NFEs. Nevertheless, our PnP-DM method significantly outperforms these methods while achieving comparable runtimes with DPS (⇡ 1.5⇥), despite its larger NFEs (⇡ 3⇥). This is primarily due to two factors: 1) PnP-DM avoids running the full diffusion process by adapting the starting noise level to ⇢ k at each iteration, and 2) the runtime is further reduced by using an annealing schedule of ⇢ k . We also note that the runtime reported for DDRM and DPS below is the time it takes to generate one sample. For the linear inverse problem experiments, where we generated 20 samples for each sampling method, PnP-DM was faster than DPS because we took 20 samples that PnP-DM generated along one Markov chain of batch size 1 (hence same runtime as below, around 50 seconds) but DPS requires running a diffusion process with batch size 20, which was significantly slower (around 330 seconds).

D Implementation details of baseline methods

PnP-ADMM We set the ADMM penalty parameter as 2 and ran for 500 iterations to ensure convergence. We used the pretrained DnCNN denoiser [81] available at the deepinv libraryfoot_6 .

DPIR We followed the annealing schedule in [80] and ran for 40 iterations. We used the pretrained DRUNet denoiser [82] available at the deepinv library.

DDRM We ran all the experiments with the default parameters: ⌘ B = 1.0, ⌘ = 0.85, and 20 steps for the DDIM sampler [59]. For the Gaussian deblur problem, we used the SVD-based forward model implementation based on separable 1D convolution. We ran it with an additional batch dimension to collect multiple samples.

DPS We followed the original paper to use a 1000-step DDPM sampler backbone. For the linear inverse problems, we used the step size given in [17], i.e. ⇣ 0 = 1. For the nonlinear inverse problems, we optimized the step size ⇣ 0 by performing a grid search, which led to ⇣ 0 = 3 for CDP and Fourier phase retrieval and ⇣ 0 = 0.001 for black hole imaging. For the synthetic prior experiments, we also optimized the step size and used ⇣ 0 = 0.1 for compressed sensing and ⇣ 0 = 1 for Gaussian deblur. We ran it with an additional batch dimension to collect multiple samples.

PnP-SGS We performed a grid search for the coupling parameter ⇢ and found that ⇢ = 0.1 worked the best for all problems. We followed the practice in [19] to have a burn-in period of 20 iterations during which the reverse diffusion is early-stopped. We ran the algorithm for 100 iterations in total and collect 20 samples in the 80 iterations after the burn-in period.

DPnP We implemented the DDS-DDPM sampler for the prior step. For fair comparison, we used the same annealing schedule for the coupling strength (denoted as ⌘ k in [78]) as PnP-DM. We ran it for the same number of iterations for each inverse problem with the same way of collecting samples as our method.

HIO We set = 0.7 and applied both the non-negative constraint and the finite support constraint.

To mitigate the instability of reconstruction depending on initialization, we first repeatedly ran the algorithm with 100 different random initializations and chose the reconstruction that has the best measurement fit. Then we ran another 10,000 iterations with the chosen reconstruction to ensure convergence and report the metrics on the final reconstruction.

E Additional related works

Image reconstruction with plug-and-play priors Plug-and-Play priors (PnP) [69] is an algorithmic framework that leverages off-the-shelf denoisers for solving imaging inverse problems.

Recognizing the equivalence between the proximal operator and finding the maximum a posteriori (MAP) solution to a denoising problem, PnP substitutes the proximal update in many optimization algorithms, such as ADMM [13,56] and HQS [82,80], with generic denoising algorithms, particularly those based on deep learning [49,82,80]. The PnP framework enjoys both convergence guarantees [65,56] and strong empirical performance [76,1] due to its compatibility with stateof-the-art learning-based denoising priors. Recent works have also proposed learning-based PnP frameworks that have direction interpretations from an optimization perspective [20,26]. See [36] for a comprehensive review on the theory and practice of PnP.

Posterior sampling with MCMC and learning priors Learning-based priors have also been considered in the Bayesian context [35], where one seeks to sample the posterior distribution defined under a learned prior. An important technique is denoising score matching (DSM) [70], which connects image denoising with learning the score function of an image distribution. Based on DSM, prior works have incorporated deep denoising priors into MCMC formulations, particularly focusing the Langevin Monte Carlo and its variants as they involve the score function of the target distribution [40,34,41,66,54]. Recently, methods based on SGS have also gained increasing popularity [53,19,6,27,78]. Unlike PnP methods based on optimization, these sampling methods possess the ability to generate diverse solutions and quantify the uncertainty of solution space.

Solving inverse problems with diffusion models

The remarkable performance of diffusion models [33,63] on modelling image distributions makes them desirable choices as images priors for solving inverse problems. One popular approach is to leverage a pretrained unconditional model and modify the reverse diffusion process during inference to enforce data consistency [73,18,17,60,39,61,62,8,44,84,58]. Despite of the promising performance of these methods, they usually involve approximations and empirically driven designs that are hard to justify theoretically and may lead to inconsistent sample distributions. Another line of work learns task-specific models, which achieves higher accuracy at the cost of re-training models for new problems [4,57,44]. Methods based on Particle Filtering and Sequential Monte Carlo are also considered to ensure asymptotic consistency [11,23,74]. Diffusion models have also been considered as a prior for variational inference [28,47] and pluy-and-play image reconstruction [32,48]. In addition to the compressed sensing experiment presented in the main paper, we show another comparison on a Gaussian deblurring problem in Figure 10. Here, the linear forward model A 2 R m⇥n is a 2D convolution matrix with a Gaussian blur kernel of size 7⇥7 and standard deviation 3.0. Similar to the compressed sensing experiment, both methods yield accurate reconstructions of the mean. However, in terms of the posterior standard deviation, DPS exhibits a notable difference from the ground truth, whereas our method achieves a significantly more accurate result.

G Licenses

We list the licenses of all the assets we used in this paper: 2.05 1.82 1.79 2.02 1.96 1.92 2.08 1.98 2.15 1.90 Indices used in teaser: 5, 11, 32, 50, 59, 88 Indices used here: 9, 14, 25, 29, 52, 54, 64, 74, 86, 99 Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We have included the code for reproducing the main experimental results of our work in the supplemental material.

Guidelines:

• The answer NA means that paper does not include experiments requiring code.

• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer.

Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.cc/public/guides/CodeSubmissionPolicy) for more details. • The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes] Justification: Due to the limited space and constraints of time and computational resources, we were unable to include error bars for our experiments over multiple runs. Alternatively, we evaluated our method and baselines based on multiple samples and various metrics. We conducted a uncertainty quantification analysis where we calculated the per-pixel statistics and confidence intervals. We also extensively included visual examples to better demonstrate the performance of the methods we considered in this paper. All of these efforts aimed to better reflect the significance of our improvement over existing methods.

Guidelines:

• The answer NA means that the paper does not include experiments.

• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean.

• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 8. Experiments Compute Resources Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: We have provided details on our compute resources in the appendix. Guidelines: • The answer NA means that the paper does not include experiments. • The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 9. Code Of Ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: We confirm that our research conform with the NeurIPS Code of Ethics. Guidelines: • The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We have added a "Broarder Impacts" section in the main manuscript. Guidelines:

• The answer NA means that there is no societal impact of the work performed.

• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] Justification: Our work does not involve models that have a high risk for misuse.

Guidelines:

• The answer NA means that the paper poses no such risks.

• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have properly credited the owners of all the assets used in the paper. Their licenses are listed in the footnotes and Appendix G.

Guidelines:

• The answer NA means that the paper does not use existing assets.

• The authors should cite the original paper that produced the code package or dataset.

• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.

• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators.