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

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.

Paper

References (75)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer 43sJ6/10 · confidence 3/52024-07-11

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.

Strengths

- The writing is mostly clear; I have no complaints on the paper organization or exposition. - In particular, the exposition of SGS in the general case followed by its application to DM priors is helpful. - As far as I know, this is an original way of solving the diffusion posterior sampling problem. - Strong results (if we are to accept that they support the claims) and application to a real imaging problem. - Code is provided.

Weaknesses

- On fairness and strength of experimental comparisons: - The results do not come with confidence intervals, which makes it hard to assess significance. - 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). - 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? - Computation cost is not sufficiently discussed. - 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. - Convergence rate: - 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. - 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. - Please also see questions below.

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. - 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. - 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.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, but it would be good to compare computation costs, see above.

Reviewer TmFe7/10 · confidence 5/52024-07-12

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.

Strengths

I highly enjoyed going through this paper and I found it to be interesting. The strengths are: - The method is original and clever; the joint distribution introduced, which stems from the SGS approach, allows almost exact sampling 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. - 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). - 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. - **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. [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.

Questions

see above

Rating

7

Confidence

5

Soundness

4

Presentation

4

Contribution

4

Limitations

see above

Reviewer G2Ty8/10 · confidence 3/52024-07-12

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.

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. - 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. - 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. - 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. 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.

Questions

1. Can the authors compare the computational resources and time required for PnP-DM versus existing DM-based methods? 2. 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?

Rating

8

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

Yes.

Reviewer G64m8/10 · confidence 4/52024-07-14

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.

Strengths

- The proposed method outperforms DPS, with satisfactory evidence provided by the authors. - 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.

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.

Questions

No questions.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

No limitations other than those mentioned in the paper.

Reviewer G64m2024-08-12

Thank you for sharing the sampling speed results. The shared results make sense to me. I strongly feel that this paper provides novel enough contributions to be considered for acceptance. I have raised my confidence score to 4. Hope this helps AC decide about the paper.

Reviewer u5tS4/10 · confidence 4/52024-07-15

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 proposed to unify various diffusion models in a common general formulation.

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. The 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? 2. **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. 3. **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.

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? - 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. - Did you employ any regularization parameters to balance the data-fidelity and regularization terms? - 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$ ?

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

No limitations.

Authorsrebuttal2024-08-09

Thank you for your response and comments. Below we answer them point by point. > **W1.1** Indeed, we have tried to apply PnP-SGS [1] to nonlinear inverse problems by combining our likelihood step with its prior step design, which was how we obtained the PnP-SGS results for Table 2. While it could still handle the coded diffraction pattern problem, it failed on the more challenging Fourier phase retrieval problem. We also tried to include the annealing strategy for PnP-SGS but found that the method diverged with large $\rho$, probably due to its heuristic design of the prior step, and thus did not benefit from annealing. Overall, our method significantly outperforms PnP-SGS by at least 1dB in PSNR for all linear problems and coded diffraction patterns, and 15dB for Fourier phase retrieval. Our experimental results indicate that PnP-SGS struggles with challenging nonlinear inverse problems, such as Fourier phase retrieval. > **W1.2** There are mainly two aspects. Here we explain using the notations in the original DDPM paper. First, to rigorously implement the prior step of SGS, one need to assume the input as an observation of the *unscaled* image $\boldsymbol{x}_0$ with additive white Gaussian noise. However, in DDPM, the mean of state $\boldsymbol{x}\_t$ is not $\boldsymbol{x}\_0$, but $\sqrt{\bar{\alpha}\_t}\boldsymbol{x}\_0$, which is a *down-scaled* version of $\boldsymbol{x}\_0$. So, it is inaccurate to directly use DDPM as a denoiser. This mismatch is particularly significant when starting from a large $t^\ast$ as $\sqrt{\bar{\alpha}\_{t^\ast}}$ is close to 0. Second, according to the SGS formulation, the noise estimation module is unnecessary, and one should always denoise $\boldsymbol{z}^{(k)}$ (the $\boldsymbol{z}$ iterate of SGS at iteration $k$) assuming a noise level $\rho$. However, PnP-SGS does not take into account the hyperparameter $\rho$ in the denoising problem. It is unclear how the generated $\boldsymbol{x}^{(k+1)}$ relates to the target conditional distribution $\pi^{X|Z=\boldsymbol{z}^{(k)}}(\boldsymbol{x})$ for each prior step of SGS. Therefore, the posterior distribution sampled from the denoiser in [1] is not the desired posterior distribution even with a perfect score function. > **W1.3** We never claimed that the idea of annealing is our contribution. Nevertheless, we believe that it is our contribution to propose a general framework that accommodates annealing in an easy yet rigorous way for SGS and leads to significant performance improvement. Our focus is more on the algorithm design and experimental validation. To the best of our knowledge, no existing theory on SGS provides a non-asymptotic convergence bound with annealing, so this is a a challenging open question. Although our current theory assumes fixed $\rho$, it still provides some theoretical insights on the algorithm behavior, such as the interpretation of $\rho$ as step size, and potentially opens up new theoretical directions for SGS based on the Fisher information. We hope to extend the analysis to the annealing case in the future. > **W2** We respectfully disagree with your last statement. DPS requires backpropagating through the entire denoiser network for each diffusion step after the forward pass of the denoiser, introducing a significant computational overhead. On the other hand, our method does not require doing so. This is the main reason why the clock time does not scale linearly with NFE for these two methods. Moreover, unlike DPS that applies a likelihood update for every function evaluation. Our method does so only every multiple function evaluations. > **W3.1** As we showed in Appendix C.3 "Annealing schedule for $\rho$" paragraph, our annealing schedule decreases to a fixed minimum level at $\rho_{\min}$, so we indeed fixed $\rho$ at small values after certain numbers of iterations. Furthermore, we did run our algorithms for the synthetic prior experiments with a fixed schedule of $\rho$ throughout the process, as shown in Table 4. Experimental results show that our method can accurately sample the posterior, which corroborates our theory. For experiments on linear inverse problems and coded diffraction patterns with FFHQ images, the results with fixed $\rho$ are on par with those with annealing, probably because the problems are not highly non-convex. For more challenging nonlinear problems, such as the Fourier phase retrieval, we empirically find that an annealing schedule is essential to overcome the non-convexity of the problem and provide accurate reconstruction. > **W3.2** By combining the above argument with those in the manuscript, we will get the following bound with score function error: $$\frac{1}{T} \int_0^{T} \mathsf{FI}\left(\pi_\tau || \nu_\tau\right) \mathrm{d} \tau \leq \frac{4\mathsf{KL}(\pi^X||\nu_0^X)}{K(1+t^\ast) \min(\rho, \delta)^2} + \frac{1}{K(1+t^\ast)\delta^2}\sum_{k=1}^K\epsilon_{\text{score}}(k).$$

Reviewer u5tS2024-08-13

Response to authors

I thank the authors for their careful answer. I think that the manuscript would benefit from a more detailed comparison with Ceurdoux et al. and from an extended theoretical analysis with approximate score. Also if should be made clear that the theory is only with fixed rho parameter. With these clarifications I am ready to recommend acceptance.

Authorsrebuttal2024-08-13

We appreciate your constructive feedback on our work. We will include a more detailed comparison with Coeurdoux et al. and the theoretical analysis with an approximate score function in the final version. We will also clarify that the current theory is only with fixed $\rho$.

Reviewer u5tS2024-08-08

Response to rebuttal

I thank the authors for their detailed answers. Here are additional remarks / questions. **W1:** 1. Although not demonstrated in the experiments, I believe the approach in [1] could be equally applied to nonlinear inverse problems. 2. Could you clarify why you consider the denoiser design based on diffusion models in [1] to be heuristic? 3. The annealing strategy, despite not being adopted in [1], is commonly used in Plug-and-Play literature. However, incorporating this strategy introduces a discrepancy between your theoretical framework and experimental results, which is concerning. 4. I acknowledge that using EDM as the diffusion model offers a more flexible and optimized framework. **W2:** I do not understand why the clock time does not scale linearly with NFE. It appears that the operations outside the denoiser evaluation in both DPS and your method are relatively insignificant in time. **W3:** 1. I strongly disagree with your statement, as you did not run the algorithms with fixed $\rho$ even for smaller values. 2. Can you explicitly state the final convergence bound analogous to Theorem 3.1, including the time integral from 0 to T, that would result from this additional theoretical consideration?

Reviewer 43sJ2024-08-12

I'm happy to recommend acceptance given the clarifications. Thank you.

Reviewer TmFe2024-08-12

Thank you for your response. Reviewer u5tS mentions that this work bears significant similarities with the work of Coeurdoux et al. and I agree that this is the case with the main difference being that they use a constant $\rho$, which is also more or less the main difference between your work and that of [1] that I am mentioning in my initial review. While reading the paper I did not notice these similarities with the work of Coeurdoux et al., this means that paper in its current form does not adequately address the related work and might mislead the reader into thinking that the paper is much more novel than what it actually is. Nonetheless, the empirical results are promising and show that the decreasing schedule is indeed a good idea. I maintain my score but i strongly recommend the authors to modify their paper and clearly state the similarities with previous works.

Authorsrebuttal2024-08-13

Thank you for your response. We will provide additional clarifications on the similarities and differences between our work and prior works in the final version. Please also see [responses W1.1--W1.3](https://openreview.net/forum?id=Xq9HQf7VNV&noteId=O9TzVFCQmB) to Reviewer u5tS for more information on other differences.

Reviewer G2Ty2024-08-12

Thanks for the clarification and additional experiments. This makes sense to me.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC