Learning in the Presence of Low-dimensional Structure: A Spiked Random Matrix Perspective
Authors: Jimmy Ba
Murat A Erdogdu
Taiji Suzuki
Zhichao Wang
Denny Wu
Venue: NeurIPS 2023 poster (2023)
Keywords: random matrix theory
high-dimensional statistics
neural network
kernel method
feature learning

Abstract: We consider the learning of a single-index target function $f_*: \mathbb{R}^d\to\mathbb{R}$ under spiked covariance data: 
$$f_*(\boldsymbol{x}) = \textstyle\sigma_*(\frac{1}{\sqrt{1+\theta}}\langle\boldsymbol{x},\boldsymbol{\mu}\rangle), ~~ \boldsymbol{x}\overset{\small\mathrm{i.i.d.}}{\sim}\mathcal{N}(0,\boldsymbol{I_d} + \theta\boldsymbol{\mu}\boldsymbol{\mu}^\top), ~~ \theta\asymp d^{\beta} \text{ for } \beta\in[0,1), $$ 
where the link function $\sigma_*:\mathbb{R}\to\mathbb{R}$ is a degree-$p$ polynomial with information exponent $k$ (defined as the lowest degree in the Hermite expansion of $\sigma_*$), and it depends on the projection of input $\boldsymbol{x}$ onto the spike (signal) direction $\boldsymbol{\mu}\in\mathbb{R}^d$. 
In the proportional asymptotic limit where the number of training examples $n$ and the dimensionality $d$ jointly diverge: $n,d\to\infty, n/d\to\psi\in(0,\infty)$, we ask the following question: how large should the spike magnitude $\theta$ (i.e., the strength of the low-dimensional component) be, in order for $(i)$ kernel methods, $(ii)$ neural networks optimized by gradient descent, to learn $f_*$? 

We show that for kernel ridge regression, $\beta\ge 1-\frac{1}{p}$ is both sufficient and necessary. Whereas for two-layer neural networks trained with gradient descent, $\beta>1-\frac{1}{k}$ suffices. Our results demonstrate that both kernel methods and neural networks benefit from low-dimensional structures in the data. Further, since $k\le p$ by definition, neural networks can adapt to such structures more effectively.

Learning in the Presence of Low-dimensional Structure: A Spiked Random Matrix Perspective
Jimmy 1..3. Murat Erdogdu 1.2 Taiji Suzuki4.5 Zhichao Wang Denny Wu?.8
University of Toronto, Vector Institute, "xAI, *University of Tokyo, RIKEEN AIP,
University of California San Diego, "New York University, Flatiron Institute
{jba, erdogdu 0cc tos.duunntottootta toronto. edu, taiji@mist. i. w.w.mmiitt...ttyoo..... zhw0360ucsd, edu, dennywu@nyu.cm
Abstract
We consider the problem of learning a single-index target function f= Rd under the spiked covariance data:
fx(x) - (ViLg (x. (1.. x ~ NN Ia + 0um1), 8x d for 3 E [o, 1),
where the link function a= exponent & (defined as the lowest degree in the Hermite expansion of a.), and it depends on the projection of input æ onto the spike (signal) direction μ E Rd. In the proportional asymptotic limit where the number of training examples n and the dimensionality d jointly diverge: n, d > 30, n/d - tb (0, oo), we ask the following question: how large should the spike magnitude 0 be, in order for (i) kernel methods, (ii) neural networks optimized by gradient descent, to learn f.? We show that for kernel ridge regression, 3 2 1- is both sufficient and nec- essary. Whereas for two-layer neural networks trained with gradient descent, 3 > 1 suffices. Our results demonstrate that both kernel methods and neural networks benefit from low-dimensional structures in the data. Further, since k S p by definition, neural networks adapt such structures more effectively.
R degree-p polynomial with information
Introduction
Learning under spiked covariance. Real-world data is often nigh-dimensional but it also ex- hibits certain low-dimensional structures. Indeed, the number of input features is exceedingly large machine learning, yet most relevant information concentrated subspace [LV07, HTFF09]. Such structures have important consequences on learning perfor- mance. In random matrix theory. high dimensionality reflected by the proportional asymp- totic limit where the numbers of training examples and input features diverge to infinity at the same rate, whereas the low-dimensional structure often described by a spiked random matrix model [Joh01, BAPO5, BS06, BGN11] in which low-rank signal ("spike") is hidden in large- dimensional noise ("bulk"). In this work, we restrict ourselves to the simple setting where the signal rank-1, and consider the following Kata-generating process:
low-dimensional
x (o, Ia + 0p41)) fa(x) = a. (V (I. (1))
(1.1)
where the goal is to estimate the target function (teacher) f., which is single-index model depend- ing on the signal direction also embedded in the spiked input data, and > 0 controls the spike magnitude; we scale x da for 3 E (0, 1), where larger indicates more prominent low-dimensional structure, and hence the learning problem becomes easier. We aim to characterize the efficiency of learning f. using kernel methods and neural networks (NNs) in relation to the strength the low-dimensional signal 6.
Rdd and R R is an unknown link function. The signal µa is
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Prior results: isotropic data. When & = 0, the imput features are isotropic and do not reveal any information about the target f.. In this setting, the performance of kernel methods and two-layer NNs have been extensively studied in the proportional asymptotic limit.
Kernel methods. For isotropic Gaussian or spherical input data, the performance of random features (RF) regression and kernel ridge regression (KRR) has been precisely characterized in the proportional limit [MM22, GLK 20, DL20, BMR21J. However, in this regime, RF models and KRR suffers from the "curse of dimensionality" [Bac17, YS19) and cannot achieve vanishing generalization error unless f. is linear [EK10, 10, GMMM21, HI.22]. More generally, for KRR to learn degree-p polynomial target, a sample complexity of = (dd)) required [GMMM21, LRZ20, DWY21, MMM21J.
Neural networks Gradient descent. While NNs can efficiently approximate a single-index model (with polynomial link o.), the optimization complexity of gradient-base based learning isotropic data is known to be governed by the information exponent, defined as the index of the first non-zero Hermite coefficient of o,, [AGJ21]. For NNs trained with gradient descent (GD) to learn target with information exponent k recent works showed that sample complexity of (dd)) suffices! [AGJ21, DLS22, BBSS22, ABAM23).. Hence, the proportional regime (n = d), learn- ability has only been established for A 1 [AAM22, BES 22, MHPG 23, BMZ23].
These prior results delineate the limitation of kernel methods and GD-trained NNs under isotropic data in the high-dimensional proportional scaling: KRR can only learn linear functions on the input, and NNs learn targets with information exponent 1, which cannot cover many important prob- lems such as phase retrieval [Fie82, CC15, TV23J. Following the intuition that low-dimensional structure in the input data make the learning problem easier, ask the following question.
If the input data contains cow-dimensional structure given by the spiked covariance model (1.1), can kernel methods and neural networks learn a larger class of f. in the proportional regime?
1.1 Our Contributions
We answer this question in the affirmative by showing that the proportional asymptotic limit when the spike magnitude 0 reaches a certain threshold, both KRR and GD-trained two-layer NN can learn the target function fo, where the link function a. is degree-p polynomial with information exponent & (see Definition 1). Our find- ings are summarized below and illustrated in Figure 1.
10a
NN KRR succeed
low-dim. signal 8
NN succeeds, KRR fails
10a
For KRR with the inner-produe kernel, we give a sharp analysis of the prediction risk and show that to learn a link function o, with degree in the proportional regime, a spike magnitude of 0 = (dd'=3 is both necessary and sufficient.
0 = 1
0 x x 1/1k
10a
10a dimensionality d
For two-layer NN with ReLU activation, we upper bound the prediction risk when representation learning is performed via gradient descent step on the first- layer parameters (analogous to [BES+ 22, DLS22]), and show that a spike magnitude of 0 w (d'-+) is sufficient to learn a. with information exponent k..
Figure 1: Spike magnitude 0 x d for KRR and GD-trained NN to learn the target f. with degree and information exponent in proportional regime. When A p. KRR requires a stronger low-dimensional struc- ture (i.c., larger spike) to learn f-.
Our analysis reveals that both KRR and NNs can benefit from the low-dimensional structure (spike) in the input data, and they can learn wider class of target functions in the proportional regime compared to the isotropic setting (i.c., p > 1 for KRR, and k > 1 for NNs). Furthermore, since for a given link function k < p by definition, we know that NNs require smaller spike magnitude 0 to learn the same target compared to KRR (see Figure 1). In other words, a GD-trained NN can utilize such a dimeeseiioo structure more effectively than KRR.
'The appearance of & in the exponent of the dimensionality is also consistent with the correlational statistical query (CSQ) lower bound [DLS22, ABAM23J, which relates to the hardness learning via gradient descent.

---

Related works. Related to our problem setting, [GMMM20] studied the sample complexity of KRR and NNs in a spiked covariance model, and showed that both models indeed benefit from a stronger low-dimensional structure (i.e., larger spike). However, their analysis assumed the number of spikes to be diverging and cannot cover the proportional regime we consider. More importantly, the NN analysis is limited to approximation, and gradient-based optimization guarantee is not given.
In [RGKZ21], the authors presented a non-rigorous analysis on classifying XOR mixtures using an RF model and two-layer NN trained with online SGD and showed that the NN can learn the task under smaller spike magnitude. Such a task is parallel to the information exponent k = 2 setting (see [ABAM23]); our results rigorously demonstrate that similar separation exists for different class of target functions with potentially higher degree and information exponent.
2 Preliminaries: Problem Setting and Assumptions
Notations. ||- denotes the Ey norm for vectors and the & - E2 operator norm for matrices, and I| || is the Frobenius norm. Oa(.) and ou(-) stand for the standard big-O and little-o notations, where the subscript highlights the asymptotic variable; we write Ō(.) when tors are ignored. (.,((.) (resp- 0a,F()) represents big-O (resp. little-o) in probability d
fac-
())),((), w-) are defined analogously. Denote 7 as the standard Gaussian distribution in R. Given
(-))((),)
f: d + R. we denote its norm with respect to the data distribution as [[
2.1 Basic Assumptions
We consider Gaussian input data with a spiked covariance:
I ~ NV Ia + +1ppp) )) 0 x d for some 3 E [0,
(2.1)
where the signal direction μ E I2d satisfies || = 1, and the spike magnitude @ is allowed to grow with the input dimensionality d as specified by the exponent B: as 3 gets larger, the spike magnitude increases and hence the data exhibits stronger low-dimensional structure. Following the terminology in [GMMM20], we may define the effective dimensionality of the input data as deff x dt-3, which captures the intuition that a larger spike renders the features more low-dimensional. As we will see, the sample complexity of KRR and two-layer NN is decided by this effective dimensionality.
Teacher model. We consider student-teaching setting, where the labels y are generated from teacher model (target function) f. Red -> TR. In the spiked covariance model (2.1), is known that input directions with large variations are often good predictors of the labels indeed, this is the main reason that principal component analysis is used in practice [HTFF09]. We therefore consider the following single-index target function where the index features align with the spike direction μr:
f- (x) = =a. Vi+o (ac, 12)
(2.2)
where the link function a. E (7) is centered such that [lo.(2)) - 0 for ~ NN0, 1) (this can be achieved by subtracting the mean from the training labels in [DLS22J). and the normalization factor (1 0)- 1/2 ensures that the L norm of f. remains constant. We remark that given the prior knowledge of low-dimensional structure, we may efficiently learn f- by first performing PCA on the input features; our goal, however, is not construct optimal learning algorithm, but understand the behavior of KRR and two-layer NN without such data preprocessing, similar to [GMMM20].
We also assume G. is a degree-p polynomial with information exponent (IE) defined below.
Definition (Information exponent). Let f denote the normalized Hermite polynomials. The
information exponent of f E (Y²,,, which we denote by k. is the index of the first non-zero Hermite coefficient of f. that is, given the Hermite expansion f == =,,0 a,,y k := min > 0 : a, # 0}.
is clear that by definition, we always have ls S p, and equality is only achieved when or, is a "pure" degree-k Hermite polynomial. Intuitively, IE measures the magnitude of information contained in the gradient, and larger & implies increased gradient descent complexity [AGJ21, BBSS22]. Such definition is also related to the recently introduced leap complexity [ABA22].
Remark. In (2.2) we do not corrupt the training labels y with ii.d. label noise, meaning that, in light of the bias-variance decomposition, we analyze the bias term describing the extent fa is learned by the student model. If label noise is introduced, the additional variance term (describing the "overfitting" " due noise) can be handled by standard concentration argument.
3

---

2.2 Learning Objective and Training Procedure
where H, L.d. NO Ia + 0up') and yi fzz,) we con- sider two student models obtained via the following training procedure in the proportional regime:
Given training examples
v:) {();,
Proportional asymptotic limit: n, d + 00, n/d > &b E (0, 00).
(2.3)
Student model I: Kernel ridge regression. We focus on the inner-product kernel: ka y)
(a.9)) for some E (n) (similar argument can also apply to rotationally invariant kernels as in
[EKI,, DWY21J. Denoting the associated reproducing kernel Hilbert space (RKHS) with H. the kernel ridge regression estimator is given by
( Σ - f(x))) + +11f1(33 = fx)(()) = k(x, x) (K + XI) y. (2.4)
FLoe= = argmin SEH
where À is the ridge parameter, and K E R is the kernel Gram matrix with entries Ki,- ka,, x, Spectral properties of the kernel matrix and the prediction risk of KRR have been exten- sively studied in various high-dimensional settings [FKI0, GMMM21, BMR21, XHM LY22J.
Since the kernel function is fixed before seeing the data, we intuitively expect that KRR cannot effectively adapt to low-dimensional structure in the learning problem.
Student model II: Two-layer neural network. We aim to learn the following width-N network:
fNN(2)) 1 Σ aao w.) + bb) = (a. N/2.N x +b)).
(2.5)
where a R + R is nonlinear activation, and the trainable parameters W E RdxN a, b E R are initialized from standard Gaussian distribution. We require the network width to be N (dd)) for some small > 0. Such a choice entails that it is sufficient to have O(d total training parameters, which is almost proportional to the sample size n x d; contrast, for KRR we need to store the kernel Gram matrix with n² entries, which is less computationally efficient at test time.
It is known that NNs can adapt to the learning problem via representation learning, in which the trained features encode relevant information of the target function. To realize this advantage, we up- date the parameters in (2.5) via the two-stage procedure introduced in [BES 22, DLS22, BEG 22], where we first learn the representation by taking one gradient descent step on the first-layer param- eters W, and then estimate the second-layer parameters a separately (which is a convex problem). This procedure is summarized as follows (see Algorithm 1 in Appendix for more details).
1.
Feature learning for the 1st layer. We optimize the representation of the NN via gradient de- scent, where the training objective is the empirical squared loss L(f) 1 >im (f(z)) )² n Specifically, denote the i-th column of the initialized weight matrix Wn Wo w, E Rd. and the initialized NN as SN. we take one GD step with learning rate 1) as follows,
(0)
w t w (0)
N/N V (o (())))
Due to the anisotropic input data, the pre-activation (x, w. may blow up if the trained param- eter wo"" aligns with the signal direction μ.. To circumvent this issue, we normalize each neuron to have (roughly) unit pre-activation, i.e., for i E [N), we perform the weight normalization:
(1) w.!
(1) プグ (æ,, (1))))) -1/2
At high level, this resembles the often-used normalization layers in deep learning [IS15, BKH16).
2. Ridge regression for the 2nd layer. After obtaining the updated first-layer parameters W, E RdxX we optimize the second-layer a by solving the ridge regression objective this can be done by using GD to solve an Ez-penalized least squares. To circumvent the dependence be- tween the training data X and the learned W 1. we follow [BES+ 22] and estimate the regression coefficients à using new set of training data [ē,, V. ). with the same sample size
Denoting the feature matrix on the fresh training set X, ( as Φ = J/ X W ++b)= E BRXN
the ridge regression estimator can be obtained by: = argmin
[ - $a 2 + A|[a|1)]

