CoSIGN: Few-Step Guidance of ConSIstency Model to Solve General INverse Problems

TL;DR

This paper introduces CoSIGN, a novel method that significantly reduces the number of inference steps needed for diffusion model-based inverse problem solving, achieving high-quality reconstructions with only 1-2 NFEs.

Contribution

The paper proposes a few-step guidance framework using a pretrained consistency model with measurement constraints, enabling high-quality inverse problem solutions with minimal inference steps.

Findings

Achieves state-of-the-art results with 1-2 NFEs.

Supports both single-step and multi-step refinement.

Balances image quality and computational cost effectively.

Abstract

Diffusion models have been demonstrated as strong priors for solving general inverse problems. Most existing Diffusion model-based Inverse Problem Solvers (DIS) employ a plug-and-play approach to guide the sampling trajectory with either projections or gradients. Though effective, these methods generally necessitate hundreds of sampling steps, posing a dilemma between inference time and reconstruction quality. In this work, we try to push the boundary of inference steps to 1-2 NFEs while still maintaining high reconstruction quality. To achieve this, we propose to leverage a pretrained distillation of diffusion model, namely consistency model, as the data prior. The key to achieving few-step guidance is to enforce two types of constraints during the sampling process of the consistency model: soft measurement constraint with ControlNet and hard measurement constraint via optimization. Supporting both single-step reconstruction and multistep refinement, the proposed framework further provides a way to trade image quality with additional computational cost. Within comparable NFEs, our method achieves new state-of-the-art in diffusion-based inverse problem solving, showcasing the significant potential of employing prior-based inverse problem solvers for real-world applications. Code is available at: https://github.com/BioMed-AI-Lab-U-Michgan/cosign.