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

TL;DR

The paper introduces SGPS, a novel method that improves diffusion model-based inverse problem solutions by reducing sampling errors with SURE-guided corrections, enabling high-quality reconstructions with fewer steps.

Contribution

SGPS is the first approach to incorporate SURE-based trajectory correction in diffusion sampling, significantly reducing the number of steps needed for accurate inverse problem solutions.

Findings

SGPS achieves high-quality reconstructions with fewer than 100 NFEs.

SGPS outperforms existing methods at low NFE counts across various inverse problems.

The method effectively mitigates noise-induced errors during sampling.

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.