Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Think Twice Before You Act: Improving Inverse Problem Solving With MCMC

Yaxuan Zhu, Zehao Dou, Haoxin Zheng, Yasi Zhang, Ying Nian Wu, Ruiqi Gao

TL;DR

This work tackles ill-posed inverse problems by leveraging diffusion priors and identifies limitations in DPS's Tweedie-based posterior approximation, especially at high noise. It proposes Diffusion Posterior MCMC (DPMC), which couples a diffusion prior with annealed Langevin MCMC over a sequence of intermediate distributions to better approximate the posterior and reduce error along the diffusion path. Empirically, DPMC outperforms DPS and remains competitive with strong baselines across linear and nonlinear image inverse tasks (e.g., 4× super-resolution, deblurring, inpainting, and phase retrieval) while requiring fewer function evaluations. The results underline the benefit of explicit exploration via MCMC within diffusion-based posteriors and suggest avenues for future work, including explicit likelihood models and HMC-based samplers, as well as considerations of potential misuse and safety.

Abstract

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.

Think Twice Before You Act: Improving Inverse Problem Solving With MCMC

TL;DR

This work tackles ill-posed inverse problems by leveraging diffusion priors and identifies limitations in DPS's Tweedie-based posterior approximation, especially at high noise. It proposes Diffusion Posterior MCMC (DPMC), which couples a diffusion prior with annealed Langevin MCMC over a sequence of intermediate distributions to better approximate the posterior and reduce error along the diffusion path. Empirically, DPMC outperforms DPS and remains competitive with strong baselines across linear and nonlinear image inverse tasks (e.g., 4× super-resolution, deblurring, inpainting, and phase retrieval) while requiring fewer function evaluations. The results underline the benefit of explicit exploration via MCMC within diffusion-based posteriors and suggest avenues for future work, including explicit likelihood models and HMC-based samplers, as well as considerations of potential misuse and safety.

Abstract

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.
Paper Structure (17 sections, 2 theorems, 35 equations, 12 figures, 7 tables, 1 algorithm)

This paper contains 17 sections, 2 theorems, 35 equations, 12 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Inverse Problem
  4. Diffusion Models
  5. Diffusion Posterior Sampling
  6. Diffusion Posterior with MCMC sampling
  7. Experiments
  8. Experiment Settings
  9. Experimental Results
  10. Ablation Study
  11. Conclusion and Future Work
  12. Experimental Details
  13. Theoretical Analysis
  14. More Experimental Results
  15. Sampling time
  16. ...and 2 more sections

Key Result

Lemma B.2

After choosing the step size $\eta = \frac{m\varepsilon^2}{32Ld^2}$, and the number of iterations the last iterate distribution $p_K := \mathrm{Law}(\mathbf{X}_K)$ holds that:

Figures (12)

  • Figure 1: The illustration of DPMC. At each step, DPMC iterates over diffusion proposal step and MCMC exploration step.
  • Figure 2: Qualitative comparison between DPS samples and our DPMC samples. Images in the top rows are from super resolution task. Images in the bottom rows are from Gaussian deblurring task.
  • Figure 3: Qualitative results of DPMC on different linear inverse problem tasks.
  • Figure 4: Inpainting results with box-shape mask using different random seed. The first column is the ground truth. The second column is the masked observation. The third to fifth columns are results get by our algorithm under different random seeds.
  • Figure 5:
  • ...and 7 more figures

Theorems & Definitions (3)

  • Lemma B.2: TV-Convergence of Langevin MCMC Algorithm
  • Theorem B.6
  • proof