DiffStateGrad: Manifold-Guided Inverse Solver

• DiffStateGrad is a module that projects measurement gradients onto locally-adaptive subspaces derived from the SVD of diffusion states, ensuring updates remain on the data manifold.
• It reduces spurious artifacts and enhances robustness to hyperparameters in both linear and nonlinear image restoration tasks.
• Empirical evaluations show improved metrics such as PSNR, LPIPS, and SSIM across diverse diffusion solvers and image restoration scenarios.

Diffusion State-Guided Projected Gradient (DiffStateGrad) is a principled module for enhancing diffusion-based inverse problem solvers by constraining measurement-gradient updates to low-dimensional, data-adaptive subspaces derived from the geometry of intermediate diffusion states. This approach selectively filters measurement-guidance steps so that updates remain aligned with the underlying data manifold of the learned diffusion prior, thereby reducing spurious artifacts, improving robustness to hyperparameters, and enhancing worst-case performance in both linear and nonlinear image restoration tasks (Zirvi et al., 2024).

1. Theoretical Foundations and Problem Setting

The central problem addressed is the reconstruction of an unknown signal $x^* \in \mathbb{R}^n$ from noisy linear or nonlinear measurements $y = \mathcal{A}(x^*) + \epsilon$, where $\mathcal{A}$ is a known (possibly nonlinear) forward operator and $\epsilon$ is noise. Diffusion models, often implemented via Stochastic Differential Equations (SDEs) or deterministic samplers (such as DDPM/DDIM), provide powerful generative priors for regularizing such ill-posed inverse problems.

Measurement-gradient guidance, wherein the gradient of a measurement-consistency loss is injected into the diffusion process at each reverse step, is a common mechanism for enforcing data fidelity. However, unconstrained application of this gradient can disrupt the generative process, pushing iterates away from the data manifold captured by the diffusion prior and thereby introducing artifacts, overfitting to noise, or instabilities—especially with large guidance steps or in high-noise regimes (Zirvi et al., 2024).

DiffStateGrad addresses this by projecting the measurement-gradient update at each diffusion step onto a low-rank subspace that locally approximates the tangent space of the data manifold at the current intermediate state, thereby preserving the learned generative geometry.

2. Construction of the Projection Subspace

At each diffusion time step $t$, the current noisy state $x_t \in \mathbb{R}^n$ is reshaped into a matrix $Z_t \in \mathbb{R}^{H \times W}$ (image or latent-image grid). To capture the local geometry, DiffStateGrad computes a truncated Singular Value Decomposition (SVD):

$Z_t = U_t S_t V_t^T$

where $U_t \in \mathbb{R}^{H \times H}$ and $V_t \in \mathbb{R}^{W \times W}$ are orthonormal matrices, $$S_t = \operatorname{diag}(s_1, \dots, s_{\min(H, W)})$$ singular values ordered decreasingly. The principal subspace is defined by selecting the smallest $r$ so that a fixed fraction $\tau$ of the total variance is retained:

$$c_k = \frac{\sum_{j=1}^k s_j^2}{\sum_{j=1}^{\min(H, W)} s_j^2}, \quad r = \min \{ k \mid c_k \geq \tau \}$$

The projection of any gradient matrix $G_t \in \mathbb{R}^{H \times W}$ onto this subspace is

$$P_t(G_t) = U_{t,r}\left(U_{t,r}^T G_t V_{t,r}\right)V_{t,r}^T,$$

where $U_{t, r}$ and $V_{t, r}$ comprise the top-$r$ singular vectors. This constructs an adaptive, data-driven basis that approximates the tangent plane of the local manifold $M_t$ at $x_t$, assuming $M_t$ has low intrinsic dimension (Zirvi et al., 2024).

3. Guided Diffusion Update with Projected Gradient

For inverse problems, data fidelity is imposed by adding a measurement-consistency gradient in the reverse diffusion step. In the general DDPM/DDIM reverse update:

$$x_{t-1} = x_t + \alpha_t s_\theta(x_t, t) + \sigma_t \varepsilon - \eta_t\,g_t$$

where $s_\theta$ is the unconditional score, $\varepsilon \sim \mathcal{N}(0, I)$, and $\eta_t$ is the step size, $g_t$ is the measurement gradient:

$$g_t = \nabla_{x_t} \| y - \mathcal{A}(\hat{x}_{0|t}) \|_2^2$$

with $\hat{x}_{0|t}$ an estimate of the clean signal from $x_t$.

DiffStateGrad replaces the raw update $-\eta_t g_t$ with its projected version $-\eta_t P_t(g_t)$:

$$x_{t-1} = x_t + \alpha_t s_\theta(x_t, t) + \sigma_t \varepsilon - \eta_t P_t(g_t)$$

This confines the update to the locally-adaptive low-rank manifold, suppressing components orthogonal to the data geometry. The result is a reverse step that maintains closer proximity to the learned data manifold, as formalized in the theoretical guarantee (see below).

4. Algorithmic Framework and Hyperparameters

DiffStateGrad is implemented as a modular wrapper around any diffusion-based inverse solver in either pixel or latent space. The essential steps are:

