Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Diffusion State-Guided Projected Gradient for Inverse Problems

Rayhan Zirvi, Bahareh Tolooshams, Anima Anandkumar

TL;DR

This work tackles inverse problems with unconditional diffusion priors by addressing the intractable measurement likelihood that can cause artifacts during diffusion-based sampling. The proposed Diffusion State-Guided Projected Gradient (DiffStateGrad) projects the measurement guidance gradient onto a low-rank subspace derived from the intermediate diffusion state, effectively enforcing tangent-space consistency with the data manifold. The method is demonstrated to boost robustness to guidance step sizes and noise, reduce failure rates, and improve reconstruction quality across both linear and nonlinear image restoration tasks, while incurring minimal computational overhead. DiffStateGrad is versatile and compatible with a range of diffusion-based solvers, offering a practical pathway to more reliable diffusion-driven inverse problem solutions with broad potential impact in imaging sciences.

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.

Diffusion State-Guided Projected Gradient for Inverse Problems

TL;DR

This work tackles inverse problems with unconditional diffusion priors by addressing the intractable measurement likelihood that can cause artifacts during diffusion-based sampling. The proposed Diffusion State-Guided Projected Gradient (DiffStateGrad) projects the measurement guidance gradient onto a low-rank subspace derived from the intermediate diffusion state, effectively enforcing tangent-space consistency with the data manifold. The method is demonstrated to boost robustness to guidance step sizes and noise, reduce failure rates, and improve reconstruction quality across both linear and nonlinear image restoration tasks, while incurring minimal computational overhead. DiffStateGrad is versatile and compatible with a range of diffusion-based solvers, offering a practical pathway to more reliable diffusion-driven inverse problem solutions with broad potential impact in imaging sciences.

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.
Paper Structure (23 sections, 2 theorems, 23 equations, 20 figures, 6 tables, 5 algorithms)

This paper contains 23 sections, 2 theorems, 23 equations, 20 figures, 6 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Background & Related Works
  3. Diffusion State-Guided Projected Gradient (DiffStateGrad)
  4. Results
  5. Conclusion
  6. Reproducibility Statement
  7. Acknowledgments
  8. Additional Results
  9. Visualizations
  10. Implementation Details
  11. Hyperparameters
  12. Efficiency Experiment
  13. PSLD
  14. ReSample
  15. DPS
  16. ...and 8 more sections

Key Result

Proposition 1

Let $\mathcal{M}$ be a smooth $m$-dimensional submanifold of a $d$-dimensional Euclidean space $\mathbb{R}^d$, where $m < d$. Assume that for each state ${\bm z}_t \in \mathcal{M}$, the tangent space $T_{{\bm z}_t} \mathcal{M}$ is well-defined, and the projection operator $\mathcal{P}_{\mathcal{S}_{ where $\eta > 0$ is a small step size. Then, for sufficiently small $\eta$, the projected update ${

Figures (20)

  • Figure 1: High-level interpretation of Diffusion State-Guided Projected Gradient (DiffStateGrad). DiffStateGrad projects the measurement gradient onto a subspace defined to capture statistics of the diffusion state at time $t$ on which the gradient guidance is applied. This helps the process stay closer to the data manifold during the diffusion process, resulting in better posterior sampling. Without such projection, the measurement gradient pushes the process off the data manifold. For when the measurement gradient guidance is applied to ${\bm z}_{0\mid t}$, the projection is defined to capture the structure of the tangent space of the clean data manifold. The dotted straight line conceptually visualizes the subspace to which the measurement gradient is projected.
  • Figure 2: Visualization of DiffStateGrad in removing artifacts. The large MG step size pushes the process away from the manifold in PSLD, while DiffStateGrad-PSLD is unaffected. The title refers to the diffusion steps.
  • Figure 3: DiffStateGrad improves the worst-case performance. The PSNR histogram for phase retrieval shows that DiffStateGrad significantly lowers the failure rate.
  • Figure 4: Runtime complexity of DiffStateGrad. The increase of runtime with DiffStateGrad is minimal.
  • Figure 5: Robustness of DiffStateGrad to MG step size. (a-b) Performance on box inpainting across various MG step sizes. (c-e) Performance on different tasks with default and large step sizes. We evaluate the performance of PSLD and DiffStateGrad-PSLD using FFHQ $256 \times 256$.
  • ...and 15 more figures

Theorems & Definitions (3)

  • Proposition 1
  • Proposition 1
  • proof