Diffusion State-Guided Projected Gradient for Inverse Problems

Authors: Rayhan Zirvi, Bahareh Tolooshams, Anima Anandkumar
Primary area: generative models
Keywords: Diffusion models, Inverse problems, Robustness, Subspace, Projection, Box inpainting, Phase retrieval

Abstract:
Recent advancements in diffusion models have been effective in learning data priors for solving inverse problems. They leverage diffusion sampling steps for inducing a data prior while using a measurement guidance gradient at each step to impose data consistency. For general inverse problems, approximations are needed when an unconditionally trained diffusion model is used since the measurement likelihood is intractable, leading to inaccurate posterior sampling. In other words, due to their approximations, these methods fail to preserve the generation process on the data manifold defined by the diffusion prior, leading to artifacts in applications such as image restoration. To enhance the performance and robustness of diffusion models in solving inverse problems, we propose Diffusion State-Guided Projected Gradient (DiffStateGrad), which projects the measurement gradient onto a subspace that is a low-rank approximation of an intermediate state of the diffusion process. DiffStateGrad, as a module, can be added to a wide range of diffusion-based inverse solvers to improve the preservation of the diffusion process on the prior manifold and filter out artifact-inducing components. We highlight that DiffStateGrad improves the robustness of diffusion models in terms of the choice of measurement guidance step size and noise while improving the worst-case performance. Finally, we demonstrate that DiffStateGrad improves upon the state-of-the-art on linear and nonlinear image restoration inverse problems. Our code is available at https://github.com/Anima-Lab/DiffStateGrad.

Decision: accept (poster)
Reviewer scores: 8, 8, 8, 3

Reviews and discussion:
<thread>
creation date:
1729507679028

thread_type:
official_review

<thread_object>
creator: reviewer_3fsw

<content>
summary:
The authors propose DiffStateGrad, a method for projecting the gradients to a low-rank subspace through SVD, where the gradients are those arising from the diffusion models for inverse problem-solving (DIS) framework. The authors propose an adaptive thresholding value, where it is automatically calculated for every timestep $t$. It is shown that DiffStateGrad improves the performance compared to when it is not used. DiffStateGrad is shown to be especially useful for latent diffusion-based methods, which is understandable as gradients are noisier due to the existence of the decoder.

strengths:
1. The method is straightforward to understand and easy to implement.

2. The proposed framework can be used regardless of the base framework. In the three frameworks that the authors tested DiffStateGrad, it always shows improved performance.

weaknesses:
1. I believe DiffStateGrad should mainly be compared to approaches like MPGD [1] or DDS [2], which are highly related to this work. The authors do mention [1] as a related work, but says that DiffStateGrad is different to [1], as it considers a projection on $\mathcal{M}_t$ and not $\mathcal{M}$. However, care should be taken because applying gradient guidance to the posterior mean or $x_t$ are equivalent up to a constant. Hence, projecting the gradient guidance on MPGD has about the same effect. The same goes for [2], which does not explicitly project the gradient but uses a Krylov subspace method. I believe DiffStateGrad should be directly compared to MPGD.

2. The fact that guidance on the posterior mean and the guidance on $x_t$ only differs by a scale factor questions if the theory presented in the paper is actually useful.

3. Most of the experiments are conduced with latent diffusion based methods. More experiments should be conducted on other widely used pixel domain methods such as DPS and $\Pi$GDM.




**References**

[1] He, Yutong, et al. "Manifold preserving guided diffusion." ICLR 2024

[2] Chung, Hyungjin, et al., "Decomposed diffusion sampler for accelerating large-scale inverse problems", ICLR 2024

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681156401

thread_type:
official_review

<thread_object>
creator: reviewer_kb1e

<content>
summary:
This paper tackles the problem of solving from a posterior distribution defined with a diffusion-based prior. The authors assume that a diffusion model has been pre-trained and then use it to solve various inverse problems while not requiring anymore training. The method proposed can be thought of as enhancement for existing solvers; it can applied on top of any existing diffusion-based posterior sampling algorithm. It basically consists in adding a projection step, where the projection operator is computed based on the current diffusion state.

