# Solving Inverse Problems with Latent Diffusion Models via Hard Data
Consistency

Bowen Song1 Corresponding authors {bowenbw<EMAIL_ADDRESS>Soo Min Kwon1∗
Zecheng Zhang2 Xinyu Hu3 Qing Qu1 Liyue Shen1
1 University of Michigan  2 Kumo.AI  3 Microsoft

###### Abstract

Latent diffusion models have been demonstrated to generate high-quality
images, while offering efficiency in model training compared to diffusion
models operating in the pixel space. However, incorporating latent diffusion
models to solve inverse problems remains a challenging problem due to the
nonlinearity of the encoder and decoder. To address these issues, we propose
ReSample, an algorithm that can solve general inverse problems with pre-
trained latent diffusion models. Our algorithm incorporates data consistency
by solving an optimization problem during the reverse sampling process, a
concept that we term as hard data consistency. Upon solving this optimization
problem, we propose a novel resampling scheme to map the measurement-
consistent sample back onto the noisy data manifold and theoretically
demonstrate its benefits. Lastly, we apply our algorithm to solve a wide range
of linear and nonlinear inverse problems in both natural and medical images,
demonstrating that our approach outperforms existing state-of-the-art
approaches, including those based on pixel-space diffusion models. Our code is
available at https://github.com/soominkwon/resample.

##  1 Introduction

Inverse problems arise from a wide range of applications across many domains,
including computational imaging [1, 2], medical imaging [3, 4], and remote
sensing [5, 6], to name a few. When solving these inverse problems, the goal
is to reconstruct an unknown signal $\bm{x}_{*}\in\mathbb{R}^{n}$ given
observed measurements $\bm{y}\in\mathbb{R}^{m}$ of the form

$\bm{y}=\mathcal{A}(\bm{x}_{*})+\bm{\bm{\eta}},$

where $\mathcal{A}(\cdot):\mathbb{R}^{n}\to\mathbb{R}^{m}$ denotes some
forward measurement operator (can be linear or nonlinear) and
$\bm{\eta}\in\mathbb{R}^{m}$ is additive noise. Usually, we are interested in
the case when $m<n$, which follows many real-world scenarios. When $m<n$, the
problem is ill-posed and some kind of regularizer (or prior) is necessary to
obtain a meaningful solution.

In the literature, the traditional approach of using hand-crafted priors (e.g.
sparsity) is slowly being replaced by rich, learned priors such as deep
generative models. Recently, there has been a lot of interests in using
diffusion models as structural priors due to their state-of-the-art
performance in image generation [7, 8, 9]. Compared to GANs, diffusion models
are generally easier and more stable to train as they do not rely on an
adversarial training scheme, making them a generative prior that is more
readily accessible [7]. The most common approach for using diffusion models as
priors is to resort to posterior sampling, which has been extensively explored
in the literature [10, 11, 12, 13, 14, 15, 16, 17]. However, despite their
remarkable success, these techniques exhibit several limitations. The primary
challenge is that the majority of existing works that employ diffusion models
as priors train these models directly in the pixel space, which requires
substantial computational resources and a large volume of training data. In
real-world applications such as computed tomography (CT) and magnetic
resonance imaging (MRI) reconstruction, where images are inherently either 3D
or even 4D [18], training diffusion models directly in the pixel space is
often infeasible.

Latent diffusion models (LDMs), which embed data in order to operate in a
lower-dimensional space, present a potential solution to this challenge, along
with considerable improvements in computational efficiency [19, 20] by
training diffusion models in a compressed latent space. They can also provide
a great amount of flexibility, as they can enable one to transfer and
generalize these models to different domains by fine-tuning on small amounts
of training data [21]. Nevertheless, using LDMs to solve inverse problems
poses a significant challenge. The main difficulty arises from the inherent
nonlinearity and nonconvexity of the decoder, making it challenging to
directly apply existing solvers designed for pixel space. To address this
issue, a concurrent work by Rout et al. [22] recently introduced a posterior
sampling algorithm operating in the latent space (PSLD), designed to solve
_linear_ inverse problems with provable guarantees. However, we observe that
PSLD may reconstruct images with artifacts in the presence of measurement
noise. Therefore, developing an efficient algorithm capable of addressing
these challenges remains an open research question.

Figure 1: Example reconstructions of our algorithm (ReSample) on two noisy
inverse problems, nonlinear deblurring and CT reconstruction, on natural and
medical images, respectively.

In this work, we introduce a novel algorithm named ReSample, which effectively
employs LDMs as priors for solving inverse problems. Our algorithm can be
viewed as a two-stage process that incorporates data consistency by (1)
solving a hard-constrained optimization problem, ensuring we obtain the
correct latent variable that is consistent with the observed measurements, and
(2) employing a carefully designed resampling scheme to map the measurement-
consistent sample back onto the correct noisy data manifold. As a result, we
show that our algorithm can achieve state-of-the-art performance on various
inverse problem tasks and different datasets, compared to existing algorithms.
Notably, owing to using the latent diffusion models as generative priors, our
algorithm achieves improved computational efficiency with a reduction in
memory complexity. Below, we highlight some of our key contributions.

  * •

We propose a novel algorithm that enables us to leverage latent diffusion
models for solving general inverse problems (linear and nonlinear) through
hard data consistency.

  * •

Particularly, we carefully design a stochastic resampling scheme that can
reliably map the measurement-consistent samples back onto the noisy data
manifold to continue the reverse sampling process. We provide a theoretical
analysis to further demonstrate the superiority of the proposed stochastic
resampling technique.

  * •

With extensive experiments on multiple tasks and various datasets,
encompassing both natural and medical images, our proposed algorithm achieves
state-of-the-art performance on a variety of linear and nonlinear inverse
problems.

##  2 Background

###  2.1 Diffusion Models

#### Denoising Diffusion Probabilistic Models.

We first briefly review the basic fundamentals of diffusion models, namely the
denoising diffusion probabilistic model (DDPM) formulation [23]. Let
$\bm{x}_{0}\sim p_{\text{data}}(\bm{x})$ denote samples from the data
distribution. Diffusion models start by progressively perturbing data to noise
via Gaussian kernels, which can be written as the variance-preserving
stochastic differential equation (VP-SDE) [24] of the form

$\displaystyle d\bm{x}=-\frac{\beta_{t}}{2}\bm{x}dt+\sqrt{\beta_{t}}d\bm{w},$
(1)

where $\beta_{t}\in(0,1)$ is the noise schedule that is a monotonically
increasing sequence of $t$ and $\bm{w}$ is the standard Wiener process. This
is generally defined such that we obtain the data distribution when $t=0$ and
obtain a Gaussian distribution when $t=T$, i.e.
$\bm{x}_{T}\sim\mathcal{N}(\bm{0},\bm{I})$. The objective of diffusion models
is to learn the corresponding reverse SDE of Equation (1), which is of the
form

$\displaystyle
d\bm{x}=\left[-\frac{\beta_{t}}{2}\bm{x}-\beta_{t}\nabla_{\bm{x}_{t}}\log
p(\bm{x}_{t})\right]dt+\sqrt{\beta_{t}}d\bar{\bm{w}},$ (2)

where $d\bar{\bm{w}}$ is the standard Wiener process running backward in time
and $\nabla_{\bm{x}_{t}}\log p(\bm{x}_{t})$ is the (Stein) score function. In
practice, we approximate the score function using a neural network
$\bm{s}_{\theta}$ parameterized by $\theta$, which can be trained via
denoising score matching [25]:

$\displaystyle\hat{\theta}=\underset{\theta}{\operatorname*{arg\,min}}\,\mathbb{E}\left[\|\bm{s}_{\theta}(\bm{x}_{t},t)-\nabla_{\bm{x}_{t}}\log
p(\bm{x}_{t}|\bm{x}_{0})\|_{2}^{2}\right],$ (3)

where $t$ is uniformly sampled from $[0,T]$ and the expectation is taken over
$t$, $\bm{x}_{t}\sim p(\bm{x}_{t}|\bm{x}_{0})$, and $\bm{x}_{0}\sim
p_{\text{data}}(\bm{x})$. Once we have access to the parameterized score
function $\bm{s}_{\theta}$, we can use it to approximate the reverse-time SDE
and simulate it using numerical solvers (e.g. Euler-Maruyama).

Figure 2: Overview of our ReSample algorithm during the reverse sampling
process conditioned on the data constraints from measurement. The entire
sampling process is conducted in the latent space upon passing the sample
through the encoder. The proposed algorithm performs hard data consistency at
some time steps $t$ via a skipped-step mechanism.

#### Denoising Diffusion Implicit Models.

As the DDPM formulation is known to have a slow sampling process, Song et al.
[26] proposed denoising diffusion implicit models (DDIMs) that defines the
diffusion process as a non-Markovian process to remedy this [23, 26, 9, 27].
This enables a faster sampling process with the sampling steps given by

$\displaystyle\bm{x}_{t-1}=\sqrt{\bar{\alpha}_{t-1}}\hat{\bm{x}}_{0}(\bm{x}_{t})+\sqrt{1-\bar{\alpha}_{t-1}-\eta\delta_{t}^{2}}\bm{s}_{\theta}(\bm{x}_{t},t)+\eta\delta_{t}\bm{\epsilon},\quad
t=T,\ldots,0,$ (4)

where $\alpha_{t}=1-\beta_{t}$, $\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$,
$\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$, $\eta$ is the temperature
parameter, $\delta_{t}$ controls the stochasticity of the update step, and
$\hat{\bm{x}}_{0}(\bm{x}_{t})$ denotes the predicted $\bm{x}_{0}$ from
$\bm{x}_{t}$ which takes the form

$\displaystyle\hat{\bm{x}}_{0}(\bm{x}_{t})=\frac{1}{\sqrt{\bar{\alpha}_{t}}}(\bm{x}_{t}+\sqrt{1-\bar{\alpha}_{t}}\bm{s}_{\theta}(\bm{x}_{t},t)),$
(5)

which is an application of Tweedie’s formula. Here, $\bm{s}_{\theta}$ is
usually trained using the epsilon-matching score objective [26]. We use DDIM
as the backbone of our algorithm and show how we can leverage these update
steps for solving inverse problems.

###  2.2 Diffusion Models for Solving Inverse Problems

#### Solving Inverse Problems with Diffusion Models.

Given measurements $\bm{y}\in\mathbb{R}^{m}$ from some forward measurement
operator $\mathcal{A}(\cdot)$, we can use diffusion models to solve inverse
problems by replacing the score function in Equation (2) with the conditional
score function $\nabla_{\bm{x}_{t}}\log p(\bm{x}_{t}|\bm{y})$. Then by Bayes
rule, notice that we can write the conditional score as

$\displaystyle\nabla_{\bm{x}_{t}}\log
p(\bm{x}_{t}|\bm{y})=\nabla_{\bm{x}_{t}}\log
p(\bm{x}_{t})+\nabla_{\bm{x}_{t}}\log p(\bm{y}|\bm{x}_{t}).$

This results in the reverse SDE of the form

$\displaystyle
d\bm{x}=\left[-\frac{\beta_{t}}{2}\bm{x}-\beta_{t}(\nabla_{\bm{x}_{t}}\log
p(\bm{x}_{t})+\nabla_{\bm{x}_{t}}\log
p(\bm{y}|\bm{x}_{t}))\right]dt+\sqrt{\beta_{t}}d\bar{\bm{w}}.$

In the literature, solving this reverse SDE is referred as _posterior
sampling_. However, the issue with posterior sampling is that there does not
exist an analytical formulation for the likelihood term
$\nabla_{\bm{x}_{t}}\log p(\bm{y}|\bm{x}_{t})$. To resolve this, there exists
two lines of work: (1) to resort to alternating projections onto the
measurement subspace to avoid using the likelihood directly [12, 13, 28] and
(2) to estimate the likelihood under some mild assumptions [11, 14]. For
example, Chung et al. [11] proposed diffusion posterior sampling (DPS) that
uses a Laplacian approximation of the likelihood, which results in the
discrete update steps

$\displaystyle\bm{x}_{t-1}$
$\displaystyle=\sqrt{\bar{\alpha}_{t-1}}\hat{\bm{x}}_{0}(\bm{x}_{t})+\sqrt{1-\bar{\alpha}_{t-1}-\eta\delta_{t}^{2}}\bm{s}_{\theta}(\bm{x}_{t},t)+\eta\delta_{t}\bm{\epsilon}$
(6) $\displaystyle\bm{x}_{t-1}$
$\displaystyle=\bm{x}^{\prime}_{t-1}-\zeta\nabla_{\bm{x}_{t}}\|\bm{y}-\mathcal{A}(\hat{\bm{x}}_{0}(\bm{x}_{t}))\|_{2}^{2},$
(7)

where $\zeta\in\mathbb{R}$ can be viewed as a tunable step-size. However, as
previously mentioned, these techniques have limited applicability for real-
world problems as they are all built on the pixel space.

#### Solving Inverse Problems with Latent Diffusion Models.

The limited applicability of pixel-based diffusion models can be tackled by
alternatively utilizing more efficient LDMs as generative priors. The setup
for LDMs is the following: given an image $\bm{x}\in\mathbb{R}^{n}$, we have
an encoder $\mathcal{E}:\mathbb{R}^{n}\to\mathbb{R}^{k}$ and a decoder
$\mathcal{D}:\mathbb{R}^{k}\to\mathbb{R}^{n}$ where $k\ll n$. Let
$\bm{z}=\mathcal{E}(\bm{x})\in\mathbb{R}^{k}$ denote the embedded samples in
the latent space. One way of incorporating LDMs to solve inverse problems
would be to replace the update steps in Equation (6) with

$\displaystyle\bm{z}^{\prime}_{t-1}$
$\displaystyle=\sqrt{\bar{\alpha}_{t-1}}\hat{\bm{z}}_{0}(\bm{z}_{t})+\sqrt{1-\bar{\alpha}_{t-1}-\eta\delta_{t}^{2}}\bm{s}_{\theta}(\bm{z}_{t},t)+\eta\delta_{t}\bm{\epsilon},$
(8) $\displaystyle\bm{z}_{t-1}$
$\displaystyle=\bm{z}^{\prime}_{t-1}-\zeta\nabla_{\bm{z}_{t}}\|\bm{y}-\mathcal{A}(\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t})))\|_{2}^{2},$
(9)

