SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Corollary 4. Assume  $\sigma_{*}$  is a degree- $p$  polynomial with  $p \geq 1$  defined in (2.2). Under the same assumptions of Theorem 3, if  $\beta > 1 - \frac{1}{p}$ , then the KRR estimator with  $\lambda = \mathcal{O}_d(1)$  learns  $f_{*}$  in the proportional regime, i.e.,  $\mathcal{R}(\hat{f}_{\mathrm{ker}}) = o_d(1)$  with probability 1 when  $n, d \to \infty$  proportionally.
We make the following remarks on Theorem 3 and Corollary 4.
- At a high level, the above theorem aligns with the conclusion of [GMMM20] which assumes spherical data with a diverging number of spikes. Recall that the effective dimensionality of our anisotropic Gaussian data is  $d_{\mathrm{eff}} \asymp d^{1 - \beta}$ , and Theorem 3 implies that to extract the signal direction  $\mu$  and learn a degree- $p$  teacher model, a sample size of  $n = \omega(d_{\mathrm{eff}}^p)$  is required; note that for  $\beta > 0$ , this implies an improvement over the isotropic setting where the sample complexity is  $n = \omega(d^p)$ .
- While we only state the result for inner-product KRR (similar to [LRZ20, BMR21]), we expect similar findings for general rotationally invariant kernels:  $k(\pmb{x}_i, \pmb{x}_j) = g(\langle \pmb{x}_i, \pmb{x}_j \rangle, \| \pmb{x}_i \|, \| \pmb{x}_j \|)$ , since the norm term has negligible contribution to the learning of a single-index target (2.2). In fact, [DWY21] established a lower bound for rotationally invariant KRR, which, when applied to our setting, gives a necessary sample complexity of  $n \asymp d_{\mathrm{eff}}^{\Omega(p)}$ . Although this lower bound may not be sharp, it illustrates that the required spike magnitude needs to scale with target degree  $p$ .
3.2 Intuition behind the Analysis
We briefly summarize the two main ingredients for the proof of Theorem 3. As a preliminary step, we can rotate the signal  $\pmb{\mu}$  to the  $e_1$  direction due to the rotational invariance of Gaussian distribution. After such transformation, we have  $([\pmb{x}_i]_1, [\pmb{x}_i]_{2:d}) \stackrel{\mathrm{i.i.d.}}{\sim} \mathcal{N}(0, 1 + \theta) \oplus \mathcal{N}(0, I_{d-1})$  with  $\theta = \Theta(d^\beta)$ .
Polynomial approximation of the kernel matrix. Firstly, we approximate the inner-product KRR with some degree- $L$  polynomial kernel in the form of
$$
k \left(\boldsymbol {x} _ {i}, \boldsymbol {x} _ {j}\right) = \sum_ {k = 1} ^ {L} c _ {k} d ^ {- k} \left\langle \boldsymbol {x} _ {i}, \boldsymbol {x} _ {j} \right\rangle^ {k}.
$$
Such approximation can be used to establish a learning lower bound for KRR. For instance in the proportional limit, [EK10] showed that for well-conditioned data, the inner-product kernel can be approximated by the linear kernel, which implies that KRR only learns linear  $\sigma_{*}$  [BMR21]. For general deterministic datasets, similar approximation has been studied in [WZ23] where the "orthogonality" of input data determines the polynomial approximation error. In our setting, the inner product  $\frac{1}{d}\pmb{x}_i^\top\pmb{x}_j = \mathcal{O}_d(d^{\beta - 1})$  depends on the spike magnitude  $\theta \asymp d^{\beta}$ . Consequently, we approximate our inner-product kernel by some polynomial kernel whose degree depends on  $\beta$ .
Orthogonal polynomial expansion. While the approximation in the preceding step simplifies the kernel, such analysis is generally not sharp. To refine the result, we expand the polynomial kernel in the Hermite bases (analogous to the analysis of spherical data in [GMMM21, MMM21]). After rotation, the teacher model (2.2) can be written as  $f_{*}(\pmb{x}) = \sigma_{*}(z)$ , for  $z \sim \mathcal{N}(0,1)$  and  $[\pmb{x}]_1 = \sqrt{(1 + \theta)} z$  independent with other entries  $[\pmb{x}]_i$  for  $2 \leq i \leq d$ . In light of the kernel trick [Ver18, Section 3.7], we write  $k(\pmb{x}_i, \pmb{x}_j) = \langle \Phi(\pmb{x}_i), \Phi(\pmb{x}_j) \rangle$  where the feature vector  $\Phi(\pmb{x})$  is composed of Hermite polynomials of  $\pmb{x}$  up to some degree. Importantly, we only need to extract the Hermite components with respect to the first entry  $[\pmb{x}]_1$  to learn a function  $\sigma_{*}(z)$ .
4 Two-layer Neural Network
During the feature learning phase of NN training, we intuitively expect the parameters to align with the (rank-1) signal direction which allows NN to overcome the "curse of dimensionality". Indeed, in this section, we show that the first gradient step on the first-layer weights  $W$  enables the model to "zoom in" to the low-dimensional structure, given that the spike magnitude  $\theta$  is sufficiently large.
4.1 Upper Bound on the Prediction Risk
We consider the following standard Gaussian initialization for the two-layer NN (2.5):
$$
\left[ \boldsymbol {W} _ {0} \right] _ {i j} \stackrel {\text {i . i . d .}} {\sim} \mathcal {N} (0, 1 / d), \quad \left[ \boldsymbol {a} _ {0} \right] _ {i} \stackrel {\text {i . i . d .}} {\sim} \mathcal {N} (0, 1 / N), \quad \left[ \boldsymbol {b} _ {0} \right] _ {i} \stackrel {\text {i . i . d .}} {\sim} \mathcal {N} (0, 1). \tag {4.1}
$$
Denote the NN optimized by the two-stage procedure outlined in Section 2.2 as  $\hat{f}_{\mathrm{NN}}$ . In the sequel, we restrict ourselves to  $\sigma(z) = \mathrm{ReLU}(z) = \max \{z, 0\}$ . We expect similar characterization to hold when the student activation  $\sigma$  has non-zero Hermite coefficients up to degree- $p$  (see [ABAM23]).
The following theorem gives a sufficient condition on the spike magnitude  $\theta \asymp d^{\beta}$  in order for  $\hat{f}_{\mathrm{NN}}$  to learn the target function  $f_{*}$  with linear sample complexity (in terms of dimension dependence).
Theorem 5. Consider the NN training procedure in Section 2.2 with Gaussian initialization defined in (4.1),  $N = \Omega(d^{\varepsilon})$ ,  $\eta = \Omega(N^{1/2 + \varepsilon})$ , and appropriately chosen  $\ell_2$  regularization  $d^{\varepsilon - 1} \lesssim \lambda \lesssim d^{-\varepsilon}$  for small  $\varepsilon > 0$ . Then, for  $\sigma = R e L U$ , we have
$$
\mathcal {R} \left(\hat {f} _ {\mathrm {N N}}\right) = o _ {d} (1) \text {w i t h p r o b a b i l i t y} 1, \text {w h e n} \beta > 1 - 1 / k,
$$
as  $n, d \to \infty, n / d \to \psi$ . That is, the two-layer ReLU NN can learn the target function with the information exponent  $k$  in the proportional regime.
The above theorem predicts that a spike magnitude of  $\theta = \omega(d^{1 - \frac{1}{k}})$  is sufficient for a GD-trained two-layer NN to learn  $f_{*}$  in the  $n \asymp d$  regime. Similar to prior works [AGJ21, BBSS22], the information exponent  $k$  also determines the complexity of the learning problem in our setting: larger  $k$  implies a more "difficult" task for GD, hence we need a larger  $\theta$  (i.e., stronger low-dimensional signal) to achieve linear sample complexity. Moreover, in light of the effective dimensionality defined in [GMMM20], Theorem 5 implies a sample size of  $n = \omega(d_{\mathrm{eff}}^k)$  suffices to learn a single-index target with information exponent  $k$ ; this contrasts the  $n = \omega(d_{\mathrm{eff}}^p)$  complexity of KRR given in Theorem 3. Such discrepancy will be highlighted and empirically validated in Section 5.
4.2 Intuition behind the Analysis
Theorem 5 is established in two steps: first, we prove that under certain conditions, gradient descent can align the first-layer parameters  $\mathbf{W}$  with the signal direction  $\pmb{\mu}$ ; then we show that after achieving such an alignment,  $f_{*}$  can be efficiently learned by optimizing the second-layer  $\pmb{a}$ .
Spiked covariance amplifies gradient signal. Due to our "mean-field" initialization (i.e., small second-layer coefficients), the initial NN has small output  $f_{\mathrm{NN}}^{(0)}(\pmb{x}) \approx 0$ . Therefore the gradient of the squared loss is dominated by the correlation between the student and teacher model. Concretely, consider the population gradient for one parameter vector  $\pmb{w} \in \mathbb{R}^d$  at initialization (for simplicity we omit the bias unit here):
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {x}} \left[ \nabla_ {\boldsymbol {w}} f _ {\mathrm {N N}} ^ {(0)} (\boldsymbol {x}) \right] \approx - \mathbb {E} _ {\boldsymbol {x}} [ \boldsymbol {x} \sigma^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {w} \rangle) f _ {*} (\boldsymbol {x}) ] \\ \stackrel {(i)} {=} - \boldsymbol {\Sigma} \mathbb {E} _ {\boldsymbol {x}} \left[ f _ {*} ^ {\prime} (\boldsymbol {x}) \sigma^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {w} \rangle) \cdot (1 + \theta) ^ {- 1 / 2} \boldsymbol {\mu} + f _ {*} (\boldsymbol {x}) \sigma^ {\prime \prime} (\langle \boldsymbol {x}, \boldsymbol {w} \rangle) \cdot \boldsymbol {w} \right] \\ \stackrel {(i i)} {=} - \sqrt {1 + \theta} \boldsymbol {\mu} \cdot \mathbb {E} _ {\boldsymbol {x}} \left[ \sigma_ {*} ^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {\mu} \rangle) \sigma^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {w} \rangle) \right] + \text {r e s i d u a l t e r m s}, \\ \end{array}
$$
where  $(i)$  is due to multivariate Stein's lemma, and  $(ii)$  follows from the definition of  $\pmb{\Sigma}$ . Utilizing the Hermite expansions of the student and teacher nonlinearities,
$$
\sigma (z) = \sum_ {i = 0} ^ {\infty} \alpha_ {i} h _ {i} (z), \quad \sigma_ {*} (z) = \sum_ {i = 0} ^ {\infty} \alpha_ {i} ^ {*} h _ {i} (z),
$$
we have the following decomposition on the strength of the correlation term,
$$
\mathbb {E} _ {\boldsymbol {x}} \left[ \sigma_ {*} ^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {\mu} \rangle) \sigma^ {\prime} (\langle \boldsymbol {x}, \boldsymbol {w} \rangle) \right] \approx \sum_ {i = 0} ^ {\infty} (i + 1) ^ {2} \alpha_ {i + 1} \alpha_ {i + 1} ^ {*} \cdot \left\langle \sqrt {1 + \theta} \boldsymbol {\mu}, \boldsymbol {w} \right\rangle^ {i}.
$$
This calculation illustrates that the (expected) first gradient step indeed contains the direction of the signal  $\mu$ . More importantly, the magnitude of the signal term is affected by two factors:  $(i)$  it is amplified by a larger spike  $\theta$ , and  $(ii)$  it vanishes when the information exponent  $k$  is large, since  $\langle \pmb{\mu}, \pmb{w} \rangle = \tilde{\Theta}_{\mathbb{P}}(d^{-1/2})$  at initialization. The gradient magnitude relates to the sample complexity because in order to establish learnability in the proportional regime, we need to achieve nontrivial gradient concentration using  $n \asymp d$  samples. Consequently, the sufficient condition in Theorem 5 requires a larger  $\theta$  when the information exponent  $k$  is large, and vice versa.
Remark. One caveat in the above derivation is that the Hermite coefficients of the student nonlinearity  $\{\alpha_i\}_{i=1}^{\infty}$  need to be non-zero up to degree- $k$ , which is not satisfied by ReLU or most commonly used activation functions. We solve this issue by taking into account the bias units  $b_i$  at initialization which "diversify" the nonlinearity, similar to [BBSS22] (see Lemma 15 for details).
Nonparametric learning with random biases. After the signal direction is identified via gradient descent, we need to learn the unknown link function  $\sigma_{*}$ . We utilize random biases as a resource for univariate approximation - similar argument has appeared in many prior works [DLS22, BBSS22, BEG+22, MHPG+23, ABAM23]. Recall that in the representation learning phase, the bias units  $b_{i}$  are not optimized. Therefore, if the weight vectors  $\boldsymbol{w}$  align with the signal direction  $\boldsymbol{w} \approx \boldsymbol{\mu}$ , then learning the second-layer coefficients  $\boldsymbol{a}$  via ridge regression can be reduced to a univariate kernel regression problem; specifically, given  $\boldsymbol{x}_{i}, \boldsymbol{x}_{j} \in \mathbb{R}^{d}$ , the univariate kernel is given by  $k(z_{i}, z_{j}) = \mathbb{E}_{a,b}[\sigma(a \cdot z_{i} + b)\sigma(a \cdot z_{j} + b)]$ , where  $z_{i} = \langle \boldsymbol{x}_{i}, \boldsymbol{\mu} \rangle$ ,  $z_{j} = \langle \boldsymbol{x}_{j}, \boldsymbol{\mu} \rangle$ . For  $\sigma = \mathrm{ReLU}$ , it is known that such kernel can efficiently learn a polynomial link function  $\sigma_{*}$  (e.g., see [DLS22, BBSS22]).
5 Experiments: Comparing KRR and NN
In this section, we empirically validate the theoretical results presented in previous sections. We construct target functions  $f_{*}$  with a link of varying degree  $p$  and information exponent  $k$  and experimentally compute the prediction risk of KRR and that of a two-layer NN.
$\boxed{k=p}$ : NN & KRR are comparable. We first consider the setting where  $k=p$ , which indicates that the link function is a pure degree- $p$  Hermite polynomial:  $\sigma_{*}(z)=h_{p}(z)$ . Recall that  $\theta \asymp d^{\beta}$ ; Theorem 3 implies that KRR learns the target function in the proportional regime when the spike magnitude satisfies  $\beta>1-\frac{1}{p}$ , which matches the sufficient condition for NN given in Theorem 5.
In the experiments of Figure 2, we compute the prediction risk of KRR and that of two-layer NN optimized via the two-stage procedure outlined in Section 2, where the link function  $\sigma_{*}$  is a pure degree-2 and degree-3 Hermite polynomial. We observe that both KRR and NN can learn  $f_{*}$  with linear sample complexity when the spike magnitude  $\theta$  exceeds the same threshold predicted by our theory; intuitively speaking, this means that for the  $k = p$  setting, NN and KRR utilize the low-dimensional structure with the same efficiency. Also, note that while Theorem 5 only provides an upper bound on the risk, the experiments demonstrate that the predicted scaling of  $\theta$  is sharp.
(a)  $p = 2, k = 2$  (KRR).
(b)  $p = 2, k = 2$  (NN).
(c)  $p = 3, k = 3$  (KRR).
(d)  $p = 3,k = 3$  (NN).
Figure 2: Prediction risk of KRR (a)&(c) and GD-trained two-layer NN (b)&(d) to learn  $f_{*}$ , where the link function  $\sigma_{*}(z) = h_{2}(z)$  in (a)&(b), and  $\sigma_{*}(z) = h_{3}(z)$  in (c)&(d). We set  $\psi = 5$ . For KRR we use the Gaussian RBF kernel and  $\lambda = 10^{-2}$ . For two-layer NN, we set  $\eta = \sqrt{N}$ . The experiments are averaged over 50 runs. Darker color corresponds to a smaller prediction risk, and the red lines are predicted by the sufficient conditions of learnability given by Theorem 3 and 5.
$\boxed{k < p}$ : NN outperforms KRR. Next, we consider a setting where the link function  $\sigma_{*}$  is a mixture of low- and high-degree components:  $\sigma_{*}(z) = \sum_{j\in \mathcal{J}}\alpha_{j}h_{j}(z)$ , with  $\min \mathcal{J} = k$ ,  $\max \mathcal{J} = p$ . In this case, Theorem 3 predicts that KRR requires a spike magnitude of  $\beta >1 - \frac{1}{p}$  to learn  $f_{*}$  with linear sample complexity, whereas Theorem 5 entails that a smaller spike magnitude  $\beta >1 - \frac{1}{k}$  is sufficient for the two-layer NN.
Figure 3 shows that the predictions from Theorem 3 and 5 once again accurately align with the experimental results. Observe that when  $k < p$ , GD-trained NN indeed outperforms KRR, in the sense that it requires a less prominent low-dimensional structure (i.e., smaller spike magnitude  $\theta$ ) to attain linear sample complexity. This illustrates the benefit of NN under structured data due to the presence of representation (feature) learning. Roughly speaking, for target function with  $k < p$ , the sample complexity of representation learning (i.e., finding the signal direction  $\mu$  via gradient-based learning) is low compared to directly estimating  $f_{*}$  on the original input features; therefore GD-trained NNs achieve better performance than KRR.
(b)  $p = 3,k = 1$  (NN).
(a)  $p = 3,k = 1$  (KRR).
(c)  $p = 3, k = 2$  (KRR).
(d)  $p = 3,k = 2$  (NN).
Figure 3: Prediction risk of KRR (a)&(c) and GD-trained two-layer NN (b)&(d) to learn  $f_{*}$ , where the link function  $\sigma_{*}(z) = h_{1}(z) + h_{3}(z)$  in (a)&(b), and  $\sigma_{*}(z) = h_{2}(z) + h_{3}(z)$  in (c)&(d). We use the same hyperparameters as in Figure 2. The red lines are predicted by the sufficient conditions in Theorem 3 and 5.
6 Conclusion and Future Directions
We investigated the performances of KRR and GD-trained NNs in the proportional asymptotic limit. We analyzed the prediction risk of these models in learning a nonlinear single-index model (2.2), where the index feature in the teacher model only relies on the spike direction  $\mu$  hidden in the anisotropic Gaussian input data (2.1). The strength  $\theta$  (signal-to-noise ratio) of this spike determines the level of low dimensionality, which affects the learnability of both NNs and KRR. Our results clearly demonstrate that both NNs and KRR benefit from this additional structure in the covariance; yet, NNs adapt to the low-dimensional patterns more efficiently.
It is worth noting that our theoretical analysis focuses on the "well-specified" scenario where the input spike perfectly aligns with the index features of  $f_{*}$ . In a companion work [MHWSE23], we investigate a more general setting where the spiked covariance data only provides partial information of  $f_{*}$  (i.e. the "misaligned" scenario), and show that feature learning (via gradient flow) also benefits from such structured data. Finally, extending our results to multi-index teacher models [CM20, DLS22, ABAM23] and considering more general training dynamics (e.g., beyond the first gradient step) applied to multiple-layer NNs are interesting directions left for future studies.
The authors thank Enric Boix-Adsera, Joan Bruna, Bruno Loureiro, and Alireza Mousavi-Hosseini for the discussions and feedback on the manuscript. MAE was partially supported by NSERC Grant [2019-06167], CIFAR AI Chairs program, CIFAR AI Catalyst grant. TS was partially supported by JSPS KAKENHI (20H00576) and JST CREST (JPMJCR2015). ZW was partially supported by NSF DMS-2055340 and NSF DMS-2154099.
[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
[ABA22] Emmanuel Abbe and Enric Boix-Adsera. On the non-universality of deep learning: quantifying the cost of symmetry. Advances in Neural Information Processing Systems, 35:17188-17201, 2022.
[ABAM23] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302.11055, 2023.
[AGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22(106):1-51, 2021.
[Bac17] Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-681, 2017.
[Bac23] Francis Bach. Learning Theory from First Principles. MIT Press, 2023.
[BAP05] Jinho Baik, Gérard Ben Arous, and Sandrine Péché. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. The Annals of Probability, 33(5):1643-1697, 2005.
[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems, 2022.
[  $\mathrm{BEG^{+}22}]$  Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21764, 2022.
$\left[\mathrm{BES}^{+}22\right]$  Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. In Advances in Neural Information Processing Systems 35, 2022.
[BGN11] Florent Benaych-Georges and Raj Rao Nadakuditi. The eigenvalues and eigenvectors of finite, low rank perturbations of large random matrices. Advances in Mathematics, 227(1):494–521, 2011.
[BKH16] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
[BMR21] Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: a statistical viewpoint. Acta numerica, 30:87-201, 2021.
[BMZ23] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055, 2023.
[BS06] Jinho Baik and Jack W Silverstein. Eigenvalues of large sample covariance matrices of spiked population models. Journal of multivariate analysis, 97(6):1382-1408, 2006.
2 Preliminaries: Problem Setting and Assumptions 3
2.1 Basic Assumptions 3
2.2 Learning Objective and Training Procedure 4
3 Kernel Ridge Regression 5
3.1 Sharp Analysis of the Prediction Risk 5
3.2 Intuition behind the Analysis 6
4 Two-layer Neural Network 6
4.1 Upper Bound on the Prediction Risk 6
4.2 Intuition behind the Analysis 7
5 Experiments: Comparing KRR and NN 8
6 Conclusion and Future Directions 9
A Analysis of Kernel Ridge Regression 14
A.1 Polynomial Approximation of Kernel Matrix 14
A.2 Hermite Expansion of Polynomial Kernel 15
A.3 Analysis of Prediction Risk 17
B Analysis of Neural Network 19
B.1 Hermite Expansion of Population Gradient 19
B.2 Finite-sample Concentration 22
B.3 Learning  $f_{*}$  with Neural Network Features 24
A Analysis of Kernel Ridge Regression
In this section, we analyze the performance of KRR defined in Section 2. Recall that
$$
\boldsymbol {x} _ {i} \sim \mathcal {N} \left(0, \boldsymbol {I} _ {d} + \theta \boldsymbol {\mu} \boldsymbol {\mu} ^ {\top}\right), \quad \| \boldsymbol {\mu} \| _ {2} = 1, \tag {A.1}
$$
where  $\theta \asymp d^{\beta}$ , and  $\beta \in [0,1)$ . Let  $\pmb{\Sigma} = \pmb{I} + \theta \pmb{\mu} \pmb{\mu}^{\top}$  and  $\pmb{x}_i = \pmb{\Sigma}^{1/2} \pmb{z}_i$ , where  $\pmb{z}_i \sim \mathcal{N}(0, \pmb{I})$ , and  $\pmb{\Sigma}^{1/2} = \pmb{I} + (\sqrt{1 + \theta} - 1) \pmb{\mu} \pmb{\mu}^{\top}$ . Then, the training label  $y_i$  can be equivalently expressed as
$$
y _ {i} = f _ {*} (\boldsymbol {x} _ {i}) = \sigma_ {*} (\langle \boldsymbol {z} _ {i}, \boldsymbol {\mu} \rangle),
$$
for  $i\in [n]$ . Recall  $\pmb{X}^{\top} = [\pmb{x},\dots \pmb{x}_n]\in \mathbb{R}^{d\times n}$ . We define the nonlinear kernel matrix by  $\pmb{K} = g(\frac{1}{d}\pmb {X}\pmb{X}^{\top})\in \mathbb{R}^{n\times n}$ , where smooth function  $g:\mathbb{R}\to \mathbb{R}$  is applied to each entry of the matrix. Additionally, we write
$$
\boldsymbol {K} (\boldsymbol {x}, \boldsymbol {X}) = \left[ g \left(\boldsymbol {x} ^ {\top} \boldsymbol {x} _ {1} / d\right), \dots , g \left(\boldsymbol {x} ^ {\top} \boldsymbol {x} _ {n} / d\right) \right].
$$
We analyze the kernel ridge regression estimator:
$$
\hat {f} _ {\ker} (\boldsymbol {x}) = \boldsymbol {K} (\boldsymbol {x}, \boldsymbol {X}) \left(\boldsymbol {K} + \lambda \boldsymbol {I}\right) ^ {- 1} \boldsymbol {y}.
$$
Specifically, we aim to compute the prediction risk in the proportional regime:
$$
\mathcal {R} (\hat {f} _ {\ker}) = \mathbb {E} _ {\boldsymbol {x}} (\hat {f} _ {\ker} (\boldsymbol {x}) - f _ {*} (\boldsymbol {x})) ^ {2},
$$
where  $\pmb{x} \in \mathbb{R}^d$  is an independent copy of  $\pmb{x}_i$  sampled from (A.1). Since the kernel  $\pmb{K}$  is rotationally invariant, i.e.,
$$
k (\boldsymbol {S} \boldsymbol {x} _ {i}, \boldsymbol {S} \boldsymbol {x} _ {j}) = k (\boldsymbol {x} _ {i}, \boldsymbol {x} _ {j}),
$$
for any orthogonal matrix  $S \in \mathbb{R}^{d \times d}$  and  $i, j \in [n]$ , we may apply an orthogonal matrix  $O \in \mathbb{R}^{d \times d}$  for each data sample such that
$$
\boldsymbol {O} \boldsymbol {x} _ {i} \sim \mathcal {N} (0, \boldsymbol {I} _ {d} + \theta \boldsymbol {e} _ {1} \boldsymbol {e} _ {1} ^ {\top}), \qquad \boldsymbol {O} \boldsymbol {x} \sim \mathcal {N} (0, \boldsymbol {I} _ {d} + \theta \boldsymbol {e} _ {1} \boldsymbol {e} _ {1} ^ {\top}),
$$
where  $e_1$  is the standard coordinate vector. It is clear that in distribution,  $\hat{f}_{\mathrm{ker}}(\pmb{x})$  and  $f_{*}(\pmb{x})$  remain unchanged before and after the transformation. Therefore, without loss of generality, we will consider the case that  $\pmb{\mu} = \pmb{e}_1$  for the proofs in this section, that is, we write