strengths:
- The paper is well written and the methodology is clearly presented. The fact that the proposed algorithm can be plugged on top of existing methods is a significant feat. The theoretical justification for the method is loose but the explanations provided are intuitive. 
- The experiments are convincing are well thought and rather extensive.

weaknesses:
- While the additional computational cost is marginal for latent diffusion models, isn't it prohibitive for pixel space diffusion models? Having to compute the SVD at each iteration is certainly a significant drawback. 
- The methodology is based on having at some point a sample on the manifold of "artifact free images" (this is loosely defined). It is unclear how this arises in practice. 
- There are some inconsistencies in the experiments which I believe are explained nowhere in the paper. First, the authors compare to DAPS in pixel space but not in latent space, although in the original paper DAPS is applied in latent space too. Second, in the only example involving imagenet, the authors compare to PSLD only and not ReSample nor PSLD. Also, no experiment on pixel space imagenet is provided. Is there any reason for this?

questions:
- While it is claimed that you test the robustness of the method to measurement noise, I am failing where this is done in the paper. 
- How would the method apply to more general methods, e.g. to methods that do not use the approximate guidance term, for example [1]. Would you then simply project the updated $x_t$ and reconstruct it, then move on to the next step? 

[1] Lugmayr, Andreas, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. "Repaint: Inpainting using denoising diffusion probabilistic models." In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11461-11471. 2022.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730683900052

thread_type:
official_review

<thread_object>
creator: reviewer_d85g

<content>
summary:
Authors propose Diffusion State-Guided Projected Gradient (DiffStateGrad), a module that increases robustness of existing diffusion based inverse problem solvers to both increase the performance and the robustness against choice of hyperparameters. DiffStateGrad projects the measurement gradients onto the low-rank approximation of intermediate states of the diffusion process (noisy manifolds for each timestep). Effectiveness of DiffStateGrad is demonstrated through multiple datasets (ImageNet and FFHQ), wide range of linear and non-linear inverse problems and applied to several SOTA posterior sampling algorithms (PSLD, ReSample and DAPS).

strengths:
* The paper is written very well. The work is contextualized well among related work and I enjoyed reading the paper.
* DiffStateGrad is formulated as a module that greatly increases robustness of existing SOTA posterior sampling approaches (such as PSLD and DAPS) with respect to the choice of step siz and measurement noise. 
* The effectiveness of DiffStateGrad is demonstrated in diverse set of forward models such as box inpainting, random inpainting, Gaussian deblur, motion deblur, etc. for linear inverse problems and phase retrieval, nonlinear deblur and HDR for nonlinear inverse problems.
* Additional cost of calculating SVD and projecting gradients seems negligible based on Figure 4.

weaknesses:
* See the questions below.

