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

## Abstract

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 D iffusion P osterior MC MC ( 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.

## 1 Introduction

Diffusion Models [, , , ] have recently achieved significant success in high-dimensional data generation, including images [, , , ], videos [, , , , ], audio [, ], text [], and and 3D generation [, , , , ]. Beyond generation, recent works have applied diffusion models to solve inverse problems in a plug-and-play fashion without the need for fine-tuning [, , , , , , , , , , , , , , , ], where the goal is to restore data ${\mathbf{x}}$ from degraded measurements ${\mathbf{y}}$. Among these, one line of works [, , , , ] proposes to modify the inference process of diffusion models with guidance that encourages samples to be consistent with the measure.
Other perspectives include variational inference [, , ], Bayesian filtering [], and solving an inner loop optimization problem during sampling [, ].

A typical challenge in this context is that the posterior distribution $p({\mathbf{x}}|{\mathbf{y}})$ is only defined for clean samples ${\mathbf{x}}$, yet during sampling, an estimation of $p({\mathbf{x}}_{t}|{\mathbf{y}})$ is needed at each diffusion step $t$. DPS [] approximates the intractable posterior through Tweedie’s formula [], enabling its application in general inverse problems. However, this approximation might be inaccurate especially on high noise levels, leading to samples of low quality. To address this, $\Pi$GDM [] attempts to improve the approximated guidance by pseudo-inverting the measurement. Additionally, [, ] aim to enhance DPS in the Latent Diffusion model [] by introducing extra guidance terms or resampling processes to the data space.

In this work, we propose to leverage annealed Markov Chain Monte Carlo (MCMC) to improve solving inverse problems with diffusion models.
As a family of effective algorithms that draw samples from complex distributions, MCMC is widely used in training Energy-Based Models (EBMs) [, , ] and sampling from diffusion models [, ]. Annealed MCMC further proposes to facilitate MCMC sampling by gradually sampling from a sequence of distributions from decreasingly lower temperature, to accelerate the mixing.
Previous work [] has leveraged Annealed MCMC in compositional generation with EBMs or diffusion models. In this work, we propose using MCMC to reduce the error of posterior approximation in solving inverse problems.

There are some earlier works [, ] that employ annealed MCMC with Langevin dynamics to solve inverse problems. They tackle the intractability of the posterior distribution by projecting the current sample onto the measurement subspace. However, these approaches might fail when measurements are noisy or the measurement process is non-linear, as discussed in []. Moreover, a single MCMC step is executed at every noise level before moving to the next noise level.
Alternatively, we propose to build the the intermediate distributions of our annealed MCMC with the approximated posterior distributions derived in []. Despite the fact that such approximation might be inaccurate for the data distributions defined by the forward diffusion, they still remain as a valid sequence of distributions that can be leveraged in annealed MCMC. We further propose to run multiple sampling steps at each noise level similar to [], that empirically improves the performance.

In summary, we make the following contributions in this work:

- •
We propose the Diffusion Posterior MCMC (DPMC) algorithm, which leverages annealed MCMC with a sequence of posterior distribution of data given measurements, whose formula is inspired by DPS [].
- •
We demonstrate that empirically DPMC outperforms DPS in terms of sample quality across various types of inverse problems in image domain, including both linear and nonlinear inverse problems. DPMC also establishes competitive performance compared with other existing approaches.
- •
Through extensive ablation study and comparison with other approaches, we highlight the effectiveness of MCMC-based approaches for solving inverse problems, and a constant improvement of performance with increasingly larger number of sampling step.

## 2 Background

### 2.1 Inverse Problem

We denote data distribution ${\mathbf{x}}\sim p_{data}({\mathbf{x}})$. In many scientific applications, instead of directly observing ${\mathbf{x}}$, we might only have a partial measurement ${\mathbf{y}}$, which is derived from ${\mathbf{x}}$, and we want to restore ${\mathbf{x}}$ from ${\mathbf{y}}$. Formally, we might assume the following mapping between ${\mathbf{x}}$ and ${\mathbf{y}}$

$$ ${\mathbf{y}}={\mathbf{A}}({\mathbf{x}})+{\mathbf{n}},\;{\mathbf{x}}\in\mathbb{ R}^{D},\;{\mathbf{y}}\in\mathbb{R}^{d},\;{\mathbf{n}}\sim\mathcal{N}(0,\sigma^ {2}{\mathbf{I}})$ (1) $$

where ${\mathbf{A}}(\cdot):\mathbb{R}^{D}\mapsto\mathbb{R}^{d}$ is called forward measurement operator and ${\mathbf{n}}$ is the measurement noise following Gaussian distribution. Therefore, we have $p({\mathbf{y}}|{\mathbf{x}})\sim\mathcal{N}({\mathbf{A}}({\mathbf{x}}),\sigma^
{2}{\mathbf{I}})$. Mapping between ${\mathbf{x}}$ to ${\mathbf{y}}$ can be many-to-one. This makes exactly restoring ${\mathbf{x}}$ become an ill-posed problem.

### 2.2 Diffusion Models

Diffusion models [, , , ] define a generative process that gradually transforms a random noise distribution into a clean data distribution. A diffusion model consists of a forward noise injection process and a backward denoising process. Let ${\mathbf{x}}_{0}\sim p_{data}$ denote clean observed samples. DDPM [] defines a Markovian forward process as follows:

$$ $\displaystyle q({\mathbf{x}}_{1:T}|{\mathbf{x}}_{t-1})$ $\displaystyle=q({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1})=\mathcal{N}(\sqrt{\alpha_ {t}}{\mathbf{x}}_{t-1},\beta_{t}{\mathbf{I}})$ $\displaystyle q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ $\displaystyle=\mathcal{N}(\sqrt{\bar{\alpha}_{t}}{\mathbf{x}}_{0},(1-\bar{ \alpha}_{t}){\mathbf{I}})$ (2) $$

where $\{\beta_{t}\}_{t=1}^{T}$ is the manually designed noise schedule that might differ from different works [, , ]. And $\alpha_{t}=1-\beta_{t}$, $\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$. In the backward process, we start from the noise distribution and gradually denoise the samples as follows:

$$ ${\mathbf{x}}_{t-1}=\mu_{\theta}({\mathbf{x}}_{t},t)+\sqrt{\tilde{\beta}_{t}}{ \mathbf{z}}_{t},\;\;{\mathbf{z}}_{t}\sim\mathcal{N}(0,{\mathbf{I}})$ $$

$$ $\mu_{\theta}({\mathbf{x}}_{t},t)=\frac{1}{\sqrt{1-\beta_{t}}}\left({\mathbf{x} }_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}({\mathbf{x} }_{t},t)\right),\;\;\tilde{\beta}_{t}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{ \alpha}_{t}}\beta_{t}$ (3) $$

where $\epsilon_{\theta}({\mathbf{x}}_{t},t)$ is parameterized by a neural network [, ]. Let the marginal distribution defined by the forward process of ${\mathbf{x}}_{t}$ be denoted as $p_{t}({\mathbf{x}}_{t})$. When trained with denoising score matching loss [, ], with sufficient data and model capacity, for almost all ${\mathbf{x}}$ and $t$, $\epsilon_{\theta}({\mathbf{x}}_{t},t)$ corresponds to the gradient field $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ as follows:

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})=-\frac{ \epsilon_{\theta}({\mathbf{x}}_{t},t)}{\sqrt{1-\bar{\alpha}_{t}}}$ (4) $$

Other than DDPM, DDIM[] defines a Non-Markovian forward process that shares the same training objective as DDPM. Thus, a model trained with DDPM can be directly applied in the DDIM sampler to accelerate the sampling process. A DDIM sampler follows:

$$ $\displaystyle{\mathbf{x}}_{t-1}=\sqrt{\bar{\alpha}_{t-1}}\left(\frac{{\mathbf{ x}}_{t}-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}({\mathbf{x}}_{t},t)}{\sqrt{ \bar{\alpha}_{t}}}\right)+\sqrt{1-\bar{\alpha}_{t-1}-\sigma_{t}^{2}}\epsilon_{ \theta}({\mathbf{x}}_{t},t)+\sigma_{t}{\mathbf{z}}_{t}$ (5) $$

where ${\mathbf{z}}_{t}\sim\mathcal{N}(0,{\mathbf{I}})$ and variance $\sigma_{t}$ can be arbitrary defined.

### 2.3 Diffusion Posterior Sampling

