SURE Guided Posterior Sampling (SGPS)

• SGPS is a trajectory-corrected inference algorithm that integrates diffusion denoising with SURE-based error correction and PCA noise estimation.
• It interleaves conditional posterior guidance with local residual measurement to achieve high-quality image reconstructions under tight computational budgets.
• The method leverages unbiased risk estimates and KL convergence theory to mitigate error accumulation, ensuring efficient correction during sampling.

SURE Guided Posterior Sampling (SGPS) is a trajectory-corrected inference algorithm for diffusion-based inverse problems that leverages Stein’s Unbiased Risk Estimate (SURE) and PCA-based noise estimation to mitigate error accumulation in the critical early and middle stages of sampling. SGPS consistently achieves high-quality reconstructions under tight computational budgets—requiring fewer than 100 Neural Function Evaluations (NFEs)—by interleaving diffusion denoising, conditional posterior guidance, local residual noise measurement, and data-guided correction steps (Kim et al., 29 Dec 2025).

1. Inverse Problem Formulation and Diffusion Priors

The core objective is to recover an unknown image $x \in \mathbb{R}^n$ from noisy linear measurements: $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$ where $A \in \mathbb{R}^{m \times n}$ is a known forward operator (e.g., for super-resolution or deblurring), and $\sigma_y$ is the measurement noise standard deviation.

A diffusion model serves as the learned prior, specified by a stochastic differential equation (SDE): $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$ with reparameterization in EDM by $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$.

The unconditional backward sampling (reverse-time ODE) is: $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$ approximated in practice via a pre-trained denoiser $D_\theta$: $$\nabla_x\log p(x;\sigma) = \frac{D_\theta(x;\sigma) - x}{\sigma^2}$$

Posterior sampling for the inverse problem uses Bayes’ rule: $$\nabla_{x_t}\log p(x_t|y) = \nabla_{x_t}\log p_t(x_t) + \nabla_{x_t}\log p(y|x_t)$$ The data-consistency term $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$0 is typically intractable and must be approximated.

2. SURE-Based Trajectory Correction

2.1 Stein's Unbiased Risk Estimate (SURE)

SURE provides an unbiased estimator of the mean squared error (MSE) for denoising under additive Gaussian noise. For $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$1, $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$2, and denoiser $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$3: $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$4 where $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$5. The trace is estimated by a Monte Carlo probe: $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$6

The SURE gradient direction is obtained by differentiating SURE w.r.t. $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$7: $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$8 The correction is: $$y = A x + \eta, \quad \eta \sim \mathcal{N}(0, \sigma_y^2 I_m)$$9 where $A \in \mathbb{R}^{m \times n}$0 is a user-chosen step size; experiments use $A \in \mathbb{R}^{m \times n}$1.

2.2 Local SURE Gradient Update

After applying conditional guidance, let $A \in \mathbb{R}^{m \times n}$2 be the resulting state. Given a residual noise estimate $A \in \mathbb{R}^{m \times n}$3 (see Section 3), the denoiser is applied: $A \in \mathbb{R}^{m \times n}$4 with SURE evaluated as: $A \in \mathbb{R}^{m \times n}$5 A correction step via autodiff follows, reducing residual noise and pulling samples toward the data manifold.

3. PCA-Based Residual Noise Estimation

Accurate SURE application requires knowledge of the residual variance $A \in \mathbb{R}^{m \times n}$6 in $A \in \mathbb{R}^{m \times n}$7. SGPS employs a patch PCA estimator:

• Decompose $A \in \mathbb{R}^{m \times n}$8 into $A \in \mathbb{R}^{m \times n}$9 overlapping patches $\sigma_y$0, compute mean $\sigma_y$1 and covariance

$\sigma_y$2

• Eigen-decompose $\sigma_y$3 to obtain eigenvalues $\sigma_y$4. For each $\sigma_y$5, define: $\sigma_y$6 The smallest $\sigma_y$7 with $\sigma_y$8 equal to the median of $\sigma_y$9 is chosen. The noise level is then

$$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$0

This estimator is efficient and requires no additional training.

4. SURE Guided Posterior Sampling Algorithm

The SGPS algorithm proceeds as follows:

1. Initialization: Sample $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$1.
2. For $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$2:
   • a) Denoising: $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$3.
   • b) Conditional Guidance: Use Langevin iterations to obtain $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$4 that balances prior and data likelihood.
   • c) PCA Noise Estimation: Estimate residual noise $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$5 from $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$6.
   • d) SURE Gradient Correction: Apply local correction using the SURE gradient to $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$7, yielding $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$8.
   • e) Sample for Next Step: $$\mathrm{d} x_t = \sqrt{2 \dot{\sigma}(t) \sigma(t)}\,\mathrm{d}w_t,\;\; x_0 \sim p_{\text{data}},\; x_T \approx \mathcal{N}(0, \sigma_T^2 I)$$9.
