Original: Learning Single-Index Models with Shallow Neural Networks
Revised: Learning Single-Index Models with Shallow Neural Networks

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress
Revised: AfﬁliationAddress

Original: Abstract
Revised: Abstract

Original: Single-index models are a class of functions given by an unknown univariate“link” function applied to an unknown one-dimensional projection of the input.
Revised: Single-index models are a class of functions given by an unknown univariate“link” function applied to an unknown one-dimensional projection of the input.

Original: These models are particularly relevant in high dimension, when the data mightpresent low-dimensional structure that learning algorithms should adapt to.
Revised: These models are particularly relevant in high dimension, when the data mightpresent low-dimensional structure that learning algorithms should adapt to.

Original: Whileseveral statistical aspects of this model, such as the sample complexity of recovering the relevant (one-dimensional) subspace, are well-understood, they rely ontailored algorithms that exploit the speciﬁc structure of the target function.
Revised: Whileseveral statistical aspects of this model, such as the sample complexity of recovering the relevant (one-dimensional) subspace, are well-understood, they rely ontailored algorithms that exploit the speciﬁc structure of the target function.

Original: In thiswork, we introduce a natural class of shallow neural networks and study its abilityto learn single-index models via gradient descent .
Revised: In thiswork, we introduce a natural class of shallow neural networks and study its abilityto learn single-index models via gradient descent .

Original: More precisely, we considershallow networks in which  biases of the neurons are frozen at random initialization.
Revised: More precisely, we considershallow networks in which the ﬁrst layer weights are tied, to mirror the singleindex model structure, and biases of the neurons are frozen at random initialization.

Original: We show that the corresponding optimization landscape is benign, which inturn leads to generalization guarantees that match the optimal sample complexityof dedicated semi-parametric methods.
Revised: We show that the corresponding optimization landscape is benign, which inturn leads to generalization guarantees that match the optimal sample complexityof dedicated semi-parametric methods.

Original:  Introduction
Revised:  Introduction

Original: High-dimensional learning with both computational and statistical guarantees remains an outstanding challenge, with even more relevance in the face of current scaling trends.
Revised: High-dimensional learning with both computational and statistical guarantees remains an outstanding challenge, with even more relevance in the face of current scaling trends.

Original: One important questionwhich has seen considerable attention recently is to understand the advantages brought by non-linearlearning models, such as neural networks, over more mature (from a theoretical standpoint) counterparts, such as kernel methods [32, 33, 49, 13], even in their simplest shallow instantiations.
Revised: One important questionwhich has seen considerable attention recently is to understand the advantages brought by non-linearlearning models, such as neural networks, over more mature (from a theoretical standpoint) counterparts, such as kernel methods [33, 34, 53, 13], even in their simplest shallow instantiations.

Original: While approximation beneﬁts of shallow neural networks over non-adaptive kernels have beenknown for decades [8, 40], an important piece of this theoretical puzzle was provided by the statistical analysis of [5], which hinted at an inherent advantage of neural networks for extracting information from high-dimensional data with a “hidden” low-dimensional structure.
Revised: While approximation beneﬁts of shallow neural networks over non-adaptive kernels have beenknown for decades [8, 41], an important piece of this theoretical puzzle was provided by the statistical analysis of [5], which hinted at an inherent advantage of neural networks for extracting information from high-dimensional data with a “hidden” low-dimensional structure.

Original: Alas, the remainingpiece of this puzzle, namely providing computational guarantees, is still mostly unresolved.
Revised: Alas, the remainingpiece of this puzzle, namely providing computational guarantees, is still mostly unresolved.

Original: Several computational hardness results for learning functions that can be efﬁciently approximatedby shallow neural networks [18, 26, 45, 17] have been established in the literature, which rule outpositive results in the general setting.
Revised: Several computational hardness results for learning functions that can be efﬁciently approximatedby shallow neural networks [18, 26, 48, 17] have been established in the literature, which rule outpositive results in the general setting.

Original: On the other hand, progress has been made on the positiveside [1, 2, 43] by focusing on function classes with strong structural properties, thereby showcasingthe adaptive representation learning capabilities of neural networks.
Revised: On the other hand, progress has been made on the positiveside [1, 2, 45] by focusing on function classes with strong structural properties, thereby showcasingthe adaptive representation learning capabilities of neural networks.

Original: This work aligns with the latter effort, and focuses on the class of single-index models .
Revised: This work aligns with the latter effort, and focuses on the class of single-index models .

Original: Single-indexmodels are high-dimensional functions F : R d !
Revised: Single-indexmodels are high-dimensional functions F : R d !

Original: R of the form F ⇤ ( x ) =
Revised: R of the form F ⇤ ( x ) =

Original: f
Revised: f

Original: ⇤ ( h ✓ ⇤ , x i ) , where boththe univariate “link” function f
Revised: ⇤ ( h ✓ ⇤ , x i ) , where boththe univariate “link” function f

Original: ⇤ : R !
Revised: ⇤ : R !

Original: R and relevant (one-dimensional) subspace ✓ ⇤ 2 S d  1are unknown.
Revised: R and relevant (one-dimensional) subspace ✓ ⇤ 2 S d  1are unknown.

Original: These models have been extensively studied in the statistics literature [16, 28, 19],leading to dedicated algorithmic procedures, and can be provably approximated with shallow neuralnetworks without incurringin a curse-of-dimensionality [5].
Revised: These models have been extensively studied in the statistics literature [16, 28, 19],leading to dedicated algorithmic procedures, and can be provably approximated with shallow neuralnetworks without incurring a curse-of-dimensionality [5].

Original: In contrast, the analysis of neural network learning using gradient-descent techniques has been more heavily focused on the so-called“Teacher-Student” setup [27, 24, 50, 51, 3], where the link function f is assumed to be known.
Revised: In contrast, the analysis of neural network learning using gradient-descent techniques has been more heavily focused on the so-called “TeacherStudent” setup [27, 24, 54, 55, 3], where the link function f is assumed to be known.

Original: The mix of a high-dimensional parametric component (the hidden direction) with a low-dimensionalnon-parametric one (the link function) in single-index models naturally suggests a shallow neuralnetwork architecture where the inner weights are “active”, while the biases are “lazy” [14].
Revised: The mix of a high-dimensional parametric component (the hidden direction) with a low-dimensionalnon-parametric one (the link function) in single-index models naturally suggests a shallow neuralnetwork architecture where the inner weights are “active”, while the biases are “lazy” [14].

Original: Weinstantiate such an architecture by freezing the biases at their random initialization, and analyzegradient descent dynamics on the remaining free parameters in the continuous-time limit.
Revised: Weinstantiate such an architecture by freezing the biases at their random initialization, and analyzegradient descent dynamics on the remaining free parameters in the continuous-time limit.

Original: Our main results establish that as soon as the width N of the network is bigger than a quantity, whichdepends solely on the smoothness properties of the link function f ⇤ 1 , gradient descent recovers theunknown direction ✓ ⇤ and approximates the univariate link function f ⇤ with optimal sample andtime complexity
Revised: Our main results establish that as soon as the width N of the network is bigger than a quantity, whichdepends solely on the smoothness properties of the link function f ⇤ 1 , gradient descent recovers theunknown direction ✓ ⇤ and approximates the univariate link function f ⇤ with optimal sample andtime complexity

Original: O ( d s ) (Theorem 5.1), where s is the so-called information-exponent of the linkfunction [3].
Revised: O ( d s ) (Theorem 5.1), where s is the so-called information-exponent of the linkfunction [3].

Original: The information exponent roughly captures the signal strength, which here refers tothe alignment between the network direction ✓ and the hidden direction ✓ ⇤ , at typical initializations.
Revised: The information exponent roughly captures the signal strength, which here refers tothe alignment between the network direction ✓ and the hidden direction ✓ ⇤ , at typical initializations.

Original: The success of gradient descent relies on the benign optimization landscape of the empirical loss,though the presence of degenerate saddles requires a careful convergence analysis leveraging uniform convergence results for the empirical landscape [36, 20].
Revised: The success of gradient descent relies on the benign optimization landscape of the empirical loss,though the presence of degenerate saddles requires a careful convergence analysis leveraging uniform convergence results for the empirical landscape [37, 20].

