Title: Gradient-Based Feature Learning under Structured Data
Venue: NeurIPS 2023
Accepted: yes
Keywords: feature learning

neural networks

single-index model

gradient descent
Abstract: Recent works have demonstrated that the sample complexity of gradient-based learning of single index models, i.e. functions that depend on a 1-dimensional projection of the input data, is governed by their information exponent. However, these results are only concerned with isotropic data, while in practice the input often contains additional structure which can implicitly guide the algorithm. In this work, we investigate the effect of a spiked covariance structure and reveal several interesting phenomena. First, we show that in the anisotropic setting, the commonly used spherical gradient dynamics may fail to recover the true direction, even when the spike is perfectly aligned with the target direction. Next, 
we show that appropriate weight normalization that is reminiscent of batch normalization can alleviate this issue. Further, by exploiting the alignment between the (spiked) input covariance and the target, we obtain improved sample complexity compared to the isotropic case. In particular, under the spiked model with a suitably large spike, the sample complexity of gradient-based training can be made independent of the information exponent while also outperforming lower bounds for 
rotationally invariant kernel methods.

Gradient-Based Feature Learning under Structured Data

Abstract

Recent works have demonstrated that the sample complexity of gradient-based learning of single index models, i.e. functions that depend on a 1-dimensional projection of the input data, is governed by their information exponent . However, these results are only concerned with isotropic data, while in practice the input often contains additional structure which can implicitly guide the algorithm. In this work, we investigate the effect of a spiked covariance structure and reveal several interesting phenomena. First, we show that in the anisotropic setting, the commonly used spherical gradient dynamics may fail to recover the true direction, even when the spike is perfectly aligned with the target direction. Next, we show that appropriate weight normalization that is reminiscent of batch normalization can alleviate this issue. Further, by exploiting the alignment between the (spiked) input covariance and the target, we obtain improved sample complexity compared to the isotropic case. In particular, under the spiked model with a suitably large spike, the sample complexity of gradient-based training can be made independent of the information exponent while also outperforming lower bounds for rotationally invariant kernel methods.

1 Introduction

A fundamental feature of neural networks is their adaptivity to learn unknown statistical models. For instance, when the learning problem exhibits certain low-dimensional structure or sparsity, it is expected that neural networks optimized by gradient-based algorithms can efficiently adapt to such structure via feature/representation learning. A considerable amount of research has been dedicated to understanding this phenomenon under various assumptions and to demonstrate the superiority of neural networks over non-adaptive methods such as kernel models [GMMM19, WLLM19, BES + 19, LMZ20, AAM22, BES + 22, DLS22, Tel23, MHPG + 23].

A particular relevant problem setting for feature learning is the estimation of single index models, where the response y ∈ R depends on the input x ∈ R d via y = g ( ⟨ u , x ⟩ ) + ϵ , where g : R → R is the nonlinear link function and u is the unit target direction. Here, learning corresponds to recovering the unknowns u and g , which requires the model to extract and adapt to the low-dimensional target direction. Recent works have shown that the sample complexity is determined by certain properties of the link function g . In particular, the complexity of gradient-based optimization is captured by the information exponent of g introduced by [BAGJ21]. Intuitively, a larger information exponent s corresponds to a more complex g (for gradient-based learning), and it has been proven that when the input is isotropic x ∼ N (0 , I d ) , gradient flow can learn the single index model with ˜ O ( d s ) sample complexity [BBSS22].

In practice, however, real data always exhibits certain structures such as low intrinsic dimensionality, and isotropic data assumptions fail to capture this fact. In statistics methodology, it is known that the directions along which the input x has high variance are often good predictors of the target y [HTFF09]; indeed, this is the main reason principal component analysis is used in pretraining [JWHT13]. A fundamental model that captures such a structure is the spiked matrix model in which x ∼ N (0 , I d + κ θθ ⊤ ) for some unit direction θ ∈ R d and κ > 0 [Joh01]. Along the direction θ , data has higher variability and predictive power. In single index models, such predictive power translates to a non-trivial alignment between the vectors u and θ -our focus is to investigate the effect of such alignment on the sample complexity of gradient-based training.

1.1 Contributions: learning single index models under spiked covariance

In this paper, we study the sample complexity of learning a single index model using a two-layer neural network and show that it is determined by an interplay between

spike-target alignment : ⟨ u , θ ⟩ ≍ d -r 1 , r 1 ∈ [0 , 1 / 2] ,

spike magnitude : κ ≍ d r 2 , for r 2 ∈ [0 , 1] .

Our contributions can be summarized as follows.

We show that even in the case of perfect spiketarget alignment ( r 1 = 0 ), the spherical gradient flow commonly employed in recent literature (see e.g. [BAGJ21, BBSS22]) cannot recover the target direction for moderate spike magnitudes in the population limit. The failure of this covariance-agnostic procedure under anisotropic structure insinuates the necessity of an appropriate covariance-aware normalization to effectively learn the single index model.

Hard

n

∝

d

s

r

2

(spike magnitude)

Intermediate

r

1

(spike-target alignment)

Easy

sample complexity

High

Low

n

∝

d

1+(

s

-

1)(1

-

r

2

)

n

∝

d

1+(

s

-

1)(1

-

2(

r

2

-

r

1

))

r

1

=

r

2

2

r

1

=

r

2

We show that a covariance-aware normalization that resembles batch normalization resolves this issue. Indeed, the resulting gradient flow can successfully recover the target direction u in this case, and depending on the amount of spike-target alignment, the sample complexity can significantly improve compared to the isotropic case.

Under the spiked covariance model, we prove a three-stage phase transition for the sample complexity depending on the quantities r 1 and r 2 . For a suitable direction and magnitude of the spike, the sample complexity can be made ˜ O ( d 3+ ν ) for any ν > 0 which is independent of the information exponent s . This should be compared against the known complexity of ˜ O ( d s ) under isotropic data.

We finally show that preconditioning the training dynamics with the inverse covariance improves the sample complexity. This is particularly significant for the spiked covariance model where ˜ O ( d 3+ ν ) samples can be reduced to ˜ O ( d 1+ ν ) for any ν > 0 , i.e. almost linear in d . The three-stage phase transition also emerges, as illustrated in Figure 1: in the 'hard' regime, the complexity remains ˜ O ( d s ) regardless of the magnitude and direction of the spike, while in the 'easy' regime the complexity only depends on the spike magnitude and not its direction. The 'intermediate' regime interpolates between these two; smaller r 1 and larger r 2 improve the sample complexity.

The rest of the paper is organized as follows. We discuss the notation and the related work in the remainder of this section. We provide preliminaries on the statistical model and the training procedure in Section 2, and provide a negative result on the covariance-agnostic gradient flow in Section 2.1. Our main sample complexity result on a single neuron is presented in Section 3.2. We provide our results on multi-neuron neural networks in Section 4 and also discuss extensions such as preconditioning and its implications. We provide a technical summary in Section 5 and conclude in Section 6.

Notation. We use ⟨· , ·⟩ and ∥·∥ to denote Euclidean inner product and norm. For matrices, ∥·∥ denotes the usual operator norm, and λ max ( · ) and λ min ( · ) denote the largest and smallest eigenvalues respectively. We reserve γ for the standard Gaussian distribution on R , and let ∥·∥ γ denote the L 2 ( γ ) norm. S d -1 is the unit d -dimensional sphere. For quantities a and b , we will use a ≲ b to convey there exists a constant C (a universal constant unless stated otherwise, in which case may depend on polylogarithmic factors of d ) such that a ≤ Cb , and a ≍ b signifies that a ≲ b and b ≲ a .

<latexi sh

1_b

64="jfyN

8

uY

zJ

ZH

+

/XDO

Bw

>A

C

c

F

M

EI

d5

r

S

g

R

q

L

G

U

K

2

m

v

W

o0TP

n

V

3

Qk

9

7

p

<latexi sh

1_b

64="

T

o

RNywfcd

C2

V

9j

J

>A

q3

H

M

E

W

X

LI

0

v

B

gpU

Q

G

K

/Z

P8

7r

O

z

m

u

Y

5

D

S

+

n

k

F

<latexi sh

1_b

64="GvK

8m5O

WDIpAXcBy

Q

0

Y

>

C

3

dVF

M2

9

fR

n

S

LJz

NZE

7

P

/

q

o

H

+

U

T

j

k

r

<latexi sh

1_b

64="

GJO3wn9

poN

2RPk uj

Y

/U

c

>A

M

EHW

r

K+

y

Q

B

g0

L

S

7Z

T

m

v8

5

I

z

f

X

D

<latexi sh

1_b

64="0

5UZ+YK

9Lw fJ

k

Ty

Qo

>A

Cqn cdVFNj

M

EHXD

u

I

r

R

p

gS

3

Wm

z

/

B

v2

7

P

G

O

8

<latexi sh

1_b

64="y

8K

Cvf

0Sq

+knYGE

Fr

U

D

>A

3

c

V

L

M

H

W

wm

I5

NJ

gd

B

R7

u

j

T

9

Q

X

2

P

O

/

p

Z

o

z

<latexi sh

1_b

64="5

qFYjz2DXH+

k/g

0

>A

C

cdV

9N

3Gr7

K

BE

SU

RW

O

Qy

TM

n

v

u

o

P

Z

L

8

<latexi sh

1_b

64="pFg

Gq0kMO

v+o

>A

Cun

H

E

W

wmX

I

RT

K

L

j

2

N

/

Z

7

3

J

Y

P

5

f

B

D

U

y

<latexi sh

1_b

64="FKgR9

YH

3

PTG0

/

rk

>A

C

n

cdV

M

E

7

+

f

Q

W

U

JuN

O

52

B

w

X

j

q

m

v

z

D

L

S

8

I

y

<latexi sh

1_b

64="QE

2Fc qYg

SB

jN

9n

7urGdU

>A

C

X

VH

O

5

p+M3

Z

T

W

P

f

v

0

z

D8

/

o

J

w

y

mK

I

R

<latexi sh

1_b

64="uSGA

O

cU

fW

0X

52j

K

p+n

g

>

MwE

D

m3

k

J

q

L

N

R

9

I

F/

T

B

Z

z

P

8o

v

7

y

Q

Y

<latexi sh

1_b

64="pm3g0z5L

N

vQM

RY

c

7

I

>A

E

W

w

F

U

K

B

PSy

n

2

/

D

J

ZOj

o

+9

u

f

G

8

X

k

<latexi sh

1_b

64="/O

5

L

CnR

8T+uW

mzS

P2

A

>

q3

cdVH

M

E

w

F

Y

Q

U

K

B

rj

k

7

yp

f

N

0

G

Z

9

JI

v

D

g

X

o

1.2 Further related work

Non-Linear Feature Learning with Neural Networks. Recently, two popular scaling regimes of neural networks have emerged for theoretical studies. A large initialization variance leads to the lazy training regime, where the weights do not move significantly, and the training dynamics is captured by the neural tangent kernel (NTK) [JGH18, COB19]. However, there are many instances of function classes that are efficiently learnable by neural networks and not efficiently learnable by the NTK [YS19, GMMM19]. Under a smaller initialization scaling, gradient descent on infinite-width neural networks becomes equivalent to Wasserstein gradient flow on the space of measures, known as the mean-field limit [CB18, RVE18, MMN18, MMM19, NWS22, Chi22], which can learn certain low-dimensional target functions efficiently [WLLM19, AAM22, HC22, ASKL23].

