# Solving Linear Inverse Problems Provably via
Posterior Sampling with Latent Diffusion Models

Litu Rout Negin Raoof Giannis Daras
Constantine Caramanis Alexandros G. Dimakis Sanjay Shakkottai

The University of Texas at Austin
litu.rout@utexas.eduneginmr@utexas.edugiannisdaras@utexas.educonstantine@utexas.edudimakis<EMAIL_ADDRESS>
###### Abstract

We present the first framework to solve linear inverse problems leveraging
pre-trained latent diffusion models. Previously proposed algorithms (such as
DPS and DDRM) only apply to pixel-space diffusion models. We theoretically
analyze our algorithm showing provable sample recovery in a linear model
setting. The algorithmic insight obtained from our analysis extends to more
general settings often considered in practice. Experimentally, we outperform
previously proposed posterior sampling algorithms in a wide variety of
problems including random inpainting, block inpainting, denoising, deblurring,
destriping, and super-resolution.

##  1 Introduction

We study the use of pre-trained latent diffusion models to solve linear
inverse problems such as denoising, inpainting, compressed sensing and super-
resolution. There are two classes of approaches for inverse problems:
supervised methods where a restoration model is trained to solve the task at
hand [35, 37, 52, 30], and unsupervised methods that use the prior learned by
a generative model to guide the restoration process [49, 38, 5, 32, 11, 26];
see also the survey of [34] and references therein.

The second family of unsupervised methods has gained popularity because: (i)
general-domain foundation generative models have become widely available, (ii)
unsupervised methods do not require any training to solve inverse problems and
leverage the massive data and compute investment of pre-trained models and
(iii) generative models sample from the posterior-distribution, mitigating
certain pitfalls of likelihood-maximization methods such as bias in the
reconstructions [33, 24] and regression to the mean [23, 22].

Figure 1: Overall pipeline of our proposed framework from left to right. Given
an image (left) and a user defined mask (center), our algorithm inpaints the
masked region (right). The known part of the images are unaltered (see
Appendix B for web demo and image sources).

Diffusion models have emerged as a powerful new approach to generative
modeling [44, 45, 46, 20, 28, 18, 51]. This family of generative models works
by first corrupting the data distribution $p_{0}({\bm{x}}_{0})$ using an Itô
Stochastic Differential Equation (SDE),
$\mathrm{d}{\bm{x}}={\bm{f}}({\bm{x}},t)\mathrm{d}t+g(t)\mathrm{d}{\bm{w}}$,
and then by learning the score-function, $\nabla_{{\bm{x}}_{t}}\log
p_{t}({\bm{x}}_{t})$, at all levels $t$, using Denoising Score Matching (DSM)
[21, 50]. The seminal result of [1] shows that we can reverse the corruption
process, i.e., start with noise and then sample from the data distribution, by
running another Itô SDE. The SDE that corrupts the data is often termed as
Forward SDE and its reverse as Reverse SDE [46]. The latter depends on the
score-function $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ that we learn
through DSM. In [8, 9], the authors provided a non-asymptotic analysis for the
sampling of diffusion models when the score-function is only learned
approximately.

The success of diffusion models sparked the interest to investigate how we can
use them to solve inverse problems. [46] showed that given measurements
${\bm{y}}=\mathcal{A}{\bm{x}}_{0}+\sigma_{y}{\bm{n}}$, we can provably sample
from the distribution $p_{0}({\bm{x}}_{0}|{\bm{y}})$ by running a modified
Reverse SDE that depends on the unconditional score $\nabla_{{\bm{x}}_{t}}\log
p_{t}({\bm{x}}_{t})$ and the term $\nabla_{{\bm{x}}_{t}}\log
p({\bm{y}}|{\bm{x}}_{t})$. The latter term captures how much the current
iterate explains the measurements and it is intractable even for linear
inverse problems without assumptions on the distribution $p_{0}(x_{0})$ [11,
14]. To deal with the intractability of the problem, a series of approximation
algorithms have been developed [22, 11, 2, 13, 26, 10, 6, 43, 12, 27] for
solving (linear and non-linear) inverse problems with diffusion models. These
algorithms use pre-trained diffusion models as flexible priors for the data
distribution to effectively solve problems such as inpainting, deblurring,
super-resolution among others.

Recently, diffusion models have been generalized to learn to invert non-
Markovian and non-linear corruption processes [16, 15, 3]. One instance of
this generalization is the family of Latent Diffusion Models (LDMs) [39]. LDMs
project the data into some latent space,
${\bm{z}}_{0}=\mathcal{E}({\bm{x}}_{0})$, perform the diffusion in the latent
space and use a decoder, $\mathcal{D}({\bm{z}}_{0})$, to move back to the
pixel space. LDMs power state-of-the-art foundation models such as Stable
Diffusion [39] and have enabled a wide-range of applications across many data
modalities including images [39], video [4], audio [29] and medical domain
distributions (e.g., for MRI and proteins) [36, 48]. Unfortunately, none of
the existing algorithms for solving inverse problems works with Latent
Diffusion Models. Hence, to use a foundation model, such as Stable Diffusion,
for some inverse problem, one needs to perform finetuning for each task of
interest.

In this paper, we present the first framework to solve general inverse
problems with pre-trained latent diffusion models. Our main idea is to extend
DPS by adding an extra gradient update step to guide the diffusion process to
sample latents for which the decoding-encoding map is not lossy. By harnessing
the power of available foundation models, we are able to outperform previous
approaches without finetuning across a wide range of problems (see Figure 1
and 2).

Our contributions are as follows:

  * (i)

We show how to use Latent Diffusion Models models (such as Stable Diffusion)
to solve linear inverse problem when the degradation operator is known.

  * (ii)

We theoretically analyze our algorithm and show provable sample recovery in a
linear model setting with two-step diffusion processes.

  * (iii)

We achieve a new state-of-the-art for solving inverse problems with latent
diffusion models, outperforming previous approaches for inpainting, block
inpainting, denoising, deblurring, destriping, and super-resolution.111The
source code is available at: https://github.com/LituRout/PSLD and a web
application for image inpainting is available at:
https://huggingface.co/spaces/PSLD/PSLD.

##  2 Background and Method

Notation: Bold lower-case ${\bm{x}}$, bold upper-case ${\bm{X}}$, and normal
lower case $x$ denote a vector, a matrix, and a scalar variable, respectively.
We denote by $\odot$ element-wise multiplication. ${\bm{D}}({\bm{x}})$
represents a diagonal matrix with entries ${\bm{x}}$. We use $\mathcal{E}(.)$
for the encoder and $\mathcal{D}(.)$ for the decoder. $\mathcal{E}\sharp p$ is
a pushforward measure of $p$, i.e., for every ${\bm{x}}\in p$, the sample
$\mathcal{E}({\bm{x}})$ is a sample from $\mathcal{E}\sharp p$. We use arrows
in Section 3 to distinguish random variables of the forward ($\rightarrow$)
and the reverse process ($\leftarrow$).

The standard diffusion modeling framework involves training a network,
${\bm{s}}_{\theta}({\bm{x}}_{t},t)$, to learn the score-function,
$\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$, at all levels $t$, of a
stochastic process described by an Itô SDE:

$\displaystyle\mathrm{d}{\bm{x}}={\bm{f}}({\bm{x}},t)\mathrm{d}t+g(t)\mathrm{d}{\bm{w}},$
(1)

where ${\bm{w}}$ is the standard Wiener process. To generate samples from the
trained model, one can run the (unconditional) Reverse SDE, where the score-
function is approximated by the trained neural network. Given measurements
${\bm{y}}=\mathcal{A}x_{0}+\sigma_{y}{\bm{n}}$, one can sample from the
distribution $p_{0}({\bm{x}}_{0}|{\bm{y}})$ by running the conditional Reverse
SDE given by:

$\mathrm{d}{\bm{x}}=\left({\bm{f}}({\bm{x}},t)-g^{2}(t)\left(\nabla_{{\bm{x}}_{t}}\log
p_{t}({\bm{x}}_{t})+\nabla_{{\bm{x}}_{t}}\log
p({\bm{y}}|{\bm{x}}_{t})\right)\right)\mathrm{d}t+g(t)\mathrm{d}{\bm{w}}.$ (2)

As mentioned, $\nabla_{{\bm{x}}_{t}}\log p({\bm{y}}|{\bm{x}}_{t})$ is
intractable for general inverse problems. One of the most effective
approximation methods is the DPS algorithm proposed by [11]. DPS assumes that:

$\displaystyle p({\bm{y}}|{\bm{x}}_{t})\approx
p\left({\bm{y}}|{\bm{x}}_{0}=\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]\right)=\mathcal{N}({\bm{y}};\mu=\mathcal{A}\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}],\Sigma=\sigma_{y}^{2}I).$
(3)

Essentially, DPS substitutes the unknown clean image ${\bm{x}}_{0}$ with its
conditional expectation given the noisy input,
$\mathbb{E}[{\bm{x}}_{0}|{\bm{x}}_{t}]$. Under this approximation, the term
$p({\bm{y}}|{\bm{x}}_{t})$ becomes tractable.

The theoretical properties of the DPS algorithm are not well understood. In
this paper, we analyze DPS in a linear model setting where the data
distribution lives in a low-dimensional subspace, and show that DPS actually
samples from $p({\bm{x}}_{0}|{\bm{y}})$ (Section 3.2). Then, we provide an
algorithm (Section 2.1) and its analysis to sample from
$p({\bm{x}}_{0}|{\bm{y}})$ using latent diffusion models (Section 3.3).
Importantly, our analysis suggests that our algorithm enjoys the same
theoretical guarantees while avoiding the curse of ambient dimension observed
in pixel-space diffusion models including DPS. Using experiments (Section 4),
we show that our algorithm allows us to use powerful foundation models and
solve linear inverse problems, outperforming previous unsupervised approaches
without the need for finetuning.

###  2.1 Method

In Latent Diffusion Models, the diffusion occurs in the latent space.
Specifically, we train a model ${\bm{s}}_{\theta}({\bm{z}}_{t},t)$ to predict
the score $\nabla_{{\bm{z}}_{t}}\log p_{t}({\bm{z}}_{t})$, of a diffusion
process:

$\displaystyle\mathrm{d}{\bm{z}}={\bm{f}}({\bm{z}},t)\mathrm{d}t+g(t)\mathrm{d}{\bm{w}},$
(4)

where ${\bm{z}}_{0}=\mathcal{E}({\bm{x}}_{0})$ for some encoder function
$\mathcal{E}(\cdot):\mathbb{R}^{d}\to\mathbb{R}^{k}$. During sampling, we
start with ${\bm{z}}_{T}$, we run the Reverse Diffusion Process and then we
obtain a clean image by passing ${\bm{z}}_{0}\sim
p_{0}({\bm{z}}_{0}|{\bm{z}}_{T})$ through a decoder
$\mathcal{D}:\mathbb{R}^{k}\to\mathbb{R}^{d}$.

Although Latent Diffusion Models underlie some of the most powerful foundation
models for image generation, existing algorithms for solving inverse problems
with diffusion models do not apply for LDMs. The most natural extension of the
DPS idea would be to approximate $p({\bm{y}}|{\bm{z}}_{t})$ with:

$\displaystyle p({\bm{y}}|{\bm{z}}_{t})\approx
p({\bm{y}}|{\bm{x}}_{0}=\mathcal{D}\left(\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]\right)),$
(5)

i.e., to approximate the unknown clean image ${\bm{x}}_{0}$ with the decoded
version of the conditional expectation of the clean latent ${\bm{z}}_{0}$
given the noisy latent ${\bm{z}}_{t}$. However, as we show experimentally in
Section 4, this idea does not work. The failure of the “vanilla” extension of
the DPS algorithm for latent diffusion models should not come as a surprise.
The fundamental reason is that the encoder is a many-to-one mapping. Simply
put, there are many latents ${\bm{z}}_{0}$ that correspond to encoded versions
of images that explain the measurements. Taking the gradient of the density
given by (5) could be pulling ${\bm{z}}_{t}$ towards any of these latents
${\bm{z}}_{0}$, potentially in different directions. On the other hand, the
score-function is pulling ${\bm{z}}_{t}$ towards a specific ${\bm{z}}_{0}$
that corresponds to the best denoised version of ${\bm{z}}_{t}$.

To address this problem, we propose an extra term that penalizes latents that
are not fixed-points of the composition of the decoder-function with the
encoder-function. Specifically, we approximate the intractable $\nabla\log
p({\bm{y}}|{\bm{z}}_{t})$ with:

$\displaystyle\nabla_{{\bm{z}}_{t}}\log
p({\bm{y}}|{\bm{z}}_{t})=\underbrace{\nabla_{{\bm{z}}_{t}}p({\bm{y}}|{\bm{x}}_{0}=\mathcal{D}\left(\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]\right))}_{\mathrm{DPS\
vanilla\
extension}}+\gamma_{t}\underbrace{\nabla_{z_{t}}\left|\left|\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]-\mathcal{E}(\mathcal{D}(\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]))\right|\right|^{2}}_{\mathrm{``goodness"\
of}\ {\bm{z}}_{0}}.$ (6)

We refer to this approximation as Goodness Modified Latent DPS (GML-DPS).
Intuitively, we guide the diffusion process towards latents such that: i) they
explain the measurements when passed through the decoder, and ii) they are
fixed points of the decoder-encoder composition. The latter is useful to make
sure that the generated sample remains on the manifold of real data. However,
it does not penalize the reverse SDE for generating other latents
${\bm{z}}_{0}$ as long as $\mathcal{D}({\bm{z}}_{0})$ lies on the manifold of
natural images. Even in the linear case (see Section 3), this can lead to
inconsistency at the boundary of the mask in the pixel space. The linear
theory in Section 3 suggests that we can circumvent this problem by
introducing the following gluing objective. In words, the gluing objective
penalizes decoded images having a discontinuity at the boundary of the mask.

$\displaystyle\nabla_{{\bm{z}}_{t}}\log p({\bm{y}}|{\bm{z}}_{t})$
$\displaystyle=\underbrace{\nabla_{{\bm{z}}_{t}}p({\bm{y}}|{\bm{x}}_{0}=\mathcal{D}\left(\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]\right))}_{\mathrm{DPS\
vanilla\ extension}}$
$\displaystyle+\gamma_{t}\underbrace{\nabla_{z_{t}}\left|\left|\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]-\mathcal{E}({\mathcal{A}}^{T}{\mathcal{A}}{\bm{x}}_{0}^{*}+({\bm{I}}-{\mathcal{A}}^{T}{\mathcal{A}})\mathcal{D}(\mathbb{E}[{\bm{z}}_{0}|{\bm{z}}_{t}]))\right|\right|^{2}}_{\mathrm{``gluing"\
of}\ {\bm{z}}_{0}}.$ (7)

The gluing objective is critical for our algorithm as it ensures that the
denoising update, measurement-matching update, and the gluing update point to
the same optima in the latent space. We refer to this approximation (7) as
Posterior Sampling with Latent Diffusion (PSLD). In the next Section 3, we
provide an analysis of these gradient updates, along with the associated
algorithms.

##  3 Theoretical Results

Input: $T$, $\bm{y}$, $\zeta_{i=1}^{T}$,
${\\{\tilde{\sigma}_{i}\\}_{i=1}^{T}},\bm{s}_{\theta}$

1 $\bm{x}_{T}\sim\mathcal{T}(\bm{0},\bm{I})$

2 for  _$i=T-1$ to $0$_ do

3 $\hat{\bm{s}}\leftarrow\bm{s}_{\theta}(\bm{x}_{i},i)$

4
$\hat{{\bm{x}}}_{0}\leftarrow\frac{1}{\sqrt{\bar{\alpha}_{i}}}({\bm{x}}_{i}+{(1-\bar{\alpha}_{i})\hat{\bm{s}})}$

5 ${\bm{z}}\sim\mathcal{N}(\bm{0},\bm{I})$

6
${\bm{x}}^{\prime}_{i-1}\leftarrow\frac{\sqrt{\alpha_{i}}(1-\bar{\alpha}_{i-1})}{1-\bar{\alpha}_{i}}{\bm{x}}_{i}+\frac{\sqrt{\bar{\alpha}_{i-1}}\beta_{i}}{1-\bar{\alpha}_{i}}\hat{{\bm{x}}}_{0}+{\tilde{\sigma}_{i}{\bm{z}}}$

7
${\bm{x}}_{i-1}\leftarrow{\bm{x}}^{\prime}_{i-1}-{\zeta_{i}}\nabla_{{\bm{x}}_{i}}\|{\bm{y}}-\mathcal{A}(\hat{{\bm{x}}}_{0})\|_{2}^{2}$

8 end for

9return $\hat{{\bm{x}}}_{0}$

Algorithm 1 DPS

Input: $T$, $\bm{y}$, {$\eta_{i}\\}_{i=1}^{T}$,
{$\gamma_{i}\\}_{i=1}^{T},{\\{\tilde{\sigma}_{i}\\}_{i=1}^{T}},\mathcal{E},\mathcal{D},{\mathcal{A}}{\bm{x}}^{*}_{0},{\mathcal{A}},\bm{s}_{\theta}$

1 $\bm{z}_{T}\sim\mathcal{N}(\bm{0},\bm{I})$

2 for  _$i=T-1$ to $0$_ do

3 $\hat{\bm{s}}\leftarrow\bm{s}_{\theta}(\bm{z}_{i},i)$

4
$\hat{{\bm{z}}}_{0}\leftarrow\frac{1}{\sqrt{\bar{\alpha}_{i}}}({\bm{z}}_{i}+{(1-\bar{\alpha}_{i})\hat{\bm{s}})}$

5 $\bm{\bm{\epsilon}}\sim\mathcal{N}(\bm{0},\bm{I})$

6
${\bm{z}}^{\prime}_{i-1}\leftarrow\frac{\sqrt{\alpha_{i}}(1-\bar{\alpha}_{i-1})}{1-\bar{\alpha}_{i}}{\bm{z}}_{i}+\frac{\sqrt{\bar{\alpha}_{i-1}}\beta_{i}}{1-\bar{\alpha}_{i}}\hat{{\bm{z}}}_{0}+{\tilde{\sigma}_{i}\bm{\bm{\epsilon}}}$

7
${\bm{z}}^{\prime\prime}_{i-1}\leftarrow{\bm{z}}^{\prime}_{i-1}-{\eta_{i}}\nabla_{{\bm{z}}_{i}}\|{\bm{y}}-\mathcal{A}(\mathcal{D}{(\hat{{\bm{z}}}_{0})})\|_{2}^{2}$

8
${\bm{z}}_{i-1}\leftarrow{\bm{z}}^{\prime\prime}_{i-1}-{\gamma_{i}}\nabla_{{\bm{z}}_{i}}\|\hat{\bm{z}}_{0}-\mathcal{E}({{\mathcal{A}}^{T}{\mathcal{A}}{\bm{x}}^{*}_{0}+({\bm{I}}-{\mathcal{A}}^{T}{\mathcal{A}})\mathcal{D}(\hat{\bm{z}}_{0}))}\|_{2}^{2}$

9 end for

return $\mathcal{D}(\hat{{\bm{z}}}_{0})$

Algorithm 2 PSLD

As discussed in Section 2, diffusion models consist of two stochastic
processes: the forward and reverse processes, each governed by Itô SDEs. For
implementation purposes, these SDEs are discretized over a finite number of
(time) steps, and the diffusion takes place using a transition kernel. The
forward process starts from $\overrightarrow{{\bm{x}}_{0}}\sim
p(\overrightarrow{{\bm{x}}_{0}})$ and gradually adds noise, i.e.,
$\overrightarrow{{\bm{x}}}_{t+1}=\sqrt{1-\beta_{t}}\overrightarrow{{\bm{x}}}_{t}+\sqrt{\beta_{t}}\bm{\epsilon}$
where $\beta_{t}\in[0,1]$ and $\beta_{t}\geq\beta_{t-1}$ for $t=0,\dots,T-1$ .
The reverse process is initialized with
$\overleftarrow{{\bm{x}}}_{T}\sim{\mathcal{N}}\left(\mathbf{0},{\bm{I}}_{d}\right)$
and generates
$\overleftarrow{{\bm{x}}}_{t-1}=\mu_{\theta}(\overleftarrow{{\bm{x}}}_{t},t)+\sqrt{\beta_{t}}\bm{\epsilon}$.
In the last step, $\mu_{\theta}(\overleftarrow{{\bm{x}}}_{1},1)$ is displayed
without the noise.

In this section, we consider the diffusion discretized to two steps
($\\{\overrightarrow{{\bm{x}}_{0}},\overrightarrow{{\bm{x}}_{1}}\\}$), and a
Gaussian transition kernel that arises from the Ornstein-Uhlenbeck (OU)
process. We choose this setup because it captures essential components of
complex diffusion processes without raising unnecessary complications in the
analysis. We provide a principled analysis of Algorithm 1 and Algorithm 2 in a
linear model setting with this two-step diffusion process under assumptions
that guarantee exact reconstruction is possible in principle. A main result of
our work is to prove that in this setting we can solve inverse problems
perfectly. As we show, this requires some novel algorithmic ideas that are
suggested by our theory. In Section 4, we then show that these algorithmic
ideas are much more general, and apply to large-scale real-world applications
of diffusion models that use multiple steps
($\\{\overrightarrow{{\bm{x}}_{0}},\overrightarrow{{\bm{x}}_{1}},\cdots,\overrightarrow{{\bm{x}}_{T}}\\}$,
where $T=1000$), and moreover do not satisfy the recoverability assumptions.
We provide post-processing details of Algorithm 2 in Appendix B.1. All proofs
are given in Appendix A.

###  3.1 Problem Setup

The goal is to show that posterior sampling algorithms (such as DPS) can
provably solve inverse problems in a perfectly recoverable setting. To show
exact recovery, we analyze two-step diffusion processes in a linear model
setting similar to [40, 7], where the images
($\overrightarrow{{\bm{x}}_{0}}\in\mathbb{R}^{d}$) reside in a linear subspace
of the form
$\overrightarrow{{\bm{x}}_{0}}={\mathcal{S}}\overrightarrow{{\bm{w}}_{0}},{\mathcal{S}}\in\mathbb{R}^{d\times
l},\overrightarrow{{\bm{w}}_{0}}\in\mathbb{R}^{l}$. Here, ${\mathcal{S}}$ is a
tall thin matrix with $rank({\mathcal{S}})=l\leq d$ that lifts any latent
vector
$\overrightarrow{{\bm{w}}_{0}}\sim\mathcal{N}\left(\mathbf{0},{\bm{I}}_{l}\right)$
to the image space with ambient dimension $d$. Given the measurements
${\bm{y}}={\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}+\sigma_{y}{\bm{n}}$,
${\mathcal{A}}\in\mathbb{R}^{l\times d},{\bm{n}}\in\mathbb{R}^{l}$, the goal
is to sample from $p_{0}(\overrightarrow{{\bm{x}}_{0}}|{\bm{y}})$ using a pre-
trained latent diffusion model. In the inpainting task, the measurement
operator ${\mathcal{A}}$ is such that ${\mathcal{A}}^{T}{\mathcal{A}}$ is a
diagonal matrix ${\bm{D}}({\bm{m}})$, where ${\bm{m}}$ is the masking vector
with elements set to 1 where data is observed and 0 where data is masked (see
Appendix A for further details). Recall that in latent diffusion models, the
diffusion takes place in the latent space of a pre-trained Variational
Autoencoder (VAE). Following the common practice [39], we consider a setting
where the latent vector of the VAE is $k$-dimensional and the latent
distribution is a standard Gaussian
${\mathcal{N}}\left(\mathbf{0},{\bm{I}}_{k}\right)$. Our analysis shows that
the proposed Algorithm 2 provably solves inverse problems under the following
assumptions.

######  Assumption 3.1.

The columns of the data generating model ${\mathcal{S}}$ are orthonormal,
i.e., ${\mathcal{S}}^{T}{\mathcal{S}}={\bm{I}}_{l}$.

######  Assumption 3.2.

The measurement operator ${\mathcal{A}}$ satisfies
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})\succ\mathbf{0}$.

These assumptions have previously appeared, e.g., [40]. While Assumption 3.1
is mild and can be relaxed at the expense of (standard) mathematical
complications, Assumption 3.2 indicates that
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$ is a positive
definite matrix. The latter ensures that there is enough energy left in the
measurements for perfect reconstruction. More precisely, any subset of $l$
coordinates exactly determines the remaining $(d-l)$ coordinates of
$\overrightarrow{{\bm{x}}_{0}}$. The underlying assumption is that there
exists a solution and it is unique [40]. Thus, the theoretical question
becomes how close the recovered sample is to this groundtruth sample from the
true posterior. Alternatively, one may consider other types of posteriors and
prove that the generated samples are close to this posterior in distribution.
However, this does not guarantee that the exact groundtruth sample is
recovered. Therefore, motivated by prior works [40, 7], we analyze posterior
sampling in a two-step diffusion model and answer a fundamental question: Can
a pre-trained latent diffusion model provably solve inverse problems in a
perfectly recoverable setting?

###  3.2 Posterior Sampling using Pixel-space Diffusion Model