Original: We show that gradient descent overour proposed neural network architecture solves two distinct problems–univariate non-parametrickernel ridge regression and non-convex optimization in high dimension– simultaneously and efﬁciently , cementing its role as a versatile algorithm for high-dimensional learning.
Revised: We show that gradient descent overour proposed neural network architecture solves two distinct problems–univariate non-parametrickernel ridge regression and non-convex optimization in high dimension– simultaneously and efﬁciently , cementing its role as a versatile algorithm for high-dimensional learning.

Original: We illustrate ourtheoretical results with experiments in Appendix D.
Revised: We illustrate ourtheoretical results with experiments in Appendix D.

Original:  Related Work
Revised:  Related Work

Original: Our work builds upon methods from different areas in statistics, optimization, and machine learning.
Revised: Our work builds upon methods from different areas in statistics, optimization, and machine learning.

Original: Single–index and multi-index models.
Revised: Single–index and multi-index models.

Original: A useful modeling assumption in high-dimensional regression is that the regression function F ( x ) =
Revised: A useful modeling assumption in high-dimensional regression is that the regression function F ( x ) =

Original: E [ y | x ] only depends on one or a few directions.
Revised: E [ y | x ] only depends on one or a few directions.

Original: This leads to the single-index model F ( x ) =
Revised: This leads to the single-index model F ( x ) =

Original: f ⇤ ( h ✓ ⇤ , x i ) , and multi-index model F ( x ) =
Revised: f ⇤ ( h ✓ ⇤ , x i ) , and multi-index model F ( x ) =

Original: f ⇤ ( h ✓ ⇤ 1 , x i , . . .
Revised: f ⇤ ( h ✓ ⇤ 1 , x i , . . .

Original: , ✓ ⇤ k , x i ) , with k typically much smaller than the dimension.
Revised: , ✓ ⇤ k , x i ) , with k typically much smaller than the dimension.

Original: Such models have along history in the statistics literature and different methods exist for various estimation problems,including projection pursuit [21], slicing [30], gradient-based estimators [31], and moment-basedestimators [16].
Revised: Such models have along history in the statistics literature and different methods exist for various estimation problems,including projection pursuit [21], slicing [31], gradient-based estimators [32], and moment-basedestimators [16].

Original: When the function f
Revised: When the function f

Original: ⇤ is also to be estimated, we face a semi-parametric probleminvolving parameter recovery of ✓ ⇤ and non-parametric estimation of f ⇤ .
Revised: ⇤ is also to be estimated, we face a semi-parametric probleminvolving parameter recovery of ✓ ⇤ and non-parametric estimation of f ⇤ .

Original: Our work is closely related to [19], which also characterizes the population landscape of certain objectives by leveraging
Revised: Our work is closely related to [19], which also characterizes the population landscape of certain objectives by leveraging

Original: Gaussian data.
Revised: Gaussian data.

Original: Single and multi-index models are also studied in [5] in the context of shallow neural networks, where it is shown that certain models of inﬁnite-width shallow networks can adapt tosuch low-dimensional structure, though no tractable algorithms are introduced.
Revised: Single and multi-index models are also studied in [5] in the context of shallow neural networks, where it is shown that certain models of inﬁnite-width shallow networks can adapt tosuch low-dimensional structure, though no tractable algorithms are introduced.

Original: Finally, [1] studiesthe learnability of functions on the hypercube by shallow neural networks and stochastic gradientdescent and introduces the merged staircase property , which provides necessary and sufﬁcient conditions for learnability with linear sample complexity n =
Revised: Finally, [1] studiesthe learnability of functions on the hypercube by shallow neural networks and stochastic gradientdescent and introduces the merged staircase property , which provides necessary and sufﬁcient conditions for learnability with linear sample complexity n =

Original: O ( d ) .
Revised: O ( d ) .

Original: While they learn a broader classof functions (including multi-index model) for a more efﬁcient sample complexity regime ( O ( d ) vs
Revised: While they learn a broader classof functions (including multi-index model) for a more efﬁcient sample complexity regime ( O ( d ) vs

Original: O ( d s ) ), our work captures the regime of semi-parametric estimation by considering Gaussian datawithout the sparsity constraints restrictions on F implied by their merged staircase property.
Revised: O ( d s ) ), our work captures the regime of semi-parametric estimation by considering Gaussian datawithout the sparsity constraints restrictions on F implied by their merged staircase property.

Original: Teacher–student models.
Revised: Teacher–student models.

Original: In the context of neural networks, several works have considered theteacher–student setting, where the target function F takes the form of a neural network with thesame activation as network used for learning [27, 24, 50, 51, 3].
Revised: In the context of neural networks, several works have considered theteacher–student setting, where the target function F takes the form of a neural network with thesame activation as network used for learning [27, 24, 54, 55, 3].

Original: In this case the problem doesnot involve non-parametric estimation as in our setup, but this line of work often involves studyingoptimization landscapes similar to ours for estimating hidden directions.
Revised: In this case the problem doesnot involve non-parametric estimation as in our setup, but this line of work often involves studyingoptimization landscapes similar to ours for estimating hidden directions.

Original: Kernels and random features
Revised: Kernels and random features.

Original: Our work builds on the kernel methods literature for approximation and non-parametric regression [44, 7], their links with neural networks [15, 5], and in particularon random feature approximation [41, 6, 42, 37].
Revised: Our work builds on the kernel methods literature for approximation and non-parametric regression [47, 7], their links with neural networks [15, 5], and in particularon random feature approximation [42, 6, 44, 38].

Original: Non-convex optimization landscapes.
Revised: Non-convex optimization landscapes.

Original: There is a vast literature studying tractable non-convexoptimization landscapes, arising from high-dimensional statistics and statistical physics [35, 34, 22,4, 46, 10, 24].
Revised: There is a vast literature studying tractable non-convexoptimization landscapes, arising from high-dimensional statistics and statistical physics [36, 35, 22,4, 43, 49, 10, 24].

Original: A particular aspect of our setup is that the optimization landscape does not have thestrict saddle property, which is often leveraged to establish global convergence [29].
Revised: A particular aspect of our setup is that the optimization landscape does not havethe strict saddle property, which is often leveraged to establish global convergence [29].

Original:  Preliminaries
Revised:  Preliminaries

Original: We focus on regression problems under a single-index model with Gaussian input data.
Revised: We focus on regression problems under a single-index model with Gaussian input data.

Original: Speciﬁcally,we assume
Revised: Speciﬁcally,we assume

Original: d -dimensional inputs x ⇠  d := N (0 , I d ) , and labels
Revised: d -dimensional inputs x ⇠  d := N (0 , I d ) , and labels

Original: [Equation]
Revised: [Equation]

