Gradient-Based Feature Learning under Structured Data
Authors: Alireza Mousavi-Hosseini
Denny Wu
Taiji Suzuki
Murat A Erdogdu
Venue: NeurIPS 2023 poster (2023)
Keywords: feature learning
neural networks
single-index model
gradient descent
TL;DR: We show that when learning single-index models with neural networks using gradient-based training, additional structure in the covariance matrix can improve the sample complexity.

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
Alireza Mousavi-Hossein', Denny Wu', Taiji Suzuki', Murat A. Erdogdu'
University of Toronto and Vector Institute, New York University and Flatiron Institute, University of Tokyo and RIKEN AIP
(mousavi, erdoddu toronto edu, dennyvu@nyu.com edu, tai maiji0mist i. u-tokyo.ae jp
Abstract
Recent works have demonstrated that the sample complexity of gradient-based learning of single index models, i.e. functions that depend on 1-dimensional projection 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. this work, 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, obtain improved sample complexity compared to the isotropic case. In particular, under the spiked model with suitably large spike, the sample complexity of gradient-basse training can be made independent of the information exponent while also outperforming lower bounds for rotationally invariant kernel methods.
Introduction
fundamental feature of neural networks is their adaptivity 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-basse based algorithms can efficiently adapt such structure via feature/repressettaiio learning. A considerable amount of research has been dedicated understanding this phenomenon under various assumptions and demonstrate the superiority of neural networks over -daapivee methods such as kernel models [GMMM19, WLLM19, BES 19, LMZ20, AAM22, + 22, DLS22, Te123, MHPG+
particular relevant problem setting for feature learning is the estimation single index models, where the response y R depends the input æ Rd via g((u, x)) + c, where g IR + R is the nonlinear link function and u is the unit target direction. Here, learning corresponds to recovering the unknowns 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, larger information exponent corresponds more complex g (for gradient-based learning), and it has been proven that when the input is isotropic (O Ia), gradient flow can learn the single index model with (dd) sample complexity [BBSS22].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 [HTFRO9]; indeed, this is the main reason principal component analysis used in pre- training [JWHT13]. A fundamental model that captures such a structure is the spiked matrix model in which æ ~ (0, La + k00") for some unit direction 0 E IR and k > 0 [Joh01]. Along the direction 0, data has higher variability and predictive power. single index models, such predictive power translates to non-trivial alignment between the vectors and 0 our focus is to investigate the effect of such alignment on the sample complexity of gradicnt- -ased 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
r spikke Sepuuii
111-12)
infermediate
spike-target alignment: (u, 0) x d d' E [0,1/2],
86
Easy
spike magnitude: d"² for r2 E [0, 1].
04
Our contributions can be summarized as follows.
Hard
1. We show that even the case of perfect spike- target alignment (r1 0), the spherical gradient flow commonly employed in recent literature (see e.g. [BAGJ21, BBSS22]) cannot recover the target direc- tion for moderate spike magnitudes in the population limit. The failure of this covariance-agnosti procedure under anisotropic structure insinuates the necessity of appropriate covariance- normalization ef- fectively learn the single index model.
d²
0.0
0.1
0.2
03
4.4
F1 (spike- ttrget alignment)
Low
sample complescity
High
Figure I: Sample complexity to learn u and under the spiked model. Smaller denotes better spike-target alignment, while larger v2 denotes a larger spike magnitude. The sample complexities are based on Corollary
2. We show that a covariance-awart 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.
3. Under the spiked covariance model, we prove a three-stage phase transition for the sample complex- ity depending on the quantities r1 and r2. For suitable direction and magnitude of the spike, the sample complexity can be made Odd++r) for any v > 0 which is independent of the information exponent This should be compared against the known complexity of O(d") under isotropic data.
4. 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 (dd++) samples can be reduced to (d)++) for any v > 0, i.e. almost linear in d. The three-stage phase transition also emerges, illustrated in Figure 1: in the "hard" regime, the complexity remains ((d") 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 r11 and larger r32 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-agnosti gradient flow in Section 2.1. Our main sample complexity result on a single neuron is presented Section 3.2. We provide our results on multi- nuutinneuron neural networks 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, ||:I1 denotes the usual operator norm, and (.ax and Amin (-) denote the largest and smallest eigenvalues respectively. We reserve y for the standard Gaussian distribution on IR, and let II-Il, denote the L²(r) norm. gd-1 is the unit d-dimensional sphere. For quantities a and b, we will use S 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 Cb, and b signifies that S b and b S

---

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, COB 19]. However, there are many instances of function classes that 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 [CB 18, RVE18, MMN18, MMM19, NWS22, Chi22], which can learn certain low-dimensional target functions efficiently [WLLM19, AAM22, HC22, ASKL233.
As for neural networks with smaller width, recent works showed that a two-stage feature learning pro- cedure can outperform the NTK when the data is sampled uniformly from the hypercube [BEG 22] or isotropic Gaussian [DLS22, BES+ 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 companion work
we zoom into the setting where the spike and target are perfectly aligned (r1 = 0), and prove learnability in the n x 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 single index model is classical in statistics [LD89], with efficient dedicated algorithms (IKKSKI1, I, 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 considered the effect of taking single gradient step on the ability of a two-layer neural network to learn a single index model, and [BBSS22, MHPG* + 23] considered training 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 as the activations are required to be monotone or have similar properties, see e.g. [So117, 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 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 æ and a link function g E L?))), consider the single index model
y 0
(u (u,x) |
+ with x E),
(2.1)
where is zero-mean noise with (11) sub-Gaussiaa norm and Sd-1, Learning the model (2.1) corresponds to approximately recovering the unknown link g and the unknown direction u. Note that normalization is needed to make this problem well-defined; without loss of generality, we write (u, / //2 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
[(: W. b) = Σ (o(((w x) + by).
(2.2)
where W- [w: is the m x d matrix whose rows corresponds to first-layer weights un, {a, { denote the second-layer weights, b = {b: pm denote the biases, and ø is the non-linear
3

---

activation function. We assume and á are weakly differentiable with weak derivatives g' and d' respectively. and g.g', o, o' (7). 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 and u. by minimizing the population risk
R(W, a, b) = [[j(;; W.a, a. y)²].
(2.3)
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; the end of this stage, the neurons w, align with u. Here, the goal is to recover only the direction. Next, using random biases and training the second 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;'s ss 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 = L, 01 = 1, bj = 0, i.c., 0(:; W, a. b) is replaced with ğ(æ; w) = 6((w, æ)) and we write R(w) == R W, a, b) for simplicity. We emphasize that unless & 9 (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, > 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 >
Characteristics of the link function play an important role in the complexity of learning the model. As such, central part of our analysis will rely on particular property based on the Hermite expansion of functions in a basis defined by the normalized Hermite polynomials (h, 3220 given as
h2(2) = (-1 d²
e- 12 dei
(2.4)
vs!
These polynomials form orthonormal basis in the space 1.2 and the resulting expansion yields the following measure of complexity for which is termed as the information exponent.
Definition 1 (Information exponent). Let = Σ520 oghy be the Hermite expansion of g. The information exponent of g is defined to be ses == inf > 0 a3 #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 Ó(d)) for s 2 3 where s 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 certain normalization in the training. In what follows, we prove that such normalization in training procedure indeed necessary.
2.1 The spiked model and limitations of covariance-agnostic training
In practice, data often exhibit certain structure which may have profound impact on the statistical procedure. A well-known model that captures such a structure the spiked model [JohO1] for which one or several large eigenvalues of the input covariance matrix Σ are separated from the bulk of the spectrum (see also [BBAP05, BS06J). Although our results hold for generic covariance matrices, they reveal interesting phenomena under the following spiked model assumption.
Assumption 1. The covariance Σ follows the (N. O)-spiked model if Σ
To0 1+*
where | 1.
In pursuit of the target (unit) direction u, the magnitude of the neuron w is immaterial; thus, recent works take advantage this and simplify the optimization trajectory by projecting ws onto unit sphere gd-1 throughout the training process [BAGJ21, BBSS22). In the sequel, 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
du = - --V where R(w))> = (((w) (VR(w)) w) w.w. dt
(2.5)
where V is the spherical gradient at the current iterate. It is straightforward see that when the initialization wD is on the unit sphere, the entire flow will remain on the unit sphere, i.e. w E Sd-1 for all & 2 O. The flow (2.5) has been proven useful for learning the direction u [BBSS22] in the isotropic case Σ = Ia when the activation ø is ReLU. In contrast, when Σ follows a spiked model, we show that it can get stuck stationary points that are almost orthogonal to u. Indeed, when the imput covariance E has a spike in the target direction u, i.c. 0 = 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', u) will be insignificant for all
Theorem 2. Let s be the information exponent of g with Eigl = O, and assume Σ follows the (is, u)-spiked model with 2(1) S K < (d For ReL.U activation, let w² denote the solution to (2.5) initialized uniformly at random over gd-1, then with probability at least 0.99,
supl (w', =)) 1/VVd 120
(2.6)
non-trivial alignment between the first-layer weights wi 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 aligmment 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. x 2 (dd). 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
R(w)= 1/2 1 (o((w, x)) -v² B[[o((w }²} - Elw((w, x))v + L8 [+*].
(2.7)
If the input was isotropic, i.e. x ~ (0, Ia), the first term in (2.7) would be equal to |0 which is independent of w. Thus, minimizing R(w) in this case is equivalent to maximizing the "correlation" term E[cp( (ur, jr) However, under the spiked model, the alignment between and u breaks the symmetry; consequently, the first term in the dccomposition grows with (w, u), creating a repulsive force that traps the dynamics around the equator where is almost orthogonal to
3 Main Results: Alignment via Normalized Dynamies
Having established that the covariance-agnostic training dynamics (2.5) is likely to fail, we consider 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 x = Σ 1/2 x. e = 1/2 /////////////// /2/ wll and similarly define प, and consider the prediction function ğι:; w) := ol (w. 2)). Due to symmetry, the second moment of
the prediction is E [ax; )² = | which independent of wn thus, the population risk reads
R(w)= = E (0(2; w) y)²| -
E [
Elo( (w, z))y]-
(3.1)
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 V.. R(w) = - V.. Elo((W, z))y]. Based on this, we use the following normalized gradient flow for training
dw
dt = =n((2²) V R(w') where (ww) = // w 1|.
(3.2)
5

---

We remark that, though not identical, this normalization is closely related to batch normalization which commonly employed in practice [IS15]. Under the invariance provided by the current normalization, minimizing R(w)) corresponds to maximizing (1, z) )yj. 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' using the gradient flow (3.2). Then wt solves the following ODE
d'o
dt = So (aaa((w' (w', (La- w'w' T )E( T' u
(3.3)
where So.g W., U) == - Bio (w, 2))g (o'ggg,,TT x))
We will investigate if the modified flow (3.3) achieves alignment: in this context, alignment corre- sponds to (w², 4) R 1. Towards that end, we make the following assumption.
Assumption 2. Let g -Ezz ash; and & = -Ezo Rygh be the Hermite decomposition of g and respectively, Let 8 be the information exponent of g. For some universal constant > 0, we assume E S-
Sōlg(w)
Ve E (0,
There are several important examples that readily satisfy Assumption 2. when the link function known as [BAGJ21], i.e. e = g. A more interesting example is when is an activation with degree s non-zero Hermite coefficient (e.g. ReI.U when s is even, see [GKK 19, Claim 1D and g a degree Hermite polynomial, which for = corresponds to the phase retrieval problem. In this case, the assumption is satisfied if a, and 8. have the same sign, which occurs with probability 0.5 we randomly choose the sign of the second layer.
The obvious example is
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 (w5", ④ > 0. Then, we have (wo T) 2 - E as soon as
1
1
w 4 T
In(1/e) Amin (
where T2(2)= =
In(1/z) (1/2)²-2
8 = 2.
(3.4)
> 2
We remark that the information exponent enters the rate in (3.4) through the function Ts, 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.d, samples (()(0, | from the single index model (2.1), we consider the flow dw
di - ( VR((')) with VR(w) == - Ei= Φ
(w.
(())
(3.5)
13' wlll
where we estimate the covariance matrix Σ using the sample mean Σ ii=1 x (1) over n i..d. samples; the above dynamics defines an empirical gradient flow. Notice that we ignored the gradient associated with the term d² 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' 2 n² by assuming access to an additional n' n unlabeled data points. Similar the previous section, we track the dynamics of normalized w by defining w $1/2w u/ (1/2w (and leave unchanged from Section 3.1). The same arguments as in Lemma 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 1 p > For & denoting either the ReLU activation or a smooth activation satisfying | (8" S 1, suppose Assumption 2 holds. For any O, suppose
11+1-P for some
the finite sample gradient flow (3.5)
with wy(w) = |/ wl? initialized such that (w u) > O. and with number of samples
n 2 dx(x)) (10) u 2(1-)) ve-2]

---

where x ( is the condition number of Σ. Then, for T T ((((T"x)))+u((1/)) Amin we have
(3.6)
with probability at least 1 c1d d for some universal constants C1,02 > 0 over the randomness of the dataset. Here, Ta is defined in (3.4) and 2 hides poly-lggaritmiic
Remark. We make the following remarks on the above theorem.
The initial condition (w) 2i) is required when we have odd information exponent. When w is initialized uniformly over Sd-1 the condition holds with probability 0.5 over the initialization. See [BAGJ21, Remark 1.8] for further discussion this condition.
Although ur is defined using the empirical covariance unlike which is defined by population co- variance, this definition the suitable choice to approximate the target function g (c.f. Theorem 9), since it ensures the arguments of & and g sufficiently close when recovers
The intuition behind the proof of Theorem 5 is presented in Section 5 with the complete proof 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 concrete example of such improvement by specializing Theorem 5 for a spiked covariance model.
Corollary 6. Consider the setting of Theorem 5 with E following the (1s, Oy-pikeed model, where (u, 0) x d= and K. x d²= with r1 E [0, 1/2) and r2 E [0, 1]. Suppose w is sampled uniformly from gd-1 Then, when conditioned on w u >
O, the sample complexity in Theorem 5 reads
22r (d V 422 ( ve
0<r2 <T
n 2
n <m < 211,
(3.7)
(d(= (
dL+2ra
2r1 r2 <1
where 2 hides poly-logarithmic factors of
Remark. We have the following observations on the above sample complexity.
Corollary 6 demonstrates that structured data can lead better sample complexity when the right normalization used during training. This complements Theorem 2 where recall that spherical training dynamics ignores the structure data and the target direction cannot be recovered.
When g is polynomial of degree p. the lower bound for rotationally invariant kernels (including the neural tangent kernel at initialization) implies complexity of at least ((((())r)p) [DWY21]] Thus the sample complexity of Corollary can always outperform the kernel lower bound when p is sufficiently large and remains constant.
Three-step phase transition. Recall that in the isotropic setting Σ - Ig. the sample complexity of learning with information exponent s using full-batch gradient flow is ((d") for 2 3 [BBSS22]. The sample complexity in Corollary 6 is strictly smaller than (d') as soon as 1)rr/(s 2) < Furthermore, for any y > 0 it is at most ((d++>) soon r2 2 1 1/v((s 3) and 2r r2, in which case the sample complexity becomes independent the information exponent. Interestingly, the complexity becomes independent of 71 when 72 > 2r; 72 < rı, i.c. the direction of the spike becomes irrelevant when the spike magnitude is sufficiently large or small.
The three-stage phase transition Corollary is due to the different behaviour of the inner product (w, u) different regimes Tı and r2- When r2 < ri, we have (w", u) x (w", u). thus the initial alignment is just as uninformative as the isotropic case providing no improvement. Moreover, potentially large condition number may hurt the sample complexity in this case. On the other hand, when r2 < 2r1 we have (w". ④. x n(u, ("w", 0), and r2 > 2r, leads to (w", प) x w 0). thus large (u, 0) in this regime may improve the sample complexity.
7

---

Implications to Neural Networks and Further Improvements
4.1 Improving Sample Complexity via Preconditioning
We now demonstrate that preconditioning the training dynamics with E-1 can remove the dependency on x (), ultimately improving the sample complexity. Consider the preconditioned gradient flow
dw
dt = (-=n(ww) 2 'VR(w") with n(w) = =||8² w
(4.1)
We have the following alignment result.
Theorem 7. Consider the same setting Theorem 5, and assume the preconditioned empirical gradient flow (4.1) with number of samples
Z d
2(1 ve (
where 2 hides poly-logarithmic factors of d. Then, for T x Ts (T(((2)) ,)) In(1/E), we have
w 1
with probability at least 1 c,d didd"2 for some universal constants C1, C2 >> 0.
Preconditioning removes the condition number dependence, which is particularly important in the spiked model case where this quantity be large.
Corollary 8. Consider the setting of Theorem and assume we the preconditioned empirical gradient flow (4.1) for the (j., 0) --p-keddd model where (u, O x d= and x = d" with rı E [0, 1/2] and r2 E [0, 1.. Suppose w is sampled uniformly from Sd-1 Then, when conditioned on (w υ) > O, the sample complexity of Theorem 7 reads
d(d"-Ive-2)
0 0<12 <T
2
v6-2)
< < 2m1,
(4.2)
d(d(o- ve-2)
2n r2 <1
where 2 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 Odd") complexity in the isotropic case [BBSS22] when 12 > T1- Further, for any v > O, we can obtain the complexity of (dd++2) (nearly linear in dimension) when r2 addition to the remarks of Corollary 6, we note that the complexity is independent of both rı and r2 when r2 < rı (cf. Figure hard regime), i.e. the spike magnitude and the spike-target alignment have no effect on the complexity unless r2 2 T1.
p/(s 1) and > 2r, r 1/2(1 v/(s 1)) 12 2r1. In
Under the spiked covariance model, one could improve the above results by instead using spectral ini- tialization, i.e. initializing at 0, which can be estimated from unlabeled data. Assuming perfect access to 0 using the statement of Theorems 5 and 7, this initialization would imply sample complexity of Odd!+ +222 221-22221------- (00-((((11)))))))) -221v1))))))) without and (d]++((>- ()(2r: -12)v0) 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 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 needed to learn the single index target which we compare against the results of [BBSS22]. We highlight that layer-wise training in Algorithm frequently employed in the literature 2BES+22, BBSS22, DLS22, MHPG 233 and in particular [BBSS22] also used gradient flow for training.

---

Algorithm Layer-wise training of a two-layer ReLU network with gradient flow (GF).
Input: w" E T, A, À E R,, and data ( JE
1: Train the first layer weights W using the GF (3.5) or the preconditioned GF (4.1).
2: Normalize the weights 3 / 2/2
WIll for every S S m.
3: Let by ldd Unif(- A) and 0j = 1/m for i m.
4: Train the second layer weights z" via the gradient flow
da'
(0(= WT, .a.b)) (0) + Alla'11
Va
dt
2n
5: return (WT a b).
Theorem 9. Let g be twice weakly differentiable with information exponent 8 and assume g" has at most polynomial growth. Suppose & is the ReLU activation, Assumption 2 holds and we run Algorithmm I with w initialized uniformly over S6-1 For any > 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 1. Then, for A x VIn(nd), some regime of X given by (E.3) and sufficiently large T" given by (E.4), we have
((a..)) (ilo: WT a) v) S b))
C2(E CC(eee)C((++1/mm)
(4.3)
conditioned on (wD, u) > O with probability at least 0.99 over the randomness of the dataset, biases, and initialization, where Ci is a universal constant and C2 hides polylog(m, n1, d) fuctors.
The next result immediately follows from the previous theorem together with Corollaries 6 & 8.
Corollary 10. In the setting Theorem 9, E follows the (R. -spiieddd model, the sample complexity is given by (3.7) if we use the empirical gradient flow and (4.2) if we the preconditioned version.
We remark that for fixed E, the sample complexity to learn in the isotropic case is Ö(d") [BBSS22]. Under the spiked model, we assume that r2 is sufficiently large and r1 is sufficiently small as discussed in the previous section, Corollary 10 improves this rate either (3.7) when the empirical gradient flow is used without preconditioning or (4.2) with preconditioning.
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 Σ = Ia. where we have the following decomposition for population loss
R(w)= E.. [(r((ww x)) ")
_ 1012 + 1 E[y"] B[o((w, (g(())((()))))) x))]
(5.1)
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
d(w' u) dt
((d'((w', ())gg'((())))))))) a)) (1 (www,,, u) = (1 (wo", ,,)² u)² Σ Jj (wo", u w)--1
j23
where the second identity follows from the Hermite expansion; see also [EDB16, EBD19]. Assume a,8, 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", u) < 1, this signal roughly of
the order (wo u) u)=-1, As the uniform concentration error over gd-1 scales with v d/nn, we need
d( "0, u) 1(= ensure the signal remains dominant and wh moves towards u. When w is initialized uniformly over gd-1 this translates to sample complexity of x d", which 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 (x, u - model and ø ReLU. Using Lemma 12, is easy to show that with the spherical gradient flow, the alignment obeys the following ODE
d(w', u) dl
Kirs, 1
(w', u) (11-(ww,,)) ) ,w), 1)
E IE(o(((T))
(TT, 2))]
where o (wo") is introduced in Lemma 12. The additional ( term creates a repulsive force towards the equator (w', u) = 0. The presence of this term is due to the fact that unlike (5.1), the term E [wo("", ,²))²] is no longer independent of to and cannot be replaced by | When wD is initialized uniformly over s and ((1) S K O(d), we have (w u) x vk(ww, u). Furthermore, at this initialization (w)) 2 1/2. Therefore,
d(w", u)
80,, (Vk(W) u)) (wo .
di
2
Hence the dynamics is trapped at (w0', u) O(L/Vd) for all t >> 0 as long
O(d
To remove the repulsive force in the spherical dynamics, 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', α) via
directly normalize the input of
d (ww, u dt w'. Thus, the strength of the signal at initialization is of order x ()) which after controlling the error in the estimate of Σ and the estimate of popula-
Eld' ((w, 2))g ((u, ( Eu'")
(5.2)
where :7
tion gradient using finitely many samples, leads to the sample complexity n x dx(E)) 2 (w", .) (1)20--
Importantly, E can incude a much stronger initial alignment (w", 4) than the isotropic case ( (w", 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 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. 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 improve upon the isotropic setting, while also outperforming lower bounds for rotationally-invariinn 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.
contrast, we showed that an appropriate
We outline few limitations of our current work and discuss directions for future research.
While studying single index models provides 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 Correlational Statistical Query (CSQ) algorithm [BF02, Rey20J, i.c. an algorithm that only accesses noisy estimates of expected correlation queries from the model. Understanding the limitations of learning single index models under 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 induces. interesting to consider cases where we train a network with multiple neurons starting from more standard initialization which can help relax Assumption 2.
10

---

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-stairease property: a necessary and nearly sufficient condition for sgd learning of sparse functioos on two-layer neural networks, Conference Learning Theory, 2022.
[ABAM23]
Emmanuel Abbe, Enric Boix-Adsers and Theodor Misiakiewicz, Sgd learning neural networks: leap complexity and saddle-to-saddle dynamics, arXiv preprint arXiv:2302.11055 (2023).
[ASKL23] Luca Amaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro, From high- dimensional & mean-field dynamics to dimensionless odes: A unifying approuch to sgd in two-layers networks, arXiv preprint atXiv:2302.05882 (2023).
[ATV22]
Pranjal Awasthi, Alex Tang, and Aravindan Vijayaraghavan, Agnostic learning of gen- eral relu activation using gradient descent, arXiv preprint arXiv:2208.0271 (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.
[BBAPO5]
Jinho Baik, Gérard Ben Arous, and Sandrine Péché, Phase transition of the largest eigenvalue for nommull 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.
[BE 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.08 088799900708799 (2022).
[BES+19]
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang, Generaliza- tion of two-layer neural networks: An asymplotic viewpoint, International Conference on Learning Representations, 2019.
[BES+22]
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang. Denny Wu, and Greg Yang. High digeediimenional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation, arXiv preprint arXiv:2205.01 22455. (2022).
[BES+23] Jimmy Ba, Murat 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, 359- 395.
[BMZ23]
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou, Learning time-scales in two-layers neural nefworks, arXiv preprint 4xXi:2303.00055 (2023).
[BPVZ23]
Joan Bruna, Loucas Pillaud-Vivii and Aaron Zweig, On single index models beyond gaussian data, arXiv preprint arXiv:2307.15800 (2023).
[BS06)
Jinho Baik and Jack Silverstein, Figenvalues 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-paraaeterized Models using Optimal Transport, Advances in Neural Information Processing Systems, 2018
[Chi22]
Lénaïc Chizat, Mean-field langevin dynamics: Exponential convergence and annealing. arXiv preprint arXiv:2202.00 0000000..000.... (2022).
[CM20]
Sitan Chen and Raghu Meka, Learning polynomials in few relevant dimensions, Con- ference 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-Woodhhe Courtney Paquette, Elliot Paquette, and Inbar Seroussi, Hitring the high-dimensions notes: An ode for sgd learning dynamics on gims and multi- index models, arXiv preprint atXiv:2308.089 (2023).
[DH18]
Rishabh Dudeja and Daniel Hsu, Learning single-index models in gaussian space, Conference On Learning Theory, PMLR, 2018, pp. 1887-1930.
[DKL++33] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan, Learning two-layer neural networks, one (giant) step at time, arXiv preprint arXiv:2305.18270 (2023).
[DKTZ22] Ilias Diakonikolas, Vasilis Kontonis, Christos Tzamos, and Nikos Zarifis, Learning 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. 10633300..33333 (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 lemma, Proceedings of Advances in Neural Information Processing Systems, 2015,
pp. 1216-1224.
[FCG20]
Spencer Frei, Yuan Cao, and Quanquan Gu, Agnostic learning of 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.
12

---

[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 loffe and Christian Szegedy, Batch normalization: Accelerating deep network training by reducing internal covariate shift, International conference machine learning, pmlr, 2015, pp. 448-456.
[JGH18]
Arthur Jacot, Franck Gabriel, and Clement Hongler, Neural Tangent Kernel: Con- vergence and Generalization in Neural Networks, Advances in Neural Information Processing Systems, 2018
[JohO1]
lain M Johnstone, On the distribution of the largest eigenvalue in principal components analysis, The Annals of statistics 29 (2001), 2, 295-327.
[JWHT13]
Gareth James, Daniela Witten, Trevor Hastie, and Robert Tibshirani, An introduction to statistical learning, vol. 112. Springer, 2013.
[KKSKII] 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 two- layer 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 two- layers neural networks: dimension-free bounds and kernel limit, Conference on Learn- ing 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 Artificial Intelligence and Statisties, PMLR, 2022, pp. 9741-9757.
[O'D14]
Ryan O' D'ounelll Analysis of boolean functions, Cambridge University Press, 2014.
[PSE22]
Sejun Park, Umut Simsekli, and Murat A. Erdogdu, Generalization Bounds for Stochas tic Gradient Descent via Localized g- Covers, arXiv preprint arXiv:2209.0888 108951 (2022).
[Rey20]
Lev Reyzin, Statistical queries and statistical algorithms: Foundations and applica- fions, arXiv preprint arXiv:2004.0000 00557 (2020).
[RGKZ21]
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová, Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed, International Conference on Machine Learning, 2021.
13

---

[RVE18]
Grant M Rotskoff and Eric Vanden-Eijndee Neural networks as Interacting Particle Systems: Asymptotic convexity of the Loss Landscape and Universal Scaling of the (2018).
Approximation Error, arXiv preprint
[Sol17]
Mahdi Soltanolkotabi, Learning relus via gradient descent, Advances in neural infor- mation 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 (NeurlPS 2023), 2023.
[Te123]
Matus Telgarsky, Feature selection and low test error in shallow low-rotation relu networks, The Eleventh International Conference 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, Cam- bridge University Press, 2019.
[WLLM19]
Colin Wei, Jason D Lee. Qiang Liu, and Tengyu Ma, Regularization matters: General- ization and optimization of neural nets vs their induced kernel, Advances in Neural Information Processing Systems 32 (2019).
[YO20]
Gilad Yehudai and Shamir Ohad, Learning single neuron with gradient methods, Proceedings Thirty Third Conference on Learning Theory, Proceedings of Machine Learning Research, vol. 125, PMLR, 2020, pp. 3756-3786.
[YS19]
Gilad Ychudai and Ohad Shamir, On the Power and Limitations of Random Features for Understanding Neural Networks, Advances in Neural Information Processing Systems, 2019.
14

---

A Concurrent Works
In this paragraph we briefly summarize few relevant results concurrent our submission. [BMZ23] provided a precise analysis of the two-timescale dynamics in learning single index model with information exponent 1. [MHD considered the learning of single index target with s using neural network in the mean-field regime. [BPVZ23] extended the information exponent- based characterization of online SGD to input data beyond Gaussian. [DNGL233 showed that gradient-smoothed dynamics can improve the sample complexity and match the CSQ lower bound. Finally, beyond the single-index setting, [DKI,*23, 23,, SWON23, CWPPS23) considered learning low-dimensional target functions supported on k > dimensions via gradient-based feature learning.
B Background on Hermite Expansion
The normalized Hermite polynomials (hib20 given by (2.4) provide an orthonormal basis for (7)) thus for every f E L2()) we have
(f,h))m
where (f, = -B--wroonl((()))((())) We will commonly invoke the following well-known properties of Hermite polynomials. ] 1, then h2 Vihy-i. where h2 stands for the derivative of hy. Furthermore, 21 and 22 are two standard Gaussian random variables with E[ p. then Blh4(z1) ōzpp where dij is the Kronecker delta. We refer the interested reader to [O (14,,4 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 [BAGI21]. their setting, they assume the true data distribution P,, is parameterized by some unit vector u $6-1, and we know the parametric family ] 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(ww, ).). 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 if
a (0) = 0 dgs d'
15 <8 d=k (0) = -<< k=s
d d2k (2) SC
k E>s,VEe|11.11
where C, c > are universal constants. To specialize the above abstract definition to the Gaussian case, consider the setting where the input data is standard Gaussian NO, L2) and the problem is to estimate u E gd-1 given response variable y f((u, x)) with known f. Via the Hermite expansion of f, one can write
1 E [(((((,xx) f((u,
R(ww, u))
x)))
(f.hy/ u)² + const.
2
J21
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 æ (O, E). The spherical symmetry assumed in [BAGJ21] no longer holds. However, after proper normalization of weights, if we consider the population loss
(w, a:) 1/12
)
(u, æx) ||
R(w) = E
(w0 Eu, then R(w) R ||2/// reveals that for their results to hold, the proper symmetry to consider is the ellipsoidal symmetry, and
Indeed, close examination of the arguments of [BAGJ21]
1//2
15

---

the proper definition of information exponent is the degree of the first non-zero term in the Hermite expansion of R. which reads
- -(=)=-- ((E + const.
321
Once again, we can consistently define the information exponent to 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 closed-form expression for the population gradient. We refer to [Erd15, MHPG 23] for similar statements.
Lemma 11. Let f.g R + IR with weakly differentiable. Suppose 2 ~ (N(0,A)) Then, for any w, u E Sd-1, we have
[(w (u,, z))= = [[f((w, 2))g z)) ( 2)) z) - g ((u, (2, w)) H2..
Proof. Consider the conditional distribution z (w =) ~ (1, Σ), where
(w, 2)w and Σ L w
Recall that Stein's 's lemma (Gaussian integration by parts) states that when & ~ NH, E), then
Elg(2) (l)(z)) [V())
Hence,
Elg((u, x))²² | (w, x)] E[g((u, ()) (w. < w)w + (Ia ww E[g' ((u, z3) (w, 2)]u.
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 is weakly differentiable. Let 2 (O,
Define w
11/2 and similarly define u. Then w
V VR(w) = E Vo + Soig(w)))
where
(w)( El S((,, 2))g((( 2)) + o' ((w, 2))g (( 2)) (u, w))
1
+ 2
Σ 22
and
Eld' 22) 2))g ( 2))]
E (w)
1/2
Proof.
Notice that the population risk is given by
( (() ( uu, x)
2
E
3
R(w) E
18 [[(((W, [B][[],
x))9]+
|/
2
2
Notice that 1//22. By the homogeneity of ReLU, we can rewrite the first term
Ew][o|| w Eww
E ([((( )² w
16

---

Then we have,
VR(w) Ew E
(u, x²)
(o(lw,x))@
||
=10(w)
Note that ('w 2)) o' ((w, z)). Then,
v(w)= = 1/2 Elo' (w 2)) L (u, 2))=
1/2 [B|ø' ( (w, 2) g ((u 2) + (lo( 2)g((u, 2)) 6'(w,, 2))g ((u, (u, w)
Rld' ( z) ))) ( (u, z))] u (w z)g( (u, z)) 0 ((w, )g' ((u, 2 w))
{
E ull
|///301
where we used Lemma and the fact that d'(2)2 )(2). Therefore,
R((w) E ( + Eo.g (un))
which concludes the proof.
Particularly, the above lemma yields the following corollary for the spherical dynamics in the population limit.
Corollary 13. Suppose (wo 1:20 is solution to the population spherical gradient flow (2.5), is the ReLU activation, and E follows the (R, ( sp-ske model. Then,
d(w', u) dt
(0.0(w")) (w0", =))
1+
of (oo 0) + Ea (wo)) 01] (u,
((0,u) (w', (od,0)(('' [
(C.1)
1
C.1 Proof of Theorem
Plugging in f = u in Corollary 13, we obtain
d(w', u)
-fccoolut + kua (w (wo")"",, u)
)( (ue )0
(C.2)
dt
+ K
To prove the statement of the theorem, we will show that whenever
C
Kwo, u) Vd
we will have d(w',a)) < 0, thus when initialized from O(IVVd), (w", u) can never escape the saddle point near the equator (w, u) =
Recall from the properties of the Hermite expansion Appendix B that
g' Σ Vioghy-1 and Φ Σ Vi8,hz-1- 121
J21
Since we additionally assume Eig] ap by the definition of Vo.g Lemma 12 and the properties of Hermite expansion discussed in Appendix B, we have
"an(w')) Ejzo(i - 1)a,,,, 1 2)
2
and similarly
(oo(w²)== =- Σ
ja, By (w' -u)-11
128
17

---

Thus we obtain
d(wo,, u)
dt = 2F2FF(W)) (1 - (w',u)² u)²),
where
k(w0,, u)
F(w')= = Ba u) 1 (w" u) -)-
_ (m)
(w0 u)
(wo,
1 +
2(1 + K)
328
We proceed by upper bounding F. To do so, first note that
+ K W, u)
I
Σ 2 =
+ K
IIK
To bound the first term of F, we have
Ea (w'', u) 1-1 (wo, u) 5 (20', ; -2) (1) -)1²= j|ar (wo, (w'
u2²
j28
12a
s 11 12)l=²11 ( (w', z2)l
s 110''11 m)) 1//21 (20 u4)
Similarly,
()-1))a,Bj(ti u u))+t((
m) s 119'1119' x))x+))/2 1))
1 $28
Hence, for k 2 1,
F(w') 1110'1l1lll' [|llllllll( [1 xx) x)(--1)/21 w))² (1 + (1 (w) (w',u)² u²²
Suppose & < d/C² 1, then
F(w²)) I (w², u) 11 (1) 13/// (w' , ---
- -1/4)
(1 [1+ k)) (1+ 1/4
S {w0,, u)²
2110'1l, 2
d
Thus, for any x such that
2-2
d=
d
1SK5
C2
(80*-211
11x))-1
d(w' u
and any wt such that (wo, u)1 S C/vd, we have as long as the above holds true at initialization.
0. hence supezol (wo, u)) S C/Vd,
dt
Finally, we will show (w", u)) < C/vV 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 (10 u)² 1/d, thus by the Markov inequality
P((w",u) u) ")2 c20///) $1/C²,
hence choice of C' 2 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, will prove the following proposition.
18

---

Proposition 14. Suppose we initialize the spherical population gradient flow w. Let be the ReLU activation and assume
( == N...-ooo.)))())) 1/12,
and K. (w0 u/" flow on the sphere satisfies
(w) for a sufficiently large constant C depending only on g. Then, the gradient
(w" u) 2 11-E
if (wo u > (w") < 0
(C.3)
(ww, u) S 5-1+
if
whenever
1 a - 1/2 In(2/6).
T
(C.4)
Before proceeding to the proof, we notice that if we uniformly initialize w" over $d-1 then the typical value for ww,, u) is of order d meaning that the above proposition asks for A (2(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= u) > and (w u) < O are symmetric, thus we only present the proof for the former. Using (C.2), we can write the dynamics the sphere more explicitly as
d(wo" u) dl
BTd'((w" 2))g ((u, 2))] n(w' u)²
k(w0 u.) E E[[(((")) Wh, 2)) 2))]
n(w', u) 2(1 + k)
(1-(w', ,u)3
VT + K / k(w0,, u2)²
Our goal is study the regime of large therefore we will bound how much B₁ and Bg can deviate from their corresponding λ = oo values. In particular, we have
1 10
11
B1 2
+ n(w', u)² + n(w0 u)
x((0'
Furthermore, assuming (wt, u) > and (Ø, ().g)> > O, let cu(2t) the Lipschitzness of d,
Then, by
VI+N 11xtt// m2)2"
Bg (ww")) (b, g)n cx(w") ( (((((w";)))) - o((u, z) 2))))(((2,
( (0,g)) (xx(w") ((o0))lllll |||g|1 1/1l =u1l
cx (w') (ww))((0.))) (), (cx(w") (1lll|
||ll|| V 2(1 - (w', u)3
ca ca (w)()(() (6) 9)n, v 2x / (wo", u2)
noww", u)²
where we used क neighborhood around initialization, and will see below holds for all >> 0), then,
/1/22" in the last step. Suppose (w', u) > (which holds at least on |/''==''|
N(w", u2)
(w()
+ kiw", u1) 2
As result, obtain
k(w", u)
11 12/11/2 V2[1e11 1
B1 + Bg
2 (0, ,), ro(we,, u²²²
2(1 + x)
V Nssw 2n)² 2
19

---

Consequently, the lower bound of the time derivative of (w", u) becomes larger as (w", u) increases. Therefore, assuming (w) u) > O, we only need to control this lower bound at initialization. Assume
(11g9l>> + vS(ggl+>²
K 2 (wo, u) u)²{(a 112)
a
1/2
From this, we conclude that when (wo u) > 0 and (0. g) 1/2, we have
d(w', u) 1/2 (10', 11)²)
dt
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
(un, x²)
R(w) = E -
g( 0((
+ 2
1l + E
/2
Then,
(Ia VR(w)= 1/2
w w
Koh
(w, w. u
(D.1)
Σ 1/ /2/22 w
where
( (w, ②) = Blo' ((w, 2))g'((( ((T,, 2))1 = Σ (, (,,) 2)
(D.2)
$28
Proof. Recall from (3.1) that
(w, x))
VR(w)= = V
-
gl (it, 2))
1/2 w
1
Eow
Id
11/2 E(p'((((u) z))g( (u, 2))2]
// woll
1201/20111
Σ 1/2 /1(
W w
(Go.g (w0, ((W,
y + Vin (w), w))
(by Lemma 11)
|| 222/2w1) w|l
uD 1/2 (I.
[G (w, u.)u
where
V W U.) = Elo' 2))g((((())))))))) (T) ((u.))) (w) 2) 6'(( )'(w,))))'' (w, x)]
(D.3)
(the above only function of (w, τ.) due the Hermite expansion).
Given the closed form the population gradient, the proof of Lemma 3 is immediate by noticing that
d dt
(I. w u
Σ 1/2 dww
(D.4)
Σ 22/2
wil
dt
Next, move to prove Proposition 4.
20

---

D.1 Proof of Proposition 4
From Lemma 3, we have
d( w τ)
Cooo (wr,, u))) u
uं 1|
dt
chmin ( (
where Vi == ū w
U w.
The above inequality and the fact that (w σ) > 0 imply that
(w', u) is non-decreasing in time. Let
Ti = suplt > 0: (w', ②) < 1/2}.
Then, on & E [0, we have
d w u)
3cdmiu
Σ w u
4
and integration yields
1/2 (w0 a
8
Ti S OV
In(1 In(1/(2(1u)
1 u
8 2.
3cdmiu Σ
w U 8-2 2
8>
Therefore, Ti S Tx((w", 21) //min (E). For tort > Tí, we have
d(w0, u) dt
cdmiul (E)
(1 (w',, 2)
2
Let
T: stpppp > 0: (un", ,n) < 1 - $}.
Once again, integration implies
2--2
T2 S Ti cdmiu Σ)
In(2/(3E))
which completes the proof.
D.2
Preliminary Lemmas for proving Theorem 5
We first introduce number of concentration (and anti-conncnntaation) lemmas that will be useful for proving Theorem 5.
Suppose {a (2 lid. (O, and g : R + R satisfies 1g(-) S C(1 + 1-1)) for some C and p 2 1, Additionally, suppose {e(t) )==1 are i.i.d. e-sub-Gaussian zero-mean noise
Lemma 16.
1i=1
independent of (=(() ym - Let ( =g(((( 21) + (8) for some Sd-1 Then, for any q
with probability at least 4d we have
(
SC C(2 + a (2ldmdd))) Z In(nd")//2,
for all i.
Proof. Notice that (u. (0) -N(0, 1), thus
P (K u. ( 2 in(d)))) < 2n d
Similarly, by the sub-Gaussian and zero-mean property of (t))
P P((e(0) 2 a 2lm((nd))
< 2n 11 d
Thus, by a union bound, we have
T. (0) s 2 In(nd)) and
(00 S In(nd9), for 1 S S n,
with probability at least - 4d Using the upper bound on [g| finishes the proof.
21

---

Lemma 17. Suppose { are i.i.d. samples drawn uniformly from gd-1. Then,
3vd 8n
(., (e.zz1)
3d (12++((8/VVa)) se
P
sup WEB
i=1
Proof. Fix some E (0, 1). Let Nr be minimal e-covering of gd-1, onto No. Notice that by the triangle inequality and the union bound
Let in be the projection of w
P
sup
Σ C1(K 1=1
x 50) 2
sup WEN.
Σ 1
(1<00,
④0, 2000(())))))
$ 26
2 a
WES
==1
$ |N:|P
④ x
2
N
1=1
$ (3/c)²P
Σ
1
w x010 S
526)
IN
Moreover, due to [BBSS22, Lemma A.7]. ([ =) 2c)] = () S 2e) 8vde.
Choose
3vd 8n
By Lemma 25
a P
E1(K w. (w, > 22c 2 3d 2 23(iv(((v/[ - (3/0))
which completes the proof.
We summarize the above statements into "good event", characterized by the following lemma. Lemma 18. Let (2 p ud NNO, and (()) = x1)) (I I for every i. We say event G occurs whenever:
1. (0) S In(ndd) for all sisn. n
2. supwega Ei= ((T, $ vd/n) z dIn(n/Vd).
3. Amax (1 ((1)((0) - -1 vd/n
4. 1 - Amin
L 2(00 (00)
S vd/n.
For n. Z d, event & occurs with probability at least 1 Odd-).
Proof. The first and second statements of the lemma follow from Lemmas 16 and 17 respectively, The third and fourth statements standard Gaussian covariance concentration bounds (sec c.g. [Wail9, Theorem 6. where both statements hold with probability at least 1 2 d) 2e-d).
D.3 Proof of Theorem 5
1/2 w 1/2 wll
We begin by recalling the definition w copy here for the reader's convenience,
and the finite-samples dynamics (3.5), which we
dw
= n(w") V w
(w", x0)))
(()
dt
71
1/2
20
where n(ww)) = //20/2 w01 113 Moreover, via chain rule, we obtain
w
1/2
dw dt
(I.
Σ
dw
(D.5)
1/2 w0ll
dt
22

---

Let ((0 == -ラ- 1122 -1/2//2 x()) Σ 1/2 1//
x00)
1/13/11// (*) Then,
dw
(w',z(') x
= (I J [L v
ø'
y
(6)
dt
| 1/2 w0
==1
To simplify the notation, define
D 2 = (Ll
w 2 )o.
Then
d n v2
v(w')
Σ のk
( w 210) )) (0))
dt
We can decompose the above dynamics into a population term and three different error terms in the following manner:
d( " ń)
=(u(w"), Bs =)))z()>
dt
+
v(w0",
1
( ( ((w0, 200 )) 1012 EzN Lo ((u',2))yz
n2
=1
Φ | (100,2200) ((( x(11) (2)0)))))))))))))))))))))))))))) (2,00
1=1
1 ø
(( (1,,((0) ( (000 10) D (
n
We will proceed in three steps. In the first, we bound εg. the concentration error. In the second, we bound & and Eg, the errors due to estimating E with Σ (i.e. replacing 200 with ( 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 & Let K x In(nd")o/2 and notice that on event g
we have |y(0 S K for all i. Let yK = yy1 |l S K). On the event G. we have yk ( for all i, & (=(w"(>) )
and
where
w. )1)) yk (1) (0) Ez,y 10' (wW, z) 2))z
Thus, our objective is to bound || Anll uniformly for all w E gd-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,
Σ Φ ( w. z(1))) (1) 20))) 2))))))) .0) Ezzloo 2))2 (z,))) w)
sup (A., v) W. 1.
sup W
1=1
Ez.y ((1,((2)))( ()) 10vl K)].
By the Cauchy-Schwarc inequality,
Exzyyo'' ((w, 2) 2))2(z, 1)1((=) > K)]l < Ez,y ( 2)) w) 1/2
E[1(b| K2)
X B [ |11 1/4 E
1/4
P(yy > K)
S 1/2
23

---

where the last inequality follows from Lemma 16. Hence,
n
sup nn v) WE
sup 40 10.
(w
(10 32) (0 .0) - -Blo 2) yyk (z, v)
+ O(d
Next, we need to establish uigh-probability bounds via covering argument. To simplify the exposition, define the stochastic process indexed by w gd-1 and v Sd-1 via
x w o ((w,, (1)
yK .0)
Fix some Ear, fu > O. Let .. and θ be e and E, coverings of sd-1 and let is and á denote the projection of onto O and of t onto O₂ respectively, then
1 W
E[X = sup Σ 1=1
[(xO
(1) w.
x-i
(1)
sup
+
w.
w.
1=1
1=1
Ex y (x - Xw + Ex,y (x
X
x( (
n ū.
We bound each of the terms using Cauchy-Schwartz, Specifically,
x
(()
(( (w, 20))) (0)
Σ (2(), 100000000 -8)- S Ken.
72
w.
72
where we used the upper bound on the operator norm of 1 Ei=1 20) from Lemma 18 together with the fact that 2 d. Similarly,
Ex (x., Xw,el S Ex,v [o'( (w, 1/2 Ex (z, -} 1/2
S KE
To bound the differences when replace with w need to make distinction between ReLU and smooth activations as the respective arguments are to some extent different. When o' is Lipschitz,
x00 m.
-x(1) 60.6
=
(2) wu
(d) - 2(000
(2))) 02) S Kew,
n
n 1=1
i=1
and
Vi ((w, 2)) - ('((0,, 2()x)= 1/2
E E[ 0) 1/2
N - X
<
E
S Kcw.
Therefore, for a smooth activation & we choose G. = Cw = v d n,, and obtain
x Ezy [x S sup T
1
x (t) X E. (X + (Vd/n)
sup 20, O
6, n
When & the ReLU activation, need to show that the sign of the preactivation changes only for small number of samples when we change the weight w to i.. Notice that
sign ((T, A sign ( @.. z0))
< ( ④ w, x (0)
S Cw.
Choose e x vd/n. n. On event 9, we know from Lemma 18 that at most Oddnn(n/ (n/Vd)) samples can satisfy the above condition. Therefore.
Recall that (+) = 200 f101 (1
x((1
(ORP(((( - )(((w,x())))
1 Σ (1) 8)²
w. 10.0
<
i=1
din(n/vd)
24

---

and
E X Xmm,)] s IE yk (1) (w 2)) o' ((ir, z))) =))) 1/2
=
B B1(e, 03/ 030]12
(iisig (w, z)) + sign( (w, 2))) 1/2
KP( (w. 2> Ew) 1/2 )
/K /2
where the last inequality follows from the anti-concentration on the sphere [BBSS222 Lemma A.7]. Thus, for ReLU we choose Ev x vd/n and em x vd/n, and once again obtain
sup 8
1
n1
sup Ez,y [m
x 8
E. Pzzy(Xmnb + Ó(vd/n).
0, m2
It remains to bound the term
n sup 0
x( i
E[
72
Notice that for fixed w, é, X is sub-Gaussian with sub-Gaussian norm O(K). Thus, via the sub-Gaussian maximal inequality [VHI6, Lemma 5.2].
x
sup t Y2
with probability at least
Consequently, we have
A S Ö(Vd/n+ t
sup
WEBd
with probability at least (d).). Assuming that n grows at most polynomially dimension
and choosing a sufficiently large q, we have supwesa-1l Anll s ((va/n) with probability at least 1 (.
Finally, by Lemma 23,
||
< Amax Σ
Amax E))
(D.6)
with probability at least 1 e-2'/2 Combining the above with the bound on || we have E -mmax (E)O(v( with probability at least - O(d T). which concludes the first step of the proof.
Step 2. Bounding the error due to the estimation of Σ, i.e. & and Eg. Recall that we are (i)
considering the event G, thus (0)) VK We control each of the error terms separately. We begin by Ez, where by Sauchy-Schward
1 i=1
Φ ((w (00) 0' ((<00', (000) (8) (00,'(()))
E
n
d' ((0)))()))) 8'( ((w,,,x0))}
1 n
yk (z(1) xz())))w'))²
(1)
n
i=1
1
Σ きt 2(00, o ((w0 (00 }²
- ||
where the last line follows from Lemma 18 and the fact that n. 2 d. When ø' is additionally Lipschitz, we have
& ス -KIp((l"
1
w 2(0
2(1) z(0)
n -=1
25

---

Moreover, for any w E gd-1
7 71
1
W. 200
=(0)²²
2
m Σ
(0 20) || --1
02/² )200)
n1
(0) (0) T
1 -5--- 1-----11 1/111111111 D2/2/22
22
where the last inequality holds with probability at least 1 - 22 on the event of Lemma 24. Hence for smooth activations we conclude
& Kllw("))ll d//.
When is the ReL.U activation, we need more involved argument to control &g. In particular, we will show that for any W. at most only (dd datapoints can have sign( Notice that
(w, z 2))) # sigu ( प x0 ))
sign w. ) * sign (N,, (w, x((1)
W (00
- 1 (0) 2(1)
w. z(1)
S La Σ 1/2 /
(D.7)
where z(1)
20 |
is distributed uniformly over gd-11 From Lemma 24 have [[
1/2 Σ 111
with probability least 1 2e d. On the other hand, from Lemma 17 we know with probability at least e-d, for any w E s most Ó(d) of the labeled samples have (w, > S vd/n. Recall that n'' 2 n² when using the ReLU activation. This is precisely why make this choice for the ReLU activation, as we need to balance the RHS of (D.7) which is of order vd/n' with the LHS of (D.7) which should at most be of order vd/n if we want to ensure only ((d) samples satisfy the bound. When n' n² we can balance these two terms, thus with probability least 1 3d- the sign change can occur for at most (d)) many samples, and
(( (0," s 0 a m
In this case, we end up with
& K )1( d/n).
Bounding &g for ReLU and Lipschitz o' is identical. In both cases, by Cauchy- Schwartz,
8 ((w ((w,2(1)) 2000 22) 2)
1 z(1) 2(0) ulwaa) 2
Eg 2
2 K
- 1/2 1/2 w(w\\\|
which holds the intersection of event G and of Lemma 23. At last, using the bound on ||((w") from (D.6), we obtain
A 2 -Amax (E)ō(Va/m)
with probability at least Od--)).
Step 3. Analyzing the Convergence. As a result of the previous steps, we have established
d (w', u) 2 (u(w"), (11'((='22))9z) Amax (d)m)))/n)
dt
26

---

Thanks to Lemma 11, we can write
Elo' (ww z)) El z) ())g((()) x))=
6 ( (w, U v u ( (ww² u )w
where So.g and vass were introduced in (D.2) and (D.3) respectively. Recall the definition of v(w"),
v(w') (Io w'w' )ā(Lo w' "w"" )
Therefore,
d 1 ④
2 Co (ww,, ) (Ia w'w )E(I "' )u Amax (E)((vd))))))
dt
2 ol w u) u)/-1 ( ui Emi) Amax ()O(vd/1
(By Assumption
cmmin Σ w s-1 (w' F) Amax (Z)O(vd/n)
where T._
u,, wh w Moreover, from Lemma 23, we have
Lr( (r(
Amin
$ 2 Amin ())
Ammn E
dx( Σ
Amia EE)
V
n'
with probability at least - - nn'88 Hence, for n' 2 dxx(2) )) we have Amann (1) 2 Amin [),
and
consequently,
d W ④ dt
Amn (w) 11)=-111 (w", u)² ) Xmax (E)(())/))
where c' is universal constant. Notice that the first term in the RHS above denotes the signal, while the second term denotes the noise. We want ensure the noise remains smaller than the signal throughout the trajectory, which leads to the convergence of w' to Z. Notice that the signal term first increasing, then decreasing for (wr' πτ [0, 1]. Thus, suffices to ensure the noise smaller than the signal on the two ends of the interval, i.e. at time = O and at time / = T where (w"" u) = 8. initialization, this condition leads to
N (dx(E))(x u. 211-))
and at time leads
n Cd)(E)²
where C hides constant depending only 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
d( 1(
Amin (E) u)
W F)
dt
where c" is universal constant. The rest of the proof follows by integration and is identical the proof of Proposition 4 in Appendix D.I.
D.4 Proof of Corollary 6
The proof follows immediately from Theorem 5 and the following lemma which describes how (w, u) bchaves under different regimes of r1 and r2-
27

---

Lemma 19. Suppose Σ follows the (k, -spiiedd model, w is sampled uniformly from Sd n2' 2 d.
and there exist universal constants C2 C2 Cz. C33 > 0 such that
C2d K S C2dd²
and
Tl (u, 0) S Cid
for ri [0 and r2 E (0, 1]. Then, conditioned on (w² u) > O, with any arbitrarily large constant probability 1 8, for sufficiently large d (that depends on 8) we have
d-1/2
0 12 <n
w u 2
d=2-111112
TI r2 21
(D.8)
11/2-1//2
2r1 < 12 < 1
Proof. By definition,
1/2 Σ
w0 /1/2
1
1/2 w w |||
1/2 u
Recall that we are conditioning our arguments on (wo, 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 ( ww", u) in terms of d, rı, and r2.
We begin by observing that
w/
S
ΣΣ 1/ x
3wll $ SE 1/ //// |//1||1
where the last inequality holds on the event of Lemma 24, which happens with probability at least 1 2e -d Consequently.
1/2 wD
(wo, En) + (w0, (2 =1/3)
M
1/2
w
1//22ll
Σ 1/2 2 || ||
Furthermore, due to the Markov inequality,
C1 (100) (x0,0)²
$ 11/OI.
Similarly (by conditioning on Σ)
PP - -1'3""2-''3'000 )² 2
C 11 1/2
1/2 uull
< 1/0
d
Additionally, on the event of Lemma 24,
|1 1/2 - 14/2 u|| < || 1/2 2-12 Eu|| S Vd/n Eu].
Therefore, on the above events, for some absolute constant C' > 0
(wo ' Eu) CI Eull EuIlv/In'
1+0,0
|Eu|
(w) u) + k2w00 0)(u) 0) C(1 no(u,, 0)1) 1/m
VI+ C.C1 + k/u, 0)2
Recall that CTd"a S < Chdla and Cad S (u, 0) < Cha (notice that changing -0 does not change the spiked model of Assumption 1, thus we can assume (w, 0) 2 0 without loss of generality). Then,
C(1 n(u, 01) 1// (wo ④) 2 ' u) + w .0 (u, 0)
(D.9)
+ C,C 1 d dr
28

---

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 r2 2r,, the denominator of constant order. If additionally 12 < ri, the dominant term in the numerator is (w) u) and (w= u) = 1/Vd, otherwise the dominant term is x(w", -0) (uu 0) and (w", , = d==- On the other hand, when r2 > 2r the denominator is of order d'2/2-ni, and once again the dominant term of the numerator is o(ww: 0) (u, 0), therefore (w), म) x d 1)/2, Using this intuition, we analyze each of the following regimes separately.
Case 1.0 < r1: In this case, by [BBSS22, Lemma A.7] we have (w u) 2 c/Vd with
probability at least 4c. On the intersection of all considered events with (wo 4) 0, and for sufficiently large and n' Z d. we must have (w) u) > (otherwise (w) 2 plugging the values in (D.9),
< 0). Thus by
VCICs - C(1 + "1
1 Vd
(D.10)
Va C2 d d
The intersection of all desired events and (w), π) > 0 happens with probability at least 1 - 4c 2/C, - 2e, thus conditioned on (w"; u) the probability is at least 1 -c 4/C, 4e Choosing sufficiently small c, large C1. and respectively large and n' 2 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. ri < <12 < 2r This time we use the fact that (w", uu) S v Wi/d d with probability at least 1 1/1, and (un0, 0) 2 c/ ad with probability at least 1 4c. By an argument similar to the previous case, for sufficiently large and n' Z d. (w², 2 > O implies (w) O) > 0, hence by (D.9)
-VG + cC2Cgd'2 - C(1 "2 /d/n
T T
d
(D.11)
vavI+ + C.C C dr
with probability at least 8c - 4/C1 - -d when conditioned on (w u >
Case 3. 2r1 < r2 < 1: Once again recall (D.9). To bound the numerator, we repeat the exact same argument in the previous case, thus
-VC1 + cC2Cyd'a - C(1 + +CC 2 d
(D.12)
V1+C 1t 2n1
which finishes the proof of the lemma.
E Proofs of Section 4
E.1 Proof of Theorem 7
We recall from (D.5) that
1/2
dw
(I.
1
Σ
dw
dt |Σ
w
dt
Furthermore, the preconditioned dynamics of w given by (4.1) reads
n(w²) ) 1/2 (I, w "w'"""
dw dt
Φ ((wr,, ;00)) y (())()
w
i=1
where we recall {0 = 5 112 (1) Plugging in n(w')
1/2 2| yields
Σ
dw
(( ( x(0) ) (0)
= (Ia t T
dt
1=1
Σ ø' (ww 20) ) i=1
(I. wh W
29

---

The rest of the analysis is identical to that of the proof of Theorem 5 in Appendix D.3. Specifically, by defining
u'_ -ww,, uiyw',
have
d 10
=(u'
Ez,y [0'((",, .2)))22)>
dt
0 (( ( x(0)) 2(0) - Ex, [o' (d'20,, ;))
u
71
i=1
=18
11
お ((T0 (1T) 0 8((1T, ,2(1)) 1 (2(1) (z))) σ
1=1
=
n
お' (( (w,, (.)""))) (0)0 =0(0) u1
As long as n 2 d, n' n for the smooth case, and n' 2 n² for the ReLU case, the first two steps of the proof of Theorem 5 in Appendix D.3 implies that
E, A A &3 N - -I<1l/((///)) -0lVd//))
Once again, we apply Lemma to obtain
Elo'((,, z))zz] S l W1 u u vo (w', u) )
with So.g and u given in (D.2) and (D.3) respectively. As a result,
d(w u) dt
(o w 22 2 - Ö(Va/n) 2 c0
(1
प प. d(vd/))
(By Assumption 2).
We need to ensure the noise term, i.e. the second term on the RHS remains smaller than the signal, ie. the first term. The signal term attains its minimum at either initialization = 0 or at the end of the trajectory & = 7 where (w=" u) = 1 - e, which imposes the following sufficient conditions on n. Namely, at initialization we require
n 2 Cd w x)2 2(1-s)
while at t t= T we require
2 Cd/e²
where C' hides constant that only depend on and polylogarithmic factors of d. Hence, we obtain
d(un' u) 2 c (w'' u) (1, - ( (1- x
df
for some universal constant d' Appendix D.1), for
0. Via integration (similar to the proof of Proposition 4 in
Ti = sup(t 0: (w', σ) 1/2}
we obtain
Ti S To T)),
and for
To supft 0: (w' π =]
we obtain T2 - Ti $ In(1/)) which completes the proof We conclude by remarking that the proof
of Corollary 8 is immediate given Theorem 7 and Lemma 19.
D
30

---

E.2 Preliminary Lemmas for Proving Theorem 9
We will adapt the following lemma from [MHPG 23], which provides non-parametric approxima- tion of via random biases.
Lemma 20. [MHPG* 23, Lemma 22] For any smooth g R - R and > 0, let g. -
be smooth function such that ī(=) = g(=) for |=| S A and 5( 0(-2A) -
0. Suppose
{b; 1 lid. Unif( 2aif(-2A,220 2A), and let A. = A supielszal (2) Then, there exist second layer
weights (a,(b;)) with ||a]| & Aw/ such that for any fixed E |- Al and any 8 0, with probability at least å over the random biases,
m
In(1/0)
((b;)((=+b) by) g(b;)
SAA.V
m
where é is the ReLU activation.
We use the above lemma to show the existence of a second layer with Ö(1 ((1//m) norm with training error of order (1/m)).
Lemma 21. For any < 1, suppose W. 2 1 - Then for any q 0, sufficiently large d, n 2 d/E², with probability at least O(d over the random biases and the dataset, there exists a second layer a with ||a]| S Ó(1/vm) described by Lemma 20 such that
2
a W. 21)) b,) -y11)
811) 210) [ +0(1/m++)
where ó is the ReLU activation.
Proof. We begin by replacing w and 2()) with it and (0) Specifically, via Jensen's inequality,
2
0 w. 2(0) +b;) (()))
ad (1))())) + (b)) =( ((TL x(0))
=1,
n
9 s((TT, (((T, 2(1)) )){
y
1=1
=
2
Σ a 2 w. 2(01 (07,2(1))+0) +b,) o((TT, x(0) bb)}}
1=1
Y22
00 w. 200 +b,) -0((T,, (,)) +b,}}
((,,,((0)+bb))
j=1
We bound each term separately. For E,, we can invoke Lemma 20 which implies that each term in the sum can be bounded by ((1mm) with probability at least 1 1/(n""), thus by a union bound, with probability at least d over the random biases.
& S (11m))
By sub-Guu Subb-uassianity of (()) (hence sub-exponentiality of e for 2 (with sufficiently large constant) we have
& E + Vd/n,
31

---

with probability at least 1 e For Eg. via the Lipschitzness of ReLU and the Cauchy-Schwartz inequlllty we can write
Õ(1) E
"
w (0)
=(()) S -O s
where we used the event of Lemma 24 which happens with probability at least - 2e and Õ(1) represents a constant that depends at most polylogarihmmcally on d.
Finally, we bound the last term. Once again via the Lipschitzness of the ReLU activation and the Cauchy-Schwarle inequality
(1)
E
S
28
U. 2(1)
S 6 Ö(E),
72
where once again used the event of Lemma 24. On the intersection of all desired events, we have
a3 w. =(0) + bb²) y0)) S [[[]]+ + + df d/n' +
We conclude the proof by noticing that 2 n² and n 2 d
Additionally, will use the following standard Lemma the Rademacher complexity of two-layer neural networks, which in particular a restatement of [MHPG* 23, Lemma 18] in way suitable for our analysis.
Lemma 22. Let FF be class of real-valued functions on (2,y). Given n samples (2(0) , define the empirical Rademacher complexity of F as
)00)
(R) = B(
sup (z)((0)) fEF n
where (st) are i..dd Rademacher random variables (i.e. +1 with equal probability). Suppose F is given by
2
m bb)) w
F (z.1)= +
2x)
AC ||ll| S ra/vm, 1byl S rb, r1, Ai m
for some fixed gd-1, Suppose [2( rid NO, Ia), and suppose |0'l 1. Then.
By (F) < 2v20 + rb)ra
Vn
Proof. See the proof of [MHPG+2] 23, Lemma 18].
E.3 Proof of Theorem 9
Throughout the proof, we will assume (w, π 2 E where we recall
1/2 w
11/2 u य =
and
1/2
/
wl
From either Theorem 5 or Theorem 7, assume (w. τ 1 with probability at least - O(d=)) 9) for any fixed g > 0. For simplicity, let
((;: = Σ a)o((W, + bb)),
32

---

and similarly define Ú(z; U). We define the following quantities, R(w) = zz (ú(2; w) - --)² and
R(w) = Ex (i(z; - )) -
(E.1)
and similarly define their empirical counterparts,
R(w) 1 Σ 0( w -0 and R 22 = ((2(00 -1(0)
(E.2)
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(T). Then, via a generalization bound, we show that the R()) can be upper bounded by R(u)). Finally, we show that R(w) can be upper bounded by the training error, i.e. R( wi), and convex optimization of the last layer can attain the near-optimal value of this training error which is bounded by Lemma 21.
Step Bounding R R(w) via R R(u). By Jensen's 's inequality.
Ez.y (0(2; w)²]
E (Ú(ž; w) )(); +3 (0(2; w) 01 0(2; +3 (0(z; -3²
Suppose |aa|| S ra/ /mm. For the first term, by Lipschitzness of and the Sennns--cccaarzzzeuuuiiiii
27
Ex (0(2; - ((:
Ex
Σ asfo ( (w, +b,) -0 (((T, +b,))),))) +bb))}
< r Ez (w0, 2z= - 2)²1
s 2 - 1 1 z r2d/n',
where the last inequality holds with probability at least For the middle term, via a similar argument,
2e 2e-d on the event of Lemma 24.
E (ú(z: w) - 0(2; u)) S r Tr (w u, 2) S 2ra8.
In what follows, we will restrict the analysis to the case where ra = Ō(1). Therefore, we have
Ez (0(2:w) -) S 3 3 (0(z:T) v)²] O(d/n'++ +
Step 2. Generalization: Bounding R(u) via R(u).
Define the event
E = [ z>| V |e| S v 2n(nd))))
and similarly define E by replacing and with z (t) and et) respectively. Via the Cauchy-Schwards and Jensen inequalities
Ez.y (((2; u) w)²]
Ez.y (b(2; 2) - (E)
+ Bz,, 11'10)
< Ex,,
(0(2; T - 1E))
+ [] 1/2 (E 1/2
Moreover, [[ S 1, E., v/ |0(z; [) < Ō((), and and ( _ 4/(dd)) (via standard sub-Gaussian
tail bound). Consequently.
Ez,y (0(z; -y)l s Ez,y -)) + (dm-9)
Let
m
((210, 1000, 9 a. b) =
Σ a
T, 2(1) + bb,) - y 1(E).
33

---

Notice that य is fixed. Then, by standard symmetrization argument (see e.g. [VH16, Lemma 7.4]) and Lemma 22
1 Σ ((2 , b) - 2B (nn(F)
B
sup
Bz.g (e,le(zz3;;, b)l
((VIn))
where (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
sup |allSr./
B (2 a, b)]
Σ
a. b) S O(vd/n)
with probability at least - 2e
Therefore, conclude this step by noticing that
Ez,y (b(z; T)) - y)² - (1200 [ - v10) ((E + m
1=1
-
Σ
[( (1) ( - y10)² + OlVd/m),
with probability at least 2e
Step 3. Bounding the training and finishing the proof. This step is similar to the proof of [MHPG*22 23, Theorem 4]. For conciseness, define
n > n
(., v0))
R(a)
Our goal is to choose suitable such that the minimizer and (a( -((a)= + +/1aa11//2.
a" = arg min R(aa) + //2
aERT
satisfies ||a"|| S r.//m while the value of the above minimization problem which we denote with R2 does not significantly exceed
min Sra vm
R(a).
|a
We argue that the suitable choice for À is
m
21
+ mE +1 = (mm + me + 1).
(E.3)
r
Let R. denote the minimizer of the regularized problem and à == arg /in R(a). From
Lemma 21, with a proper choice ra @(1). we have
R( S E[e2| + Õ(1/m +s).
with probability least (d 9) over the biases and the dataset. Note that a" is the minimizer of Rr, we have
A|1a*112
R(a)) + A||ā||2
R(aa)
2
and in particular
X \|aa*11
\|āa|
< R(ā)) +
||00| S 5(11/m)
2
2
and
\|
R(a") < R(ā))
< E + 0(11m +
34

---

Let {a' be the solution to the gradient flow of a. Then,
alla' a²l 2a2aa'''aa a""a'aa", VRi(a")),
dt
and by the first-order condition of strong convexity
a', a' aa' VRi(a')) 2 Allaa a a=,
therefore
aa" a"'l se 2XT -AT
l|a" - a²12
As the training error (of the regularized problem) is X-strongly convex in a, by applying the standard condition, gradient flow for training a obtains
Ri(a²) R S (Ra(a") Ri) RRi)e= 22T e
Furthermore, since
||a0| 112 - ||a" 112 S I||aa" T' aall 11a" - -a" 22l1aa a"" -a"Illa"||,
we have
R(a" - R s 22 0 a "a|||a lle 22" + (Rs(aa) - -Ri) e 222""
Consequently, choosing
aa aa" a
In 4
In ) --i))
T
(E.4)
X
2X
implies
R(aa") S RD
+ and || s 2| - Õ(1/vm)
Therefore
Ria ) Õ(1m +
Recall that
R(a" )(a-)= Σ ( ()) y/t) 2
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(w)).
Thus the last step is to upper bound R((u) via R(w). To that end, via Jensen's 's inequality
)0)
2
3
(E20) w (0)
7
(1)00 w) -D(2(1) w)) 2
+
i=1
Σ
((000) w) 0(2(()) ui)
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,
) ()))
-
1/2
/2
1
Σ
1 O(d/n))
71 where the last inequality holds with probability at least Similarly,
2e -d (over the event of Lemma 24).
0(2 (1) w ;(=(0;)))
2 x))
5 ra wo w < Ō(s).
71
35

---

Putting the bounds back together (recall n' 2 2 de )), we arrive at
R(T) S R( + Ö(s + d/ni - R w + O(e).
Combining the result of this step with the two previous steps implies
R(w) v -| + (11mm+
with probability at least 1 - O(d (when conditioned on (w", प) > 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.
Lid N( E). Ler E E xl0x(0 x0) Lemma 23. [Wai19, Theorem 6.1]. Suppose (x(0)n' Ji=1 Then, for n' 2 tr(E)/ )
tr Σ
Amax Σ
< Amax (
5 nt' rmxmax E
with probability at least e Furthermore, for for 2 d,
tr(2) (2) nimin (2)
Amin D
Amin (Σ)
4
with probability at least 1 1 e-n118
Lemma 24. Suppose (2(0) mn ltd N((O let (+)
1/// 21) for some invertible E, and define
n'
Σ = n'
(0) (
x
i=1
Then
[1 - 1/2 Σ 1/2
V I[a - E 1/1
1//2
S
with probability at least 1 - 2c -d
Proof. We have
I - Σ
Σ 1//
Amax
Σ
1/2 E1/2 -1
V
{
Amin Σ -1/2x112)
Amax
1/25 22/2)
x
1/2
Amin
1125 1/122 (1//351'>1122)
Amin
Σ 1/1/255 1/2)
-1/2
V {
--max
E 1/2/)/=-111 -1/2
-1/2
-1/2
n'' M
Amin
1
(1)
1 V
1
Amax
7
1
(1)
(1)
i=1
i=1
Similarly,
1/2
1/2
n' > (1))
[ - 3
1/2 Σ 1/2
Amax
T
V
- Amiu
1
(4)
(6)
n2'
i=1
1=1
36

---

Moreover, by [Wai19, Example 6.2], we have with probability at least 1 - 2e 2e-
a' n Σ=
n' (1) (())
1
Amin $1+(2+1 क. and
2 11-(v22 \ m.
Amax
11' ==1
Thus, for n' Z d (with a sufficiently large absolute constant), we have
[1 - 1/2 V I - -
n1
with probability at least 1-2e 2e-d
Lemma 25 (Chernoff's Inequality). Suppose X1...., Xn are i.i.d. Bernoulli random variables, and further assume that E Xi) S
Then, for any 8 2 1.
P Σ
X; x u(1+ Se-u0/3
(F.I)
Proof. The proof follows from standard Chernoff bound. From [Ver18, Theorem 2.3.1]
TP Σ Ex2 s ((+))()))))(())))))
(notice that the statement [er188 Theorem 2.3. holds true even when EE X,] = y is replaced with E, X,] pi). We conclude by remarking that 8 (1 + 8) In(1 + 6) S -6/3 for 2 1. O
37


Reviews:
[{"review_id": "liXNYTfXIg", "reviewer": "Reviewer_XbBV", "summary": "This 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.", "questions": "$\\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?", "limitations": "Not applicable", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "${\\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", "weaknesses": "${\\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and positive evaluation. We address the reviewer's concerns below.\n***\n**Impractical training procedure**:\nWhile we agree with the reviewer's statement, we emphasize that two-step procedures have become mainstream in the theory of neural networks, mainly because they allow for relatively more in-depth analysis (e.g. [DLS22,BES+22,BBSS22,MHPG+23,ABAM23]). Developing theory when both layers are trained simultaneously is certainly interesting and more practical; yet, existing works in this setting (e.g. large-width mean-field regime) are not able to characterize the generalization behavior to the extent we did in our paper. However, extending our results to such regimes is an interesting direction for future research.\n***\n\n**Difference from prior works**:\nOur paper distinguishes itself from prior work by focusing on structured input covariance rather than isotropic, which allows us to obtain a sample complexity that, unlike prior works [BAGJ21,BES+22,BBSS22], can be *independent of information exponent* and almost linear in the input dimension $d$ under a spiked model. \n\nMore generally, by focusing on the spiked model, we reveal the existence of a *three-stage phase transition* based on the spike-target alignment and the spike magnitude, with each stage representing a separate regime of statistical complexity of learning the target. Furthermore, we establish how preconditioning can help generalization, an observation absent in the non-linear feature learning literature due to concentrating on isotropic data. \n\nNote that we also establish the *necessity of weight normalization* when learning under such covariance structures, which demonstrates that the covariance-agnostic training procedures used in prior work will provably fail in this setting. We discuss these differences in the paragraph starting at line 170 and throughout the introduction (lines 55 to 78), and we will further emphasize the distinctions with the existing works in the revision.\n***\n\n**Practical implications**: While our primary objective in this paper was entirely theoretical, there are still immediate practical implications. First, note that our results can be interpreted as “*neural networks benefit from additional structure in the input data*”. This is particularly interesting given that, in practice, principal directions of the input covariance contain non-trivial information about the target [HTFF09]; thus, *data whitening* as pre-processing may hurt generalization. Second, our work proves the necessity of proper weight normalization when the data is anisotropic, which provides theoretical justification for weight normalization practices used in deep learning.\n***\n\n**Differentiable $\\phi$ and $g$**: For recovering the target direction, smoothness does not affect *supervised sample complexity*, which is the main quantity of interest. However, smooth activations allow for a better *unsupervised sample complexity* (number of samples used to estimate the covariance), as discussed in line 252.\n\nIn section 4.2, we simply use ReLU since its piecewise linearity simplifies the function approximation needed to learn the link function $g$. In principle, Theorem 9 can be restated for any activation that is a universal function approximator, without any change in the sample complexity.\n\nAs for the link function $g$, further differentiability will not provide any benefits in theory since the main property of $g$ that determines the sample complexity is its information exponent.\n***\n\n**Extending the analysis to multiple-index models**: This is a very interesting direction for future work. We believe that in principle, our results should generalize to $k$-index models with $k = \\mathcal{O}(1)$ with roughly the same sample complexity. To learn such a model, one needs to show that training the first layer recovers the entire span of $k$ target directions (prior work has shown convergence to this subspace, but recovering the entire subspace is more challenging [MHPG+23]). Even in the simpler isotropic case, [ABAM23] only provides general conjectures in this setting, and their theorems are valid only for specific examples of multiple-index targets. Therefore, while conceptually simple, such generalizations can be highly non-trivial from a technical perspective, mainly due to the complex interactions that arise between the neurons.\n***\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n\n[ABAM23] Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz, “Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics”, COLT 2023.\n\n[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath, “Online stochastic gradient descent on non-convex losses from high-dimensional inference”, JMLR 2021.\n\n[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song, “Learning single-index models with shallow neural networks”, NeurIPS 2022.\n\n[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”, NeurIPS 2022.\n\n[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi, “Neural Networks can Learn Representations with Gradient Descent”, COLT 2022.\n\n[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.\n\n[MHPG+23] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu, “Neural networks efficiently learn low-dimensional representations with SGD”, ICLR 2023.\n"}, {"review_id": "vNsnVwZXsu", "reviewer": "Reviewer_GGcV", "summary": "In 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. ", "questions": "1) 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. ", "limitations": "I do not see any potential negative social impact of this paper.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The 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.", "weaknesses": "The main weakness of the manuscript is the writing, sometimes not entirely clear. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s thoughtful evaluation and detailed comments. Following the reviewer’s suggestions, we will improve the general readability in the final version. We address the reviewer’s questions below.\n***\n**Q1**. We appreciate the reviewer's feedback and we will rephrase the sentence in the final version as follows. \n\n*“We emphasize that unless $\\phi = g$ (i.e. the link function is known), the first stage of training only recovers the relevant direction $\\boldsymbol{u}$ and is not able to approximate the non-linear link function $g$; thus, a second stage of training with $m > 1$ is required to learn the unknown $g$.”*\n***\n\n**Q2**. We provide a technical overview in Appendix A and we will move this section to the main text in the final version.\nTo provide a quick intuition between the connection of Assumption 2 and Proposition 4, consider the following ODE (see (A.2) for derivation), which follows from Lemma 3. The below inequality can be used to describe the evolution of the alignment between $\\overline{\\boldsymbol{w}}^t$ and $\\overline{\\boldsymbol{u}}$\n\n$$\\frac{\\mathrm{d}\\langle \\overline{\\boldsymbol{w}}^t, \\overline{\\boldsymbol{u}}\\rangle}{\\mathrm{d}t} = A_t\\mathbb{E}\\left[\\phi'(\\langle \\overline{\\boldsymbol{w}}^t, \\boldsymbol{z}\\rangle)g'(\\langle\\overline{\\boldsymbol{u}}, \\boldsymbol{z}\\rangle )\\right] = -A_t\\zeta_{\\phi,g}(\\langle \\overline{\\boldsymbol{w}}^t, \\overline{\\boldsymbol{u}}\\rangle) \\geq cA_t\\langle \\overline{\\boldsymbol{w}}^t, \\overline{\\boldsymbol{u}}\\rangle^{s-1},$$\nwhere $A_t \\geq \\lambda_{\\text{min}}(\\boldsymbol{\\Sigma})(1 - \\langle \\overline{\\boldsymbol{w}}^t, \\overline{\\boldsymbol{u}}\\rangle^2)$ and we used Assumption 2 in the last step.\nIntegrating this equation will yield the statement of Proposition 4. See Appendix D.1 for a precise derivation.\n\nRegarding the class of admissible activations, the population limit (Proposition 4) can cover any weakly differentiable activation $\\phi$ with $\\phi,\\phi' \\in L^2(\\gamma)$ (pointed out in line 140). Our finite-sample analysis (Theorems 5 and 7) however is more specific and can accommodate the following activations: 1-ReLU (non-smooth) and 2- Lipschitz activations with a Lipschitz derivative (smooth). Examples of the second class include popular activations like the sigmoid and tanh. We will further clarify this in the final version.\n***\n\n**Q3**. For Theorem 9, in addition to recovering the relevant direction, we need to approximate the function $g$ using a linear combination of the activation $\\phi$. The piecewise-linear structure of ReLU makes the approximation analysis easier, which is why we (along with several recent works in the literature [DLS22,BBSS22,MHPG+23]) focus on function approximation via ReLU units. However, in principle, Theorem 9 can be restated for any activation which is capable of universal function approximation.\n***\n\n**Q4**. The intermediate regime interpolates between the hard and easy regimes, in which larger $r_2$ and smaller $r_1$ continuously improve the sample complexity. This is in contrast to the hard regime where the complexity is not affected by either $r_1$ or $r_2$, and the easy regime where only $r_2$ determines the sample complexity. We will add a discussion of this regime to the final version of the manuscript.\n***\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n\n[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song, “Learning single-index models with shallow neural networks”, NeurIPS 2022.\n\n[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi, “Neural Networks can Learn Representations with Gradient Descent”, COLT 2022.\n\n[MHPG+23] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu, “Neural networks efficiently learn low-dimensional representations with SGD”, ICLR 2023."}, {"review_id": "gdFlYkgK94", "reviewer": "Reviewer_oCpw", "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.\n", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n", "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.\n\n[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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful feedback and positive assessment. \n\n**Initialization;** We remark that our main objective was to investigate the different complexity regimes arising from structural variations in data, rather than designing optimal algorithms for the learning task at hand. Nevertheless, the reviewer's suggestion to initialize at the spike direction is in fact very useful and will be discussed in the main text; the resulting sample complexity after spectral initialization can be directly obtained from Theorems 5 and 7 as they hold for general initialization. In particular, up to polylog factors, the sample complexity from Theorem 5 (without preconditioning) would read $d^{1 + 2r_2 + ((s-1)(2r_1 - r_2) \\lor 0)}$ and that from Theorem 7 (with preconditioning) would read $d^{1 + ((s-1)(2r_1 - r_2) \\lor 0)}$. As expected, these complexities are strictly better than what is achieved under random initialization (Corollaries 6 and 8), and notably, the sample complexity for the preconditioned algorithm becomes linear in $d$ (up to polylog factors) as soon as $r_2 > 2r_1$. \n\n**Further discussion:** We remark that while spectral initialization does provide some benefit under the spiked model, it does not always trivialize the problem. More specifically, unless we have $r_2 > 2r_1$, the information exponent still affects the sample complexity. Furthermore, there are settings where uniform initialization has smaller sample and time complexity than initializing at the spike direction, e.g. if the spike-target alignment is worse than uniformly at random (i.e. $r_1 > 0.5$). This phenomenon can more naturally occur under a general covariance structure, where e.g. if the target direction aligns well with some other eigenvector of $\\boldsymbol{\\Sigma}$, Theorems 5 and 7 show that the convergence time and sample complexity of initializing from the first eigenvector can be arbitrarily large.\n\nWe thank the reviewer for pointing to this and we will add a detailed discussion about this observation in the final version of the manuscript.\n***\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period."}]

Author rebuttal:

Meta-review: The paper analyzes neural network learning in the single index model, with structured data. In this setup, the observed label y is a function of a one dimensional projection <u,x> of the features x. Previous works assumed an isotropic Gaussian distribution for x. This submission considers the situation in which x follows a Gaussian distribution with spiked covariance. It performs a comprehensive analysis of this setup, identifying phase transitions in terms of (i) the alignment between u and the covariance spike and (ii) the information exponent of the link function. An interesting byproduct of this analysis is the following observation: in the spiked model, SGD does not accurately learn u, but a normalized variant (somewhat reminiscent of batch normalization) does. 

Reviewers positively evaluated the manuscript’s technical quality, noting that extending results on the single index models to more realistic (and structured) distributions is an important direction for theory. The paper’s phase transition results give a comprehensive picture of this model. Concerns included the realism of the model, as well as certain artificial aspects of the algorithm (this is a two-stage theoretical algorithm, not SGD or batch normalized SGD). The review consensus is that the paper’s merits (high technical quality, important direction for theory) outweigh these concerns; the AC concurs, and recommends acceptance.
Decision: Accept (poster)