---

Learning in the proportional limit. Given a target function fa and learned (student) model f, we evaluate the model performance using the prediction risk:
R(f) = Em - f. (z))) -f f- f-1122.
where the expectation is taken over the (anisotropic) Gaussian data distribution. To benchmark the model performance in the proportional limit (2.3), we introduce the following notion of learnability. Definition 2 (Learnability in the proportional regime). We say algoritlum learns the target func- tion fa in the proportional regime if, for any (small) constant e >> 0, there exists some constant 2b. > 0 such that when n/ d + 2b 2 the the algorithm constructs f using n i.i.d. training examples, and achieves R(f) IIf f-lli= S with probability 1 when n.d d + oo proportionally.
Based on this definition, in order for a student model (e.g.. KRR or two-layer NN) to learn the teacher f. in the proportional regime, it needs to achieve arbitrarily small prediction risk with high probability, and with linear sample complexity (i.e., using n = training samples).
3 Kernel Ridge Regression
For an anner-product kernel, the performance of KRR depends on the behavior of the pairwise inner- product of the training examples vy). Intuitively speaking, since the inner anner-poouuts concentrate in high dimensions, the kernel matrix can be approximated in operator norm via low-degree Taylor expansion, and thus KRR can only learn f. that is low-degree (e.g., see (DWY21, W723). In- troducing low-dimensional structure to the data is one way to counteract such near-orthogonality in high dimensions. Hence we expect that KRR can achieve better performance when the spike magnitude 0 becomes larger- this intuition will be made rigorous in the following subsections.
3.1 Sharp Analysis of the Prediction Risk
In this section, we study the prediction risk of KRR defined in (2.4) with positive semidefinite kernel given by k(z,, xy) x()=dx,))/d) i.e., (0) 0 for all & (see [SS02, Section 13.1)) Recall that the strength of the low-dimensional component is determined by the spike magnitude 0 d in the following, KRR.
characterize how the exponent 3 E [0. 1) affects the performance of
We denote by Pse: 1,2 > 1.2 the orthogonal projector onto the subspace of polynomials of degree at most l for any l E N. Similarly, we denote by Pst = Id Pge the projector to the orthogonal complement where Id is the identity operator. To illustrate this notion, consider = 1, then we have
(ax) PaiSa (xx + P>)f=(2)) a" (b",, .²) (s))
where a", ..' = argmin,,6 E. (f (x) a - (b. x)) Note that any polynomial f. with degree at most satisfies Psef. Denoting the KRR model as feors we have the following sharp analysis of the asymptotic prediction risk.
Theorem 3. Given any fixed E N. Suppose that g (0) > for all S l, and the spike magnitude scales as d for B (1 - 2, 2 Then as n, d + 0o, n/d + ab E (0, oo), the prediction risk of the KRR estimator fuer with X (1) satisfies the following with probability 1, R(fher) - 1| = oa(1).
As corollary of Theorem 3, based on Definition 2, we can obtain a sharp threshold of 3 for learning degree-p polynomial link function in the proportional regime (Definition 2).
Corollary 4. Assume a degree-p polynomial with p 2 defined in (2.2). Under the same assumptions of Theorem 3, if 9 - then the KRR estimator with X Oa(1) learns fa in the
proportional regime, i.e., R(fxer) = (1( with probability I when n., d + co proportionally.
We make the following remarks on Theorem and Corollary 4.
At high level, the above theorem aligns with the conclusion of [GMMM20] which assumes spherical data with diverging number spikes. Recall that the effective dimensionality of our anisotropic Gaussian data is deir x d=-3 and Theorem 3 implies that to extract the signal direction I and learn a degree-p teacher model, sample size of = w(PP required; note that for 3 > 0, this implies an improvement over the isotropic setting where the sample complexity is wdP).
5

---

While we only state the result for inner- amer--roductt KRR (similar to [LRZ200, BMR21J, we expect similar findings for general rotationally invariant kernels: 2( g xz)) | ll. | since the norm term has negligible contribution to the learning of single-index target (2.2). In fact, [DWY2I] established lower bound for rotationally invariant KRR, which, when applied to
our setting, gives a necessary sample complexity of n (0(p) Although this lower bound may not be sharp. illustrates that the required spike magnitude needs 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, can rotate the signal j to the ej direction due to the rotational invariance of Gaussian distribution.
After such transformation, we have ([x;)1, (x,]2.d) ~ (0, @V IF-1) with 0 = O(d
L.d.
Polynomial approximation of the kernel matrix. Firstly, we approximate the inner-product KRR with some degree-L L polynomial kernel in the form of
k(æ;, 22)) = Σk= (2,,)=
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 a. [BMR21]. For general deterministic datasets, similar approximation has been studied in [WZ23] where the "or- thogonality" of input data determines the polynomial approximation error. In our setting, the inner product j æj Oa(d3-1) depends on the spike magnitude 0 x d8 Consequently, we approxi- mate our inner-product kemne by some polynomial kernel whose degree depends on B.
Orthogonal polynomial expansion. While the approximation in the preceding step simplifies the kernel, such analysis is generally not sharp. To refine the result, we expend the polynomial kernel in the Hermite bases (analogous to the analysis of spherical data in [GMMM21, M- After rotation, the teacher model (2.2) can be written as fx (2),, for Z _N 1) and
[a]1 = /(1 + ()) independent with other entries [z), for S S In light of the kernel trick [Ver18 8, Section 3.7]. we write (,)
(.(x)), ((x,)) where the feature vector ((x) is composed of Hermite polynomials of x up to some degree. Importantly, we only need to extract the
Hermite components with respect to the first entry [zz]] to learn function a. (z).
4 Two-layer Neural Network
During the feature learning phase of NN training, we intuitively expect the parameters 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- -iw-dimesioona structure, given that the spike magnitude @ 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):
[ ooku 1.d.d. (0,11/d) [ao] lid. (N 1/N),, [bo], tid. N/ 1).
(4.1)
Denote the NN optimized by the two-stage procedure outlined in Section 2.2 as fNN. In the sequel, we restrict ourselves to (2) ReLU(z) = max(2,0) 0). We expect similar characterization hold when the student activation o has non-zero Hermite coefficients up to degree-p (see [ABAM23]).
The following theorem gives a sufficient condition on the spike magnitude 0 x d in order for î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 = ()dF), 1/ = (2N2/2+*)), and appropriately chosen lg regularization de S X d= for small 0. Then, for o ReLU, we have
R(fNN) = o1(1) with probability I, when B 1 1>1-1/B 1/k:
as n1, + so,n/d + tb. 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 A = w(d' ) is sufficient for a GD-trerrr -D-triined two-layer NN learn f. in the n x d regime. Similar to prior works [AGJ21, BBSS22], the infor- mation exponent k also determines the complexity of the learning problem in our setting: larger l implies more "difficult" task for GD, hence we need a larger 0 (i.e., stronger low-dimensional sig- nal) to achieve linear sample complexity. Moreover, in light of the effective dimensionality defined in [GMMM20], Theorem 5 implies a sample size n = w(d suffices to learn a single-index target with information exponent k; this contrasts the Theorem 3, Such discrepancy will be highlighted and empirically validated Section 5.
w(den) complexity of KRR given in
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 W with the signal direction μe; then we show that after achieving such an alignment, fo can be efficiently learned by optimizing the second-layer a.
Spiked covariance amplifies gradient signal. Due to our "mean-field" initialization (i.e., small second-layee coefficients), the initial NN has small output f(0) (a) O. 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 ur E TRd at initialization (for simplicity we omit the bias unit here):
Eg V NN (w)) LEx (x)]
(
EEs (z) ()) (1 + 0) - 1/2 F f=(x))" u)) w
vI+ 0p En lo (ax u) Jo' ((z, w))] + residual terms,
where (i) is due to multivariate Stein's 's lemma, and (ii) follows from the definition of Σ. Utilizing the Hermite expansions of the student and teacher nonlinearities,
a (2) =
a.(=)= = E o
we have the following decomposition on the strength of the correlation term,
Ex [o ((x, ())oo w))l s
Σ Ci (i++)) 1)9aa1aa w)) ==0
This calculation illustrates that the (expected) first gradient step indeed contains the direction of the signal μ.. More importantly, the magnitude of the signal term is affected by two factors: (i) it is amplified by larger spike 0, and (ii) it vanishes when the information exponent l: is large, since (p,, w) (d 1/2) initialization. The gradient magnitude relates to the sample complexity because in order to establish leamability in the proportional regime, we need to achieve nontrivial gradient concentration using n x d samples. Consequently, the sufficient condition in Theorem requires larger 0 when the information exponent l is large, and vice versa.
Remark. One caveat in the above derivation is that the Hermite coefficients of the student nonlin- earity { need to be non-zero up to degree-k, which not satisfied by ReLU or most commonly- used activation functions. We solve this issue by taking into account the bias units b, 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 learn the unknown link function a_. We utilize random biases resource for univariate approximation similar argument has appeared in many prior works [DLS22, BBSS22, B2G+22, MHPG+ ABAM23]. Recall that in the representation learning phase, the bias units b; are not optimized. Therefore, if the weight vectors w align with the signal direction w 2 μ., then learning the second-layer coefficients a via ridge regression can be reduced to a univariate kernel regression problem; specifically, given T;, E RRd the univariate kernel is given by k(z,, 2z) Ba,,(ar(u z + (a) 2, + b)j, where 2; = (a., 2, = (I. 11). For a = ReI.U, it is known that such kernel can efficiently learn a polynomial link function a. (c.g., see [DLS22, BBSS22)).
7

---

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 and information exponent & and exper- imentally compute the prediction risk of KRR and that of two-layer NN.
k= p NN & KRR comparable. We first consider the setting where k: p. which indicates
that the link function is pure degree-p Hermite polynomial: o.(=) (2) Recall that = d Theorem 3 implies that KRR learns target function the proportional regime when the spike p which matches the sufficient condition for NN given in Theorem 5.
magnitude satisfies B > 1 1.
In the experiments of Figure 2, we compute the prediction risk of KRR and that of two-layer NN optimized via the tto-stage procedure outlined in Section 2, where the link function σ. is a pure 2 and degree-3 Hermite polynomial. We observe that both KRR and NN can learn f. with linear sample complexity when the spike magnitude 0 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 is sharp.
10
1.0
0.8
0.6
spike 0
101
0.6
spike
0.6
0.4
0.2
10a
10a
En7
109
10a dimensionality
Er7
dimensionality
(a) 2, k = (KRR).
(b) p = 2, == (NN).
1.2
2/
2
0.8
spike 0
0.8
10
104
0.6
0.4
0.4
0.2
10a
10a dimensionality d
dimensionality d
(c) 3, k 3 (KRR).
(d) p t: (NN).
Figure 2: Prediction of KRR (a)&(c) and GD-trained two-layer (b)& learn f., where link function (z) h2(2) in (b). and (2) = ha(z) ) in (c)&(d). We set t = 5. For KRR we use the Gaussian RBF kernel and À 10 For two-layer NN, we set 7 = /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.
k <<p: NN outperforms KRR. Next, we consider a setting where the link function is mixture low- and high-degree components: o. Σjer (z), with min J k, max = In this case, Theorem 3 predicts that KRR requires a spike magnitude of B > to learn f- with linear sample complexity, whereas Theorem 5 entails that a smaller spike magnitude 3 > 1 is sufficient for the two-l NN.

---

Figure 3 shows that the predictions from Theorem and 5 once again accurately align with the experimental results. Observe that when < p. GD-trained NN indeed outperforms KRR, in the sense that it requires less prominent low-dimensional structure (i.e., smaller spike magnitude 0) 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 & < p. the sample complexity of representation learning (i.e., finding the signal direction μ via gradient- based learning) low compared to directly estimating f. the original input features; therefore GD-trained NNs achieve better performance than KRR.
1.0
d20
0.7
0.6
0.6
101
0.5
spike
spike
10a
0.4
0.3
0.2
0.1
10a
10a
10a dimensionality
dimensionality
(a) p (KRR).
p 3, k (NN).
10?
20
1.2
= d
1.0
1.0
0.8
0.0
Splke
Spike
0.6
101
0.4
0.4
0.2
0.2
109
10 dimensionality d
dimensionality d
(c)p- 3, k = 2 (KRR).
(d) 3, k = 2 (NN).
Figure 3: Prediction risk of KRR (a)&(c) and GD-trained two-layer NN (b)& learn f.. where the link function (2) (2) hz(2) (b))(b), and (2)2) ha(z) + hs in (c)&(d). We use the same hyperparameters Figure 2. The red lines are predicted by the sufficient conditions Theorem and 5.
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 nonlinear single-index model (2.2), where the index feature in the teacher model only relies on the spike direction hidden in the anisotropic Gaussian input data (2.1). The strength (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 fa. In a companion work [MHWSE23), we investigate 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, D1, 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.

---

Acknowledgement
The authors thank Enric Boix-Adsera Joan Bruna, Bruno Loureiro, and Alireza Mousavi-Hosseiní for the discussions and feedback on the manuseript. MAE was partially supported by NSERC Grant [2019-06167]. CIFAR AI Chairs program, CIFAR AI Catalyst grant. TS partially supported by JSPS KAKENHI (20H00576) and JST CREST (JPMJCR2015). ZW was partially supported by NSF DMS-2055340 and NSF DMS-2154099.
References
[AAM22]
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged- staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. 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 Sys- tems, 35:17188 1188-182111 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. 1055, 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): 2021.
[Bac17]
Francis Bach. Breaking the curse of dimensionality with convex neural networks The Journal of Machine Learning Research, 18(1):629-68 1811229-61.222 2017.
[Bac23]
Francis Bach. Learning Theory 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 Proba- biliry, 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 Pro- cessing Systems, 2022.
[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.
[BES+22] 22]
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- and Raj Rao Nadakuditi. The eigenvalues and eigenvectors of finite, low rank perturbations of large random matrices. Advances in Mathematics, 227(1):494-52) 2011.
[ВKН16]
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607.0000000 06450, 2016
[BMR21]
Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: a statis- tical 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 preprini arXiv:2303.00m 00055 2023.
[BS06]
Jinho Baik and Jack W Silverstein. Eigenvalues of large sample covariance matrices of spiked population models. Journal of multivariate analysis, 1766):132-11400 2006.
10

