SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left| T _ {1 3} \right| \leq \| \boldsymbol {y} \| \left\| (\boldsymbol {K} + \lambda \boldsymbol {I}) ^ {- 1} \right\| \| \boldsymbol {E} _ {2} \| \\ \leq C n \| f _ {*} \| _ {L ^ {2}} ^ {2} \frac {1}{d _ {\text {e f f}} ^ {\ell + 1}} = o _ {d, \mathbb {P}} (1), \\ \end{array}
$$
based on the effective dimension  $d_{\mathrm{eff}} = d^{1 - \beta}$ . This completes the proof of this lemma.
Combining all the above lemmas together, Theorem 3 can be established by considering (A.3) and the following approximation
$$
\begin{array}{l} \mathcal {R} \left(f _ {\ker}\right) = \| f _ {*} \| _ {L ^ {2}} - 2 T _ {1} + T _ {2} = \| P _ {> \ell} f _ {*} \| _ {L ^ {2}} ^ {2} + o _ {d, \mathbb {P}} (1) \\ = \| P _ {> \ell} f _ {*} \| _ {L ^ {2}} ^ {2} + o _ {d, \mathbb {P}} (1), \\ \end{array}
$$
where  $P_{> \ell}$  is defined by in Section 3.
B Analysis of Neural Network
For the two-layer neural network, we perform the gradient-based optimization procedure outlined in Algorithm 1. This algorithm is almost identical to the two-stage representation learning procedure analyzed in recent works [BES+22, DLS22, BEG+22], with the only difference being the additional normalization step (B.2) to handle the anisotropy of the input data.
B.1 Hermite Expansion of Population Gradient
We first consider one parameter vector at random initialization  $\boldsymbol{w}_i \stackrel{\mathrm{i.i.d.}}{\sim} \mathcal{N}(0, d^{-1} \boldsymbol{I}_d)$ . In the following, we restrict ourselves to parameters in a subset defined as
$$
\mathcal {E} _ {\tau} = \left\{i \in [ N ] \left| \frac {1}{\tau \sqrt {d}} \leq | \langle \boldsymbol {\mu}, \boldsymbol {w} _ {i} \rangle | \leq \frac {\tau}{\sqrt {d}} \right. \right\},
$$
where the positive scalar  $\tau = \mathcal{O}(\mathrm{polylog}(d))$  will be appropriately selected in the sequel. Note that due to the Gaussian initialization, setting  $\tau = \Theta(1)$  ensures that a constant fraction of neurons  $\boldsymbol{w}_i$  falls into the set  $\mathcal{E}_{\tau}$ . The gradient of  $\boldsymbol{w}_i$  with respect to the population squared loss can be written as
$$
\nabla_ {\boldsymbol {w} _ {i}} \Big [ \mathbb {E} _ {\boldsymbol {x}} \big (f _ {*} (\boldsymbol {x}) - f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x}) \big) ^ {2} \Big ] = \mathbb {E} _ {\boldsymbol {x}} \Big [ \frac {1}{\sqrt {N}} a _ {i} \big (f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x}) - f _ {*} (\boldsymbol {x}) \big) \sigma^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}) \boldsymbol {x} \Big ].
$$
Due to our specific choice of parameter initialization (4.1), we know that the output of the student model at initialization  $f_{\mathrm{NN}}^{(0)}$  can be ignored (see Lemma 18 for justification). We, therefore, focus on
Algorithm 1: Gradient-based training for two-layer ReLU neural network.
Hyperparameters: learning rate  $\eta$ ,  $\ell_2$  regularization strength  $\lambda$
Normalized gradient descent on 1st layer: Given i.i.d. training examples  $\{(x_{j},y_{j})\}_{j = 1}^{n}$
$$
\boldsymbol {w} _ {i} ^ {(1)} \leftarrow \boldsymbol {w} _ {i} ^ {(0)} - \eta \sqrt {N} \cdot \nabla_ {\boldsymbol {w} _ {i} ^ {(0)}} \left(\frac {1}{n} \sum_ {j = 1} ^ {n} \left(y _ {j} - f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x} _ {j})\right) ^ {2}\right) \tag {B.1}
$$
$$
\boldsymbol {w} _ {i} ^ {(1)} \leftarrow \frac {1}{\sqrt {\frac {1}{n} \sum_ {j = 1} ^ {n} \left\langle \boldsymbol {x} _ {j} , \boldsymbol {w} _ {i} ^ {(1)} \right\rangle^ {2}}} \cdot \boldsymbol {w} _ {i} ^ {(1)} \tag {B.2}
$$
Ridge regression for 2nd layer: Given i.i.d. training examples  $\{(\tilde{\pmb{x}}_j,\tilde{y}_j)\}_{j = 1}^n$
$$
\hat {\boldsymbol {a}} \leftarrow \operatorname {a r g m i n} _ {\boldsymbol {a}} \left\{\frac {1}{n} \sum_ {j = 1} ^ {n} \left(\tilde {y} _ {j} - \langle \phi_ {j}, \boldsymbol {a} \rangle\right) ^ {2} + \lambda \| \boldsymbol {a} \| ^ {2} \right\},
$$
where  $[\phi_j]_i\coloneqq \frac{1}{\sqrt{N}}\sigma \big(\big{\langle}\tilde{\pmb{x}}_j,\pmb {w}_i^{(1)}\big{\rangle} + \tilde{b}_i\big)$
return prediction function  $f_{\mathrm{NN}}(\pmb {x}) = \frac{1}{\sqrt{N}}\sum_{i = 1}^{N}\hat{a}_{i}\sigma \bigl {(}\bigl {\langle}\pmb {x},\pmb{w}_{i}^{(1)}\bigr{\rangle} + b_{i}\bigr)$
the following term capturing the correlation between the  $i$ -th student neuron and the teacher model:
$$
\begin{array}{l} g \left(\boldsymbol {w} _ {i}\right) := \mathbb {E} _ {\boldsymbol {x}} \left[ f _ {*} (\boldsymbol {x}) \sigma^ {\prime} \left(\left\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \right\rangle + b _ {i}\right) \boldsymbol {x} \right] \tag {B.3} \\ \stackrel {(i)} {=} \Sigma \mathbb {E} _ {\boldsymbol {x}} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}) \cdot \frac {1}{\sqrt {1 + \theta}} \boldsymbol {\mu} + f _ {*} (\boldsymbol {x}) \sigma^ {\prime \prime} (\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}) \cdot \boldsymbol {w} _ {i} \right], \\ = \underbrace {\sqrt {1 + \theta} \boldsymbol {\mu} \cdot \mathbb {E} [ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} (\langle \boldsymbol {x} , \boldsymbol {w} _ {i} \rangle + b _ {i}) ]} _ {(I)} + \underbrace {(\boldsymbol {I} + \theta \boldsymbol {\mu} \boldsymbol {\mu} ^ {\top}) \boldsymbol {w} _ {i} \cdot \mathbb {E} [ f _ {*} (\boldsymbol {x}) \sigma^ {\prime \prime} (\langle \boldsymbol {x} , \boldsymbol {w} _ {i} \rangle + b _ {i}) ]} _ {(I I)}. \\ \end{array}
$$
where  $(i)$  is due to Stein's lemma. Note that  $\left\| \frac{1}{\sqrt{1 + \theta}} \pmb{\Sigma}^{1/2} \pmb{\mu} \right\| = 1$ . Therefore, if  $\|\pmb{\Sigma}^{1/2} \pmb{w}_i\| \approx 1$  (which is entailed by the defined event  $\mathcal{E}_{\tau}$ ), we may decompose the first term in the last equation via the Hermite expansion as follows (e.g., see [BBSS22]):
$$
\begin{array}{l} (I) = \sqrt {1 + \theta} \boldsymbol {\mu} \cdot \mathbb {E} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} \left(\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}\right) \right] \\ \stackrel {(i)} {\approx} \sqrt {1 + \theta} \boldsymbol {\mu} \cdot \sum_ {j = 0} ^ {\infty} (j + 1) ^ {2} \alpha_ {j + 1} ^ {b _ {i}} \alpha_ {j + 1} ^ {*} \cdot \left\langle \sqrt {1 + \theta} \boldsymbol {\mu}, \boldsymbol {w} _ {i} \right\rangle^ {j}, \tag {B.4} \\ \end{array}
$$
where  $\alpha^{b_i},\alpha^*$  are the Hermite coefficients of the (shifted) student and teacher nonlinearities:
$$
\sigma (z + b _ {i}) = \sum_ {j = 0} ^ {\infty} \alpha_ {j} ^ {b _ {i}} h _ {j} (z), \quad \sigma_ {*} (z) = \sum_ {j = 0} ^ {\infty} \alpha_ {j} ^ {*} h _ {j} (z).
$$
The following lemma controls the error in  $(i)$  due to the norm fluctuation of  $w_{i}$ .
Lemma 14. Define  $\tilde{\pmb{w}}_i = \pmb {w}_i / \| \pmb{\Sigma}^{1 / 2}\pmb {w}_i\|$  for  $i\in \mathcal{E}(\tau)$ , then we have
$$
\Delta := \left| \mathbb {E} _ {\boldsymbol {x}} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} \left(\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}\right) \right] - \mathbb {E} _ {\boldsymbol {x}} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} \left(\langle \boldsymbol {x}, \tilde {\boldsymbol {w}} _ {i} \rangle + b _ {i}\right) \right] \right| \lesssim \frac {1}{\sqrt {d}} + \frac {\tau}{d ^ {1 - \beta}}.
$$
Proof. By Cauchy-Schwarz inequality, we have
$$
\begin{array}{l} \left. \right.\left| \mathbb {E} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} \left(\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}\right)\right] - \mathbb {E} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} \left(\langle \boldsymbol {x}, \tilde {\boldsymbol {w}} _ {i} \rangle + b _ {i}\right)\right]\right| \\ \leq \sqrt {\| f _ {*} ^ {\prime} (\boldsymbol {x}) \| _ {L ^ {2}} ^ {2} \cdot \| \sigma^ {\prime} (\langle \boldsymbol {x} , \boldsymbol {w} _ {i} \rangle + b _ {i}) - \sigma^ {\prime} (\langle \boldsymbol {x} , \tilde {\boldsymbol {w}} _ {i} \rangle + b _ {i}) \| _ {L ^ {2}} ^ {2}}. \\ \end{array}
$$
Note that  $\| f_{*}^{\prime}(\pmb {x})\|_{L^{2}} = \mathbb{E}[\sigma_{*}^{\prime}(\xi)^{2}]^{1 / 2} = \mathcal{O}(1)$ , where  $\xi \sim \mathcal{N}(0,1)$ .
$$
\left\| \sigma^ {\prime} \left(\left\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \right\rangle + b _ {i}\right) - \sigma^ {\prime} \left(\left\langle \boldsymbol {x}, \tilde {\boldsymbol {w}} _ {i} \right\rangle + b _ {i}\right) \right\| _ {L ^ {2}} \leq \left\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {w} _ {i} \right\| - 1 | \leq \frac {1}{\sqrt {d}} + \frac {\tau}{d ^ {1 - \beta}},
$$
Similarly, for the second term, we have
$$
\begin{array}{l} (I I) = \left(\boldsymbol {I} _ {d} + \theta \boldsymbol {\mu} \boldsymbol {\mu} ^ {\top}\right) \boldsymbol {w} _ {i} \cdot \mathbb {E} \left[ f _ {*} (\boldsymbol {x}) \sigma^ {\prime \prime} (\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle) \right] \\ \stackrel {(i)} {\approx} \theta \langle \boldsymbol {\mu}, \boldsymbol {w} _ {i} \rangle \boldsymbol {\mu} \cdot \sum_ {j = 0} ^ {\infty} (j + 1) (j + 2) \alpha_ {j + 2} \alpha_ {i} ^ {*} \cdot \left\langle \sqrt {1 + \theta} \boldsymbol {\mu}, \boldsymbol {w} _ {i} \right\rangle^ {j} \\ + \boldsymbol {w} _ {i} \cdot \sum_ {j = 0} ^ {\infty} (j + 1) (j + 2) \alpha_ {j + 2} \alpha_ {i} ^ {*} \cdot \left\langle \sqrt {1 + \theta} \boldsymbol {\mu}, \boldsymbol {w} _ {i} \right\rangle^ {j}, \tag {B.5} \\ \end{array}
$$
where  $(i)$  can be justified using the same argument from Lemma 14, the details of which we omit.
Random bias units "diversify" the gradient update. Importantly, the shift introduced by the bias term  $b_{i}$  ensures almost all neurons will have a non-vanishing Hermite coefficient at the given degree, as shown by the following lemma.
Lemma 15. Given degree  $p \in \mathbb{N}$ , for all  $b \in \mathbb{R}$  such that  $h_{p-2}(b) \neq 0$  when  $p \geq 2$ , we have that the  $p$ -th Hermite coefficient of  $f(z) = \sigma(z + b)$  is non-zero, i.e.,  $\left| \mathbb{E}_{z \sim \mathcal{N}(0,1)}[\sigma(z + b)h_p(z)] \right| > 0$ .
Proof. First note that when  $p = 1$ , we have
$$
\int_ {\mathbb {R}} \operatorname {R e L U} (x + b) x e ^ {- \frac {x ^ {2}}{2}} d x = \frac {1}{2} \left(\operatorname {e r f} \left(\frac {b}{\sqrt {2}}\right) + 1\right) > 0.
$$
For  $b \geq 2$ , we divide the computation based on the parity of  $p$ . When  $p$  is even, i.e.,  $p = 2s$ , for any  $b \in \mathbb{R}$ , the Hermite coefficients are given by
$$
\int_ {\mathbb {R}} \mathrm {R e L U} (x + b) h _ {2 s} (x) e ^ {- \frac {x ^ {2}}{2}} d x = e ^ {- \frac {b ^ {2}}{2}} h _ {2 s - 2} (b),
$$
for all  $s \in \mathbb{N}$ . As for odd  $p = 2s + 1$ , the odd-order Hermite coefficients are given by
$$
\int_ {\mathbb {R}} \mathrm {R e L U} (x + b) h _ {2 s + 1} (x) e ^ {- \frac {x ^ {2}}{2}} d x = - e ^ {- \frac {b ^ {2}}{2}} h _ {2 s - 1} (b).
$$
Therefore, as long as  $b$  is not the zero of the Hermite polynomial of degree  $p - 2$ , the Hermite coefficients of  $\mathrm{ReLU}(x + b)$  would be nonzero.
Approximation of population gradient. Note that for fixed  $p \in \mathbb{N}$ , the condition  $h_{p - 2}(b) \neq 0$  excludes less than  $p$  possible values of bias  $b_{i}$ . Combining (B.4) and (B.5), we have the following characterization of the correlation term (B.3). Now given  $k \in \mathbb{N}$  as the information exponent of  $\sigma_*$ , we restrict ourselves to  $b_{i}$  such that the corresponding degree-  $k$  Hermite coefficient of the shifted ReLU  $\alpha_{k}^{b_{i}} = \mathbb{E}_{z \sim \mathcal{N}(0,1)}[\sigma(z + b_{i})h_{k}(z)]$  is at least of order  $\frac{1}{\log^C d}$  for some constant  $C > 0$ . In Section B.3.1 we show that such a condition is satisfied by a sufficiently large set of  $b_{i}$ .
Proposition 16. Given fixed  $k \in \mathbb{N}$  and  $\theta \asymp d^{\beta}$  where  $\beta \in [0,1)$ , consider neurons  $\boldsymbol{w}_i$  where the index  $i \in \mathcal{E}_{\tau}$  for  $\tau = \mathcal{O}\left(\log^{C}d\right)$  ensures that  $\left|\alpha_k^{b_i}\right| = \Omega \left(\frac{1}{\log^C d}\right)$ . Then we have
$$
\begin{array}{l} \left\| (I) \right\| ^ {2} = \Theta \left(\alpha_ {k} ^ {b _ {i}} \langle \boldsymbol {\mu}, \boldsymbol {w} _ {i} \rangle^ {2 (k - 1)} \cdot (1 + \theta) ^ {k}\right) + \mathcal {O} (\Delta). \\ \left\| (I I) \right\| ^ {2} = \mathcal {O} \left(\operatorname {p o l y} (\tau) \cdot \left(d ^ {- k - 1} (1 + \theta) ^ {k + 2} + d ^ {- k} (1 + \theta) ^ {k}\right)\right) + \mathcal {O} (\Delta). \\ \end{array}
$$
Where  $\Delta \lesssim \frac{1}{d^{1/2}} + \frac{\tau}{d^{1-\beta}}$  is defined in Lemma 14.
Proof. Since the link function  $\sigma_{*}$  has information exponent  $k$ , by definition, we know that  $\alpha_{i}^{*} = 0$  for all  $i < k$ . Also, since  $\alpha_{k}^{b_{i}} = \Omega \left(\frac{1}{\log^{C}d}\right)$ , under event  $\mathcal{E}_i(\tau)$  with  $\tau = \mathcal{O}\left(\log^{C}d\right)$ , the expected
gradient is dominated by the term involving  $\alpha_{k}^{*}$  in the summation. For  $(I)$  we have
$$
\begin{array}{l} \left\| \sqrt {1 + \theta} \boldsymbol {\mu} \cdot \sum_ {j = 0} ^ {\infty} (j + 1) ^ {2} \alpha_ {j + 1} ^ {b _ {i}} \alpha_ {j + 1} ^ {*} \cdot \left\langle \sqrt {1 + \theta} \boldsymbol {\mu}, \boldsymbol {w} _ {i} \right\rangle^ {j} \right\| \\ = \left(1 + o _ {d} (1)\right) \left\| \sqrt {1 + \theta} \boldsymbol {\mu} \cdot \alpha_ {k} ^ {b _ {i}} \alpha_ {k} ^ {*} \cdot \left\langle \sqrt {1 + \theta} \boldsymbol {\mu}, \boldsymbol {w} _ {i} \right\rangle^ {k - 1} \right\| \\ = \Theta \left(\alpha_ {k} ^ {b _ {i}} \langle \boldsymbol {\mu}, \boldsymbol {w} _ {i} \rangle^ {k - 1} (1 + \theta) ^ {\frac {k}{2}}\right). \\ \end{array}
$$
The computation of  $\| (II)\|$  follows from the exact same procedure. Finally, the proof is complete by including the norm fluctuation error  $\Delta$  in Lemma 14.
Proposition 16 establishes that  $(I)$  is the dominating term in the correlation  $g(\pmb{w}_i)$  for any  $\beta > 1 - \frac{1}{k}$  as shown in the following corollary.
Corollary 17. Under the same assumptions as Proposition 16, if  $\beta > 1 - \frac{1}{k}$ , we have
$$
\| g (\boldsymbol {w} _ {i}) - (I) \| \lesssim d ^ {- \varepsilon} \| g (\boldsymbol {w} _ {i}) \|,
$$
for some small constant  $\varepsilon > 0$ .
B.2 Finite-sample Concentration
Given  $n$  training examples  $\{(x_i, y_i)\}_{i=1}^n$ , denote the matrix of training data  $\mathbf{X} \in \mathbb{R}^{n \times d}$ , and the vector of labels  $\mathbf{y} \in \mathbb{R}^n$ . Similarly, we denote the initial first-layer parameters as  $\mathbf{W}_0 \in \mathbb{R}^{d \times N}$ ,  $\mathbf{b}_0 \in \mathbb{R}^N$ ,  $\mathbf{a}_0 \in \mathbb{R}^n$ . We first show that the initial gradient (B.1) is dominated by the correlation term (B.3).
Lemma 18. When  $n, d, N \to \infty$  proportionally, there exist some constants  $c, C > 0$  such that
$$
\mathbb {P} \left(\left\| \frac {1}{n} \sum_ {j = 0} ^ {n} f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x} _ {j}) \sigma^ {\prime} \left(\langle \boldsymbol {x} _ {j}, \boldsymbol {w} _ {i} \rangle + b _ {i}\right) \boldsymbol {x} _ {j} \right\| \geq C \frac {\sqrt {1 + \theta}}{d}\right) \leq \exp \left(- c \log^ {2} d\right).
$$
Proof. Using the matrix & vector notation,
$$
\frac {1}{n} \sum_ {j = 0} ^ {n} f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x} _ {j}) \sigma^ {\prime} (\langle \boldsymbol {x} _ {j}, \boldsymbol {w} _ {i} \rangle + b _ {i}) \boldsymbol {x} _ {j} = \frac {1}{n} \boldsymbol {X} ^ {\top} (f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {X}) \odot \sigma^ {\prime} (\boldsymbol {X} \boldsymbol {w} _ {i} + b _ {i} \boldsymbol {1} _ {n})).
$$
Since  $\sigma^{\prime}$  is bounded, we know that
$$
\left\| \sum_ {j = 0} ^ {n} f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x} _ {j}) \sigma^ {\prime} (\langle \boldsymbol {x} _ {j}, \boldsymbol {w} _ {i} \rangle + b _ {i}) \boldsymbol {x} _ {j} \right\| \lesssim \| \boldsymbol {X} \| \| f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {X}) \| = \frac {1}{\sqrt {N}} \| \boldsymbol {X} \| \| \sigma (\boldsymbol {X} \boldsymbol {W} _ {0} + \boldsymbol {b} _ {0}) \boldsymbol {a} _ {0} \|.
$$
To control the RHS, we first recall that the operator norm of  $X$  is bounded as follows,
$$
\mathbb {P} \left(\| \boldsymbol {X} \| \geq C \sqrt {(1 + \theta) d}\right) \leq \exp (- c d), \tag {B.6}
$$
for some constants  $c, C > 0$ . Next, we consider the event that the norm of each training example  $\pmb{x}_j$  can be controlled as follows:  $\mathcal{B} := \left\{\left|\| \pmb{x}_j\| / \sqrt{d} - 1\right| \leq 1/2, j \in [n]\right\}$ . Recall that  $\pmb{x}_j = \pmb{\Sigma}^{1/2}\pmb{z}_j$ ; hence  $\| \pmb{x}_j\|^2 \leq \| \pmb{z}_j\|^2 + (1 + \theta)|\langle \pmb{z}_j, \pmb{\mu} \rangle|^2$ , which is clearly dominated by  $\| \pmb{z}_j\|^2$  when  $\beta < 1$ . Consequently, [Ver18, Theorem 3.1.1] implies that
$$
\mathbb {P} (\mathcal {B}) \geq 1 - n \exp (- c d). \tag {B.7}
$$
Under event  $\mathcal{B}$ , we know that  $f_{j} := \langle \pmb{a}_{0}, \sigma(\pmb{W}_{0}^{\top} \pmb{x}_{j} + \pmb{b}_{0}) \rangle = \sum_{k=0}^{N} a_{j} \sigma(\langle \pmb{x}_{j}, \pmb{w}_{k} \rangle + b_{k})$  is a sum of independent sub-exponential random variables, where the sub-exponential norm can be bounded as  $\| f_{j} \|_{\psi_{1}} \leq \| a_{j} \|_{\psi_{2}} \| \sigma(\langle \pmb{x}_{j}, \pmb{w}_{k} \rangle + b_{k}) \|_{\psi_{2}} \lesssim N^{-1/2}$  conditioning on event  $\mathcal{B}$ . Therefore, for each  $j \in [n]$ , we have
$$
\mathbb {P} \left(\left| f _ {j} \right| \geq C \log d\right) \leq \exp \left(- c \log^ {2} d\right). \tag {B.8}
$$
Combining (B.6), (B.7) and taking a union bound over all  $f_{j}$  in (B.8) establishes the lemma.
The following proposition establishes the finite-sample concentration of the correlation term  $g(\pmb{w}_i)$ .
Proposition 19. Define the empirical correlation term
$$
g _ {n} (\boldsymbol {w} _ {i}) := \frac {1}{n} \sum_ {j = 1} ^ {n} \boldsymbol {x} _ {j} f _ {*} (\boldsymbol {x} _ {j}) \sigma^ {\prime} (\langle \boldsymbol {x} _ {j}, \boldsymbol {w} _ {i} \rangle + b _ {i}),
$$
and recall the population counterpart
$$
g (\boldsymbol {w} _ {i}) = \mathbb {E} _ {\boldsymbol {x}} \left[ f _ {*} (\boldsymbol {x}) \sigma^ {\prime} \left(\langle \boldsymbol {x}, \boldsymbol {w} _ {i} \rangle + b _ {i}\right) \boldsymbol {x} \right].
$$
Then for any fixed  $\boldsymbol{w}_i$  and  $b_i$ , and any  $p$ -degree polynomial  $\sigma_*$  with information exponent  $k$ , if  $\beta \in [0,1)$ , we have
$$
\mathbb {P} \left(\left\| g _ {n} (\boldsymbol {w} _ {i}) - g (\boldsymbol {w} _ {i}) \right\| \geq \log^ {4 p} (n) \sqrt {\frac {d}{n}}\right) \leq n \exp (- c \log^ {2} n), \tag {B.9}
$$
for some constant  $c > 0$ . Moreover, when  $i \in \mathcal{E}_{\tau}$  and  $\alpha_{k}^{b_{i}} = \Omega \left( \frac{1}{\log^{C} d} \right)$ , we have
$$
\mathbb {P} \left(\| g _ {n} (\boldsymbol {w} _ {i}) - g (\boldsymbol {w} _ {i}) \| \geq C \| g (\boldsymbol {w} _ {i}) \| \cdot \operatorname {p o l y l o g} (n, d) \sqrt {\frac {d _ {\text {e f f}} ^ {k}}{n}}\right) \leq n \exp \left(- c \log^ {2} n\right), \tag {B.10}
$$
for any  $1 > \beta > 1 - \frac{1}{k}$ , where  $d_{\mathrm{eff}} \coloneqq d^{1 - \beta}$ .
Proof. Define the projection  $P \coloneqq I_d - \mu \mu^\top$ , we have the following decomposition
$$
\boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {z} _ {i} = \boldsymbol {P} \boldsymbol {z} _ {i} + \sqrt {1 + \theta} \boldsymbol {\mu} \boldsymbol {\mu} ^ {\top} \boldsymbol {z} _ {i},
$$
where Gaussian random vectors  $P\mathbf{z}_i$  and  $(\boldsymbol{\mu}^{\top}\mathbf{z}_{i})\boldsymbol{\mu}$  are centered and independent with each other. With this decomposition, we further decompose  $g_{n}(\pmb{w}_{i})$  by
$$
\begin{array}{l} g _ {n} (\pmb {w} _ {i}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \pmb {\Sigma} ^ {1 / 2} \pmb {z} _ {i} \sigma_ {*} (\pmb {\mu} ^ {\top} \pmb {z} _ {i}) \sigma^ {\prime} (\pmb {w} ^ {\top} \pmb {\Sigma} ^ {1 / 2} \pmb {z} _ {i} + b) \\ = \underbrace {\left(\frac {1}{n} \sum_ {i = 1} ^ {n} \boldsymbol {P} \boldsymbol {z} _ {i} \sigma_ {*} (\boldsymbol {\mu} ^ {\top} \boldsymbol {z} _ {i}) \sigma^ {\prime} (\boldsymbol {w} ^ {\top} \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {z} _ {i} + b)\right)} _ {I _ {1}} \\ + \underbrace {\boldsymbol {\mu} \left(\frac {\sqrt {1 + \theta}}{n} \sum_ {i = 1} ^ {n} \boldsymbol {\mu} ^ {\top} \boldsymbol {z} _ {i} \sigma_ {*} (\boldsymbol {\mu} ^ {\top} \boldsymbol {z} _ {i}) \sigma^ {\prime} (\boldsymbol {w} ^ {\top} \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {z} _ {i} + b)\right)} _ {I _ {2}}. \\ \end{array}
$$
Therefore, for any  $t > 0$  and  $q\in (0,1)$ , we have
$$
\begin{array}{l} \mathbb {P} \left(\left\| g _ {n} \left(\boldsymbol {w} _ {i}\right) - g \left(\boldsymbol {w} _ {i}\right) \right\| > t\right) \\ \leq \inf  _ {q \in (0, 1)} \left\{\mathbb {P} \left(\| I _ {1} - \mathbb {E} I _ {1} \| \geq q t\right) + \mathbb {P} \left(\| I _ {2} - \mathbb {E} I _ {2} \| \geq (1 - q) t\right) \right\}. \tag {B.11} \\ \end{array}
$$
In the following, we present the concentration inequalities of  $I_{1}$  and  $I_{2}$  separately.
Now we first consider the concentration of  $I_{2}$ . Notice that for any  $s > 0$ ,
$$
\mathbb {P} (\| \boldsymbol {I} _ {2} - \mathbb {E} \boldsymbol {I} _ {2} \| \geq s) = \mathbb {P} \Big (\sqrt {1 + \theta} \big | \frac {1}{n} \sum_ {i = 1} ^ {n} f (\boldsymbol {z} _ {i}) - \mathbb {E} [ f (\boldsymbol {z} _ {i}) ] \big | \geq s \Big),
$$
where function  $f: \mathbb{R}^d \to \mathbb{R}$  is defined by  $f(z) := \pmb{\mu}^\top z \sigma_*(\pmb{\mu}^\top z) \sigma'(\pmb{w}^\top \pmb{\Sigma}^{1/2} z + b)$ . Based on [Sam23], we next estimate the Orlicz norm of  $f(z_i)$ . Recall that the  $\alpha$ -exponential Orlicz norm of a random variable is defined by
$$
\left\| X \right\| _ {\psi_ {\alpha}} = \inf  _ {t > 0} \left\{t: \mathbb {E} \left[ \exp \left(| X | ^ {\alpha} / t ^ {\alpha}\right) \right] \leq 2 \right\}.
$$
For any  $m \in \mathbb{N}$  and  $\boldsymbol{z} \sim \mathcal{N}(0, I_d)$ , since  $\sigma_*$  is any polynomial with degree at most  $p$ , we have
$$
\begin{array}{l} \left\| f (\boldsymbol {z}) \right\| _ {L ^ {m}} = \mathbb {E} [ | f (\boldsymbol {z}) | ^ {m} ] ^ {1 / m} \leq \mathbb {E} _ {\xi \sim \mathcal {N} (0, 1)} [ | \xi \sigma_ {*} (\xi) | ^ {m} ] ^ {1 / m} \\ \leq m ^ {\frac {p + 1}{2}} \| \xi \sigma_ {*} (\xi) \| _ {L ^ {2}}, \\ \end{array}
$$
where the last inequality is due to Gaussian hypercontractivity [MMM21, Lemma 20]. Thus, random variable  $f(z)$  has a finite Orlicz norm of order  $\alpha \coloneqq \frac{2}{p + 1}$ : there exists some universal constant  $C > 0$  such that  $\| f(z)\|_{\psi_\alpha}\leq C$ . Then by Theorem 1.5 of [GSS21], we can obtain that
$$
\mathbb {P} \left(\left\| \boldsymbol {I} _ {2} - \mathbb {E} \boldsymbol {I} _ {2} \right\| \geq s\right) \leq 2 \exp \left(- c \left(s \sqrt {\frac {n}{1 + \theta}}\right) ^ {\frac {2}{p + 1}}\right), \tag {B.12}
$$
for any  $s > 0$  and some universal constant  $c > 0$ .
Next, we prove the concentration of  $I_{1}$ . Let  $\pmb{u}_{i} := \pmb{P}\pmb{z}_{i}\sigma_{*}(\pmb{\mu}^{\top}\pmb{z}_{i})\sigma^{\prime}(\pmb{w}^{\top}\pmb{\Sigma}^{1 / 2}\pmb{z}_{i} + b)$  for  $i \in [n]$ . Consider a matrix  $\pmb{A} \in \mathbb{R}^{d \times n}$  where the  $i$ -th column of  $\pmb{A}$  is  $\bar{\pmb{u}}_i := \pmb{u}_i - \mathbb{E}\pmb{u}_i$  for  $i \in [n]$ . Then, for any  $s > 0$  we have,
$$
\mathbb {P} \left(\left\| \boldsymbol {I} _ {1} - \mathbb {E} \boldsymbol {I} _ {1} \right\| \geq s\right) \leq \mathbb {P} \left(\frac {1}{\sqrt {n}} \| \boldsymbol {A} \| \geq s\right)
$$
which implies that it is sufficient to analyze the concentration of  $\| \mathbf{A} \|$ . Notice that  $\mathbf{A}$  is composed of independent centered columns  $\bar{\mathbf{u}}_i$  with a population covariance  $S \coloneqq \mathbb{E}[\bar{\mathbf{u}}_i\bar{\mathbf{u}}_i^\top]$ . We can check that