Title: Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Inverse Problems

Posterior Sampling

Latent Diffusion Model

Stable Diffusion

Sample Recovery
Abstract: We present the first framework to solve linear inverse problems leveraging pre-trained \textit{latent} diffusion models.  Previously proposed algorithms (such as DPS and DDRM) only apply to \textit{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.

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

Abstract

We present the first framework to solve linear inverse problems leveraging pretrained 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 [37, 39, 56, 31], and unsupervised methods that use the prior learned by a generative model to guide the restoration process [52, 40, 5, 33, 11, 26]; see also the survey of Ongie et al. [36] 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 [35, 24] and regression to the mean [23, 22].

Diffusion models have emerged as a powerful new approach to generative modeling [47, 48, 49, 20, 29, 18, 54]. This family of generative models works by first corrupting the data distribution p 0 ( x 0 ) using an Itô Stochastic Differential Equation (SDE), d x = f ( x , t )d t + g ( t )d w , and then by learning the score-function, ∇ x t log p t ( x t ) , at all levels t , using Denoising Score Matching (DSM) [21, 53]. The seminal result of Anderson [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 [49]. The latter depends on the score-function ∇ x t log p t ( 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. Song et al. [49] showed that given measurements y = A x 0 + σ y n , we can

SOP

provably sample from the distribution p 0 ( x 0 | y ) by running a modified Reverse SDE that depends on the unconditional score ∇ x t log p t ( x t ) and the term ∇ x t log p ( y | 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, 46, 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) [41]. LDMs project the data into some latent space, z 0 = E ( x 0 ) , perform the

diffusion in the latent space and use a decoder, D ( z 0 ) , to move back to the pixel space. LDMs power state-of-the-art foundation models such as Stable Diffusion [41] and have enabled a widerange of applications across many data modalities including images [41], video [4], audio [30] and medical domain distributions (e.g., for MRI and proteins) [38, 51]. 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. 2

2 Background and Method

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

The standard diffusion modeling framework involves training a network, s θ ( x t , t ) , to learn the score-function, ∇ x t log p t ( x t ) , at all levels t , of a stochastic process described by an Itô SDE:

where 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 y = A x 0 + σ y n , one can sample from the distribution p 0 ( x 0 | y ) by running the conditional Reverse SDE given by:

As mentioned, ∇ x t log p ( y | x t ) is intractable for general inverse problems. One of the most effective approximation methods is the DPS algorithm proposed by Chung et al. [11]. DPS assumes that:

Essentially, DPS substitutes the unknown clean image x 0 with its conditional expectation given the noisy input, E [ x 0 | x t ] . Under this approximation, the term p ( y | 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 ( x 0 | y ) (Section A.1). Then, we provide an algorithm (Section 2.1) and its analysis to sample from p ( x 0 | y ) using latent diffusion models (Section 3.2). 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 s θ ( z t , t ) to predict the score ∇ z t log p t ( z t ) , of a diffusion process:

where z 0 = E ( x 0 ) for some encoder function E ( · ) : R d → R k . During sampling, we start with z T , we run the Reverse Diffusion Process and then we obtain a clean image by passing z 0 ∼ p 0 ( z 0 | z T ) through a decoder D : R k → 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 ( y | z t ) with:

i.e., to approximate the unknown clean image x 0 with the decoded version of the conditional expectation of the clean latent z 0 given the noisy latent 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 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 z t towards any of these latents z 0 , potentially in different directions. On the other hand, the score-function is pulling z t towards a specific z 0 that corresponds to the best denoised version of 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 ∇ log p ( y | z t ) with:

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 z 0 as long as D ( 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.

The gluing objective is critical for our algorithm as it ensures that the denoising update, measurementmatching 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.

Remark 2.1. Consider the optimization problem of projecting onto the measurements:

In the linear setting, the optimal solution is given by x ∗ 0 = A T ( AA T ) -1 y + (ˆ x 0 -A T ( AA T ) -1 ( A ˆ x 0 )) . Now further suppose that the measurement rows are orthogonal, i.e. AA T = I l . This condition holds for some natural linear inverse problems like inpainting. Suppose that we want to update the latent vector z t such that E [ z 0 | z t ] = E ( x ∗ 0 ) ; this ensures that the gradients

Algorithm 1: DPS

Input:

T

,

y

,

ζ

T

i

=1

,

{

˜

σ

i

}

T

i

=1

,

s

θ

1

x

T

∼ N

(

0

,

I

)

2

for

i

=

T

-

1

to

0

do

3

ˆ

s

s

θ

(

x

i

, i

)

4

ˆ

x

0

1

√

¯

α i

(

x

i

+(1

-

¯

α

i

)ˆ

s

)

5

z

∼ N

(

0

,

I

)

6

x

′

i

-

1

√

α i

(1

-

¯

α i

-

1 )

1

-

¯

α i

x

i

+

√

¯

α i

-

1

β i

1

-

¯

α i

ˆ

x

0

+ ˜

σ

i

z

7

x

i

-

1

x

′

i

-

1

-

ζ

i

∇

x

i

∥

y

-A

(ˆ

x

0

)

∥

2

2

8

end

9

return

ˆ

x

0

Algorithm 2: PSLD

Input:

T

,

y

, {

η

i

}

T

i

=1

, {

γ

i

}

T

i

=1

,

{

˜

σ

i

}

T

i

=1

,

E

,

D

,

A

,

s

θ

1

z

T

∼ N

(

0

,

I

)

2

for

i

=

T

-

1

to

0

do

3

ˆ

s

s

θ

(

z

i

, i

)

4

ˆ

z

0

1

√

¯

α i

(

z

i

+(1

-

¯

α

i

)ˆ

s

)

5

ϵ

∼ N

(

0

,

I

)

6

z

′

i

-

1

√

α i

(1

-

¯

α i

-

1 )

1

-

¯

α i

z

i

+

√

¯

α i

-

1

β i

1

-

¯

α i

ˆ

z

0

+ ˜

σ

i

ϵ

7

z

′′

i

-

1

z

′

i

-

1

-

η

i

∇

z

i

∥

y

-A

(

D

(ˆ

z

0

))

∥

2

2

8

z

i

-

1

z

′′

i

-

1

-

γ

i

∇

z

i

∥

ˆ

z

0

-E

(

A

T

y

+(

I

-A

T

A

)

D

(ˆ

z

0

))

∥

2

2

9

end

10

return

D

(ˆ

z

0

)

resulting from the two terms in (7) both point to the same optima in the latent space. Equivalently, we want to solve the following minimization problem: min z t ∥ E [ z 0 | z t ] -E ( x ∗ 0 ) ∥ 2 2 . Substituting E ( x ∗ 0 ) = E ( A T y +(ˆ x 0 -A T A ˆ x 0 )) = E ( A T y +( I -A T A )ˆ x 0 ) , and ˆ x 0 = D ( E [ z 0 | z t ]) , we can thus interpret the gluing objective in (7) as a one step of gradient descent of this loss ∥ E [ z 0 | z t ] -E ( x ∗ 0 ) ∥ 2 2 with respect to z t . Note that, if there was no latent space, our gluing would be equivalent to a projection on the measurements, but now because of the encoder and decoder, it is not.

3 Theoretical Results

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 - → x 0 ∼ p ( - → x 0 ) and gradually adds noise, i.e., - → x t +1 = √ 1 -β t - → x t + √ β t ϵ where β t ∈ [0 , 1] and β t ≥ β t -1 for t = 0 , . . . , T -1 . The reverse process is initialized with ← -x T ∼ N ( 0 , I d ) and generates ← -x t -1 = µ θ ( ← -x t , t ) + √ β t ϵ . In the last step, µ θ ( ← -x 1 , 1) is displayed without the noise.

In this section, we consider the diffusion discretized to two steps ( { - → x 0 , - → 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 largescale real-world applications of diffusion models that use multiple steps ( { - → x 0 , - → x 1 , · · · , - → x T } , where T = 1000 ), and moreover do not satisfy the recoverability assumptions. We provide post-processing details of Algorithm 2 in Appendix C.1. All proofs are given in Appendix B.

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 [42, 7], where the images ( - → x 0 ∈ R d ) reside in a linear subspace of the form - → x 0 = S - → w 0 , S ∈ R d × l , - → w 0 ∈ R l , and σ y = 0 . Here, S is a tall thin matrix with rank ( S ) = l ≤ d that lifts any latent vector - → w 0 ∼ N ( 0 , I l ) to the image space with ambient dimension d . Given the measurements y = A - → x 0 + σ y n , A ∈ R l × d , n ∈ R l , the goal is to sample from p 0 ( - → x 0 | y ) using a pre-trained latent diffusion model. In the inpainting task, the measurement operator A is such that A T A is a diagonal matrix D ( m ) , where m is the masking vector with elements set to 1 where data is observed and 0 where data is masked (see Appendix B 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 [41], we consider a setting where the latent vector of the V AE is k -dimensional and the latent distribution is a standard Gaussian N ( 0 , I k ) . 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 S are orthonormal, i.e., S T S = I

l .

Assumption 3.2. The measurement operator A satisfies ( AS ) T ( AS ) ≻ 0 .

These assumptions have previously appeared, e.g., [42]. While Assumption 3.1 is mild and can be relaxed at the expense of (standard) mathematical complications, Assumption 3.2 indicates that ( AS ) T ( AS ) 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 - → x 0 . The underlying assumption is that there exists a solution and it is unique [42]. 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 [42, 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 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 V AE is in the same space as the latent distribution of the data generating model, i.e., k = l , and normalize γ i = 1 (as this is immaterial in the linear setting). In Proposition 3.3 , we provide analytical solutions for the encoder and the decoder of the VAE.

Proposition 3.3 (Variational Autoencoder) . Suppose Assumption 3.1 holds. For an encoder E : R d → R k and a decoder D : R k → R d , denote by L ( ϕ, ω ) the training objective of VAE:

then the combination of E ( - → x 0 ; ϕ ) = S T - → x 0 and D ( ← -z 0 ; ω ) = S ← z 0 is a minimizer of L ( ϕ, ω ) .

Using the encoder E ( - → x 0 ; ϕ ) = S T - → x 0 , we can use the analytical solution θ ∗ of the LDM obtained in Theorem A.1 . To verify that θ ∗ recovers the true subspace p ( - → x 0 ) , we compose the decoder D ( ← -z 0 ; ω ) = S ← z 0 with the generator of the LDM, i.e., ← -x 0 = D ( θ ∗ ← -z 1 ) = D ( I k ← -z 1 ) = S ← z 1 . Since ← -z 1 ∼ N ( 0 , I k ) and S is the data generating model, this shows that ← -x 0 is a sample from p ( - → x 0 ) . Thus we have the following.

Theorem 3.4 (Generative Modeling using Diffusion in Latent Space) . Suppose Assumption 3.1 holds. Let the optimal solution of the latent diffusion model be

For a fixed variance β > 0 , if µ θ ( - → z 1 ( - → z 0 , - → ϵ )) := θ - → z 1 ( - → z 0 , - → ϵ ) , then the closed-form solution is θ ∗ = √ 1 -β I k , which after normalization by 1 √ 1 -β and composition with the decoder D ( ← -z 0 ; ω ) = S ← z 0 recovers the true subspace of p ( - → x 0 ) .

With this optimal θ ∗ , we can now prove exact sample recovery using GML-DPS (6).

Theorem 3.5 (Posterior Sampling using Goodness Modified Latent DPS) . Let Assumptions 3.1 and 3.2 hold. Let σ j , ∀ j = 1 , . . . , r , denote the singular values of ( AS ) T ( AS ) , and let

Given a partially known image - → x 0 ∼ p ( - → x 0 ) , any fixed variance β ∈ (0 , 1) , then with the (unique) step size η j i = 1 / 2 σ j , j = 1 , 2 , . . . , r , the GML-DPS Algorithm (6) samples from the true posterior p ( - → x 0 | y ) and exactly recovers the groundtruth sample, i.e., ← -x 0 = - → x 0 .

Theorem 3.5 shows that GML-DPS (6) recovers the true sample using an LDM. This approach, however, requires the step size η 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 ← -z 1 for which ∇ ← -z 1 ∣ ∣ ∣ ∣ ← -z 0 ( ← -z 1 )] -E ( D ( ← -z 0 ( ← -z 1 ))) ∣ ∣ ∣ ∣ 2 = 0 . These challenges motivate the gluing objective in Theorem 3.6 . This is crucial for two reasons. First, we show that it helps recover the true sample even when the step size η is chosen arbitrarily. Second, it assigns all the probability mass to the desired (unique) solution in the pixel space.

Theorem 3.6 (Posterior Sampling using Diffusion in Latent Space) . Let Assumptions 3.1 and 3.2 hold. Let σ j , ∀ j = 1 , . . . , r denote the singular values of ( AS ) T ( AS ) and let

Given a partially known image - → x 0 ∼ p ( - → x 0 ) , any fixed variance β ∈ (0 , 1) , and any positive step sizes η j i , j = 1 , 2 , . . . , r , the PSLD Algorithm 2 samples from the true posterior p ( - → x 0 | y ) and exactly recovers the groundtruth sample, i.e., ← -x 0 = - → x 0 .

The important distinction between Theorem 3.5 and Theorem 3.6 is that the former requires the exact step size while the latter works for any finite step size. Combining denoising, measurementconsistency (with a scalar η ), and gluing updates, we have

When η is chosen arbitrarily, then the third term guides the reverse SDE towards the optimal solution - → z 0 . When the reverse SDE generates the exact same groundtruth sample, i.e., D ( ← -z 1 ( ← -z 0 )) = - → 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 η .

We draw the following key insights from our Theorem 3.6 : 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 η . Furthermore, it allows the same (scalar) step size across all the coordinates of - → x 0 .

4 Experimental Evaluation

We experiment with in-distribution and out-ofdistribution datasets. For in-distribution, we conduct our experiments on a subset of the FFHQ dataset [25] (downscaled to 256 × 256 3 , denoted by FFHQ 256). For out-of-distribution, we use images from the web and ImageNet dataset [17] (resized to 256 × 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 [44, 45]; and (ii) the Latent Diffusion model (LDM-VQ-4) trained on the FFHQ 256 dataset [41]. 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 × 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 [32]: (vi) Super-resolution downsamples images at 2 × , 3 × , and 4 × scale. (vii) Denoising has Gaussian noise with σ = 0 . 05 . (viii) Destriping has vertical and horizontal stripes in the input images.

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 indistribution 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 × 512 , we apply the forward operator after upsampling inputs to 512 × 512 , run posterior sampling at 512 × 512 , and then downsample images to the original 256 × 256 resolution for a fair comparison with DPS. We observed a similar performance while applying the masking operator at 256 × 256 and upscaling to 512 × 512 before running PSLD. More implementation details are provided in Appendix C.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 Distance 4 (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.

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 6 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 1 shows the super-resolution results using nearestneighbor kernels from [32] 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 × 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

(a) Input

(c) Comm. Serv. 1

(d) Comm. Serv. 2

(e) PSLD (Ours)

reconstruction-see Table 6. 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.6 ), 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 C.

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

CO

e

e

e

e

(a) Input

(c) DPS [9]

(d) PSLD (Ours)

(a) Input

(c) DPS [9]

(d) PSLD (Ours)

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

0.30

DPS

32

0.90

0.25

PSLD

SdI

PSNR

30

SSIM

0.85

0.20

LPI

0.15

28

DPS

0.80

DPS

PSLD

PSLD

0.10

26

0.75

20

40

60

80

20

40

60

80

20

40

60

80

Percentage ofdroppedpixels

Percentageofdroppedpixels

Percentageofdroppedpixels

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.

References

Brian D.O. Anderson. 'Reverse-time diffusion equation models'. In: Stochastic Processes and their Applications 12.3 (1982), pp. 313-326 (page 1).

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 (page 2).

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) (page 2).

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 (page 3).

Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G Dimakis. 'Compressed sensing using generative models'. In: International Conference on Machine Learning . PMLR. 2017, pp. 537546 (page 1).

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 (2016), pp. 84-98 (pages 2, 8, 24, 27).

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) (pages 5, 6).

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) (page 1).

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) (page 1).

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) (page 2).

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 (pages 1-3, 7-10, 15, 18, 20, 22, 24-31).

Hyungjin Chung, Jeongsol Kim, and Jong Chul Ye. 'Direct Diffusion Bridge using Data Consistency for Inverse Problems'. In: arXiv preprint arXiv:2305.19809 (2023) (page 2).

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 . Ed. by Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho. 2022. URL: https://openreview.net/forum?id=nJJjv0JDJju (pages 2, 8, 24, 27, 31).

Giannis Daras, Yuval Dagan, Alexandros G Dimakis, and Constantinos Daskalakis. 'Scoreguided intermediate layer optimization: Fast langevin mixing for inverse problem'. In: arXiv preprint arXiv:2206.09104 (2022) (page 2).

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) (page 2).

Mauricio Delbracio and Peyman Milanfar. 'Inversion by direct iteration: An alternative to denoising diffusion for image restoration'. In: arXiv preprint arXiv:2303.11435 (2023) (page 2).

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 'Imagenet: A largescale hierarchical image database'. In: 2009 IEEE conference on computer vision and pattern recognition . Ieee. 2009, pp. 248-255 (pages 7, 21, 22, 27-29).

Prafulla Dhariwal and Alexander Nichol. 'Diffusion models beat gans on image synthesis'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 8780-8794 (page 1).

Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, et al. 'DataComp: In search of the next generation of multimodal datasets'. In: arXiv preprint arXiv:2304.14108 (2023) (page 10).

Jonathan Ho, Ajay Jain, and Pieter Abbeel. 'Denoising diffusion probabilistic models'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 6840-6851 (page 1).

Aapo Hyvärinen and Peter Dayan. 'Estimation of non-normalized statistical models by score matching.' In: Journal of Machine Learning Research 6.4 (2005) (page 1).

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 (pages 1, 2).

Ajil Jalal, Sushrut Karmalkar, Alexandros G Dimakis, and Eric Price. 'Instance-optimal compressed sensing via posterior sampling'. In: arXiv preprint arXiv:2106.11438 (2021) (page 1).

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 . Ed. by Marina Meila and Tong Zhang. Vol. 139. Proceedings of Machine Learning Research. PMLR, 18-24 Jul 2021, pp. 4721-4732. URL: https://proceedings.mlr.press/v139/jalal21b.html (page 1).

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 (pages 7-10, 24-28, 31).

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. 'Denoising Diffusion Restoration Models'. In: Advances in Neural Information Processing Systems (pages 1, 2, 8, 24, 27, 31).

Bahjat Kawar, Noam Elata, Tomer Michaeli, and Michael Elad. 'GSURE-Based Diffusion Model Training with Corrupted Data'. In: arXiv preprint arXiv:2305.13128 (2023) (page 2).

Bahjat Kawar, Gregory Vaksman, and Michael Elad. 'SNIPS: Solving noisy inverse problems stochastically'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 2175721769 (page 31).

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 . PMLR. 2022, pp. 1120111228 (page 1).

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) (page 3).

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) (Oct. 2019). DOI: 10.1109/iccv.2019.00427 . URL: http://dx.doi.org/10.1109/ICCV. 2019.00427 (page 1).

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. 1146111471 (pages 7, 8, 30).

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 (page 1).

Xiangming Meng and Yoshiyuki Kabashima. 'Diffusion model based posterior sampling for noisy linear inverse problems'. In: arXiv preprint arXiv:2211.12343 (2022) (pages 27, 31).

Sachit Menon, Alexandru Damian, Shijia Hu, Nikhil Ravi, and Cynthia Rudin. 'Pulse: Selfsupervised photo upsampling via latent space exploration of generative models'. In: Proceedings of the ieee/cvf conference on computer vision and pattern recognition . 2020, pp. 24372445 (page 1).

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 (2020), pp. 39-56 (page 1).

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 (page 1).

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 . Springer. 2022, pp. 117-126 (page 3).

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) (page 1).