After incorporating LDMs, this can be viewed as a non-linear inverse problem
due to the non-linearity of the decoder $\mathcal{D}(\cdot)$. As this builds
upon the idea behind DPS, we refer to this algorithm as _Latent-DPS_. While
this formulation seems to work, we empirically observe that Latent-DPS often
produces reconstructions that are often noisy or blurry and inconsistent with
the measurements. We conjecture that since the forward operator involving the
decoder is highly nonconvex, the gradient update may lead towards a local
minimum. We provide more insights in Appendix (Section D).

##  3 ReSample: Inverse Problems using Latent Diffusion Models

###  3.1 Proposed Method

#### Hard Data Consistency.

Similar to Latent-DPS, our algorithm involves incorporating data consistency
into the reverse sampling process of LDMs. However, instead of a gradient-like
update as shown in Equation (9), we propose to solve an optimization problem
on _some time steps $t$_ of the reverse process:

$\displaystyle\hat{\bm{z}}_{0}(\bm{y})\in\underset{\bm{z}}{\operatorname*{arg\,min}}\,\frac{1}{2}\|\bm{y}-\mathcal{A}(\mathcal{D}(\bm{z}))\|^{2}_{2},$
(10)

where we denote $\hat{\bm{z}}_{0}(\bm{y})$ as the sample consistent with the
measurements $\bm{y}\in\mathbb{R}^{m}$. This optimization problem has been
previously explored in other works that use GANs for solving inverse problems,
and can be efficiently solved using iterative solvers such as gradient descent
[29, 30, 31, 32]. However, it is well known that solving this problem starting
from a random initial point may lead to unfavorable local minima [29]. To
address this, we solve Equation (10) starting from an initial point
$\hat{\bm{z}}_{0}(\bm{z}_{t+1})$, where $\hat{\bm{z}}_{0}(\bm{z}_{t+1})$ is
the estimate of ground-truth latent vector at time $0$ based on the sample at
time $t+1$. The intuition behind this initialization is that we want to start
the optimization process within local proximity of the global solution of
Equation (10), to prevent resulting in a local minimum. We term this overall
concept as _hard data consistency_ , as we strictly enforce the measurements
via optimization, rather than a “soft” approach through gradient update like
Latent-DPS. To obtain $\hat{\bm{z}}_{0}(\bm{z}_{t+1})$, we use Tweedie’s
formula [33] that gives us an approximation of the posterior mean which takes
the following formula:

$\displaystyle\hat{\bm{z}}_{0}(\bm{z}_{t})=\mathbb{E}[\bm{z}_{0}|\bm{z}_{t}]=\frac{1}{\sqrt{\bar{\alpha}_{t}}}(\bm{z}_{t}+(1-\bar{\alpha}_{t})\nabla\log
p(\bm{z}_{t})).$ (11)

However, we would like to note that performing hard data consistency on every
reverse sampling iteration $t$ may be very costly. To address this, we first
observe that as we approach $t=T$, the estimated
$\hat{\bm{z}}_{0}(\bm{z}_{t})$ can deviate significantly from the ground
truth. In this regime, we find that hard data consistency provides only
marginal benefits. Additionally, in the literature, existing works point out
the existence of a three-stage phenomenon [34], where they demonstrate that
data consistency is primarily beneficial for the semantic and refinement
stages (the latter two stages when $t$ is closer to $0$) of the sampling
process. Following this reasoning, we divide $T$ into three sub-intervals and
only apply the optimization in the latter two intervals. This approach
provides both computational efficiency and accurate estimates of
$\hat{\bm{z}}_{0}(\bm{z}_{t})$.

Furthermore, even during these two intervals, we observe that we do not need
to solve the optimization problem on every iteration $t$. Because of the
continuity of the sampling process, after each data-consistency optimization
step, the samples in the following steps can still remain similar semantic or
structural information to some extent. Thus, we “reinforce” the data
consistency constraint during the sampling process via a skipped-step
mechanism. Empirically, we see that it is sufficient to perform this on every
$10$ (or so) iterations of $t$. One can think of hard data consistency as
guiding the sampling process towards the ground truth signal $\bm{x}^{*}$ (or
respectively $\bm{z}^{*}$) such that it is consistent with the given
measurements. Lastly, in the presence of measurement noise, minimizing
Equation (10) to zero loss can lead to overfitting the noise. To remedy this,
we perform early stopping, where we only minimize up to a threshold $\tau$
based on the noise level. We will discuss the details of the optimization
process in the Appendix. We also observe that an additional Latent-DPS step
after unconditional sampling can (sometimes) marginally increase the overall
performance. We perform an ablation study on the performance of including
Latent-DPS in the Appendix.

Algorithm 1 ReSample: Solving Inverse Problems with Latent Diffusion Models

Measurements $\bm{y}$, $\mathcal{A}(\cdot)$, Encoder $\mathcal{E}(\cdot)$,
Decoder $\mathcal{D}(\cdot)$, Score function $\bm{s}_{\theta}(\cdot,t)$,
Pretrained LDM Parameters $\beta_{t}$, $\bar{\alpha}_{t}$, $\eta$, $\delta$,
Hyperparameter $\gamma$ to control $\sigma^{2}_{t}$, Time steps to perform
resample $C$

$\bm{z}_{T}\sim\mathcal{N}(\bm{0},\bm{I})$ $\triangleright$ Initial noise
vector

for $t=T-1,\ldots,0$ do

$\boldsymbol{\epsilon}_{1}\sim\mathcal{N}(\bm{0},\bm{I})$

$\hat{\bm{\epsilon}}_{t+1}=\bm{s}_{\theta}(\bm{z}_{t+1},{t+1})$
$\triangleright$ Compute the score

$\hat{\bm{z}}_{0}(\bm{z}_{t+1})=\frac{1}{\sqrt{\bar{\alpha}_{t+1}}}(\bm{z}_{t+1}-\sqrt{1-\bar{\alpha}_{t+1}}\hat{\bm{\epsilon}}_{t+1})$
$\triangleright$ Predict $\hat{\bm{z}}_{0}$ using Tweedie’s formula

$\bm{z}^{\prime}_{t}=\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}(\bm{z}_{t+1})+\sqrt{1-\bar{\alpha}_{t}-\eta\delta^{2}}\hat{\bm{\epsilon}}_{t+1}+\eta\delta\boldsymbol{\epsilon}_{1}$
$\triangleright$ Unconditional DDIM step

if $t\in C$ then $\triangleright$ ReSample time step

$\hat{\bm{z}}_{0}(\bm{y})\in\underset{\bm{z}}{\operatorname*{arg\,min}}\,\frac{1}{2}\|\bm{y}-\mathcal{A}(\mathcal{D}(\bm{z}))\|_{2}^{2}$
$\triangleright$ Solve with initial point $\hat{\bm{z}}_{0}(\bm{z}_{t+1})$

$\bm{z}_{t}=\text{StochasticResample}(\hat{\bm{z}}_{0}(\bm{y}),\bm{z}^{\prime}_{t},\gamma)$
$\triangleright$ Map back to $t$

else

$\bm{z}_{t}=\bm{z}^{\prime}_{t}$ $\triangleright$ Unconditional sampling if
not resampling

$\bm{x}_{0}=\mathcal{D}(\bm{z}_{0})$ $\triangleright$ Output reconstructed
image

#### Remapping Back to $\bm{z}_{t}$.

Following the flowchart in Figure 2, the next step is to map the measurement-
consistent sample $\hat{\bm{z}}_{0}(\bm{y})$ back onto the data manifold
defined by the noisy samples at time $t$ to continue the reverse sampling
process. Doing so would be equivalent to computing the posterior distribution
$p(\bm{z}_{t}|\bm{y})$. To incorporate $\hat{\bm{z}}_{0}(\bm{y})$ into the
posterior, we propose to construct an auxiliary distribution
$p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ in replace of
$p(\bm{z}_{t}|\bm{y})$. Here, $\hat{\bm{z}}_{t}$ denotes the remapped sample
and $\bm{z}^{\prime}_{t}$ denotes the unconditional sample before remapping.
One simple way of computing this distribution to obtain $\hat{\bm{z}}_{t}$ is
shown in Proposition 1.

######  Proposition 1 (Stochastic Encoding).

Since the sample $\hat{\bm{z}}_{t}$ given $\hat{\bm{z}}_{0}(\bm{y})$ and
measurement $\bm{y}$ is conditionally independent of $\bm{y}$, we have that

$\displaystyle
p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}))=\mathcal{N}(\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}(\bm{y}),(1-\bar{\alpha}_{t})\bm{I}).$
(12)