Original: where ✓ ⇤ 2 S d  1 and ✏ is random noise such that E [ ✏ | x ] = 0 , E [ ✏
Revised: where ✓ ⇤ 2 S d  1 and ✏ is random noise such that E [ ✏ [Equation] x [Equation]

Revised: = 0 , E

Original: [Equation]
Revised: [Equation]

Revised: ✏

Revised: [Equation]

Original: | x ]   2 .
Revised: x [Equation]   [Equation]

Original: Shallow networks and random features.
Revised: Shallow networks and random features.

Original: We consider learning algorithms based on shallow neural networks of the form
Revised: We consider learning algorithms based on shallow neural networks of the form

Original: [Equation]
Revised: [Equation]

Original: with ✓ 2 S d  1 , where  ( u ) = max { 0 , u } is the ReLU activation, b i ⇠ N (0 , ⌧ 2 ) (we assume ⌧ > 1 )are random “bias” scalars that are frozen throughout training, and  ( u ) =
Revised: with ✓ 2 S d  1 , where  ( u ) = max { 0 , u } is the ReLU activation, b i ⇠ N (0 , ⌧ 2 ) (we assume ⌧ > 1 )are random “bias” scalars that are frozen throughout training, and  ( u ) =

Original: 1 p N (  ( u  b i ))
Revised: 1 p N (  ( u  b i ))

Original: i 2 [ N ]is the vector of random features.
Revised: i 2 [ N ]is the vector of random features.

Original: If we also ﬁx the direction ✓ , learning c alone may be seen as arandom feature model [41] that approximates a kernel method with the following kernel:
Revised: If we also ﬁx the direction ✓ , learning c alone may be seen as arandom feature model [42] that approximates a kernel method with the following kernel:

Original: [Equation]
Revised: [Equation]

Original: where  b ( u ) :=  ( u  b ) and  ⌧ = N (0 , ⌧ 2 ) is the Gaussian measure on R with mean zero andvariance ⌧ 2 .
Revised: where  b ( u ) :=  ( u  b ) and  ⌧ = N (0 , ⌧ 2 ) is the Gaussian measure on R with mean zero andvariance ⌧ 2 .

Original: The kernel  corresponds to an RKHS H , which we precisely deﬁne as
Revised: The kernel  corresponds to a Reproducing Kernel Hilbert Space (RKHS) H , whichwe precisely deﬁne as

Original: [Equation]
Revised: [Equation]

Original: The following lemma characterizes the corresponding RKHS norm k · k H .
Revised: The following lemma characterizes the corresponding RKHS norm k · k H .

Original: Lemma 3.1 (RKHS Sobolev Representation, [48, Theorem 5]) .
Revised: Lemma 3.1 (RKHS Sobolev Representation, [52, Theorem 5]) .

Original: The RKHS norm in H is given by
Revised: The RKHS norm in H is given by

Original: [Equation]
Revised: [Equation]

Original: where f 00 is deﬁned in the generalised (Sobolev) sense.
Revised: where f 00 is deﬁned in the generalised (Sobolev) sense.

Original: Remark 3.2.
Revised: Remark 3.2.

Original: The boundary conditions for this RKHS can be relaxed by symmetrizing the kernel, ie s ( u, v ) = 12 (  ( u, v ) +  (  u,  v )) , but we will keep the original boundary conditions and assumethem on f ⇤ for simplicity and w.l.o.g.
Revised: The boundary conditions for this RKHS can be relaxed by symmetrizing the kernel, ie s ( u, v ) = 12 (  ( u, v ) +  (  u,  v )) , but we will keep the original boundary conditions and assumethem on f ⇤ for simplicity and w.l.o.g.

Original: Hermite decompositions.
Revised: Hermite decompositions.

Original: Given that our data is normally distributed, we consider the familyof (normalized) Hermite polynomials { h j } j 2 N , which form an orthonormal basis of L 2 (  ) , thespace of squared-integrable function under the Gaussian measure  := N (0 , 1) .
Revised: Given that our data is normally distributed, we consider the familyof (normalized) Hermite polynomials { h j } j 2 N , which form an orthonormal basis of L 2 (  ) , thespace of squared-integrable function under the Gaussian measure  := N (0 , 1) .

Original: We will denoteby f ⇤ =
Revised: We will denoteby f ⇤ =

Original: P j ↵ j h j the Hermite decomposition of the target link function.
Revised: P j ↵ j h j the Hermite decomposition of the target link function.

Original: We apply the following useful properties of Hermite polynomials [39]:
Revised: We apply the following useful properties of Hermite polynomials [40]:

Original: [Equation]
Revised: [Equation]

Original: where h · , · i  d is the inner product in L 2 (  d ) and  the Kronecker delta.
Revised: where h · , · i  d is the inner product in L 2 (  d ) and  the Kronecker delta.

Original: We will assume throughoutthat k f ⇤ k 2  =
Revised: We will assume throughoutthat k f ⇤ k 2  =

Original: P [Equation]{+j ↵ 2 j , k f 
+}
Revised: P {+j ↵ 2 j , k f 
+}

Revised: ⇤ k

Original: [Equation]
Revised: [Equation]

Original: and
Revised:  = P

Original: [Equation]
Revised: [Equation]

Original: {+↵ 2 j , and k f 
 ⇤ k 2  = P j j ( j +} 1) ↵ 2 j are all ﬁnite (see Assumption 4.1).
Revised: {+↵ 2 j , and k f 
 ⇤ k 2  = P j j ( j +} 1) ↵ 2 j are all ﬁnite (see Assumption 4.1).

Original: Deﬁnition 3.3 (Population loss) .
Revised: Deﬁnition 3.3 (Population loss) .

Original: We deﬁne the L 2 -regularized population loss by
Revised: We deﬁne the L 2 -regularized population loss by

Original: [Equation]
Revised: [Equation]

Original: From features to Hermite coefﬁcients.
Revised: From features to Hermite coefﬁcients.

Original: To precisely characterize the landscape of L ( c, ✓ ) , weintroduce notation to represent the relationship each student feature  b in terms of Hermite polynomials h j .
Revised: To precisely characterize the landscape of L ( c, ✓ ) , weintroduce notation to represent the relationship between features  b in terms of Hermite polynomialsh j .

Original: We deﬁne the linear integral operator T :
Revised: We deﬁne the linear integral operator T :

Original: L 2 (  ) !
Revised: L 2 (  ) !

Original: R N by
Revised: R N by

Original: [Equation]
Revised: [Equation]

Original: Note that the operator T has rank N almost surely.
Revised: Note that the operator T has rank N almost surely.

Original: We abuse notation to equivalently let T ( j, b i ) :
Revised: We abuse notation to equivalently let T ( j, b i ) :

Original: =1 p N h h j ,  b
Revised: =1 p N h h j ,  b

Original: j i  =
Revised: i i  =

Original: ( T h j ) i and observe that ( T f ) i =
Revised: ( T h j ) i and observe that ( T f ) i =

Original: P j T ( j, b i ) .
Revised: P j T ( j, b i ) .

Original: Observe that T has an adjoint
Revised: Observe that T has an adjoint

Original: T ⇤ : R N !
Revised: T ⇤ : R N !

Original: L 2 (  ) deﬁned as ( T ⇤ c )( u ) = 1 p N
Revised: L 2 (  ) deﬁned as ( T ⇤ c )( u ) = 1 p N

Original: P
Revised: P

Original: Ni =1 c i  ( u  b i ) =
Revised: Ni =1 c i  ( u  b i ) =

Original: c >  ( u ) .
Revised: c >  ( u ) .

Original: Let T j :=
Revised: Let T j :=

Original: T h j =( T ( j, b i ))
Revised: T h j =( T ( j, b i ))

Original: i =1 ...
Revised: i =1 ...

Original: N 2 R N .
Revised: N 2 R N .

Original: We can then write down the Hermite expansion of thestudent network:
Revised: We can then write down the Hermite expansion of the network:

Original: [Equation]
Revised: [Equation]

Original: Denoting m := h ✓ , ✓ ⇤ i , this allows us to rewrite the regularized objective as
Revised: Denoting m := h ✓ , ✓ ⇤ i , this allows us to rewrite the regularized objective as

Original: [Equation]
Revised: [Equation]

Original: where the term
Revised: where the term

Original: P j ↵ 2 j is a constant that can be ignored.
Revised: P j ↵ 2 j is a constant that can be ignored.

Original: Let Q :=
Revised: Let Q :=

Original: T T ⇤ 2 R N ⇥ N be a featurecovariance matrix and Q  = Q +  I .
Revised: T T ⇤ 2 R N ⇥ N be a featurecovariance matrix and Q  = Q +  I .

Original: Note that Q  is positive semi-deﬁnite for  > 0 .
Revised: Note that Q  is positive semi-deﬁnite for  > 0 .

Original: We deﬁnethe regularized projection
Revised: We deﬁnethe regularized projection

Original: ˆ P  = ˆ ⌃ (ˆ ⌃ +  I )  1 onto the random feature space ˆ H for ˆ ⌃ =
Revised: ˆ P  = ˆ ⌃ (ˆ ⌃ +  I )  1 onto the random feature space ˆ H for ˆ ⌃ =

Original: T ⇤ T .
Revised: T ⇤ T .

Original: Observe that ( ˆ P  f )( u ) =
Revised: Observe that ( ˆ P  f )( u ) =

Original: c ⇤>  ( u ) , where c
Revised: c ⇤>  ( u ) , where c

Original: ⇤ 2 R N is the solution to the following objective:
Revised: ⇤ 2 R N is the solution to the following objective:

Original: [Equation]
Revised: [Equation]

Original: Geometry on the sphere.
Revised: Geometry on the sphere.

Original: Because the direction ✓ is constrained to lie on the sphere, our optimization algorithms rely on spherical (Riemannian) gradients, which are deﬁned as follows:
Revised: Because the direction ✓ is constrained to lie on the sphere, our optimization algorithms rely on spherical (Riemannian) gradients, which are deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: where Proj ✓ ?
Revised: where Proj ✓ ?

Original: v = v  h ✓ , v i ✓ .
Revised: v = v  h ✓ , v i ✓ .

Original: We say that ( c, ✓ ) is a critical point of L if r S d  1 ✓
Revised: We say that ( c, ✓ ) is a critical point of L if r S d  1 ✓

Original: L ( c, ✓ ) = 0 andr c L ( c, ✓ ) = 0 .
Revised: L ( c, ✓ ) = 0 andr c L ( c, ✓ ) = 0 .

Original:  Population Landscape under Frozen Random Biases
Revised:  Population Landscape under Frozen Random Biases

Original: In order to understand the optimization and generalization properties of applying gradient descent tothe empirical loss L n ( c, ✓ ) , we ﬁrstthe study  optimization landscape of the population loss L ( c, ✓ ) .
Revised: In order to understand the optimization and generalization properties of applying gradient descent tothe empirical loss L n ( c, ✓ ) , we ﬁrst study the optimization landscape of the population loss L ( c, ✓ ) .

Original: We precisely characterize the critical points of the population loss in order to show that a gradientbased optimizer of astudent neural network of sufﬁcient width N converges only when its directionvector ✓ is either parallel or orthogonal to the teacher direction ✓ ⇤ .
Revised: We precisely characterize the critical points of the population loss in order to show that a gradientbased optimizer of a neural network of sufﬁcient width N converges only when its direction vector✓ is either parallel or orthogonal to the target direction ✓ ⇤ .

Original: In the subsequent section, weshow that having sufﬁciently many training samples n provides a similar optimization landscape andhas favorable generalization properties for most initializations.
Revised: In the subsequent section, we showthat having sufﬁciently many training samples n provides a similar optimization landscape and hasfavorable generalization properties for most initializations.

Original: We will make the following regularity assumptions on the teacher link function f
Revised: We will make the following regularity assumptions on the target link function f

Original: ⇤ , in order to ensuresmall approximation error from random features, and a benign empirical landscape.
Revised: ⇤ , in order to ensuresmall approximation error from random features, and a benign empirical landscape.

Original: Assumption 4.1 (Regularity of f ⇤ ) .
Revised: Assumption 4.1 (Regularity of f ⇤ ) .

Original: Let f ( p ) ⇤ be the p th derivative of f ⇤ .
Revised: Let f ( p ) ⇤ be the p th derivative of f ⇤ .

Original: We assume that f ( p ) ⇤ 2 141 L 2 (  ) (i.e. k f ( p )
Revised: We assume that f ( p ) ⇤ 2 143 L 2 (  ) (i.e. k f ( p )

Original: ⇤ k
Revised: ⇤ k

Original:  < 1 ) for all p  6 , that f ⇤ is Lipschitz and that f ⇤ is polynomially bounded.
Revised:  < 1 ) for all p  6 , that f ⇤ is Lipschitz and that f ⇤ is polynomially bounded.

Original: The index of the ﬁrst non-zero Hermite coefﬁcient ↵ j of f
Revised: The index of the ﬁrst non-zero Hermite coefﬁcient ↵ j of f

Original: ⇤ is referred to as the information exponentand is given by s :
Revised: ⇤ is referred to as the information exponentand is given by s :

Original: = min { j : ↵ j 6 = 0 } .
Revised: = min { j : ↵ j [Equation] = 0 } .

Original: To consider the critical points of L ( c, ✓ ) , we introduce the projected population loss at a givenstudent direction ✓ :
Revised: To consider the critical points of L ( c, ✓ ) , we introduce the projected population loss at a givendirection ✓ :

Original: [Equation]
Revised: [Equation]

Original: Crucially, the sufﬁcient number of features N depends on the regularization parameter  2 (0 , 1) ,the variance ⌧ 2 > 1 of the random biases, and the teacher link function f ⇤ , but not on the inputdimension d .
Revised: Crucially, the sufﬁcient number of features N depends on the regularization parameter  2 (0 , 1) ,the variance ⌧ 2 > 1 of the random biases, and the target link function f ⇤ , but not on the inputdimension d .

Original: Theorem 4.2 (Critical Points of Population Loss) .
Revised: Theorem 4.2 (Critical Points of Population Loss) .

Original: For ⌧ > 1 , and  2 (0 , 1) , there exist constants
Revised: For ⌧ > 1 , and  2 (0 , 1) , there exist constants

Original: C ⌧ and C f ⇤ depending only on ⌧ and the teacher link function f
Revised: C ⌧ and C f ⇤ depending only on ⌧ and the target link function f

Original: ⇤ respectively such that if < C  (
Revised: ⇤ respectively such that if

Original: [Equation]
Revised: [Equation]

Revised: ⇤ and

Original: [Equation]
Revised: [Equation]

Original: 2 ) f ⇤ and

Original: N

Original: [Equation]
Revised: [Equation]

Original:  ◆ (5) then with probability 1   over the random biases b 1 , . . .
Revised: (5) then with probability 1   over the random biases b 1 , . . .

Original: , b N , the set of ﬁrst-order critical points
Revised: , b N , the set of ﬁrst-order critical points

Original: C L := { ( c, ✓ ) : r S d  1 ✓ L
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: c
Revised: c

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]

Original: satisﬁes:1.
Revised: satisﬁes:1.

Original: (orientation relative to ✓ ⇤ )
Revised: (orientation relative to ✓ ⇤ )

Original: if ( c, ✓ ) 2 C L , then either ✓ 2 {± ✓ ⇤ } or h ✓ , ✓ ⇤ i = 0 ; and2.
Revised: if ( c, ✓ ) 2 C L , then either ✓ 2 {± ✓ ⇤ } or h ✓ , ✓ ⇤ i = 0 ; and2.

Original: (existence and uniqueness of c ) if r S d  1 ✓ ¯ L ( ✓ ) = 0 , then there exists a unique c such that
Revised: (existence and uniqueness of c ) if r S d  1 ✓ ¯ L ( ✓ ) = 0 , then there exists a unique c such that

Original: [Equation]
Revised: [Equation]

Original: , ✓ [Equation] L . Theorem 4.2 thus establishes a benign optimization landscape in the population limit, rejoiningseveral known non-convex objectives with similar behavior, such as tensor decomposition [25] ormatrix completion [23].
Revised: Theorem 4.2 thus establishes a benign optimization landscape in the population limit, rejoiningseveral known non-convex objectives with similar behavior, such as tensor decomposition [25] ormatrix completion [23].

Revised: It implies that all critical points are either at the equator or the poles andthus a gradient ﬂow has to either avoid or escape the “spurious” critical points in the equator, i.e., ✓ ’ssatisfying h ✓ ⇤ , ✓ i = 0 to succeed.

Original: Importantly, this optimization landscape has the same topology as if onewere using Hermite basis functions rather than random scalar features [19, Theorem 5], the tailoredchoice that exploits the knowledge that the data is being generated by a single-index model.
Revised: Importantly, this optimization landscape has the same topology asif one were using Hermite basis functions rather than random scalar features [19, Theorem 5], thetailored choice that exploits the knowledge that the data is being generated by a single-index modelon Gaussian samples.

Original: Weview this as an interesting robustness property of shallow neural networks – at least in the regimewhere biases are randomly frozen.
Revised: We view this as an interesting robustness property of shallow neural networks– at least in the regime where biases are randomly frozen.

Original: We now outline the analysis that leads to the proof of the theorem before proving it formally in
Revised: We now outline the analysis that leads to the proof of the theorem before proving it formally in

Original: Section 4.3.
Revised: Section 4.3.

Original: 4.1 Exact expression of ﬁrst-order critical points
Revised: 4.1 Exact expression of ﬁrst-order critical points

Original: The ﬁrst step is to characterize the critical points of the regularized population loss.
Revised: The ﬁrst step is to characterize the critical points of the regularized population loss.

Original: This is describedin the following result.
Revised: This is describedin the following result.

Original: Claim 4.3 (Critical point equation) .
Revised: Claim 4.3 (Critical point equation) .

Original: Let g m , ¯ g m 2 L 2 (  ) be deﬁned by
Revised: Let g m , ¯ g m 2 L 2 (  ) be deﬁned by

Original: [Equation]
Revised: [Equation]

Original: where m = h ✓ , ✓ ⇤ i .
Revised: where m = h ✓ , ✓ ⇤ i .

Original: Then, the critical points of L ( c, ✓ ) satisfy the following equations:
Revised: Then, the critical points of L ( c, ✓ ) satisfy the following equations:

Original: [Equation]
Revised: [Equation]

Original: Remark 4.4 (Inﬁnite-widthstudent networks) .
Revised: Remark 4.4 (Inﬁnite-width networks) .

Original: If the kernel is universal and N ! 1 , then  P  isthe identity operator when  = 0 .
Revised: If the kernel is universal and N ! 1 , then ˆ P  is theidentity operator when  = 0 .

Original: Hence, h h j , P  h
Revised: Hence, h h j , P  h

Original: j 0 i =
Revised: j 0 i =

Original:  j,j 0 in this limit and the teacher can be172 exactly  represented when m = 1 .
Revised:  j,j 0 in this limit and the target can be exactly 176 represented when m = 1 .

Original: We prove Claim 4.3 in Appendix B.1 by exactly computing the gradients and relating the criticalpoints of ¯ L ( ✓ ) and L ( c, ✓ ) .
Revised: We prove Claim 4.3 in Appendix B.1 by exactly computing the gradients and relating the criticalpoints of ¯ L ( ✓ ) and L ( c, ✓ ) .

Original: 4.2 Controlling random feature approximation
Revised: 4.2 Controlling random feature approximation

Original: We now consider (ﬁnite) random feature approximations to functions in the RKHS.
Revised: We now consider (ﬁnite) random feature approximations to functions in the RKHS.

Original: We show thatthe best possible loss of a linear combination of sufﬁciently many ﬁnite features is bounded aboveby the best approximation with inﬁnitely many features with high probability.
Revised: We show thatthe best possible loss of a linear combination of sufﬁciently many ﬁnite features is bounded aboveby the best approximation with inﬁnitely many features with high probability.

Original: Let A ( f,  ) be the (regularized)
Revised: Let A ( f,  ) be the (regularized)

Original: L 2 approximation error for functions in the space H with respect tothe target function f and measure  .
Revised: L 2 approximation error for functions in the space H with respect tothe target function f and measure  .

Original: Formally,
Revised: Formally,

Original: [Equation]
Revised: [Equation]

Original: Lemma 4.5 (Random features approximation, adapted from [6]) .
Revised: Lemma 4.5 (Random features approximation, adapted from [6]) .

Original: Let  2 (0 , 1) , ⌧ > 1 , and f 

Revised: Let  2 (0 , 1) , ⌧ > 1 , and f

Original: L 2 (  ) .
Revised: L 2 (  ) .

Original: There exists a constant C ⌧ > 0 such that if N  C ⌧  ln 1  , then with probability at least 1   ,
Revised: There exists a constant C ⌧ > 0 such that if N  C ⌧  ln 1  , then with probability at least 1   ,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The claim is immediate from Lemmas B.5 and B.6 in Appendix B.2.
Revised: The claim is immediate from Lemmas B.5 and B.6 in Appendix B.2.

Original: Next, we show that the approximation error of the RKHS corresponding to an inﬁnite number ofrandom features can be bounded in terms of the regularization  and the  -norm of the secondderivative of the target function.
Revised: Next, we show that the approximation error of the RKHS corresponding to an inﬁnite number ofrandom features can be bounded in terms of the regularization  and the  -norm of the secondderivative of the target function.

Original: The bound is independent of d , as we consider ⌧ and  ﬁxed as dgrows.
Revised: The bound is independent of d , as we consider ⌧ and  ﬁxed as dgrows.

Original: Lemma 4.6 (RKHS approximation error) .
Revised: Lemma 4.6 (RKHS approximation error) .

Original: Assume f 2 L 2 (  ) is in the weighted Sobolev space
Revised: Assume f 2 L 2 (  ) is in the weighted Sobolev space

Original: H 2 (  ) and has polynomial growth (i.e. there exist C > 0 and M > 0 such that | f ( u ) |  C ( u M +u  M ) ).
Revised: H 2 (  ) and has polynomial growth (i.e. there exist C > 0 and M > 0 such that | f ( u ) |  C ( u M +u  M ) ).

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: with  := ⌧ 2 ⌧ 2 +2 .
Revised: with  := ⌧ 2 ⌧ 2 +2 .

Original: The proof appears in Appendix B.3.
Revised: The proof appears in Appendix B.3.

Original: 4.3 Proof of Theorem4.
Revised: 4.3 Proof of Theorem

Original: Proof.
Revised: Proof.

Original: For ( c, ✓ ) to be a critical point of L , it must satisfy r S d  1 ✓L
Revised: For ( c, ✓ ) to be a critical point of L , it must satisfy r S d  1 ✓

Original: [Equation] c, ✓ ) = 0 and r c L ( ✓ , c ) = 0 .
Revised: L ( c, ✓ ) = 0 and r c L ( ✓ , c ) = 0 .

