Abstract
Diffusion models have found extensive use in solving numerous inverse problems. Such diffusion inverse problem solvers aim to sample from the posterior distribution of data given the measurements, using a combination of the unconditional score function and an approximation of the posterior related to the forward process. 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 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, not amenable to large-scale problems. In this work, we reinterpret CMs as proximal operators of a prior, enabling their integration into plug-and-play (PnP) frameworks. We propose a solver based on PnP-ADMM, which enables us to leverage the fast convergence of conjugate gradient method. We further accelerate this with noise injection and momentum, dubbed PnP-CM, and show it maintains the convergence properties of the baseline PnP-ADMM. We evaluate our approach on a variety of inverse problems, including inpainting, super-resolution, Gaussian deblurring, and magnetic resonance imaging (MRI) reconstruction. To the best of our knowledge, this is the first CM trained for MRI datasets. Our results show that PnP-CM achieves high-quality reconstructions in as few as 4 NFEs, and can produce meaningful results in 2 steps, highlighting its effectiveness in real-world inverse problems while outperforming comparable CM-based approaches.
Abstract (translated)
扩散模型在解决各种逆向问题中得到了广泛应用。这些扩散逆向问题求解器旨在利用无条件评分函数和正向过程相关的后验分布的近似值来从给定测量数据的数据后验分布采样。最近，一致性模型（CMs）被提出直接预测扩散ODE轨迹上任一点的最终输出，从而仅通过少量的NFE（网络前向传递次数）就能实现高质量的采样。CMs也被应用于逆向问题求解，但现有的基于CM的方法要么需要额外的任务特定训练，要么使用收敛缓慢的数据保真度操作，不适合大规模的问题。 在本工作中，我们重新解释了CMs作为先验概率的proximal算子，使其能够被集成到即插即用（PnP）框架中。我们提出了一个基于PnP-ADMM的方法，使我们可以利用共轭梯度法快速收敛的优点。进一步地，通过引入噪声注入和动量加速这一方法，并将其命名为PnP-CM，证明其保持了基础的PnP-ADMM收敛特性。 我们在包括图像修复、超分辨率重建、高斯去模糊以及磁共振成像（MRI）重构等不同类型的逆向问题上评估了我们的方法。据我们所知，这是首个专门用于MRI数据集训练的一致性模型。实验结果显示，PnP-CM能够在仅4次NFE的情况下实现高质量的图像重建，并在2步内生成有意义的结果，突出显示其在实际逆向问题中的高效性和优越性，超过了同类CM方法的表现。
URL
https://arxiv.org/abs/2509.22736
