Learning single-index models with shallow neural networks
Authors: Alberto Bietti
Joan Bruna
Clayton Sanford
Min Jae Song
Venue: NeurIPS 2022 Accept (2022)
Keywords: single-index models
gradient descent
landscape analysis
random feature approximation
shallow neural networks

Abstract: 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. These models are particularly relevant in high dimension, when the data might present low-dimensional structure that learning algorithms should adapt to. While several statistical aspects of this model, such as the sample complexity of recovering the relevant (one-dimensional) subspace, are well-understood, they rely on tailored algorithms that exploit the specific structure of the target function. In this work, we introduce a natural class of shallow neural networks and study its ability to learn single-index models via gradient flow. More precisely, we consider shallow networks in which biases of the neurons are frozen at random initialization. We show that the corresponding optimization landscape is benign, which in turn leads to generalization guarantees that match the near-optimal sample complexity of dedicated semi-parametric methods.

Learning Single-Index Models with Shallow Neural Networks
Alberto Bietti
Joan Bruna
Clayton Sanford Columbia University
Min Jae Song New York University
New York University
New York University
Meta Al
Abstract
Single-index models are class of functions given by an unknown univariate "link" function applied unknown one-dimensional projection of the input. These models are particularly relevant in high dimension, when the data might present low-dimensional structure that learning algorithms should adapt to. While several statistical aspects of this model, such as the sample complexity of recov- ering the relevant (one-dimensional) subspace, are well-understood, they rely on tailored algorithms that exploit the specific structure of the target function. In this work, we introduce natural class of shallow neural networks and study its ability to learn single- index models via gradient More precisely, we consider shal- low networks in which biases of the neurons are frozen at random initialization. We show that the corresponding optimization landscape is benign, which in turn leads to generalization guarantees that match the near-optimal sample complexity of dedicated semi-pant parametrratticccc methods.
Introduction
High-dimensional learning with both computational and statistical guarantees, which particularly relevant given the current scaling trends, remains an outstanding challenge. One important question which has received considerable attention is on understanding the advantages of using non-linear learning models, such neural networks, over more mature (from theoretical standpoint) coun- terparts, such as kernel methods [56, 57, 84, 22J. Perhaps surprisingly, the question remains largely open even for shallow neural networks.
While approximation benefits of shallow neural networks over non-adaptive kernels have been known for decades [9, 69], another important piece of the theoretical puzzle was provided by [5], whose analysis hinted at an inherent statistical advantage of neural networks for extracting infor- mation from high-diiensiooal data with a "hidden" low-dimensional structure. Providing computa- tional guarantees, the remaining piece of this puzzle, still mostly unresolved.
Several computational hardness results for learning functions that can be efficiently approximated
by shallow neural networks have been established in the literature [30, 41, 28, 78, 20], ruling out positive results in the general setting. On the other hand, progress has been made on the positive side [1, 2, 73] by focusing on function classes with strong structural properties, thereby showcasing the adaptive representation learning capabilities of neural networks.
This work aligns with the latter effort, and focuses on the class of single-index models. Single- index models are high-dimensional functions F Rd + IR of the form F.(x) = f= x)), where both the univariate "link" function fa : R R and relevant (one-dimensional) subspace 9t 1d-1 are unknown. These models have been extensively studied in the statistics literature [45, 48, 26, 46, 33], leading to dedicated algorithmic procedures, and can be provably approximated with shallow neural networks without incurring in curse-of-ding d--eemnssionattty [5]. In contrast, the analysis
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

