Diffusion State-Guided Projected Gradient for Inverse Problems
Authors: Rayhan Zirvi
Bahareh Tolooshams
Anima Anandkumar
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion models
Inverse problems
Robustness
Subspace
Projection
Box inpainting
Phase retrieval
Primary area: generative models
TL;DR: We can improve the performance and robustness of diffusion-based models in solving inverse problems by adding a Diffusion State-Guided Projected Gradient step.

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.

Published as a conference paper at ICLR 2025
DIFFUSION STATE-GUDED PROJECTED GRADIENT FOR INVERSE PROBLEMS
Rayhan Zirvi", Bahareh Tolooshams", & Anima Anandkumar Computing and Mathematical Sciences California Institute of Technology (rayhanzi rvi, otol oosh, animal Equal contribution
ABSTRACT
Recent advancements in diffusion models have been effective in learning data priors for solving inverse problems. They leverage diffusion sampling steps for inducing data prior while using measurement guidance gradient at each step impose data consistency. For general inverse problems, approximations are needed when unconditionally trained diffuusion model is used since the measurement likelihood is intractable, leading to inaccurate posterior sampling. In other words, due to their approximations, these methods fail preserve the generation process on the data manifold defined by the diffusion prior, leading to artifacts in appli- cations such as image restoration. To enhance the performance and robustness of diffusion models in solving inverse problems, we propose Diffussio Projected Gradient (DiffStateGrad), which projects the measurement gradient onto subspace that is low-rank approximation of an intermediate state of the diffusion process. Dira as 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 measurement guidance step size and noise while improving the worst-case performance. Finally, demonstrate that DiffStateGrad improves upon the state- of-the-ar art on linear and nonlinear image restoration inverse problems. Our code is available https: / thub. com Ani ma Lab/D Di f FStateGradd
INTRODUCTION
Inverse problems are ubiquitous in science and engineering, playing crucial role in simulation-based scientific discovery and real-world applications (Groetsch & Groetsch, 1993). They arise in fields such as medical imaging, remote sensing, astrophysics, computational neuroscience, molecular dynamics simulations, systems biology, and generally solving partial differential equations (PDEs). Inverse problems aim recover unknown signal x: R from noisy observations
y A(x") + n R
where A denotes the measurement operator, and n is the noise. Inverse problems are ill- ill-posed, i.c., in the absence a structure governing the underlying desired signal x, many solutions can explain the measurements y. In the Bayesian framework, this structure is translated into prior p(x=), which can be combined with the likelihood term p( ty(ar xx) to define a posterior distribution p(x) y) xx p(ylaa p(x).. Hence, solving the inverse problem translates into performing a Maximum Posteriori (MAP) estimation or drawing high-probability samples from the posterior (Stuart, 2010). Given the forward model p(z) (z)) the critical step to choose the prior p(z), which often challenging: one needs domain knowledge to define prior or large amount of data to learn it.
Prior works consider sparse priors and provide a theoretical analysis of conditions for the unique recovery of data, problem known compressed sensing (Donoho, 2006; Candès et al., 2006). Sparse priors have shown usefulness medical imaging (Lustig al., 2007), computational neuro- science (Olshausen & Field, 1997), and engineering applications. This approach is categorized into model-based priors where structure is assumed on the signal instead of being learned.

---

Published as a conference paper at ICLR 2025
Diffusion Image Generation
Diffusion for Inverse Problems
Diffusion Inverse Problems w// DiffStateGrad
4
M m+1
%-
M,
20
80 1-1
M 1-]
%0
Denoising Gradient
MG Dred
Measurement Gradient
Comparing measurement gradients
DifSstatGGad Measurement Gradient
pilisiutCrred
Proj P MG
Diffusion Manifold
Figure 1: High-level interpretation of Diffusion State-Guided Projected Gradient (DiffState- Grad). DiffSStateGrad projects the measurement gradient onto subspace defined to capture statistics of the diffusion state at time 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 Zuje. 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 which the measurement gradient is projected.
Recent literature goes beyond such model-based pri- and leverages information from data. The lat- est works employ generative diffusion models (Song & Ermon, 2019; Kadkhodaic & Simoncelli, 2021),
Originll
Meaugrement
PUD
which implicitly learn the data prior p(x:) by learning process that transforms noise into samples from a complex data distribution. For inverse problems, this reverse generation process is guided by the likeli- hood p(y|=), forming denoiser posterior, to gener- data-cone fata-consistent samples. While diffusion models are state-of-the-arte they still face challenges in solv- ing inverse problems.
(a)
Image visualizations.
999
990
980
970
960
950
2014
21
The main challenge arises from the fact that the denoiser posterior, specifically the likelihood com- ponent p(y ()), is intractable since the diffusion is trained unconditionally (Song et al., 2021). Prior work addresses this challenge by proposing various approximations or projections to the gradient related to the measurement likelihood p(yl az) to achieve likely solutions (Kawar et al., 2022); when these ap- proximations are not valid, it results in inaccurate posterior sampling and the introduction of "artifacts" in the reconstructed data (Chung et al., 2023). Latent diffusion models (LDMs) (Rombach et al., 2022), due the nonlinearity of the intent-to-pixe deco der, further exacerbate this challenge. Besides this approximation, the lack of robustness of diffusion models to the measurement gradient step size (Peng et al., 2024) and the measurement noise, and the lack of guarantees for worst-case performance limits their practical applications for inverse problems.
91
(b) PSLD diffusion dynamics.
999
990
980
970
960
950
g
(c) DiffStateGrad-PSII diffusson dynamics.
Figure 2: Visualization of DiflStateGrad in removing artiffacts. The large MG step size pushes the process away from the manifold in PSLD, while DiffStateGrad- fected. The title refers to the diffusion steps.
is unaf-
We propose Diffussin State- Guided Projected Gradient (DiffStateGrad) to address the challenge of staying on the data manifold in solving inverse problems. We focus on gradient-based measurement guidance approaches that use the measurement guidance to move
Our contributions:
2

---

Published as a conference paper at ICLR 2025
the intermediate diffusion state E; toward high-probability regions of the posterior. DiffStateGrad projects the measurement guidance gradient onto a low-rank subspace, capturing the data statistics of the learned prior (Figure 1). We visualize how the diffusion process is pushed off the manifold when the measurement step size is relatively large in a diffusion model and how the incorporation of DiffStateGrad alleviates this challenge (Figure 2). We define projection step to preserve the measurement gradient on the tangent space of the state manifold. We achieve this projection by performing singular value decomposition (SVD) on the diffusion state of image which guidance is applied and use the highest contributing singular vectors as choice of our projection matrix; by projecting the measurement gradient onto our proposed subspace, we aim to remove the directions orthogonal to the local manifold structure.
We show that the crucial factor is the choice of the subspace, not the low-rank nature of the subspace projection. We find that our DiffStateGrad enhances performance. Our projection defines the subspace based on the structure of the state to which the measurement guidance is applied. This is in contrast to random subspace projections or low-rank approximations (Table 1).
We theoretically prove how DiffStateGrad helps the samples remain on or close the manifold, hence improving reconstruction quality (Proposition 1).
We demonstrate that Diff increases the robustness of diffusion models to the measure- ment guidance gradient step size (Figure 5, Table 5) and the measurement noise (Figure 6). For example, for large step size, DiffStateGrad drastically improves the LPIPS of PSLD (Rout et al.. 2023) from 0.463 to 0.165 on random inpainting. For large measurement noise, DiffStaaeGrad improves the SSIM of DAPS (Zhang et al., 2024) from 0.436 to 0.705 on box inpainting.
We empirically show that DiffStateGrad improves the worst-case performance of the diffusion model, e.g., significantly reducing the failure rate (PSNR < 20) from 26% to 4% on the phase retrieval task, increasing their reliability (Figure 3). DiffStateGrad consistently shows lower standard deviation across the test datasets than state-of-the-ar methods.
We demonstrate that DiffStateeGrad significantly improves the performance of state-of-the-art (SOTA) methods, especially in challenging tasks such as phase retrieval and high dynamic range reconstruction. For example, DiffStateGrad improves the PSNR of ReSample (Song et al., 2023a) from 27.61 (8.61)).07) a wide range of linear inverse problems of box inpainting, random inpainting, Gaussian deblur, motion deblur, and super-resolution (Tables 3 and 4) and nonlinear inverse problems of phase retrieval, nonlinear deblur, and high dynamic range (HDR) (Table 3) for image restoration tasks.
31. (9..33) for phase retrieval, reporting mean (std). Our experiments cover
BACKGROUND & RELATED WORKS
Learning-based priors. These methods leverage data structures captured by a pre-trained denoiser (Romano et al., 2017) as plug-and-play priors (Venkatakrishnan al., 2013), or deep generative models such variational autoencoders (VAEs) (Kingma, 2013) and generative ad- versarial networks (GANs) (Goodfellow et al., 2014) to solve inverse problems (Bora et al., 2017: Ulyanov et al., 2018). The -tate-of-the-art is based on generative diffu- sion models, which have shown promising performance in generating high-quality samples in computer vision (Song et al., 2023b), solving PDEs (Shu al., 2023), and high-energ, physics (Shmakov et al., 2024).
Projection Subspace
LPIPS
SSIM
PSNR1
No Projection
0.246
0.809
29.05
Random matrix
0.299
0.753
27.30
Mcasurement gradient
0.242
0.808
29.21
DiftStateGrad (ours)
0.165
0.898
31.68
Table 1: Advantage of diffusion state- guided projection. Results are from ran- dom inpainting on FFHQ 256 x 256.
Diffusion models. Diffusion models conceptualize the generation of data as the reverse of a noising process, where data sample x, at time / within the interval [0, T] follows specified stochastic differential equation (SDE). This SDE (Song et al., 2021) for the data noising process is described by
da = - (Be/2) x dt + V dw,
where B. E (0, 1) is a positive, monotonically increasing function of time t, and w represents standard Wiener process. The process begins with an initial data distribution Ey Pdata and transitions to an approximately Gaussian distribution Vy ~ (O, I) by time T. The objective of regenerating the original data distribution from this Gaussian distribution involves reversing the
3