Yaniv Romano, Michael Elad, and Peyman Milanfar. 'The little engine that could: Regularization by denoising (RED)'. In: SIAM Journal on Imaging Sciences 10.4 (2017), pp. 1804-1844 (pages 1, 31).

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 (pages 2, 3, 5, 7).

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) (pages 5, 6, 15).

Chitwan Saharia, William Chan, Huiwen Chang, Chris A. Lee, Jonathan Ho, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Palette: Image-to-Image Diffusion Models . 2022. arXiv: 2111.05826 [cs.CV] (page 31).

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] (page 7).

Christoph Schuhmann et al. LAION-5B: An open large-scale dataset for training next generation image-text models . 2022. arXiv: 2210.08402 [cs.CV] (page 7).

Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. 'Pseudoinverse-guided diffusion models for inverse problems'. In: International Conference on Learning Representations . 2023 (pages 2, 31).

Yang Song and Stefano Ermon. 'Generative modeling by estimating gradients of the data distribution'. In: Advances in Neural Information Processing Systems 32 (2019) (page 1).

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 (page 1).

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 (page 1).

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 (pages 8, 24, 27).

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 (page 3).

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 . IEEE. 2013, pp. 945-948 (page 1).

Pascal Vincent. 'A connection between score matching and denoising autoencoders'. In: Neural computation 23.7 (2011), pp. 1661-1674 (page 1).