---

[CC15]
Yuxin Chen and Emmanuel Candes. Solving random quadratic systems of equations nearly as easy as solving linear systems. Advances Neural Information Process- ing Systems, 28, 2015.
[Chil1]
Theodore S Chihara. An introduction orthogonal polynomials. Courier Corpora- tion, 2011.
[CM20)
Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference Learning Theory, pages 1161-1227. PMLR, 2020.
[CS13]
Xiuyuan Cheng and Amit Singer. The spectrum of random inner-product kernel ma- trices. Random Matrices: Theory and Applications, 2(04):1350010, 2013.
[DL20]
Oussama Dhifallah and Yue M Lu. A precise performance analysis of learning with random features. arXiv preprint arXiv:2008. 1904, 2020.
[DLS22]
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, pages 5413-5452. PMLR, 2022.
[DWY21]
Konstantin Donhauser, Mingqi Wu, and Fanny Yang. How rotational invariance of common kernels prevents generalization in high dimensions. In International Con- ference on Machine Learning, pages 2804-2814. PMLR, 2021.
[ЕK10]
Noureddine El Karoui. The spectrum of kernel random matrices. The Annals of
Statistics, 38(1):1-50, 2010.
[Fie82]
James R Fienup. Phase retrieval algorithms: a comparison. 21(15):2758-2773 1982.
Applied optics,
[GLK+20] Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc Mézard, and Lenka Zde- borová. Generalisation error in learning with random features and the hidden man- ifold model. In International Conference Machine Learning, pages 3452-3002 -462.. PMLR, 2020.
[GMMM20] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems, 33:14820-14830, 2020.
[GMMM21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Lin- earized two-layers neural networks in high dimension. The Annals of Statistics, 1925):102-11044 2021.
[GSS21] Friedrich Götze, Holger Sambale, and Arthur Sinulis. Concentration inequalities for polynomials a-subbexponential random variables. Electronic Journal of Probabil- iry, 20(2on)))) 22, 2021.
[HL22]
Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with ran- dom features. IEEE Transactions on Information Theory, 2022.
[HTFF09]
Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman. The elements of statistical learning: data mining, inference, and prediction, volume 2. Springer, 2009.
[IS15]
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pages 448-456. pmlr, 2015.
[JohO1]
Iain M Johnstone. On the distribution of the largest eigenvalue in principal compo- nents analysis. The Annals of statistics, 29(2):295-33) 3-27, 2001.
[LRZ20]
Tengyuan Liang, Alexander Rakhlin, and Xiyu Zhai. On the multiple descent of minimum-norm interpolants and restricted lower isometry of kernels. In Conference on Learning Theory, pages 2683-2711. PMLR, 2020.
[LV07]
John A Lee and Michel Verleysen. Nonlinear dimensionality reduction, volume 1. Springer, 2007.
[IYY22]
Yue and Horng-Tzer Yau. An equivalence principle the spectrum random inner-product kernel matrices. arXiv preprint arXiv:2205.0m6 063000 2022.

---

[MHPG* 23] Alireza Mousavi-Hosseinn Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Mu- rat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. International Conference on Learning Representations, 2023.
[MHWSE23] Alireza Mousavi- Nousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A. Erdogdu. Gradient-based feature learning under structured data. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
[Min17] Stanislav Minsker. On some extensions of Bernstein's inequality for self-adjoint op- erators. Statistics & Probability Letters, 127:111-119, 2017.
[MM22]
Song Mei and Andrea Montanari. The generalization
of random features
gression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4):667-766, 2022.
[MMM21]
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Generalization error of random feature and kernel methods: hypercontractivity and kernel matrix concentra- tion. Applied and Computational Harmonic Analysis, 2021.
[Rah17]
Sharif Rahman. Wiener-hermite polynomial expansion for multivariate gaus- sian probability measures. Journal of Mathematical Analysis and Applications,
454(1):033-334, 2017.
[RGKZ21] Maria Refinetti. Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classify- ing high-dimensional gaussian mixtures: Where kemel methods fail and neural net- works succeed. In International Conference on Machine Learning, pages 9936-8947. PMLR, 2021.
[Sam23]
Holger Sambale. Some notes concentration for a-subexponential random vari- ables. In High Dimensional Probability IX: The Ethereal Volume, pages 167-192. Springer, 2023.
[SS02]
Bernhard Schölkopf and Alexander J Smola. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2002.
[Sze39]
Gabor Szeg. Orthogonal polynomials, volume 23. American Mathematical Soc., 1939.
[TV23]
Yan Shuo Tan and Roman Vershynin. Online stochastic gradient descent with arbi- trary initialization solves non-smooth, non-convex phase retrieval. Journal of Ma- chine Learning Research, 14458):1-47,
[Ver10]
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv:1011.3027. 2010.
[Ver18]
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
[WZ21]
Zhichao Wang and Yizhe Zhu. Deformed semicirele law and concentration of nonlinear random matrices for ultra-wide neural networks. arXiv preprint arXiv:2109.09304 2021.
[W723] Zhichao Wang and Yizhe Zhu. Overparameterizzed random feature regression with nearly orthogonal data. In International Conference on Artificial Intelligence and Statistics, pages 8463-8493. PMLR, 2023.
[XHM+22] Lechao Xiao, Hong Hu, Theodor Misiakiewicz, Yue Lu, and Jeffrey Pennington. Precise learning curves and higher-order scalings for dot-product kernel regression. Advances in Neural Information Processing Systems, 35:4558 4570, 2022.
[YS19]
Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[Zhi21]
Nikita Zhivotovskiy. Dimension-Fe freee bounds for sums of independent matrices and simple tensors via the variational principle. arXiv preprint arXiv:2108.088 2021.
12

---

Table of Contents
1
Introduction
1.1
Our Contributions
2
Preliminaries: Problem Setting and Assumptions
2.1
Basic Assumptions
2.2
Learning Objective and Training Procedure
3
Kernel Ridge Regression
3.1
Sharp Analysis of the Prediction Risk
3.2
Intuition behind the Analysis
6
4
Two-layer Neural Network
4.1
Upper Bound on the Prediction Risk
4.2
Intuition behind the Analysis
5
Experiments: Comparing KRR and NN
6 Conclusion and Future Directions
Analysis of Kernel Ridge Regression
14
A.I Polynomial Approximation of Kernel Matrix
14
A.2 Hermite Expansion of Polynomial Kernel
15
A.3 Analysis of Prediction Risk
17
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
13

---