We first consider the reverse process, starting with
$\overleftarrow{{\bm{x}}_{1}}\sim{\mathcal{N}}\left(\mathbf{0},{\bm{I}}_{d}\right)$,
and borrow a result from [40] to show that the sample
$\overleftarrow{{\bm{x}}_{0}}$ generated by the reverse process is a valid
image from $p(\overrightarrow{{\bm{x}}_{0}})$.

######  Theorem 3.3 (Generative Modeling using Diffusion in Pixel Space,
[40]).

Suppose Assumption 3.1 holds. Let

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{x}}_{1}}(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{x}}_{0}}\right)-\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{x}}_{1}}\left(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

For a fixed variance $\beta>0$, if
$\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{x}}_{1}}\left(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\coloneqq{\bm{\theta}}\overrightarrow{{\bm{x}}_{1}}\left(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}\right)$,
then the closed-form solution ${\bm{\theta}}^{*}$ is
$\sqrt{1-\beta}{\bm{S}}{\bm{S}}^{T}$, which after normalization by
$1/\sqrt{1-\beta}$ recovers the true subspace of
$p\left(\overrightarrow{{\bm{x}}_{0}}\right)$.

Though this establishes that $\overleftarrow{{\bm{x}}_{0}}$ generated by the
reverse process is a valid image from $p(\overrightarrow{{\bm{x}}_{0}})$, it
is not necessarily a sample from the posterior
$p(\overrightarrow{{\bm{x}}_{0}}|{\bm{y}})$ that satisfies the measurements.
To accomplish this we perform one additional step of gradient descent for
every step of the reverse process. This gives us Algorithm 1, the DPS
algorithm. The next theorem shows that the reverse SDE guided by these
measurements (3) recovers the true underlying sample222While the DPS Algorithm
[11] uses a scalar step size $\zeta_{i}$ at each step, this does not suffice
for exact recovery. However, by generalizing to allow a different step size
per coordinate, we can show sample recovery. Thus, in this section, we denote
$\zeta_{i}^{j}$ to be the step size at step $i$ and coordinate $j$, $1\leq
j\leq r$. Also note that the step index $i$ is vacuous in this section, as we
consider a two-step diffusion process (i.e., $i$ is always ’1’)..

######  Theorem 3.4 (Posterior Sampling using Diffusion in Pixel Space).

Suppose Assumption 3.1 and Assumption 3.2 hold. Let us denote by
$\sigma_{j},\forall j=1,\dots,r$, the singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$ and

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{x}}_{1}}(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{x}}_{0}}\right)-\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{x}}_{1}}\left(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

Given a partially known image $\overrightarrow{{\bm{x}}_{0}}\sim
p(\overrightarrow{{\bm{x}}_{0}})$, a fixed variance $\beta>0$, there exists a
step size $\zeta_{i}^{j}=1/2\sigma_{j}$ for all the coordinates of
$\overrightarrow{{\bm{x}}_{0}}$ such that Algorithm 1 samples from the true
posterior $p(\overrightarrow{{\bm{x}}_{0}}|y)$ and exactly recovers the
groundtruth sample, i.e.,
$\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

###  3.3 Posterior Sampling using Latent Diffusion Model

In this section, we analyze two approximations: GML-DPS based on (6), and PSLD
based on (7), displayed in Algorithm 2. We consider the case where the latent
distribution of the VAE is in the same space as the latent distribution of the
data generating model, i.e., $k=l$, and normalize $\gamma_{i}=1$ (as this is
immaterial in the linear setting). In Proposition 3.5, we provide analytical
solutions for the encoder and the decoder of the VAE.

######  Proposition 3.5 (Variational Autoencoder).

Suppose Assumption 3.1 holds. For an encoder
$\mathcal{E}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ and a decoder
$\mathcal{D}:\mathbb{R}^{k}\rightarrow\mathbb{R}^{d}$, denote by
$\mathcal{L}\left(\phi,\omega\right)$ the training objective of VAE:

$\displaystyle\arg\min_{\phi,\omega}\mathcal{L}\left(\phi,\omega\right)\coloneqq\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}}\sim
p}\left[\left\|\mathcal{D}(\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi);\omega)-\overrightarrow{{\bm{x}}_{0}}\right\|_{2}^{2}\right]+\lambda
KL\left(\mathcal{E}\sharp p,\mathcal{N}(\mathbf{0},{\bm{I}}_{k})\right),$

then the combination of
$\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi)={\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}}$
and
$\mathcal{D}(\overleftarrow{{\bm{z}}_{0}};\omega)={\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$
is a minimizer of $\mathcal{L}\left(\phi,\omega\right)$.

Using the encoder
$\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi)={\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}}$,
we can use the analytical solution ${\bm{\theta}}^{*}$ of the LDM obtained in
Theorem 3.3. To verify that ${\bm{\theta}}^{*}$ recovers the true subspace
$p\left(\overrightarrow{{\bm{x}}_{0}}\right)$, we compose the decoder
$\mathcal{D}(\overleftarrow{{\bm{z}}_{0}};\omega)={\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$
with the generator of the LDM, i.e.,
$\overleftarrow{{\bm{x}}_{0}}=\mathcal{D}\left({\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}\right)=\mathcal{D}\left({\bm{I}}_{k}\overleftarrow{{\bm{z}}_{1}}\right)={\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}$.
Since
$\overleftarrow{{\bm{z}}_{1}}\sim{\mathcal{N}}\left(\mathbf{0},{\bm{I}}_{k}\right)$
and ${\mathcal{S}}$ is the data generating model, this shows that
$\overleftarrow{{\bm{x}}_{0}}$ is a sample from
$p(\overrightarrow{{\bm{x}}_{0}})$. Thus we have the following.

######  Theorem 3.6 (Generative Modeling using Diffusion in Latent Space).

Suppose Assumption 3.1 holds. Let the optimal solution of the latent diffusion
model be

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{z}}_{1}}(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

For a fixed variance $\beta>0$, if
$\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\coloneqq{\bm{\theta}}\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)$,
then the closed-form solution is
${\bm{\theta}}^{*}=\sqrt{1-\beta}{\bm{I}}_{k}$, which after normalization by
$\frac{1}{\sqrt{1-\beta}}$ and composition with the decoder
${\mathcal{D}}\left(\overleftarrow{{\bm{z}}_{0}};\omega\right)={\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$
recovers the true subspace of $p\left(\overrightarrow{{\bm{x}}_{0}}\right)$.

With this optimal ${\bm{\theta}}^{*}$, we can now prove exact sample recovery
using GML-DPS (6).

######  Theorem 3.7 (Posterior Sampling using Goodness Modified Latent DPS).

Let Assumptions 3.1 and 3.2 hold. Let $\sigma_{j},\forall j=1,\dots,r$, denote
the singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$, and let

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{z}}_{1}}(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

Given a partially known image $\overrightarrow{{\bm{x}}_{0}}\sim
p(\overrightarrow{{\bm{x}}_{0}})$, any fixed variance $\beta\in(0,1)$, then
with the (unique) step size $\eta_{i}^{j}=1/2\sigma_{j},j=1,2,\ldots,r$, the
GML-DPS Algorithm (6) samples from the true posterior
$p(\overrightarrow{{\bm{x}}_{0}}|y)$ and exactly recovers the groundtruth
sample, i.e., $\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

Theorem 3.7 shows that GML-DPS (6) recovers the true sample using an LDM. This
approach, however, requires the step size $\eta$ to be chosen coordinate-wise
in a specific manner. Also, multiple natural images could have the same
measurements in the pixel space. This is a reasonable concern for LDMs due to
one-to-many mappings of the decoder. Note that the goodness objective (Section
2.1) cannot help in this scenario because it assigns uniform probability to
many of these latents $\overleftarrow{{\bm{z}}_{1}}$ for which
$\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left|\left|\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}})]-\mathcal{E}(\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}})))\right|\right|^{2}=0$.
These challenges motivate the gluing objective in Theorem 3.8. This is crucial
for two reasons. First, we show that it helps recover the true sample even
when the step size $\eta$ is chosen arbitrarily. Second, it assigns all the
probability mass to the desired (unique) solution in the pixel space.

######  Theorem 3.8 (Posterior Sampling using Diffusion in Latent Space).

Let Assumptions 3.1 and 3.2 hold. Let $\sigma_{j},\forall j=1,\dots,r$ denote
the singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$ and let

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{z}}_{1}}(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

Given a partially known image $\overrightarrow{{\bm{x}}_{0}}\sim
p(\overrightarrow{{\bm{x}}_{0}})$, any fixed variance $\beta\in(0,1)$, and any
positive step sizes $\eta_{i}^{j},j=1,2,\ldots,r$, the PSLD Algorithm 2
samples from the true posterior $p(\overrightarrow{{\bm{x}}_{0}}|y)$ and
exactly recovers the groundtruth sample, i.e.,
$\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

The important distinction between Theorem 3.7 and Theorem 3.8 is that the
former requires the exact step size while the latter works for any finite step
size. Combining denoising, measurement-consistency (with a scalar $\eta$), and
gluing updates, we have

$\displaystyle\overleftarrow{{\bm{z}}_{0}}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}-\eta\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}}))-{\bm{y}}\right\|_{2}^{2}-\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}})-\mathcal{E}({\mathcal{A}}^{T}{\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}+({\bm{I}}_{d}-{\mathcal{A}}^{T}{\mathcal{A}})\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}})))\right\|_{2}^{2}.$

When $\eta$ is chosen arbitrarily, then the third term guides the reverse SDE
towards the optimal solution $\overrightarrow{{\bm{z}}_{0}}$. When the reverse
SDE generates the exact same groundtruth sample, i.e.,
$\mathcal{D}(\overleftarrow{{\bm{z}}_{1}}(\overleftarrow{{\bm{z}}_{0}}))=\overrightarrow{{\bm{x}}_{0}}$,
then the third term becomes zero. For all other samples, it penalizes the
reverse SDE. Thus, it forces the reverse SDE to recover the true underlying
sample irrespective of the value of $\eta$.

We draw the following key insights from our Theorem 3.8: Curse of ambient
dimension: In order to run posterior sampling using diffusion in the pixel
space, the gradient of the measurement error needs to be computed in the
$d$-dimensional ambient space. Therefore, DPS algorithm suffers from the curse
of ambient dimension. On the other hand, our algorithm uses diffusion in the
latent space, and therefore avoids the curse of ambient dimension. Large-scale
foundation model: We propose a posterior sampling algorithm which offers the
provision to use large-scale foundation models, and it provably solves general
linear inverse problems. Robustness to measurement step: The gluing objective
makes our algorithm robust to the choice of step size $\eta$. Furthermore, it
allows the same (scalar) step size across all the coordinates of
$\overrightarrow{{\bm{x}}_{0}}$.

##  4 Experimental Evaluation

We experiment with in-distribution and out-of-distribution datasets. For in-
distribution, we conduct our experiments on a subset of the FFHQ dataset [25]
(downscaled to $256\times
256$333https://www.kaggle.com/datasets/denislukovnikov/ffhq256-images-only,
denoted by FFHQ 256). For out-of-distribution, we use images from the web and
ImageNet dataset [17] (resized to $256\times 256$, denoted by ImageNet 256).
To make a fair comparison, we use the same validation subset and follow the
same masking strategy as the baseline DPS [11]. It is important to note that
our main contribution is an algorithm that can leverage any latent diffusion
model. We test our algorithm with two pre-trained latent diffusion models: (i)
the Stable Diffusion model that is trained on multiple subsets of the LAION
dataset [41, 42]; and (ii) the Latent Diffusion model (LDM-VQ-4) trained on
the FFHQ $256$ dataset [39]. The DPS model is similarly trained from scratch
for 1M steps using 49k FFHQ $256$ images, which excludes the first 1K images
used as validation set.

Inverse Problems. We experiment with the following task-specific measurement
operators from the baseline DPS [11]:  (i) Box inpainting uses a mask of size
128×128 at the center.  (ii) Random inpainting chooses a drop probability
uniformly at random between $(0.2,0.8)$ and applies this drop probability to
all the pixels.  (iii) Super-resolution downsamples images at $4\times$ scale.
(iv) Gaussian blur convolves images with a Gaussian blur kernel.  (v) Motion
blur convolves images with a motion blur kernel. We also experiment with these
additional operators from RePaint [31]:  (vi) Super-resolution downsamples
images at $2\times$, $3\times$, and $4\times$ scale.  (vii) Denoising has
Gaussian noise with $\sigma=0.05$.  (viii) Destriping has vertical and
horizontal stripes in the input images.

Table 1: Quantitative inpainting results on FFHQ $256$ validation set [25,
11]. We use Stable Diffusion v-1.5 and the measurement operators as in DPS
[11]. As shown, our PSLD model outperforms DPS since it is able to leverage
the power of the Stable Diffusion foundation model.