We defer all of the proofs to the Appendix. Proposition 1 provides us a way of
computing $\hat{\bm{z}}_{t}$, which we refer to as _stochastic encoding_.
However, we observe that using stochastic encoding can incur a high variance
when $t$ is farther away from $t=0$, where the ground truth signal exists.
This large variance can often lead to noisy image reconstructions. To address
this issue, we propose a posterior sampling technique that reduces the
variance by additionally conditioning on $\bm{z}^{\prime}_{t}$, the
unconditional sample at time $t$. Here, the intuition is that by using
information of $\bm{z}^{\prime}_{t}$, we can get closer to the ground truth
$\bm{z}_{t}$, which effectively reduces the variance. In Lemma 2, under some
mild assumptions, we show that this new distribution
$p(\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ is a
tractable Gaussian distribution.

######  Proposition 2 (Stochastic Resampling).

Suppose that
$p(\bm{z}^{\prime}_{t}|\hat{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ is
normally distributed such that
$p(\bm{z}^{\prime}_{t}|\hat{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=\mathcal{N}(\bm{\mu}_{t},\sigma_{t}^{2})$.
If we let $p(\hat{\bm{z}_{t}}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ be a prior for
$\bm{\mu}_{t}$, then the posterior distribution
$p(\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ is
given by

$\displaystyle
p(\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=\mathcal{N}\left(\frac{\sigma_{t}^{2}\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}(\bm{y})+(1-\bar{\alpha}_{t})\bm{z}^{\prime}_{t}}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})},\frac{\sigma_{t}^{2}(1-\bar{\alpha}_{t})}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})}\bm{I}\right).$
(13)

We refer to this new mapping technique as _stochastic resampling_. Since we do
not have access to $\sigma_{t}^{2}$, it serves as a hyperparameter that we
tune in our algorithm. The choice of $\sigma_{t}^{2}$ plays a role of
controlling the tradeoff between prior consistency and data consistency. If
$\sigma_{t}^{2}\to 0$, then we recover unconditional sampling, and if
$\sigma_{t}^{2}\to\infty$, we recover stochastic encoding. We observe that
this new technique also has several desirable properties, for which we
rigorously prove in the next section.

###  3.2 Theoretical Results

In Section 3.1, we discussed that stochastic resampling induces less variance
than stochastic encoding. Here, we aim to rigorously prove the validity of
this statement.

######  Lemma 1.

Let $\tilde{\bm{z}}_{t}$ and $\hat{\bm{z}}_{t}$ denote the stochastically
encoded and resampled image of $\hat{\bm{z}}_{0}(\bm{y})$, respectively. If
$\text{VAR}(\bm{z}^{\prime}_{t})>0$, then we have that
$\text{VAR}(\hat{\bm{z}}_{t})<\text{VAR}(\tilde{\bm{z}}_{t})$.

######  Theorem 1.

If ${\hat{\bm{z}}}_{0}(\bm{y})$ is measurement-consistent such that
$\bm{y}=\mathcal{A}(\mathcal{D}({\hat{\bm{z}}}_{0}(\bm{y})))$, i.e.
$\hat{\bm{z}}_{0}=\hat{\bm{z}}_{0}(\bm{z}_{t})=\hat{\bm{z}}_{0}(\bm{y})$, then
stochastic resample is unbiased such that
$\mathbb{E}[\hat{\bm{z}}_{t}|\bm{y}]=\mathbb{E}[\bm{z}^{\prime}_{t}]$.

These two results, Lemma 1 and Theorem 1, prove the benefits of stochastic
resampling. At a high-level, these proofs rely on the fact the posterior
distributions of both stochastic encoding and resampling are Gaussian and
compare their respective means and variances. In the following result, we
characterize the variance induced by stochastic resampling, and show that as
$t\to 0$, the variance decreases, giving us a reconstructed image that is of
better quality.

######  Theorem 2.

Let $\bm{z}_{0}$ denote a sample from the data distribution and $\bm{z}_{t}$
be a sample from the noisy perturbed distribution at time $t$. Then,

$\displaystyle\text{Cov}(\bm{z}_{0}|\bm{z}_{t})=\frac{(1-\bar{\alpha}_{t})^{2}}{\bar{\alpha}_{t}}\nabla^{2}_{\bm{z}_{t}}\log
p_{\bm{z}_{t}}(\bm{z}_{t})+\frac{1-\bar{\alpha}_{t}}{\bar{\alpha}_{t}}\bm{I}.$

By Theorem 2, noice that since as $\alpha_{t}$ is an increasing sequence that
converges to 1 as $t$ decreases, the variance between the ground truth
$\bm{z}_{0}$ and the estimated $\hat{\bm{z}}_{0}$ decreases to 0 as $t\to 0$,
assuming that $\nabla^{2}_{\bm{z}_{t}}\log p_{\bm{z}_{t}}(\bm{z}_{t})<\infty$.
Following our theory, we empirically show that stochastic resampling can
reconstruct signals that are less noisy than stochastic encoding, as shown in
the next section.††footnotetext: *We have updated the baseline results for
PSLD. More details are provided in the Appendix (Section A.4).

##  4 Experiments

Method | Super Resolution $4\times$ | Inpainting (Random $70\%$)
---|---|---
LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
DPS [11] | $0.173\pm 0.04$ | $28.41\pm 2.20$ | $0.782\pm 0.06$ | $\underline{0.102}\pm 0.02$ | $\underline{32.48}\pm 2.30$ | $\underline{0.899}\pm 0.03$
MCG [12] | $0.193\pm 0.03$ | $25.92\pm 2.35$ | $0.740\pm 0.05$ | $0.134\pm 0.03$ | $29.53\pm 2.70$ | $0.847\pm 0.05$
ADMM-PnP [35] | $0.304\pm 0.04$ | $21.08\pm 3.13$ | $0.631\pm 0.11$ | $0.627\pm 0.07$ | $15.40\pm 2.09$ | $0.342\pm 0.09$
DDRM [13] | $0.151\pm 0.03$ | $\underline{29.49}\pm 1.93$ | $0.817\pm 0.05$ | $0.166\pm 0.03$ | $27.69\pm 1.54$ | $0.798\pm 0.04$
DMPS [16] | $\underline{0.147}\pm 0.03$ | $28.48\pm 1.92$ | $\underline{0.811}\pm 0.05$ | $0.175\pm 0.03$ | $28.84\pm 1.65$ | $0.826\pm 0.03$
Latent-DPS | $0.272\pm 0.05$ | $26.83\pm 2.00$ | $0.690\pm 0.07$ | $0.226\pm 0.04$ | $26.23\pm 1.84$ | $0.703\pm 0.07$
PSLD-LDM∗ [22] | $0.209\pm 0.10$ | $27.61\pm 2.95$ | $0.704\pm 0.17$ | $0.260\pm 0.08$ | $27.07\pm 2.45$ | $0.689\pm 0.11$
ReSample (Ours) | $\bm{0.144}\pm 0.029$ | $\bm{30.45}\pm 2.09$ | $\bm{0.832}\pm 0.05$ | $\bm{0.082}\pm 0.02$ | $\bm{32.77}\pm 2.23$ | $\bm{0.903}\pm 0.03$

Table 1: Quantitative results of super resolution and inpainting on the
CelebA-HQ dataset. Input images have an additive Gaussian noise with
$\sigma_{\bm{y}}=0.01$. Best results are in bold and second best results are
underlined.

Method | Nonlinear Deblurring | Gaussian Deblurring
---|---|---
LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
DPS [11] | $0.230\pm 0.065$ | $\underline{26.81}\pm 2.84$ | $\underline{0.720}\pm 0.077$ | $0.175\pm 0.03$ | $28.36\pm 2.12$ | $0.772\pm 0.07$
MCG [12] | - | - | - | $0.517\pm 0.06$ | $15.85\pm 1.08$ | $0.536\pm 0.08$
ADMM-PnP [35] | $0.499\pm 0.073$ | $16.17\pm 4.01$ | $0.359\pm 0.140$ | $0.289\pm 0.04$ | $20.98\pm 4.51$ | $0.602\pm 0.15$
DDRM [13] | - | - | - | $0.193\pm 0.04$ | $26.88\pm 1.96$ | $0.747\pm 0.07$
DMPS [16] | - | - | - | $0.206\pm 0.04$ | $26.45\pm 1.83$ | $0.726\pm 0.07$
Latent-DPS | $\underline{0.225}\pm 0.04$ | $26.18\pm 1.73$ | $0.703\pm 0.07$ | $0.205\pm 0.04$ | $27.42\pm 1.84$ | $0.729\pm 0.07$
PSLD-LDM∗ [22] | - | - | - | $0.323\pm 0.09$ | $24.21\pm 2.79$ | $0.548\pm 0.12$
ReSample (Ours) | $\bm{0.153}\pm 0.03$ | $\bm{30.18}\pm 2.21$ | $\bm{0.828}\pm 0.05$ | $\bm{0.148}\pm 0.04$ | $\bm{30.69}\pm 2.14$ | $\bm{0.832}\pm 0.05$

Table 2: Quantitative results of Gaussian and nonlinear deblurring on the
CelebA-HQ dataset. Input images have an additive Gaussian noise with
$\sigma_{\bm{y}}=0.01$. Best results are in bold and second best results are
underlined. For nonlinear deblurring, some baselines are omitted, as they can
only solve _linear_ inverse problems.

We conduct experiments to solve both linear and nonlinear inverse problems on
natural and medical images. We compare our algorithm to several state-of-the-
art methods that directly use the pixel space: DPS [11], Manifold Constrained
Gradients (MCG) [12], Denoising Diffusion Destoration Models (DDRM) [13],
Diffusion Model Posterior Sampling (DMPS) [16], and plug-and-play using ADMM
(ADMM-PnP) [35]. We also compare our algorithm to Latent-DPS and Posterior
Sampling with Latent Diffusion (PSLD) [22], a concurrent work we recently
notice also tackling latent diffusion models. Various quantitative metrics are
used for evaluation including Learned Perceptual Image Patch Similarity
(LPIPS) distance, peak signal-to-noise-ratio (PSNR), and structural similarity
index (SSIM). Lastly, we conduct ablation study to compare the performance of
stochastic encoding and our proposed stochastic resampling technique as
mentioned in Section 3.1, and also demonstrate the memory efficiency gained by
leveraging LDMs.

###  4.1 Experimental Results

#### Experiments on Natural Images.

For the experiments on natural images, we use datasets FFHQ [36], CelebA-HQ
[37], and LSUN-Bedroom [38] with the image resolution of $256\times 256\times
3$. We take pre-trained latent diffusion models LDM-VQ4 trained on FFHQ and
CelebA-HQ provided by Rombach et al. [19] with autoencoders that yield images
of size $64\times 64\times 3$, and DDPMs [23] also trained on FFHQ and CelebA-
HQ training sets. Then, we sample $100$ images from both the FFHQ and CelebA-
HQ validation sets for testing evaluation. For computing quantitative results,
all images are normalized to the range $[0,1]$. All experiments had Gaussian
measurement noise with standard deviation $\sigma_{\bm{y}}=0.01$. Due to
limited space, we put the results on FFHQ and details of the hyperparameters
to the Appendix.

For linear inverse problems, we consider the following tasks: (1) Gaussian
deblurring, (2) inpainting (with a random mask), and (3) super resolution. For
Gaussian deblurring, we use a kernel with size $61\times 61$ with standard
deviation $3.0$. For super resolution, we use bicubic downsampling and a
random mask with varying levels of missing pixels for inpainting. For
nonlinear inverse problems, we consider nonlinear deblurring as proposed by
Chung et al. [11]. The quantitative results are displayed in Tables 1 and 2,
with qualitative results in Figure 3. In Tables 1 and 2, we can see that
ReSample significantly outperforms all of the baselines across all three
metrics on the CelebA-HQ dataset. We also observe that ReSample performs
better than or comparable to all baselines on the FFHQ dataset as shown in the
Appendix. Interestingly, we observe that PSLD [22] seems to struggle in the
presence of measurement noise, often resulting in reconstructions that
underperform compared to pixel-based methods. Remarkably, our method excels in
handling nonlinear inverse problems, further demonstrating the flexibility of
our algorithm. We further demonstrate the superiority of ReSample for handling
nonlinear inverse problems in Figure 6, where we show that we can consistently
outperform DPS.

Figure 3: Qualitative results of multiple tasks on the LSUN-Bedroom and
CelebA-HQ datasets. All inverse problems have Gaussian measurement noise with
variance $\sigma_{\bm{y}}=0.01$. Figure 4: Qualitative results of CT
reconstruction on the LDCT dataset. We annotate the critical image structures
in a red box, and zoom in below the image.

Method | Abdominal | Head | Chest
---|---|---|---
PSNR$\uparrow$ | SSIM$\uparrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
Latent-DPS | 26.80 $\pm 1.09$ | 0.870 $\pm 0.026$ | 28.64 $\pm 5.38$ | 0.893 $\pm 0.058$ | 25.67$\pm 1.14$ | 0.822 $\pm 0.033$
MCG [12] | 29.41 $\pm 3.14$ | 0.857 $\pm 0.041$ | 28.28 $\pm 3.08$ | 0.795 $\pm 0.116$ | 27.92 $\pm 2.48$ | 0.842 $\pm 0.036$
DPS [11] | 27.33 $\pm 2.68$ | $0.715\pm 0.031$ | 24.51 $\pm 2.77$ | 0.665 $\pm 0.058$ | 24.73 $\pm 1.84$ | 0.682 $\pm 0.113$
PnP-UNet [39] |  32.84 $\pm 1.29$ |  0.942 $\pm 0.008$ |  33.45 $\pm 3.25$ |  0.945 $\pm 0.023$ | 29.67 $\pm 1.14$ |  0.891 $\pm 0.011$
FBP | 26.29 $\pm 1.24$ | 0.727 $\pm 0.036$ | 26.71 $\pm 5.02$ | 0.725 $\pm 0.106$ | 24.12 $\pm 1.14$ | 0.655 $\pm 0.033$
FBP-UNet [40] | 32.77 $\pm 1.21$ | 0.937 $\pm 0.013$ | 31.95 $\pm 3.32$ | 0.917 $\pm 0.048$ |  29.78 $\pm 1.12$ | 0.885 $\pm 0.016$
ReSample (Ours) |  35.91 $\pm 1.22$ |  0.965 $\pm 0.007$ |  37.82 $\pm 5.31$ |  0.978 $\pm 0.014$ |  31.72 $\pm 0.912$ |  0.922 $\pm 0.011$

Table 3: Quantitative results of CT reconstruction on the LDCT dataset. Best
results are in bold and second best results are underlined.

#### Experiments on Medical Images.

For experiments on medical images, we fine-tune LDMs on 2000 2D computed
tomography (CT) images with image resolution of $256\times 256$, randomly
sampled from the AAPM LDCT dataset [41] of 40 patients, and test on the 300 2D
CT images from the remaining 10 patients. We take the FFHQ LDM-VQ4 model
provided by Rombach et al. [19] as pre-trained latent diffusion model in this
experimental setting followed by 100K fine-tuning iterations. Then, we
simulate CT measurements (sinograms) with a parallel-beam geometry using 25
projection angles equally distributed across 180 degrees using the torch-radon
package [42]. Following the natural images, the sinograms were perturbed with
additive Gaussian measurement noise with $\sigma_{\bm{y}}=0.01$. Along with
the baselines used in the natural image experiments, we also compare to the
following methods: (1) Filtered Backprojection (FBP), which is a standard
conventional CT reconstruction technique, and (2) FBP-UNet, a supervised
learning method that maps the FBP reconstructed images to ground truth images
using a UNet model. For FBP-Unet and PnP-UNet, we trained a model on 3480 2D
CT from the training set. For MCG and DPS, we used the pre-trained checkpoints
provided by Chung et al. [12]. We present visual reconstructions with
highlighted critical anatomic structures in Figure 4 and quantitative results
in Table 3. Here, we observe that our algorithm outperforms all of the
baselines in terms of PSNR and SSIM. Visually, we also observe that our
algorithm is able to reconstruct smoother images that have more accurate and
sharper details.

###  4.2 Ablation Studies

#### Effectiveness of the Resampling Technique.

Here, we validate our theoretical results by conducting ablation studies on
stochastic resampling. Specifically, we perform experiments on the LSUN-
Bedroom and CelebA-HQ datasets with tasks of Gaussian deblurring and super-
resolution. As shown in Figure 5, we observe that stochastic resampling
reconstructs smoother images with higher PSNRs compared to stochastic
encoding, corroborating our theory.

#### Effectiveness of Hard Data Consistency.

In Figure 6, we perform an ablation study on the ReSample frequency on CT
reconstruction. This observation is in line with what we expect intuitively,
as more ReSample time steps (i.e., more data consistency) lead to more
accurate reconstructions. More specifically, we observe in Figure 6 (right)
that if we perform more iterations of ReSample, we obtain reconstructions with
higher PSNR.

Table 4: Memory Usage of Different Methods for Gaussian Deblurring on the FFHQ
Dataset

Model | Algorithm | Model Only | Memory Increment | Total
---|---|---|---|---
DDPM | DPS | 1953MB | +3416MB (175%) | 5369MB
| MCG |  | +3421MB (175%) | 5374MB
| DMPS |  | +5215MB (267 %) | 7168MB
| DDRM |  | +18833MB (964 %) | 20786MB
LDM | PSLD | 3969MB | +5516MB (140%) | 9485MB
| ReSample |  | +1040MB (26.2%) | 5009MB

Figure 5: Effectiveness of our resampling technique compared to stochastic
encoding. Results were demonstrated on the LSUN-Bedroom and CelebA-HQ datasets
with measurement noise of $\sigma_{\bm{y}}=0.05$ to highlight the
effectiveness of stochastic resample. Figure 6: Additional results on the
performance of ReSample. Left: Results on nonlinear deblurring highlighting
the performance of ReSample. Right: Ablation study on the ReSample frequency
on the performance.

#### Memory Efficiency.

To demonstrate the memory efficiency of our algorithm, we use the command
nvidia-smi to monitor the memory consumption during solving an inverse
problem, in comparison to our baselines. We present the memory usage for
Gaussian deblurring on the FFHQ dataset in Table 4. Although the entire LDM
models occupy more memory due to the autoencoders, our algorithm itself
exhibits memory efficiency, resulting in lower overall memory usage. This
underscores its potential in domains like medical imaging, where memory
complexity plays a crucial role in feasibility.

##  5 Conclusion

In this paper, we propose ReSample, an algorithm that can effectively leverage
LDMs to solve _general_ inverse problems. We demonstrated that (1) our
algorithm can reconstruct high-quality images compared to many baselines and
(2) we can fine-tune existing models for medical imaging tasks and achieve an
additional reduction in memory usage. One fundamental limitation of our method
lies in the computational overhead of hard data consistency, which we leave as
a significant challenge for future work to address and improve upon.

##  6 Acknowledgement

BS and LS acknowledges support from U-M MIDAS PODS Grant and U-M MICDE
Catalyst Grant, and computing resource support from NSF ACCESS Program and
Google Cloud Research Credits Program. This work used NCSA Delta GPU through
allocation CIS230133 and ELE230011 from the Advanced Cyberinfrastructure
Coordination Ecosystem: Services & Support (ACCESS) program, which is
supported by National Science Foundation grants #2138259, #2138286, #2138307,
#2137603, and #2138296. SMK and QQ acknowledge support from U-M START & PODS
grants, NSF CAREER CCF-2143904, NSF CCF-2212066, NSF CCF-2212326, ONR
N00014-22-1-2529, AWS AI Award, and a gift grant from KLA.

## References

  * [1] A. Beck and M. Teboulle, “A fast iterative shrinkage-thresholding algorithm for linear inverse problems,” SIAM Journal on Imaging Sciences, vol. 2, no. 1, pp. 183–202, 2009.
  * [2] M. V. Afonso, J. M. Bioucas-Dias, and M. A. T. Figueiredo, “An augmented lagrangian approach to the constrained optimization formulation of imaging inverse problems,” IEEE Transactions on Image Processing, vol. 20, no. 3, pp. 681–695, 2011.
  * [3] P. Suetens, Fundamentals of medical imaging.  Cambridge university press, 2017.
  * [4] S. Ravishankar, J. C. Ye, and J. A. Fessler, “Image reconstruction: From sparsity to data-adaptive methods and machine learning,” Proceedings of the IEEE, vol. 108, no. 1, pp. 86–109, 2019.
  * [5] W. Liu, X. Xia, L. Xiong, Y. Lu, L. Gao, and Z. Yu, “Automated vehicle sideslip angle estimation considering signal measurement characteristic,” IEEE Sensors Journal, vol. 21, no. 19, pp. 21675–21687, 2021.
  * [6] W. Liu, K. Quijano, and M. M. Crawford, “Yolov5-tassel: detecting tassels in rgb uav imagery with improved yolov5 based on transfer learning,” IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, vol. 15, pp. 8085–8094, 2022.
  * [7] P. Dhariwal and A. Nichol, “Diffusion models beat gans on image synthesis,” Advances in Neural Information Processing Systems, vol. 34, pp. 8780–8794, 2021.
  * [8] T. Karras, M. Aittala, T. Aila, and S. Laine, “Elucidating the design space of diffusion-based generative models,” arXiv preprint arXiv:2206.00364, 2022.
  * [9] Y. Song, P. Dhariwal, M. Chen, and I. Sutskever, “Consistency models,” arXiv preprint arXiv:2303.01469, 2023.
  * [10] Y. Song, L. Shen, L. Xing, and S. Ermon, “Solving inverse problems in medical imaging with score-based generative models,” International Conference on Learning Representations, 2022.
  * [11] H. Chung, J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye, “Diffusion posterior sampling for general noisy inverse problems,” in The Eleventh International Conference on Learning Representations, 2023.
  * [12] H. Chung, B. Sim, D. Ryu, and J. C. Ye, “Improving diffusion models for inverse problems using manifold constraints,” arXiv preprint arXiv:2206.00941, 2022.
  * [13] B. Kawar, M. Elad, S. Ermon, and J. Song, “Denoising diffusion restoration models,” arXiv preprint arXiv:2201.11793, 2022.
  * [14] J. Song, A. Vahdat, M. Mardani, and J. Kautz, “Pseudoinverse-guided diffusion models for inverse problems,” in International Conference on Learning Representations, 2023.
  * [15] H. Chung, S. Lee, and J. C. Ye, “Fast diffusion sampler for inverse problems by geometric decomposition,” arXiv preprint arXiv:2303.05754, 2023.
  * [16] X. Meng and Y. Kabashima, “Diffusion model based posterior sampling for noisy linear inverse problems,” arXiv preprint arXiv:2211.12343, 2022.
  * [17] D. Zhang and F. Zhou, “Self-supervised image denoising for real-world images with context-aware transformer,” IEEE Access, vol. 11, pp. 14340–14349, 2023.
  * [18] M. Lustig, D. Donoho, and J. M. Pauly, “Sparse mri: The application of compressed sensing for rapid mr imaging,” Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, vol. 58, no. 6, pp. 1182–1195, 2007.
  * [19] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, “High-resolution image synthesis with latent diffusion models,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10684–10695, 2022.
  * [20] A. Vahdat, K. Kreis, and J. Kautz, “Score-based generative modeling in latent space,” in Advances in Neural Information Processing Systems (M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, eds.), vol. 34, pp. 11287–11302, Curran Associates, Inc., 2021.
  * [21] N. Ruiz, Y. Li, V. Jampani, Y. Pritch, M. Rubinstein, and K. Aberman, “Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22500–22510, 2023.
  * [22] L. Rout, N. Raoof, G. Daras, C. Caramanis, A. G. Dimakis, and S. Shakkottai, “Solving linear inverse problems provably via posterior sampling with latent diffusion models,” arXiv preprint arXiv:2307.00619, 2023.
  * [23] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” Advances in Neural Information Processing Systems, vol. 33, pp. 6840–6851, 2020.
  * [24] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, “Score-based generative modeling through stochastic differential equations,” arXiv preprint arXiv:2011.13456, 2021.
  * [25] P. Vincent, “A connection between score matching and denoising autoencoders,” Neural Computation, vol. 23, no. 7, pp. 1661–1674, 2011.
  * [26] J. Song, C. Meng, and S. Ermon, “Denoising diffusion implicit models,” arXiv preprint arXiv:2010.02502, 2020.
  * [27] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu, “Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps,” arXiv preprint arXiv:2206.00927, 2022.
  * [28] Y. Wang, J. Yu, and J. Zhang, “Zero-shot image restoration using denoising diffusion null-space model,” arXiv preprint arXiv:2212.00490, 2022.
  * [29] A. Bora, A. Jalal, E. Price, and A. G. Dimakis, “Compressed sensing using generative models,” in International Conference on Machine Learning, pp. 537–546, PMLR, 2017.
  * [30] A. Jalal, M. Arvinte, G. Daras, E. Price, A. G. Dimakis, and J. Tamir, “Robust compressed sensing mri with deep generative priors,” in Advances in Neural Information Processing Systems (M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, eds.), vol. 34, pp. 14938–14954, Curran Associates, Inc., 2021.
  * [31] V. Shah, R. Hyder, M. S. Asif, and C. Hegde, “Provably convergent algorithms for solving inverse problems using generative models,” arXiv preprint arXiv:2105.06371, 2021.
  * [32] V. Lempitsky, A. Vedaldi, and D. Ulyanov, “Deep image prior,” in 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9446–9454, 2018.
  * [33] B. Efron, “Tweedie’s formula and selection bias,” Journal of the American Statistical Association, vol. 106, no. 496, pp. 1602–1614, 2011.
  * [34] J. Yu, Y. Wang, C. Zhao, B. Ghanem, and J. Zhang, “Freedom: Training-free energy-guided conditional diffusion model,” arXiv preprint arXiv:2303.09833, 2023.
  * [35] R. Ahmad, C. A. Bouman, G. T. Buzzard, S. H. Chan, E. T. Reehorst, and P. Schniter, “Plug and play methods for magnetic resonance imaging.,” 2019.
  * [36] V. Kazemi and J. Sullivan, “One millisecond face alignment with an ensemble of regression trees,” in 2014 IEEE Conference on Computer Vision and Pattern Recognition, pp. 1867–1874, 2014.
  * [37] Z. Liu, P. Luo, X. Wang, and X. Tang, “Deep learning face attributes in the wild,” in Proceedings of International Conference on Computer Vision (ICCV), 2015.
  * [38] F. Yu, A. Seff, Y. Zhang, S. Song, T. Funkhouser, and J. Xiao, “Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop,” arXiv preprint arXiv:1506.03365, 2016.
  * [39] D. Gilton, G. Ongie, and R. Willett, “Model adaptation for inverse problems in imaging,” IEEE Transactions on Computational Imaging, vol. 7, pp. 661–674, 2021.
  * [40] K. H. Jin, M. T. McCann, E. Froustey, and M. Unser, “Deep convolutional neural network for inverse problems in imaging,” IEEE Transactions on Image Processing, vol. 26, no. 9, pp. 4509–4522, 2017.
  * [41] T. R. Moen, B. Chen, D. R. Holmes III, X. Duan, Z. Yu, L. Yu, S. Leng, J. G. Fletcher, and C. H. McCollough, “Low-dose ct image and projection dataset,” Medical physics, vol. 48, no. 2, pp. 902–911, 2021.
  * [42] M. Ronchetti, “Torchradon: Fast differentiable routines for computed tomography,” arXiv preprint arXiv:2009.14788, 2020.
  * [43] A. Q. Nichol and P. Dhariwal, “Improved denoising diffusion probabilistic models,” in International Conference on Machine Learning, pp. 8162–8171, PMLR, 2021.
  * [44] K. Zhang, W. Zuo, Y. Chen, D. Meng, and L. Zhang, “Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising,” IEEE transactions on image processing, vol. 26, no. 7, pp. 3142–3155, 2017.
  * [45] Y. Han and J. C. Ye, “Framing u-net via deep convolutional framelets: Application to sparse-view ct,” IEEE transactions on medical imaging, vol. 37, no. 6, pp. 1418–1429, 2018.
  * [46] B. Zhu, J. Z. Liu, S. F. Cauley, B. R. Rosen, and M. S. Rosen, “Image reconstruction by domain-transform manifold learning,” Nature, vol. 555, no. 7697, pp. 487–492, 2018.
  * [47] H. Gupta, K. H. Jin, H. Q. Nguyen, M. T. McCann, and M. Unser, “Cnn-based projected gradient descent for consistent ct image reconstruction,” IEEE transactions on medical imaging, vol. 37, no. 6, pp. 1440–1453, 2018.
  * [48] H. Wei, F. Schiffers, T. Würfl, D. Shen, D. Kim, A. K. Katsaggelos, and O. Cossairt, “2-step sparse-view ct reconstruction with a domain-specific perceptual network,” arXiv preprint arXiv:2012.04743, 2020.
  * [49] J. Liang, J. Cao, G. Sun, K. Zhang, L. Van Gool, and R. Timofte, “Swinir: Image restoration using swin transformer,” in Proceedings of the IEEE/CVF international conference on computer vision, pp. 1833–1844, 2021.
  * [50] S. A. Hussein, T. Tirer, and R. Giryes, “Image-adaptive gan based reconstruction,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, pp. 3121–3129, 2020.
  * [51] G.-H. Liu, T. Chen, E. A. Theodorou, and M. Tao, “Mirror diffusion models for constrained and watermarked generation,” arXiv preprint arXiv:2310.01236, 2023.
  * [52] A. Lou and S. Ermon, “Reflected diffusion models,” arXiv preprint arXiv:2304.04740, 2023.
  * [53] N. Fishman, L. Klarner, V. De Bortoli, E. Mathieu, and M. Hutchinson, “Diffusion models for constrained domains,” arXiv preprint arXiv:2304.05364, 2023.
  * [54] H. Chung, J. C. Ye, P. Milanfar, and M. Delbracio, “Prompt-tuning latent diffusion models for inverse problems,” arXiv preprint arXiv:2310.01110, 2023.
  * [55] P. A. Bromiley, “Products and convolutions of gaussian probability density functions,” 2013.

Appendix

In this section, we present additional experimental results to supplement
those presented in the main paper. In Section A, we present additional
qualitative and quantitative results to highlight the performance of ReSample.
In Section B, we briefly discuss some of the implementation details regarding
hard data consistency. In Section C, we outline all of the hyperparameters
used to produce our results. In Section D, we discuss some more reasons to why
Latent-DPS fails to consistently accurately recover the underlying image.
Lastly, in Section F, we present our deferred proofs.

##  Appendix A Additional Results

###  A.1 Results on FFHQ

Here, we provide additional quantitative results on the FFHQ dataset. Similar
to the results in the main text on the CelebA-HQ dataset, ReSample outperforms
the baselines across many different tasks.

Method | Super Resolution $4\times$ | Inpainting (Random $70\%$)
---|---|---
LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
DPS | $\underline{0.175}\pm 0.04$ | $\underline{28.47}\pm 2.09$ | $\underline{0.793}\pm 0.05$ | $\underline{0.106}\pm 0.03$ | $\bm{32.32}\pm 2.21$ | $\bm{0.897}\pm 0.04$
MCG | $0.223\pm 0.04$ | $23.74\pm 3.18$ | $0.673\pm 0.07$ | $0.178\pm 0.03$ | $24.89\pm 2.57$ | $0.731\pm 0.06$
ADMM-PnP | $0.303\pm 0.03$ | $21.30\pm 3.98$ | $0.760\pm 0.04$ | $0.308\pm 0.07$ | $15.87\pm 2.09$ | $0.608\pm 0.09$
DDRM | $0.257\pm 0.03$ | $27.51\pm 1.79$ | $0.753\pm 0.05$ | $0.287\pm 0.03$ | $24.97\pm 1.51$ | $0.680\pm 0.05$
DMPS | $0.181\pm 0.03$ | $27.21\pm 1.92$ | $0.766\pm 0.06$ | $0.150\pm 0.03$ | $28.17\pm 1.58$ | $0.814\pm 0.04$
Latent-DPS | $0.344\pm 0.05$ | $24.65\pm 1.77$ | $0.609\pm 0.07$ | $0.270\pm 0.05$ | $27.08\pm 1.97$ | $0.727\pm 0.06$
PSLD-LDM | $0.267\pm 0.10$ | $27.22\pm 2.99$ | $0.705\pm 0.13$ | $0.270\pm 0.06$ | $25.61\pm 2.02$ | $0.630\pm 0.09$
ReSample (Ours) | $\bm{0.164}\pm 0.03$ | $\bm{28.90}\pm 1.86$ | $\bm{0.804}\pm 0.05$ | $\bm{0.099}\pm 0.02$ | $\underline{31.34}\pm 2.11$ | $\underline{0.890}\pm 0.03$

Table 5: Comparison of quantitative results on inverse problems on the FFHQ
dataset. Input images have an additive Gaussian noise with
$\sigma_{\bm{y}}=0.01$. Best results are in bold and second best results are
underlined.

Method | Nonlinear Deblurring | Gaussian Deblurring
---|---|---
LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
DPS | $\underline{0.197}\pm 0.06$ | $\underline{27.13}\pm 3.11$ | $\underline{0.762}\pm 0.08$ | $\bm{0.169}\pm 0.04$ | $\underline{27.70}\pm 2.04$ | $\underline{0.774}\pm 0.05$
MCG | - | - | - | $0.371\pm 0.04$ | $25.33\pm 1.74$ | $0.668\pm 0.06$
ADMM-PnP | $0.424\pm 0.03$ | $21.80\pm 1.42$ | $0.497\pm 0.05$ | $0.399\pm 0.03$ | $21.23\pm 3.01$ | $0.675\pm 0.06$
DDRM | - | - | - | $0.299\pm 0.039$ | $26.51\pm 1.67$ | $0.702\pm 0.06$
DMPS | - | - | - | $0.227\pm 0.036$ | $26.04\pm 1.75$ | $0.699\pm 0.066$
Latent-DPS | $0.319\pm 0.05$ | $24.52\pm 1.75$ | $0.637\pm 0.07$ | $0.258\pm 0.05$ | $25.98\pm 1.66$ | $0.704\pm 0.06$
PSLD-LDM | - | - | - | $0.422\pm 0.11$ | $20.08\pm 2.97$ | $0.400\pm 0.16$
ReSample (Ours) | $\bm{0.171}\pm 0.04$ | $\bm{30.03}\pm 1.78$ | $\bm{0.838}\pm 0.04$ | $\underline{0.201}\pm 0.04$ | $\bm{28.73}\pm 1.87$ | $\bm{0.794}\pm 0.05$

Table 6: Comparison of quantitative results on inverse problems on the FFHQ
dataset. Input images have an additive Gaussian noise with
$\sigma_{\bm{y}}=0.01$. Best results are in bold and second best results are
underlined.

We observe that ReSample outperforms all baselines in nonlinear deblurring and
super resolution, while demonstrating comparable performance to DPS [11] in
Gaussian deblurring and inpainting. An interesting observation is that
ReSample exhibits the largest performance gap from DPS in nonlinear deblurring
and the smallest performance gap in random inpainting, mirroring the pattern
we observed in the results of the CelebA-HQ dataset. This may suggest that
ReSample performs better than baselines when the forward operator is more
complex.

###  A.2 Additional Results on Box Inpainting

In this section, we present both qualitative and quantitative results on a
more challenging image inpainting setting. More specifically, we consider the
“box” inpainting setting, where the goal is to recover a ground truth image in
which large regions of pixels are missing. We present our results in Figure 7
and Table 7, where we compare the performance of our algorithm to PSLD [22]
(latent-space algorithm) and DPS [11] (pixel-space algorithm). Even in this
challenging setting, we observe that ReSample can outperform the baselines,
highlighting the effectiveness of our method.

Method | Box Inpainting
---|---
LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
DPS | $0.127$ | $22.85$ | $\underline{0.861}$
DDRM | $\underline{0.120}$ | $\underline{24.33}$ | $0.860$
DMPS | $0.233$ | $22.01$ | $0.803$
Latent-DPS | $0.199$ | $23.14$ | $0.784$
PSLD-LDM | $0.201$ | $23.99$ | $0.787$
ReSample (Ours) | $\mathbf{0.093}$ | $\mathbf{24.67}$ | $\mathbf{0.892}$

Table 7: Comparison of quantitative results for box inpainting on CelebA-HQ
dataset. Input images have an additive Gaussian noise with
$\sigma_{\bm{y}}=0.01$. Best results are in bold and second best results are
underlined.

###  A.3 Additional Preliminary Results on High-Resolution Images

In this section, we present techniques for extending our algorithm to address
inverse problems with high-resolution images. One straightforward method
involves using arbitrary-resolution random noise as the initial input for the
LDM. The convolutional architecture shared by the UNet in both the LDM and the
autoencoder enables the generation of images at arbitrary resolutions, as
demonstrated by [19]. To illustrate the validity of this approach, we
conducted a random-inpainting experiment with images of dimensions $(512\times
512\times 3)$ and report the results in Figure 8. In Figure 8, we observe that
our method can achieve excellent performance even in this high-resolution
setting. Other possible methods include obtaining an accurate pair of encoders
and decoders that can effectively map higher-resolution images to the low-
dimensional space, and then running our algorithm using the latent diffusion
model. We believe that these preliminary results represent an important step
towards using LDMs as generative priors for solving inverse problems for high-
resolution images.

Figure 7: Qualitative results on box inpainting with measurement noise
$\sigma_{\bm{y}}=0.01$. These results highlight the effectiveness of ReSample
on more difficult inverse problem tasks. Figure 8: Additional results on 70%
random inpainting with high-resolution images ($512\times 512\times 3$) with
measurement noise $\sigma_{\bm{y}}=0.01$.

###  A.4 Discussion on PSLD Baseline

To better reflect the performance of PSLD [22], we re-ran several experiments
with PSLD while collaborating with the original authors to fine-tune
hyperparameters. Furthermore, as the work by Rout et al. [22] initially
employed the stable diffusion (SD) model rather than the LDM model, we have
relabeled the baselines as “PSLD-LDM”. We report the hyperparameters used to
generate the new results, as well as the ones used to generate the previous
results in Table 8. The previous hyperparameters were chosen based on the
implementation details provided by Rout et al. in Section B.1 [22], where they
stated that they used $\gamma=0.1$ and that the hyperparameters were available
in the codebase. However, it has been brought to our attention that this was
not the optimal hyperparameter. Therefore, we additionally tuned the parameter
so that each baseline could obtain the best results. We have also conducted
additional experiments directly comparing to PSLD-LDM for inpainting tasks on
the FFHQ 1K validation set and present the results in Table 9. Throughout
these results, we still observe that our algorithm largely outperforms the
baselines, including PSLD.

Hyperparameter [22] | SR (FFHQ) | SR (CelebA-HQ) | Gaussian Deblurring (CelebA-HQ)
---|---|---|---
Previous ($\gamma$) | $0.1$ | $0.1$ | $0.1$
New ($\gamma$) | $1.0$ | $1.1$ | $0.15$
Table 8: Different hyperparameters used to generate PSLD results. SR refers to
super resolution $(4\times)$, “previous” and “new” refer to the hyperparameter
$\gamma$ used to generate the old and new results, respectively.

Method | Inpainting (Random) | Inpainting (Box)
---|---|---
LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$ | LPIPS$\downarrow$ | PSNR$\uparrow$ | SSIM$\uparrow$
DPS | $\underline{0.212}$ | $29.49$ | $0.844$ | $0.214$ | $23.39$ | $0.798$
PSLD-LDM | $0.221$ | $\underline{30.31}$ | $\underline{0.851}$ | $\underline{0.158}$ | $\mathbf{24.22}$ | $\underline{0.819}$
ReSample (Ours) | $\mathbf{0.135}$ | $\mathbf{31.64}$ | $\mathbf{0.906}$ | $\mathbf{0.128}$ | $\underline{23.81}$ | $\mathbf{0.868}$

Table 9: Comparison of quantitative results for different inpainting tasks on
the FFHQ 1k validation set. Best results are in bold and second best results
are underlined.

###  A.5 Additional Qualitative Results

Figure 9: Additional results on super resolution ($4\times$) on the CelebA-HQ
dataset with Gaussian measurement noise of variance $\sigma_{\bm{y}}=0.01$.
Figure 10: Additional results on Gaussian deblurring on the CelebA-HQ dataset
with Gaussian measurement noise of variance $\sigma_{\bm{y}}=0.01$. Figure 11:
Additional results on inpainting with a random mask ($70\%$) on the LSUN-
Bedroom dataset with Gaussian measurement noise of variance
$\sigma_{\bm{y}}=0.01$. Figure 12: Comparison of algorithms on Gaussian
deblurring on the FFHQ dataset with Gaussian measurement noise of variance
$\sigma_{\bm{y}}=0.01$. Figure 13: Additional results on Gaussian deblurring
with additive Gaussian noise $\sigma_{\bm{y}}=0.05$. Figure 14: Additional
results on super resolution $4\times$ with additive Gaussian noise
$\sigma_{\bm{y}}=0.05$. Figure 15: Additional results on CT reconstruction
with additive Gaussian noise $\sigma_{\bm{y}}=0.01$. Figure 16: Additional
results on CT reconstruction with additive Gaussian noise
$\sigma_{\bm{y}}=0.01$.

###  A.6 Ablation Studies

Here, we present some ablation studies regarding computational efficiency and
stochastic resampling, amongst others.

Figure 17: An ablation study on a few of the hyperparameters associated with
ReSample. Left: Performance of different values of $\gamma$ (hyperparameter in
stochastic resampling) on the CelebA-HQ dataset. Right: CT reconstruction
performance as a function of the number of timesteps to perform optimization.

#### Effectiveness of Stochastic Resampling.

Recall that in stochastic resampling, we have one hyperparameter
$\sigma_{t}^{2}$. In Section C, we discuss that our choice for this
hyperparameter is

$\sigma_{t}^{2}=\gamma\left(\frac{1-\bar{\alpha}_{t-1}}{\bar{\alpha}_{t}}\right)\left(1-\frac{\bar{\alpha}_{t}}{\bar{\alpha}_{t-1}}\right),$

where $\bar{\alpha}$ is an adaptive parameter associated to the diffusion
model process. Thus, the only parameter that we need to choose here is
$\gamma$. To this end, we perform a study on how $\gamma$ affects the image
reconstruction quality. The $\gamma$ term can be interpreted as a parameter
that balances the prior consistency with the measurement consistency. In
Figure 17 (left) observe that performance increases a lot when $\gamma$
increases from a small value, but plateaus afterwards. We also observe that
the larger $\gamma$ gives more fine details on the images, but may introduce
additional noise. In Figure 18, we provide visual representations
corresponding to the choices of $\gamma$ in order.

Figure 18: Visual representation of varying $\gamma$ on the CelebA-HQ dataset.

Table 10: Average inference of different algorithms for Gaussian Deblurring

Algorithm | Inference Time
---|---
DPS | 1:12
MCG | 1:20
DMPS | 1:25
DDRM | 10 seconds
PSLD | 2:11
ReSample (Ours) | 1:54

#### Effect of the Skip Step Size.

In this section, we conduct an ablation study to investigate the impact of the
skip step size in our algorithm. The skip step size denotes the frequency at
which we ’skip’ before applying the next hard data consistency. For instance,
a skip step size of $10$ implies that we apply hard data consistency every
$10$ iterations of the reverse sampling process. Intuitively, one might expect
that fewer skip steps would lead to better reconstructions, indicating more
frequent hard data consistency steps. To validate this intuition, we conducted
an experiment on the effect of the skip step size on CT reconstruction, and
the results are displayed in Figure 19, with corresponding inference times in
Table 16. We observe that for skip step sizes ranging from $1$ to $10$, the
results are very similar. However, considering the significantly reduced time
required for a skip step size of $10$ as shown in Table 16, we chose a skip
step size of $10$ in our experimental setting in order to balance the trade-
off between reconstruction quality and inference time. Lastly, we would like
to note that in Figure 19, a skip step size of $4$ exhibits a (very) slight
improvement over a skip step of $1$. Due to the reverse sampling process
initiating with a _random_ noise vector, there is a minor variability in the
results, with the average outcomes for skip step sizes of $1$ and $4$ being
very similar.

Algorithm | Skip Step Size | Inference Time | PSNR | SSIM
---|---|---|---|---
ReSample (Ours) | $1$ | 31:20 | 31.74 | 0.922
| $4$ | 8:24 | 31.74 | 0.923
| $10$ | 3:52 | 31.72 | 0.922
| $20$ | 2:16 | 31.23 | 0.918
| $50$ | 1:20 | 30.86 | 0.915
| $100$ | 1:08 | 30.09 | 0.906
| $200$ | 0:54 | 28.74 | 0.869
Table 11: Inference times and corresponding performance for performing hard
data consistency on with varying skip step sizes for chest CT reconstruction.
Figure 19: Observing the effect of the skip step size of hard data consistency
optimization on CT reconstruction. The numbers above the image correspond to
the number of hard data consistency steps per iteration of the reverse
sampling process (e.g., $1$ refers to hard data consistency on every step),
and the red boxes outline the differences in reconstructions.

#### Computational Efficiency.

Here, we present the memory usage and inference times of ReSample compared to
our baselines in Tables 10 and 12. We observe that the memory gap between
ReSample and other algorithms widens when the pre-trained diffusion model
increases in size, as the hard data consistency step of ReSample requires
minimal memory. Conversely, since we only apply hard data consistency in a
subset of all sampling steps (which differs from existing methods that perform
consistency updates at every time step), we note a slight increase in
inference time compared to algorithms such as DPS, DMPS, and MCG. ReSample
outperforms PSLD in terms of inference time.

Table 12: Memory Usage of different algorithms with different pretrained
models for Gaussian Deblurring on FFHQ256 and ImageNet512 datasets

Pretrained Model | Algorithm | Model Only | Memory Increment | Total Memory
---|---|---|---|---
DDPM(FFHQ) | DPS | 1953MB | +3416MB (175%) | 5369MB
| MCG |  | +3421MB (175%) | 5374MB
| DMPS |  | +5215MB (267 %) | 7168MB
| DDRM |  | +18833MB (964 %) | 20786MB
LDM(FFHQ) | PSLD | 3969MB | +5516MB (140%) | 9485MB
| ReSample (Ours) |  | +1040MB (26.2%) | 5009MB
DDPM(ImageNet) | DPS | 4394MB | +6637MB (151%) | 11031MB
| MCG |  | +6637MB (151%) | 11031MB
| DMPS |  | +8731MB (199 %) | 13125MB
| DDRM |  | +4530MB (103 %) | 8924MB
LDM(ImageNet) | PSLD | 5669MB | +5943MB (105%) | 11612MB
| ReSample (Ours) |  | +1322MB (30.1%) | 7002MB

#### Discussion on Latent-DPS.

In the main text, we briefly discussed how adding the Latent-DPS gradient term
into our ReSample algorithm _can_ improve the overall reconstruction quality.
Generally, we observe that adding the Latent-DPS gradient will cause a very
marginal boost in the PSNR, but only when the learning rate scale is chosen
“correctly”. Interestingly, we observe that choosing $\bar{\alpha}_{t}$ is
_critical_ in the performance of Latent-DPS. To this end, we perform a brief
ablation study on the learning rate for Latent-DPS, where we choose the
learning rate to be $k\bar{\alpha}_{t}$ for some $k>0$. We vary $k$ and test
the performance on 50 images of chest CT images and display the results in
Table 13. In Table 13, we observe that $k=0.5$ returns the best results, and
should be chosen if one were to adopt the method of adding Latent-DPS into
ReSample.

$k$ | $0$ | $0.5$ | $1.0$ | $1.5$ | $2.0$ | $2.5$
---|---|---|---|---|---|---
PSNR | $31.64$ | $\mathbf{31.73}$ | $31.58$ | $31.51$ | $31.27$ | $31.04$
Measurement Loss | $\mathbf{9.82e-5}$ | $1.21e-4$ | $5.03e-4$ | $1.35e-3$ | $1.96e-3$ | $2.63e-3$
Table 13: The effect of the Latent-DPS gradient scale (learning rate) as a
function of $k>0$. We study the PSNR and measurement loss (objective function)
changes for varying $k$ for chest CT reconstruction on 50 test samples. The
best results are in bold. Note that $k=0$ here refers to no Latent-DPS and
only using ReSample.

#### Training Efficiency of LDMs.

To underscore the significance of our work, we conducted an ablation study
focusing on the training efficiency of LDMs. This study demonstrates that LDMs
require significantly fewer computational resources for training, a trait
particularly beneficial for downstream applications like medical imaging. To
support this assertion, we present the training time and performance of LDMs,
comparing them to DDPMs [23] in the context of medical image synthesis. For
LDMs, we utilized the pretrained autoencoder and LDM architecture provided by
[19], training them on 9000 2D CT slices from three organs across 40 patients
sourced from the LDCT training set [41]. In the case of DDPMs, we employed the
codebase provided by [43] to train pixel-based diffusion models on the same
set of 2D CT training slices. The results are presented in comparison with 300
2D slices from 10 patients in the validation set, detailed in Table 14.
Observing the results in Table 14, we note that by utilizing LDMs, we can
significantly reduce training time and memory usage while achieving better
generation quality, as measured by the FID score.

Model type | Training Time | Training Memory | FID
---|---|---|---
LDM (Ours) | 133 mins | 7772MB | 100.48
DDPM | 229 mins | 10417MB | 119.50
Table 14: Comparison of training efficiency and performance of LDMs and DDPMs
on CT images, computed on a V100 GPU.

##  Appendix B Discussion on Hard Data Consistency

Recall that the hard data consistency step involves solving the following
optimization problem:

$\displaystyle\hat{\bm{z}}_{0}(\bm{y})\in\underset{\bm{z}}{\operatorname*{arg\,min}}\,\frac{1}{2}\|\bm{y}-\mathcal{A}(\mathcal{D}(\bm{z}))\|^{2}_{2},$
(14)

where we initialize using $\hat{\bm{z}}_{0}(\bm{z}_{t})$. Instead of solving
for the latent variable $\bm{z}$ directly, notice that one possible technique
would be to instead solve for vanilla least squares:

$\displaystyle\hat{\bm{x}}_{0}(\bm{y})\in\underset{\bm{x}}{\operatorname*{arg\,min}}\,\frac{1}{2}\|\bm{y}-\mathcal{A}(\bm{x})\|^{2}_{2},$
(15)

where we instead initialize using $\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))$,
where $\mathcal{D}(\cdot)$ denotes the decoder. Similarly, our goal here is to
find the $\hat{\bm{x}}_{0}(\bm{y})$ that is close to
$\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))$that satisfies the measurement
consistency: $\|\bm{y}-\mathcal{A}(\bm{x})\|^{2}_{2}<\sigma^{2}$, where
$\sigma^{2}$ is an estimated noise level. Throughout the rest of the Appendix,
we refer to the former optimization process as _latent optimization_ and the
latter as _pixel optimization_.