To solve the ill-posed inverse problem, Diffusion Posterior Sampling (DPS) [] recruits pretrained diffusion models [, ] as prior and propose an iterative optimization algorithm. According to Bayes rule, the gradient field of posterior distribution $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t}|y)=\nabla_{{\mathbf{x}}_{t}}
\log p_{t}({\mathbf{x}}_{t})+\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{y}}|{
\mathbf{x}}_{t})$ where the term $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ is estimated by the pretrained diffusion model using equation 4. In DPS, the authors adapt the approximation $p({\mathbf{y}}|{\mathbf{x}}_{t})\approx p({\mathbf{y}}|\hat{{\mathbf{x}}}_{0})$, where $\hat{{\mathbf{x}}}_{0}({\mathbf{x}}_{t})=E_{{\mathbf{x}}_{0}\sim p({\mathbf{x}
}_{0}|{\mathbf{x}}_{t})}[{\mathbf{x}}_{0}]$ can be estimated through Tweedie’s formula []:

$$ $\displaystyle\hat{{\mathbf{x}}}_{0}=\frac{{\mathbf{x}}_{t}-\sqrt{1-\bar{\alpha }_{t}}\epsilon_{\theta}({\mathbf{x}}_{t},t)}{\sqrt{\bar{\alpha}_{t}}}$ (6) $$

Given noisy observation data $x_{t}$, DPS makes the following updates

$$ $\displaystyle{\mathbf{x}}_{t-1}^{\prime}$ $\displaystyle=\frac{1}{\sqrt{1-\beta_{t}}}({\mathbf{x}}_{t}-\frac{\beta_{t}}{ \sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}({\mathbf{x}}_{t},t))+\sigma_{t}{ \mathbf{z}}_{t},\;{\mathbf{z}}_{t}\sim\mathcal{N}(0,{\mathbf{I}})$ $\displaystyle{\mathbf{x}}_{t-1}={\mathbf{x}}_{t-1}^{\prime}-\zeta_{t}\nabla_{{ \mathbf{x}}_{t}}\|{\mathbf{y}}-{\mathbf{A}}(\hat{{\mathbf{x}}}_{0})\|^{2}_{2}$ (7) $$

In practice, DPS employs the following adaptive step size parameter $\zeta_{t}$:

$$ $\zeta_{t}=\frac{\zeta}{\|{\mathbf{y}}-{\mathbf{A}}(\hat{{\mathbf{x}}}_{0})\|_{ 2}}$ (8) $$

where $\zeta$ is a fix constant.

## 3 Diffusion Posterior with MCMC sampling