| Inpaint (random) | Inpaint (box) | SR ($4\times$) | Gaussian Deblur |
---|---|---|---|---|---
Method | FID ($\downarrow$) | LPIPS ($\downarrow$) | FID ($\downarrow$) | LPIPS ($\downarrow$) | FID ($\downarrow$) | LPIPS ($\downarrow$) | FID ($\downarrow$) | LPIPS ($\downarrow$) |
PSLD (Ours) | 21.34 | 0.096 | 43.11 | 0.167 | 34.28 | 0.201 | 41.53 | 0.221 |
DPS [11] | 33.48 | 0.212 | 35.14 | 0.216 | 39.35 | 0.214 | 44.05 | 0.257 |
DDRM [26] | 69.71 | 0.587 | 42.93 | 0.204 | 62.15 | 0.294 | 74.92 | 0.332 |
MCG [13] | 29.26 | 0.286 | 40.11 | 0.309 | 87.64 | 0.520 | 101.2 | 0.340 |
PnP-ADMM [6] | 123.6 | 0.692 | 151.9 | 0.406 | 66.52 | 0.353 | 90.42 | 0.441 |
Score-SDE [47] | 76.54 | 0.612 | 60.06 | 0.331 | 96.72 | 0.563 | 109.0 | 0.403 |
ADMM-TV | 181.5 | 0.463 | 68.94 | 0.322 | 110.6 | 0.428 | 186.7 | 0.507 |

Table 2: Quantitative super-resolution (using measurement operator from [31]) results on FFHQ $256$ validation samples [25, 11]. We use PSLD with Stable Diffusion. Table shows LPIPS ($\downarrow$).  Method | PSLD (Ours) | DPS [11]
---|---|---
$2\times$ | 0.185 | 0.220
$3\times$ | 0.220 | 0.247
$4\times$ | 0.233 | 0.291

Evaluation. We compare the performance of our PSLD algorithm with the state-
of-the-art DPS algorithm [11] on random inpainting, box inpainting, denoising,
Gaussian deblur, motion deblur, arbitrary masking, and super-resolution tasks.
We show that PSLD outperforms DPS, both in-distribution and out-of-
distribution datasets, using the Stable Diffusion v-1.5 model pre-trained on
the LAION dataset. We also test PSLD with LDM-VQ-4 trained on FFHQ $256$, to
compare with DPS trained on the same data distribution. Note that the LDM-v4
is a latent-based model released prior to Stable Diffusion. Therefore, it does
not match the performance of Stable Diffusion in solving inverse problems.
However, it shows the general applicability of our framework to leverage an
LDM in posterior sampling. Since Stable Diffusion v-1.5 is trained with an
image resolution of $512\times 512$, we apply the forward operator after
upsampling inputs to $512\times 512$, run posterior sampling at $512\times
512$, and then downsample images to the original $256\times 256$ resolution
for a fair comparison with DPS. We observed a similar performance while
applying the masking operator at $256\times 256$ and upscaling to $512\times
512$ before running PSLD. More implementation details are provided in Appendix
B.1.

Metrics.  We use the commonly used Learned Perceptual Image Patch Similarity
(LPIPS), Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index Metric
(SSIM), and Fréchet Inception Distance444https://github.com/mseitzer/pytorch-
fid (FID) metrics for quantitative evaluation.

Results. Figure 2 shows the inpainting results on out-of-distribution samples.
This experiment was performed on commercial platforms that use (to the best of
our knowledge) Stable diffusion and additional proprietary models. This
evaluation was performed on models deployed in May 2023 and may change as
commercial providers improve their platforms.

(a) Input

(b) Groundtruth

(c) Comm. Serv. 1

(d) Comm. Serv. 2

(e) PSLD (Ours)

Figure 2: Inpainting results in general domain images from the web (see
Appendix B for image sources). Our model compared to state-of-art commercial
inpainting services that leverage the same foundation model (Stable Diffusion
v-1.5).

The qualitative advantage of PSLD is clearly demonstrated in Figures 2, 3, 4,
15 and 16. In Figure 5, we compare PSLD and DPS in random inpainting task for
varying percentage of dropped pixels. Quantitatively, PSLD outperforms DPS in
commonly used metrics: LPIPS, PSNR, and SSIM.

In our PSLD algorithm, we use Stable Diffusion v1.5 model and (zero-shot) test
it on inverse problems. Table 1 compares the quantitative results of PSLD with
related works on random inpainting, box inpainting, super-resolution, and
Gaussian deblur tasks. PSLD significantly outperforms previous approaches on
the relatively easier random inpainting task, and it is better or comparable
on harder tasks. Table 4 draws a comparison between PSLD and the strongest
baseline (among the compared methods) on out-of-distribution images. Table 2
shows the super-resolution results using nearest-neighbor kernels from [31] on
FFHQ 256 validation dataset. Observe that PSLD outperforms state-of-the-art
methods across diverse tasks and standard evaluation metrics.

In Table 3, we compare PSLD (using LDM-VQ-4) and DPS on random and box
inpainting tasks with the same operating resolution ($256\times 256$) and
training distributions (FFHQ 256). Although the LDM model exceeds DPS
performance in box inpainting, it is comparable in random inpainting. As
expected, using a more powerful pre-trained model such as Stable Diffusion is
beneficial in reconstruction–see Table 1. This highlights the significance of
our PSLD algorithm that has the provision to incorporate a powerful foundation
model with no extra training costs for solving inverse problems. Importantly,
PSLD uses latent-based diffusion, and thus it avoids the curse of ambient
dimension (Theorem 3.8), while still achieving comparable results to the
state-of-the-art method DPS [11] that has been trained on the same dataset.
Additional experimental evaluation is provided in Appendix B.

Table 3: Quantitative inpainting results on FFHQ $256$ validation set [25, 11]. We use the latent diffusion (LDM-VQ-4) trained on FFHQ $256$. Note that in this experiment PSLD and DPS use diffusion models trained on the same dataset. As shown, PSLD with LDM-VQ-4 as diffusion model outperforms DPS in box inpainting and has comparable performance in random inpainting.  | Inpaint (random) | Inpaint (box)
---|---|---
Method |  PSNR ($\uparrow$) |  SSIM ($\uparrow$) |  LPIPS ($\downarrow$) |  PSNR ($\uparrow$) |  SSIM ($\uparrow$) |  LPIPS ($\downarrow$)
PSLD (Ours) | 30.31 | 0.851 | 0.221 | 24.22 | 0.819 | 0.158
DPS [11] | 29.49 | 0.844 | 0.212 | 23.39 | 0.798 | 0.214
Table 4: Quantitative results of random inpainting and denoising on FFHQ $256$ [25, 11] using Stable Diffusion v-1.5. Note that DPS is trained on FFHQ $256$. The results show that our method PSLD generalizes well to out-of-distribution samples even without finetuning.  |  Random inpaint + denoise $\sigma=0.00$ |  Random inpaint + denoise $\sigma=0.05$
---|---|---
Method |  PSNR ($\uparrow$) |  SSIM ($\uparrow$) |  LPIPS ($\downarrow$) |  PSNR ($\uparrow$) |  SSIM ($\uparrow$) |  LPIPS ($\downarrow$)
PSLD (Ours) | 34.02 | 0.951 | 0.083 | 33.71 | 0.943 | 0.096
DPS [11] | 31.41 | 0.884 | 0.171 | 29.49 | 0.844 | 0.212
Figure 3: Left panel: Random Inpainting on images from FFHQ 256 [25] using
PSLD with Stable Diffusion v-1.5. Notice the text in the top row and the
facial expression in the bottom row.  Right panel: Block ($128\times 128$)
inpainting, using the LDM-VQ-4 model trained on FFHQ $256$ [25]. Notice the
glasses in the top row and eyes in the bottom row.

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 4:  Inpainting (random and box) results on out-of-distribution samples,
$256\times 256$ (see Appendix B for image sources). We use PSLD with Stable
Diffusion v-1.5 as generative foundation model.

Figure 5:  Comparing DPS and PSLD performance in random inpainting on FFHQ 256
[25, 11], as the percentage of masked pixels increases. PSLD with Stable
Diffusion outperforms DPS.

##  5 Conclusion

In this paper, we leverage latent diffusion models to solve general linear
inverse problems. While previously proposed approaches only apply to pixel-
space diffusion models, our algorithm allows us to use the image prior learned
by latent-based foundation generative models. We provide a principled analysis
of our algorithm in a linear two-step diffusion setting, and use insights from
this analysis to design a modified objective (goodness and gluing). This leads
to our algorithm – Posterior Sampling with Latent Diffusion (PSLD) – that
experimentally outperforms state-of-art baselines on a wide variety of tasks
including random inpainting, block inpainting, denoising, destriping, and
super-resolution.

Limitations. Our evaluation is based on Stable Diffusion which was trained on
the LAION dataset. Biases in this dataset and foundation model will be
implicitly affecting our algorithm. Our method can work with any LDM and we
expect new foundation models trained on better datasets like [19] to mitigate
these issues. Second, we have not explored how to use latent-based foundation
models to solve non-linear inverse problems. Our method builds on the DPS
approximation (which performs well on non-linear inverse problems), and hence
we believe our method can also be similarly extended.

## Acknowledgements

This research has been supported by NSF Grants 2019844, 2112471, AF 1901292,
CNS 2148141, Tripods CCF 1934932, the Texas Advanced Computing Center (TACC)
and research gifts by Western Digital, Wireless Networking and Communications
Group (WNCG) Industrial Affiliates Program, UT Austin Machine Learning Lab
(MLL), Cisco and the Stanly P. Finch Centennial Professorship in Engineering.
Litu Rout has been supported by the Ju-Nam and Pearl Chew Endowed Presidential
Fellowship in Engineering. Giannis Daras has been supported by the Onassis
Fellowship (Scholarship ID: F ZS 012-1/2022-2023), the Bodossaki Fellowship
and the Leventis Fellowship. We thank the HuggingFace team for providing us
GPU support for the demo of our work.