A Analysis of Kernel Ridge Regression
In this section, we analyze the performance of KRR defined in Section 2. Recall that
2; ~ (O Ia + +1pp)) || = 1,
(A.1)
0pp and æ; 1/2 zi. where 2, NO I), and g112 I (vI+0 1)p Then, the training label yı can be equivalently expressed
where 0 x d and 3 E [0, Let Σ
f. (z;) = (z((2,)14)
for i E [n]. Recall X = [........,, .,. E Rdx= We define the nonlinear kernel matrix by K 0 xx E Rux", where smooth function g R - IR is applied to each entry of the matrix. Additionally, we write
K (z. X) [o( 21 ..///)...,,... g( an/d))
We analyze the kernel ridge regression estimator:
fx) = (xx, (K y.
Specifically, we aim compute the prediction risk in the proportional regime:
R(fier) E. (z) f=)))
where a E Rd is an independent copy of æ; sampled from (A.1). Since the kernel K is rotationally invariant, i.e..
k k(Sa; Sa.)) k(at,,)))
for any orthogonal matrix S E Rad and i.j [n], we may apply an orthogonal matrix O E RRdxd for each data sample such that
Oz: NN I. Dele]), Oz ~ DeleT):
where e₁ is the standard coordinate vector. is clear that in distribution, (x) and f.(x) re- main unchanged before and after the transformation. Therefore, without loss of generality, we will consider the case that μ e for the proofs in this section, that is, we write
Σ Ia de
(A.2)
and training label ys = fxx) only depends on the first entry the data æ; for i E [n]]
A.1 Polynomial Approximation of Kernel Matrix
Denote the concatenation of and X as X (ar, X Rdx(++1), and the enlarged nonlinear kernel matrix by K x 1(=+1)x(==I), First present preliminary approximation of K by a proper polynomial kernel.
Lemma 6. Suppose that B < n, d + do proportionally, have
for some p N. Under the assumptions of Theorem 3, as
HH + o Kl o0,r(1),
where
H = =P{x /d),
and
(1k) s(1) ())
op :
(1)
Proof.
For simplicity, we denote x = 20g and x.. = |x: x. T1 where a.! is the first entry of æs and
Z; the remaining entries for 0 < is < n. Notice that 1 x; 1 2 Ezy. for any I i.j $ n. By Bernstein inequality [Ver18, Corollary 2.8.3], we have
<
x æi
> 1/2
+ PP
> 1/2
x cxp ed min (t² (t, (2,t)) + exp exxp(cctll (-ctat--))
14

---

for any > and > 0. Therefore, _ æ5 oupp(d log d) for 3 [1/2. 1) and | x
od.rdd11/2lg d) for B E [0, 1/2). Then have
|oofdig H K)1l = ||offdig H K)
Σ
Σ i 2=2
1
k : 2)
(d
Here, offdig(-) denotes setting the diagonal entries of the matrix to be zero. In addition, following the same proof Theorem of [EK10], we can obtain a similar approximation for the diagonal part dig( H K), the details of which we omit.
A.2 Hermite Expansion of Polynomial Kernel
Equipped with the above polynomial expansion of the kernel matrix, we now introduce further approximation using orthogonal polynomials. Recall that
R(Frer)
||f=l1.= 2y 2yJ(K
(K + A E + (K - + "
where
E = Eall((r) x)] E R
M == (x x | X)] E g"x*
From Lemma we know that for some p E
K H, AR.
k(x, X)= = H2( X) + InIn+1 + Ak's
for A (1) and A k o1,,(1), where on is defined in Lemma 6, Hp = P )), and p (2,, X) = ( x))
is row vector. Then following from Lemmas A.3, A.8 and A.13 in [BMR21], we can conclude that R(frer) R2 - od,((1), where
R, = 11 2y (H, + AI) E, + y (K, + AT) - M,(K, +1)
(A.3)
and
E, = E. x H (x, X) E ER.
M, == E p X x H (z. X)) =R
Now we analyze the asymptotic behavior of Rp by utilizing the Hermitian polynomials to reorganize kernel Hy. p. Since Σ now diagonal matrix in (A.2), from [Rah17], we know the multivariate Hermite polynomials with respect (0, Σ in Rdd are strongly orthogonal, each of which the product of univariate Hermite polynomials on each coordinate. More precisely, for any k E and multi-index I = (i1, 42...,...,44 ia) E [A1d with Hermite polynomial with respect to I by 11, (x)) = II (1,) where æ = (.i,..., ..a) and ha, (-) is the the univariate Hermite polynomials (note that the variance for the first coordinate and the remaining coordinates are different). Hence, any L-degree polynomial of x E Rd has the form
i2 + is k. define the k-th multivariate
f(x)
Σ 11
ch.
(A.4)
K=0
I=
for some coefficients C E TR. Similar expansions using Gegenbauer polynomials for uniform unit sphere distribution have been used in [GMMM21, MMM21, GMMM20J. The decomposition consider is specifically related to the Hermite expansions of a.:
fL(x) Σ
Σ az))(("),
(A.S)
1=
1=
15

---

where (-) is the j-th Hermite polynomial with 2 (0, I). Notice that (2)),, δR, J for æ sampled from the spiked model (2.1).
By the kernel trick (e.g., Section 3.7 of [Ver18]); k-th monomial kernel satisfies
(a,, x)) { 80l
where each entry of o E IRd a k-th degree monomial a. Hence we can rewrite the above inner product based on Hermite expansions of each entry of x; and az Therefore, for kernel H ( R. there exists Hermite polynomial feature map +b Rd - R& where
2p
1
+ K ) k(d,k) = d dd-k- 1
L
k
such that
H (æ sy/d) (<../V/VVDDDD /e (x.
where each row vector (.)) the concatenation orthogonal polynomials h.) for S k < L and D is the diagonal matrix where the diagonal entries the Hermite coefficients corresponding to H, which will be specified in the sequel. We define
$ 11 (21/Va)
+
(22/VV
RRLL
(A.6)
4 (2x//Vd)
which gives Hp from the other entries and the label y; only depends on the first entry, we will separate the Hermite polynomials depending on xf and the terms depending on the remaining entries a, E Rd-1 That is, we further decompose each (xx:/v //d/ = [Dolæs/ ((x./V where Po(a;// Vd) E go and
4 Furthermore, since the variance of the first entry of æ; is different
2p
d dd+k-2
Bold, d k: k- k
Lo where (ddd, kt) is the total number of multivariate Hermite polynomials of æ; with degree k whose components with respect x! is not vanishing. Based this decomposition, also write
k-
( dd)
$ @(12/Va
$ =
(A.7)
( /d)
Lemma 7. Let dar dl-B, Under the assumptions of Theorem we know that
H ( w;/d) (<<r//
where 2 is diagonal matrix with entries
g (0)
((d, m) = d dm
+ oa - k
(A.8)
for m E [k] and [2p]. Furthermore, based on the definition of &o as above, we have
H, D0Di + Hpp
where each entry of H (T. 22)
2 2,), and D is diagonal matrix defined analogously 2 but with 0 in (A.8). Note that H is inner-product polynomial kernel matrix of X pi...... Tpp] (d-1)xn
Proof. This lemma follows from properties of multivariate Hermite polynomials [Chil1]. Recall Var(=;)) + 0 ((d)).
By re-normalizing the variables, we obtain
///d) ΣΣ Σ
S(d, k, mihh, r he (xr))
where I = (ii,..., id). As for Po. the diagonals of Do correspond m above.
16

---

A.3 Analysis of Prediction Risk
The next lemma shows that H p is approximated by linear kernel which does not aid the learning of fa: this entails that we may focus on the term involving Do- Notice that H, is a rotationally invariant kernel on isotropic Gaussian data X whose limiting global law has been analyzed in [CS13]. Lemma 8. Under the assumptions of Theorem 3, as n, d + so proportionally, we have
H + I Kinl = 0d,F
where
Kin c1 x x caln 2 11
and || K in = Odr(1). Here, c1, 02 and C2 are some positive constants that only depend on g(*)()) for k = 0, 1inand2 2.
The above lemma is consequence of Theorem of [EKI0], and hence omit the detailed proof; for the formulae of constants C1, C2 and C32 we refer to Equation (57) of [BMR21]. We remark that in the following concentration analysis, the exact limiting spectral distribution of Kinn does not appear, since we do not consider the exact setting of B =
for l E N.
Lemma 9, Under the assumptions of f. Theorem 3, we have that
y 4 Pa,
for some vector =0,..- 10,...,0, Okk Opp 0, T]T onto the Hermite components with (11, 0, 0).
E R22, where P Loxx2p is projector
Here a represents the Hermite coefficients of the function f.. Lemma 9 directly follows from (A.5) since y only depends on at for E [d]. Notice that here P selects (xx)) for i [d] and k [2p] from all multivariate Hermite polynomials.
Next, we prove the following concentration statements, parallel to existing results in [GMMM21, MMM21] for spherical data. For any l [2p], we decompose t f Lo.>el where o e composed of all Hermite polynomials hy.r with & < l, and Lo,> collects all Hermite polynomials ha.r with & f. Similarly, we define Do,se and Do,z corresponding
$o Raxm and
Po.>E [ax(Lo- respectively. Here we denote by m the number of all Hermite polynomials hu.r with k < l where m >> n.
Lemma 10. Under the assumptions of Theorem we have
Po. Im + (1.)
This lemma shows the concentration of the Gram matrix since E[ Po.se] I The proof of
this lemma can be derived from the concentration result [Zhi21]; see also Lemma 11 in [GMMM21]. The following two lemmas the crucial ingredients of the proof of Theorem which have been established in Lemmas 13 and 14 [GMMM21] for the spherical setting.
Lemma 11. Under the assumptions of Theorem 3 with n, d + so proportionally, we have
<
<
1
E+1 for some N, as
1 o
n H + +\, MM, M
2 P
0.5l
o1.)(1).
Proof. Recall the definition of in (A.6). From Lemma 7, we can obtain that Mp $ DD' due to the orthogonal polynomial decomposition in (A.4). Hence,
M = uu.ce D J.se + 2 91
where - is defined similarly as Po.se which is formed by all orthogonal components of Po.5e- In general, > is composed by Po.>e and & defined in (A.7). Therefore,
n nHH, + AI M H AI,))
H + AI,) = =n(H, + AI, 3 0 0,51
+ nnHH, + >I
D
H *I
17

---

Recall that H. = =D By Lemma 7 we know that with high probability
|n(H AI)
$ D
H + - S Cnd = od(
For the second part, let A Φ > D formula to obtain
+ AI and apply the
n(Hp p + AI 16
$ 0.
Do
0,
0.50
H
+
-2
= - 1 0.
D
+ 0,5e
A 1
$ 0,
0
>{
A P B 1.
where
B : (nD $ T
0.5e4
A
Po.S
Following the proof of Lemma 13 in [GMMM21], we can show that
|| | In = =.,p(1) and
|B I, (1) due assumption that - } < l and Lemma 10.
Lemma 12. Recall the definition of Po above. Under the assumptions of Theorem 3 with 1 1- << 1 E+1 for some E N, when n, Do proportionally, we have
(H + AI Po Do,se Im o1.
The proof of Lemma 12 is analogous to Lemma 11, the details of which we omit.
From (A.3), let us denote by
T = y HH AT)-EE, Ep
T: = y MM,
Now aim to control Ti and T₂ based Lemmas 11 and 12. We define the projection of y onto the Hermite components with degree less than or equals to and components with degree larger than as: y yı + y/2 where y/ Po,<e Pe P/ax and y- = Po.> Here Pse projects onto the Hermite components with I = 0, (i.,0....,0) and degree no greater e, and P> P Pst. Similarly, we decompose Ep as Ep Ep,1 Ep.,, where we denote
pp, =
Ep.2 = To Do.se
Lemma 13. Under the assumptions of Theorem 3, we have
Ti = -PseS+ll + o1,,(1),
Tr = PPels + 0a.(1))
Proof. For simplicity, we first control Tr and then apply some of these results to bound T,. Let Tr Teı 2T22 + T23, where
T2 y4 (H,
T22 = MM, +X) - y/2
T23 = yI (H, + - XT) -
Based on Lemmas 9, 11 and (A.8), we know that
T2 = n² P
T 2 P Psi 0.5e
=PPesfll (1))
For T23. following from (23) in [GMMM21], when dft < 1, by Markov's inequality, we have
T23 a
T TT
TT 0.
Pu Do sea os.(1),,
n²
d++1
dll
-PPssfff + -1p(11 od.
(A.9)
18

---

since polynomials in Po.ge are orthogonal to polynomials in Po,>e. Notice that
2222 S 2 (..((1)
This finishes the proof for T2.
Now let us consider Tj = T11 + T12 + T13, where
T1 = Vi (H + XI) Ei.p.
T12 V2 H AI) E1.pp
Tis = y (H, + AI) E2..p-
From Lemmas 9 and 12, we directly have that
Ti = a T Ps + our(1)
-1PPef=llie od.p(1).
Applying (A.9) from above, we have Ti2
od,o Lastly, we know that
[iil ||l||| K + AI) - |22l
S Cnllfa
od,((1),
del
based on the effective dimension der = d=-3 This completes the proof of this lemma.
Combining all the above lemmas together, Theorem 3 can be established by considering (A.3) and following approximation
R(Srar) = || 27; + T2 = 1P>seselle odx(11)
= ||PseJalizz +ar((1),
where Pse defined by Section
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 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 random initialization we Lidd ing, we restrict ourselves to parameters in subset defined as
A(O0 d= In the follow-
1
E, (iENN
TVd - (,),
Val
where the positive scalar - O(polylog(d) will be appropriately selected in the sequel. Note that due to the Gaussian initialization, setting 7 @(1) ensures that constant fraction of neurons w; falls into the set E,. The gradient of w; with respect to the population squared loss can be written as
1
up Ex (f.()) - /| ( = Ex
a: (1 f. (x))o'((xx w(x, + b,)az
Due to our specific choice of parameter initialization (4.1), we know that the output of the student model initialization { can be ignored (see Lemma 18 for justification). We, therefore, focus
19

---

Algorithm 1: Gradient-based training for two-layer ReLU neural network. Hyperparameters: learning rate 1/, E2 regularization strength À Normalized gradient descent on 1st layer: Given i.i.d. training examples ((ar, Uy }z))
w!!h! w (0) -VVN-
10) ( E (vj
INN
[N(())))))
(B.1)
w (1)
w (1)
(B.2)
- (a), w 2)
Ridge regression for 2nd layer: Given i.i.d. training examples {(ā}, D;);
à t argmin E (ū) a)) + }||a|l)]
where 16 == J 0 ( w (0) + b.).
return prediction function fNN(())= = 1
VN Ei= ā ((x,, (x, (.w w!!)) + b.)
the following term capturing the correlation between the i-th student neuron and the teacher model:
g(w:)) == E. (x)) ((x,, w.) + b:)arl
(B.3)
(4) En (x)) ((c, uxs) + b:) VI + ((x))"(((0; ws) + bi) wi
vI+ 0p ((f((a) ((2, w.) bi)] + (I + Opp4 )wn, (x)) wss + b,]]
(I))
where (i) is due to Stein's lemma. Note that - 122 4 1. Therefore, if E 1/2 Besll R
(which is entailed by the defined event E,), we may decompose the first term in the last equation via the Hermite expansion as follows (e.g., see [BBSS22]):
(T) = v + Op (()((z) + b,)l
Σ0+18 (V1+6u..)) w)
VI++pu 0p
(B.4)
3-0
where a a" are the Hermite coefficients of the (shifted) student and teacher nonlincarities:
a(= b:) Σ hg(s), a.(2)= -a)3(3))3))
j=0
The following lemma controls the error in (i) due to the norm fluctuation of w,.
Lemma 14. Define w = "w (2 for i E E(T), then we have
will
1 A : Eng (x) ((x, w + b,)] (x) ir;) + b:)]l S
+ d
Vd
Proof. By (x) ((2,, un) + b.)]
inequality, we have
wi) + b,)]
VIIf((>)) ||o'(( w.) b.) o' ((x, wi + by) b,)11
Note that [=)] E|o (5) 0(1), where NN 1).
1 T
[o'( + b.) a' d'((x, w. + 6:)1l I 1/2 will 11 <
Vd
d
20

---

which concludes the proof.
Similarly, for the second term, have
(IT) = (Ia + Opp
ww: B2ff(())) ((x))
fi, wi) p
Σ (j + 1)j + 2)a5+20 (VI+ Uye, wi)
1=0
+ w ( 2)03+203 (VI+01, wi)
(B.5)
1=0
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, 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 N. for all R such that hp-2(b) + when p 2 we have that the p-th Hermite coefficient of f(2) of + b) non-zero, i.e., 10 (b)p(()) >
Proof. First note that when = 1, we have
1 / ReLU(x + b)re dx
/ (ent( (/2) + +1) >0 0. 2
For 2 2, divide the computation based on the parity of p. When p is even, i.e,, p 2s, for any b E IR, the Hermite coefficients are given by
ReLU(x + b)), 2 dx = e- h25-2(b),
for all N. As for odd p 2s + 1, the odd-order Hermite coefficients given by
/ RcLU(x b))12+11()))))) o dr - h22- (b).
Therefore, as long as not the zero of the Hermite polynomial of degree - 2, the Hermite coefficients of ReLU(x + b) would nonzero.
Approximation of population gradient. Note that for fixed p N, the condition hp-2(b) * excludes less than possible values of bias b;. Combining (B.4) and (B.5), we have the following characterization of the correlation term (B.3). Now given k N the information exponent of o.. we restrict ourselves to b; such that the corresponding degree-k Hermite coefficient of the shifted b, ReLU Og (== + (z)) is at least of order क Section B.3.1 we show that such condition is satisfied by sufficiently large set of b;.
1 for some constant C > 0. In
Proposition 16. Given fixed N and x d where 3 E [0, 1), consider neurons w; where the index i E for = O (log" d) ensures that akk = Ω log
Then we have
|1(1)11 = @ (( (M,, w. 2(k- (1+ 0)) - +OAA).
11 o (poly(T) (d-*~1 h) (1+ - *11 +)+)) + ((A).
Where A S 1 + पه defined in Lemma 14.
Proof. Since the link function a. has information exponent k; by definition, we know that of = under event E7(7) with т = -0log d). the expected
for all k. Also, since a = Ω log
21

---

gradient is dominated by the term involving of in the summation. For (I) we have
VI+ 0u
E 1) a3z1 (VI+ + 0u,, wi)
[VI+0uuuu
vI vI+0pp 0u ax ah . [v1+0u,,0) k-1
(1 + oa(1))
= =0(a (Me, (1+
The computation of [I]]]]]] follows from the exact same procedure. Finally, the proof is complete by including the norm fluctuation error A in Lemma 14.
Proposition 16 establishes that (1) is the dominating term the correlation g(w;) for any B 1- as shown in the following corollary.
Corollary 17. Under the same assumptions Proposition 16, if 1. we have
([g(w;; (1)ll X (19()()
for some small constant E > 0.
B.2 Finite-sample Concentration
Given n training examples {(a;, y:)} denote the matrix of training data X E Raxd, and the vector of labels E IR" Similarly, we denote the initial first-layer parameters as Wo E gdxN bu E R.N, ag E RR. We first show that the initial gradient (B.1) 1) dominated by the correlation term (B.3).
Lemma 18. When n, d, N + proportionally. there exist some constants C' > 0 such that (( NN (z)))o ( + b.)xzz
2 C cx S exp(- exp(-clog' d).
Proof. Using the matrix & vector notation.
(0) ( )o', ((zr,, w.) + by) b:)00, = x _ (X) (Xw, + b,1,)))
Since o' is bounded, we know that
(0) (n, )o' ((a,,, un)) bi)ar,
||
NN X
X (XW. + bo)aoll-
To control the RHS, we first recall that the operator norm of X is bounded as follows,
P (IXI C cV(1+0)d) exp(-cd),
(B.6)
for some constants C > O. Next, we consider the event that the norm of each training example æ; can be controlled as follows: B = ||
< 1/2, j E [n]]. Recall that æ;
1/2 zjj
hence ||a,;| S ||, + (1 + 0) (z,; n)|², which is clearly dominated by ||=,|| when < 1. Consequently, [Ver18, Theorem 3. implies that
P(B)= 2 1 eexp(-ctt
(B.7)
Under event B, we know that f, == (as. o (W x, + bo)) = Σ ( wx) + by) is sum of independent sub-exponential random variables, where the sub-exponential norm can be bounded as || ||a, I| [0((2, wx) b 1/2 conditioning event B. Therefore, for each j E [n], we have
P(fsl 2 C' Co d) S exp(-clog d).
(B.8)
Combining (B.6), (B.7) and taking a union bound over all fj in (B.8) establishes the lemma. I1
The following proposition establishes the finite-sample concentration of the correlation term g(w.).
22

---

Proposition 19. Define the empirical correlation term
9n (ws)
Σ wjS+ (a))(o ((2))) w.) + b:),
and recall the population counterpart
g(w,; Th [S= (2))0a((() wi) + b.)a).
Then for any fixed w; and by, and any p-degree polynomial o, with information exponent k. if B E 1), we have
n)) [[mn(ww) g(w:))l 2 log 40 (n) S n
(B.9)
for some constant c 0. Moreover, when i & and ax = S2 logtd) we have
d
P [1gn(w.) g(w.))l 2 Cllg(w:))l polylog(n, d) S <nexp_-clog' n),
(B.10)
for any 1> B> E. where doff d!-3.
Proof.
Define the projection P := Ia - pe we have the following decomposition
1/2 2i Pz; + v1 + 0puzz' 2;:
where Gaussian random vectors Pz; and (fx 2; 2:)p are centered and independent with each other. With this decomposition, we further decompose 9n (w:) by
In (w.: = 1/2 21.22,0(11 21) 2i)o'(w" Σ2//2 b) 1
n 1=1
Pzia. (1 2i)o T
T=1/22 +b)
n
0
n
k 2,0, (H z:)o'(w 121/221 + b)
+ μ
7
Therefore, for any > 0 and g (0, 1), we have
P(lgg, (w;) g (w:) || 1)
inf [[P([]] ET qt) + P E12l 2 (1 g)t)} (E(0,1)
(B.11)
In the following, present the concentration inequalities of I, and 1a separately.
Now we first consider the concentration of Iz. Notice that for any 0,
1( - EI2l 2 s)
/I 0 ΣE f(z.) - [[ 2 s).
where function f Id + IR is defined by f(z) = 20+ 2) T + b). Based on [Sam23], we next estimate the Orlicz norm of f(z,). Recall that the a-exponential Orlicz norm of random variable is defined by
|X inf inf Elexp (1X((11) 2).
For any E N and N'0, Ia), since o is any polynomial with degree most p, have
[[f(z)) [l < (l=
< m [50 (5) 12,
23

---

where the last inequality is due to Gaussian hypercontractivity [MMM21, Lemma 20]. Thus, random variable f(z) has a finite Orliez norm of order == 2 there exists some universal constant C' p+I
such that ll) I C. Then by Theorem 1.5 [GSS21], we can obtain that
P+1
P( EI21 2 8) S 2exp
-c (
(B.12)
1+0
for any 8 > O and some universal constant > O.
Next, we prove the concentration of I1. Let u, == P2,0, (n 2,)o'(w" T " 1//2g 2; + b) for i [n]. Consider matrix A Rdxx" where the i-th column of A is un = u; Eu, for i [n]. Then, for any > we have,
P( EL 8) S
A 2
which implies that it is sufficient to analyze the concentration of A Notice that A is composed of independent centered columns ü, with population covariance S E|ū,ü, We check that
|||| sup
Var(w (ooT)))) ux)
sup
Var(w" (o)))
|e||=1
||||==1, wif
sup E[(u v |+||=1, vlu
Ei(v) (()
sup
|e||=1,
vl
for some universal constant C > Meanwhile,
|| ||
< Pz,ll 1o, (n [2,)) 2.) + B||| | (n 22)]
I || (p T 2;) |||| 2,10 [E)
< || (,.(#' z:)1 + C" Vd,
(B.13)
of the Gaussian random vector. Applying Theorem 3.1.1 of [Ver18], know that || c'vd almost surely, for some universal constant é' > 0. Since a. is polynomial with degree p. we know that ((o is uniformly bounded by some constant for a 1/p and 5 ~ -N(00, 1). Hence, (o-(E)) I log tog2((n) almost surely. Together with (B.13), we can conclude that ||ú.]|| S C'log*((n) almost surely as n + co. Therefore, by Theorem 5.44 in [Ver10], we have that with probability at least exp (-cs?),
where the inequality is due to the sub-Gaussian tail of the
d | $ (C + s log (m)) -.
n
Combining (B.12) and choosing 1/ and / log 27 (m) for any E [0, 1). As a corollary of (B.9), applying Proposition 16, we can directly obtain (B. This completes the proof of the proposition.
in (B.11), one can easily obtain (B.9)
B.3 Learning f, with Neural Network Features
Univariate approximation. We first adapt an argument from [DLS22] showing that ReLU non- linearity can approximate a polynomial function in one dimension with the aid of random bias units. Lemma 20 (Corollary 3 in [DLS22], adapted). Let a()) ReLU(=), b ~ (O, 1). Then given wifce-differentiable f + R, there exists some v(s, b) such that for any |x| < C with some constant C 0,
Unif(1-1, 1}) and
Elu(s,b)o + b)] = f(z), and sup,, ( $ exp(C).
2/2 362 12 be(e.2)) any > Then we have
Proof. Let
()o(s) + b)] =
24

---

for all |2| S c. Now let b(s = Jioj 2e) for any > O, where f(c) := -e
1 V=
1/- /2
For
similar reasons, may verify that
[ ()z 6]] =,
10E(0,c| for all |=| v c. Let v(s,)) 2v |" (-sb)e /2 for any O, where ʃ E (RR).. Then we can easily check that
([ 6)0 + f G(c)z,
for all |2| S where
F(c) - - + - 2f(c),
G(c) (00)
Notice that this construction of e(s, 6) approximates the function f(z) using some constant and linear function perturbations. Now recalling the first two cases, we know that we can approximate any constant and linear functions on this interval, and hence we obtain the desired claim by combining all three cases together.
B.3.1 Analysis of the Learned Feature Map
Recall the training procedure in Algorithm 1. The following lemma shows that after the first-layer parameters optimized for one (normalized) gradient step, there exist some second-layer coeffi- cients that achieve small prediction risk with high probability.
Proposition 21. Given fixed & E N and B > 1 - E. let W 1 be the weight matrix optimized via one normalized gradient descent step defined in Algorithm I, there exists some second-layer coefficients ā E R such that f(x) ~ x + b)ā VN Σ āso (ax, w! + bi) achieves
llf - fll 1/ d, and ||ā|| S log d.
for some large constant C 0, with probability exp(_clog" N) when n, d oo proportionally and N ).
Proof. We follow argument similar to [BES+22, Appendix DJ by selecting a suitable set of neurons that can approximate fa- Throughout this proof, we take e >> to be some small non- vanishing constant and C' 0 to be some large constant. At high level, we select the neurons following the univariate approximation result in Lemma 20 to polynomial target function. To this end, we need to show that after one gradient descent step, there exists sufficiently large set of neurons w (1) that align with the signal direction, i.e., v p ~= 1.
We first restrict ourselves to the "average-case" neurons around the equator w, where the indices i E,. Due to the Gaussian initialization, setting r = logg d entails IP & E,) S exp(- log² d). Now we consider the subset of bias units that gives an "informative" gradient update, i.e., the k-th Henmite coefficient is not vanishing: a Ω log l for some large constant C' 0. Recall that Lemma 15 implies that for bias unit satisfying - (b:) +f O, the degree-k: Hermite coefficient of f(=) = oz + by) is non-zero. Therefore, we define the index set
B, = (i E [N] 2bb)) 2 e, if 2), 2}.
(B.14)
Notice that all the zeros of Hermite polynomials Additionally, from Theorem 3.3.2 in [Szc39], the zeros of hp+1(z) are interlacing between the zeros of (x)). From the recurrence relation of Hermite polynomials, namely,
real and distinct, see Theorem 3.3.1 1 in [Sze39].
h (a) xx)( h,,x)):
we know that (x) + Of is a root of hp- Thus, applying the implicit function theorem, we can deduce that by choosing x T for some large C >> O, we attain o bb = Ω Tog for any E B. Due to the Gaussian initialization of b, we know what the expected size of the set E.B.J = 2(N1-=') for some small e' > 0. Therefore, Hoeffding's 's inequality implies that for large enough N = 2(d'), ) we have
B
< N-1/2+=' with probability eepp( -cllollo' N).
25

---

Gradient approximation. For the subset of neurons with indices i E Be, we introduce a sequence of approximations to the gradient update given as
In (w
x 72
(J - [x())) o (
(0) + +b,1))
where we dropped the factor initial NN output follows:
= for convenience. First recall that by Lemma 18, ignore the VN
Im (w (0)) -
11_ x x ( o b,11,)) x
VT+8 d'
(B.15)
with probability 1 - exp(-clog' d) for some constant > O. Note that since 3 < 1, the RHS can always be upper- ooper-bounded by dt for some small O. Next, we take into account the concentration and truncation error in Propositions and 16, respectively. Define the population gradient (again omitting the scalar --))
g(w(") E. (f) w/")) + b.)]
(B.10) in Proposition 19 implics that for any fixed p E N. 3 1 +.2 and large enough n, with probability 1 - exp( caclollog d), we have
x (
+ b,11,) s(wp:0) S s oo
|
(B.16)
for some e 0. As for the population counterpart, Corollary 17 implies that for neurons with E. then we have
indices i E & nB. if 3
>
((ww,)) ) a a v @( 8u,, w
R-1 H 15 [o
(B.17)
where [((p))) x (1 + 0) (10
(0) 14) = =
due to the condition i E &
for r logo d. Combining (B.15), (B.,6), and (B.17), we know that the gradient vector go (w, achieves significant alignment with the target direction μ when B > E-
Normalization step. Now we consider the normalization procedure (B.2). When B > 1 - k (B.15), (B.16), and (B.17) imply that for neurons w, in the subset E & n Be, we have
w (1)
w (0)
X (N x)= - fx(X)) 2 (Xw; (0) + 111
ija
n
4 w + na;
(0)
by aj
v 8 [vI+ [VI+ 0u, w; w0 k-1 (mt (u+
where || A
S d= for some > 0, with probability 1 - cxppp(c d). Now recall that we select
sufficiently large learning rate 17 = ((N1/2+) for some small > 0. Since ~ NN00 1/N),
W kapel S d
(B.18)
by aj /1 0 VI 1V+++pp mo 100)= where K; na,a logo d and a x log a- Define the normalization factor
k-1 (1
0)
by the choice of T
X 1)=((113 (1))))) 1/2
1
xw (1)
Vn
by 18.18) we know that
Ix k;g X X
1
X k
(1)
x
vn
1
x
w
(1)
xit1
5 + 10 (1))))))))
26