Original: According to Claim 4.3, every ﬁxed ✓ has a unique complementary c such that r c L ( ✓ , c ) = 0 , andr S d  1 ✓
Revised: According to Claim 4.3, every ﬁxed ✓ has a unique complementary c such that r c L ( ✓ , c ) = 0 , andr S d  1 ✓

Original: L ( c, ✓ ) = 

Revised: L ( c, ✓ ) =

Original: if and only if ✓ 2 {  ✓ ⇤ , ✓ ⇤ } or
Revised: if and only if ✓ 2 {  ✓ ⇤ , ✓ ⇤ } or

Original: [Equation]
Revised: [Equation]

Original: The remainder of the proof involves showing that the latter condition is true only if m = 0 .
Revised: The remainder of the proof involves showing that the latter condition is true only if m = 0 .

Original: Wederive a contradiction whenever m 6 = 0 .
Revised: Wederive a contradiction whenever m 6 = 0 .

Original: Let  :=
Revised: Let  :=

Original: ⌧ 2 ⌧ 2 +

Revised: ⌧ 2 ⌧ 2 +

Original: and take N  C ⌧  ln 1  and
Revised: and take N  C ⌧  ln 1  and

Original: [Equation]
Revised: [Equation]

Original: for sufﬁciently large C 0 and C ⌧ .
Revised: for sufﬁciently large C 0 and C ⌧ .