In this section, we introduce our DPMC model, which combines diffusion models and Markov Chain Monte Carlo (MCMC) sampling. The former one is for the progressive denoising to provide an unconditional proposal distribution. The latter is for the conditional guidance by the measurement ${\mathbf{y}}$. According to Theorem 1 in [], error of the approximation $p({\mathbf{y}}|{\mathbf{x}}_{t})\approx p({\mathbf{y}}|\hat{{\mathbf{x}}}_{0})$ is bounded by the estimation error between ${\mathbf{x}}_{0}$ and $\hat{{\mathbf{x}}}_{0}$. When the noise level is low, $p({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ can be single-modal and the estimation $\hat{{\mathbf{x}}}_{0}$ and the posterior approximation might be accurate. However, at higher noise levels, where $p({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ can be indeed multi-modal and $\hat{{\mathbf{x}}}_{0}$ can be far from ${\mathbf{x}}_{0}$, this approximation can be too loose, potentially leading to inferior results. As illustrated in Figure 2, while the samples of DPS are valid, they might fail to capture vivid details. To address these problems, we instead resort to MCMC sampling.

Figure: Figure 1: The illustration of DPMC. At each step, DPMC iterates over diffusion proposal step and MCMC exploration step.
Refer to caption: extracted/5852151/images/diagram.png

Unlike diffusion models, which explicitly define all the noise distributions $p_{t}({\mathbf{x}}_{t})$ through the forward process, MCMC sampling does not require knowledge of the sample distributions at each intermediate MCMC step. As long as a valid transition kernel and a sufficient number of MCMC sampling steps are used, MCMC is known to sample from arbitrary underlying distributions.

Figure: Algorithm 1 DPMC Algorithm

Annealed MCMC [, , ] is a widely used technique to accelerate the MCMC sampling process for highly multi-modal data. In the annealed sampling process, samples gradually progress through a series of intermediate distributions with different temperatures. MCMC is applied at each intermediate distribution to enable samples to transition from the previous distribution to the current one. In this work, we apply annealed MCMC to solve the inverse problem. We introduce a series of intermediate distributions $\{\tilde{p}_{t}({\mathbf{x}}_{t}|{\mathbf{y}})\}_{t=1}^{T}$. Note that we do not expect $\tilde{p}_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ to be close to the true posterior distribution $p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ at every intermediate distribution. Instead, we only require $\tilde{p}_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ to agree with $p_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ at the clean image distribution. The key idea is that MCMC sampling can bridge the gap between different intermediate distributions. And the intermediate distributions only need to form a trajectory that enables MCMC to smoothly transition from the noise distribution to the target distribution in the clean data space.

Following [], we define the intermediate distributions $\tilde{p}_{t}({\mathbf{x}}_{t}|y)$ as

$$ $\displaystyle\tilde{p}_{t}({\mathbf{x}}_{t}|{\mathbf{y}})\propto p_{t}({ \mathbf{x}}_{t})\exp(-\rho\|{\mathbf{y}}-{\mathbf{A}}(\hat{{\mathbf{x}}}_{0}) \|_{2}^{2})$ (9) $$

where $\rho=1/\sigma^{2}$. $p_{t}({\mathbf{x}}_{t})$ denotes the diffusion prior at noise level $t$ and $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ can be estimated using equation 4. However, instead of using 7 for sampling, which requires the intermediate distribution to be sufficiently close to the ground truth posterior at each noise level, we propose a new proposal-and-update algorithm based on annealed MCMC sampling.

Proposal Stage: Given samples ${\mathbf{x}}_{t}$ that follow the intermediate distribution $\tilde{p}_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$, we first denoise them to $t-1$ following the standard diffusion step without considering the extra guidance. We denote the proposed sample as $\tilde{{\mathbf{x}}}_{t-1}$. The proposal step aims to provide a good initialization for the intermediate distribution at $\tilde{p}_{t-1}({\mathbf{x}}_{t-1}|{\mathbf{y}})$ with the help of the diffusion model. The samples $\tilde{{\mathbf{x}}}_{t-1}$ might not fully adhere to the target distribution $\tilde{p}_{t-1}({\mathbf{x}}_{t-1}|{\mathbf{y}})$, but we hope they are close enough to the target distribution so that we can obtain good samples with a few MCMC sampling steps.

Exploration Stage: We then encourage the samples to explore the landscape at noise level $t-1$ and follow $\tilde{p}_{t-1}({\mathbf{x}}_{t-1}|{\mathbf{y}})$ with MCMC updates. We employ Langevin Dynamics [] as the transition kernel. Specifically, starting from ${\mathbf{x}}_{t-1}^{0}=\tilde{{\mathbf{x}}}_{t-1}$, we iterate the following updates:

$$ $\displaystyle{\mathbf{x}}_{t-1}^{k+1}={\mathbf{x}}_{t-1}^{k}+\eta_{t-1}\nabla_ {{\mathbf{x}}_{t-1}^{k}}\log\tilde{p}_{t-1}({\mathbf{x}}_{t-1}^{k}|{\mathbf{y} })+\sqrt{2\eta_{t-1}}{\mathbf{w}},\;{\mathbf{w}}\sim\mathcal{N}(0,{\mathbf{I}})$ (10) $$

We named our algorithm Diffusion Posterior MCMC (DPMC). We illustrate DPMC Figure 1 and in Algorithm 1. We also provide a theoretical analysis which shows that the conditional distribution derived by DPMC is $\varepsilon$-close to the ground truth, i.e.

$$ $D_{\mathrm{TV}}(q_{0}({\mathbf{x}}_{0}\mid{\mathbf{y}})~{}\|~{}p^{*}({\mathbf{ x}}_{0}\mid{\mathbf{y}}))\leq\varepsilon$ $$

with a carefully chosen step size $\eta_{t}$ and inner loop $K$ of the Langevin MCMC algorithm, under several assumptions on the score estimation error, conditional probability approximation error as well as some convexity and Lipschitz continuity conditions. Detailed assumptions and main theorem for the convergence of our DPMC algorithm is deferred to Appendix B.

Figure: 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.
Refer to caption: extracted/5852151/images/DPS_ours.png

Implementation Details
We apply the same implementation as [] in equation 8 by introducing a step size parameter $\zeta_{t}$ to be inverse proportional to optimization distance $\|{\mathbf{y}}-{\mathbf{A}}(\hat{{\mathbf{x}}}_{0})\|_{2}$. This setting equals to sampling according to the prior distribution tilted by an exponential distribution of the optimization distance. In practice, we calculate $\nabla_{{\mathbf{x}}_{t}}\tilde{p}_{t}({\mathbf{x}}_{t}|{\mathbf{y}})$ at each step using the following equation:

$$ $\displaystyle\nabla_{{\mathbf{x}}_{t}}\log\tilde{p}_{t}({\mathbf{x}}_{t}|{ \mathbf{y}})=\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})-\xi_{t} \nabla_{{\mathbf{x}}_{t}}\|{\mathbf{y}}-{\mathbf{A}}(\hat{{\mathbf{x}}}_{0})\| _{2}$ (11) $$

where $\xi_{t}$ represents a constant or variable weighting schedule. Note that we replace square $l_{2}$ norm by $l_{2}$ norm itself. According to [, ], using $l_{2}$ norm might make the optimization to be more robust to outliers and can achieve more stable results than square $l_{2}$ norm as it imposes a smaller penalty for large errors.

While adding MCMC exploration might inevitably introduce more sampling steps at each intermediate distribution and thus increase the sampling time, we find that with DPMC, we can counteract this by reducing the number of intermediate distributions needed. In fact, we find that using 200 intermediate distributions is sufficient for DPMC. Moreover, during the early stage of diffusion, the samples are very close to Gaussian noise. Considering that the reason for inserting intermediate distributions is to provide a good initialization for sampling more complex distributions that are very different than Gaussian, we can skip those early stages with high noise levels and start MCMC sampling directly from a moderate noise level. On the other hand, at very low noise levels, $p({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ almost becomes single-modal, and the approximation $p({\mathbf{y}}|{\mathbf{x}}_{t})\sim p({\mathbf{y}}|{\mathbf{x}}_{0})$ becomes sufficiently accurate. Thus, we can directly follow DPS without the need for additional MCMC sampling. In fact, we find that DPMC performs well when we apply the proposal and exploration steps only to the middle steps of the sampling process, while following the original DPS setting at both ends. This further reduces the number of evaluation steps needed. In practice, DPMC can achieve much better results compared to DPS with an even smaller number of evaluations (NFE).

## 4 Experiments

In this section, we conduct experiments on the DPMC method proposed by us, which outperforms existing baselines on diffusion posterior sampling. To start with, we thoroughly introduce the models, datasets and settings of our experiments.

### 4.1 Experiment Settings

Datasets and Pretrained Model: Following [], we test our algorithm on FFHQ $256\times 256$ dataset [] and ImageNet $256\times 256$ dataset []. Same as [, ], we use 1k validation images for each dataset. For FFHQ dataset, we use the pretrained model provided by []. For ImageNet dataset, we use the pretrained model provided by [].

Figure: Figure 3: Qualitative results of DPMC on different linear inverse problem tasks.
Refer to caption: extracted/5852151/images/all_samples.png

Inverse Problems: We evaluate the effectiveness of our DPMC algorithm on the following inverse problems: (i) Super-resolution with 4x bicubic downsampling as the forward measurement; (ii) Random inpainting using both box masks and random masks; (iii) Deblurring with Gaussian blur kernels and motion blur kernels (^1^11[https://github.com/LeviBorodenko/motionblur](https://github.com/LeviBorodenko/motionblur)); (iv) Phase retrieval, where we perform a Fourier transformation on each image and use only the Fourier magnitude as the measurement. Among these inverse problems, (i), (ii), and (iii) are linear inverse problems, while (iv) is a nonlinear inverse problem. We use Gaussian noise with $\sigma=0.05$ for all tasks and set task-related parameters according to [, ] to ensure a fair comparison.

Hyper-Parameter Setting:
We use the DDIM sampler with default variance $\sigma_{t}=0$ as our diffusion proposal sampler. We use $T=200$ intermediate distributions and $K=4$ MCMC steps at each intermediate distribution. Empirically, we set $\eta_{t}=\eta\beta_{t}$ and $\xi_{t}=\xi\bar{\alpha}^{3}_{t}$, where $\eta$ and $\xi$ are task-related constants. We find these schedules work well across all settings. The detailed parameter settings can be found in Section A in the Appendix. As discussed in 3, we apply the proposal-exploration step in the middle 60% of the total sampling steps and resort to original DPS sampler step at the first 30% and the last 10% noise levels. An exception is the inpainting task with a box-type mask on the ImageNet dataset, where we find that applying the proposal-exploration step until clean images yields better results. Considering both the initial proposal step and the MCMC sampling steps, the NFE of this setting is around 700, which is smaller than the 1000 NFE of DPS. All of the experiments are carried out on a single Nvidia A100 GPU. We report the running clock time of DPMC in Table 6 in Appendix C.

Baselines:
On linear inverse problems, we compare DPMC algorithm with the original DPS[], filter posterior sampling (FPS) with and without sequential Monte Carlo sampling[], denoising diffusion
restoration models (DDRM) [], manifold constrained gradients (MCG) [], plug-and-play alternating direction method of multiplier (PnP-ADMM) [], Score-Based SDE [, ] and alternating direction method of multiplier with total-variation (TV) sparsity regularized optimization (ADMM-TV). On phase retrieval, we compare DPMC with DPS[], oversampling
smoothness (OSS)[], Hybrid input-output (HIO) [] and error reduction (ER) [] algorithm.

### 4.2 Experimental Results

Linear Inverse Problems: We carried out experiments on FFHQ and ImageNet datasets. We show qualitative samples of each task in Figure 3. DPMC is able to generate valid samples given the noisy, degraded input. Note that, as discussed in 2.1, given the information loss in the image degradation process, exactly restoring the original ${\mathbf{x}}$ is ill-posed. Instead, an effective algorithm should be able to fill in meaningful content that agrees with the noisy observation. This is particularly true for tasks with large information loss, such as inpainting large areas in the image with random or box masks. We show an example in Figure 4, where we present samples generated with different random seeds using the same noisy observation ${\mathbf{y}}$ in the inpainting task with a box-shaped mask. DPMC is capable of generating various samples that agree with the observed part.

Qualitative results are shown in Table 1 and Table 2. Following [, ], we report the LPIPS score [] and FID [] score. The LPIPS score measures the similarity of predicted samples with the ground truth at the single image level. As discussed in [], unlike PSNR and SSIM, which capture shallow, low-level features and might fail to account for many nuances of human perception, LPIPS focuses more on structured information related to human perception. On the other hand, FID measures the distribution differences between generated samples and observations. We believe these two metrics are well-suited to reflect sample performance given the many-to-one mapping nature between ${\mathbf{y}}$ and ${\mathbf{x}}$. As we can see, DPMC achieves similar or better results than DPS on both datasets across all tasks using less NFE. This is also evident in the qualitative comparisons in Figure 2, where DPS provides blurred samples and our DPMC fills in more vivid details. Compared to other strong baselines, such as FPS or FPS-SMC [], DPMC achieves superior results on most tasks, especially in terms of FID. We have included more qualitative and quantitative comparison with strong baselines in Appendix C.3. These results demonstrate the effectiveness of introducing MCMC in the sampling process.

Nonlinear Inverse Problem: We conducted the phase retrieval experiment on the FFHQ dataset. For this, we utilized $T=200$ intermediate distributions and $K=6$ MCMC sampling steps at each intermediate distribution. The proposal-exploration step was still applied to the middle 60% of the sampling steps, corresponding to 920 NFE. Similar to DPS, we observed that the quality of final samples depend on the initialization. Therefore, we followed DPS by generating four different samples and selecting the best one. Our qualitative results are shown in Figure 5, and the quantitative results are reported in Table 4.2. Compared to DPS, our DPMC achieved better LPIPS and a similar FID score.

**Table 1: Quantitative results of various linear inverse problems on FFHQ $256\times 256$-1k validation set. Bold denotes the best result for each task and underline denotes the second best result.**
| Methods | Super Resolution | Inpainting (box) | Gaussian Deblur | Inpainting (random) | Motion Deblur |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS |
| DPMC (Ours) | 21.93 | 0.212 | 19.59 | 0.160 | 21.34 | 0.210 | 21.26 | 0.205 | 20.73 | 0.213 |
| FPS [] | 26.66 | 0.212 | 26.13 | 0.141 | 30.03 | 0.248 | 35.21 | 0.265 | 26.18 | 0.221 |
| FPS-SMC [] | 26.62 | 0.210 | 26.51 | 0.150 | 29.97 | 0.253 | 33.10 | 0.275 | 26.12 | 0.227 |
| DPS [] | 39.35 | 0.214 | 33.12 | 0.168 | 44.05 | 0.257 | 21.19 | 0.212 | 39.92 | 0.242 |
| DDRM [] | 62.15 | 0.294 | 42.93 | 0.204 | 74.92 | 0.332 | 69.71 | 0.587 | - | - |
| MCG [] | 87.64 | 0.520 | 40.11 | 0.309 | 101.2 | 0.340 | 29.26 | 0.286 | - | - |
| PnP-ADMM [] | 66.52 | 0.353 | 151.9 | 0.406 | 90.42 | 0.441 | 123.6 | 0.692 | - | - |
| Score-SDE [, ] | 96.72 | 0.563 | 60.06 | 0.331 | 109.0 | 0.403 | 76.54 | 0.612 | - | - |
| ADMM-TV | 110.6 | 0.428 | 68.94 | 0.322 | 186.7 | 0.507 | 181.5 | 0.463 | - | - |

**Table 2: Quantitative results of various linear inverse problems on ImageNet $256\times 256$-1k validation set. Bold denotes the best result for each task and underline denotes the second best result.**
|  | Super Resolution | Inpainting (box) | Gaussian Deblur | Inpainting (random) | Motion Deblur |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS |
| DPMC (Ours) | 31.74 | 0.307 | 30.55 | 0.221 | 33.62 | 0.318 | 30.25 | 0.292 | 30.88 | 0.303 |
| FPS [] | 47.32 | 0.329 | 33.19 | 0.204 | 54.41 | 0.396 | 42.68 | 0.325 | 52.22 | 0.370 |
| FPS-SMC [] | 47.30 | 0.316 | 33.24 | 0.212 | 54.21 | 0.403 | 42.77 | 0.328 | 52.16 | 0.365 |
| DPS [] | 50.66 | 0.337 | 38.82 | 0.262 | 62.72 | 0.444 | 35.87 | 0.303 | 56.08 | 0.389 |
| DDRM [] | 59.57 | 0.339 | 45.95 | 0.245 | 63.02 | 0.427 | 114.9 | 0.665 | - | - |
| MCG [] | 144.5 | 0.637 | 39.74 | 0.330 | 95.04 | 0.550 | 39.19 | 0.414 | - | - |
| PnP-ADMM [] | 97.27 | 0.433 | 78.24 | 0.367 | 100.6 | 0.519 | 114.7 | 0.677 | - | - |
| Score-SDE [, ] | 170.7 | 0.701 | 54.07 | 0.354 | 120.3 | 0.667 | 127.1 | 0.659 | - | - |
| ADMM-TV | 130.9 | 0.523 | 87.69 | 0.319 | 155.7 | 0.588 | 189.3 | 0.510 | - | - |

Figure: 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.
Refer to caption: extracted/5852151/images/inpainting_seed2.png

[Uncaptioned image]: extracted/5852151/images/phase_retrieval.png

### 4.3 Ablation Study

We conducted ablation studies on critical parameters of DPMC using the Gaussian deblur task on the FFHQ dataset. Table 4(a) examines the impact of the number of MCMC sampling steps. Table 4(b) explores the influence of the number of intermediate distributions. Table 4(c) evaluates different weighting schedules by setting $\xi_{t}\propto f(\bar{\alpha}_{t})$, where $f(\cdot)$ ranges from a constant function to $\bar{\alpha}_{t}^{4}$. For each ablation setting, we adjusted the task-related parameters $\eta$ and $\xi$ to optimize the current configuration. The results demonstrate that using more sampling steps or intermediate distributions improves performance, with $T=200$ and $K=4$ being a good choice for balancing performance and sample efficiency. Additionally, setting $\xi_{t}=\xi\bar{\alpha}_{t}^{3}$ yields better results compared to constant scheduling or other alternatives.

**Table 4: Ablation studies on number of MCMC sampling step $K$, number of intermediate distribution $T$ and $\epsilon_{t}$ schedule. In 4(a), we keep the same number of intermediate distributions and change the number of MCMC sampling steps; in 4(b), we keep the same MCMC sampling steps and change the number of intermediate distributions; in 4(c), we keep the number of MCMC steps and the number of intermediate distributions and change different weighted schedule $\xi_{t}$.**
|  |  |
| --- | --- |

## 5 Conclusion and Future Work

In this study, we propose DPMC, an algorithm based on a specific formula of the posterior distribution and Annealed MCMC sampling to solve inverse problems. We demonstrate that DPMC results in superior sample quality compared to DPS across various inverse problems with fewer functional evaluations. Our study underscores the benefit of incorporating more sampling steps into each intermediate distribution to encourage exploration. Additionally, it is beneficial to trade in the number of intermediate distributions with the number of MCMC exploration steps. One potential limitation of our current work is the necessity to manually tune the weighting schedule and other hyper-parameters, whose optimal values can vary for different tasks. An interesting future direction is to explore scenarios where an explicit estimation $p({\mathbf{x}})$ is provided by an EBM or another likelihood-based estimation technique, enabling the use of more advanced samplers such as Hamiltonian Monte Carlo (HMC) with adaptive step sizes. As a powerful algorithm that can process images, DPMC might have the potential to cause negative social consequences, including deepfakes, misinformation, and privacy breaches. We believe that more research and resources are needed to mitigate these risks.

## Appendix A Experimental Details

In this study, for linear inverse problem, we use $T=200$ intermediate distributions and $K=4$ MCMC sampling step at each intermediate distributions as our default setting. For the nonlinear inverse problem, phase retrieval, we use T = 200 intermediate distributions and K = 6 MCMC sampling step at each intermediate distributions. Our main task specific parameters are the Guidance weight $\xi$ and Langevin step size coefficient $\eta$. We show the parameter setting for each task in Table 5.
All of our experiments are carried on a single Nvidia A100 GPU.

**Table 5: Hyper-parameter settings for various linear inverse problems. $\xi$: Guidance Weight; $\eta$: Langevin Step Size Coefficient**
|  |  |
| --- | --- |

## Appendix B Theoretical Analysis

In this section, we provide some theoretical analysis on the DPMC model proposed by us. First, we start with some existing results on the KL-convergence of the Langevin MCMC algorithm. Denote $p^{*}$ as the target distribution over $\mathbb{R}^{d}$, and $s^{*}(\mathbf{x}):=\nabla\log p^{*}(\mathbf{x})$ as its score function. The Langevin MCMC algorithm with step size $\eta$ is given by:

$$ $\mathbf{X}_{0}\sim p_{0},~{}~{}\mathbf{X}_{i+1}=\mathbf{X}_{i}+\eta s^{*}( \mathbf{X}_{i})+\sqrt{2\eta}\cdot\xi_{i}$ $$

where $\xi_{i}\sim\mathcal{N}(0,\bm{I}_{d})$. Before we state the convergence rate, we propose the strong convexity and Lipschitz continuity assumption that the target distribution $p^{*}$ needs to satisfy.

###### Assumption B.1 (Strong Convexity and Lipschitz Continuity of Potential Function) .

Let $U(\mathbf{X})=-\log p^{*}(\mathbf{X})$ be the potential function. It has $m$-strong convexity and its gradient has $L$-Lipschitz continuous, i.e.:

$$ $m\bm{I}_{d}\preceq\nabla^{2}U(\mathbf{X})\preceq L\bm{I}_{d}~{}~{}~{}\text{for }~{}\forall\mathbf{X}\in\mathbb{R}^{d}.$ $$

Under this assumption, [] propose the following total variation convergence result of the Langevin MCMC algorithm.

###### Lemma B.2 (TV-Convergence of Langevin MCMC Algorithm) .

After choosing the step size $\eta=\frac{m\varepsilon^{2}}{32Ld^{2}}$, and the number of iterations

$$ $K=\frac{32L^{2}d\log(D_{\mathrm{TV}}(p_{0}\|p^{*})/\varepsilon)}{m^{2} \varepsilon^{2}},$ $$

the last iterate distribution $p_{K}:=\mathrm{Law}(\mathbf{X}_{K})$ holds that:

$$ $D_{\mathrm{TV}}\left(p_{K}\|p_{0}\right)\leq\varepsilon.$ $$

Notice that in the original Theorem 1 of [], the authors set $p_{0}=\mathcal{N}(0,\frac{1}{m}\bm{I}_{d})$. However, the proof does not actually rely on this choice. Besides, we use the corollary of TV-distance convergence instead of the original KL-divergence result since TV distance holds triangular inequality, which benefits our analysis.

Next, we study the distribution estimation error and provide a convergence guarantee for DPMC proposed by us. Our main result requires the following assumptions on the conditional distribution $p({\mathbf{x}}_{t}\mid{\mathbf{y}})$ as well as the unconditional score estimation error.

###### Assumption B.3 (Unconditional Score Estimation Error) .

For all $k=1,2,\ldots,N$, it holds that:

$$ $\mathbb{E}_{\mathbf{x}\sim p_{kh}}\|\hat{s}_{\theta}(\mathbf{x},kh)-\nabla\log p _{kh}(\mathbf{x})\|^{2}\leq\varepsilon_{\mathrm{score}}^{2}$ $$

where $\hat{s}_{\theta}(\mathbf{x},t)$ is the pretrained score estimator we plug into our algorithm, and $h=T/N$ is the time step.

###### Assumption B.4 (Conditional Score Approximation) .

For all $k=1,2,\ldots,N$, we have an upper bound for the TV-distance between the true conditional distribution $p_{kh}({\mathbf{x}}_{kh}\mid{\mathbf{y}})$ and our DPS-type approximation $\bar{p}_{kh}({\mathbf{x}}_{kh}\mid{\mathbf{y}})$ as follows:

$$ $D_{\mathrm{TV}}\left(\bar{p}_{kh}(\mathbf{x}_{kh}\mid{\mathbf{y}})~{}\|~{}p_{ kh}(\mathbf{x}_{kh}\mid{\mathbf{y}})\right)\leq\varepsilon_{\mathrm{cond}}.$ $$

###### Assumption B.5 (Lipschitz Continuity, Strongly Convexity and Bounded Moment of Conditional Score) .

For all $t$, the conditional score $\nabla\log p_{t}(\mathbf{x}_{t}\mid{\mathbf{y}})$ is $L$-Lipschitz continuous, and the potential function $\log p_{t}(\mathbf{x}_{t}\mid{\mathbf{y}})$ is $m$-strongly convex, which enables us to obtain exponential convergence of Langevin MCMC algorithm. We also assume that

$$ $\mathbb{E}_{\mathbf{x}_{t}\sim p_{t}(\cdot\mid{\mathbf{y}})}\|\nabla\log p_{t} ({\mathbf{y}}\mid\mathbf{x}_{t})\|^{2}\leq U_{\mathrm{cond}}^{2}.$ $$

For $t=0$, $p_{0}(\cdot\mid{\mathbf{y}})$ has bounded second-order moment, i.e.

$$ $m_{2}^{2}:=\mathbb{E}_{{\mathbf{x}}_{0}\sim p_{0}(\cdot\mid{\mathbf{y}})}\|{ \mathbf{x}}_{0}\|^{2}<\infty.$ $$

It also leads to the conclusion that $p_{0}({\mathbf{x}}_{0}\mid{\mathbf{y}})$ has a bounded KL-divergence from the standard Gaussian distribution, i.e.

$$ $D_{\mathrm{KL}}\left(p_{0}({\mathbf{x}}_{0}\mid y)\|\mathcal{N}(0,\bm{I}_{d}) \right)\leq\mathrm{poly}(d).$ $$

Now, we state our main theorem as follows:

###### Theorem B.6 .

Under Assumptions B.3-B.5, once our time step $h<1/L\wedge 1$, we can guarantee that our last-iterate conditional distribution $q_{0}(\cdot\mid{\mathbf{y}})$ derived from DPMC is $(\varepsilon+\varepsilon_{\mathrm{cond}})$-close from ground truth with respect to TV-distance:

$$ $D_{\mathrm{TV}}(q_{0}({\mathbf{x}}_{0}\mid{\mathbf{y}}),p_{0}({\mathbf{x}}_{0} \mid{\mathbf{y}}))\leq\varepsilon+\varepsilon_{\mathrm{cond}}$ $$

after choosing the step size $\eta$ and the number of inner loops (denoted by $K$) of the Langevin MCMC algorithm at all time steps as follows:

$$ $\eta=\frac{m\varepsilon^{2}}{32Ld^{2}},~{}~{}K=\frac{32L^{2}d\cdot\log((\sqrt{ \mathrm{poly}(d)\cdot\exp(-T)}+\varepsilon_{\mathrm{cond}})/\varepsilon~{}\vee ~{}\varepsilon_{\mathrm{inter}}/\varepsilon)}{m^{2}\varepsilon^{2}}.$ $$

Here,

$$ $\varepsilon_{\mathrm{inter}}:=\varepsilon+\varepsilon_{\mathrm{cond}}+C\sqrt{h }\cdot(L\sqrt{dh}+Lm_{2}h)+C\sqrt{h}(\varepsilon_{\mathrm{score}}+U_{\mathrm{ cond}}),$ $$

and $C$ is a universal constant.

###### Proof.

Denote $q_{t}(\cdot)$ as the probability measure derived by the backward process of diffusion model. For the initial step of backward process, we have $q_{T}=\mathcal{N}(0,\bm{I}_{d})$, whose distance from $p_{T}$ shows the convergence of forward process. According to [], the variance-preserving framework leads to exponential convergence of forward process, i.e.

$$ $D_{\mathrm{KL}}\left(p_{T}({\mathbf{x}}_{T}\mid y)\|\mathcal{N}(0,\bm{I}_{d}) \right)\leq D_{\mathrm{KL}}\left(p_{0}({\mathbf{x}}_{0}\mid y)\|\mathcal{N}(0, \bm{I}_{d})\right)\cdot\exp(-T)\leq\mathrm{poly}(d)\cdot\exp(-T).$ $$

By using Pinsker’s Inequality and Assumption B.4, we have:

$$ $D_{\mathrm{TV}}\left(\bar{p}_{T}({\mathbf{x}}_{T}\mid y)\|\mathcal{N}(0,\bm{I} _{d})\right)\leq\sqrt{\mathrm{poly}(d)\cdot\exp(-T)}+\varepsilon_{\mathrm{cond }}.$ $$

Starting from $\tilde{{\mathbf{x}}}_{T}\sim q_{T}=\mathcal{N}(0,\bm{I}_{d})$, we apply Langevin MCMC algorithm with regard to the score function $\nabla\log p_{T}({\mathbf{x}}_{T}\mid{\mathbf{y}})$. By using Lemma B.2, we can make $D_{\mathrm{TV}}(q_{T}({\mathbf{x}}_{T}\mid{\mathbf{y}}),\bar{p}_{T}({\mathbf{x
}}_{T}\mid{\mathbf{y}}))\leq\varepsilon$ by using step size $\eta=\frac{m\varepsilon^{2}}{32Ld^{2}}$, and the number of iterations

$$ $K=\frac{32L^{2}d\log(D_{\mathrm{TV}}(q_{T}\|\bar{p}_{T}({\mathbf{x}}_{T}\mid{ \mathbf{y}}))/\varepsilon)}{m^{2}\varepsilon^{2}}.$ $$

Next, we apply unconditional backward step as well as Langevin MCMC inner loops to make sure $D_{\mathrm{TV}}(q_{kh}({\mathbf{x}}_{kh}\mid{\mathbf{y}}),\bar{p}_{kh}({
\mathbf{x}}_{kh}\mid{\mathbf{y}}))\leq\varepsilon$ holds for $\forall k\in[N]$, including the last iterate $k=0$. We prove it by the method of induction. Assume it holds that $D_{\mathrm{TV}}\left(q_{(k+1)h}({\mathbf{x}}_{(k+1)h}\mid{\mathbf{y}})~{}\|~{}
\bar{p}_{(k+1)h}({\mathbf{x}}_{(k+1)h}\mid{\mathbf{y}})\right)\leq\varepsilon$, then we immediately have

$$ $D_{\mathrm{TV}}\left(q_{(k+1)h}({\mathbf{x}}_{(k+1)h}\mid{\mathbf{y}})~{}\|~{} p_{(k+1)h}({\mathbf{x}}_{(k+1)h}\mid{\mathbf{y}})\right)\leq\varepsilon+ \varepsilon_{\mathrm{cond}}$ $$

by using Assumption B.4 as well as the triangular inequality of TV distance. After applying a backward diffusion step with unconditional score estimator $\hat{s}_{\theta}(\cdot,t)$ plugged in and obtain $\bar{{\mathbf{x}}}_{kh}$. During this backward step, the score estimation error is actually

$$ $\displaystyle\mathbb{E}_{{\mathbf{x}}_{kh}\sim p_{kh}({\mathbf{x}}_{kh}\mid{ \mathbf{y}})}\|\hat{s}_{\theta}({\mathbf{x}}_{kh})-\nabla\log p_{kh}({\mathbf{ x}}_{kh}\mid{\mathbf{y}})\|^{2}$ $\displaystyle~{}~{}\leq\mathbb{E}_{{\mathbf{x}}_{kh}}\|\hat{s}_{\theta}({ \mathbf{x}}_{kh})-\nabla\log p_{kh}({\mathbf{x}}_{kh})-\nabla\log p_{kh}({ \mathbf{y}}\mid{\mathbf{x}}_{kh})\|^{2}$ $\displaystyle~{}~{}\leq 2\mathbb{E}_{{\mathbf{x}}_{kh}}\|\hat{s}_{\theta}({ \mathbf{x}}_{kh})-\nabla\log p_{kh}({\mathbf{x}}_{kh})\|^{2}+2\mathbb{E}_{{ \mathbf{x}}_{kh}}\|\nabla\log p_{kh}({\mathbf{y}}\mid{\mathbf{x}}_{kh})\|^{2}$ $\displaystyle~{}~{}\leq 2(\varepsilon_{\mathrm{score}}^{2}+U_{\mathrm{cond}}^{ 2})\leq 2(\varepsilon_{\mathrm{score}}+U_{\mathrm{cond}})^{2}.$ $$

Therefore, we substitute the score estimation error $\varepsilon_{\mathrm{score}}$ in Theorem 2 of [] with $\varepsilon_{\mathrm{score}}+U_{\mathrm{cond}}$.
We use the results in [] as well as Girsanov Theorem, and conclude that:

$$ $\displaystyle D_{\mathrm{TV}}\left(\mathrm{Law}(\bar{{\mathbf{x}}}_{kh})~{}\|~ {}p_{kh}({\mathbf{x}}_{kh}\mid{\mathbf{y}})\right)-D_{\mathrm{TV}}\left(q_{(k+ 1)h}({\mathbf{x}}_{(k+1)h}\mid{\mathbf{y}})~{}\|~{}\bar{p}_{(k+1)h}({\mathbf{x }}_{(k+1)h}\mid{\mathbf{y}})\right)$ $\displaystyle~{}~{}\leq\underbrace{C\sqrt{h}\cdot(L\sqrt{dh}+Lm_{2}h)}_{\text{ discretization error}}+\underbrace{C\sqrt{h}(\varepsilon_{\mathrm{score}}+U_{ \mathrm{cond}})}_{\text{score estimation error}}$ $$

where $C$ is a universal constant, which leads to

$$ $D_{\mathrm{TV}}\left(\mathrm{Law}(\bar{{\mathbf{x}}}_{kh})~{}\|~{}p_{kh}({ \mathbf{x}}_{kh}\mid{\mathbf{y}})\right)\leq\varepsilon+\varepsilon_{\mathrm{ cond}}+C\sqrt{h}\cdot(L\sqrt{dh}+Lm_{2}h)+C\sqrt{h}(\varepsilon_{\mathrm{score }}+U_{\mathrm{cond}}):=\varepsilon_{\mathrm{inter}}.$ $$

As the initial step of Langevin MCMC inner loops at time $t=kh$, we apply the exponential convergence (Lemma B.2) and show that we can make $D_{\mathrm{TV}}(q_{kh}({\mathbf{x}}_{kh}\mid{\mathbf{y}}),\bar{p}_{kh}({
\mathbf{x}}_{kh}\mid{\mathbf{y}}))\leq\varepsilon$ and complete the induction by letting the step size $\eta=\frac{m\varepsilon^{2}}{32Ld^{2}}$ and the number of iterations

$$ $K=\frac{32L^{2}d\cdot\log(\varepsilon_{\mathrm{inter}}/\varepsilon)}{m^{2} \varepsilon^{2}}.$ $$

To sum up, we can guarantee that $D_{\mathrm{TV}}(q_{0}({\mathbf{x}}_{0}\mid{\mathbf{y}}),\bar{p}_{0}({\mathbf{x
}}_{0}\mid{\mathbf{y}}))\leq\varepsilon$ by choosing the step size $\eta$ and the number of inner loop $K$ of Langevin MCMC algorithm as follows:

$$ $\eta=\frac{m\varepsilon^{2}}{32Ld^{2}},~{}~{}K=\frac{32L^{2}d\cdot\log((\sqrt{ \mathrm{poly}(d)\cdot\exp(-T)}+\varepsilon_{\mathrm{cond}})/\varepsilon~{}\vee ~{}\varepsilon_{\mathrm{inter}}/\varepsilon)}{m^{2}\varepsilon^{2}}$ $$

where

$$ $\varepsilon_{\mathrm{inter}}:=\varepsilon+\varepsilon_{\mathrm{cond}}+C\sqrt{h }\cdot(L\sqrt{dh}+Lm_{2}h)+C\sqrt{h}(\varepsilon_{\mathrm{score}}+U_{\mathrm{ cond}}).$ $$

It finally comes to our conclusion as $D_{\mathrm{TV}}(\bar{p}_{0}({\mathbf{x}}_{0}\mid{\mathbf{y}})~{}\|~{}p_{0}({
\mathbf{x}}_{0}\mid{\mathbf{y}}))\leq\varepsilon+{\mathrm{cond}}$.
∎

###### Assumption B.1 (Strong Convexity and Lipschitz Continuity of Potential Function) .

###### Lemma B.2 (TV-Convergence of Langevin MCMC Algorithm) .

###### Assumption B.3 (Unconditional Score Estimation Error) .

###### Assumption B.4 (Conditional Score Approximation) .

###### Assumption B.5 (Lipschitz Continuity, Strongly Convexity and Bounded Moment of Conditional Score) .

###### Theorem B.6 .

###### Proof.

## Appendix C More Experimental Results

### C.1 Sampling time

In Table 6, we report the time for generating one sample with DPMC default setting, which sets $T=200$, $K=4$ and applies MCMC sampling steps in the middle 60% intermediate distributions. (See section 4.1 for more details.) The time is tested on the FFHQ dataset using a single Nvidia A100 GPU.

**Table 6: Running Times of Different Methods for Generating one Sample on FFHQ-1k Validation Dataset**
| Model | Running Time (Seconds) |
| --- | --- |
| DPMC(ours) | 54.63 |
| FPS [] | 33.07 |
| FPS-SMC (with $M=5$) [] | 57.88 |
| DPS [] | 70.42 |
| Score-SDE [] | 32.93 |
| DDRM [] | 2.034 |
| MCG [] | 73.16 |
| PnP-ADMM [] | 3.595 |
| $\Pi$GDM [] | 33.18 |

### C.2 More Samples for Each Inverse Problems

In Figure 6, 8, 7, 9, 10, we show more samples for each inverse problem. DPMC works well across different tasks and datasets.

Figure: Figure 6: More Results for Super Resolution Task
Refer to caption: extracted/5852151/images/SP.png

Figure: Figure 7: More Results for Motion Deblur Task
Refer to caption: extracted/5852151/images/motion_deblur.png

Figure: Figure 8: More Results for Gaussian Deblur Task
Refer to caption: extracted/5852151/images/gaussian_deblur.png

Figure: Figure 9: More Results for Inpainting (Box) Task
Refer to caption: extracted/5852151/images/inpainting_box.png

Figure: Figure 10: More Results for Inpainting (Random) Task
Refer to caption: extracted/5852151/images/inpainting_random.png

### C.3 More baselines

In this section we compare DPMC with more strong baselines. More precisely, we compare our model with $\Pi$GDM [], DDNM[], RED-Diff [], DiffPir [] and CCDF []. We conducted experiments for RED-diff, $\Pi$GDM and DDNM. And we obtained the reported results of DiffPir and CCDF from their original papers. The quantitative results are shown in Table 7 and Table 8.

We have observed that the results of RED-diff and $\Pi$GDM are significantly influenced by whether noise is added to the degradation process. We observed that if noise is added, the optimized samples might exhibit noise-like artifacts and sometimes completely fail. Therefore, we include both the noisy version (similar to our setting) and the noiseless version (denoted as $\sigma=0$ ) here. Note that due to setting differences, the noiseless results are listed here for reference only.

From the quantitative results, we can see that our DPMC demonstrates comparable or better performance compared to strong baselines in the noisy measurement case ($\sigma=0.05$), particularly in terms of FID scores. Our results remain robust even when compared with the noiseless results of RED-diff and $\Pi$GDM.

We also include qualitative comparision in Figure 11. The main advantage of our DPMC is that our samples contain more vivid details and can accept noisy settings ($\sigma>0$). The samples from RED-diff and $\Pi$GDM might contain noise-like artifacts in noisy settings. The results of FPS-SMC, DDNM, $\Pi$GDM ($\sigma=0.0$) and RED-diff ($\sigma=0.0$) don’t contain noise-like textures; however, these samples tend to be blurrier than ours.

**Table 7: Quantitative results of various linear inverse problems on FFHQ $256\times 256$-1k validation set with more strong baselines.**
| Methods | Super Resolution | Inpainting (box) | Gaussian Deblur | Inpainting (random) | Motion Deblur |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS |
| DPMC (Ours) | 21.93 | 0.212 | 19.59 | 0.160 | 21.34 | 0.210 | 21.26 | 0.205 | 20.73 | 0.213 |
| DDNM[] | 26.64 | 0.214 | 25.97 | 0.150 | 28.69 | 0.212 | 28.71 | 0.201 | - | - |
| RED-Diff [] | 89.13 | 0.435 | - | - | 37.35 | 0.255 | - | - | - | - |
| $\Pi$GDM [] | 29.59 | 0.214 | - | - | 431.83 | 0.887 | - | - | - | - |
| CCDF [] | 60.90 | - | 49.77 | - | - | - | - | - | - |  |
| DiffPir [] | - | 0.260 | - | 0.236 | - | - | - | - | - | 0.255 |
| RED-Diff $(\sigma=0.0)$ [] | 39.68 | 0.185 | - | - | 30.54 | 0.161 | - | - | - | - |
| $\Pi$GDM $(\sigma=0.0)$ [] | 39.61 | 0.207 | - | - | 34.52 | 0.140 | - | - | - | - |

**Table 8: Quantitative results of various linear inverse problems on ImageNet $256\times 256$-1k validation set with more strong baselines.**
|  | Super Resolution | Inpainting (box) | Gaussian Deblur | Inpainting (random) | Motion Deblur |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Methods | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS | FID | LPIPS |
| DPMC (Ours) | 31.74 | 0.307 | 30.55 | 0.221 | 33.62 | 0.318 | 30.25 | 0.292 | 30.88 | 0.303 |
| RED-Diff [] | 82.62 | 0.471 | - | - | 39.11 | 0.319 | - | - | - | - |
| $\Pi$GDM [] | 43.55 | 0.343 | - | - | 371.33 | 0.813 | - | - | - | - |
| DiffPir [] | - | 0.371 | - | 0.355 | - | - | - | - | - | 0.366 |
| RED-Diff $(\sigma=0.0)$ [] | 45.17 | 0.304 | - | - | 32.29 | 0.232 | - | - | - | - |
| $\Pi$GDM $(\sigma=0.0)$ [] | 50.21 | 0.342 | - | - | 32.99 | 0.200 | - | - | - | - |

Figure: Figure 11: Qualitative comparison of different methods on ImageNet $256\times 256$-1k dataset.
Refer to caption: extracted/5852151/rebuttal_images/super_res_all.png

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: See also weaknesses above.  Explicit questions:

1) It is not clear whether this method performs better for high corruption or for low corruption. Could the authors provide a plot where the horizontal axis is the corruption level and the vertical axis the performance and show how their method compares with DPS across these different regimes?

2) Could the authors provide comparisons of what happens if both DPS and the proposed method have big enough computational budget? If NFEs are not a concern, does the method still outperform DPS?

3) Could the authors clarify whether the theoretical results would also hold for DPS under similar assumptions and learning rate tuning?

4) Could the authors compare with other methods that perform similar inner loop optimizations (such as Resample and Beyond First Order Tweedie)?
soundness: 3
strengths: 1) The empirical results are quite strong. The authors demonstrate that their method outperforms natural baselines across different inverse problems such as mask inpainting, box inpainting, super-resolution, deblurring, and phase retrieval.

2) The algorithmic modification is simple, which might allow for wide adoption from the community.

