Diffusion State-Guided Projected Gradient for Inverse Problems

We thank the reviewer for their thoughtful comments. First, we would like to refer the reviewer to our general summary of response post. The post summarizes the key points and additional experiments that we have now included in the paper. Below, we provide our response to each of the reviewer's comments.

1. [Comparison of DiffStateGrad projection to other projection methods]

We thank the reviewer for raising this point. We have included additional explanations and experimental results on explaining how our method differs from MPGD [1] and MCG [2].

• MPGD [1] employs a manifold-preserving approach to enhance the efficiency of diffusion generation and solve inverse problems. While this method aligns with a similar motivation as our proposed framework, it is specifically designed for cases where the measurement gradient is applied to $x_{0|t}$ or $z_{0|t}$. In this scenario, MPGD suggests projecting the gradient onto the tangent space of the clean data manifold. Its most effective implementation, MPGD-AE, utilizes an autoencoder (AE) for this manifold projection. However, this approach is limited by the need for an autoencoder, and its performance heavily depends on the expressive power of the autoencoder used. In contrast, DiffStateGrad is more versatile, as it can be applied to methods that use guidance on $z_{0|t}$/$x_{0|t}$ (e.g., ReSample [7] and DAPS [6]) as well as methods that apply the guidance to $x_{t}$/$z_{t}$ (e.g., PSLD [8] and DPS [1]). Notably, MPGD does not support the latter scenario. For the former scenario, we agree with a reviewer that both approaches of MPGD and DiffStateGrad can be compared and the comparison can highlight the effectiveness of the choice of projection step. Our approach constructs the projection using SVD to capture the data structure and does not need a trained AE. We have clarified this distinction in the revised manuscript and included additional explanations in the Related Works section. We conclude that our results from new Tables 3 and 4 confirm that among pixel-based methods using the guidance to update $x_{0|t}$, DiffStateGrad-DAPS outperforms DAPS and MPGD-AE. Moreover, while MPGD achieves LPIPS metrics that are comparable to or slightly better than DAPS for FFHQ, its PSNR performance is significantly lacking. To further highlight MPGD-AE's dependency on the expressive power of the autoencoder (AE) used, we refer the reviewer to our new results on MPGD-AE for inpainting tasks on ImageNet, where MPGD-AE performs very poorly compared to both DAPS and our DiffStateGrad-DAPS.

• We agree with the reviewer that projecting onto the clean data manifold for $x_{0|t}$ is equivalent to projecting onto $M_t$ of $x_t$ up to a constant. We clarify that the key aspect of our paper is that the projection should be defined based on the variable to which the gradient guidance is being applied. For example, if the guidance is applied to $x_{0|t}$, the projection should be defined based on $x_{0|t}$; this is similar to the intuition behind MPGD, and is applicable to methods such as ReSample and DAPS. Moreover, if the guidance is applied to $x_t$, the projection should be defined based on $x_t$. This latter scenario is applicable to approaches like DPS and PSLD. Importantly, our DiffStateGrad framework offers performance improvements in both cases.

• Lastly, we would like to refer the reviewer to our general response G2, where we include additional explanations on how our method differs from MCG, which also uses a manifold projection concept.

2. [Theory]

We thank the reviewer for raising this important point, as elaborating on it helps to improve the clarity and quality of the paper. The theoretical analysis in this work supports the argument that the measurement gradient should be projected onto the tangent space of the manifold at time $t$ (if the gradient update is applied to $z_t$ or $x_t$) or at time $0$, corresponding to the clean data manifold (if the gradient update is applied to $z_{0|t}$ or $x_{0|t}$).

This intuition aligns with the reasoning discussed in MPGD [1], which advocates for projecting the gradient onto the tangent space of the clean data manifold. While our analysis primarily formulates the problem for scenarios where the measurement gradient is applied to update $z_t$, we have now clarified in the paper that, in cases where the gradient update is applied to $z_{0|t}$, the projection should be onto the tangent space of $z_{0|t}$. This latter scenario is indeed used for DiffStateGrad-ReSample and DiffStateGrad-DAPS, and it is analogous to the approach taken in MPGD. We hope that this clarification has addressed the reviewer's concern.