---

with high probabilityy. where we used the proportional scaling of n. d in the last inequality, Now note that Jo xull + vI+ almost surely when n. d + oc proportionally, we have
x (1+
Xi - [AVV (mVV++)
(1+)
(xxx
Combining the above calculations and recalling N = 22(d"), we arrive at the following characteri- zation the normalized weight vector:
n2 1ー==(11 (ag 20; (0
1/2
w (1) sign(k;)( 1/21/21
$ Xi w (2) ki (kgn(k:) (K,) 1)- kxx;")
2 <d (1+
(B.19)
with probability 1 exp(-clog' d). Observe that (B.19) entails that after the normalization step:
with + ( E (1)))) 112
(1)
the updated weight vector with index E &, n B, approximates the target direction in the sense that | 0 w.!") si x d with high probability, for some small 0.
Approximation of link function. Now we construct certain second-layer coefficients ā; on sub- set of neurons with indices i E I such that the (sub)network approximates the link function a.. Roughly speaking. we wish find some collection of neurons that utilizes the univariate approxi- mation result in Lemma 20:
f(=):= = EiEz Ai ((ar, ( + b,)) R Bly(s, (s(ar, (1 + (0++0)111 (4) + b)].
First, we consider the ideal setting where w; exactly in the direction of μa to sign flip, that is, we may write w; = s1(1 + 0) -1/2 for 8; ~ (nif((+1,,1)) This allows us to reparameterize (ar, (1 + 0) 1/2 =: IR. and choose the second-layer coefficients according to Lemma 20 on domain |2| log log d. This implies that [u(s, b) O(exp(log log d)) O(log d). The corresponding "infinite-widhh width" neural network is written as
(z) = Un- -,t---ii((((+110)))))-))))- (1).
b) ass= + b)].
From Lemma 20 we know that for |x| I log log d, we have f=) = (x))))
Now we control the difference between the ideal model and the finite-wide network of interest. First we restrict ourselves to b E B¹, defined as "continuous" analogue of (B. 14): B = {b (bp-((b) > E, if 2 2}. with log d for large C > and write the corresponding infinite-widhe model as
(z) E. (u,,(((s, ()z + +) 1ses:]
By the Lipschitz property of a, we have (2) fz)( & sup ju(s, b) log oa()) Next, bound the emror due truncation the interval |a| I loglog d. Recall that (x, (1 + 0/21/2 u) =: N 1); hence we have the following evaluation
E Ex-N1o.1 [0) 1x(()) S E. -ox + Eg (5))))))))=)))
EE (z)) |=22og
(B.20)
logod d
To upper- bound the above expectation, observe that
[=(o.(2) | al S /
1
(2)2)² exp(-22/2) dz
(B.21)
|=/>2og logd
log
for large > O, where we used the fact that o, a degree-p polynomial with bounded 1 norm in the last inequality. On the other hand, by Cauchy-Schwand we have
[x)( $
E E[ b) Blo(sz sz b) & suplu(s, b)l
EEo(ss + b)?
27

---

Therefore, by the Lipschitz property of a and the fact that sup/u(s, b)) = log d,
1
(2) |=/2 for al - [Bf |1/2/o log .gdd d
(B.22)
Combining (B.20), (B.21), and (B.22) yields E- - (á()))) 2 log d (a(1) for
some large C² > 0. Now consider the finite-width model
f(=) Eiez a, / (s:2 + LieB.
where 8, LLd. Unif((+1,-1)) (+), and by 1.d. NN 1). Similar to above, we set the second-layer according to Lemma 20, that is, 0; VN u(s,
To establish concentration of around f we first restrict [z| S log Let cls b) - o(sz + b). Then || Vo(s, b)ll I log d. Hence, applying Lipschitz concentration inequality for and b, we have (lo(s, b) ++ log d. Since uniformly [[v(s, bl I log d, we know that the sub- exponential norm b))(s, k)l,,, S 2og" d. Thus, utilizing the Bernstein's inequality [Ver18. Theorem 2.8.1], we know that
Eier u(s,, (,)0 + b,)li.B - fāc > - log lo प.j
2
(- dzi
min(f?
log'
for any O. By construction, we can construct a subset neurons E with [I| ((V), with probability at least exp(-clog' d). Since we choose N 2(d'). we may take log
for large C, and union bound yields f()) f (z) log with probability 1 exp sep(-clog log d) as d, N - oc, under the event that |a| s log d. Then by the tail bound of the standard Gaussian random variable z. we have
E.~Nion1 (f(x) f (2))
E (f(=)) f (z)) |s| + Ez (5()))))
S log d' (B.23)
with probability 1 - 1-exp(-cclog d) when d, N + co.
To control the deviation from the ideal setting, we consider neurons w. (1) where i E & n Bg for log" d and x 2 Recall that for this subset of neurons, with at least probability 1
exp(-clog" d) we have VI+0 w sign(r=)) E d= where sign(N.) ~ Unif( [[+1 -1}) due to the Gaussian initialization of a, (note that wi, a, b, are independent). Therefore,
Eg
Σ
āro ((x, (h))) + bb))
1/211)1 aso (sign(s:)); (1 0) 01 + b:)
VN IEI
NET
(3)
||a||10
Ex EiEz x w sign(k;) (1+ 1/2 <
(B.24)
N
Nd
for small constant E > O. where we used the Lipschitz property of a in (i).
Finally, recall that with high probability implies the following norm control of the chosen coefficients,
construct I such that [[]
18 n B&l = N. This
ΣiEr a S [Z] ((Z) log dvN)² = N log² d [ S log c d,
(B.25)
This entails that the RHS of (B.24) be bounded d log N oa(1).
Combining (B.20), (B.23). (B.24), and (B.25), we know that there exists some second-layer coeffi- cients such that as n, d, N -> oc, there exists some some large constant C > such that
Fg (z) [(x)))
log c d' where f(x)
VN Σ àgo (æx, ( + b.), ||ā||| S log d,
with probability exp(-clog* d). Note here here we simply set à; for $ & 0 B..
28

---

B.3.2 Ridge Regression Estimator
Proposition 21 suggests there exists some second-layer coefficients with small l2 norm can achieve small prediction risk. Based a standard concentration argument, we can show that the ridge regression estimator with appropriately chosen À also achieves low risk (e.g., see [Bac23]). Here we briefly sketch argument following [BES+ 22, Appendix D.3].
Given feature map x 1 o ( x + b) conditioned on first-layer weights W1.b, b, we de- VN
note the associated Hilbert space as H. We define the optimal predictor in the RKHS as := argminseH (f(x)) (x))). which takes the form of f(x) = ( or (W b), a) for some
RN,. We have the orthogonal decomposition of fx in L
f- = f(x) + / (æ)).
Given fresh training data ((;; 0c))= it is known that the prediction risk of ridge regression admits the following decomposition:
4 (T8) ))
R Ex (f.) f(x)) + -
Es + =1)
71
$ Ex (f((()) - f(x))) Ex [(x) Φ
Ei + 11)
1
R1
-1 EEg + X1) Ex f + E 21)
T TfL1 1
n
Rs
where &. == VN 0 (W (W
b) for E Rd. which gives +
[ ..... $ T
..... P2. where
ã; is the i-th row of X; we denoted En == L 4., Eg == 1 (W 1 + (W x + T.
Also, the i-th entry of vector f and f_ are given by [f]: = f(x.), [fil (ā,). respectively.
Concentration of feature covariance. We first establish the concentration of the kernel matrix in the proportional regime. The following lemma provides estimate of the norm of the features.
Lemma 22. Given the weight matrix after one gradient step W1 1 as specified in Algorithm 1, then have Σ W, S VN, with probability 1 - eexp(-clog? d), as d, N + n/d + sb > 1.
Proof.
Write n; /1/2 (w (0) na;; ga(w! (o))) by definition æ =1122 z for NO, Ia), and
12 201 10 Z Zu
( 2
Amin
|
"
m // (1)
S 1 almost surely as n, co, n/d > 1. x ( z z)
which implies that
Summing over w yields the desired claim.
(1)
Observe that the feature vector &=
/ o [w VN ( is sub-Gaussian
vector with parameter VN
< 1 due the Lipschitzity of hence we may apply the following sub-Gaussian concentration result.
Σ 1
W
Lemma 23. Under the same assumption as Lemma 22, as d + oc, we have || S log d almost surely.
Proof. Notice that ||-I| and (-) are I-Lipschitz convex functions. Following [WZ21, Section 3), we can utilize the convex concentration inequality (see [Ver18, Chapter 5]) to obtain
P( ||
Ex ll O |III > t) < 2 exp (-ot²),
(B.26)
29

---

for any t conditioning on the event that Σ we know this event that Furthermore, under this event, we have
w
Σ
w
S VN. From Lemma 22,
F.
Σ W F
2 VN occurs with probability at least exp (-clog² d).
E ||| @ |II
Bl 1|
N Σ
W,
Thus, we can take t = log and combine Lemma 22 and (B.26) to attain the assertion.
Lemma 24 ([Min17], 7],, adapted). Under the same assumption as Lemma 22, given nd 2 clog* for some constant > O, we have
ndt
[E + A1) 1/2 1/ (E + 1/2 2t) 1 Cnexp
2(1 + 1/33) 1/3))oo log
if > log" d. (1+t/3)/
Proof. Let us denote
G.
E +
0z EE) A1)-11/2
72
Notice that E[G,] = 0. From Lemma 23, we know that ||[G;| S log" d/ An almost surely. More-
over, we may verify that [[ < 1 log² 4 N and An
1 log" d n An
log" An
Tr E
Tr(Es +I) Eg
Here we use the fact that all the eigenvalues of (E+ + XI) Eip are between 0 and 1. Thus, we can apply Matrix Bernstein bound (see (3.9) in [Min conclude this concentration result. O
With Lemma 24 we have the following useful corollary.
Corollary 25. Given nd > log* d, for any E (0, for sufficiently large d.
|| [Eo AI) 1/2 (E Et) + A1) 1/2 oa((),
122
with probability at least exp(-clopg d).
Prediction risk of ridge regression. We upper- the prediction risk the ridge regression estimator by controlling the decomposed error terms separately.
Proposition 26. Under the training procedure described in Algorithm with N 2(d and d=1 S S de for small > 0, 3 > 1 - then we have R(fxx) = oa(1) with probability when n. dl no,n/d v..
Proof. First note that R1 B. (f. - (xx)) = || by definition. For R2, [BES+ 22, Appendix D.3] showed that R2 S Ilʃ. - -1li2 + Alss 1 Finally, for R3 due to the concentration in Corollary 25, we have the following evaluation,
R₂
1
fi
Eg 4
+ XI)
E.
{
+1
-1 TT
n
1
+
Eo +A1)
-1/2
|
Eo ~1)
-1/2
(Es - Eo) (Eo Ar) -1/2
Vn
L
+
fT
E. + XI)
T
<
f 11
Recall that in Proposition we constructed some that achieves ll.. - f11,2- S 1/log d and Ilf S log d. This implies that I I 1/ log d and hence R1 o1.((1). Similarly, for = Od-=), we have R2 = -1r((). The proof is completed by Markov's inequality on Rs.
30


Reviews:
[{"review_id": "NggvAlpHov", "reviewer": "Reviewer_sXLf", "summary": "This paper studies a problem of parameter recovery under a generalized linear model-like setting. The novelty of this problem is revisiting this setting in the proportional limit setting: The sample size divided by the feature dimension converges to a fixed constant as both increases to infinity. This is a high-dimensional setting that received interest in the statistics theory community for many machine learning problems. This paper aims to study this setting with two learning methods: The kernel ridge regression and two-layer neural networks trained by gradient descent.\n\nFor both settings, the paper contributes tight bounds on recovering the unknown signal vector, given feature vector and label vectors in the proportional limit setting. Empirical studies validate the theoretical results.\n\n", "questions": "- Would the result hold for anisotropic input data?\n- Does the result apply to classification problems?\n- What if $y$ is contaminated by some random noise (which is a common assumption for this type of setting), or is it the case that the results already apply to this case?", "limitations": "Limitations (and extensions) have been discussed in the conclusion section.\n\nPotential negative societal impact would be unexpected given the technical nature of the paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1. A solid paper with strong technical contributions. Working out the asymptotic rates for kernel ridge regression and two-layer neural nets is a highly nontrivial problem, and the results shed some light on when one is better than the other, depending on the relation between signal and dimension.\n\nS2. The paper is overall nicely written. The discussion with respect to related work is also clear to read.", "weaknesses": "W1. Although the results are strong, there is a lack of insight that one can take away from the analysis. I consider this a limitation because such discussions have been commonly stated in existing random matrix theory / machine learning papers, e.g.,\n\nHastie, T., Montanari, A., Rosset, S., & Tibshirani, R. J. (2022). Surprises in high-dimensional ridgeless least squares interpolation. Annals of statistics, 50(2), 949.\n\nWei, A., Hu, W., & Steinhardt, J. (2022, June). More than a toy: Random matrix models predict how real-world neural representations generalize. In International Conference on Machine Learning (pp. 23549-23588). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the comments and questions. We address the technical points below. \n\n**Q: Lack of insight** \n\nHere we briefly summarize the takeaways of our analysis. \n\n1. By introducing a low-dimensional structure to the input data in the form of a spiked covariance (which is a common model in high-dimensional statistics), both kernel ridge regression (KRR) and neural network trained by gradient descent (NN+GD) can achieve better performance in the high-dimensional asymptotic limit. \nIn other words, both methods can benefit from the anisotropy of input data, provided that the target function (ground truth) focuses on the dominant feature directions. \n\n2. Similar to the isotropic setting, the performance of KRR and NN+GD depends on different notions of target complexity: for KRR the complexity of learning is measured by the *degree* of $f_*$, whereas for NN+GD it is the *information exponent*, which determines the difficulty of gradient-based optimization (note that this differs from prior results such as [GMMM20], where the optimization complexity of NN is not studied). \nSince the information exponent is by definition no larger than the degree of $f_*$, our analysis indicates that NN+GD can adapt to low-dimensional structures more effectively than KRR. \nThese theoretical predictions are supported by experiments. \n\n\n**Q: Anisotropic input data** \n\nNote that the spiked covariance input data we consider is by definition anisotropic. We, therefore, interpret the reviewer's comment as the extension to *arbitrary* covariance, which our current analysis does not cover. As remarked in the Introduction, most existing results on this topic only handled the isotropic setting, whereas we show that the addition of one single spike (arguably the simplest anisotropic setting) already gives rise to\nrich phenomenologies, such as learnability of a wider range of target functions in the *proportional regime*, and a new separation between the performance of KRR and NN+GD. We believe that our analysis serves as an important first step toward the understanding of learning behavior under more complicated and realistic structured data. \n\n\n**Q: Extension to classification problems**\n\nIn this submission we focus on the regression setting, for which the performance of KRR and NN has been extensively studied under high-dimensional isotropic input data (see Introduction); this provides a useful benchmark for us to compare our findings against (e.g., the difficulty of gradient-based learning measured by the information exponent). \nIt is definitely interesting to extend our current analysis to the classification setting, where the spike structure can be encoded in different ways, such as a Gaussian mixture model. \n\n**Q: Handling of label noise** \n\nRecall that the prediction risk admits a bias-variance decomposition, where the bias term comes from the learning of target function $f_*$, and the variance term is due to contamination of label noise $\\varepsilon$. Our current analysis focuses on the bias term (i.e., the noiseless setting), which is more challenging to characterize. When i.i.d. label noise is added, it can be shown via standard arguments that under an appropriate choice of regularization $\\lambda$, the variance term in both KRR and NN+GD vanishes when $n/d\\to\\psi$ becomes large. \nWe will include a discussion on the handling of label noise in the revised manuscript. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period. "}, {"review_id": "ooscQvxjMU", "reviewer": "Reviewer_GruC", "summary": "The paper studies the learning capacity of Kernel Ridge Regression and 2-Layer Neural Network to approximate a target function $f^\\star$ (denoted teacher) under a low dimensional structure. In particular, starting from a generative process where the covariance of the data has a low dimensional structure ($\\mathbf{I}_d +\\theta\\mu\\mu^\\top$), the target function $f^\\star$ depends on the signal $\\mu$ and the objective is to approximate the target function as best as possible with a model known as a student. The authors investigate two models for the student in particular (Kernel Ridge Regression and 2-Layer Neural Network with one gradient step).\n\nIn these two models, the authors ask the question of the condition on the strength of the low dimensional structure ($\\theta$) to approximate the function target $f^\\star$. The conclusions of the study make it possible to highlight a scaling $\\mathcal{O}(d^{1-1/d})$ and $\\mathcal{O}(d^{1-1/k})$ respectively for the KRR and the NN where the link function of the target is a $p$-degree polynomial with Information exponent $k$, which makes it possible to highlight the ability of these two models to find the low dimensional structure and which also highlights the superiority in some cases of the NN. Experiments have been provided to corroborate the various theoretical conclusions.", "questions": "1- Can the authors motivate why they choose the convention for the low dimensional structure both for the data and the target function $f^\\star$?\n2- Can the authors comment on the independence of the result on the kernel used which has generally huge impact?\n3- These conclusions are valid on real data as well or at least some insights on the actual performance for real data?", "limitations": "1- The main theoretical limitation is the lack of motivation of the choice of the teacher model. Why to model low dimensional structure as done by authors? Does this reflect the actual low dimensional structure encountered in real applications? The loose bounds (not depending on the kernel) also is a bit problematic\n\n2- The main empirical limitation is the lack of insights for real world application, the lack of experimentation on different kernels, ...", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1- Understanding the influence of low dimensional structures on the learning and generalization capabilities of machine models is an interesting question. The conclusions drawn by the authors seem to provide part of the answer although limited to a simple setting.\n\n2- The paper is quite well written and easy to follow and although not having checked all the calculations, the theoretical results seem new and rigorous.\n\n3- The experiments although in a synthetic setting corroborate the theoretical conclusions", "weaknesses": "1- The motivation for designing the target function $f^\\star$ seems to be vague and has very little motivation. It is clear that we want to introduce a low dimension structure but why introduce it both in the covariance of the data and also as a weights (oracle) for the teacher network. One could choose weights that have nothing to do with the structure of the data and see how the structure of the data impacts the learning capacity. Also, the choice to put the low dimensional structure on the covariance matrix of the data deserves to be explained in accordance with the type of low dimensional structure generally encountered in real data. In particular the low dimensional structure can be found for example in the mean of the data in the form of sparsity, ...\n\n2- The results obtained on the kernel seem to be a little too pessimistic or too loose. Indeed, one would have expected that the bounds depend on the type of kernel used and its proximity to the link function oracle $\\sigma^\\star$. Indeed, it is clear that for this problem where the means are zero, a linear kernel would give catastrophic performance compared to a quadratic kernel for example. I'm not sure that the boundary is tight enough. It is probably valid for the best choice of kernel. For example, I would be curious to see how Figure 2 and Figure 3 behave if instead of a Gaussian kernel one would take a quadratic kernel. Indeed, [1] proved that in the setting of zero mean data considered, the optimal kernel is not the Gaussian RBF but instead the quadratic kernel (see also other references in [2, 3]).\n\n3- Experiments on real data would have been useful to illustrate the conclusions of the paper, for example by taking a pre-trained network as a teacher and learning a KRR and a 2-NN to approximate the teacher. The low dimension structure can be added manually or modelled.\n\nReferences\n[1] Covariance discriminative power of kernel clustering methods, A Kammoun, R Couillet - Electronic Journal of Statistics, 2023\n[2] R. Couillet, F. Benaych-Georges, \"Kernel Spectral Clustering of Large Dimensional Data\", Electronic Journal of Statistics, vol. 10, no. 1, pp. 1393-1454, 2016\n[3] H. Tiomoko Ali, A. Kammoun, R. Couillet, Random Matrix Asymptotics of Inner Product Spectral Clustering\", IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP'18), Calgary, Canada, 2018.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful feedback. We address the technical points below. \n\n**Q: low dimensional structure both for the data and the target function $f_*$**\n\nAs remarked in the Introduction, the spiked covariance model is a classical statistical model that captures low-dimensional structure in high-dimensional data. The underlying intuition is that input directions with high variance are often good predictors of the labels -- indeed, this is the main reason that principal component analysis is used in practice.  \nIn this paper, we focus on the simplest spiked covariance setting with only one spike $\\boldsymbol{\\mu}\\in\\mathbb{R}^d$ and the target function $f_*$ is a *single-index model*. Note that this single-index setting has been extensively studied in many recent works as a canonical example of learning low-dimensional signals from high-dimensional data (e.g., [BAJ21] [BBSS22]), where the advantage of representation learning in neural networks can be demonstrated. \n \nIn order for the spikes to have higher predictive power, the target function $f_*$ needs to align with these large directions -- such alignment between the features and target function is analogous to the *source condition* in the nonparametric literature. \nFor this reason, we chose the single-index $f_*$ to be in the direction of $\\boldsymbol{\\mu}$, so that a larger spike corresponds to a stronger learning signal from the input data. \nIn contrast, when the index features are orthogonal to spike direction, then we would not expect the low-dimensional structure to improve the learning performance: imagine the extreme setting where the spike completely dominates the remaining directions (i.e., the input is almost one-dimensional), then the learning problem becomes ``misspecified'' in the sense that the input features contain no information about the target function. \nAs mentioned in the Conclusion section, an interesting future direction is to study the intermediate regime where $f_*$ is *partially* aligned with the spike direction. \n\n\n\n**Q: Pessimistic bounds for KRR**\n\nWe emphasize that our characterization of KRR prediction risk is  *sharp* for inner-product kernels we consider, and does not depend on the particular choice of the kernel function, similar to [GMMM19] [GMMM20].  \nThe main difference between our results and the references [1][2][3] mentioned by the reviewer is that we allow the spectral norm of the covariance to *grow with the dimensionality $d$*. In this setting, when the spike magnitude is sufficiently large, the kernel matrix cannot be approximated via a low-degree Taylor expansion; consequently, it becomes possible for KRR to go beyond learning a low-degree $f_*$.  \nMore specifically, the references [1][2][3] considered spectral clustering by approximating the kernel as a linear kernel plus some low-rank spikes and using these low-rank spikes for clustering. These results does not handle the divergent spike in the dataset, and thus the kernel model cannot represent high-degree $f_*$. For instance, in the kernel clustering setting, [1] proved the optimal kernel is a quadratic kernel in classifying zero-mean datasets with different covariances; but in our case, it is clear that a quadratic kernel cannot achieve vanishing error when the target function is beyond degree-2.  \nWe will include a more detailed discussion in the revision. \n\n**Q: Insights on real-world data**\n\nWhile the focus of our work is mostly theoretical, the following high-level takeaways may transfer to more complicated and realistic datasets. \n\n1. By introducing a low-dimensional structure to the input data in the form of a spiked covariance, both kernel ridge regression (KRR) and neural network trained by gradient descent (NN+GD) can achieve better performance in the high-dimensions. In other words, both methods can benefit from the anisotropy of input data, provided that the target function (ground truth) focuses on the dominant feature directions. \n\n2. The performance of KRR and NN+GD depends on different notions of target complexity: for KRR the complexity of learning is measured by the *degree* of $f_*$, whereas for NN+GD it is the *information exponent*, which determines the difficulty of gradient-based optimization (note that this differs from prior results such as [GMMM20], where the optimization complexity of NN is not studied). Since the information exponent is by definition no larger than the degree of $f_*$, we know that NN+GD can adapt to low-dimensional structures more effectively than KRR. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period."}, {"review_id": "wZaE6goX7M", "reviewer": "Reviewer_b5dX", "summary": "This paper considers the problem of learning a single-index target function under the spiked covariance data in the proportional asymptotic limit.\nThe results show the spike magnitude required for both kernel ridge regression (KRR) and two-layer neural networks (NN) trained with gradient descent and reveal that the NN can adapt to low-dimensional structures more effectively than KRR.\nNumerical experiments on synthetic data are also supportive.", "questions": "1. Can the results in this paper be applied to NN trained only by GD rather than the two-stage procedure considered in this paper?\n", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper is clear and easy to follow.\n\n* This paper shows theoretically the advantage of NN over KRR under a spiked convariance setting, which may provide more insight into the generalization of NN.", "weaknesses": "* It would be better to justify assumption of the single-index model considered in this paper.\n\n* Although both upper and lower bounds for KRR are established, this paper only provides upper bound for NN.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the comments and questions. We address the technical points below. \n\n**Q: Single-index assumption** \n\nThe single-index model is a canonical setting of learning low-dimensional signals from high-dimensional data; it covers many statistical learning problems such as phase retrieval and generalized linear models. Importantly, the low dimensionality of the single-index target may be exploited by *adaptive* learning procedures (as opposed to fixed bases models such as kernel methods), and thus this model class has been extensively studied in recent works to demonstrate the advantage of representation learning in neural networks -- see [AAM22] [BBSS22] and references therein. We will include a more thorough discussion in the revised manuscript. \n\n\n**Q: Only upper bound for NN**\n\nIndeed we only provide an upper bound for the performance of NN+GD. This being said, for the purpose of establishing a separation between KRR and NN+GD, our upper bound (together with the sharp analysis of KRR) is already sufficient. Moreover, from the experiments in Section 5, we empirically observe that our derived upper bound is tight in terms of the required spike magnitude. \n\nOptimization lower bound for neural networks is generally difficult to establish. One possibility is to study the correlational statistical query (CSQ) lower bound, which relates to the gradient descent complexity (for the squared loss) -- see [DLS22] and [ABAM23] for the isotropic setting. We leave this direction as future work.\n\n**Q: Beyond the two-stage procedure**\n\nOur two-stage feature learning paradigm has been employed in many recent works, including [AAM22], [BBSS22], [BES+22], [DLS22], and [ABAM23], but to our knowledge, we are the first to analyze the effectiveness of such procedure under anisotropic input data. Roughly speaking, this two-stage process realizes the representation learning advantage via gradient descent training of the first layer parameters, whereas the second-layer parameters can approximate the unknown link function. We believe that insights of our current results may be extended to general NNs trained via standard GD, but this likely requires new proof techniques and a more involved analysis. \n\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n\n"}, {"review_id": "DTsKjbM5NP", "reviewer": "Reviewer_qX6D", "summary": "This manuscript presents an exploration of high-dimensional regression using a theoretical approach. It focuses on a model characterized by:\n- A single-index model where the output is dependent on a one-dimensional projection of the input in a specific direction, $\\mu$.\n- Spiked covariance data with Gaussian input that exhibits variance significantly larger in the $\\mu$ direction.\n\nThe authors build upon this theory in an asymptotic framework where the dimension $d$, the number of samples $n$, and the variance of the spike approach infinity. They examine the performance of kernel methods and one hidden-layer neural networks in this context, particularly when n and d are proportional. The networks are simplified for the sake of analysis, having been trained with only one gradient step in the inner layer.\n\nThe key theoretical findings suggest that neural networks efficiently learn high-degree components, provided there are low-degree components in the same direction. Conversely, it is more challenging for kernel methods to learn high-degree components under similar conditions. This observation underscores an advantage of neural networks over kernel methods. These results are backed by numerical simulations in the same idealized setting, indicating a strong agreement with the theory.", "questions": "Comments:\n- l.104-105: This sentence is unclear. Could you elaborate?\n- Eq. (2.3): It seems the notations $X$ or $k(x,X)$ have not been introduced at this stage.\n- l.165: Could you clarify the \"i.e.\" part?\n- In Section 5: How would the neural network training plots appear if you used plain GD/SGD? Would the predicted scalings remain valid?\n- l.475-476: This sentence seems to lack clarity and precision. A better wording could be, \"The distribution of the random variable $(f_{ker}(x), f_*(x))$ remains unaltered through this rotation.\"\n\nPotential Mathematical Corrections:\n- Display of Theorem 3: It appears there is a square missing on the norm. Also, could we prove the stronger statement $\\Vert f_{ker} - P_{>l} f_* \\Vert \\to 0$?\n- Second equality of l.255: Considering $x$ is multidimensional, shouldn't $f'_*(x)$ be a vector?\n- Display of l.258: Shouldn't there be (i+1) factors in the sum?\n- l.483: p=0 seems to be an invalid choice as 1/p appears.\n- Display of l.489: Based on Prop. 2.7.1(a) of [Ver18], shouldn't there be 2 factors?\n- Display of l.491: Should there be constants in the inequality?\n- Second equation of l.499: The dimensions do not seem to match.\n- l.506: Considering the variance in one direction is not equal to 1, shouldn't the formula for the Hermite polynomials be different?\n\nTypographical Errors:\n- l.100: The correct term should be \"teacher\" instead of \"teaching\".\n- l.206: Please replace \"expend\" with \"expand\".\n- l.207: \"bases\" should be replaced with \"basis\".\n- l.467: Considering the regression setting, it would be more accurate to use \"output\" instead of \"label\".\n- l.480-481: This does not form a complete sentence.\n- l.484: Remove one \"+\" sign.\n- Second equation of l.499: It should be \\Delta_K' (not \\Delta_{K'}).\n- l.501: Shouldn't it be \"Hermite\" instead of \"Hermitian\"?\n", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This manuscript is well-articulated and provides transparent results. To the best of my understanding, it is the first work to explore this specific model, providing insights into the limitations of neural networks and kernel methods.", "weaknesses": "While this study builds upon previous works that examined single-index models and/or spiked covariance data, I encourage the authors to delve deeper into the following areas:\n\n(1) Why is the considered model better than those considered before? Could the authors identify a dataset that is likely to present a single-index model and spiked covariance data in the same direction, even approximately? The [GMMM20] work furnished an example to support their research, however, their example seemed to advocate for a multi-index model. I understand that your study deviates from theirs in the focus on the single-index scenario.\n\n(2) While the results are appropriately compared with related findings, there is a shortage of comparison in terms of the methods of analysis. Can you highlight the novelty of your analyses compared to previous works, particularly in contrast to [GMMM20], which seems closely related to your study? How harder was it to tackle the single-index case rather than multi-index with a diverging number of directions?\n\nFurthermore, I found it challenging to navigate through the appendix of this paper due to a plethora of typos and inaccuracies. Several sections were validated by aligning with another paper's proof (e.g., l.492-493, l.550, l.568-569, l.584-585). This raises questions about the originality of the proof techniques. As a consequence of all of this, I cannot confirm the accuracy of the proofs.\n\nI provide detailed comments and potential corrections below to enhance the clarity and correctness of the work. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and thoughtful feedback. We address the technical points below. \n\n**Q: Motivation of single-index $f_*$**\n\nThe single-index model is a canonical setting of learning low-dimensional signals from high-dimensional data; it covers many statistical learning problems such as phase retrieval and generalized linear models. Importantly, the low dimensionality of the single-index target may be exploited by *adaptive* learning procedures (as opposed to fixed bases models such as kernel methods), and thus this model class has been extensively studied in recent works to demonstrate the benefit of representation learning in neural networks trained with gradient descent -- see [AAM22] [BBSS22] and references therein. Note that such optimization-based advantage was not established in [GMMM20]. \n\nAs for why we do not consider the multiple-index setting, the main reason is that the optimization of neural network is less understood and more challenging to analyze. For instance, in order for the first gradient step to reveal information of the entire support of $f_*$, one may need to subtract the low-degree Hermite terms as done in [DLS22]. We therefore focus on the single-index target, which has been extensively studied in the isotropic setting (e.g., [ABJ21], [BES+22], [BBSS22], ...), and hence our results can be directly compared against existing benchmarks to illustrate the benefit of anisotropy in both KRR and NN+GD. \n\n**Q: Novelty of KRR analysis**\n\nWe make the following remarks. \n\n1. Our analysis of KRR builds upon previous results in [BMR21], [GMMM21], and [EK10]. However, our analysis is non-trivial as these prior works cannot be directly applied to KRR on an anisotropic Gaussian dataset with a divergent spike. The novelty of our proof comes from a combination of the proof techniques in [EK10, BMR21] and [GMMM21]. The former addresses the case when a nonlinear kernel with general data distribution is equivalent to a linear kernel; while the latter shows that a nonlinear kernel on unit sphere can be approximated by a polynomial kernel. In our case, we separate the data distribution into divergent spike and isotropic Gaussian; we then apply [EK10, BMR21] to show that the isotropic Gaussian part is equivalent to a linear kernel and will not help us to learn the target function; then, analogously to [GMMM21], we use the general Hermite polynomial expansion to exactly compute the performance of KRR on the spike direction. In the last step, since we are dealing with a polynomial kernel in one dimension similar to prior works, we omit some of the details by citing previous results. For completeness and convenience for readers, we will expand the proof and computations in our revision. \n\n2. As the reviewer mentioned, the main difference between our kernel analysis and [GMMM20] is that we consider an anisotropic Gaussian dataset with one divergent spike in the covariance, whereas [GMMM20] considered a multi-index target with diverging number of directions. The proof strategies are also different: [GMMM20] directly used spherical harmonics to truncate the kernel due to the spherical setting; while in our setting, we need to first apply a rotation and a polynomial truncation to the kernel matrix to separate the isotropic Gaussian component and the spike. \nWe note that our current approach cannot precisely characterize the performance of KRR under anisotropic Gaussian data with diverging number of spikes (as in [GMMM20] for spherical data), as we do not have a precise description of the orthogonal polynomial expansion for kernels under Gaussian data. We leave this direction as future work. \n\n**Q: Clarifications** \n\n* *lines 104-105:* We apply a normalization $\\frac{1}{\\sqrt{1+\\theta}}$ in the teacher model to make sure that the order of the argument inside the link function is unchanged as we increase the signal-to-noise ratio $\\theta$, and therefore the $L^2$ norm of $f_*$ also remains constant. \n* *line 165:* The \"i.e.\" part is a sufficient condition of positive semidefiniteness for the kernel $k(x_i,x_j)=g(x_i^\\top x_j)$. When $g$ is real analytic with nonnegative $g^{(k)}(0)$ for all $k\\in\\mathbb{N}$, we obtain a positive semidefinite inner product kernel. \n* *Scaling for plain GD/SGD:* In this work, we focus on the two-stage feature learning paradigm employed in recent works, including [AAM22], [BBSS22], [BES+22], [DLS22], and [ABAM23]; to our knowledge, we are the first to analyze the effectiveness of such procedure under anisotropic input data. Roughly speaking, we realize the representation learning advantage via gradient descent training of the first layer parameters, whereas the second-layer parameters can approximate the unknown link function. \nWe believe that insights of our current results may be extended to NNs trained via standard GD, but this likely requires new proof techniques and a more involved analysis.  \n\n**Q: Typos and corrections**\n\nThank you for the close reading and detailed feedback. We have revised manuscript accordingly, and below we summarize a few noticeable corrections. \n\n* In the display of Theorem 3, we missed a square on the norm when stating the prediction risk of KRR. \n \n* Below line 258, there should be (i+1) factors in the sum since we consider the Hermite expansion of $\\sigma'$. \n\n* In line 483, we need to exclude the case $p=0$ which corresponds to constant functions (that are easily learnable by KRR, e.g., see [BMR21]). In fact for simple presentation, in Corollary 4 we only consider teacher model with degree $p\\ge 1$. \n\n* In line 506, the variance of the first coordinate is not equal to 1 and hence different from the other coordinates. We will use a different notation to distinguish between Hermite polynomial expansions with different variances (for the first coordinate vs. the remaining). \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period. "}]

Author rebuttal:

Meta-review: In this paper, the authors study the problem of learning a scalar valued target function when data is provided with a spiked covariance structure.  The authors provide theoretical justification of the spike size necessary for detection by Kernel Ridge Regression and a two-layer neural network trained using gradient descent.  The results show that the performance of these methods depends on different notions of complexity: one depends on the degree the function and the other on its information exponent.  The strength of the paper is the novel theoretical results that identify relevant notions of complexity under which learning guarantees can be proved under data with a particular spiked covariance structure.  The results notably provide qualitative insight that neural networks can adapt to low-dimensional structures more effectively than Kernel Ridge Regression can under this structured data.  Several reviewers suggested that the paper could be improved by better justification of the single-index model being considered.
Decision: Accept (poster)