3) The research direction of developing algorithms for solving inverse problems with diffusion models is relevant and interesting.
summary: The authors propose a new algorithm for solving inverse problems with pre-trained diffusion models. The proposed algorithm is a variation of Diffusion Posterior Sampling (DPS) where instead of taking one gradient step to match the measurements, the authors propose running Langevin dynamics to sample from the approximate conditional distribution. The authors demonstrate the empirical performance of their algorithm for linear and non-linear inverse problems.
weaknesses: 1) I believe that the presentation of the paper could be improved. There are several examples. In Line 43, I believe that measure should be changed to measurement. In Lines 46-47, I believe that the authors want to say that the posterior is intractable (it is always defined, but it is intractable to sample from it or write down an explicit formula for its density). I also believe that it is a little confusing to what former and later refer to in lines 158-159. The theoretical result is not clearly stated, there is no mention in the main paper of the assumptions and on what epsilon depends to and how. 

2) Apart from not being clearly stated, the theoretical result is very weak (it has very strong assumptions such as strong convexity) and I do not understand the value it offers to the paper. My guess is that under all these strong assumptions on the data distribution, DPS would also converge to the true distribution (given properly tuned learning rates and a small enough discretization step). 

3) I believe that the authors overcomplicate the presentation of their algorithm and as far as I understand what's happening is an inner loop of Langevin Dynamics that is getting interleaved between the denoising steps. Langevin Dynamics is the simplest algorithm in the MCMC family. Having an inner loop that enforces measurement consistency in the diffusion sampling has been proposed previously in the literature. For example, the papers Beyond First Order Tweedie and Resample also seem to solve an inner optimization problem between each denoising step. 