As for neural networks with smaller width, recent works showed that a two-stage feature learning procedure can outperform the NTK when the data is sampled uniformly from the hypercube [BEG + 22] or isotropic Gaussian [DLS22, BES + 22, BBSS22, MHPG + 23, ABAM23]. However, these results do not take into account the additional structure that might be present in the covariance matrix of the input data. Two notable exceptions are [GMMM20, RGKZ21], where the authors analyzed a spiked covariance and Gaussian mixture data, respectively. Our setting is closer to [GMMM20], however, they do not provide optimization guarantees through gradient-based training. Furthermore, in a companion work [BES + 23], we zoom into the setting where the spike and target are perfectly aligned ( r 1 = 0 ), and prove learnability in the n ≍ d regime for both kernel regression and two-layer neural network. Finally, we go over some results concurrent to our work in Appendix A.

Learning Single Index Models. The problem of estimating the relevant direction in a single index model is classical in statistics [LD89], with efficient dedicated algorithms ([KKSK11, CM20] among others). However, these algorithms are non-standard and instead, we are concerned with standard iterative algorithms like training neural networks with gradient descent. Recently, [DH18] considered an iterative optimization procedure for learning such models with a polynomial sample complexity that is controlled by the smoothness of the link function. [BES + 22] considered the effect of taking a single gradient step on the ability of a two-layer neural network to learn a single index model, and [BBSS22, MHPG + 23] considered training a special two-layer neural network architecture where all neurons share the same weight with gradient flow or online SGD. However, these works only consider the isotropic Gaussian input, and the effect of anisotropy in the covariance matrix when training a neural network to learn a single index model has remained unclear.

Training a Single Neuron with Gradient Descent. When training the first layer, we consider a setting where there is only one effective neuron. A large body of works exists on training a single neuron using variants of gradient descent. In the realizable setting (i.e. identical link and activation), the typical assumptions on the activation correspond to information exponent 1 as the activations are required to be monotone or have similar properties, see e.g. [Sol17, YO20, DKTZ22]. In the agnostic setting, [FCG20] considered initializing from the origin which is a saddle point for information exponent larger than 1. [ATV22] also considered the agnostic learning of a ReLU activation, albeit their sample complexity is not explicit other than being polynomial in dimension.

2 Preliminaries: Statistical Model and Training Procedure

For a d -dimensional input x and a link function g ∈ L 2 ( γ ) , consider the single index model

where ϵ is a zero-mean noise with O (1) sub-Gaussian norm and u ∈ S d -1 . Learning the model (2.1) corresponds to approximately recovering the unknown link g and the unknown direction u . Note that a normalization is needed to make this problem well-defined; without loss of generality, we write ⟨ u , x ⟩ / ∥ Σ 1 / 2 u ∥ to ensure that the input variance and the scaling of g both remain independent of the conditioning of Σ . For this learning task, we will use a two-layer neural network of the form

where W = { w i } i m =1 is the m × d matrix whose rows corresponds to first-layer weights w i , a = { a i } i m =1 denote the second-layer weights, b = { b i } i m =1 denote the biases, and ϕ is the non-linear

activation function. We assume g and ϕ are weakly differentiable with weak derivatives g ′ and ϕ ′ respectively, and g, g ′ , ϕ, ϕ ′ ∈ L 2 ( γ ) . We are interested in the high-dimensional regime; thus, d is assumed to be sufficiently large throughout the paper. Our ultimate goal is to learn both unknowns g and u by minimizing the population risk

using a gradient-based training method such as gradient flow.

We follow the two-step training procedure employed in recent works [BES + 22, MHPG + 23, BBSS22, DLS22]: First, we train the first-layer weights W to learn the unknown direction u ; at the end of this stage, the neurons w i align with u . Here, the goal is to recover only the direction. Next, using random biases and training the second-layer weights, we obtain a good approximation for the unknown link function g . In the majority of this work, we focus on the first part of this two-stage procedure as the alignment between w i 's and u essentially determines the sample complexity of the overall procedure. This problem is somewhat equivalent to the simplified problem of minimizing (2.3) with m = 1 , a 1 = 1 , b 1 = 0 , i.e., ˆ y ( x ; W , a , b ) is replaced with ˆ y ( x ; w ) := ϕ ( ⟨ w , x ⟩ ) and we write R ( w ) := R ( W , a , b ) for simplicity. We emphasize that unless ϕ = g (i.e. the link function is known), the first stage of training only recovers the relevant direction u and is not able to approximate g . Indeed, m> 1 is often needed to learn the non-linear link function; this is the focus of Section 4.2 where we derive a complete learnability result for a two-layer neural network with m> 1 .

Characteristics of the link function play an important role in the complexity of learning the model. As such, a central part of our analysis will rely on a particular property based on the Hermite expansion of functions in a basis defined by the normalized Hermite polynomials { h j } j ≥ 0 given as

These polynomials form an orthonormal basis in the space L 2 ( γ ) , and the resulting expansion yields the following measure of complexity for g , which is termed as the information exponent .

Definition 1 (Information exponent) . Let g = ∑ j ≥ 0 α j h j be the Hermite expansion of g . The information exponent of g is defined to be s := inf { j > 0 : α j = 0 } .

̸

This concept was introduced in [BAGJ21] in a more general framework, and our definition is more in line with the setting in [BBSS22]. We remark that the definition of [BAGJ21] can be modified to handle anisotropy in which case one arrives at Definition 1. We provide a detailed discussion on this concept together with some properties of the Hermite expansion in Appendix B. Throughout the paper, we assume that the information exponent does not grow with dimension.

In the case where the d -dimensional input data is isotropic, [BBSS22] showed that learning a single index target with full-batch gradient flow requires a sample complexity of ˜ O ( d s ) for s ≥ 3 where s is the information exponent of g . We will show that this sample complexity can be improved under anisotropy. More specifically, if the input covariance Σ has non-trivial alignment with the unknown direction u , we prove in Section 3 that the resulting sample complexity can be even made independent of the information exponent if we use a certain normalization in the training. In what follows, we prove that such a normalization in training procedure is indeed necessary.

2.1 The spiked model and limitations of covariance-agnostic training

In practice, data often exhibit a certain structure which may have a profound impact on the statistical procedure. A well-known model that captures such a structure is the spiked model [Joh01] for which one or several large eigenvalues of the input covariance matrix Σ are separated from the bulk of the spectrum (see also [BBAP05, BS06]). Although our results hold for generic covariance matrices, they reveal interesting phenomena under the following spiked model assumption.

Assumption 1. The covariance Σ follows the ( κ, θ ) -spiked model if Σ = I d + κ θθ ⊤ 1+ κ where ∥ θ ∥ = 1 .

In pursuit of the target (unit) direction u , the magnitude of the neuron w is immaterial; thus, recent works take advantage of this and simplify the optimization trajectory by projecting w onto unit sphere S d -1 throughout the training process [BAGJ21, BBSS22]. In the sequel, we study the same dynamics

which is agnostic to the input covariance in order to motivate our investigation of normalized gradient flow in Section 3. More specifically, we consider the spherical population gradient flow

where ∇ S is the spherical gradient at the current iterate. It is straightforward to see that when the initialization w 0 is on the unit sphere, the entire flow will remain on the unit sphere, i.e. w t ∈ S d -1 for all t ≥ 0 . The flow (2.5) has been proven useful for learning the direction u [BBSS22] in the isotropic case Σ = I d when the activation ϕ is ReLU. In contrast, when Σ follows a spiked model, we show that it can get stuck at stationary points that are almost orthogonal to u . Indeed, when the input covariance Σ has a spike in the target direction u , i.e. θ = u , one expects that the training procedure benefits from this as the input x contains information about the sought unknown u without even querying the response y . The following result proves the contrary; for moderate spike magnitudes, the alignment between the first-layer weights and target ⟨ w t , u ⟩ will be insignificant for all t .

Theorem 2. Let s > 2 be the information exponent of g with E [ g ] = 0 , and assume Σ follows the ( κ, u ) -spiked model with Ω(1) ≤ κ ≤ O ( d s -2 s -1 ) . For ReLU activation, let w t denote the solution to (2.5) initialized uniformly at random over S d -1 , then with probability at least 0 . 99 ,

A non-trivial alignment between the first-layer weights w t and the target direction u is required to learn the single index model (2.1). However, the above result implies that in high dimensions when d ≫ 1 , the alignment is negligible in the population limit (when the number of samples goes to infinity). We remark that when the spike magnitude is large, i.e. κ ≥ Ω( d ) , the flow (2.5) can achieve alignment as the problem essentially becomes one-dimensional, as we demonstrate in Appendix C.

To see why the flow (2.5) gets stuck at saddle points and fails to recover the true direction, notice that

If the input was isotropic, i.e. x ∼ N (0 , I d ) , the first term in (2.7) would be equal to ∥ ϕ ∥ 2 γ , which is independent of w . Thus, minimizing R ( w ) in this case is equivalent to maximizing the 'correlation' term E [ ϕ ( ⟨ w , x ⟩ ) y ] . However, under the spiked model, the alignment between w and u breaks the symmetry; consequently, the first term in the decomposition grows with ⟨ w , u ⟩ , creating a repulsive force that traps the dynamics around the equator where w is almost orthogonal to u .

3 Main Results: Alignment via Normalized Dynamics

Having established that the covariance-agnostic training dynamics (2.5) is likely to fail, we consider a covariance-aware normalized flow in this section and show that it can achieve alignment with the unknown target and enjoy better sample complexity compared to the existing results [BAGJ21, BBSS22] in the isotropic case. We start with the population dynamics.

3.1 Warm-up: Population dynamics

To simplify the exposition, we define z := Σ -1 / 2 x , w := Σ 1 / 2 w / ∥ Σ 1 / 2 w ∥ and similarly define u , and consider the prediction function ˆ y ( x ; w ) := ϕ ( ⟨ w , z ⟩ ) . Due to symmetry, the second moment of the prediction is E [ ˆ y ( x ; w ) 2 ] = ∥ ϕ ∥ 2 γ which is independent of w ; thus, the population risk reads

In (3.1), the only term that depends on the weights w is the correlation term and the source of the repulsive force in (2.7) is eliminated; we have ∇ w R ( w ) = -∇ w E [ ϕ ( ⟨ w , z ⟩ ) y ] . Based on this, we use the following normalized gradient flow for training

We remark that, though not identical, this normalization is closely related to batch normalization which is commonly employed in practice [IS15]. Under the invariance provided by the current normalization, minimizing R ( w ) corresponds to maximizing E [ ϕ ( ⟨ w , z ⟩ ) y ] . Thus, instead of w , it will be more useful to track the dynamics of its normalized counterpart w , which is made possible by the following intermediary result that follows from Stein's lemma; also see e.g. [EDB16, MHPG + 23].

Lemma 3. Suppose we train w t using the gradient flow (3.2) . Then w t solves the following ODE

We will investigate if the modified flow (3.3) achieves alignment; in this context, alignment corresponds to 〈 w t , u 〉 ≈ 1 . Towards that end, we make the following assumption.

Assumption 2. Let g = ∑ j ≥ 0 α j h j and ϕ = ∑ j ≥ 0 β j h j be the Hermite decomposition of g and ϕ respectively. Let s be the information exponent of g . For some universal constant c > 0 , we assume

