---
title: SURE Guided Posterior Sampling (SGPS)
url: https://www.emergentmind.com/topics/sure-guided-posterior-sampling-sgps
type: topic
---

# SURE Guided Posterior Sampling (SGPS)

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 [2512.23232].

## 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 $\nabla_{x_t}\log p(y|x_t)$ 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 $x_{\text{noisy}} = x_0 + z$, $z \sim \mathcal{N}(0, \sigma^2 I)$, and denoiser $f$:
\[
\text{SURE}(x_{\text{noisy}}) = -n \sigma^2 + \|x_{\text{noisy}} - f(x_{\text{noisy}})\|^2 + 2\sigma^2 \operatorname{tr}(J_f(x_{\text{noisy}}))
\]
where $J_f = \partial f / \partial x_{\text{noisy}}$. The trace is estimated by a Monte Carlo probe:
\[
\operatorname{tr} J_f(x) \approx \frac{b^\top(f(x + \epsilon b) - f(x))}{\epsilon},\;\; b \sim \mathcal{N}(0, I)
\]

The SURE gradient direction is obtained by differentiating SURE w.r.t. $x$:
\[
\nabla_x\,\mathrm{SURE}(x) = 2(x-f(x)) - 2\sigma^2 \nabla_x [\operatorname{tr} J_f(x)]
\]
The correction is:
\[
x_{\text{corrected}} = x_{\text{noisy}} - \alpha \nabla_{x_{\text{noisy}}} \mathrm{SURE}(x_{\text{noisy}})
\]
where $\alpha$ is a user-chosen step size; experiments use $\alpha=0.5$.

### 2.2 Local SURE Gradient Update

After applying conditional guidance, let $x_{\text{noisy}}$ be the resulting state. Given a residual noise estimate $\hat{\sigma}_0$ (see Section 3), the denoiser is applied:
\[
\hat{x} = f(x_{\text{noisy}}; \hat{\sigma}_0)
\]
with SURE evaluated as:
\[
\mathrm{SURE} = -n\hat{\sigma}_0^2 + \|x_{\text{noisy}} - \hat{x}\|^2 + 2\hat{\sigma}_0^2 \frac{b^\top(f(x_{\text{noisy}} + \epsilon b) - \hat{x})}{\epsilon}
\]
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 $\hat{\sigma}_0^2$ in $x_{\text{noisy}}$. SGPS employs a patch PCA estimator:

- Decompose $x_{\text{noisy}}$ into $s$ overlapping patches $\{p_i\}_{i=1}^s$, compute mean $\mu$ and covariance
\[
\Sigma = \frac{1}{s} \sum_{i=1}^s (p_i - \mu)(p_i - \mu)^\top
\]
- Eigen-decompose $\Sigma$ to obtain eigenvalues $\lambda_1 \ge \lambda_2 \ge \ldots \ge \lambda_r$. For each $i$, define:
\[
\tau_i = \frac{1}{r - i + 1} \sum_{j=i}^r \lambda_j
\]
The smallest $i$ with $\tau_i$ equal to the median of $\{\lambda_i, ..., \lambda_r\}$ is chosen. The noise level is then
\[
\hat{\sigma}_0 = \sqrt{\tau_i}
\]

This estimator is efficient and requires no additional training.

## 4. SURE Guided Posterior Sampling Algorithm

The SGPS algorithm proceeds as follows:

1. **Initialization:** Sample $x_T \sim \mathcal{N}(0, \sigma_T^2 I)$.
2. **For $t = T, \dots, 1$:**
   - a) **Denoising:** $\hat{x}_{0|t} = D_\theta(x_t, \sigma_t)$.
   - b) **Conditional Guidance:** Use Langevin iterations to obtain $x_{0|t,y}$ that balances prior and data likelihood.
   - c) **PCA Noise Estimation:** Estimate residual noise $\hat{\sigma}_0$ from $x_{0|t,y}$.
   - d) **SURE Gradient Correction:** Apply local correction using the SURE gradient to $x_{0|t,y}$, yielding $x^*_{0|t,y}$.
   - e) **Sample for Next Step:** $x_{t-1} \sim \mathcal{N}(x^*_{0|t,y}, \sigma_{t-1}^2 I)$.
3. **Return $x_0$.**

Distinctive features:
- Estimated, not assumed, noise levels at each step ($\hat{\sigma}_0$ 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 $O(\eta^2 n \sigma_t^2)$, justifying the use of SURE at each iteration.
- **KL-Convergence with SURE Correction (Theorem 2):** Under local strong convexity of $-\log p(x|y)$ and bounded SURE bias/variance, each correction step reduces the KL divergence to the true posterior, up to $O(\beta_t^2)$ error, where $\beta_t = \alpha \hat{\sigma}_0^2$:
\[
D_{\mathrm{KL}}(q_t^* || p) \leq (1 - \beta_t\mu) D_{\mathrm{KL}}(q_t || p) + \beta_t^2 C + \Delta_t
\]
- **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 $<100$ NFEs.

## 6. Empirical Performance and Cost Analysis

### 6.1 Benchmark Domains

SGPS was evaluated on linear (FFHQ256 super-resolution $4\times$, box inpainting, random inpainting, Gaussian and motion deblurring) and nonlinear (phase retrieval, nonlinear deblurring, HDR recovery) inverse problems.

### 6.2 Quantitative Results

Performance with $T=16$ ($\approx 48$ NFE) and $T=33$ ($\approx 99$ 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 $\approx$ 4.13 s/image, 99 NFE $\approx$ 8.46 s/image.
- In competitive SR4 settings at comparable runtime (4 s), SGPS achieves PSNR $\approx$ 29.06 dB versus DDNM's 29.09 dB.
- Overhead breakdown (for 48 NFE): SURE update (denoiser $\times2$ + 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 $t_{\max} = T$ to $t_{\min} = 0.02$, with $\rho = 7$ (Karras et al.).
- **Sampling Steps:** $T=16$ (48 NFE), $T=33$ (99 NFE).
- **Langevin Conditional Guidance:** 100 iterations per outer step, step size $\eta \approx 0.1$.
- **PCA:** Patch size $8 \times 8$, stride 4, $s \approx 1000$ patches per image.
- **SURE Hyperparameters:** $\epsilon = \max(x_{\text{noisy}})/1000 \approx 10^{-3}$, $\alpha = 0.5$.
- **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 $A$ (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 [2512.23232].

Source: https://www.emergentmind.com/topics/sure-guided-posterior-sampling-sgps