We sincerely thank the reviewer for their thoughtful feedback. We kindly refer the reviewer to our general comment addressed to all reviewers, where we summarize the new experiments we conducted and highlight the key points of our work. Below, we provide a detailed response to the reviewer's specific concerns, addressing issues related to computational memory requirements for gradient computation, comparisons with other projection-based methods, additional pixel-based methods, and the adoption of strategies for selecting a good initial diffusion state or noise.

1. [Computational memory]

DiffStateGrad targets a class of diffusion models that leverage measurement gradient guidance to solve inverse problems. These methods, such as DPS [1], PSLD [2], and DAPS [3], were state-of-the-art at the time of their release, and DiffStateGrad is directly applicable to these gradient-based guidance approaches. We recognize that in certain applications, the memory required for backpropagating gradients in PSLD or DPS might limit their applicability. In such cases, our method can be applied to DAPS [3], a diffusion-based approach where the measurement gradient is derived based on $x_{0|t}$ rather than $x_t$, offering a more computationally efficient alternative. In this setting, our results demonstrate that DiffStateGrad-DAPS consistently outperforms DAPS. This is evident in our initial comparisons for FFHQ (Table 3) as well as in the new comparisons on ImageNet (Table 4), where DiffStateGrad-DAPS exhibits superior performance. For the set of new results, please see our Tables in our general response (G1) above.

Finally, we note that our implementation allows DiffStateGrad to be applied less frequently during diffusion iterations. For instance, in ReSample, DiffStateGrad is applied every five iterations. To further illustrate the impact of this choice, we conducted an additional ablation study, varying the projection frequency of DiffStateGrad-DAPS for random inpainting tasks on FFHQ. The results which we included in the general response (G3) indicate that reducing the frequency of projection still leads to performance improvements, demonstrating the flexibility and efficiency of our approach.

2. [Comparison to other projection-based diffusion models]

To emphasize the subspace choice, recommended by the reviewer, we have included additional comparison results with MCG [4] and MPGD [5] (see G1 in our general response). Below, we re-state our general response. The above-mentioned methods differ as follows in their respective choices of projection:

• MPGD [4] employs a manifold-preserving approach to enhance the efficiency of diffusion generation and solve inverse problems. While this method aligns with a similar motivation as our proposed framework, it is specifically designed for cases where the measurement gradient is applied to $x_{0|t}$ or $z_{0|t}$. In this scenario, MPGD suggests projecting the gradient onto the tangent space of the clean data manifold. Its most effective implementation, MPGD-AE, utilizes an autoencoder (AE) for this manifold projection. However, this approach is limited by the need for an autoencoder, and its performance heavily depends on the expressive power of the autoencoder used. In contrast, DiffStateGrad is more versatile, as it can be applied to methods that use guidance on $x_{0|t}$/$z_{0|t}$ (e.g., ReSample [7] and DAPS [6]) as well as methods that apply the guidance to $x_{t}$/$z_{t}$ (e.g., PSLD [8] and DPS [1]). Notably, MPGD does not support the latter scenario. Moreover, our approach constructs the projection using SVD to capture the data structure and does not need a trained AE. We have clarified this distinction in the revised manuscript and included additional explanations in the Related Works section. We conclude that our results from new Tables 3 and 4 confirm that among pixel-based methods using the guidance to update $x_{0|t}$, DiffStateGrad-DAPS outperforms DAPS and MPGD-AE.

• MCG [5] introduces a manifold constraint that projects the measurement gradient onto the data manifold $M_0$ corresponding to $x_{0|t}$ while applying the guidance to update $x_t$. In essence, MCG uses the concept of projection to map the gradient onto the data manifold related to $x_0$ but subsequently uses this projected gradient to update $x_t$. In contrast, our work argues that when the gradient of the measurement fidelity is used to update $x_t$, it should be projected onto a manifold that represents the noisy data at time $t$, i.e., ${\mathcal{M}}_t$. This ensures that the updates respect the state of the diffusion process at $t$. Our method is more effective as it projects the measurement gradient—used to update $x_t$—onto the noisy diffusion state associated with $x_t$, thereby preserving the $t$-state on ${\mathcal{M}}_t$ rather than ${\mathcal{M}}_0$.