4) The fact that the authors only use the proposed method for a subset of the steps makes the evaluation harder. 

5) I believe that the paper would benefit from some latent diffusion experiments since there are powerful latent diffusion models available that can/should be leveraged to solve inverse problems.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see my comments above.
soundness: 3
strengths: 1. The theoretical bound presented in the appendix is insightful.
2. The paper is well-structured and easy to follow (altough there are some problems with writing as I explained in the weakness part).
3. The experimental results demonstrate the superiority of the proposed method over state-of-the-art alternatives, both quantitatively and qualitatively.
4. The experiments are thorough and consider relevant datasets and settings.
summary: The paper addresses an interesting topic: solving inverse problems using diffusion models. The authors argue that the performance of the popular DPS algorithm is limited by inaccuracies in the posterior approximation, particularly at high noise levels. To address this, they propose an MCMC algorithm to sample more accurately from the posterior distribution.
weaknesses: 1. **Marginal Contribution**: The contribution of the paper appears limited, as it primarily applies the established MCMC algorithm within diffusion models for solving inverse problems. Specifically, the novelty seems to lie in the use of Equation (10) atop standard diffusion models, where Equation (10) serves as the MCMC update step. 

2. **Method Clarity**: The authors should dedicate more space to clearly explaining their method in Section 3, especially detailing the deployment of MCMC in posterior sampling. Currently, the description on page 5 is divided into two parts (Proposal Stage and Exploration Stage), which makes the algorithm challenging to follow. Mathematical expressions should be used alongside explanations to clarify each step. Specific areas for improvement include:
   
   a) **Proposal Stage Description**: The description of the proposal stage is vague. For instance, the phrase "..we first denoise them to \( t-1 \) following..." is unclear. Does \( t-1 \) refer to \( x_{t-1} \)?
   
   b) **Standard Diffusion Step**: When the authors mention the "standard diffusion step," are they referring to using a pre-trained diffusion model to estimate \( x_t \)?
   
   c) **Understanding of the Sampling Process**: My understanding is that the method first uses an unconditional pre-trained diffusion model to estimate \( x_{t-1} \) and then modifies \( x_{t-1} \) to sample from the posterior. If this is correct, the writing should be adjusted to make this process explicitly clear in the paper.

