Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors
Authors: Zihui Wu
Yu Sun
Yifan Chen
Bingliang Zhang
Yisong Yue
Katherine Bouman
Venue: NeurIPS 2024 poster (2024)
Keywords: Computational imaging
Inverse problems
Diffusion models
Primary area: machine_learning_for_other_sciences_and_fields
TL;DR: We propose a new method for sampling the posterior distributions of inverse problems using diffusion models in a rigorous way.

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.

Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors
Zihui Wu Yu Sun Yifan Chen Bingliang Zhang' Yisong Yue' Katherine L. Bouman 1,2.3
Department of Computing and Mathematical Sciences, Caltech
Department of Electrical Engineering, Caltech
Department of Astronomy, Caltech Department of Electrical and Computer Engineering, Johns Hopkins University Courant Institute of Mathematical Sciences, New York University
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 the generative process to 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 Markov chain Monte Carlo algo- rithm that performs posterior sampling for general inverse problems by leverage a general DM formulation as a unified interface that allows for rigorously solving the denoising problem with range of state-of-the-arr DMs. We demonstrate the effectiveness of the proposed method on six inverse problems (three linear and three nonlinear), including real-world black hole imaging problem. Experimental results indicate that our proposed method offers more accurate reconstructions and posterior estimation compared existing DM-based imaging inverse methods.
to sampling the posterior of Gaussian denoising problem. Crucially, reducing it
Introduction
Inverse problems arise in many computational imaging applications, where the goal is to recover an image R from set of sparse and noisy measurements y E IRm", The relationship between and y can described by
y A(ar) n,
where A(-): 12 R the forward operator (linear or nonlinear) and the random measure- ment noise in RT Since the sparsity and noisiness of y often lead to significant uncertainty in it is preferable to sample the posterior distribution p(æly) over all possible solutions based on some prior distribution p(æ), rather than finding 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 powerful tool for modeling highly complex image distributions B3] 6s1 Nevertheless, it remains a challenge to turn DMs into reliable imaging inverse solvers, which motivates to develop principled Bayesian method that leverages DMs as priors for posterior sampling.
Diffusion models generate samples from distribution by reversing diffusion process from the target distribution to a simple (usually Gaussian) distribution [33] 63). In particular, it estimates a clean image Co from a noise image ØT by successively denoising noisy images, where æt Pt the intermediate noisy image at time [0, Reversing diffusion requires one to estimate the time-ng gradient log density (score function) V log (arg) along the diffusion process, v log pe (xx (y) in the case of sampling the posterior p(zly).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Real-world imaging system
Posterfor samples
Mean Image
Officil image by EHT
Praposed method
Diffusion image prior
PriM-MM
Experiment performed with dieta MB7 black hole
Figure 1: Demonstration of the proposed method, PnP-DM, for posterior sampling using the real data for the M87 black hole from April 6t, 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 real-world imaging system with expressive image prior in the form of a diffusion model, which was trained with images from the GRMHD black hole simulation [2 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 approximate the time-varying gradient log density V lo (2x ( [1 73
39 588
60,
44
[5
[8]
55
In particular they first apply Bayes' rule to separate the forward operator from an unconditional prior over the intermediate noisy image zz:
log (2)+)) =
+ log log
(2)
By instead aiming to evaluate the right hand side, one can leverage the existing pre-trained DMs for the unconditional term V llog (z,). However, the main challenge in this case is that V log Pe (z,) is intractable to compute in general, as pt (y(xz) involves an integral over all possible 2p 's that could give rise to x; [77. Various methods have been proposed circumvent the intractability and can mostly be categorized into two groups. One group of methods explicitly approximate log Pr (y (ar)) by making simplifying assumptions [62
60 However, even for arguably the finest approximation to date proposed in the recent work 88 is exact only when the prior distribution p(ar) Gaussian. For general prior distributions beyond Gaussian, these methods do not sample the true posterior p(æ) y). The other group of methods do not make explicit approximations but instead substitute V log Pe y (x,) with empirically designed updates where y is treated as a guidance signal 39
[5) 777 18 155. 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 guidance methods, where the generative process is guided towards the regions where the measurement y is more likely to be observed, not posterior sampling methods
773 84
58
44
[8) We also note that some recent work considered combing DMs with Sequential Monte Carlo to ensure asymptotic consistency in posterior sampling but the investigation has been limited to linear imaging inverse problems.
11 23
Our contributions In this work, we pursue a different path towards posterior sampling with DM priors by proposing 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 the Split Gibbs Sampler (71) that alternates between two sampling steps that separately involve the likelihood and prior. While the likelihood step 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 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 black hole interferometric imaging problem involving real data that is both nonlinear and severely ill ios-posed (see Figure T. Overall, PnP-DM outperforms existing bascline methods, achieving higher accuracy in posterior estimation.

---

2 Preliminaries
Split Gibbs Sampler (SGS) is an MCMC approach developed for Bayesian inference [71]. It is also related to the Proximal Sampler 42
and serves as the backbone for the Generative Plug Plug-and-Pul and-Play (GPnP) [6] and Diffusion Plug-and- and-Play (DPnP) [78) frameworks in computational
14
25 79
imaging. The goal of SGS is to sample the posterior distribution
p oc p(y) = exp(- y) g(x))
where (x; log plyla))))))))))))))))))))()))( and g(xx) == 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-spliting strategy and considers sampling an alternative distribution
π ((, z) x exp (-5(z;))- - g(x)) 2pd 1
=
where 2 E RR is an augmented variable and p > 0 is a hyperparameter that controls the strength of the coupling between and z. We denote the æ- and z-marginal distributions of (4) as x (at := J <( z)dz and n² <<(x) (x) := f ,(. z)dæ, respectively. As + 0, π converges to the target posterior p(æ) yy in terms of total variation distance [7T], 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 (o) and, for iteration k- 0, .... K 1, alternates between
1. Likelihood step: sample (k)) n Z1X==(")) (z) exp
(- (-s(2; y) 2/2
2p [1x)(k) =11
2. Prior step: sample (k-1) n xI (x) exp (-9(x) 2/2 1| 2/x) 113)
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 modular way. A similar variable- splitting strategy 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 78). The key distinction among these methods lies in their approaches to the prior step. For instance, the works [53] 6 (2) applied Langevin-base updates for sampling π x 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 A concurrent work 78 improved the implementation by devising two diffusion processes that rigorously solve the prior step. Our method
inverse problems based on SGS [53 [9
6 27
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 optimization- based variablc-splittig formulation of HQS and utilized general DMs as image priors. We instead considers the SGS formulation from 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, Appendix E
3 Method
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 p. We start with our implementations of the first step for solving both lincar and nonlinear inverse problems.
3

---

horation
Itorttion
Fit--
serevong
aiert from 6:45 notaa
strt frn 1 noise Imaet
Prior stop: danolaing with nolon invel Pt.
[oa]
the acheoule
A schematic diagram our method. Our method alternates between likelihood step that enforces data consistency and prior step that solves denoising posterior sampling problem by leveraging the Split Gibbs Sampler [7I]. 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 key connection to a general diffusion model framework called the EDM 37 This connection allows to easily incorporate family of state-of-the-an diffusion models as priors to conduct posterior sampling in a principled way without additional training. Our method demonstrates strong performance variety of linear and nonlinear inverse problems.
Figure 2:
3.1 Likelihood step: enforcing data consistency
For the likelihood step at iteration k we sample
(z) exp (-1(2:))
2p2 ((xe(
=
(S)
Linear forward model and Gaussian noise We first consider simple yet common where the forward model .A is linear and the noise distribution is zero-mean Gaussian, i.e. A := E R x and
(0, Σ). In this case, the potential function of the likelihood term is f(x;; y) 3 ly Aar
(up to an additive constant that does not depend on and y) where || |13 == (., Σ --1.). It is then straightforward to show that
= π = =
where := ATΣ-'AS 0² and m(x) := AA' + - x). The problem of sampling
from Gaussian distributions has been systematically studied [72]. We refer readers to Appendix C.1 for a more detailed discussion.
For general nonlinear inverse problems, the likelihood step is not sampling from Gaussian distribution anymore. Nevertheless, since have access to 2 in closed form up to multiplicative factor, we can use Monte Carlo methods hased on Langevin dynamics to draw samples from it long the likelihood potential is differentiable. Specifically, first set up the following Langevin SDE that admits x X= as the stationary distribution
General case
dæ = V log π z|x
(zt)dt + + 2dw: = - V fz;
- (z x)) dt + V2dw.-
We then initialize the SDE 2g x and run it with Euler discretization. The pseudocode provided in Appendix C.I
3.2
Prior step: denoising via the EDM framework
For the prior step at iteration k. we sample
x (k+1) N X (k) (x) x exp -g(x)
2p |
z 2(() 113)
2x
(6)
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 (A), the noise level is p. and the prior distribution is p(æ) exp(-g(#)).
We tackle this denoising posterior sampling problem within SGS using DMs as image priors. In particular, 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 ( consider a family of mollified distributions p(æ; σt) given by adding i.idd Gaussian noise of standard deviation a to the prior distribution p(a:), i.e. x + de ~ p(x; g). The core idea of the EDM framework that variety of state-of-the-art DMs can be unified into the following reverse SDE:
dx.- = 8(() s(t)
(<R) ((0)) dt s(t) (()))((((())))))
logg
(7)
where Wy 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), s(t) are their time derivatives. As shown in [37]. the defining property of 7 is that 21/s(() ~ p(æ; (t)) ) for any time f. Therefore, solving this SDE backward in time allows to travel from any noise level (t) to
the clean image distribution at t = O. This means that we can use p to solve (6) with arbitrary noise level p as long as p is within the range of a (t). Indeed, the distribution of Zy conditioned 2; is
pl (ra)) p(a)e/=o)p(ar
(N((() aoo (L)2 (1) exp(-g(mo))
exp -g(
(t) no ==/=(()]3
We highlight that the last line exactly matches 6 when T2 = =(t)2(()) and (t) p. Therefore, can naturally design a practical algorithm that samples (6) by following these three steps: (1) find t" such that (t*) = p. (2) initialize at Øt- = =(t))2 and (3) solve 77) 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 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 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
3.3 Putting it all together
The pscudocode of PnP-DM in complete form is presented in Algorithm [ PnP-DM alternates between the two sampling steps with an annealing schedule {px} } 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 a common practice both Langevin-based improve the empirical performance in solving inverse problems.
40
34 651
and SGS-based (6) 78) MCMC algorithms
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 [99]. 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 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 [9] is heuristic (which is also pointed out by (78) and not rigorously designed to sample ( Third, unlike PnP-SGS [19] that uses constant p. we consider annealing schedule {pk} } for the coupling parameter, which is important for highly ill ill-posed inverse problems.
Algorithm 1 Plug-and Play Diffusion Models (PnP-DM)
Input: initialization 2p E I",", total number of iterations K > O, coupling strength schedule {pr > 0)K==' likelihood potential f(: y) with measurements y E IR"", pretrained model
Do(:; that approximates V log (x; o) with (Do(z;o) - x)/a².
1: for k = 0. K do
2: 2(k) f- Likelikkokoee iodeelhoodSteppp!!! Pkr f(:1g)
Section 3.1
3: x (k+1) < PriorStep(# 2 (k), Pk. Do(::.))
>> Section 3.2
4: end for
5: return x (k+1)
5

---

3,4 Theoretical insights
Non process
We provide some theoretical insights on the non asymptotic behavior of PnP-DM. We start with the following definitions. For two probability measures and j such that l & μt, the Kullback-Leibler (KI.) divergence and Fisher information (or Fisher divergence) of with respect to jú are defined, respec- tively, as
Stationary pracess
sap
Pror soop with
Figure 3: conceptual illustration of the non-stationary and stationary time-continoous processes as inter- polations of K discretize iterations of PnP-DM.
(L((llm) := / ulog 1 p
1 and FI(uIlm) := 1 log
2
Both divergences are equal to zero if and only if μL pt. KL divergence 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
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 8, s V log Pt. == V Voggp( (oogp((;ee -ppppooattt :t))) Let vg be the distribution of the
initialization (0). Let of and u be the distributions of (k) and (k++1) at the hthh iteration. Recall that the stationary distributions are x and n² Our analysis is concerned with two continuous-time processes: (1) the non-stationary process from vy a non-stationary initialization, to v where (7 is run with the approximated score function 8, and (2) the stationary process that alternates between stationary distributions x and 7², These two processes are the interpolation PnP-DM in con-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:
Theorem 3.1. Consider running K iterations of PnP-DM with Pk = p > and a score estimate 8, & V Vogg Pt == V log (t)) Let t" > 0 be such that aft)) p and == inf v(t) where ut) == (t) Define Vy and n, as the distributions at time T of the non statiiaaio and stationary process, respectively. Then, for over K iterations of PnP-DM- equivalently over r E [0,T] with ;= K(t + 1), we have
4KL(TT (1/6))
Ac
T (T+l14)) s K(I* + 1) min(p, 8)2
(8)
L + 1)82
convergence from initiolization
souve appresimation error
where we assume that the score estimation error Escore == S +1 "T. Ve--Voogp.. 12d7 < co.
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 (11/K) under certain conditions up to the score approximation error. Note that theory only requires aaccurratee score estimate under the measure 7,, which is a relatively weaker condition than the common L 00 -accurate score estimate assumption in prior analysis of sampling methods involving score estimates (5, (66) This result resem- bles the first-order stationarity for Langevin Monte
Compressed sensing
GT image
GT posterior
DPS
Ours (EDM)
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 77: see further discussions on the role of 8 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 pro- cess that alternates between the two target stationary distributions.
Moun
GT prior
Stil dovaatin (SD)
mess and dec
Figure 4: Results on a synthetic problem with the ground truth posterior available. PnP-DM can sample it more accurately that DPS [7

---

Table 1: Quantitative comparison three noisy linear inverse problems for 100 FFHQ color test
images. Bold: best; Underline: second best.
Method
Canssion dehlur
Mexion dohlur
Super-rsolution (4x)
PSNR (1)
SSIM (1)
LPIPS (3)
PSNR (1)
SSIM (1)
LPIES (4)
PSNR (1)
SSIM (1)
LPIPS (U)
PiI-ADMM
0.7855
0.3472
26.55
0.7655
0.3600
26.61
0.7834
0.3766
DPPR X0
28.74
0.8348
0.2677
20.97
0,8529
0.2404
28.75
0.8378
0.2577
DDRM 39
27.05
0,7819
0.2570
29.47
0,8437
0.2322
0.8035
0.2542
0.2445
28. 98
0.8447
0.2190
0.8349
DPaP [8]
29:21
0.8360
0,2098
30.21
0,8527
0.2010
29.32
0,8407
0.1277
PiP-M
29.46
0.2255
0.2302
30.06
0.8330
0.2069
29.40
0.2219
0.8547
29:57
PaP-MM (DDPM)
29.00
0,833
0.2203
30.26
0.8507
0.2100
29.55
0.8404
0.2213
PaP (EDM)
29.66
0,8411
0,2170
30.35
0.8547
0.2002
29.60
0.8335
0.2191
Messurement
DPIR
DPS
PnP-SGS
Ours: PnP DM (VE)
Ours: PnP DM (EDM)
Ground trush
Figure 5: Visual examples for the motion deblur problem (ay 0.05). We visualize one sample generated by each sampling algorithm.
Experiments
4.1
Validation with ground truth posterior
We first demonstrate the accuracy of PnP-DM for posterior sampling on 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
the bottom left of Figure 4 The linear forward model A E R=*x is a Gaussian matrix (m = n n2/2), i... Aij ~ ~((,, 1). test image is randomly generated from the prior (see top left of Figure 4 and the measurement is calculated according to (T) with ~ (O 0.01²I). We compare our method with the popular DM-based method DPS [7]. 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 terms of both absolute magnitude and spatial distribution. These results highlight the accuracy of our method over DPS by taking more principled Bayesian approach.
4.2 Benchmark experiments
Dataset and inverse problems We test our proposed algorithm and several baseline methods on 100 images from the validation set the FFHQ dataset (38) for five inverse problems: (1) Gaussian deblur with kernel size 61 x61 and standard deviation 3.0, (2) Motion deblur with kernel size 61 x61 and intensity of 0.5, (3) Super-resolution with 4x downsampling ratio, (4) the coded diffraction patterns (CDP) reconstruction problem (nonlinear) in 1o 51l (phase retrieval with phase mask), and (5) the Fourier phase retrieval (nonlinear) with 4x oversampling. We add i.i.d. Gaussian noise to all the simulated measurements y. In particular, i.e. for Fourier phase retrieval, the noise standard deviation is set as dy = 0.05. Due to the severe ill- ill-posedness of Fourier phase retrieval. we consider smaller noise standard deviation Jy 0.01.
(00 o For all problems except
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
methods: PnP-ADMM [3], DPIR 80 (2) conditional DMs: DDRM [39]. DPS [18] and (3) SGS- SGS-based method: PnP-SGS 19). DPnP [78). For fair comparison, we use the same pre-trained score function checkpoint for
our method with various baselines, including (1) optimization-
all DM-based methods. Since the pre-trained score function was trained with the DDPM formu-
7

---

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 com- parison. 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
Alx eevor ((F zil Std. deviction (std)
1
/2od
Se Voc std
Results: linear problems A quantitative com-
parison is provided in Table [ PnP-DM gener- ally outperforms the baseline methods and that the VE and EDM variants consistently outper- form the other two variants on these linear prob-
3,80 %
DPS
91 200 %
lems. Figure [] contains visual examples for the motion deblur problem (see Appendix F.2 for the other two linear problems). PnP-DM
PPPDM PhP--SS
provides high-quality reconstructions that are both sharp and consistent with the ground truth image. We also provide an uncertainty quantifi-
cation analysis based on rixel-wise statistics in Figure 6 In the left three columns, we visualize the absolute error (i z/), standard deviation (std), and absolute z-score (|r =l/std). In the third column, red pixels highlight locations where the ground truth pixel values are outliers of the s-iimma credible interval (CI) under the estimated posterior uncertainty. The fourth col- umn contains scatter plots of [ Ø versus std for each pixel of the reconstructions, where red boxes show the percentages of outliers (outside of 3-sigma CT) and gray boxes indicate the per- centages within the 3-sigma CL. Similar to the synthetic prior experiment, DPS tends to have larger standard deviation estimations, as shown by the less concentrated distribution of gray points around the origin. Compared with baselines, especially PnP-SGS, our approach captures a higher percentage (97.46%) ) of ground truth pixels than the base- lines (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 DM-based prior, we do not necessarily expect to reach 99% coverage.
57 -7.466 95
Figure 6: Comparison of uncertainty quantifica- tion (UQ) for the motion deblur. Left 3 columns: absolute error (a= æ1), standard deviation (std), and absolute z-score (- - x /std) with the outlier pixels in red. Right column: scatter plot of | x versus std. Note that PnP-DM leads to better UQ performance than the baselines by having the low- est percentage of outliers while avoiding having overestimated per-pixel standard deviations.
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 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 1809-rotated images, that equally fit the measurement. To increase the chance of getting properly-oriented reconstructions, 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 posed inverse problem. As shown in Figure (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 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 (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-mode posterior distributions.
4.3 Experiments on black hole imaging
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
Problem setup
8

---

Table 2: Quantitative evaluation two noisy nonlinear inverse problems for 100 FFHQ grayscale test images. Bold: best; Underline: second best.
Method
Coded Codedd difliccionn pateerss
Fourier pluse netrivval
PSNR (t)
SSIM (1)
LPIPS (1)
PSNR (1)
SSIM (1)
LPIPS (L)
HIO
20.66
0,4308
0.6669
DPS 1z
33.43
0 0.9049
0.1374
23.60 60
0.6804
0.3126
PaP-SGS [1]
92.19
0.8889
0.2010
15.36
0,3659
0,5750
0,8653
PaP (VP)
32,99
0,8846
0.1906
30.36
0.8553
0.2115
PuP-DM (VE)
35.13
0.8971
0.1663
29.88
0.8464
0.2186
PPP-UM (IDPPM)
33:35
0,9083
0.1471
30.11
0.1975
(EDM)
0,9050
0.1386
31.14
0.2024
(a)
DPS
PnP-SGS
Ours: PnP (EDM)
Ground truth
(b)
Rotated
Fourier PR
Roratee
81.1%
Upright
Figure 7: Results of the Fourier phase retrieval problem. (a) PnP-DM provides both upright and rotated reconstructions (two modes given by the invariance of the forward model to 180 rotation) with high fidelity, while the baseline methods cannot. (b) We visualize the percentages of upright and rotated reconstructions out of 90 runs for test image with two samples for each orientation.
in Figure 8 (a). This BHI inverse problem is severely ill ill-posed. Even with an 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 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.
telescope,
Results on simulated data We use the simulated data from 664] 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 1-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. 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 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 6th 2017 211, with the results shown in Figure [ 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 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 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 state-of-the-an DMs as priors through a unified interface. Experimental results demonstrate that method outperforms existing DM-based methods across

---

(a)
Graund truth
(b)
Mode
Mode
Made
Made
Mode 3
Data mismatchh
3.12
4.19
3.24
41.17
788. 44
Mean image
Fourier spectrum
2.60
2.60
3.24
19.79
5.58
2.83
45.79
1007.7
cestenir samples
Zoom-in
2.92
filly sampled low-lee))
Ours: PnP-DM (EDM)
DPS
Figure 8: Results a nonlinear and severely ill-posed black hole imaging problem. Our method, PnP-DM, is compared with the conditional diffusion model baseline DPS. A metric quantifying the mismatch with the observed measurements is labeled for each sample. which should be around for ideal measurement fit. Samples generated by PnP-DM exhibit two distinct modes with sharp details and consistent ring structure, while samples given by DPS display inconsistent ring sizes and sometimes fail to capture the black hole structure entirely with samples having poor measurement fit.
both linear and nonlinear inverse problems, including 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 (c.g. 3D imaging) or those with expensive likelihood evaluation (c.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 need facilitate image struction with expressive image priors and quantify uncertainty, which could lead to better imaging systems that enables further understanding of the imaging target. Nonetheless, DMs as priors into the imaging process, it is inevitable to inherent the potential bias of these models.
are introducing
Acknowledgments
The authors thank Charles Gammie, Ben Prather, Abhishek Joshi, Vedant Dhruv, and Chi-kwan Chan for providing the black hole simulations. The authors also thank the generous funding from Schmidt Sciences and the Heritage Medical Research Fellowship. Z.W W. was supported by an Amazon Al4Science Fellowship. Y.S. was supported by a Computing, Data, and Society Fellowship. B.Z. was supported by Kortschak Fellowship.
10

---

References
[1] Rizwan Ahmad, Charles A. Bouman, Gregery T. Buzzard, Stanley Chan, Sizhuo Liu, Edward T. Reehorst, and Philip Schniter. Plug-and-play methods for magnetic resonance imaging: Using denoisers for image recovery. IEEE Signal Processing Magazine, 37(1):105-116, 2020.
[2] Brian. D. O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applica- tions, 12:313-322 -26,, 1982.
[3] Krishna Balasubramanian, Sinho Chewi, Murat A Erdogdu, Adil Salim, and Shunshi Zhang. Towards a theory of non-lop non-log-conean sampling:firs--ore Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 2896-2923, PMLR, 02-05 Jul 2022.
stationarity guarantees for langevin monte carlo, In
[4] Georgios Batzolis, Jan Stanczuk, generation with score-based diffusion models. ArXiv, abs/2111.1300l 2021.
Schonlich, and Christian Ermann. Conditional image
[5] Valentin De Bortoli, James Thorton, Jeremy Heng, and Arnau Doucet. Diffusion schrödinger bridge with applications cosee generative modeling. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[6] Charles A. Bouman and Gregery T. Buzzard. Generative plug and play: Posterior sampling for inverse problems, 2023.
[7] Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends in Machine Learning, 3:1-122, 01 2011.
[8] Benjamin Boys, Mark Girolami, Jakiw Pidstrigach, Sebastian Reich, Alan Mosca, and Deniz Akyildiz. Tweedie moment projected diffusions for inverse problems, 2023.
[9] Andrew Brock, Jeff Donahuc, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019.
[10] Emmanuel Candès, Xiaodong Li, and Mahdi Soltanolkotabi. Phase retrieval from coded diffraction patterns. Applied and Computational Harmonic Analysis, 39, 10 2013.
[11] Gabriel Victorino Cardoso, Yazid Janati, Sylvain Le Corff, and Éric Moulines. Monte carlo guided diffusion for bayesian linear inverse problems, ArXiv, abs/2308 009883, 2023.
[12] Andrew Chael, Michael Johnson, Katherine Bouman, Lindy L. Blackburn, Kazunori Akiyama, and Ramesh Narayan. Interferometric imaging directly with closure phases and closure amplitudes. The Astrophysical Journal, 857(1): (1)3, apr 2018.
[13] Stanley Chan, Xiran Wang. and Omar Elgendy. Plug-and and-play admm for image restoration: Fixed point convergence and applications. IEEE Transactions on Computational Imaging, PP, 05 2016.
[14] Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for a proximal algorithm for sampling. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 2984-300 3014. PMLR, 05 Jul 2022.
[15] Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Conditioning method for denoising diffusion probabilistic models. 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 14347-14356, 2021.
[16] Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon. Perception prioritized training of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 1472-11481, June 2022.
[17] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Representctions, 2023.
[18] Hyungjin Chung and Jong Chul Yc. Score-based diffusion models for accelcrated mri. Medical Image Analysis, page 102479, 2022.
[19] Florentin Coeurdoux, Nicolas Dobigeon, and Pierre Chainais. Plug-and and-play split gibbs sampler: embed- ding deep generative priors in bayesian inference, 2023.

---

[20] Regev Cohen, Yochai Blau, Daniel Freedman, and Ehud Rivlin. It has potential: Gradient-driven denoisers for convergent solutions inverse problems. In Neural Information Processing Systems, 2021.
[21] The Event Horizon Telescope Collaboration. First m87 event horizon telescope results. iv, imaging the central supermassive black hole. The Astrophysical Journal Letters, 875(1):L4, apr 2019.
[22] The Event Horizon Telescope Collaboration. First m87 event horizon telescope results. v. physical origin of the asymmetric ring. The Astrophysical Journal Letters, 875(1):LS5, apr 2019.
[23] Zehao Dou and Yang Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations, 2024.
[24] Bradlcy Efron, Tweedic's 's formula and selection bias. Journal of the American Statistical Association, 106:1602-166 1614, 12 2011.
[25] Jiaojiao Fan, Bo Yuan, and Yongxin Chen. Improved dimension dependence of a proximal algorithm for sampling. In Gergely Neu and Lorenzo Rosasco, editors. Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 1473- PMLR, 12-15 Jul 2023.
[26] Zhenghan Fang, Sam Buchanan, and Jeremias Sulam. What's in prior? learned proximal networks for inverse problems. The Twelfth International Conference on Learning Representations, 2024.
[27] Elhadji C. Faye, Mame Diarra Fall, and Nicolas Dobigeon. Regularization by denoising: Bayesian model and langevin-within-se gibbs sampling, 2024.
[28] Berthy T. Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L. Bouman, and William T. Freeman. Score- Score-based diffusion models as principled priors for inverse imaging. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), pages 10520 10531, October 2023.
[29] James R. Fienup. Phase retrieval algorithms: comparison. Applied optics, 21 15:2758-69 -9, 1982.
[30] Daniel Gabay and Bertrand Mercier. A dual algorithm for the solution of nonlinear variational problems via finite element approximation. Computers & Mathematics With Applications, 2:17-40, 40, 1976.
[31] Donald Geman and Chengda Yang. Nonlincar image recovery with hulf-quadriticc regularization. IEEE transactions on image processing: a publication of the IEEE Signal Processing Sociery, 4 7:932-46, 1995.
[32] Alexandros Graikos, Nikolay Malkin, Nehojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and- play priors. ArXiv, abs/2206.09912,, 2022.
[33] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle. M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1840-6851. Curran Associates, Inc., 2020.
[34] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir. Robust compressed sensing mri with deep generative priors. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, volume 34, pages 44938-14954 Curran Associates, Inc., 2021.
[35] Zahra Kadkhodaic and Eero Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in a denoiser. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[36] Ulugbek Kamilov, Charles Bouman, Gregery Buzzard, and Brendt Wohlberg. Plug-and-play methods for integrating physical and leamed models in computational imaging: Theory, algorithms, and applications. IEEE Signal Processing Magazine, 40:85-97, 01 2023.
[37] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion feffusioo-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[38] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. 2019 JEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 4396-4405. 2018
[39] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems, 2022.
12

---

[40] Bahjat Kawar, Gregory Vaksman, and Michael Elad. SNIPS: Solving noisy inverse problems stochastically. Advances in Neural Information Processing Systems, 34:21757-2176,, 2021.
[41] Rémi Laumont, Valentin De Bortoli, Andrés Almansa, Julie Delon, Alain Durmus, and Marcelo Pereyra. Bayesian imaging using plug & play priors: When langevin meets tweedie. SIAM Journal on Imaging Sciences, 15(2)7701-377, 2022.
[42] Yin Tat Lee, Ruoqi Shen, and Kevin Tian. Structured logconcave sampling with a restricted gaussian oracle. In Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages
PMLR, 15 19 Aug 2021.
[43] Xiang Li, Soo Min Kwon, Ismail R. Alkhouri, Saiprasad Ravishankar, and Qing Qu Decoupled data consistency with diffusion purification for image restoration, 2024.
[44] Jiaming Liu, Rushil Anirudh, Jayaraman J. Thiagarajan, Stewart He, Aditya Mohan, Ulugbek Kamilov, and Hyojin Kim. Dolce: A model-based probabilistic diffusion framework for limited-angle et reconstruction. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 10498 10508, October 2023.
[45] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In P Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[46] Ilya Loshchilov Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[47] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective solving inverse problems with diffusion models. In The Twelfth International Conference on Learning Representations. 2024.
[48] Ségolène Martin, Anne Gagneux, Paul Hagemann, and Gabriele Steidl. Pnp-flow: Plug-and-play image restoration with flow matching, 2024.
[49] Tim Meinhardt, Michael Möller, Caner Hazirbas, and Daniel Cremers. Learning proximal operators: Using denoising networks for regularizing inverse imaging problems. 2017 IEEE International Conference on Computer Vision (ICCV), pages 1799-1808, 2017.
[50] Chenlin Meng. Yutong He, Yang Song, Jiaming Song. Jiajun Wu. Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Reprexentations, 2021.
[51] Christopher A. Metzler, Philip Schniter, Ashok Veeraraghavan, and Richard Baraniuk. prdeep: Robust phase retrieval with a flexible deep network. In International Conference Machine Learning. 2018
[52] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 8162-817 -171. PMLR,, 18 24 2021.
[53] Marcelo Pereyra, Luis A. Vargas- Mieles, and Konstantinos C. Zygalakis. The split gibbs sampler revisited: Improvements its algorithmic structure and augmented target distribution. SIAM Journal Imaging Sciences, 16(4):2040 2014000-7711--0 2023. 2023.
[54] B. Remy, F. Lanusse, N. Jeffrey, J. Liu, Jean-Luc Starck, K. Osato, and T. Schrabback. Probabilistic mass mapsing with neural score estimation. Astronomy & &Astrophysics, 672, 12 2022.
[55] Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis. Alex Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffiusion models. In Thiriy- seventh Conference Neural Information Processing Systems, 2023.
[56] Ernest K. Ryu, Jialin Liu, Sicheng Wang. Xiaohan Chen, Zhangyang Wang. and Wotao Yin. Plug-and and-play methods provably converge with properly trained denoisers. In International Conference on Machine Learning, 2019.
[57] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence. PP:1-14, 09 2022.
[58] Bowen Song. Soo Min Kwon, Zecheng Zhang. Xinyu Hu, Qing Qu, and Liyue Shen. Solving inverse problems with latent diffusion models via hard data consistency. In The Tweifth International Conference Learning Representations, 2024.
13

---

[59] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[60] Jiaming Song, Arash Vahdat, Morteza Mardami, and Jan Kautz. Pscudoinverse-guiddd diffusion models for inverse problems. In International Conference on Learnnng Representations, 2023.
[61] Jiamning Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and and-play controllable generation. In Proceedings of the 40th International Conference Machine Learning, ICML'23, JMLR.org. 2023.
[62] Yang Song, Liyue Shen, Lei Xing. and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations, 2022.
[63] Yang Song, Jascha Sihl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[64] He Sun and Katherine L. Bouman. Deep probabilistic imaging: Uncertainty quantification and multi -oodal solution characterization for computational imaging. In AAAI Conference on Artificial Intellignnce (AAAI), 2021.
[65] Yu Sun, Brendt Wohlberg, and Ulugbek Kamilov, An online plug-and-play algorithm for regularized image reconstruction. IEEE Transactions on Computational Imaging, PP:1-1, 01 2019.
[66] Yu Sun, Zihui Wu, Yifan Chen, Berthy T. Feng, and Katherine L. Bouman. Provable probabilistic imaging using cosre-based generative priors. IEEE Transactions on Computational Imaging, 10:1290-1300 1305, 2024.
[67] Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Rexearch, 9(86):2579-22 2600,,,,,.66,..2. 2008
[68] Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Neural Information Processing Systems, 2019.
[69] Singanallur V. Venkatakrishnan, Charles A. Bouman, and Brendt Wohlberg. Plug-and-play priors for model based reconstruction. In 2013 IEEE Global Conference on Signal and Information Processing. pages 945 948, 2013.
[70] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 13:161111 2011.
[71] Maxime Vono, Nicolas Dobigeon, and Pierre Chainais. Split-and-augmeette gibbs sampler-applicatienn to large-scale inference problems. IEEE Transactions on Signal Processing, (7(6):1645 1661,16611112100 2019.
[72] Maxime Vono, Nicolas Dobigeon, and Pierre Chainais. High-dimensional gaussian sampling: A review and unifying approach based stochastic proximal point algorithm. SIAM Review, 64(1):3-56, 2022.
[73] Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. The Eleventh International Conference on Learning Representations, 2023.
[74] Luhuan Wu. Brian L Trippe, Christian A. Naesseth, David M Blei. and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. arXiv preprint MXXX::3306.17755 2023.
[75] Zihui Wu, Sun, Jiaming Liu, and Ulughek Kamilov. Online regularization by denoising with applications to phase retrieval. In 2019 IEEF/CVF International Conference on Computer Vision Workshop (ICCWW), pages 3887-3895, 2019.
[76] Zihui Wu, Yu Sun, Alex Matlock, Jiaming Liu, Lci Tian, and Ulugbek S. Kamilov, Simba: Scalable inversion in optical tomography using deep denoising priors. JEEE Journal of Selected Topics in Signal Processing, 14(6):1163 1175, 2020.
[77] Bin Xia, Yulun Zhang, Shiyin Wang, Yitong Wang. Xinglong Wu, Yapeng Tian, Wenming Yang, and Luc Van Gool. Diffir: Efficient diffusion model for image restoration. ICCV, 2023.
[78] Xingyu Xu and Yuejie Chi. Provably robust score-based diffusion posterior sampling for plug-and-plar image reconstruction, 2024.
[79] Bo Yuan, Jiaojiao Fan, Jiaming Liang, Andre Wibisono, and Yongxin Chen. On a class of gibbs sampling over networks. In Gergely Neu and Lorenzo Rosasco, editors, Proceedings of Thirtyy Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 5754-5780. PMLR, 12- 15 Jul 2023.
14

---

[80] Kai Zhang, Yawei Li, Wangmeng Zuo, Lei Zhang, Luc Gool, and Radu Timofte. Plug-and-plan image restoration with deep denoiser prior. IEEE Transactions Pattern Analysis and Machine Intelligence. PP:1-1, 06 2021.
[81] Kai Zhang. Wangmeng Zuo. Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond Gaussian denoiser: Residual learning of deep CNN for image denoising. IEEE Transactions on Image Processing. 26(7):3142- 3155, 2017.
[82] Kai Zhang, Wangmeng Zuo, Shuhang Gu, and Lei Zhang. Learning deep enn denoiser prior for image restoration. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 2808 2817, 2017.
[83] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator: In The Eleventh International Conference on Learning Representations, 2023.
[84] Yuanzhi Zhu, Kai Zhang. Jingyun Liang Jiezhang Cao, Bihan Wen, Radu Timofte, and Luc Van Gool. Denoising diffusion models for plug-and-play image restoration. In JEEE Conference on Computer Vision and Partern Recognition Workshops (NTIRE), 2023.
15

---

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 p is constant, i.e. Pk = p and make the following assumption. Assumption A.I. There exists a unique 1* such that ((()) =
This assumption is satisfied for common diffusion models. Popular choices of the noise level schedule include (t) = or (t) = vt. 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 & with iterate x**), the likelihood step of SGS is equivalent to solving the following from t to
dæ; p²V log de dt + pdw,
(9)
where zy = x (k) and (r(x) := /eepp= f z: y)
x = dz.
Proof. This proposition due to the Brownian bridge construction presented in Lemma of [79..
This SDE satisfies that p(z1 exp (-f 2p² x| - =113)). Therefore, solving
1
from f == 0 to t = 1 is equivalent to taking likelihood step.
Proposition A.3 (EDM reverse diffusion for the prior step). For iteration k with iterate (8)), the prior step of SGS equivalent to solving the following SDE from
to
dæ; [u(l)ar, (1)) logpt Pr (a:+] di + v(()dw.
(10)
where Ex- (#)) u(t)
s(t) v(t) == st)) (2) and pt is the distribution of s(t)
a(t)
(t)) with x following the prior distribution p(æ) x exp(-g(x)) and ~ (O,
Proof. First note that (10) is exactly 7 written in terms of u(t) and v()). We know that the (10 is
the reverse SDE the following SDE
dæ; u(t) + v(t)du;;
(11)
where To p(z)) and Pt is the marginal distribution of æ,, As we showed in the main text, it holds for that
p( (xe) exp (9(20)) 20 1| st/s(()133
As (10) is the time-reversed process of [1]. they share the same path distribution and thus the same conditional distribution p( p(zplaz (z,). So, if we set 2,- =((t) s(t*)2((e)))) (1)), we have that
p( m x exp
(g(()))
2a (t=))² -00 20x) x exp -g(#x)
2p2 ar 2xx)
1
which is the desired conditional distribution of the prior step. Therefore, solving (T)) from = 1" to t == 0 is equivalent taking 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 DM, the prior step involves network that approximates the score function of the prior distribution, i.e. 8, Re V log Pt, so the continuous-time process for the actual update is
dæ: [u(()a; v(()3s, 8, (a:)) di v()ddw;
(12)
We can then interpolate PnP-DM by considering dynamic that alternates between running ( and (12).
Since each likelihood step takes unit of time and each prior step takes [* unit 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 ()) We define {ur} and { } as the distributions at time 7 of the non-stationary process initialized at (0) ~ rg (Figure top) stationary process initialized at (0) ~ xx (Figure bottom), respectively. Therefore, we have
top) and the
16

---

* k/ TT = Λ for = k(t" + 1) with à = 0,... K, and
X.
V = =T for r = k(t + 1) + with & 0,... K - 1.
A.2 Proof of Theorem 3.1
Before proving our main result, we present key lemma for 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
da b( t)dt + c(C)du,, dă b(ã,, t) + c(t)dw:;
(13)
(14)
where b(.,)) x b(., IR" x I + R", and cl.) : R. + R. Let μ4 be the distribution
of Φ, initialized with zu po for 13 and let pi be the distribution of कn initialized with Bo po for (1).. Then we have
c(t)2
(Hs (Tte + c(t)² / 1ī. be mr.
a.KL(Heliid) - 4
(15)
Proof of LemmaaA.4. Writing b(.tt) be and b(., t) be. by the Fokker- equations of and 14, we have that
13
dep = div [
((t)2
log ue - bz) b1)(²)) mr and 0.m = div [( V
V log - 5z p
2
Defining (5xx)= == x: log x and d'(x)
de (aa) = log + 1, we can calculate
d
Pe lite
(Mr [it) a,
Mt
He Oepi Oepec
+ ff0
H1 1+t
Orj
p
iile
c(()2
[c(1)
F
div
Voo ut be
14
Mt div
V log j4:
i36
Tee
2
2
c((t)2
+
div
V log pre
The
14c
2
c(1)
+ /
c(1)²
V
14
V
Vlog 14+
b2
Nt
V
Vog Tix
Tha
2
T
2
ct))
V
log
be 1/6
P1
2
c(12
c(12
V
Voog
U4
b
+ de
+
Me
V log Jt+
Φ
Tha
2
P
2
i
V
Mt
c(t)2
V
log Th
be Φ
ct))
#1
Pt /
He
V
log
v
log
br
2
14
c(1)2
/
14
/
17 b/
12 #1
V log
M1 + c(1)²
4
14
c(1)2
/
/ |7, b 14
log
14t
+ c(t)2
4
Jh
c(t)2
Fl (l
+
[/2 - by 1|2 Ftz
a(t)2
where we used the fact that
Ja²2 ab < aaa2 + b², Va, E R for the inequality.
17

---

Now we are ready to prove Theorem 3.1
We first consider the likelihood steps over iterations of PnP-DM. Applying Lemma 2 of 779 to the likelihood steps (9) of the non-stationary and stationary processes, we have
Proof of Theorem 3.7
that
0² 0 0,KL(T+|lu5) = FI( [14)) -
0²
2
4
for r E [R<L6- + +), + 1) + 1] with h: = =,,.. K Integrating both sides over T E [k(L"*
1),(((* + +1) + 11. we get
kel +1)+1
4|KL(T 11 KL ( [f]
(T+
(16)
+()++1)
p²
for k 0,
Then, applying Lemma A.4
to the prior steps 12 with
b( t) = u(t) rtt)² v log P: (x)))
b(2,, t) == u(t) u(t)2s, (xx,)
c(t) v(t)
8 = inf v()),
AE(0.t*)
we have that
u(T)2
1
/ ||0(7) (s, V pr pr)
0.KL(T+ $
Fl (m=112>) + u(T)2
Vlogp>>
u 2
Fl [1<=||11 (1+) + n(r)2 / ||8, Vlog
S
8²
$ Fl (Tr|| ( ||8, V 1ogpr112.
for E [k((t + 1) + (k: 1)((" + 1)] with k 0,... K 1. Integrating both sides over
T E [k(L" +1) + (k 1)(( (1" + 1)], we get
(++1(0+1))
4 4|KL(mZ||u
KL( x 1 (u++)) )
4c 82
JXC-40441
Fl ( dr
(17)
82
where
(++1)(+++)
EE, ||8, - V log pr dr = ( 1'+11 2(T) ||sT - V log Porg--ll 112.7. dr.
Excore
4(1+11++1
Finally, combining
(6) and 17 for h: 0, ..... K 1, we obtain
4[KL(T |1 KL(T* 11x21
4K Souor
FI ( (.) dr
min(p, 8)2
82
4KL(T> (1/²))
+ 4K
min 8)2
82
The proof is concluded by dividing = K(L" + 1) on both sides.
A.3 Discussion
To facilitate the discussion, we first present the following proposition. Proposition A.5. Define a weighting function >(T) over 7 E [0, T] such that for h: = O,... K 1,
A(F) S02 is E [k(("* + 1), k(t + 1) + 1), (0(7)²
if [k(t" + + (k: + 1) + 1)]-
18

---

Then, under the same settings of Theorem 3.7 we have
1
/ X(T) (1,) dr 4KL(T 110)
4c
(18)
T where Excore ;= S1* u(T) |s, V log 1Bdd.
K(t + +1
+1
Proof. With the definition of A(T), we can apply Lemma 2 of (79 to the likelihood steps and obtain
+1+++
(1r X()) (Tr |u)) dr 4[KL(T |
KL (1
(19)
(((+1)
for 1- = O, ...,.. K - 1. Similarly, we can apply Lemma A.4 to the prior steps and obtain
(1++)0++1))
1((++)++1 (r) dr <4|LL( KL(x 1 + Aescore
(20)
where
+1+1)(-++1
u(T)?E,, ||,, Vlogp- dr f
(++1 / u ||st - V Vlog Plogp-ll P-112 dr.
Cecore
+(1++1)+1
Together, for [0, TJ. We can then get 18 by combining dividing by := K (t* 1).
199 and 20 for k 0,... K and
Unlike Theorem 3.T this proposition calculates the weighted average of the Fisher information along
the two processes with the weighting function X(r). 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 X(T) over E [0, 1". Given this observation, we can see the role of ò in Theorem 3.1 With a strictly positive ó, the weighting function A(T) 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 8 O, the Fisher information Fl (T> [12)) may be increasingly large as A(T) gets closer to 0. For iDDMM and EDM, this could happen near t = 0 in the reverse diffusion at (0) = 0 for these diffusion processes. Nevertheless, we can instead consider a slightly adjusted diffusion coefficient ū(()= == v(() + with > O. Using the relation between scores and diffusions div log Ap, we get the following reverse SDE which has the same law as at each t:
dæ;
s(t) a(1)
Lt
2
2s(1))(()0()) v log
It (()) di+ (s(()(0)2)()))) + dw.
In this case, ū(t) = s(t) (t) + is strictly positive, so the convergence on the unweighted average Fisher information is also guaranteed.
B Inverse problem setup
Data usage
We list the data we have used for our experiments:
For the synthetic prior experiment, we took images from the CelebA dataset 45). turned them into grayscale, rescaled them to [-1, 1], and resized them to 32 x 32 pixels for efficient computation. We then found the empirical mean and covariance the images to construct the Gaussian image prior. The test image was randomly drawn from this Gaussian prior.
For the benchmark experiments, we used the first 100 images (index 00000 to 00099) in the FFHQ dataset 38). For all linear inverse problems, the test images were in RGB and normalized to range [- 1, 1]. For all nonlinear problems, the test images were in grayscale and normalized to range [0, 1].
used the simulated data used in [64] and the publicly available EHT 2017 date that was used to produce the first image of the M87 black hole.
For the black hole experiments,
horizon- wwwtrooooeeeeeeopee.e..oo www.eecoopee11--oooseeeee 2017-observations
19

---

Gaussian and motion deblur The forward model is defined as
N (Bc, o
where B 12 is circulant matrix that effectively implements convolution with kernel &s under the circular boundary condition. For the Gaussian deblurring problem, we fixed the kernel as a Gaussian kernel with standard deviation 3.0 and size 61 x 61. For the motion deblurring problem, we randomly generated the kernel & for each test image using the code with with intensity of 0.5 and size 61 61. For fair comparison, the blur kernel for each test image set the same for all compared methods.
Super-resolution
The forward model is defined
y ~ N(Pp2.00)1)
where Pj E IR xm is matrix that implements a block averaging filter to downscale the images by factor of f. Specifically, we set f and used the SVD implementation from the code of [39
Coded diffraction patterns (CDP) CDP is measurement model originally proposed in [T]. The target a illuminated by a coherent source and modulated by a phase mask D. The light field then undergoes the far far-field Fraunhofer diffraction and is measured by standard camera. Mathematically, the forward model of CDP is defined as
y N(IFDa)
where F denotes the 2D Fourier transform. We followed 75
D diagonal matrix with
entries drawn randomly from the complex unit circle.
Fourier phase retrieval We adopted similar setting defined as
18
In particular, the forward model is
y ~ (((FP))
where denotes the oversampling matrix that effectively pads æ in 2D matrix form with zeros. We considered 4x oversampling ratio for grayscale images of size 256 x 256, so Pæ 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
Vt_
-999 (00 -5) F (T, + Ja,b E CC,
(21)
where a and b denote a pair of telescopes, t represents the time of measurement acquisition, i is the imaginary unit, and F is the Fourier component of the image corresponding to the baseline between telescopes and b time In practice, there are three main sources of noise in (2T): gain error Ia and 9b at the telescopes, phase error di and of. and baseline- additive white Gaussian noise 1), 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 these errors, which called closure phase and log closure amplitude measurements
cph (t,ba.).)) <(Va, VV,.VV.aVa,e
Apph t,(a,b,e) (z))
logramp Vt.(a.b.e.d)
log
A logcamp ((a.bbe,,)) (æ),
|a.cl| VV.
where 7 computes the angle of a complex number. Given a total of M telescopes, there are in total (M-1)(M-2) closure phase and M(M-3) log closure amplitude measurements at time t, after eliminating repetitive measurements. In our experiments, we used a 9-telescope array (M 9) from
(license unknown)
bttps//gtthibbcco/bbbjaaaaaattt kww.r//ddm (MIT license)
20

---

calculated using the 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 constraint on the total flux the likelihood term. The overall potential function of the likelihood is given by
||.A,c (x)
Mt.c
cph
2
|ogcamm 1.d
(2) 4z.d
logcamp
Σ, a. y lux
f(;; y)
(22)
20 cph
20 logcamp
20
f,d
In this equation, y is the total flux of the underlying black hole, which can be accurately measured. We use y/ = (z/oph g g fllx) to denote all the measurements and c, as the indices for the closure phase and log closure amplitude measurements. Parameters deph. Ologeamp were given by the telescope system and duux The data mismatch metric reported in Figure g is defined as the sum of the reduced x² values for the closure phase and log closure amplitude measurements, which ehtim. obsdata. Obsdata. chiss function of the ehtim package") Both x² values should ideally be around for data with high signal-to-noise ratio (SNR). Therefore, data mismatch value around 2 to 3 is considered as fitting the measurements well.
set to v2 in experiments to constrain the total flux.
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
= n 2====== (
where A = -- I and ((x) = A- } z). The bottleneck here is that both
the mean and the covariance involve the matrix inverse A-1, which can be prohibitive to compute directly for high-dimensional problems. Nevertheless, the computational cost be significantly alleviated when the noise is i.i.d. Gaussian, i.e. Σ = o and A can be efficiently decomposed. For
example, if one can efficiently calculate the SVD of the forward model A, i.e. A = USVF, one can find the Cholesky decomposition of -1 as
A LL" where L: (
1 +2++111 7 11)
1/2
Since S is diagonal matrix, the second term can calculated with only (nn) complexity. Then, leveraging the property of multivariate Gaussian distribution, can sample r ~ A/(0, I) and calcu- late m(x)+++ + Ln as sample that exactly follows the target Gaussian distribution A(m (m (x), A-). An analogous derivation with Fourier transform can be done when A is 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
Algorithm 2
Langevin Monte Carlo for the likelihood step under general A
Input: state æ, coupling strength p likelihood potential f(:iy) with measurements y Hyperparameter: step size 7 > 0, number of iterations >
1: Mg <
2: for j 0,
do
3: Uj+1 + Uj - (u, x) /27E, where Ej ~(0,(
4: end for
5: return UJ
Table summarizes the hyperparameters used for solving the nonlinear inverse problems consid- ered in this work.
https//gghhhooccmmcchhhtttttttti (GPL-3.0 license)
21

---

Table 3: List of hyperparameters for the likelihood step of PnP-DM
Inverse problem
Step (t)
Number of iterations (J)
Coded diffraction pallerss
1.0e3
100
Fourier phase retrieval
1.0e4 Black bole imaging
100
1.0-55
200
C.2 Prior step
The EDM framework We formally introduce the EDM formulation 37 forward diffusion process is defined as the following linear Itō SDE
using our notations.
The
dæ;
u())a,dl + =(()dum;
(23)
where u(() : IR + R, v((): R R are the drift and diffusion coefficients. The generative process is the time-reversed version of 23 According to (2) is another Itô SDE of the form
dæy [(() "11) log dt + =())du,
(24)
where (a,)) is the marginal distribution of æt. There also exists reverse probability flow ODE
[((t)x: vt)) lo dt,
dæ.
(25)
which shares the same marginal distributions 24 Based on (33 have
p(a, (20) =N((s(() s(1) (1)
where s() exp (Io w(()d)) and o1)) =- So u(E) क(द) dg. We also have ar/s(() ~ plæ; (())
where p(2; (t)) is the distribution obtained by adding i.i.d. Gaussian noise of standard deviation (t)) the prior data. The idea of the EDM formulation to write the reverse diffuison directly in terms of the scaling and noise level of a; with respect to Co, which are more important than the drift and diffusion coefficients. With the relations between u(t), (t). Pt and st), (t), p(: (t)). we can
rewrite (24) and (25) as
s()) ((1)
Lt. ((1)) dt s(t) (6 26(()2 logp
dæ.-
(26)
and
[t()) 8c - ((1))(()) logp [x(
Zt x(1) (0) 0
dæ;
dt.
(27)
Note that 26 is precisely the SDE we considered for the prior step. Finally, due to the Tweedie's formula (24), we can approximate V æ, log (:: by denoiser ( ()) 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,
further rewrite (26 and (7 as
(( (20(0) (t) s(t))
$(1)
20(()s(1) Do o(t)
( xc ((0)):(0))) dt +tt(t) ((t))/2(())) (28)
dx.
and
dæ; - [[(o()= 8(1)
((t)s(t) o(t)
( x a(t)) s(t)
Do
dt.
(29)
s(t)
Pseudocode We provide the pseudocode for our prior step in Algorithm 3 Note that the update rule precisely the Euler discretization of (28) and (29). The discretization time steps [ scaling schedule s-)), and noise schedule (-). are kept the same as in Table of [37]. For all experiments, we the total number of time steps 100, i.c. N = 100. We note that this does not imply that each prior step has a number of function evaluations (NFE) equal to 100. Since p is to small value as the algorithm runs, the number of steps 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
(Creative Commons Attributio-Nonnomonnnn NonCommerciallSharre ShareAlike 4.0
International Public license)
22

---

Algorithm 3 EDM for the prior step (Bayesian denoising with noise level p)
Input: noisy image
R", assumed noise level
O, pretrained model Dø +3 that approxi-
mates V log p (z; or) with (Do(æ; σ) - x
Hyperparameter: discretization time steps [t (monotonically decreasing to in = 0), scaling schedule s(-), noise schedule (-), solver (SDE or pf-ODE)
i" < smallest such that (t;;) P
Find the starting point of the reverse diffusion > Initialize at time 4.-
2: ei- t a(1;-)z
3: for = -i"," ,..... NN do
4:
À +-2 2 solver is SDE else 1
d: f Ad())) k(t))
(o(C.)((t.) Do 377,3 (4i) (t:))
5:
o(4)) (t,)
6:
U--1 + + (t++1 1;
D Drift
7:
if + N and solver is SDE then
8:
U1+1 1t+1 + ((t))/0(()))()) t411) where E, ~~,((0,1)
Difffusion
9:
end if
end for
11: return
Model checkpoint For experiments with FFHQ color images, we used the pretrained checkpoint from [16] available at the repository [ 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 [ residual blocks, multi-resolution attention, and multi-head attention with fixed channels per head. See the appendix of [16] for architecture details. Specifically, changed the input and output channels 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:
and 2, respectively, to accommodate
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 mugnetolydrodynamic simulation 22) See
Figure for some example training images. We applied data augmentation with random flipping and resizing, that the flux spin rotation and the ring diameter vary from sample to sample.
Figure Example images from the dataset for training the black hole diffusion model prior.
Preconditioning Since the checkpoints used are all trained based the DDPM (or VP-SDE)
formulation [33 we converted them to the denoiser Dg under the EDM formulation via the VP
preconditioning 37 Specifically, if denote the pretrained model Fo ; .). the model we used for Algorithm is
Do(#; 0) == Caip(() (o)x + Cour(a) (a) Fo (Cin( )az; Craise (A)),
(30)
where Caip (0) = 1, Com (0) = -0, (a) 1/ 1, and Croise (a) = (p Here op (-)
is the inverse of the VP-SDE noise schedule defined as (v Ve 2ot2
with 3g 19.9
and Bmin 0.1. This adaption allows us make fair comparison with other DM-based methods using the same pretrained models. One can also incorporate DMs trained with other formulations into PnP-DM by properly setting the preconditionng parameters.
'bttps//github..om/yyhhht 1itp:hhhi11/////////geeee weighting (MIT license)
23

---

Connection to DDS-DDPM in 78
A concurrent work 78 introduced rigorous implementation by converting the DDPM B3 (or VP-SDE (3] sampler into
of the prior step, called
reverse diffusion based on the VE-SDE 63) The diffusion process after the conversion can be used 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 SvE(t) = 1, OVE vt, becomes
(vE(t) 0, and (VE(t) 1. So 28
1 Do (r vii de + drū;
dæ
at
(31)
t
Applying the VP preconditioning
(30 (3I.. we obtain
dæ.
Fe
9990v
(v dt + dee,
(32)
Vi+1
We can then rescale the time range from [0, 1] to [0, 1000), discretize (32) backward time over the time steps {Te} } from (78). and apply the exponential integrator 83 to the drift term, resulting in the following update rule:
- = .. 2(VT= T(-1) Fo
xt
ovp (VTr)) + Tr= T--1 where e <~ ~~00 I)
VT: Ti
Based on the definition Ti == a - 1 1= (t)( in DDS-DDPM, we get
: = -; - 2 1) Fo (Varūr:) + V Tr- where ~ I)
(33)
This exactly the update rule of DDS-DDPM with :o(-;t) denoting the noise estimate a-)) of
DDPM. One can also verify that the initialization in DDS-DDPM is equivalent to ours by checking that 0i 2 भनम equivalent to T. = (t) S 77 where r) = p is the assumed noise level in (6). As one can see, DDS-DDPH is equivalent to our prior step by choosing the VP-preconditioning, VE reverse diffusion, and particular integration scheme. In fact, our prior step allows for more general definitions of diffusion processes and includes both the ODE and SDE solvers.
C.3 Others
Annealing schedule for p In this work, we considered an exponential annealing schedule for the coupling strength p. We note that the schedule can be more general than exponential decay and we leave the investigation of other decay schedules in future work. Specifically, we specified a starting level Po. decay rate a, and a minimum value Pmin. Then we set
Pk max(a" po, Pmin)
for k 0, ... K Table summarizes the annealing hyperparameters that we used for all the inverse problems considered in this work.
Table 4: List of hyperparameters for the annealing schedule of p in PnP-DM
Inverse problem
Starting level (po)
Minimam level (Peis)
Decay rate (a)
Synthetic prior experiments
0.03
0.03
Geussian deblur
10
0,3
0.9
Motion deblur
0.3
0.9
Super resolution
10
0.3
0.9
Coded diffraction patterns
0.9
Fourier phase retrieval
0.1
0.9
Black hole imaging
0.02
0.93
Initialization For the linear inverse problems, we used the zero initialization, i.c. x = 0 e R". For the CDP and Fourier phase retrieval problems, we used the Gaussian initialization, i.c. (0))) ~ (O, I). For black hole imaging experiments, we used the uniform random initialization between 0 and 1 for each pixel. We found that PnP-DM, as an MCMC algorithm, is insensitive to the initialization. Except for the black hole experiments where we found the negative values would cause problems, any reasonable initialization would lead to comparable results. This observation corroborates convergence result, which holds for any initialization 1/.
24

---

Number of iterations We ran 500 iterations for the synthetic prior experiments, 200 iterations for the black hole experiments, and 100 for all other experiments. The numbers were chosen so that the algorithm was fully converged.
Sample collection To collect multiple samples using our method, there are two main approaches: (I) 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 small effective sample size. The second approach ensures all samples are fully independent but takes longer to run. In our experiments, used the first approach for all tests involving 256 x 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 A 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 {px} } (as it takes more network evaluations for larger p for our EDM-based denoiser). In our actual experiments, we ran each image for at least 100 iterations to ensure convergence, which took around I minute for a single Markov chain. Here we present comparison of computational efficiency with the major baselines on a linear super-resolution and a nonlinear coded diffraction patterns problem in Table 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 I and Table Sableein in the manuscript. As expected, DM-based approaches (DDRM & DPS) generally yield shorter runtimes due to their lower NFBs. Nevertheless, our PnP-DM method significantly outperforms these methods while achieving comparable runtimes with DPS ( 1.5x), despite its larger NFEs (Rs 3x). This is primarily due to two factors: 1) PnP-DM avoids rumning the full diffusion process by adapting the starting noise level to Pk at each iteration, and 2) the runtime is further reduced by using an annealing schedule of Pk- 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 (hence same runtime as below, around 50 seconds) but DPS requires running diffusion process with batch size 20, which was significantly slower (around 330 seconds).
Table 5: Comparison of computational efficiency between PnP-DM and other baseline methods
Inverse problem
DDRM
DPS
PGP-SGS
DPuP
PnP- (ours)
Super-resolution
Clock time (s)
0.4
39
20
322
NFE
20
1000
1030
18372
3032
Clock time (s)
37
54
261
50
Coded diffraction patterns
D Implementation details of baseline methods
PmP-ADMM
We set the ADMM penalty parameter as 2 and ran for 500 iterations to ensure convergence. We used the pretrained DrCNN denoiser [8 available at the deepinv library"
DPIR We followed the annealing schedule in 88 and ran for 40 iterations. We used the pretrained DRUNet denoiser [22 available at the deepinv library.
DDRM We ran all the experiments with the default parameters: na 1.0, 77 = 0.85, and 20 steps for the DDIM sampler [59]. For the Gaussian deblur problem, we used the -VDebased forward model implementation based on separable 1D convolution. We ran with an additional batch dimension to collect multiple samples.
httgss//ithhb.ccm//deee///eeeege (BSD- BOSD-3-Clause license)
25

---

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 [7]. i.c. <' 1. For the nonlinear inverse problems, we optimized the step size Ç by performing grid search, which led to Ç" 3 for CDP and Fourier phase retrieval and Ç" = 0.001 for black hole imaging. For the synthetic prior experiments, we also optimized the step size and used Ç" 0.1 for compressed sensing and Ç = 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 p= 0.1 worked the best for all problems. We followed the practice in [19] to have 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 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 3 = 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 metries on the final reconstruction.
E
Additional related works
Image reconstruction with plug-and-play priors Plug-and-Play priors (PnP) [69] is an algo- rithmic 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 denoising problem, PnP substitutes the proximal update in many optimization algorithms, such as ADMM 13 56) and HQS 82
80), with generic denoising algorithms, par- The PnP framework enjoys both convergence guarantees [655] (56) and strong empirical performance [76, [ due to its compatibility with state- of-the-art learning-basee based denoising priors. Recent works have also proposed learning-based PnP frameworks that have direction interpretations from an optimization perspective (20 26
ticularly those based on deep learning 49
82
80
See (36
for a comprehensive review on the theory and practice of PnP.
Posterior sampling with MCMC and learning priors Learning- bcaseing-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 they involve the score function of the target Recently, methods based on SGS have also gained increasing 78 Unlike PnP methods based on optimization, these sampling methods
distribution 40
34
41
66
54].
popularity 5 19
6
27
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 pretrained unconditional model and modify
the reverse diffusion process during inference to enforce data consistency 773 18 17 60
39 61]
62 g
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
84
58
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
I1 23 74.. Diffusion models have also been considered as a prior for variational inference (28 47
and pluy-and-play image reconstruction
32 48
26

---

Additional experimental results
F.1 Synthetic prior experiment
In addition to the compressed sensing exper- iment presented in the main paper, show another comparison on a Gaussian deblurring problem in Figure I0 Here, the linear forward model A Rm** is a 2D convolution matrix with Gaussian blur kernel of size 7 7 and stan- dard deviation 3.0. Similar to the compressed sensing experiment, both methods yield accu- rate reconstructions of the mean. However, in
Gaussian deblur
mage
posterior
DPS
Ours (EDM)
Mean
GT prior
Std. devittioon (S0)
terms of the posterior standard deviation, DPS exhibits notable difference from the ground truth, whereas our method achieves signifi- cantly more accurate result.
Figure 10: Comparison of method and DPS 17] on estimating the posterior distribution of Gaussian deblurring problem under Gaussian prior. While the mean estimations of the two meth- ods are of roughly the same quality, our approach provides a much more accurate estimation of the posterior per-pixel standard deviation than DPS.
F.2 Linear inverse problems
We provide visual comparisons for the Gaus- sian deblur and super-resolution problems in Figure
Meesurement
DPIR
DPS
PnP-SGS
Ours: PnP-DM (VE)
Ground truch
Gousion deblur
Superreeslutoon
Figure 11: Visual comparison between our method and baselines on solving the Gaussian deblurring and super-resolution problems with i.i.d. Gaussian noise (o, 0.05). We visualize one sample generated by each algorithm.
Additional visual examples are provided in Figure 12 (Gaussian deblurring), Figure (motion deblurring), and Figure (superrresollution).
F.3 Nonlinear inverse problems
We provide visual comparisons for the CDP reconstruction problem in Figure 15 where visualize one sample for each method. As shown by the red zoom-in boxes, PnP-DM can recover fine-grained features such the hair threads that are missing in the reconstructions by the baselines. Additional reconstruction examples are given in Figure 16 for the CDP reconstruction problem.
We then show some additional reconstruction examples with comparison to DPS in Figure 17 For each method, we visualize best reconstruction out of four runs for each test image according to the PSNR value. While DPS failed on around half of the test images, our proposed method provided high-fidelity reconstructions on almost all test images. This comparison highlights the better robustness of our method over DPS.
27

---

Messurement
Messurement
Ground
Figure 12: Additional visual examples for the Gaussian deblurring problem.
Messurement
(VE)
Ours (EDM)
truth
Meesunement
Ours (VE)
Ours (FDM)
Ground truth
Pors
Pors
Porsi
Figure 13: Additional visual examples for the motion deblurring problem.
E4
Black hole imaging
Finally, we present visual examples from the black hole imaging experiments. Samples generated by PnP-DM (EDM) and DPS using the simulated data are shown in Figure 18 with the data mismatch metric labeled at the top right corner of each sample. Consistent with the results in Figure X DPS can only capture one of the two posterior modes. DPS samples from Mode 2 and Mode significantly deviate from the measurements and lack the expected black hole structure. contrast, PnP-DM successfully samples both posterior modes and consistently produces samples that fit the measurements well. Additionally, Figure presents more samples obtained by applying PnP-DM to the real M87 black hole data. The generated samples are not only diverse but also fit the measurements
28

---

Messurement
tuuhh
Messuremment
DVEL
S
Figure 14: Additional visual examples for the 4x super-resolution problem.
PaP-5G5
DPaP
Ours: PwR-DM (DDPM)
PnP-DM (EDM)
Ground
Codad diffraction
Figure 15: Visual comparison between method and baselines solving the coded diffraction pattern (CDP) reconstruction problems with i.i.d. Gaussian noise (ay 0.05). We visualize one sample generated by each algorithm.
well with data mismatch values around These samples exhibit a ring diameter consistent with the official EHT reconstruction in Figure T and share common bright spot location at the lower half of the ring.
F.5 Further analysis
Sensitivity analysis on the annealing schedule {pt} Figure 20 we present a sensitivity analysis the annealing schedule {Pk]. In particular, show the PSNR curves of æ& with different exponential decay rates (left) and minimum coupling levels Pmin (right) for one linear (super-resolution) and nonlinear (coded diffraction patterns) problem. We have the following conclusions based on the results. First, different decay rates lead to different rates convergence, which corroborates with our theoretical insights that P plays the same role as the step size. The final level of PSNR is not sensitive to different decay rates, as all curves converge to the same level. Second, Pmin decreases, the final PSNR becomes higher. This is expected because the stationary distribution of the Et π should converge to the true target posterior, p( a (y), as decreases.
we show some visual examples of intermediate Fa and iterates (left) and convergence plots of PSNR, SSIM, and LPIPS for æg (right) on the super-resolution problem. Ps decreases, a4 becomes closer to the ground truth and ø gets less noisy. Both the visual quality and metric curves stabilize after the minimum
Convergence curves with intermediate visual examples Figure 21
29

---

(iDPMM)
buthh
(CDPM)
Ground
100
103
MOO
Figure 16: Additional visual examples for the Fourier phase retrieval problem.
DPS
(DDPM)
Ground truth
DP5
(DDPM)
Ground
Figure 17: Additional visual examples for the Fourier phase retrieval problem.
coupling strength achieved. Despite being run for 100 iterations in total, our method generates good images in around 40 iterations, which is around 30 seconds and 1600 NFEs.
G Licenses
We list the licenses of all the assets we used in this paper:
Data
CelebA 45]:
Unknown
30

---

Mode
Mode 2
Mode
Mode
Mode 3
3.08
2.83
4.03
21.28
842.39
2.43
3.84
3.23
39.95
505.21
2.91
3.02
4.47
37
654.74
Ours: PnP-DM (EDM)
DPS
Figure 18: Additional visual examples given by PnP-DM and DPS using the simulated black hole data.
2.15
1.82
1.79
1.90
2.08
1.92
2.02
1.96
Figure 19: Additional visual examples given by PnP-DM using the real M87 black hole data.
Creative Commons BY-NC-SA 4.0 license Simulated and real black hole data: Unknown
FFHQ
Code
The license for each code repository that have used is listed in the footnote after the repository link.
Pretrained model
FFHQ model by Choi al. 16 MIT license
31

---

decay rate
minimum coupling strength Puiin
Supar-nesalution
Coded diffraction patterna
Supar-naiolution
Coded diffraction patterns
32
32
PSNR (o8)
P3NR (0B)
0.1
Paiin
Pain 0.5
0.50
0.96
Paine 0.9
Iteration 10)
100
0
lteration (4)
100
Iteration (A)
100
0
Iteration (
100
Figure 20: Sensitivity analysis on the annealing schedule Ps. with different decay rates a (left)
and minimum coupling strength Pmin (right) for linear (suner-resolution) and nonlinear (coded diffraction patterns) inverse problem. Recall from Appendix C.3 that Pk = max(o" (ax Pos Pmii), where
Po 10 for this experiment.
27g
Convergence
PONR 690
SSIM
Inttllzztion
IPPS
30
210
Z50
230
Z
Figure 21: Visual examples of intermediate 24 and 2 iterates (left) and convergence plots of PSNR, SSIM, and LPIPS for I iterates (right) on the super-resolution problem. The vertical dashed lines show the iterations which the æ and z iterates visualized.
32

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have made sure that the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We have added a paragraph on the limitations of our method to the conclusion
section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiitiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
33

---

Answer: [Yes]
We have presented our theoretical result in the main paper with the proof in the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We have provided all the necessary details to reproduce the main results in our paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
34

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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 NeurlPS code and data submission guidelines (https://nips.cm cc/
public/guides http://oooeeeaeossnooooooooooooe
for more details.
While we encourage the release of code and data, we understand that this might 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results, See the NeurIPS code and data submission guidelines (https //nip.cc. wttp:ccccppcccccceeeeec gww /eodedeee/ooeebbiieeeeee caatimissinooooiiiieico
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 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 supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have provided all the training and test details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 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 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.
35

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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 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 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 figures symmetric error bars that would yield results that are out range (e.g. negative 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have provided details 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 compute required for each of the individual experimental runs as well 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 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
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations 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 "Broarder Impacts" section in the main manuscript.
Guidelines:
36

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt 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 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 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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 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 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
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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
37

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.coo//aaaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of 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 s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We have documented the code we provide in the supplement.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Our research does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Our research does involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
38

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
39


Reviews:
[{"review_id": "ZHQDRKyq3A", "reviewer": "Reviewer_u5tS", "summary": "In this paper, the authors treat the problem of sampling from posterior of image inverse problems. Their formulation is based on the Split Gibbs Sampler (SGS) algorithm, which alternates between sampling from Moreau regularized versions of the prior and of the likelihood. The main contribution of the paper is to use the EDM formulation of diffusion models for sampling the prior term. The EDM formulation was \nproposed to unify various diffusion models in a common general formulation.", "questions": "- In Figure 4, the diffusion process is represented in a closed-form. Did you train a denoiser, or did you use the closed-form score instead?\n\n- How did you tune the hyperparameters for your method and for each method in the comparison? Did you use their default hyperparameters? For example, DPIR provides hyperparameters fine-tuned for various inverse problems. However, the SR x4 task is not one of the inverse problems for which DPIR's authors fine-tuned their hyperparameters, likely making the default hyperparameters sub-optimal for this task.\n\n- Did you employ any regularization parameters to balance the data-fidelity and regularization terms?\n\n- For the backward diffusion process, did you use the same number of iterations for all values of  $\\rho_k$ ? Specifically, is the step size of the diffusion process consistent, or is it adjusted according to the value of $\\rho_k$ ?", "limitations": "No limitations.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Firstly the paper is well-written, provides clear and concise explanations of the methods and contributions. The authors have also included some theoretical insights that guarantee the convergence of each sampling step of the algorithm to the right stationary distributions when the number of iterations tends to infinity.  \n\nThe strength of the paper is mainly its experimental section, which covers a range of both linear and non-linear inverse problems. Additionally, I find interesting the inclusion of a toy experiment with a Gaussian prior, which allows comparing to ground-truth posterior distributions.", "weaknesses": "1. **Significance of Contribution**: My primary concern is the overall significance of the contribution. When compared to the works of Coeurdoux et al. (2023) and Xu & Chi (2024), the advancements appear to be relatively minor. Specifically, the contribution seems to involve using a more general formulation of diffusion models (EDM) instead of the more specific DDPM or DDIM. Are there additional conceptual differences? Furthermore, given this similarity, what accounts for the significant performance disparity with these methods?\n\n2. **Algorithm Efficiency**: The algorithms seem to require considerable computational time, necessitating the diffusion process to be run at each iteration. While this is mathematically sensible, it appears excessive. Given the annealing process chosen, how many calls to the denoiser are typically made in practice? Additionally, how does this number compare to methods that perform a single diffusion using an approximation of p(y|x_k), such as DPS? Including a comparative table in the paper would be valuable.\n\n3. **Theoretical Results**: The theoretical results presented, while interesting, are somewhat limited as they assume a fixed rho parameter, which is not the case in practical applications. Additionally, the algorithm is presumed to run with the true score, which is also not realistic. In addition, although it may be beyond the scope of the current paper, providing theoretical insights into the convergence of the proposed Gibbs Sampler (with approximate score) to the true posterior would be beneficial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your efforts on reviewing our paper. Below we provide point-by-point responses to your comments:\n\n> Weakness 1\n\nWe believe that our work has several significant contributions over the existing works [1] and [2]. Moreover, we highlight that according to [NeurIPS 2024 policy](https://neurips.cc/Conferences/2024/CallForPapers), our work should be considered as a concurrent work with [2] and \"not be rejected on the basis of the comparison to contemporaneous work\" as [2] was first published on arXiv on March 25th, 2024, which is within two months of the submission deadline of NeurIPS 2024.\n\n - Compared to [1], the significance of our work lies in three aspects:\n     - Nonlinear inverse problems are beyond the scope of [1], while we have investigated three different nonlinear inverse problems. \n     - We have adopted a more rigorous formulation than [1]. As also pointed out by [2], the denoiser design based on diffusion models in [1] is heuristic. \n     - Unlike PnP-SGS [1], we consider an annealing strategy, which is important for ill-posed inverse problems.\n\n- Compared to [2], our work have the following three contributions:\n     - Our work proposes a more general formulation based on EDM. To the best of our knowledge, using EDM as a rigorous Bayesian denoiser for solving inverse problems has not been explored in prior literature. \n     - Our proposed method has a significant improvement in performance over DPnP [2]. Response Table 2 in the response to reviewer 43sJ below further shows the statistical significance of the improvement. One reason for the performance gain is that we inherited the optimized design choices from EDM [3], which provides better image quality with less diffusion steps. Another reason is that the flexibility of the framework provides a larger design space with more flexible parameter choices. \n\n> Weakness 2\n\nWe provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.\n\n> Weakness 3\n\n- **Fixed $\\rho$**. Our current analysis considers a fixed $\\rho$ across different iterations, but we find that a variable schedule for $\\rho$ is more practical. Theoretically, our analysis could be extended to the setting with a variable $\\rho$ by viewing the iterations with larger $\\rho$ values as warm-up iterations that produce favorable initial conditions for the algorithm with smaller $\\rho$ values. \n- **Approximate score bound**. We can extend our convergence bound to the case where the learned score has error, which we can include in our revised manuscript. More precisely, assume we implement the EDM reverse diffusion step (as given by equation (10) in the manuscript) as: $$\\mathrm{d} \\boldsymbol{x}\\_t=\\left[u(t) \\boldsymbol{x}\\_t- v(t)^2 s\\_t\\left(\\boldsymbol{x}\\_t\\right)\\right] \\mathrm{d} t + v(t)\\mathrm{d}\\bar{\\boldsymbol{w}}\\_t,$$ \n\n    where $s\\_t$ is the approximate score. Then, by differentiating the KL divergence along the two dynamics or using the Girsanov theorem (in the spirit of the proof in [4]), we can prove the following: for $\\tau \\in [k(1+t^\\ast)+1, (k+1)(1+t^\\ast)]$, which corresponds to the prior step in the Split Gibbs Sampler, we have:\n    $$\\frac{1}{4}\\int\\_{k(1+t^\\ast)+1}^{(k+1)(1+t^\\ast)} \\lambda(\\tau)\\text{FI}(\\pi\\_\\tau\\Vert \\nu\\_\\tau)\\mathrm{d}\\tau \\leq \\epsilon\\_{\\text{score}}(k) + \\text{KL}(\\pi\\_{(k+1)(1+t^\\ast)}\\Vert \\nu\\_{(k+1)(1+t^\\ast)}) - \\text{KL}(\\pi\\_{k(1+t^\\ast)+1}\\Vert \\nu\\_{k(1+t^\\ast)+1})$$\n\n    where $\\epsilon\\_{\\text{score}}(k) := \\int\\_{0}^{t^\\ast} v(t)^2\\mathbb{E}||s\\_t(\\boldsymbol{x}\\_t^{(k)}) - \\nabla \\log p\\_t(\\boldsymbol{x}\\_t^{(k)})||^2 \\mathrm{d} t$ and $\\boldsymbol{x}\\_t^{(k)}, 0\\leq t \\leq t^\\ast$ is the exact EDM reverse process, which satisfies equation (10) in the manuscript, starting at $\\boldsymbol{x}\\_0^{(k)} \\sim \\nu\\_{k}^Z$ and ending up at $\\boldsymbol{x}\\_{t^\\ast}^{(k)} \\sim \\nu\\_{k+1}^X$. Therefore, up to score approximation errors, a convergence bound analogous to Theorem 3.1 will hold.    \n- **Convergence to the true posterior**. Indeed, the asymptotic convergence to the true posterior as $\\rho\\to 0$ is difficult to obtain (see Section 5 of [5]). We agree that this is out of the scope of the paper, so we leave it to future work.\n\n> Question 1\n\nWe trained a denoiser for the Gaussian image prior. More details are presented in Appendix C.2 \"Model checkpoint\" paragraph.\n\n> Question 2\n\nPlease refer to Appendix D for how we chose the hyperparameters for all the methods. For our method, we fine-tuned the parameters in the annealing schedule $\\rho_0$, $\\alpha$, and $\\rho_\\min$ (notations from Appendix C.3) by performing a grid search over 20 FFHQ images outside of the test set we used for final reporting (referred to as the validation set hereafter). For DDRM and DPS on linear inverse problems, we used the default parameters in their official repositories. For all other cases, we performed a grid search of the main parameter(s) of each method on the validation set. Specifically, for DPIR, we used the same annealing function as that in the official repository of [6] but fine-tuned the starting/ending noise level and number of iterations.\n\n> Question 3\n\nWe did not employ any explicit regularization parameters. Within the Bayesian framework, the data-fidelity (likelihood) and regularization (prior) should be automatically balanced based on the noise distribution. \n\n> Question 4\n\nThe discretization time steps are the same for all values $\\rho_k$. The number of iterations that are actually run depends on the specific value of $\\rho_k$. See Appendix C.2 “Pseudocode” paragraph and Algorithm 3 for more details.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3T108H4wPJ", "reviewer": "Reviewer_G64m", "summary": "This paper proposes a Markov Chain Monte Carlo algorithm for posterior sampling in both linear and non-linear inverse problems. The core of the proposed method is based on a Split Gibbs Sampler that alternates between two steps: one involving the likelihood and the other the prior. Additionally, the paper connects the Bayesian denoising problem with unconditional generation using Diffusion Models. The proposed method is validated on a range of linear and non-linear inverse problems, with an additional real-world application in black hole imaging.", "questions": "No questions.", "limitations": "No limitations other than those mentioned in the paper.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The proposed method outperforms DPS, with satisfactory evidence provided by the authors.\n- Unlike some previous works, the proposed method is effective for both linear and non-linear inverse problems.\n- The paper provides sufficient theoretical analysis to support the proposed method.", "weaknesses": "There are no apparent weaknesses. However, I am curious about the reconstruction speed comparison (seconds/image) between DPS and the proposed method. It would be practically very attractive for the community to use this algorithm if a rigorous comparison of reconstruction speed is made.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your evaluation of our paper. We respond to your comment below.\n\n> There are no apparent weaknesses. However, I am curious about the reconstruction speed comparison (seconds/image) between DPS and the proposed method. It would be practically very attractive for the community to use this algorithm if a rigorous comparison of reconstruction speed is made.\n\nWe provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aNCpgcbGkD", "reviewer": "Reviewer_G2Ty", "summary": "The paper introduces a method to address Bayesian inverse problems in computational imaging. It leverages the generative capabilities of diffusion models (DMs) to sample the posterior distribution over all possible solutions from noisy and sparse measurements. The method combines a Markov chain Monte Carlo (MCMC) algorithm with a general DM formulation to perform rigorous posterior sampling, effectively integrating state-of-the-art DMs as expressive image priors. The method deviates from current methods that rely on approximating the intractable posterior via separating the forward operator from an unconditional prior over the intermediate noisy image. The approach is validated on six inverse problems, demonstrating superior accuracy and posterior estimation compared to existing DM-based methods.", "questions": "1. Can the authors compare the computational resources and time required for PnP-DM versus existing DM-based methods?\n      \n2. How sensitive is the performance of PnP-DM to the choice of the annealing schedule for the coupling parameter ρ? Is there an optimal range or strategy for selecting these parameters?", "limitations": "Yes.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper presents a novel combination of MCMC and diffusion models to solve inverse problems. This use of DMs as priors in a Bayesian framework can be integrated into SOTA diffusion models as plug-n-play expressive image priors for Bayesian inference.\n   \n- The presented method is rigorous. The use of the Split Gibbs Sampler and the EDM formulation for the prior step is well-executed. The theoretical insights provided, including the stationarity guarantee in terms of average Fisher information, add depth to the method.\n\n- The paper is well-structured and supported by diagrams and pseudocode to enhance understanding. The distinction between existing methods and the proposed approach is clearly articulated.\n\n- Experiments show the efficacy of the proposed method compared to existing methods on a diverse set of real-world problems, highlighting its potential broader impact on computational imaging applications.", "weaknesses": "- While the method demonstrates superior performance, the computational cost and efficiency are not thoroughly discussed. A comparison of computational resources required compared to other methods would be beneficial.\n\nThe impact of various parameters, such as the annealing schedule for the coupling parameter ρ, on the method's performance is not deeply investigated. A sensitivity analysis could provide insights into the method's robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback on our work. Below we provide point-by-point responses to your comments.\n\n> While the method demonstrates superior performance, the computational cost and efficiency are not thoroughly discussed. A comparison of computational resources required compared to other methods would be beneficial. Can the authors compare the computational resources and time required for PnP-DM versus existing DM-based methods?\n\nWe provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.\n\n> The impact of various parameters, such as the annealing schedule for the coupling parameter ρ, on the method's performance is not deeply investigated. A sensitivity analysis could provide insights into the method's robustness. How sensitive is the performance of PnP-DM to the choice of the annealing schedule for the coupling parameter ρ? Is there an optimal range or strategy for selecting these parameters?\n\nFollowing your suggestion, we have included a sensitivity analysis on the annealing schedule. In Figure 1 of the attached PDF file, we show the PSNR curves of different exponential decay rates $\\alpha$ (left) and minimum coupling levels $\\rho_{\\min}$ (rate) for one linear (super-resolution) and one nonlinear (coded diffraction patterns) problem. \n\nWe have the following conclusions based on the results. First, different decay rates lead to different rates of convergence, which corroborates with our theoretical insights that $\\rho$ plays the same role as the step size. The final level of PSNR is not sensitive to different decay rates, as all curves converge to the same level. Second, as $\\rho_{\\min}$ decreases, the final PSNR becomes higher as the stationary distribution converges to the true target posterior. \n\nOur strategy for choosing the annealing parameters is as follows. The starting coupling strength $\\rho_0$ should be large to overcome the ill-posedness of the problem (usually around 5 to 10 is sufficient). The minimum coupling strength $\\rho_{\\min}$ should be small to ensure that the stationary distribution is close to the target posterior; empirically around 0.1 to 0.3 works the best. The range of decay rate $\\alpha$ is generally flexible; a value around 0.9 usually leads to good results.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WtufxyczwO", "reviewer": "Reviewer_TmFe", "summary": "This paper proposes a novel posterior sampling algorithm for solving inverse problems with Diffusion models. The proposed algorithm is based on the Split Gibbs sampling scheme and consists in, first introducing an extended distribution $\\pi_\\rho(x, z)$ that admits the original posterior $\\pi$ as $x$ marginal after letting $\\rho$ tend to $0$. Then, for a fixed $\\rho$ this joint distribution can be sampled via Gibbs sampling; the form of the joint distribution $\\pi_\\rho(x, z)$ allows sampling sampling both $X$ and $Z$ easily. For $Z$, this can be done exactly in the case of linear inverse problems, or with Langevin Monte Carlo otherwise. For $X$, the particular structure of the joint distribution is exploited and as it turns out, it can be sampled by simply running the learned backward Diffusion. Notably, this is the only step that requires using the learned denoiser. The theoretical properties of the ideal sampler are inverstigated and extensive experiments are considered.", "questions": "see above", "limitations": "see above", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I highly enjoyed going through this paper and I found it to be interesting. The strengths are: \n- The method is original and clever; the joint distribution introduced, which stems from the SGS approach, allows almost exact \nsampling of both the conditional distributions. Furthermore, the considered framework allows having actual theoretical guarantees for the sampler (although in an ideal case). This is still quite informative and contrasts with existing approaches (besides the SMC-based ones which also allow more straightforward theoretical guarantees). The main perk of this method wrt to the latter is that it has theoretical guarantees while not requiring absurd memory requirements due to the storage of a large number of particles. \n\n- The writing is extremely clear. The paper is written in a way that It is straightforward to understand while still being quite detailed. This is a significant perk (Still, i think that section 3.4 could be improved as I had to go through the supplementary material to actually understand what the proof technique was). \n\n- Finally, the experiments are very good; the authors first consider a toy example in which the posterior has a closed form. This serves as a nice introduction to the experimental section and showcases the perks of the present method in comparison with DPS. Next, the imaging experiments are rigorously executed, with a precise look at the uncertainty quantification. Finally, an original black hole experiment is considered and I found it to be very interesting.", "weaknesses": "- **Computational cost**:  I believe that the proposed method is much slower than existing alternatives. This is not that much of an issue for me (but of course this depends on the reader and practitioner). In my opinion the compute time is not properly addressed; although the authors discuss it in the bottom of page 21, I would like to see a proper runtime comparison with existing methods. I would like to emphasize that  this is only a minor weakness.\n\n- **Originality**: I believe that the proposed method is related to SDEdit [1]; it can be thought of as a variant of SDEdit. Indeed, to see why this is the case, when the present algorithm is applied to inpainting, the first iteration, the likelihood step, will result in a noisy image with the observation already present In it. This image is then denoised using the backward process during the prior step. This is in fact very similar to sdedit which starts with some initial image, say the pseudo-inverse of the observation, noises it to some given noise level and denoises it using the backward Diffusion process. See Algorithm 3, [1]. I believe that the main difference with the algorithm presented in your paper is the use of the decreasing noise schedule, which makes sense theoretically. For a fair treatment I suggest you mention this in your paper. \n\n[1] Meng, C., He, Y., Song, Y., Song, J., Wu, J., Zhu, J.Y. and Ermon, S., 2021. Sdedit: Guided image synthesis and editing with stochastic differential equations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your evaluation of our work. Below we provide point-by-point responses to your comments.\n\n> **Computational cost**: I believe that the proposed method is much slower than existing alternatives. This is not that much of an issue for me (but of course this depends on the reader and practitioner). In my opinion the compute time is not properly addressed; although the authors discuss it in the bottom of page 21, I would like to see a proper runtime comparison with existing methods. I would like to emphasize that this is only a minor weakness.\n\nWe provide a comparison on computational efficiency for all methods in Response Table 1 above. Our method PnP-DM achieves superior performance while remaining comparable in runtime with DPS, with a runtime no more than 1.5 times that of DPS. In fact, for the linear inverse problem experiments, PnP-DM was faster than DPS because we were able to take multiple samples along one Markov chain generated by PnP-DM but DPS requires running multiple diffusion processes.\n\n> **Originality**: I believe that the proposed method is related to SDEdit [7]; it can be thought of as a variant of SDEdit. Indeed, to see why this is the case, when the present algorithm is applied to inpainting, the first iteration, the likelihood step, will result in a noisy image with the observation already present In it. This image is then denoised using the backward process during the prior step. This is in fact very similar to sdedit which starts with some initial image, say the pseudo-inverse of the observation, noises it to some given noise level and denoises it using the backward Diffusion process. See Algorithm 3, [7]. I believe that the main difference with the algorithm presented in your paper is the use of the decreasing noise schedule, which makes sense theoretically. For a fair treatment I suggest you mention this in your paper.\n\nThank you for bringing this work to our attention. This work indeed shares some similarity to our own in that it employs a process of noising followed by denoising. Nevertheless, our method is more theoretically grounded and meant for rigorously estimating the posterior of an inverse problem, whereas SDEdit is more empircally driven and designed for solving image synthesis/editing problems. We will cite and discuss [7] in the related work section in the final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "c6Za8jcrWS", "reviewer": "Reviewer_43sJ", "summary": "It is proposed to solve image inverse problems with diffusion priors by sliced Gibbs sampling (SGS). To sample from the product of prior and likelihood distributions, SGS assumes a pair of variables $x,z$ coupled by a Gaussian and alternates conditional sampling steps: on $x$ using the prior density multiplied with the conditional Gaussian and on $z$ using the likelihood multiplied with the conditional Gaussian. When the prior is a diffusion model, the two steps can be implemented by modifying the given pretrained score model and by Langevin dynamics (or exactly in the case of a linear problem), respectively. This procedure is evaluated on a few standard linear and nonlinear IPs and on a black home imaging problems, and the reconstructions found by the proposed method are typically visually and quantitatively better than the ones found by algorithms from prior work.", "questions": "- After equation (1), writing $y=A(x)+n$ where $n\\in\\mathbb{R}^n$ means that $y$ depends deterministically on $x$. In fact $n$ should be a random variable taking values in $\\mathbb{R}^n$, which makes $y$ also a random variable.\n- I am surprised not to see any reference to the paper whose title is a suffix of this one's, [\"Diffusion models as plug-and-play priors\", NeurIPS'22], also abbreviated \"PnP\" in some subsequent work. There a stochastic optimisation is used to find modes of the posterior.\n  - Related to this, \"PnP-DM\" is not a very informative name for the proposed algorithm when so many of the baseline methods already have \"PnP\" in the name in some other configuration of letters.", "limitations": "Yes, but it would be good to compare computation costs, see above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The writing is mostly clear; I have no complaints on the paper organization or exposition.\n  - In particular, the exposition of SGS in the general case followed by its application to DM priors is helpful.\n- As far as I know, this is an original way of solving the diffusion posterior sampling problem.\n- Strong results (if we are to accept that they support the claims) and application to a real imaging problem.\n- Code is provided.", "weaknesses": "- On fairness and strength of experimental comparisons:\n  - The results do not come with confidence intervals, which makes it hard to assess significance.\n  - I may have missed/misunderstood this, but are all methods using the same pretrained prior? Appendix D seems to suggest otherwise, which would be a problem (it isn't fair to compare posterior sampling methods with different priors). \n    - Similarly, different methods seem to be assuming different numbers of sampling steps and otherwise incompatible choices. How were hyperparameters, such as the $\\rho_k$ annealing schedule, chosen?\n  - Computation cost is not sufficiently discussed.\n    - The method is run for at least a hundred Gibbs iterations in each experiment, and each one involves diffusion sampling initialized from some intermediate noise level. This makes the number of model evaluations quite large. It would be necessary to compare the number of function calls, as well as wall time, for all the methods.\n- Convergence rate:\n  - Theorem 3.1 tells us the rate of convergence to the stationary process for a fixed $\\rho$, but not how close the stationary marginal is to the true posterior. They are not equal, in general, for positive $\\rho$, and the rate of convergence shown has a $\\frac1\\rho$ factor, which means the convergence guarantees with small $\\rho$ are weaker.\n  - To understand the method, it would be important to show the dependence on the number of iterations, both by showing examples at different iterations in an illustrative case and by tracking the convergence of some metric at a function of the number of steps.\n- Please also see questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper. Below we provide point-by-point responses to your comments. \n\n> Weakness 1.1\n\nWe provide the confidence intervals of all the results in Tables 1 and 2 in the attached PDF. Based on these results, we claim that our method achieves comparable or superior outcomes on most tasks, particularly for complex nonlinear inverse problems such as Fourier phase retrieval. Compared to the baseline methods, our standard deviation is slightly lower, indicating comparable and slightly better robustness of our approach.\n\nTo better show the significance of the improvement of our method over baselines, we further conducted a sample-wise PSNR comparison between PnP-DM (EDM) and the most competitive baseline DPnP. The table below shows the sample-wise improvement rate PnP-DM over DPnP and $p$-value of the one-sided t-test for each inverse problem. Note that our method improves on *all samples in the test set* and the $p$-values indicate high statistical significance.\n\n**Response Table 2:** Sample-wise comparison between our method and DPnP [2].\n\n|                  | Gaussian deblur | Motion deblur | Super-resolution | Coded diffraction patterns | Fourier phase retrieval |\n|------------------|-----------------|---------------|------------------|----------------------------|-------------------------|\n| Improvement rate | 100%            | 100%          | 100%             | 100%                       | 100%                    |\n| t-test $p$-value | 8.18e-8         | 1.15e-6       | 9.40e-11         | 3.43e-70                   | 2.28e-2                 |\n\n> Weakness 1.2\n\nAs mentioned in Page 7 line 244-246, we used the same model checkpoints for all diffusion model-based methods, which includes DDRM, DPS, PnP-SGS, DPnP, and our PnP-DM. This suggests that these sampling methods should share the same target posterior distribution.\nFor our method, we fine-tuned the parameters in the annealing schedule $\\rho_0$, $\\alpha$, and $\\rho_\\min$ (notations from Appendix C.3) by performing a grid search over 20 FFHQ images outside of the test set we used for final reporting. Please also see the response to Question 2 by Reviewer u5tS.\n\n> Weakness 1.3\n\nPlease refer to the response to all reviewers above for a comparison on computational efficiency. To clarify, as we mentioned in Appendix C.2 \"Pseudocode\" paragraph, the total number of time steps for the *entire* diffusion process is 100, but we only start running from some intermediate noise level. Therefore, the number of NFEs for each iteration is less than 100. In fact, according to the table, the average number of NFEs is only around 30 per iteration given the current annealing schedule.\n\n> Weakness 2.1\n\nAlthough we agree that there is a gap between the stationary process and the true posterior, we would like to argue that the convergence guarantees with small $\\rho$ are not weaker and that the $\\frac{1}{\\rho}$ factor is in fact expected. For SGS-base methods, $\\rho$ can be interpreted as the \"step size.\" The step size has a similar presence in the bound of prior non-asymptotic analysis on the convergence of Langevin-based MCMC algorithm [8], [9]. For these algorithms, it is common to have an $O(1/T)$ convergence rate where $T$ is the total diffusion time and scales with the step size, so there will be a $\\frac{1}{\\text{step size}}$ factor on the right hand side. This also justifies the annealing strategy, which starts with large $\\rho$ (hence large step size) at the beginning when the distribution is far away from the target and gradually decreases.\n\n> Weakness 2.2\n\nThank you for your suggestion. In Figure 2 of the attached PDF file, we show some visual examples of intermediate $\\boldsymbol{x}$ and $\\boldsymbol{z}$ iterates (left) and convergence plots of PSNR, SSIM, and LPIPS for $\\boldsymbol{x}$ iterates (right) on the super-resolution problem. As $\\rho_k$ decreases, the $\\boldsymbol{x}$ iterate becomes closer to the ground truth and the $\\boldsymbol{z}$ iterate gets less noisy. Both the visual quality and metric curves stabilize after the minimum coupling strength $\\rho_{\\min}$ is achieved. Despite being run for 100 iterations in total, our method generates good images in around 40 iterations, which is around 30 seconds and 1600 NFEs. We can include these results in our supplemental material of the revised manuscript. \n\n> Question 1\n\nWe will change the writing on this part in the final version of the paper.\n\n> Question 2\n\nThank you for bringing this paper to our attention. We notice that [one reviewer's comment](https://openreview.net/forum?id=yhlMZ3iR7Pu&noteId=IRfkHFMGG6#:~:text=for%20downstream%20tasks.-,The%20naming%20of%20the%20framework%20is%20confusing.%20Plug%2Dand%2Dplay%20priors%20(PnP)%20is%20a%20well%2Dknown%20framework%20in%20the%20literature%20of%20image%20restoration%20and%20imaging%20inverse%20problems.%20The%20title%20confuses%20me%20in%20the%20first%20place.,-The%20naming%20is) of this work also touches on the naming of the method. According to the authors' response, this happens to be a clash of terminology. However, we will include this paper in the related work. As for the \"PnP-DM\" name, we understand that \"PnP\" is so commonly used that it may cause some ambiguity. We used it mainly to highlight its similarity to deterministic PnP methods that incorporate an image prior via denoising.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## **Response to all reviewers**
We thank all the reviewers for their careful reviews and constructive feedback. We are glad that our method was recognized as "rigorous", "well-executed" (reviewer G2Ty), and “unlike some previous works...effective for both linear and non-linear inverse problems” (reviewer G64m), including “an application to a real imaging problem” (reviewer 43sJ). We are also encouraged that the reviewers found our paper "well-written" (reviewer u5tS) and the writing "extremely clear" (reviewer TmFe). 

In the responses below, we address the reviewers' comments individually. Additional experiments are presented in the attached PDF file. 

--------------------------------------------------------------------------------

### **Common response to computational efficiency**

We present a comparison of computational efficiency with the major baselines on a linear super-resolution and a nonlinear coded diffraction patterns problem in Response Table 1 below. 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 ($\approx 1.5\times$), despite its larger NFEs ($\approx 3\times$). This is primarily due to two factors: 1) PnP-DM avoids running the full diffusion process by adapting the starting noise level to $\rho_k$ at each iteration, and 2) the runtime is further reduced by using an annealing schedule of $\rho_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).

**Response Table 1**: Comparison of runtime in seconds and number of function evaluations (NFE) for our method and baselines.

|                            | Metric         | DDRM | DPS | PnP-SGS | DPnP  | PnP-DM (ours) |
|----------------------------|----------------|------|-----|---------|-------|---------------|
| Super-resolution           | Clock time (s) | 0.4             | 39             | 20      | 322   | 55           |
|                            | NFE            | 20              | 1000           | 1030             | 18372 | 3032         |
| Coded diffraction patterns | Clock time (s) | --              | 37             | 54      | 261   | 50           |
|                            | NFE            | --              | 1000           | 2572    | 14596 | 2482         |

--------------------------------------------------------------------------------

### **References for all responses**

[1] Coeurdoux et al., 2023. Plug-and-play split Gibbs sampler: embedding deep generative priors in Bayesian inference.

[2] Xu and Chi, 2024. Provably robust score-based diffusion posterior sampling for plug-and-play image reconstruction.

[3] Karras, et al., 2022. Elucidating the design space of diffusion-based generative models.

[4] Chen et al., 2023. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions.

[5] Yuan et al., 2023. On a class of Gibbs sampling over networks.

[6] Zhang et al., 2022. Plug-and-play image restoration with deep denoiser prior.

[7] Meng, C., He, Y., Song, Y., Song, J., Wu, J., Zhu, J.Y. and Ermon, S., 2021. SDEdit: Guided image synthesis and editing with stochastic differential equations.

[8] Balasubramanian et al., 2022. Towards a theory of non-log-concave sampling: first-order stationarity guarantees for Langevin Monte Carlo.

[9] Sun et al., 2023. Provable probabilistic imaging using score-based generative priors.

Meta-review: This paper received 5 reviews with balanced views but overall positive. The discussion period has enabled to reach a consensus for acceptance.

All reviewers recognized the quality of the work :
- "the paper is well-written, provides clear and concise explanations of the methods and contributions"
- "The strength of the paper is mainly its experimental section, which covers a range of both linear and non-linear inverse problems."
- "Unlike some previous works, the proposed method is effective for both linear and non-linear inverse problems."
- "The paper provides sufficient theoretical analysis to support the proposed method."
- "The theoretical insights provided, including the stationarity guarantee in terms of average Fisher information, add depth to the method."
- "Experiments show the efficacy of the proposed method compared to existing methods"
- "the experiments are very good; the authors first consider a toy example in which the posterior has a closed form. This serves as a nice introduction to the experimental section ..."
- "this is an original way of solving the diffusion posterior sampling problem."

The authors/reviewers discussion period was useful for resolving the reviewers concerns. Yet, it is required for the final version of the paper that the authors :
- provide a more detailed comparison with Ceurdoux et. al. and clearly state the similarities/differences,
- provide an extended theoretical analysis with approximate score (at least in supplementary material).
Decision: Accept (poster)