Diffusion Models for Solving Inverse Problems via Posterior Sampling with Piecewise Guidance

Abstract

Diffusion models are powerful tools for sampling from high-dimensional distributions by progressively transforming pure noise into structured data through a denoising process. When equipped with a guidance mechanism, these models can also generate samples from conditional distributions. In this paper, a novel diffusion-based framework is introduced for solving inverse problems using a piecewise guidance scheme. The guidance term is defined as a piecewise function of the diffusion timestep, facilitating the use of different approximations during high-noise and low-noise phases. This design is shown to effectively balance computational efficiency with the accuracy of the guidance term. Unlike task-specific approaches that require retraining for each problem, the proposed method is problem-agnostic and readily adaptable to a variety of inverse problems. Additionally, it explicitly incorporates measurement noise into the reconstruction process. The effectiveness of the proposed framework is demonstrated through extensive experiments on image restoration tasks, specifically image inpainting and super-resolution. Using a class conditional diffusion model for recovery, compared to the pseudoinverse-guided diffusion model ($\Pi$GDM) baseline, the proposed framework achieves a reduction in inference time of $25\%$ for inpainting with both random and center masks, and $23\%$ and $24\%$ for $4\times$ and $8\times$ super-resolution tasks, respectively, while incurring only negligible loss in PSNR and SSIM.