## References

  * [1] Brian D.O. Anderson  “Reverse-time diffusion equation models”  In _Stochastic Processes and their Applications_ 12.3 Elsevier, 1982, pp. 313–326
  * [2] Marius Arvinte, Ajil Jalal, Giannis Daras, Eric Price, Alex Dimakis and Jonathan I Tamir  “Single-Shot Adaptation using Score-Based Models for MRI Reconstruction”  In _International Society for Magnetic Resonance in Medicine, Annual Meeting_ , 2022
  * [3] Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping and Tom Goldstein  “Cold Diffusion: Inverting arbitrary image transforms without noise”  In _arXiv preprint arXiv:2208.09392_ , 2022
  * [4] Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler and Karsten Kreis  “Align your latents: High-resolution video synthesis with latent diffusion models”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2023, pp. 22563–22575
  * [5] Ashish Bora, Ajil Jalal, Eric Price and Alexandros G Dimakis  “Compressed sensing using generative models”  In _International Conference on Machine Learning_ , 2017, pp. 537–546  PMLR
  * [6] Stanley H Chan, Xiran Wang and Omar A Elgendy  “Plug-and-play ADMM for image restoration: Fixed-point convergence and applications”  In _IEEE Transactions on Computational Imaging_ 3.1 IEEE, 2016, pp. 84–98
  * [7] Minshuo Chen, Kaixuan Huang, Tuo Zhao and Mengdi Wang  “Score Approximation, Estimation and Distribution Recovery of Diffusion Models on Low-Dimensional Data”  In _arXiv preprint arXiv:2302.07194_ , 2023
  * [8] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim and Anru R Zhang  “Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions”  In _arXiv preprint arXiv:2209.11215_ , 2022
  * [9] Sitan Chen, Giannis Daras and Alexandros G Dimakis  “Restoration-Degradation Beyond Linear Diffusions: A Non-Asymptotic Analysis For DDIM-Type Samplers”  In _arXiv preprint arXiv:2303.03384_ , 2023
  * [10] Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon and Sungroh Yoon  “Ilvr: Conditioning method for denoising diffusion probabilistic models”  In _arXiv preprint arXiv:2108.02938_ , 2021
  * [11] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky and Jong Chul Ye  “Diffusion Posterior Sampling for General Noisy Inverse Problems”  In _The Eleventh International Conference on Learning Representations_ , 2023  URL: https://openreview.net/forum?id=OnD9zGAGT0k
  * [12] Hyungjin Chung, Jeongsol Kim and Jong Chul Ye  “Direct Diffusion Bridge using Data Consistency for Inverse Problems”  In _arXiv preprint arXiv:2305.19809_ , 2023
  * [13] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu and Jong Chul Ye  “Improving Diffusion Models for Inverse Problems using Manifold Constraints”  In _Advances in Neural Information Processing Systems_ , 2022  URL: https://openreview.net/forum?id=nJJjv0JDJju
  * [14] Giannis Daras, Yuval Dagan, Alexandros G Dimakis and Constantinos Daskalakis  “Score-guided intermediate layer optimization: Fast langevin mixing for inverse problem”  In _arXiv preprint arXiv:2206.09104_ , 2022
  * [15] Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alexandros G Dimakis and Peyman Milanfar  “Soft diffusion: Score matching for general corruptions”  In _arXiv preprint arXiv:2209.05442_ , 2022
  * [16] Mauricio Delbracio and Peyman Milanfar  “Inversion by direct iteration: An alternative to denoising diffusion for image restoration”  In _arXiv preprint arXiv:2303.11435_ , 2023
  * [17] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li and Li Fei-Fei  “Imagenet: A large-scale hierarchical image database”  In _2009 IEEE conference on computer vision and pattern recognition_ , 2009, pp. 248–255  Ieee
  * [18] Prafulla Dhariwal and Alexander Nichol  “Diffusion models beat gans on image synthesis”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 8780–8794
  * [19] Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh and Jieyu Zhang  “DataComp: In search of the next generation of multimodal datasets”  In _arXiv preprint arXiv:2304.14108_ , 2023
  * [20] Jonathan Ho, Ajay Jain and Pieter Abbeel  “Denoising diffusion probabilistic models”  In _Advances in Neural Information Processing Systems_ 33, 2020, pp. 6840–6851
  * [21] Aapo Hyvärinen and Peter Dayan  “Estimation of non-normalized statistical models by score matching.”  In _Journal of Machine Learning Research_ 6.4, 2005
  * [22] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis and Jon Tamir  “Robust compressed sensing mri with deep generative priors”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 14938–14954
  * [23] Ajil Jalal, Sushrut Karmalkar, Alexandros G Dimakis and Eric Price  “Instance-optimal compressed sensing via posterior sampling”  In _arXiv preprint arXiv:2106.11438_ , 2021
  * [24] Ajil Jalal, Sushrut Karmalkar, Jessica Hoffmann, Alex Dimakis and Eric Price  “Fairness for Image Generation with Uncertain Sensitive Attributes”  In _Proceedings of the 38th International Conference on Machine Learning_ 139, Proceedings of Machine Learning Research  PMLR, 2021, pp. 4721–4732  URL: https://proceedings.mlr.press/v139/jalal21b.html
  * [25] Tero Karras, Samuli Laine and Timo Aila  “A style-based generator architecture for generative adversarial networks”  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2019, pp. 4401–4410
  * [26] Bahjat Kawar, Michael Elad, Stefano Ermon and Jiaming Song  “Denoising Diffusion Restoration Models”  In _Advances in Neural Information Processing Systems_
  * [27] Bahjat Kawar, Noam Elata, Tomer Michaeli and Michael Elad  “GSURE-Based Diffusion Model Training with Corrupted Data”  In _arXiv preprint arXiv:2305.13128_ , 2023
  * [28] Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang and Il-Chul Moon  “Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation”  In _International Conference on Machine Learning_ , 2022, pp. 11201–11228  PMLR
  * [29] Haohe Liu, Zehua Chen, Yi Yuan, Xinhao Mei, Xubo Liu, Danilo Mandic, Wenwu Wang and Mark D Plumbley  “Audioldm: Text-to-audio generation with latent diffusion models”  In _arXiv preprint arXiv:2301.12503_ , 2023
  * [30] Hongyu Liu, Bin Jiang, Yi Xiao and Chao Yang  “Coherent Semantic Attention for Image Inpainting”  In _2019 IEEE/CVF International Conference on Computer Vision (ICCV)_ IEEE, 2019  DOI: 10.1109/iccv.2019.00427
  * [31] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte and Luc Van Gool  “Repaint: Inpainting using denoising diffusion probabilistic models”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 11461–11471
  * [32] Gary Mataev, Peyman Milanfar and Michael Elad  “DeepRED: Deep image prior powered by RED”  In _Proceedings of the IEEE/CVF International Conference on Computer Vision Workshops_ , 2019, pp. 0–0
  * [33] Sachit Menon, Alexandru Damian, Shijia Hu, Nikhil Ravi and Cynthia Rudin  “Pulse: Self-supervised photo upsampling via latent space exploration of generative models”  In _Proceedings of the ieee/cvf conference on computer vision and pattern recognition_ , 2020, pp. 2437–2445
  * [34] Gregory Ongie, Ajil Jalal, Christopher A Metzler, Richard G Baraniuk, Alexandros G Dimakis and Rebecca Willett  “Deep learning techniques for inverse problems in imaging”  In _IEEE Journal on Selected Areas in Information Theory_ 1.1 IEEE, 2020, pp. 39–56
  * [35] Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell and Alexei A Efros  “Context encoders: Feature learning by inpainting”  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 2536–2544
  * [36] Walter HL Pinaya, Petru-Daniel Tudosiu, Jessica Dafflon, Pedro F Da Costa, Virginia Fernandez, Parashkev Nachev, Sebastien Ourselin and M Jorge Cardoso  “Brain imaging generation with latent diffusion models”  In _Deep Generative Models: Second MICCAI Workshop, DGM4MICCAI 2022, Held in Conjunction with MICCAI 2022, Singapore, September 22, 2022, Proceedings_ , 2022, pp. 117–126  Springer
  * [37] Elad Richardson, Yuval Alaluf, Or Patashnik, Yotam Nitzan, Yaniv Azar, Stav Shapiro and Daniel Cohen-Or  “Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation”  In _arXiv preprint arXiv:2008.00951_ , 2020
  * [38] Yaniv Romano, Michael Elad and Peyman Milanfar  “The little engine that could: Regularization by denoising (RED)”  In _SIAM Journal on Imaging Sciences_ 10.4 SIAM, 2017, pp. 1804–1844
  * [39] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser and Björn Ommer  “High-resolution image synthesis with latent diffusion models”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 10684–10695
  * [40] Litu Rout, Advait Parulekar, Constantine Caramanis and Sanjay Shakkottai  “A Theoretical Justification for Image Inpainting using Denoising Diffusion Probabilistic Models”  In _arXiv preprint arXiv:2302.01217_ , 2023
  * [41] Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev and Aran Komatsuzaki  “LAION-400M: Open Dataset of CLIP-Filtered 400 Million Image-Text Pairs”, 2021  arXiv:2111.02114 [cs.CV]
  * [42] Christoph Schuhmann et al.  “LAION-5B: An open large-scale dataset for training next generation image-text models”, 2022  arXiv:2210.08402 [cs.CV]
  * [43] Jiaming Song, Arash Vahdat, Morteza Mardani and Jan Kautz  “Pseudoinverse-guided diffusion models for inverse problems”  In _International Conference on Learning Representations_ , 2023
  * [44] Yang Song and Stefano Ermon  “Generative modeling by estimating gradients of the data distribution”  In _Advances in Neural Information Processing Systems_ 32, 2019
  * [45] Yang Song and Stefano Ermon  “Improved techniques for training score-based generative models”  In _Advances in neural information processing systems_ 33, 2020, pp. 12438–12448
  * [46] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon and Ben Poole  “Score-Based Generative Modeling through Stochastic Differential Equations”  In _International Conference on Learning Representations_ , 2021
  * [47] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon and Ben Poole  “Score-Based Generative Modeling through Stochastic Differential Equations”  In _International Conference on Learning Representations_
  * [48] Yu Takagi and Shinji Nishimoto  “High-resolution image reconstruction with latent diffusion models from human brain activity”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2023, pp. 14453–14463
  * [49] Singanallur V Venkatakrishnan, Charles A Bouman and Brendt Wohlberg  “Plug-and-play priors for model based reconstruction”  In _2013 IEEE Global Conference on Signal and Information Processing_ , 2013, pp. 945–948  IEEE
  * [50] Pascal Vincent  “A connection between score matching and denoising autoencoders”  In _Neural computation_ 23.7 MIT Press, 2011, pp. 1661–1674
  * [51] Su Wang, Chitwan Saharia, Ceslee Montgomery, Jordi Pont-Tuset, Shai Noy, Stefano Pellegrini, Yasumasa Onoe, Sarah Laszlo, David J Fleet and Radu Soricut  “Imagen editor and editbench: Advancing and evaluating text-guided image inpainting”  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2023, pp. 18359–18369
  * [52] Jiahui Yu, Zhe Lin, Jimei Yang, Xiaohui Shen, Xin Lu and Thomas Huang  “Free-Form Image Inpainting With Gated Convolution”  In _2019 IEEE/CVF International Conference on Computer Vision (ICCV)_ IEEE, 2019  DOI: 10.1109/iccv.2019.00457

##  Appendix A Technical Proofs

Notation and Measurement Matrix. We elaborate on the structure of the
measurement matrix ${\mathcal{A}}\in\mathbb{R}^{l\times d}.$ In our setting,
we are considering linear inverse problems. Thus, this matrix is a pixel
selector and consists of a subset of the rows from the $d\times d$ identity
matrix (the rows that are present correspond to the indices of the selected
pixels from the image $\overrightarrow{{\bm{x}}_{0}}\in\mathbb{R}^{d}$). Given
this structure, it immediately follows that ${\mathcal{A}}^{T}{\mathcal{A}}$
is a $d\times d$ matrix that has the interpretation of a pixel selection mask.
Specifically, ${\mathcal{A}}^{T}{\mathcal{A}}$ is a $d\times d$ diagonal
matrix ${\bm{D}}({\bm{m}})$, where the elements of ${\bm{m}}$ are set to 1
where data (pixel) is observed and 0 where data (pixel) is masked. Without the
loss of generality, we suppose that the first $k$ coordinates are known.

The rest of this section contains proofs of all the theorems and propositions
presented in the main body of the paper. For clarity, we restate the theorems
more formally with precise mathematical details.

###  A.1 Proof of Theorem 3.4

######  Theorem A.1 (Posterior Sampling using Diffusion in Pixel Space).