In our experiments, we actually found that these two different formulations
yield different results, in the sense that performing latent optimization
gives reconstructions that are “noisy”, yet much sharper with fine details,
whereas pixel optimization gives results that are “smoother” yet blurry with
high-level semantic information. Here, the intuition is that pixel
optimization does not directly change the latent variable where as latent
optimization directly optimizes over the latent variable. Moreover, the
encoder $\mathcal{E}(\cdot)$ can add additional errors to the estimated
$\hat{\bm{z}}_{0}(\bm{y})$, perhaps throwing the sample off the data manifold,
yielding images that are blurry as a result.

There is also a significant difference in time complexity between these two
methods. Since latent optimization needs to backpropagate through the whole
network of the decoder $\mathcal{D}(\cdot)$ on every gradient step, it takes
much longer to obtain a local minimum (or converge). Empirically, to balance
the trade-off between reconstruction speed and image quality, we see that
using _both_ of these formulations for hard data consistency can not only
yield the best results, but also speed up the optimization process. Since
pixel optimization is easy to get a global optimum, we use it first during the
reverse sampling process and then use latent optimization when we are closer
to $t=0$ to refine the images with the finer details.

Lastly, we would like to remark that for pixel optimization, there is a
closed-form solution that could be leveraged under specific settings [28]. If
the forward operator $\mathcal{A}$ is linear and can take the matrix form
$\mathbf{A}$ and the measurements are noiseless (i.e.,
$\bm{y}=\mathcal{A}(\hat{\bm{x}}_{0}(\bm{y}))$), then we can pose the
optimization problem as