3. **Terminology on Line 248**: On line 248, the authors state that "..The samples \( \tilde{x}_{t-1} \) might not fully adhere to the target distribution \( \tilde{p}_{t-1}(x_{t-1}|y) \)." Given that samples drawn from the unconditional distribution would indeed differ significantly from those conditioned on \( y \), the use of "might not fully adhere" is imprecise. The terminology should be refined here to convey a more scientifically rigorous statement.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the weaknesses above.
soundness: 2
strengths: 1. This paper aims to reduce the bias in diffusion posterior sampling by incorporating MCMC sampling into the intermediate steps of the reverse diffusion process.

2. The experimental results are promising across a range of inverse tasks, including super-resolution (4x), random inpainting, motion deblurring, Gaussian deblurring, and box inpainting.
summary: Diffusion Posterior Sampling (DPS) is a versatile algorithm for solving inverse problems with pre-trained diffusion models. However, DPS suffers from bias due to the Dirac delta approximation of the noise-conditional posterior $p(x_0|x_t) \approx \delta(x_0 - E[x_0|x_t]),$ which can lead to large errors even for a simple mixture of Gaussians. To address this, the authors propose Diffusion Posterior MCMC (DPMC), which reduces bias by defining a series of intermediate distributions inspired by the approximated conditional distributions in DPS.
In its simplest form, DPMC can be viewed as performing multiple DPS steps with added noise. Through annealed MCMC sampling, the intermediate latents closely follow the carefully constructed intermediate distributions before moving to the next noise level. Extensive experiments on FFHQ 256x256 and ImageNet 256x256 demonstrate DPMC’s effectiveness on five inverse problems: super-resolution, box inpainting, random inpainting, Gaussian deblurring, and motion deblurring.
weaknesses: 1. In the Line 6 of Algorithm 1, the update follows Eq 10. The second term in Eq 10 requires $\nabla \log p_t(x_t)$ and the gradient of the measurement error. How is the score function $\nabla \log p_t(x_t)$ computed using the pre-trained neural network? Since the NN expects the state and time as inputs, and time changes after every update, how is the time step (or equivalently the noise level) chosen while computing the score?

