We present the first framework to solve linear inverse problems leveraging pre-trained latent diffusion models. Previously proposed algorithms (such as DPS and DDRM) only apply to pixel-space diffusion models. We theoretically analyze our algorithm showing provable sample recovery in a linear model setting. The algorithmic insight obtained from our analysis extends to more general settings often considered in practice. Experimentally, we outperform previously proposed posterior sampling algorithms in a wide variety of problems including random inpainting, block inpainting, denoising, deblurring, destriping, and super-resolution.
Paper
Similar papers
Peer review
Summary
This paper proposes a method to deal with linear inverse problems using pre-trained latent diffusion models (LDM) as a prior. The main idea is to extend the original diffusion posterior sampling (DPS) to the case of LDM by approximating the gradient term of the intractable likelihood. Two approximations (GLM-DPS), and (PSLD) are proposed for such approximation, and it turns out that PSLD achieves better performances. This paper also theoretically analyzes the DPS and PSLD in a toy linear setting. Experiment results show superior performances compared with the original DPS in a variety of tasks, such as random/block inpainting, denoising, super-resolution.
Strengths
1. It is the first work, to my best knowledge, that LDM is used to address the linear inverse problems. 2. Compared to the original DPS, the proposed PSLD achieves apparently better performances in most cases. 3. Some theoretical analysis is also provided, although it is based on some toy linear model.
Weaknesses
1. There is a lack of analysis of the complexity, or running time, of the proposed PSLD method. Adding details of the running time of PSLD and other comparison methods is suggested. 2. Part 3 contains some theoretical analysis using a simple toy linear model. For example, Theorem 3.4 states that DPS exactly recovers the groundtruth sample. It is known that DPS uses a Laplace approximation to approximate the gradient term of the likelihood. As a result, how can it exactly recover the posterior samples with the crude approximation? This is also the case with Theorem 3.7 and Theorem 3.8. To what extent can the results of Theorem 3.4, 3.7, and 3.8 be generalized to the real diffusion models? 3. In the experiments part, while PSLD shows superior performances over the original DPS, there is a lack of comparison with other latest diffusion-model-based algorithms that were available long before NeurIPS submission, e.g., [A] Wang, Yinhuai, Jiwen Yu, and Jian Zhang. "Zero-shot image restoration using denoising diffusion null-space model." arXiv preprint arXiv:2212.00490 (2022). [B] Meng, Xiangming, and Yoshiyuki Kabashima. "Diffusion model based posterior sampling for noisy linear inverse problems." arXiv preprint arXiv:2211.12343 (2022). [C] Song, Jiaming, Arash Vahdat, Morteza Mardani, and Jan Kautz. "Pseudoinverse-guided diffusion models for inverse problems." In International Conference on Learning Representations. 2022. 4. Most previous image restoration methods using diffusion models also consider colorization, can the authors also add some results in the case of colorization? Besides, only Table 2 and Table 5 show quantitative comparisons with other methods but only for inpainting. It is suggested to add quantitative comparisons with other methods in other tasks such as super-resolution, denoising, colorization, etc.
Questions
Three additional questions: 1. How do you add the additive Gaussian noise with standard deviation \sigma_y? 2. Can I understand that the improved performance of PSLD over DPS comes solely from the improvement of diffusion models in the latent space over that in the pixel space? In other words, if DPS is implemented with more powerful diffusion models, is it possible to outperform PSLD? In this sense, how to ensure a fair comparison between PSLD and DPS? 3. Are the results of PSLD sensitive to the choice of hyper-parameters? Compared to DPS which has only one scaling hyperparameter, there are two scaling hyperparameters, namely, \eta and \gamma. It is suggested to evaluate the sensitivity of these hyperparameters.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Please refer to the weakness and questions parts.
Summary
This paper focuses on solving inverse problems using diffusion based probabilistic models without **retraining**. The authors build on "Diffusion posterior sampling" which basically builds a diffusion model for the posterior using only the score of the prior distribution. Indeed, the score of the posterior $\nabla \log p_t(x_t | y)$ is equal to $\nabla \log p(y|x_t) + \nabla \log p_t(x_t)$ where the second term is given by the prior diffusion model. The first term is however intractable but can be approximated by noting that $p(y | x_t) = \int p(y|x_0) p(x_0 | x_t) dx_0$. This integral is then approximated by simply taking the mean a posteriori i.e. $p(y | x_t) \approx p(y | E(x_0 | x_t))$ which itself can be approximated using Tweedie's formula and the prior score. The original DPS paper only considers the diffusion in the pixel space and this paper extends the methodology to latent diffusion models. While the direct extension is straightforward, the authors claim that it does not work in practice and they provide a rather intuitive modified posterior diffusion model that ensures consistency at the borders of the mask.
Strengths
The paper tackles an important problem and provides a sound methodology for diffusion posterior sampling for latent space diffusion models. Furthermore, the numerical experiments suggest that this method outperforms other existing methods.
Weaknesses
The fact that the authors show recovery in a 2 step diffusion model in a very basic setting proves nothing. First, the results hold only when the optimal solution is normalized. In realistic settings and with multi step diffusion one does not know in advance how to rescale (if even rescaling is the right to do?) the drift so that the diffusion model matches the target distribution. Next, the authors prove that DPS "provably" recovers the posterior but this to me is also meaningless since one can easily show numerically in a noisy inpainting setting with Gaussian mixtures (the posterior is tractable in this case) that DPS **does not** in fact sample the posterior. It samples outside the support of the posterior even in the simplest settings. I believe that the authors should provide a numerical experiment where the target posterior is known (again, Gaussian mixtures) and show that their method does sample from the posterior. To me there is no reason why it should exactly sample from the posterior and not outside the support like DPS. Indeed, this issue is not related to the quality of the generative model but more to the fact that DPS (and truthfully many of the difussion posterior sampling methods) are not exact samplers since there is no correction like in MCMC algorithms. I believe that the authors should drop the "provably" in the title as it is misleading.
Questions
- What do the authors mean by the *curse of ambient dimension*, line 226-228? In DPS the gradient is indeed computed in the $d$-dimensional space but this is also the case of PSLD by the chain rule (and in fact the computational cost of PSLD is slightly larger). - In Figure 2 right panel, an inpainting example is considered. In this case p(y|x_0) is a dirac delta so I am not sure how the authors manage to take the gradient. Do the authors approximate the dirac delta with a Gaussian with small variance? Was the same courtesy applied to DPS? - It is quite surprising that in Algorithm 2 (but also Algo1) nothing depends on the measurement variance. Why is it not taken into account? - While I could guess what it respresents, $x_0 *$ is not defined anywhere.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
see below.
Summary
While the preivous methods have focused on solving linear inverse problems based on diffusion models, the paper presents a first extension to latent diffusion models (LDM). The core idea is developed upon the existing DPS method, which forms an approximation to p(y|xt) by using the denoising score estimate. The key challenges of the extension to LDM are (1) the non-bijective mapping between encoder and decoder, and (2) the inconsistency at the boundary of mask (for inpainting problems). To address these challenges, the authors propose an additional gluing term that penalizes the discontinuity at the mask boundary. The main practical contribution of this paper is that if unlocks the potential of large pre-trained LDM (e.g. stable diffusion), and hence the proposed based based on more powerful pre-trained models enjoys superior performance compared to its counterpart with ambient / pixel space settings. In addition, the paper also presents a theoretical analysis in a toy setting (linear two-step diffusion models) to motivate the proposed objective.
Strengths
- While the extension of DPS to LDMs seems natural, I appreciate the practical impact of the proposed methodology given the popularity and superior power of existing pre-trained LDMs. - If I understand it correctly, doing conditional inference in latent space would also be more efficient than doing inference in the ambient space (after controlling the model complexity though in practice I wonder if that comparison can be made) - The proposed method is also claimed to be robustness of the choice of step-size, at least in toy settings, though I am not sure how this would extend to general case.
Weaknesses
- The theoretical analysis seems to follow reference 30 closely, takes up a large portion of the paper. However, I am not sure how much insights can be carried over to the general cases. In my opinion, the proposed method stays well-motivated without this analysis. Hence. I believe either discussing the extension to general cases, or trimming this part would make a better paper. - The evaluation part can be misleading given that pre-trained LDMs are "trained on much more data compared to the one used by DPS" (ref: appendix, caption above table 5). I think this information should be disclosed in the main paper instead of the appendix. In addition, it would be helpful to also disclose the unconditional generation performance of all backbone diffusion models, as well as the size of the models, the number of training data, etc.
Questions
- Theorem 3.8: what is the value for gamma_i in Alg 2? - Theorem 3.8 and its surrounding text: can you elaborate more on why it is robust to different values of step size eta? In the equation between line 221 and line 222 on page 7, apparently if I choose different value of eta, the resulting z0 would be different (as long as the grad norm after eta is non-zero). Am I missing anything here? In addition, in this equation, should there be a multiplicative factor \gamma before the last term, or is it set to 1 for the theorem and the statement to be correct? - Page 7. line 228 "DPS algorithm suffers from the curse of ambient dimension". Can you elaborate on the issue? Is it mainly the efficiency issue or the approximation accuracy issue? - Would be helpful to accomplish the theoretical analysis with toy data empirical verification. - (If possible), would be helpful to provide results where the base diffusion models for DPS and PSLD have similar performance.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Yes.
Summary
This paper investigates the use of diffusion models for solving inverse problems. While many recent papers have explored diffusion models for inverse problems, to the best of my knowledge, this study is the first to propose a method for inverse problems with latent diffusion. One advantage of latent diffusion is its lower computational demand. The main challenge in solving inverse problems with diffusion models lies in computing the intractable likelihood term $\nabla \log p_{x_t}( y | x_t )$, where $x_t$ represents a noisy point in pixel-space. Recently, the DPS algorithm was introduced to approximate this term as $\nabla \log p_{x_t}( y | \mathbb{E}[x_0| x_t] )$. The first contribution of this paper is to extend the DPS framework to a latent diffusion model, approximating the intractable likelihood $\nabla \log p_{x_t}( y | z_t )$, where $z_t$ denotes a noisy point in latent space, with $\nabla \log p_{z_t}( y | \mathbb{E}[z_0| z_t] )$ along with an additional term that enforces $\mathbb{E}[z_0| z_t]$ to be a fixed point of the autoencoder (GML-DPS). Furthermore, the latter term is modified to enforce consistency in the measurements (PSLD). The empirical performance of the PSLD algorithm is compared with the standard DPS in high-dimensional reconstruction tasks, and improvements over DPS are demonstrated. These three algorithms (DPS, GML-DPS, PSLD) are theoretically studied in the case of a data distribution corresponding to a Gaussian supported on a low-dimensional subspace, perfect knowledge of the subspace, zero noise in the measurements, and a measurement matrix that is bijective over the subspace. Specifically, this implies that the inverse problem can be exactly solved given the measurements. Finally, the sampling processes are two-step diffusion processes. The authors demonstrate that all three algorithms precisely recover the target signal, and PSLD exhibits robustness to variations in the specification of the step sizes.
Strengths
Solving inverse problems with (large) latent diffusion models is a relatively underexplored area, and this paper fills a gap in the existing literature. In particular, the authors demonstrate how their methods can be used with SOTA latent-based foundation models such as stable diffusion. The empirical results presented are promising and demonstrate the potential of this approach. Additionally, the authors offer some theoretical guarantees, albeit in a limited setting. These guarantees are valuable and noteworthy, considering the field's scarcity of rigorous results, even in simple toy examples.
Weaknesses
The main weaknesses of the paper lie in the experimental evaluations and results. Specifically, all experiments are conducted on subsets of the FFHQ dataset and images sourced from the internet. The authors evaluate the PSLD algorithm on two latent models: LDM-VQ-4, trained on FFHQ, and Stable Diffusion, trained on a significantly larger dataset (LAION). However, the comparisons are made between PSLD on these two latent models (proposed) and DPS (baseline) on a standard diffusion model trained *exclusively* on FFHQ. The authors demonstrate improved performance of PSLD on the (latent) Stable Diffusion model compared to DPS on the standard diffusion model. However, when employing the latent model LDM-VQ-4 trained solely on FFHQ, the improvements are marginal (Table 3). As a result, it becomes challenging to ascertain whether the observed improvements stem from algorithmic enhancements or simply the utilization of superior models trained on larger datasets, which may also account for the enhanced out-of-distribution performance. Furthermore, the experiments do not shed light on the improvement attributable to the addition of the "gluing" term in comparison to the DPS's vanilla extension term. One might hypothesize that the "gluing" term could be employed in DPS (by removing the encoder-decoder), potentially enhancing the algorithm's performance.
Questions
Some typos and questions: - In Section 2, the notation can be confusing. To clarify, it is suggested to use $x_0^*$ to denote the target signal and distinguish it from the first sample of the backward diffusion ($x_0$). - For Equations (3), should it be $\hat{x}_0 = \mathbb{E}[x_0 | x_t]$? - In Algorithm 1: line 1 what is $\mathcal{T}$? - Equations (7) should include a $\log$ term for the DPS vanilla extension. Additionally, as shown in the supplementary material, $\mathcal{A}^\top$ should be added to (7) and line 8 of Algorithm 2. - In Algorithm 2, it appears that the unknown target $x_0^*$ is needed, but in reality, the algorithm only requires the measurements $y$. It is advisable to modify the algorithm to reflect this (as in Algorithm 1). - Line 142 of the Theoretical section introduces the noisy inverse problem as $y = A x_0 + \sigma_y n$. However, it should be noted that exact recovery is not possible in this case and it seems that the theoretical results are proven under the assumption of $\sigma_y = 0$, where the problem can be solved exactly. The authors should explicitly address these points to ensure clarity.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
2 fair
Contribution
3 good
Limitations
Yes
I appreciate the thorough response provided by the authors. I believe that this paper constitutes a valuable contribution and increased my score. However, I would like to discuss the phrasing and structure of the experiments in Section 5. I find it slightly misleading, which not only led me astray but possibly other reviewers as well. I understand the enthusiasm surrounding the outcomes achieved with the stable diffusion model. However, I believe it would be more equitable for the readers if the authors first presented and discussed the results for PSLD on the latent diffusion model LDM-V Q-4. Comparing these outcomes with those of DPS for a standard diffusion model makes sense, as these two models are trained using the same volume of data. This approach would effectively establish the authors' assertion that the primary innovation lies in the proposed algorithm for latent diffusion, which can perform comparably well, if not slightly better, than the state-of-the-art DPS for standard diffusion models. Following these initial experiments, presenting the results for the Stable diffusion model would be logical and would further reinforce the authors' argument. Finally, I am assuming that the authors have tested equation (5), subsequently moving on to (6), before eventually arriving at the proposed (7). I am curious whether the authors have conducted experiments to demonstrate that using (6) indeed outperforms the basic (5), and also to showcase that adopting (7) results in improvements over (6). It is possible that this information is present in the appendices, but I may have overlooked it.
Discussion with Reviewer vbWT
Dear Reviewer vbWT, Thank you for reading our response and increasing the score. We greatly appreciate your timely feedback and active participation during the discussion phase. Below, we respond to your new comments. (1) **Discussion on the phrasing and structure of the experiments in Section 5.** We thank the reviewer for the suggestion on how to make the flow of the experimental Section 5 more logical and further reinforce our arguments. As suggested by the reviewer, we will rephrase and restructure Section 5 in the revised version of the paper. (2) **Have the authors tested equation (5), (6) and (7)?** We have tested equation (5), (6) and (7) leading up to the main idea of PSLD. In our inpainting experiment, we found that equation (5) failed to generate coherent images at the boundary, equation (6) made the boundary smooth at the cost of extensive parameter tuning, and finally equation (7) mitigated these issues that resulted in high PSNR and SSIM. Quantitatively, we have already demonstrated the improvement of equation (7) over equation (6) in the Appendix B.2 (Table 5). Also, the newly added experiments in the attached PDF (Table 1) shows this improvement of equation (7) over equation (6). In the revised version, we will add experiments with equation (5) as another baseline. ### Concluding Remark We are happy to address any remaining points during the discussion phase.
I would like to thank your for your response. I still have some disagreements on certain points: **DPS does not sample the posterior in noisy setting - Many diffusion posterior sampling algorithms are not exact samplers.** What I meant is that the result proven for DPS in the paper does not have any practical implication; as soon as one departs from the assumptions made in the paper, DPS fails to sample from the simplest posterior. The authors can perhaps try it on simple toy examples; take a $d$-dimensional Gaussian mixture and observe only one coordinate for example. DPS will sample inside but also outside of the support of the posterior and there is no evidence on why this shouldn't be the case of the proposed algorithm. Hence, the assumptions made in the paper are very strong and yield results with questionnable meaning. Regarding the comment on MCMC, I strongly disagree with the conclusion. Assuming that MCMC algorithms only work on toy algorithms, methods such as DPS do not work on such toy examples and as such, there is no reason to believe that they sample approximately from the correct posterior in high dimensional examples if they fail on the simplest examples. The quantitative metrics that are usually used like the FID and LPIPS mean nothing in practice; they do not assert the quality of posterior sampling. They quantify how coherent the images are but not if they're accurate samples from the posterior. It is practically impossible to know if methods such as DPS or PSLD sample from the posterior as we do not have access to samples from the said posterior. The only way to get a rough idea is to use these methods on examples for which the posterior is available in closed form. I have one remaining question; from my experience DPS is very sensitive to the "step-size" $\zeta_i$. It can yield very bad results when chosen inappropriately. In the main paper there does not seem to be any discussion on this matter for your methods and I believe that this is quite important. Did the authors tune the parameter $\eta$? Furthermore, I understand that in the derivations the noise std is factored into $\eta$. But what I meant is that intuitively $\eta$ should depend, in the experiments, on the noise std. It seems to me that the noise std is the same on all experiments and that this matter is not discussed.
Discussion with Reviewer S3JW
Dear Reviewer S3JW, (1) **DPS does not sample the posterior in noisy setting.** Thanks for the comment, we now better understand your intent. We agree with your comment on DPS, and our result only holds in the linear noiseless case under restrictive assumptions. We will emphasize this in the paper. However, this stylized theory did produce one practically useful result: We wanted to solve linear inverse problems in the latent space to leverage the power of Stable diffusion and other pre-trained foundation models. A standard way of using DPS in the latent space failed, and this led us to design the gluing objective. Specifically, even in the linear noiseless case, (latent) DPS did not recover the ground-truth (due to lack of contraction to the unique fixed point), motivating the gluing objective to fix this issue. As we responded to reviewer yVXi, we will remove the pixel-space analysis (move it to the appendix), shorten the latent-space analysis, and emphasize it only holds for the noiseless case. We will further highlight that the analysis goal is to provide intuition on the gluing objective. (2) **Sampling from true posterior and metrics such as LPIPS and FID.** This is a fair point. There is no theoretical justification to assert that DPS or PSLD sample from the true posterior. However, in our paper we study the solution of inverse problems where the ground truth is known. This helps in validating the reconstruction produced by our algorithm, and thus pairwise error-metrics like MSE/PSNR and LPIPS can be used. These metrics are used, e.g. for MRI quality assessment and numerous other inverse problems, as non-perfect but still reasonable metrics to compare inverse problem solvers. Finally, as you rightly point out, there is no theory showing that metrics like FID or Inception score say anything about the quality of posterior sampling and this remains as an important research direction. (3) **Did the authors tune the parameter $\eta$? In the main paper there does not seem to be any discussion on this matter for your methods and I believe that this is quite important.** We tuned the hyper-parameter $\eta$ to the extent possible with our available computing resources. The reported results are with the best hyper-parameter we could find. We will revise the discussion to make this clear. (4) **It seems to me that the noise standard deviation is the same on all experiments and this matter is not discussed.** As rightly pointed out by the reviewer, the noise standard deviation is the same on all experiments. This is the same setting as the baseline DPS. We will add this discussion in the revised version. ### Concluding Remark We are happy to address any remaining points during the discussion phase.
Thank the authors for their thorough and thoughtful response. All my concerns have been touched and addressed.
Thanks for the rebuttal
I thank the authors for the rebuttal, especially the added comparisons with other methods, which further show the superiority of the proposed method. Please revise the future version accordingly. I have raised my score after reading the rebuttal.
Decision
Accept (poster)