SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Litu Rout Negin Raoof Giannis Daras
Constantine Caramanis Alexandros G. Dimakis Sanjay Shakkottai
The University of Texas at Austin*
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.
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(\pmb{x}_0)$  using an Itô Stochastic Differential Equation (SDE),  $\mathrm{d}\pmb{x} = \pmb{f}(\pmb{x},t)\mathrm{d}t + g(t)\mathrm{d}\pmb{w}$ , and then by learning the score-function,  $\nabla_{\pmb{x}_t}\log p_t(\pmb{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  $\nabla_{\pmb{x}_t}\log p_t(\pmb{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  $\pmb{y} = \mathcal{A}\pmb{x}_0 + \sigma_y\pmb{n}$ , we can
*Email: {litu.rout,neginmr,giannisdaras,constantine,sanjay.shakkottai}utexas.edu, dimakis@austin.utexas.edu
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 C for web demo and image sources).
provably sample from the distribution  $p_0(\pmb{x}_0|\pmb{y})$  by running a modified Reverse SDE that depends on the unconditional score  $\nabla_{\pmb{x}_t}\log p_t(\pmb{x}_t)$  and the term  $\nabla_{\pmb{x}_t}\log p(\pmb{y}|\pmb{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 = \mathcal{E}(\boldsymbol{x}_0)$ , perform the
diffusion in the latent space and use a decoder,  $\mathcal{D}(\pmb{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 wide range 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).
(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  $\pmb{x}$ , bold upper-case  $\pmb{X}$ , and normal lower case  $x$  denote a vector, a matrix, and a scalar variable, respectively. We denote by  $\odot$  element-wise multiplication.  $D(\pmb{x})$  represents a diagonal matrix with entries  $\pmb{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  $\pmb{x} \in p$ , the sample  $\mathcal{E}(\pmb{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,  $s_{\theta}(\pmb{x}_t,t)$ , to learn the score-function,  $\nabla_{\pmb{x}_t}\log p_t(\pmb{x}_t)$ , at all levels  $t$ , of a stochastic process described by an Itô SDE:
$$
\mathrm {d} \boldsymbol {x} = \boldsymbol {f} (\boldsymbol {x}, t) \mathrm {d} t + g (t) \mathrm {d} \boldsymbol {w}, \tag {1}
$$
where  $\boldsymbol{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  $\boldsymbol{y} = \mathcal{A}x_0 + \sigma_y\boldsymbol{n}$ , one can sample from the distribution  $p_0(\boldsymbol{x}_0|\boldsymbol{y})$  by running the conditional Reverse SDE given by:
$$
\mathrm {d} \boldsymbol {x} = (\boldsymbol {f} (\boldsymbol {x}, t) - g ^ {2} (t) (\nabla_ {\boldsymbol {x} _ {t}} \log p _ {t} (\boldsymbol {x} _ {t}) + \nabla_ {\boldsymbol {x} _ {t}} \log p (\boldsymbol {y} | \boldsymbol {x} _ {t})) \mathrm {d} t + g (t) \mathrm {d} \boldsymbol {w}. \tag {2}
$$
As mentioned,  $\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol {y}|\boldsymbol {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:
$$
p (\boldsymbol {y} \mid \boldsymbol {x} _ {t}) \approx p \left(\boldsymbol {y} \mid \hat {\boldsymbol {x}} _ {0} := \mathbb {E} [ \boldsymbol {x} _ {0} \mid \boldsymbol {x} _ {t} ]\right) = \mathcal {N} (\boldsymbol {y}; \mu = \mathcal {A} \mathbb {E} [ \boldsymbol {x} _ {0} \mid \boldsymbol {x} _ {t}, \Sigma = \sigma_ {y} ^ {2} I). \tag {3}
$$
Essentially, DPS substitutes the unknown clean image  $\pmb{x}_0$  with its conditional expectation given the noisy input,  $\mathbb{E}[\pmb{x}_0|\pmb{x}_t]$ . Under this approximation, the term  $p(\pmb{y}|\pmb{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(\pmb{x}_0|\pmb{y})$  (Section A.1). Then, we provide an algorithm (Section 2.1) and its analysis to sample from  $p(\pmb{x}_0|\pmb{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.
The 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.
In Latent Diffusion Models, the diffusion occurs in the latent space. Specifically, we train a model  $s_{\theta}(z_t,t)$  to predict the score  $\nabla_{z_t}\log p_t(z_t)$ , of a diffusion process:
$$
\mathrm {d} \boldsymbol {z} = \boldsymbol {f} (\boldsymbol {z}, t) \mathrm {d} t + g (t) \mathrm {d} \boldsymbol {w}, \tag {4}
$$
where  $\boldsymbol{z}_0 = \mathcal{E}(\boldsymbol{x}_0)$  for some encoder function  $\mathcal{E}(\cdot): \mathbb{R}^d \to \mathbb{R}^k$ . During sampling, we start with  $\boldsymbol{z}_T$ , we run the Reverse Diffusion Process and then we obtain a clean image by passing  $\boldsymbol{z}_0 \sim p_0(\boldsymbol{z}_0 | \boldsymbol{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(\pmb{y}|\pmb{z}_t)$  with:
$$
p (\boldsymbol {y} | \boldsymbol {z} _ {t}) \approx p (\boldsymbol {y} | \boldsymbol {x} _ {0} = \mathcal {D} \left(\mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ]\right)), \tag {5}
$$
i.e., to approximate the unknown clean image  $\boldsymbol{x}_0$  with the decoded version of the conditional expectation of the clean latent  $\boldsymbol{z}_0$  given the noisy latent  $\boldsymbol{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  $\boldsymbol{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  $\boldsymbol{z}_t$  towards any of these latents  $\boldsymbol{z}_0$ , potentially in different directions. On the other hand, the score-function is pulling  $\boldsymbol{z}_t$  towards a specific  $\boldsymbol{z}_0$  that corresponds to the best denoised version of  $\boldsymbol{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(\pmb{y}|\pmb{z}_t)$  with:
$$
\nabla_ {\boldsymbol {z} _ {t}} \log p (\boldsymbol {y} | \boldsymbol {z} _ {t}) = \underbrace {\nabla_ {\boldsymbol {z} _ {t}} \log p (\boldsymbol {y} | \hat {\boldsymbol {x}} _ {0} = \mathcal {D} \left(\mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ]\right))} _ {\text {D P S v a n i l l a e x t e n s i o n}} + \gamma_ {t} \underbrace {\nabla_ {\boldsymbol {z} _ {t}} \left| \left| \mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ] - \mathcal {E} \left(\mathcal {D} (\mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ])\right) \right| \right| ^ {2}} _ {\text {" g o o d n e s s " o f z _ {0}}} \tag {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  $z_0$  as long as  $\mathcal{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.
$$
\begin{array}{l} \nabla_ {\boldsymbol {z} _ {t}} \log p (\boldsymbol {y} | \boldsymbol {z} _ {t}) = \underbrace {\nabla_ {\boldsymbol {z} _ {t}} \log p (\boldsymbol {y} | \boldsymbol {x} _ {0} = \mathcal {D} \left(\mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ]\right))} _ {\text {D P S v a n i l l a e x t e n s i o n}} \\ + \gamma_ {t} \underbrace {\nabla_ {z _ {t}} \left| \left| \mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ] - \mathcal {E} \left(\mathcal {A} ^ {T} \boldsymbol {y} + (\boldsymbol {I} - \mathcal {A} ^ {T} \mathcal {A}) \mathcal {D} \left(\mathbb {E} [ \boldsymbol {z} _ {0} | \boldsymbol {z} _ {t} ]\right)\right) \right| \right| ^ {2}} _ {\text {" g l u i n g " o f } z _ {0}}. \tag {7} \\ \end{array}
$$
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.
Remark 2.1. Consider the optimization problem of projecting onto the measurements:
$$
\min  _ {\boldsymbol {x} _ {0}} \left\| \hat {\boldsymbol {x}} _ {0} - \boldsymbol {x} _ {0} \right\| _ {2} ^ {2}
$$
$$
\text {s u b j e c t} \mathcal {A} x _ {0} = y,
$$
In the linear setting, the optimal solution is given by  $x_0^* = \mathcal{A}^T(\mathcal{A}\mathcal{A}^T)^{-1}\mathbf{y} + (\hat{\mathbf{x}}_0 - \mathcal{A}^T(\mathcal{A}\mathcal{A}^T)^{-1}(\mathcal{A}\hat{\mathbf{x}}_0))$ . Now further suppose that the measurement rows are orthogonal, i.e.  $\mathcal{A}\mathcal{A}^T = \mathbf{I}_t$ . This condition holds for some natural linear inverse problems like inpainting. Suppose that we want to update the latent vector  $\mathbf{z}_t$  such that  $\mathbb{E}[z_0|\mathbf{z}_t] = \mathcal{E}(x_0^*)$ ; this ensures that the gradients
Algorithm 1: DPS
Input:  $T,\pmb{y},\zeta_{i = 1}^{T},\{\tilde{\sigma}_{i}\}_{i = 1}^{T},\pmb{s}_{\theta}$    
1  $\pmb{x}_T\sim \mathcal{N}(\pmb {0},\pmb {I})$    
2 for  $i = T - 1$  to 0 do   
3  $\hat{\pmb{s}}\gets \pmb{s}_{\theta}(\pmb {x}_i,i)$    
4  $\hat{\pmb{x}}_0\gets \frac{1}{\sqrt{\bar{\alpha}_i}} (\pmb {x}_i + (1 - \bar{\alpha}_i)\hat{\pmb{s}})$    
5  $\pmb {z}\sim \mathcal{N}(\pmb {0},\pmb {I})$    
6  $\pmb{x}_{i - 1}^{\prime}\gets$ $\begin{array}{r}\frac{\sqrt{\alpha_i}(1 - \bar{\alpha}_i - 1)}{1 - \bar{\alpha}_i}\pmb {x}_i + \frac{\sqrt{\bar{\alpha}_i - 1}\beta_i}{1 - \bar{\alpha}_i}\hat{\pmb{x}}_0 + \tilde{\sigma}_iz\\ \pmb{x}_{i - 1}\gets \pmb{x}_{i - 1}^{\prime} - \zeta_i\nabla_{\pmb{x}_i}\| \pmb {y} - \mathcal{A}(\hat{\pmb{x}}_0)\| _2^2 \end{array}$    
7   
8 end   
9 return  $\hat{\pmb{x}}_0$
Algorithm 2: PSLD
Input:  $T, y$ .  $\{\eta_i\}_{i=1}^T$ ,  $\{\gamma_i\}_{i=1}^T$ ,  $\{\tilde{\sigma}_i\}_{i=1}^T$ ,  $\mathcal{E}$ ,  $\mathcal{D}$ ,  $\mathcal{A}$ ,  $\boldsymbol{s}_{\theta}$   
1  $z_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$   
2 for  $i = T - 1$  to 0 do  
3  $\hat{\boldsymbol{s}} \gets \boldsymbol{s}_{\theta}(\boldsymbol{z}_i, i)$   
4  $\hat{\boldsymbol{z}}_0 \gets \frac{1}{\sqrt{\bar{\alpha}_i}} (\boldsymbol{z}_i + (1 - \bar{\alpha}_i) \hat{\boldsymbol{s}})$   
5  $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$   
6  $z_{i-1}' \gets \frac{\sqrt{\alpha_i} (1 - \bar{\alpha}_i - 1)}{1 - \bar{\alpha}_i} \boldsymbol{z}_i + \frac{\sqrt{\bar{\alpha}_i - 1} \beta_i}{1 - \bar{\alpha}_i} \hat{\boldsymbol{z}}_0 + \tilde{\sigma}_i \epsilon$   
7  $z_{i-1}'' \gets z_{i-1}' - \eta_i \nabla_{\boldsymbol{z}_i} \| \boldsymbol{y} - \mathcal{A}(\mathcal{D}(\hat{\boldsymbol{z}}_0)) \|_2^2$   
8  $z_{i-1} \gets z_{i-1}'' - \gamma_i \nabla_{\boldsymbol{z}_i} \| \hat{\boldsymbol{z}}_0 - \mathcal{E}(\mathcal{A}^T \boldsymbol{y} + (\boldsymbol{I} - \mathcal{A}^T \mathcal{A}) \mathcal{D}(\hat{\boldsymbol{z}}_0)) \|_2^2$   
9 end  
10 return  $\mathcal{D}(\hat{\boldsymbol{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_{\boldsymbol{z}_t} \| \mathbb{E}[z_0 | \boldsymbol{z}_t] - \mathcal{E}(x_0^*) \|_2^2$ . Substituting  $\mathcal{E}(x_0^*) = \mathcal{E}(\mathcal{A}^T \boldsymbol{y} + (\hat{\boldsymbol{x}}_0 - \mathcal{A}^T \mathcal{A} \hat{\boldsymbol{x}}_0)) = \mathcal{E}(\mathcal{A}^T \boldsymbol{y} + (\boldsymbol{I} - \mathcal{A}^T \mathcal{A}) \hat{\boldsymbol{x}}_0)$ , and  $\hat{\boldsymbol{x}}_0 = \mathcal{D}(\mathbb{E}[z_0 | \boldsymbol{z}_t])$ , we can thus interpret the gluing objective in (7) as a one-step gradient descent of this loss  $\| \mathbb{E}[z_0 | \boldsymbol{z}_t] - \mathcal{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.
As discussed in Section 2, diffusion models consist of two stochastic processes: the forward and reverse processes, each governed by Ito 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{\pmb{x}_0} \sim p(\overrightarrow{\pmb{x}_0})$  and gradually adds noise, i.e.,  $\overrightarrow{\pmb{x}}_{t+1} = \sqrt{1 - \beta_t} \overrightarrow{\pmb{x}}_t + \sqrt{\beta_t} \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{\pmb{x}}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}_d)$  and generates  $\overleftarrow{\pmb{x}}_{t-1} = \mu_\theta (\overleftarrow{\pmb{x}}_t, t) + \sqrt{\beta_t} \epsilon$ . In the last step,  $\mu_\theta (\overleftarrow{\pmb{x}}_1, 1)$  is displayed without the noise.
In this section, we consider the diffusion discretized to two steps  $(\{\overrightarrow{x_0},\overrightarrow{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{x_0},\overrightarrow{x_1},\dots ,\overrightarrow{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.
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  $(\overrightarrow{\pmb{x}_0} \in \mathbb{R}^d)$  reside in a linear subspace of the form  $\overrightarrow{x_0} = S\overrightarrow{w_0}$ ,  $S \in \mathbb{R}^{d \times l}$ ,  $\overrightarrow{w_0} \in \mathbb{R}^l$ , and  $\sigma_y = 0$ . Here,  $S$  is a tall thin matrix with rank  $(S) = l \leq d$  that lifts any latent vector  $\overrightarrow{w_0} \sim \mathcal{N}(0, I_l)$  to the image space with ambient dimension  $d$ . Given the measurements  $\pmb{y} = \mathcal{A}\overrightarrow{x_0} + \sigma_y\pmb{n}$ ,  $\mathcal{A} \in \mathbb{R}^{l \times d}$ ,  $\pmb{n} \in \mathbb{R}^l$ , the goal is to sample from  $p_0(\overrightarrow{x_0}|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  $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 VAE is  $k$ -dimensional and the latent distribution is a standard Gaussian  $\mathcal{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.,  $\boldsymbol{S}^T\boldsymbol {S} = \boldsymbol {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., [42]. While Assumption 3.1 is mild and can be relaxed at the expense of (standard) mathematical complications, Assumption 3.2 indicates that  $(\mathcal{AS})^T (\mathcal{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  $\overrightarrow{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