Neural Networks Efficiently Learn Low-Dimensional Representations with SGD
Authors: Alireza Mousavi-Hosseini
Sejun Park
Manuela Girotti
Ioannis Mitliagkas
Murat A Erdogdu
Venue: ICLR 2023 notable top 25% (2023)
Keywords: feature learning
generalization
compressibility
sgd
neural networks

Abstract: We study the problem of training a two-layer neural network (NN) of arbitrary width using stochastic gradient descent (SGD) where the input $\boldsymbol{x}\in \mathbb{R}^d$ is Gaussian and the target $y \in \mathbb{R}$ follows a multiple-index model, i.e., $y=g(\langle\boldsymbol{u_1},\boldsymbol{x}\rangle,...,\langle\boldsymbol{u_k},\boldsymbol{x}\rangle)$ with a noisy link function $g$. We prove that the first-layer weights in the NN converge to the $k$-dimensional principal subspace spanned by the vectors $\boldsymbol{u_1},...,\boldsymbol{u_k}$ of the true model, when online SGD with weight decay is used for training. This phenomenon has several important consequences when $k \ll d$. First, by employing uniform convergence on this smaller subspace, we establish a generalization error bound of $\mathcal{O}(\sqrt{{kd}/{T}})$ after $T$ iterations of SGD, which is independent of the width of the NN. We further demonstrate that, by recovering the principal direction, SGD-trained ReLU NNs can learn a single-index target of the form $y=f(\langle\boldsymbol{u},\boldsymbol{x}\rangle) + \epsilon$ with a sample complexity linear in $d$ (up to log factors), where $f$ is a monotonic function with at most polynomial growth, and $\epsilon$ is the noise. This is in contrast to the known $d^{\Omega(p)}$ samples required to learn any degree $p$ polynomial in the kernel regime, and shows that SGD-trained NNs can outperform the Neural Tangent Kernel at initialization. Finally, we establish compressibility guarantees for NNs using that SGD produces an approximately rank-$k$ first-layer weight matrix.

Published as a conference paper at ICLR 2023
NEURAL NETWORKS EFFICIENTLY LEARN LOW-DIMENSIONAL REPRESENTATIONS WITH SGD
Alireza MousaviiHoossenn", Park', Manuela Girotti', Ioannis Mitliagkas Murat A. Erdogdu University of Toronto, 'Vector Insitute. 'Korea University, "Saint Mary's University, Université de Montréal, MMila Institute
{mousavi, erdogdu @ toronto, edu, par..ppar0000eemal com manuela, aareea.ggrooiieee.. eerroottismu.....m smu. ca, loannisliro.. um..t..rrr..n.aall......... ca
ABSTRACT
We study the problem of training a two-layer neural network (NN) of arbitrary width using stochastic gradient descent (SGD) where the input E H is Gaussian and the target E IR follows a multiple-index model, i.c., y = g (w1, z).... (Ug x)) with noisy link function g. We prove that the first-layer weights of the NN con- verge to the k-dimensional principal subspace spanned by the vectors u1...., un of the true model, when online SGD with weight decay is used for training. This phenomenon has several important consequences when A: & d. First, by employ- ing uniform convergence on this smaller subspace, we establish generalization error bound of Ol /d/TT) after iterations of SGD, which is independent of the width of the NN. We further demonstrate that, SGD-trained ReLU NNs can learn single-index target of the form y = f( x)) + by recovering the principal direction, with a sample complexity linear in d (up to log factors), where f is a monotonic function with at most polynomial growth, and e the noise. This contrast to the known d()) sample requirement to learn any degree p polynomial in the kernel regime, and it shows that NNs trained with SGD can outperform the neural tangent kernel at initialization. Finally, we also provide compressibility guarantees for NNs using the approximate low-rank structure produced by SGD.
INTRODUCTION
The task of learning an unknown statistical (teacher) model using data is fundamental in many areas of learning theory. There has been a considerable amount of research dedicated to this task, especially when the trained (student) model is neural network (NN), providing precise and non-asymptotic guarantees in various settings (Zhong et al., 2017; Goldt et al., 2019; Ba et al., 2019; Sarao Mannelli al., 2020; Zhou et al., 2021; Akiyama & Suzuki, 2021: Abbe et al., 2022; Ba et al., 2022; Damian et al., 2022; Veiga et al., 2022). As evident from these works, explaining the remarkable learning capabilities of NNs requires arguments beyond the classical learning theory (Zhang al., 2021).
The connection between NNs and kernel methods has been particularly useful towards this expedi- tion (Jacot et al., 2018: Chizat al., 2019). In particular, two-layer NN with randomly initialized and untrained weights is an example of a random features model (Rahimi & Recht, 2007), and regres- sion on the second layer captures several interesting phenomena that NNs exhibit in practice (Louart et al., 2018: 8: Mei & Montanari, 2022), e.g. cusp in the learning curve. However, NNs also inherit favorable characteristics from the optimization procedure (Ghorbani et al.. 2019: Allen-Zhi & Li, 2019; Yehudai & Shamir, 2019; Li et al., 2020; Refinetti et al., 2021), which cannot be captured by associating NNs with regression on random features. Indeed, recent works have established a separation between NNs and kernel methods, relying the emergence of representation learning as a consequence of gradient- training (Abbe et al., 2022; Ba et al., 2022: Barak et al., 2022; Damian et al., 2022), which often exhibits natural bias towards low-complexity models.
theme that has emerged repeatedly in modern learning theory is the implicit regularization effect provided by the training dynamics (Neyshabur al., 2014). The work by Soudry al. (2018) has inspired an abundance of recent works focusing on the implicit bias of gradient descent favoring, in some sense, low-complexity models, e.g. by achieving min-norm and/or max-margin solutions

---

Published as a conference paper at ICLR 2023
despite the lack of any explicit regularization (Gunasekar et al., 2018: Li et al., 2018; Ji & Telgarsky, 2019; Gidel et al., 2019; Chizat & Bach, 2020; Pesme et al., 2021). However, these works mainly consider linear models unrealistically wide NNs, and the notion of reduced complexity its implications on generalization varies. A concrete example in this domain is compressiblity and its connection to generalization (Arora al., 2018; Suzuki et al., 2020). Indeed, when a trained NN can be compressed into a smaller NN with similar prediction behavior, the resulting models exhibit similar generalization performance. Thus, the model complexity of the original NN can be explained by the smaller complexity of the compressed one, which is classically linked to better generalization.
well as
In this paper, we demonstrate the emergence of low-complexity structures during the training procedure. More specifically. we consider training a two-layer student NN with arbitrary width m where the input æ: E Rdd is Gaussian and the target y E R follows a multiple-index teacher model, i.e. y g( (M1, ..,.... (us,, 2;); t) with link function g and a noise independent of the input. In this setting, we prove that the first-layer weights trained by online stochastic gradient descent (SGD) with weight decay converge to the k-dimensional subspace spanned by the weights of the teacher model, span(u1 U1. "I...... Uz), which we refer to as the principal subspace. Our primary focus the case where the target values depend only on few important directions along the input, i.e. k < d. which induces a low-dimensional structure on the SGD-trained first-layer weights, whose impact on generalization is profound. First, convergence to the principal subspace leads to an improved bound on the generalization gap for SGD, independent of the width of the NN. In the specific case of learning a single-index target with a ReLU student network, we show that this convergence leads to useful features that improve upon the initial random features. Hence we prove that NNs can learn certain degree-p polyno- mials with number of samples (almost) linear in d using online SGD, while learning a degree polyno- mial with any rotationally invariant kernel, including the neural tangent kernel (NTK) at initialization, re- quires (p( samples (Donhauser et al., 2021). We summarize contributions follows.
-2
spontut
Inttalized
Tained
-3 -2
Figure 1: Two-layer ReLU network with m
1000, d = 2 is trained to recover a lanh single- index model via SGD with weight decay. Initial neurons (red) converge to the principal subspace. 10% of student neurons are visualized.
We show in Theorem 3 that NNs learn low-dimensional representations by proving that the iterates of online SGD on the first layer of a two-layer NN with width converge to /mm neighborhood of the principal subspace after Odd/=)) iterations, with high probability. The error tolerance of /ms is sufficient to guarantee that the risk of SGD iterates and that of its orthogonal projection to the principal subspace are within ((s) distance.
We demonstrate the impact of learning low-dimensional representations with three applications.
For a single-index target y = f(uu, x)) + c with a monotonic link function f where f" has at most polynomial growth, we prove in Theorem 4 that ReLU networks of width m can learn this target after T iterations of SGD with an excess risk estimate of (V//T + 1/m), with high probability (see the illustration in Figure 1). In particular, the number of iterations is linear in the input dimension even when f is polynomial of any (fixed) degree p.
Based on a uniform convergence argument on the principal subspace, we prove in Theorem 5 that T iterations of SGD will produce model with generalization error of O (k/kTT)) with high probability. Remarkably, this rate is independent of the width m of the NN, even in the case k x d where the target is any function of the input, and not necessarily low-dimensional. Finally, we provide a compressiblity result directly following from the wow-dimensionality of the principal subspace. We prove that T' iterations of SGD produce first-layer weights that are compressible rank-k with risk deviation of O(d/TT), with high probability.
The rest of the paper is organized as follows. We discuss the notation and the related work in the remainder of this section. We describe the problem formulation and preliminaries in Section 2, and provide an analysis for the warm-up case of population gradient descent in Section 2.1. Our main result SGD is presented in Section 3. We discuss three implications of our main theorem in Section 4, where we provide results on lcarnability, generalization gap. and compressibility in Sections 4.1, 4.2, and 4.3, respectively. We finally conclude with brief discussion in Section 5.
2

---