1. Unconditional Sampling: Perform the standard reverse update for $x_t$.
2. Estimate Clean State: Compute $\hat{x}_{0|t}$, either directly (pixel solvers) or via encoder-decoder in latent solvers.
3. Compute Measurement Gradient: $$g_t = \nabla_{x_t}\|y - \mathcal{A}(\hat{x}_{0|t})\|^2$$.
4. Project Gradient: Form $Z_t$ from $x_t$, compute its SVD, select $r$ via variance threshold $\tau$ (e.g., 0.99), then apply $P_t$ to $g_t$.
5. Update: $x_{t-1} \leftarrow$ unconditional step minus $\eta_t P_t(g_t)$.

Hyperparameter guidelines:

• $\tau$: Robust in $[0.6, 0.99]$, default $0.99$.
• Projection period $P$: Default every step ($P=1$); $P=5$ effective for inner-loop algorithms (e.g., ReSample).
• Step size $\eta_t$: Insensitive, coarse grid search suffices; DiffStateGrad improves robustness, admitting an order-of-magnitude range.
• SVD performed on $H \times W$; overhead is amortized and does not critically impact runtime (Zirvi et al., 2024).

5. Theoretical Guarantees and Geometric Interpretation

Let $M \subset \mathbb{R}^d$ be a smooth $m$-dimensional manifold $(m < d)$. The tangent space at $x_t$ is $T_{x_t} M$. If $P_t$ projects onto a subspace approximating $T_{x_t} M$, then for small enough $\eta$ the projected update

$x_{t-1}' = x_t - \eta P_t(g_t)$

remains strictly closer to $M$ (in Euclidean distance) than the unprojected update $x_{t-1} = x_t - \eta g_t$. The update decomposition $g_t = g_t^\parallel + g_t^\perp$ shows that without projection, the full step includes a normal component ($g_t^\perp$) that moves off-manifold, whereas projection removes most of this normal component with only a small approximation error. A first-order analysis of the local manifold projection map underpins this guarantee, demonstrating superior manifold adherence and reduced risk of artifact generation (Zirvi et al., 2024).

6. Empirical Evaluation and Comparative Results

Extensive experiments on FFHQ 256$\times$256 and ImageNet 256$\times$256 validate the improvements of DiffStateGrad across both pixel-space (DPS, DAPS) and latent-space (PSLD, ReSample) diffusion solvers. Tasks encompass box inpainting, random inpainting, Gaussian deblurring, motion deblurring, $4\times$ super-resolution, nonlinear phase retrieval, nonlinear deblur, and high-dynamic-range (HDR) reconstruction.

Key findings:

• LPIPS improvements, e.g., DiffStateGrad-PSLD: 0.158 $\rightarrow$ 0.092 (box inpainting), 0.246 $\rightarrow$ 0.165 (random inpainting).
• PSNR gains: up to +2–3 dB for many tasks.
• Nonlinear phase retrieval: PSNR $27.61 \pm 8.07 \rightarrow 31.19 \pm 4.33$ (ReSample), failure rate (PSNR$<20$) reduced from 26% to 4%.
• Robustness: Wide range of step sizes $\eta_t$ and added measurement noise $\sigma_\text{noise}$; e.g., DAPS SSIM with noise drops from 0.80$\rightarrow$0.44 without DiffStateGrad, but only to 0.70 with DiffStateGrad.
• Projection period and subspace ablations: Full projection ($P=1$) and state-guided subspaces produce maximal gain. Performance is stable across $\tau \in [0.6, 0.99]$.

Qualitative analysis shows substantial reduction in spurious artifacts and hallucinations, especially under aggressive measurement gradients or high-noise conditions (Zirvi et al., 2024).

7. Integration Guidelines and Applicability

DiffStateGrad is architected for compatibility with a broad array of diffusion-guided inverse solvers, serving as a drop-in module just after the unconditional reverse step and prior to state update. Practical steps for integration include:

• Performing SVD on intermediate state, selecting variance threshold, and projecting the measurement gradient before applying it.
• Using default $\tau=0.99$ and matching the guidance step size scale of the base solver.
• Selecting projection period $P=1$ for standard solvers, $P=5$ for those with intensive inner loops.
• Adjustment requires minimal code changes and incurs modest computational overhead.
• Empirical evidence supports significant improvements to robustness, reliability, and artifact mitigation over a range of datasets, tasks, and solver architectures.

Summary Table: Core Components of DiffStateGrad

Component	Mathematical Operation	Practical Role
Projection Subspace	Truncated SVD of $Z_t$ at each step	Data-driven tangent plane
Projection Operator	$$P_t(G_t) = U_{t, r} (U_{t, r}^T G_t V_{t, r}) V_{t, r}^T$$	Filters measurement guidance
Diffusion Update	$x_{t-1} \leftarrow$ uncon. step $- \eta_t P_t(g_t)$	Manifold-adherent, artifact-suppressing

DiffStateGrad systematically leverages the instantaneous low-rank geometry of the diffusion state to ensure that the generative process stays closely aligned to the prior data manifold throughout inverse problem solving. This approach is supported by theoretical analysis, robust empirical gains in standard metrics (LPIPS, PSNR, SSIM), and practical ease of adoption within existing diffusion-based ecosystems (Zirvi et al., 2024).