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

TL;DR

This paper introduces a versatile diffusion-based framework with piecewise guidance for efficiently solving various inverse problems like image inpainting and super-resolution, without retraining for each task.

Contribution

The proposed method is problem-agnostic, incorporates measurement noise explicitly, and reduces inference time significantly while maintaining high reconstruction quality.

Findings

Achieves 25% faster inference in image inpainting.

Reduces inference time by 23-24% in super-resolution tasks.

Maintains comparable PSNR and SSIM to baseline methods.

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 \blue{pseudoinverse-guided diffusion model (\textrm{\(\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.