There are several important examples that readily satisfy Assumption 2. The obvious example is when the link function is known as in [BAGJ21], i.e. ϕ = g . A more interesting example is when ϕ is an activation with degree s non-zero Hermite coefficient (e.g. ReLU when s is even, see [GKK19, Claim 1]) and g is a degree s Hermite polynomial, which for s = 2 corresponds to the phase retrieval problem. In this case, the assumption is satisfied if α s and β s have the same sign, which occurs with probability 0 . 5 if we randomly choose the sign of the second layer.

Under this condition, the following result shows that the population flow (3.3) can achieve alignment.

Proposition 4. Suppose Assumption 2 holds and consider the gradient flow given by (3.3) with initialization satisfying 〈 w 0 , u 〉 > 0 . Then, we have 〈 w T , u 〉 ≥ 1 -ε as soon as

We remark that the information exponent enters the rate in (3.4) through the function τ s , and time needed to achieve ε alignment gets worse with larger information exponent. Indeed, it is understood that this quantity serves as a measure of complexity for the target function being learned.

3.2 Empirical dynamics and sample complexity

Given n i.i.d. samples { ( x ( i ) , y ( i ) ) } n i =1 from the single index model (2.1), we consider the flow

where we estimate the covariance matrix Σ using the sample mean ˆ Σ := 1 n ′ ∑ n ′ i =1 x ( i ) x ( i ) ⊤ over n ′ i.i.d. samples; the above dynamics defines an empirical gradient flow. Notice that we ignored the gradient associated with the term ϕ 2 since the population dynamics ensures that its gradient will concentrate around zero; thus, it is redundant to estimate this term. Below, we will use n ′ = n for smooth activations, i.e. the same dataset can be used for covariance estimation; For ReLU, we require a more accurate covariance estimator, thus, we use n ′ ≳ n 2 by assuming access to an additional n ′ -n unlabeled data points. Similar to the previous section, we track the dynamics of normalized w by defining w := ˆ Σ 1 / 2 w / ∥ ˆ Σ 1 / 2 w ∥ (and leave u unchanged from Section 3.1). The same arguments as in Lemma 3 allow us to track the evolution of w , which ultimately yields the following alignment result under general covariance structure.

Theorem 5. Let s be the information exponent of g , and assume it satisfies | g ( · ) | ≲ 1 + |·| p for some p > 0 . For ϕ denoting either the ReLU activation or a smooth activation satisfying | ϕ ′ | ∨ | ϕ ′′ | ≲ 1 , suppose Assumption 2 holds. For any ε > 0 , suppose we run the finite sample gradient flow (3.5) with η ( w ) = ∥ ˆ Σ 1 / 2 w ∥ 2 , initialized such that 〈 w 0 , u 〉 > 0 , and with number of samples

where κ ( Σ ) is the condition number of Σ . Then, for T τ s ( ⟨ w 0 , u ⟩ ) +ln(1 /ε ) λ ( Σ ) , we have

with probability at least 1 -c 1 d -c 2 for some universal constants c 1 , c 2 > 0 over the randomness of the dataset. Here, τ s is defined in (3.4) and ≳ hides poly-logarithmic factors.

Remark. We make the following remarks on the above theorem.

The initial condition 〈 w 0 , u 〉 > 0 is required when we have odd information exponent. When w 0 is initialized uniformly over S d -1 , the condition holds with probability 0 . 5 over the initialization. See [BAGJ21, Remark 1.8] for further discussion on this condition.

Although w is defined using the empirical covariance unlike u which is defined by population covariance, this definition is the suitable choice to approximate the target function g (c.f. Theorem 9), since it ensures the arguments of ϕ and g are sufficiently close when w recovers u .

The intuition behind the proof of Theorem 5 is presented in Section 5 with the complete proof in the appendix. We highlight that the improvement in the sample complexity compared to the isotropic setting occurs whenever the covariance structure induces a stronger initial alignment and consequently stronger signal. The following corollary demonstrates a concrete example of such improvement by specializing Theorem 5 for a spiked covariance model.

Corollary 6. Consider the setting of Theorem 5 with Σ following the ( κ, θ ) -spiked model, where ⟨ u , θ ⟩ ≍ d -r 1 and κ ≍ d r 2 with r 1 ∈ [0 , 1 / 2] and r 2 ∈ [0 , 1] . Suppose w 0 is sampled uniformly from S d -1 . Then, when conditioned on 〈 w 0 , u 〉 > 0 , the sample complexity in Theorem 5 reads

where ≳ hides poly-logarithmic factors of d .

Remark. We have the following observations on the above sample complexity.

Corollary 6 demonstrates that structured data can lead to better sample complexity when the right normalization is used during training. This complements Theorem 2 where we recall that spherical training dynamics ignores the structure in data and the target direction cannot be recovered.

When g is a polynomial of degree p , the lower bound for rotationally invariant kernels (including the neural tangent kernel at initialization) implies a complexity of at least d Ω((1 -r 2 ) p ) [DWY21]. Thus the sample complexity of Corollary 6 can always outperform the kernel lower bound when p is sufficiently large and s remains constant.

Three-step phase transition. Recall that in the isotropic setting Σ = I d , the sample complexity of learning g with information exponent s using full-batch gradient flow is ˜ O ( d s ) for s ≥ 3 [BBSS22]. The sample complexity in Corollary 6 is strictly smaller than ˜ O ( d s ) as soon as ( s -1) r 1 / ( s -2) < r 2 . Furthermore, for any ν > 0 it is at most ˜ O ( d 3+ ν ) as soon as r 2 ≥ 1 -ν/ ( s -3) and 2 r 1 < r 2 , in which case the sample complexity becomes independent of the information exponent. Interestingly, the complexity becomes independent of r 1 when r 2 > 2 r 1 or r 2 < r 1 , i.e. the direction of the spike becomes irrelevant when the spike magnitude is sufficiently large or small.

The three-stage phase transition of Corollary 6 is due to the different behaviour of the inner product 〈 w 0 , u 〉 in different regimes of r 1 and r 2 . When r 2 < r 1 , we have 〈 w 0 , u 〉 ≍ 〈 w 0 , u 〉 , thus the initial alignment is just as uninformative as the isotropic case providing no improvement. Moreover, a potentially large condition number may hurt the sample complexity in this case. On the other hand, when r 1 < r 2 < 2 r 1 we have 〈 w 0 , u 〉 ≍ κ ⟨ u , θ ⟩ 〈 w 0 , θ 〉 , and r 2 > 2 r 1 leads to 〈 w 0 , u 〉 ≍ √ κ 〈 w 0 , θ 〉 , thus large κ or ⟨ u , θ ⟩ in this regime may improve the sample complexity.

4 Implications to Neural Networks and Further Improvements

4.1 Improving Sample Complexity via Preconditioning

Wenowdemonstrate that preconditioning the training dynamics with ˆ Σ -1 can remove the dependency on κ ( Σ ) , ultimately improving the sample complexity. Consider the preconditioned gradient flow

We have the following alignment result.

Theorem 7. Consider the same setting as Theorem 5, and assume we run the preconditioned empirical gradient flow (4.1) with number of samples

where ≳ hides poly-logarithmic factors of d . Then, for T ≍ τ s (〈 w 0 , u 〉) +ln(1 /ε ) , we have

with probability at least 1 -c 1 d -c 2 for some universal constants c 1 , c 2 > 0 .

Preconditioning removes the condition number dependence, which is particularly important in the spiked model case where this quantity can be large.

Corollary 8. Consider the setting of Theorem 7, and assume we run the preconditioned empirical gradient flow (4.1) for the ( κ, θ ) -spiked model where ⟨ u , θ ⟩ ≍ d -r 1 and κ ≍ d r 2 with r 1 ∈ [0 , 1 / 2] and r 2 ∈ [0 , 1] . Suppose w 0 is sampled uniformly from S d -1 . Then, when conditioned on 〈 w 0 , u 〉 > 0 , the sample complexity of Theorem 7 reads

where ≳ hides poly-logarithmic factors of d .

The above result improves upon Corollary 6; thus, making a case for preconditioning in practice. The complexity results also strictly improve upon the ˜ O ( d s ) complexity in the isotropic case [BBSS22] when r 2 > r 1 . Further, for any ν > 0 , we can obtain the complexity of ˜ O ( d 1+ ν ) (nearly linear in dimension) when r 2 > 1 -ν/ ( s -1) and r 2 > 2 r 1 or r 1 +1 / 2(1 -ν/ ( s -1)) < r 2 < 2 r 1 . In addition to the remarks of Corollary 6, we note that the complexity is independent of both r 1 and r 2 when r 2 < r 1 (cf. Figure 1 hard regime), i.e. the spike magnitude and the spike-target alignment have no effect on the complexity unless r 2 ≥ r 1 .

Under the spiked covariance model, one could improve the above results by instead using spectral initialization, i.e. initializing at θ , which can be estimated from unlabeled data. Assuming perfect access to θ , using the statement of Theorems 5 and 7, this initialization would imply a sample complexity of ˜ O ( d 1+2 r 2 +(( s -1)(2 r 1 -r 2 ) ∨ 0) ) without and ˜ O ( d 1+(( s -1)(2 r 1 -r 2 ) ∨ 0) ) with preconditioning.

4.2 Two-layer neural networks and learning the link function

Our main focus so far was learning the target direction u . Next, we consider learning the unknown link function with a neural network, providing a complete learnability result for single index models.

We use Algorithm 1 and train the first-layer of the neural network with either the empirical gradient flow (3.5) or the preconditioned version (4.1). Then, we randomly choose the bias units and minimize the second layer weights using another gradient flow. Our goal is to track the sample complexity n needed to learn the single index target which we compare against the results of [BBSS22]. We highlight that layer-wise training in Algorithm 1 is frequently employed in the literature [BES + 22, BBSS22, DLS22, MHPG + 23] and in particular [BBSS22] also used gradient flow for training.

Algorithm 1 Layer-wise training of a two-layer ReLU network with gradient flow (GF).

Input: w 0 ∈ R d , T, T ′ , ∆ , λ ∈ R + and data { ( x ( i ) , y ( i ) ) } n i =1 .

1: Train the first layer weights W T j using the GF (3.5) or the preconditioned GF (4.1).

2: Normalize the weights W T j := W T j / ∥ ˆ Σ 1 / 2 W T j ∥ for every 1 ≤ j ≤ m .

4: Train the second layer weights a T ′ via the gradient flow

3: Let b j i.i.d. ∼ Unif ( -∆ , ∆) and a 0 j = 1 /m for 1 ≤ j ≤ m .

5: return ( W T , a T ′ , b ) .

Theorem 9. Let g be twice weakly differentiable with information exponent s and assume g ′′ has at most polynomial growth. Suppose ϕ is the ReLU activation, Assumption 2 holds and we run Algorithm 1 with w 0 initialized uniformly over S d -1 . For any ε > 0 , let n and T be chosen according to Theorem 5 when we run the gradient flow (3.5) and Theorem 7 when we run the preconditioned gradient flow (4.1) . Then, for ∆ ≍ √ ln( nd ) , some regime of λ given by (E.3) and sufficiently large T ′ given by (E.4) , we have

conditioned on 〈 w 0 , u 〉 > 0 with probability at least 0 . 99 over the randomness of the dataset, biases, and initialization, where C 1 is a universal constant and C 2 hides polylog( m,n,d ) factors.