Suppose Assumption 3.1 and Assumption 3.2 hold. Let us denote by
$\bm{\sigma}=\\{\sigma_{j}\\}_{j=1}^{k}$ the singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$, i.e.
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})={\bm{U}}\Sigma{\bm{V}}^{T}\coloneqq{\bm{U}}{\bm{D}}(\bm{\sigma}){\bm{V}}^{T},{\bm{U}}\in\mathbb{R}^{k\times
k},{\bm{V}}\in\mathbb{R}^{k\times k}$ and

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{x}}_{1}}(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{x}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{x}}_{1}}\left(\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

Suppose $\overrightarrow{{\bm{x}}_{0}}\sim p(\overrightarrow{{\bm{x}}_{0}})$.
Given measurements $y={\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}$ and a fixed
variance $\beta\in(0,1)$, there exists a matrix step size555We use the term
‘step size’ in a more general way than is normally used. In this case, the
step size is a ‘pre-conditioning’ positive definite matrix, whose eigenvalue
magnitudes correspond to the scalar step sizes per coordinate along an
appropriately rotated basis. This general form is needed and with carefully
selected (unique) eigenvalues; otherwise the DPS algorithm fails to converge
to the groundtruth sample. We will later see that for our PSLD Algorithm in
Theorem 3.8, we can revert to the commonly used notion of step size (a single
scalar), as any finite step size (including a single scalar common across all
coordinates) suffices for proving recovery.
$\bm{\zeta}=(1/2)({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i})({\mathcal{S}}{\bm{U}})^{T},\bm{\zeta}_{i}=\\{\zeta_{i}^{j}=1/\sigma_{j}\\}^{k}_{j=1}$
for all the coordinates of $\overrightarrow{{\bm{x}}_{0}}$ such that Algorithm
1 samples from the true posterior $p(\overrightarrow{{\bm{x}}_{0}}|y)$ and
exactly recovers the groundtruth sample, i.e.,
$\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

Proof. Our goal is to show that
$\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$, where
$\overleftarrow{{\bm{x}}_{0}}$ is returned by Algorithm 1. Recall that the
reverse process starts with
$\overleftarrow{{\bm{x}}_{1}}\sim{\mathcal{N}}\left(\mathbf{0},{\bm{I}}_{d}\right)$
and generates the following:

$\displaystyle\overleftarrow{{\bm{x}}_{0}}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{x}}_{1}}-\bm{\zeta}\nabla_{\overleftarrow{{\bm{x}}_{1}}}\left\|{\mathcal{A}}\overleftarrow{{\bm{x}}_{0}}(\overleftarrow{{\bm{x}}_{1}})-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{x}}_{1}}-\bm{\zeta}\nabla_{\overleftarrow{{\bm{x}}_{1}}}\left\|{\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}\left({\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\right)^{T}\left({\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-{\bm{y}}\right)$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}\left({\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-{\bm{y}}\right)$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\bm{y}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2\bm{\zeta}{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}.$

Now, we use the singular value decomposition of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$ with left
singular vectors in ${\bm{U}}\in\mathbb{R}^{k\times k}$, right singular
vectors in ${\bm{V}}\in\mathbb{R}^{k\times k}$, and singular values
$\bm{\sigma}=[\sigma_{1},\dots,\sigma_{k}]$ in $\Sigma={\bm{D}}(\bm{\sigma})$.
Thus, the above expression becomes

$\displaystyle\overleftarrow{{\bm{x}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2\bm{\zeta}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2\bm{\zeta}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2\bm{\zeta}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i})({\mathcal{S}}{\bm{U}})^{T}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i})({\mathcal{S}}{\bm{U}})^{T}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i}){\bm{U}}^{T}{\bm{S}}^{T}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i}){\bm{U}}^{T}{\bm{S}}^{T}{\mathcal{S}}{\bm{U}}\Sigma{\bm{V}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle\stackrel{{\scriptstyle(ii)}}{{=}}{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i}){\bm{U}}^{T}{\bm{U}}\Sigma{\bm{U}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i}){\bm{U}}^{T}{\bm{U}}\Sigma{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i})\Sigma{\bm{U}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2({\mathcal{S}}{\bm{U}}){\bm{D}}(\bm{\zeta}_{i})\Sigma{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2{\mathcal{S}}{\bm{U}}{\bm{D}}(\bm{\zeta}_{i}){\bm{D}}(\bm{\sigma}){\bm{U}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2{\mathcal{S}}{\bm{U}}{\bm{D}}(\bm{\zeta}_{i}){\bm{D}}(\bm{\sigma}){\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-2{\mathcal{S}}{\bm{U}}{\bm{D}}(\bm{\zeta}_{i}\odot\bm{\sigma}){\bm{U}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+2{\mathcal{S}}{\bm{U}}{\bm{D}}(\bm{\zeta}_{i}\odot\bm{\sigma}){\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}},$

where (i) is due to Assumption 3.1 and (ii) uses Assumption 3.2. By choosing
$\zeta_{i}^{j}$ as half the inverse of the non-zero singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$, i.e.,
$\zeta_{i}^{j}=1/2\sigma_{i}~{}\forall i=1,\dots,k$, we obtain

$\displaystyle\overleftarrow{{\bm{x}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-{\mathcal{S}}{\bm{U}}{\bm{U}}^{T}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+{\mathcal{S}}{\bm{U}}{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle={\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}-{\mathcal{S}}{\mathcal{S}}^{T}\overleftarrow{{\bm{x}}_{1}}+{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}=\overrightarrow{{\bm{x}}_{0}},$

which completes the statement of the theorem. $\square$

###  A.2 Proof of Proposition 3.5

######  Proposition A.2 (Variational Autoencoder).

Suppose Assumption 3.1 holds. For an encoder
$\mathcal{E}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{k}$ and a decoder
$\mathcal{D}:\mathbb{R}^{k}\rightarrow\mathbb{R}^{d}$, denote by
$\mathcal{L}\left(\phi,\omega\right)$ the training objective of VAE:

$\displaystyle\arg\min_{\phi,\omega}\mathcal{L}\left(\phi,\omega\right)\coloneqq\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}}\sim
p}\left[\left\|\mathcal{D}(\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi);\omega)-\overrightarrow{{\bm{x}}_{0}}\right\|_{2}^{2}\right]+\lambda
KL\left(\mathcal{E}\sharp p,\mathcal{N}(\mathbf{0},{\bm{I}}_{k})\right),$

then the combination of
$\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi)={\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}}$
and
$\mathcal{D}(\overleftarrow{{\bm{z}}_{0}};\omega)={\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$
is a minimizer of $\mathcal{L}\left(\phi,\omega\right)$.

###### Proof.

To show that the encoder
$\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi)={\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}}$
and the decoder
$\mathcal{D}(\overleftarrow{{\bm{z}}_{0}};\omega)={\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$
minimize the VAE training objective $\mathcal{L}\left(\phi,\omega\right)$, we
begin with the first part of the loss, which is also called reconstruction
error $\mathcal{L}_{recon}\left(\phi,\omega\right)$. Substituting
$\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi)={\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}}$
and
$\mathcal{D}(\overleftarrow{{\bm{z}}_{0}};\omega)={\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$,
we have

$\displaystyle\mathcal{L}_{recon}\left(\phi,\omega\right)$
$\displaystyle\coloneqq\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}}\sim
p}\left[\left\|\mathcal{D}(\mathcal{E}(\overrightarrow{{\bm{x}}_{0}};\phi);\omega)-\overrightarrow{{\bm{x}}_{0}}\right\|_{2}^{2}\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}}\sim
p}\left[\left\|\mathcal{D}({\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}};\omega)-\overrightarrow{{\bm{x}}_{0}}\right\|_{2}^{2}\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}}\sim
p}\left[\left\|{\mathcal{S}}{\mathcal{S}}^{T}\overrightarrow{{\bm{x}}_{0}}-\overrightarrow{{\bm{x}}_{0}}\right\|_{2}^{2}\right]$

Using the fact that $\overrightarrow{{\bm{x}}_{0}}$ lives in a linear
subspace, we arrive at

$\displaystyle\mathcal{L}_{recon}\left(\phi,\omega\right)$
$\displaystyle=\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}}\sim
p}\left[\left\|{\mathcal{S}}{\mathcal{S}}^{T}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}-{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}\right\|_{2}^{2}\right]$
$\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}}\sim\mathcal{N}\left(\mathbf{0},{\bm{I}}_{k}\right)}\left[\left\|{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}-{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}\right\|_{2}^{2}\right]=0,$

where (i) is due to Assumption 3.1. Now, we analyze the distribution loss.
Note that the KL-divergence between two Gaussian distributions with moments
$(\mu_{1}$, $\sigma_{1})$ and $(\mu_{2},\sigma_{2})$ is given by

$\displaystyle
KL\left(\mathcal{N}(\mu_{1},\sigma_{1}),\mathcal{N}(\mu_{2},\sigma_{2})\right)=\log\left(\frac{\sigma_{2}}{\sigma_{1}}\right)+\frac{\sigma_{1}^{2}+\left(\mu_{1}-\mu_{2}\right)^{2}}{2\sigma_{2}^{2}}-\frac{1}{2}.$

Since
$\mathcal{E}\left({\bm{x}}_{0}\right)={\mathcal{S}}^{T}{\bm{x}}_{0}={\mathcal{S}}^{T}{\mathcal{S}}{\bm{z}}_{0}={\bm{z}}_{0}$,
the distribution loss becomes:

$\mathcal{L}_{dist}\left(\phi\right)\coloneqq KL\left(\mathcal{E}\sharp
p,\mathcal{N}(\mathbf{0},{\bm{I}}_{k})\right)=KL\left(\mathcal{N}(\mathbf{0},{\bm{I}}_{k}),\mathcal{N}(\mathbf{0},{\bm{I}}_{k})\right)=0.$

∎

###  A.3 Proof of Theorem 3.6

######  Theorem A.3 (Generative Modeling using Diffusion in Latent Space).

Suppose Assumption 3.1 holds. Let the optimal solution of the latent diffusion
model be

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{z}}_{1}}(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

For a fixed variance $\beta>0$, if
$\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\coloneqq{\bm{\theta}}\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)$,
then the closed-form solution is
${\bm{\theta}}^{*}=\sqrt{1-\beta}{\bm{I}}_{k}$, which after normalization by
$\frac{1}{\sqrt{1-\beta}}$ and composition with the decoder
${\mathcal{D}}\left(\overleftarrow{{\bm{z}}_{0}};\omega\right)\coloneqq{\mathcal{S}}\overleftarrow{{\bm{z}}_{0}}$
recovers the true subspace of $p\left(\overrightarrow{{\bm{x}}_{0}}\right)$.

Proof. In latent diffusion models, the training is performed in the latent
space of a pre-trained VAE. If the VAE is chosen from Proposition 3.5, then
the training objective becomes:

$\displaystyle\min_{{\bm{\theta}}}$
$\displaystyle\mathbb{E}_{\overrightarrow{{\bm{x}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}(\overrightarrow{{\bm{z}}_{1}}\left(\mathcal{E}(\overrightarrow{{\bm{x}}_{0}}),\overrightarrow{\bm{\epsilon}}),\mathcal{E}(\overrightarrow{{\bm{x}}_{0}})\right)-\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{z}}_{1}}\left(\mathcal{E}(\overrightarrow{{\bm{x}}_{0}}),\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\overrightarrow{{\bm{z}}_{0}}-\mu_{\bm{\theta}}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right]=\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\overrightarrow{{\bm{z}}_{0}}-{\bm{\theta}}\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right\|^{2}\right]$
$\displaystyle=\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\overrightarrow{{\bm{z}}_{0}}-{\bm{\theta}}\left(\overrightarrow{{\bm{z}}_{0}}\sqrt{1-\beta}+\sqrt{\beta}\overrightarrow{\bm{\epsilon}}\right)\right\|^{2}\right]$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}\overrightarrow{{\bm{z}}_{0}}\sim
p\\\
\overrightarrow{\bm{\epsilon}}\sim\mathcal{N}\left(\mathbf{0},{\bm{I}}_{k}\right)\end{subarray}}\left[\sum_{i=1}^{k}\left(\overrightarrow{{\bm{z}}_{0,i}}-{\bm{\theta}}_{i}^{T}\left(\overrightarrow{{\bm{z}}_{0}}\sqrt{1-\beta}+\overrightarrow{\bm{\epsilon}}\sqrt{\beta}\right)\right)^{2}\right],$

where ${\bm{\theta}}_{i}^{T}$ denotes the $i^{th}$ row of matrix
${\bm{\theta}}$. The solution of this regression problem is given by666For
ease of notation, we drop the forward arrow in the rest of this proof.

