Title: Learning in the Presence of Low-dimensional Structure: A Spiked Random Matrix Perspective
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

We consider the problem of learning a single-index target function f ∗ : R d → R under the spiked covariance data:

where the link function σ ∗ : R → R is a degreep polynomial with information exponent k (defined as the lowest degree in the Hermite expansion of σ ∗ ), and it depends on the projection of input x onto the spike (signal) direction µ ∈ R d . In the proportional asymptotic limit where the number of training examples n and the dimensionality d jointly diverge: n, d → ∞ , n/d → ψ ∈ (0 , ∞ ) , we ask the following question: how large should the spike magnitude θ be, in order for ( i ) kernel methods, ( ii ) neural networks optimized by gradient descent, to learn f ∗ ? We show that for kernel ridge regression, β ≥ 1 -1 p is both sufficient and necessary. Whereas for two-layer neural networks trained with gradient descent, β > 1 -1 k suffices. Our results demonstrate that both kernel methods and neural networks benefit from low-dimensional structures in the data. Further, since k ≤ p by definition, neural networks can adapt to such structures more effectively.

1 Introduction

Learning under spiked covariance. Real-world data is often high-dimensional but it also exhibits certain low-dimensional structures. Indeed, the number of input features is exceedingly large in machine learning, yet most relevant information is concentrated in a low-dimensional subspace [LV07, HTFF09]. Such structures have important consequences on learning performance. In random matrix theory, high dimensionality is reflected by the proportional asymptotic limit where the numbers of training examples and input features diverge to infinity at the same rate, whereas the low-dimensional structure is often described by a spiked random matrix model [Joh01, BAP05, BS06, BGN11] in which a low-rank signal ('spike') is hidden in a largedimensional noise ('bulk'). In this work, we restrict ourselves to the simple setting where the signal is rank-1, and consider the following data-generating process:

where the goal is to estimate the target function (teacher) f ∗ , which is a single-index model depending on the signal direction µ ∈ R d , and σ ∗ : R → R is an unknown link function . The signal µ is also embedded in the spiked input data, and θ > 0 controls the spike magnitude; we scale θ ≍ d β for β ∈ [0 , 1) , where larger β indicates a 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 of the low-dimensional signal θ .

Prior results: isotropic data. When θ = 0 , the input 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, BMR21]. 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, GMMM21, HL22]. More generally, for KRR to learn a degreep polynomial target, a sample complexity of n = Ω( d p ) is required [GMMM21, LRZ20, DWY21, MMM21].

Neural networks + Gradient descent. While NNs can efficiently approximate a single-index model (with polynomial link σ ∗ ), the optimization complexity of gradient-based learning on isotropic data is known to be governed by the information exponent , defined as the index of the first non-zero Hermite coefficient of σ ∗ [AGJ21]. For NNs trained with gradient descent (GD) to learn a target with information exponent k , recent works showed that a sample complexity of n = ˜ Θ( d k ) suffices 1 [AGJ21, DLS22, BBSS22, ABAM23]. Hence, in the proportional regime ( n ≍ d ), learnability has only been established for k = 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 k = 1 , which cannot cover many important problems such as phase retrieval [Fie82, CC15, TV23]. Following the intuition that low-dimensional structure in the input data can make the learning problem easier, we ask the following question.

If the input data contains low-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 in the proportional asymptotic limit when the spike magnitude θ reaches a certain threshold, both KRR and GD-trained two-layer NN can learn the target function f ∗ , where the link function σ ∗ is a degreep polynomial with information exponent k (see Definition 1). Our findings are summarized below and illustrated in Figure 1.

For KRR with the inner-product kernel, we give a sharp analysis of the prediction risk and show that to learn a link function σ ∗ with degree p in the proportional regime, a spike magnitude of θ = Ω ( d 1 -1 p ) is both necessary and sufficient .

For a two-layer NN with ReLU activation, we upper bound the prediction risk when representation learning is performed via one gradient descent step on the firstlayer parameters (analogous to [BES + 22, DLS22]), and show that a spike magnitude of θ = ω ( d 1 -1 k ) is sufficient to learn σ ∗ with information exponent k .

10 2

10 3

dimensionality

d

10 1

10 2

low-dim. signal

NN succeeds,

KRR fails

NN & KRR succeed

d

1

1/

p

d

1

1/

k

̸

Our analysis reveals that both KRR and NNs can benefit from the low-dimensional structure (spike) in the input data, and they can learn a wider class of target functions in the proportional regime compared to the isotropic setting (i.e., 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 a smaller spike magnitude θ to learn the same target compared to KRR (see Figure 1). In other words, a GD-trained NN can utilize such a low-dimensional structure more effectively than KRR.

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 a 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 a different class of target functions with potentially higher degree and information exponent.

2 Preliminaries: Problem Setting and Assumptions

Notations. ∥ · ∥ denotes the ℓ 2 norm for vectors and the ℓ 2 → ℓ 2 operator norm for matrices, and ∥ · ∥ F is the Frobenius norm. O d ( · ) and o d ( · ) stand for the standard big-O and little-o notations, where the subscript highlights the asymptotic variable; we write ˜ O ( · ) when (poly-)logarithmic factors are ignored. O d, P ( · ) (resp. o d, P ( · ) ) represents big-O (resp. little-o) in probability as d → ∞ . Ω( · ) , Θ( · ) , ω ( · ) are defined analogously. Denote γ as the standard Gaussian distribution in R . Given f : R d → R , we denote its L p -norm with respect to the data distribution as ∥ f ∥ L p .

2.1 Basic Assumptions

We consider Gaussian input data with a spiked covariance:

where the signal direction µ ∈ R d satisfies ∥ µ ∥ = 1 , and the spike magnitude θ is allowed to grow with the input dimensionality d as specified by the exponent β : as β 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 d eff ≍ d 1 -β , 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 a student-teaching setting , where the labels y are generated from a teacher model (target function) f ∗ : R d → R . In the spiked covariance model (2.1), it 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 µ :

where the link function σ ∗ ∈ L 2 ( γ ) is centered such that E [ σ ∗ ( z )] = 0 for z ∼ N (0 , 1) (this can be achieved by subtracting the mean from the training labels as in [DLS22]), and the normalization factor (1 + θ ) -1 / 2 ensures that the L 2 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 to construct an optimal learning algorithm, but to understand the behavior of KRR and two-layer NN without such data preprocessing, similar to [GMMM20].

We also assume σ ∗ is a degreep polynomial with information exponent (IE) k defined below.

̸

Definition 1 (Information exponent) . Let { h j } ∞ j =0 denote the normalized Hermite polynomials. The information exponent of f ∈ L 2 ( γ ) , which we denote by k , is the index of the first non-zero Hermite coefficient of f , that is, given the Hermite expansion f = ∑ ∞ j =0 α j h j , k := min { j > 0 : α j = 0 } .

It is clear that by definition, we always have k ≤ p , and equality is only achieved when σ ∗ is a 'pure' degreek Hermite polynomial. Intuitively, IE measures the magnitude of information contained in the gradient, and larger k 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 i.i.d. label noise, meaning that, in light of the bias-variance decomposition, we analyze the bias term describing the extent f ∗ is learned by the student model. If label noise is introduced, the additional variance term (describing the 'overfitting' due to noise) can be handled by standard concentration argument.

2.2 Learning Objective and Training Procedure

Given training examples { ( x i , y i ) } n i =1 , where x i i . i . d . ∼ N (0 , I d + θ µµ ⊤ ) and y i = f ∗ ( x i ) , we consider two student models obtained via the following training procedure in the proportional regime:

Student model I: Kernel ridge regression. We focus on the inner-product kernel: k ( x , y ) = g ( ⟨ x , y ⟩ d ) for some g ∈ L 2 ( γ ) (similar argument can also apply to rotationally invariant kernels as in [EK10, DWY21]). Denoting the associated reproducing kernel Hilbert space (RKHS) with H , the kernel ridge regression estimator is given by

where λ > 0 is the ridge parameter, and K ∈ R n × n is the kernel Gram matrix with entries K i,j = k ( x i , x j ) . Spectral properties of the kernel matrix and the prediction risk of KRR have been extensively studied in various high-dimensional settings [EK10, GMMM21, BMR21, XHM + 22, LY22]. Since the kernel function is fixed before seeing the data, we intuitively expect that KRR cannot effectively adapt to a low-dimensional structure in the learning problem.

Student model II: Two-layer neural network. We aim to learn the following widthN network:

where σ : R → R is a nonlinear activation, and the trainable parameters W ∈ R d × N , a , b ∈ R N are initialized from standard Gaussian distribution. We require the network width to be N = Ω( d ε ) for some small ε > 0 . Such a choice entails that it is sufficient to have Θ( d 1+ ε ) total training parameters, which is almost proportional to the sample size n ≍ d ; in contrast, for KRR we need to store the kernel Gram matrix with n 2 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 update 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 parameters 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).

