SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma E.2. Suppose  $P_{1}$  and  $P_{2}$  are two distributions over  $\mathbb{R}^d$  and  $m$  is a function defined on  $\mathbb{R}^d$ , then  $|\mathbb{E}_{x\sim P_1}[m(z)] - \mathbb{E}_{z\sim P_2}[m(z)]|$  can be bounded in terms of  $\mathrm{d}_{\mathrm{TV}}(P_1,P_2)$ , specifically when  $P_{1}$  and  $P_{2}$  are Gaussians and  $m(z) = \| z\| _2^2$ :
$$
\mathbb {E} _ {x \sim P _ {1}} [ \| z \| _ {2} ^ {2} ] = \mathcal {O} \left(\mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ] \left(1 + \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right)\right)\right). \tag {E.2}
$$
When  $P_{1}$  and  $P_{2}$  are Gaussians and  $m(z) = \| z\| _2$
$$
\left. \right.\left| \mathbb {E} _ {z \sim P _ {1}} \left[ \| z \| _ {2} \right] - \mathbb {E} _ {z \sim P _ {2}} \left[ \| z \| _ {2} \right]\right| = \mathcal {O} \left(\left(\sqrt {\mathbb {E} _ {z \sim P _ {1}} \left[ \| z \| _ {2} ^ {2} \right]} + \sqrt {\mathbb {E} _ {z \sim P _ {2}} \left[ \| z \| _ {2} ^ {2} \right]}\right) \cdot \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right)\right). \tag {E.3}
$$
Proof. Proof is in §E.7.
Lemma E.3. We compute  $\mathbb{E}_{z\sim P^{LD}(a)}[\| z\| _2^2 ],\mathbb{E}_{z\sim P_{t_0}^{LD}(a)}[\| z\| _2^2 ],\mathbb{E}_{x\sim \widehat{P}_a}[\| x\| _2^2 ],\mathbb{E}_{z\sim (U^\top V^\top)_\#}\widehat{P}_a[\| z\| _2^2 ]$  in this Lemma.
$$
\mathbb {E} _ {z \sim P ^ {L D} (a)} [ \| z \| _ {2} ^ {2} ] = \frac {\widehat {\beta} ^ {\top} \Sigma^ {2} \widehat {\beta}}{\left(\| \widehat {\beta} \| _ {\Sigma} ^ {2} + \nu^ {2}\right) ^ {2}} a ^ {2} + \operatorname {t r a c e} (\Sigma - \Sigma \widehat {\beta} \left(\widehat {\beta} ^ {\top} \Sigma \widehat {\beta} + \nu^ {2}\right) ^ {- 1} \widehat {\beta} ^ {\top} \Sigma). \tag {E.4}
$$
Let  $M(a) \coloneqq \mathbb{E}_{z \sim P^{LD}(a)}[\| z\|_2^2]$ , which has an upper bound  $M(a) = O\left(\frac{a^2}{\|\beta^*\|_{\Sigma}} + d\right)$ .
$$
\mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \| z \| _ {2} ^ {2} ] \leq M (a) + t _ {0} d. \tag {E.5}
$$
$$
\mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| x \| _ {2} ^ {2} \right] \leq \mathcal {O} \left(c t _ {0} D + M (a) \cdot \left(1 + T V (\widehat {P} _ {a})\right)\right). \tag {E.6}
$$
Proof. Proof is in §E.8.
E.2 Proof of Lemma D.4
Proof. The first two assertions (D.2) and (D.3) are consequences of [8, Theorem 3, item 1 and 3]. To show (D.4), we first have the conditional score matching error under distribution shift being
$$
\mathcal {T} (P (x, \widehat {y} = a), P _ {x \widehat {y}}; \bar {\mathcal {S}}) \cdot \epsilon_ {d i f f} ^ {2},
$$
where  $\mathcal{T}(P(x,\widehat{y} = a),P_{x\overline{y}};\bar{S})$  accounts for the distribution shift as in the parametric case (Lemma D.7). Then we apply [8, Theorem 3, item 2] to conclude
$$
T V (\widehat {P} _ {a}) = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} (P (x , \widehat {y} = a) , P _ {x \widehat {y}} ; \bar {\mathcal {S}})}{c _ {0}}} \cdot \epsilon_ {d i f f}\right).
$$
E.3 Proof of Lemma D.5
Proof. Given
$$
\mathcal {E} _ {1} = \mathbb {E} _ {\widehat {P} _ {a}} \left| x ^ {\top} (\theta^ {*} - \widehat {\theta}) \right| \leq \mathbb {E} _ {\widehat {P} _ {a}} \| x \| _ {V _ {\lambda} ^ {- 1}} \cdot \| \theta^ {*} - \widehat {\theta} \| _ {V _ {\lambda}},
$$
then things to prove are
$$
\left. \mathbb {E} _ {\widehat {P} _ {a}} \| x \| _ {V _ {\lambda} ^ {- 1}} = \sqrt {\operatorname {t r a c e} \left(V _ {\lambda} ^ {- 1} \Sigma_ {\widehat {P} _ {a}}\right)}; \right. \tag {E.7}
$$
$$
\left\| \theta^ {*} - \widehat {\theta} \right\| _ {V _ {\lambda}} \leq \mathcal {O} \left(\sqrt {d \log n _ {2}}\right), \tag {E.8}
$$
where the second inequality is to be proven with high probability w.r.t the randomness in  $\mathcal{D}_{\text{label}}$ . For (E.7),  $\mathbb{E}_{\widehat{P}_a}\|x\|_{V_\lambda^{-1}} \leq \sqrt{\mathbb{E}_{\widehat{P}_a}x^\top V_\lambda^{-1}x} = \sqrt{\mathbb{E}_{\widehat{P}_a}\mathrm{trace}(V_\lambda^{-1}xx^\top)} = \sqrt{\mathrm{trace}(V_\lambda^{-1}\mathbb{E}_{\widehat{P}_a}xx^\top)}$ .
For (E.8), what's new to prove compared to a classic bandit derivation is its  $d$  dependency instead of  $D$ , due to the linear subspace structure in  $x$ . From the closed form solution of  $\widehat{\theta}$ , we have
$$
\widehat {\theta} - \theta^ {*} = V _ {\lambda} ^ {- 1} X ^ {\top} \eta - \lambda V _ {\lambda} ^ {- 1} \theta^ {*}. \tag {E.9}
$$
$$
\left\| \theta^ {*} - \widehat {\theta} \right\| _ {V _ {\lambda}} \leq \left\| X ^ {\top} \eta \right\| _ {V _ {\lambda} ^ {- 1}} + \lambda \left\| \theta^ {*} \right\| _ {V _ {\lambda} ^ {- 1}}, \tag {E.10}
$$
where  $\lambda \| \theta^{*}\|_{V_{\lambda}^{-1}}\leq \sqrt{\lambda}\| \theta^{*}\|_{2}\leq \sqrt{\lambda}$  and
$$
\begin{array}{l} \| X ^ {\top} \eta \| _ {V _ {\lambda} ^ {- 1}} ^ {2} = \eta^ {\top} X (X ^ {\top} X + \lambda I _ {D}) ^ {- 1} X ^ {\top} \eta \\ = \eta^ {\top} X X ^ {\top} \left(X X ^ {\top} + \lambda I _ {n _ {2}}\right) ^ {- 1} \eta . \\ \end{array}
$$
Let  $Z^{\top} = (z_{1}, \dots, z_{i}, \dots, z_{n_{2}})$  s.t.  $A z_{i} = x_{i}$ , then it holds that  $X = Z A^{\top}$ , and  $X X^{\top} = Z A^{\top} A Z^{\top} = Z Z^{\top}$ , thus
$$
\begin{array}{l} \| X ^ {\top} \eta \| _ {V _ {\lambda} ^ {- 1}} ^ {2} = \eta^ {\top} X X ^ {\top} \left(X X ^ {\top} + \lambda I _ {n _ {2}}\right) ^ {- 1} \eta \\ = \eta^ {\top} Z Z ^ {\top} \left(Z Z ^ {\top} + \lambda I _ {n _ {2}}\right) ^ {- 1} \eta \\ = \eta^ {\top} Z \left(Z ^ {\top} Z + \lambda I _ {d}\right) ^ {- 1} Z ^ {\top} \eta \\ = \left\| Z ^ {\top} \eta \right\| _ {\left(Z ^ {\top} Z + \lambda I _ {d}\right) ^ {- 1}}. \\ \end{array}
$$
With probability  $1 - \delta$ ,  $\|z_i\|^2 \leq d + \sqrt{d \log \left(\frac{2n_2}{\delta}\right)} \coloneqq L^2, \forall i \in [n_2]$ . Then Theorem 1 in "Improved algorithms for linear stochastic bandits" (by Yasin Abbasi-Yadkori, David Pal, and Csaba Szepesvari) gives rise to
$$
\left\| Z ^ {\top} \eta \right\| _ {\left(Z ^ {\top} Z + \lambda I _ {d}\right) ^ {- 1}} \leq \sqrt {2 \log (2 / \delta) + d \log \left(1 + n _ {2} L ^ {2} / (\lambda d)\right)}
$$
with probability  $1 - \delta / 2$ . Combine things together and plugging in  $\lambda = 1$ ,  $L^2 = d + \sqrt{d \log \left(\frac{2n_2}{\delta}\right)}$ , we have with high probability
$$
\| \theta^ {*} - \widehat {\theta} \| _ {V _ {\lambda}} = \mathcal {O} \left(\sqrt {d \log (n _ {2} \sqrt {\log (n _ {2})})}\right) = \mathcal {O} \left(\sqrt {d \log n _ {2} + \frac {1}{2} d \log (\log n _ {2})}\right) = \mathcal {O} \left(\sqrt {d \log n _ {2}}\right).
$$
E.4 Proof of Lemma D.6
Proof. Recall the definition of  $\widehat{\Sigma}_{\lambda}$  and  $\Sigma_{P_a}$  that
$$
\widehat {\Sigma} _ {\lambda} = \frac {1}{n _ {2}} X ^ {\top} X + \frac {\lambda}{n _ {2}} I _ {D},
$$
$$
\Sigma_ {P _ {a}} = \mathbb {E} _ {x \sim P _ {a}} \left[ x x ^ {\top} \right],
$$
where  $X$  are stack matrix of data supported on  $\mathcal{A}$  and  $P_{a}$  is also supported on  $\mathcal{A}$ ,  $\mathcal{A}$  is the subspace encoded by matrix  $A$ . The following lemma shows it is equivalent to measure trace  $(\widehat{\Sigma}_{\lambda}^{-1}\Sigma_{P_a})$  on  $\mathcal{A}$  subspace.
Lemma E.4. For any P.S.D. matrices  $\Sigma_1, \Sigma_2 \in \mathbb{R}^{d \times d}$  and  $A \in \mathbb{R}^{D \times d}$  such that  $A^\top A = I_d$ , we have
$$
\operatorname {T r} \left(\left(\lambda I _ {D} + A \Sigma_ {1} A ^ {\top}\right) ^ {- 1} A \Sigma_ {2} A ^ {\top}\right) = \operatorname {T r} \left(\left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {2}\right).
$$
The lemma above allows us to abuse notations  $\widehat{\Sigma}_{\lambda}$  and  $\Sigma_{P_a}$  in the following way while keeping the same trace  $(\widehat{\Sigma}_{\lambda}^{-1}\Sigma_{P_a})$  value:
$$
\widehat {\Sigma} _ {\lambda} = \frac {1}{n _ {2}} Z ^ {\top} Z + \frac {\lambda}{n _ {2}} I _ {d},
$$
$$
\Sigma_ {P _ {a}} = \mathbb {E} _ {z \sim P ^ {L D} (a)} \left[ z z ^ {\top} \right],
$$
where  $Z^{\top} = (z_{1},\dots ,z_{i},\dots ,z_{n_{2}})$  s.t.  $A z_{i} = x_{i}$  and recall notataion  $P^{LD}(a) = P_z\left(z\mid \widehat{f} (Az)\right)$ .
Given  $z \sim \mathsf{N}(\mu, \Sigma)$ , as a proof artifact, let  $\widehat{f}(x) = \widehat{\theta}^\top x + \xi, \xi \sim \mathsf{N}(0, \nu^2)$  where we will let  $\nu \to 0$  in the end, then let  $\widehat{\beta} = A^\top \widehat{\theta} \in \mathbb{R}^d$ ,  $(z, \widehat{f}(Az))$  has a joint distribution
$$
\left(z, \widehat {f}\right) \sim \mathbb {N} \left(\left[ \begin{array}{l} \mu \\ \widehat {\beta} ^ {\top} \mu \end{array} \right], \left[ \begin{array}{c c} \Sigma & \Sigma \widehat {\beta} \\ \widehat {\beta} ^ {\top} \Sigma & \widehat {\beta} ^ {\top} \Sigma \widehat {\beta} + \nu^ {2} \end{array} \right]\right). \tag {E.11}
$$
Then we have the conditional distribution  $z \mid \widehat{f}(Az) = a$  following
$$
P _ {z} (z \mid \widehat {f} (A z) = a) = \mathbb {N} (\mu + \Sigma \widehat {\beta} (\widehat {\beta} ^ {\top} \Sigma \widehat {\beta} + \nu^ {2}) ^ {- 1} (a - \widehat {\beta} ^ {\top} \mu), \Gamma) \tag {E.12}
$$
with  $\Gamma \coloneqq \Sigma -\Sigma \widehat{\beta}\left(\widehat{\beta}^{\top}\Sigma \widehat{\beta} +\nu^{2}\right)^{-1}\widehat{\beta}^{\top}\Sigma .$
When  $\mu = 0$ , we compute  $\mathrm{trace}(\widehat{\Sigma}_{\lambda}^{-1}\Sigma_{P_a})$  as
$$
\begin{array}{l} \mathrm {t r a c e} (\widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma_ {P _ {a}}) = \mathrm {t r a c e} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \frac {\Sigma \widehat {\beta} \widehat {\beta} ^ {\top} \Sigma}{\left(\| \widehat {\beta} \| _ {\Sigma} ^ {2} + \nu^ {2}\right) ^ {2}} a ^ {2}\right) + \mathrm {t r a c e} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \Gamma\right) \\ = \operatorname {t r a c e} \left(\frac {\widehat {\beta} ^ {\top} \Sigma \widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma \widehat {\beta}}{\left(\| \widehat {\beta} \| _ {\Sigma} ^ {2} + \nu^ {2}\right) ^ {2}} a ^ {2}\right) + \operatorname {t r a c e} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma\right) - \operatorname {t r a c e} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \frac {\Sigma \widehat {\beta} \widehat {\beta} ^ {\top} \Sigma}{\| \widehat {\beta} \| _ {\Sigma} ^ {2} + \nu^ {2}}\right) \\ = \operatorname {t r a c e} \left(\frac {\Sigma^ {1 / 2} \widehat {\beta} \widehat {\beta} ^ {\top} \Sigma^ {1 / 2} \Sigma^ {1 / 2} \widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma^ {1 / 2}}{\left(\| \widehat {\beta} \| _ {\Sigma} ^ {2} + \nu^ {2}\right) ^ {2}} a ^ {2}\right) \\ \leq \frac {\left\| \Sigma \widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma \right\| _ {o p} \cdot \left\| \widehat {\beta} \right\| _ {\Sigma} ^ {2}}{\left(\left\| \widehat {\beta} \right\| _ {\Sigma} ^ {2} + \nu^ {2}\right) ^ {2}} \cdot a ^ {2} + \operatorname {t r a c e} \left(\Sigma^ {\frac {1}{2}} \widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma^ {\frac {1}{2}}\right) \\ \end{array}
$$
By the Lemma 3 in [7], it holds that
$$
\left\| \Sigma^ {\frac {1}{2}} \widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma^ {\frac {1}{2}} - I _ {d} \right\| _ {2} \leq O \left(\frac {1}{\sqrt {\lambda_ {\operatorname* {m i n}} n _ {2}}}\right). \tag {E.13}
$$
$$
\operatorname {t r a c e} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma_ {P _ {a}}\right) \leq \frac {1 + \frac {1}{\sqrt {\lambda_ {\operatorname* {m i n}} n _ {2}}}}{\| \widehat {\beta} \| _ {\Sigma} ^ {2}} \cdot a ^ {2} + O \left(d \left(1 + \frac {1}{\sqrt {\lambda_ {\operatorname* {m i n}} n _ {2}}}\right)\right).
$$
Then, what left is to bound  $\| \widehat{\beta}\|_{\Sigma} = \| \widehat{\theta}\|_{A\Sigma A^{\top}}\geq \| \theta^{*}\|_{A\Sigma A^{\top}} - \| \widehat{\theta} -\theta^{*}\|_{A\Sigma A^{\top}}$  by triangle inequality. On one hand,
$$
\left\| \theta^ {*} \right\| _ {A \Sigma A ^ {\top}} = \left\| \beta^ {*} \right\| _ {\Sigma}. \tag {E.14}
$$
$$
\begin{array}{l} \left\| \widehat {\theta} - \theta^ {*} \right\| _ {A \Sigma A ^ {\top}} = \mathcal {O} \left(\left\| \widehat {\theta} - \theta^ {*} \right\| _ {\widehat {\Sigma} _ {\lambda}}\right) \\ = \mathcal {O} \left(\frac {\| \widehat {\theta} - \theta^ {*} \| _ {V _ {\lambda}}}{\sqrt {n _ {2}}}\right) \\ = \mathcal {O} \left(\sqrt {\frac {d \log (n _ {2})}{n _ {2}}}\right). \\ \end{array}
$$
with high probability. Thus when  $n_2 = \Omega \left(\frac{d}{\|\beta^*\|_\Sigma^2}\right)$
$$
\| \widehat {\beta} \| _ {\Sigma} \geq \frac {1}{2} \| \beta^ {*} \| _ {\Sigma}.
$$
$$
\operatorname {t r a c e} \left(\widehat {\Sigma} _ {\lambda} ^ {- 1} \Sigma_ {P _ {a}}\right) \leq \mathcal {O} \left(\frac {1 + \frac {1}{\sqrt {\lambda_ {\operatorname* {m i n}} n _ {2}}}}{\| \beta^ {*} \| _ {\Sigma}} \cdot a ^ {2} + d \left(1 + \frac {1}{\sqrt {\lambda_ {\operatorname* {m i n}} n _ {2}}}\right)\right) = \mathcal {O} \left(\frac {a ^ {2}}{\| \beta^ {*} \| _ {\Sigma}} + d\right). \tag {E.15}
$$
E.5 Proof of lemma D.7
Proof. Recall the definition of  $g(x)$  that
$$
g (x) = \theta^ {* ^ {\top}} A A ^ {\top} x,
$$
note that  $g(x) = \theta^{*^{\top}}x$  when  $x$  is supported on  $\mathcal{A}$ . Thus,
$$
\begin{array}{l} \left| \mathbb {E} _ {x \sim P _ {a}} [ g (x) ] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ g (x) ] \right| \\ = \left| \mathbb {E} _ {x \sim P _ {a}} \left[ \theta^ {* ^ {\top}} x \right] - \mathbb {E} _ {x \sim \hat {P} _ {a}} \left[ \theta^ {* ^ {\top}} A A ^ {\top} x \right] \right| \\ \leq \left| \mathbb {E} _ {x \sim P _ {a}} [ \theta^ {* ^ {\top}} x ] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \theta^ {* ^ {\top}} V V ^ {\top} x ] \right| + \underbrace {\left| \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \theta^ {* ^ {\top}} V V ^ {\top} x ] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \theta^ {* ^ {\top}} A A ^ {\top} x ] \right|} _ {e _ {1}}, \\ \end{array}
$$
$$
\begin{array}{l} e _ {1} = \left| \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \theta^ {* ^ {\top}} \left(V V ^ {\top} - A A ^ {\top}\right) x \right] \right| \\ \leq \mathbb {E} _ {x \sim \hat {P} _ {a}} \left[ \left(\left\| \left(V V ^ {\top} - A A ^ {\top}\right) x \right\| \right] \right. \\ \leq \| V V ^ {\top} - A A ^ {\top} \| _ {F} \cdot \sqrt {\mathbb {E} _ {x \sim \hat {P} _ {a}} \left[ \| x \| _ {2} ^ {2} \right]}. \\ \end{array}
$$
Use notation  $P^{LD}(a) = P(z \mid \widehat{f}(Az) = a), P_{t_0}^{LD}(a) = P_{t_0}(z \mid \widehat{f}(Az) = a)$ .
$$
\begin{array}{l} \left| \mathbb {E} _ {x \sim P _ {a}} [ \theta^ {* ^ {\top}} x ] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \theta^ {* ^ {\top}} V V ^ {\top} x ] \right| \\ = \left| \mathbb {E} _ {z \sim P (z | \widehat {f} (A z) = a)} \left[ \theta^ {* ^ {\top}} A z \right] - \mathbb {E} _ {z \sim (U ^ {\top} V ^ {\top}) _ {\#}} \widehat {P} _ {a} \left[ \theta^ {* ^ {\top}} V U z \right] \right| \\ \leq \left| \mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \theta^ {* ^ {\top}} A z ] - \mathbb {E} _ {z \sim (U ^ {\top} V ^ {\top}) _ {\#}} \widehat {P} _ {a} [ \theta^ {* ^ {\top}} V U z ] \right| + \underbrace {\left| \mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \theta^ {* ^ {\top}} A z ] - \mathbb {E} _ {z \sim P ^ {L D} (a)} [ \theta^ {* ^ {\top}} A z ] \right|} _ {e _ {2}}, \\ \end{array}
$$
$$
\begin{array}{l} e _ {2} = \left| \alpha (t _ {0}) \mathbb {E} _ {z \sim P ^ {L D} (a)} [ \theta^ {* ^ {\top}} A z ] + h (t _ {0}) \mathbb {E} _ {u \sim \mathsf {N} (0, 1 _ {\mathrm {d}})} [ \theta^ {* ^ {\top}} A u ] - \mathbb {E} _ {z \sim P ^ {L D} (a)} [ \theta^ {* ^ {\top}} A z ] \right| \\ \leq \left(1 - \alpha \left(t _ {0}\right)\right) \left| \mathbb {E} _ {z \sim P ^ {L D} (a)} \left[ \theta^ {* ^ {\top}} A z \right] \right| \\ \leq t _ {0} \cdot \mathbb {E} _ {z \sim P ^ {L D} (a)} [ \| z \| _ {2} ]. \\ \end{array}
$$
Then what is left to bound is
$$
\begin{array}{l} \left| \mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \theta^ {* ^ {\top}} A z ] - \mathbb {E} _ {z \sim (U ^ {\top} V ^ {\top}) _ {\#}} \widehat {P} _ {a} [ \theta^ {* ^ {\top}} V U z ] \right| \\ \leq \underbrace{\left|\mathbb{E}_{z\sim P_{t_{0}}^{LD}(a)}\bigl[\theta^{*^{\top}}VUz\bigr] - \mathbb{E}_{z\sim(U^{\top}V^{\top})_{\#}}\widehat{P}_{a}\bigl[\theta^{*^{\top}}VUz\bigr]\right|}_{e_{3}} + \underbrace{\left|\mathbb{E}_{z\sim P_{t_{0}}^{LD}(a)}\bigl[\theta^{*^{\top}}VUz\bigr] - \mathbb{E}_{z\sim P_{t_{0}}^{LD}(a)}\bigl[\theta^{*^{\top}}A z\bigr]\right|}_{e_{4}}. \\ \end{array}
$$
Then for term  $e_3$ , by Lemma E.2, we get
$$
\begin{array}{l} e _ {3} \leq \left| \mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \| z \| _ {2} ] - \mathbb {E} _ {z \sim (U ^ {\top} V ^ {\top}) _ {\#}} \widehat {P} _ {a} [ \| z \| _ {2} ] \right| \\ = \mathcal {O} \left(T V (\widehat {P} _ {a}) \cdot \left(\sqrt {\mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \| z \| _ {2} ^ {2} ]} + \sqrt {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x \| _ {2} ^ {2} ]}\right)\right), \\ \end{array}
$$
where we use  $\mathbb{E}_{z\sim (U^{\top}V^{\top})_{\#}}\widehat{P}_a[\| z\| _2]\leq \mathbb{E}_{x\sim \widehat{P}_a}\left[\| x\| _2\right]$
For  $e_4$ , we have
$$
\begin{array}{l} e _ {4} = \left| \mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} [ \theta^ {* ^ {\top}} (V U - A) z ] \right| \\ = \alpha (t _ {0}) \left| \mathbb {E} _ {z \sim P (a)} \left[ \theta^ {* ^ {\top}} (V U - A) z \right] \right| \\ \leq \| V U - A \| _ {F} \cdot \mathbb {E} _ {z \sim P ^ {L D} (a)} [ \| z \| _ {2} ]. \\ \end{array}
$$
Therefore, by combining things together, we have
$$
\begin{array}{l} \mathcal {E} _ {2} \leq e _ {1} + e _ {2} + e _ {3} + e _ {4} \\ \leq \| V V ^ {\top} - A A ^ {\top} \| _ {F} \cdot \sqrt {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x \| _ {2} ^ {2} ]} + (\| V U - A \| _ {F} + t _ {0}) \cdot \sqrt {M (a)} \\ + \mathcal {O} \left(T V (\widehat {P} _ {a}) \cdot \left(\sqrt {M (a) + t _ {0} d} + \sqrt {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x \| _ {2} ^ {2} ]}\right)\right). \\ \end{array}
$$
By Lemma D.4 and Lemma E.1, we have
$$
T V (\widehat {P} _ {a}) = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} (P (x , \widehat {y} = a) , P _ {x \widehat {y}} ; \bar {S})}{\lambda_ {\min}}} \cdot \epsilon_ {d i f f}\right),
$$
$$
\| V V ^ {\top} - A A ^ {\top} \| _ {\mathrm {F}} = \widetilde {\mathcal {O}} \left(\frac {\sqrt {t _ {0}}}{\sqrt {\lambda_ {\operatorname* {m i n}}}} \cdot \epsilon_ {d i f f}\right),
$$
$$
\| V U - A \| _ {F} = \mathcal {O} \left(d ^ {\frac {3}{2}} \| V V ^ {\top} - A A ^ {\top} \| _ {F}\right).
$$
And by Lemma E.3
$$
\mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| x \| _ {2} ^ {2} \right] = \mathcal {O} \left(c t _ {0} D + M (a) \cdot \left(1 + T V (\widehat {P} _ {a})\right)\right).
$$
Therefore, leading term in  $\mathcal{E}_2$  is
$$
\mathcal {E} _ {2} = \mathcal {O} \left(\left(T V \left(\widehat {P} _ {a}\right) + t _ {0}\right) \sqrt {M (a)}\right).
$$
By plugging in score matching error  $\epsilon_{diff}^2 = \widetilde{\mathcal{O}}\left(\frac{1}{t_0}\sqrt{\frac{Dd^2 + D^2d}{n_1}}\right)$ , we have
$$
T V (\widehat {P} _ {a}) = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} (P (x , \widehat {y} = a) , P _ {x \widehat {y}} ; \bar {\mathcal {S}})}{\lambda_ {\mathrm {m i n}}} \cdot \left(\frac {D d ^ {2} + D ^ {2} d}{n _ {1}}\right) ^ {\frac {1}{4}} \cdot \frac {1}{\sqrt {t _ {0}}}}\right).
$$
When  $t_0 = \left(\frac{Dd^2 + D^2d}{n_1}\right)^{1/6}$ , it admits the best trade off in  $\mathcal{E}_2$  and  $\mathcal{E}_2$  is bounded by
$$
\mathcal {E} _ {2} = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} (P (x , \widehat {y} = a) , P _ {x \widehat {y}} ; \bar {\mathcal {S}})}{\lambda_ {\mathrm {m i n}}} \cdot \left(\frac {D d ^ {2} + D ^ {2} d}{n _ {1}}\right) ^ {\frac {1}{6}} \cdot a}\right).
$$
E.6 Proof of Lemma E.1
Proof. From Lemma 17 in [8], we have
$$
\| U - V ^ {\top} A \| _ {F} = \mathcal {O} (\| V V ^ {\top} - A A ^ {\top} \| _ {F}).
$$
Then it suffices to bound
$$
\left| \left\| V U - A \right\| _ {F} ^ {2} - \left\| U - V ^ {\top} A \right\| _ {F} ^ {2} \right|,
$$
$$
\begin{array}{l} \left\| V U - A \right\| _ {F} ^ {2} = 2 d - \operatorname {t r a c e} \left(U ^ {\top} V ^ {\top} A + A ^ {\top} V U\right) \\ \| U - V ^ {\top} A \| _ {F} ^ {2} = d + \operatorname {t r a c e} \left(A ^ {\top} V V ^ {\top} A\right) - \operatorname {t r a c e} \left(U ^ {\top} V ^ {\top} A + A ^ {\top} V U\right). \\ \end{array}
$$
$$
\left| \| V U - A \| _ {F} ^ {2} - \| U - V ^ {\top} A \| _ {F} ^ {2} \right| = \left| d - \operatorname {t r a c e} \left(A ^ {\top} V V ^ {\top} A\right) \right| = \left| \operatorname {t r a c e} \left(A A ^ {\top} \left(V V ^ {\top} - A A ^ {\top}\right)\right) \right|,
$$
which is because trace  $\left(A^{\top}VV^{\top}A\right)$  is calculated as
$$
\begin{array}{l} \operatorname {t r a c e} \left(A ^ {\top} V V ^ {\top} A\right) = \operatorname {t r a c e} \left(A A ^ {\top} V V ^ {\top}\right) \\ = \operatorname {t r a c e} \left(A A ^ {\top} A A ^ {\top}\right) + \operatorname {t r a c e} \left(A A ^ {\top} \left(V V ^ {\top} - A A ^ {\top}\right)\right) \\ = d + \operatorname {t r a c e} \left(A A ^ {\top} \left(V V ^ {\top} - A A ^ {\top}\right)\right). \\ \end{array}
$$
Then we will bound  $\left|\operatorname{trace}\left(AA^{\top}(VV^{\top} - AA^{\top})\right)\right|$  by  $\|VV^{\top} - AA^{\top}\|_{F}$ ,
$$
\begin{array}{l} \left| \operatorname {t r a c e} \left(A A ^ {\top} \left(V V ^ {\top} - A A ^ {\top}\right)\right) \right| \leq \operatorname {t r a c e} \left(A A ^ {\top}\right) \operatorname {t r a c e} \left(\left| V V ^ {\top} - A A ^ {\top} \right|\right) \\ \leq d \cdot \operatorname {t r a c e} \left(\left| V V ^ {\top} - A A ^ {\top} \right|\right) \\ \leq d \cdot \sqrt {2 d \| V V ^ {\top} - A A ^ {\top} \| _ {F} ^ {2}}. \\ \end{array}
$$
Thus,  $\| VU - A\| _F = \mathcal{O}\left(d^{\frac{3}{2}}\sqrt{\angle(V,A)}\right)$
E.7 Proof of Lemma E.2
Proof. When  $P_{1}$  and  $P_{2}$  are Gaussian,  $m(z) = \| z\| _2^2$
$$
\begin{array}{l} \left. \left| \mathbb {E} _ {z \sim P _ {1}} [ m (z) ] - \mathbb {E} _ {z \sim P _ {2}} [ m (z) ] \right| \right. \\ = \left| \int m (z) \left(p _ {1} (z) - p _ {2} (z)\right) \mathrm {d} z \right| \\ \leq \left| \int_ {\| z \| _ {2} \leq R} \| z \| _ {2} ^ {2} \left(p _ {1} (z) - p _ {2} (z)\right) \mathrm {d} z \right| + \int_ {\| z \| _ {2} > R} \| z \| _ {2} ^ {2} p _ {1} (z) \mathrm {d} z + \int_ {\| z \| _ {2} > R} \| z \| _ {2} ^ {2} p _ {2} (z) \mathrm {d} z \\ \leq R ^ {2} \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right) + \int_ {\| z \| _ {2} > R} \| z \| _ {2} ^ {2} p _ {1} (z) \mathrm {d} z + \int_ {\| z \| _ {2} > R} \| z \| _ {2} ^ {2} p _ {2} (z) \mathrm {d} z. \\ \end{array}
$$
Since  $P_{1}$  and  $P_{2}$  are Gaussians,  $\int_{\|z\|_2 > R} \|z\|_2^2 p_1(z) \mathrm{d}z$  and  $\int_{\|z\|_2 > R} \|z\|_2^2 p_2(z) \mathrm{d}z$  are bounded by some constant  $C_{1}$  when  $R^2 \geq C_2 \max \{\mathbb{E}_{z \sim P_1}[\|z\|_2^2], \mathbb{E}_{z \sim P_2}[\|z\|_2^2]\}$  as suggested by Lemma 16 in [8].
$$
\begin{array}{l} \mathbb {E} _ {z \sim P _ {1}} [ \| z \| _ {2} ^ {2} ] \leq \mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ] + C _ {2} \max  \left\{\mathbb {E} _ {P _ {1}} [ \| z \| _ {2} ^ {2} ], \mathbb {E} _ {P _ {2}} [ \| z \| _ {2} ^ {2} ] \right\} \cdot \mathrm {d} _ {\mathrm {T V}} (P _ {1}, P _ {2}) + 2 C _ {1} \\ \leq \mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ] + C _ {2} (\mathbb {E} _ {z \sim P _ {1}} [ \| z \| _ {2} ^ {2} ] + \mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2}) \cdot \mathrm {d} _ {\mathrm {T V}} (P _ {1}, P _ {2}) + 2 C _ {1}. \\ \end{array}
$$
$$
\mathbb {E} _ {z \sim P _ {1}} [ \| z \| _ {2} ^ {2} ] = \mathcal {O} \left(\mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ] + \mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ] \cdot \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right)\right)
$$
since  $\mathsf{d}_{\mathrm{TV}}(P_1,P_2)$  decays with  $n_1$
Similarly, when  $m(z) = \| z\| _2$