Think Twice Before You Act: Improving Inverse Problem Solving With MCMC
This work addresses a bottleneck in using diffusion models for inverse problems, offering a more efficient and accurate method for applications such as image restoration, though it is incremental over existing approaches.
The paper tackles the inaccuracy of Diffusion Posterior Sampling (DPS) in solving inverse problems at high noise levels by proposing DPMC, an Annealed MCMC-based algorithm that reduces accumulated error through intermediate distributions, outperforming DPS with fewer evaluations across tasks like super-resolution and deblurring.
Recent studies demonstrate that diffusion models can serve as a strong prior for solving inverse problems. A prominent example is Diffusion Posterior Sampling (DPS), which approximates the posterior distribution of data given the measure using Tweedie's formula. Despite the merits of being versatile in solving various inverse problems without re-training, the performance of DPS is hindered by the fact that this posterior approximation can be inaccurate especially for high noise levels. Therefore, we propose \textbf{D}iffusion \textbf{P}osterior \textbf{MC}MC (\textbf{DPMC}), a novel inference algorithm based on Annealed MCMC to solve inverse problems with pretrained diffusion models. We define a series of intermediate distributions inspired by the approximated conditional distributions used by DPS. Through annealed MCMC sampling, we encourage the samples to follow each intermediate distribution more closely before moving to the next distribution at a lower noise level, and therefore reduce the accumulated error along the path. We test our algorithm in various inverse problems, including super resolution, Gaussian deblurring, motion deblurring, inpainting, and phase retrieval. Our algorithm outperforms DPS with less number of evaluations across nearly all tasks, and is competitive among existing approaches.