Original: We contradict the equality with probability at least 1   over therandomly sampled biases b 1 , . . .
Revised: We contradict the equality with probability at least 1   over therandomly sampled biases b 1 , . . .

Original: , b N :
Revised: , b N :

Original: [Equation]
Revised: [Equation]

Original:  Lemma 4.5 ]
Revised: [ Lemma 4.5 ]

Original: [Equation]
Revised: [Equation]

Original: [ Lemma 4.6 ]
Revised: [ Lemma 4.6 ]

Revised: [Equation]

Revised: [ Lemma

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  Empirical Landscape and Generalization Guarantees
Revised:  Empirical Landscape and Generalization Guarantees

Original: Now that we have established that the population landscape L ( c, ✓ ) has a particularly simplestructure given sufﬁciently many random features N & O (1) , we study the optimization properties of its ﬁnite-sample counterpart.
Revised: Now that we have established that the population landscape L ( c, ✓ ) has a particularly simplestructure given sufﬁciently many random features N & O (1) , we study the optimization properties of its ﬁnite-sample counterpart.

Original: For training samples { ( x i , y i ) } i 2 [ n ] generated by independently drawn x i ⇠ N (0 , I ) and y i =
Revised: For training samples { ( x i , y i ) } i 2 [ n ] generated by independently drawn x i ⇠ N (0 , I ) and y i =

Original: f
Revised: f

Original: ⇤ ( h x i , ✓ ⇤ i ) 2 and random features  :
Revised: ⇤ ( h x i , ✓ ⇤ i ) 2 and random features  :

Original: R !
Revised: R !

Original: {+R N with
+}
Revised: {+R N with
+}

Revised: For simplicity of the analysis and exposition, we assume clean labels without additive noise, but note we may incorporate additive noise to our analysis in a straightforward manner.

Original: R N with ( u ) j = N  1 / 2  ( u  b j ) , we deﬁne the empirical risk as
Revised:  ( u ) j = N  1 / 2  ( u  b j ) , we deﬁne the empirical risk as

Original: [Equation]
Revised: [Equation]

Original: We speciﬁcally consider the estimator ˆ F ( x ) :
Revised: We speciﬁcally consider the estimator ˆ F ( x ) :

Original: = ˆ f ( h x, ˆ ✓ i ) , where ( ˆ f, ˆ ✓ ) are obtained by using gradient descent on c and ✓ to minimize L n ( c, ✓ ) .
Revised: = ˆ f ( h x, ˆ ✓ i ) , where ( ˆ f, ˆ ✓ ) are obtained by using gradient descent on c and ✓ to minimize L n ( c, ✓ ) .

Original: Such strategy appears to be reasonable, in light of theproperties of the population landscape studied in Section 4, namely that its only local minimisers arealso global, and correspond to ˜ f
Revised: Such strategy appears to be reasonable, in light of theproperties of the population landscape studied in Section 4, namely that its only local minimisers arealso global, and correspond to ˜ f

Original: = ˆ P  f ⇤ and ˜ ✓ = ✓ ⇤ .
Revised: = ˆ P  f ⇤ and ˜ ✓ = ✓ ⇤ .

Original: Given a sufﬁciently large sample size n , oneexpects the empirical landscape L n to concentrate around its expectation L and inherit its benignoptimization properties.
Revised: Given a sufﬁciently large sample size n , oneexpects the empirical landscape L n to concentrate around its expectation L and inherit its benignoptimization properties.

Original: However, the presence of a degenerate saddle at ( c, m ) =
Revised: However, the presence of a degenerate saddle at ( c, m ) =

Original: (0 , 0) ﬂattensthe landscape around the “equator” { ✓ : h ✓ , ✓ ⇤ i = 0 } , and a sufﬁciently large sample complexity is 214 ultimately required to ensure that gradient descent iterates can escape from it, given its dangerouslyclose initialization | h ✓ 0 , ✓ ⇤ i | = ⇥ (1 / p d ) .
Revised: (0 , 0) ﬂattensthe landscape around the “equator” { ✓ : h ✓ , ✓ ⇤ i = 0 } , and a sufﬁciently large sample complexity is 218 ultimately required to ensure that gradient descent iterates can escape from it, given its dangerouslyclose initialization | h ✓ 0 , ✓ ⇤ i | = ⇥ (1 / p d ) .

Original: Prior works have obtained the optimal sample complexity of n =
Revised: Prior works have obtained the optimal sample complexity of n =

Original: O ( d s ) (recall that s is the information exponent of the teacher function f ⇤ ) for recovering the teacher direction ✓ ⇤ either by employinga learning algorithm that explicitly learns individual harmonic polynomials [19] or by assuming thatf ⇤ is known a priori [3].
Revised: O ( d s ) (recall that s is the information exponent of the target function f ⇤ ) for recovering the target direction ✓ ⇤ either by employinga learning algorithm that explicitly learns individual harmonic polynomials [19] or by assuming thatf ⇤ is known a priori [3].

Original: Their argument goes essentially as follows:
Revised: Their argument goes essentially as follows:

Original: • The empirical optimization landscape (when regarded as a function only of the direction ✓ ) near the equator | m | ⌧ 1 is of the form L ( ✓ ) = ⇥ ( m s ) .
Revised: • The empirical optimization landscape (when regarded as a function only of the direction ✓ )near the equator | m | ⌧ 1 is of the form L ( ✓ ) = ⇥ ( m s ) .

Original: • In order to certify that the optimization does not converge to a suboptimal critical point onthe equator (i.e. kr L ( ✓ ) k & ✏ ) one must require that m & ✏ 1 / ( s  1) .
Revised: • In order to certify that the optimization does not converge to a suboptimal critical point onthe equator (i.e. kr L ( ✓ ) k & ✏ ) one must require that m & ✏ 1 / ( s  1) .

Original: • A standard uniform concentration argument of the form kr L ( ✓ ) r
Revised: • A standard uniform concentration argument of the form kr L ( ✓ ) r

Original: L n ( ✓ ) k .
Revised: L n ( ✓ ) k .

Original: q dn and the
Revised: q dn and the

Original: fact that m ⇠ d  1 / 2 at initialization together imply that n
Revised: fact that m ⇠ d  1 / 2 at initialization together imply that n

Original: = O ( d s ) samples are sufﬁcient to escape from the “inﬂuence” of the equator.
Revised: = O ( d s ) samples are sufﬁcientto escape from the “inﬂuence” of the equator.

Original: In order to repurpose these arguments to our setting, the relative scaling of the top-layer weights crelative to the direction vector ✓ will turn out to be crucial, as has also been observed in the literatureon lazy -vs- rich regimes [14, 49] in the context of overparametrized neural networks.
Revised: In order to repurpose these arguments to our setting, the relative scaling of the top-layer weights crelative to the direction vector ✓ will turn out to be crucial, as has also been observed in the literatureon lazy -vs- rich regimes [14, 53] in the context of overparametrized neural networks.

Original: We consider an idealized version of a Gradient Descent algorithm over the empirical loss L n in theinﬁnitesimally small learning rate regime, resulting in the gradient ﬂow ODE:
Revised: We consider an idealized version of a Gradient Descent algorithm over the empirical loss L n in theinﬁnitesimally small learning rate regime, resulting in the gradient ﬂow ODE:

Original: [Equation]
Revised: [Equation]

Original: where ⇣ is the relative time-scale between ﬁrst- and second-layer weights, and r S d  1 ✓ is the Riemannian gradient.
Revised: where ⇣ is the relative time-scale between ﬁrst- and second-layer weights, and r S d  1 ✓ is the Riemannian gradient.

Revised: The learning rate schedule ⇣ ( t ) is set so that only ✓ is trained ﬁrst and after aﬁxed amount of time, which we denote by T 0 , then c and ✓ is jointly optimized.

Revised: This is requiredto guarantee that the gradient ﬂow escapes the equator (see Lemma C.4), which is crucial since the(random) starting point for the ﬂow is close to the equator with high probability.

Original: Our main result establishes that this gradient ﬂow, under a simple schedule for ⇣ , efﬁciently ﬁnds anapproximate minimiser of the population loss, with an error explicitly quantiﬁed as a function of n :
Revised: Our main result, proved in Appendix C, establishes that this gradient ﬂow, under a simple schedulefor ⇣ , efﬁciently ﬁnds an approximate minimiser of the population loss, with an error explicitlyquantiﬁed as a function of n :

Original: Theorem 5.1 (Gradient ﬂow ﬁnds approximate minimizers) .
Revised: Theorem 5.1 (Gradient ﬂow ﬁnds approximate minimizers) .

Original: Let 0 <  < 1 / 8 and assume (i)n  C   1 d s log   1  , (ii)N  C   1 log (    1  1 ) , (iii) C f ⇤   / 2 < s 2 ↵ 2 s / 2 , and (iii) a targetf ⇤ satisfying the regularity assumptions 4.1.
Revised: For  > 0 and f ⇤ satisfying Assumption 4.1, suppose the following are true: (i) d = ⌦ (log 1  ) , (ii)  = O (1) (iii) n = ⌦ ( d s  4 log 1  ) , (iv)

Original: Then with probability (over the empirical data, therandomness of initialization and the draw of the random features) greater than 1/ 4  2  , the outputof the procedure 1 with ⇢ = ⇥ (N  1   s  ) ,  N 0 =
Revised: N = ⌦ ( 1  log 1  ) & N = O ( 1  (  4 nd ) 1 / (2( s +1)) ) , (v) N 0 = ⌦ (log 1  ) & N 0 =

Original: O    nd  2 s (2 s   1)( s +1) and T [-
-] = e ⇥   d s/ 2 returns( c T , ✓ T ) after time T > T 0 satisfying
Revised: O ((  4 nd ) 1 / (2( s +1)) ) ,(vi) ⇢ = ⇥ ( p NN  (2+ s ) / 2) , (vii) T [-
-]0 = ˜ ⇥ ( N 1+ s/ 2d  1+ s/ 2 ) , and (viii) T 1 = ⇥

Original: [Equation]
Revised: (  4 nd ) .

Original: Corollary 5.2 (Excess risk of Algorithm 1) . Under the assumptions of Theorem 5. [Equation] after optimizingthe regularisation parameter in the previous estimator
Revised: Then, if werun Procedure 1 for T = T 0 + T 1 time steps with the above parameters, its outcome (ˆ c, ˆ ✓ ) satisﬁeswith probability at least 12   the following:

Original: [Equation]
Revised: [Equation]

Original: one obtains an excess risk of the form

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Require: N 0 , ⇢ , T 0 , T 1 , N, 
Revised: Require: N 0 , ⇢ , T 0 , T 1 , N,

Original: Initialize ✓ (0)
Revised: Initialize ✓ (0)

Original: ⇠ Unif( S d  1 ) , c (0) ⇠ Unif( { c 2 R N ; k c k 2 = ⇢ ; k c k 0 = N 0 } ) .
Revised: ⇠ Unif( S d  1 ) , c (0) ⇠ Unif( { c 2 R N ; k c k 2 = ⇢ ; k c k 0 = N 0 } ) .

Original: Run Gradient Flow (11) with ⇣ ( t ) = 1 ( t > T 0 ) up to time T
Revised: Run Gradient Flow (11) with ⇣ ( t ) = 1 ( t > T 0 ) up to time T

Original: = T 0 + T 1 .
Revised: = T 0 + T 1 .

Original: Set ˆ ✓ = ✓ ( T ) , ˆ c = c ( T ) .
Revised: Set ˆ ✓ = ✓ ( T ) , ˆ c = c ( T ) .

Original: Require:  n 0 , ˆ ✓ from Algorithm 1.
Revised: Require:  n 0 , ˆ ✓ from Procedure 1.

Original: Set ˆ c = arg
Revised: Set ˆ c = arg

Original: min c L 0 n 0 ( c, ˆ ✓ ) as in (15).
Revised: min c L 0 n 0 ( c, ˆ ✓ ) as in (14).

Revised: Corollary 5.2 (Excess risk of Procedure 1) .

Revised: Under the assumptions of Theorem 5.1, after optimizingthe regularization parameter  , one obtains an excess risk of the form

Revised: [Equation]

Revised: [Equation]

Original: Fine-tuning the second layer.
Revised: Fine-tuning the second layer.

Original: After running Algorithm 1, we may include a ﬁnal ﬁne-tuningphase of training for second layer weights c alone, using a separate training sample ( x 0 i , y 0 i ) , i =1 , . . . , n 0 and a possibly different regularization parameter  n 0 .
Revised: After running Procedure 1, we may include a ﬁnal ﬁne-tuningphase of training for second layer weights c alone, using a separate training sample ( x 0 i , y 0 i ) , i =1 , . . . , n 0 and a possibly different regularization parameter  n 0 .

Original: More precisely, we set
Revised: More precisely, we set

Original: [Equation]
Revised: [Equation]

Original: where ˆ ✓ denotes the output of the previous gradient descent phase.
Revised: where ˆ ✓ denotes the output of the previous gradient descent phase.

Original: Note that this is a strongly convexoptimization problem, and can thus be optimized efﬁciently using gradient methods or by solvinga linear system.
Revised: Note that this is a strongly convexoptimization problem, and can thus be optimized efﬁciently using gradient methods or by solvinga linear system.

Original: While this may not be needed in practice, we use a different training sample fortechnical reasons, namely to break the dependence between the data and the kernel, which dependson the initial training sample through ˆ ✓ .
Revised: While this may not be needed in practice, we use a different training sample fortechnical reasons, namely to break the dependence between the data and the kernel, which dependson the initial training sample through ˆ ✓ .

Original: We note that such sample splitting strategies are commonlyused in other contexts in the statistics literature (e.g., [9, 12]).
Revised: We note that such sample splitting strategies are commonlyused in other contexts in the statistics literature (e.g., [9, 12]).

Original: We obtain the following guarantee.
Revised: We obtain the following guarantee.

Original: Proposition 5.3 (Excess risk of ﬁne-tuning) .
Revised: Proposition 5.3 (Excess risk of ﬁne-tuning) .

Original: Let  > 0 .
Revised: Let  > 0 .

Original: Let m =
Revised: Let m =

Original: h ˆ ✓ , ✓ ⇤ i , where ˆ ✓ is obtained from the previous gradient descent phase, and let ˆ c be the ridge regression estimator obtained from a fresh dataset D 0 of n 0 samples, N random features, 3 and regularization parameter  n 0 :=
Revised: h ˆ ✓ , ✓ ⇤ i , where ˆ ✓ is obtained from the previous gradient descent phase, and let ˆ c be the ridge regression estimator obtained from a fresh dataset D 0 of n 0 samples, N random features, 3 and regularization parameter  n 0 :=

Original: (  2 ⌧ 2 / k f 

Revised: (  2 ⌧ 2 / k f

Original: ⇤ k 2  n 0 ) 1 / (  +1) , and let ˆ F ( x ) =
Revised: ⇤ k 2  n 0 ) 1 / (  +1) , and let ˆ F ( x ) =

Original: ˆ c >  ( h ˆ ✓ , x i ) .
Revised: ˆ c >  ( h ˆ ✓ , x i ) .

Original: Assume
Revised: Assume

Original: [Equation]

Original: 0 & max

Original: [Equation]
Revised: [Equation]

Original: ⌘ With probability at least 1   over the random features, we have
Revised: ⌘ With probability at least 1   over the random features, we have

Original: [Equation]
Revised: [Equation]

Original: where the expectation is over the n 0 fresh samples, and is conditioned on the previously obtained ˆ ✓ .
Revised: where the expectation is over the n 0 fresh samples, and is conditioned on the previously obtained ˆ ✓ .

Original: Combining the two phases leads to the following generalization guarantee on the excess risk for
Revised: Combining the two phases leads to the following generalization guarantee on the excess risk for

Original: Algorithm 2.
Revised: Procedure 2.

Original: Corollary 5.4 (Excess risk of Algorithm 2) . Let0 <  < 1 / 8 and let  =
Revised: Corollary 5.4. Excess risk of Algorithm 2 Let  > 0 .

Original: ( s 2 ↵ 2 s / 2 C f ⇤) 2 /  ,assuming 4.1, with  C f ⇤that depends on regularity of f ⇤ .
Revised: As in Theorem 5.1, assume d  C log(   1 ) ,µ s = h h s , ⌃ h s i > 0 , and assumption 4.1 on f ⇤ with a constant C f ⇤ .

Original: Assume the following:
Revised: Let  = ( s 2 ↵ 2 s / 2 C f ⇤ ) 2 /  ,and assume the following on the sample sizes and number of random features for the ﬁrst phase( n, N, N 0 ) and ﬁne-tuning phase ( n 0 , N 0 ):

Original: [Equation]
Revised: [Equation]

Revised: and let ⇢ be as in Theorem 5.1.

Original: With probability at least 1 / 4  2  over the initial n samples, initialization, random features, we have
Revised: With probability at least 1 / 2  2  over the initial n samples,initialization, random features, we have

Original: [Equation]
Revised: [Equation]

Original: where C 1 and C 2 only depend on regularity of f ⇤ , ignoring the optimization error.
Revised: where C 1 and C 2 only depend on regularity of f ⇤ , ignoring the optimization error, and logarithmicfactors in n , d , and N .

Original: We can make the following remarks:
Revised: We can make the following remarks:

Original: • The empirical gradient ﬂow therefore escapes the inﬂuence of the degenerate saddle withoptimal sample complexity n = ⇥ ( d s ) [3].
Revised: • The empirical gradient ﬂow therefore escapes the inﬂuence of the degenerate saddle withoptimal sample complexity n = ⇥ ( d s ) [3].

Original: This is an instance of gradient descent suc- cessfully optimizing a non-convex objective without the strict saddle property as soon as s > 2 , building from the simpler optimization landscapes of [19, 3].
Revised: This is an instance of gradient descent successfully optimizing a non-convex objective without the strict saddle property as soon ass > 2 , building from the simpler optimization landscapes of [19, 3].

Original: This is in contrast with
Revised: This is in contrast withe.g.

Original: e.g. spiked tensor recovery problems [4] where the signal strength is substantially weaker, leading to complexity in the optimization landscape.
Revised: spiked tensor recovery problems [4] where the signal strength is substantially weaker,leading to complexity in the optimization landscape.

Original: • While the recovery of the high-dimensional direction is optimal in statistical terms,thejoint recovery of the link function is sub-optimal, by comparing Corollaries 5.2 and 5.4.In this sense, the ﬁne-tuning stage recovers the optimal sample complexity, and
Revised: • While the recovery of the high-dimensional direction is optimal in statistical terms,

Original: where the
Revised: the

Original: non-parametric rate no longer depends on the ambient dimension d .•
Revised: joint recovery of the link function is sub-optimal, by comparing Corollaries 5.2 and 5.4.In this sense, the ﬁne-tuning stage recovers the optimal sample complexity, and where thenon-parametric rate no longer depends on the ambient dimension d .•

Original: The time-scale separation scheduling for ⇣ is sufﬁcient but possibly not necessary.
Revised: The time-scale separation scheduling for ⇣ is sufﬁcient but possibly not necessary.

Original: Theanalysis of vanilla dynamics ( ⇣ ( t ) ⌘ 1 ) is challenging, since during the initial phase oftraining there may be adverse interaction effects between c and ✓ , which under naive analysis lead to sub-optimal sample complexity of n  O ( d 2 s ) .•
Revised: The analysis of vanilla dynamics ( ⇣ ( t ) ⌘ 1 ) is challenging, since during the initial phase oftraining there may be adverse interaction effects between c and ✓ , which under naive analysis lead to sub-optimal sample complexity of n  O ( d 2 s ) .•

Original: The time discretisation to turn Procedure 1 into a proper algorithm should follow fromstandardEuler time discretisation arguments, since L n has Lipschitz gradients.
Revised: The time discretisation to turn Procedure 1 into a proper algorithm should follow fromstandard time discretisation arguments, although the case where  = ReLU requires specialcare due to the non-smoothness of the loss. In such setting, such discretisation arguments do not hold for vanilla gradient descent in the worst-case [30], although may be recoveredby appropriately smoothing the objective prior to computing the gradient.

Original:  Conclusions
Revised:  Conclusions

Original: This work studies the ability of shallow neural networks to learn single-index models with gradientdescent.
Revised: This work studies the ability of shallow neural networks to learn single-index models with gradientdescent.

Original: Our main results are positive, and demonstrate their ability to solve a semi-parametric problem with nearly optimal guarantees.
Revised: Our main results are positive, and demonstrate their ability to solve a semi-parametric problem with nearly optimal guarantees.

Original: Interestingly, this success story combines elements from thefeature-learning regime —the ability to efﬁciently identify the hidden direction in high-dimensionsunder a non-convex objective, with ingredients from the lazy-regime, which offer better computational tradeoffs to solve low-dimensional non-parametric problems.
Revised: Interestingly, this success story combines elements from thefeature-learning regime —the ability to efﬁciently identify the hidden direction in high-dimensionsunder a non-convex objective, with ingredients from the lazy-regime, which offer better computational tradeoffs to solve low-dimensional non-parametric problems.

Original: Our technical analysis leverages tools from high-dimensional probability (such as uniform gradient concentration) and RKHSapproximation theory, and complements the growing body of theoretical work certifying the efﬁciency of gradient methods on non-convex objectives.
Revised: Our technical analysis leverages tools from high-dimensional probability (such as uniform gradient concentration) and RKHSapproximation theory, and complements the growing body of theoretical work certifying the efﬁciency of gradient methods on non-convex objectives.

Original: We have followed the standard approach ofﬁrst establishing benign topological properties of the population loss, and then extend them to theempirical loss.
Revised: We have followed the standard approach ofﬁrst establishing benign topological properties of the population loss, and then extend them to theempirical loss.

Original: That said, there are several unanswered questions that our work has not addressed, such as:
Revised: That said, there are several unanswered questions that our work has not addressed, such as:

Original: • Extension to Multi-Index Models.
Revised: • Extension to Multi-Index Models.

Original: Multi-index models are natural extensions of singleindex models where the hidden direction is replaced by a hidden low-dimensional subspace.
Revised: Multi-index models are natural extensions of single- index models where the hidden direction is replaced by a hidden low-dimensional subspace.

Original: They enjoy similar statistical guarantees as single-index models [19, 5], and thus a naturalquestion is whether the same algorithmic tools developed here extend to the multi-indexsetting.• Gradient dynamics without warmstart.
Revised: They enjoy similar statistical guarantees as single-index models [19, 5] and thus a natural question is whether the same algorithmic tools developed here extend to the multi-index setting.

Revised: • Gradient dynamics without warmstart.

Original: An unsatisfactory aspect of our results is the requirement that the algorithm starts by only optimizing ✓ for t < T 0 .
Revised: An unsatisfactory aspect of our results is the re- quirement that the algorithm starts by only optimizing ✓ for t < T 0 .

Original: It would be interestingto understand whether the vanilla dynamics could also succeed.• Trainable biases and untied directions.
Revised: It would be interestingto understand whether the vanilla dynamics could also succeed. • Trainable biases and untied directions.

Original: Our proposed neural network architecture isnon-standard, in the sense that its biases are frozen at initialization and all neurons sharethe same inner weight.
Revised: Our proposed neural network architecture isnon-standard, in the sense that its biases are frozen at initialization and all neurons sharethe same inner weight.

Original: For the purposes of learning single-index models, removing theserestrictions would not bring any statistical beneﬁts.
Revised: For the purposes of learning single-index models, removing theserestrictions would not bring any statistical beneﬁts.

Original: However, it would be interesting toextend our analysis to the general setting.
Revised: However, it would be interesting toextend our analysis to the general setting.

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.
Revised: The checklist follows the references.

Original: Please read the checklist guidelines carefully for informationon how to answer these questions.
Revised: Please read the checklist guidelines carefully for informationon how to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
Revised: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
Revised: You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:
Revised: For example:

Original: • Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .
Revised: [Yes] See Section ?? .

Original: • Did you include the license to the code and datasets?
Revised: • Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.
Revised: [No] The code and the data areproprietary.•

Original: • Did you include the license to the code and datasets?
Revised: Did you include the license to the code and datasets?

Original: [N/A]
Revised: [N/A]

Original: Please do not modify the questions and only use the provided macros for your answers.
Revised: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the
Revised: Note that the

Original: Checklist section does not count towards the page limit.
Revised: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
Revised: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] See Conclusions.
Revised: [Yes] See Conclusions.(c)

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: Did you discuss any potential negative societal impacts of your work?

