Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SURE Guided Posterior Sampling: Trajectory Correction for Diffusion-Based Inverse Problems

Minwoo Kim, Hongki Lim

TL;DR

This work tackles the inefficiency of diffusion-model–based inverse problem solvers, which typically require many Neural Function Evaluations (NFEs) due to error accumulation. It introduces SGPS, a method that applies SURE-gradient updates and PCA-based noise estimation to correct sampling trajectories after conditional guidance, thereby aligning samples with the true data manifold. The authors provide theoretical KL-convergence support for the SURE-based corrections and demonstrate empirically that SGPS delivers state-of-the-art results across linear and nonlinear inverse problems at low NFEs on FFHQ, with favorable runtime trade-offs. Practically, SGPS enables high-quality reconstructions under tight compute budgets, broadening the applicability of diffusion priors to real-time or resource-constrained imaging tasks.

Abstract

Diffusion models have emerged as powerful learned priors for solving inverse problems. However, current iterative solving approaches which alternate between diffusion sampling and data consistency steps typically require hundreds or thousands of steps to achieve high quality reconstruction due to accumulated errors. We address this challenge with SURE Guided Posterior Sampling (SGPS), a method that corrects sampling trajectory deviations using Stein's Unbiased Risk Estimate (SURE) gradient updates and PCA based noise estimation. By mitigating noise induced errors during the critical early and middle sampling stages, SGPS enables more accurate posterior sampling and reduces error accumulation. This allows our method to maintain high reconstruction quality with fewer than 100 Neural Function Evaluations (NFEs). Our extensive evaluation across diverse inverse problems demonstrates that SGPS consistently outperforms existing methods at low NFE counts.

SURE Guided Posterior Sampling: Trajectory Correction for Diffusion-Based Inverse Problems

TL;DR

This work tackles the inefficiency of diffusion-model–based inverse problem solvers, which typically require many Neural Function Evaluations (NFEs) due to error accumulation. It introduces SGPS, a method that applies SURE-gradient updates and PCA-based noise estimation to correct sampling trajectories after conditional guidance, thereby aligning samples with the true data manifold. The authors provide theoretical KL-convergence support for the SURE-based corrections and demonstrate empirically that SGPS delivers state-of-the-art results across linear and nonlinear inverse problems at low NFEs on FFHQ, with favorable runtime trade-offs. Practically, SGPS enables high-quality reconstructions under tight compute budgets, broadening the applicability of diffusion priors to real-time or resource-constrained imaging tasks.

Abstract

Diffusion models have emerged as powerful learned priors for solving inverse problems. However, current iterative solving approaches which alternate between diffusion sampling and data consistency steps typically require hundreds or thousands of steps to achieve high quality reconstruction due to accumulated errors. We address this challenge with SURE Guided Posterior Sampling (SGPS), a method that corrects sampling trajectory deviations using Stein's Unbiased Risk Estimate (SURE) gradient updates and PCA based noise estimation. By mitigating noise induced errors during the critical early and middle sampling stages, SGPS enables more accurate posterior sampling and reduces error accumulation. This allows our method to maintain high reconstruction quality with fewer than 100 Neural Function Evaluations (NFEs). Our extensive evaluation across diverse inverse problems demonstrates that SGPS consistently outperforms existing methods at low NFE counts.
Paper Structure (65 sections, 5 theorems, 77 equations, 12 figures, 6 tables, 1 algorithm)

This paper contains 65 sections, 5 theorems, 77 equations, 12 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Diffusion Models and Inverse Problems
  4. EDM
  5. DAPS
  6. Stein's Unbiased Risk Estimate (SURE)
  7. PCA-based Noise Level Estimation
  8. Method
  9. SGPS Algorithm Components
  10. Denoising with Diffusion Denoiser
  11. Conditional Guidance
  12. Noise Level Estimation
  13. SURE Gradient Update
  14. Experiments
  15. Experimental setup
  16. ...and 50 more sections

Key Result

Theorem 3.1

Under the following assumptions: For a single step from $t$ to $t-1$: Let $\hat{\mathbf{x}}_{0|t,\mathbf{y}} := \hat{\mathbf{x}}_{0|t} + \eta\,\nabla\log p(\mathbf{y}|\hat{\mathbf{x}}_{0|t}) + \sqrt{2\eta}\,\sigma_t\,\boldsymbol{\xi}$, and $R_t := \hat{\mathbf{x}}_{0|t,\mathbf{y}} - \mathbf{m}_t$. Here $\mathcal{L}(R_t)$ denotes th For the multi-step case with $K$ steps: Let $\mathbf{x}_{t-k} \si

Figures (12)

  • Figure 1: Geometric illustration of SGPS. Green curves represent data manifolds $\mathcal{M}_t$ at varying noise levels. Practical sampling deviates from ideal paths: diffusion denoising (red arrow) produces $\hat{\mathbf{x}}_{0|t}$ with errors from high initial noise, while conditional guidance (yellow arrow) creates further deviation by moving to $\hat{\mathbf{x}}_{0|t,\mathbf{y}}$ for measurement consistency but away from the true manifold $\mathcal{M}_0$. SGPS applies a SURE gradient update (purple arrow) to correct these deviations, yielding $\hat{\mathbf{x}}^*_{0|t,\mathbf{y}}$ closer to $\mathcal{M}_0$. Adding noise $\sigma_{t-1}$ (blue arrow) produces $\mathbf{x}_{t-1}$ with reduced accumulated error.
  • Figure 2: Overview of the SGPS sampling process. Starting with noisy sample $\mathbf{x}_t$, the process applies: (1) diffusion denoising to produce $\hat{\mathbf{x}}_{0|t}$, (2) conditional guidance using measurement $\mathbf{y}$ to yield $\hat{\mathbf{x}}_{0|t,\mathbf{y}}$, (3) PCA-based noise level estimation to determine $\hat{\sigma}_0$, (4) SURE gradient update to correct the trajectory to $\hat{\mathbf{x}}^*_{0|t,\mathbf{y}}$, and (5) addition of noise level $\sigma_{t-1}$ to obtain $\mathbf{x}_{t-1}$ for the next step. Image examples show the evolution from noise to clean reconstruction.
  • Figure 3: Qualitative results of SGPS and baseline methods on different general inverse problems. The figure compares the top-3 performing methods for each task, as reported in the NFE 50 comparisons in Table \ref{['linear_combined']}, \ref{['nonlinear_combined']}.
  • Figure 4: Average noise levels (left) and average PSNR (right) across sampling steps in SGPS with and without the SURE gradient update, evaluated on the SR$\times 4$ task with 33 sampling steps (T=33, NFE=99) over 100 samples, using PCA-based noise level estimation.
  • Figure 5: Sensitivity Analysis of the MC-SURE Perturbation Scale ($\epsilon$) on PSNR and LPIPS Performance for the SR4 Task.
  • ...and 7 more figures

Theorems & Definitions (5)

  • Theorem 3.1: Gaussian Preservation in Diffusion Sampling
  • Theorem 3.2: KL Convergence under Biased SURE Gradients
  • Theorem A.3: Single-Step Gaussian Preservation
  • Theorem A.4: Multi-Step Extension
  • Theorem A.8: KL Convergence under Biased SURE Gradients