Feature learning for the 1st layer. We optimize the representation of the NN via gradient descent, where the training objective is the empirical squared loss L ( f ) = 1 n ∑ n i =1 ( f ( x i ) -y i ) 2 . Specifically, denote the i -th column of the initialized weight matrix W 0 as w (0) i ∈ R d , and the initialized NN as f 0 NN , we take one GD step with learning rate η as follows,

Due to the anisotropic input data, the pre-activation ⟨ x , w (1) i ⟩ may blow up if the trained parameter w (1) i aligns with the signal direction µ . To circumvent this issue, we normalize each neuron to have (roughly) unit pre-activation, i.e., for i ∈ [ N ] , we perform the weight normalization:

At high level, this resembles the often-used normalization layers in deep learning [IS15, BKH16].

Ridge regression for the 2nd layer . After obtaining the updated first-layer parameters W 1 ∈ R d × N , we optimize the second-layer a by solving the ridge regression objective - this can be done by using GD to solve an ℓ 2 -penalized least squares. To circumvent the dependence between the training data X and the learned W 1 , we follow [BES + 22] and estimate the regression coefficients ˆ a using a new set of training data { ˜ x i , ˜ y i } n i =1 with the same sample size n .

Denoting the feature matrix on the fresh training set { ˜ X , ˜ y } as Φ := 1 √ N σ ( ˜ XW 1 + b ) ∈ R n × N the ridge regression estimator can be obtained by: ˆ a = argmin a { 1 n ∥ ˜ y -Φ a ∥ 2 + λ ∥ a ∥ 2 } .

,

Learning in the proportional limit. Given a target function f ∗ and a learned (student) model ˆ f , we evaluate the model performance using the prediction risk:

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 an algorithm learns the target function f ∗ in the proportional regime if, for any (small) constant ϵ > 0 , there exists some constant ψ ∗ > 0 such that when n/d → ψ ≥ ψ ∗ , the algorithm constructs ˆ f using n i.i.d. training examples, and achieves R ( ˆ f ) = ∥ ˆ f -f ∗ ∥ 2 L 2 ≤ ϵ with probability 1 when n, d →∞ proportionally.

Based on this definition, in order for a student model ˆ f (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 ≍ d training samples).

3 Kernel Ridge Regression

For an inner-product kernel, the performance of KRR depends on the behavior of the pairwise innerproduct of the training examples ⟨ x i , x j ⟩ . Intuitively speaking, since the inner-products 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, WZ23]). Introducing a 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 θ 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 a positive semidefinite kernel given by k ( x i , x j ) = g ( ⟨ x i , x j ⟩ /d ) , i.e., g ( k ) (0) ≥ 0 for all k (see [SS02, Section 13.1]). Recall that the strength of the low-dimensional component is determined by the spike magnitude θ ≍ d β ; in the following, we characterize how the exponent β ∈ [0 , 1) affects the performance of KRR.

We denote by P ≤ ℓ : L 2 → L 2 the orthogonal projector onto the subspace of polynomials of degree at most ℓ for any ℓ ∈ N . Similarly, we denote by P >ℓ = Id -P ≤ ℓ the projector to the orthogonal complement where Id is the identity operator. To illustrate this notion, consider ℓ = 1 , then we have

where a ∗ , b ∗ = argmin a, b E x [( f ∗ ( x ) -a -⟨ b , x ⟩ ) 2 ] . Note that any polynomial f ∗ with degree at most ℓ satisfies P >ℓ f ∗ = 0 . Denoting the KRR model as ˆ f ker , we have the following sharp analysis of the asymptotic prediction risk.

Theorem 3. Given any fixed ℓ ∈ N . Suppose that g ( k ) (0) > 0 for all k ≤ ℓ , and the spike magnitude scales as θ ≍ d β for β ∈ ( 1 -1 ℓ , 1 -1 ℓ +1 ) . Then as n, d →∞ , n/d → ψ ∈ (0 , ∞ ) , the prediction risk of the KRR estimator ˆ f ker with λ = O d (1) satisfies the following with probability 1,

As a corollary of Theorem 3, based on Definition 2, we can obtain a sharp threshold of β for learning a degreep polynomial link function in the proportional regime (Definition 2).

Corollary 4. Assume σ ∗ is a degreep polynomial with p ≥ 1 defined in (2.2) . Under the same assumptions of Theorem 3, if β > 1 -1 p , then the KRR estimator with λ = O d (1) learns f ∗ in the proportional regime, i.e., R ( ˆ f ker ) = o d (1) with probability 1 when n, d →∞ proportionally.

We make the following remarks on Theorem 3 and Corollary 4.

At a high level, the above theorem aligns with the conclusion of [GMMM20] which assumes spherical data with a diverging number of spikes. Recall that the effective dimensionality of our anisotropic Gaussian data is d eff ≍ d 1 -β , and Theorem 3 implies that to extract the signal direction µ and learn a degreep teacher model, a sample size of n = ω ( d p eff ) is required; note that for β > 0 , this implies an improvement over the isotropic setting where the sample complexity is n = ω ( d p ) .

While we only state the result for inner-product KRR (similar to [LRZ20, BMR21]), we expect similar findings for general rotationally invariant kernels: k ( x i , x j ) = g ( ⟨ x i , x j ⟩ , ∥ x i ∥ , ∥ x j ∥ ) , since the norm term has negligible contribution to the learning of a single-index target (2.2). In fact, [DWY21] established a lower bound for rotationally invariant KRR, which, when applied to our setting, gives a necessary sample complexity of n ≍ d Ω( p ) eff . Although this lower bound may not be sharp, it illustrates that the required spike magnitude needs to scale with target degree p .

3.2 Intuition behind the Analysis

We briefly summarize the two main ingredients for the proof of Theorem 3. As a preliminary step, we can rotate the signal µ to the e 1 direction due to the rotational invariance of Gaussian distribution. After such transformation, we have ([ x i ] 1 , [ x i ] 2: d ) i . i . d . ∼ N (0 , 1+ θ ) ⊕N (0 , I d -1 ) with θ = Θ( d β ) .

Polynomial approximation of the kernel matrix. Firstly, we approximate the inner-product KRR with some degreeL polynomial kernel in the form of

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 σ ∗ [BMR21]. For general deterministic datasets, similar approximation has been studied in [WZ23] where the 'orthogonality' of input data determines the polynomial approximation error. In our setting, the inner product 1 d x ⊤ i x j = O d ( d β -1 ) depends on the spike magnitude θ ≍ d β . Consequently, we approximate our inner-product kernel by some polynomial kernel whose degree depends on β .

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, MMM21]). After rotation, the teacher model (2.2) can be written as f ∗ ( x ) = σ ∗ ( z ) , for z ∼ N (0 , 1) and [ x ] 1 = √ (1 + θ ) z independent with other entries [ x ] i for 2 ≤ i ≤ d . In light of the kernel trick [Ver18, Section 3.7], we write k ( x i , x j ) = ⟨ Φ( x i ) , Φ( x j ) ⟩ 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 [ x ] 1 to learn a function σ ∗ ( z ) .

4 Two-layer Neural Network

During the feature learning phase of NN training, we intuitively expect the parameters to align with the (rank-1) signal direction which allows NN to overcome the 'curse of dimensionality'. Indeed, in this section, we show that the first gradient step on the first-layer weights W enables the model to 'zoom in' to the low-dimensional structure, given that the spike magnitude θ 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):

Denote the NN optimized by the two-stage procedure outlined in Section 2.2 as ˆ f NN . In the sequel, we restrict ourselves to σ ( z ) = ReLU( z ) = max { z, 0 } . We expect similar characterization to hold when the student activation σ has non-zero Hermite coefficients up to degreep (see [ABAM23]).

The following theorem gives a sufficient condition on the spike magnitude θ ≍ d β in order for ˆ f 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 = Ω( d ε ) , η = Ω( N 1 / 2+ ε ) , and appropriately chosen ℓ 2 regularization d ε -1 ≲ λ ≲ d -ε for small ε > 0 . Then, for σ = ReLU, we have

as n, d → ∞ , n/d → ψ . 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 θ = ω ( d 1 -1 k ) is sufficient for a GD-trained two-layer NN to learn f ∗ in the n ≍ d regime. Similar to prior works [AGJ21, BBSS22], the information exponent k also determines the complexity of the learning problem in our setting: larger k implies a more 'difficult' task for GD, hence we need a larger θ (i.e., stronger low-dimensional signal) to achieve linear sample complexity. Moreover, in light of the effective dimensionality defined in [GMMM20], Theorem 5 implies a sample size of n = ω ( d k eff ) suffices to learn a single-index target with information exponent k ; this contrasts the n = ω ( d p eff ) complexity of KRR given in Theorem 3. Such discrepancy will be highlighted and empirically validated in Section 5.