Please see G2 or Tables 3 and 4 for additional experiments showing that DiffStateGrad is superior to MCG and MPGD-AE.

The authors address my questions well, the comparison with pixel diffusion models (e.g. + DPS/+DAPS) demonstrates that this method can generalize to both LDM and pixelDM. I am looking forward to seeing whether this method can be combined with inversion techniques / other types of models such as flow-based models. I will raise my score.

We thank the reviewer for their thoughtful comments and appreciation of our extensive experimental results and intuitions on the proposed method. We kindly refer the reviewer to our general comment addressed to all reviewers, where we summarize the new experiments we conducted and highlight the key points of our work. Below, we provide a detailed response to the reviewer's specific concerns, addressing issues related to more experiments on pixel-based diffusion models, the intuition on the artifact-free images, latent DAPS, robustness to measurement noise, and comparison to gradient-free methods such as RePaint.

1. [Pixel-based Diffusion Models]

We thank the reviewer for raising this concern.

• We now include extensive additional pixel-based diffusion models in Tables 3 and 4 (please see our G1 general response in the comment to all reviewers). The results confirm that DiffStateGrad improves performance. This supports the scalability and applicability of DiffStateGrad to pixel-based diffusion models.

• Moreover, we note that our implementation includes a projection period parameter, which allows SVD computations and projections to be performed less frequently. For instance, while our experiments with DAPS apply projections at every iteration, the ReSample experiments use projections every five iterations. This information has been included in the Appendix for reference. We have conducted an additional set of experiments, where DiffStateGrad is applied to DAPS every 5 or 10 iterations. The results which we report below suggest that DiffStateGrad can be applied at less frequent manner for efficiency purposes and still improve performance upon the DAPS without projection. See our response in G3 for more detailed information. These results can be found in the new Figure 10 in the Appendix. Overall, the results, summarized below, demonstrate the flexibility and efficiency of our approach.

where P denotes the projection period.

• Lastly, we kindly refer the reviewer to our efficiency analysis presented in Figure 4, where we demonstrate that the additional runtime complexity introduced by DiffStateGrad is minimal relative to the overall runtime of the diffusion solver. In this analysis, DAPS, a pixel-based method where the projection happens as every iteration (P=1). We reason that for P>1, this runtime should even be more negligible.

Overall, we deeply appreciate the reviewer's insightful concern and acknowledge that future work on DiffStateGrad could explore more efficient implementations of SVD. Additionally, while we currently use SVD to identify a low-rank structure of the diffusion state ($x_{t}$ for DPS, $x_{0|t}$ for DAPS), future research could investigate alternative approaches for identifying subspaces that effectively represent the diffusion state and the local manifold structure for projection purposes.

2. [Artifact-free images]

The literature on using diffusion models to solve inverse problems operates under the assumption that there exists a diffusion model unconditionally trained on a dataset (e.g., clean images). The term "artifact-free" refers to the generation process of such an unconditional diffusion model, trained on clean data samples, which provides an artifact-free trajectory from $\mathcal{M}_T$ to $\mathcal{M}_0$. For further intuition on projection onto the tangent space of the clean data manifold, we refer the reviewer to MPGD [4].

In our work, the reference to artifact-free images pertains specifically to this learned path, which captures general "clean and artifact-free" images, rather than blurred images or images with missing pixels. The intuition behind our framework is to define a projection step that minimizes movements caused by the measurement gradient that are orthogonal to the manifold locally, either at $x_t$ (if the measurement guidance is applied to $x_t$) or at $x_{0|t}$ (if the measurement guidance is applied to $x_{0|t}$).