3. Return $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$0.

Distinctive features:

• Estimated, not assumed, noise levels at each step ($$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$1 via PCA).
• Local SURE-based correction at every iteration directly addresses sampling trajectory deviations.

5. Theoretical Properties

• Gaussian-Preservation (Theorem 1): Small-step Langevin guidance ensures the output of the denoiser remains nearly Gaussian in Wasserstein-2 distance $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$2, justifying the use of SURE at each iteration.
• KL-Convergence with SURE Correction (Theorem 2): Under local strong convexity of $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$3 and bounded SURE bias/variance, each correction step reduces the KL divergence to the true posterior, up to $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$4 error, where $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$5: $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$6
• Error-Cascade Mitigation: By removing residual noise at each iteration, SGPS avoids error accumulation characteristic of earlier-stage high-noise samples, enabling accurate inference with $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$7 NFEs.

6. Empirical Performance and Cost Analysis

6.1 Benchmark Domains

SGPS was evaluated on linear (FFHQ256 super-resolution $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$8, box inpainting, random inpainting, Gaussian and motion deblurring) and nonlinear (phase retrieval, nonlinear deblurring, HDR recovery) inverse problems.

6.2 Quantitative Results

Performance with $$\sigma(t) = t,\, t \in [\sigma_{\min} \approx 0, \sigma_{\max} = T]$$9 ($$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$0 NFE) and $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$1 ($$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$2 NFE) is reported using PSNR (higher is better) and LPIPS (lower is better):

Method	NFE	SR4 PSNR / LPIPS	InpaintBox PSNR / LPIPS	InpaintRnd PSNR / LPIPS	GaussDebl PSNR / LPIPS	MotDebl PSNR / LPIPS
SGPS	99	29.38 / 0.179	24.23 / 0.133	30.47 / 0.116	29.35 / 0.179	31.24 / 0.148
DAPS	100	27.69 / 0.230	22.51 / 0.192	26.64 / 0.238	27.77 / 0.220	29.84 / 0.167

Method	NFE	PhaseRet PSNR / LPIPS	NonlinDebl PSNR / LPIPS	HDR PSNR / LPIPS
SGPS	99	24.08 / 0.268	27.33 / 0.197	24.87 / 0.179
DAPS	100	20.83 / 0.402	25.56 / 0.255	24.09 / 0.199

6.3 Computational Cost

• On an RTX 4090: 48 NFE $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$3 4.13 s/image, 99 NFE $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$4 8.46 s/image.
• In competitive SR4 settings at comparable runtime (4 s), SGPS achieves PSNR $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$5 29.06 dB versus DDNM's 29.09 dB.
• Overhead breakdown (for 48 NFE): SURE update (denoiser $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$6 + autograd) 51.2%, Langevin guidance 35.5%, forward denoise 11.2%, PCA 1.8%.

7. Implementation Considerations and Limitations

• Denoiser: U-Net in VP-DDPM/EDM configuration, trained on FFHQ256 images.
• Noise schedule: Geometric, from $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$7 to $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$8, with $$\frac{\mathrm{d}x}{\mathrm{d}t} = -\dot{\sigma}(t)\sigma(t)\nabla_x\log p(x;\sigma(t))$$9 (Karras et al.).
• Sampling Steps: $D_\theta$0 (48 NFE), $D_\theta$1 (99 NFE).
• Langevin Conditional Guidance: 100 iterations per outer step, step size $D_\theta$2.
• PCA: Patch size $D_\theta$3, stride 4, $D_\theta$4 patches per image.
• SURE Hyperparameters: $D_\theta$5, $D_\theta$6.
• Trace Vectors: One random vector per step; additional vectors confer no empirical benefit.

Principal limitations include restriction to pixel-space diffusion samplers, an assumption of known $D_\theta$7 (forward operator), and requirement of local strong convexity for convergence theory. PCA noise estimation may fail for images with little self-similarity; alternative estimators (e.g., spectral) are potential directions. The SURE update uses backpropagation; forward-mode JVP or SPSA could reduce cost. Blind or partially unknown forward operators, non-Gaussian noise, and adaptation to latent-diffusion models remain open areas.

---

For detailed derivations and algorithmic implementations, see (Kim et al., 29 Dec 2025).