Think Twice Before You Act: Improving Inverse Problem Solving With MCMC (2409.08551v1)

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.

• The paper introduces DPMC, a novel algorithm that fuses MCMC techniques with diffusion models to achieve robust posterior approximations.
• It employs an annealed MCMC strategy with intermediate distributions to progressively reduce estimation errors and improve sample quality.
• Empirical evaluations demonstrate that DPMC outperforms DPS in tasks like image inpainting and super-resolution while requiring fewer functional evaluations.

Advancements in Inverse Problem Solving via Diffusion Posterior MCMC

The paper "Think Twice Before You Act: Improving Inverse Problem Solving With MCMC" introduces a novel algorithm, Diffusion Posterior MCMC (DPMC), to enhance the capability of solving inverse problems using pretrained diffusion models. Building upon previous methodologies like Diffusion Posterior Sampling (DPS), this work presents significant methodological progress by integrating Markov Chain Monte Carlo (MCMC) techniques. The paper combines the generative strength of diffusion models with the sampling robustness of MCMC, aiming to address the limitations of posterior approximation in high-noise circumstances.

Context and Challenges

Diffusion models have garnered notable attention due to their proficiency in high-dimensional data generation across modalities like images, videos, audio, and text. In parallel, these models have found utility in solving inverse problems where direct observations of data are unattainable or corrupted by noise. Inverse problems, typically framed as reconstructing original data from degraded measurements, often present computational challenges due to their ill-posed nature. DPS presents a notable framework within this context, yet it struggles with fidelity under noisy conditions due to approximate posterior estimation limitations using Tweedie's formula.

The DPMC Methodology

DPMC innovates by employing annealed MCMC to refine the sampling process, contingent upon a sequence of approximated posterior distributions. While DPS hinges on a direct but inaccurate posterior approximation, DPMC systematically transitions through intermediate distributions, inspired by DPS's approximations, using MCMC to achieve better convergence to the target posterior distribution. Key contributions of the DPMC algorithm include:

1. Implementation of a proposal-exploration sampling strategy where diffusion serves as the proposal step, and MCMC facilitates exploration at each noise level.
2. Introduction of a series of intermediate distributions, which regulate the transition from high noise to cleaner data representations.
3. Utilization of annealed MCMC to accommodate the discrepancies in distributional approximations at each stage, thereby reducing cumulative errors.

Empirical evaluations across a spectrum of inverse problem tasks such as super-resolution, deblurring, inpainting, and phase retrieval demonstrate that DPMC requires a reduced number of functional evaluations (NFE) compared to DPS while yielding superior sample quality. This performance improvement is particularly evident in tasks characterized by high levels of noise or significant information loss, such as image inpainting with large mask intervals.

Implications and Future Directions

DPMC's convergence in reducing estimation noise and improving sample quality has practical implications for scientific fields demanding precise data reconstructions from imperfect measurements. By refining the process of inverse problem-solving, DPMC reinforces the potential of pretrained diffusion models to serve as robust priors in complex applications.

Theoretically, DPMC suggests a broader avenue for integrating advanced sampling strategies into existing diffusion-based frameworks, opening pathways for further research into adaptive tuning of hyperparameters to enhance sample efficiency. Moreover, the work prompts exploration into more advanced MCMC variants, such as Hamiltonian Monte Carlo, which may leverage exact probability estimations for augmented performance.

While the paper underscores the improvements over traditional DPS approaches, future work may focus on automating hyperparameter selection to tailor the sampler dynamically across varied tasks. This refinement would further fortify DPMC’s adaptability and robustness, potentially impacting a broader range of applications in artificial intelligence and beyond. Overall, the paper provides a significant step forward in the domain of inverse problem-solving, proposing a method that combines theoretical rigor with empirical superiority.