SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} (Q R) v = Q (R v) \\ = Q (Q w) = Q ^ {2} w \\ = Q w = R v. \\ \end{array}
$$
Since the choice of  $v$  is arbitrary, it must be the case that  $R\Sigma w_{i}^{*} = 0$  and  $QR = R$ . For any  $v\in \mathbb{R}^d$ , we can check
$$
(R Q) v = R (Q v) = Q v
$$
by the definition of orthogonal projection. Therefore, it must be the case that  $RQ = Q$ . Finally, we use  $R = QR = RQ^T$  to show that
$$
\begin{array}{l} \operatorname {T r} (R \Sigma R) = \operatorname {T r} (R Q ^ {T} \Sigma Q R) = \operatorname {T r} (Q ^ {T} \Sigma Q R) \\ = \operatorname {T r} (Q ^ {T} \Sigma Q) - \operatorname {T r} (Q ^ {T} \Sigma Q (I - R)) \\ \geq \operatorname {T r} \left(Q ^ {T} \Sigma Q\right) - \sqrt {\operatorname {T r} \left(\left(Q ^ {T} \Sigma Q\right) ^ {2}\right) \operatorname {T r} \left(\left(I - R\right) ^ {2}\right)} \\ = \operatorname {T r} \left(Q ^ {T} \Sigma Q\right) \left(1 - \sqrt {\frac {k}{R \left(Q ^ {T} \Sigma Q\right)}}\right) \\ = \operatorname {T r} \left(Q ^ {T} \Sigma Q\right) \left(1 - \frac {k}{n} - \frac {n}{R \left(Q ^ {T} \Sigma Q\right)}\right) \\ \end{array}
$$
$$
\begin{array}{l} \operatorname {T r} \left(\left(R \Sigma R\right) ^ {2}\right) = \operatorname {T r} (\Sigma R \Sigma R) \\ = \operatorname {T r} \left(\Sigma Q R Q ^ {T} \Sigma Q R Q ^ {T}\right) \\ = \operatorname {T r} ((R Q ^ {T} \Sigma Q) R (Q ^ {T} \Sigma Q R)) \\ \leq \operatorname {T r} \left(\left(R Q ^ {T} \Sigma Q\right) \left(Q ^ {T} \Sigma Q R\right)\right) = \operatorname {T r} \left(\left(Q ^ {T} \Sigma Q\right) ^ {2} R\right) \\ \leq \operatorname {T r} \left(\left(Q ^ {T} \Sigma Q\right) ^ {2}\right). \\ \end{array}
$$
Rearranging concludes the proof.
D.1 Phase Retrieval
Proof. Without loss of generality, we assume that  $\mu$  lies in the span of  $\{\Sigma w_1^*, \dots, \Sigma w_k^*\}$  because otherwise we can simply increase  $k$  by one. Moreover, we can assume that  $\{\Sigma^{1/2}w_1^*, \dots, \Sigma^{1/2}w_k^*\}$  are orthonormal because otherwise we let  $\tilde{W} = W(W^T\Sigma W)^{-1}$  and conditioning on  $W^T(x - \mu)$  is the same as conditioning on  $\tilde{W}^T(x - \mu)$ . By Lemma 5, conditioned on
$$
\left( \begin{array}{c} \eta_ {1} ^ {T} \\ \dots \\ \eta_ {k} ^ {T} \end{array} \right) = [ W ^ {T} (x _ {1} - \mu), \dots , W ^ {T} (x _ {n} - \mu) ]
$$
the distribution of  $X$  is the same as
$$
X = 1 \mu^ {T} + \sum_ {i = 1} ^ {k} \eta_ {i} \left(\Sigma w _ {i} ^ {*}\right) ^ {T} + Z \Sigma^ {1 / 2} Q
$$
where  $Z$  has i.i.d. standard normal entries. Furthermore, conditioned on  $W^{T}(x - \mu)$  and the noise of variable in  $y$  (which is independent of  $x$ ), by the multi-index assumption (B), the label  $y$  is non-random. Since  $Qw^{\sharp} = 0$ , we have  $w^{\sharp} = \sum_{i=1}^{k}\langle w_i^*, \Sigma w^{\sharp}\rangle w_i^*$  and so
$$
\langle w ^ {\sharp}, x \rangle = \langle w ^ {\sharp}, \mu \rangle + \sum_ {i = 1} ^ {k} \langle w _ {i} ^ {*}, \Sigma w ^ {\sharp} \rangle \langle w _ {i} ^ {*}, x - \mu \rangle .
$$
Therefore,  $\langle w^{\sharp},x\rangle$  also becomes non-random after conditioning. We can let  $I = \{i\in [n]:\langle w^{\sharp},x_i\rangle \geq 0\}$  and define  $\xi \in \mathbb{R}^n$  by
$$
\xi_ {i} = \left\{ \begin{array}{l l} y _ {i} - | \langle w ^ {\sharp}, x _ {i} \rangle | & \text {i f} i \in I \\ | \langle w ^ {\sharp}, x _ {i} \rangle | - y _ {i} & \text {i f} i \notin I \end{array} \right.
$$
and  $\xi$  is non-random after conditioning. Following the construction discussed in the main text, for any  $w^{\sharp} \in \mathbb{R}^{d}$ , the predictor  $w = w^{\sharp} + w^{\perp}$  satisfies  $|\langle w, x_{i} \rangle| = y_{i}$  where
$$
w ^ {\perp} = \arg \min  \| w \| _ {2}
$$
$$
\begin{array}{c} w \in \mathbb {R} ^ {d}: \\ X w = \xi \end{array}
$$
by the definition of  $\xi$ . Hence, we have
$$
\min  _ {w \in \mathbb {R} ^ {d}: \forall i \in [ n ], \langle w, x _ {i} \rangle^ {2} = y _ {i} ^ {2}} \| w \| _ {2} \leq \| w ^ {\sharp} \| _ {2} + \| w ^ {\perp} \| _ {2}
$$
and it suffices to control  $\| w^{\perp}\| _2$
Let  $R$  be the orthogonal projection matrix onto the image of  $Q$  and we consider  $w$  of the form  $Rw$  to upper bound  $\| w^{\perp}\|_{2}$ . By Lemma 13, we know  $QR = R$  and  $R\Sigma w_{i}^{*} = 0$ . By the assumption that  $\mu$  lies in the span of  $\{\Sigma w_1^*,\dots,\Sigma w_k^*\}$ , we have
$$
\left(1 \mu^ {T} + \sum_ {i = 1} ^ {k} \eta_ {i} \left(\Sigma w _ {i} ^ {*}\right) ^ {T} + Z \Sigma^ {1 / 2} Q\right) R w = Z \Sigma^ {1 / 2} R w.
$$
Since  $R$  is an orthogonal projection, it holds that  $\| Rw\|_2 \leq \| w\|_2$ . Finally, we observe that the distribution of  $Z\Sigma^{1/2}R$  is the same as  $Z(R\Sigma R)^{1/2}$  and so
$$
\| w^{\perp}\|_{2}\leq \min_{\substack{w\in \mathbb{R}^{d}:\\ Z(R\Sigma R)^{1 / 2}w = \xi}}\| w\|_{2}.
$$
We are now ready to apply Lemma 12 to the covariance  $R\Sigma R$ . We are allowed to replace the dependence on  $R\Sigma R$  by the dependence on  $\Sigma^{\perp}$  by the last two inequalities of Lemma 13. The desired conclusion follows by the observation that  $\| \xi \|_2^2 = n\hat{L}_f(w^\sharp)$  and the assumption that  $\hat{L}_f(w^\sharp) \leq (1 + \rho)L_f(w^\sharp)$ .
D.2 ReLU Regression
The proof of Theorem 3 will closely follow the proof of Theorem 2.
Proof. We let  $I = \{i\in [n]:y_i > 0\}$  and for any  $(w^{\sharp},b^{\sharp})\in \mathbb{R}^{d + 1}$ , we define  $\xi \in \mathbb{R}^n$  by
$$
\xi_ {i} = \left\{ \begin{array}{l l} y _ {i} - \langle w ^ {\sharp}, x _ {i} \rangle - b ^ {\sharp} & \text {i f} i \in I \\ - \sigma (\langle w ^ {\sharp}, x _ {i} \rangle + b ^ {\sharp}) & \text {i f} i \notin I. \end{array} \right.
$$
By the definition of  $\xi$ , the predictor  $(w, b) = (w^{\sharp} + w^{\perp}, b^{\sharp})$  satisfies  $\sigma(\langle w, x_i \rangle + b) = y_i$  where
$$
w^{\perp} = \operatorname *{arg  min}_{\substack{w\in \mathbb{R}^{d}:\\ Xw = \xi}}\| w\|_{2}.
$$
$$
\min_{\substack{(w,b)\in \mathbb{R}^{d + 1}:\\ \forall i\in [n],\sigma (\langle w,x_{i}\rangle +b) = y_{i}}}\| w\|_{2}\leq \| w^{\sharp}\|_{2} + \| w^{\perp}\|_{2}
$$
Similar to the proof of Theorem 2, we make the simplifying assumption that  $\mu$  lies in the span of  $\{\Sigma w_1^*, \dots, \Sigma w_k^*\}$  and  $\{\Sigma^{1/2}w_1^*, \dots, \Sigma^{1/2}w_k^*\}$  are orthonormal. Conditioned on  $W^T(x_i - \mu)$  and the noise variable in  $y_i$ , both  $y_i$  and  $\langle w^\sharp, x_i \rangle$  are non-random, and so  $\xi$  is also non-random. The distribution of  $X$  is the same as
$$
X = 1 \mu^ {T} + \sum_ {i = 1} ^ {k} \eta_ {i} (\Sigma w _ {i} ^ {*}) ^ {T} + Z \Sigma^ {1 / 2} Q.
$$
If we consider  $w$  of the form  $Rw$ , then we have
We are now ready to apply Lemma 12 to the covariance  $R\Sigma R$ . We are allowed to replace the dependence on  $R\Sigma R$  by the dependence on  $\Sigma^{\perp}$  by the last two inequalities of Lemma 13. The desired conclusion follows by the observation that  $\| \xi \|_2^2 = n\hat{L}_f(w^\sharp, b^\sharp)$  due to the definition (13) and the assumption that  $\hat{L}_f(w^\sharp) \leq (1 + \rho)L_f(w^\sharp, b^\sharp)$ .
D.3 Low-rank Matrix Sensing
Proof. Without loss of generality, we will assume that  $d_{1} \leq d_{2}$ . We will vectorize the measurement matrices and estimator  $A_{1}, \ldots, A_{n}, X \in \mathbb{R}^{d_{1} \times d_{2}}$  as  $a_{1}, \ldots, a_{n}, x \in \mathbb{R}^{d_{1} d_{2}}$  and define  $\| x \|_{*} = \| X \|_{*}$ . Denote  $A = [a_{1}, \ldots, a_{n}]^{T} \in \mathbb{R}^{n \times d_{1} d_{2}}$ . We define the primary problem  $\Phi$  by
$$
\Phi := \min  _ {\forall i \in [ n ], \langle A _ {i}, X \rangle = \xi} \| X \| _ {*} = \min  _ {A x = \xi} \| x \| _ {*}.
$$
By Lemma 11, it suffices to consider the auxiliary problem
$$
\Psi := \min  _ {\| G \| x \| _ {2} - \xi \| _ {2} \leq - \langle H, x \rangle} \| x \| _ {*}.
$$
We will pick  $x$  of the form  $x = -\alpha H$  for some  $\alpha \geq 0$ , which needs to satisfy  $\alpha \| H \|_2^2 \geq \| \alpha G \| H \|_2 - \xi \|_2$ . By a union bound, the following events occur simultaneously with probability at least  $1 - \delta / 2$ :
1. by Lemma 3, it holds that
$$
\begin{array}{l} \left\| G \right\| _ {2} \leq \sqrt {n} + 2 \sqrt {\log (3 2 / \delta)} \\ \frac {\| \xi \| _ {2}}{\sigma} \leq \sqrt {n} + 2 \sqrt {\log (3 2 / \delta)} \\ \| H \| _ {2} \leq \sqrt {d _ {1} d _ {2}} + 2 \sqrt {\log (3 2 / \delta)} \\ \end{array}
$$
2. Condition on  $\xi$ , we have  $\frac{1}{\|\xi\|}\langle G,\xi\rangle \sim \mathcal{N}(0,1)$  and so by standard Gaussian tail bound  $\operatorname*{Pr}(|Z| > t) \leq 2e^{-t^2 /2}$
$$
\frac {| \langle G , \xi \rangle |}{\| \xi \|} \leq \sqrt {2 \log (1 6 / \delta)}
$$
Then we can use AM-GM inequality to show for sufficiently large  $n$
$$
\begin{array}{l} \| \alpha G \| H \| _ {2} - \xi \| _ {2} ^ {2} \\ = \alpha^ {2} \| G \| _ {2} ^ {2} \| H \| _ {2} ^ {2} + \| \xi \| ^ {2} - 2 \alpha \| H \| _ {2} \langle G, \xi \rangle \\ \leq n \alpha^ {2} \| H \| _ {2} ^ {2} \left(1 + 2 \sqrt {\frac {\log (3 2 / \delta)}{n}}\right) ^ {2} + \| \xi \| ^ {2} + 2 \sqrt {n} \alpha \| H \| _ {2} \| \xi \| _ {2} \sqrt {\frac {2 \log (1 6 / \delta)}{n}} \\ \leq n \alpha^ {2} \| H \| _ {2} ^ {2} \left(1 + 1 0 \sqrt {\frac {\log (3 2 / \delta)}{n}}\right) + \left(1 + \sqrt {\frac {2 \log (1 6 / \delta)}{n}}\right) \| \xi \| _ {2} ^ {2} \\ \end{array}
$$
and it suffices to let
$$
\alpha^ {2} \| H \| _ {2} ^ {4} \geq n \alpha^ {2} \| H \| _ {2} ^ {2} \left(1 + 1 0 \sqrt {\frac {\log (3 2 / \delta)}{n}}\right) + \left(1 + \sqrt {\frac {2 \log (1 6 / \delta)}{n}}\right) \| \xi \| _ {2} ^ {2}.
$$
Rearranging the above inequality, we can choose
$$
\alpha = \left(\frac {1 + 1 0 \sqrt {\frac {\log (3 2 / \delta)}{n}}}{1 - \frac {n}{d _ {1} d _ {2}} \left(1 + 1 0 \sqrt {\frac {\log (3 2 / \delta)}{n}}\right) \left(1 + 2 \sqrt {\frac {\log (3 2 / \delta)}{d _ {1} d _ {2}}}\right) ^ {2}}\right) ^ {1 / 2} \frac {\sqrt {n \sigma^ {2}}}{\| H \| _ {2} ^ {2}}
$$
and since  $H$  as a matrix can have at most rank  $d_{1}$ , by Cauchy-Schwarz inequality on the singular values of  $H$ , we have  $\| H \|_{*} \leq \sqrt{d_{1}} \| H \|_{2}$  and
$$
\| x \| _ {*} = \alpha \| H \| _ {*} \leq \alpha \sqrt {d _ {1}} \| H \| _ {2} \leq (1 + \epsilon) \sqrt {\frac {d _ {1} \left(n \sigma^ {2}\right)}{d _ {1} d _ {2}}} = (1 + \epsilon) \sqrt {\frac {n \sigma^ {2}}{d _ {2}}}
$$
for some  $\epsilon \lesssim \sqrt{\frac{\log(32 / \delta)}{n}} +\frac{n}{d_1d_2}$ . The desired conclusion follows by the observation that  $\| X^{*}\|_{*}\leq \sqrt{r}\| X^{*}\|_{F}$  because  $X^{*}$  has rank  $r$ .
Proof. Note that  $\langle A,X^{*}\rangle \sim \mathcal{N}(0,\| X^{*}\|_{F}^{2})$  and so by the standard Gaussian tail bound  $\operatorname *{Pr}(|Z|\geq t)\leq 2e^{-t^2 /2}$ , Theorem 9 and a union bound, it holds with probability at least  $1 - \delta /8$  that
$$
\begin{array}{l} | \langle A, X ^ {*} \rangle | \leq \sqrt {2 \log (3 2 / \delta)} \| X ^ {*} \| _ {F} \\ \left\| A \right\| _ {o p} \leq \sqrt {d _ {1}} + \sqrt {d _ {2}} + \sqrt {2 \log (3 2 / \delta)}. \\ \end{array}
$$
$$
\begin{array}{l} \left\| A - \frac {\langle A , X ^ {*} \rangle}{\| X ^ {*} \| _ {F} ^ {2}} X ^ {*} \right\| _ {o p} \leq \| A \| _ {o p} + \frac {| \langle A , X ^ {*} \rangle |}{\| X ^ {*} \| _ {F} ^ {2}} \| X ^ {*} \| _ {o p} \\ \leq \sqrt {d _ {1}} + \sqrt {d _ {2}} + \sqrt {2 \log (3 2 / \delta)} + \frac {\| X ^ {*} \| _ {o p}}{\| X ^ {*} \| _ {F}} \sqrt {2 \log (3 2 / \delta)} \\ \leq \sqrt {d _ {1}} + \sqrt {d _ {2}} + \sqrt {8 \log (3 2 / \delta)}. \\ \end{array}
$$
Therefore, we can choose  $C_{\delta}$  in Theorem 1 by
$$
C _ {\delta} (X) := \left(\sqrt {d _ {1}} + \sqrt {d _ {2}} + \sqrt {8 \log (3 2 / \delta)}\right) \| X \| _ {*}
$$
and applying Theorem 1 and Theorem 4, we have
$$
\begin{array}{l} (1 - \epsilon) L (\hat {X}) \leq \frac {C _ {\delta} (X) ^ {2}}{n} \\ \leq \frac {\left(\sqrt {d _ {1}} + \sqrt {d _ {2}} + \sqrt {8 \log (3 2 / \delta)}\right) ^ {2}}{n} \left(\sqrt {r} \| X ^ {*} \| _ {F} + (1 + \epsilon) \sqrt {\frac {n \sigma^ {2}}{d _ {1} \vee d _ {2}}}\right) ^ {2} \\ = \left(\sqrt {\frac {d _ {1}}{d _ {1} \vee d _ {2}}} + \sqrt {\frac {d _ {2}}{d _ {1} \vee d _ {2}}} + \sqrt {\frac {8 \log (3 2 / \delta)}{d _ {1} \vee d _ {2}}}\right) ^ {2} \left(\sqrt {\frac {r (d _ {1} \vee d _ {2})}{n}} + (1 + \epsilon) \frac {\sigma}{\| X ^ {*} \| _ {F}}\right) ^ {2} \| X ^ {*} \| _ {F} ^ {2} \\ \end{array}
$$
where  $\epsilon$  is the maximum of the two  $\epsilon$  in Theorem 1 and Theorem 4. Finally, recall that
$$
L (\hat {X}) = \sigma^ {2} + \| \hat {X} - X ^ {*} \| _ {F} ^ {2}.
$$
Assuming that  $d_{1} \leq d_{2}$ , then the above implies that
$$
\begin{array}{l} \frac {\| \hat {X} - X ^ {*} \| _ {F} ^ {2}}{\| X ^ {*} \| _ {F} ^ {2}} \\ \leq (1 - \epsilon) ^ {- 1} (1 + \epsilon) ^ {2} \left(1 + \sqrt {\frac {d _ {1}}{d _ {2}}} + \sqrt {\frac {8 \log (3 2 / \delta)}{d _ {2}}}\right) ^ {2} \left(\sqrt {\frac {r (d _ {1} + d _ {2})}{n}} + \frac {\sigma}{\| X ^ {*} \| _ {F}}\right) ^ {2} - \frac {\sigma^ {2}}{\| X ^ {*} \| _ {F} ^ {2}} \\ \lesssim \frac {r (d _ {1} + d _ {2})}{n} + \sqrt {\frac {r (d _ {1} + d _ {2})}{n}} \frac {\sigma}{\| X ^ {*} \| _ {F}} + \left(\sqrt {\frac {d _ {1}}{d _ {2}}} + \frac {n}{d _ {1} d _ {2}}\right) \frac {\sigma^ {2}}{\| X ^ {*} \| _ {F} ^ {2}} \\ \end{array}
$$
E Counterexample to Gaussian Universality
By assumption (G), we can write  $x_{i|d - k} = h(x_{i|k}) \cdot \Sigma_{|d - k}^{1/2} z_i$  where  $z_i \sim \mathcal{N}(0, I_{d - k})$ . We will denote the matrix  $Z = [z_1, \dots, z_n]^T \in \mathbb{R}^{n \times (d - k)}$ . Following the notation in section 7, we will also write  $X = [X_{|k}, X_{|d - k}]$  where  $X_{|k} \in \mathbb{R}^{n \times k}$  and  $X_{|d - k} \in \mathbb{R}^{n \times (d - k)}$ . The proofs in this section closely follow the proof of Theorem 6.
Theorem 15. Consider dataset  $(X,Y)$  drawn i.i.d. from the data distribution  $\mathcal{D}$  according to  $(G)$  and  $(H)$ , and fix any  $f:\mathbb{R}\times \mathcal{V}\to \mathbb{R}_{\geq 0}$  such that  $\sqrt{f}$  is I-Lipschitz for any  $y\in \mathcal{V}$ . Fix any  $\delta >0$  and suppose there exists  $\epsilon_{\delta} < 1$  and  $C_\delta :\mathbb{R}^{d - k}\rightarrow [0,\infty ]$  such that
(i) with probability at least  $1 - \delta / 2$  over  $(X, Y)$  and  $G \sim \mathcal{N}(0, I_n)$ , it holds uniformly over all  $w_{|k} \in \mathbb{R}^k$  and  $\| w_{|d - k} \|_{\Sigma_{|d - k}} \in \mathbb{R}_{\geq 0}$  that
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \frac {f (\langle w _ {| k} , x _ {i | k} \rangle + h (x _ {i | k}) \| w _ {| d - k} \| _ {\Sigma_ {| d - k}} G _ {i} , y _ {i})}{h (x _ {i | k}) ^ {2}} \geq (1 - \epsilon_ {\delta}) \mathbb {E} _ {\mathcal {D}} \left[ \frac {f (\langle w , x \rangle , y)}{h (x _ {| k}) ^ {2}} \right]
$$
(ii) with probability at least  $1 - \delta /2$  over  $z_{|d - k}\sim \mathcal{N}(0,\Sigma_{|d - k})$ , it holds uniformly over all  $w_{|d - k}\in \mathbb{R}^{d - k}$  that
$$
\langle w _ {| d - k}, z _ {| d - k} \rangle \leq C _ {\delta} (w _ {| d - k}) \tag {75}
$$
then with probability at least  $1 - \delta$ , it holds uniformly over all  $w \in \mathbb{R}^d$  that
$$
(1 - \epsilon_ {\delta}) \mathbb {E} \left[ \frac {f (\langle w , x \rangle , y)}{h (x _ {| k}) ^ {2}} \right] \leq \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \frac {f (\langle w , x _ {i} \rangle , y _ {i})}{h (x _ {| k}) ^ {2}} + \frac {C _ {\delta} (w _ {| d - k})}{\sqrt {n}}\right) ^ {2}. \tag {76}
$$
$$
\langle w _ {| d - k}, x _ {i | d - k} \rangle = h (x _ {i | k}) \cdot \langle w _ {| d - k}, \Sigma_ {| d - k} ^ {1 / 2} z _ {i} \rangle
$$
and so for any  $f:\mathbb{R}\times \mathcal{Y}\times \mathbb{R}^k\to \mathbb{R}$  we can write
$$
\begin{array}{l} \Phi := \sup  _ {w \in \mathbb {R} ^ {d}} F (w) - \frac {1}{n} \sum_ {i = 1} ^ {n} f (\langle w, x _ {i} \rangle , y _ {i}, x _ {i | k}) \\ = \sup_{\substack{w\in \mathbb{R}^{d},u\in \mathbb{R}^{n}\\ u = Z\Sigma_{|d - k}^{1 / 2}w_{|d - k}}}F(w) - \frac{1}{n}\sum_{i = 1}^{n}f(\langle w_{|k},x_{i|k}\rangle +h(x_{i|k})u_{i},y_{i},x_{i|k}) \\ = \sup  _ {w \in \mathbb {R} ^ {d}, u \in \mathbb {R} ^ {n}} \inf  _ {\lambda \in \mathbb {R} ^ {n}} \langle \lambda , Z \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} - u \rangle + F (w) - \frac {1}{n} \sum_ {i = 1} ^ {n} f (\langle w _ {| k}, x _ {i | k} \rangle + h (x _ {i | k}) u _ {i}, y _ {i}, x _ {i | k}). \\ \end{array}
$$
By the same truncation argument used in Lemma 7, it suffices to consider the auxiliary problem:
$$
\begin{array}{l} \Psi := \sup  _ {w \in \mathbb {R} ^ {d}, u \in \mathbb {R} ^ {n}} \inf  _ {\lambda \in \mathbb {R} ^ {n}} \| \lambda \| _ {2} \langle H, \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \rangle + \langle G \| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \| _ {2} - u, \lambda \rangle \\ + F (w) - \frac {1}{n} \sum_ {i = 1} ^ {n} f \left(\left\langle w _ {| k}, x _ {i | k} \right\rangle + h \left(x _ {i | k}\right) u _ {i}, y _ {i}, x _ {i | k}\right) \\ = \sup  _ {w \in \mathbb {R} ^ {d}, u \in \mathbb {R} ^ {n}} \inf  _ {\lambda \geq 0} \lambda \left(\langle H, \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \rangle - \left\| G \| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \| _ {2} - u \right\| _ {2}\right) \\ + F (w) - \frac {1}{n} \sum_ {i = 1} ^ {n} f (\langle w _ {| k}, x _ {i | k} \rangle + h (x _ {i | k}) u _ {i}, y _ {i}, x _ {i | k}) \\ \end{array}
$$
$$
\begin{array}{l} \Psi = \sup_{\substack{w\in \mathbb{R}^{d},u\in \mathbb{R}^{n}\\ \langle H,\Sigma_{|d - k}^{1 / 2}w_{|d - k}\rangle \geq \left\| G\| \Sigma_{|d - k}^{1 / 2}w_{|d - k}\|_{2} - u\right\|_{2}}}F(w) - \frac{1}{n}\sum_{i = 1}^{n}f(\langle w_{|k},x_{i|k}\rangle +h(x_{i|k})u_{i},y_{i},x_{i|k}) \\ = \sup_{w\in \mathbb{R}^{d}}F(w) - \frac{1}{n}\inf_{\substack{u\in \mathbb{R}^{n}\\ \langle H,\Sigma_{|d - k}^{1 / 2}w_{|d - k}\rangle \geq \left\| G\| \Sigma_{|d - k}^{1 / 2}w_{|d - k}\|_{2} - u\right\|_{2}}}\sum_{i = 1}^{n}f(\langle w_{|k},x_{i|k}\rangle +h(x_{i|k})u_{i},y_{i},x_{i|k}). \\ \end{array}
$$
Next, we analyze the infimum term:
$$
\begin{array}{l} \inf_{u\in \mathbb{R}^{n}}\sum_{i = 1}^{n}f(\langle w_{|k},x_{i|k}\rangle +h(x_{i|k})u_{i},y_{i},x_{i|k}) \\ = \inf_{\substack{u\in \mathbb{R}^{n}\\ \| u\|_{2}\leq \langle H,\Sigma^{1 / 2}_{|d - k}w_{|d - k}\rangle}}\sum_{i = 1}^{n}f(\langle w_{|k},x_{i|k}\rangle +h(x_{i|k})\left(u_{i} + \| \Sigma^{1 / 2}_{|d - k}w_{|d - k}\|_{2}G_{i}\right),y_{i},x_{i|k}) \\ = \inf  _ {u \in \mathbb {R} ^ {n}} \sup  _ {\lambda \geq 0} \lambda (\| u \| ^ {2} - \langle H, \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \rangle^ {2}) \\ + \sum_ {i = 1} ^ {n} f \left(\left\langle w _ {| k}, x _ {i | k} \right\rangle + h \left(x _ {i | k}\right) \left(u _ {i} + \left\| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \right\| _ {2} G _ {i}\right), y _ {i}, x _ {i | k}\right) \\ \geq \sup_{\lambda \geq 0}\inf_{u\in \mathbb{R}^{n}}\lambda (\| u\|^{2} - \langle H,\Sigma_{|d - k}^{1 / 2}w_{|d - k}\rangle^{2}) \\ + \sum_ {i = 1} ^ {n} f \left(\left\langle w _ {| k}, x _ {i | k} \right\rangle + h \left(x _ {i | k}\right) \left(u _ {i} + \left\| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \right\| _ {2} G _ {i}\right), y _ {i}, x _ {i | k}\right) \\ = \sup_{\lambda \geq 0} - \lambda \langle H,\Sigma_{|d - k}^{1 / 2}w_{|d - k}\rangle^{2} \\ + \sum_ {i = 1} ^ {n} \inf  _ {u _ {i} \in \mathbb {R}} f (\langle w _ {| k}, x _ {i | k} \rangle + u _ {i} + \| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \| _ {2} h (x _ {i | k}) G _ {i}, y _ {i}, x _ {i | k}) + \frac {\lambda}{h (x _ {i | k}) ^ {2}} u _ {i} ^ {2}. \\ \end{array}
$$
Now suppose that  $f$  takes the form  $f(\hat{y}, y, x_{|k}) = \frac{1}{h(x_{|k})^2} \tilde{f}(\hat{y}, y)$  for some 1 square-root Lipschitz  $\tilde{f}$  and by a union bound, it holds with probability at least  $1 - \delta$  that
$$
\langle \Sigma_ {| d - k} ^ {1 / 2} H, w _ {| d - k} \rangle^ {2} \leq C _ {\delta} (w _ {| d - k}) ^ {2}
$$
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w _ {| k}, x _ {| k} \rangle + \| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \| _ {2} h (x _ {| k}) G _ {i}, y _ {i}) \geq (1 - \epsilon_ {\delta}) \mathbb {E} \left[ \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w, x \rangle , y) \right],
$$
then the above becomes
$$
\begin{array}{l} \sup  _ {\lambda \geq 0} - \lambda \langle \Sigma_ {| d - k} ^ {1 / 2} H, w _ {| d - k} \rangle^ {2} + \sum_ {i = 1} ^ {n} \frac {1}{h (x _ {i | k}) ^ {2}} \tilde {f} _ {\lambda} (\langle w _ {| k}, x _ {i | k} \rangle + \| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \| _ {2} h (x _ {i | k}) G _ {i}, y _ {i}) \\ \geq \sup _ {\lambda \geq 0} - \lambda \langle \Sigma_ {| d - k} ^ {1 / 2} H, w _ {| d - k} \rangle^ {2} + \frac {\lambda}{\lambda + 1} \sum_ {i = 1} ^ {n} \frac {1}{h (x _ {i | k}) ^ {2}} \tilde {f} (\langle w _ {| k}, x _ {i | k} \rangle + \| \Sigma_ {| d - k} ^ {1 / 2} w _ {| d - k} \| _ {2} h (x _ {i | k}) G _ {i}, y _ {i}) \\ \geq \sup  _ {\lambda \geq 0} - \lambda C _ {\delta} (w _ {| d - k}) ^ {2} + \frac {\lambda}{\lambda + 1} (1 - \epsilon) n \mathbb {E} \left[ \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w, x \rangle , y) \right] \\ \geq n \left(\sqrt {(1 - \epsilon_ {\delta}) \mathbb {E} \left[ \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w , x \rangle , y) \right]} - \frac {C _ {\delta} (w _ {| d - k})}{\sqrt {n}}\right) _ {+} ^ {2} \\ \end{array}
$$
where we apply Lemma 8 in the last step. Then if we take
$$
F (w) = \left(\sqrt {(1 - \epsilon_ {\delta}) \mathbb {E} \left[ \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w , x \rangle , y) \right]} - \frac {C _ {\delta} (w _ {| d - k})}{\sqrt {n}}\right) _ {+} ^ {2}
$$
then we have  $\Psi \leq 0$ . To summarize, we have shown
$$
\left(\sqrt {(1 - \epsilon_ {\delta}) \mathbb {E} \left[ \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w , x \rangle , y) \right]} - \frac {C _ {\delta} (w _ {| d - k})}{\sqrt {n}}\right) _ {+} ^ {2} - \frac {1}{n} \sum_ {i = 1} ^ {n} \frac {1}{h (x _ {i | k}) ^ {2}} \tilde {f} (\langle w, x _ {i} \rangle , y _ {i}) \leq 0
$$
$$
\mathbb {E} \left[ \frac {1}{h (x _ {| k}) ^ {2}} \tilde {f} (\langle w, x \rangle , y) \right] \leq (1 - \epsilon_ {\delta}) ^ {- 1} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \frac {1}{h (x _ {i | k}) ^ {2}} \tilde {f} (\langle w, x _ {i} \rangle , y _ {i}) + \frac {C _ {\delta} (w _ {| d - k})}{\sqrt {n}}\right) ^ {2}.
$$
Theorem 16. Under assumptions  $(G)$  and  $(H)$ , fix any  $w_{|k}^{*} \in \mathbb{R}^{k}$  and suppose for some  $\rho \in (0,1)$ , it holds with probability at least  $1 - \delta /8$
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \left(\frac {y _ {i} - \langle w _ {| k} ^ {*} , x _ {i | k} \rangle}{h (x _ {i | k})}\right) ^ {2} \leq (1 + \rho) \cdot \mathbb {E} \left[ \left(\frac {y - \langle w _ {| k} ^ {*} , x _ {| k} \rangle}{h (x _ {| k})}\right) ^ {2} \right]. \tag {77}
$$
Then with probability at least  $1 - \delta$ , for some  $\epsilon \lesssim \rho + \log \left(\frac{1}{\delta}\right)\left(\frac{1}{\sqrt{n}} + \frac{1}{\sqrt{R(\Sigma_{|d - k})}} + \frac{n}{R(\Sigma_{|d - k})}\right)$ , it holds that
$$
\min  _ {w \in \mathbb {R} ^ {d}: \forall i, \langle w, x _ {i} \rangle = y _ {i}} \| w \| _ {2} ^ {2} \leq \| w _ {| k} ^ {*} \| _ {2} ^ {2} + (1 + \epsilon) \frac {n \mathbb {E} \left[ \left(\frac {y - \langle w _ {| k} ^ {*} , x _ {| k} \rangle}{h (x _ {| k})}\right) ^ {2} \right]}{\operatorname {T r} (\Sigma_ {| d - k})} \tag {78}
$$
Proof. Fix any  $w_{|k}^{*}\in \mathbb{R}^{k}$ , we observe that
$$
\begin{array}{l} \min_{w\in \mathbb{R}^{d};\forall i,\langle w,x_{i}\rangle = y_{i}}\| w\|_{2}^{2} = \min_{w\in \mathbb{R}^{d};\forall i,\langle w_{|k},x_{i|k}\rangle +\langle w_{|d - k},x_{i|d - k}\rangle = y_{i}}\| w_{|k}\|_{2}^{2} + \| w_{|d - k}\|_{2}^{2} \\ \leq \| w_{|k}^{*}\|_{2}^{2} + \min_{w_{|d - k}\in \mathbb{R}^{d - k}:}\| w_{|d - k}\|_{2}^{2}. \\ \forall i, \left\langle w _ {| d - k}, x _ {i | d - k} \right\rangle = y _ {i} - \left\langle w _ {| k}, x _ {i | k} \right\rangle \\ \end{array}
$$
Therefore, it is enough analyze