---

Published as a conference paper at ICLR 2025
noising process through reverse SDE of the form
dx 1 (f+/2)x V. log (e + BBdwa,
where indicates time moving backward and up is the reversed Wiener process. To approximate Vz, log Pe(x). ), a neural network 89 trained via denoising score matching (Vincent, 2011) is used.
Solving inverse problems with diffusion models. Diffusion-based approaches to inverse problems seek to reconstruct the data Zo from the measurement y A(zo) via the following reverse SDE
dx - (P//2)) B.(Va, (xx))
log (y) (x:
V.dw..
Conceptually, the learned score function Vz, Pe (x²) guides the reverse diffusion process from noise to the data distribution, and the likelihood-related term Va, log Po 222) ensures measurement consistency. When the model is trained unconditionallyy the main challenge is the intractable denoiser posterior due the lack of an explicit analytical expression for Vas log (r x;)); the exact relationship between y and intermediate states æ, is not well-defined, except at the initial state Zo-
Solving inverse problems with latent diffusion models. For complex scenarios where direct ap- plication of pixel-based models is computationally expensive or ineffective, latent diffusion models (LDMs) offer a promising alternative (Rombach al., 2022). Given data E IR", the LDM framework utilizes an encoder & : R + I& and a decoder R", with A < to work in compressed latent space. Zr is encoded into a latent representa- tion æτ = E(TT) and serves the starting point for the reverse diffusion process. Then, 2g is decoded to 2t =(zo), the final image reconstruction. Using a latent diffusion model introduces an additional com- plexity to solving inverse problems. The challenge arises from the nonlinear nature and non-uniqueness mapping of the encoder/decoder Rout et al. (2023); tateGrad significantly lowers the failure rate. PSLD proposed to improve performance by enforcing fixed-point properties on representations.
ReSemple
0.20
SeSSaaterrad-ReSappee
Fallure Rate (PSNR < 20): ReSemple: 26%
Density
0.15
DifSinneCrad-ReSampee 4%
0.10
D: P +
0.05
0.00
10
15
20
25
30
35
PSNR
Figure 3: DiflStateGrad improves the worst-case performance. The PSNR his- togram for phase retrieval shows that DiffS-
Diffusion-based inverse problems addressing challenges of intractable denoiser posterior. To address the intractability of the gradient for the reverse diffusion, Diffusion Posterior Sampling (DPS) (Chung et al., 2023), approximates the probability p(y xx.)) R p'y ( == TE a ax.)) using the conditional expectation of the data. Extending to the latent case, Latent-ID DPS uses p,)]=;) p(y (2o := ( o =<l)) (Song et al., 2023a). Two intuitive drawbacks of this approach are that a) the image estimate to is reconstructed using an expectation, which results in inaccurate estimations for multi-mon mudti-modal complex distributions, and b) the measurement gradient directly updates the noisy state 1t, which may push away the state from the desired noise level at t.
=,) R
Prior works aim to address the first challenge by going beyond first-order statistics (Rout et al., 2024) or incorporating posterior covariance into the maximum likelihood estimation step (Peng al., 2024). Other lines of work address the second issue by decoupling the measurement guidance from the sampling process; they update the data estimate æo at time t using the measurement gradient guidance before resampling it to the noisy manifold at time f - 1 (Song et al., 2023a; Zhang et al., 2024). The above-discussed approaches are still highly sensitive to the measurement gradient step size (Peng et al.. 2024). Indeed, balancing the measurement gradient with the unconditional score function remains significant challenge to solving inverse problems using measurement -uidde generation. Wu et al. (2024) avoids the discussed approximation and samples from the posterior directly to resolve the need to find a balance between measurement guidance and the prior process.
Projections in diffusion models. Manifold and subspace projections are used in various contexts in diffusion models. MPGD (He et al., 2024) uses a hanifold-preserving approach to improve the efficiency of diffusion generation and solving inverse problems. While this method follows a similar sentiment as our proposed framework, is only applicable when the measurement gradient is applied to -n Moreover, it requires the existence of autoencoder for achieving manifold projection, and its performance is heavily dependent on the expressive power of the autoencoder. Unlike MPGD, Diff'StateGrad is applicable to methods that apply the guidance / (i.e., ReSample and
4

---