$\displaystyle{\bm{\theta}}_{i}^{*}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{x}}_{0},\bm{\epsilon}\end{subarray}}\left[\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)^{T}\right]^{-1}\mathbb{E}_{{\bm{x}}_{0},\bm{\epsilon}}\left[{\bm{z}}_{0,i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{x}}_{0},\bm{\epsilon}\end{subarray}}\left[\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)^{T}\right]^{-1}\mathbb{E}_{{\bm{x}}_{0},\bm{\epsilon}}\left[{\bm{z}}_{0,i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{x}}_{0},\bm{\epsilon}\end{subarray}}\left[\left(\mathcal{E}({\bm{x}}_{0})\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\left(\mathcal{E}({\bm{x}}_{0})\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)^{T}\right]^{-1}\mathbb{E}_{{\bm{x}}_{0},\bm{\epsilon}}\left[\mathcal{E}({\bm{x}}_{0})_{i}\left(\mathcal{E}({\bm{x}}_{0})\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[\left(\mathcal{E}({\mathcal{S}}{\bm{z}}_{0})\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\left(\mathcal{E}({\mathcal{S}}{\bm{z}}_{0})\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)^{T}\right]^{-1}\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[\mathcal{E}({\mathcal{S}}{\bm{z}}_{0})_{i}\left(\mathcal{E}({\mathcal{S}}{\bm{z}}_{0})\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[\left({\mathcal{S}}^{T}{\mathcal{S}}{\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\left({\mathcal{S}}^{T}{\mathcal{S}}{\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)^{T}\right]^{-1}\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\mathcal{S}}^{T}{\mathcal{S}}{\bm{z}}_{0})_{i}\left({\mathcal{S}}^{T}{\mathcal{S}}{\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$

Using Assumption 3.1, the above expression simplifies to

$\displaystyle{\bm{\theta}}_{i}^{*}$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)^{T}\right]^{-1}\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[(1-\beta){\bm{z}}_{0}{\bm{z}}_{0}^{T}+{\bm{z}}_{0}\bm{\epsilon}^{T}\sqrt{\beta(1-\beta)}+\bm{\epsilon}{\bm{z}}_{0}^{T}\sqrt{\beta(1-\beta)}+\beta\bm{\epsilon}\bm{\epsilon}^{T}\right]^{-1}\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\left[\left((1-\beta)\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[{\bm{z}}_{0}{\bm{z}}_{0}^{T}\right]+\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[{\bm{z}}_{0}\bm{\epsilon}^{T}\right]\sqrt{\beta(1-\beta)}+\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[\bm{\epsilon}{\bm{z}}_{0}^{T}\right]\sqrt{\beta(1-\beta)}+\beta\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0},\bm{\epsilon}\end{subarray}}\left[\bm{\epsilon}\bm{\epsilon}^{T}\right]\right)\right]^{-1}$
$\displaystyle\hskip
170.71652pt\times\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\left[\left((1-\beta){\bm{I}}_{k}+\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0}\end{subarray}}\left[{\bm{z}}_{0}\right]\mathop{\mathbb{E}}_{\begin{subarray}{c}\bm{\epsilon}\end{subarray}}\left[\bm{\epsilon}\right]^{T}\sqrt{\beta(1-\beta)}+\mathop{\mathbb{E}}_{\begin{subarray}{c}\bm{\epsilon}\end{subarray}}\left[\bm{\epsilon}\right]\mathop{\mathbb{E}}_{\begin{subarray}{c}{\bm{z}}_{0}\end{subarray}}\left[{\bm{z}}_{0}\right]^{T}\sqrt{\beta(1-\beta)}+\beta{\bm{I}}_{k}\right)\right]^{-1}$
$\displaystyle\hskip
170.71652pt\times\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right],$

where the last step uses the fact that ${\bm{z}}_{0}$ and $\bm{\epsilon}$ are
independent Gaussian random vectors with zero mean and unit covariance.
Simplifying further, we arrive at

$\displaystyle{\bm{\theta}}_{i}^{*}$
$\displaystyle=\left[(1-\beta){\bm{I}}_{k}+\beta{\bm{I}}_{k}\right]^{-1}\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\left({\bm{z}}_{0}\sqrt{1-\beta}+\bm{\epsilon}\sqrt{\beta}\right)\right]$
$\displaystyle=\mathbb{E}_{{\bm{z}}_{0}}\left[({\bm{z}}_{0})_{i}{\bm{z}}_{0}\sqrt{1-\beta}\right]+\mathbb{E}_{{\bm{z}}_{0},\bm{\epsilon}}\left[({\bm{z}}_{0})_{i}\bm{\epsilon}\sqrt{\beta}\right]$
$\displaystyle=\mathbb{E}_{{\bm{z}}_{0}}\left[({\bm{z}}_{0})_{i}{\bm{z}}_{0}\sqrt{1-\beta}\right]+\mathbb{E}_{{\bm{z}}_{0}}\left[({\bm{z}}_{0})_{i}\right]\mathbb{E}_{\bm{\epsilon}}\left[\bm{\epsilon}\right]\sqrt{\beta}.$

The final step follows from independence of ${\bm{z}}_{0}$ and
$\bm{\epsilon}$. Since ${\bm{z}}_{0}$ and $\bm{\epsilon}$ are also
$\mathcal{N}\left(\mathbf{0},{\bm{I}}_{k}\right)$, we get

$\displaystyle{\bm{\theta}}_{i}^{*}$
$\displaystyle=\mathbb{E}_{{\bm{z}}_{0}}\left[({\bm{z}}_{0})_{i}{\bm{z}}_{0}\sqrt{1-\beta}\right]=\left[0,\dots,0,\sqrt{1-\beta},0,\dots,0\right]^{T},$

where the $i^{th}$ coordinate is $\sqrt{1-\beta}$ and zero everywhere else.
Therefore, stacking all the rows together, we get
${\bm{\theta}}^{*}=\sqrt{1-\beta}{\bm{I}}_{k}$, which after normalization by
$1/\sqrt{1-\beta}$ gives the desired result.

Next, we show that $\bm{\theta}^{*}$ recovers the true subspace of
$\overrightarrow{{\bm{x}}_{0}}\sim
p\left(\overrightarrow{{\bm{x}}_{0}}\right)$. When composed with the decoder
of VAE, the generator of the LDM gives
$\overleftarrow{{\bm{x}}_{0}}=\mathcal{D}\left({\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}\right)=\mathcal{D}\left({\bm{I}}_{k}\overleftarrow{{\bm{z}}_{1}}\right)={\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}$.
Since
$\overleftarrow{{\bm{z}}_{1}}\sim{\mathcal{N}}\left(\mathbf{0},{\bm{I}}_{k}\right)$,
this completes the statement of the theorem. $\square$

###  A.4 Proof of Theorem 3.7

Recall that the the latent-space GML-DPS (6) algorithm (based on the pixel-
space DPS algorithm [11]) has three key steps. In the first step, it uses the
normalized closed-form solution obtained in Theorem 3.6 to perform one step of
denoising by the reverse SDE. In the second step, it runs one step of gradient
descent to satisfy the measurements in the pixel space. Finally, it takes one
step of gradient descent on the goodness objective, which acts as a
regularizer to ensure that the reconstructed image lies on the data manifold.

This can be formalized as:

$\displaystyle\overleftarrow{{\bm{z}}^{\prime}_{0}}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}-\bm{\eta}\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}}))-{\bm{y}}\right\|_{2}^{2};$
(8) $\displaystyle\overleftarrow{{\bm{z}}_{0}}$
$\displaystyle=\arg\min_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left|\left|\overleftarrow{{\bm{z}}^{\prime}_{0}}-\mathcal{E}(\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))\right|\right|_{2}^{2},$
(9)

In practice, solving (9) can be difficult, and can be approximated via
gradient descent. In our analysis however, we analyze the exact system of
equations above, as (9) has a closed-form solution in the linear setting.

######  Theorem A.4 (Posterior Sampling using Goodness Modified Latent DPS).

Suppose Assumptions 3.1 and Assumption 3.2 hold. Denote by
$\bm{\sigma}=\\{\sigma_{j}\\}_{j=1}^{k}$ the singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$, i.e.,
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})={\bm{U}}\Sigma{\bm{U}}^{T}\coloneqq{\bm{U}}{\bm{D}}(\bm{\sigma}){\bm{U}}^{T},{\bm{U}}\in\mathbb{R}^{k\times
k}$, and let

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{z}}_{1}}(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|_{2}^{2}\right].$

Suppose $\overrightarrow{{\bm{x}}_{0}}\sim p(\overrightarrow{{\bm{x}}_{0}})$.
Given measurements ${\bm{y}}={\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}$ and
any fixed variance $\beta\in(0,1)$, then with the (unique) step size
$\bm{\eta}=(1/2){\bm{U}}{\bm{D}}(\bm{\eta}_{i}){\bm{U}}^{T},\bm{\eta}_{i}=\\{\eta_{i}^{j}=1/2\sigma_{j}\\}_{j=1}^{k}$,
the GML-DPS algorithm (6) samples from the true posterior
$p(\overrightarrow{{\bm{x}}_{0}}|y)$ and exactly recovers the groundtruth
sample, i.e., $\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

Proof. We start with the measurement consistency update (8) and then show that
the solution obtained from (8) is already a minimizer of (9). Therefore, we
have

$\displaystyle\overleftarrow{{\bm{z}}^{\prime}_{0}}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}-\bm{\eta}\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}}))-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle={\bm{I}}_{k}\overleftarrow{{\bm{z}}_{1}}-\bm{\eta}\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}({\bm{I}}_{k}\overleftarrow{{\bm{z}}_{1}})-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-\bm{\eta}\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}})-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-\bm{\eta}\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}})-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle\stackrel{{\scriptstyle(i)}}{{=}}\overleftarrow{{\bm{z}}_{1}}-\bm{\eta}\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}\left({\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}-{\bm{y}}\right)$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}+2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\bm{y}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}+2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}+2\bm{\eta}{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}},$

where (i) is due to Assumption 3.1. By Assumption 3.2,
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$ is a positive
definite matrix and can be written as ${\bm{U}}\Sigma{\bm{U}}^{T}$:

$\displaystyle\overleftarrow{{\bm{z}}^{\prime}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\bm{\eta}{\bm{U}}\Sigma{\bm{U}}^{T}\overleftarrow{{\bm{z}}_{1}}+2\bm{\eta}{\bm{U}}\Sigma{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2{\bm{U}}{\bm{D}}(\bm{\eta}_{i}){\bm{U}}^{T}{\bm{U}}\Sigma{\bm{U}}^{T}\overleftarrow{{\bm{z}}_{1}}+2{\bm{U}}{\bm{D}}(\bm{\eta}_{i}){\bm{U}}^{T}{\bm{U}}\Sigma{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2{\bm{U}}{\bm{D}}(\bm{\eta}_{i})\Sigma{\bm{U}}^{T}\overleftarrow{{\bm{z}}_{1}}+2{\bm{U}}{\bm{D}}(\bm{\eta}_{i})\Sigma{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2{\bm{U}}{\bm{D}}(\bm{\eta}_{i}){\bm{D}}(\bm{\sigma}){\bm{U}}^{T}\overleftarrow{{\bm{z}}_{1}}+2{\bm{U}}{\bm{D}}(\bm{\eta}_{i}){\bm{D}}(\bm{\sigma}){\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2{\bm{U}}{\bm{D}}(\bm{\eta}_{i}\odot\bm{\sigma}){\bm{U}}^{T}\overleftarrow{{\bm{z}}_{1}}+2{\bm{U}}{\bm{D}}(\bm{\eta}_{i}\odot\bm{\sigma}){\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}.$

Since $\eta_{j}^{i}=1/2\sigma_{j}$, the above expression further simplifies to

$\displaystyle\overleftarrow{{\bm{z}}^{\prime}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-{\bm{U}}{\bm{U}}^{T}\overleftarrow{{\bm{z}}_{1}}+{\bm{U}}{\bm{U}}^{T}\overrightarrow{{\bm{z}}_{0}}=\overrightarrow{{\bm{z}}_{0}}.$

Next, we show that $\overleftarrow{{\bm{z}}^{\prime}_{0}}$ is already a
minimizer of (9). This is a direct consequence of the encoder-decoder
architecture of the VAE:
$\mathcal{E}(\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))={\mathcal{S}}^{T}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}}=\overleftarrow{{\bm{z}}^{\prime}_{0}}$.
Hence,
$\left|\left|\overleftarrow{{\bm{z}}^{\prime}_{0}}-\mathcal{E}(\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))\right|\right|^{2}=0$,
and consequently
$\overleftarrow{{\bm{z}}_{0}}=\overleftarrow{{\bm{z}}^{\prime}_{0}}-\gamma\nabla_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left|\left|\overleftarrow{{\bm{z}}^{\prime}_{0}}-\mathcal{E}(\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))\right|\right|^{2}=\overrightarrow{{\bm{z}}_{0}}$.
Thus, the reconstructed sample becomes
$\overleftarrow{{\bm{x}}_{0}}=\mathcal{D}(\overleftarrow{{\bm{z}}_{0}})={\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

Furthermore, as
$\left|\left|\overleftarrow{{\bm{z}}^{\prime}_{0}}-\mathcal{E}(\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))\right|\right|^{2}=0$
for all $\overleftarrow{{\bm{z}}^{\prime}_{0}}$, it is evident that the
goodness objective cannot rectify the error incurred in the measurement update
(8). For this reason, GML-DPS algorithm (6) requires the exact step size to
sample from the posterior. $\square$

Beyond the linear setting, we also refer to Table 5 for experiments supporting
this result.

###  A.5 Proof of Theorem 3.8

Different from GML-DPS, PSLD Algorithm 2 replaces the goodness objective (6)
with the gluing objective (7), which can be formalized as:

$\displaystyle\overleftarrow{{\bm{z}}^{\prime}_{0}}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}-\eta\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}}))-{\bm{y}}\right\|_{2}^{2};$
(10) $\displaystyle\overleftarrow{{\bm{z}}_{0}}$
$\displaystyle=\arg\min_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left|\left|\overleftarrow{{\bm{z}}^{\prime}_{0}}-\mathcal{E}({\mathcal{A}}^{T}{\mathcal{A}}\overrightarrow{{\bm{z}}_{0}}+({\bm{I}}_{d}-{\mathcal{A}}^{T}{\mathcal{A}})\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))\right|\right|_{2}^{2}.$
(11)

We again remind that solving the minimization problem (11) is hard in general,
and can be approximated by gradient descent as typically followed in practice
[11]. However, in a linear model setting, (11) has a closed-form solution
which we derive to prove exact recovery.

######  Theorem A.5 (Posterior Sampling using Diffusion in Latent Space).

Let Assumptions 3.1 and 3.2 hold. Let $\sigma_{j},\forall j=1,\dots,r$ denote
the singular values of
$({\mathcal{A}}{\mathcal{S}})^{T}({\mathcal{A}}{\mathcal{S}})$ and let

$\displaystyle{\bm{\theta}}^{*}=\arg\min_{{\bm{\theta}}}\mathbb{E}_{\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}}\left[\left\|\tilde{\mu}_{1}\left(\overrightarrow{{\bm{z}}_{1}}(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}),\overrightarrow{{\bm{z}}_{0}}\right)-\mu_{\theta}\left(\overrightarrow{{\bm{z}}_{1}}\left(\overrightarrow{{\bm{z}}_{0}},\overrightarrow{\bm{\epsilon}}\right)\right)\right\|^{2}\right].$