Published as a conference paper at ICLR 2023
Notation. For a loss function 2:² - R. let 0,e and 0. denote its partial derivatives with respect to ith and jth inputs fori,J i,j E {1, For quantities and b, a S b implies S. Cb for some absolute constant C > 0, and a x b implies both a 2 b and a S b. Finally, Unif( A) denotes the uniform distribution over set A and (0, I2) denotes the d-dimensional isotropic Gaussian distribution.
1.1 RELATED WORK
Training dynamies of NNs. Several works have demonstrated learnability in a special case of teacher- student setting where the teacher model is similar to the student NN being trained (Zhong et al., 2017; Brutzkus & Globerson, 2017: Li Yuan, 2017: Zhang et al., 2019; Zhou et al., 2021). This setting has also been studied through the lens of loss landscape (Safran et al., 2021) and optimization over measures (Akiyama & Suzuki, 2021). We stress that our results work under misspecification and hold for generic teacher models that are not necessarily NNs with similar architecture to the student.
Two scaling regimes of analysis have seen a surge of recent interest. In the regime of lazy train- ing (Chizat et al., 2019), the parameters hardly move from initialization and the NN does not learn useful features, behaving like a kernel method (Jacot et al., 2018; Du et al., 2019; Allen-Zhu et al., 2019; Arora et al., 2019; Oymak & Soltanolkotabi, 2020). However, many works have shown that deep learning is more powerful than kernel models (Yehudai & Shamir, 2019; Ghorbani et al., 2020; Geiger et al., 2020), establishing a clear separation between them; thus, several important characteristics of NNs cannot be captured with lazy training (Ghorbani et al., 2019), even though might still perform better than feature learning in certain low-dimensional settings (Petrini et al., 2022). In the other scaling regime, gradient descent on infinitely wide NNs reduces to Wasserstein gradient flow, known as the mean-field regime where feature learning is possible (Chizat & Bach, 2018; Rotskoff & Vanden-Eijnden 2018; Mei et al., 2019; Nitanda et al., 2022; Chizat, 2022). Closer to our results, the concurrent work of Hajjar& & Chizat (2022) shows that low-dimensional targets induce low-dimensional dynamics on mean-field NNs. However, these results mostly hold for infinite or very wide NNs, and quantitative guarantees are difficult to obtain. Our setting is different from both of these regimes, we allow for NNs of urbitrary width without excessive overparameterizzaton.
Feature learning with multiple-index teacher models. The task of learning a target of an unknown low-dimensional function of the input is fundamental in statistics (Li & Duan, 1989). Several recent works in the learning theory literature have also focused on this problem, with an aim to demonstrate NNs can learn useful feature representations, outperforming kernel methods (Bauer & Kohler, 2019; Ghorbani et al., 2020). In particular, Abbe et al. (2022) studies the necessary and sufficient conditions for learning with sample complexity linear in d with inputs on the hypercube, in the mean-field limit. Closer to our setting are the recent works Ba et al. (2022); Damian et al. (2022); Barak et al. (2022) which demonstrate a clear separation between NNs and kernel methods, leveraging the effect of representation learning. However, their analysis considers a single (full) gradient step on the first-layer weights followed by training the second-layer parameters. In contrast, in our learnability result, we consider training both layers with SGD, which induces essentially different learning dynamics.
Generalization bounds for SGD. popular algorithm-dependen approach for studying general- ization is through algorithmic stability (Bousquet & Elisseeff, 2002; Feldman & Vondrak, 2018; Bousquet et al., 2020), which has been used to study the generalization behavior of gradient-based methods in various settings (Hardt et al., 2016; Bassily et al., 2020; Farghly & Rebeschini, 2021; Kozachkov et al., 2022). Other approaches include studying the low-dimensional structure of the trajectory (Simsekli et al., 2020; Park al., 2022) or the invariant measure of continuous-tin time approximations of SGD (Camuto et al., 2021), and employing information-theoree 2021). Among these works, Barsbey et al. (2021) show that SGD is able to learn compressible networks. However, they require large width for the mean-field approximation and assume that the SGD iterates converge to heavy-tailed distribution, while we do not make either of the assumptions.
tools (Neu et al.,
2 PRELIMINARIES: NEURAL NETWORKS AND THE PRINCIPAL SUBSPACE
For an input æ E Rd, we consider training a two-layer neural network (NN) with m neurons
((:; W, a)
Σ a, ((us, x) + bb),
(2.1)
3

---

Published as a conference paper at ICLR 2023
where a is the activation function, (wisissm are the first-layer weights collected in the rows of the matrix W Rmxd be E R is the bias, and R" is the second-layer weights. We assume a ~ (0, Ia) and the target is generated from a multiple-index (teacher) model given by
y = (u1, .)....,
...... (uA: c);
(2.2)
for weakly differentiable link function g R++1 + R and a noise e. Throughout the paper, the noise e is assumed to be independent from the input and our framework covers the special noiseless case where
While our results remain valid regardless of how k and d compare, they are most insightful when k < d; thus, we specifically consider this regime when interpreting the results. We also collect the teacher weights {u. Igigk in the rows of the matrix U E RRxx and use y =UUx; e) for simplicity.
For a given loss function E(), y). we consider the population and the empirical risks
T-1 E (j((=)) W. a. b).
R(W, a. b) (z W, a. v)l and R(W, a. b)
where the expectation is over the data distribution. Similarly, for some + 2 1, the truncated loss is defined as é (9. y) == {ū. y) A T with the corresponding risks R, and R, both of which are used in Section 4 obtain sharp high probability statements. In the warm-up case, we consider the L--regularized population risk with a penalty parameter À 2 0, defined as
X I|WI
(w b) = R(W, b)
(2.3)
To minimize (2.3), we use stochastic gradient descent (SGD) over the first-layer weights, where we interested in the convergence of iterates to the principal subspace defined by the teacher weights
S(U)= == span(u1...... U2) ) {CU : C E m=
Notice that the principal subspace satisfies S(U) S R. and its dimension is mk as opposed to the ambient dimension of md, with any matrix in this subspace having rank at most k. For any vector E Rd we let ln denote the orthogonal projection of onto span(u).. (......... Ugg and v e BI|- Similarly, for a matrix W E mxd we define WI || and W1 T by applying the projection to each row.
We make the following assumption on the data generating process.
Assumption 1 (Student-teacher setup). The student model is a two-layer NN (2.1) trained over the data set (x) )o21 where the target values y() are generated according to the teacher model (2.2) and the inputs satisfy (i) iid ~ (0, The link function g(.,..,, is weakly differentiable (see e.g. Evans (2010, Sec. 5.2.1) for definition) for any fixed
The Gaussian input is a rather standard assumption in the literature, especially in recent works that consider the student-teacher setup; see e.g. Safran et al. (2021); Zhou al. (2021); Damian al. (2022). The multiple-index teacher model (2.2) can encode a broad class of input-output relations through the non-line non-linear link function, including a multi-layer fulne-connected NN with arbitrary depth and width and weakly differentiable activations. The smoothness properties of the activation a play an important role in our analysis. As such, we consider two scenarios, with different requirements on the loss function.
Assumption 2.A (Smooth activation). The activation functiono a satisfies [a(2),, (o)( (o"))),o"())) (z) < 1 for all z E IR, the loss is (ÿ. y) = [/ y)² for simplicity, and y satisfies [y/ K almost surely. Assumption 2.B (ReLU activation). The activation function a is a(2) = max( 0) for E R. The loss satisfies 0 < (7 3) < (8, 3)1 S 1, and (8 g)1 < L.
Commonly used activations such as sigmoid and tanh satisfy Assumption 2.A. For ReI.U activation in Assumption 2.B we choose d2(2) = ((2 2 0) as its weak derivative. We highlight that Assumption 2.B is satisfied by common Lipschitz and convex loss functions such as the Huber loss
6H - y) = (ύ - 3)
if ーー 1
(2.4)
1 yl -
if
as well the logistic loss (ū y) == log(1 + e- BM), up appropriate scaling constants.
4

---

Published as a conference paper at ICLR 2023
2.1 WPRM-UP: POPULATION GRADIENT DESCENT
In this section, we study the dynamics of population gradient descent (PGD) to motivate our investi- gation of the more practically relevant case of SGD. When initialized from W PGD with a fixed step size y will update the current iterate W' according to the update rule
W W "VVRRx("),
(2.5)
We use the following initialization throughout the paper.
Assumption 3 (Initialization). For all 1 S i S m, 1 S j S d, we initialize the NN weights and biases with Vdw? ~ (O, 1), ma? ~ Unif([-1, 1D), and by ~ Uni
iid
iid
While this initialization is standard in the mean-field regime, only use it to simplify the exposition. Indeed, we can initialize W and a with any scheme that guarantees | WI S ym and ||a||so S m -1 with high probability. Further, initialization of b mostly matters in the analysis of ReLU activation.
Next, we show that the population gradient admits a certain decomposition which plays a central role in our analysis. For smooth activations, the below result is remarkable consequence of Stein's lemma, which provides a certain alignment between the true statistical model (teacher) and the model being trained (student), which has profound impact the learning dynamics. We generalize this result for ReLU through a sequence of smooth approximations (see Appendix A.1 for details).
Lemma 1. Under Assumptions 1&2.A 1&2.B, the gradient of the population risk can be written D(W),,
W + VWRx = (HL(W) + > for some HW) E Rmxx and D(W) E kxd (with explicit forms provided in Appendix A.
(2.6)
Notice that the subset of critical points W. VR.(W.) 03 for which H(W.)+ AI is invertible belongs to the principle subspace, i.e. W. E &SUU). Further, if initialize PGD (2.5) within the principal subspace, i.e. W E &(U), the subsequent iterates for t > 0 remain in this subspace.
In statistics literature, the setting we consider is often termed as model misspecification, i.e. the teacher model generating the data and the student model being trained are different. Proposition states a general result that, as long as the target depends on certain directions. PGD will produce weights their span, despite the possible mismatch between the two models. We highlight that the classical results on dimension reduction, e.g. Li & Duan (1989); Li (1991), rely on a similar principle, which was also used for designing optimization algorithms, see e.g. Erdogdu al. (2019). However, we are interested in the implications of this phenomenon in the setting of NNs trained with SGD.
The following result, proved in Appendix A.2, demonstrates the algorithmic implications of Lemma 1 for the simplistic case of PGD, and shows that the iterates converges the principal subspace. Proposition Consider running T PGD iterations (2.5) with an initialization satisfying Assump- tion 3 and a step size 1j > O. For any > O, choose X and 1 = mij based on the following.
m
1. Smooth activation. Under Assumptions 1&2.A, let X 1+ + V + 27 + 2R(W9) and
i S (x+ de +s) where == E[||U Vg(Ux; c) [|2].
2. ReLU activation. Under Assumptions 1&2.B, let Ā 2 7 + 2/ 2 and vj i X-1
Then, with probability least 1 e -Cmd over the initialization, where C is absolute constant, the iterates of PGD satisfy
S1 (T)
(2.7)
A few remarks are in order. First and most importantly, PGD iterates converge to the principal subspace as T + oc and this phenomenon is mainly due to the alignment provided by Lemma 1. Indeed in the limit, (2.7) provides sparsity in the basis of principal subspace, and it is widely known that Lz-regularizzatin does not have similar sparsity effect (unless À - oc) in contrast to its L₁ counterpart. Thus, the choice of À in Proposition 2 will lead to non-trivial orthogonal projections W in general, as we will demonstrate in Section 4.1 with a learnability result. However, without T.2- regularization, i.e. x 0, it is possible to converge to a critical point W for which H H(W") + Alm. is not invertible; hence, the weights are likely to be outside of the principal subspace in this case (cf. Figures 1&2). It is also worth emphasizing that the penalty level used in the above proposition still allows for non-convexity as we demonstrate with an example in Appendix D. We finally remark that, as evident from the proof, Proposition 2 remains valid even with unbounded smooth activations.
5

---

for Published as a conference paper at ICLR 2023
3 CONVERGENCE OF STOCHASTIC GRADIENT DESCENT
We now consider stochastic gradient descent (SGD) in the online setting where at each iteration I, we have access to new data point (xx*), (ft) drawn independently of the previous samples from the same distribution. We update the first-layer weights W with (possibly) time varying step size f/t and weight decay, according to the update rule
(3.1) The above algorithm can be used to minimize the population risk (2.3) in practice (Polyak & Juditsky, 1992), even in certain non-convex landscapes (Yu et al., 2021). As we demonstrate next, SGD still preserves several important characteristics of its population counterpart, PGD.
W = (1 - > - 7/ ( W a, (,,),,²²))
Theorem Consider running T SGD iterations (3.1) over samples satisfying Assumption 1, with an initialization satisfying Assumption 3, and using the following step size schedules.
1. Constant step size. Under Assumption 2.A. choose the constant step size π) 2m log(T), For any TT
7 let À + + 27 + +2R(W*) +
Clog((T)dd where C is an absolute constant, and T2T suppose 2 (1/6)c/d 1 log 5. Then, for a penalty m
A with probability at least 1
|
(d log(1/6))
(3.2)
Vm
72T
2. Decaying step size. Let ( := Elyl 1 under Assumption 2.A and Ç 2//2/ET under 2(+++)++1 À 2 y + S, and t" = 3
Assumption 2.B. Choose the decreasing step size The- = 7[++++1]2
any y > 0. Then, for X X m. with probability at least 1- 5,
+ log(1/0)
(3.3)
VI
22
whenever m E 200((1/)) and T 24
X
Remark. Our results are most insightful when 7 = 1 with respective rates of D(Va/T) and O(Vd/T) in the constant and decreasing step size settings; this scaling allows efficient learning of certain targets (see Section 4.).. Indeed, choosing large y may significantly restrict the learn- ability properties and will result in underfitting. However, if we ignore the underfitting issue, one can get the fastest convergence rate by choosing x V + log(1/5)) from which we obtain [WIIlle/i & 1/T. We also note that the convergence rate is stated in the normalized distance to the principal subspace, i.c. ||WT Il/vm S E, this is sufficient to guarantee that the risk of W and its orthogonal projection WT are within O(e) distance.
The above result states that, with a number of samples linear in d, SGD is able to produce iterates that are in close proximity to the principal subspace; thus, it efficiently learns (approximately) low-dimensional weights, exhibiting an implicit bias towards low- low-complyxity models. While prior works have established that NNs adapt to low-dimensional manifold structures in the data in some contexts (Chen et al., 2019; Buchanan et al., 2021; Wang et al., 2021), our result has a different nature. More specifically, the interplay between two forces is in effect here. The most important one the linear relationship between the first-layer weights and the input in both student and teacher models together with the input distribution. The alignment described in Lemma I yields sparsified weights in basis defined by the teacher network, effectively reducing the dimension from d to k. The second force is the explicit Le-reuularization. We emphasize that Le-reullarization does not play the main role in this sparsification; even though it may provide shrinkage to zero, La penalty will in general produce non-sparse solutions. However, it is still required to ensure that SGD avoids critical points outside of the principal subspace.
Although Theorem 3 does not have any implications on the convergence behavior of W I|. in the next section we show that the implied low-dimensional structure is sufficient to provide guarantees on the generalization error, learnability, and compressibility of SGD. The proof of this Theorem is provided in Appendix B, and is based on a recursion the moment generating function of [[', l|.. We note that, as in Proposition 2, regularization in Theorem does not imply (strong) convexity in general, which we demonstrate in non-convex example in Appendix D.

---

accommodate any Published as a conference paper at ICLR 2023
4 IMPLICATIONS OF LOW-DMENNIONALIITY
4.1 LEARNING SINGLE-INDEX TARGETS
An essential characteristic of NNs is their ability to learn useful representations, which allows them to adapt to the underlying misspecified statistical model. Although this fundamental property has been the guiding principle in all empirical studies, it was mathematically proven only recently for gradient-based training (Abbe et al., 2022: Ba et al., 2022; Barak et al., 2022; Damian et al., 2022; Frei al., 2022); see also survey of prior works in Malach et al. (2021). Our results in the previous section are in the same spirit, establishing the convergence of SGD to the principal subspace which indeed a span of useful directions associated with the target function being learned, As such, we leverage the learned fow-dimensional representations to demonstrate that SGD is capable of learning target function of the form y = f(uu, æ)) + with a number of samples linear in d (up to logarithmic work with the Huber loss below; however,
factors). For simplicity, Lipschitz and convex loss at the expense of a more detailed analysis.
Algorithm 1 Training two-layer ReLU network with SGD.
Input: a 6¹ E R=, D
Rmxd {(x(1), (1)) JOSeST-1. (n.) (16120 A, X.
1: for = 0, T 1 do
2: W (1 n, W 1/ V V W aa.abb)),. (0),())))))
3: end for
4: Let by ~ UnifU - A, A) for 1 S m. m.
5: for .... T' do
6: Sample i, ~ Unif(0 .... T - 1}.
7: a++1= = -1-1/X)aa
V. ai(((()))) W ,a',b), (())
8: end for
9: return (W"", a" ,))
In the sequel, we use Algorithm and train the first layer of the NN with online SGD using T data samples. Then, we randomly choose the biases and run 7" SGD iterations on the second layer using the same data samples used to train the first layer. Thus, the overall sample complexity is T whereas the total number of SGD iterations performed (2022); Barak et al. (2022); Damian et al. (2022) perform only one gradient step on the first layer weights, whereas Algorithm 1, we train the entire NN with SGD.
+ T'. We highlight that the recent works Ba et al.
Theorem 4. Suppose that the data is from a single-inder model f((uu, x)) + with monotone differentiable ff and v-sub-Gaun Gaussian noise c, and Assumptions 1&2.B hold. Further, let |uul| = [0)) < 1, and consider the Huber loss (2.4) for simplicity. Consider running Algorithm 1 with
the initialization O < a) a S 1/m, 0 < b9 bb S 1, and wn. = w ~ ~N(0,, 4La).or for all j with for any x 1,
the hyper-parameters X X
2a 2
m "("00+ +141" with 1" x 7
2(t" +1)+1
en
21+1 gl. x(t+iş= and A x v [T)(TT Then, for T Z (d log( V ( log( m )) some X' > 0 (see (C.8)), and sufficiently large T (see (C.9)), with probability at least 5,
Tog(T/5)
d + log(1/5) + +V
(1g(1
R- a"" b)- E[((() A2
+
(4.1)
T
T
where A., defined in (C.7), is poly( (T) when f" has at most polynomial growth.
This result implies that a ReLU NN trained with SGD can learn any monotone polynomial with a sample complexity linear in the input dimension d, up to logarithmic factors. Indeed, this is consistent with the work of Ben Arous et al. (2021); they establish sharp sample complexity of (IIVVII))) to learn a target with online SGD using the same activation f in the student network, where I the information exponent (I 1 in the above case due to the monotonocity of f). Despite assuming the link function f is known, we highlight that their setting covers I 2 1, whereas Theorem 4 is a proof of concept to demonstrate the learnabllity implications of convergence to the principal subspace, even when if is unknown. Building on their work, the concurrent work of Bietti et al. (2022) also proves learnability for unkown single-index targets with I > 1, albeit with sample complexity of d2 for
7

---

Published as a conference paper at ICLR 2023
I= when training ReLU students. Nonparametric regression with NNs has also been considered within the NTK framework (Hu et al., 2021; Kuzborskij & Szepesvári, 2021), but our result holds beyond this regime
grows with poly (T)), in contrast to the poly(T) requirement of the NTK regime. Additionally, learning any degree p polynomial using rotationally invariant kernels requires df(p) samples for a variety of input distributions including isotropic Gaussian (Donhauser al., 2021); thus, our result shows that SGD is able to efficiently learn target function where kernel methods cannot. For polynomial targets, Damian et al. (2022) consider training the first-layer weights with one gradient descent step with carefully chosen weight decay, and obtain a sample complexity of d² to learn any degree polynomial depending on a few directions. Finally, Chen & Meka (2020) propose method that can train NNs to learn such polynomials with sample complexity linear in d; yet, their algorithm is not a simple variant of SGD and requires a non-trivial warm-start initialization.
The proof of Theorem 4, detailed in Appendix C.2, relies on the fact that after training the first layer, the weights will align with the true direction u. Then, similarly to Damian et al. (2022) we construct an optimal with |a.| x m for every i with small empirical risk, and employ the generalization bound of Theorem 5 to achieve a rate estimate on the population risk.
4.2 GENERALIZATION GAP
For given learning algorithm, the gap between its empirical and population risks is termed as the generalization gap (not be confused with excess risk), and establishing convergence estimates for this quantity is a fundamental problem in learning theory. Classical results rely on uniform convergence over the feasible domain containing the weights; thus, they apply to any learning algorithm including SGD (Neyshabur et al., 2019). However, these bounds often diverge with the width of the NN, yielding vacuous estimates in the overparameterized regime (Zhang et al., 2021). To alleviate this, recent works considered establishing estimates for specific learning algorithm; see e.g. Hardt et al. (2016); Soudry et al. (2018); Yun et al. (2021); Park et al. (2022).
Here, we are interested in deriving an estimate for the generalization gap over the SGD-trained first- layer weights, which holds uniformly over the second layer weights and biases. More specifically, we study, after T' iterations of SGD (3.1) initialized with (W", aa. 6"). the following quantity
<(W")==uppa (W T. a. b) R)- (WT,a, b) with S = aa b E RT: ||al| S - S rb):
where the scaling ensures ý = (1) when [/w x 1, which is the setting considered in Theorem 4. We state the following bound on E(W"); covering argument over the smaller dimensional principal subspace implied by Theorem 3.
); the proof is provided in Appendix C.I, and it is based on a
Theorem 5. Consider the setting of Theorem 3 with either decreasing constant step size. For any
8> $>0 0, if 2 (d + log(1/5)) V u log(m/6)). then with probability at least
E(W") S STryy {
(x(d + log(1/6) + -rr+ + x
dk T
(4.2)
72T
where we let k for decreasing step size and = log(T") ( for constant step size.
The above bound is independent of the width m of the NN, and only grows with the dimension of the input space d and that of the principal subspace k; thus, producing non-vacuous estimates in the overparametrized regime where m is large. Further, the bound is stable in the number of SGD iterations T, that is, it converges to that rely on algorithmic stability are optimal for strongly convex objectives (Hardt et al., 2016); yet, they lead to unstable diverging bounds non-convex settings as T + co. As such, these techniques often require early stopping, which is clearly not needed in our result.
T - oc. We remark that generalization bounds for SGD
4.3 COMPRESSIBILITY
NNs exhibit compressiblity features in empirical studies, which is known to be associated with better generalization. Under the assumption that the trained network is compressible, several works established generalization bounds, see e.g. Arora et al. (2018); Suzuki et al. (2020). However, a theoretical justification of this assumption, specifically for NN trained with SGD, missing.
8

---

Published as a conference paper at ICLR 2023
Indeed, Theorem provides concrete answer this question; since the SGD iterate T converges to a low-rank matrix, the resulting weights are compressible. More precisely, my Rmxk + Rm>* be the low-rank approximation operator defined by TT& (W) := arg (W'rank( As Wi lies in the principal subspace, it has rank at most k. Thus, we can write
W W
I T Tx (WT)l: - || w Il = [|WT.e
and the following is an immediate consequence of the bound in Theorem 3, that is [ ll///mm v Vd/T, combined with the Lipschitzness of R, (W), which we prove in Lemma 17.
Proposition 6. Consider the setting of Theorem with either decreasing or constant step size. Then, with probability at least
R& (Tx( a,, Rr (W" a. b)
2 TK 7 T
d log(1/6)
(4.3)
where we let for decreasing step size and VTo for constant step size.
This result demonstrates that the low-dimensionality exhibited by the trained NN provides rate of O / / for the gap between its population risk and that of its compressed version. The bound is independent of both the width m and the dimension of the principal subspace k. Finally, we highlight that Suzuki et al. (2020) provide generalization bounds by assuming a near low-rank structure for the weight matrix, namely that jth singular value decays proportional j- for some a > 1/2. However, this condition imposes a structure quite different than what we proved in Theorem 3.
5 CONCLUSION
We studied the dynamics of SGD with weight decay on two-layer NNs, and proved that under multiple-index teacher model, the first-layer weights converge to the prin- cipal subspace, i.e. the span of the weights of the teacher. This phenomenon is of particular interest when the target depends the input along few important directions. In this setting, we proved novel generalization bounds for SGD via uniform convergence on the low-dimensional principal subspace. Further, we proved that two-layer ReLU networks can learn single-index target with a monotone link that has at most polynomial growth, using online SGD, with number of samples almost linear in d. Thus, as implication of low-dimensionality, we estab- lished separation between kernel methods and trained NNs where the former suffers from the curse of dimensionality.
2
spantul
Inttaiized
Tained
Figure 2: Neurons fail to converge the principal subspace without weight decay, in the same experimental setup of Figure I.
Two principal forces are responsible for the emergence of the fow-dimensional structure. The main one is the linear interaction between the Gaussian input and the first-layer weights in both student and teacher models. The secondary one is the weight decay which allows SGD to avoid critical points outside of principal subspace. Figure 2 shows the convergence behavior in absence of weight decay. Understanding more precisely the range of that implies convergence to the principal subspace, as well as investigating the possibility of learning multiple-index models using this convergence, left as important directions for future studies.
ACKNOWLEDGMENTS
The authors would like to thank Denny Wu for generating the figures, and both DW and Matthew Zhang for valuable feedback on the manuscript. This project was mainly funded by the CIFAR AI Catalyst grant. The authors also acknowledge the following funding sources: SP was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University). MG supported by NSERC Grant [2022-04106), IM was supported by Samsung grant and CIFAR AI Chairs program. Finally, MAE was supported by NSERC Grant [2019-06167] and CIFAR AI Chairs program.

---

Published as a conference paper at ICLR 2023
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, 2022.
Shunta Akiyama and Taiji Suzuki. On Learnabllityy via Gradient Method for Two-Layer ReLU Neural Networks in Teacher-Student Setting. In International Conference Machine Learning, 2021.
Zeyuan Allen-Zhu and Yuanzhi Li. What Can ResNet Learn Efficiently, Going Beyond Kernels? In Advances in Neural Information Processing Systems, 2019.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A Convergence Theory for Deep Learning via pher-Parameterization.. In International Conference on Machine Learning. 2019.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger Generalization Bounds for Deep Nets via a Compression Approach. In International Conference on Machine Learning, 2018.
Sanjeev Arora, Simon Du Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks. In International Conference Machine Learning, 2019.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Generalization of two-layer neural networks: An asymptotic viewpoint. In International Conference Machine Learning, 2019.
Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representa- tion. arXiv preprint wrXiv:2205.01444 2022.
Boaz Barak, Benjamin I. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit. arXiv preprint arXiv:2207.08799 2022.
Melih Barsbey, Milad Sefidgaran, Murat A Erdogdu, Gael Richard, and Umut Simsekli. Heavy tails in SGD and compressibility of overparametrized neural networks. In Advances in Neural Information Processing Systems, 2021.
Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of stochastic gradient descent on nonsmooth convex losses. In Advances in Neural Information Processing Systems, 2020.
Benedikt Bauer and Michael Kohler. On deep learning as remedy for the curse of dimensionality in nonparametric regression. The Annals of Statistics, 2019.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent non-convex losses from high-dimensional inference. J. Mach. Learn. Res., 22:106-1, 2021.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems, 2022.
Olivier Bousquet and André Elisseeff. Stability and generalization. Journal of Machine Learning Research, 2002.
Olivier Bousquet, Yegor Klochkov, and Nikita Zhivotovskiy. Sharper bounds for uniformly stable algorithms. In Conference on Learning Theory, 2020.
Alon Brutzkus and Amir Globerson. Globally Optimal Gradient Descent for a ConvNet with Gaussian Inputs. In International Conference on Machine Learning, 2017.
Sam Buchanan, Dar Gilboa, and John Wright. Deep networks and the multiple manifold problem. In International Conference Learning Representations, 2021.
10

---

Published as a conference paper at ICLR 2023
Alexander Camuto, George Deligiannidis, Murat A Erdogdu, Mert Gurbuzbalaban, Umut Simsekli, and Lingjiong Zhu. Fractal structure and generalization properties of stochastic optimization algorithms. In Advances in Neural Information Processing Systems, 2021.
Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Nonparametrie regression on low- dimensional manifolds using deep relu networks: Function approximation and statistical recovery. arXiv preprint 1908v:908.000 00844 2019.
Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory, 2020.
Lénaic Chizat. Mean-field langevin dynamics: Exponential convergence and annealing, arXiv preprint arXiv:2202.01009, 2022.
Lenaic Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. In Advances in Neural Information Processing Systems, 2018.
Lénaic Chizat and Francis Bach. Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss. In Conference on Learning Theory, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming. In Advances Neural Information Processing Systems, 2019.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural Networks can Learn Representa- tions with Gradient Descent. In Conference on Learning Theory, 2022.
Konstantin Donhauser, Mingqi Wu, and Fanny Yang. How rotational invariance of common kernels prevents generalization in high dimensions. In International Conference on Machine Learning, 2021.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient Descent Provably Optimizes Over-parameterizze Neural Networks. In International Conference on Learning Representations, 2019.
Murat A. Erdogdu, Mohsen Bayati, and Lee H. Dicker. Scalable Approximations to Generalized Linear Problems. Journal of Machine Learning Research, 2019.
Lawrence C. Evans. Partial Differential Equations. Graduate Studies in Mathematics. American Mathematical Society, 2010.
Tyler Farghly and Patrick Rebeschini. Time-independent generalization bounds for SGLD in non- convex settings. In Advances in Neural Information Processing Systems, 2021.
Vitaly Feldman and Jan Vondrak. Generalization bounds for uniformly stable algorithms. Advances in Neural Information Processing Systems, 31, 2018.
Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202.07626. 2022.
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy learning in deep neural networks: an empirical study. arXiv preprint arXiv: (xXxv:/906.00 190669903333 2020.
B. Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of Lazy Training of Two-layers Neural Networks. In Advances in Neural Information Processing Systems, 2019.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When Do Neural Networks Outperform Kernel Methods? In Advances in Neural Information Processing Systems, 2020.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
11

---

Published as a conference paper at ICLR 2023
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamies of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In Advances in Neural Information Processing Systems, 2019.
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. In International Conference on Machine Learning, 2019.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit Bias of Gradient Descent on Linear Convolutional Networks. In Advances in Neural Information Processing Systems, 2018.
Karl Hajjar and Lenaic Chizat, On the symmetries in the dynamics of wide two-layer neural networks. arXiv preprint arXiv:2211.088 1.0771,, 2022.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, 2016.
Tianyang Hu, Wenjia Wang, Cong Lin, and Guang Cheng. Regularization matters: A nonparamet- ric perspective on overparametrized neural network. In International Conference on Artificial Intelligence and Statistics, 2021.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernll: Convergence and Generalization in Neural Networks. Advances in Neural Information Processing Systems, 2018.
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory, 2019.
Leo Kozachkov, Patrick M Wensing, and Jean-Jacques Slotine. Generalization in supervised learning through riemannian contraction. arXiv preprint arXiv:2201.066 26666666,,6666 2022.
Ilja Kuzborskij and Csaba Szepesvári. Nonparametric regression with shallow overparameterized neural networks trained by gd with early stopping. In Conference Learning Theory, 2021.
Ker-Chau Li. Sliced inverse regression for dimension reduction. Journal of the American Statistical Association, 1991.
Ker-Chau Li and Naihua Duan. Regression Analysis Under Link Violation. The Annals of Statistics, 1989.
Yuanzhi Li and Yang Yuan. Convergence Analysis of Two-layer Neural Networks with ReLU Activation. In Advances in Neural Information Processing Systems, 2017.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic Regularization in Over-parameterized Matrix Sensing and Neural Networks with Quadratic Activations. In Conference on Learning Theory, 2018.
Yuanzhi Li, Tengyu Ma. and Hongyang R Zhang. Learning wwrrpaaraetrrized two-layer neural networks beyond NTK. In Conference on Learning Theory, 2020.
Cosme Louart, Zhenyu Liao, and Romain Couillet. A random matrix approach to neural networks. The Annals of Applied Probability, 2018
Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. International Conference on Machine Learning, 2021.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 2022.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, 2019.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018.
12

---

Published as a conference paper at ICLR 2023
Gergely Neu, Gintare Karolina Dziugaite, Mahdi Haghifam, and Daniel M. Roy. Information- theoretic generalization bounds for stochastic gradient descent. Conference on Learning Theory, 2021.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.660 141..6612, 2014.
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. The role of preepparametrization in generalization of neural networks. In International Conference on Learning Representations, 2019.
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 9741-9757. PMLR, 2022.
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparamettrivaatoon: Global con- vergence guarantees for training shallow neural networks. IEEE Journal Selected Areas in Information Theory, 2020.
Sejun Park, Umut Şimşekli, and Murat A. Erdogdu. Generalization Bounds for Stochastic Gradient Descent via Localized e-Covers. arXiv preprint arXiv:2209.0895/ 2022.
Scott Pesme, Loucas Pillaud- Dillaud-Vivien, and Nicolas Flammarion. Implicit Bias of SGD for Diago- nal Linear Networks: a Provable Benefit of Stochasticity, In Advances in Neural Information Processing Systems, 2021.
Leonardo Petrini, Francesco Cagnetta, Eric Vanden-Eijnden, and Matthieu Wyart. Learning sparse features can lead to overfitting in neural networks. arXiv preprint arXiv:2206. 22t.::2206..3114,, 2022.
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, (0(4):838-855 1992.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems, 2007.
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high- dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning, 2021.
Grant M Rotskoff and Eric Vanden-Eijnden. Neural networks as Interacting Particle Systems: Asymptotic convexity of the Loss Landscape and Universal Scaling of the Approximation Error. stat, 1050:22, 2018.
Itay M Safran, Gilad Yehudai, and Ohad Shamir. The Effects of Mild Overrparameterrizaton on the Optimization Landscape of Shallow ReLU Neural Networks. Conference on Learning Theory, 2021.
Stefano Sarao Mannelli, Eric Vanden-Eijnden and Lenka Zdeborová. Optimization and generalization of shallow neural networks with quadratic activation functions. Advances in Neural Information Processing Systems, 2020.
Umut Simsekli, Ozan Sener, George Deligiannidis, and Murat A Erdogdu. Hausdorff dimension, heavy tails, and generalization neural networks. In Advances in Neural Information Processing Systems, 2020.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The Implicit Bias of Gradient Descent on Separable Data. Journal of Machine Learning Research, 2018.
Taiji Suzuki, Hiroshi Abe, and Tomoaki Nishimura. Compression based bound for non-compressed network: unified generalization error analysis of large compressible deep neural network. In International Conference Learning Representations, 2020.
Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of Stochastic Gradient Descent in high-dimenssonal two-layer neural networks. arXiv preprint 2:x:::2202..00220 2022.
13

---

Published as a conference paper at ICLR 2023
Roman Vershynin. Hiihh-Dimensionee Probability: An Introduction with Applications in Data Science. Cambridge University Press, 2018.
Martin J. Wainwright. Iligh-Dimensional Statistics: A Non-Asympiotic Viewpoint. Cambridge
University Press, 2019.
Tingran Wang, Sam Buchanan, Dar Gilboa, and John Wright. Deep networks provably classify data on curves. Advances in Neural Information Processing Systems, 14:29900-28955 2021.
Gilad Ychudai and Ohad Shamir. On the Power and Limitations of Random Features for Understand- ing Neural Networks. In Advances Neural Information Processing Systems, 2019.
Lu Yu, Krishna Balasubramanian, Stanislav Volgushev, and Murat A Erdogdu. An Analysis of Constant Step Size SGD in the Non-convex Regime: Asymptotic Normality and Bias. In Advances in Neural Information Processing Systems, 2021.
Chulhee Yun, Shankar Krishman, and Hossein Mobahi. A Unifying View on Implicit Bias in Training Linear Neural Networks. In International Conference on Learning Representations, 2021.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 2021.
Xiao Zhang, Yaodong Yu, Lingxiao Wang, and Quanquan Gu. Learning One-hidden-layer ReLU Networks via Gradient Descent. In International Conference on Artificial Intelligence and Statistics, 2019.
Kai Zhong, Zhao Song, Prateek Jain, Peter L.. Bartlett, and Inderjit Dhillon. Recovery Guarantees for One-hidden-layer Neural Networks. In International Conference on Machine Learning, 2017.
Mo Zhou, Rong Ge, and Chi Jin. A Local Convergence Theory for Mildly Over-Parameterized Two-Layer Neural Network. In Conference on Learning Theory, 2021.
14

---

Published as a conference paper at ICLR 2023
A PROOFS OF SECTION 2.1
Additional Notation. We employ the following notation throughout the appendix. For vectors v and u, we use (v, u) and v o u denote their Euclidean inner product and the element-wise product, and use / p and diag(v) to denote the Lp norm and the diagonal matrix whose diagonal entries are V. For matrices V and W, we use (V, w)p V le. and [2 to denote the Frobenius inner product, Frobenius norm, and the operator norm, respectively. For activation function IR - R, a' and a" denote its first and second (weak) derivatives, which are applied element-wise for vector V to denote V w when it is clear from the context. We the shorthand
inputs. We frequently notation aa.b b W x)) to denote o o x + b), and similarly for da, (Was x) and o (W). We use vec( E IR to denote the vectorized representation of matrix A E IR=xx, and A & B for the Kronecker product two matrices A Rm and B E RPx;; we recall that the Kronecker product is mp X nq block matrix comprised of m X blocks of shape p q, where block (i. is given by A. B.
In the appendix, particular, for smooth activations, we assume suplo'| S 31 and suplo"| - 3g for some $1, B2 E R and we denote suplo o I Bo. Bo (0, ooj. We will consider the following general case for the bias vector b E mR: by Da. such that 16y| 2 b' > O, for some b" O. This setting clearly covers the case of bj = =1 from the initialization Assumption 3. Throughout the appendix, C will denote generic positive absolute constant (e.g. 10), whose value may change from line to line.
will prove the statements of the main text in a more general formulation. In
A.1 PROOF OF LEMMA 1
In what follows, V T is the Jacobian matrix and V the transpose of Jacobian for vector valued functions, which is the same as gradient for real-valued functions.
When twice differentiable (Assumption 2.A), standard matrix calculations yield
W [(((W)
E[V W),
W)) W)]
E[D,
E[ES(A.((je::W)):j)ddaaaaa''"' [e]]
= "FBB::/ooriiiee....uuobodd...vid [c]
=[(0ffod + diag(o",
+ ( (W=))) w|+
= H(W)W D(WJU,
(A.1)
where (a) follows from the dominated convergence theorem and (b) follows from the Stein's 's lemma, and Vg. the weak derivative of ge w.r.t. its inputs. Combining the above calculations with the gradient of the regularization term, with
D(W) = E Eo (a d' o' v))
(A.2)
where Vg. is the weak derivative of ge w.r.t. its inputs, and
H(W) (a.o b)) -aa
b)) b + E y) d (W b)))],
(A.3)
the proof is complete for smooth activations.
For ReLU activations and l satisfying Assumption 2.B, we introduce the following smooth approxi- mation
(2 = log(1 + c <> 0.
15

---

Published as a conference paper at ICLR 2023
Then have
(W)
+ Blae diag(a (Wo + +b)]
= - a oo
max mm
E (ww,, x) + by) Im
As o 2 0, the critical step is to show lim E ((w, x) + 6)] < cc, uniformly for all w. Let
(w, x) + b. Then ~ (b, (2wl) (2) and
((-h) 2
(--n)
000
w
of
dz S
dz
2T w
2
0
2T
w
62
(1wll22-
1
|| ||a
+c||
2llwl1
dz.
v
|2012
(1wll2
b
2|
))-
| |
62
64)
21/
V FT
20 112
[0
I/3
2
21l
7
2
IAG
[6]
em
where follows from the Gaussian tail bound 1- (xx) S
where & is the standard Gaussian CDF; (b) holds for large enough 4; and (c) holds by considering supremum over [ [22 Thus [lo" x)) + b,,] S
V2rz
2
2 and consequently,
16,1
2||alloc
V
2 Im 3 H. (W) en
2||a|l b
2 I.
|al12 +
b
where 6* mini smm 1byl. Moreover, o' (b)+b) converges a.s. (i.e. except when (w;, x) +b; for some j) a (Wa x + b), by the dominated convergence theorem,
VR(W) lim H. (W) + lim D.(W)U 1-000
1-000
We can immediately observe from the dominated convergence theorem that (W + D(W) as 1 + oo with D(W) given in (A.2). Moreover, we let H(W) lim, (W), and observe that
-2||a| b+
2 Im H(W) - ||a|
2 2||a||
2
I
(A.4)
b
This finishes the proof of Lemma
In the of smooth activations (Assumption 2.A), we have the following bounds.
Lemma 7. Let R(W) == E[((((;; W, a. b). be the unregularized population risk. Under Assumptions 1&2.A we have
32lla| WI, = H(W) = (f 8ellalloc V2R(W) Im.
(A.5)
Proof. Assumption 2.A requires ((j.)) = }/4 (ū - 2 Hence by definition H.
H(W) [[C',. T E [úi(;; W) y) diag(o"" (W (Wa)) )))]
16

---

Published as a conference paper at ICLR 2023
The first term is positive semi-definite and it can be easily bounded:
0 I E[ W Wa)oa, (W "] 0 s Blllo (Wx)² 1 [1))lll]]lll
for an arbitrary vector R For the second term, we have
- Bll - y Im N (ū W) - y) diag(o", W [ = 3/ 50 y Lm
and Blli(a; W) yll - / W))(W) by Jensen's 's inequality.
A.2 PROOF OF PROPOSITION 2
In order to present the proof of Proposition 2, we need uniform control over the eigenspectrum of H(W) In the case of ReLU (Assumption 2.B), this follows from (A.4). For smooth activations (Assumption 2.A), we need to establish the boundedness of R ((")) along the trajectory. The first step towards achieving this goal is to obtain an estimate of R2 (W++1) (W which depends on the local smoothness of (WV
We denote by (W) the full Hessian of the risk function, an md x md matrix comprised of
x d blocks (V w.y (x(W)) isidsmm where (V ( (W)))) = 8 (W)
Lemma 8. Let R(W) E[e(j(æ;; W, a, b), be the unregularized population risk. Under Assumptions 1&2.A, we have the following estimate for the eigenspectrun of the Hessian
- -2llaaalla (OR(W)) Ind I (RR(WW = (> Bllal13 3211aa/> /BR(W)
Ind-
(A.6)
Proof. By the chain rule for derivatives, we have
V2 w, R(W)- E (lw,,,o'((ww,, x)) + (oo x)) + b,) (0(æ; W) y) ((wo, x)) + ba))
where 6ig the Kronecker delta. As result, in matrix form, the Hessian reads
V'R(W) = E o'
æa [(((;; W) - y) diag("",, (wx)) @ xx "
The first term is a positive semi-definite matrix with bounded spectral norm; indeed, for any V Rmxd xd
0 s vec(V) B[o', (Wz)o',,, (oa. (Wx). Vz
x=
vec(V)
= (o,, v)
"] - E || (Wx)
V= < s
The second term is bounded by the following:
10 T E [(((:; W) y) ((1n ))))
"[o E (;(r;; W) y) (0 ((10)) m)) + by) e)
< 32 || E. [(ú(:; W) -) E. (<x o & 3z ||1a v 2R W 3110
for all 1 J S m and for any v E Rd, which completes the proof.
Lemma' 9. In the same setting as the previous Lemma, for any W, W' pmxd we have
Rw(")) < ( (VR> w' wW) (A 87 l
282
00 3R(W)) W W 12
2
V6 1a0B1laalool111
W W
2
Proof. By Taylor's theorem,
1 (ve W), V'RA(W.)... Mec(W' W))
RA Rw(W") Rw) (VRA(W), W' W)p +
(A.7)
17

---

Published as a conference paper at ICLR 2023
for some Wa W + +a(W" W), E 10. The last term can be estimated using Lemma 8:
(vec(W' - W), V2R (W vec W' - w)) S (W))lel w <
A + .
a B2 | 6R( W
(A.8)
Next, we provide an upper bound for R W.): a): 2(R(WW) .) R(W))
E ((((; W a) - y)² (0(2; W) y)
E (0(2; W - Væ: w))²
+ +0 Wa) - 0(x: w))( W) y)
+ 2B (0(x; W.) - ú(æ; (()]:))) 22(W)
< E (0(2; W.) - ú(:; w))²
_ 8ilaa / W. a W 1 + /3/ W1l+V/2(())
$ | W' w W +223/|a||2ll
2R(W)
S <
W' + 22(WW,
(A.9)
where the last inequality follows from Young's inequality and (+) is due to the estimate below:
27
E [ú((;; Wo)) - f(;; w)) = E
Σ (( +b)) - -0((00, x)) + bز)}
of IE (a(((wa)j, bb) o o(w + b;))
<
07
Σ ((wa); - w.; x²)²
j=1
= 3 a
W
WIl.
(A.10)
Plugging (A.9) into (A.8) completes the proof.
In order to prove Proposition 2 we will additionally need bound on the norm of the iterates (W') J(20 of the trajectory.
Lemma 10. Let (W") be the sequence of PGD iterates (2.5). Suppose that there exists T > such that Rx (W') is non-increasing in t 0, ,....,TT T. Under Assumptions 1&2.A, for mij
A 3211a|| 00 / (2x(W0))
and n < + 3allall + 3211a|l (2x
(A.11)
we have
| || S - B U1l2]
1IS
(A.12)
where r/m- A - -2lla|1 v 2R, (W
Proof. The update rule of PGD reads
W (Im 17 (W + AI -WW'-""DW''Ju
(A.13)
Since R(W*) S RA(W")) s Rx(W"), for all < T, we obtain from Lemma 7
A (A- AA 32lla|lox 1 I 3 <H(W") AI 3 (A+BP1alB flaa11 + Ballalloo 2wx I
(A.14)
and for v as in (A.11) we have
0 x I. "(HL(WW)) AI (1 Tm
18

---

Published as a conference paper at ICLR 2023
Therefore,
W |1e S (1 8 )
(A.15)
and easily bound the last term
|D(W")U = (W' Vg. U| |e (W')) U|F]]
S B1 E Vg. U112].
(A.16)
The statement of the lemma then follows by plugging the above bound back into (A.15) and expanding the recursion.
We are now ready to prove Proposition 2.
Proof. [Proposition 2] Along the proof, set 0 : A+ 3 B2lla]l> / ( (W"). Notice that in the setting of Proposition 2, we have e x || S 22m, which happens with probability at least exp(-Cmd).
We will consider the event where
Smooth activations. We begin by considering the following condition
A 7 + f2 2
(A.17)
m
Solving the quadratic equation to find the range of À where the above condition is satisfied yields
3 Ba I
1
33 Balla
w
+ a | + 233]
m
4
In the setting of Proposition 2 with ||llo S mm-1 the above simplifies to
1 1 + 27 + 2R(W")
m
which is satisfied in Proposition 2. Thus we will assume (A.17) holds in the rest of the proof for smooth activations.
We will use induction on i to show that RA (W") is non-increasing. The base case is trivial and assuming the claim holds up to time i, Lemma 9 implies
(W S RA("))
7/ "l
(
+ C3 oZ | | (W")
(A.18)
Moreover, we have the following upper bound on gradient norm
2/3 IΛg
VRA I|
||H((') + AI ||
W D(W')Ul
|F 3is||al|a
mg +
(A.19)
where (a) follows from the closed form of the gradient (2.6) and (b) follows from (A.12).
(A.14) and (A.16). Thus with choice of 1) S [ Bis Bis||all(( + +1) -1 n||a/|2|| VR, [F] < 1. Consequently, with Rx < RA nl VR || C + 31 (R Therefore, with choice of 1) & (e + (0)22||||>>0))
have
< 1,
will have
RA R +1 < (W - C
As ||a|| $ m 11/2
I || S m 1/2 and = A _ we can simplify the two conditions to S m(A + X -1 and 77 S mA respectively, hence proof of the induction is complete.
Finally, recall the update rule of PGD (A.13):
W (Im (HH(WW AL AI ,DD(W')U.
19

---

Published as a conference paper at ICLR 2023
By projecting each row of this recursion onto the orthogonal complement of the principal subspace, we have
WT 1 (I.. ()H((W")) + W
Again from (A.14), we have that for < 2
Im (W) AIm) (1 -
therefore
|Wi w |F S (1 77)) [ l|r-
We have shown (2.7) and the proof for smooth activations is complete.
ReLU activation. By (A.4), for À 2 +
2||a/|
2 and < (A a
21/a
2 )-1
we have
b
b
1 m 7y ( Am) 3 (1 m Im.
Notice that in the setting of Proposition 2, ||a ||||| S X, thus S X suffices for the above inequality hold. The rest of the proof follows similarly to the smooth case.
B PROOFS OF SECTION 3
We begin by characterizing the tail behavior of the stochastic gradient noise in the SGD updates (3. I) through the following lemma.
Lemma 11. For any fixed W E mR , les
r = Vdf((;; W), B(V((((
-iViic, W).
denote the zero-mean stochastic noise in the gradient of the loss function when (x, y) are generated according to Assumption and recall that
Vilū(æ; W), = 0,E(5(a;; W), (W
Suppose supg,y, is (33xx|aa) ( ----- Gan-Saasiian
1,, EÚ v)1 S x. Then for any V E Rmxd, the zero-mean random variable (V, D)
V
Proof. We use the shorthand notation VE Vwl(Ww, and VR Vw R(W) We compute the following
V ve P
(V, P BIKV P|
(V [
22x L.. W ( 2
where (a) and (b) follow from the Minkowski and Jensen inequalities respectively. Furthermore, have
(v VVV. (Wa)ar > ppp] = B[[ V x, da. (W))))=))
+
$ 3 E V x: 1
s 3111|||12(11 V + V
where the last inequality follows from Gaussianity of Væ and Lemma 32. Hence
EKKV V "') < C, x ||||||||||| IFVp.
Invoking Lemma 28 implies sub- sub-Gaussianity of (V, V RR) and completes the proof.
We proceed by presenting a lemma which constitutes the main part of the proof of Theorem via establishing recursive bound the moment generating function (MGF) of W l1 which will in turn be used to prove high probability statements for [
20

---

Published as a conference paper at ICLR 2023
Lemma 12. Consider running the iterates of SGD (3.1), under either Assumptions 1&2.A 1&2.B,
+0)++ with stepsize sequence tne that is either constant T 1] or decreasing 1t m 1)+
(cf. (Gower et al., 2019, Theorem 3.2)). Let x = Au y)1. = Bi|lal|2x, and @ == À + Ilaal + B2 Suppose тo Σ δ Let F, denote the sigma algebra generated by {W") and ler {Ae be a sequence of decreasing events (i.e. A4+1 S At). such that A, E Fe and on A, we have H(W') + ALm = I Then, for every 2 0, with probability at least P(At) 8.
mink² (d + log(1/6))
22 m
(B.1)
Proof. Let Fe denote the sigma algebra generated by (W') 3ー= Recall from Lemma 11 that we define
r = V((j(x!)); w,))) (0) E Ve(j(> w'),
with
Vi(ū(=!) w y, (0)) - 0 (()(x( W y0) )o (W's (0)
Then for the SGD updates we have
W W T/ VR -.
By projecting the iterates onto the orthogonal complement of the principal subspace,
W (Im ((H(W")) ALm) W r
Let M, = Im N1 (WL(W() A.)) Then, by observing that 1A. S L for any 0 S S _ح mnnek
we have
Layss wi Fo
S E n, M.wi Eotdiiii M. w.w..ir.. ri), [Fo]
E ell c II (-
M. M. r1), Fi [Fo
- E 1 -l
B 2 Fi]] E W M.
Fi
Fo
(B.2)
where the last inequality follows from the Cauchy-Schwartz inequality for conditional expectation.
Moreover, it is straightforward to observe that ( W'), ) y ( s x² | hence
E [V W W'), S kdd.
Note that by Jensen's 's inequality
| l2 < 2 (0) + 2E Ve((W' (0, "0)
Consequently
E [lexp(2sn) [1)
| Fi] S xidd) E E[exp(Asnf x² || | Fi]
S Eexp(4s7f x² 2d) exp( (8sm) (d)),
where the second inequality follows from Lemma 33 for 4s7 < 1/4. Since 8 S mo in order to satisfy the condition of Lemma 33 need to ensure m///m < 1, which is guaranteed by our qe S 1 assumption for a suitably small absolute constant, as 7 //m= & S Ø.
Next, we bound the last term (B.2). Let V == -4s, W Then by Lemma we have
E E[[xp((V, (v)))()) Fi] S exp(Css1 x W lFF)
Putting things back together in (B.2) and using the tower property of expectation, we have
E L e
wi
112 [Fo
< S
1 < [[tMM...11/Cooffppfr 11 Fo].
(B.3)
21

---

Published as a conference paper at ICLR 2023
Next, we bound || By definition of A,, we can already ensure H(W") ž 2 Im in (B.3). Recall the definition of H(W")
H(W") = = W"), (W' (W'xz)
TT + 01((((x;; W"), y) diag(o" (W' x))]
Notice that 0 0 v. I under either Assumption 2.A or Assumption 2.B. Moreover we have, [0, S x. Thus,
H(W*) AI 3 (> / a 822 102x @T.
Therefore,
Im ne (H") + AIm) 31 267 Im. As result ||M2|12 S 1 n=Σ Combined with (B.3) we have
m
E L44 exp(s
WT +1 (()) |oo] = E E[LA, exp(s(1 + Csn x2 -1 W 1 Csnf [oo
m
f exp (xp(Csnfr x² d) E 1A, exp(s(1- m:) W ()) Fil (B.4)
where the second inequality holds by the fact that Csni x² < nt/mm, which in turn holds when small enough absolute constant is chosen
S maxx? Also notice that for decreasing stepsize,
(t++-")²
(++ (1+1 112 S
+1+
T
(B.5)
11+
(1+71)²
(and the above holds trivially for constant step size), thus when S C1 we have s(1 - 77:7) S CCty with the same absolute constant. Hence we are allowed to expand the recursion (B.4), which implies
City for some absolute constant noK
E[1 [[1A
exp(>) exIxxexp(s])
W' (1F) Fo] S exp
- Csk 2d
- 2:)
II
(1
1o7 m
j=0
for all 182 Moreover, direct calculation implies that with both constant and decreasing step sizes of Lemma 12, we have E ni [
1/3
Crann
(with C for constant
m
?
stepsize). Thus, for all 0 < s S 1
E LA exp(<) Foo
exp
II
20 m
Common
Finally, we can apply Chernoff bound to obtain
P(
Fo) s exp
-
I[ 7- |/ W
Cmmex3d
By choosing
1-1 E =
Cmmpa 2 (d + log(1 (10))
7:7 W w
?
==0
and the largest possible
mn R we obtain
P( W | 2 Fo) P(A. n W II 2)) +))++P(9))) 56+
Taking another expectation to remove conditioning on initialization completes the proof.
The proof of Theorem for decreasing stepsize follows by a direct computation the quantities in Lemma 12 and is presented below. On the other hand, in order to get a better dependence on X, choosing the events A, for constat stepsize is more subtle and is presented in Section B.2.
22

---

Published as a conference paper at ICLR 2023
B.1 PROOF OF THEOREM 3 FOR DECREASING STEPSIZE
This part is directly implied by Lemma 12. The following argument holds on an event where [ |F S /m. which happens with probability least ( In order to see this connection, we will first present an improved statement over Lemma 7 for the case of smooth activations, Recall the definition of (W) ) for the squared error loss E(j. y) ("
H W ) ( W Wa)oaa (Wx) |] [(((: W. .. y) diag(o", (iag(""((W=))]
Notice that under Assumption 2.A we have lı/ - Bolla]]1. Then basic matrix algebra similar to that of Lemma 7 along with the triangle inequality shows
82llall (Bolaal): + [ < H(W) = (Allaa + 8ellall + [[(l)))
Therefore, with 2 7/m B2 (|||]]
1 9o El[y])), have H(W) AIm 7/mIm for all W. In addition, 0 (v. v)1 & :Bolaal1 + by the triangle inequality. Thus we can invoke Lemma 12 with 17c = m
x = Bolla|l: ||| + K, and 14. = 1. Recall that in the statement the theorem, 8g = 31 = 32 = 1, K $ 1, ||a lcc I 1/m, |alla I 1/vm, and ||a|| I 1, hence I x X, and with t* x desired result follows.
((*+))2
(1+++1))
can guarantee niA - 1. As the step size condition of Lemma 12 is satisfied, the
Similarly, for ReLU we have [0 v)1 S by Assumption 2.B, and for 7/m ب|اااا (recall b* 1 in the statement of the theorem), we have H(W) + AIM 2 2//mIm- Hence this time, we can invoke Lemma 12 with the same decreasing Tr. 14, = 1, and x
5" en
B.2 PROOF OF THEOREM 3 FOR CONSTANT STEPSIZE
In order to improve the condition on X, we will specifically look at the events A, on which maxo</<t Rx(W5) is bounded. The following lemma indicates that these events occur with high probability.
Lemma 13. Under Assumptions 1&2AA, consider the setting of Lemma 12 with constant stepsize 7 S 2- Then we have with probability at least - eexp(CCTnéd))
(w < RA(W CTT2 0StST
max
(B.6)
Proof. First, recall from Lemma 8 together with / /Ra W 2 /.llaaa1 + + x, that
( & A + |||| 2 + 32 =
We will first prove that for any t S and any s S ( we have
E (
_  (
== = Recall that F, is the sigma algebra generated by (W'} and and Taylor's theorem and Young's inequality
V -V(W")-VR(W)) By
RA (w < RA(W (VV) VV
" +
< RA(W"') - - (
(R, 1' कد ||
and
(w Fo
-
1r' Fo
18
RR
-
+ |
F Fo]
M/ E
- (
[F
B B[
|F:
1Fo
23

---

Published as a conference paper at ICLR 2023
where (a) follows from the Cauchy-Schwand inequality for conditional expectation. Moreover, in this setting we have [ y) = 1 yl & Bollall1 + K, thus letting x =
=Bo0laa Bullaal + K Lemma
11, by the sub-Gaussianity of I we have
e 2an (W")))) Fr S
Furthermore, we have the following upper bound
E Caan VeE
Fi
S E
Capo wd
1|
Fi
Cxan
where the second inequality holds for s 2 0 with sufficiently small absolute constant by
Lemma 33. Similar the argument in Lemma 12, we choose S (17²)=11, order satisfy the
condition of Lemma 12 it suffices to have nē S for sufficiently small absolute constant, Putting the above bounds back together, we have
E e R
RR -n((1 -n(1 CSnn (
+Canp +Con
Fo
Expanding the recursion yields, for any S (17²)-1 (with sufficiently small absolute constant chosen)
(W')) Fo S (WCCen"
As result, by applying Markov's inequality at time I s T, we have
(P(R 2 Rx(W + CTT)R220) (d) ou S CTTnd Consequently, with a union bound we have
RA(W') 2 Rx(W") + CTT22 Ød S Te Te- -Trnd
max
OSEST
which completes the proof.
As depicted by the following proposition, the rest of the proof analogous to the decreasing stepsize case.
Proposition 14. Consider the setting of Lemma 13 with constant stepsize S ē- and X sufficiently
large such that À 2 7/m B2 lal 2 2(Rx(W)) CTr @d). Then with probability at least 1 Te-CTrèe we have
m 1 Cmmr? (d + log(1/6))
(B.7)
?
Proof. Let A, (maxosige Rx (W)) S RW + CTr22 ūd). Notice that A, is ( (W)
measurable and Art1 A, By the bound established on H(W) in Lemma 7, on A, have H(W') Alm 7/mIm for all < i < T. Moreover, from Lemma 13, we have IP (AC))
<
Te -Tnnd Invoking Lemma 12 finishes the proof.
LI
The above proposition immediately implies the statement of Theorem 3 for constant stepsize, which we repeat here as a corollary of Proposition 14.
Corollary 15 (Proof of Theorem 3 for constant stepsize). Consider the setting of Lemma 13 with X given in Proposition 14, for constant stepsize 2m (Tg(T) with T (1/5)C1d, Then with
probability at least 8 we have
| |F S
||e
mk
(Tg())d + log(1/5))
(B.8)
T
T
24

---

Published as a conference paper at ICLR 2023
C PROOFS OF SECTION 4
C.1 PROOF OF THEOREM 5
As our arguments are based on the Rademacher complexity two-layer neural network, we require the knowledge of the norm of W'. We prove high probability bound for this norm in the following lemma.
Lemma 16. Under Assumptions 1&2.A or 1&2.B with either decreasing or constant stepsize as in Theorem 3, let x = supg.y 3.y
0 f(g y)1 < oo and Kio = 3111allo0 Then for any 2 1, with
probability at least m exp
gytd mpn
we have for all 1 j S m
r-1
3 v
|l2
(1- 2
(C.1)
==0
where ip = for decreasing step size and = log(T) for constant step size.
Proof. First, we prove that for any t >> 0 and 0 S s
2vd
we have
t-1
m1lll + 28 vd (1
E[exp(s
| w I exp
$ III
(C.2)
1=0
The base case of / 0 trivial, and for the induction step, for any 0 S S Ric n
2v have
E [exp (xx) l12)) W
E [exp(<s((1 nA))w V (ū((;; W W)), v))2)) w
()) (l) wo] S [[exp(s(1 m)) sn V ((((:; w))
sy
[[exp(s(11 -
- > sp || |1) W
E [[[ep(s((1 (anos((1=>)]llto/ (2) B/exp(s77 oo || W w wo
(a)
8² 22 m
S eexp(s(1 ,) exp s7/ vd+
2
(b) exp
23 Vd
s II > X
where (a) holds since ||>|| is 1-Lipschitz function of a standard Gaussian random vector, thus it is sub-Gaussian with parameter (Lemma 29) and additionally [[[[=||2] S Vd, and (b) holds by the induction hypothesis (notice that for decreasing stepsize s(1 TTA) S 2va - apply the following Chernoff bound,
by (B.5)). Next, we
1-1
(1 - 71Aillo91le + 3100 Vd w I exp
Sti0o vd
PP
112
II 1
A
which holds for any 0
Choosing the largest possible and noting that Mh-1 S upper bound on the conditional probability, which followed by taking expectation removes the randomness of conditioning wT. Finally applying union bound gives us the desired bound.
2 Kiso Tpt1
2mg yields expl
-yid mpxX
In addition, we would like to approximate R, (WT) and R. (WT) with R, (W(W) and R, (WT) respectively. As result, will investigate the Lipschitzness of the population and empirical risk in the next lemma
Lemma 17. Under either Assumptions 1&2.A or 1&2.B, the truncated risk W + R,(W) is V ||a||l Lipschii, Moreover, for T 2 d + log(1/8) with probability at least 1 - 5 over the stochasticity of (x(t) JoSEST-> the truncated empirical risk W + R, (W) ) is CT || Lipschitt for some absolute constant C.
25

---

Published as a conference paper at ICLR 2023
Proof. We begin by the simple observation that + E(g. y) A /2r-Lipschitz when E(0. y) 1/2(ý y) and L-Lipschitz when (0, v)1 - 1. As T 2 can consider both of them as /
Lipschitz. Thus by Jensen's inequality
(W) R. (W /2
Blū(a; أ(:; W
27
m
V E
aj (Ww,, x)) + b;) Σ a a + bb,)
V2T||a|12,
E (o((uu) by) o (( 2²) + b,))
V2T
E (n0, w/) 2
(C.3)
$ /2 W
where (a) follows from the
inequality. Note that Equation (C.3) also holds for |R, (W) - R (W') when expectation over the empirical distribution given by the training samples,
meaning
m
(W) - Rr (W')l V2
Σ
w - -
T
Σ
2002)(1) (w, - wwj). (C.4)
By Lemma 30, with probability at least
8, we have
T 1
T
Σ x
(t)
(1)
- Ls
11,
2
which completes the proof.
Lemma 18. Suppose either Assumptions (1,2.A) or (1,2.B) hold. Denote the loss with E(g. y)
l(j y).
5
w E m ba. |aal2
Fa
|b|lo
10,l S rw A1 S m Srb,
Vm
and
G - { (w.v) (0) E(j(;; W. a. (W.a,b) E $}
for x R²th and y E R. Let R(G) denote the Rademacher complexity of the function class G (see Lemma 18 for definition). Then with æ NV. for some U Rexd we have
2 R(G) < 27 U| ( T
where T is the number of samples.
(ū. v) : (W, a b) E 5) for fa w (w. v) (( W. a b) y- Proof. Let {(t.)) v) Ja,w(.
Define g(2) = ((2) A 5, and notice G = () + (2, y)) : fa w E F}. and that @ V2r-Lipschitz (thus V2r-Lipschitz well, for - 1) function. Then by Talagrand's contraction principle we have (G)) < V2TR(F). Moreover, let (EJooeeT-1 be sequence of i.i.d.
26

---

Published as a conference paper at ICLR 2023
Rademacher random variables. Then similar to the Rademacher bound of Damian et al. (2022)
RiF) = E
sup (W a),)
& (a oWWx +b) (o(())
T
1=
T-1
La (Wax +b)
sup
(..EE
Ta T
E
sup
Go (W + b)llo0
(W.N)E
Ta E
Σ S.o ( i0, 26(10) (0000(00+00) +B)1
sup
T
|| 18/5r 1=0
(b) 2.17a
T-1
E
sup
Ee ū0, x1)) +b)
n
[ 1b/
T
2 T
E
sup
Ex क) + sup
Σ Ecbl
ib lla Srm
18/5ro
2/ ra
E
E. ll2
+ TbvT
T
2 (rai |U |F + rb)ra
VT
where (a) holds by Hölder's inequality and the fact that ||a|]a S /m S ra. and (b) follows from the fact that is 31 Lipschitz, thus another application of Talagrand's contraction principle.
Proof. (Proof of Theorem 5] Let & denote the event of Lemma 17. We begin with the following decomposition for generalization error which holds on E7. R, - R, (WT) = R, (W") - R (WT) + R S CTB (T W || + R2 (W R2 (W
R. (WT) + R. (WT) (W
where the upper bound follows from Lemma 17. Consequently,
sup R2 (W" .. R, a. b) < a.b
Crêine Vm
|W + sup R, (WT a. R (WT. a. b).
a.b
(C.5)
We begin by upper bounding the first term. From Theorem 3, event &g we have with probability at least 1 - (8)
W IF S mn
d + log(1/5)
2T
Next, we bound the second term in (C.5). For each W, define W = U' where Ul the Moore-Penrose pseudo-inverse of U. Then, since we have the representation W MU for some M Rmnk
WU = W U MUU = MU = W
Thus, Wx = Wa and ((j(:; W. b).0) = <(i(;; W, a. b). y) for ã = Uz when W is in the principal subspace, i.e. W WI. ||- Let Eg denote the event of Lemma 16, on which
3xo vd
| ||2 < II
+
==
and consequently
3rizo Vdl A
| 1|2 < |Ut||
II 2=0
(1
n m
110 +
27

---

Published as a conference paper at ICLR 2023
for any 1 1 S m. Define " the RHS bound above. Then on Eg
sup R. (W R (WF)) &
R a, b) Rr a.
sup (Wa,.)es
a.b
where we recall
5 = w E m a. a.bER" ||aa22 S
Ts vm
|||||| rbb ا S .1
v و m)
Additionally define
g= ( v) + E(ū(;; W, a. a) Ar (W)
Then Lemma 31 and Lemma 18 yield
sup R ( R2 (Wr) S 23(()) S T8i(r + + -
(W b)
Besides, as the loss is bounded by t, by McDiarmid's inequality, on an event &y which happens with probability at least 1 - ((6) we have
RS (W) (W) SE
B(W) + CC
log(1/8)
sup (W...
sup ()
T
and consequently on E
log (1/6)
sup R2 (W a, b) ft. (WT a. b) TBi( + +r)raa|lU1 + a.b
T
T
Finally, observe that ||a|| S /milaal S Ta. and without loss of generality assume is orthonormal, hence |Ut|l2 = and | = Vk, thus with probability at least 1 - o(8),
d + log(1/5)
sup R. R. (WT b) R.-WW" .) STBirany a.b
2T
Koo
dk
+ TBr
Tw +
+
Tb
1+ +T
T
log(1/8)
+
(C.6)
T
We remark that in the setting of Theorem 3 which is adapted in Theorem 5, ||a|| m- thus
Noo < m- Finally, observe that r.. /2m with probability at least - (8) over initialization, which completes the proof.
C.2 PROOF OF THEOREM 4
Note that due to the special symmetry in the initialization of Algorithm 1, while training the first layer, all neurons have an identical value, i.c. wi wt for all j. and that the stochastic gradient with respect any neuron can be denote by Ve ab ((i, y)o' ((w, æ)) + b)æ. Furthermore, V.,, R (W) will also be identical for all j. which due to the population gradient formula (2.6), we denote by
VRA(w) = (h(w) + A)w + d(w)u,
where h(w)) = En Hig (W) and d(w) = a I [oP (() y)o x) + b)) x))] Additionally, via the arguments in the proof of Lemma 1, it is not difficult to observe 7/m  < h(w) + X m m- Furthermore, similar to the arguments of Lemma 11, (V,, v) is Callu v Callvl|2-sub ll-suu-Gauussian Gaussian for any E Rd. Next, will derive lower bound for (w', 2)) to argue that useful features are learned, which first requires obtaining a sharper upper bound on | 1 than that of Lemma 16. This improvement is possible due considering the special case of w w here.
28

---

Published as a conference paper at ICLR 2023
Lemma 19. Suppose 2 d. Then,
1200 12 S to (x+t) w /lll
t"
Cma
with probability at least 1 exp( ((t + +().. In particular, using the union bound, we have
Cma
sup u 1l2 S 1|20 +
7 2 1
124
with probability at least 1 - exp(- -(( + to)) exp(-Co
Proof. Let e' = Vwl V R. Then we have
Te v Rr 1/.e". w ++| = w'
Recall that (c', v) is Callee|a--ub--Gasssia and F; is the sigma algebra generated by {wr Josige. Let w' = w² - T/ V w Rx. Then, for any 0 S s R mm
E B[exp(>|le0*** (1x Fol = E eeep w 11
2sT (estiew" 5²) + 87h ( Fol
E exp (alp(ss|l+)
E eexp( 1aap--4sypp 48mh (w', e¹)) Fil E (eeex(22sn?lle Fr] 1Fo
By sub-Gaussianity of (wo e²) we have E[eep(( 4smp (wo, e¹)) F v exp(Cs a More-
over, as || / I |a|||æ]||2 by Jensen's inequality
113 < v + 2E < 2a (1| + dd).
Thus EE [exp(20n)]1c Fi 3 exp(Cinnfadd) for s s n. (which holds by s $ mn see the proof of Lemma 12 for more details), i.e. we have
E [exxp(s|l +1 (12) Fu
E [ex((s(1 + Csn w Csntt a² |Fo
Recall that by our choice of Mr. S (1 Te (x((")))) + A)) - 1 Ye2 (cf. proof of Lemma 12), and
w² =1 (h(w"))) + A)) ( u. As ||ul|2 and (D(rr) } S [a]. we have
||
(1 7:
CaFni 22,, 2:2 )|(=0² ||2
m
WV@
40ma² 7 (1 nu mx7) 1 112 + Ca2ng
(1 7: 2
M
7
4m
IA@
Cmmiia?
(1
3 ||
Caan?
2m
where (a) holds by Young's inequality and (b) holds for ma/m S 1 with a sufficiently small absolute constant. Therefore, for m
|20*11 112 + Csmmaa
exp w0 +1 112) Fo] s
exp s(1
2x2)
C's a²d
|
Notice that we can expand the recursion since s(1 n+?) S Eq. (B.5)). Expanding the recursion yields,
(cf. proof of Lemma 12,
1a
(5( 1"
2
Csm²aa( d)
R [[cxp(s|lur
(3) Fol S exp
| ||
(++
2 + t)
Finally, we apply Chernoff bound with the maximum choice of s S mmjaz. and combine it with the fact that || with probability at least exp(-Cd).
Lemma 20. Suppose mab < (f)) Then, we have (ww", u)) 2 with probability at least
1 2 2exp(-Ct) exp(-Cd).
29

---

Published as a conference paper at ICLR 2023
Proof. We will only prove for the case where f is increasing as the case for decreasing f is similar. We begin by proving an upper bound for ((w) when
|| | 222 S 1. By the triangle inequality,
1 yl - | + [f())) + ( x)) - f(0) + lcl
Furthermore, |ú| I maa((w, x²) + b). Thus, for
|s(0)] - mab) mab
(wr, -
A
2ma
xu, S and |e| X 1 for sufficiently small absolute constants, we have vl S 1 hence 0f2 = -1. Then we have,
d(w) E[ y) ((w, .>) + ((u, x -a [[]((e] $ () s²) S ()), 11) B[1((c) $ B] =)| K (),(((u,))) x)) / (), x\)]
S -a.
where the last line is obtained by considering supremum over || S 1.
Let A,- = (supt,sr ||2 2 }. Then,
[exp( (-s =<(((000 (wo't! 1
< E eepp( s w ++ 1A.]
= R [exp(
-s (ww" "o',, u) + 8m) (vl+) + Xwo u))
< 18 [expp( (s(" u) + + (VR,, u) + Cs2 a² 1A.]
[[[ex((=<( - (h(ww)) + X)) w'' u) + sm (d(w²) ) + CssRaa)) Atd]
exp(-Csmma) E [exp( (-s(L - n (h(w") + ))) u))
where (a) follows from the sub-uuuuuuseee uub-Gausianity of the stochastic noise in the gradient, (b) follows since (VRx("") u) = d(w") ) by definition, and (c) holds for s S (1)a) with a sufficiently small absolute constant. Notice that by the condition on !" inherited from Theorem 3, - 1--mmhh (nhh(u)))) +)) >
7+7) 0, and since s(1 - (h(w")) + A)) S &1( m ), we can expand the recursion,
E B[exp(-s((u) (((' LA.]
S E
exp
-Osa
(1- (1-
n= II
I
j=i+1
i=fo
t" + to 2)
+ )
1' to ) + C's
I
< E
cxp
-Cs
1
+t
where in the second incquality we used a x m 1 and y x 1. Applying the Chernoff bound implics
that (wo u) 2 1 with probability at least P(49) exp(-C's 2 exp(-C(t" + to)) exp(-Cd)- exx((CU) Finally the result follows by letting to Ct for a sufficiently small absolute constant C.
We have proven that Kw", u) 2 1 while ||w 1l2 + 0. This fact shows that the features learned in the first layer are useful. What remains to be shown is an approximation result, such that for carefully constructed second layer, the network can approximate polynomials of the desired type. This type of approximation using random biases has been adopted from Damian et al. (2022). We first present an approximation result using infinite neurons.
Lemma 21. Let 0 < |a| < r and b ~ Unif(- -2r A, 2r For any smooth f: - R, les fa R -
R be smooth function such that fz( f(=) for
|z|
< r and Fa( 2r fa 2nA
0.
Then, for |a| S A we have
A
b
E
fu
(azz + b)
= =(()).
a
α
30

---

Published as a conference paper at ICLR 2023
Proof. Using integration by parts, we have
B
4r a
Ja
b
o + 6)
=
2nd
Fa(
(2 +
db
-as
a
=- 2r
)2 + 2r
+
(b)
a
2r
= =f((=) = f(2).
Now, by a concentration argument, we state an approximation result with finitely many neurons. Lemma 22. Let r* & lays S r and by ~ Unif( 2r 2r Let
A. A sup
I (=)1,
(C.7)
2
where Jay is the extension of Sa; introduced in Lemma 21. Then there exists a(az. by) such that for any fixed E A), with probability at least - 8 over the choice of (bz), we have
A xx
log(1/6)
Σ b(aa,b)) f(z)
Moreover. ||a||2 2 /m
Proof. Let (2) be a candidate in Lemma 21, which can be obtained by e.g. extending f with suitable polynomials (notice that fa only needs to be twice differentiable its domain). Now choose aj 4 -- fu ( 4.). Then Lemma 21 ensures that a3m
Es, [a(aa, by)o((zz + by)) f(2).
It immediately follows that ||all S Cra. and (a,o(azbb;) < inequality finishes the proof.
Cr'ss. Applying the Hoeffding's pm
pd ym
In the following lemma, we will briefly record useful properties of W which will be of help for invoking the above approximation results and providing guarantees when the second layer is optimized by SGD. Through the rest of the proof. will add the mild assumption that log(1/ Otherwise, we need to add e-Cd to the probability of failure in Theorem 4.
Lemma 23. Suppose T Z d + log(1/6). Then with probability at least - ô over the choice of (by)isssm and ((x((),)) (²) TT- the following statements hold:
1. /w7/l2 (w), u) x for all 1 S S m.
2. DT (00) (0) ( T 1
3. [WT]lFS mid+log(1/8))
T
4. (u, x (z) S A for all 0 stsT
5. I 2001 S mml ( for all StST-
Proof. We will show that each of the events holds with probability (w.p.) at least 1 1- 06)) Recall from Lemma 19 that |/ l|2 S for all j w.p. 2 1 (0), which implies the same for (w. u). On the other hand, from Lemma 20, (wo u)1 ス for all j w.p. 2 - 0()). Combining these events implies that kwf, u)) x 1. The fact that | Dd- (0)) T ||2 S 1 w.p. 0()) for T 2
31

---

Published as a conference paper at ICLR 2023
d + log(1/)) follows from the statement of Lemma 30. Furthermore, [ X
mld+log(1/00)
w.p. 1 1 (6) follows from Theorem 3. Note that as (u, .z!)) ~ N0, 1), by the choice of A, we have (u, 2 Z A w.p. (6//T), thus Ku, az(t))| x A for any 0 S T w.p.
by union bound. Finally, we have
m(d + log 1/5)
W x " 12 S W
x
122 + W
x 10 X Vmll
x x(()))
+
||x
T
- m
u. x.()
m /mll=((1) 112
The first term is already bounded by vmA with probability at least (6). Moreover, recall that thus by the union bound (+) ||2 - /d S Vlog(T/8) 2 A for all 0 < t S T - 1. Thus w.p. 2 1 (6) we have [ (x()) 2l S (m((v++) A) which completes
(|=(*) ||2 E || (012 [|2] is
the proof.
From this point onwards, will denote the Huber loss with (n(ū.y) l(ÿ y) Notice that En is 1-Lischitz.
Lemma 24. Recall
T-1
R(WT, b) T
1
Σ Eaja((u3 w. xll)) + bbz u, (((u,z(1)) - c01) -c())
t=0
the empirical risk of T given by Algorithm 1. Let A logg()) A. as desined in (C.7), and
by lid Unif(- A). Then, with probability at least 1 8 (over the randomness of (b))sssmm and
(x T hence WT), for T 2 d log(1/)), there exists with |a 1l2 S A such that vm
R(W a" b) - E(en(e) S A. V log (T)))
d + log(1 (10)
+ log(1/5)
+ A.
m
T
T
Proof. We will condition the following discussion on the event of Lemma 23, Let = (w u). and let a" be constructed according to Lemma 22. By the Lipschitzness of the Huber loss, for an inividual sample (x, y) we have
(u(((:; W aa, b) - -) S En(e) + V5 W aa, b) - [((uu, x;)1
S (H(E) + 10(x; W a" b) - -i(æ; WT, b)
+ |0(;; WT. a" b) f((u,z)))
Moreover, by the Mauchy-Schwarze inequality
m
10(:; WT. b). 0(>; WT: a², b) s ||a0|2
0 (aw(w , xx) + b,) o(((10)) 11 2)) + b,))
$ ||a0|12, 11a
Σ (wiz ي (wof) 2.)
Additionally, since + ETT (x(())
1/2 S 1, by Jensen's inequality,
Σ Σ H(: WT a" b)) 0(2: WT, a- b)1 ||0||2
T Σ 1
|| ||2 W |F
d + log (1/1) 1/8)
T
32

---

Published as a conference paper at ICLR 2023
On the other hand, let 2(t) == (u, x))) S A. Then, we can apply Lemma 22 along with union bound, which states that with probability
(()) over the choice of (b,)ısism
T-1 Σ 1) WT, a" b) 5( u, x T
T-I
ajalazz (1) + bb) )(
T
t=
1=1
log(T/0)
SAA.
m
Combining the events above, we have with probability at least 1 6,
R(WT, aa. b)
+ Σ ()) و A.
Vv log(T/6)
d log (1)))
m
T
The final step is to apply a concentration bound for ΣT-1 (H(E())) Note that as CH(C) S |c|, if |e| is Gaussaa,, then (c()) [[[n(e)] is also Cv-sub- Gaussian (can be verified e.g. by Lemma 28).
Cv-sub-Gaussian
Then sub-Guun Gaussian concentration bound implies that R[(u(e)] Ei (n(e S v (og(1/8) which finishes the proof.
T
Let [-][-] denote expectation w.r.t. the random sampling of SGD used train o, hence conditioned on (al y 101 Also, define the stochastic noise in the gradient w.r.t. a as
c¹ Va (al((( WT, aa, b) - y()) V.R. aa b).
Notice that Es(ea] = O.
Lemma 25. On the event of Lemma 23 and with (b,) itd
Unif( A, A). consider the mapping a - Rx(a), Then, (a) Z mm²++ + X, and [|ealla 2 vml Vd A).
Proof. For v R((a), and any E RM with ||0|| = 1, we have the following computation:
v. (2R
Pels, ToW"" (0) + b)o(() (1) +b) o
To + X
lo (1) + +]1l+ x
||: + |1b12 +X'
3 2 +X
where (a) holds since 1 1+ (0 lla S 1. Thus (a) 3 mA2 + X. On the other
hand, as | ( 112 2 vm(V++ A) for all 0 <t - 1, we have
|c s 2 la s 2 T ( + bll2 x vmm(V+A)
Now we can analyze the SGD run on the second layer a to give a high probability statement for (ax. As (a) is a smooth and strongly convex function of a, we will state the following well-known elementary convergence result of SGD for smooth and strongly convex functions with
bounded noise, which we present in a high-probability framework suitable for our analysis.
Lemma 26. Let R R. - TR. be ju-strongly convex function satisfying Imm I VaR(() 1 LI
Suppose we run the SGD iterates a'+1 a' n/'' with Elg" a'l = VaR(a') and I 2 S G.
2++1 Then with probability at least 1 1 Choose T p(t-
R T
CLG²
CG2 log(/)))/))
R(a") - R* <
+
where R* arg min, R(a).
33

---

Published as a conference paper at ICLR 2023
Proof. Let c' g¹ VaR(a") denote the stochastic noise. By the smoothness property of R, have
R²= R((a') - -R* - V. + dt) + Lnh
R(a"t1) -
2
S R((a') - R - - - (V. ,e)) c')+ Ln
2
v G. thus ||e' v 2G and the zero-mean random variable (V (aa)) is 2G R(a (222a))l2200)))) gill--ubb-Gaussiin conditioned a'. Now, the following recursive bound on the MGF of R = R(a') - R'. For 0 S 5 1g.
Notice that by Jensen's inequality, [a
establish
we have
E
R
exp
sR
- snel V Ro ST (.. e.)
ni LG
s7I(1 G²) ( IG ]
IAB W/
exp
sR'
exp s(1 - nH)RT* + 1G
2
where (a) follows since R(aa is strongly convex thus satisfies the Polyak-L Lojyak--jjsiiiiizz inequality
(p(R(a)) R²) < R(()) [3. As s(1 7(p) $ 4m-1Gu (cf. (B.5)), we can expand the recursion
and have
t'
16LG
B (exp(s < exp
s
to +t
R
1² +
Finally, applying a Chernoff bound using 8 (4m) -1 concludes the proof.
We are finally in a position to complete the proof of Theorem 4.
Proof. [Proof of Theorem 4] We will consider the event of Lemma 23 which from Lemma 24 know with probabilility at least over the dataset and (6,)ısısm we have
log(T/8) m
d log(1/6)
log(1/6)
min R(WT b), [() A3
a: a:laalaas
T
Notice that a + Re( RWW a, b) is a convex function. Thus by strong duality, there exists X >> 0 such that the value of the above constrained minimization problem is equal to the value of the following regularized minimization problem,
( (11001770)
log(T/0)
d + log(1/0)
log (1g(1/))
min RxxWW' ,a, b) [((n(e)) & A3
m
T
Explicitly, this X' can be chosen such that the unique solution to
VaR(W" aa" b) + X'a =0
(C.8)
has ||a"||2 S A.. Notice that this a" is the unique solution arg min, (WT a,
Moreover, from Lemma 26 we have
R(WT ,a",b) (d+ AA) X//m+log(11) log(1/)))
Rx ( a" b) (T. aa, and by strong convexity
T12
(N/m)2T'
2 (w "n"T"aaT.aa _aT,b)= - Rx aa' (ax"""a',,b) ,a',
|la" a
1 -
Thus, with sufficiently large T" such that
RT a. b) + (d + AA? + 2/m + (116)) (2/m)2T'
(1g81
XA.
$ A?
(C.9)
T
Vm
34

---

Published as a conference paper at ICLR 2023
have ||a 1l2 x vm
A.
and
(
d log (1))
Rx [((n(()) A?
m
Finally, we invoke Theorem 5, to close the generalization gap and get
log(T/5)
d + log 1/8)
log(1/8) (6))
(W aa" .b) B((((())) A.
+ VV
m
T
D EXAMPLE OF NON-CONVER (W
Here, outline examples for which Rx(W) is non-convex neighborhood around W while å satisfies condition Proposition 2 or Theorem 3. For simplicity of exposition, in m
both examples we fix a 1m where 1, is the vector of all ones. It is easy to observe that the results hold with high probability when follows the initialization of Assumption 3 well. Furthermore, work the event where |W S 2/m, which happens with probability at least 1 cxp(-mdd/2)
We begin by constructing a non-convex example for Proposition 2. For this example, we choose o such that $1 S 1, (1) o--1) 0, o' (1) (-(1) O, and o"(-1) o"(1) B2 1. An example of such a function is (z))
c06(52)+1
Then, using the computations of Lemma we have x²
ì m
(W = (W
(g(x; W, a, b) - y) diag(o" (WWx)) @ xx
mmd-
Which simplifies to
V W (0) -Im
(- m
E yaæ
i m
& E
yyx
AI = Im @
Id
Therefore, (0) is not positive semi-definite (PSD) if and only if =1 PSD. Moreover, by Jensen's inequality
E yov
+
L is not
1 <
(c; W a. b)
(a(w a) + b;) abb)))
hence
2R W < 22 +2 y
$ 8 + y²
Now, let K(w, 2 for some w with ||w|| = 1. Then E[y"]
3K? and choosing
1 + 9 19+66* 6K for arbitrarily small & suffices satisfy the condition of Proposition 2. Then have
E y "]
+
- B K(w,
+ X
V Ry = T
w
w
m
3K + 1+
v9 6AC
+
where the above inequality holds for sufficiently large KC, hence Rx(W) non-convex at least on neighborhood around zero.
Next, we construct a non-convex example for the smooth and decaying step size case of Theorem 3. This time, we require o -B0 -1 (which automatically implies d'#11) 0 as a attains
its minimum) and o((+1) = 32 = For instance, we can choose a (2)) cos(az) Then simplifying V²Rx (0) yields
Im 3 (-I, yex
{A-1
E yarx
RW
Lmd Im &
35

---

Published as a conference paper at ICLR 2023
Thus, need show that 1 yauw
is not PSD. Let y } (anh(ww,x)* (1))
and X Ely] + 7 (notice that y 2 O thus X indeed satisfies the assumption in Theorem 3) with
IE (120, x²)² - -1/w0 tanh( x²)² 1| 0.
7
Then we have
7|12012 + E ((1lerlé - (wo xx)²)
1 1
ya
w
d
171
m
m
((w, x²²² 1/
| (
tanh ((w, xx² -|| 0.
m
Therefore, once again we have shown that Rx(W) is not convex neighborhood around zero
E AUXILIARY LEMMAS
In order be explicit, we state the following definitions and lemmas that will be used in the proof of Theorem 5. We only state the next definitions and lemmas and refer the reader to Wainwright (2019) and Vershynin (2018) for proof and more details.
Definition 27. (Wainwright, 2019, Definitions 2.2 and 2.7) A real-valued random variable is said to u-sub-Gaussin if for all 8 R we have E[exp[s=)] S exp(s B[z] [z]+ + s262 ). and is said to be u-sub-ewt u-sub-exxonene exponential for all [s] S 1 we have Escxp(s2)] S exp(s E[2) +
Lemma 28. (Vershynin, 2018, Propositions 2.5.2 and 2.7.1) Suppose = is a zero-mean random variable and [||=||| T. for all p 2 1. Then = is cl-sub-GGussie for absolute constant O, i.e. E[exp(sz)] S exp(="ee 2 ) for all R. Similarly, suppose E||2|1] I Lp. Then is eJ-subb-xponential for an absolute constant e. > 0.
Lemma 29. (Wainwright, 2019, Theorem 2.26) Let x (o) If Rd - R is L-Liscchitz then f(x) is sub-Gaussian with parameter L.
Lemma 30. (Wainwright, 2019, Example 6.3) Let (x.1))ssss Jisisn be a sequence of i.i.d. standard Gaussian random vectors (0, Ia). It holds with probability at least 1 ô that
Σ
Tall2 s c cc(
log(1/5) d + log(1/8) n
n
where C an absolute constant.
The next lemma is the well-known symmetrization argument that upper bounds the expected value of an empirical process with Rademacher complexity.
Lemma 31. (Mohri et al., 2018, Theorem 3.3) Let F be a class functions f RR - lR for some p O. For number of samples T and probability distribution P IR, define the Rademacher complexity of F as
T-1 Σ 1=
R(F) = E
sup fEF
((f(()))))
(E.1)
where (x(1) lid P and ( are independent Rademacher random variables (i.e. #1 M2D equiprobably). Then the following holds,
E
sup
E(ad)) -E[(())]]] <
2R(FF.
JEF
Furthermore, we have the following fact for standard normal random vectors.
36

---

Published as a conference paper at ICLR 2023
Lemma 32. Let (O, Ia). There exists an absolute constant C such that for any V E Rm and p 2 we have
E V
+
Proof. First of all, VV is a V -l-pssccii function of x, thus Lemma 29 applies and | |Vll2 sub-Gaussian. Furthermore, by applying Lemma 28 to V ||]] Væ [|]] and Minkowski's inequality, we have
C 2 E V 2:
B V
E V
|
where the last inequality follows from Jensen's is inequality.
Lemma 33. Let a (N) Then E [exp(c||<(13) I exp(2cd) for S 1/4.
Proof. Gaussian integration yields E E(exp(caat) v1-2c S exp((2c).
Furthermore, for S 1 have
Vc-2c
37


Reviews:
[{"review_id": "kIEeMDh9q_", "reviewer": "Reviewer_vyfX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_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": "This paper  proves  that SGD on neural networks can learn low-dimensional features in certain settings, and uses this to derive novel generalization and excess risk bounds.", "summary_of_the_review": "A nice work. And it can be better if  the author will study the nonparametric estimation error. ", "strength_and_weaknesses": "\nThis paper study the convergence of SGD (applied  to two layer networks)  to the principle subspaces of the teacher model.\nStrength:\nStudying the convergence of SGD sequence to the estimation target (not only to the minimizer at the optimization level) is an important perspective to understand SGD in deep learning.\nSince the the target model is a multi-index model which is easier to due to the low dimensional intrinsics structure. So it is not surprise the final   results   reducing  the curse of dimensionality.\n\nWeakness:  In the perspective of nonparametric regression, we also need bound the estimation error, i.e., the estimated network function and \ntarget function g.  But the authors never mentioned  about this.   See for example  \n1:  I. Kuzborskij and Cs. Szepesvári. Nonparametric regression with shallow overparameterized neural networks trained by gd with early stopping. COLT, 2021.\n2:  T. Hu, W. Wang, C. Lin, and G. Cheng. Regularization matters: A nonparametric perspective\non overparametrized neural network. AISTATS, 2021.\n3: S Frei, NS Chatterji, PL Bartlett， Benign Overfitting without Linearity: Neural Network Classifiers Trained by Gradient Descent for Noisy Linear Data， COLT,  2022.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Fhn7K_gcJh6", "reviewer": "Reviewer_qmRw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_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": "This paper studies layer-wise training by SGD for two-layer neural networks under teacher-student setups. Specifically, for the multiple-index models as teachers, the paper shows the convergence to the principal subspace spanned by the true parameters of the teacher. Moreover, several applications of this result are provided: (i) the learnability of the single-index model, (ii) an efficient generalization error bound, and (iii) compressibility guarantees.", "summary_of_the_review": "Although there are still several limitations in the theory, this paper certainly contributes to the context by providing an improvement over the existing studies.", "strength_and_weaknesses": "**Strengths**:\n\nThe feature learning property of neural networks trained with SGD has become a challenging research interest. Recently, many studies have been devoted to this line of research, but there are still several limitations. In this sense, the paper certainly contributes to the context by studying the SGD with multiple steps.\n\n- This paper provides an efficient sample complexity linear in the input dimension for learning the single-index model, which improves the existing results in related studies which treat one-step gradient descent. This improvement would be nice.\n\n- Moreover, the connection with an efficient generalization bound (Theorem 5) and the compressibility guarantee (Theorem 6) are also interesting, although these results are straightforward applications of Theorem 3.\n\n**Weaknesses**:\n\n- For the case $k>1$, the theory (Theorem 3) does not guarantee the convergence to the true parameter. In other words, the optimization accuracy is unclear in general, and thus the theory cannot estimate how small the expected risk can be.\n- A learnable class (i.e., $k=1$ and $f$ is monotone) is still rather limited compared to [Abbe et al. (2022)] which shows the learnability of staircase functions.\n\nHowever, I also acknowledge the difficulty of overcoming these limitations.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "8KC7TvKk3n", "reviewer": "Reviewer_k8J3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_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": "The authors consider the representation learning ability of two layer neural\nnetworks trained with gradient descent to approximate the outputs of a function\nin a nonparametric class, applied to a few projections of a high-dimensional\ninput vector -- in this setting, kernel methods and other rotation-invariant\npredictors generally suffer from a curse of dimensionality, but methods such as\nneural networks can potentially avoid the curse of dimensionality by learning\nthe relevant directions and then fitting the target in this lower-dimensional\nspace. They consider the setting of gaussian inputs, standard (mean-field-type)\nrandom initialization, and training with variants of stochastic gradient\ndescent on the square loss and smooth activations, and on convex Lipschitz losses\nand the ReLU, with large weight decay regularization. Their main technical result\nconsists of showing that online SGD (samples from the population loss) on the\nfirst layer weights leads to fairly rapid convergence to weights that are\nsuitably close to the span of the relevant $k$ directions of the teacher model\n-- the rates for the norm of the irrelevant components of the trained weights\nare on the order $\\sqrt{dm/T}$, where $d$ is the ambient dimension, $m$ is the\nnumber of neurons, and $T$ is the number of iterations of SGD, which is\nsufficient to bound the risk given that the losses are locally Lipschitz. The\nauthors apply this result to obtain a guarantee for learning noisy single-index\nteacher models (single neuron) up to the noise level with a two-stage procedure\n(after SGD on the first-layer weights to learn the right representation for the\ninput, fit the readout layer weights with essentially ridge regression), and\ngeneralization gap control (no empirical risk control) in more general settings\ndue to the low-dimensional representation learned. \n\n", "summary_of_the_review": "Similar to other quasi-concurrent works, this work presents an interesting\nmathematical study of a single level of representation learning in nonlinear\nneural networks trained in function approximation-type tasks with imperfectly\nspecified models (beyond the typical teacher-student). The results here are\nsimilar to those in concurrent works, but the excellent presentation makes this\nwork a valuable entry point to this evolving topic in the literature.\n\n", "strength_and_weaknesses": "## Strengths\n\n- The paper is very well written. The coverage of background material and\n  related work seems essentially exhaustive. The notation and technical\n  discussions are precise and easy to parse at multiple levels of detail. The\n  authors provide interesting intuitions and connections to classical work in\n  statistics (e.g.  model misspecification) that enhances the level of insight\n  of the presentation.\n\n- The results concern the important problem of representation learning in\n  nonlinear neural networks. The authors seem to have taken some care to\n  present their results in some generality (different losses and activations,\n  step size regimes for SGD, etc.). The conclusions for learning the first\n  layer weights are qualitatively interesting.\n\n## Weaknesses\n\n- The algorithmic results for neural network training + generalization are\n  relatively weaker than the results for recovery of the first layer weights --\n  just concerning single index models or the generalization gap.\n\n- It would be helpful for the sake of comparison if the theorems did not hide\n  properties of the link $g$ -- especially to compare with the results of\n  Damian et al., which expose specific rates when the target is a polynomial\n  with a certain neural-net-like structure.\n\n## Minor / Questions / Etc.\n\n- The role of the noise $\\epsilon$ in the link function does not seem clear in\n  the theorems in the main body, prior to section 4 -- what dependence does it\n  play in the results?  Can it be zero? In section 4, why is this noise\n  necessary? (It would be helpful in understanding this if Theorem 4 exposed\n  the dependence on the subgaussian rate of the noise $\\sigma^2$, say.)\n\n- Some discussion of why the risk needs to be truncated at level $\\tau$ in\n  section 4 would be helpful.\n\n- With regards to the similarity to previous works, some discussion in the\n  rebuttal might be helpful to understand the distinct insights in this work\n  (especially, on the technical side, the differences that the authors allude\n  to in the related work section between analysis of the SGD trajectory, as\n  here, with the one-step gradient trajectory, given that both are unnatural\n  two-stage (\"layerwise training\") algorithms).\n\n- The claim at the start of section 4.1, suggesting that a key aspect of the\n  neural network training/generalization problem revolves around coping with\n  model misspecification (positing the existence of such an underlying model),\n  is interesting and thought-provoking. Another line of theoretical work on\n  neural nets focuses on the ability of neural networks to learn to adapt to\n  specific low-dimensional structures in data itself -- e.g. in terms of\n  representation capacity [1], or algorithmic aspects [2-5]. Adapting to this\n  type of structure, and the analysis it entails, seems somewhat different from\n  what the authors study here, and I would be curious to hear the thoughts of\n  the authors on whether there are similarities, whether they are in tension,\n  etc.\n\n[1] http://dx.doi.org/10.1093/imaiai/iaac001\n\n[2] https://proceedings.neurips.cc/paper/2020/file/a9df2255ad642b923d95503b9a7958d8-Paper.pdf\n\n[3] https://openreview.net/forum?id=O-6Pm_d_Q-\n\n[4] https://papers.nips.cc/paper/2021/hash/f26df67e8110ee2b44923db775e3e47f-Abstract.html\n\n[5] http://arxiv.org/abs/2206.12314\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%