SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\boldsymbol {\theta} ^ {*} = \arg \min  _ {\boldsymbol {\theta}} \mathbb {E} _ {\overrightarrow {\boldsymbol {x} _ {0}}, \overrightarrow {\epsilon}} \left[ \left\| \tilde {\mu} _ {1} \left(\overrightarrow {\boldsymbol {x} _ {1}} (\overrightarrow {\boldsymbol {x} _ {0}}, \overrightarrow {\epsilon}), \overrightarrow {\boldsymbol {x} _ {0}}\right) - \mu_ {\boldsymbol {\theta}} \left(\overrightarrow {\boldsymbol {x} _ {1}} \left(\overrightarrow {\boldsymbol {x} _ {0}}, \overrightarrow {\epsilon}\right)\right) \right\| ^ {2} \right].
$$
Given a partially known image  $\overrightarrow{\pmb{x}_0}\sim p(\overrightarrow{\pmb{x}_0})$ , a fixed variance  $\beta >0$ , there exists a step size  $\zeta_i^j = 1 / 2\sigma_j$  for all the coordinates of  $\overrightarrow{\pmb{x}_0}$  such that Algorithm 1 samples from the true posterior  $p(\overrightarrow{\pmb{x}_0} |y)$  and exactly recovers the groundtruth sample, i.e.,  $\overleftarrow{\pmb{x}}_0 = \overrightarrow{\pmb{x}}_0$ .
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  $\sigma = \{\sigma_j\}_{j=1}^k$  the singular values of  $(\mathcal{A}\mathcal{S})^T(\mathcal{A}\mathcal{S})$ , i.e.  $(\mathcal{A}\mathcal{S})^T(\mathcal{A}\mathcal{S}) = U\Sigma V^T := U D(\sigma)V^T$ ,  $U \in \mathbb{R}^{k \times k}$ ,  $V \in \mathbb{R}^{k \times k}$  and
$$
\boldsymbol {\theta} ^ {*} = \arg \min  _ {\boldsymbol {\theta}} \mathbb {E} _ {\vec {x} _ {0} ^ {\rightarrow}, \vec {\epsilon}} \left[\left\| \tilde {\mu} _ {1} \left(\vec {x} _ {1} ^ {\rightarrow} (\vec {x} _ {0} ^ {\rightarrow}, \vec {\epsilon}), \vec {x} _ {0} ^ {\rightarrow}\right) - \mu_ {\theta} \left(\vec {x} _ {1} ^ {\rightarrow} \left(\vec {x} _ {0} ^ {\rightarrow}, \vec {\epsilon}\right)\right)\right\| ^ {2} \right].
$$
While the DPS Algorithm [11] uses a scalar step size  $\zeta_{i}$  at each step, this does not suffice for exact recovery. However, by generalizing to allow a different step size per coordinate, we can show sample recovery. Thus, in this section, we denote  $\zeta_i^j$  to be the step size at step  $i$  and coordinate  $j$ ,  $1 \leq j \leq r$ . Also note that the step index  $i$  is vacuous in this section, as we consider a two-step diffusion process (i.e.,  $i$  is always '1').
Suppose  $\overrightarrow{\pmb{x}_0}\sim p(\overrightarrow{\pmb{x}_0})$ . Given measurements  $y = \mathcal{A}\overrightarrow{\pmb{x}_0}$  and a fixed variance  $\beta \in (0,1)$ , there exists a matrix step size $^6$ $\zeta = (1 / 2)(\mathcal{SU})D(\zeta_i)(\mathcal{SU})^T$ ,  $\zeta_{i} = \{\zeta_{i}^{j} = 1 / \sigma_{j}\}_{j = 1}^{k}$  for all the coordinates of  $\overrightarrow{\pmb{x}_0}$  such that Algorithm 1 samples from the true posterior  $p(\overrightarrow{\pmb{x}_0} |y)$  and exactly recovers the groundtruth sample, i.e.,  $\overleftarrow{\pmb{x}}_0 = \overrightarrow{\pmb{x}_0}$ .
Proof. Our goal is to show that  $\overleftarrow{\pmb{x}_0} = \overrightarrow{\pmb{x}_0}$ , where  $\overleftarrow{\pmb{x}_0}$  is returned by Algorithm 1. Recall that the reverse process starts with  $\overleftarrow{\pmb{x}_1} \sim \mathcal{N}(\pmb{0}, \pmb{I}_d)$  and generates the following:
$$
\begin{array}{l} \overleftarrow {\boldsymbol {x} _ {0}} = \boldsymbol {\theta} ^ {*} \overleftarrow {\boldsymbol {x} _ {1}} - \zeta \nabla_ {\underline {{\boldsymbol {x} _ {1}}}} \left\| \mathcal {A} \overleftarrow {\boldsymbol {x} _ {0}} (\overleftarrow {\boldsymbol {x} _ {1}}) - \boldsymbol {y} \right\| _ {2} ^ {2} \\ = \boldsymbol {\theta} ^ {*} \overleftarrow {\boldsymbol {x}} _ {1} - \zeta \nabla_ {\overleftarrow {\boldsymbol {x}} _ {1}} \left\| \mathcal {A S S} ^ {T} \overleftarrow {\boldsymbol {x}} _ {1} - \boldsymbol {y} \right\| _ {2} ^ {2} \\ = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x}} _ {1} - 2 \boldsymbol {\zeta} \left(\mathcal {A} \mathcal {S} \mathcal {S} ^ {T}\right) ^ {T} \left(\mathcal {A} \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x}} _ {1} - \mathbf {y}\right) \\ = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x}} _ {1} - 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \left(\mathcal {A} \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x}} _ {1} - \mathbf {y}\right) \\ = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\boldsymbol {x}} _ {1} - 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \mathcal {A} \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\boldsymbol {x}} _ {1} + 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \boldsymbol {y} \\ = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\boldsymbol {x}} _ {1} - 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \mathcal {A} \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\boldsymbol {x}} _ {1} + 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \mathcal {A} \overrightarrow {\boldsymbol {x}} _ {0} \\ = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x}} _ {1} - 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \mathcal {A} \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x}} _ {1} + 2 \zeta \mathcal {S} \mathcal {S} ^ {T} \mathcal {A} ^ {T} \mathcal {A} \mathcal {S} \overrightarrow {\mathbf {z}} _ {0}. \\ \end{array}
$$
Now, we use the singular value decomposition of  $(\mathcal{AS})^T(\mathcal{AS})$  with left singular vectors in  $\pmb{U} \in \mathbb{R}^{k \times k}$ , right singular vectors in  $\pmb{V} \in \mathbb{R}^{k \times k}$ , and singular values  $\pmb{\sigma} = [\sigma_1, \dots, \sigma_k]$  in  $\Sigma = D(\pmb{\sigma})$ . Thus, the above expression becomes
where (i) is due to Assumption 3.1 and (ii) uses Assumption 3.2. By choosing  $\zeta_i^j$  as half the inverse of the non-zero singular values of  $(\mathcal{AS})^T (\mathcal{AS})$ , i.e.,  $\zeta_i^j = 1 / 2\sigma_i\forall i = 1,\ldots ,k$ , we obtain
$$
\begin{array}{l} \overleftarrow {\mathbf {x} _ {0}} = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x} _ {1}} - \mathcal {S} \mathbf {U} \mathbf {U} ^ {T} \mathcal {S} ^ {T} \overleftarrow {\mathbf {x} _ {1}} + \mathcal {S} \mathbf {U} \mathbf {U} ^ {T} \overrightarrow {\mathbf {z} _ {0}} \\ = \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\vec {x} _ {1}} - \mathcal {S} \mathcal {S} ^ {T} \overleftarrow {\vec {x} _ {1}} + \mathcal {S} \overrightarrow {z _ {0}} = \overrightarrow {x _ {0}}, \\ \end{array}
$$
which completes the statement of the theorem.
Proposition B.2 (Variational Autoencoder). Suppose Assumption 3.1 holds. For an encoder  $\mathcal{E}$  :  $\mathbb{R}^d\to \mathbb{R}^k$  and a decoder  $\mathcal{D}:\mathbb{R}^k\to \mathbb{R}^d$  , denote by  $\mathcal{L}(\phi ,\omega)$  the training objective of VAE:
$$
\arg \min _ {\phi , \omega} \mathcal {L} (\phi , \omega) := \mathbb {E} _ {\overrightarrow {\boldsymbol {x} _ {0}} \sim p} \left[ \left\| \mathcal {D} (\mathcal {E} (\overrightarrow {\boldsymbol {x} _ {0}}; \phi); \omega) - \overrightarrow {\boldsymbol {x} _ {0}} \right\| _ {2} ^ {2} \right] + \lambda K L (\mathcal {E} \sharp_ {p}, \mathcal {N} (\boldsymbol {0}, \boldsymbol {I} _ {k}))
$$
<sup>6</sup>We use the term 'step size' in a more general way than is normally used. In this case, the step size is a 'pre-conditioning' positive definite matrix, whose eigenvalue magnitudes correspond to the scalar step sizes per coordinate along an appropriately rotated basis. This general form is needed and with carefully selected (unique) eigenvalues; otherwise the DPS algorithm fails to converge to the groundtruth sample. We will later see that for our PSLD Algorithm in Theorem 3.6, we can revert to the commonly used notion of step size (a single scalar), as any finite step size (including a single scalar common across all coordinates) suffices for proving recovery.
Proof. To show that the encoder  $\mathcal{E}(\overrightarrow{\pmb{x}_0};\phi) = S^T\overrightarrow{\pmb{x}_0}$  and the decoder  $\mathcal{D}(\overleftarrow{z_0};\omega) = S\overleftarrow{z_0}$  minimize the VAE training objective  $\mathcal{L}(\phi ,\omega)$ , we begin with the first part of the loss, which is also called reconstruction error  $\mathcal{L}_{recon}(\phi ,\omega)$ . Substituting  $\mathcal{E}(\overrightarrow{\pmb{x}_0};\phi) = S^T\overrightarrow{\pmb{x}_0}$  and  $\mathcal{D}(\overleftarrow{z_0};\omega) = S\overleftarrow{z_0}$ , we have
$$
\begin{array}{l} \mathcal {L} _ {\text {r e c o n}} (\phi , \omega) := \mathbb {E} _ {\overrightarrow {\boldsymbol {x} _ {0}} \sim p} \left[ \| \mathcal {D} \left(\mathcal {E} \left(\overrightarrow {\boldsymbol {x} _ {0}}; \phi\right); \omega\right) - \overrightarrow {\boldsymbol {x} _ {0}} \| _ {2} ^ {2} \right] \\ = \mathbb {E} _ {\overrightarrow {{\pmb x} _ {0}} \sim p} \left[\left\| \mathcal {D} (\mathcal {S} ^ {T} \overrightarrow {{\pmb x} _ {0} ^ {\rightarrow}}; \omega) - \overrightarrow {{\pmb x} _ {0} ^ {\rightarrow}} \right\| _ {2} ^ {2} \right] \\ = \mathbb {E} _ {\overrightarrow {{\pmb x} _ {0}} \sim p} \left[ \left\| \mathcal {S S} ^ {T} \overrightarrow {{\pmb x} _ {0}} - \overrightarrow {{\pmb x} _ {0}} \right\| _ {2} ^ {2} \right] \\ \end{array}
$$
Using the fact that  $\overrightarrow{x_0}$  lives in a linear subspace, we arrive at
$$
\begin{array}{l} \mathcal {L} _ {r e c o n} (\phi , \omega) = \mathbb {E} _ {\overrightarrow {\boldsymbol {x} _ {0} ^ {\prime}} \sim p} \left[ \left\| \mathcal {S S} ^ {T} \mathcal {S} \overrightarrow {\boldsymbol {z} _ {0}} - \mathcal {S} \overrightarrow {\boldsymbol {z} _ {0}} \right\| _ {2} ^ {2} \right] \\ \stackrel {(i)} {=} \mathbb {E} _ {\overrightarrow {{\boldsymbol {z}}} _ {0} \sim \mathcal {N} (\mathbf {0}, \boldsymbol {I} _ {k})} \left[ \left\| \mathcal {S} \overrightarrow {{\boldsymbol {z}}} _ {0} - \mathcal {S} \overrightarrow {{\boldsymbol {z}}} _ {0} \right\| _ {2} ^ {2} \right] = 0, \\ \end{array}
$$
where (i) is due to Assumption 3.1. Now, we analyze the distribution loss. Note that the KL-divergence between two Gaussian distributions with moments  $(\mu_1,\sigma_1)$  and  $(\mu_{2},\sigma_{2})$  is given by
$$
K L \left(\mathcal {N} \left(\mu_ {1}, \sigma_ {1}\right), \mathcal {N} \left(\mu_ {2}, \sigma_ {2}\right)\right) = \log \left(\frac {\sigma_ {2}}{\sigma_ {1}}\right) + \frac {\sigma_ {1} ^ {2} + \left(\mu_ {1} - \mu_ {2}\right) ^ {2}}{2 \sigma_ {2} ^ {2}} - \frac {1}{2}.
$$
Since  $\mathcal{E}(\pmb{x}_0) = \mathcal{S}^T\pmb{x}_0 = \mathcal{S}^T\mathcal{S}\pmb{z}_0 = \pmb{z}_0$ , the distribution loss becomes:
$$
\mathcal {L} _ {\text {d i s t}} (\phi) := K L \left(\mathcal {E} \sharp p, \mathcal {N} \left(\boldsymbol {0}, \boldsymbol {I} _ {k}\right)\right) = K L \left(\mathcal {N} \left(\boldsymbol {0}, \boldsymbol {I} _ {k}\right), \mathcal {N} \left(\boldsymbol {0}, \boldsymbol {I} _ {k}\right)\right) = 0.
$$
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  $\beta > 0$ , if  $\mu_{\theta}\left(\overrightarrow{z_1} (\overrightarrow{z_0},\overrightarrow{\epsilon})\right) := \theta \overrightarrow{z_1} (\overrightarrow{z_0},\overrightarrow{\epsilon})$ , then the closed-form solution is  $\pmb{\theta}^{*} = \sqrt{1 - \beta}\pmb{I}_{k}$ , which after normalization by  $\frac{1}{\sqrt{1 - \beta}}$  and composition with the decoder  $\mathcal{D}\left(\overline{z_0};\omega\right) := S\overleftarrow{z}_0$  recovers the true subspace of  $p(\overrightarrow{x_0})$ .
Proof. In latent diffusion models, the training is performed in the latent space of a pre-trained VAE. If the VAE is chosen from Proposition 3.3, then the training objective becomes:
$$
\begin{array}{l} \operatorname * {m i n} _ {\boldsymbol {\theta}} \mathbb {E} _ {\overrightarrow {{\boldsymbol {x}}} _ {0}, \overrightarrow {{\boldsymbol {\epsilon}}}} \left[ \right.\left\| \right. \tilde {\mu} _ {1} (\overrightarrow {{\boldsymbol {z}}} _ {1} \left(\mathcal {E} (\overrightarrow {{\boldsymbol {x}}} _ {0} ^ {\prime}), \overrightarrow {{\boldsymbol {\epsilon}}}\right), \mathcal {E} (\overrightarrow {{\boldsymbol {x}}} _ {0})\left. \right) - \mu_ {\boldsymbol {\theta}} \left(\overrightarrow {{\boldsymbol {z}}} _ {1} ^ {\prime} \left(\mathcal {E} (\overrightarrow {{\boldsymbol {x}}} _ {0}), \overrightarrow {{\boldsymbol {\epsilon}}}\right)\right)\left. \right\| ^ {2} \left. \right] \\ = \mathbb {E} _ {\overrightarrow {z _ {0}}, \overrightarrow {\epsilon}} \left[ \left\| \tilde {\mu} _ {1} (\overrightarrow {z _ {1}} (\overrightarrow {z _ {0}}, \overrightarrow {\epsilon}), \overrightarrow {z _ {0}}) - \mu_ {\theta} (\overrightarrow {z _ {1}} (\overrightarrow {z _ {0}}, \overrightarrow {\epsilon})) \right\| ^ {2} \right] \\ = \mathbb {E} _ {\overrightarrow {z _ {0}}, \overrightarrow {\epsilon}} \left[ \left\| \overrightarrow {z _ {0}} - \mu_ {\theta} \left(\overrightarrow {z _ {1}} \left(\overrightarrow {z _ {0}}, \overrightarrow {\epsilon}\right)\right) \right\| ^ {2} \right] = \mathbb {E} _ {\overrightarrow {z _ {0}}, \overrightarrow {\epsilon}} \left[ \left\| \overrightarrow {z _ {0}} - \theta \overrightarrow {z _ {1}} \left(\overrightarrow {z _ {0}}, \overrightarrow {\epsilon}\right) \right\| ^ {2} \right] \\ = \mathbb {E} _ {\vec {z} _ {0}, \vec {\epsilon}} \left[\left\| \vec {z} _ {0} ^ {\rightarrow} - \boldsymbol {\theta} \left(\vec {z} _ {0} ^ {\rightarrow} \sqrt {1 - \beta} + \sqrt {\beta} \vec {\epsilon}\right)\right\| ^ {2} \right] \\ = \underset { \begin{array}{c} \overrightarrow {\boldsymbol {z}} _ {0} \sim p \\ \overrightarrow {\boldsymbol {\epsilon}} \sim \mathcal {N} (\mathbf {0}, I _ {k}) \end{array} } {\mathbb {E}} \left[ \sum_ {i = 1} ^ {k} \left(\overrightarrow {\boldsymbol {z} _ {0 , i}} - \boldsymbol {\theta} _ {i} ^ {T} \left(\overrightarrow {\boldsymbol {z} _ {0}} \sqrt {1 - \beta} + \overrightarrow {\boldsymbol {\epsilon}} \sqrt {\beta}\right)\right) ^ {2} \right], \\ \end{array}
$$
where  $\pmb{\theta}_i^T$  denotes the  $i^{th}$  row of matrix  $\pmb{\theta}$ . The solution of this regression problem is given by
$$
\begin{array}{l} \boldsymbol {\theta} _ {i} ^ {*} = \underset {\boldsymbol {x} _ {0}, \epsilon} {\mathbb {E}} \left[ \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) ^ {T} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {x} _ {0}, \epsilon} \left[ \boldsymbol {z} _ {0, i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \underset {\boldsymbol {x} _ {0}, \epsilon} {\mathbb {E}} \left[ \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) ^ {T} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {x} _ {0}, \epsilon} \left[ \boldsymbol {z} _ {0, i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \underset {\boldsymbol {x} _ {0}, \epsilon} {\mathbb {E}} \left[ \left(\mathcal {E} (\boldsymbol {x} _ {0}) \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \left(\mathcal {E} (\boldsymbol {x} _ {0}) \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) ^ {T} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {x} _ {0}, \epsilon} \left[ \mathcal {E} (\boldsymbol {x} _ {0}) _ {i} \left(\mathcal {E} (\boldsymbol {x} _ {0}) \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \underset {z _ {0}, \epsilon} {\mathbb {E}} \left[ \left(\mathcal {E} (\mathcal {S} z _ {0}) \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \left(\mathcal {E} (\mathcal {S} z _ {0}) \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) ^ {T} \right] ^ {- 1} \mathbb {E} _ {z _ {0}, \epsilon} \left[ \mathcal {E} (\mathcal {S} z _ {0}) _ {i} \left(\mathcal {E} (\mathcal {S} z _ {0}) \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ \left(\mathcal {S} ^ {T} \mathcal {S} \boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \left(\mathcal {S} ^ {T} \mathcal {S} \boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) ^ {T} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ (\mathcal {S} ^ {T} \mathcal {S} \boldsymbol {z} _ {0}) _ {i} \left(\mathcal {S} ^ {T} \mathcal {S} \boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ \end{array}
$$
Using Assumption 3.1, the above expression simplifies to
$$
\begin{array}{l} \boldsymbol {\theta} _ {i} ^ {*} = \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) ^ {T} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ (\boldsymbol {z} _ {0}) _ {i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ (1 - \beta) \boldsymbol {z} _ {0} \boldsymbol {z} _ {0} ^ {T} + \boldsymbol {z} _ {0} \boldsymbol {\epsilon} ^ {T} \sqrt {\beta (1 - \beta)} + \boldsymbol {\epsilon} \boldsymbol {z} _ {0} ^ {T} \sqrt {\beta (1 - \beta)} + \beta \boldsymbol {\epsilon} \boldsymbol {\epsilon} ^ {T} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ (\boldsymbol {z} _ {0}) _ {i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \boldsymbol {\epsilon} \sqrt {\beta}\right) \right] \\ = \left[ \left(\left(1 - \beta\right) \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ \boldsymbol {z} _ {0} \boldsymbol {z} _ {0} ^ {T} \right] + \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ \boldsymbol {z} _ {0} \boldsymbol {\epsilon} ^ {T} \right] \sqrt {\beta (1 - \beta)} + \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ \boldsymbol {\epsilon} \boldsymbol {z} _ {0} ^ {T} \right] \sqrt {\beta (1 - \beta)} + \beta \underset {\boldsymbol {z} _ {0}, \epsilon} {\mathbb {E}} \left[ \boldsymbol {\epsilon} \boldsymbol {\epsilon} ^ {T} \right]\right) \right] ^ {- 1} \\ \times \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ \left(\boldsymbol {z} _ {0}\right) _ {i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \left[ \left((1 - \beta) \boldsymbol {I} _ {k} + \underset {\boldsymbol {z} _ {0}} {\mathbb {E}} [ \boldsymbol {z} _ {0} ] \underset {\boldsymbol {\epsilon}} {\mathbb {E}} [ \boldsymbol {\epsilon} ] ^ {T} \sqrt {\beta (1 - \beta)} + \underset {\boldsymbol {\epsilon}} {\mathbb {E}} [ \boldsymbol {\epsilon} ] \underset {\boldsymbol {z} _ {0}} {\mathbb {E}} [ \boldsymbol {z} _ {0} ] ^ {T} \sqrt {\beta (1 - \beta)} + \beta \boldsymbol {I} _ {k}\right) \right] ^ {- 1} \\ \times \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ \left(\boldsymbol {z} _ {0}\right) _ {i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right], \\ \end{array}
$$
where the last step uses the fact that  $z_0$  and  $\epsilon$  are independent Gaussian random vectors with zero mean and unit covariance. Simplifying further, we arrive at
$$
\begin{array}{l} \boldsymbol {\theta} _ {i} ^ {*} = \left[ (1 - \beta) \boldsymbol {I} _ {k} + \beta \boldsymbol {I} _ {k} \right] ^ {- 1} \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ (\boldsymbol {z} _ {0}) _ {i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ \left(\boldsymbol {z} _ {0}\right) _ {i} \left(\boldsymbol {z} _ {0} \sqrt {1 - \beta} + \epsilon \sqrt {\beta}\right) \right] \\ = \mathbb {E} _ {\boldsymbol {z} _ {0}} \left[ \left(\boldsymbol {z} _ {0}\right) _ {i} \boldsymbol {z} _ {0} \sqrt {1 - \beta} \right] + \mathbb {E} _ {\boldsymbol {z} _ {0}, \epsilon} \left[ \left(\boldsymbol {z} _ {0}\right) _ {i} \epsilon \sqrt {\beta} \right] \\ = \mathbb {E} _ {\boldsymbol {z} _ {0}} \left[ (\boldsymbol {z} _ {0}) _ {i} \boldsymbol {z} _ {0} \sqrt {1 - \beta} \right] + \mathbb {E} _ {\boldsymbol {z} _ {0}} \left[ (\boldsymbol {z} _ {0}) _ {i} \right] \mathbb {E} _ {\boldsymbol {\epsilon}} [ \epsilon ] \sqrt {\beta}. \\ \end{array}
$$
The final step follows from independence of  $z_0$  and  $\epsilon$ . Since  $z_0$  and  $\epsilon$  are also  $\mathcal{N}(\mathbf{0},\mathbf{I}_k)$ , we get
$$
\boldsymbol {\theta} _ {i} ^ {*} = \mathbb {E} _ {\boldsymbol {z} _ {0}} \left[ (\boldsymbol {z} _ {0}) _ {i} \boldsymbol {z} _ {0} \sqrt {1 - \beta} \right] = \left[ 0, \dots , 0, \sqrt {1 - \beta}, 0, \dots , 0 \right] ^ {T},
$$
where the  $i^{th}$  coordinate is  $\sqrt{1 - \beta}$  and zero everywhere else. Therefore, stacking all the rows together, we get  $\pmb{\theta}^{*} = \sqrt{1 - \beta}\pmb{I}_{k}$ , which after normalization by  $1 / \sqrt{1 - \beta}$  gives the desired result.
Next, we show that  $\pmb{\theta}^{*}$  recovers the true subspace of  $\overrightarrow{x_0} \sim p(\overrightarrow{x_0})$ . When composed with the decoder of VAE, the generator of the LDM gives  $\overleftarrow{\pmb{x}_0} = \mathcal{D}(\pmb{\theta}^*\overline{\pmb{z}_1}) = \mathcal{D}(\pmb{I}_k\overleftarrow{\pmb{z}_1}) = \mathcal{S}\overleftarrow{\pmb{z}_1}$ . Since  $\overleftarrow{\pmb{z}_1} \sim \mathcal{N}(\mathbf{0},\mathbf{I}_k)$ , this completes the statement of the theorem.
B.4 Proof of Theorem 3.5
Recall that 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
For ease of notation, we drop the forward arrow in the rest of this proof.
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:
$$
\overleftarrow {z _ {0} ^ {\prime}} = \boldsymbol {\theta} ^ {*} \overleftarrow {z _ {1}} - \boldsymbol {\eta} \nabla_ {\overleftarrow {z} _ {1}} \left\| \mathcal {A D} \left(\overleftarrow {z _ {0}} \left(\overleftarrow {z _ {1}}\right)\right) - \boldsymbol {y} \right\| _ {2} ^ {2}; \tag {8}
$$
$$
\left. \overleftarrow {z} _ {0} = \arg \min  _ {\overleftarrow {z} _ {0} ^ {\prime}} \left| \left| \overleftarrow {z} _ {0} ^ {\prime} - \mathcal {E} \left(\mathcal {D} \left(\overleftarrow {z} _ {0} ^ {\prime}\right)\right) \right| \right| _ {2} ^ {2}, \right. \tag {9}
$$
In practice, solving (9) can be difficult, and can be approximated via gradient descent. In our analysis however, we analyze the exact system of equations above, as (9) has a closed-form solution in the linear setting.
Theorem B.4 (Posterior Sampling using Goodness Modified Latent DPS). Suppose Assumptions 3.1 and Assumption 3.2 hold. Denote by  $\sigma = \{\sigma_{j}\}_{j=1}^{k}$  the singular values of  $(\mathcal{A}\mathcal{S})^T(\mathcal{A}\mathcal{S})$ , i.e.,  $(\mathcal{A}\mathcal{S})^T(\mathcal{A}\mathcal{S}) = U\Sigma U^T := U D(\sigma)U^T$ ,  $U \in \mathbb{R}^{k \times k}$ , and let
$$
\boldsymbol {\theta} ^ {*} = \arg \min  _ {\boldsymbol {\theta}} \mathbb {E} _ {\vec {z _ {0}}, \vec {\epsilon}} \left[ \left\| \tilde {\mu} _ {1} \left(\vec {z _ {1}} (\vec {z _ {0}}, \vec {\epsilon}), \vec {z _ {0}}\right) - \mu_ {\theta} \left(\vec {z _ {1}} (\vec {z _ {0}}, \vec {\epsilon})\right) \right\| _ {2} ^ {2} \right].
$$
Suppose  $\overrightarrow{\pmb{x}_0} \sim p(\overrightarrow{\pmb{x}_0})$ . Given measurements  $\pmb{y} = \mathcal{A}\overrightarrow{\pmb{x}_0}$  and any fixed variance  $\beta \in (0,1)$ , then with the (unique) step size  $\pmb{\eta} = (1/2)\pmb{U}\pmb{D}(\pmb{\eta}_i)\pmb{U}^T$ ,  $\pmb{\eta}_i = \{\eta_i^j = 1/2\sigma_j\}_{j=1}^k$ , the GML-DPS algorithm (6) samples from the true posterior  $p(\overrightarrow{\pmb{x}_0} | y)$  and exactly recovers the groundtruth sample, i.e.,  $\overleftarrow{\pmb{x}_0} = \overrightarrow{\pmb{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