Diffusion models have found extensive use in solving inverse problems, by sampling from an approximate posterior distribution of data given the measurements. Recently, consistency models (CMs) have been proposed to directly predict the final output from any point on the diffusion ODE trajectory, enabling high-quality sampling in just a few neural function evaluations (NFEs). CMs have also been utilized for inverse problems, but existing CM-based solvers either require additional task-specific training or utilize data fidelity operations with slow convergence, limiting their applicability to large-scale problems and making them difficult to extend to nonlinear settings. In this work, we reinterpret CMs as proximal operators of a prior, enabling their integration into plug-and-play (PnP) frameworks. Specifically, we propose PnP-CM, an ADMM-based PnP solver that provides a unified framework for solving a wide range of inverse problems, and incorporates noise perturbations and momentum-based updates to improve performance in the low-NFE regime. We evaluate our approach on a diverse set of linear and nonlinear inverse problems. We also train and apply CMs to MRI data for the first time. Our results show that PnP-CM achieves high-quality reconstructions in as few as 4 NFEs, and produces meaningful results in 2 steps, highlighting its effectiveness in real-world inverse problems while outperforming existing CM-based approaches.
翻译:扩散模型通过从给定测量数据的近似后验分布中采样,已在解决反问题中得到广泛应用。近期,一致性模型(CMs)被提出,其能够直接从扩散ODE轨迹上的任意点预测最终输出,从而仅需少量神经函数评估(NFEs)即可实现高质量采样。CMs也被用于反问题,但现有基于CM的求解器要么需要额外的任务特定训练,要么采用收敛缓慢的数据保真操作,这限制了其在大规模问题中的适用性,且难以推广至非线性场景。本文重新将CM解释为先验的近端算子,使其能够集成到即插即用(PnP)框架中。具体而言,我们提出PnP-CM,一种基于ADMM的PnP求解器,为求解多种反问题提供统一框架,并通过引入噪声扰动和基于动量的更新策略,提升低NFE模式下的性能。我们在多样化的线性和非线性反问题上评估该方法,并首次将CM训练应用于MRI数据。结果表明,PnP-CM在仅需4次NFE时即可实现高质量重建,且在2步迭代中产生有意义的结果,突显了其在真实世界反问题中的有效性,同时优于现有基于CM的方法。