$\displaystyle\hat{\bm{x}}_{0}(\bm{y})\in\underset{\bm{x}}{\operatorname*{arg\,min}}\,\frac{1}{2}\|\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))-\bm{x}\|^{2}_{2},\quad\text{s.t.}\quad\mathbf{A}\bm{x}=\bm{y}.$
(16)

Then, the solution to this optimization problem is given by

$\displaystyle\hat{\bm{x}}_{0}(\bm{y})=\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))-(\mathbf{A}^{+}\mathbf{A}\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))-\mathbf{A}^{+}\bm{y}),$
(17)

where by employing the encoder, we obtain

$\displaystyle\hat{\bm{z}}_{0}(\bm{y})=\mathcal{E}(\hat{\bm{x}}_{0}(\bm{y}))=\mathcal{E}(\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))-(\mathbf{A}^{+}\mathbf{A}\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))-\mathbf{A}^{+}\bm{y})).$
(18)

This optimization technique does not require iterative solvers and offers
great computational efficiency. In the following section, we provide ways in
which we can compute a closed-form solution in the case in which the
measurements may be noisy.

###  B.1 Accelerated Pixel Optimization by Conjugate Gradient

Notice that since pixel optimization directly operates in the pixel space, we
can use solvers such as conjugate gradients least squares for linear inverse
problems. Let $\mathbf{A}$ be the matrix form of the linear operator
$\mathcal{A}(\cdot)$. Then, as discussed in the previous subsection, the
solution to the optimization problem in the noiseless setting is given by