Focusing on the former case (without loss of generality), given a single sample $x_t$ from the manifold $\mathcal{M}_t$, we use $x_t$ to approximate and capture the local structures of $\mathcal{M}_t$. These points have now been included in the revised manuscript, and we hope this explanation clarifies the reviewer’s comment.

We hope that our response has adequately addressed the reviewer's concerns. We look forward to hearing from the reviewer on whether their concerns have been resolved or if there are additional ways in which this paper can be further improved.

[1] Chung, H.d , Kim, J., Mccann, M. T., Klasky, M. L., & Ye, J. C. Diffusion Posterior Sampling for General Noisy Inverse Problems. In The Eleventh International Conference on Learning Representations.

[2] Rout, L., Raoof, N., Daras, G., Caramanis, C., Dimakis, A., & Shakkottai, S. (2024). Solving linear inverse problems provably via posterior sampling with latent diffusion models. Advances in Neural Information Processing Systems, 36.

[3] Zhang, B., Chu, W., Berner, J., Meng, C., Anandkumar, A., & Song, Y. (2024). Improving diffusion inverse problem solving with decoupled noise annealing. arXiv preprint arXiv:2407.01521.

We sincerely thank the reviewer for their thoughtful and positive feedback. We kindly refer the reviewer to our general comment addressed to all reviewers, where we summarize the new experiments we conducted and highlight the key points of our work. Below, we provide a detailed response to the reviewer's specific concerns, addressing issues related to comparisons with other projection-based methods, initialization strategies, and robustness to measurement noise.

1. [Comparison to MCG]

We thank the reviewer for their comment. Below, we discuss the similarities and differences between MCG and the proposed frameworks.

• MCG introduces a manifold constraint that projects the measurement gradient onto the data manifold $M_0$ corresponding to $x_{0|t}$ while applying the guidance to update $x_t$. In essence, MCG uses the concept of projection to map the gradient onto the data manifold related to $x_0$ but subsequently uses this projected gradient to update $x_t$. In contrast, our work argues that when the gradient of the measurement fidelity is used to update $x_t$, it should be projected onto a manifold that represents the noisy data at time $t$, i.e., ${\mathcal{M}}_t$. This ensures that the updates respect the state of the diffusion process at $t$. Our method is more effective as it projects the measurement gradient—used to update $x_t$—onto the noisy diffusion state associated with $x_t$, thereby preserving the $t$-state on ${\mathcal{M}}_t$ rather than ${\mathcal{M}}_0$.

Please see G2 in our general response or Tables 3 and 4 for additional experiments showing that DiffStateGrad is superior to MCG. We refer the reviewer to the comparison between MCG, DPS [1], and DiffStateGrad-DPS to evaluate the impact of our proposed projection (see DPS [1] for a detailed discussion on how DPS is a generalization of the MCG method). Results show that DiffStateGrad is superior to MCG.

2. [DiffStateGrad when added to approaches using initialization strategies]

We thank the reviewer for highlighting the relevant literature on improving performance through better initialization strategies. We agree that such approaches can enhance the performance of solving inverse problems; however, we deliberately chose not to include these steps—such as "choosing a good initial point with the help of the measurement data"—in our experiments. This decision was made to ensure a fair comparison between each method and its corresponding DiffStateGrad without introducing additional variables. We now cite these references as an avenue to further improve DiffStateGrad for solving inverse problems.

3. [Robustness to measurement noise]

We clarify that the experiments in Tables 3-4 are conducted when there is a small measurement noise of 0.05 std. This noise level is chosen to follow a similar experiment setup to DPS [1], PSLD [2], and DAPS [3]. This is to show the impact of DiffStateGrad projection when there is little noise and the Signal-to-Noise ratio (SNR) is high. Moreover, our experiments in Figure 6 show the impact of measurement noise on the performance and the performance in general in low SNR settings. The figure shows that DiffStateGrad results in less performance decline (blue solid line) than no projection (red dashed line). We have now updated Figure 6 to include an additional point on the x-axis where the measurement noise is 0. The figure, which we show its results in numbers below, suggests that as the measurement noise level increases, the performance degradation of methods can be ameliorated by adding DiffStateGrad.