Suppose $\overrightarrow{{\bm{x}}_{0}}\sim p(\overrightarrow{{\bm{x}}_{0}})$.
Given measurements ${\bm{y}}={\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}$, any
fixed variance $\beta\in(0,1)$, and any positive step sizes
$\eta_{i}^{j},j=1,2,\ldots,r$, the PSLD Algorithm 2 samples from the true
posterior $p(\overrightarrow{{\bm{x}}_{0}}|y)$ and exactly recovers the
groundtruth sample, i.e.,
$\overleftarrow{{\bm{x}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.

Proof. Following the proof in Appendix A.4, we have

$\displaystyle\overleftarrow{{\bm{z}}^{\prime}_{0}}$
$\displaystyle={\bm{\theta}}^{*}\overleftarrow{{\bm{z}}_{1}}-\eta\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}(\overleftarrow{{\bm{z}}_{0}}(\overleftarrow{{\bm{z}}_{1}}))-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle={\bm{I}}_{k}\overleftarrow{{\bm{z}}_{1}}-\eta\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}\mathcal{D}(\overleftarrow{{\bm{z}}_{1}})-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-\eta\nabla_{\overleftarrow{{\bm{z}}_{1}}}\left\|{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}-{\bm{y}}\right\|_{2}^{2}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\eta{\mathcal{S}}^{T}{\mathcal{A}}^{T}({\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}-{\bm{y}})$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\eta{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}+2\eta{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}$
$\displaystyle=\overleftarrow{{\bm{z}}_{1}}-2\eta{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}_{1}}+2\eta{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}.$

We use the above expression to derive a closed-form solution to the
minimization problem (11):

$\displaystyle\mathbf{0}$
$\displaystyle=\nabla_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left\|\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}({\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}+({\bm{I}}_{d}-{\mathcal{A}}^{T}{\mathcal{A}}){\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}})\right\|_{2}^{2}$
$\displaystyle=\nabla_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left\|\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}-{\mathcal{S}}^{T}({\bm{I}}_{d}-{\mathcal{A}}^{T}{\mathcal{A}}){\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}})\right\|_{2}^{2}$
$\displaystyle=\nabla_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left\|\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}-{\mathcal{S}}^{T}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}})\right\|_{2}^{2}$
$\displaystyle=\nabla_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left\|\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}-{\mathcal{S}}^{T}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}}+{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}})\right\|_{2}^{2}$
$\displaystyle=2\left({\bm{I}}_{k}-{\mathcal{S}}^{T}{\mathcal{S}}+{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\right)\left(\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}-{\mathcal{S}}^{T}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}}+{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}})\right)$
$\displaystyle=2{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\left({\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overleftarrow{{\bm{z}}^{\prime}_{0}}-{\mathcal{S}}^{T}{\mathcal{A}}^{T}{\mathcal{A}}{\mathcal{S}}\overrightarrow{{\bm{z}}_{0}})\right),$

where the last step is due to Assumption 3.1. Thus, we have

$\displaystyle\overleftarrow{{\bm{z}}_{0}}=\arg\min_{\overleftarrow{{\bm{z}}^{\prime}_{0}}}\left|\left|\overleftarrow{{\bm{z}}^{\prime}_{0}}-\mathcal{E}({\mathcal{A}}^{T}{\mathcal{A}}\overrightarrow{{\bm{z}}_{0}}+({\bm{I}}_{d}-{\mathcal{A}}^{T}{\mathcal{A}})\mathcal{D}(\overleftarrow{{\bm{z}}^{\prime}_{0}}))\right|\right|_{2}^{2}=\overrightarrow{{\bm{z}}_{0}},$

which produces
$\overleftarrow{{\bm{x}}_{0}}=\mathcal{D}(\overleftarrow{{\bm{z}}_{0}})=\mathcal{D}(\overrightarrow{{\bm{z}}_{0}})={\mathcal{S}}\overrightarrow{{\bm{z}}_{0}}=\overrightarrow{{\bm{x}}_{0}}$.
$\square$

It is worth highlighting that PSLD exactly recovers the groundtruth sample
irrespective of the choice of the step size $\eta$, whereas GML-DPS requires
the step size to be exactly
$\bm{\eta}=(1/2){\bm{U}}{\bm{D}}(\bm{\eta}_{i}){\bm{U}}^{T}$.

##  Appendix B Additional Experiments

###  B.1 Implementation Details

For inpainting tasks, we note that the PSLD sampler generates missing parts
(by design of our gluing objective) that are consistent with the known
portions of the image, i.e.,
$\overleftarrow{{\bm{x}}_{0}}={\mathcal{A}}^{T}{\mathcal{A}}\overrightarrow{{\bm{x}}_{0}}+({\bm{I}}_{d}-{\mathcal{A}}^{T}{\mathcal{A}})\mathcal{D}(\overleftarrow{{\bm{z}}_{0}})$.
This is different from the DPS sampler, which generates the whole image which
may not match the observations exactly. In other words, in the last of step of
our algorithm, the observations are glued onto the corresponding parts of the
generated image, leaving the unmasked portions untouched [51]. This sometimes
creates edge effects which are then removed by post-processing the glued image
through the encoder and decoder of the SD model, i.e. running one last step of
our algorithm. Figure 2 illustrates that gluing the observations in commercial
services still leads to visually inconsistent results (e.g. head in top row)
unlike our method.

For all other tasks, such as motion deblur, Gaussian deblur, and super-
resolution, this last step is not needed, as there is no box inpainting, i.e.,
$\overleftarrow{{\bm{x}}_{0}}=\mathcal{D}(\overleftarrow{{\bm{z}}_{0}})$.
Furthermore, we use the same measurement operator ${\mathcal{A}}$ and its
transpose ${\mathcal{A}}^{T}$ as provided by the DPS code
repository777https://github.com/DPS2022/diffusion-posterior-
sampling/blob/main/guided_diffusion/measurements.py. However, since Stable
Diffusion v1.5 generates images of size $512\times 512$ resolution and DPS
operates at $256\times 256$, we adjust the size of the kernels used in PSLD to
ensure that both the methods use the same amount of information while sampling
from the posterior. During evaluation, we downsample PSLD generated images
from $512\times 512$ to $256\times 256$ to compare with DPS at the same
resolution.

PSLD (Stable Diffusion-V1.5 ): We run Algorithm 2 with Stable Diffusion
version 1.5 as the foundation model888https://huggingface.co/runwayml/stable-
diffusion-v1-5. We use a fixed $\eta=1$ and $\gamma=0.1$. Since we study
posterior sampling of images without conditioning on text inputs, we pass an
empty string to the Stable Diffusion foundation model, which accepts texts as
an input argument. For better performance, we recommend using the latest
pretrained weights.

PSLD (LDM-VQ-4 ): This is the same sampling algorithm as before but with a
different latent diffusion model, LDM-
VQ-4999https://github.com/CompVis/latent-diffusion , which contains pretrained
weights for FFHQ 256101010https://ommer-lab.com/files/latent-
diffusion/ffhq.zip and large-scale text-to-image generative
model111111https://ommer-lab.com/files/latent-
diffusion/nitro/txt2img-f8-large/model.ckpt. We keep the hyperparameters same
($\eta=1$ and $\gamma=0.1$). For each task, we provide hyper-parameter details
in our codebase121212https://github.com/LituRout/PSLD. Although we have tested
our framework with these two latent-diffusion-models, one may experiment with
other latent-diffusion-models available in the same repository.

DPS: We use the original source code provided by the
authors131313https://github.com/DPS2022/diffusion-posterior-sampling.

OOD images are sourced online:

  1. 1.

Figure 1: the original images are generated by Stable Diffusion
v-2.1141414https://huggingface.co/spaces/stabilityai/stable-diffusion.

  2. 2.

Figure 2 first row: Walking example from the web.

  3. 3.

Figure 2 second row, Obama-Biden image from the web.

  4. 4.

Figure 2 third row, Fisherman from ImageNet 256 [17].

  5. 5.

Figure 4 first row: Racoon image from the web.

  6. 6.

Figure 4 second row: Fisherman from ImageNet 256 [17].

  7. 7.

Figure 15: Celebrity face from the web.

###  B.2 Additional Experimental Evaluation

Here, we provide additional results to support our theoretical claims on
various inverse problems.

Figures 6, 7, 8, and 9 show the inpainting results of user defined masks
obtained from our PSLD inpainting web demo. Note that the foundation model
used in this demo is a generic model. For better performance on specific
images, we recommend finetuning the foundation model on this class and then
running posterior sampling using our web demo:
https://huggingface.co/spaces/PSLD/PSLD.

Figure 6: Results from the web application of our PSLD algorithm, $512\times
512$. The original image (1) is generated by Stable Diffusion v-2.1 with the
prompt,“A dinner date between a robot couple during sunset”. Figure 7: Results
from the web application of our PSLD algorithm, $512\times 512$. The original
image (1) is generated by Stable Diffusion v-2.1 with the prompt,“A panda
wearing a spiderman costume”. Figure 8: Results from the web application of
our PSLD algorithm, $512\times 512$. The original image (1) is generated by
Stable Diffusion v-2.1 with the prompt,“A teddy bear showing stop sign at the
traffic”. Figure 9: Results from the web application of our PSLD algorithm,
$512\times 512$. The original image (1) is generated by Stable Diffusion v-2.1
with the prompt,“A cute dog playing with a toy teddy bear on the lawn”.

Figure 10 and 11 illustrate super-resolution (4$\times$) of in-distribution
samples from the validation set of FFHQ 256. Observe that the samples
generated by DPS are far from the groundtruth sample. On the other hand, the
samples generated by PSLD closely capture the perceptual quality of the
groundtruth sample. In other words, one may identify (b) and (c) as images of
two different individuals, whereas (b) and (d) of the same individual. We
attribute this photorealism of our method to the power of Stable Diffusion
foundation model and the ability to use the knowledge of the VAE encoder-
decoder in the gluing objective.

In addition, we test on out-of-distribution samples from ImageNet [17]
validation set. Figure 12 and Figure 13 show the results in motion deblur and
Gaussian deblur, respectively. By leveraging the foundation model Stable
Diffusion v1.5, our PSLD method clearly outperforms DPS [11] in the general
domain. Further, Figures 14, 15, and 16 show reconstruction of general domain
samples for random inpainting, super-resolution, and destriping tasks,
respectively. In all these tasks, the samples generated by PSLD are closer to
the groundtruth sample than the ones generated by DPS. Table 5 shows the
quantitative results.

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 10:  Super-resolution results on images from FFHQ 256 [25, 11] (in
distribution).

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 11:  Super-resolution results on FFHQ 256 [25, 11] (in distribution).

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 12: Motion deblur results on ImageNet 256 [17] (out-of-distribution).

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 13:  Gaussian deblur results on ImageNet 256 [17] (out-of-distribution). Table 5: Quantitative random inpainting results on FFHQ $256$ validation set [25, 11]. We use Stable Diffusion (v1.5) trained on LAION.  | Inpaint (random) | SR ($4\times$) | Gaussian Deblur
---|---|---|---
Method | PSNR ($\uparrow$) | SSIM ($\uparrow$) | PSNR ($\uparrow$) | SSIM ($\uparrow$) | PSNR ($\uparrow$) | SSIM ($\uparrow$)
PSLD (Ours) | 30.31 | 0.851 | 30.73 | 0.867 | 30.10 | 0.843
GML-DPS (Ours) | 29.49 | 0.844 | 29.77 | 0.860 | 29.21 | 0.820
DPS [11] | 25.23 | 0.851 | 25.67 | 0.852 | 24.25 | 0.811
DDRM [26] | 9.19 | 0.319 | 25.36 | 0.835 | 23.36 | 0.767
MCG [13] | 21.57 | 0.751 | 20.05 | 0.559 | 6.72 | 0.051
PnP-ADMM [6] | 8.41 | 0.325 | 26.55 | 0.865 | 24.93 | 0.812
Score-SDE [47] | 13.52 | 0.437 | 17.62 | 0.617 | 7.12 | 0.109
ADMM-TV | 22.03 | 0.784 | 23.86 | 0.803 | 22.37 | 0.801

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 14:  Random inpainting results on ImageNet 256 [17] (out-of-
distribution).

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 15:  Super-resolution (using nearest neighbor kernel from [31]) results
on out-of-distribution samples from the web, $256\times 256$ (see Table 2 for
LPIPS of these images).

(a) Input

(b) Groundtruth

(c) DPS [11]

(d) PSLD (Ours)

Figure 16:  Destriping results on out-of-distribution samples from the web,
$256\times 256$. (Top row) Horizontal destriping: LPIPS of PSLD=0.244 and DPS
[11]=0.613. (Bottom row) Vertical destriping: LPIPS of PSLD=0.255, DPS
[11]=0.597.