questions:
* In line 214, there is a brief comparison between MCG and DiffStateGrad, stating that one of them enforces iterates to stay close near $\mathcal{M}_t$ and the other one enforces closeness to $\mathcal{M}_0$. The fact that these methods are similar in terms of premise (measurement gradients throwing iterates off the manifold) and intuition, I think it deserves a more in depth comparison/discussion about their similarities and differences.
* How would DiffStateGrad interact with diffusion-based solvers that adopt earlier initialization strategies such as (CCDF [1] or it's adaptive version Adapt-and-Diffuse [2])? I would expect bigger step sizes are more detrimental in the "chaotic" regimes of reverse diffusion process where SNR is low. Would DiffStateGrad + CCDF combination should perform similar or better (albeit with less margin) than vanilla CCDF?
* Except for the robustness experiments, measurement noise level is not specified for the experiments (Table 2, 3, etc.). If noise was present, what was the variance and based on line 365 was it added to the images in the range [0,1] (this information is very useful for reproducibility in the future)? Does the presense/absense of measurement noise have any effect on the performance of DiffStateGrad?
  * On a related note, how does the measurement noise robustness experiment figures look when noise level is $<0.05$? 

***

[1] Chung, Hyungjin, Byeongsu Sim, and Jong Chul Ye. "Come-closer-diffuse-faster: Accelerating conditional diffusion models for inverse problems through stochastic contraction." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[2] Fabian, Zalan, Berk Tinaz, and Mahdi Soltanolkotabi. "Adapt and Diffuse: Sample-adaptive Reconstruction via Latent Diffusion Models." Proceedings of machine learning research 235 (2024): 12723.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711125183

thread_type:
official_review

<thread_object>
creator: reviewer_yapa

<content>
summary:
This paper present a novel method for adding an add-on of projected gradient during solving inverse problems with (latent) diffusion models. The method can be applied for an arbitrary inverse problem solver. Results show improvement over baselines such as PSLD and ReSample for solving inverse problems with LDMs.

strengths:
1. Well written
2. Easy to Understand
3. The idea of projecting the gradient to the manifold of intermediate noise is novel and making sense to me. This method supposes to suppress artifacts that arises with hard optimization.

weaknesses:
1. In some cases where gradient computation may incur some additional burdens (for example when PSLD takes a lot of memory), this method may not be feasible.
2. There are some other works that try to project the restoration gradient onto the prior manifold (for example DreamClean [1], and MCG [2]).
3. The baselines with LDMs are sufficient in my view, but this paper could benefits more with pixel diffusion baselines such as DDNM [3], DDRM and so on.


[1] Xiao, Jie, et al. "DreamClean: Restoring Clean Image Using Deep Diffusion Prior." The Twelfth International Conference on Learning Representations.

[2] Chung, Hyungjin, et al. "Improving diffusion models for inverse problems using manifold constraints." Advances in Neural Information Processing Systems 35 (2022): 25683-25696

[3] Wang, Yinhuai, Jiwen Yu, and Jian Zhang. "Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model." The Eleventh International Conference on Learning Representations.

questions:
1. Could authors provide more baselines results especially for pixel diffusion
2. I am curious about whether the authors consider choosing a good initial noise. Like [4, 5]


[4] Chung, Hyungjin, Byeongsu Sim, and Jong Chul Ye. "Come-closer-diffuse-faster: Accelerating conditional diffusion models for inverse problems through stochastic contraction." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[5] Fabian, Zalan, Berk Tinaz, and Mahdi Soltanolkotabi. "Adapt and Diffuse: Sample-adaptive Reconstruction via Latent Diffusion Models." Proceedings of machine learning research 235 (2024): 12723.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732209897616

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Authors' Summary of General Response and Modifications to the Manuscript (1)

comment:
We sincerely thank all the reviewers for their thoughtful and constructive feedback. We have carefully addressed each reviewer’s comments individually and incorporated their suggestions into the revised manuscript. Below, we provide a summary of the key modifications and additional analyses introduced in the paper. We believe these enhancements have strengthened our work and adequately addressed the reviewers' concerns. Finally, our changes to the manuscript can be found in color.

**G1. [Pixel-based comparison]**

In response to the reviewers' valuable suggestions, we have conducted additional experiments with "pixel-based" diffusion models and included the results in the revised manuscript. These experiments demonstrate that DiffStateGrad is not only applicable but also scalable to both latent- and pixel-based diffusion models. The results confirm that DiffStateGrad consistently enhances performance and outperforms prior work in the literature when applied to both pixel-based and latent-based models that use a guidance term to solve inverse problems. While we highlight the new pixel-based comparisons here, we kindly refer the reviewers to the manuscript for details on our initial experiments with latent diffusion models.

We have combined the previous Tables 2, 3, and 4 into a new Table 3, which now includes results for all linear and nonlinear inverse problems. In addition, we conducted new experiments for DPS [1], DiffStateGrad-DPS, and MPGD-AE [2], and reported results for DDNM [3], DDRM [4], and MCG [5] on linear inverse problems using the FFHQ dataset. 

Furthermore, we included new pixel-based quantitative results in Table 4a, showcasing comparisons among DAPS [6], DiffStateGrad-DAPS, DPS, DiffStateGrad-DPS, and MPGD-AE [2] for box inpainting and random inpainting on ImageNet. Lastly, Table 4b presents additional results for DAPS, DiffStateGrad-DAPS, and MCG on Gaussian deblurring, motion deblurring, and super-resolution inverse problems conducted on ImageNet.

Below, we provide a summary of the results.

**New Results in Table 3: Performance comparison of linear tasks on FFHQ 256 x 256.**

| Method                        | LPIPS ↓ (Box-In) | PSNR ↑ (Box-In) | LPIPS ↓ (Random-In) | PSNR ↑ (Random-In) | LPIPS ↓ (Gaussian-Db) | PSNR ↑ (Gaussian-Db) | LPIPS ↓ (Motion-Db) | PSNR ↑ (Motion-Db) | LPIPS ↓ (SR) | PSNR ↑ (SR) |
|-------------------------------|---------------|--------------|------------------|-----------------|--------------------|-------------------|------------------|-----------------|--------------|-------------|
| **Pixel-based**               |               |              |                  |                 |                    |                   |                  |                 |              |             |
| DAPS                          | 0.136         | 24.57        | 0.130            | 30.79           | 0.216              | 27.92             | 0.154            | 30.13           | 0.197        | 28.64       |
| DiffStateGrad-DAPS (ours)     | **0.113**     | **24.78**    | **0.099**        | **32.04**       | 0.180          | **29.02**         | 0.119        | **31.74**       | 0.181    | **29.35**   |
| DPS                           | 0.127         | 23.91        | 0.130            | 28.67           | 0.145              | 25.48             | 0.132            | 26.75           | 0.191        | 24.38       |
| DiffStateGrad-DPS (ours)      | 0.114     | 24.10    | 0.107        | 30.15      | **0.128**          | 26.29         | **0.118**        | 27.61      | 0.186    | 24.65   |
| DDNM                          | 0.235         | 24.47        | 0.121            | 29.91           | 0.216              | 28.20             | -                | -               | 0.197        | 28.03       |
| DDRM                          | 0.159         | 22.37        | 0.218            | 25.75           | 0.236              | 23.36             | -                | -               | 0.210        | 27.65       |
| MCG                           | 0.309         | 19.97        | 0.286            | 21.57           | 0.340              | 6.72              | 0.702            | 6.72            | 0.520        | 20.05       |
| MPGD-AE                       | 0.138         | 21.59        | 0.172            | 25.22           | 0.150              | 24.42             | 0.120            | 25.72           | **0.168**    | 24.01   |
</content>
</thread_object>

</thread>
<thread>
creation date:
1732211779084

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Authors' Summary of Manuscript Modifications

comment:
We thank all reviewers for their constructive feedback. Given our extensive new analyses, we have made several changes to the manuscript.  We refer the reviewers to the changes in color green in the manuscript, and summarize below the key changes

- New Table 2: We have created a new Table 2 that categorizes prior works based on gradient guidance computation, incorporation, and projection. This table highlights that DiffStateGrad is applicable to both $x_t$/$z_t$ and $x_{0|t}$/$z_{0|t}$.

- Revised Table 3: Tables 2, 3, and 4 have been merged into a single Table 3. Table 3a now includes both pixel-based (6 new rows) and latent-based diffusion models for all linear tasks on FFHQ. Table 3b reports our previous results for nonlinear tasks on FFHQ. This table now focuses on LPIPS and PSNR metrics, with SSIM values moved to the Appendix (see Table 6).

- New Table 4a: This table contains results for both pixel-based (5 new rows) and latent-based diffusion models for inpainting tasks on ImageNet.

- New Table 4b: This table includes results for pixel-based (3 new rows) and latent-based diffusion models for deblurring and super-resolution tasks on ImageNet.

- Figure Reorganization: To accommodate the new results, we have moved the original Figure 7 (“Performance of DiffStateGrad-DAPS for different variance retention thresholds”) to the Appendix, where it is now labeled as Figure 9.

 - Expanded Related Works: The Related Works section has been updated with additional information on MCG and MPGD, clarifying how our proposed framework differs from these methods.

- Reproducibility: We have included additional details on the new experiments in the Appendix to support reproducibility.

- Additional Citations: We have added citations for methods suggested by reviewers: [1-2]: Approaches that accelerate or improve diffusion-based inverse solvers through initialization strategies. [3]: A method focused on inpainting inverse problems using a gradient-free masking approach to generate only missing pixels.

- Visualization Adjustments: To save space for additional experimental results, half of the visualizations from Figure 7 have been moved to the Appendix.

- Figure 10: We have included an additional figure for the numbers we report on the impact of the projection period (See G3).

We hope the reviewers find this summary helpful, and we refer the reviewer to the manuscript for more details.

[1] Chung, H., Sim, B., & Ye, J. C. (2022). Come-closer-diffuse-faster: Accelerating conditional diffusion models for inverse problems through stochastic contraction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 12413-12422).