Published as a conference paper at ICLR 2025
DAPS) and methods that apply the guidance to 2,/2, (i.e., PSLD and DPS). Chung al. (2022a) proposes a manifold constraint to project the measurement gradient into the data manifold 2y while the guidance updates 2; our proposal is more effective since project the measurement gradient on the noisy diffusion state related to dt or 2t, preserving the / state on M₁ rather than Mo-
Gradient guidance incorporation. Prior works differ from one another in two key aspects: (a) how the loss for the gradient is computed and (b) how the gradient used to update the diffusion state. Table 2 categorizes prior works based on these character- istics. While few approaches, such as diffusion- -ased MRI (Chung & Ye, 2022), compute the gradient using 2t, most recent literature has shifted to- ward using Zote 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 time I, whereas ReSample, DAPS, and MPGD apply the guidance to the conditional state at 0jL before resampling. Additionally, while the projections in MCG (Chung et al., 2022a) and MPGD (He al., 2024) are restricted to Zojz- DiffSttateGrd applies to both types of methods
Method
Gradient computation Gradient incorporation
Projection
/
Zoy /uje
DPS
PSLD
x
X
X
ReSamplc
X
X
DAPS
/
X
MCG
X
K
MPGD
/
X
DitSSitteGrad (ours)
Table 2: Gradient guidance computation, incorporation, and projection for diffusion- inverse problems.
Conditional diffusion models for inverse problems. We focus on unconditional diffusion models as learned priors to solve general inverse problems. This approach leverages already trained diffusion models, which is useful for domains with abundant data. Another approach is to train conditional diffusion models where Væ, log Pt (y xx) is directly captured by the score function, or where the diffusion directly transforms the measurement into the underlying data (e.g., image-to-ingg -o-ggee diffusion) (Saharia et al., 2022; Liu et al., 2023; Chung et al., 2024). This latter approach problem- specific; hence, it is not generalizable across inverse tasks. Finally, we note that while this work focuses on gradient- based guidance. prior work such as RePaint (Lugmayr et al., 2022) introduces a gradient-free free masking strategy to solve inverse problems. Although RePaint is appealing, it is limited to inpainting tasks and scenarios where measurement noise is negligible.
DIFFUSION STATE-GUIDED PROJECTED GRADIENT (DIFFSTATEGRAD)
We propose a Diffusion State-Guided Projected Gradient (DiffStateGrad) to solve inverse problems. DiffStateGrad can be incorporated into wide range of diffusion models to improve guidance-bart based diffusion models. Without loss of generality, we explain DiffStateGrad in the context of Latent- DPS (Chung et al., 2023) hich applies the measurement guidance to z,. We note that DiffStateGrad applies to wide range of pixel and latent diffusion- based inverse solvers (see Section 4).
Given 1,+11 we sample from the unconditional reverse process, and then compute the estimate Eo(2,) = (EE
20 2,]). Then, the data-consistency guidance term can be incorporated follows.
(5) where 9: VV = p(y εo(2))) is the measurement gradient (MG), Te is the step size, and Ps, is a projection step onto the low-rank subspace S. The main contribution of this paper is a) to highlight that the measurement gradient should be projected onto a subspace imposed by the state being updated by the gradient (see gradient incorporation column in Table 2) and b) to define this subspace so it results in better posterior sampling; in other words, to define a subspace such that when the measurement gradient is projected onto, the diffusion process is not disturbed and pushed away from the data manifold. In Table 1, we show for PSLD that indeed the subspace S,, defined by the intermediate diffusion state, results in an improved posterior sampling, unlike subspace that is constructed based on random matrix or the low-rank structure of the measurement gradient. Hence, we choose the diffusion state 2, define S,. Finally, for methods where the measurement gradient guidance is being applied to Z0Jr, we define the low-rank subspace based on Z0 (see Table 2).
2, 2, PsP (g:).
We focus on images as our data modality and implement the projection Ps, by computing the SVD of in its image matrix form, denoted by Z, (i.c., U, S, V + SVD(Z,)). Then, we compute an adaptive rank r + arg minn 2 1 o3 2 T leveraging a fixed variance retention threshold т.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Diffusion State-Guided Projected Gradient (DiffSttateGrad) for Latent Diffusion-based Inverse Problems (Image Restoration Tasks)
Require: Normal imput + variance retention threshold
1: Let I number of total iterations of sampling algorithm and assume we calculate latent image representation Z, for each iteration. Note that Z. is a matrix.
2: for i T 1 to do
3:
Compute measurement gradient G: according sampling algorithm U.S.V. V t ZZ)
4:
Perform SVD on current diffusion state b Calculate eigenvalues
5:
Ay + sy (where 8, are the singular values of S)
2 Ck Σn
6:
Cumulative sum of eigenvalues
7:
mig 2 TT
> Determine rank based on threshold
8:
A: U,
Get first r left singular vectors D Get first right singular vectors Project gradient > Reconstruct approximated gradient
9:
B. V.
10:
R. + A. G.
11:
A.
12:
Use updated gradient G;, in sampling algorithm
13:
end for
14:
return D(2e)
The gradient It, which takes matrix form for images, is projected onto subspace defined by the highest r singular values of Z, as follows:
G. + U.V...VVV...
(6)
where G the measurement gradient in image matrix form, and U, and V, contain the first left and right singular vectors, respectively (Section 3). While we use the full SVD projection (i.e., combining both left and right projection). in practice, one may choose to either left or right projection. Next, provide mathematical intuitions (Proposition 1) the effectiveness of subspace projections in preserving 2t. particularly for high-dimensional data with low-rank structure, after the MG update on the manifold M,. Finally, note that while DiffStateGGad the runtime and computational efficiency of diffusion frameworks that use Adam optimizers for data consistency (Song et al., 2023a; Zhao al., 2024), the current implementation and this paper does not explore this aspect and, instead, focuses on the property of the proposed subspace.
significantly improve
Proposition 1. Let M be a smooth m-dimensional submanifold of a d-dimensional Euclidean space Rd where < d. Assume that for each state E M,, the tangent space Tz,M is well-defined, and the projection operator Psz, onto an approximate subspace Sz, closely approximates the projection onto M. For the state E . and measurement gradient 9t RRd, consider two update rules:
2:-1 2: 79:
(standard update),
2/111 2t npss, (gt) (projected update).
where nj 0 is small step size. Then, for sufficiently small n, the projected update 21 stays closer to the manifold M than the standard update 21-1. That is,
dist(z!-1,, < dis((zz-1,M).
(8)
The remainder of this section provides intuition on how DiffStateGrad improves solving inverse prob- lems in the presence of a suitable learned prior. Let the initial latent state z, be on the manifold M, (c.g.. being urtifact-free). The term "artifact-fred" refers to the generation process of an unconditional diffusion model that is trained on clean data samples and provides an artifact- trajectory from Mr to Mo. We observe that pushing away from the manifold process (e.g., introducing artifacts) can only be introduced via the guidance by the data-consistency gradient step, as this is the sole mechanism by which information from the measurement process enters the latent space. Consider the manifold M of artifact-free latent representations. Each 2, lies on this manifold, and the tangent space T represents the directions of "allowable" updates that maintain the artifact-free property staying on the current manifold. Finally, note that this motivates project the gradient onto the tangent space of the data manifold where the guidance is applied. Alternatively, when the measurement gradient guidance is applied ZaJt> we define the projection step based So (see also (He et al., 2024)).

---

Published as a conference paper at ICLR 2025
Our DiffStateGrad method, through the projection operator Ps.. approximates this tangent space. The effectiveness of DifiStateGrad depends on how well Ps. approximates the projection onto M. Hence, discuss rationale how the approximated projection is sufficient for performance; we, accordingly, support this by experimental results in Section 4. First, the SVD captures the principal directions of variation in 2,, which are likely align with the local structure of the manifold when the data is high-dimensional. Second, by adaptively choosing the rank based on a variance retention threshold, we ensure that the projection preserves the most significant state-related structural information while filtering out potential noise or artifact artifact-inducing components from the measurement gradient. Finally, the low-rank nature of our approximation aligns with the assumption that the manifold of representations has lower intrinsic dimensionality than the ambient space.
Hence, by projecting the measurement gradient onto this subspace defined by the current latent state zt. we effectively filter the directions orthogonal to the local manifold structure, and hence, remove indfucit-induuing components. This projection ensures that updates to 2₁ remain closer the manifold M than unprojected updates would, as stated in Proposition 1. Consequently, DiffStateGrad relies on the reliability of the learned prior and helps to provide high-probability posterior samples. This creates an inductive process: if 2, is artifact-free, and we only allow updates that align with its structure (i.e., updates that stay close to the manifold M), subsequent latent representation 2₁ will likely be samples from the nigh-probability regions of the posterior.
Figure demonstrates the effectiveness of DiffStateGrad in removing artifacts when the MG step size is large; artifacts introduced onto the measurement gradient and stay within the latent representation in PSLD (Rout et al., 2023). On the other hand, the reverse process via DiffStateGrad- PSLD (our method applied to PSLD) stays artifact-free, consistent with the mathematical analysis in Proposition and the practical efficacy of the proposed SVD-based subspace projection. Finally, we note that the most significant improvements appear in challenging tasks such as phase retrieval, HDR, and inpainting. We attribute the effectiveness of DiffStateGrad, particularly in challenging tasks, to reduced rate of failure cases (Figure 3). By constraining solutions closer to the data manifold, DiffStateGrad minimizes extreme failures, enhances consistency in reconstruction quality.
Efficiency. DiffStateGrad introduces minimal computational overhead. We perform SVD at most once per iteration, and for latent diffusion solvers, this occurs in the latent space on 64 x 64 matrices. By selecting low rank based a variance threshold, subsequent projection and reconstruction operations are performed on reduced matrices, further decreasing computational complexity.
Unmodified With DiffStateGrad
00.78
PSLD
2:38 .0
ReSemple
248.55
23.18
DAPS
Figure 4 illustrates the runtime of PSLD (Rout et al., 2023), ReSample (Song et al., 2023a), and DAPS (Zhang et al., 2024) with and without DiffStateGradd The figure shows that the additional computational cost of incorporating DiffStateGraa is marginal, typically adding only a few seconds to the total runtime (see C.2 for further details). We note that our method is not intended to improve efficiency, but rather to enhance performance and robustness.
50
100
250
Runtime (seconds)
Figure 4: Runtime complexity of DiffStateGrad. The increase of runtime with DiffStateGrad is minimal.
4 RESULTS
This section provides extensive experimental results on the effectiveness of DiffStateGrad for image- based inverse problems. We show that DiffStateGrra significantly improves (1) the robustness of diffusion-bassee methods to the choice of measurement gradient step size and measurement noise, and (2) the overall posterior sampling performance of diffusion.
Experimental setup. We evaluate the performance of DiffStaaeGGad applied to four SOTA diffusion methods of PSLD (Rout et al., 2023), ReSample (Song et al., 2023a), DPS (Chung et al., 2023), and DAPS (Zhang et al., 2024). These methods span both latent solvers (PSLD and ReSample) and pixel-based solvers (DPS and DAPS). We also directly compare against other methods includ- ing DDNM (Wang et al., 2023), DDRM (Kawar al., 2022), MCG (Chung al., 2022a), and
7

---

Published as a conference paper at ICLR 2025
0.40
0.35
DiffStateGrad-PSLD
0.30
PSLD
PSNR (1)
18
NMSE
0.20
0.15
0.10
DiffstateGrad-PSI
0.05
PSLD
0.2
0.3
02
0.3
0.4
MG Step Size
MG Step Size
(a) Box Inpainting (NMSE).
(b) Box Inpainting (PSNR).
0.8
+
DiffStateGrad PSLD
0.8
DiffStateGred PSLD
0.8
DifistateGred PSLD
PSLD
PSLD
0.6
0.6
0.6
LPIPS
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.1 (default)
0.4
0.10
0.15
0.9 Step
1.5
(default)
(cofault]
Step Size
Step Size
(c) Random Inpainting (LPIPS).
(d) Motion Deblur (LPIPS).
(c) Super- Resol-Resolttion (LPIPS).
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 DiffStateGGra-PS5 PSLD using FFHQ 256 256.
DiffStateGrad-PSLD PSLD
0.8
0.65
0.60
SSIM
SSIM
0.6
0.55
DiffStateGrad DAPS
0.50
0.4
DAPS
0.00
0.05
0.10
0.15
0.20
0.00
0.05
0.10
0.15
0.20
Measurement Noise Std
Measurement Noise Std (b) Motion Deblur.
(a) Inpainting (Box).
Figure 6: Robustness of Dird to measurement noise. We evaluate the performance of DiffStateGGrad-SI.. and DiffStateGrad-DAA with their respective counterparts on different tasks across a range of measurement noise levels (std of 0 to 0.2) using FFHQ 256 x 256.
MPGD-AE (He et al., 2024). We evaluate performance based on key quantitative metrics, including LPIPS (Learned Perceptual Image Patch Similarity), PSNR (Peak Signal-to-Noise Ratio), and SSIM (Structural Similarity Index) (Wang et al., 2004). We demonstrate the effectiveness of DiffStateGrad on two datasets: a) the FFHQ 256 x
256 validation dataset (Karras et al., 2021), and b) the ImageNet 256 x 256 validation dataset (Deng et al., 2009). For pixel-based experiments, we use (i) the pre- trained diffusion model from (Chung et al., 2023) for the FFHQ dataset, and (ii) the pre-trained model from (Dhariwal & Nichol, 2021) for the ImageNet dataset. For latent diffusion experiments, we use (i) the unconditional LDM-VQ-4 model trained on FFHQ (Rombach et al., 2022) for the FFHQ dataset, and (ii) the Stable Diffusion v1.5 (Rombach al., 2022) model for the ImageNet dataset.
We consider both linear and nonlinear inverse problems for natural images. For evaluation, we sample a fixed set of 100 images from the FFHQ and ImageNet validation sets. Images are normalized to the range [0, 1]. We use the default settings for all experiments (see Appendix C for more details). For linear inverse problems, we consider (1) box inpainting, (2) random inpainting, (3) Gaussian deblur, (4) motion deblur, and (5) super-resolution. In the box inpainting task, random 128 x 128 box is used, while the random inpainting task employs a 70% random mask. Gaussian and motion deblurring tasks utilize kernels of size 61 x 61, with standard deviations of 3.0 and 0.5, respectively, For super- resolutess, images are downscaled by a factor of using a bicubic resizer. For nonlinear inverse problems, we consider (1) phase retrieval, (2) nonlinear deblur, and (3) high dynamic range (HDR). For phase retrieval, we use an oversampling rate of 2.0, and due to the instability and non- uniqueness of reconstruction, we adopt the strategy from DPS (Chung al., 2023) and DAPS (Zhang
8

---

Published as a conference paper at ICLR 2025
Table 3: Performance comparison for linear and nonlinear tasks on FFHQ 256 x 256.
Method
Inpaint (Box)
Inpaint (Random)
Gaussian deblur
Motion deblur
SR (x4)
LPIPSL PSNRT
LPIPSI PSNRT LPIPSL PSNRT
LPIPS PSNR? LPIPS) PSNRT
Pixel-based
DAPS
0.136
24.57
0.130
30.79
0.216
27.92
0.154
30.13
0.197
28.64
DiffShateGrad-DAP (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-De (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
Latent
PSLD
0.158
24.22
0.246
29.05
0.357
22.87
0.322
24.25
0.313
24.51
DiffStateGrdd-SSL. PSLD (ours)
0.092
24.32
0.165
31.68
0.355
22.95
0.319
24.31
0.320
24.56
ReSample
0.198
19.91
0.115
31.27
0.253
27.78
0.160
30.55
0.204
28.02
DiftStateGrad -raeteeaGdd--eepplleee (ours)
0.156
23.59
0.106
31.91
0.245
28.04
0.153
30.82
0.200
28.27
(a) Linear inverse problems.
Method
Phase retrieval
Nonlinear deblur
High dynamic range
LPIPS
PSNRT
LPIPS
PSNR1
LPIPS4
PSNR1
Pixel-based
DAPS
0.139 (0.026) 30.52 (2.61) 0.184 (0.032)
27.80 (1.97) 0.170 (0.075)
26.91 (3.94)
DiffStateGrad-DDa (ours)
0.105 (0.023) 32.25 (1.34)
0.145 (0.027)
29.51 (2.16) 0.143 (0.070)
27.76 (3.18)
Latent
ReSample
0.237 (0.189) 27.61 (8.07) 0.188 (0.037) 29.54 (1.89)
0.190 (0.067)
24.88 (3.46)
SerStattGGaddReeSamppl (ours)
0.154 (0.104) 31.19 (4.33) 0.185 (0.035) 29.91 (1.60) 0.164 (0.041) 25.50 (3.07)
(b) Nonlinear inverse problems.
Original
Measurement
DIfiStatGGrad
Original
Measurement
(b) Random inpainting. Figure 7: Qualitative comparison of DiftStateGrad-SSI. and PSLD for their best-performing MG step size. DiffStateGraa-PSII reliable reconstructions. Images are chosen at random for visualization.
(a) Box inpainting.
remove artifacts and reduce failure cases, producing more
et al., 2024), generating four separate reconstructions and reporting the best result. Like DAPS (Zhang et al., 2024), normalize the data to lie in the range [0, 1] before applying the discrete Fourier transform. For nonlinear deblur, we use the default setting from (Tran et al.. 2021). For HDR, we use scale factor of We that PSLD not designed handle nonlinear inverse problems. We also conduct additional experiment for Magnetic Resonance Imaging (MRI) (see Appendix E).
Robustness. Figure exhibits the sensitivity of PSLD to the choice of MG step size; the perfor- mance of PSLD significantly deteriorates when relatively large MG step size is used, leading to poor results across all tasks. In contrast, DiffStateGrad PSLD shows superior robustness, maintain- ing high performance over wide range of MG step sizes. Figure 6 demonstrates the robustness of DiffStateGrad methods compared to their non- non-DiffS DiffSttateGrd counterparts when faced with in- creasing measurement noise. For both inpainting and Gaussian deblur tasks, the performance of DAPS and PSLD deteriorates significantly as noise levels rise. In contrast, DiffStateGrad DAPS and DiffStateGrad PSLD D exhibit superior resilience across the range of noise levels tested.

---

Published as a conference paper at ICLR 2025
Table 4: Performance comparison for linear tasks on ImageNet 256 256.
Method
Inpaint (Box)
Inpaint (Random)
LPIPS
SSIMI
PSNRT
LPIPS
SSIMT
PSNR
Pixel-based
DAPS
0.217 (0.043) 0.762 (0.041) 20.90 (3:69) 0.158 (0.039) 0.794 (0.067) 28.34 (3,65)
DifSSateGrrd-DAAe (ours)
0.191 (0.044) 0.801 (0.056) 21.07 (3.77) 0.107 (0.037) 0.856 (0.067) 29.78 (4.17)
DPS
0.257 (0.086) 0.718 (0.097) 19.85 (3:54) 0.256 (0.133) 0.728 (0.143) 26.25 (4.15)
DiftStateGrad-DD (ours)
0.243 (0.093) 0.731 (0. (0.100) 19.87 (3.61) 0.233 (0.138) 0.754 (0.150) 27.28 (4.88)
MPGD-AE
0.295 (0.057) 0.621 (0.053)
16.12 (2.26) 0.554 (0.148) 0.388 (0.112) 17.91 (3.25)
Latent
PSLD
0.182 (0.033) 0.780 (0.044) 16.28 (3.49) 0.217 (0.073) 0.846 (0.070)
26.56 (2.98)
DiffstatGGradPPLLD (ours)
0.176 (0.030) 0.803 (0.045) 18.90 (3.82) 0.169 (0.050) 0.878 (0.051) 28.48 (4.04)
(a) Inpainting inverse problems.
Method
Gaussian deblur
Motion deblur
SR (x4)
LPIPS|
PSNRT
LPIPS
PSNRT
LPIPS
PSNR1
Pixel-bused
DAPS
0.266 (0.087) 25.27 (3.56)
0.166 (0.058)
28.85 (3.64) 0.259 (0.073)
25.67 (3.40)
DiffStateGrad-Der DAPS (ours)
0.243 (0.075)
25.87 (3.56)
0.143 (0.050)
29.71 (3.54)
0.229 (0.057)
26.40 (3.44)
MCG
0,550 (-)
16.32 (-)
0.758 (-)
5.89 (-)
0.637 (-)
13.39 (-)
Latent
PSLD
0.466 (0.085) 20.70 (3.01) 0.435 (0. 102) 21.26 (3.44) 0.416 (0.063)
22,29 (3.08)
DiflStateGrad-PSL. (ours)
0.446 (0.076) 22.34 (3.19) 0.399 (0.060) 23.80 (3.27) 0.370 (0.081) 23.53 (3.52)
(b) Deblurring and super-resolution inverse problems.
Performance. We provide quantitative results Tables 4, and qualitative results in Figure These results demonstrate the substantial improvement in the performance of the methods with DiffStateGrad against their respective SOTA counterparts across wide variety of linear and nonlinear tasks on both FFHQ and ImageNet. For example, DiftStateGrad significantly improves performance and increases reconstruction consistency in both pixel-based and latent solvers for phase retrieval (Table 3). Our results also highlight the superiority of our choice of subspace compared to prior works. For example, DifiStateGrad-DDS manifold constrained diffusion. Furthermore, the outperformance of SerrSttatGrad-DAAPP against DAPS and MPGD-AE emphasizes the effectiveness the proposed SVD-based subspace projection. We additionally show that Diff does not impact the diversity of the posterior sampling (Figure 16). Finally, we note that DiffStateGrad is robust to the choice of subspace rank (Figures and 9). We refer the reader Appendix B for the extensive qualitative performance of DiftStateGrad.
superior to MCG, which is the special case of DPS for
CONCLUSION
We introduce Diffusion State-Guided Projected Gradient (DiffStateGradd) to enhance the perfor- mance and robustness of diffusion models in solving inverse problems. DiffStateGrad addresses the introduction of artifacts and deviations from the data manifold by constraining gradient updates to a subspace approximating the manifold. DiffStateGrad is versatile, applicable across various diffusion models and sampling algorithms, and includes an adaptive rank that dynamically adjusts the gradi- ent's 's complexity. Overall, DiffStateGrad reduces the need for excessive tuning of hyperparameters and significantly boosts performance for more challenging inverse problems. We note that DiffSttate- Grad assumes that the learned prior is relatively good prior for the task at hand. Since Diff'StateGrad encourages the process to stay close to the manifold structure captured by the generative prior, may introduce the prior's biases into image restoration tasks. Hence, DiffStateGrad may not be recommended for certain inverse problems such as black hole imaging (Feng et al., 2024). We finally note that DiffStateGraa can be combined with prior works that adopt initialization strategies for the diffusion process to further accelerate and improve performance (Fabian et al., 2024; Chung et al., 2022b). We leave this for future work.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our results, we thoroughly detail the hyperparameterss used in our experiments in C.I. We also provide specific implementation and configuration details of all the baselines used in C.3, C.4, C.5, and C.6. Moreover, we use easily accessible pre-trained diffusion mod- els throughout our experiments. PSLD (https://github com/Litiun wtt.tttttttttsst ReSam- ple (https: htap::://gghhhie http////oooooowwo////eeeaaaaaaaa and DAPS (https:/ htt.bb.....uuggg.o/gg///.t...... com/ charggginnglllangg219999/442329 all have publicly available code. We also make our code available
DPS (https://githin
com/
DPS2022/ DPSS02//dffffaao0rr-oeeeeera-aa0
ACKNOWLEDGMENTS
R.Z. did the project under the Summer Undergraduate Research Fellowships (SURF) program at Caltech. R.Z. thanks the SURF program for funding this project. B.T. was supported by the Swartz Foundation Fellowship for Postdoctoral Research in Theoretical Neuroscience at Caltech. A.A. was supported Bren Chair Professor and A12050 Senior Fellow.
B.T. proposed and designed the project. R.Z. implemented the method and executed all the experi- ments, excluding the MRI experiment. B.T. ran the MRI experiment. R.Z and B.T. wrote the paper. A.A. supervised the project, provided valuable feedback, and offered editorial comments. R.Z. and B.T. prepared the rebuttal together. B.T. led the discussion during the review process.

---

Published as a conference paper at ICLR 2025
REFERENCES
Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G Dimakis. Compressed Sensing using Generative Models. International conference on machine learning, pp. 537-546. PMLR, 2017.
Emmanuel J Candès, Justin Romberg, and Terence Tao. Robust Uncertainty Principles: Exact Signal Reconstruction from Highly Incomplete Frequency Information. IEEE Transactions on information theory, 52(2):489-509, 2006.
Hyungjin Chung and Jong Chul Ye. Score-Based Diffusion Models for Accelerated MRI. Medical Image Analysis, pp. 102479, 2022.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving Diffusion Models for Inverse Problems using Manifold Constraints. Advances in Neural Information Processing Systems, 35:25683- 25696, 2022a.
Hyungjin Chung, Byeongsu Sim, and Jong Chul Ye. come--loser-ddiffus-fasstts 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, 2022b.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion Posterior Sampling for General Noisy Inverse Problems. In The Eleventh International Conference on Learning Representations, 2023. URL https: /ppp/eeeeeeeeeeeeeeeeeeeeee net forum?
92GAGTOK
Hyungjin Chung, Jeongsol Kim, and Jong Chul Ye. Direct diffusion bridge using data consistency for inverse problems. Advances in Neural Information Processing Systems, 36, 2024.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. ImageNet: A Large- Scale Hierarchical Image Database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, IEEE, 2009.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion Models Beat GANs Image Synthesis. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems. Curran Associates, Inc., 2021.
David L Donoho. Compressed Sensing. IEEE Transactions on information theory, 2006.
Zalan Fabian, Berk Tinaz, and Mahdi Soltanolkotabi. Adapt and diffuse: Sample-adaptive recon- struction via latent diffusion models. Proceedings of machine learning research, 235:12723, 2024.
Berthy T Feng, Katherine Bouman, and William T Freeman. Event-horizon--call imaging m87* under different assumptions via deep generative image priors. arXiv preprint arXiv:2406.02785. 2024.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Nets. Advances in neural information processing systems, 27, 2014.
Charles Groetsch and CW Groetsch. Inverse problems in the mathematical sciences, volume 52. Springer, 1993.
Yutong He, Naoki Murata, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Dongjun Kim, Wei- Hsiang Liao, Yuki Mitsufuji, J Zico Kolter, Ruslan Salakhutdinov, et al. Manifold preserving guided diffusion. In The Twelfth International Conference on Learning Representations, 2024.
Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir. Robust Compressed Sensing MRI with Deep Generative Priors. In Advances in Neural Information Processing Systems, volume 34, pp. 14938-144 14954. Curran Associates, Inc., 2021.
12

---

Published as a conference paper at ICLR 2025
Zahra Kadkhodaic and Eero Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in a denoiser. Advances in Neural Information Processing Systems, 34:13242-1333 -3254, 2021.
Tero Karras, Samuli Laine, and Timo Aila. Style- Basee Generator Architecture for Generative Adversarial Networks. IEEE Transactions on Pallern Analysis & Muchine Intelligence, 43(12): 4217-4228, Dec 2021.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593-23606, 2022.
Diederik Kingma, Auto-encoding Variational Bayes. arXiv preprint arXiv: 2013,
Florian Knoll, Jure Zbontar, Anuroop Sriram, Matthew J Muckley, Mary Bruno, Aaron Defazio, Mare Parente, Krzysztof Geras, Joe Katsnelson, Hersh Chandarana, et al. fastmri: A publicly available raw k-space and dicom dataset of knee images for accelerated mr image reconstruction using machine learning. Radiology: Artificial Intelligence, ((1):e19000)) 2020.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. 12sb: Image-to-image schrödinger bridge. In International Conference on Machine PMLR, 2023.
Learning, pp.
Andreas Lugmayr, 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 computer vision and pattern recognition, pp. 11461-11471, 2022.
Michael Lustig, David Donoho, and John M Pauly. Sparse mri: The application of compressed sensing for rapid mr imaging. Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, (3(6):1822-181 2007.
Bruno A Olshausen and David Field, Sparse coding with an overcomplete basis set: A strategy employed by v1? Vision research, 37(23):3311-333
1997.
Xinyu Peng. Ziyang Zheng, Wenrui Dai, Nuoqian Xiao, Chenglin Li. Junni Zou, and Hongkai Xiong. Improving diffusion models for inverse problems using optimal posterior covariance. In Forty-first International Conference on Machine Learning, 2024.
Yaniv Romano, Michael Flad, and Peyman Milanfar. denoising (red). SIAM Journal on Imaging Sciences, 10(4):18044 1844, 2017.
The little engine that could: Regularization by
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution Image Synthesis with Latent Diffussio Models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkottai. Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023. URL https: //popreevveww.. net/ forum? 1d=
Litu Rout, Yujia Chen, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkottai, and Wen-Sheng Chu. Beyond first-order tweedie: Solving inverse problems using latent diffusion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9472-9481, 2024.
Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: image-to-image diffusion models. In ACM SIGGRAPH 2022 conference proceedings, pp. 1-10, 2022.
Alexander Shmakov, Kevin Greif, Michael Fenton, Aishik Ghosh, Pierre Baldi, and Daniel White- son. End-to-end latent variational diffusion models for inverse problems in high energy physics. Advances in Neural Information Processing Systems, 36, 2024.
Dule Shu, Zijie Li, and Amir Barati Farimani. A Physics-informed Diffusion Model for fidh-fiielity Flow Field Reconstruction. Journal of Computational Physics, 478:1 1972,, 2023.
13

---

Published as a conference paper at ICLR 2025
Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen. Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency. In Conference on Parsimony and Learning (Recent Spotligtt Track), 2023a. URL id=i A.d=iHcarbe rDCZIn.
formm
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-Cuuided Diffusion Models for Inverse Problems. In International Conference on Learning Representations, 2023b.
URL httpps://ppenrvvette net.eeeeeeeweeeeeeeeeeeeeeeeeeee htt.:/ee////oooeeee Eorum? idd (iddd____MMMMMKKK
Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. In Advances in Neural Information Prucessing Systems, volume 32, 2019.
Yang Song and Stefano Ermon. Improved Techniques for Training Score-Based Generative Models. 2020.
In Advances in Neural Information Processing Systems, volume 33, pp.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based Generative Modeling through Stochastic Differential Equations. In The International Conference on Learning Representations, 2021. URL https: /ttLhtts////peeeeee /openrevie. net/pdf/
11c.B11c144490 pdf.
Andrew M Stuart. Inverse problems: bayesian perspective. Acta numerica, 19:451-559, 2010.
Phong Tran, Anh Tran, Quynh Phung, and Minh Hoai. Explore Image Deblurring via Encoded Blur Kernel Space. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2021.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Deep Image Prior. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 9446-9454, 2018.
Singanallur V Venkatakrishnan, Charles A Bouman, and Brendt Wohlberg. Plug-and-play priors for model based reconstruction. In 2013 IEEE global conference on signal and information processing, pp. 945 948. IEEE, 2013.
Pascal Vincent. A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7):1661-1674, 2011. doi: 10. 10.1162/NECO_a_000
Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference on Learning Representations, 2023.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from eπror visibility to structural similarity. IEEE transactions image processing. 13(4):600-612, 2004.
Zihui Wu, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, and Katherine L Bouman. Prin- cipled probabilistic imaging using diffusion models as plug-and-play priors. arXiv preprint arXiv:2405.18782, 2024.
Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew J. Muckley, Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, Marc Parente, Krzysztof Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzal, Adriana Romero, Michael Rabbat, Pascal Vincent, Nafissa Yakubova, James Pinkerton, Duo Wang, Erich Owens, C. Lawrence Zitnick, Michael P. Recht, Daniel K. Sodickson, and Yvonne W. Lui. fastMRI: An open dataset and benchmarks for accelerated MRI. 2018.
Bingliang Zhang, Wenda Chu, Julius Berner, Chenlin Meng, Anima Anandkumar, and Yang Song. Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing, arXiv preprint arXiv:2407.01521, 2024.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficien 11m training by gradient low-rank projection. In Forty-first International Conference Machine Learning, 2024.
14

---

Published as a conference paper at ICLR 2025
A
ADDITIONAL RESULTS
Table 5: Robustness comparison of PSLD and DiffStateGGad-PSS.DD linear tasks under different MG step sizes on FFHQ 256 x 256.
Method
Inpaint (Box) Inpaint (Random) Gaussian deblur
Motion deblur
SR (x4)
LPIPS
PSNRT LPIPS
PSNRT LPIPS PSNRT
LPIPSI PSNRT
LPIPS. PSNR
Defautt MG step
PSLD
24.22
0.246
0.357
22.87
0.322
24.25
0.313
24.51
DitSSuteCGrad-PLLD (ours)
0.095
23.76
0.265
28.14
0.366
22.24
0.335
23.34
0.392
22.12
Large MG step
PSLD
0.252
11.99
0.463
20.62
0.549
17.47
0.514
18.81
0.697
7.700
DilSStatGGra-PSS. (ours)
0.092
24.32
0.165
31.68
0.355
22.95
0.319
24.31
0.320
24.56
GSS 1-05
0.30
WININOneFLLD 0-95
PAL 0.5
POL 2073
3.00
20
0.00
NASSE
0.23
0.05
0.07
PSNR
0.00
16
USSunsrrs PSLE:
EENAYO PMDI
Des PMD 0.555
0.55
2.32
0.15
5.25
0.30
0.33
0.0
0.2
0.25
0.30
0.40
022
B25 0.25 0.30
MG Step Size
MG Step Sice
Step Size
(a) NMSE
(b) NMSE (zoomed)
(c) PSNR
Figure Robustness comparison PSLD and DiffStateGrad-PS.. for different variance retention thresholds T. We evaluate images from FFHQ 256 256 on box inpainting.
PSNA 25.0
SSWS (11
LAIPS
Vaniance Retention Threshold (T)
Variacce Retention Threshold (t)
Var Jari Retention Threshold (t)
(a) PSNR
(b) SSIM
(c) LPIPS
Figure 9: Performance of DiffStateGrad-DAPS for different variance retention thresholds т. DiffStateGrad is robust to the choice of r, values > 0.6 show similar performance in this figure. In our main experiments, we use T = 0.99 (see C.I for further details). We evaluate images from FFHQ 256 x 256 box inpainting.
15

---

Published as a conference paper at ICLR 2025
34
0.92
0.14
DAPS
DAPS
DiffStateGrad DAPS
DiffStateGrad DAPS
0.13
SSIM (1)
0.88
0.11
LPIPS
62 31
0.10
0.84
0.09
DAPS
0.08
DifIStateGrad DAPS
0.80
0.07
10
Projection Period
Projection Period
Projection Period
(a) PSNR
(b) SSIM
(c) LPIPS
Figure 10: Performance of DiffStateGrad-DAPS for different projection periods P. As period of projection increases, DiffStateGrad still outperforms DAPS without projection, which is the SOTA. See C.1 for details of P values used in experiments. We evaluate images from FFHQ 256 256 on random inpainting.
Table 6: SSIM comparison on FFHQ 256 x 256.
(a) Linear tasks.
Method
Inpaint (Box) Inpaint (Random)
Gaussian deblur Motion deblur
SR (x4)
SSIM
SSIM
SSIM1
SSIM1
SSIM
Pixel-based
DAPS
0.806 (0.028)
0.829 (0.022)
0.786 (0.051)
0.837 (0.040)
0.797 (0.044)
DiffStateGradDDa. (ours)
0.849 (0.029)
(0.023)
0.803 (0.044)
(0.028)
0.801 (0.039)
DPS
0,810 (0.039)
0,815 (0.045)
0.709 (0.062)
0.754 (0.056)
0.675 (0.071)
DiffStateGrad- (ours)
(0.039)
0.852 (0.046)
0.739 (0.062)
0.782 (0.056)
0.683 (0.073)
MPGD-AE
0.753 (0.029)
0.731 (0.050)
0.664 (0.071)
0.723 (0.061)
0.670 (0.070)
Latent
PSLD
0.819 (0.031)
0.809 (0.049)
0,537 (0.094)
0.615 (0.075)
0.650 (0.140)
DifSStateGrad-PS... (ours)
0.880 (0.028)
0.898 (0.024)
0.542 (0.077)
0.620 (0.065)
0.640 (0.123)
ReSample
0.807 (0.036)
0.892 (0.030)
0.757 (0.049)
0.854 (0.034)
0.790 (0.048)
DiffStateGrad RaSmpee (ours)
0.841 (0.032)
0.913 (0.023)
0.767 (0.041)
0.860 (0.031)
0.795 (0.044)
(b) Nonlinear tasks.
Method
Phase retrieval Nonlincar deblur High dynamic range
SSIM
SSIM
SSIMI
Pixel-based
DAPS
0.823 (0.033)
0.723 (0.034)
0.817 (0. 109)
DitlStateGrad-DD (ours)
0.868 (0.026)
(0.035)
(0.098)
Latent
ReSample
0.750 (0.246)
0.842 (0.038)
0.819 (0.109)
BritSaalGrud-ReeSummee (ours)
0.855 (0.130)
0.847 (0.035)
(0.059)
16

---

Published as a conference paper at ICLR 2025
B VISUALIZATIONS
Original
Neesurement
Measurement DifSsaterra--SSL
(a) Box inpainting.
(b) Random inpainting.
Original
Measurement
DIfSSatteGrad-PSLD
PSLO
Original
Measurement
DifStateGGrad-P5L
PSUD
(c) Gaussian deblur.
(d) Super-resolution (x
Figure 11: Qualitative comparison of DifSStateGrad-PSLD and PSLD for large MG step size. Images are chosen at random for visualization.
Originel
eemeuremment
PSLD
Original
Meesanement
DiffSiateGrrd-PSLD
PSLD
(a) Box inpainting.
(b) Random inpainting.
Measurement
DifSSateGrad-PSUU
Measurement DifttateGGadd-PSL
BLO
BLO
BLO
BLO
B
(c) Gaussian deblur.
(d) Super-resolution (x1).
Figure 12: Qualitative comparison of DiTIStateCrad-PSLD and PSLD for their best-performing
MG step size. Images are chosen random for visualization.
17

---

Published as a conference paper at ICLR 2025
Originl
Odiahal
Figure 13: Qualitative comparison of SrrStattGrad-ReSammee and ReSample for phase re- trieval. Whereas the performance of ReSample is inconsistent, DiffStateGrad-ReSS consistently produces accurate reconstructions. Images are chosen random for visualization.
Original
Messurement
ReSample
Oriiinll
Measurement
Dratteerrr ReSample
ReSample
Figure 14: Qualitative comparison of Baataaterraa--eRanappp and ReSample for high dynamic
range (HDR). Images are chosen at random for visualization.
18

---

Published as a conference paper at ICLR 2025
Original
Measurement
DMfSsateCCrad-SLL
PSLD
Original
Measurement
DifiSiateG ac-
PSLD
(a) Gaussian deblur.
(b) Super-resolution (x
Figure 15: Qualitative comparison of DiflStateGrad-PSI and PSLD for their pest-performing MG step size. DiffStateGrad-PSLD can remove artifacts and reduce failure cases, producing more reliable reconstructions. Images are chosen at random for visualization.
Original
Measurement
Sample
Sample 2
Sample 3
Original
Measurement
Sample
Sample
Sample
(a) Box inpainting (192 192 box).
(b) Gaussian deblur (81 81 kemel,, SD of 7).
Figure 16: Reconstruction diversity of DiftStateGrad. DiflStateGrad-PSLD with large MG step size can produce diverse range of images from multimodal posteriors. Generated images have distinctive facial features.
C IMPLEMENTATION DETAILS
C.1 HYPERPARAMETERS
For all main experiments across all four methods, we use the variance retention threshold r = 0.99. For all experiments involving PSLD, DPS, and DAPS, we perform the DiffStateeGrad projection step every iteration (P 1). For all experiments involving ReSample, we perform the step every five iterations (P = 5). See the sections dedicated to each method for further implementation details. We reiterate that various values of r and P are reasonable options for optimal performance (see Figures to 10).
C.2 EFFICIENCY EXPERIMENT
We evaluate the computational overhead introduced by DiffStateGrad across three diffusion- based methods: PSLD, ReSample, and DAPS. We conduct these experiments on the box inpainting task using an NVIDIA GeForce RTX 4090 GPU with 24GB of VRAM. Each method is run with its default settings on set of 100 images from FFHQ 256 256, and we measure the average runtime in seconds per image.
C.3 PSLD
Our DiffStateGrad-PSLD algorithm integrates the state-guided projected gradient directly into the PSLD update process. For each iteration (i.e., period = 1) of the main loop. after computing the standard PSLD update 2 we introduce DiffStateGrad method. First, calculate the full gradient G, according to PSLD, combining both the measurement consistency term and the fixed-point constraint. We then perform SVD on the current latent representation (or diffusion state)
19

---

Algorithm 2 DiffStateCGrad-PSL.D for Image Restoration Tasks Published as a conference paper at ICLR 2025
Zi (z, in image matrix form). Using the variance retention threshold determine the appropriate rank for our projection. We construct projection matrices from the truncated singular vectors and use these to approximate the gradient. This approximated gradient G¹, is then used for the final update step, replacing the separate gradient updates in standard PSLD. This process is repeated at every iteration, allowing for adaptive, low-rank updates throughout the entire diffusion process.
For experiments, we use the official implementation of PSLD (Rout et al., 2023) with default configurations (i.e., noise, step size, etc.) for reproducing baselines.
Require: T, y, )T th f (
Require: E, A.A.A A, variance retention threshold r
1: 2T ~O,
2: for f = T to 0 do
3: â - No (2,
4:
20
(Ze
+ (1 (8))
Jor
5:
I
z(=1 yarll-as (ar(1- + va, 1
6:
I- 20 ō.E
1-6,
7:
@. g *.-V..ll
(E)
+ Va,leu E(A ArT + +(I AT ATD))(())))
8:
U.S.V + -SVD(Z,) 2.)
9:
A + of (8, are the singular values of S) - Σ E
10:
11:
T arg r
12:
A U.
B V
14:
R. + R
15:
16:
21-1
# 9t
end for
18: return D(2))
C.4 RESAMPLE
Our DiffStateGrad- ReSaale algorithm integrates the state- guided projected gradient into the op- timization process of ReSample (Song et al., 2023a). We introduce two new hyperparameters: the variance retention threshold t and period P for applying our DiffStateGrad step. During each ReSample step, we first perform SVD the current latent representation (or diffusion state) Z (*) in image matrix form). Note that we do not perform SVD within the gradient descent loop itself, meaning that we only perform SVD most once per iteration of the sampling algorithm. We then determine the appropriate rank based on and construct projection matrices. Then, within the gradient descent loop for solving zn(y), (y). we approximate the gradient in the diffusion state subspace using our projection matrices every P = 5 steps. On steps where DiffStateeGrd is not applied, we use the standard gradient. This adaptive. periodic application of DiffStateGrad allows for a balance between the benefits of low-rank approximation and the potential need for full gradient information. The rest of the ReSample algorithm, including the stochastic resampling step, remains unchanged.
We note that the ReSample algorithm employs two-stage approach for its hard data consistency step. Initially, it performs pixel-space optimization. This step is computationally efficient and produces smoother, albeit potentially blurrier, results with high-level semantic information. As the diffusion process approaches the image with finer details. Our tiffStateGrad method is specifically integrated into this latter, latent-space optimization stage. By applying DiffStateGrad to the latent optimization, we aim to mitigate the potential introduction of artifacts and off-manifold deviations that occur due to the direct manipulation of latent variables. This application of DiffStateGrad allows us to benefit from the computational efficiency of initial pixel-space optimization while enhancing the robustness and
O. ReSample transitions to latent-space optimization to refine
20

---

Published as a conference paper at ICLR 2025
quality of the final latent-space refinement. Importantly, DiffStateGrad is not applied during the pixel-space optimization phase, as this stage already tends to produce smoother results and less prone to artifact introduction.
For experiments, we use the official implementation of ReSample (Song et al., 2023a) with default configurations (i.e., noise, step size, etc.) for reproducing baselines.
Algorithm 3 DiffStateGraddReeaanii for Image Restoration Tasks
Require: Measurements A(-), Encoder E(-), Decoder D(-), Score function so(: t), Pretrained LDM
Parameters 3,, (z. 1/. 6, Hyperparameter to control o?, Time steps to perform resample C, Variance retention threshold Period P
zr (O,
Initial noise vector
2: for -or=- --..-,
3:
(0,1)
4:
êr+1 -0(2+1,,1+11 ++
Compute the score Predict using Tweedie's 's formula
5:
fo(2(+1) Vall1 (Zi+1 v1 a
6:
zi = vārzo(Zz41) + vI - āe - n E njb1
D Unconditional DDIM step ReSammll time step
7:
E
8:
Initialize Eo(y) with (2
9:
U.S. +- SVD(Z))
Perform SVD on current diffusion state b Calculate eigenvalues
10:
Aj - (where are the singular values of S)
Ck E
11:
b Cumulative sum of eigenvalues Determine rank based on threshold
12:
arg min 7}
13:
A 4-1 U.
Get first r left singular vectors b first right singular vectors
B x V,
15:
for each step in gradient descent do if step number mod P = 0 then
16:
17:
g V 4olvi b1y A(D(
D Compute gradient Project gradient b Reconstruct approximated gradient
18:
R < A ABGB"
19:
G' ARB
else
21:
g' + 4- an(n)) A(D(
> Compute gradient without modification
22:
end if
23:
Update eo(y) using gradient g' end for
24:
2; StochasticResample (y)) zí,
Map back
26:
else
27:
21 x
Unconditional sampling if not resampling
28:
end
end for
30: 3o = D(zo)
Output reconstructed image
31:
return Zo
C.5 DPS
Like DiffStateGrad-PSI.DD our DiftStateGrad-DPS algorithm integrates the state-guided projected gradient into each iteration of the DPS process. The key difference lies in operating directly pixel-based diffusion states rather than latent representations. After computing the standard DPS update z/-1, follow the same projection strategy: performing SVD the current diffusion state X,, determining projection rank using threshold т, and constructing an approximated gradient G² that replaces the standard DPS gradient update.
For experiments, we use the official implementation of DPS (Chung et al., 2023) with default configurations (i.e., noise, step size, etc.) for reproducing baselines.
21

---

Published as a conference paper at ICLR 2025
Algorithm 4 DiffStateGrad-DPS for Image Restoration Tasks
Require: T. T.y. [()) (ōx Y vervaiience retention threshold ær (0. I)
2: for = t do
3: 8 - ss (x),l)
4: to + Vac
(æ: (1 (8)
5:
(O
6:
x111 (1
V dC 18t
to
0,
1-01
I-a:
7:
g: 4 Vv2ly A(
8:
U,S, V (V)(X,)
9:
A + of (8, are the singular values S)
C Σ 22
10:
11:
min T}
12:
A U
13:
B. - V.
R. + AT G.
15:
C
16:
#1_1 x/--1 96
17:
end for
18: return Ø0
C.6 DAPS
We improve upon DAPS by incorporating a state-guided projected gradient. We introduce a variance retention threshold 7 to determine the projection rank. For each noise level the annealing loop, DAPS computes the initial estimate 2 by solving the probability flow ODE using the score model 8p- This estimate represents guess of the clean image given the current noisy sample æt,- We then perform SVD on this estimate in image matrix form (0) (using it as our diffusion state), determine the appropriate rank based on т, and construct projection matrices. Within the Langevin dynamics loop, we calculate the full gradient, combining both the prior term log p(g((())) (ax;) and
likelihood term log pgy) ( ). For each step (i.e., period P 1), we project this gradient using pre-computed matrices and use this projected gradient for the update step. This process repeated for each noise level. progressively refining our estimate of the clean image.
For experiments, we use the official implementation of DAPS (Zhang et al., 2024) with default configurations (i.e., noise, step size, etc.) for reproducing baselines.
22

---

Published as a conference paper at ICLR 2025
Algorithm 5 DiffStateGrad-DAN -DPS for Image Restoration Tasks
Require: Score model 8g, measurement y, noise schedule or, {fa (o. N.) variance retention threshold Sample æT ~ (N
2: for NA, NA - 1.,..... do
3:
Compute To = (o( by solving the probability flow ODE in Eq. (39) with 86 U,S.V V - o "
4:
Perform SVD on initial estimate b Calculate cigenvalues
5:
A - (where s, are the singular values of S)
6:
Cumulative sum of cigenvalues
7:
r arg 2T
Determine rank based on threshold
8:
A U,
Get first left singular vectors b Get first right singular vectors
9:
B+ V,
10:
for = O....,
do
9 Ván log p(a ax.)) Vao 1
B Compute gradient Project gradient
12:
R t- A GB"
G' ARB.
Reconstruct approximated gradient
14:
20+1)
0 nig + Vinie, ~ 100,
for
16:
Sample Ct;
of !I) N (N)
(N)
17:
end for
18:
return 20
C.7 DDNM
We report numbers from DAPS (Zhang et al., 2024), which conducts experiments under identical settings to ours.
C.8 DDRM
We report numbers from DAPS (Zhang et al., 2024), which conducts experiments under identical settings to ours.
C.9
MCG
We report numbers from DPS (Chung et al., 2023), which conducts experiments under identical settings ours.
C.
MPGD-AE
For experiments, we use the official implementation of MPGD-AE (He et al., 2024). For accurate comparison, 1000 DDIM steps and the guidance weight parameter of 0.5 (which find through fine-tuning).
D PROOFS
Proposition 1. Let M be a smooth m-dimensional submanifold of a d-dimensional Euclidean space Rdd where m < d. Assume that for each state E M,, the tangent space T2.M is well-defined, and the projection operator Psz, onto an approximate subspace Sz, closely approximates the projection onto T2, M. For the state E MM and measurement gradient gt E Rd, consider two update rules:
2 = - (standard update),
2 = =z:- 2 (gt) (projected update),
23

---

Published as a conference paper at ICLR 2025
where ni 0 is small step size. Then, for sufficiently small 1, the projected update 2/-1 stays closer to the manifold M than the standard update 21-1. That is,
dist( dist(z/11, < dist(Zt-1,
Proof. Let 21 E M and gr E RR. Decompose the gradient gr into components tangent and normal to M at zz:
9. g gi.
where of M, and gf E N,,M, the normal space 2,.
We have two projection operators:
Mi the exact orthogonal projection onto the tangent space T,
Ps., : an approximate projection operator onto subspace Sz, that closely approximates M.
Assuming that Ps., approximates Pr.,Ms Ms we have:
Ps.,, (gi) = g!و + E,
where e = Ps., (gi) g is the approximation error, which small.
The standard update is:
2t-1 = 2: 11: = 2: -2--7(9 + g.t).
The projected update is:
2=1= -1iss,, (gi) 21 - cc..
Let T (2)) denote the orthogonal projection of onto M. For points close to 2t. approximate
(=)) using the tangent space projection, which comes from the first-order Taylor expansion of M at Zc:
((2) Zt + Pla,M(z 21).
Here, the higher-order terms for T(2;-1) are of order 0( 2+| = (( = O(7)).
=
Similarly, the higher-order terms for T (2 1) are of order (1 2z 21 O(nn (g)))))))
O(7²). We will address these terms the end of the proof.
First, compute the distance from - to M:
dist(z;-1,,M)
x π (z
2 n(. + g. - (z1 + Pra,u(ln(de + gi ))
n(g! + g.tt) + TP + +g)
- + nyPT MP
|| -Pr ((
Since g. E N and Pr.,((9 ( = 0, we have:
dist(z(-1,M) M) = |1
24

---

Published as a conference paper at ICLR 2025
Now, compute the distance from z to M:
dist (ist(z((1,
2 x((1-1)]
2 2 -ng: + +c) - (2+ Pr.,M(-10) + k))
nng + + +TTPo, +c)
= ||- + TPPT=
= n1/lc Pr M()
n IT[(T Pr., Mell
Since ||+ = [|(1 Pro M)ll, we have:
dist dist(2/-1,, =n|[c=]]
(10)
Because small, we can bound ||+||| s g for some small constant d'
Therefore,
dist(z/_11,M) S d'n 1|
Comparing the distances, have:
dist(z--1, dist(z; 2 n g dn 9
= (1 - c' d')m |
= en || ||,
where c c= 1- 0.
Including higher-order terms, we can write:
dist(z/-1,,M) S dist(z--1. - en || + 0
(11)
Therefore, for sufficiently small the linear term dominates the higher-order terms, meaning that the projected update z stays closer to the manifold M than the standard update zt-1:
dist(z dist(z(-1,,M) < 1ist(z)-1,
E MRI EXPERIMENTS
To demonstrate the applicability of DiffStateGrad beyond natural images and the discussed diffusion- based inverse problems, we conduct an additional experiment on 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 MRI in the image space. We utilize the Compressed Sensing Generative Model (Jalal al., 2021), which employs Langevin dynamics for MRI reconstruc- tion (available at ht tttp gitbbb.. com/ ut csi lab/ cssmmmr wariiiaanneiiiiii This method, which we refer to as CSGM-Langevin, is built upon prior works (Song Ermon, 2020; 2019).
We incorporate our proposed DiffStateGrad into the measurement gradient guidance. Specifically, given the complex nature of MRI data, we apply DiftStateCGaa to the magnitude of the complex- valued data while preserving the phase information. This framework uses x; for both measurement gradient computation and incorporation. The unconditional diffusion model was trained on T2- weighted brain datasets from the NYU fastMRI dataset (Zbontar al., 2018; Knoll et al., 2020), and the reported results were averaged over 30 test examples (reporting avg (std)). The measurement operators give samples vertically equispaced in k-space at undersampling rate of R = 4.
While the default optimized measurement gradient step size from (Jalal et al., 2021) is 5.0, the results confirm that DiffStateCira enhances the robustness of CSGM-Lannevin when the step size deviates
25

---

Published as a conference paper at ICLR 2025
40.0
1.0
0.4
DrSttaterredCC Langevis
DifStatedred
CarakLangeetn
0.9
0.3
CSGMManneetn
35.0
0.8
32.5
0.7
0.2
30.0
0.6
Difilannacied Langeaih
27.5
0.5
Langennn
5.0 (default)
7.5
(default)
(default)
MG Step Size
MG Step Size
MG Step Size
(a) PSNR.
(b) SSIM.
(c) LPIPS.
Robustness of DiffStateGrad to MG step size MRI. Comparing CSGM- Langevin (Jalal et al., 2021) with DiffSateGradCCSM-Laggge of the T2-weighted brain fastMRI dataset (Zbontar al., 2018; Knoll et al., 2020). The measurement operator gives vertically and equispaced samples in the k-space with the undersampling rate of R 4
Figure 17:
on 30 examples from the test set
from this optimal value. For example, CSGM-Langevin's performance is significantly reduced when the measurement gradient step size increases to 7.5 or 10, while DiffStateGrad's performance decline less. This suggests that DiffStateGrad can a practical tool to reduce the risk of drastic failure in real-world applications. Finally, given the small sample size and the preliminary nature of this analysis, we leave extensive analysis for future work.
E.1 NYU FASTMRI DATASET
The data used in the above experiment was obtained from the NYU fastMRI Initiative database (fastmri. i med. nyu. edu) (Zbontar et al., 2018; Knoll et al., 2020). We note that NYU fastMRI investigators provided data but did not participate the analysis or writing of this report. See their website for more information on the primary goal of fastMRI to advance machine learning research for the reconstruction of medical images.
26


Reviews:
[{"review_id": "ecS2DKI0Z3", "reviewer": "Reviewer_yaPA", "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.", "questions": "1. Could authors provide more baselines results especially for pixel diffusion\n2. I am curious about whether the authors consider choosing a good initial noise. Like [4, 5]\n\n\n[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.\n\n[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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Well written\n2. Easy to Understand\n3. 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.\n2. There are some other works that try to project the restoration gradient onto the prior manifold (for example DreamClean [1], and MCG [2]).\n3. 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.\n\n\n[1] Xiao, Jie, et al. \"DreamClean: Restoring Clean Image Using Deep Diffusion Prior.\" The Twelfth International Conference on Learning Representations.\n\n[2] Chung, Hyungjin, et al. \"Improving diffusion models for inverse problems using manifold constraints.\" Advances in Neural Information Processing Systems 35 (2022): 25683-25696\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LbkUBnrmyw", "reviewer": "Reviewer_D85G", "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).", "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.\n* 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?\n* 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?\n  * On a related note, how does the measurement noise robustness experiment figures look when noise level is $<0.05$? \n\n***\n\n[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.\n\n[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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The paper is written very well. The work is contextualized well among related work and I enjoyed reading the paper.\n* 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. \n* 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.\n* Additional cost of calculating SVD and projecting gradients seems negligible based on Figure 4.", "weaknesses": "* See the questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FP7FRGoH8W", "reviewer": "Reviewer_kB1E", "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.", "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. \n- 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? \n\n[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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n- 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. \n- 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. \n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iHWJAfcPJo", "reviewer": "Reviewer_3FsW", "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.", "questions": "Please see weaknesses.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The method is straightforward to understand and easy to implement.\n\n2. 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.\n\n2. 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.\n\n3. 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.\n\n\n\n\n**References**\n\n[1] He, Yutong, et al. \"Manifold preserving guided diffusion.\" ICLR 2024\n\n[2] Chung, Hyungjin, et al., \"Decomposed diffusion sampler for accelerating large-scale inverse problems\", ICLR 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)