SlowGuess/NeurIPS_2023_OCR huggingface dataset:
for small constant  $\varepsilon > 0$ , where we used the Lipschitz property of  $\sigma$  in  $(i)$ .
Finally, recall that with high probability we can construct  $\mathcal{I}$  such that  $|\mathcal{I}| = |\mathcal{E}_{\tau} \cap \mathcal{B}_{\varepsilon}| \asymp N$ . This implies the following  $\ell_{2}$ -norm control of the chosen coefficients,
$$
\sum_ {i \in \mathcal {I}} \tilde {a} _ {i} ^ {2} \lesssim | \mathcal {I} | \cdot \left(| \mathcal {I} | ^ {- 1} \log d \sqrt {N}\right) ^ {2} = N \log^ {2} d \cdot | \mathcal {I} | ^ {- 1} \lesssim \log^ {C} d, \tag {B.25}
$$
This entails that the RHS of (B.24) can be bounded as  $d^{-\varepsilon}\log^{C}N = o_{d}(1)$ .
Combining (B.20), (B.23), (B.24), and (B.25), we know that there exists some second-layer coefficients  $\tilde{a}$  such that as  $n, d, N \to \infty$ , there exists some large constant  $C > 0$  such that
$$
\mathbb {E} _ {\boldsymbol {x}} \left[ \left(f _ {*} (\boldsymbol {x}) - \tilde {f} (\boldsymbol {x})\right) ^ {2} \right] \lesssim \frac {1}{\log^ {C} d}, \text {w h e r e} \tilde {f} (\boldsymbol {x}) = \frac {1}{\sqrt {N}} \sum_ {i = 1} ^ {N} \tilde {a} _ {i} \sigma \left(\left\langle \boldsymbol {x}, \boldsymbol {w} _ {i} ^ {(1)} \right\rangle + b _ {i}\right), \| \tilde {\boldsymbol {a}} \| \lesssim \log^ {C} d,
$$
with probability  $1 - \exp \left(-c\log^2 d\right)$ . Note here we simply set  $\tilde{a}_i = 0$  for  $i\notin \mathcal{E}_{\tau}\cap \mathcal{B}_{\varepsilon}$
B.3.2 Ridge Regression Estimator
Proposition 21 suggests there exists some second-layer coefficients with small  $\ell_2$  norm can achieve small prediction risk. Based on a standard concentration argument, we can show that the ridge regression estimator with appropriately chosen  $\lambda$  also achieves low risk (e.g., see [Bac23]). Here we briefly sketch an argument following [BES+22, Appendix D.3].
Given feature map  $\pmb{x} \rightarrow \frac{1}{\sqrt{N}} \sigma(\pmb{W}_1^\top \pmb{x} + \pmb{b})$  conditioned on first-layer weights  $\pmb{W}_1, \pmb{b}$ , we denote the associated Hilbert space as  $\mathcal{H}$ . We define the optimal predictor in the RKHS as  $\check{f} := \operatorname*{argmin}_{f \in \mathcal{H}} \mathbb{E}_\pmb{x}(f(\pmb{x}) - f_*(\pmb{x}))^2$ , which takes the form of  $\check{f}(\pmb{x}) = \left\langle \frac{1}{\sqrt{N}} \sigma(\pmb{W}_1^\top \pmb{x} + \pmb{b}), \check{\pmb{a}} \right\rangle$  for some  $\check{\pmb{a}} \in \mathbb{R}^N$ . We have the orthogonal decomposition of  $f_*$  in  $L^2$ :
$$
f _ {*} (\boldsymbol {x}) = \tilde {f} (\boldsymbol {x}) + f _ {\perp} (\boldsymbol {x}).
$$
Given fresh training data  $\{(\tilde{x}_i,\tilde{y}_i)\}_{i = 1}^n$ , it is known that the prediction risk of ridge regression admits the following decomposition:
$$
\begin{array}{l} \mathcal {R} = \mathbb {E} _ {\boldsymbol {x}} \left(\left(f _ {*} (\boldsymbol {x}) - \check {f} (\boldsymbol {x})\right) + \left(\check {f} (\boldsymbol {x}) - \frac {1}{n} \phi_ {x} \left(\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Phi} ^ {\top} \tilde {\boldsymbol {y}}\right)\right) ^ {2} \\ \lesssim \underbrace {\mathbb {E} _ {\boldsymbol {x}} \left(f _ {*} (\boldsymbol {x}) - \check {f} (\boldsymbol {x})\right) ^ {2}} _ {R _ {1}} + \underbrace {\mathbb {E} _ {\boldsymbol {x}} \left(\check {f} (\boldsymbol {x}) - \frac {1}{n} \phi_ {x} \left(\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Phi} ^ {\top} \check {\boldsymbol {f}}\right) ^ {2}} _ {R _ {2}} \\ + \underbrace {\frac {1}{n ^ {2}} \boldsymbol {f} _ {\perp} ^ {\top} \Phi \left(\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Sigma} _ {\Phi} \left(\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Phi} ^ {\top} \boldsymbol {f} _ {\perp}} _ {R _ {3}}, \\ \end{array}
$$
where  $\phi_{\pmb{x}} := \frac{1}{\sqrt{N}} \sigma(\pmb{W}_{1}^{\top} \pmb{x} + \pmb{b})$  for  $\pmb{x} \in \mathbb{R}^{d}$ , which gives  $\Phi^{\top} = [\phi_{\tilde{\pmb{x}}_{1}}^{\top}, \dots, \phi_{\tilde{\pmb{x}}_{i}}^{\top}, \dots, \phi_{\tilde{\pmb{x}}_{n}}^{\top}]$ , where  $\tilde{\pmb{x}}_{i}$  is the  $i$ -th row of  $\tilde{\pmb{X}}$ ; we denoted  $\widehat{\Sigma}_{\Phi} := \frac{1}{n} \Phi^{\top} \Phi$ ,  $\Sigma_{\Phi} := \frac{1}{N} \mathbb{E}_{\pmb{x}}\left[\sigma(\pmb{W}_{1}^{\top} \pmb{x} + \pmb{b}) \sigma(\pmb{W}_{1}^{\top} \pmb{x} + \pmb{b})^{\top}\right]$ . Also, the  $i$ -th entry of vector  $\check{f}$  and  $f_{\perp}$  are given by  $[\check{f}]_i = \check{f}(\tilde{\pmb{x}}_i)$ ,  $[\pmb{f}_{\perp}]_i = f_{\perp}(\tilde{\pmb{x}}_i)$ , respectively.
Concentration of feature covariance. We first establish the concentration of the kernel matrix in the proportional regime. The following lemma provides an estimate of the norm of the features.
Lemma 22. Given the weight matrix after one gradient step  $\mathbf{W}_1$  as specified in Algorithm 1, then we have  $\left\| \boldsymbol{\Sigma}^{1/2} \mathbf{W}_1 \right\|_F \lesssim \sqrt{N}$ , with probability  $1 - \exp(-c \log^2 d)$ , as  $n, d, N \to \infty, n/d \to \psi > 1$ .
Proof. Write  $\boldsymbol{v}_i = \boldsymbol{\Sigma}^{1/2}(\boldsymbol{w}_i^{(0)} - \eta a_i \cdot g_n(\boldsymbol{w}_i^{(0)}))$ , by definition  $\boldsymbol{x} = \boldsymbol{\Sigma}^{1/2}\boldsymbol{z}$  for  $z \sim \mathcal{N}(0, I_d)$ , and
$$
\left\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} _ {i} ^ {(1)} \right\| ^ {- 2} = \frac {\frac {1}{n} \boldsymbol {v} ^ {\top} \boldsymbol {Z} ^ {\top} \boldsymbol {Z} \boldsymbol {v}}{\| \boldsymbol {v} \| ^ {2}} \geq \lambda_ {\min } \left(\frac {1}{n} \boldsymbol {Z} ^ {\top} \boldsymbol {Z}\right),
$$
which implies that  $\left\| \boldsymbol{\Sigma}^{1/2}\boldsymbol{w}_i^{(1)}\right\|^2 \leq \lambda_{\min}^{-1}\left(\frac{1}{n}\boldsymbol{Z}^\top\boldsymbol{Z}\right) \lesssim 1$  almost surely as  $n, d \to \infty$ ,  $n/d > 1$ . Summing over  $\boldsymbol{w}_i^{(1)}$  yields the desired claim.
Observe that the feature vector  $\phi_{\pmb{x}} = \frac{1}{\sqrt{N}}\sigma(\pmb{W}_{1}^{\top}\pmb{x} + \pmb{b}) = \frac{1}{\sqrt{N}}\sigma(\pmb{W}_{1}^{\top}\pmb{\Sigma}^{1/2}\pmb{z} + \pmb{b})$  is a sub-Gaussian vector with parameter  $\left.\frac{1}{\sqrt{N}}\right\| \pmb{\Sigma}^{1/2}\pmb{W}_{1}\bigg\|_{F} \lesssim 1$  due to the Lipschitzity of  $\sigma$ ; hence we may apply the following sub-Gaussian concentration result.
Lemma 23. Under the same assumption as Lemma 22, as  $d \to \infty$ , we have  $\| \phi_{\pmb{x}} \| \leq \log d$  almost surely.
Proof. Notice that  $\| \cdot \|$  and  $\sigma(\cdot)$  are 1-Lipschitz convex functions. Following [WZ21, Section 3], we can utilize the convex concentration inequality (see [Ver18, Chapter 5]) to obtain
$$
\mathbb {P} \left(\left| \left| \phi_ {\boldsymbol {x}} \right| \right| - \mathbb {E} _ {\boldsymbol {x}} \left[ \left| \left| \phi_ {\boldsymbol {x}} \right| \right| \right] \mid > t\right) \leq 2 \exp \left(- c t ^ {2}\right), \tag {B.26}
$$
for any  $t > 0$ , conditioning on the event that  $\left\| \boldsymbol{\Sigma}^{1/2}\mathbf{W}_1\right\| \leq \left\| \boldsymbol{\Sigma}^{1/2}\mathbf{W}_1\right\|_F \lesssim \sqrt{N}$ . From Lemma 22, we know this event that  $\left\| \boldsymbol{\Sigma}^{1/2}\mathbf{W}_1\right\|_F \lesssim \sqrt{N}$  occurs with probability at least  $1 - \exp(-c\log^2 d)$ . Furthermore, under this event, we have
$$
\mathbb {E} _ {\boldsymbol {x}} [ \| \phi_ {\boldsymbol {x}} \| ] ^ {2} \leq \mathbb {E} [ \| \phi_ {\boldsymbol {x}} \| ^ {2} ] = \frac {1}{N} \left\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {W} _ {1} \right\| _ {F} ^ {2} \mathbb {E} [ \sigma (\xi) ^ {2} ] \lesssim 1.
$$
Thus, we can take  $t = \log d$  and combine Lemma 22 and (B.26) to attain the assertion.
Lemma 24 ([Min17], adapted). Under the same assumption as Lemma 22, given  $n\lambda \geq c\log^2 d$  for some constant  $c > 0$ , we have
$$
\mathbb {P} \Big (\big \| (\boldsymbol {\Sigma} _ {\Phi} + \lambda \boldsymbol {I}) ^ {- 1 / 2} (\boldsymbol {\Sigma} _ {\Phi} - \widehat {\boldsymbol {\Sigma}} _ {\Phi}) (\boldsymbol {\Sigma} _ {\Phi} + \lambda \boldsymbol {I}) ^ {- 1 / 2} \big \| \geq t \Big) \leq C n \exp \Big (- \frac {n \lambda t ^ {2}}{2 (1 + t / 3) \log^ {2} d} \Big),
$$
if  $t^2 \geq \log^2 d \cdot (1 + t / 3) / \lambda n$ .
Proof. Let us denote
$$
\boldsymbol {G} _ {i} = \frac {1}{n} \left(\boldsymbol {\Sigma} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1 / 2} \left(\phi_ {\boldsymbol {x} _ {i}} \phi_ {\boldsymbol {x} _ {i}} ^ {\top} - \boldsymbol {\Sigma} _ {\Phi}\right) \left(\boldsymbol {\Sigma} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1 / 2}.
$$
Notice that  $\mathbb{E}[G_i] = 0$ . From Lemma 23, we know that  $\| [G_i\| \leq \log^2 d / \lambda n$  almost surely. Moreover, we may verify that  $\mathbb{E}[G_i^2 ]\preccurlyeq \frac{1}{n}\frac{\log^2d}{\lambda n} I_N$  and
$$
\operatorname {T r} \mathbb {E} [ G _ {i} ^ {2} ] \leq \frac {1}{n} \frac {\log^ {2} d}{\lambda n} \operatorname {T r} (\boldsymbol {\Sigma} _ {\Phi} + \lambda I) ^ {- 1} \boldsymbol {\Sigma} _ {\Phi} \leq \frac {\log^ {2} d}{\lambda n}.
$$
Here we use the fact that all the eigenvalues of  $(\Sigma_{\Phi} + \lambda I)^{-1}\Sigma_{\Phi}$  are between 0 and 1. Thus, we can apply Matrix Bernstein bound (see (3.9) in [Min17]) to conclude this concentration result.
With Lemma 24 we have the following useful corollary.
Corollary 25. Given  $n\lambda \gg \log^4 d$ , for any  $t \in (0,1)$ , for sufficiently large  $d$
$$
\left\| \left(\boldsymbol {\Sigma} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1 / 2} \left(\boldsymbol {\Sigma} _ {\Phi} - \widehat {\boldsymbol {\Sigma}} _ {\Phi}\right) \left(\boldsymbol {\Sigma} _ {\Phi} + \lambda \boldsymbol {I}\right) ^ {- 1 / 2} \right\| = o _ {d} (1),
$$
with probability at least  $1 - n\exp \bigl (-c\log^2 d\bigr)$
Prediction risk of ridge regression. We upper-bound the prediction risk of the ridge regression estimator by controlling the decomposed error terms separately.
Proposition 26. Under the training procedure described in Algorithm 1 with  $N = \Omega(d^{\varepsilon})$  and  $d^{\varepsilon - 1} \lesssim \lambda \lesssim d^{\varepsilon}$  for small  $\varepsilon > 0$ , if  $\beta > 1 - \frac{1}{k}$ , then we have  $\mathcal{R}(\hat{f}_{\mathrm{NN}}) = o_d(1)$  with probability 1 when  $n, d \to \infty, n / d \to \psi$ .
Proof. First note that  $R_{1} = \mathbb{E}_{\boldsymbol{x}}\big(f_{*}(\boldsymbol {x}) - \check{f} (\boldsymbol {x})\big)^{2} = \| f_{\perp}\|_{L^{2}}^{2}$  by definition. For  $R_{2}$ , [BES+22, Appendix D.3] showed that  $R_{2}\lesssim \| f_{*} - \tilde{f}\|_{L^{2}}^{2} + \lambda \| \tilde{f}\|_{\mathcal{H}}^{2}$ . Finally, for  $R_{3}$  due to the concentration in Corollary 25, we have the following evaluation,
$$
\begin{array}{l} R _ {3} = \frac {1}{n ^ {2}} \boldsymbol {f} _ {\perp} ^ {\top} \boldsymbol {\Phi} \left(\widehat {\boldsymbol {\Sigma}} _ {\boldsymbol {\Phi}} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Sigma} _ {\boldsymbol {\Phi}} \left(\widehat {\boldsymbol {\Sigma}} _ {\boldsymbol {\Phi}} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Phi} ^ {\top} \boldsymbol {f} _ {\perp} \\ \leq \frac {1}{n} \left\| \frac {1}{\sqrt {n}} \boldsymbol {f} _ {\perp} ^ {\top} \Phi (\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}) ^ {- 1 / 2} \right\| ^ {2} \cdot \left\| (\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}) ^ {- 1 / 2} (\boldsymbol {\Sigma} _ {\Phi} - \widehat {\boldsymbol {\Sigma}} _ {\Phi}) (\widehat {\boldsymbol {\Sigma}} _ {\Phi} + \lambda \boldsymbol {I}) ^ {- 1 / 2} \right\| \\ + \frac {1}{n} \left\| \frac {1}{n} \boldsymbol {f} _ {\perp} ^ {\top} \boldsymbol {\Phi} \left(\widehat {\boldsymbol {\Sigma}} _ {\boldsymbol {\Phi}} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Phi} ^ {\top} \right\| ^ {2} \leq \frac {1}{n} \| \boldsymbol {f} _ {\perp} \| ^ {2}, \\ \end{array}
$$
Recall that in Proposition 21 we constructed some  $\tilde{f}$  that achieves  $\| f_{*} - \tilde{f}\|_{L^{2}}^{2}\lesssim 1 / \log^{C}d$  and  $\| \tilde{f}\|_{\mathcal{H}}^2\lesssim \log^C d$ . This implies that  $\| f_{\perp}\|_{L^2}^2\lesssim 1 / \log^C d$  and hence  $R_{1} = o_{d,\mathbb{P}}(1)$ . Similarly, for  $\lambda = \mathcal{O}(d^{-\varepsilon})$ , we have  $R_{2} = o_{d,\mathbb{P}}(1)$ . The proof is completed by Markov's inequality on  $R_{3}$ .