$\displaystyle\hat{\bm{x}}=\bm{x}_{0}-(\mathbf{A}^{+}\mathbf{A}\bm{x}_{0}-\mathbf{A}^{+}\bm{y}),$
(19)

where $\mathbf{A}^{+}=\mathbf{A}^{\top}(\mathbf{AA}^{\top})^{-1}$ and
$(\mathbf{AA}^{\top})^{-1}$ can be implemented by conjugate gradients. In the
presence of measurement noise, we can relax this solution to

$\displaystyle\hat{\bm{x}}=\bm{x}_{0}-\kappa(\mathbf{A}^{+}\mathbf{A}\bm{x}_{0}-\mathbf{A}^{+}\bm{y}),$
(20)

$\kappa\in(0,1)$ is a hyperparameter that can reduce the impact between the
noisy component of the measurements and $\bm{x}_{0}$ (the initial image before
optimization, for which we use $\mathcal{D}(\hat{\bm{z}}_{0}(\bm{z}_{t}))$).

We use this technique for CT reconstruction, where the forward operator
$\mathbf{A}$ is the radon transform and $\mathbf{A}^{\top}$ is the non-
filtered back projection. However, we can use this conjugate gradient
technique for any linear inverse problem where the matrix $\mathbf{A}$ is
available.

##  Appendix C Implementation Details

In this section, we discuss the choices of the hyperparameters used in all of
our experiments for our algorithm. All experiments are implemented in PyTorch
on NVIDIA GPUs (A100 and A40).

Notation | Definition
---|---
$\tau$ | Hyperparameter for early stopping in hard data consistency
$\sigma_{t}$ | Variance scheduling for the resampling technique
$T$ | Number of DDIM or DDPM steps
Table 15: Summary of the hyperparameters with their respective notations for
ReSample.

###  C.1 Linear and Nonlinear Inverse Problems on Natural Images

For organizational purposes, we tabulate all of the hyperparameters associated
to ReSample in Table 15. The parameter for the number of times to perform hard
data consistency is not included in the table, as we did not have any explicit
notation for it. For experiments across all natural images datasets (LSUN-
Bedroom, FFHQ, CelebA-HQ), we used the same hyperparameters as they seemed to
all empirically give the best results.

For $T$, we used $T=500$ DDIM steps. For hard data consistency, we first split
$T$ into three even sub-intervals. The first stage refers to the sub-interval
closest to $t=T$ and the third stage refers to the interval closest to $t=0$.
During the second stage, we performed pixel optimization, whereas in the third
stage we performed latent optimization for hard data consistency as described
in Section B. We performed this optimization on every $10$ iterations of $t$.

We set $\tau=10^{-4}$, which seemed to give us the best results for noisy
inverse problems, with a maximum number of iterations of $2000$ for pixel
optimization and $500$ for latent optimization (whichever convergence criteria
came first). For the variance hyperparameter $\sigma_{t}$ in the stochastic
resample step, we chose an adaptive schedule of

$\sigma_{t}^{2}=\gamma\left(\frac{1-\bar{\alpha}_{t-1}}{\bar{\alpha}_{t}}\right)\left(1-\frac{\bar{\alpha}_{t}}{\bar{\alpha}_{t-1}}\right),$

as discussed in Section A. Generally, we see that $\gamma=40$ returns the best
results for experiments on natural images.

###  C.2 Linear Inverse Problems on Medical Images

Since the LDMs for medical images is not as readily available as compared to
natural images, we largely had to fine-tune existing models. We discuss these
in more detail in this section.

Backbone Models. For the backbone latent diffusion model, we use the pre-
trained model from latent diffusion [19]. We select the VQ-4 autoencoder and
the FFHQ-LDM with CelebA-LDM as our backbone LDM.For inferencing, upon taking
pre-trained checkpoints provided by [19], we fine-tuned the models on 2000 CT
images with 100K iterations and a learning rate of $10^{-5}$.

#### Inferencing.

For $T$, we used a total of $T=1000$ DDIM steps. For hard data consistency, we
split $T$ into three sub-intervals: $t>750$, $300<t\leq 750$, and $t\leq 300$.
During the second stage, we performed pixel optimization by using conjugate
gradients as discussed previously, with 50 iterations with $\kappa=0.9$. In
the third stage, we performed latent optimization with $\tau$ as the estimated
noise level $\tau=10^{-4}$. We set skip step size to be $10$ and $\gamma=40$,
with $\sigma_{t}$ as the same as the experiments for the natural images.

###  C.3 Implementations of Baselines