Su Wang, Chitwan Saharia, Ceslee Montgomery, Jordi Pont-Tuset, Shai Noy, Stefano Pellegrini, Yasumasa Onoe, Sarah Laszlo, David J Fleet, Radu Soricut, et al. '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 (pages 1, 21).

Yinhuai Wang, Jiwen Yu, and Jian Zhang. 'Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model'. In: The Eleventh International Conference on Learning Representations . 2023. URL: https://openreview.net/forum?id=mRieQgMtNTQ (page 31).

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) (Oct. 2019). DOI: 10.1109/iccv.2019.00457 . URL: http: //dx.doi.org/10.1109/ICCV.2019.00457 (page 1).

A Additional Theoretical Results

Notation and Measurement Matrix. We elaborate on the structure of the measurement matrix A ∈ R l × 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 × d identity matrix (the rows that are present correspond to the indices of the selected pixels from the image - → x 0 ∈ R d ). Given this structure, it immediately follows that A T A is a d × d matrix that has the interpretation of a pixel selection mask . Specifically, A T A is a d × d diagonal matrix D ( m ) , where the elements of 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.

A.1 Posterior Sampling using Pixel-space Diffusion Model

We first consider the reverse process, starting with ← -x 1 ∼ N ( 0 , I d ) , and borrow a result from [42] to show that the sample ← -x 0 generated by the reverse process is a valid image from p ( - → x 0 ) .

