Recent advancements in diffusion models have been effective in learning data priors for solving inverse problems. They leverage diffusion sampling steps for inducing a data prior while using a measurement guidance gradient at each step to impose data consistency. For general inverse problems, approximations are needed when an unconditionally trained diffusion model is used since the measurement likelihood is intractable, leading to inaccurate posterior sampling. In other words, due to their approximations, these methods fail to preserve the generation process on the data manifold defined by the diffusion prior, leading to artifacts in applications such as image restoration. To enhance the performance and robustness of diffusion models in solving inverse problems, we propose Diffusion State-Guided Projected Gradient (DiffStateGrad), which projects the measurement gradient onto a subspace that is a low-rank approximation of an intermediate state of the diffusion process. DiffStateGrad, as a module, can be added to a wide range of diffusion-based inverse solvers to improve the preservation of the diffusion process on the prior manifold and filter out artifact-inducing components. We highlight that DiffStateGrad improves the robustness of diffusion models in terms of the choice of measurement guidance step size and noise while improving the worst-case performance. Finally, we demonstrate that DiffStateGrad improves upon the state-of-the-art on linear and nonlinear image restoration inverse problems. Our code is available at https://github.com/Anima-Lab/DiffStateGrad.
翻译:近年来,扩散模型在学习数据先验以解决逆问题方面取得了显著进展。这些方法利用扩散采样步骤诱导数据先验,同时在每一步使用测量引导梯度来强制数据一致性。对于一般逆问题,当使用无条件训练的扩散模型时,由于测量似然难以处理,需要进行近似处理,这会导致后验采样不准确。换言之,由于这些近似方法的存在,它们无法保持扩散先验所定义的数据流形上的生成过程,从而在图像恢复等应用中产生伪影。为提升扩散模型在解决逆问题中的性能与鲁棒性,我们提出扩散状态引导的投影梯度方法,该方法将测量梯度投影到一个子空间上,该子空间是扩散过程中间状态的低秩近似。DiffStateGrad作为一个模块,可集成到多种基于扩散的逆问题求解器中,以改善先验流形上扩散过程的保持能力,并滤除导致伪影的成分。我们强调,DiffStateGrad通过优化测量引导步长和噪声的选择,提升了扩散模型的鲁棒性,同时改善了最坏情况下的性能表现。最后,我们在线性与非线性图像恢复逆问题上验证了DiffStateGrad相较于现有最优方法的性能提升。代码已开源:https://github.com/Anima-Lab/DiffStateGrad。