2. The proposed DPMC sampler is a stochastic equivalent of RB-Modulation (Algorithm 1). Substituting the gradient of the proposal distribution in Eq. (10):
$$x_{t-1}^{k+1} = x_{t-1}^k + \eta_{t-1} \nabla_{x_{t-1}^k} \log p_{t-1}(x_{t-1}^k) -  \rho  \eta_{t-1} \nabla_{x_{t-1}^k} || y - A \hat{x}(x_{t-1}^k) ||^2 + \sqrt{2 \eta_{t-1}} w.$$
Without the noise term $w$, DPMC is equivalent to RB-Modulation (Algorithm 1, https://arxiv.org/pdf/2405.17401 ). While $w$ helps in exploration, it slows down the inference process due to large mixing time. On the other hand, RB-Modulation aims for the maximum likelihood estimate and converges much faster. Both the algorithms make certain approximations to reduce inference time but the core idea remains the same. The authors are encouraged to clarify the distinction between these two algorithms and highlight the algorithmic novelty of DPMC over prior work RB-Modulation. 

3. Contributions 2 and 3 could be combined, as they convey a similar message.

4. The first equality in Eq. (2) is incorrect. It should be properly factorized to accurately represent the probabilities.

5. Inconsistent notation used in Line 136 for $p_t(x_t)$. With Bayes rule, it should just be a marginal of $p$.

6. Typo: Missing log in Line 269.

7. The primary motivation of this paper is that $p(x_0|x_t)$ may be multi-modal; hence, replacing  $x_0 \sim p(x_0|x_t)$ with its conditional expectation $\hat{x}_0$ at higher noise levels could produce misleading results. However, the authors still use $\hat{x}_0$ in the first 30\% of the reverse process, which corresponds to high noise levels – contradicting the main motivation.

8. What happens if these initial 30\% steps mislead the generation process? Does the proposed sampling mechanism rectify this issue in subsequent steps? If not, the technical benefits of these MCMC steps must be clearly explained, rather than better storytelling via MCMC sampling.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: See the above section.
soundness: 3
strengths: **Originality**: While the proposed idea of using annealed MCMC has been utilized for unconditional diffusion sampling in several prior works [Song et al.], its application to solving inverse problems combined with existing posterior approximation methods is novel.

**Quality and Clarity**: The paper is written clearly, and the choice of baselines is adequate. The empirical results in Tables 1 and 2 are comprehensive regarding the selection of datasets and inverse problems. The ablation studies take into account the impact of different design choices in the method.

**Significance**: Though the computational overhead of introducing MCMC corrections as an inner loop in iterative diffusion sampling seems much, the method has significance in terms of yielding better quality generations at similar sampling budgets as DPS, which could be useful in scenarios when a large sampling budget is available.
summary: The authors present a method based on annealed MCMC for solving inverse problems using diffusion models. The motivation for this work stems from the fact that existing methods like DPS [Chung et al.] rely on overly simplistic approximations of the diffusion posterior $p(x_0|x_t)$ throughout the diffusion sampling procedure when solving inverse problems. In this work, alternatively, the authors propose to sample from intermediate distributions by introducing additional MCMC correction steps. Empirical results are presented on different linear and non-linear inverse problems on the FFHQ and ImageNet datasets at 256x256 resolution.
weaknesses: **Requested Changes/Questions**

1. **On the choice of the intermediate distribution in Eq. 9**: The authors use the posterior approximation from DPS to define the intermediate distributions as stated in Eq. 9. In principle, the approximation from Pi-GDM can also be used to formulate the same. What is the rationale behind sticking with the approximation in DPS since Pi-GDM can perform better with fewer NFE? It would be great if the authors could clarify this choice in the main text.

2. **On the choice of the proposal**: The authors propose to initialize MCMC sampling by generating a sample from a single denoising step from the diffusion model without guidance. Intuitively, it looks like initializing the MCMC chain with denoising + guidance could yield better samples since the latter might be closer to the target intermediate distribution. Is this choice motivated by reducing the computational overhead incurred by initializing using denoising + guidance, as the guidance term requires computing a Jacobian-vector product? It would be great if the authors could clarify more on this or share some experimental evidence that justifies the choice of their current proposal over the one mentioned here.

3. In Eq. 11, the authors suggest to use the  l2-norm over the squared l2-norm. How critical is this choice? Moreover, it is unclear if the same choice is also used for the vanilla DPS baseline for a fair comparison since the authors specify using a squared l2 norm in Eq. 7. It would be nice if the authors could present a small ablation for one inverse problem on the dataset of their choice.

4. **Computational Overhead**: By introducing an inner MCMC loop, the computational overhead of this method increases significantly, especially since each MCMC correction step requires a Jacobian-vector product through the score network. The authors propose to counteract this by reducing the number of intermediate distributions to around 200, which results in a total number of NFE to around 700. Though this is smaller than DPS, it is quite a lot compared to some recent work [1,2] in solving inverse problems in diffusion. While I won't stress on comparing DPMC to these methods, it would be nice to discuss this aspect of inverse problems in the paper (maybe in the main text or the appendix) as a possible extension of this work. Moreover, can the authors extend the results in Table 4b to say T=25 or 50 to assess the impact on performance at lower sampling budgets? Lastly, I would expect K to decrease as the signal-to-noise ratio increases in diffusion sampling since lower MCMC corrections should be required as sampling progresses. Did the authors experiment with a schedule on K, or is the choice of a fixed K in the paper primarily motivated by simplicity constraints? Perhaps the last section of the paper can benefit from this discussion.

[1] Fast Samplers for Inverse Problems in Iterative Refinement Models, Pandey et al.

[2] Accelerating Diffusion Models for Inverse Problems through Shortcut Sampling, Liu et al.

**Minor Comments/Corrections**

1. In some places, like Line 354, method names and citations are hard to decipher. I think readability can be improved by using a different color for citation links.

2. Tables 1 and 2 can benefit from indicating that lower is better for lpips and fid.

3. Line 369, the link to the appendix is broken.

4. There are a lot of missing entries in Table 1. Can the authors clarify in the caption if the results were unavailable for these cases or if the methods were unstable?