#### Latent-DPS.

For the Latent-DPS baseline, we use $T=1000$ DDIM steps for CT reconstruction
and $T=500$ DDIM steps for natural image experiments. Let $\zeta_{t}$ denote
the learning rate. For medical images, we use $\zeta_{t}=2.5\bar{\alpha}_{t}$
and $\zeta_{t}=0.5\bar{\alpha}_{t}$ for natural images. Empirically, we
observe that our proposed $\zeta_{t}$ step size schedule gives the best
performance and is robust to scale change, as previously discussed.

#### DPS and MCG.

For DPS, we use the original DPS codebase provided by [11] and pre-trained
models trained on CelebA and FFHQ training sets for natural images. For
medical images, we use the pretrained checkpoint from [12] on the dataset
provided by [41]. For MCG [12], we modified the MCG codebase by deleting the
projection term tuning the gradient term for running DPS experiments on CT
reconstruction. Otherwise, we directly used the codes provided by [12] for
both natural and medical image experiments.

#### DDRM.

For DDRM, we follow the original code provided by [13] withDDPM models trained
on FFHQ and CelebA training sets adopted from the repository provided by [7].
We use the default parameters as displayed by [13].

#### DMPS.

We follow the original code from the repository of [16] with the DDPM models
trained on FFHQ and CelebA training sets adopted from the repository of [7].
We use the default parameters as displaye d by [16].

#### PSLD-LDM.

We follow the original code from the repository provided by [22] with the
pretrained LDMs on CelebA and FFHQ datasets provided by [19]. We use the
default hyperparameters as implied in [22]. For some tasks, the
hyperparameters were tuned by grid search, and the results according to the
optimal hyperparameters are reported.

#### ADMM-PnP and Other (Supervised) Baselines.

For ADMM-PnP we use the 10 iterations with $\tau$ tuned for different inverse
problems. We use $\tau=5$ for CT reconstruction, $\tau=0.1$ for linear inverse
problems, and $\tau=0.075$ for nonlinear deblurring. We use the pre-trained
model from original DnCNN repository provided by [44]. We observe that ADMM-
PnP tends to over-smooth the images with more iterations, which causes
performance degradation. For FBP-UNet, we trained a UNet that maps FBP images
to ground truth images. The UNet network architecture is the same as the one
explained by [40].

##  Appendix D More Discussion on the Failure Cases of Latent-DPS

In this section, we provide a further explanation to which why Latent-DPS
often fails to give accurate reconstructions.

###  D.1 Failure of Measurement Consistency

Previously, we claimed that by using Latent-DPS, it is likely that we converge
to a local minimum of the function
$\|\mathbf{y}-\mathcal{A}(\mathcal{D}(\hat{\bm{z}}_{0}))\|^{2}_{2}$ and hence
cannot achieve accurate measurement consistency. Here, we validate this claim
by comparing the measurement consistency loss between Latent-DPS and ReSample.
We observe that ReSample is able to achieve better measurement consistency
than Latent-DPS. This observation validates our motivation of using hard data
consistency to improve reconstruction quality.

Anatomical Site | ReSample (Ours) | Latent-DPS
---|---|---
Chest | 5.36e-5 | 2.99e-3
Abdominal | 6.28e-5 | 3.92e-3
Head | 1.73e-5 | 2.03e-3
Table 16: Comparison of average measurement consistency loss for CT
reconstruction between ReSample and Latent-DPS

###  D.2 Failure of Linear Manifold Assumption

We hypothesize that one reason that Latent-DPS fails to give accurate
reconstructions could due to the nonlinearity of the decoder
$\mathcal{D}(\cdot)$. More specifically, the derivation of DPS provided by
[11] relies on a _linear manifold assumption_. For our case, since our forward
model can be viewed as the form $\mathcal{A}(\mathcal{D}(\cdot))$, where the
decoder $\mathcal{D}(\cdot)$ is a highly nonlinear neural network, this linear
manifold assumption fails to hold. For example, if two images $\bm{z}^{(1)}$
and $\bm{z}^{(2)}$ lie on the clean data distribution manifold $\mathcal{M}$
at $t=0$, a linear combination of them $a\bm{z}^{(1)}+b\bm{z}^{(2)}$ for some
constants $a$ and $b$, may not belong to $\mathcal{M}$ since
$\mathcal{D}(a\bm{z}^{(1)}+b\bm{z}^{(2)})$ may not give us an image that is
realistic. Thus, in practice, we observe that DPS reconstructions tend to be
more blurry, which implies that the reverse sampling path falls out of this
data manifold. We demonstrate this in Figure 20, where we show that the
average of two latent vectors gives a blurry and unrealistic image.

We would like to point out that this reasoning may also explain why our
algorithm outperforms DPS on nonlinear inverse tasks.

Figure 20: Depiction of the violation of the linear manifold assumption.

###  D.3 Inaccurate Estimates of the Posterior Mean

In this section, we discuss how inaccurate estimates of
$\hat{\bm{z}}_{0}(\bm{z}_{t})$ (i.e., the posterior mean via Tweedie’s
formula) can be one of the reasons why Latent-DPS returns image
reconstructions that are noisy. This was mainly because for values of $t$
closer to $t=T$, the estimate of $\hat{\bm{z}}_{0}(\bm{z}_{t})$ may be
inaccurate, leading us to images that are noisy at $t=0$. Generally, we find
that the estimation of $\hat{\bm{z}}_{0}(\bm{z}_{t})$ is inaccurate in the
early timesteps (e.g. $t>0.5T$) and vary a lot when $t$ decreases. This
implies that the gradient update may not point to a consistent direction when
$t$ is large. We further demonstrate this observation in Figure 21.

Figure 21: Comparison of the prediction of the ground truth signal
$\hat{\bm{z}}_{0}(\bm{z}_{t})$ for different values of $t$. Left:
$\hat{\bm{z}}_{0}(\bm{z}_{t})$ when $t=0.5T$. Right:
$\hat{\bm{z}}_{0}(\bm{z}_{t})$ when $t=0$. This serves to show the estimation
error of the posterior mean for large values of $t$ (i.e. when $t$ is closer
to pure noise).

##  Appendix E Related Works

Deep neural networks have been extensively employed as priors for solving
inverse problems [45, 29, 46, 47]. Numerous works focus on learning the
mapping between measurements and clean images, which we term as supervised
methods [45, 46, 48, 49]. Supervised methods necessitate training on pairs of
measurements and clean images, requiring model retraining for each new task.
On the other hand, another line of research aims at learning the prior
distribution of ground truth images, solving inverse problems at inference
time without retraining by using the pre-trained prior distributions [29, 30,
50, 32]. We categorize this as unsupervised methods.

Until the advent of diffusion models, unsupervised methods struggled to
achieve satisfactory reconstruction quality compared to supervised methods
[30]. These methods rely on optimization within a constrained space [29],
while constructing a space that accurately encodes the ground truth data
distribution to solve the optimization problem is challenging. However, with
the accurate approximation of the prior distribution now provided by diffusion
models, unsupervised methods can outperform supervised methods for solving
inverse problems more efficiently without retraining the model [24, 30].

For unsupervised approaches using diffusion models as priors, the plug-and-
play approach has been widely applied for solving linear inverse problems [24,
13, 28]. These methods inject the measurement on the noisy manifold or
incorporate an estimate of the ground truth image into the reverse sampling
procedure. We refer to these methods as hard data consistency approaches, as
they directly inject the measurements into the reverse sampling process.
However, to the best of our knowledge, only a few of these methods extend to
nonlinear inverse problems or even latent diffusion models [19]. Other
approaches focus on approximating the conditional score [7, 11, 12] under some
mild assumptions using gradient methods. While these methods can achieve
excellent performance and be extended to nonlinear inverse problems,
measurement consistency may be compromised, as demonstrated in our paper.
Since these methods apply data consistency only through a gradient in the
reverse sampling process, we refer to these methods as soft data consistency
approaches.

Recently, there has been a growing interest in modeling the data distribution
by diffusion models in a latent space or a constraint domain [51, 19, 52, 53].
There has also been a growing interest in solving inverse problems using
latent diffusion models. In particular, [22] proposed PSLD, an unsupervised
soft approach that adds a gradient term at each reverse sampling step to solve
linear inverse problems with LDMs. However, there are works that observe that
this approach may suffer from instability due to estimating the gradient term
[54]. Our work proposes an unsupervised hard approach that involves
“resampling” and enforcing hard data consistency for solving general inverse
problems (both linear and nonlinear) using latent diffusion models. By using a
hard data consistency approach, we can obtain much better reconstructions,
highlighting the effectiveness of our algorithm.

##  Appendix F Deferred Proofs for ReSample

Here, we provide the proofs regarding the theory behind our resampling
technique. To make this section self-contained, we first restate all of our
results and discuss notation that will be used throughout this section.

#### Notation.

We denote scalars with under-case letters (e.g. $\alpha$) and vectors with
bold under-case letters (e.g. $\bm{x}$). Recall that in the main body of the
paper, $\bm{z}\in\mathbb{R}^{k}$ denotes a sample in the latent space,
$\bm{z}^{\prime}_{t}$ denotes an unconditional sample at time step $t$,
$\hat{\bm{z}}_{0}(\bm{z}_{t})$ denotes a prediction of the ground truth signal
$\bm{z}_{0}$ at time step $t$, $\hat{\bm{z}}_{0}(\bm{y})$ denotes the
measurement-consistent sample of $\hat{\bm{z}}_{0}(\bm{z}_{t})$ using hard
data consistency, and $\hat{\bm{z}}_{t}$ denotes the re-mapped sample from
$\hat{\bm{z}}_{0}(\bm{y})$ onto the data manifold at time step $t$. We use
$\hat{\bm{z}}_{t}$ as the next sample to resume the reverse diffusion process.

######  Proposition 1 (Stochastic Encoding).

Since the sample $\hat{\bm{z}}_{t}$ given $\hat{\bm{z}}_{0}(\bm{y})$ and
measurement $\bm{y}$ is conditionally independent of $\bm{y}$, we have that

$\displaystyle
p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}))=\mathcal{N}(\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}(\bm{y}),(1-\bar{\alpha}_{t})\bm{I}).$
(21)

###### Proof.

By Tweedie’s formula, we have that $\hat{\bm{z}}_{0}(\bm{y})$ is the estimated
mean of the ground truth signal $\bm{z}_{0}$. By the forward process of the
DDPM formulation [23], we also have that

$p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}))=\mathcal{N}(\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}(\bm{y}),(1-\bar{\alpha}_{t})\bm{I}).$

Then, since $\bm{y}$ is a measurement of $\bm{z}_{0}$ at $t=0$, we have
$p(\bm{y}|\hat{\bm{z}}_{0}(\bm{y}),\hat{\bm{z}}_{t})=p(\bm{y}|\hat{\bm{z}}_{0}(\bm{y}))$.
Finally, we get

$\displaystyle p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$
$\displaystyle=\frac{p(\bm{y}|\hat{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}))p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}))p(\hat{\bm{z}}_{0}(\bm{y}))}{p(\bm{y},\hat{\bm{z}}_{0}(\bm{y}))}$
(22)
$\displaystyle=\frac{p(\bm{y},\hat{\bm{z}}_{0}(\bm{y}))p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y}))}{p(\bm{y},\hat{\bm{z}}_{0}(\bm{y}))}$
(23) $\displaystyle=p(\hat{\bm{z}}_{t}|\hat{\bm{z}}_{0}(\bm{y})).$ (24)

∎

######  Proposition 2 (Stochastic Resampling).

Suppose that
$p(\bm{z}^{\prime}_{t}|\hat{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ is
normally distributed such that
$p(\bm{z}^{\prime}_{t}|\hat{\bm{z}}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=\mathcal{N}(\mathbf{\mu}_{t},\sigma_{t}^{2})$.
If we let $p(\hat{\bm{z}_{t}}|\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ be a prior for
$\bm{\mu}_{t}$, then the posterior distribution
$p(\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})$ is
given by

$\displaystyle
p(\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y})=\mathcal{N}\left(\frac{\sigma_{t}^{2}\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}(\bm{y})+(1-\bar{\alpha}_{t})\bm{z}^{\prime}_{t}}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})},\frac{\sigma_{t}^{2}(1-\bar{\alpha}_{t})}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})}\bm{I}\right).$
(25)

###### Proof.

We have that

$\displaystyle
p(\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0},\bm{y})=\frac{p(\bm{z}^{\prime}_{t}|\hat{\bm{z}}_{t},\hat{\bm{z}}_{0},\bm{y})p(\hat{\bm{z}_{t}}|\hat{\bm{z}}_{0},\bm{y})p(\hat{\bm{z}}_{0},\bm{y})}{p(\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0},\bm{y})},$
(26)

where both $p(\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0},\bm{y})$ and
$(\hat{\bm{z}}_{0},\bm{y})$ are normalizing constants. Then by Lemma 1, we
have
$p(\hat{\bm{z}_{t}}|\hat{\bm{z}}_{0},\bm{y})=\mathcal{N}(\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0},(1-\bar{\alpha}_{t})\bm{I})$.
Now, we can compute the posterior distribution:

$\displaystyle
p(\hat{\bm{z}}_{t}=\bm{k}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0},\bm{y})$
$\displaystyle\propto
p(\bm{z}^{\prime}_{t}|\hat{\bm{z}}_{t}=\bm{k},\hat{\bm{z}}_{0},\bm{y})p(\hat{\bm{z}}_{t}=\bm{k}|\hat{\bm{z}}_{0},\bm{y})$
(27)
$\displaystyle\propto\exp{\left\\{-\left(\frac{(\bm{k}-\bm{z}^{\prime}_{t})^{2}}{\sigma_{t}^{2}}\right)+\left(\frac{(\bm{k}-\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0})^{2}}{1-\bar{\alpha}_{t}}\right)\right\\}}$
(28)
$\displaystyle\propto\exp{\left\\{\frac{-\left(\bm{k}-\frac{(1-\bar{\alpha}_{t})\bm{z}^{\prime}_{t}+\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}\sigma_{t}^{2}}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})}\right)^{2}}{\frac{1}{\sigma^{2}_{t}}+\frac{1}{1-\bar{\alpha}_{t}}}\right\\}}.$
(29)

This as a Gaussian distribution, which can be easily shown using moment-
generating functions [55]. ∎

######  Theorem 1.

If ${\hat{\bm{z}}}_{0}(\bm{y})$ is measurement-consistent such that
$\bm{y}=\mathcal{A}(\mathcal{D}({\hat{\bm{z}}}_{0}(\bm{y})))$, i.e.
$\hat{\bm{z}}_{0}=\hat{\bm{z}}_{0}(\bm{z}_{t})=\hat{\bm{z}}_{0}(\bm{y})$, then
stochastic resample is unbiased such that
$\mathbb{E}[\hat{\bm{z}}_{t}|\bm{y}]=\mathbb{E}[\bm{z}^{\prime}_{t}]$.

###### Proof.

We have that
$\mathbb{E}[\hat{\bm{z}}_{t}|\bm{y}]=\mathbb{E}_{\bm{z}^{\prime}_{t}}[\mathbb{E}_{\hat{\bm{z}}_{0}(\bm{y})}[\mathbb{E}_{\hat{\bm{z}}_{t}}[\hat{\bm{z}}_{t}|\bm{z}^{\prime}_{t},\hat{\bm{z}}_{0}(\bm{y}),\bm{y}]]]$.

Let $\gamma=\frac{\sigma_{t}^{2}}{\sigma_{t}^{2}+1-\bar{\alpha}_{t}}$. By
using Proposition 2, we have
$\mathbb{E}[\hat{\bm{z}}_{t}|\bm{y}]=\mathbb{E}_{\bm{z}^{\prime}_{t}}[\mathbb{E}_{\hat{\bm{z}}_{0}}[(\gamma\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}+(1-\gamma)\bm{z}^{\prime}_{t})|\hat{\bm{z}}_{0},\bm{z}^{\prime}_{t},\bm{y}]]$.

Since $\hat{\bm{z}}_{0}$ is measurement-consistent such that
$\bm{y}=\mathcal{A}(\mathcal{D}({\hat{\bm{z}}}_{0}(\bm{y})))$, let $k=0$, we
have

$\displaystyle\hat{\bm{z}}_{0}(\bm{y})=\hat{\bm{z}}_{0}=\frac{1}{\sqrt{\bar{\alpha}_{t-k}}}(\bm{z}^{\prime}_{t-k}-(1-\bar{\alpha}_{t-k})\nabla\log
p(\bm{z}^{\prime}_{t-k}))$ (30)

Then,we have that

$\displaystyle\mathbb{E}[\hat{\bm{z}}_{t}|\bm{y}]$
$\displaystyle=\mathbb{E}_{\bm{z}^{\prime}_{t}}[\mathbb{E}_{\hat{\bm{z}}_{0}}[(\gamma\sqrt{\bar{\alpha}_{t}}\hat{\bm{z}}_{0}+(1-\gamma)\bm{z}^{\prime}_{t})|\hat{\bm{z}}_{0},\bm{z}^{\prime}_{t},\bm{y}]]$
(31)
$\displaystyle=\gamma\sqrt{\frac{\bar{\alpha}_{t}}{\bar{\alpha}_{t-k}}}\mathbb{E}_{\bm{z}^{\prime}_{t}}[\mathbb{E}[\bm{z}^{\prime}_{t-k}-(1-\bar{\alpha}_{t-k})\nabla\log
p(\bm{z}^{\prime}_{t-k})|\bm{z}^{\prime}_{t}]]+(1-\gamma)\mathbb{E}[\bm{z}^{\prime}_{t}],$
(32)

as both $\bm{z}^{\prime}_{t}$ and $\bm{z}^{\prime}_{t-k}$ are unconditional
samples and independent of $\bm{y}$. Now, we have

$\displaystyle\mathbb{E}_{\bm{z}^{\prime}_{t}}[\mathbb{E}[\bm{z}^{\prime}_{t-k}-(1-\bar{\alpha}_{t-k})\log
p(\bm{z}^{\prime}_{t-k})|\bm{z}^{\prime}_{t}]]=\mathbb{E}_{\bm{z}^{\prime}_{t-k}}[\bm{z}^{\prime}_{t-k}-(1-\bar{\alpha}_{t-k})\nabla\log
p(\bm{z}^{\prime}_{t-k})]$ (33)

Since $\bm{z}^{\prime}_{t-k}$ is the unconditional reverse sample of
$\bm{z}^{\prime}_{t}$, we have
$\mathbb{E}[\bm{z}^{\prime}_{t}]=\sqrt{\frac{\bar{\alpha}_{t}}{\bar{\alpha}_{t-k}}}\mathbb{E}[\bm{z}^{\prime}_{t-k}]$,
and then

$\displaystyle\mathbb{E}_{\bm{z}^{\prime}_{t-k}}[\nabla\log
p(\bm{z}^{\prime}_{t-k})]$ $\displaystyle=\int\nabla\log
p(\bm{z}^{\prime}_{t-k})p(\bm{z}^{\prime}_{t-k})d\bm{z}^{\prime}_{t-k}$ (34)
$\displaystyle=\int\frac{p^{\prime}(\bm{z}^{\prime}_{t-k})}{p(\bm{z}^{\prime}_{t-k})}p(\bm{z}^{\prime}_{t-k})d\bm{z}^{\prime}_{t-k}$
(35) $\displaystyle=\frac{\partial(1)}{\partial\bm{z}^{\prime}_{t-k}}=0$ (36)

Finally, we have
$\mathbb{E}[\hat{\bm{z}}_{t}|\bm{y}]=\gamma\mathbb{E}[\bm{z}^{\prime}_{t}]+(1-\gamma)\mathbb{E}[\bm{z}^{\prime}_{t}]=\mathbb{E}[\bm{z}^{\prime}_{t}]$.
∎

######  Lemma 1.

Let $\tilde{\bm{z}}_{t}$ and $\hat{\bm{z}}_{t}$ denote the stochastically
encoded and resampled image of $\hat{\bm{z}}_{0}(\bm{y})$, respectively. If
$\text{VAR}(\bm{z}^{\prime}_{t})>0$, then we have that
$\text{VAR}(\hat{\bm{z}}_{t})<\text{VAR}(\tilde{\bm{z}}_{t})$.

###### Proof.

Recall that $\hat{\bm{z}}_{t}$ and $\tilde{\bm{z}}_{t}$ are both normally
distributed, with

$\displaystyle\text{VAR}(\hat{\bm{z}}_{t})$ $\displaystyle=1-\bar{\alpha}_{t}$
(37) $\displaystyle\text{VAR}(\tilde{\bm{z}}_{t})$
$\displaystyle=\frac{1}{\frac{1}{1-\bar{\alpha}_{t}}+\frac{1}{\sigma_{t}^{2}}}$
(38)
$\displaystyle=\frac{\sigma_{t}^{2}(1-\bar{\alpha}_{t})}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})}.$
(39)

For all $\sigma_{t}^{2}\geq 0$, we have

$\displaystyle\frac{\sigma_{t}^{2}(1-\bar{\alpha}_{t})}{\sigma_{t}^{2}+(1-\bar{\alpha}_{t})}$
$\displaystyle<1-\bar{\alpha}_{t}$ (40)
$\displaystyle\implies\text{VAR}(\hat{\bm{z}}_{t})$
$\displaystyle<\text{VAR}(\tilde{\bm{z}}_{t}).$ (41)

∎

######  Theorem 2.

Let $\bm{z}_{0}$ denote a sample from the data distribution and $\bm{z}_{t}$
be a sample from the noisy perturbed distribution at time $t$. Given that the
score function $\nabla_{\bm{z}_{t}}\log p_{\bm{z}_{t}}(\bm{z}_{t})$ is
bounded,

$\displaystyle\text{Cov}(\bm{z}_{0}|\bm{z}_{t})=\frac{(1-\bar{\alpha}_{t})^{2}}{\bar{\alpha}_{t}}\nabla^{2}_{\bm{z}_{t}}\log
p_{\bm{z}_{t}}(\bm{z}_{t})+\frac{1-\bar{\alpha}_{t}}{\bar{\alpha}_{t}}\bm{I},$

where $\alpha_{t}\in(0,1)$ is an decreasing sequence in $t$.

###### Proof.

By the LDM forward process, we have

$\displaystyle
p(\bm{z}_{t}|\bm{z}_{0})=\mathcal{N}(\sqrt{\bar{\alpha}_{t}},1-\bar{\alpha}_{t}).$

Consider the following variable:
$\hat{\bm{z}}_{t}=\frac{\sqrt{\bar{\alpha}_{t}}}{1-\bar{\alpha}_{t}}\bm{z}_{t}$,
which is simply a scaled version of $\bm{z}_{t}$. Then, consider the
distribution

$\displaystyle
p(\bar{\bm{z}}_{t}|\bm{z}_{0})=\frac{1}{\left(2\pi\left(\frac{\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}\right)\right)^{d/2}}\exp\left\\{-\frac{\|\bar{\bm{z}}_{t}-\frac{\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}\bm{z}_{0}\|_{2}^{2}}{2\left(\frac{\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}\right)}\right\\}.$

Now, we want to separate the term only with $\bar{\bm{z}}_{t}$ and the term
only with $\bm{z}_{0}$ for applying Tweedie’s formula. Hence, let
$p_{0}\left(\bar{\bm{z}}_{t}\right)=\frac{1}{\left(2\pi\left(\frac{\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}\right)\right)^{d/2}}\exp\left(-\frac{||\bar{\bm{z}}_{t}||^{2}}{2\left(\frac{\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}\right)}\right)$,
we achieve

$\displaystyle
p(\bar{\bm{z}}_{t}|\bm{z}_{0})=p_{0}(\bar{\bm{z}}_{t})\exp\left\\{{\bar{\bm{z}}_{t}}^{\top}\bm{z}_{0}-\frac{\bar{\alpha}_{t}}{2(1-\bar{\alpha}_{t})}\|\bm{z}_{0}\|^{2}\right\\},$

which separates the distribution interaction term
$\bar{\bm{z}}_{t}^{\top}\bm{z}_{0}$ from $\|\bm{z}_{0}\|^{2}$ term.
Let $\lambda(\bar{\bm{z}}_{t})=\log p(\bar{\bm{z}}_{t})-\log
p_{0}(\bar{\bm{z}}_{t})$. Then by Tweedie’s formula, we have
$\mathbb{E}[\bm{z}_{0}|\bar{\bm{z}}_{t}]=\nabla\lambda(\bar{\bm{z}}_{t})$ and
$\text{Cov}(\bm{z}_{0}|\bar{\bm{z}}_{t})=\nabla^{2}\lambda(\bar{\bm{z}}_{t})$.
Since $p_{0}(\bar{\bm{z}}_{t})$ is a Gaussian distribution with mean at $0$
and variance equal to $\frac{\bar{\alpha}_{t}}{1-\bar{\alpha_{t}}}$, we obtain

$\displaystyle\nabla\lambda(\bar{\bm{z}}_{t})=\nabla\log
p(\bar{\bm{z}}_{t})+\frac{1-\bar{\alpha}_{t}}{\bar{\alpha}_{t}}\bar{\bm{z}}_{t}.$

We observe that since $\hat{\bm{z}}$ is a scaled version of $\bm{z}_{t}$, we
can obtain the distribution of $\hat{\bm{z}}$ as

$p(\bar{\bm{z}}_{t})=\frac{1-\bar{\alpha}_{t}}{\sqrt{\bar{\alpha}_{t}}}p\left(\frac{1-\bar{\alpha}_{t}}{\sqrt{\bar{\alpha}_{t}}}\cdot\bm{z}_{t}\right).$

Then, we can apply chain rule to first compute then gradient on $\bm{z}_{t}$,
and then account for $\bm{\bar{z}}_{t}$. As a result we get

$\displaystyle\nabla_{\bar{\bm{z}}_{t}}\log
p(\bar{\bm{z}}_{t})=\frac{1-\bar{\alpha}_{t}}{\sqrt{\bar{\alpha}_{t}}}\nabla_{\bm{z_{t}}}\log
p({\bm{z}}_{t})$

and then

$\displaystyle\nabla\lambda(\bar{\bm{z}}_{t})=\mathbb{E}[\bm{z}_{0}|\bm{z}_{t}]=\frac{1-\bar{\alpha}_{t}}{\sqrt{\bar{\alpha}_{t}}}\nabla\log
p(\bm{z}_{t})+\frac{1}{\sqrt{\bar{\alpha}_{t}}}\bm{z}_{t},$

which is consistent with the score function of [11]. Afterwards, we take the
gradient again with respective to the $\bm{\bar{z}}_{t}$, and apply the chain
rule again, then we have

$\displaystyle\nabla^{2}\lambda(\bar{\bm{z}}_{t})=\text{Cov}(\bm{z}_{0}|\bm{z}_{t})=\frac{(1-\bar{\alpha}_{t})^{2}}{\bar{\alpha}_{t}}\nabla^{2}\log
p(\bm{z}_{t})+\frac{1-\bar{\alpha}_{t}}{\bar{\alpha}_{t}}\bm{I}.$

This gives the desired result. ∎