SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We do not see outstanding limitations in our analysis. The linear subspace assumption initiates the study of conditional diffusion models on low-dimensional data. We expect to stimulate more sophisticated analyses under general assumptions such as manifold data.
C Omitted Proof in Section 3
Proof. For any  $t \geq 0$ , it holds that  $\nabla_{x_t} \log p_t(x_t \mid y) = \nabla_{x_t} \log p_t(x_t, y)$  since the gradient is taken w.r.t.  $x_t$  only. Then plugging in this equation and expanding the norm square on the LHS gives
$$
\begin{array}{l} \mathbb {E} _ {(x _ {t}, y) \sim P _ {t}} \left[ \| \nabla_ {x _ {t}} \log p _ {t} (x _ {t}, y) - s (x _ {t}, y, t) \| _ {2} ^ {2} \right] = \mathbb {E} _ {(x _ {t}, y) \sim P _ {t}} \left[ \| s (x _ {t}, y, t) \| _ {2} ^ {2} \right. \\ \left. - 2 \langle \nabla_ {x _ {t}} \log p _ {t} (x _ {t}, y), s (x _ {t}, y, t) \rangle \right] + C. \\ \end{array}
$$
Then it suffices to prove
$$
\mathbb {E} _ {(x _ {t}, y) \sim P _ {t}} \left[ \langle \nabla_ {x _ {t}} \log p _ {t} (x _ {t}, y), s (x _ {t}, y, t) \rangle \right] = \mathbb {E} _ {(x, y) \sim P _ {x \hat {y}}} \mathbb {E} _ {x ^ {\prime} \sim \mathsf {N} (\alpha (t) x, h (t) I)} \left[ \langle \nabla_ {x ^ {\prime}} \phi_ {t} (x ^ {\prime} \mid x), s (x ^ {\prime}, y, t) \rangle \right]
$$
Using integration by parts to rewrite the inner product we have
$$
\begin{array}{l} \mathbb {E} _ {(x _ {t}, y) \sim P _ {t}} [ \langle \nabla_ {x _ {t}} \log p _ {t} (x _ {t}, y), s (x _ {t}, y, t) \rangle ] = \int p _ {t} (x _ {t}, y) \langle \nabla_ {x _ {t}} \log p _ {t} (x _ {t}, y), s (x _ {t}, y, t) \rangle d x _ {t} d y \\ = \int \langle \nabla_ {x _ {t}} p _ {t} (x _ {t}, y), s (x _ {t}, y, t) \rangle d x _ {t} d y \\ = - \int p _ {t} (x _ {t}, y) \operatorname {d i v} (s (x _ {t}, y, t)) d x _ {t} d y, \\ \end{array}
$$
where denote by  $\phi_t(x'|x)$  the density of  $\mathsf{N}(\alpha(t)x, h(t)I_D)$  with  $\alpha(t) = \exp(-t/2)$  and  $h(t) = 1 - \exp(-t)$ , then
$$
\begin{array}{l} - \int p _ {t} (x _ {t}, y) \operatorname {d i v} (s (x _ {t}, y, t)) d x _ {t} d y = - \mathbb {E} _ {(x, y) \sim P _ {x \hat {y}}} \int \phi_ {t} (x ^ {\prime} \mid x) \operatorname {d i v} (s (x ^ {\prime}, y, t)) d x ^ {\prime} \\ = \mathbb {E} _ {(x, y) \sim P _ {x \hat {y}}} \int \langle \nabla_ {x ^ {\prime}} \phi_ {t} (x ^ {\prime} \mid x), s (x ^ {\prime}, y, t) \rangle d x ^ {\prime} \\ = \mathbb {E} _ {(x, y) \sim P _ {x \hat {y}}} \mathbb {E} _ {x ^ {\prime} \sim \mathsf {N} (\alpha (t) x, h (t) I)} \left[ \left\langle \nabla_ {x ^ {\prime}} \phi_ {t} \left(x ^ {\prime} \mid x\right), s \left(x ^ {\prime}, y, t\right) \right\rangle \right]. \\ \end{array}
$$
D Omitted Proofs in Section 4
Additional Notations: We follow the notations in the main paper along with some additional ones. Use  $P_{t}^{LD}(z)$  to denote the low-dimensional distribution on  $z$  corrupted by diffusion noise. Formally,  $p_{t}^{LD}(z) = \int \phi_{t}(z'|z)p_{z}(z)\mathrm{d}z$  with  $\phi_t(\cdot |z)$  being the density of  $\mathsf{N}(\alpha (t)z,h(t)I_d)$ .  $P_{t_0}^{LD}(z|\widehat{f} (Az) = a)$  the corresponding conditional distribution on  $\widehat{f} (Az) = a$  at  $t_0$ , with shorthand as  $P_{t_0}^{LD}(a)$ . Also give  $P_{z}(z|\widehat{f} (Az) = a)$  a shorthand as  $P^{LD}(a)$ . In our theorems,  $\mathcal{O}$  hides constant factors and higher order terms in  $n_1^{-1}$  and  $n_2^{-1}$  and,  $\widetilde{\mathcal{O}}$  further hides logarithmic terms and can also hide factors in  $d$ .
D.1 Parametric Conditional Score Matching Error
Theorems presented in Section 4 are established upon the conditional score estimation error, which has been studied in [8] for general distributions, but in Lemma D.1 we provide a new one specific to our setting where the true score is linear in input  $(x_{t},\widehat{y})$  due to the Gaussian design. Despite the linearity of score in Gaussian case, we emphasize matching score in (3.2) is not simply linear regression as  $S$  consists of an encoder-decoder structure for estimating matrix  $A$  to reduce dimension (see  $\S \mathrm{H}$  for  $S$  construction and more proof details).
In the following lemma, we first present a general result for the case where the true score is within  $S$ , which is constructed as a parametric function class. Then the score matching error is bounded in terms of  $\mathcal{N}(S, 1/n_1)$ , the log covering number of  $S$ , recall  $n_1$  is the size of  $\mathcal{D}_{\mathrm{unlabel}}$ . Instantiating this general result, we derive score matching error for Gaussian case by upper bounding  $\mathcal{N}(S, 1/n_1)$  in this special case.
Lemma D.1. Under Assumption 4.1, if  $\nabla \log p_t(x\mid y)\in S$  , where
$$
\mathcal {S} = \left\{\mathbf {s} _ {V, \psi} (x, y, t) = \frac {1}{h (t)} \left(V \cdot \psi \left(V ^ {\top} x, y, t\right) - x\right): V \in \mathbb {R} ^ {D \times d}, \psi \in \Psi : \mathbb {R} ^ {d + 1} \times \left[ t _ {0}, T \right]\rightarrow \mathbb {R} ^ {d} \right\}, \tag {(3.8) revisited}
$$
with  $\Psi$  parametric. Then for  $\delta \geq 0$ , with probability  $1 - \delta$ , the square score matching error is bounded by  $\epsilon_{diff}^2 = \mathcal{O}\left(\frac{1}{t_0}\sqrt{\frac{\mathcal{N}(S,1 / n_1)(d^2\vee D)\log\frac{1}{\delta}}{n_1}}\right)$ , i.e.,
$$
\frac {1}{T - t _ {0}} \int_ {t _ {0}} ^ {T} \mathbb {E} _ {\left(x _ {t}, y\right) \sim P _ {t}} \left[ \| \nabla \log p _ {t} \left(x _ {t} \mid y\right) - \widehat {s} \left(x _ {t}, y, t\right) \| _ {2} ^ {2} \right] d t \leq \epsilon_ {d i f f} ^ {2}, \tag {D.1}
$$
recall  $P_{t}$  comes from  $P_{x\widehat{y}}$  by noising  $x$  at  $t$  in the forward process. Under Assumption 4.4 and given  $\widehat{f} (x) = \widehat{\theta}^{\top}x$  and  $\widehat{y} = \widehat{f} (x) + \xi ,\xi \sim \mathsf{N}(0,\nu^{2})$ , the score function  $\nabla \log p_t(x\mid \widehat{y})$  to approximate is linear in  $x$  and  $\widehat{y}$ . When approximated by  $S$  with  $\Psi$  linear,  $\mathcal{N}(\mathcal{S},1 / n_1) = \mathcal{O}((d^2 +Dd)\log (Ddn_1))$ .
Proof. Proof is in  $\S H$ .
To provide fidelity and reward guarantees of  $\widehat{P}_a$ : the generated distribution of  $x$  given condition  $\widehat{y} = a$ , we will need the following lemma. It provides a subspace recovery guarantee between  $V$  (score matching output) and  $A$  (ground truth), as well as a distance measure between distributions  $P_a$  and  $\widehat{P}_a$ , given score matching error  $\epsilon_{diff}$ .
Note  $P_{a}$  and  $\widehat{P}_{a}$  are over  $x$ , which admits an underlying low-dimensional structure  $x = A z$ . Thus we measure distance between  $P_{a}$  and  $\widehat{P}_{a}$  by defining
Definition D.2.  $TV(\widehat{P}_a) \coloneqq \mathsf{d}_{\mathrm{TV}} \left( P_{t_0}^{LD}(z \mid \widehat{f}(Az) = a), (U^\top V^\top)_\# \widehat{P}_a \right)$  with notations:
-  $\mathrm{d}_{\mathrm{TV}}(\cdot, \cdot)$  is the TV distance between two distributions.
-  $f_{\sharp}P$  denotes a push-forward measure, i.e., for any measurable  $\Omega$ ,  $(f_{\sharp}P)(\Omega) = P(f^{-1}(\Omega))$
-  $(V^{\top})_{\#} \widehat{P}_a$  pushes generated  $\widehat{P}_a$  forward to the low dimensional subspace using learned subspace matrix  $V$ .  $U$  is an orthonormal matrix of dimension  $d$ .
-  $P_{t_0}^{LD}(z \mid \widehat{f}(Az) = a)$  is close to  $(A^\top)_{\#} P_a$ , with  $t_0$  taking account for the early stopping in backward process.
We note that there is a distribution shift between the training and the generated data, which has a profound impact on the generative performance. We quantify the influence of distribution shift by the following class restricted divergence measure.
Definition D.3. Distribution shift between two arbitrary distributions  $P_{1}$  and  $P_{2}$  restricted under function class  $\mathcal{L}$  is defined as
$\mathcal{T}(P_1,P_2;\mathcal{L}) = \sup_{l\in \mathcal{L}}\mathbb{E}_{x\sim P_1}[l(x)] / \mathbb{E}_{x\sim P_2}[l(x)]$  with arbitrary two distributions  $P_{1},P_{2}$
Definition D.3 is well perceived in bandit and RL literature [30, 29, 6, 12].
Lemma D.4. Given the square score matching error (D.1) upper bounded by  $\epsilon_{diff}^{2}$ , and when  $P_z$  satisfying Assumption 4.3 with  $c_0I_d \preceq \mathbb{E}_{z \sim P_z}[zz^\top]$ , it guarantees on for  $x \sim \widehat{P}_a$  and  $\angle(V, A) := \|VV^\top - AA^\top\|_F^2$  that
$$
\left(I _ {D} - V V ^ {\top}\right) x \sim \mathsf {N} (0, \Lambda), \quad \Lambda \prec c t _ {0} I _ {D}, \tag {D.2}
$$
$$
\angle (V, A) = \widetilde {\mathcal {O}} \left(\frac {t _ {0}}{c _ {0}} \cdot \epsilon_ {d i f f} ^ {2}\right). \tag {D.3}
$$
$$
T V \left(\widehat {P} _ {a}\right) = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} \left(P \left(x , \widehat {y} = a\right) , P _ {x \widehat {y}} ; \bar {\mathcal {S}}\right)}{c _ {0}}} \cdot \epsilon_ {d i f f}\right). \tag {D.4}
$$
with  $\bar{S} = \left\{\frac{1}{T - t_0}\int_{t_0}^T\mathbb{E}_{x_t|x}\| \nabla \log p_t(x_t\mid y) - s(x_t,y,t)\| _2^2\mathrm{d}t:s\in S\right\}$ .  $TV(\widehat{P}_a)$  and  $\mathcal{T}(P(x,\widehat{y} = a),P_{x\widehat{y}};\bar{S})$  are defined in Definition D.2 and D.3.
Proof. Proof is in §E.2.
Proof. Proof of  $\angle (V,A)$ . By Lemma 3 of [8], we have
$$
\angle (V, A) = \mathcal {O} \left(\frac {t _ {0}}{c _ {0}} \cdot \epsilon_ {d i f f} ^ {2}\right)
$$
when the latent  $z$  satisfying Assumption 4.3 and  $c_{0}I_{d}\preceq \mathbb{E}_{z\sim P_{z}}[zz^{\top}]$ . Therefore, by (D.1), we have with high probability that
$$
\angle (V, A) = \widetilde {\mathcal {O}} \left(\frac {1}{c _ {0}} \sqrt {\frac {\mathcal {N} (\mathcal {S} , 1 / n _ {1}) (D \vee d ^ {2})}{n _ {1}}}\right).
$$
When Assumption 4.4 holds, plugging in  $c_{0} = \lambda_{\min}$  and  $\mathcal{N}(\mathcal{S},1 / n_1) = \mathcal{O}((d^2 +Dd)\log (Ddn_1))$  it gives
$$
\angle (V, A) = \widetilde {\mathcal {O}} \left(\frac {1}{\lambda_ {\mathrm {m i n}}} \sqrt {\frac {(D \lor d ^ {2}) d ^ {2} + (D \lor d ^ {2}) D d}{n _ {1}}}\right),
$$
where  $\widetilde{\mathcal{O}}$  hides logarithmic terms. When  $D > d^2$ , which is often the case in practical applications, we have
$$
\angle (V, A) = \widetilde {\mathcal {O}} \left(\frac {1}{\lambda_ {\mathrm {m i n}}} \sqrt {\frac {D d ^ {2} + D ^ {2} d}{n _ {1}}}\right).
$$
Proof of  $\mathbb{E}_{x\sim \widehat{P}_a}[\| x_\perp \| _2]$ . By the definition of  $x_{\perp}$  that  $x_{\perp} = (I_D - AA^\top)x$
$$
\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x ^ {\perp} \| _ {2} ] = \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (I _ {D} - A A ^ {\top}) x \| _ {2} ] \leq \sqrt {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (I _ {D} - A A ^ {\top}) x \| _ {2} ^ {2} ]}.
$$
Score matching returns  $V$  as an approximation of  $A$ , then
$$
\left\| \left(I _ {D} - A A ^ {\top}\right) x \right\| _ {2} \leq \left\| \left(I _ {D} - V V ^ {\top}\right) x \right\| _ {2} + \left\| \left(V V ^ {\top} - A A ^ {\top}\right) x \right\| _ {2},
$$
$$
\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (I _ {D} - A A ^ {\top}) x \| _ {2} ^ {2} ] \leq 2 \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (I _ {D} - V V ^ {\top}) x \| _ {2} ^ {2} ] + 2 \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (V V ^ {\top} - A A ^ {\top}) x \| _ {2} ^ {2} ],
$$
where by (D.2) in Lemma D.4 we have
$$
(I _ {D} - V V ^ {\top}) x \sim \mathsf {N} (0, \Lambda), \quad \Lambda \prec \operatorname {c t} _ {0} \mathsf {I}
$$
for some constant  $c \geq 0$ . Thus
$$
\mathbb {E} _ {x \sim \hat {P} _ {a}} \left[ \| \left(I _ {D} - V V ^ {\top}\right) x \| _ {2} ^ {2} \right] = \operatorname {T r} (\Lambda) \leq c t _ {0} D. \tag {D.5}
$$
$$
\| (V V ^ {\top} - A A ^ {\top}) x \| _ {2} ^ {2} \leq \| V V ^ {\top} - A A ^ {\top} \| _ {o p} ^ {2} \| x \| _ {2} ^ {2} \leq \| V V ^ {\top} - A A ^ {\top} \| _ {F} ^ {2} \| x \| _ {2} ^ {2},
$$
where  $\|VV^{\top} - AA^{\top}\|_{F}^{2}$  has an upper bound as in (D.3) and  $\mathbb{E}_{x\sim \widehat{P}_a}\left[\| x\| _2^2\right]$  is bounded in Lemma E.3 by
$$
\mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| x \| _ {2} ^ {2} \right] = \mathcal {O} \left(c t _ {0} D + M (a) \cdot (1 + T V (\widehat {P} _ {a})\right).
$$
with  $M(a) = O\left(\frac{a^2}{\|\beta^*\|_{\Sigma}} + d\right)$ .
Therefore, to combine things together, we have
$$
\begin{array}{l} \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x ^ {\perp} \| _ {2} ] \leq \sqrt {2 \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (I _ {D} - V V ^ {\top}) x \| _ {2} ^ {2} ] + 2 \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| (V V ^ {\top} - A A ^ {\top}) x \| _ {2} ^ {2} ]} \\ \leq c ^ {\prime} \sqrt {t _ {0} D} + 2 \sqrt {\angle (V , A)} \cdot \sqrt {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x \| _ {2} ^ {2} ]} \\ = \mathcal {O} \left(\sqrt {t _ {0} D} + \sqrt {\angle (V , A)} \cdot \sqrt {M (a)}\right). \\ \end{array}
$$
$\mathcal{O}$  hides multiplicative constant and  $\sqrt{\angle(V,A)t_0D}$ ,  $\sqrt{\angle(V,A)M(a)TV(\widehat{P}_a)}$ , which are terms with higher power of  $n_1^{-1}$  than the leading term.
Remark of Theorem 4.5. 1. Guarantee (4.2) applies to general distributions with light tail as assumed in Assumption 4.3.
2. Guarantee (4.3) guarantees high fidelity of generated data in terms of staying in the subspace when we have access to a large unlabeled dataset.
3. Guarantee (4.3) shows that  $\mathbb{E}_{x\sim \widehat{P}_a}[\| x_\perp \|_2]$  scales up when  $t_0$  goes up, which aligns with the dynamic in backward process that samples are concentrating to the learned subspace as  $t_0$  goes to 0. Taking  $t_0 \to 0$ ,  $\mathbb{E}_{x\sim \widehat{P}_a}[\| x_\perp \|]$  has the decay in  $O(n_1^{-\frac{1}{4}})$ . However, taking  $t_0 \to 0$  is not ideal for the sake of high reward of  $x$ , we take the best trade-off of  $t_0$  later in Theorem 4.6.
D.3 Proof of Theorem 4.6
Proof of Theorem 4.6 and that of some results in "Implications and Discussions" following the theorem in main paper are provided in this section. This section breaks down into three parts: Suboptimality Decomposition, Bounding  $\mathcal{E}_1$  Relating to Offline Bandits, Bounding  $\mathcal{E}_2$  and the Distribution Shift in Diffusion.
D.3.1 SubOpt  $(\widehat{P}_a; y^* = a)$  Decomposition
Proof. Recall notations  $\widehat{P}_a\coloneqq \widehat{P} (\cdot |\widehat{y} = a)$  (generated distribution) and  $P_{a}\coloneqq P(\cdot |\widehat{y} = a)$  (target distribution) and  $f^{*}(x) = g^{*}(x_{\parallel}) - h^{*}(x_{\perp})$ $\mathbb{E}_{x\sim \widehat{P}_a}[f^\star (x)]$  can be decomposed into 3 terms:
$$
\begin{array}{l} \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ f ^ {\star} (x) ] \geq \mathbb {E} _ {x \sim P _ {a}} [ f ^ {\ast} (x) ] - \left| \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ f ^ {\ast} (x) ] - \mathbb {E} _ {x \sim P _ {a}} [ f ^ {\ast} (x) ] \right| \\ \geq \mathbb {E} _ {x \sim P _ {a}} [ \widehat {f} (x) ] - \mathbb {E} _ {x \sim P _ {a}} \left[ \left| \widehat {f} (x) - f ^ {*} (x) \right| \right] - \left| \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ f ^ {*} (x) ] - \mathbb {E} _ {x \sim P _ {a}} [ f ^ {*} (x) ] \right| \\ \geq \mathbb {E} _ {x \sim P _ {a}} \left[ \widehat {f} (x) \right] - \underbrace {\mathbb {E} _ {x \sim P _ {a}} \left[ \left| \widehat {f} (x) - g ^ {*} (x) \right| \right]} _ {\mathcal {E} _ {1}} \\ - \underbrace {\left| \mathbb {E} _ {x \sim P _ {a}} \left[ g ^ {*} (x _ {\parallel}) \right] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ g ^ {*} (x _ {\parallel}) \right] \right|} _ {\mathcal {E} _ {2}} - \underbrace {\mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ h ^ {*} (x _ {\perp}) \right]} _ {\mathcal {E} _ {3}}, \\ \end{array}
$$
where  $\mathbb{E}_{x\sim P_a}[\widehat{f} (x)] = \mathbb{E}_{a\sim q}[a]$  and we use  $x = x_{\parallel}$ ,  $f^{*}(x) = g^{*}(x)$  when  $x\sim P_{a}$ . Therefore
$$
\begin{array}{l} \operatorname {S u b O p t} \left(\widehat {P} _ {a}; y ^ {*} = a\right) = a - \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ f ^ {\star} (x) \right] \\ \leq \underbrace {\mathbb {E} _ {x \sim P _ {a}} \left[ \left| (\widehat {\theta} - \theta^ {*}) ^ {\top} x \right| \right]} _ {\mathcal {E} _ {1}} + \underbrace {\left| \mathbb {E} _ {x \sim P _ {a}} [ g ^ {*} (x _ {\parallel}) ] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ g ^ {*} (x _ {\parallel}) ] \right|} _ {\mathcal {E} _ {2}} \\ + \underbrace {\mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ h ^ {*} (x _ {\perp}) \right]} _ {\mathcal {E} _ {3}}. \\ \end{array}
$$
$\mathcal{E}_1$  comes from regression: prediction/generalization error onto  $P_{a}$ , which is independent from any error of distribution estimation that occurs in diffusion.  $\mathcal{E}_2$  and  $\mathcal{E}_3$  do not measure regression-predicted  $\widehat{f}$ , thus they are independent from the prediction error in  $\widehat{f}$  for pseudo-labeling.  $\mathcal{E}_2$  measures the disparity between  $\widehat{P}_{a}$  and  $P_{a}$  on the subspace support and  $\mathcal{E}_3$  measures the off-subspace component in generated  $\widehat{P}_{a}$ .
D.3.2 Bounding  $\mathcal{E}_1$  Relating to Offline Bandits
For all  $x_{i} \in \mathcal{D}_{\mathrm{label}}$ ,  $y_{i} = f^{*}(x_{i}) + \epsilon_{i} = g(x_{i}) + \epsilon_{i}$ . Thus, trained on  $\mathcal{D}_{\mathrm{label}}$  the prediction model  $\widehat{f}$  is essentially approximating  $g$ . By estimating  $\theta^{*}$  with ridge regression on  $\mathcal{D}_{\mathrm{label}}$ , we have  $\widehat{f}(x) = \widehat{\theta}^{\top}x$  with
$$
\widehat {\theta} = \left(X ^ {\top} X + \lambda I\right) ^ {- 1} X ^ {\top} \left(X \theta^ {*} + \eta\right), \tag {D.6}
$$
where  $X^{\top} = (x_{1},\dots ,x_{i},\dots ,x_{n_{2}})$  and  $\eta = (\epsilon_1,\dots ,\epsilon_i,\dots ,\epsilon_{n_2})$
Lemma D.5. Under Assumption 4.1 and 4.2 and given  $\epsilon_i\sim \mathsf{N}(0,\sigma^2)$ , define  $V_{\lambda}\coloneqq X^{\top}X + \lambda I$ ,  $\widehat{\Sigma}_{\lambda}\coloneqq \frac{1}{n_2} V_{\lambda}$  and  $\Sigma_{P_a}\coloneqq \mathbb{E}_{x\sim P_a}xx^\top$  the covariance matrix (uncentered) of  $P_{a}$ , and take  $\lambda = 1$ , then with high probability
$$
\mathcal {E} _ {1} \leq \sqrt {\operatorname {T r} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma_ {P _ {a}}\right)} \cdot \frac {\mathcal {O} \left(\sqrt {d \log n _ {2}}\right)}{\sqrt {n _ {2}}}. \tag {D.7}
$$
Proof. Proof is in §E.3.
Lemma D.6. Under Assumption 4.1, 4.2 and 4.4, when  $\lambda = 1$ ,  $P_{a}$  has a shift from the empirical marginal of  $x$  in dataset by
$$
\operatorname {T r} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma_ {P _ {a}}\right) \leq \mathcal {O} \left(\frac {a ^ {2}}{\| \beta^ {*} \| _ {\Sigma}} + d\right). \tag {D.8}
$$
when  $n_2 = \Omega (\max \{\frac{1}{\lambda_{\min}},\frac{d}{\|\beta^*\|_\Sigma^2}\})$
Proof. Proof is in §E.4.
D.3.3 Bounding  $\mathcal{E}_2$  and the Distribution Shift in Diffusion
Lemma D.7. Under Assumption 4.1, 4.2 and 4.4, when  $t_0 = \left( \frac{Dd^2 + D^2d}{n_1} \right)^{1/6}$
$$
\mathcal {E} _ {2} = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} (P (x , \widehat {y} = a) , P _ {x \widehat {y}} ; \bar {\mathcal {S}})}{\lambda_ {\min}}} \cdot \left(\frac {D d ^ {2} + D ^ {2} d}{n _ {1}}\right) ^ {\frac {1}{6}} \cdot a\right).
$$
Proof. Proof is in §E.5.
Note that  $\mathcal{T}(P(x,\widehat{y} = a),P_{x\widehat{y}};\bar{S})$  depends on  $a$  and measures the distribution shift between the desired distribution  $P(x,\widehat{y} = a)$  and the data distribution  $P_{x\widehat{y}}$ . To understand this distribution's dependency on  $a$ , it what follows we give  $\mathcal{T}(P(x,\widehat{y} = a),P_{x\widehat{y}};\bar{S})$  a shorthand as  $\text{DistroShift}^2(a)$  and give it an upper bound in one special case of the problem.
Distribution Shift In the special case of covariance  $\Sigma$  of  $z$  is known and  $\| A - V\| _2^2 = \mathcal{O}\left(\| AA^\top -VV^\top \| _F^2\right)$ , we showcase a bound on the distribution shift in  $\mathcal{E}_2$ , as promised in the discussion following Theorem 4.6. We have
$$
\text {D i s t r o S h i f t} ^ {2} (a) = \frac {\mathbb {E} _ {P _ {x , \widehat {y}} = a} [ \ell (x , y ; \widehat {s}) ]}{\mathbb {E} _ {P _ {x \widehat {y}}} [ \ell (x , y ; \widehat {s}) ]},
$$
where  $\ell (x,y;\widehat{s}) = \frac{1}{T - t_0}\int_{t_0}^T\mathbb{E}_{x'\mid x}\| \nabla_{x'}\log \phi_t(x'|x) - \widehat{s} (x',y,t)\| _2^2\mathrm{d}t.$  By Proposition 3.1, it suffices to bound
$$
\operatorname {D i s t r o S h i f t} ^ {2} (a) = \frac {\mathbb {E} _ {P _ {x , \widehat {y} = a}} [ \int_ {t _ {0}} ^ {T} \| \nabla \log p _ {t} (x , y) - \widehat {s} (x , y , t) \| _ {2} ^ {2} \mathrm {d} t ]}{\mathbb {E} _ {P _ {x \widehat {y}}} [ \int_ {t _ {0}} ^ {T} \| \nabla \log p _ {t} (x , y) - \widehat {s} (x , y , t) \| _ {2} ^ {2} \mathrm {d} t ]}.
$$
We expand the difference  $\| \nabla \log p_t(x,y) - \widehat{s} (x,y,t)\| _2^2$  by
$$
\begin{array}{l} \| \nabla \log p _ {t} (x, y) - \widehat {s} (x, y, t) \| _ {2} ^ {2} \leq \frac {2}{h ^ {2} (t)} \left[ \| (A - V) B _ {t} (A ^ {\top} x + \nu^ {- 2} y \theta) \| _ {2} ^ {2} + \| V B _ {t} (A - V) ^ {\top} x \| _ {2} ^ {2} \right] \\ \leq \frac {2}{h ^ {2} (t)} \Big [ \| A - V \| _ {2} ^ {2} \| B _ {t} (A ^ {\top} x + \nu^ {- 2} y \theta) \| _ {2} ^ {2} + \| A - V \| _ {2} ^ {2} \| x \| _ {2} ^ {2} \Big ] \\ \leq \frac {2}{h ^ {2} (t)} \| A - V \| _ {2} ^ {2} (3 \| x \| _ {2} ^ {2} + y ^ {2}), \\ \end{array}
$$
where we recall  $B_{t}$  is defined in (H.2) and in the last inequality, we use  $(a + b)^{2} \leq 2a^{2} + 2b^{2}$ . In the case of covariance matrix  $\Sigma$  is known, i.e.,  $B_{t}$  is known, we also consider matrix  $V$  directly matches  $A$  without rotation. Then by [8, Lemma 3 and 17], we have  $\| A - V\| _2^2 = \mathcal{O}\left(\| AA^\top -VV^\top \| _F^2\right) =$
$\mathcal{O}\left(t_0 / c_0\mathbb{E}_{P_{x\widehat{y}}}[\ell (x,y;\widehat{s})]\right)$ . To this end, we only need to find  $\mathbb{E}_{P_{x|\widehat{y} = a}}[\| x\| _2^2 ]$ . Since we consider on-support  $x$ , which can be represented as  $x = A z$ , we have  $\| x\| _2 = \| z\| _2$ . Thus, we only need to find the conditional distribution of  $z|\widehat{y} = a$ . Fortunately, we know  $(z,\widehat{y})$  is jointly Gaussian, with mean 0 and covariance
$$
\left[ \begin{array}{c c} \Sigma & \Sigma \widehat {\beta} \\ \widehat {\beta} ^ {\top} \Sigma & \widehat {\beta} ^ {\top} \Sigma \widehat {\beta} + \nu^ {2} \end{array} \right].
$$
Consequently, the conditional distribution of  $z|\widehat{y} = a$  is still Gaussian, with mean  $\Sigma \widehat{\beta} a / (\widehat{\beta}^{\top}\Sigma \widehat{\beta} +\nu^{2})$  and covariance  $\Sigma -\Sigma \widehat{\beta}\widehat{\beta}^{\top}\Sigma /( \widehat{\beta}^{\top}\Sigma \widehat{\beta} +\nu^{2})$ . Hence, we have
$$
\mathbb {E} _ {P _ {z | \widehat {\gamma} = a}} [ \| z \| _ {2} ^ {2} ] = \frac {1}{(\widehat {\beta} ^ {\top} \Sigma \widehat {\beta} + \nu^ {2}) ^ {2}} \left((a ^ {2} - \widehat {\beta} ^ {\top} \Sigma \widehat {\beta} - \nu^ {2}) \widehat {\beta} ^ {\top} \Sigma^ {2} \widehat {\beta}\right) + \mathrm {T r} (\Sigma) = \mathcal {O} (a ^ {2} \vee d).
$$
We integrate over  $t$  for the numerator in  $\text{DistroShift}(a)$  to obtain  $\mathbb{E}_{P_{x|\hat{y}} = a}\left[\int_{t_0}^T\|\nabla\log p_t(x,y)- \widehat{s}(x,y,t)\|_2^2\mathrm{d}t\right] = \mathcal{O}\left((a^2 \vee d)\frac{1}{c_0}\mathbb{E}_{P_{x\hat{y}}}[\ell(x,y;\widehat{s})]\right)$ . Note the cancellation between the numerator and denominator, we conclude
$$
\operatorname {D i s t r o S h i f t} (a) = \mathcal {O} \left(\frac {1}{c _ {0}} (a \vee \sqrt {d})\right).
$$
As  $d$  is a natural upper bound of  $\sqrt{d}$  and viewing  $c_{0}$  as a constant, we have  $\text{DistroShift}(a) = \mathcal{O}(a \vee d)$  as desired.
E Supporting lemmas and proofs
E.1 Supporting lemmas
Lemma E.1. The estimated subspace  $V$  satisfies
$$
\left\| V U - A \right\| _ {F} = \mathcal {O} \left(d ^ {\frac {3}{2}} \sqrt {\angle (V , A)}\right) \tag {E.1}
$$
for some orthogonal matrix  $U\in \mathbb{R}^{d\times d}$
Proof. Proof is in §E.6.