Original: [N/A]
Revised: [N/A]

Original: This is apurely theoretical work with no practical applications.
Revised: This is apurely theoretical work with no practical applications.(d)

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] See Sections and 5.
Revised: [Yes] See Sections 4and 5.(b)

Original: (b) Did you include complete proofs of all theoretical results?
Revised: Did you include complete proofs of all theoretical results?

Original: [Yes] See Section 4.3.
Revised: [Yes] See Section 4.3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes] Supplementalzip ﬁle.
Revised: [Yes] Supplementalzip ﬁle.

Original: (b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes] See Section D.
Revised: [Yes] See Section D.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See Section D.
Revised: [Yes] See Section D.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Section D.
Revised: [Yes] See Section D.

Original: If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A]
Revised: [N/A]

Original: (d) Did you discuss whether and how consent was obtained from people whose data
Revised: (d) Did you discuss whether and how consent was obtained from people whose data

Original: you’re using/curating?
Revised: you’re using/curating?

Original: [N/A]
Revised: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁ-
Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁ-

Original: able information or offensive content?
Revised: able information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A](
Revised: [N/A]

Original: c) Did you include the estimated hourly wage paid to participants and the 
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: total

Original: amountspent on participant compensation
Revised: spent on participant compensation?

Original: [N/A
Revised: [N/A]