of neural network learning using gradient-basee methods has focused on the so-called "Teacher- Student" setup [43, 39, 85, 86, 10], where the link function f. is assumed to be known and used as the activation function for the student.
The mix of high-dimensional parametric component (the hidden direction) with non-parametric one in low dimension (the link function) in single-index models naturally suggests a shallow neu- ral network architecture where the inner weights are shared and "active", while the biases are "lazy" [23]. We instantiate such an architecture by freezing the biases at random initialization, and analyze gradient descent on the free parameters in the continuous-tin limit.
Our main results establish that as soon as the width N of the network is larger than a quantity which depends solely smoothness properties of the (univariate) link function fx, gradient flow recovers the unknown direction 0" with near optimal sample complexity (dd)), where 8 is the so- called enformation-eppoeen of the link function [10] (at least when 2 see Theorem 1 for the formal result), and approximates the univariate link function f. near-optimally (see Corollary 6.4). The information exponent roughly captures the signal strength, which here refers the alignment between the network direction 0 and the hidden direction 0*, at typical initializations.
The success of gradient flow relies on the benign optimization landscape of the empirical loss, though the presence of degenerate saddles necessitates a careful analysis leveraging uniform con- vergence of the empirical landscape [60, 35J. We show that gradient flow over our proposed neural network architecture solves two distinct problems- sion and non-convex optimization in high uniiviit-uivariatt non parammetrii kernel ridge regres- and efficiently, cementing its role as a versatile algorithm for high-dimensional learning. We illustrate our theoretical results with experiments in Section A.
2 Related Work
Single-index and multi-index models. A useful modeling assumption in high-dimensional re- gression is that the regression function F(x) Ely]ar] only depends on one or few direc- tions. This leads to the single-index model F(z) = ()),,, and multi-index model F(x) = x)) (0% x)). with & typically much smaller than the dimension. Such models have a long history in the statistics literature and different methods exist for various estimation problems, including projection pursuit [36, 47], slicing [52], gradient-based estimators [53], and moment- estimators [26]. When the function f- is also to be estimated, we face a
problem involving parameter recovery of 6" and non- non-parametric estimation of f.. Our work closely re- lated to [33], which also characterizes the population landscape of certain objectives by leveraging Gaussian data. Multi-index models are also studied in [5] in the context of shallow neural net- works, where it is shown that certain models of infinite-width shallow networks can adapt to such low-dimensional structure, though no tractable algorithms are introduced.
The works [8, 19, 66] show that certain neural networks trained close to initialization can learn certain sparse polynomials which take the form of multi-index models, but such networks do not directly aim to learn target directions. Recently, [1] studied the Icarnability of functions on the hypercube by shallow neural networks with stochastic gradient descent and introduces the merged stairease property, which provides necessary and sufficient conditions for learnability with linear sample complexity n = Od). While they learn a broader class of functions (including multi-index model) for a more efficient sample complexity regime (O(d) vs O(d") )), their setup is restricted to simple discrete data distributions, while our work captures the regime of semi-parametric estima- tion by considering Gaussian data without the sparsity requirements on F implied by their merged staircase property.
Concurrently to our work, [4] and [27] studied the learnability of certain single and multi-index mod- els on Gaussian data with shallow networks, by performing a single gradient step on the first layer before fitting the second layer. While the single step is sufficient to provide a separation from kernel methods in these works, we show that optimizing both layers jointly until convergence (for a more simplistic architecture) can significantly improve the rates, by fully decoupling the non- -arramaett learning part from the high-dimensional inference of the hidden direction. Finally, recently [64] studied the ability of shallow neural networks to learn certain single and multi-index models, show- ing in particular that SGD-trained ReL.U networks can learn single-index functions with monotonic index function (corresponding in our setting to s = 1) with linear (up to logarithmic factors) sam-
2

---

ple complexity. Our results therefore extend such positive guarantees to broader class of index functions with arbitrary information exponent.
Teacher-student models. In the context of neural networks, several works have considered the stuueen-ttudet setting [34], where the target function takes the form of neural network with the same activation as the network used for learning [43, 39, 85, 77, 86, 10, 82]. In this case the problem does not involve non-parne non-parametric estimation as in our setup, but this line of work often involves studying optimization landscapes similar to ours for estimating hidden directions. In particular, the population landscape appearing [10] is similar to ours, based on Hermite coefficients of link functions. The follow-up work [11] extends this to multiple student neuron directions, but still focuses on parametric rather than non-parametric statistical problems.
Kernels and random features. In order to obtain non-parametric estimation guarantees for learn- ing the target function f. the single-index model, our work builds on the kernel methods literature for approximation and con-parametric regression [76, 13, 7J, their links with neural networks [24, 5], and in particular on random feature approximation [70, 6, 72, 61].
Non-convex and non-smooth optimization landscapes. There is vast literature studying tractable non-convex optimization landscapes, arising from high-dimensional statistics and statis- tical physics [59, 58, 37, 12, 79, 17, 39, 71, 55J. A particular aspect of our setup is that the opti- mization landscape does not have the strict saddle property, which is often leveraged to establish global convergence [50]. [60, 35] study concentration properties of the empirical landscape to the population one for non-convex problems including generalized linear models. Our results rely on similar concentration analyses, but depart from these previous work by also allowing optimization of the link function, and by supporting the non-smoothness arising from ReLU activations. On the algorithmic side, we consider gradient flows non-convex and non-smooth landscapes, which quire careful technical treatment, but have been studied by previous works [32, 29, 49]. We refer the interested reader to Appendix for more details on this technical issue.
3 Preliminaries
We focus on regression problems under single-index model with Gaussian input data. Specifically, we assume d-dimensional inputs x "/a := (0, Ia), and labels
y == F (x) + (
f=0 (((0"..)) x)) +
where 6* gdd1 and & ~~ (O, o²) is independent, additive Gaussian noise. The normalization O" gd-1 is to ensure that both f. and 0" are identifiable.
Shallow networks and random features. We consider learning algorithms based on shallow neu- ral networks of the form
((r,,,6)=c-- ((((,))))
VN Σ
with 0 E gd-1, where d(u) = max (0, y is the ReLU activation. b, NN(0, 7²) (we assume T > 1) are random "bias" scalars that are frozen throughout training, and E; are random signs with Rademacher distribution (i.c., uniform over {+1]), independent from 6, and also frozen during training. The resulting vector of random features is thus (uu) 1
vN (o)
The choice of ReLU activation is motivated by its popularity among practitioners. As we shall see, the fact that é is non-smooth introduces some technical challenges, but its piece-wise linear structure enables dedicated arguments both in terms of approximation as well as in the study of the optimization landscape. In Appendix G we discuss how our main results are affected when replacing the ReLU by a smooth activation, especially when choosing it such that ø' is Lipschitz.
Empirical risk minimization. The supervised learning task is to estimate F* (and therefore both f. and 0*) from samples {(r;, y/ )=)).. We will focus on mean-squared error with Tychonov regularisation, determined by the following losses.
3

---

Definition 3.1 (Population risk). We define the -Fregularrzed population loss by [(c, 0) = E [ G(a; c, 0)) + A|c|| = -E - [ 2.y
+2 A/|c|/
Definition 3.2 (Empirical risk). We define the f-regularized erpppriiiieddmmiiiaa empirical loss by
[n(c, Σ (c"0((0,z))) - A||e||2.
9.)² 9.))
(I)
Hermite decomposition. Given that our data is normally distributed, we consider the family of (normalized) Hermite polynomials {hyljen, which form orthonormal basis of (7), the space of squared-integrable function under the Gaussian measure := (O 1). We will denote by f- E3 aryh the Hermite decomposition of the target link function, which assume is in (7) henceforth.
We apply the following useful properties of Hermite polynomials [67, Chapter 1 11.2]:
h/j Vihj-1 and (hy( +), hy'( (()' ,))) = و5.r (0,00000)) 0'),,
where )s the inner product in 1,2 (74) and 5 the Kronecker delta. We will assume throughout that || D3 az. || = Ej jaJ. and || = ΣØ - 1)az are finite (see As- sumption 5.2). We will also consider the weighted Sobolev space 2 (7)) which contains functions 00.
f= E, ayh3 (7) such that
Random features to Hermite coefficients. To precisely characterize the landscape of L(c, 0), we introduce notation to represent each random feature function in terms of Hermite polynomials hj. We define the linear integral operator T : (T) 7 RV by
1
(Tf): = (f. (f = VN 2 2
E [f())((e) b,)l ie [N].
(2)
Note that T has rank N almost surely.
The operator T has adjoint T NR - (7) defined as (T 1 E GØle,u-bi)
VN c. (u). Finally, for any E N. let TT E R defined as T, Thy. We can then write down the Hermite expansion of the student network:
G( c ((1,,)))
Σ (c. (T),(((0,x)).
120
Denoting m = (0,0"), the regularized population objective can expressed as
I(c,0) Σ 3 + Σ (c T, -2 2 - Ty)m + A||c||?
(3)
where the term Σ, is constant that can be ignored. Let Q = TT E RNNN be feature covariance matrix and Q2 = Q + AT. Note that Qx is positive definite for X >> 0. We define the regularized projection Pa E + -1 onto the random feature space for E = T-T. Observe
that (P) TT where ex IR. is the solution the following objective:
min || - - 1 AlIdIE CERN
(4)
Geometry on the sphere. Because the direction constrained to on the sphere, our opti- mization algorithms rely spherical (Riemannian) gradients, which are defined as follows:
L(c, = 1Ig= VoL(c,
where NT v w - 00, We say that (c. 0) critical poins of I if V T(e, and V V.L(c, 0) 0.

---

Univariate Approximation using Random Features
Before addressing the high-dimensionality of the learning problem, first focus the parametric approximation aspects of the univariate link function. As usual, we start by deriving approximation rates of the infinitely-wide model, given by RKHS, and then establish approxima- tion rates for our random feature model.
Univariate RKHS. If we fix the direction 8, learning alone may be model [70] that approximates a kernel method with the following kernel.
random feature
k(u, v) Ea- (o(uu b)Φ(Ev 6)] u, 1 DEE.
(5)
where 7r = (O, 7²) is the Gaussian measure on R with mean zero and variance 7², and e is random sign with Rademacher distribution. The kernel x corresponds an RKHS H. given by
H 1
f(u) +/2 [/|1+0000 [0) b) cb)ob( b)] dy, Ct. e_ (1))))
The following lemma characterizes the corresponding RKHS norm || lly. and follows from The- orem B.8, by noting that ro(u, v) = (V(u),
with (b(x) v2 (ō(u :), 0- - .)..
Lemma 4.1 (RKHS norm). The RKHS norm in H given by
I1 = inf ||ct + lle- : f(u) -2 f [c- (b)(() - b) + + (b)o(-u d d7)
The choice of ReLU for the activation function gives more explicit control over the RKHS norm, based on Sobolev representations, as already exploited by several works [68, 5, 75].
Lemma 4.2 (RKHS norm bound). Let J E HT) n C²RR) and T 1. If J and I' both have polynomial growth and f' (f di < co, then f E H with 7+(1)
I S 6T ( (11" di + | - 611 2(J,
7+(1)
The proof is in Appendix C.I.
RKHS approximation properties. Let A(f,A) A) be the (regularized) I,2 approximation error for functions the space H. with respect the target function f and measure ?. Formally,
A(f.A) = min If g|le; + Algg1/ gEH
We will now show that the approximation eиror of the RKHS corresponding to an infinite number of random features can be bounded in terms of the regularization À and the 7-norm of the second derivative of the target function. For that purpose, we consider the following 'source' condition to ensure polynomial approximation error in A.
Assumption 4.3 (Containment in (*)).. and define K := inf [B 2 [If")) I BIJ I(
effffeeee H²()) | f" e L*())) We assume f E F
Assumption 4.3 provides a sufficient condition for approximating f with functions in the RKHS. The family of approximants (hus H M > 0} use in Lemma 4.4 are exactly equal to f on [-M, M and are linear outside of [-M, M). The ItA assumption on f" ensures control over the RKHS norm of has Note that by Jensen's inequality, (*(r) C I2 I? K is always well-defined for f" & L L'(r). Sigmoidal functions, compactly supported smooth functions, and, more generally, functions with polynomial growth satisfy Assumption 4.3.
Lemma 4.4 (RKHS approximation error). Let X E (0, 1) and f E F. Then, there exists a universal constant C 0 such that
A(f, & o ( 2 AC7).
(8)
1-1//² where B = 3+1777 and C2 max ||
5

---

The proof appears in Appendix C.2. This lemma allows us to control the RKHS approximation error of a target function in terms of their Hermite decompositions. The main technical difficulty is that the RKHS integral operator E does not diagonalise in the Hermite basis; we address this with a dedicated argument exploiting the RKHS Sobolev representation of Lemma 4.2. The assumption that f E (7) (Assumption 4.3) is sufficient for our purposes but not necessary for polynomial- in-À approximation rates. In Section H, we show that the ReLU function Φ(t) = max(0, t), which Lipschitz but not (1), as the targel satisfies A(o, A) S -²\2/3 using a direct argument. Extending the class of functions approximable by H with polynomial-in-X rate is an interesting future direction.
Random feature approximation. We now consider (finite) random feature approximations to functions in the RKHS. Lemma C.2 shows that the best possible loss of a linear combination of sufficiently many finite features is bounded above by the best approximation with infinitely many features with high probability. More specifically, as long as N Z A-1, the random feature ap- proximation error behaves like the RKHS approximation error. The proof leverages the 'degrees of freedom' the kernel and closely tracks [6].
Population Landscape under Frozen Random Biases
To understand optimization and generalization properties of gradient flow on the empirical loss L,, (c, 0), we first the study optimization landscape of the population loss L(c, 0). We characterize the critical points of the population loss and show that gradient flow on a shallow neural network of sufficient width N converges only when its direction vector either parallel or orthogonal to the
the target direction 0- Importantly, the sufficient number of random features N depends (e-regularization parameter À E (0, but not on the input dimension In Section 6, we further show that sufficiently large the number of training samples, guarantees similar properties for the empirical landscape and thus has favorable generalization properties for most initializations.
One of the main measures of complexity for the target link function fa is its information exponent (see e.g., [10]), defined as follows.
Definition 5.1 (Information exponent). Let f R + R be any function such that f E L²( The information exponent of J, which we denote by s, is the index of the first non-zero Hermite coefficient. That is, == min E N:: a; 03.
We make the following regularity assumptions on the target link function f- to ensure small approx- imation error by random features, and benign population and empirical landscape.
Assumption 5.2 (Regularity of f.). We consider f. E I,2(T), with f. E, ayhh- Assume 1) f. Lipschitz, 2) Ej ja |oy oc, and 3) (²) E, + 1)as+2h5(2) is in IA(r)) (Assumption 4.3)
We also suppose w.l.o.g. that f. is normalized so that Il,f = 1. To analyze the critical points of [(c, 0), we introduce the projected population loss L()), which can be seen a semiparametric least squares (SLS) objective [48].
L(0) = min L),
Theorem 5.3 (Critical points of the population loss). Assume fa satisfies Assumption 5.2 and has information exponent s 2 1. For r > 1, and 5 E (0, 1), there exists X" < 1 depending only on and the target link function f. and universal constant C. > 0 such that if
A X and
N2 NZİ log
99)
X
A5
then with probability 1 1-6 8 over the random biases by and signs Ej. j 1... N, the set of first
critical points = {cc 0): V L(c,, 0) = O, VeI(c, 0) = 0} satisfies:
1. (orientation relative to 0*) if (c, () E Ω, then either 0 E {+0r} or (0,0r) = 0.
2. (existence and uniqueness of (e)( if v T.(0) O, then there exists a unique e E R such
that () E Ω.

---

Theorem 5.3 thus establishes a benign optimization landscape in the population limit, rejoining several known non-convex objectives with similar behavior, such as tensor decomposition [40] or matrix completion [38]. Importantly, this optimization landscape has the same topology as the one that arises from using the Hermite basis, the tailored choice for data generated by a single-index model in Gaussian space [33, Theorem 51, instead of random scalar features. We view this as an interesting robustness property of shallow neural networks, at least in the regime where biases are randomly frozen.
Empirical Landscape and Generalization Guarantees
Section 5 shows that the population landscape has a relatively simple structure given N = 9a((1) random features. We now study the optimization properties of its finite-sample counterpart.
We consider the estimator F(x)= == f((,, 0)), where (f, 0) are obtained by running gradient flow on c and 0 to minimize (c, 8). Such strategy appears to be reasonable in light of the properties of the population landscape, since its local minimizers are also global and correspond to f = PA." and 0" (Theorem 5.3). For a sufficiently large sample size n. one expects the empirical landscape Ly concentrate around its expectation L and inherit its benign optimization properties. However, the presence of degenerate saddle at (c, m) = (0, 0) for m = (0, 6*) flattens the landscape around the "equator" {0: (0, 0*) = 0}. Thus, more samples are required ensure that gradient flow escapes from the equatorial region despite its dangerously close random initialization (60, 8')| = (()/V)).
Prior works have obtained sample complexity of n = Odd)), where we recall that s is the informa- tion exponent of the target function fe. for recovering 0* either by employing a learning algorithm that explicitly learns individual Hermite polynomials [33] or by assuming that f. is known a priori [10].' ! The intuition behind this sample complexity is roughly as follows.
The empirical optimization landscape (when regarded function only of the direction 0) near the equator (Iml < 1) is of the form L()) x m
In order certify that the optimization algorithm does not converge to suboptimal critical point (i.e., [ I <) on the equator, one requires that m 2 c!/(s-1).
A uniform gradient convergence bound of the form |VI(() V VI.,, (0)1 d/vd/n) and the fact that m @ (1)vd) initialization together imply that (dd) samples are sufficient to escape from the "influence" of the equator.
In order to repurpose these arguments to our setting, the relative scaling of the top-layer weights c relative to the direction vector is crucial, has also been observed in the literature on lazy-vs-rich regimes [23, 84] in the context of overparametrized neural networks.
We consider an idealized version of Gradient Descent over the empirical loss L, in the infinitesi- mally small learning rate regime. This results in a gradient flow ODE of the form:
ē(t) = - ccn((e)) -1-1
0(1) = - - Ln(c,
(10)
where ( is the relative scale between and 0, and ad 0 is the Riemannian gradient.
Specifically, we study setting where ((t) 1(t To) for an appropriately chosen time To- This choice produces a two-stage gradient-flow. During the first phase, up until time To. we only optimize the first-layer parameter from random initialization. In the second phase, the parameters and / are jointly optimized. Additionally, the first phase only utilizes small fraction No < of the random features employed in the second phase. Our procedure implements this by randomly initializing <(0) N sparse vector with No non-zero components. The overall approach is described in Procedure 1.
'Actually, in [10] the authors obtain a slightly improved sample complexity of O(d for s 2 3 by directly analyzing SGD with fixed step-size, as well as a matching lower bound (for SGD in the small step-size regime) up to polylogarithmic factors.
7

---

Our main result, proved in Appendix E, establishes that this gradient flow efficiently finds an ap- proximate minimizer of the population loss, with an error (explicitly quantified as function of n) that reveals the fundamental role of the information exponent of f.. On top of the regularity con- ditions on the target link function of Assumption 5.2, the upper bound on X, and the lower bound on N from Eq. (9) of Theorem 5.3, the main result imposes (compatible) upper bound on N and an appropriate choice of initial norm (p) and sparsity (No/N) for c(O). In this section, we are interested in behavior as d, and N grow asymptotically and hence treat the target function f. and terms de- rived from it (including Hermite coefficients aj and information exponent 8), along with the bias parameter T and regularity parameter 3, as constants and omit them from asymptotic notation.
Theorem 6.1 (Gradient flow finds approximate minimizers). For E (0, 1/4) and fa satisfying where Dena == }), (iii) N = =0 log tt) &
Assumption 5.2, suppose the following are true: (i) X = ((1) and =
max (V d+NN ( (2)22/222-1)) (ii) n = M(max (d+N)d'- 1///31/2 X
X
N = Ö(A (iv) No (log( (v) p == @(VNN, (2+0)/22 (7² + XN//No) (vi) To ((d=/2-1), and (vii) Ti = Ô 6(2xx) Then, ie run Procedure for T = To + Ti time steps with the above parameters, with probability least - 5 we have
1- 1-1(rr 01) = Ō Ö max
d [d+N N d') n2
11
(11)
The empirical gradient flow therefore escapes the influence of the degenerate saddle with sample complexity = ((d") when À @(1) and > 2. This is an instance of gradient flow successfully optimizing a non-convex objective without the strict saddle property as soon as s > 2, building from the simpler optimization landscapes of [33, 10]. This is in contrast, for example, with spiked tensor recovery problems [12] where the signal strength is substantially weaker, leading to complexity in the optimization landscape. This sample complexity nearly matches the tight lower bound n >> d=-1 of [10], obtained in the case >>> 2 and applies to SGD rather than batch gradient descent, as is our case. For s E {1.2}, the sample complexity becomes df and d2.5, respectively, but we note that these may be improved to d' when using smooth activation (see Appendix G). For s = 2, this is comparable [27]. which requires (dd)) samples, but still above the n >> dloog of [10].
We emphasize that the "near-optimality" of our sample complexity n O(d)) only pertains to gradient-based methods in small learning rate regimes [10]. In fact, alternative methods have been shown to achieve a better sample complexity of Ö(al=/2) ) in the setting where f. is a certain degree- s polynomial with information exponent s [19], leveraging tensor factorization tools. We leave it as an interesting open question to further understand the nature of this gap.
We note that the dependence on d + N in the recovery guarantee can likely be improved to d using a more refined nbrm- landscape concentration analysis. We also remark that if we chose the number of random features N O 1 log ). then the requirement À 12 for large enough sample size n imposes À >> n 1/5 This lower bound on À guarantees that critical points near initialization obtain an excess risk that vanishes with n, with a rate independent of d, as we now show.
be escaped, but may slow down learning. Nevertheless, this is sufficient to
Corollary 6.2 (Excess risk of Algorithm I). Under the assumptions of Theorem 6.1, and further assuming n Z d3, an appropriate choice of X yields an excess risk guarantee of the form
(( n
1
|
1 =
F
+
(12)
where B is defined as in Lemma 4.4.
This result indicates that the joint training is consistent, with excess risk that vanishes with a rate with explicit dependence on the ambient dimension d and the non non-paramttric exponent S. However,
to ensure enough gradient it requires a regularisation strength À = e maxx { (() デ (*) ~ोス
concentration, which happens to be larger than the optimal regularisation of the univariate kernel ridge regression associated with learning f.. We are thus 'over-regularisingg as a consequence of the joint training, resulting in a slower rate than what would be dictated by estimating 0. and f. separately. A simple mechanism to break this inefficiency by considering a fine-tuning step of the second-layer terms.
8

---

Procedure Gradient Flow
Require: No, P. To, T1, N., and X. Initialize 8() ~ Unif(Se-1) c(0) ~ Unif E ; ||e|l2 = p: I|clo = No)).
Run Gradient Flow (10) with ((t) 1(t > To) up time T = To T1.
Set 0 (TT). c(T).
Procedure 2 Fine-Tuning
Require: 0 from Procedure 1, and Ant Set è = arg min, Li (c, 0) in (13).
Fine-tuning the second layer. After running Algorithm 1, may include final fine-tuning phase of training for second layer weights с alone, using separate training sample (22 (6), i .,..., nt' and a possibly different regularization parameter Aye, More precisely, we set
ē= arg min L., 0)
min
T Σ
(c 4 (10, v/)² + An [|cll]]]
(13)
where 9 denotes the output of the previous gradient descent phase. Note that this is a strongly convex optimization problem, and can thus be optimized efficiently using gradient methods by solving a linear system. While this may not be needed practice, we use a different training sample for technical reasons, namely to break the dependence between the data and the kernel, which depends on the initial training sample through Ò. We note that such sample splitting strategies are commonly used in other contexts in the statistics literature (e.g., [15, 21J). We obtain the following guarantee.
Proposition 6.3 (Excess risk of fine-tuning). Let E 1/4). Let m (0", where Ô is obtained from the previous gradient descent phase, and let è be the ridge regression estimator obtained from fresh dataset 'D' of n' samples, N random features, and regularization parame- ((0+)), and let F(x) ēT (10,
ter Am' = (a²,²
Assume
n' 2 max { (0²>2/115" /1,/11 (I. /0 11..11e// (02,2//1012) and (n
N 2 (n' M1 /0>>) p.ं log
Then with probability at least 1 over the random features, we have
0 n
E - F1 1 D
+ (1-
(14)
where the expectation is over the n' fresh samples, and is conditioned on the previously obtained ò.
Decoupling the regularization parameters of the two phases (along with number of random fea- tures N) allows us keep large in the first phase, leading to fast recovery as per Theorem 6.1, while obtaining vanishing excess risk through decreasing Awt. This is illustrated in the result on the excess risk for Algorithm 2.
Corollary 6.4 (Excess risk of Algorithm 2). Let 8 E 1/41. As in Theorem 6.1, let Pss
(h.; Eha) > O, and let f. satisfy Assumption 5.2. Let = ((1), and assume the following on the sample sizes and number of random features (n', NN);
featurse for the first phase (n, N, No) and fine-tuning phase
No (2 log >
1 11:
= max d(o+3)/2)) N' = (n'==t) =('
and lei p be as Theorem 6.1. With probability at least 1/2 28 over the initial n samples, initialization, random features, we have
E. E Fo- 11-1 1
S
max
n n²
d d
8+1
(15)
where the constants in Ô do not depend on d other than through logarithmic factors.

---

Comparing Corollaries 6.2 and 6.4, we observe that the fine-tuning stage recovers the optimal sample complexity, where the non -on--aamttticc rate fully independent of the ambient dimension d, while in Corollary 6.2 there is still dependence in the constants. We make the following additional remarks:
The time-scale separation schedule for Ç in Theorem 6. 1 is sufficient but possibly not necessary. The analysis of vanilla dynamics (C(t) = 1) is challenging, since during the initial phase of training there may be adverse interaction between e and 0, which under naive analysis lead to sub-optimal sample complexity of n 2 ((d)). Observe that this separate analysis of 'weak' and 'strong' recovery phases of learning appears in most contemporary related work [27, 1, 4, 10, 11].
The time discretization to turn Procedure into a proper algorithm should follow from standard time discretization arguments, although the case where & = ReLU requires special care due to the non-smoothness of the loss (see Appendix F for further discussion). In such setting, such discretization arguments do not hold for vanilla gradient descent in the worst-case [51], although these may be recovered by appropriately smoothing the objective prior to computing the gradient, or by using instead a smooth activation function (see Appendix G).
7 Conclusion
This work studies the ability of shallow neural networks to learn single-index models with gradi- ent descent. Our main results positive, and demonstrate their ability solve a semi-parametric problem with nearly optimal guarantees. Interestingly, this success story combines elements from the feature-learning regime, i.e., the ability to efficienlly identify the hidden direction in high- dimensions under a non-convex objective. with ingredients from the lazy-regime. Our technical analysis leverages tools from high-dimensional probability (such as uniform gradient concentration) and RKHS approximation, and complements the growing body of theoretical work on the efficacy of gradient methods for non-convex objectives. We have followed the standard approach of first estab- lishing benign topological properties of the population loss, and then extending them to the empirical loss. There are nonetheless several unanswered questions that our work has not addressed.
Weaker regularity and discrete-time analysis. Our approximation rate for ReLU as the larget (see Appendix H) suggests that the polynomial-in-A approximation rate may be extended to func- tion classes beyond F c (+))), such as Lipschitz. functions with smooth tail behavior. Thus, it would be interesting to extend our empirical landscape concentration results to such functions satisfying weaker regularity assumptions, which currently rely on certain polynomial decay of the Hermite coefficients (see Assumption 5.2). Additionally, by using a smooth activation function (see Appendix G). our GF dynamics can be discretized and turned into GD with analogous sample and time complexity. In that context, a natural goal to compare quantitatively the differences between GD with multiple passes over the training data and SGD by adapting tools from [10, 11].
Trainable biases and untied directions. Our proposed neural network architecture is non-standard, in the sense that its biases are frozen at initialization and all neurons share the same inner weight. For the purposes of learning single-index models, removing these restrictions would not bring any statistical benefits. However, it would be interesting to extend our analysis to the general setting where the first layer weights are not tied and biases are not frozen.
Extension to multi-index models. Multi-index models are natural extensions in which the hidden direction 4* is replaced by hidden low- dow-dimensional subspace. Typically, multi-index models enjoy similar statistical guarantees as single-index models [33, 5], and thus a natural question is whether the same algorithmic tools developed here extend to the multi-indes setting.
Gradient dynamics without warm-start. An unsatisfactory aspect of our results the requirement that the algorithm starts by only optimizing 0 for t < To- It would be interesting to understand whether the vanilla dynamics can also succeed provably.
Acknowledgments. We are thankful to Enric Adix-A, Alex Damian, Cédric Gerbelot, Daniel Hsu, Jason Lee, Theodor Misiakiewicz, Matus Telgarsky, Eric Vanden-Eijnden, and Denny Wu for useful discussions. We also thank the anonymous NeurlPS reviewers and area chair for helpful feedback. JB. AB and MJ are partially supported by NSF RI-1816753, NSF CAREER CIF 1845360, NSF CHS-190100 1, NSF Scale MoDL DMS 2134216, Capital One and Samsung Electronics. CS is supported by an NSF GRFP and by NSF grants CCF-18 181773 and IIS-1838154.
10

---

References
[1] E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. The merged property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks.
arXiv preprint wxXi:2202.08653. 2022. 1, 2, 10
[2] Z. Allen-Zhu and Y. Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04411: 2020.
[3] L. Ambrosio, N. Gigli, and G. Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005. 60
[4] J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymp- totics of feature learning: How one gradient step improves representation. arXiv preprint 2022. 2, 10
[5] F. Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629 (811,,62-6881, 2017. 1, 2, 3, 10
[6] Bach. On the equivalence between kernel quadrature rules and random feature expansions. The Journal of Machine Learning Research, 18(1):714-751, 2017.
25, 26
[7] Bach. Learning theory from first principles, 2021. 3, 26, 54, 55, 62
[8] Y. Bai and J. D. Lee. Beyond linearization: quadratic and higher-order approximation of wide neural networks. In International Conference Learning Repreventations, 2019. 2
[9] A. R. Barron. Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions on Information theory, 39(3)-930-94) 1993.
[10] G. Ben Arous, R. Gheissari, and A. Jagannath. Online stochastic gradient descent on non- convex losses from high-dimensional inference. Journal of Machine Learning Research (JMLR), 22:106-1, 2021. 2, 3, 6, 7, 8, 10, 18
[11] G. Ben Arous, R. Gheissari, and A. Jagannath. High- digh-dimensional limit theorems for sgd: Effective dynamics and critical scaling. arXiv preprint arXiv:2206. 04030, 2022. 3, 10
[12] G. Ben Arous, S. Mei, A. Montanari, and M. Nica. The landscape of the spiked tensor model. Communications on Pure and Applied Mathematics, 72(11):2222-2330 2019. 3,
[13] A. Berlinet and C. Thomas-Agnan. Reproducing kernel Hilbert spaces in probability and statistics. Springer Science & Business Media, 2011.
[14] D. Bertoin, Bolte, S. Gerchinovitz, and E. Pauwels. Numerical influence of relu' (0) on backpropagation. Advances in Neural Information Processing Systems, 34:468-479, 2021. 36
[15] X. Bing, F. Bunea, Strimas-Mackes and M. Wegkamp. Prediction under latent factor re- gression: Adaptive per, interpolating predictors and beyond. Journal of Machine Learning Research, 22(177):1-50, 2021.
[16] Borwein and A. Lewis. Convex Analysis. Springer, 2006. 59,
[17] N. Boumal, Voroninski, and A. Bandeira. The non-convex burer-monteiro approach works smooth semidefinite programs. Advances in Neural Information Processing Systems, 29, 2016.
[18] S.-H. Chang, P. C. Cosman, and L. B. Milstein, Chernoff-type bounds for the gaussian error function. IEEE Transactions on Communications, (9(11))2939-294, 2011. 24
[19] M. Chen, Y. Bai, Lee, T. Zhao, H. Wang, Xiong, and R. Socher. Towards understanding hierarchical learning: Benefits of neural representations. Advances in Neural Information Processing Systems, 2020. 2,
[20] S. Chen, A. Gollakota, R. Klivans, and R. Meka. Hardness noise-free learning for two- hidden-layer neural networks. arXiv preprint arXiv:2202.05555: 2022.

---

[21] V. CChhrmozhukov, D. Chetverikov. M. Demirer, E. Duño, C. Hansen, W. Newey, and J. Robins. Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21(1):11-66, 2018. 9
[22] L. Chizat and F. Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, pages 1305-1338. PMLR, 2020.
[23] L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming, 2018. 2,
[24] Y. Cho and L. Saul. Kernel methods for deep learning. Advances in neural information pro- cessing systems, 22, 2009. 3
[25] D.-A. Clevert, T. Unterthiner, and S. Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07250 2015. 61
[26] A. S. Dalalyan, A. Juditsky, and V. Spokoiny. A new algorithm for estimating the effective dimension-reduction subspace. The Journal of Machine Learning Research, 9:1647-1678, 2008. 1,
[27] A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, 2022. 2, 8, 10
[28] A. Daniely and G. Vardi. Hardness of learning neural networks with natural weights. Advances in Neural Information Processing Systems, 33:930-990 -000,, 2020.
[29] D. Davis, D. Drusvyatskiy, S. Kakade, and J. D. Lee. Stochastic subgradient method converges on tame functions. Foundations of computational mathematics, 20(1):119-154, 2020. 3, 60
[30] I. Diakonikolas, D. M. Kane, V. Kontonis, and N. Zarifis. Algorithms and sq lower bounds for pac learning onehhidden-layer relu networks. J. Abernethy and S. Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 1514-1539, PMLR, 09-12 Jul 2020.
[31] NIST Digital Library of Mathematical Functions. http://dlmf..ist.gvvv. Release 1.1.6 of 2022- 06-30. F. W. J. Olver, A. B. Olde Daalhuis, D. W. Lozier, B. I. Schneider, R. Boisvert, W. Clark, B. R. Miller, B. V. Saunders, H. S. Cohl, and M. A. McClain, eds. 20
[32] D. Drusvyatskiy, A. D. Ioffe, and A. S. Lewis. Curves of descent. SIAM Journal on Control and Optimization, 53(1):
2015. 3, 60
[33] R. Dudeja and D. Hsu. Learning single-index models in gaussian space. In S. Bubeck, V. Perchet, and P. Rigollet, editors, Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 1887-1930. PMLR, 06-09
Jul 2018. 1, 2, 7, 8, 10, 34
[34] A. Engel and C. Van den Broeck. Statistical mechanics of learning. Cambridge University Press, 2001. 3
[35] D. J. Foster, A. Sekhari, and K. Sridharan. Uniform convergence of gradients for non-convex learning and optimization. In Advances in Neural Information Processing Systems, 2018. 2, 3
[36] J. H. Friedman and J. W. Tukey. A projection pursuit algorithm for exploratory data analysis. IEEE Transactions on computers, 100(9):881-88))) 1974. 2
[37] R. Ge, F. Huang, C. Jin, and Y. Yuan. Escaping from saddle points-online stochastic gradient for tensor decomposition. In Conference on learning theory, pages 797-842. PMLR, 2015.
[38] R. Ge. J. D. Lee, and T. Ma. Matrix completion has no spurious local minimum. Advances in neural information processing systems, 29, 2016. 7
[39] R. Ge, J. Lee, and Ma. Learning me-hiddden-laye layer neural networks with landscape design. arXiv preprint 10X00:171..00001, 2017. 2.
12

---

[40] R. Ge and T. Ma. On the optimization landscape of tensor decompositions. Advances in Neural Information Processing Systems, 30, 2017.
[41] S. Goel, A. Gollakota, Z. Jin, S. Karmalkar, and A. Klivans. Superpolynomial lower bounds for learning one-layer neural networks using gradient descent. In International Conference on Machine Learning, pages 3587-3596. PMLR, 2020.
[42] S. Goel, S. Karmalkar, and A. Klivans. Time/accuracy tradcoff's for learning relu with respect to gaussian marginals. Advances in Neural Information Processing Systems, 32, 2019. 65
[43] S. Goldt, M. Advani, A. M. Saxe, F. Krzakala, and L. Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019. 2,3
[44] L. Gross. Logarithmic sobolev inequalities. American Journal of Mathematics, 97(4))1061- 1083, 1975. 61
[45] W. Härdle, M. Müller, S. Sperlich, and A. Werwatz. Nonparametric and semiparametric mod- els, volume 1. Springer, 2004.
[46] M. Hristache, A. Juditsky, and V. Spokoiny. Direct estimation of the index coefficient in single-index model. The Annals of Statistics, 29(3):595-623, 2001. I
[47] P. J. Huber. Projection pursuit. The annals Statistics, pages 435 475, 1985. 2
[48] H. Ichimura. Semiparametric least squares (sls) and weighted sls estimation of single-index models. Journal of econometrics, 58(1-2):71-120, 1993. 1,
[49] Z. Ji and M. Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 11117776-17186, 2020. 3, 60
[50] C. Jin, R. Ge, P. Netrapalli, S. M. Kakade, and M. I. Jordan. How to escape saddle points efficiently, 2017. 3
[51] G. Kornowski and O. Shamir. Oracle complexity in nonsmooth nonconvex optimization. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 324-334. Curran Associates, Inc., 2021. 10
reduutionrrdduction. [52] K.-C. Li. Sliced inverse regression for dimension reduction. Journal of the American Statistical Association, 86(414):316 027,, 1991. 2
[53] K.-C. Li. On principal hessian directions for data visualization and dimension reduction: Another application of stein's lemma. Journal of the American Statistical Association, 87(420):1025-1033 1992. 2
[54] K. Lyu and J. Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint 1906.009000009000 2019. 60
[55] A. Maillard, G. B. Arous, and G. Biroli. Landscape complexity for the empirical risk of generalized linear models. arXiv preprint arXiv:1912.02143, 2019. 3
[56] E. Malach, P. Kamath, E. Abbe, and N. Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. In International Conference on Machine Learning, pages 7379- 7389. PMLR, 2021. 1
[57] E. Malach, G. Ychudai, S. Shalev-Schwartz, and O. Shamir. The connection between ap- proximation, depth separation and learnailiity in neural networks. In Conference Learning Theory, pages 3265-3295. PMLR, 2021. 1
[58] S. S. Mannelli, G. Biroli, C. Cammarota, F. Krzakala, P. Urbani, and L. Zdeborová. Marvels and pitfalls of the langevin algorithm in noisy high-dimensional inference. Physical Review X, 001))011007, 2020. 3
13

---

[59] S. S. Mannelli, F. Krzakala, P. Urbani, and L. Zdeborova. Passed & spurious: Descent al- gorithms and local minima spiked matrix-tensor models. In international conference on machine learning, pages
PMLR, 2019. 3
[60] Mei, Y. Bai, and Montanari. The landscape of empirical risk for non-convex losses. arXiv preprint arXiv:1607.06554 2016. 2, 3, 20, 34, 63
[61] S. Mei, T. Misiakiewicz, and A. Montanari. Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration. Applied and Computational
Harmonic Analysis, 59:3-84, 2022. 3
[62] S. Minsker. On some extensions of bernstein's inequality for self-adjoint operators. Statistics & Probability Letters,
2017. 56,57 57
[63] M. Mitzenmacher and E. Upfal. Probability and computing: Randomization and probabilistic techniques in algorithms and data analysis. Cambridge university press, 2017. 37
[64] A. Mousavi-Hooseini, S. Park, M. Girotti, I. Mitliagkas, and M. A. Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint 2r0.1:2209.14863, 2022.
[65] E. Nelson. The free markoff field. Journal of Functional Analysis, 12(2):21 1973. 31
[66] E. Nichani, Y. Bai, and J. D. Lee. Identifying good directions to escape the ntk regime and efficiently learn low-degree plus sparse polynomials. arXiv preprint arXiv:2206.0388 03688, 2022.
[67] R. O' O'Donnell. Analysis of boolean boolean,uunction. Cambridge University Press, 2014. 4, 30, 31, 62
[68] G. Ongie, R. Willett, D. Soudry, and N. Srebro. A function space view of bounded norm infinite width relu nets: The multivariate case. arXiv preprint arXiv:1910.0166 2019.
[69] A. Pinkus. Approximation theory of the mip model in neural networks. Acta numerica, 8:143- 195, 1999.
[70] A. Rahimi and B. Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007. 3,
[71] V. Ros, G. B. Arous, G. Biroli, and C. Cammarota. Complex energy landscapes in spiked- tensor and simple glassy models: Ruggedness, arrangements of local minima, and phase tran- sitions. Physical Review X, 9()):01 1003, 2019, 3
[72] A. Rudi and L. Rosasco. Generalization properties of learning with random features. Advances in neural information processing systems, 30, 2017. 3
[73] I. Safran and J. D. Lee. Optimization-basee based separations for neural networks. arXiv preprint 12Xiv:2112.02393 2021.
[74] S. Saitoh. Integral transforms, reproducing kernels and their applications, volume 369. CRC Press, 1997. 21
[75] Savarese, I. Evron, D. Soudry, and N. Srebro. How do infinite width bounded norm networks look in function space? In Conference on Learning Theory, pages 2667-2690. PMLR, 2019.
5
[76] B. Schölkopf and A. Smola. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2002. 3
[77] M. Soltanolkotabi. Learning relus via gradient descent. Advances in neural information pro- cessing systems, 2017. 3
[78] M. J. Song, I. 7adik, and Bruna. On the cryptographic hardness of learning single periodic neurons. Advances in Neural Processing Systems (NeurIPS), 2021. 1
14

---

[79] J. Sun, Q. Qu, and J. Wright. A geometric analysis of phase retrieval. Foundations of Compu- tational Mathematics, 18(5):1131-118,, 2018.
[80] A. Tropp. User-friendly tail bounds for sums random matrices. Foundations of computa-
tional mathematics, 12(4):389-434 2012. 56, 57
[81] L. Van den Dries and C. Miller. Geometric categories and o-minimal structures. Duke Mathe- matical Journal, 84(2):497-540, 1996. 60
[82] R. Veiga, I.. Stephan, B. Loureiro, F. Krzakala, and L.. Zdehorová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. arXiv preprint arXiv:2202.00293, 2022. 3
[83] R. Vershynin. ence, volume 47. Cambridge university press, 2018. 19, 20, 52, 58
probability: An introduction with applications in data sci-
[84] B. Woodworth, Gunasekar, J. D. Lee, E. Moroshko, P. Savarese, I. Golan, D. Soudry, and
N. Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673 PMLR, 2020. 1,7
[85] K. Zhong, Z. Song, P. Jain, P. L. Bartlett, and I. S. Dhillon. Recovery guarantees for one- hidden-layer neural networks. In International conference on machine learning, pages 4140- 4149. PMLR, 2017. 2.
[86] M. Zhou, R. Ge, and C. Jin. A local convergence theory for mildly over-parameterized two- layer neural network. Conference on Learning Theory, pages 4577-4632. PMLR, 2021. 2, 3
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes) See Conclusions.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] This is purely theoretical work with no practical applications.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set assumptions all theoretical results? (Yes) See Sections and 6.
(b) Did you include complete proofs of all theoretical results? [Yes] See Section D.2.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] Supplemental zip file.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section A.
(c) Did you report enror bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Section A.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Section A.
4. If you using existing assets (e.g., code, data, models) curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you'l re using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifi- able information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "fqMOFkZpwFn", "reviewer": "Reviewer_uZAC", "summary": "This manuscript studies the supervised problem of learning a single-index model with a proxy model for a two-layer neural network where the first layer weights are tied and the biases are random. Different from previous works on single-index models, the link function of the target single-index model are not matched with the network architecture. Both layers are trained. The key result is to show that gradient flow on the empirical loss can achieve good approximation error of the target. More specifically, the main theoretical results are:\n\n1. To show that the critical points of the population loss are benign, i.e. that all the local minimisers are global and correspond to the target weights.\n\n2.  To show that gradient flow on the empirical loss with a particular choice of relative time-scaling for the update of the first and second layer weights finds an approximate minimiser of the population loss. Although recovery is optimal with respect to the first layer weights, it is sub-optimal in terms of the target function.\n\n3. To show that the optimal statistical sample complexity for the recovery of the target can be achieved by re-training the last layer on fresh samples.", "questions": "- **[Q1]**: Can the authors clarify why they believe the setting considered is natural when learning a single-index model (discussion in L38-L40)?\n\n- **[Q2]**: Why the variance of the bias needs to be larger than one in the analysis (L94)? What would change in the analysis for $\\tau<1$?\n\n- **[Q3]**: Why the $\\ell_2$ penalty parameter $\\lambda$ is restricted to $(0,1)$ (L47)? Is this always satisfied by the optimal $\\lambda$ leading to corollary 5.2?\n\n- **[Q4]**: What role does the choice of $\\phi = relu$ plays in the results? Would it be possible to derive similar guarantees for other activations?", "limitations": "The authors provide a nice discussion of the limitations of their work, as well as a list of possible improvements and future directions (L296-L308).", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This manuscript adds to the increasing literature of theoretical works showing the separation between neural networks and random features model and the benefits of overparametrisation for optimisation in non-convex landscapes [1, 43, VSL+22, SVL22, BES+22]. Although the setting considered is artificial in many ways (tied weights, random bias, delayed gradient flow), the fact that  gradient flow on this network can approximate \"well\" a mismatched single-index model with is an interesting result.\n\n**Strengths**: Despite the density of technical details, the paper is well written and the main thread is easy to follow. The main result is significant and timely.\n\n**Weaknesses**: The setting considered is artificial in many ways, and therefore it is hard to say it is directly relevant to our understanding of the approximation and optimisation properties of two-layer neural networks. Rather, I see it as a constrained first step with desirable extensions.\n\n**References** (numbered refs. are from the bibliography in the paper)\n\n[VSL+22], R. Veiga, L. Stephan, B. Loureiro, F. Krzakala, L. Zdeborová, \"Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks\", arXiv: 2202.00293 [stat.ML]\n\n[SVL22] I. Safran, G. Vardi, J.D. Lee, On the Effective Number of Linear Regions in Shallow Univariate ReLU Networks: Convergence Guarantees and Implicit Bias\n\n[BES+22] J. Ba, M.A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, G. Yang, \"High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation\", arXiv: 2205.01445 [stat.ML].", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "aTLHJBImPRS", "reviewer": "Reviewer_z2Sa", "summary": "This paper studies the problem of learning a single-index function with a shallow network trained via gradient descent. The function to be learned is of the form $f_{\\star}\\left(\\theta^{\\star}\\cdot x\\right)$, where both $f_{\\star}:\\mathbb{R}\\rightarrow\\mathbb{R}$ and $\\theta^{\\star}\\in\\mathbb{S}^{d-1}$ are unknown, and the data $x$ is Gaussian. The authors establish that the training landscape of this problem has no bad local minima, and then prove that $(f_{\\star}, \\theta^{\\star})$ can be learned efficiently, modulo conditions on the width of the network and a two-part structure during training, where only $\\theta$ is learned during the first part.", "questions": "I’m curious why the authors feel that the teacher-student terminology is necessary in this paper? Can $F_{\\theta^{\\star}}$ not simply be referred to as ‘ground truth’ or the ‘true’ data generating process, as is standard in learning theory? If they can, then I suggest eliminating nonstandard terminology.\n\nI think it would be helpful to the reader to explicitly state the following implications of Theorem 4.2 for the learning problem right after the theorem is presented: (1) all minima are global minima and (2) gradient descent will have to either avoid or escape the minima corresponding to $\\langle\\theta^{\\star}, \\theta\\rangle=0$.\n\nA little more discussion of $T_0$ and $T_1$ in the main text would also be useful. At present, $T_1$ appears in the main text for the first time in Algorithm 1, and I think there isn’t enough information in the main text about it to understand the algorithm. With regards to $T_0$, perhaps the logical place to first introduce it might be line 234? This will help the reader really understand the meaning of $\\zeta(t)$, as well, without first having to go through the proofs. This will also make it clear \\textit{before} the math is presented that the algorithm constructed by the authors has two different learning phases.\n\n**A few minor things:**\n\nTypo in line 35. Remove the word ‘in’ between ‘incurring’ and ‘a’.\n\nLine 116. Maybe insert the word ‘between’ between ‘relationship’ and ‘each’?\n\nThe first time the acronym RKHS (line 99?) is used, its full form should be given.\n\nA little bit of attention needs to be paid to capitalization in the references. For example, in line 311, ‘sgd’ should be ‘SGD’, and so on.", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is a very nice paper. It provides a positive answer to the (previously open) question of whether or not a shallow neural net can learn single-index functions via a constructive proof. The result contributes to our knowledge of the class of functions that can be learned efficiently with shallow networks and is therefore relevant to a broad machine learning audience. The paper is clearly written.\n\nI do not find significant weaknesses.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YL-25Y5wARH", "reviewer": "Reviewer_9Zca", "summary": "In this paper, the authors study the problem of learning a single index model using a specific class of two-layer neural networks. The network structure is nonstandard, because its bias terms are random initialized and fixed (not trained), and weights are tied across all neurons. The first main result is a population landscape theorem, which states that the first-order critical points are either in the poles ($\\pm \\theta^*$) or on the equator ($\\langle \\theta, \\theta^* \\rangle = 0$). Building on this result, the authors then study the empirical loss. They prove that the (nonstandard) gradient flow converges to a solution with desired generalization error.", "questions": "- Q1: In the statement of Theorem 5.1, the probability is \"greater than $1/4 - 2\\delta$\". Is it possible to strengthen this to a high probability statement? If not, why?\n- Q2: The authors mention that they student the critical points of $L(c, \\theta)$ through projected population loss. Based on that, I expected that during the initial warmup stage ($t \\le T_0$) one trains $c$ first, but Eqn. (11) is the contrary. \n- Q3: I do not understand why we cannot achieve the optimal generalization error without fine-tuning. Is this an artifact of the proof? If so, where is the bottleneck and how fine-tuning improves the rate? If not, what is the implication for the practice?\n\nSome typos:\n- Line 119: $\\varphi_{b_i}$ not $\\varphi_{b_j}$\n- Line 133: \"first the study\" --> \"first study\"", "limitations": "An evident limitation is that the model is rather simplified, but I believe this is on par with the current level of neural network theories. There are tweaks such as certain nonstandard $\\zeta$ scheduling, and these limitations are discussed in the paper.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, I find this paper very clear and solid. The model under study is simplified but more complicated than existing ones (i.e., random features model). While the proof strategy is mostly built upon existing ideas, the new challenge is the interaction between parameters $c$ and $\\theta$. Below is a more detailed assessment.\n\n**Pros**:\n- The landscape result is very clear. It shows that under essentially minimal overparametrization the critical points are either the true directions or orthogonal to the true directions. \n- It is proved that gradient flow (with certain warmup and fine-tuning) achieves the optimal generation error under optimal sample complexity $n \\gtrsim \\Omega(d^s)$.\n- The proof seems to be a nice combination of the landscape analysis that is popular for analyzing many nonconvex problems and gradient flow analysis for lazy training. \n\n**Cons**:\n- The neural network structure is rather simplified. While I do not complain about useful simplification (which is standard in the literature), the title and abstract are rather misleading; for instance, it is not mentioned that the weights are tied, which is crucial for this work.\n- It is not clear whether certain tweaks such as warmup and fine-tuning are necessary in practice. Perhaps the authors should discuss how these tweaks would affect the practical aspect of training neural nets.\n- Some further questions about theorem statement (see below).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "LBKbLKYC4rw", "reviewer": "Reviewer_WRiK", "summary": "The paper's focus is the learning of single-index models, in which the label $y$ of a datapoint $x\\sim\\mathcal N(0,I_d)\\in\\mathbb R^d$ is obtained as $y=f_\\star(\\langle \\theta_\\star,x\\rangle)$, i.e., by projecting the data in a direction $\\theta^\\star\\in\\mathbb S^{d-1}$ and then feeding the result in a \"link function\" $f_\\star\\colon\\mathbb R\\to\\mathbb R$. Both $\\theta^\\star$ and $f_\\star$ have to be learned. The authors consider a shallow neural network of the type $$G(x;c,\\theta)=\\frac{1}{\\sqrt N}\\sum_{i=1}^Nc_i\\phi(\\langle\\theta,x\\rangle-b_i),\\qquad \\phi(u)=\\max(0,u)$$\nwith frozen biases $b_i$ and depending on $c\\in\\mathbb R^N$ and $\\theta\\in\\mathbb R^d$. They show that, in the considered setting, the critical points of the population landscape can be fully characterised: for example, for sufficiently small regularisation and sufficiently large width of the network, the critical points $\\hat\\theta$ of the population loss are such that $\\langle\\hat\\theta,\\theta^\\star\\rangle\\in\\\\{-1,0,1\\\\}$.  They also provide a bound for the random feature approximation error in terms of the RKHS error. \n\nThe paper moves then to analyze the empirical landscape, extracting the sample complexity scaling required to escape from the uninformative equator $\\\\{\\theta\\colon \\langle \\theta,\\theta^\\star\\rangle=0\\\\}$, showing that the gradient flow on the empirical loss finds an approximate minimizer for a number of samples $n=O(d^s)$, where $s$, called information exponent, is related to the Hermite polynomial expansion of the target function $f_\\star$. The theoretical understanding of the (empirical) landscape properties leads to the proposal of a gradient descent algorithm for $c$ and $\\theta$, whose flow efficiently finds an approximate minimizer of the population loss.", "questions": "As the authors themselves pointed out, and as mentioned above, the considered model might look somehow rigid in its parametrization. I would like the authors to further comment on the robustness of their results. In particular, the Gaussian dataset leads the authors to choose Hermite polynomials as a basis of their decomposition and, e.g., the expansion of $f_\\star$ on this basis determines relevant quantities, such as the information exponent. The authors might further comment on what they expect in the case of different data set distributions, maybe expanding the discussion on lines 157-163.\n\nIf compatible with the length restrictions, I think it might be useful to add to the paper a brief section discussing the implementation of the algorithm (e.g., moving the material in the Appendix in the main text) and some details on the numerical results.\n\n*Minor typographical comments.*\n- On page 3 (and in many other points, in particular in Secion 5) $F^*$ is sometimes called $F_{\\theta^*}$ (see Definition 3).\n- On page 4, line 119 has $\\phi_{b_j}$ in place of $\\phi_{b_i}$. Also, I believe that the coefficients $\\alpha_j$ missing are missing in the sum. In line 133, \"we first the study optimization\" > \"we first study the optimization\".\n- On page 5, in Remark 4.4, $\\hat P_\\lambda$ misses a hat.\n- On page 6 line 206 subscript d missing on $I_d$.\n- On page 7 in Theorem 5.1, it is maybe more convenient to write \"Algorithm 1\" in place of \"procedure 1\" for reasons of clarity (similarly at a later stage in line 281 Procedure 1 refers to the box named Algorithm 1). Also, there is an extra parenthesis in eq. 12 and, in line 241, an exponent 'outside' the O-notation.\n- In Appendix eq 15, $\\sigma$ should probably be replaced by $\\phi$.\n\n*Additional references*\nThe authors might consider adding the reference\nV Ros, G Ben Arous, G Biroli, C Cammarota, Physical Review X 9 (1), 011003 (2019)\nin line 86, as in this paper the quenched entropy of the spiked tensor model is computed, complementing the rigorous results in [4].", "limitations": "The paper is of theoretical nature and there are no potential negative societal effects and/or limitations of a sort in this respect. The weaknesses of the contributions are discussed above.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Originality* The considered setting can be seen as a special one within the large literature on single- and multi-index models that the authors properly cite in the paper; the landscape analysis is very inspired by the one presented in [19] (where a similar landscape topology is found, as acknowledged by the authors). All results are supported by rigorous analysis built upon the available literature on random features and kernel methods. To my understanding, the originality of the paper consists in constructing a simple and completely characterizable example of a learning task performed by training both layers of a shallow network using gradient descent in a nonconvex yet \"benign\" landscape.\n\n*Quality* All results are supported by rigorous analysis and precise analytical bounds to the performances are provided. On the other hand, the proposed algorithm is not explored extensively from the numerical point of view (very few numerical experiments are given in the Appendix, and without implementation of the second fine-tuning step which is discussed in the main text).\n\n*Clarity* The paper is clearly written and proceeds in sequential order in studying the population landscape and the empirical landscape, up to the proposal of an efficient gradient descent algorithm for the solution of the task.\n\n*Significance* The main result of the paper is the full characterisation of the performances of a shallow neural network architecture (of a specific type) in learning a single-index model by means of gradient descent. The main significance comes, in my understanding, from the hybrid nature of the learning task (aiming at obtaining both $\\theta$ and, non-parametrically, the link function) via a gradient descent scheme whose performances are clarified in terms of the benign landscape. As somehow the authors themselves stress in the conclusion, a weakness of the contribution is the rigidity of the investigated setting: the architecture consists of a shallow neural network whose hidden layer consists of $N$ ReLU neurons parametrised by the same $\\theta$, differing only by their quenched bias. Also, assumptions on the dataset seem to be crucial for the analysis (see below). These assumptions might look a little bit restrictive/unrealistic, and it is not very clear to me how much the results can be considered robust: in particular, I wonder how much the simple properties of the landscape, which are pivotal for the results, are due to this combination of assumptions. Despite these observations, I personally found the paper interesting and possibly the first step for the study of more complex set-ups to investigate both the breaking down and the persistence/generalization of the phenomenology observed here.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3I3VqHmGgTM", "reviewer": "Reviewer_5fyA", "summary": "The authors consider the problem of regression in high dimensions, with Gaussian i.i.d. inputs and a target function that is an arbitrary function that only depends on the projection of the input along a special direction, θ*. They study the performance of a two-layer neural network with many neurons, which however all share the same weight θ (the biases are different for each neuron, but they are random and fixed). The second-layer weights are free and optimised.\n\nThe authors first show that at small regularisation, the only first-order critical points of the population loss of the model correspond to a magnetisation of the model m =〈 θ, θ*〉= 0 at initialisation, and m=1 at perfect recovery. (Theorem 4.2). In other words, the optimization landscape in the population limit does not contain spurious minima or other obstacles.\n\nThey then study the generalisation of a model trained on the empirical loss. Considering an algorithm where they first train the first layer, then fine-tune the second layer while keeping the direction θ fixed, they obtain a bound on the excess risk (eq. 17) that has two parts: one that depends on the dimension, and which relates to the cost of escaping the saddle-point at initialisation, and a second part which comes from the non-parametric fine-tuning.", "questions": "n/a", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This is a clean and well-written analysis of feature learning dynamics. The authors propose an original model for the two-layer network that allows for a precise analysis, and I believe this model can be a nice addition to the set of models where these questions can be studied, in addition to the teacher-student setup. The authors do a good job imho of putting their work into the context on other recent results, such as recent work on \"information exponents\" Ref. [3]. All in all, I believe this paper will be of interest and use to the theory community at NeurIPS and should therefore be published.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept