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

TL;DR

CoSIGN introduces a few-step framework that leverages a pretrained Consistency Model (CM) as a strong data prior for general inverse problems. By adding a soft measurement constraint through a ControlNet and an optional hard measurement constraint with optimization, the method achieves high-fidelity reconstructions in 1–2 NFEs, with multistep refinement providing further gains. The approach demonstrates state-of-the-art performance on natural and medical imaging tasks in low-NFE regimes and offers real-time inference capabilities, while acknowledging task-specific ControlNet training as a limitation for zero-shot generalization. Overall, CoSIGN showcases how CM priors can be effectively steered for measurement-consistent inverse problem solving with substantial speed advantages. This work highlights the practical potential of few-step, prior-based solvers for real-world imaging applications.

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.

Figures (11)

• Figure 1: Illustration of the sampling trajectory and reconstruction results. In low NFE region, typical DIS methods like DPS (grey arrow) fail to guide original trajectory of CM (blue arrow) towards high-fidelity results. Instead, our method (orange arrow) can guide the trajectory in a single step with ControlNet, and further refine the single-step result with hard measurement constraint and multistep sampling.
• Figure 2: Overview of our proposed CoSIGN method. In Stage 1, measurement is projected onto the signal space through initial reconstruction. In Stage 2, we input the initial reconstruction into the ControlNet as a condition, and guide the pretrained CM with soft measurement constraint. In Stage 3, we further guarantee measurement consistency with hard measurement constraint. Both the first and the third stage are optional and training-free. In the multistep sampling scheme, the single-step reconstruction result can be sent back to the second stage for further refinement by adding a lower level of noise and denoising with CM for a second time.
• Figure 3: Visual results of two linear inverse problems on LSUN bedroom validation set. Zoom in to get a better view.
• Figure 4: Visual results of nonlinear deblurring on LSUN bedroom validation set. Zoom in to get a better view.
• Figure 5: Visual results of sparse-view CT reconstruction with 23 angles on LDCT validation set. A zoomed-in patch with details is provided in the corner.