[2] Fabian, Z., Tinaz, B., & Soltanolkotabi, M. (2024). Adapt and Diffuse: Sample-adaptive Reconstruction via Latent Diffusion Models. Proceedings of machine learning research, 235, 12723.

[3] Lugmayr, A., Danelljan, M., Romero, A., Yu, F., Timofte, R., & Van Gool, L. (2022). Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 11461-11471).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732389604699

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional experiment on MRI inverse problem: showing applicability beyond natural images when Langevin-based model is used

comment:
We would like to thank the reviewer yaPA for their thoughtful comment and for raising their scores. As suggested by reviewer yaPA, we have conducted additional experiments. This is to demonstrate the applicability of DiffStateGrad beyond natural image tasks and to showcase its compatibility with other guidance-based Langevin dynamics used for solving inverse problems.

We conducted additional experiments on Compressed-Sensing Magnetic Resonance Imaging (MRI). This represents an inverse problem where undersampled measurements are observed in the frequency domain, and the task is to reconstruct high-quality MR images in the image space.

For these experiments, we utilized the Compressed Sensing Generative Model (CSGM) proposed in [3], which employs Langevin dynamics for MRI reconstruction (available at https://github.com/utcsilab/csgm-mri-langevin). We incorporated our proposed DiffStateGrad into the measurement gradient guidance. Given the complex nature of MRI data, we applied DiffStateGrad to the magnitude of the complex-valued data while preserving the phase information. This framework uses $x_t$ for both measurement gradient computation and incorporation. The unconditional diffusion model was trained on T2-weighted brain datasets from the NYU fastMRI dataset [1-2], and the reported results were averaged over 30 test examples (reporting avg (std)). The measurements were sampled vertically and equispaced at an undersampling rate of $R=4$.

While the default optimized measurement gradient step size from [3] is $5.0$, our results confirm that DiffStateGrad enhances the robustness of CSGM-Langevin when the step size deviates from this optimal value. For example, CSGM-Langevin's performance is significantly reduced when measurement gradient step size increases (see step size = 10), while DiffStateGrad's performance decline is less. This suggests that DiffStateGrad can be a practical tool to reduce the risk of drastic failure in real-world applications.

We will update the manuscript shortly to include these new results. Given the small sample size and the preliminary nature of this analysis, we will include these results in the Appendix, provide a brief reference to them in the main text and leave extensive analysis for future work.


**Step size = 2.5**
| Method                        | PSNR ↑ | SSIM ↑ | NMSE  ↓ |  LPIPS ↓ |
|-------------------------------|---------------|--------------|------------------|------------------|
CSGM-Langevin | 34.85 ( 4.4739 ) | 0.8345 ( 0.1308 ) | 0.03072 ( 0.0709 ) | 0.1474 ( 0.0618 )|
DiffStateGrad-CSGM-Langevin (Ours) | 34.78 ( 4.3892 ) | 0.8423 ( 0.1277 )| 0.02881 ( 0.0637 ) | 0.153 ( 0.0745 ) |


**Step size = 5.0** (the default value chosen by CSMRI-Langevin) 
| Method                        | PSNR ↑ | SSIM ↑ | NMSE  ↓ |  LPIPS ↓ |
|-------------------------------|---------------|--------------|------------------|------------------|
CSGM-Langevin | 34.31 ( 3.5739 ) | 0.796 ( 0.1119 ) | 0.02922 ( 0.0701 ) | 0.1797 ( 0.0546 ) |
DiffStateGrad-CSGM-Langevin (Ours) | 33.75 ( 3.9638 ) | 0.8021 ( 0.1256 ) | 0.03241 ( 0.0643 ) | 0.1716 ( 0.0611 )|


**Step size = 7.5**
| Method                        | PSNR ↑ | SSIM ↑ | NMSE  ↓ |  LPIPS ↓ |
|-------------------------------|---------------|--------------|------------------|------------------|
CSGM-Langevin | 32.88 ( 2.9101 ) | 0.723 ( 0.1023 ) | 0.03324 ( 0.0723 ) | 0.2465 ( 0.0593 )|
DiffStateGrad-CSGM-Langevin (Ours) | 33.21 ( 3.3072 ) | 0.7768 ( 0.1156 ) | 0.03043 ( 0.053 ) | 0.1882 ( 0.051 ) |


**Step size = 10.0**
| Method                        | PSNR ↑ | SSIM ↑ | NMSE  ↓ |  LPIPS ↓ |
|-------------------------------|---------------|--------------|------------------|------------------|
CSGM-Langevin | 31.38 ( 2.5123 ) | 0.6552 ( 0.0948 ) | 0.04107 ( 0.0769 ) | 0.3014 ( 0.0599 ) |
DiffStateGrad-CSGM-Langevin (Ours) | 32.49 ( 3.0165 ) | 0.7386 ( 0.1136 ) | 0.03181 ( 0.0478 ) | 0.2111 ( 0.0491 )|


[1] Knoll, F., Zbontar, J., Sriram, A., Muckley, M. J., Bruno, M., Defazio, A., ... & Lui, Y. W. (2020). fastMRI: A publicly available raw k-space and DICOM dataset of knee images for accelerated MR image reconstruction using machine learning. Radiology: Artificial Intelligence, 2(1), e190007.

[2] Zbontar, J., Knoll, F., Sriram, A., Murrell, T., Huang, Z., Muckley, M. J., ... & Lui, Y. W. (2018). fastMRI: An open dataset and benchmarks for accelerated MRI. arXiv preprint arXiv:1811.08839.

[3] Jalal, A., Arvinte, M., Daras, G., Price, E., Dimakis, A. G., & Tamir, J. (2021). Robust compressed sensing mri with deep generative priors. Advances in Neural Information Processing Systems, 34, 14938-14954.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The starting point of the paper is very reasonable. Specifically, the authors aim to relieve the issue that previous methods fail to preserve the generation process on the data manifold defined by the diffusion prior when solving the inverse problem. The authors empirically find that the crucial factor is to choose the low-rank subspace of the latent image representation for containing the gradient, which is meaningful. Three of the reviewers provided very positive feedback, while one has doubts about the paper’s contributions. 
The AC tends to agree with the opinions of the first three reviewers that the overall quality of this paper is good. However, after discussing with the reviewers, we believe that the assumption of Proposition 1 may be too strong, which may be misleading, and the claimed theoretical contribution to be reasonable. Thus, I only recommend the paper to be accepted as a  poster paper.