Theorem A.1 (Generative Modeling using Diffusion in Pixel Space, [42]) . Suppose Assumption 3.1 holds. Let

For a fixed variance β > 0 , if µ θ ( - → x 1 ( - → x 0 , - → ϵ )) := θ - → x 1 ( - → x 0 , - → ϵ ) , then the closed-form solution θ ∗ is √ 1 -β SS T , which after normalization by 1 / √ 1 -β recovers the true subspace of p ( - → x 0 ) .

Though this establishes that ← -x 0 generated by the reverse process is a valid image from p ( - → x 0 ) , it is not necessarily a sample from the posterior p ( - → x 0 | 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 sample 5 .

Theorem A.2 (Posterior Sampling using Diffusion in Pixel Space) . Suppose Assumption 3.1 and Assumption 3.2 hold. Let us denote by σ j , ∀ j = 1 , . . . , r , the singular values of ( AS ) T ( AS ) and

Given a partially known image - → x 0 ∼ p ( - → x 0 ) , a fixed variance β > 0 , there exists a step size ζ j i = 1 / 2 σ j for all the coordinates of - → x 0 such that Algorithm 1 samples from the true posterior p ( - → x 0 | y ) and exactly recovers the groundtruth sample, i.e., ← -x 0 = - → x 0 .

B Technical Proofs

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.

B.1 Proof of Theorem A.2

Theorem B.1 (Posterior Sampling using Diffusion in Pixel Space) . Suppose Assumption 3.1 and Assumption 3.2 hold. Let us denote by σ = { σ j } k j =1 the singular values of ( AS ) T ( AS ) , i.e. ( AS ) T ( AS ) = U Σ V T := UD ( σ ) V T , U ∈ R k × k , V ∈ R k × k and

Suppose - → x 0 ∼ p ( - → x 0 ) . Given measurements y = A - → x 0 and a fixed variance β ∈ (0 , 1) , there exists a matrix step size 6 ζ = (1 / 2)( S U ) D ( ζ i )( S U ) T , ζ i = { ζ j i = 1 /σ j } k j =1 for all the coordinates of - → x 0 such that Algorithm 1 samples from the true posterior p ( - → x 0 | y ) and exactly recovers the groundtruth sample, i.e., ← -x 0 = - → x 0 .

Proof. Our goal is to show that ← -x 0 = - → x 0 , where ← -x 0 is returned by Algorithm 1 . Recall that the reverse process starts with ← -x 1 ∼ N ( 0 , I d ) and generates the following:

Now, we use the singular value decomposition of ( AS ) T ( AS ) with left singular vectors in U ∈ R k × k , right singular vectors in V ∈ R k × k , and singular values σ = [ σ 1 , . . . , σ k ] in Σ = D ( σ ) . Thus, the above expression becomes

where (i) is due to Assumption 3.1 and (ii) uses Assumption 3.2 . By choosing ζ j i as half the inverse of the non-zero singular values of ( AS ) T ( AS ) , i.e., ζ j i = 1 / 2 σ i ∀ i = 1 , . . . , k , we obtain

which completes the statement of the theorem.

B.2 Proof of Proposition 3.3

Proposition B.2 (Variational Autoencoder) . Suppose Assumption 3.1 holds. For an encoder E : R d → R k and a decoder D : R k → R d , denote by L ( ϕ, ω ) the training objective of VAE:

□

Proof. To show that the encoder E ( - → x 0 ; ϕ ) = S T - → x 0 and the decoder D ( ← -z 0 ; ω ) = S ← z 0 minimize the VAE training objective L ( ϕ, ω ) , we begin with the first part of the loss, which is also called reconstruction error L recon ( ϕ, ω ) . Substituting E ( - → x 0 ; ϕ ) = S T - → x 0 and D ( ← -z 0 ; ω ) = S ← z 0 , we have

Using the fact that - → x 0 lives in a linear subspace, we arrive at

where (i) is due to Assumption 3.1 . Now, we analyze the distribution loss. Note that the KLdivergence between two Gaussian distributions with moments ( µ 1 , σ 1 ) and ( µ 2 , σ 2 ) is given by

Since E ( x 0 ) = S T x 0 = S T S z 0 = z 0 , the distribution loss becomes:

B.3 Proof of Theorem 3.4

Theorem B.3 (Generative Modeling using Diffusion in Latent Space) . Suppose Assumption 3.1 holds. Let the optimal solution of the latent diffusion model be

For a fixed variance β > 0 , if µ θ ( - → z 1 ( - → z 0 , - → ϵ )) := θ - → z 1 ( - → z 0 , - → ϵ ) , then the closed-form solution is θ ∗ = √ 1 -β I k , which after normalization by 1 √ 1 -β and composition with the decoder D ( ← -z 0 ; ω ) := S ← z 0 recovers the true subspace of p ( - → x 0 ) .

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

where θ T i denotes the i th row of matrix θ . The solution of this regression problem is given by 7

Using Assumption 3.1 , the above expression simplifies to

where the last step uses the fact that z 0 and ϵ are independent Gaussian random vectors with zero mean and unit covariance. Simplifying further, we arrive at

The final step follows from independence of z 0 and ϵ . Since z 0 and ϵ are also N ( 0 , I k ) , we get

where the i th coordinate is √ 1 -β and zero everywhere else. Therefore, stacking all the rows together, we get θ ∗ = √ 1 -β I k , which after normalization by 1 / √ 1 -β gives the desired result.

Next, we show that θ ∗ recovers the true subspace of - → x 0 ∼ p ( - → x 0 ) . When composed with the decoder of VAE, the generator of the LDM gives ← -x 0 = D ( θ ∗ ← -z 1 ) = D ( I k ← -z 1 ) = S ← z 1 . Since ← -z 1 ∼ N ( 0 , I k ) , this completes the statement of the theorem. □

B.4 Proof of Theorem 3.5

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.4 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:

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 B.4 (Posterior Sampling using Goodness Modified Latent DPS) . Suppose Assumptions 3.1 and Assumption 3.2 hold. Denote by σ = { σ j } k j =1 the singular values of ( AS ) T ( AS ) , i.e., ( AS ) T ( AS ) = U Σ U T := UD ( σ ) U T , U ∈ R k × k , and let

Suppose - → x 0 ∼ p ( - → x 0 ) . Given measurements y = A - → x 0 and any fixed variance β ∈ (0 , 1) , then with the (unique) step size η = (1 / 2) UD ( η i ) U T , η i = { η j i = 1 / 2 σ j } k j =1 , the GML-DPS algorithm (6) samples from the true posterior p ( - → x 0 | y ) and exactly recovers the groundtruth sample, i.e., ← -x 0 = - → 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

where (i) is due to Assumption 3.1 . By Assumption 3.2 , ( AS ) T ( AS ) is a positive definite matrix and can be written as U Σ U T :

Since η i j = 1 / 2 σ j , the above expression further simplifies to

Next, we show that ← -z ′ 0 is already a minimizer of (9). This is a direct consequence of the encoderdecoder architecture of the V AE: E ( D ( ← -z ′ 0 )) = S T S ← -z ′ 0 = ← -z ′ 0 . Hence, ∣ ∣ ∣ ∣ ∣ ∣ ← -z ′ 0 -E ( D ( ← -z ′ 0 )) ∣ ∣ ∣ ∣ ∣ ∣ 2 = 0 , and

consequently ← -z 0 = ← -z ′ 0 -γ ∇ ← -z ′ 0 ∣ ∣ ∣ ∣ ∣ ∣ ← -z ′ 0 -E ( D ( ← -z ′ 0 )) ∣ ∣ ∣ ∣ ∣ ∣ 2 = - → z 0 . Thus, the reconstructed sample becomes ← -x 0 = D ( ← -z 0 ) = S - → z 0 = - → x 0 .

Furthermore, as ∣ ∣ ∣ ∣ ∣ ∣ ← -z ′ 0 -E ( D ( ← -z ′ 0 )) ∣ ∣ ∣ ∣ ∣ ∣ 2 = 0 for all ← -z ′ 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. □

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

B.5 Proof of Theorem 3.6

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

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 B.5 (Posterior Sampling using Diffusion in Latent Space) . Let Assumptions 3.1 and 3.2 hold. Let σ j , ∀ j = 1 , . . . , r denote the singular values of ( AS ) T ( AS ) and let

Suppose - → x 0 ∼ p ( - → x 0 ) . Given measurements y = A - → x 0 , any fixed variance β ∈ (0 , 1) , and any positive step sizes η j i , j = 1 , 2 , . . . , r , the PSLD Algorithm 2 samples from the true posterior p ( - → x 0 | y ) and exactly recovers the groundtruth sample, i.e., ← -x 0 = - → x 0 .

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

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

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

□

It is worth highlighting that PSLD exactly recovers the groundtruth sample irrespective of the choice of the step size η , whereas GML-DPS requires the step size to be exactly η = (1 / 2) UD ( η i ) U T .

C Additional Experiments

C.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., ← -x 0 = A T A - → x 0 + ( I d -A T A ) D ( ← -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 [54]. This sometimes creates edge effects which are then removed by postprocessing 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., ← -x 0 = D ( ← -z 0 ) . Furthermore, we use the same measurement operator A and its transpose A T as provided by the DPS code repository 8 . However, since Stable Diffusion v1.5 generates images of size 512 × 512 resolution and DPS operates at 256 × 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 × 512 to 256 × 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 model 9 . We use a fixed η = 1 and γ = 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-4 10 , which contains pretrained weights for FFHQ 256 11 and large-scale text-to-image generative model 12 . We keep the hyperparameters same ( η = 1 and γ = 0 . 1 ). For each task, we provide hyper-parameter details in our codebase 13 . 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 authors 14 .

OODimages are sourced online:

Figure 1: the original images are generated by Stable Diffusion v-2.1 15 .

Figure 2 first row: Walking example from the web.

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

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

Figure 4 first row: Racoon image from the web.

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

Figure 15: Celebrity face from the web.

PSLD Image Inpainting

Image inpainting by Posterior Sampling with Latent Diffusion (PSLD)

sed uss ajeauas ouedu ux ysew pauap asn e pue paajad azis aen ewue ua

Upload

Output1

Output2

Number of difusion steps (e.g. 20o)

200

0.1

Gluing kernel size (e.g. 15)

Gluing kernel sigma (e.g.7)

Measurement factor (e.g. 1)

Your prompt (leave empty for posterior sampling)

Inpaint!

C.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 10 and 11 illustrate super-resolution (4 × ) 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 V AE 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. Figure 17 shows the results on image colorization. Table 5 and Table 6 show the quantitative results. Table 7 draws a comparison between the latent-DPS and PSLD algorithms, and shows that the PSLD objective enhances the reconstruction performance.

In Table 8, we compare the runtime and NFEs of PSLD with prior works. PSLD-SD (trained on LAION-5B) takes 776 s to generate 512x512 images. To compare with other methods that generate 256x256 images, we divide our runtime by 4. All the other methods use diffusion models trained on FFHQ and produce 256x256 images.

PSLDImageInpainting

Image inpainting by Posterior Sampling with Latent Diffusion (PSLD)

sed gussw ajeauaa o ujedu uo x ysew pauyap jasn e pue (pauajaud azis auenbs) aewi ue uani

Upload

Output1

Number of diffusion steps (e.g. 20o)

200

Output2

Gluing factor (e.g. 1e-1)

0.1

Gluing kernel size (e.g 15)

15

Gluing kernel sigma (e.g. 7)

Measurement factor (e.g. 1)

Your prompt (leave empty for posterior sampling)

Inpaint!

PSLD Image Inpainting

Image inpainting by Posterior Sampling with Latent Diffusion (PSLD)

Given an image (square size preferred) and a user defined mask, click on Inpaint to generate missing parts.

Upload

Output1

zandno

Number of diffusion steps (e.g. 20o)

200

Gluing factor (e.g. 1e-1)

0.1

Gluing kernel size (e.g 15)

15

Gluing kernel sigma (e.g. 7)

Measurement factor (e.g. 1)

Your prompt (leave empty for posterior sampling)

PSLDImageInpainting

Image inpainting by Posterior Sampling with Latent Diffusion (PSLD)

Given an mage (square size preferred) and a user defined mask, click on Inpaint to generate missing parts.

Upload

Output1

Number of diffusion steps (e.g. 20o)

200

Output2

Gluing factor (e.g. 1e-1)

0.1

Gluing kernel size (e.g 15)

15

Gluing kernel sigma (e.g. 7)

Measurement factor (e.g. 1)

Your prompt (leave empty for posterior sampling)

edu

D Additional Discussion

Curse of ambient dimension: DPS [11] suffers from the curse of ambient dimension because in this method, gradients are computed in the pixel space with dimension d . However, latent-based methods such as PSLD compute gradients in the latent dimension k , and hence the computation is more efficient. Furthermore, applying the chain rule on V AE and running diffusion in the latent space is less expensive than running diffusion in pixel space directly. In practice, the computational complexity of Stable Diffusion model ( ∼ 4 GB ) is higher (roughly 6 times) than the computational complexity of the encoder-decoder model ( ∼ 700 MB ). Therefore, applying the chain rule in the encoder-decoder and running diffusion in the latent space is less expensive than applying diffusion models in the pixel space directly.

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

(c) DPS [11]

(d) PSLD (Ours)

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

www.hnwhite.com

wwww.inwhite.com

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

(a) Input

(c) DPS [11]

(d) PSLD (Ours)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method to deal with linear inverse problems using pre-trained latent diffusion models (LDM) as a prior. The main idea is to extend the original diffusion posterior sampling (DPS) to the case of LDM by approximating the gradient term of the intractable likelihood.  Two approximations (GLM-DPS), and (PSLD) are proposed for such approximation, and it turns out that PSLD achieves better performances. This paper also theoretically analyzes the DPS and PSLD in a toy linear setting. Experiment results show superior performances compared with the original DPS in a variety of tasks, such as random/block inpainting, denoising, super-resolution.\n\nSTRENGTHS:\n1. It is the first work, to my best knowledge, that LDM is used to address the linear inverse problems. \n\n2. Compared to the original DPS, the proposed PSLD achieves apparently better performances in most cases. \n\n3. Some theoretical analysis is also provided, although it is based on some toy linear model.\n\nWEAKNESSES:\n1.  There is a lack of analysis of the complexity, or running time, of the proposed PSLD method. Adding details of the running time of PSLD and other comparison methods is suggested. \n\n2. Part 3 contains some theoretical analysis using a simple toy linear model. For example, Theorem 3.4 states that DPS exactly recovers the groundtruth sample. It is known that DPS uses a Laplace approximation to approximate the gradient term of the likelihood. As a result, how can it exactly recover the posterior samples with the crude approximation? This is also the case with Theorem 3.7 and Theorem 3.8.  To what extent can the results of Theorem 3.4, 3.7, and 3.8 be generalized to the real diffusion models?\n\n3. In the experiments part, while PSLD shows superior performances over the original DPS, there is a lack of comparison with other latest diffusion-model-based algorithms that were available long before NeurIPS submission, e.g., \n\n[A] Wang, Yinhuai, Jiwen Yu, and Jian Zhang. \"Zero-shot image restoration using denoising diffusion null-space model.\" arXiv preprint arXiv:2212.00490 (2022).\n\n[B] Meng, Xiangming, and Yoshiyuki Kabashima. \"Diffusion model based posterior sampling for noisy linear inverse problems.\" arXiv preprint arXiv:2211.12343 (2022).\n\n[C] Song, Jiaming, Arash Vahdat, Morteza Mardani, and Jan Kautz. \"Pseudoinverse-guided diffusion models for inverse problems.\" In International Conference on Learning Representations. 2022.\n\n4. Most previous image restoration methods using diffusion models also consider colorization, can the authors also add some results in the case of colorization? Besides, only Table 2 and Table 5 show quantitative comparisons with other methods but only for inpainting. It is suggested to add quantitative comparisons with other methods in other tasks such as super-resolution, denoising, colorization, etc.\n\nQUESTIONS:\nThree additional questions:\n\n1. How do you add the additive Gaussian noise with standard deviation \\sigma_y? \n\n\n2. Can I understand that the improved performance of PSLD over DPS comes solely from the improvement of diffusion models in the latent space over that in the pixel space? In other words, if DPS is implemented with more powerful diffusion models, is it possible to outperform PSLD? In this sense, how to ensure a fair comparison between PSLD and DPS? \n\n3. Are the results of PSLD sensitive to the choice of hyper-parameters? Compared to DPS which has only one scaling hyperparameter, there are two scaling hyperparameters, namely, \\eta and \\gamma. It is suggested to evaluate the sensitivity of these hyperparameters.\n\nLIMITATIONS:\nPlease refer to the weakness and questions parts.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on solving inverse problems using diffusion based probabilistic models without **retraining**. The authors build on \"Diffusion posterior sampling\" which basically builds a diffusion model for the posterior using only the score of the prior distribution. Indeed, the score of the posterior $\\nabla \\log p_t(x_t | y)$ is equal to $\\nabla \\log p(y|x_t) + \\nabla \\log p_t(x_t)$ where the second term is given by the prior diffusion model. The first term is however intractable but can be approximated by noting that $p(y | x_t) = \\int p(y|x_0) p(x_0 | x_t) dx_0$. This integral is then approximated by simply taking the mean a posteriori i.e. $p(y | x_t) \\approx p(y | E(x_0 | x_t))$ which  itself can be approximated using Tweedie's formula and the prior score. The original DPS paper only considers the diffusion in the pixel space and this paper extends the methodology to latent diffusion models. While the direct extension is straightforward, the authors claim that it does not work in practice and they provide a rather intuitive modified posterior diffusion model that ensures consistency at the borders of the mask.\n\nSTRENGTHS:\nThe paper tackles an important problem and provides a sound methodology for diffusion posterior sampling for latent space diffusion models. Furthermore, the numerical experiments suggest that this method outperforms other existing methods.\n\nWEAKNESSES:\nThe fact that the authors show recovery in a 2 step diffusion model in a very basic setting proves nothing. First, the results hold only when the optimal solution is normalized. In realistic settings and with multi step diffusion one does not know in advance how to rescale (if even rescaling is the right to do?) the drift so that the diffusion model matches the target distribution. Next, the authors prove that DPS \"provably\" recovers the posterior but this to me is also meaningless since one can easily show numerically in a noisy inpainting setting with Gaussian mixtures (the posterior is tractable in this case) that DPS **does not** in fact sample the posterior. It samples outside the support of the posterior even in the simplest settings. I believe that the authors should provide a numerical experiment where the target posterior is known (again, Gaussian mixtures) and show that their method does sample from the posterior. To me there is no reason why it should exactly sample from the posterior and not outside the support like DPS. Indeed, this issue is not related to the quality of the generative model but more to the fact that DPS (and truthfully many of the difussion posterior sampling methods) are not exact samplers since there is no correction like in MCMC algorithms. I believe that the authors should drop the \"provably\" in the title as it is misleading.\n\nQUESTIONS:\n- What do the authors mean by the *curse of ambient dimension*, line 226-228? In DPS the gradient is indeed computed in the $d$-dimensional space but this is also the case of PSLD by the chain rule (and in fact the computational cost of PSLD is slightly larger).\n\n- In Figure 2 right panel, an inpainting example is considered. In this case p(y|x_0) is a dirac delta so I am not sure how the authors manage to take the gradient. Do the authors approximate the dirac delta with a Gaussian with small variance? Was the same courtesy applied to DPS?  \n\n- It is quite surprising that in Algorithm 2 (but also Algo1) nothing depends on the measurement variance. Why is it not taken into account? \n\n- While I could guess what it respresents, $x_0 *$ is not defined anywhere.\n\nLIMITATIONS:\nsee below.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhile the preivous methods have focused on solving linear inverse problems based on diffusion models, the paper presents a first extension to latent diffusion models (LDM). The core idea is developed upon the existing DPS method, which forms an approximation to p(y|xt) by using the denoising score estimate. The key challenges of the extension to LDM are (1) the non-bijective mapping between encoder and decoder, and (2) the inconsistency at the boundary of mask (for inpainting problems). To address these challenges, the authors propose an additional gluing term that penalizes the discontinuity at the mask boundary.\n\nThe main practical contribution of this paper is that if unlocks the potential of large pre-trained LDM (e.g. stable diffusion), and hence the proposed based based on more powerful pre-trained models enjoys superior performance compared to its counterpart with ambient / pixel space settings. \n\nIn addition, the paper also presents a theoretical analysis in a toy setting (linear two-step diffusion models) to motivate the proposed objective.\n\nSTRENGTHS:\n- While the extension of DPS to LDMs seems natural, I appreciate the practical impact of the proposed methodology given the popularity and superior power of existing pre-trained LDMs. \n\n- If I understand it correctly, doing conditional inference in latent space would also be more efficient than doing inference in the ambient space (after controlling the model complexity though in practice I wonder if that comparison can be made) \n\n- The proposed method is also claimed to be robustness of the choice of step-size, at least in toy settings, though I am not sure how this would extend to general case.\n\nWEAKNESSES:\n- The theoretical analysis seems to follow reference 30 closely, takes up a large portion of the paper. However, I am not sure how much insights can be carried over to the general cases. In my opinion, the proposed method stays well-motivated without this analysis. Hence. I believe either discussing the extension to general cases, or trimming this part would make a better paper. \n\n- The evaluation part can be misleading given that pre-trained LDMs are \"trained on much more data compared to the one used by DPS\" (ref: appendix, caption above table 5). I think this information should be disclosed in the main paper instead of the appendix. In addition, it would be helpful to also disclose the unconditional generation performance of all backbone diffusion models, as well as the size of the models, the number of training data, etc.\n\nQUESTIONS:\n- Theorem 3.8: what is the value for gamma_i in Alg 2? \n\n- Theorem 3.8 and its surrounding text: can you elaborate more on why it is robust to different values of step size eta? In the equation between line 221 and line 222 on page 7, apparently if I choose different value of eta, the resulting z0 would be different (as long as the grad norm after eta is non-zero). Am I missing anything here? In addition, in this equation, should there be a multiplicative factor \\gamma before the last term, or is it set to 1 for the theorem and the statement to be correct? \n\n- Page 7. line 228 \"DPS algorithm suffers from the curse of ambient dimension\". Can you elaborate on the issue? Is it mainly the efficiency issue or the approximation accuracy issue?\n\n- Would be helpful to accomplish the theoretical analysis with toy data empirical verification. \n\n- (If possible), would be helpful to provide results where the base diffusion models for DPS and PSLD have similar performance.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the use of diffusion models for solving inverse problems. While many recent papers have explored diffusion models for inverse problems, to the best of my knowledge, this study is the first to propose a method for inverse problems with latent diffusion. One advantage of latent diffusion is its lower computational demand.\n\nThe main challenge in solving inverse problems with diffusion models lies in computing the intractable likelihood term $\\nabla \\log p_{x_t}( y | x_t )$, where $x_t$ represents a noisy point in pixel-space. Recently, the DPS algorithm was introduced to approximate this term as $\\nabla \\log p_{x_t}( y | \\mathbb{E}[x_0| x_t] )$. The first contribution of this paper is to extend the DPS framework to a latent diffusion model, approximating the intractable likelihood $\\nabla \\log p_{x_t}( y | z_t )$, where $z_t$ denotes a noisy point in latent space, with $\\nabla \\log p_{z_t}( y | \\mathbb{E}[z_0| z_t] )$ along with an additional term that enforces $\\mathbb{E}[z_0| z_t]$ to be a fixed point of the autoencoder (GML-DPS). Furthermore, the latter term is modified to enforce consistency in the measurements (PSLD). The empirical performance of the PSLD algorithm is compared with the standard DPS in high-dimensional reconstruction tasks, and improvements over DPS are demonstrated.\n\nThese three algorithms (DPS, GML-DPS, PSLD) are theoretically studied in the case of a data distribution corresponding to a Gaussian supported on a low-dimensional subspace, perfect knowledge of the subspace, zero noise in the measurements, and a measurement matrix that is bijective over the subspace. Specifically, this implies that the inverse problem can be exactly solved given the measurements. Finally, the sampling processes are two-step diffusion processes. The authors demonstrate that all three algorithms precisely recover the target signal, and PSLD exhibits robustness to variations in the specification of the step sizes.\n\nSTRENGTHS:\nSolving inverse problems with (large) latent diffusion models is a relatively underexplored area, and this paper fills a gap in the existing literature. In particular, the authors demonstrate how their methods can be used with SOTA latent-based foundation models such as stable diffusion. The empirical results presented are promising and demonstrate the potential of this approach. \n\nAdditionally, the authors offer some theoretical guarantees, albeit in a limited setting. These guarantees are valuable and noteworthy, considering the field's scarcity of rigorous results, even in simple toy examples.\n\nWEAKNESSES:\nThe main weaknesses of the paper lie in the experimental evaluations and results. Specifically, all experiments are conducted on subsets of the FFHQ dataset and images sourced from the internet. The authors evaluate the PSLD algorithm on two latent models: LDM-VQ-4, trained on FFHQ, and Stable Diffusion, trained on a significantly larger dataset (LAION). However, the comparisons are made between PSLD on these two latent models  (proposed) and DPS (baseline) on a standard diffusion model trained *exclusively* on FFHQ.\n\nThe authors demonstrate improved performance of PSLD on the (latent) Stable Diffusion model compared to DPS on the standard diffusion model. However, when employing the latent model LDM-VQ-4 trained solely on FFHQ, the improvements are marginal (Table 3).\n\nAs a result, it becomes challenging to ascertain whether the observed improvements stem from algorithmic enhancements or simply the utilization of superior models trained on larger datasets, which may also account for the enhanced out-of-distribution performance.\n\nFurthermore, the experiments do not shed light on the improvement attributable to the addition of the \"gluing\" term in comparison to the DPS's vanilla extension term. One might hypothesize that the \"gluing\" term could be employed in DPS (by removing the encoder-decoder), potentially enhancing the algorithm's performance.\n\nQUESTIONS:\nSome typos and questions:\n\n- In Section 2, the notation can be confusing. To clarify, it is suggested to use $x_0^*$ to denote the target signal and distinguish it from the first sample of the backward diffusion ($x_0$).\n\n- For Equations (3), should it be $\\hat{x}_0 = \\mathbb{E}[x_0 | x_t]$?\n\n- In Algorithm 1: line 1 what is $\\mathcal{T}$?\n\n- Equations (7) should include a $\\log$ term for the DPS vanilla extension. Additionally, as shown in the supplementary material, $\\mathcal{A}^\\top$ should be added to (7) and line 8 of Algorithm 2.\n\n- In Algorithm 2, it appears that the unknown target $x_0^*$ is needed, but in reality, the algorithm only requires the measurements $y$. It is advisable to modify the algorithm to reflect this (as in Algorithm 1).\n\n- Line 142 of the Theoretical section introduces the noisy inverse problem as $y = A x_0 + \\sigma_y n$.  However, it should be noted that exact recovery is not possible in this case and it seems that the theoretical results are proven under the assumption of $\\sigma_y = 0$, where the problem can be solved exactly. The authors should explicitly address these points to ensure clarity.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}