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

扩散模型已被证明是解决一般反问题的强先验知识。大多数现有的基于扩散模型的反问题求解器(DIS)采用即插即用的方法，用投影或梯度来指导采样轨迹。虽然有效，但这些方法通常需要数百个采样步骤，从而在推理时间和重建质量之间形成了两难。在这项工作中，我们试图将推理步骤推到1-2 NFEs的边界，同时仍然保持高重建质量。为了实现这一目标，我们提出利用预训练的扩散模型蒸馏，即一致性模型，作为数据先验知识。实现少量步骤指导的关键在于在一致性模型的采样过程中强制执行两种类型的约束：通过ControlNet的软测量约束和通过优化的硬测量约束。所提出的框架支持单步重建和多步细化，还提供了一种用额外计算成本换取图像质量的方法。在可比的NFEs下，我们的方法在基于扩散的反问题求解中实现了新的最先进水平，展示了利用基于先验知识的反问题求解器进行实际应用的巨大潜力。代码可在https://github.com/BioMed-AI-Lab-U-Michgan/cosign获得。

本文旨在解决扩散模型在解决逆问题时需要大量采样步骤的问题，提出了一种基于预训练扩散模型的一步或两步采样方案，以在保持重建质量的同时减少推理时间。

通过使用预训练的一致性模型作为数据先验，通过控制网络和优化来实现软测量约束和硬测量约束，以在少量采样步骤下实现高质量的重建。

该方法在比较少的采样步骤下实现了新的扩散模型逆问题求解的最新成果。该框架支持单步重建和多步细化，并提供了一种以额外计算成本为代价来交换图像质量的方法。

最近的相关研究包括：1）Diffusion Probabilistic Models for Inverse Problems: A Variational Perspective；2）Plug-and-Play Methods Provably Converge with Properly Trained Denoisers；3）Deep Image Prior