We thank the reviewers for their constructive feedback. We hope that the above discussion, along with our detailed responses to the reviewers' feedback, have adequately addressed their concerns. We look forward to hearing from the reviewers on whether their concerns have been resolved or if there are additional ways in which this paper can be further improved.

"Prior works differ from one another in two key aspects: (a) how the gradient is computed and (b) how the gradient is used to update the diffusion state. Table 2 categorizes prior works based on these characteristics. While a few approaches, such as diffusion-based MRI, compute the measurement gradient using $x_{t}$, most recent literature has shifted toward using $x_{0|t}$ for gradient computation. Regarding gradient incorporation, the literature is further subdivided. For instance, methods like DPS and PSLD use the measurement gradient to update the state at time $t$, whereas other approaches, such as ReSample, DAPS, and MPGD, apply the guidance to the conditional state at $0|t$ before resampling. Additionally, while the projections in MCG and MPGD are restricted to $x_{0|t}$, DiffStateGrad stands out by being applicable to both types of methods."

New Table 2

[1] Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., & Ye, J. C. Diffusion Posterior Sampling for General Noisy Inverse Problems. In The Eleventh International Conference on Learning Representations.

[2] He, Y., Murata, N., Lai, C. H., Takida, Y., Uesaka, T., Kim, D., ... & Ermon, S. Manifold Preserving Guided Diffusion. In The Twelfth International Conference on Learning Representations.

[3] Wang, Y., Yu, J., & Zhang, J. Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model. In The Eleventh International Conference on Learning Representations.

[4] Kawar, B., Elad, M., Ermon, S., & Song, J. (2022). Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35, 23593-23606.

[5] Chung, H., Sim, B., Ryu, D., & Ye, J. C. (2022). Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35, 25683-25696.

[6] Zhang, B., Chu, W., Berner, J., Meng, C., Anandkumar, A., & Song, Y. (2024). Improving diffusion inverse problem solving with decoupled noise annealing. arXiv preprint arXiv:2407.01521.

[7] Song, B., Kwon, S. M., Zhang, Z., Hu, X., Qu, Q., & Shen, L. Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency.

[8] Rout, L., Raoof, N., Daras, G., Caramanis, C., Dimakis, A., & Shakkottai, S. (2024). Solving linear inverse problems provably via posterior sampling with latent diffusion models. Advances in Neural Information Processing Systems, 36.

G3. [Efficiency of the proposed method]

In response to concerns regarding the efficiency of our projection method as image size increases, we would like to highlight the following points. First, our pixel-based analysis demonstrates that DiffStateGrad applies to both pixel-based diffusion models (e.g., DPS and DAPS) and latent-based diffusion models (e.g., PSLD and ReSample). Additionally, our implementation includes a projection period parameter, which allows projections to be performed less frequently. For instance, while our experiments with DAPS apply projections at every iteration, the ReSample experiments use projections every five iterations. This information has been included in the Appendix for reference.

We also conducted a new ablation study on DiffStateGrad-DAPS for random inpainting tasks on FFHQ where we find that increasing the projection period (decreasing the frequency) still leads to performance improvement. We include these results in the new Figure 10 in the Appendix. The results, summarized below, demonstrate the flexibility and efficiency of our approach.

where P denotes the projection period.

G4. [Increasing Robustness of Diffusion-based Inverse Solvers]

Finally, we would like to re-emphasize one of the key practical benefits of DiffStateGrad: its ability to increase the robustness of diffusion-based models to the choice of gradient step size as well as measurement noise. While the literature continues to propose new and more sophisticated approaches, their complexity and the extensive hyperparameter tuning they require pose significant challenges for experimentalists and scientists looking to adopt these frameworks in their fields. In contrast, DiffStateGrad represents a practical step towards making diffusion models more accessible and user-friendly.