4.2 Intuition behind the Analysis

Theorem 5 is established in two steps: first, we prove that under certain conditions, gradient descent can align the first-layer parameters W with the signal direction µ ; then we show that after achieving such an alignment, f ∗ 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-layer coefficients), the initial NN has small output f (0) NN ( x ) ≈ 0 . Therefore the gradient of the squared loss is dominated by the correlation between the student and teacher model. Concretely, consider the population gradient for one parameter vector w ∈ R d at initialization (for simplicity we omit the bias unit here):

where ( i ) is due to multivariate Stein's lemma, and ( ii ) follows from the definition of Σ . Utilizing the Hermite expansions of the student and teacher nonlinearities,

we have the following decomposition on the strength of the correlation term,

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 a larger spike θ , and ( ii ) it vanishes when the information exponent k is large, since ⟨ µ , w ⟩ = ˜ Θ P ( d -1 / 2 ) at initialization. The gradient magnitude relates to the sample complexity because in order to establish learnability in the proportional regime, we need to achieve nontrivial gradient concentration using n ≍ d samples. Consequently, the sufficient condition in Theorem 5 requires a larger θ when the information exponent k is large, and vice versa.

Remark. One caveat in the above derivation is that the Hermite coefficients of the student nonlinearity { α i } ∞ i =1 need to be non-zero up to degreek , which is not satisfied by ReLU or most commonlyused activation functions. We solve this issue by taking into account the bias units b i at initialization which 'diversify' the nonlinearity, similar to [BBSS22] (see Lemma 15 for details).

Nonparametric learning with random biases. After the signal direction is identified via gradient descent, we need to learn the unknown link function σ ∗ . We utilize random biases as a resource for univariate approximation - similar argument has appeared in many prior works [DLS22, BBSS22, BEG + 22, MHPG + 23, ABAM23]. Recall that in the representation learning phase, the bias units b i are not optimized. Therefore, if the weight vectors w align with the signal direction w ≈ µ , then learning the second-layer coefficients a via ridge regression can be reduced to a univariate kernel regression problem; specifically, given x i , x j ∈ R d , the univariate kernel is given by k ( z i , z j ) = E a,b [ σ ( a · z i + b ) σ ( a · z j + b )] , where z i = ⟨ x i , µ ⟩ , z j = ⟨ x j , µ ⟩ . For σ = ReLU , it is known that such kernel can efficiently learn a polynomial link function σ ∗ (e.g., see [DLS22, BBSS22]).

5 Experiments: Comparing KRR and NN

In this section, we empirically validate the theoretical results presented in previous sections. We construct target functions f ∗ with a link of varying degree p and information exponent k and experimentally compute the prediction risk of KRR and that of a two-layer NN.

k = p : NN & KRR are comparable. We first consider the setting where k = p , which indicates that the link function is a pure degreep Hermite polynomial: σ ∗ ( z ) = h p ( z ) . Recall that θ ≍ d β ; Theorem 3 implies that KRR learns the target function in the proportional regime when the spike magnitude satisfies β > 1 -1 p , which matches the sufficient condition for NN given in Theorem 5.

In the experiments of Figure 2, we compute the prediction risk of KRR and that of two-layer NN optimized via the two-stage procedure outlined in Section 2, where the link function σ ∗ is a pure degree-2 and degree-3 Hermite polynomial. We observe that both KRR and NN can learn f ∗ with linear sample complexity when the spike magnitude θ exceeds the same threshold predicted by our theory; intuitively speaking, this means that for the k = p setting, NN and KRR utilize the lowdimensional 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 2

10 3

dimensionality

d

10 1

10 2

spike

d

1/2

0.2

0.4

0.6

0.8

1.0

(a)

p

= 2

, k

= 2

(KRR).

10 3

dimensionality

d

10 0

10 1

spike

d

1/2

0.2

0.4

0.6

0.8

1.0

(b)

p

= 2

, k

= 2

(NN).

10 2

10 3

dimensionality

d

10 1

10 2

spike

d

2/3

0.2

0.4

0.6

0.8

1.0

1.2

(c)

p

= 3

, k

= 3

(KRR).

10 3

dimensionality

d

10 0

10 1

spike

d

2/3

0.2

0.4

0.6

0.8

1.0

1.2

(d)

p

= 3

, k

= 3

(NN).

k < p : NN outperforms KRR. Next, we consider a setting where the link function σ ∗ is a mixture of low- and high-degree components: σ ∗ ( z ) = ∑ j ∈J α j h j ( z ) , with min J = k, max J = p . In this case, Theorem 3 predicts that KRR requires a spike magnitude of β > 1 -1 p to learn f ∗ with linear sample complexity, whereas Theorem 5 entails that a smaller spike magnitude β > 1 -1 k is sufficient for the two-layer NN.

Figure 3 shows that the predictions from Theorem 3 and 5 once again accurately align with the experimental results. Observe that when k < p , GD-trained NN indeed outperforms KRR, in the sense that it requires a less prominent low-dimensional structure (i.e., smaller spike magnitude θ ) to attain linear sample complexity. This illustrates the benefit of NN under structured data due to the presence of representation (feature) learning. Roughly speaking, for target function with k < p , the sample complexity of representation learning (i.e., finding the signal direction µ via gradientbased learning) is low compared to directly estimating f ∗ on the original input features; therefore GD-trained NNs achieve better performance than KRR.

10 2

10 3

dimensionality

d

10 1

10 2

spike

d

2/3

0.2

0.4

0.6

0.8

1.0

(a)

p

= 3

, k

= 1

(KRR).

10 3

dimensionality

d

10 0

10 1

spike

1

0.1

0.2

0.3

0.4

0.5

0.6

0.7

(b)

p

= 3

, k

= 1

(NN).

10 2

10 3

dimensionality

d

10 1

10 2

spike

d

2/3

0.2

0.4

0.6

0.8

1.0

1.2

(c)

p

= 3

, k

= 2

(KRR).

10 3

dimensionality

d

10 0

10 1

spike

d

1/2

0.2

0.4

0.6

0.8

1.0

(d)

p

= 3

, k

= 2

(NN).

6 Conclusion and Future Directions

We investigated the performances of KRR and GD-trained NNs in the proportional asymptotic limit. We analyzed the prediction risk of these models in learning a nonlinear single-index model (2.2), where the index feature in the teacher model only relies on the spike direction µ 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 f ∗ . In a companion work [MHWSE23], we investigate a more general setting where the spiked covariance data only provides partial information of f ∗ (i.e. the 'misaligned' scenario), and show that feature learning (via gradient flow) also benefits from such structured data. Finally, extending our results to multi-index teacher models [CM20, DLS22, ABAM23] and considering more general training dynamics (e.g., beyond the first gradient step) applied to multiple-layer NNs are interesting directions left for future studies.

References

[CC15] Yuxin Chen and Emmanuel Candes. Solving random quadratic systems of equations is nearly as easy as solving linear systems. Advances in Neural Information Processing Systems , 28, 2015.

[Chi11] Theodore S Chihara. An introduction to orthogonal polynomials . Courier Corporation, 2011.

[CM20] Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory , pages 1161-1227. PMLR, 2020.

[CS13] Xiuyuan Cheng and Amit Singer. The spectrum of random inner-product kernel matrices. 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.11904 , 2020.

[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on 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 Conference on Machine Learning , pages 2804-2814. PMLR, 2021.

[EK10] 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. Applied optics , 21(15):2758-2769, 1982.

[GLK + 20] Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc M´ ezard, and Lenka Zdeborov´ a. Generalisation error in learning with random features and the hidden manifold model. In International Conference on Machine Learning , pages 3452-3462. 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. Linearized two-layers neural networks in high dimension. The Annals of Statistics , 49(2):1029-1054, 2021.

[GSS21] Friedrich G¨ otze, Holger Sambale, and Arthur Sinulis. Concentration inequalities for polynomials in α -sub-exponential random variables. Electronic Journal of Probability , 26(none):1 - 22, 2021.

[HL22] Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random 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 Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning , pages 448-456. pmlr, 2015.

[Joh01] Iain M Johnstone. On the distribution of the largest eigenvalue in principal components analysis. The Annals of statistics , 29(2):295-327, 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.

[LY22] Yue M Lu and Horng-Tzer Yau. An equivalence principle for the spectrum of random inner-product kernel matrices. arXiv preprint arXiv:2205.06308 , 2022.

[MHPG + 23] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. International Conference on Learning Representations , 2023.

[MHWSE23] Alireza Mousavi-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 operators. Statistics & Probability Letters , 127:111-119, 2017.

[MM22] Song Mei and Andrea Montanari. The generalization error of random features regression: 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 concentration. Applied and Computational Harmonic Analysis , 2021.

[Rah17] Sharif Rahman. Wiener-hermite polynomial expansion for multivariate gaussian probability measures. Journal of Mathematical Analysis and Applications , 454(1):303-334, 2017.

[RGKZ21] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborov´ a. Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning , pages 8936-8947. PMLR, 2021.

[Sam23] Holger Sambale. Some notes on concentration for α -subexponential random variables. In High Dimensional Probability IX: The Ethereal Volume , pages 167-192. Springer, 2023.

[SS02] Bernhard Sch¨ olkopf 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 arbitrary initialization solves non-smooth, non-convex phase retrieval. Journal of Machine Learning Research , 24(58):1-47, 2023.

[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 semicircle law and concentration of nonlinear random matrices for ultra-wide neural networks. arXiv preprint arXiv:2109.09304 , 2021.

[WZ23] Zhichao Wang and Yizhe Zhu. Overparameterized 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-free bounds for sums of independent matrices and simple tensors via the variational principle. arXiv preprint arXiv:2108.08198 , 2021.

Table of Contents

A Analysis of Kernel Ridge Regression

In this section, we analyze the performance of KRR defined in Section 2. Recall that

where θ ≍ d β , and β ∈ [0 , 1) . Let Σ = I + θ µµ ⊤ and x i = Σ 1 / 2 z i , where z i ∼ N (0 , I ) , and Σ 1 / 2 = I +( √ 1 + θ -1) µµ ⊤ . Then, the training label y i can be equivalently expressed as

for i ∈ [ n ] . Recall X ⊤ = [ x , . . . x n ] ∈ R d × n . We define the nonlinear kernel matrix by K = g ( 1 d XX ⊤ ) ∈ R n × n , where smooth function g : R → R is applied to each entry of the matrix. Additionally, we write

We analyze the kernel ridge regression estimator:

Specifically, we aim to compute the prediction risk in the proportional regime:

where x ∈ R d is an independent copy of x i sampled from (A.1). Since the kernel K is rotationally invariant, i.e.,

for any orthogonal matrix S ∈ R d × d and i, j ∈ [ n ] , we may apply an orthogonal matrix O ∈ R d × d for each data sample such that

where e 1 is the standard coordinate vector. It is clear that in distribution, ˆ f ker ( x ) and f ∗ ( x ) remain unchanged before and after the transformation. Therefore, without loss of generality, we will consider the case that µ = e 1 for the proofs in this section, that is, we write

and training label y i = f ∗ ( x i ) only depends on the first entry of the data x i for i ∈ [ n ] .

A.1 Polynomial Approximation of Kernel Matrix

Denote the concatenation of x and X as ˜ X ⊤ := [ x , X ⊤ ] ∈ R d × ( n +1) , and the enlarged nonlinear kernel matrix by ˜ K = g ( 1 d ˜ X ˜ X ⊤ ) ∈ R ( n +1) × ( n +1) . First we present a preliminary approximation of ˜ K by a proper polynomial kernel.

Lemma 6. Suppose that β ≤ 1 -1 p for some p ∈ N . Under the assumptions of Theorem 3, as n, d →∞ proportionally, we have

where

Proof. For simplicity, we denote x = x 0 and x ⊤ i = [ x 1 i , ¯ x ⊤ i ] , where x 1 i is the first entry of x i , and ¯ x i the remaining entries for 0 ≤ i ≤ n . Notice that 1 d x ⊤ i x j = 1 d z i Σ z j , for any 0 ≤ i, j ≤ n . By Bernstein inequality [Ver18, Corollary 2.8.3], we have

̸

for any i = j and t > 0 . Therefore, ∣ ∣ 1 d x ⊤ i x j ∣ ∣ = o d, P ( d β -1 log d ) for β ∈ [1 / 2 , 1) and ∣ ∣ 1 d x ⊤ i x j ∣ ∣ = o d, P ( d -1 / 2 log d ) for β ∈ [0 , 1 / 2] . Then we have

̸

Here, offdig ( · ) denotes setting the diagonal entries of the matrix to be zero. In addition, following the same proof as Theorem 1 of [EK10], we can obtain a similar approximation for the diagonal part dig ( ˜ H p -˜ 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 a further approximation using orthogonal polynomials. Recall that

where

From Lemma 6 we know that for some p ∈ N ,

for ∥ ∆ K ∥ = o d, P (1) and ∥ ∥ ∆ ′ K ∥ ∥ = o d, P (1) , where σ p is defined in Lemma 6, H p := P p ( 1 d XX ⊤ ) , and H p ( x , X ) := P p ( 1 d Xx ) ⊤ is a row vector. Then following from Lemmas A.3, A.8 and A.13 in [BMR21], we can conclude that R ( ˆ f ker ) = R p + o d, P (1) , where

and

Nowweanalyze the asymptotic behavior of R p by utilizing the Hermitian polynomials to reorganize kernel H p . Since Σ now is a diagonal matrix in (A.2), from [Rah17], we know the multivariate Hermite polynomials with respect to N (0 , 1 d Σ ) in R d are strongly orthogonal, each of which is the product of univariate Hermite polynomials on each coordinate. More precisely, for any k ∈ N and multi-index I := ( i 1 , i 2 , . . . , i d ) ∈ [ k ] d with i 1 + i 2 + · · · + i d = k , we define the k -th multivariate Hermite polynomial with respect to I by h k, I ( x ) := ∏ d j =1 h i j ( x j ) where x = ( x 1 , . . . , x d ) and h i j ( · ) is the i j -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 ∈ R d has the form

for some coefficients c k, I ∈ R . Similar expansions using Gegenbauer polynomials for uniform unit sphere distribution have been used in [GMMM21, MMM21, GMMM20]. The decomposition we consider is specifically related to the Hermite expansions of σ ∗ :

where h j ( · ) is the j -th Hermite polynomial with z ∼ N (0 , I ) . Notice that E [ h k, I ( x ) h l, J ( x )] = δ k,l δ I , J for x sampled from the spiked model (2.1).

By the kernel trick (e.g., see Section 3.7 of [Ver18]), a k -th monomial kernel satisfies

where each entry of x ⊗ k i ∈ R d k is a k -th degree monomial of x i . Hence we can rewrite the above inner product based on Hermite expansions of each entry of x ⊗ k i and x ⊗ k j . Therefore, for kernel H p ( x i , x j ) , there exists a Hermite polynomial feature map Φ : R d → R L where

such that

where each row vector Φ ( · ) is the concatenation of orthogonal polynomials h k, I ( · ) for 0 ≤ k ≤ L and D is the diagonal matrix where the diagonal entries are the Hermite coefficients corresponding to H p which will be specified in the sequel. We define √

which gives H p = ΦΦ ⊤ . Furthermore, since the variance of the first entry of x i is different from the other entries and the label y i only depends on the first entry, we will separate the Hermite polynomials depending on x 1 i and the terms depending on the remaining entries ¯ x i ∈ R d -1 . That is, we further decompose each Φ ( x i / √ d ) = [ Φ 0 ( x i / √ d ) , Φ (¯ x i / √ d )] , where Φ 0 ( x i / √ d ) ∈ R L 0 and

where B 0 ( d, k ) is the total number of multivariate Hermite polynomials of x i with degree k , whose components with respect to x 1 i is not vanishing. Based on this decomposition, we also write

Lemma 7. Let d eff := d 1 -β . Under the assumptions of Theorem 3, we know that √ √

where D 2 is a diagonal matrix with entries

for m ∈ [ k ] and k ∈ [2 p ] . Furthermore, based on the definition of Φ 0 as above, we have

where each entry of ¯ H p (¯ x i , ¯ x j ) := P p ( 1 d ¯ x ⊤ i ¯ x j ) , and D 2 0 is a diagonal matrix defined analogously as D 2 but with m > 0 in (A.8) . Note that ¯ H p is an inner-product polynomial kernel matrix of ¯ X := [¯ x 1 , . . . , ¯ x n ] ∈ R ( d -1) × n .

This lemma follows from properties of multivariate Hermite polynomials [Chi11]. Recall

Proof. Var ( x 1 ) = 1 + θ = Θ( d β ) .

By re-normalizing the variables, we obtain

̸

where I = ( i 1 , . . . , i d ) . As for Φ 0 , the diagonals of D 0 correspond to m = 0 above.

A.3 Analysis of Prediction Risk

The next lemma shows that ¯ H p is approximated by a linear kernel which does not aid the learning of f ∗ ; this entails that we may focus on the term involving Φ 0 . Notice that ¯ H p 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 →∞ proportionally, we have

where and ∥ K lin ∥ = O d, P (1) . Here, c 1 , c 2 and c 3 are some positive constants that only depend on g ( k ) (0) for k = 0 , 1 and 2 .

The above lemma is a consequence of Theorem 1 of [EK10], and hence we omit the detailed proof; for the formulae of constants c 1 , c 2 and c 3 , we refer to Equation (57) of [BMR21]. We remark that in the following concentration analysis, the exact limiting spectral distribution of K lin does not appear, since we do not consider the exact setting of β = 1 -1 ℓ for ℓ ∈ N .

Lemma 9. Under the assumptions of f ∗ in Theorem 3, we have that

for some vector α = [0 , . . . , 0 , α k , . . . , α p , 0 , . . . , 0] ⊤ ∈ R 2 p , where P ∈ R L 0 × 2 p is a projector onto the Hermite components with I = ( i 1 , 0 , . . . , 0) .

Here α represents the Hermite coefficients of the function f ∗ . Lemma 9 directly follows from (A.5) since y only depends on x 1 i for i ∈ [ d ] . Notice that here P selects h k ( x 1 i ) for i ∈ [ d ] and k ∈ [2 p ] from all multivariate Hermite polynomials.

Next, we prove the following concentration statements, parallel to existing results in [GMMM21, MMM21]forspherical data. For any ℓ ∈ [2 p ] , we decompose Φ 0 = [ Φ 0 , ≤ ℓ , Φ 0 ,>ℓ ] where Φ 0 , ≤ ℓ are composed of all Hermite polynomials h k, I with k ≤ ℓ , and Φ 0 ,>ℓ collects all Hermite polynomials h k, I with k > ℓ . Similarly, we define D 0 , ≤ ℓ and D 0 ,>ℓ corresponding to Φ 0 , ≤ ℓ ∈ R n × m and Φ 0 ,>ℓ ∈ R n × ( L 0 -m ) , respectively. Here we denote by m the number of all Hermite polynomials h k, I with k ≤ ℓ where m ≫ n .

Lemma 10. Under the assumptions of Theorem 3, we have

This lemma shows the concentration of the Gram matrix since E [ Φ ⊤ 0 , ≤ ℓ Φ 0 , ≤ ℓ ] = I m . The proof of this lemma can be derived from the concentration result [Zhi21]; see also Lemma 11 in [GMMM21]. The following two lemmas are the crucial ingredients of the proof of Theorem 3, which have been established in Lemmas 13 and 14 in [GMMM21] for the spherical setting.

Lemma 11. Under the assumptions of Theorem 3 with 1 -1 ℓ < β < 1 -1 ℓ +1 for some ℓ ∈ N , as n, d →∞ proportionally, we have

Proof. Recall the definition of Φ in (A.6). From Lemma 7, we can obtain that M p = Φ D 4 Φ ⊤ due to the orthogonal polynomial decomposition in (A.4). Hence,

where Φ >ℓ is defined similarly as Φ 0 ,>ℓ which is formed by all orthogonal components of Φ 0 , ≤ ℓ . In general, Φ >ℓ is composed by Φ 0 ,>ℓ and ¯ Φ defined in (A.7). Therefore,

Recall that H p = Φ D 2 Φ ⊤ . By Lemma 7 we know that with high probability

For the second part, we let A := Φ >ℓ D 4 >ℓ Φ ⊤ >ℓ + λ I n and apply the Sherman-Morrison-Woodbury formula to obtain

where

Following the proof of Lemma 13 in [GMMM21], we can show that ∥ ∥ A -1 -I n ∥ ∥ = o d, P (1) and ∥ B -I n ∥ = o d, P (1) due to assumption that 1 -1 ℓ < ℓ and Lemma 10.

Lemma12. Recall the definition of Φ 0 , ≤ ℓ above. Under the assumptions of Theorem 3 with 1 -1 ℓ < β < 1 -1 ℓ +1 for some ℓ ∈ N , when n, d →∞ proportionally, we have

The proof of Lemma 12 is analogous to Lemma 11, the details of which we omit. From (A.3), let us denote by

Now we aim to control T 1 and T 2 based on 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 1 + y 2 , where y 1 = Φ 0 , ≤ ℓ P ≤ ℓ α and y 2 = Φ 0 ,>ℓ P >ℓ α . Here P ≤ ℓ projects onto the Hermite components with I = ( i 1 , 0 , . . . , 0) and degree no greater ℓ , and P >ℓ = P -P ≤ ℓ . Similarly, we decompose E p as E p = E p, 1 + E p, 2 , where we denote

Lemma 13. Under the assumptions of Theorem 3, we have

Proof. For simplicity, we first control T 2 and then apply some of these results to bound T 1 . Let T 2 = T 21 +2 T 22 + T 23 , where

Based on Lemmas 9, 11 and (A.8), we know that

For T 23 , following from (23) in [GMMM21], when d ℓ +1 eff < 1 , by Markov's inequality, we have

since polynomials in Φ 0 , ≤ ℓ are orthogonal to polynomials in Φ 0 ,>ℓ . Notice that

This finishes the proof for T 2 .

Now let us consider T 1 = T 11 + T 12 + T 13 , where

From Lemmas 9 and 12, we directly have that

Applying (A.9) from above, we have T 12 ≤ T 23 ∥ P ≤ ℓ f ∗ ∥ 2 L 2 = o d, P (1) . Lastly, we know that

based on the effective dimension d eff = d 1 -β . This completes the proof of this lemma.

Combining all the above lemmas together, Theorem 3 can be established by considering (A.3) and the following approximation

where P >ℓ is defined by in Section 3.

B Analysis of Neural Network

For the two-layer neural network, we perform the gradient-based optimization procedure outlined in Algorithm 1. This algorithm is almost identical to the two-stage representation learning procedure analyzed in recent works [BES + 22, DLS22, BEG + 22], with the only difference being the additional normalization step (B.2) to handle the anisotropy of the input data.

B.1 Hermite Expansion of Population Gradient

We first consider one parameter vector at random initialization w i i . i . d . ∼ N (0 , d -1 I d ) . In the following, we restrict ourselves to parameters in a subset defined as

where the positive scalar τ = O (polylog( d )) will be appropriately selected in the sequel. Note that due to the Gaussian initialization, setting τ = Θ(1) ensures that a constant fraction of neurons w i falls into the set E τ . The gradient of w i with respect to the population squared loss can be written as

Due to our specific choice of parameter initialization (4.1), we know that the output of the student model at initialization f (0) NN can be ignored (see Lemma 18 for justification). We, therefore, focus on

Algorithm 1: Gradient-based training for two-layer ReLU neural network.

Hyperparameters: learning rate η , ℓ regularization strength λ

Normalized gradient descent on 1st layer: Given i.i.d. training examples { ( x , y ) } n

2 j j j =1 .

Ridge regression for 2nd layer: Given i.i.d. training examples { (˜ x j , ˜ y j ) } n j =1 .

the following term capturing the correlation between the i -th student neuron and the teacher model:

where ( i ) is due to Stein's lemma. Note that ∥ ∥ ∥ 1 √ 1+ θ Σ 1 / 2 µ ∥ ∥ ∥ = 1 . Therefore, if ∥ Σ 1 / 2 w i ∥ ≈ 1 (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]):

where α b i , α ∗ are the Hermite coefficients of the (shifted) student and teacher nonlinearities:

The following lemma controls the error in ( i ) due to the norm fluctuation of w i .

Lemma 14. Define ˜ w i = w i / ∥ Σ 1 / 2 w i ∥ for i ∈ E ( τ ) , then we have

Note that ∥ f ′ ∗ ( x ) ∥ L 2 = E [ σ ′ ∗ ( ξ ) 2 ] 1 / 2 = O (1) , where ξ ∼ N (0 , 1) .

which concludes the proof.

Similarly, for the second term, we have

where ( i ) can be justified using the same argument from Lemma 14, the details of which we omit.

Random bias units 'diversify' the gradient update. Importantly, the shift introduced by the bias term b i ensures almost all neurons will have a non-vanishing Hermite coefficient at the given degree, as shown by the following lemma.

̸

Lemma 15. Given degree p ∈ N , for all b ∈ R such that h p -2 ( b ) = 0 when p ≥ 2 , we have that the p -th Hermite coefficient of f ( z ) = σ ( z + b ) is non-zero, i.e., ∣ ∣ E z ∼N (0 , 1) [ σ ( z + b ) h p ( z )] ∣ ∣ > 0 .

Proof. First note that when p = 1 , we have

For b ≥ 2 , we divide the computation based on the parity of p . When p is even, i.e., p = 2 s , for any b ∈ R , the Hermite coefficients are given by

for all s ∈ N . As for odd p = 2 s +1 , the odd-order Hermite coefficients are given by

Therefore, as long as b is not the zero of the Hermite polynomial of degree p -2 , the Hermite coefficients of ReLU( x + b ) would be nonzero.

̸

Approximation of population gradient. Note that for fixed p ∈ N , the condition h p -2 ( b ) = 0 excludes less than p possible values of bias b i . Combining (B.4) and (B.5), we have the following characterization of the correlation term (B.3). Now given k ∈ N as the information exponent of σ ∗ , we restrict ourselves to b i such that the corresponding degreek Hermite coefficient of the shifted ReLU α b i k = E z ∼N (0 , 1) [ σ ( z + b i ) h k ( z )] is at least of order 1 log C d for some constant C > 0 . In Section B.3.1 we show that such a condition is satisfied by a sufficiently large set of b i .

Proposition 16. Given fixed k ∈ N and θ ≍ d β where β ∈ [0 , 1) , consider neurons w i where the index i ∈ E τ for τ = O ( log C d ) ensures that ∣ ∣ ∣ α b i k ∣ ∣ ∣ = Ω ( 1 log C d ) . Then we have

Where ∆ ≲ 1 d 1 / 2 + τ d 1 -β is defined in Lemma 14.

Proof. Since the link function σ ∗ has information exponent k , by definition, we know that α ∗ i = 0 for all i < k . Also, since α b i k = Ω ( 1 log C d ) , under event E i ( τ ) with τ = O ( log C d ) , the expected

gradient is dominated by the term involving α ∗ k in the summation. For ( I ) we have

The computation of ∥ ( II ) ∥ follows from the exact same procedure. Finally, the proof is complete by including the norm fluctuation error ∆ in Lemma 14.

Proposition 16 establishes that ( I ) is the dominating term in the correlation g ( w i ) for any β > 1 -1 k as shown in the following corollary.

Corollary 17. Under the same assumptions as Proposition 16, if β > 1 -1 k , we have

for some small constant ε > 0 .

B.2 Finite-sample Concentration

Given n training examples { ( x i , y i ) } n i =1 , denote the matrix of training data X ∈ R n × d , and the vector of labels y ∈ R n . Similarly, we denote the initial first-layer parameters as W 0 ∈ R d × N , b 0 ∈ R N , a 0 ∈ R n . Wefirst show that the initial gradient (B.1) is dominated by the correlation term (B.3).

Lemma 18. When n, d, N →∞ proportionally, there exist some constants c, C > 0 such that

Proof. Using the matrix & vector notation,

Since σ ′ is bounded, we know that

To control the RHS, we first recall that the operator norm of X is bounded as follows,

for some constants c, C > 0 . Next, we consider the event that the norm of each training example x j can be controlled as follows: B := {∣ ∣ ∣ ∥ x j ∥ / √ d -1 ∣ ∣ ∣ ≤ 1 / 2 , j ∈ [ n ] } . Recall that x j = Σ 1 / 2 z j ; hence ∥ x j ∥ 2 ≤ ∥ z j ∥ 2 + (1 + θ ) |⟨ z j , µ ⟩| 2 , which is clearly dominated by ∥ z j ∥ 2 when β < 1 . Consequently, [Ver18, Theorem 3.1.1] implies that

Under event B , we know that f j := ⟨ a 0 , σ ( W ⊤ 0 x j + b 0 ) ⟩ = ∑ N k =0 a j σ ( ⟨ x j , w k ⟩ + b k ) is a sum of independent sub-exponential random variables, where the sub-exponential norm can be bounded as ∥ f j ∥ ψ 1 ≤ ∥ a j ∥ ψ 2 ∥ σ ( ⟨ x j , w k ⟩ + b k ) ∥ ψ 2 ≲ N -1 / 2 conditioning on event B . Therefore, for each j ∈ [ n ] , we have

Combining (B.6), (B.7) and taking a union bound over all f j in (B.8) establishes the lemma.

The following proposition establishes the finite-sample concentration of the correlation term g ( w i ) .

Proposition 19. Define the empirical correlation term

and recall the population counterpart

Then for any fixed w i and b i , and any p -degree polynomial σ ∗ with information exponent k , if β ∈ [0 , 1) , we have

for some constant c > 0 . Moreover, when i ∈ E τ and α b i k = Ω ( 1 log C d ) , we have

for any 1 > β > 1 -1 k , where d eff := d 1 -β .

Proof. Define the projection P := I d -µµ ⊤ , we have the following decomposition

where Gaussian random vectors Pz i and ( µ ⊤ z i ) µ are centered and independent with each other. With this decomposition, we further decompose g n ( w i ) by

Therefore, for any t > 0 and q ∈ (0 , 1) , we have

In the following, we present the concentration inequalities of I 1 and I 2 separately.

Now we first consider the concentration of I 2 . Notice that for any s > 0 ,

where function f : R d → R is defined by f ( z ) := µ ⊤ z σ ∗ ( µ ⊤ z ) σ ′ ( w ⊤ Σ 1 / 2 z + b ) . Based on [Sam23], we next estimate the Orlicz norm of f ( z i ) . Recall that the α -exponential Orlicz norm of a random variable is defined by

For any m ∈ N and z ∼ N (0 , I d ) , since σ ∗ is any polynomial with degree at most p , we have

where the last inequality is due to Gaussian hypercontractivity [MMM21, Lemma 20]. Thus, random variable f ( z ) has a finite Orlicz norm of order α := 2 p +1 : there exists some universal constant C > 0 such that ∥ f ( z ) ∥ ψ α ≤ C . Then by Theorem 1.5 of [GSS21], we can obtain that

for any s > 0 and some universal constant c > 0 .

Next, we prove the concentration of I 1 . Let u i := Pz i σ ∗ ( µ ⊤ z i ) σ ′ ( w ⊤ Σ 1 / 2 z i + b ) for i ∈ [ n ] . Consider a matrix A ∈ R d × n where the i -th column of A is ¯ u i := u i -E u i for i ∈ [ n ] . Then, for any s > 0 we have,

which implies that it is sufficient to analyze the concentration of ∥ A ∥ . Notice that A is composed of independent centered columns ¯ u i with a population covariance S := E [¯ u i ¯ u ⊤ i ] . We can check that

for some universal constant C > 0 . Meanwhile,

where the inequality is due to the sub-Gaussian tail of the norm of the Gaussian random vector. Applying Theorem 3.1.1 of [Ver18], we know that ∥ z i ∥ ≤ c ′ √ d almost surely, for some universal constant c ′ > 0 . Since σ ∗ is a polynomial with degree p , we know that ∥ σ ∗ ( ξ ) ∥ ψ α is uniformly bounded by some constant for α = 1 /p and ξ ∼ N (0 , 1) . Hence, | σ ∗ ( ξ ) | ≤ log 2 p ( n ) almost surely. Together with (B.13), we can conclude that ∥ ¯ u i ∥ ≤ C log 2 p ( n ) √ d almost surely as n →∞ . Therefore, by Theorem 5.44 in [Ver10], we have that with probability at least 1 -n exp( -cs 2 ) ,

Combining (B.12) and choosing q = 1 2 and t = log 4 p ( n ) √ d n in (B.11), one can easily obtain (B.9) for any β ∈ [0 , 1) . As a corollary of (B.9), applying Proposition 16, we can directly obtain (B.10). This completes the proof of the proposition.

B.3 Learning f ∗ with Neural Network Features

Univariate approximation. We first adapt an argument from [DLS22] showing that a ReLU nonlinearity can approximate a polynomial function in one dimension with the aid of random bias units. Lemma 20 (Corollary 3 in [DLS22], adapted) . Let σ ( z ) = ReLU( z ) , s ∼ Unif( {-1 , 1 } ) and b ∼ N (0 , 1) . Then given twice-differentiable f : R → R , there exists some υ ( s, b ) such that for any | z | ≤ C with some constant C > 0 ,

Proof. Let v ( s, b ) = e c 2 / 2 √ 2 π 1 -e -3 c 2 / 2 1 b ∈ [ c, 2 c ] , for any c > 0 . Then we have

for all | z | ≤ c . Now let v ( s, b ) = s f ( c ) 1 b ∈ [ c, 2 c ] , for any c > 0 , where f ( c ) := ∫ 2 c c 1 √ 2 π e -b 2 / 2 db . For similar reasons, we may verify that

for all | z | ≤ c . Let v ( s, b ) = 2 √ 2 πf ′′ ( -sb ) e b 2 / 2 1 b ∈ [0 ,c ] , for any c > 0 , where f ∈ C 2 ( R ) . Then we can easily check that

Notice that this construction of v ( s, b ) 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 are optimized for one (normalized) gradient step, there exist some second-layer coefficients that achieve small prediction risk with high probability.

Proposition 21. Given fixed k ∈ N and β > 1 -1 k , let W 1 be the weight matrix optimized via one normalized gradient descent step defined in Algorithm 1, there exists some second-layer coefficients ˜ a ∈ R N such that ˜ f ( x ) = 1 √ N σ ( W ⊤ 1 x + b )˜ a = 1 √ N ∑ N i =1 ˜ a i σ ( a i ⟨ x , w (1) i ⟩ + b i ) achieves

for some large constant C > 0 , with probability 1 -exp ( -c log 2 N ) when n, d →∞ proportionally and N = Ω( d ϵ ) .

Proof. We follow an argument similar to [BES + 22, Appendix D] by selecting a suitable set of neurons that can approximate f ∗ . Throughout this proof, we take ε > 0 to be some small nonvanishing constant and C > 0 to be some large constant. At a high level, we select the neurons following the univariate approximation result in Lemma 20 to fit a polynomial target function. To this end, we need to show that after one gradient descent step, there exists a sufficiently large set of neurons w (1) i that align with the signal direction, i.e., ∣ ∣ ∣ √ 1 + θ ⟨ w (1) i , µ ⟩ ∣ ∣ ∣ ≈ 1 .

We first restrict ourselves to the 'average-case' neurons around the equator w i where the indices i ∈ E τ . Due to the Gaussian initialization, setting τ = log 2 d entails P ( i / ∈ E τ ) ≲ exp ( -log 2 d ) .

̸

Now we consider the subset of bias units that gives an 'informative' gradient update, i.e., the k -th Hermite coefficient is not vanishing: α b i k = Ω ( 1 log C d ) for some large constant C > 0 . Recall that Lemma 15 implies that for bias unit satisfying h p -2 ( b i ) = 0 , the degreek Hermite coefficient of f ( z ) = σ ( z + b i ) is non-zero. Therefore, we define the index set

Notice that all the zeros of Hermite polynomials are real and distinct, see Theorem 3.3.1 in [Sze39]. Additionally, from Theorem 3.3.2 in [Sze39], the zeros of h p +1 ( x ) are interlacing between the zeros of h p ( x ) . From the recurrence relation of Hermite polynomials, namely,

we know that h ′ p ( x ) = 0 if x is a root of h p . Thus, applying the implicit function theorem, we can deduce that by choosing ε ≍ 1 log C d for some large C > 0 , we attain α b i k = Ω ( 1 log C d ) , for any i ∈ B ε . Due to the Gaussian initialization of b i , we know what the expected size of the set E |B ε | = Ω( N 1 -ε ′ ) for some small ε ′ > 0 . Therefore, Hoeffding's inequality implies that for large enough N = Ω( d ϵ ) , we have ∣ ∣ ∣ |B ε | E |B ε | -1 ∣ ∣ ∣ ≤ N -1 / 2+ ε ′ with probability 1 -exp ( -c log 2 N ) .

for all | z | ≤ c , where

̸

Gradient approximation. For the subset of neurons with indices i ∈ B ε , we introduce a sequence of approximations to the gradient update given as

where we dropped the factor a i √ N for convenience. First recall that by Lemma 18, we can ignore the initial NN output as follows:

with probability 1 -exp ( -c log 2 d ) for some constant c > 0 . Note that since β < 1 , the RHS can always be upper-bounded by d -ε for some small ε > 0 . Next, we take into account the concentration and truncation error in Propositions 19 and 16, respectively. Define the population gradient (again omitting the scalar a i √ N ),

(B.10) in Proposition 19 implies that for any fixed p ∈ N , β > 1 -1 k , and large enough n, d , with probability 1 -exp ( -c log 2 d ) , we have

for some ε > 0 . As for the population counterpart, Corollary 17 implies that for neurons with indices i ∈ E τ ∩ B ε , if β > 1 -1 k , then we have

where ∥ ∥ ∥ g ( w (0) i ) ∥ ∥ ∥ ≍ (1 + θ ) k 2 ∣ ∣ ∣ ∣ 〈 w (0) i , µ 〉 k -1 ∣ ∣ ∣ ∣ = ˜ Θ ( (1 + θ ) k 2 d 1 -k 2 ) due to the condition i ∈ E τ for τ = log C d . Combining (B.15), (B.16), and (B.17), we know that the gradient vector g n ( w (0) i ) achieves significant alignment with the target direction µ when β > 1 -1 k .

Normalization step. Now we consider the normalization procedure (B.2). When β > 1 -1 k , (B.15), (B.16), and (B.17) imply that for neurons w i in the subset i ∈ E τ ∩ B ε , we have

where ∥ ∆ ∥ ≲ d -ε for some ε > 0 , with probability 1 -exp ( -c log 2 d ) . Now recall that we select a sufficiently large learning rate η = Ω( N 1 / 2+ ε ) for some small ε > 0 . Since a i ∼ N (0 , 1 /N ) ,

where κ i = ηa i α b i k α ∗ k √ 1 + θ 〈 √ 1 + θ µ , w (0) i 〉 k -1 = ˜ Θ ( N ε (1 + θ ) k 2 d 1 -k 2 ) by the choice of τ ≍ log C d and α b i k ≍ 1 log C d . Define the normalization factor

by (B.18) we know that

with high probability, where we used the proportional scaling of n, d in the last inequality. Now note that 1 √ n ∥ Xµ ∥ → √ 1 + θ almost surely when n, d →∞ proportionally, we have

Combining the above calculations and recalling N = Ω( d ϵ ) , we arrive at the following characterization of the normalized weight vector:

with probability 1 -exp ( -c log 2 d ) . Observe that (B.19) entails that after the normalization step:

the updated weight vector with index i ∈ E τ ∩B ε approximates the target direction in the sense that ∥ ∥ ∥ √ 1 + θ · w (1) i -sign( κ i ) µ ∥ ∥ ∥ ≲ d -ε with high probability, for some small ε > 0 .

Approximation of link function. Now we construct certain second-layer coefficients ˜ a i on a subset of neurons with indices i ∈ I such that the (sub)network approximates the link function σ ∗ . Roughly speaking, we wish to find some collection of neurons that utilizes the univariate approximation result in Lemma 20:

First, we consider the ideal setting where w i is exactly in the direction of µ up to sign flip, that is, we may write w i = s i (1 + θ ) -1 / 2 µ for s i ∼ Unif ( { +1 , -1 } ) . This allows us to reparameterize ⟨ x , (1 + θ ) -1 / 2 µ ⟩ =: z ∈ R , and choose the second-layer coefficients ˜ a according to Lemma 20 on domain | z | ≤ log log d . This implies that | υ ( s, b ) | = O (exp(log log d )) = O (log d ) . The corresponding 'infinite-width' neural network is written as

From Lemma 20 we know that for | z | ≤ log log d , we have f ∞ ( z ) = σ ∗ ( z ) .

Nowwecontrol the difference between the ideal model and the finite-width network of interest. First we restrict ourselves to b ∈ B ′ ε defined as a 'continuous' analogue of (B.14): B ′ ε = { b ∣ ∣ | h p -2 ( b ) | ≥ ε, if p ≥ 2 } , with ε ≍ 1 log C d for large C > 0 , and write the corresponding infinite-width model as

By the Lipschitz property of σ , we have | f ∞ ( z ) -f ε ∞ ( z ) | ≲ sup | υ ( s, b ) | · 1 log C d = o d (1) .

Next, we bound the error due to truncation at the interval | z | ≤ log log d . Recall that ⟨ x , (1 + θ ) -1 / 2 µ ⟩ =: z ∼ N (0 , 1) ; hence we have the following evaluation

To upper-bound the above expectation, observe that

for large C > 0 , where we used the fact that σ ∗ is a degreep polynomial with bounded L 2 norm in the last inequality. On the other hand, by Cauchy-Schwarz we have

Therefore, by the Lipschitz property of σ and the fact that sup | υ ( s, b ) | ≍ log d ,

Combining (B.20), (B.21), and (B.22) yields E z ∼N (0 , 1) [ ( σ ∗ ( z ) -f ε ∞ ( z ) ) 2 ] ≲ 1 log C d = o d (1) for some large C > 0 . Now consider the finite-width model

where s i i . i . d . ∼ Unif ( { +1 , -1 } ) and b i i . i . d . ∼ N (0 , 1) . Similar to above, we set the second-layer ˜ a according to Lemma 20, that is, ˜ a i = √ N |I| υ ( s i , b i ) .

To establish the concentration of ˆ f around f ε ∞ , we first restrict | z | ≤ log d . Let ϕ ( s, b ) = σ ( sz + b ) . Then ∥∇ ϕ ( s, b ) ∥ ∞ ≤ 1 + log d . Hence, applying Lipschitz concentration inequality for s and b , we have ∥ ϕ ( s, b ) ∥ ψ 2 ≤ 1 + log d . Since uniformly ∥ v ( s, b ) ∥ ∞ ≤ log d , we know that the subexponential norm ∥ v ( s, b ) ϕ ( s, b ) ∥ ψ 1 ≤ 2 log 2 d . Thus, utilizing the Bernstein's inequality [Ver18, Theorem 2.8.1], we know that

for any t > 0 . By construction, we can construct a subset of neurons I ⊂ E τ ∩B ε with |I| = Θ( N ) , with probability at least 1 -exp ( -c log 2 d ) . Since we choose N = Ω( d ϵ ) , we may take t ≍ 1 log C d for large C , and a union bound yields ∣ ∣ ∣ ˆ f ( z ) -f ε ∞ ( z ) ∣ ∣ ∣ ≲ 1 log C d with probability 1 -exp ( -c log 2 d ) as d, N → ∞ , under the event that | z | ≤ log d . Then by the tail bound of the standard Gaussian random variable z , we have

with probability 1 -exp ( -c log 2 d ) when d, N →∞ .

To control the deviation from the ideal setting, we consider neurons w (1) i where i ∈ E τ ∩ B ε for τ ≍ log C d and ε ≍ 1 log C d . Recall that for this subset of neurons, with at least probability 1 -exp ( -c log 2 d ) we have ∥ ∥ ∥ √ 1 + θ · w (1) i -sign( κ i ) µ ∥ ∥ ∥ ≲ d -ε , where sign ( κ i ) ∼ Unif ( { +1 , -1 } ) due to the Gaussian initialization of a i (note that w i , a i , b i are independent). Therefore,

for small constant ε > 0 , where we used the Lipschitz property of σ in ( i ) .

Finally, recall that with high probability we can construct I such that |I| = |E τ ∩ B ε | ≍ N . This implies the following ℓ 2 -norm control of the chosen coefficients,

This entails that the RHS of (B.24) can be bounded as d log N = o (1) .

Combining (B.20), (B.23), (B.24), and (B.25), we know that there exists some second-layer coefficients ˜ a such that as n, d, N →∞ , there exists some some large constant C > 0 such that

with probability 1 -exp ( -c log 2 d ) . Note here here we simply set ˜ a i = 0 for i / ∈ E τ ∩ B ε .

B.3.2 Ridge Regression Estimator

Proposition 21 suggests there exists some second-layer coefficients with small ℓ 2 norm can achieve small prediction risk. Based on 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 an argument following [BES + 22, Appendix D.3].

Given feature map x → 1 √ N σ ( W ⊤ 1 x + b ) conditioned on first-layer weights W 1 , b , we denote the associated Hilbert space as H . We define the optimal predictor in the RKHS as ˇ f := argmin f ∈H E x ( f ( x ) -f ∗ ( x )) 2 , which takes the form of ˇ f ( x ) = 〈 1 √ N σ ( W ⊤ 1 x + b ) , ˇ a 〉 for some ˇ a ∈ R N . We have the orthogonal decomposition of f ∗ in L 2 :

Given fresh training data { (˜ x i , ˜ y i ) } n i =1 , it is known that the prediction risk of ridge regression admits the following decomposition:

where ϕ x := 1 √ N σ ( W ⊤ 1 x + b ) for x ∈ R d , which gives Φ ⊤ = [ ϕ ⊤ ˜ x 1 , . . . , ϕ ⊤ ˜ x i , . . . , ϕ ⊤ ˜ x n ] , where ˜ x i is the i -th row of ˜ X ; we denoted ̂ Σ Φ := 1 n Φ ⊤ Φ , Σ Φ := 1 N E x [ σ ( W ⊤ 1 x + b ) σ ( W ⊤ 1 x + b ) ⊤ ] . Also, the i -th entry of vector ˇ f and f ⊥ are given by [ ˇ f ] i = ˇ f (˜ x i ) , [ f ⊥ ] i = f ⊥ (˜ x i ) , respectively.

Concentration of feature covariance. We first establish the concentration of the kernel matrix in the proportional regime. The following lemma provides an estimate of the norm of the features.

Lemma22. Given the weight matrix after one gradient step W 1 as specified in Algorithm 1, then we have ∥ ∥ ∥ Σ 1 / 2 W 1 ∥ ∥ ∥ F ≲ √ N, with probability 1 -exp ( -c log 2 d ) , as n, d, N →∞ , n/d → ψ > 1 .

Proof. Write v i = Σ 1 / 2 ( w (0) i -ηa i · g n ( w (0) i )) , by definition x = Σ 1 / 2 z for z ∼ N (0 , I d ) , and

which implies that ∥ ∥ ∥ Σ 1 / 2 w (1) i ∥ ∥ ∥ 2 ≤ λ -1 min ( 1 n Z ⊤ Z ) ≲ 1 almost surely as n, d → ∞ , n/d > 1 . Summing over w (1) i yields the desired claim.

Observe that the feature vector ϕ x = 1 √ N σ ( W ⊤ 1 x + b ) = 1 √ N σ ( W ⊤ 1 Σ 1 / 2 z + b ) is a sub-Gaussian vector with parameter 1 √ N ∥ ∥ ∥ Σ 1 / 2 W 1 ∥ ∥ ∥ F ≲ 1 due to the Lipschitzity of σ ; hence we may apply the following sub-Gaussian concentration result.

Lemma 23. Under the same assumption as Lemma 22, as d →∞ , we have ∥ ϕ x ∥ ≤ log d almost surely.

Proof. Notice that ∥·∥ and σ ( · ) are 1-Lipschitz convex functions. Following [WZ21, Section 3], we can utilize the convex concentration inequality (see [Ver18, Chapter 5]) to obtain

for any t > 0 , conditioning on the event that ∥ ∥ ∥ Σ 1 / 2 W 1 ∥ ∥ ∥ ≤ ∥ ∥ ∥ Σ 1 / 2 W 1 ∥ ∥ ∥ F ≲ √ N . From Lemma 22, we know this event that ∥ ∥ ∥ Σ 1 / 2 W 1 ∥ ∥ ∥ F ≲ √ N occurs with probability at least 1 -exp( -c log 2 d ) . Furthermore, under this event, we have

Thus, we can take t = log d and combine Lemma 22 and (B.26) to attain the assertion.

Lemma 24 ([Min17], adapted) . Under the same assumption as Lemma 22, given nλ ≥ c log 2 d for some constant c > 0 , we have

Proof. Let us denote

Notice that E [ G i ] = 0 . From Lemma 23, we know that ∥ [ G i ∥ ≤ log 2 d/λn almost surely. Moreover, we may verify that E [ G 2 i ] ≼ 1 n log 2 d λn I N and

Here we use the fact that all the eigenvalues of ( Σ Φ + λ I ) -1 Σ Φ are between 0 and 1. Thus, we can apply Matrix Bernstein bound (see (3.9) in [Min17]) to conclude this concentration result.

With Lemma 24 we have the following useful corollary.

Corollary 25. Given nλ ≫ log 4 d , for any t ∈ (0 , 1) , for sufficiently large d ,

with probability at least 1 -n exp ( -c log 2 d ) .

Prediction risk of ridge regression. We upper-bound the prediction risk of the ridge regression estimator by controlling the decomposed error terms separately.

Proposition 26. Under the training procedure described in Algorithm 1 with N = Ω( d ε ) and d ε -1 ≲ λ ≲ d ε for small ε > 0 , if β > 1 -1 k , then we have R ( ˆ f NN ) = o d (1) with probability 1 when n, d →∞ , n/d → ψ .

Proof. First note that R 1 = E x ( f ∗ ( x ) -ˇ f ( x ) ) 2 = ∥ f ⊥ ∥ 2 L 2 by definition. For R 2 , [BES + 22, Appendix D.3] showed that R 2 ≲ ∥ f ∗ -˜ f ∥ 2 L 2 + λ ∥ ˜ f ∥ 2 H . Finally, for R 3 due to the concentration in Corollary 25, we have the following evaluation,

Recall that in Proposition 21 we constructed some ˜ f that achieves ∥ f ∗ -˜ f ∥ 2 L 2 ≲ 1 / log C d and ∥ ˜ f ∥ 2 H ≲ log C d . This implies that ∥ f ⊥ ∥ 2 L 2 ≲ 1 / log C d and hence R 1 = o d, P (1) . Similarly, for λ = O ( d -ε ) , we have R 2 = o d, P (1) . The proof is completed by Markov's inequality on R 3 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nComments:\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\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-21"}

comments: SUMMARY:
This paper considers the problem of learning a single-index target function under the spiked covariance data in the proportional asymptotic limit.
The 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.
Numerical experiments on synthetic data are also supportive.

STRENGTHS:
* The paper is clear and easy to follow.

* 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.

* Although both upper and lower bounds for KRR are established, this paper only provides upper bound for NN.

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?

LIMITATIONS:
The authors adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1- 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\n\nWEAKNESSES:\n1- 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.\n\nQUESTIONS:\n1- 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?\n\nLIMITATIONS:\n1- 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, ...", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nS1. 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.\n\nWEAKNESSES:\nW1. 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations (and extensions) have been discussed in the conclusion section.\n\nPotential negative societal impact would be unexpected given the technical nature of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}