The next result immediately follows from the previous theorem together with Corollaries 6 & 8.

Corollary 10. In the setting of Theorem 9, if Σ follows the ( κ, θ ) -spiked model, the sample complexity n is given by (3.7) if we use the empirical gradient flow and (4.2) if we use the preconditioned version.

We remark that for fixed ε , the sample complexity to learn g in the isotropic case is ˜ O ( d s ) [BBSS22]. Under the spiked model, if we assume that r 2 is sufficiently large and r 1 is sufficiently small as discussed in the previous section, Corollary 10 improves this rate to either (3.7) when the empirical gradient flow is used without preconditioning or to (4.2) with preconditioning.

5 Technical Overview

In this section, we briefly discuss the key intuitions that lead to the proof of our main results. We first review the case Σ = I d , where we have the following decomposition for population loss

Notice that the only term contributing to the population gradient is the last term which measures the correlation between ϕ and g . Following the gradient flow and applying Stein's lemma yields

where the second identity follows from the Hermite expansion; see also [EDB16, EBD19]. Assume α s β s > 0 to ensure that the population dynamics will move towards u at least near initialization. When replacing the population gradient with a full-batch gradient, we need the estimation noise to be smaller than the signal existing in the gradient. When 〈 w 0 , u 〉 ≪ 1 , this signal is roughly of the order 〈 w 0 , u 〉 s -1 . As the uniform concentration error over S d -1 scales with √ d/n , we need n ≍ d 〈 w 0 , u 〉 2( s -1) to ensure the signal remains dominant and w t moves towards u . When w 0 is initialized uniformly over S d -1 this translates to a sample complexity of n ≍ d s , which is indeed obtained by [BBSS22] via similar arguments.

However, the behavior of the spherical dynamics entirely changes when we move to the anisotropic case. Suppose Σ follows a ( κ, u ) -spiked model and ϕ is ReLU. Using Lemma 12, it is easy to show that with the spherical gradient flow, the alignment obeys the following ODE

where ˜ ψ ϕ,g ( w t ) is introduced in Lemma 12. The additional ˜ ψ ϕ,g ( w t ) term creates a repulsive force towards the equator ⟨ w t , u ⟩ = 0 . The presence of this term is due to the fact that unlike (5.1), the term E [ ϕ ( ⟨ w t , u ⟩ ) 2 ] is no longer independent of w and cannot be replaced by ∥ ϕ ∥ 2 γ . When w 0 is initialized uniformly over S d -1 and Ω(1) ≤ κ ≤ O ( d ) , we have 〈 w 0 , u 〉 ≍ √ κ 〈 w 0 , u 〉 . Furthermore, at this initialization ˜ ψ ϕ,g ( w 0 ) ≈ 1 / 2 . Therefore,

Hence the dynamics is trapped at |⟨ w t , u ⟩| = O (1 / √ d ) for all t > 0 as long as κ = O ( d 1 -1 / ( s -1) ) . To remove the repulsive force in the spherical dynamics, we can directly normalize the input of ϕ . As demonstrated by (3.1), once again the only term that varies with w would be the correlation loss. Specifically, using the result of Lemma 3, in the population limit we can track 〈 w t , u 〉 via

where u t ⊥ := u -〈 w t , u 〉 w t . Thus, the strength of the signal at initialization is of order 〈 w 0 , u 〉 s -1 / κ ( Σ ) , which after controlling the error in the estimate of ˆ Σ and in the estimate of population gradient using finitely many samples, leads to the sample complexity n ≍ d κ ( Σ ) 2 〈 w 0 , u 〉 2(1 -s ) . Importantly, Σ can incude a much stronger initial alignment 〈 w 0 , u 〉 than the isotropic case 〈 w 0 , u 〉 , which is emphasized in Corollary 6. Using preconditioning will further remove the dependency on the condition number of Σ .

6 Conclusion

We studied the dynamics of gradient flow to learn single index models when the input data covariance may contain additional structure. Under a spiked model for the covariance matrix, we showed that using spherical gradient flow, as an example of a covariance-agnostic training mechanism employed in the recent literature, is unable to learn the target direction of the single index model even when the spike and the target directions are identical. In contrast, we showed that an appropriate weight normalization removes this problem and successfully recovers the target direction. Moreover, depending on the alignment between the covariance structure and the target direction, the sample complexity can improve upon the isotropic setting, while also outperforming lower bounds for rotationally-invariant kernels. This phenomenon is due to the additional information about the target direction contained in the covariance matrix which improves the effective alignment at initialization. Additionally, we showed that a simple preconditioning of the gradient flow using the inverse empirical covariance can improve the sample complexity, achieving almost linear rate in certain settings.

We outline a few limitations of our current work and discuss directions for future research.

While studying single index models provides a pathway to a general understanding of feature learning with structured covariance, considering multi-index models can provide a more complete picture [PSE22], e.g. by establishing incremental learning dynamics [ABAM23]. We leave the problem of learning multi-index models under structured input as an interesting future direction.

Gradient flow under squared loss can be seen as an example of a Correlational Statistical Query (CSQ) algorithm [BF02, Rey20], i.e. an algorithm that only accesses noisy estimates of expected correlation queries from the model. Understanding the limitations of learning single index models under a structured input through a CSQ lower bound perspective is another important direction that would complement our results in this paper.

When training the first layer, we considered a somewhat unconventional initialization and relied on the symmetry it induces. It is interesting to consider cases where we train a network with multiple neurons starting from a more standard initialization which can help relax Assumption 2.

Acknowledgments

The authors thank Alberto Bietti and Zhichao Wang for discussions and feedback on the manuscript. TS was partially supported by JSPS KAKENHI (20H00576) and JST CREST (JPMJCR2015). MAE was partially supported by NSERC Grant [2019-06167], CIFAR AI Chairs program, CIFAR AI Catalyst grant.

References

[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz, The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks , Conference on Learning Theory, 2022.

[ABAM23] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz, Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics , arXiv preprint arXiv:2302.11055 (2023).

[ASKL23] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro, From highdimensional & mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks , arXiv preprint arXiv:2302.05882 (2023).

[ATV22] Pranjal Awasthi, Alex Tang, and Aravindan Vijayaraghavan, Agnostic learning of general relu activation using gradient descent , arXiv preprint arXiv:2208.02711 (2022).

[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath, Online stochastic gradient descent on non-convex losses from high-dimensional inference. , J. Mach. Learn. Res. 22 (2021), 106-1.

[BBAP05] Jinho Baik, G´ erard Ben Arous, and Sandrine P´ ech´ e, Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices .

[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song, Learning single-index models with shallow neural networks , Advances in Neural Information Processing Systems, 2022.

[BEG + 22] Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang, Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit , arXiv preprint arXiv:2207.08799 (2022).

[BES + 19] Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang, Generalization of two-layer neural networks: An asymptotic viewpoint , International Conference on Learning Representations, 2019.

[BES + 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 , arXiv preprint arXiv:2205.01445 (2022).

[BES + 23] Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu, Learning in the presence of low-dimensional structure: a spiked random matrix perspective , Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.

[BF02] Nader H Bshouty and Vitaly Feldman, On using extended statistical queries to avoid membership queries , Journal of Machine Learning Research 2 (2002), no. Feb, 359395.

[BMZ23] Rapha¨ el Berthier, Andrea Montanari, and Kangjie Zhou, Learning time-scales in two-layers neural networks , arXiv preprint arXiv:2303.00055 (2023).

[BPVZ23] Joan Bruna, Loucas Pillaud-Vivien, and Aaron Zweig, On single index models beyond gaussian data , arXiv preprint arXiv:2307.15804 (2023).

[BS06] Jinho Baik and Jack W Silverstein, Eigenvalues of large sample covariance matrices of spiked population models , Journal of multivariate analysis 97 (2006), no. 6, 1382-1408.

[CB18] Lenaic Chizat and Francis Bach, On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport , Advances in Neural Information Processing Systems, 2018.

[Chi22] L´ ena¨ ıc Chizat, Mean-field langevin dynamics: Exponential convergence and annealing , arXiv preprint arXiv:2202.01009 (2022).

[CM20] Sitan Chen and Raghu Meka, Learning polynomials in few relevant dimensions , Conference on Learning Theory, 2020.

[COB19] Lenaic Chizat, Edouard Oyallon, and Francis Bach, On Lazy Training in Differentiable Programming , Advances in Neural Information Processing Systems, 2019.

[CWPPS23] Elizabeth Collins-Woodfin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi, Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models , arXiv preprint arXiv:2308.08977 (2023).

[DH18] Rishabh Dudeja and Daniel Hsu, Learning single-index models in gaussian space , Conference On Learning Theory, PMLR, 2018, pp. 1887-1930.

[DKL + 23] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan, Learning two-layer neural networks, one (giant) step at a time , arXiv preprint arXiv:2305.18270 (2023).

[DKTZ22] Ilias Diakonikolas, Vasilis Kontonis, Christos Tzamos, and Nikos Zarifis, Learning a single neuron with adversarial label noise via gradient descent , Conference on Learning Theory, PMLR, 2022, pp. 4313-4361.

[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi, Neural Networks can Learn Representations with Gradient Descent , Conference on Learning Theory, 2022.

[DNGL23] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee, Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models , arXiv preprint arXiv:2305.10633 (2023).

[DWY21] Konstantin Donhauser, Mingqi Wu, and Fanny Yang, How rotational invariance of common kernels prevents generalization in high dimensions , International Conference on Machine Learning, 2021.

[EBD19] Murat A. Erdogdu, Mohsen Bayati, and Lee H. Dicker, Scalable Approximations to Generalized Linear Problems , Journal of Machine Learning Research (2019).

[EDB16] Murat A Erdogdu, Lee H Dicker, and Mohsen Bayati, Scaled least squares estimator for glms in large-scale problems , Advances in Neural Information Processing Systems 29 (2016).

[Erd15] Murat A Erdogdu, Newton-stein method: a second order method for glms via stein's lemma , Proceedings of Advances in Neural Information Processing Systems, 2015, pp. 1216-1224.

[FCG20] Spencer Frei, Yuan Cao, and Quanquan Gu, Agnostic learning of a single neuron with gradient descent , Advances in Neural Information Processing Systems, vol. 33, Curran Associates, Inc., 2020, pp. 5417-5428.

[GKK19] Surbhi Goel, Sushrut Karmalkar, and Adam Klivans, Time/accuracy tradeoffs for learning a relu with respect to gaussian marginals , Advances in neural information processing systems 32 (2019).

[GMMM19] B. Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari, Limitations of Lazy Training of Two-layers Neural Networks , Advances in Neural Information Processing Systems, 2019.

[GMMM20] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari, When Do Neural Networks Outperform Kernel Methods? , Advances in Neural Information Processing Systems, 2020.

[HC22] Karl Hajjar and Lenaic Chizat, On the symmetries in the dynamics of wide two-layer neural networks , arXiv preprint arXiv:2211.08771 (2022).

[HTFF09] Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman, The elements of statistical learning: data mining, inference, and prediction , vol. 2, Springer, 2009.

[IS15] Sergey Ioffe and Christian Szegedy, Batch normalization: Accelerating deep network training by reducing internal covariate shift , International conference on machine learning, pmlr, 2015, pp. 448-456.

[JGH18] Arthur Jacot, Franck Gabriel, and Clement Hongler, Neural Tangent Kernel: Convergence and Generalization in Neural Networks , Advances in Neural Information Processing Systems, 2018.

[Joh01] Iain M Johnstone, On the distribution of the largest eigenvalue in principal components analysis , The Annals of statistics 29 (2001), no. 2, 295-327.

[JWHT13] Gareth James, Daniela Witten, Trevor Hastie, and Robert Tibshirani, An introduction to statistical learning , vol. 112, Springer, 2013.

[KKSK11] Sham M Kakade, Varun Kanade, Ohad Shamir, and Adam Kalai, Efficient learning of generalized linear and single index models with isotonic regression , Advances in Neural Information Processing Systems 24 (2011).

[LD89] Ker-Chau Li and Naihua Duan, Regression Analysis Under Link Violation , The Annals of Statistics (1989).

[LMZ20] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang, Learning over-parametrized twolayer neural networks beyond NTK , Conference on Learning Theory, 2020.

[MHD + 23] Arvind Mahankali, Jeff Z Haochen, Kefan Dong, Margalit Glasgow, and Tengyu Ma, Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time , arXiv preprint arXiv:2306.16361 (2023).

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

[MMM19] Song Mei, Theodor Misiakiewicz, and Andrea Montanari, Mean-field theory of twolayers neural networks: dimension-free bounds and kernel limit , Conference on Learning Theory, 2019.

[MMN18] Song Mei, Andrea Montanari, and Phan-Minh Nguyen, A mean field view of the landscape of two-layer neural networks , Proceedings of the National Academy of Sciences 115 (2018), no. 33, E7665-E7671.

[NWS22] Atsushi Nitanda, Denny Wu, and Taiji Suzuki, Convex analysis of the mean field langevin dynamics , International Conference on Artificial Intelligence and Statistics, PMLR, 2022, pp. 9741-9757.

[O'D14] Ryan O'Donnell, Analysis of boolean functions , Cambridge University Press, 2014.

[PSE22] Sejun Park, Umut Simsekli, and Murat A. Erdogdu, Generalization Bounds for Stochastic Gradient Descent via Localized ε -Covers , arXiv preprint arXiv:2209.08951 (2022).

[Rey20] Lev Reyzin, Statistical queries and statistical algorithms: Foundations and applications , arXiv preprint arXiv:2004.00557 (2020).

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

[RVE18] Grant M Rotskoff and Eric Vanden-Eijnden, Neural networks as Interacting Particle Systems: Asymptotic convexity of the Loss Landscape and Universal Scaling of the Approximation Error , arXiv preprint arXiv:1805.00915 (2018).

[Sol17] Mahdi Soltanolkotabi, Learning relus via gradient descent , Advances in neural information processing systems 30 (2017).

[SWON23] Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda, Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond , Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.

[Tel23] Matus Telgarsky, Feature selection and low test error in shallow low-rotation relu networks , The Eleventh International Conference on Learning Representations, 2023.

[Ver18] Roman Vershynin, High-dimensional probability: An introduction with applications in data science , Cambridge University Press, 2018.

[VH16] Ramon Van Handel, Probability in high dimension , 2016.

[Wai19] Martin J. Wainwright, High-dimensional statistics: A non-asymptotic viewpoint , Cambridge University Press, 2019.

[WLLM19] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma, Regularization matters: Generalization and optimization of neural nets vs their induced kernel , Advances in Neural Information Processing Systems 32 (2019).

[YO20] Gilad Yehudai and Shamir Ohad, Learning a single neuron with gradient methods , Proceedings of Thirty Third Conference on Learning Theory, Proceedings of Machine Learning Research, vol. 125, PMLR, 2020, pp. 3756-3786.

[YS19] Gilad Yehudai and Ohad Shamir, On the Power and Limitations of Random Features for Understanding Neural Networks , Advances in Neural Information Processing Systems, 2019.

A Concurrent Works

In this paragraph we briefly summarize a few relevant results concurrent to our submission. [BMZ23] provided a precise analysis of the two-timescale dynamics in learning a single index model with information exponent s = 1 . [MHD + 23] considered the learning of a single index target with s = 2 using a neural network in the mean-field regime. [BPVZ23] extended the information exponentbased characterization of online SGD to input data beyond Gaussian. [DNGL23] showed that a gradient-smoothed dynamics can improve the sample complexity and match the CSQ lower bound. Finally, beyond the single-index setting, [DKL + 23, SWON23, CWPPS23] considered learning low-dimensional target functions supported on k > 1 dimensions via gradient-based feature learning.

B Background on Hermite Expansion

The normalized Hermite polynomials { h j } j ≥ 0 given by (2.4) provide an orthonormal basis for L 2 ( γ ) , thus for every f ∈ L 2 ( γ ) we have

where ⟨ f, h j ⟩ γ := E z ∼N (0 , 1) [ f ( z ) h j ( z )] . We will commonly invoke the following well-known properties of Hermite polynomials. If j ≥ 1 , then h ′ j = √ jh j -1 , where h ′ j stands for the derivative of h j . Furthermore, if z 1 and z 2 are two standard Gaussian random variables with E [ z 1 z 2 ] = ρ , then E [ h i ( z 1 ) h j ( z 2 )] = δ ij ρ j where δ ij is the Kronecker delta. We refer the interested reader to [O'D14, Chapter 11.2] for additional discussions and properties of these polynomials.

We will now discuss how our Definition 1 relates to the original definition of information exponent of [BAGJ21]. In their setting, they assume the true data distribution P u is parameterized by some unit vector u ∈ S d -1 , and we know the parametric family { P w } w ∈ S d -1 ; thus the problem is to estimate the direction u . Furthermore, they assume the population loss, which is the expectation of some per-sample loss, has spherical symmetry, i.e. the population loss R ( w ) can be written as R ( w ) = ˜ R ( ⟨ w , u ⟩ ) . Then, [BAGJ21, Definition 1.2] defines the information exponent to be the degree of the first non-zero coefficient of ˜ R in its Taylor expansion around the origin. In other words, we say R has information exponent s if

where C, c > 0 are universal constants. To specialize the above abstract definition to the Gaussian case, consider the setting where the input data is standard Gaussian x ∼ N (0 , I d ) and the problem is to estimate u ∈ S d -1 given a response variable y = f ( ⟨ u , x ⟩ ) with known f . Via the Hermite expansion of f , one can write

Thus, the information exponent of ˜ R is indeed the degree of the first non-zero term in the Hermite expansion of f .

Now consider the general case where x ∼ N (0 , Σ ) . The spherical symmetry assumed in [BAGJ21] no longer holds. However, after proper normalization of weights, if we consider the population loss

then R ( w ) = ˜ R ( ⟨ w , Σ u ⟩ ∥ Σ 1 / 2 w ∥∥ Σ 1 / 2 u ∥ ) . Indeed, a close examination of the arguments of [BAGJ21] reveals that for their results to hold, the proper symmetry to consider is the ellipsoidal symmetry, and

the proper definition of information exponent is the degree of the first non-zero term in the Hermite expansion of ˜ R , which reads

Once again, we can consistently define the information exponent to be the degree of the first non-zero term in the Hermite expansion of f , as long as the input is Gaussian (potentially anisotropic).

C Proofs of Section 2.1

Before beginning our main discussions, we state the following lemma which is a generalization of Stein's lemma (Gaussian integration by parts), and will help obtain a closed-form expression for the population gradient. We refer to [Erd15, MHPG + 23] for similar statements.

Lemma 11. Let f, g : R → R with g weakly differentiable. Suppose z ∼ N (0 , I d ) . Then, for any w , u ∈ S d -1 , we have

Proof. Consider the conditional distribution z |⟨ w , z ⟩ ∼ N ( µ , Σ ) , where

Recall that Stein's lemma (Gaussian integration by parts) states that when z ∼ N ( µ , Σ ) , then

Hence,

Applying the tower property of conditional expectation and rearranging the terms yields the desired result.

We are now ready to state and prove the expression for the population gradient when using the ReLU activation.

Lemma 12. Suppose ϕ is the ReLU activation and g is weakly differentiable. Let z ∼ N (0 , I d ) . Define w := Σ 1 / 2 w ∥ Σ 1 / 2 w ∥ and similarly define u . Then

where and

Proof. Notice that the population risk is given by

Notice that x = Σ 1 / 2 z . By the homogeneity of ReLU, we can rewrite the first term as

Then we have,

where we used Lemma 11 and the fact that ϕ ′ ( z ) z = ϕ ( z ) . Therefore,

which concludes the proof.

Particularly, the above lemma yields the following corollary for the spherical dynamics in the population limit.

Corollary 13. Suppose { w t } t ≥ 0 is a solution to the population spherical gradient flow (2.5) , ϕ is the ReLU activation, and Σ follows the ( κ, θ ) -spiked model. Then,

C.1 Proof of Theorem 2

Plugging in θ = u in Corollary 13, we obtain

To prove the statement of the theorem, we will show that whenever

we will have d ⟨ w t , u ⟩ 2 d t < 0 , thus when initialized from O (1 / √ d ) , ⟨ w t , u ⟩ can never escape the saddle point near the equator ⟨ w , u ⟩ = 0 .

Recall from the properties of the Hermite expansion in Appendix B that

Since we additionally assume E [ g ] = α 0 = 0 , by the definition of ˜ ψ ϕ,g in Lemma 12 and the properties of Hermite expansion discussed in Appendix B, we have

and similarly

Thus we obtain where

We proceed by upper bounding F . To do so, first note that

To bound the first term of F , we have

Similarly,

Hence, for κ ≥ 1 ,

Suppose κ < d/C 2 -1 , then

Thus, for any κ such that

Finally, we will show ∣ ∣ 〈 w 0 , u 〉∣ ∣ ≤ C/ √ d with probability at least 0 . 99 for a suitable choice of constant C . Indeed, this is an elementary concentration of measure result on the unit sphere. For simplicity, we avoid performing sharp probability of failure analysis and only remark that E [ 〈 w 0 , u 〉 2 ] = 1 /d , thus by the Markov inequality and any w t such that |⟨ w t , u ⟩| ≤ C/ √ d , we have d ⟨ w t , u ⟩ 2 d t ≤ 0 , hence sup t ≥ 0 |⟨ w t , u ⟩| ≤ C/ √ d, as long as the above holds true at initialization.

hence a choice of C ≥ 10 suffices, and the proof is complete.

C.2 Extremely Large Spike

In this section, we will show that under extremely large spike, the spherical gradient flow (2.5) can potentially recover the true direction. Namely, we will prove the following proposition.

Proposition 14. Suppose we initialize the spherical population gradient flow (2.5) from w 0 . Let ϕ be the ReLU activation and assume

and κ ≥ C ⟨ w 0 , u ⟩ 2 for a sufficiently large constant C > 0 depending only on g . Then, the gradient flow on the sphere satisfies

Before proceeding to the proof, we notice that if we uniformly initialize w 0 over S d -1 , then the typical value for 〈 w 0 , u 〉 is of order d -1 / 2 , meaning that the above proposition asks for κ = Ω( d ) . This is a regime where lower bounds for the sample complexity of kernel methods are Ω(1) [DWY21], thus no meaningful separation in terms of dimension dependency of the sample complexity between neural networks and kernel methods is possible, as the problem becomes effectively one-dimensional.

Proof. The cases where 〈 w 0 , u 〉 > 0 and 〈 w 0 , u 〉 < 0 are symmetric, thus we only present the proof for the former. Using (C.2), we can write the dynamics on the sphere more explicitly as

Our goal is to study the regime of large κ , therefore we will bound how much B 1 and B 2 can deviate from their corresponding κ = ∞ values. In particular, we have

Furthermore, assuming ⟨ w t , u ⟩ > 0 and ⟨ ϕ, g ⟩ γ > 0 , let c κ ( w t ) := κ ⟨ w t , u ⟩ √ 1+ κ √ 1+ κ ⟨ w t , u ⟩ 2 . Then, by the Lipschitzness of ϕ ,

where we used w t = Σ 1 / 2 w t ∥ Σ 1 / 2 w t ∥ in the last step. Suppose ⟨ w t , u ⟩ > 0 (which holds at least on a neighborhood around initialization, and as we will see below holds for all t > 0 ), then,

whenever

As a result, we obtain

Consequently, the lower bound of the time derivative of ⟨ w t , u ⟩ becomes larger as ⟨ w t , u ⟩ increases. Therefore, assuming 〈 w 0 , u 〉 > 0 , we only need to control this lower bound at initialization. Assume

From this, we conclude that when ⟨ w t , u ⟩ > 0 and ⟨ ϕ, g ⟩ γ = α > 1 / 2 , we have

integration yields the desired result.

D Proofs of Section 3

We begin by stating the closed-form expression for the population gradient, i.e. the counterpart of Lemma 12 in the normalized setting.

Lemma 15. Consider the population risk R ( w ) defined by (3.1) , recall that

Then,

where

Proof. Recall from (3.1) that

where

(the above is only a function of w u due to the Hermite expansion).

⟨ , ⟩

Given the closed form of the population gradient, the proof of Lemma 3 is immediate by noticing that

Next, we move on to prove Proposition 4.

D.1 Proof of Proposition 4

where u t ⊥ := u -〈 w t , u 〉 w t . The above inequality and the fact that 〈 w 0 , u 〉 > 0 imply that 〈 w t , u 〉 is non-decreasing in time. Let

Then, on t ∈ [0 , T 1 ] , we have and integration yields

Therefore, T 1 ≲ τ s ( 〈 w 0 , u 〉 ) /λ min ( Σ ) . For t > T 1 , we have

Let

Once again, integration implies

which completes the proof.

D.2 Preliminary Lemmas for proving Theorem 5

We first introduce a number of concentration (and anti-concentration) lemmas that will be useful for proving Theorem 5.

Lemma 16. Suppose { z ( i ) } n i =1 i.i.d. ∼ N (0 , I d ) , and g : R → R satisfies | g ( · ) | ≤ C (1 + |·| p ) for some C > 0 and p ≥ 1 . Additionally, suppose { ϵ ( i ) } i =1 are i.i.d. σ -sub-Gaussian zero-mean noise independent of { z ( i ) } n i =1 . Let y ( i ) := g ( 〈 u , z ( i ) 〉 ) + ϵ ( i ) for some u ∈ S d -1 . Then, for any q > 0 , with probability at least 1 -4 d -q , we have

for all i .

Proof. Notice that 〈 u , z ( i ) 〉 ∼ N (0 , 1) , thus

Similarly, by the sub-Gaussian and zero-mean property of ϵ ( i ) ,

Thus, by a union bound, we have

with probability at least 1 -4 d -q . Using the upper bound on | g | finishes the proof.

Lemma 17. Suppose { z ( i ) } n i =1 are i.i.d. samples drawn uniformly from S d -1 . Then,

Proof. Fix some ϵ ∈ (0 , 1) . Let N ϵ be a minimal ϵ -covering of S d -1 . Let ˆ w be the projection of w onto N ϵ . Notice that by the triangle inequality and the union bound

Moreover, due to [BBSS22, Lemma A.7],

Choose ϵ = 3 √ d 8 n . By Lemma 25

which completes the proof.

We summarize the above statements into a 'good event', as characterized by the following lemma.

Lemma 18. Let { z ( i ) } n i =1 i.i.d. ∼ N (0 , I d ) , and z ( i ) := z ( i ) / ∥ z ( i ) ∥ for every i . We say event G occurs whenever:

For n ≳ d , event G occurs with probability at least 1 -O ( d -q ) .

Proof. The first and second statements of the lemma follow from Lemmas 16 and 17 respectively. The third and fourth statements are standard Gaussian covariance concentration bounds (see e.g. [Wai19, Theorem 6.1] where both statements hold with probability at least 1 -2 e -d ).

D.3 Proof of Theorem 5

We begin by recalling the definition w := ˆ Σ 1 / 2 w ∥ ˆ Σ 1 / 2 w ∥ and the finite-samples dynamics (3.5), which we copy here for the reader's convenience,

where η ( w t ) = ∥ ˆ Σ 1 / 2 w t ∥ 2 . Moreover, via chain rule, we obtain

To simplify the notation, define

Then

We can decompose the above dynamics into a population term and three different error terms in the following manner:

We will proceed in three steps. In the first, we bound E 1 , the concentration error. In the second, we bound E 2 and E 3 , the errors due to estimating Σ with ˆ Σ (i.e. replacing z ( i ) with ˜ z ( i ) ). Finally, we will analyze the convergence time similar to that of Proposition 4. Throughout the proof, we will assume that the event G of Lemma 18 occurs.

Step 1. Controlling the concentration error E 1 . Let K ≍ ln( nd q ) p/ 2 , and notice that on event G we have ∣ ∣ y ( i ) ∣ ∣ ≲ K for all i . Let y K := y 1 ( | y | ≤ K ) . On the event G , we have y ( i ) K = y ( i ) for all i , and where

Thus, our objective is to bound ∥ ∆ n ∥ uniformly for all w ∈ S d -1 . To that end, we first modify the expectation in the above definition so that the empirical average and expected value match in terms of their random variables. Specifically,

By the Cauchy-Schwartz inequality,

where the last inequality follows from Lemma 16. Hence,

Next, we need to establish high-probability bounds via a covering argument. To simplify the exposition, define the stochastic process indexed by w ∈ S d -1 and v ∈ S d -1 via

Fix some ϵ w , ϵ v > 0 . Let Θ w and Θ v be ϵ w and ϵ v coverings of S d -1 , and let ˆ w and ˆ v denote the projection of w onto Θ w and of v onto Θ v respectively, then

We bound each of the terms using Cauchy-Schwartz. Specifically,

where we used the upper bound on the operator norm of 1 n ∑ n i =1 z ( i ) z ( i ) ⊤ from Lemma 18 together with the fact that n ≳ d . Similarly,

To bound the differences when we replace w with ˆ w , we need to make a distinction between ReLU and smooth activations as the respective arguments are to some extent different. When ϕ ′ is Lipschitz,

and

Therefore, for a smooth activation ϕ we choose ϵ v = ϵ w = √ d/n , and obtain

When ϕ is the ReLU activation, we need to show that the sign of the preactivation changes only for a small number of samples when we change the weight w to ˆ w . Notice that

̸

Recall that z ( i ) := z ( i ) / ∥ z ( i ) ∥ . Choose ϵ w ≍ √ d/n . On event G , we know from Lemma 18 that at most O ( d ln( n/ √ d )) samples can satisfy the above condition. Therefore,

and

̸

where the last inequality follows from the anti-concentration on the sphere [BBSS22, Lemma A.7]. Thus, for ReLU we choose ϵ v ≍ √ d/n and ϵ w ≍ √ d/n , and once again obtain

It remains to bound the term

Notice that for fixed ˆ w , ˆ v , X ˆ w , ˆ v is sub-Gaussian with sub-Gaussian norm O ( K ) . Thus, via the sub-Gaussian maximal inequality [VH16, Lemma 5.2],

with probability at least 1 -e -d . Consequently, we have

with probability at least 1 -O ( d -q ) . Assuming that n grows at most polynomially in dimension and choosing a sufficiently large q , we have sup w ∈ S d -1 ∥ ∆ n ∥ ≤ ˜ O ( √ d/n ) with probability at least 1 -O ( d -q ) .

with probability at least 1 -e -n ′ / 2 . Combining the above with the bound on ∥ ∆ n ∥ , we have E 1 ≥ -λ max ( Σ ) ˜ O ( √ d/n ) with probability at least 1 -O ( d -q ) , which concludes the first step of the proof.

Step 2. Bounding the error due to the estimation of Σ , i.e. E 2 and E 3 . Recall that we are considering the event G , thus y ( i ) = y ( i ) K . We can control each of the error terms separately. We begin by E 2 , where by Cauchy-Schwartz

where the last line follows from Lemma 18 and the fact that n ≳ d . When ϕ ′ is additionally Lipschitz, we have

Moreover, for any w ∈ S d -1 ,

where the last inequality holds with probability at least 1 -2 e -d on the event of Lemma 24. Hence for smooth activations we conclude

When ϕ is the ReLU activation, we need a more involved argument to control E 2 . In particular, we will show that for any w , at most only ˜ O ( d ) datapoints can have sign (〈 w , z ( i ) 〉) = sign (〈 w , ˜ z ( i ) 〉) . Notice that

̸

̸

where z ( i ) := z ( i ) ∥ z ( i ) ∥ is distributed uniformly over S d -1 . From Lemma 24 we have ∥ I d -ˆ Σ -1 / 2 Σ 1 / 2 ∥ ≲ √ d n ′ with probability at least 1 -2 e -d . On the other hand, from Lemma 17 we know with probability at least 1 -e -d , for any w ∈ S d -1 at most ˜ O ( d ) of the labeled samples have ∣ ∣ 〈 w , z ( i ) 〉∣ ∣ ≲ √ d/n . Recall that n ′ ≳ n 2 when using the ReLU activation. This is precisely why we make this choice for the ReLU activation, as we need to balance the RHS of (D.7) which is of order √ d/n ′ with the LHS of (D.7) which should at most be of order √ d/n if we want to ensure only ˜ O ( d ) samples satisfy the bound. When n ′ = n 2 we can balance these two terms, thus with probability at least 1 -3 e -d the sign change can occur for at most ˜ O ( d ) many samples, and

In this case, we end up with

Bounding E 3 for ReLU and Lipschitz ϕ ′ is identical. In both cases, by Cauchy-Schwartz,

which holds on the intersection of event G and of Lemma 23. At last, using the bound on ∥ ν ( w t ) ∥ from (D.6), we obtain

with probability at least 1 -O ( d -q ) .

Step 3. Analyzing the Convergence. As a result of the previous steps, we have established

Thanks to Lemma 11, we can write

where ζ ϕ,g and ψ ϕ,g were introduced in (D.2) and (D.3) respectively. Recall the definition of ν ( w t ) ,

Therefore,

where u t ⊥ := u -〈 u , w t 〉 w t . Moreover, from Lemma 23, we have

with probability at least 1 -e -n ′ / 8 . Hence, for n ′ ≳ d κ ( Σ ) we have λ min ( ˆ Σ ) ≳ λ min ( Σ ) , and consequently,

where c ′ is a universal constant. Notice that the first term in the RHS above denotes the signal, while the second term denotes the noise. We want to ensure the noise remains smaller than the signal throughout the trajectory, which leads to the convergence of w t to u . Notice that the signal term is first increasing, then decreasing for 〈 w t , u 〉 ∈ [0 , 1] . Thus, it suffices to ensure the noise is smaller than the signal on the two ends of the interval, i.e. at time t = 0 and at time t = T where 〈 w T , u 〉 = 1 -ε . At initialization, this condition leads to

where C hides constant depending only on s and at most polylogarithmic factors of d . Thus, we have established the sample complexity as presented by Theorem 5.

It remains to obtain the convergence time. With the above sample complexity, we have

where c ′′ is a universal constant. The rest of the proof follows by integration and is identical to the proof of Proposition 4 in Appendix D.1.

D.4 Proof of Corollary 6

The proof follows immediately from Theorem 5 and the following lemma which describes how 〈 w 0 , u 〉 behaves under different regimes of r 1 and r 2 .

and at time t = T , leads to

Lemma 19. Suppose Σ follows the ( κ, θ ) -spiked model, w 0 is sampled uniformly from S d -1 , n ′ ≳ d , and there exist universal constants C 2 , C ′ 2 , C 3 , C ′ 3 > 0 such that

for r 1 ∈ [0 , 1 / 2] and r 2 ∈ [0 , 1] . Then, conditioned on 〈 w 0 , u 〉 > 0 , with any arbitrarily large constant probability 1 -δ , for sufficiently large d (that depends on δ ) we have

Proof. By definition,

Recall that we are conditioning our arguments on 〈 w 0 , u 〉 > 0 , hence the numerator of the above fraction is positive. To translate the sample complexities of Theorems 5 and 7 to the spiked model, our goal is to lower bound 〈 w 0 , u 〉 in terms of d , r 1 , and r 2 .

We begin by observing that

where the last inequality holds on the event of Lemma 24, which happens with probability at least 1 -2 e -d . Consequently,

Furthermore, due to the Markov inequality,

Similarly (by conditioning on ˆ Σ )

Additionally, on the event of Lemma 24,

Therefore, on the above events, for some absolute constant C > 0

Recall that C 2 d r 2 ≤ κ ≤ C ′ 2 d r 2 and C 3 d -r 1 ≤ ⟨ u , θ ⟩ ≤ C ′ 3 d -r 1 (notice that changing θ to -θ does not change the spiked model of Assumption 1, thus we can assume ⟨ u , θ ⟩ ≥ 0 without loss of generality). Then,

The last term in the numerator can be made arbitrarily small by sufficiently large n ′ , hence we focus on other terms for now. Intuitively, when r 2 < 2 r 1 , the denominator is of constant order. If additionally r 2 < r 1 , the dominant term in the numerator is 〈 w 0 , u 〉 and 〈 w 0 , u 〉 ≍ 1 / √ d , otherwise the dominant term is κ 〈 w 0 , θ 〉 ⟨ u , θ ⟩ and 〈 w 0 , u 〉 ≍ d r 2 -r 1 -1 / 2 . On the other hand, when r 2 > 2 r 1 , the denominator is of order d r 2 / 2 -r 1 , and once again the dominant term of the numerator is κ 〈 w 0 , θ 〉 ⟨ u , θ ⟩ , therefore 〈 w 0 , u 〉 ≍ d ( r 2 -1) / 2 . Using this intuition, we analyze each of the following regimes separately.

Case 1. 0 < r 2 < r 1 : In this case, by [BBSS22, Lemma A.7] we have ∣ ∣ 〈 w 0 , u 〉∣ ∣ ≥ c/ √ d with probability at least 1 -4 c . On the intersection of all considered events with 〈 w 0 , u 〉 > 0 , and for sufficiently large d and n ′ ≳ d , we must have 〈 w 0 , u 〉 > 0 (otherwise 〈 w 0 , u 〉 < 0 ). Thus by plugging the values in (D.9),

The intersection of all desired events and 〈 w 0 , u 〉 > 0 happens with probability at least 1 2 -4 c -2 /C 1 -2 e -d , thus conditioned on 〈 w 0 , u 〉 the probability is at least 1 -8 c -4 /C 1 -4 e -d . Choosing sufficiently small c , large C 1 , and respectively large d and n ′ ≳ d with sufficiently large absolute constant, we can arbitrarily increase the (constant) probability of success. Thus the analysis of this regime is complete.

Case 2. r 1 < r 2 < 2 r 1 : This time we use the fact that ∣ ∣ 〈 w 0 , u 〉∣ ∣ ≤ √ C 1 /d with probability at least 1 -1 /C 1 , and ∣ ∣ 〈 w 0 , θ 〉∣ ∣ ≥ c/ √ d with probability at least 1 -4 c . By an argument similar to the previous case, for sufficiently large d and n ′ ≳ d , 〈 w 0 , u 〉 > 0 implies 〈 w 0 , θ 〉 > 0 , hence by (D.9)

with probability at least 1 -8 c -4 /C 1 -4 e -d when conditioned on 〈 w 0 , u 〉 > 0 .

Case 3. 2 r 1 < r 2 < 1 : Once again recall (D.9). To bound the numerator, we repeat the exact same argument as in the previous case, thus which finishes the proof of the lemma.

E Proofs of Section 4

E.1 Proof of Theorem 7

We recall from (D.5) that

Furthermore, the preconditioned dynamics of w t given by (4.1) reads

The rest of the analysis is identical to that of the proof of Theorem 5 in Appendix D.3. Specifically, by defining we have

As long as n ≳ d , n ′ = n for the smooth case, and n ′ ≳ n 2 for the ReLU case, the first two steps of the proof of Theorem 5 in Appendix D.3 implies that

Once again, we apply Lemma 11 to obtain

with ζ ϕ,g and ψ ϕ,g given in (D.2) and (D.3) respectively. As a result,

We need to ensure the noise term, i.e. the second term on the RHS remains smaller than the signal, i.e. the first term. The signal term attains its minimum at either initialization t = 0 or at the end of the trajectory t = T where 〈 w T , u 〉 = 1 -ε , which imposes the following sufficient conditions on n . Namely, at initialization we require

while at t = T we require where C hides constant that only depend on s and polylogarithmic factors of d . Hence, we obtain

for some universal constant c ′ > 0 . Via integration (similar to the proof of Proposition 4 in Appendix D.1), for we obtain

and for we obtain T 2 -T 1 ≲ ln(1 /ε ) , which completes the proof. We conclude by remarking that the proof of Corollary 8 is immediate given Theorem 7 and Lemma 19.

E.2 Preliminary Lemmas for Proving Theorem 9

We will adapt the following lemma from [MHPG + 23], which provides a non-parametric approximation of g via random biases.

Lemma 20. [MHPG + 23, Lemma 22] For any smooth g : R → R and ∆ > 0 , let ˜ g : R → R be a smooth function such that ˜ g ( z ) = g ( z ) for | z | ≤ ∆ and ˜ g ( -2∆) = ˜ g ′ ( -2∆) = 0 . Suppose { b j } j m =1 i.i.d. ∼ Unif ( -2∆ , 2∆) , and let ∆ ∗ := ∆sup | z |≤ 2∆ | ˜ g ′′ ( z ) | . Then, there exist second layer weights { a j ( b j ) } j m =1 with ∥ a ∥ ≲ ∆ ∗ / √ m , such that for any fixed z ∈ [ -∆ , ∆] and any δ > 0 , with probability at least 1 -δ over the random biases,

where ϕ is the ReLU activation.

We use the above lemma to show the existence of a second layer with ˜ O (1 / √ m ) norm with training error of order ˜ O (1 /m ) .

Lemma 21. For any ε < 1 , suppose ⟨ w , u ⟩ ≥ 1 -ε . Then for any q > 0 , sufficiently large d , n ≳ d/ε 2 , with probability at least 1 -O ( d -q ) over the random biases and the dataset, there exists a second layer a with ∥ a ∥ ≤ ˜ O (1 / √ m ) described by Lemma 20 such that

where ϕ is the ReLU activation.

Proof. We begin by replacing w and ˜ z ( i ) with u and z ( i ) . Specifically, via Jensen's inequality,

We bound each term separately. For E 1 , we can invoke Lemma 20 which implies that each term in the sum can be bounded by ˜ O (1 /m ) with probability at least 1 -1 / ( nd q ) , thus by a union bound, with probability at least 1 -d -q over the random biases,

By sub-Guassianity of ϵ ( i ) (hence sub-exponentiality of ϵ ( i ) 2 ), for n ≳ d (with a sufficiently large constant) we have

with probability at least 1 -e -d .

For E 3 , via the Lipschitzness of ReLU and the Cauchy-Schwartz inequality we can write where we used the event of Lemma 24 which happens with probability at least 1 -2 e -d , and ˜ O (1) represents a constant that depends at most polylogarithmically on d .

Finally, we bound the last term. Once again via the Lipschitzness of the ReLU activation and the Cauchy-Schwartz inequality

where once again we used the event of Lemma 24. On the intersection of all desired events, we have

We conclude the proof by noticing that n ′ ≳ n 2 and n ≳ dε -2 .

Additionally, we will use the following standard Lemma on the Rademacher complexity of two-layer neural networks, which in particular is a restatement of [MHPG + 23, Lemma 18] in a way suitable for our analysis.

Lemma 22. Let F be a class of real-valued functions on ( z , y ) . Given n samples { z ( i ) , y } n i =1 , define the empirical Rademacher complexity of F as

where ( ς i ) are i.i.d. Rademacher random variables (i.e. ± 1 with equal probability). Suppose F is given by

for some fixed u ∈ S d -1 . Suppose { z ( i ) } n i =1 i.i.d. ∼ N (0 , I d ) , and suppose | ϕ ′ | ≤ 1 . Then,

Proof. See the proof of [MHPG + 23, Lemma 18].

E.3 Proof of Theorem 9

Throughout the proof, we will assume ⟨ w , u ⟩ ≥ 1 -ε where we recall

From either Theorem 5 or Theorem 7, we can assume ⟨ w , u ⟩ ≥ 1 -ε with probability at least 1 -O ( d -q ) for any fixed q > 0 . For simplicity, let

and similarly define ˆ y ( z ; u ) . We define the following quantities,

and similarly define their empirical counterparts,

Notice that ultimately, we are interested in bounding R ( w ) . We break down the proof into three steps. In the first step, we show that R ( w ) can be upper bounded by R ( u ) . Then, via a generalization bound, we show that the R ( u ) can be upper bounded by ˆ R ( u ) . Finally, we show that ˆ R ( u ) can be upper bounded by the training error, i.e. ˆ R ( w ) , and convex optimization of the last layer can attain the near-optimal value of this training error which is bounded by Lemma 21.

Step 1. Bounding R ( w ) via R ( u ) . By Jensen's inequality,

Suppose ∥ a ∥ ≤ r a / √ m . For the first term, by Lipschitzness of ϕ and the Cauchy-Schwartz inequality

where the last inequality holds with probability at least 1 -2 e -d on the event of Lemma 24. For the middle term, via a similar argument,

In what follows, we will restrict the analysis to the case where r a = ˜ O (1) . Therefore, we have

Step 2. Generalization: Bounding R ( u ) via ˆ R ( u ) .

Define the event

and similarly define E ( i ) by replacing z and ϵ with z ( i ) and ϵ ( i ) respectively. Via the Cauchy-Schwartz and Jensen inequalities

Moreover, E [ y 4 ] ≲ 1 , E z ,y [ ˆ y ( z ; u ) 4 ] ≤ ˜ O (1) , and P ( E C ) ≤ 4 / ( nd q ) (via a standard sub-Gaussian tail bound). Consequently,

Let

Notice that u is fixed. Then, by a standard symmetrization argument (see e.g. [VH16, Lemma 7.4]) and Lemma 22

where ˆ R n ( F ) . As the loss is bounded, we can apply McDiarmid's inequality to turn the above bound in expectation into a bound in probability, in particular

with probability at least 1 -2 e -d . Therefore, we conclude this step by noticing that

with probability at least 1 -2 e -d .

Step 3. Bounding the training error and finishing the proof. This step is similar to the proof of [MHPG + 23, Theorem 4]. For conciseness, define

and ˆ R λ ( a ) := ˆ R ( a ) + λ ∥ a ∥ 2 / 2 . Our goal is to choose suitable λ such that the minimizer

satisfies ∥ a ∗ ∥ ≤ r a / √ m while the value of the above minimization problem which we denote with ˆ R ∗ λ does not significantly exceed

We argue that the suitable choice for λ is

Let ˆ R ∗ denote the minimizer of the regularized problem and ˜ a := arg min ∥ a ∥≤ r a / √ m ˆ R ( a ) . From Lemma 21, with a proper choice of r a = ˜ Θ(1) , we have

with probability at least 1 -O ( d -q ) over the biases and the dataset. Note that as a ∗ is the minimizer of ˆ R λ , we have and in particular

and

Let { a t } t ≥ 0 be the solution to the gradient flow of a . Then,

and by the first-order condition of strong convexity

therefore

As the training error (of the regularized problem) is λ -strongly convex in a , by applying the standard Polyak-Łojasiewicz condition, gradient flow for training a obtains

Furthermore, since

we have

Consequently, choosing

implies

Therefore

Recall that

is the final training error which we also denoted by ˆ R ( w ) earlier in this section when were not focusing on the second layer. From the previous two steps, we know how to bound R ( w ) via ˆ R ( u ) . Thus the last step is to upper bound ˆ R ( u ) via ˆ R ( w ) . To that end, via Jensen's inequality

The first term on the RHS is ˆ R ( w ) for which we developed a bound earlier in this step. Bounding the latter two terms can be performed similarly to the arguments in the previous sections. In particular,

where the last inequality holds with probability at least 1 -2 e -d (over the event of Lemma 24). Similarly,

Putting the bounds back together (recall n ′ ≥ n ≳ dε -2 ), we arrive at

Combining the result of this step with the two previous steps implies

with probability at least 1 -O ( d -q ) (when conditioned on 〈 w 0 , u 〉 > 0 ) which completes the proof of Theorem 9.

F Auxiliary Lemmas

In this section, we recall a number of standard lemmas which we employ in various parts of our proofs.

Lemma 23. [Wai19, Theorem 6.1]. Suppose { x ( i ) } n ′ i =1 i.i.d. ∼ N (0 , Σ ) . Let ˆ Σ := 1 n ′ ∑ n ′ i =1 x ( i ) x ( i ) ⊤ . Then, for n ′ ≥ tr( Σ ) /λ max ( Σ ) ,

with probability at least 1 -e -n ′ / 2 . Furthermore, for n ′ ≥ d ,

with probability at least 1 -e -n ′ / 8 .

Then

Lemma 24. Suppose { z ( i ) } n ′ i =1 i.i.d. ∼ N (0 , I d ) , let x ( i ) := Σ 1 / 2 z ( i ) for some invertible Σ , and define

with probability at least 1 -2 e -d .

Proof. We have

Similarly,

Moreover, by [Wai19, Example 6.2], we have with probability at least 1 -2 e -d ,

Thus, for n ′ ≳ d (with a sufficiently large absolute constant), we have

with probability at least 1 2 e -d .

-

Lemma 25 (Chernoff's Inequality) . Suppose X 1 , . . . , X n are i.i.d. Bernoulli random variables, and further assume that E [ ∑ i X i ] ≤ µ . Then, for any δ ≥ 1 ,

Proof. The proof follows from a standard Chernoff bound. From [Ver18, Theorem 2.3.1]

(notice that the statement of [Ver18, Theorem 2.3.1] holds true even when E [ ∑ i X i ] = µ is replaced with E [ ∑ i X i ] ≤ µ ). We conclude by remarking that δ -(1 + δ ) ln(1 + δ ) ≤ -δ/ 3 for δ ≥ 1 .

Peer reviews:
comments: SUMMARY:
Summary: The authors consider single index learning with a spiked input distribution, rather than the Gaussian typical in most works in this setting, showing that a modified SGD can learn where vanilla SGD fails and showing better sample complexity under the assumption that the spike is aligned with the true hidden signal.

STRENGTHS:
Strengths: The question of moving beyond the Gaussian input distribution is an interesting one, and the results seem quite interesting.  In particular, the transitions in behavior depending on the alignment of the spike provide a rather full picture of the setting.

WEAKNESSES:
Weaknesses: The setting is somewhat artificial, as the spike strongly correlating with the true signal is a much easier setting than the typical single index model.  Furthermore, it’s not instantly clear how this version of SGD performs according to a somewhat silly baseline that’s aware of the possible correlation of the spike and true signal.  From the intuition of [1], most of the effort of the learning process is spent trying to find non-negligible correlation between the learned parameter and the true direction.  If you suspect the spike of your input distribution is already meaningfully correlated, you could simply initialize there.  Perhaps this is a naive idea ([1] requires the link function to be known for instance), but a little discussion of why the spiked covariance setting is not too easy would strengthen the paper I feel.

[1] Arous, Gerard Ben, Reza Gheissari, and Aukosh Jagannath. "Online stochastic gradient descent on non-convex losses from high-dimensional inference." The Journal of Machine Learning Research 22.1 (2021): 4788-4838.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors investigate the learning of single-index models with a gradient-based algorithm. They analyze structured inputs considering a spiked matrix model for the data. They show that an appropriate normalization of the weights in the training algorithm can help recover the target direction. They study the behavior of the sample complexity to learn the target direction as a function of the magnitude of the spike in the data, and the alignment of the spike with the target direction. They identify three learning phases, characterizing the phase transitions between them.\n\nSTRENGTHS:\nThe precise characterization of feature learning for single-index models is an important topic in machine learning theory. The analysis of structured input is a relevant step toward a better understanding of these models. The characterization of the sample complexity under the spiked model given in terms of a phase diagram is nice.\n\nWEAKNESSES:\nThe main weakness of the manuscript is the writing, sometimes not entirely clear.\n\nQUESTIONS:\n1) The overview of the relevant literature on single-neuron training is exhaustive and nice. However, when introducing the single-neuron modification after eq. (2.3), I think it is confusing for the reader to introduce a general link $\\phi$, stating assumptions on its form (e.g. Assumption 2), and then write \"there is little to no hope for learning the entire single-index model unless $\\phi = g$\".  I would rephrase this sentence in the introduction of the single-neuron simplification, given that the following results analyze non-necessarily strictly realizable settings. \n2) The authors do a good job in general in trying to explain their results. I think it would be nice to expand on Proposition 4, e.g. by explaining what is the role of Assumption 2 in a sketch of proof. On a similar note, it is not clear how far from ReLU $\\phi$ can be. If an intuitive understanding can be given on the class of $\\phi$ to be considered, is more than welcome. \n3) In Section 4.2 extends the findings to two-layer neural networks,  in Thm 9 is written\"suppose $\\phi$ is ReLU and Assumption 2 holds\", why there is a need to restrict strictly to a ReLU network in this case and not just assuming that Assumption 2 holds?\n4) I like the illustrative power of Figure 1. The hard and easy phases are explained in simple terms on page 2, however, the intermediate phase is not introduced.\n\nLIMITATIONS:\nI do not see any potential negative social impact of this paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the learning dynamics of single index models when the input data has a structured covariance. It demonstrates that the commonly used spherical gradient flow fails to learn the target direction even when the spike and target directions are identical. However, the paper introduces an appropriate weight normalization technique that overcomes this limitation and successfully recovers the target direction. It further reveals that leveraging the alignment between the covariance structure and the target direction improves the sample complexity compared to isotropic cases and outperforms lower bounds for rotationally-invariant kernel methods. The paper also suggests future directions, such as studying multi-index models, understanding the limitations from a Correlational Statistical Query (CSQ) perspective, and exploring the effects of different initializations in training networks with multiple neurons.\n\nSTRENGTHS:\n${\\bf Originality}$: This paper introduces a new investigation into learning single-index models with a spiked covariance structure, expanding upon prior research that primarily focused on isotropic data. This unique focus on the impact of additional structure in the covariance matrix sets it apart from existing works. Furthermore, the introduction of weight normalization techniques and the exploration of their effects in anisotropic scenarios demonstrate originality in addressing the limitations of conventional spherical gradient dynamics.\n\n$\\bf Quality$: This paper demonstrates a high level of quality in terms of its theoretical analysis, rigorous mathematical proofs, and well-supported findings. The authors delve into the dynamics of gradient flow, leveraging mathematical techniques to derive insights into the behavior of the learning algorithm. The use of appropriate assumptions and the incorporation of empirical covariance preconditioning showcase the paper's thoughtful methodology and attention to detail.\n\n$\\bf Clarity$: The paper is written in a clear and accessible manner, making complex concepts and mathematical formulations understandable to the reader. The abstract, introduction, and conclusion concisely summarize the objectives, methods, and key findings of the study. The organization of the paper and the logical flow of ideas contribute to its overall clarity, enhancing the reader's comprehension of the research.\n\n$\\bf Significance$: The paper holds good implications for the understanding of learning dynamics in the presence of structured covariance matrices. By uncovering phenomena where conventional approaches fail, such as the inability of spherical gradient dynamics to recover the target direction, the paper sheds light on the limitations and challenges associated with anisotropic data. Moreover, the proposed weight normalization techniques, insights into improved sample complexity, and the comparison to rotationally invariant kernel methods highlight the practical relevance and potential impact of the findings.\n\nWEAKNESSES:\n${\\bf (1)}\\ \\textbf{The training procedure considered in the paper is not practical:}$ The authors focus on a two-step training procedure that deviates from the standard gradient descent (GD) commonly used in practice. It is challenging to assess the extent to which the results obtained from this training procedure translate to impacts on standard training algorithms.\n\n${\\bf (2)} \\ \\textbf{Differences from the existing works:}$ Some existing works, such as Refs [BBSS22] [BAGJ21], have established some results on learning the single-index model using neural networks. It seems that the authors don't clarify the difference of this work from the existing ones. \n\n${\\bf (3)}\\ \\textbf{Practical implications and applications:}$ Although the authors have successfully established strong theoretical results for learning single-index models using two-layer neural networks, the practical implications of this work on real-world applications of deep learning remain somewhat unclear. Further clarification regarding the practical relevance and potential impact of the findings would greatly benefit readers seeking to understand the practical implications of the research in the field of deep learning.\n\nQUESTIONS:\n$\\bf Q1.$ What happens if assuming $g$ and $\\phi$ to be differentiable? Can you get stronger results?\n\n$\\bf Q2.$ Is it possible to extend the current analysis to the multiple-index models? What are the main technical difficulties?\n\nLIMITATIONS:\nNot applicable", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}