Figure 5 illustrates the increased robustness to measurement gradient step size in PSLD introduced by DiffStateGrad. Furthermore, Figure 6a highlights the improved robustness of DAPS when DiffStateGrad is applied, particularly against measurement noise levels for box inpainting. Similarly, Figure 6b demonstrates the enhanced robustness of PSLD against measurement noise levels for Gaussian deblurring when using DiffStateGrad. Given one reviewer's suggestion, Figure 6 is now modified to include results when the measurement noise is 0. The figure, which we show its results in numbers below, suggests that as the measurement noise level increases, the performance degradation of methods can be ameliorated by adding DiffStateGrad.

G5. [Prior works guidance computation and guidance incorporation]

To provide further clarity on how our method relates to prior works, we have added an additional paragraph to the Background and Related Works section. This paragraph categorizes prior works based on three key criteria: (a) how the measurement guidance gradient is computed (GC), (b) how the gradient guidance is incorporated into the diffusion process (GI), and (c) whether the method utilizes a projection.

For convenience, we restate this new paragraph in the next comment, along with the updated Table 2.

• MPGD [2] employs a manifold-preserving approach to enhance the efficiency of diffusion generation and solve inverse problems. While this method aligns with a similar motivation as our proposed framework, it is specifically designed for cases where the measurement gradient is applied to $x_{0|t}$ or $z_{0|t}$. In this scenario, MPGD suggests projecting the gradient onto the tangent space of the clean data manifold. Its most effective implementation, MPGD-AE, utilizes an autoencoder (AE) for this manifold projection. However, this approach is limited by the need for an autoencoder, and its performance heavily depends on the expressive power of the autoencoder used. In contrast, DiffStateGrad is more versatile, as it can be applied to methods that use guidance on $x_{0|t}$/$z_{0|t}$ (e.g., ReSample [7] and DAPS [6]) as well as methods that apply the guidance to $x_{t}$/$z_{t}$ (e.g., PSLD [8] and DPS [1]). Notably, MPGD does not support the latter scenario. Moreover, our approach constructs the projection using SVD to capture the data structure and does not need a trained AE. We have clarified this distinction in the revised manuscript and included additional explanations in the Related Works section. We conclude that our results from new Tables 3 and 4 confirm that among pixel-based methods using the guidance to update $x_{0|t}$, DiffStateGrad-DAPS outperforms DAPS and MPGD-AE.

• MCG [5] introduces a manifold constraint that projects the measurement gradient onto the data manifold $M_0$ corresponding to $x_{0|t}$ while applying the guidance to update $x_t$. In essence, MCG uses the concept of projection to map the gradient onto the data manifold related to $x_0$ but subsequently uses this projected gradient to update $x_t$. In contrast, our work argues that when the gradient of the measurement fidelity is used to update $x_t$, it should be projected onto a manifold that represents the noisy data at time $t$, i.e., ${\mathcal{M}}_t$. This ensures that the updates respect the state of the diffusion process at $t$. Thus, our method is more effective as it projects the measurement gradient—used to update $x_t$—onto the noisy diffusion state associated with $x_t$, thereby preserving the $t$-state on ${\mathcal{M}}_t$ rather than ${\mathcal{M}}_0$ (See G5 and the new Table 2).

We also refer readers to DPS [1], which discusses how DPS generalizes MCG. We finally conclude that DiffStateGrad-DPS consistently outperforms both DPS and MCG, as evidenced by the new results presented in Table 3.

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

New Results in Table 4: Performance comparison of linear tasks on ImageNet 256 x 256.

Inpainting tasks.

Deblurring and super-resolution tasks.

New Results in Table 6: SSIM comparison of linear tasks on FFHQ 256 x 256.

G2. [How our method differs from other projection method]

Our initial Table 1 provides a brief study on the impact of subspace selection. The results indicate that the improvement in performance is not solely due to the low-dimensional nature of the gradient but rather the specific choice of projection. The table demonstrates that projecting onto the state of the diffusion process, to which the gradient guidance is applied, significantly enhances performance compared to random projections or low-rank approximations of the gradient.

To further highlight the importance of subspace selection, we have included additional comparative results with MCG [5] and MPGD [2] (see above tables). These methods differ as follows in their respective choices of projection:

We have now incorporated the above into the manuscript. Please see Appendix E.