Provable Guarantees for Neural Networks via Gradient Feature Learning
Authors: Zhenmei Shi
Junyi Wei
Yingyu Liang
Venue: NeurIPS 2023 poster (2023)
Keywords: neural networks
gradient descent
feature learning
provable guarantees
theoretical analysis
TL;DR: This work proposes a unified analysis framework for two-layer networks trained by gradient descent, which is centered around the principle of feature learning from gradients.

Abstract: Neural networks have achieved remarkable empirical performance, while the current theoretical analysis is not adequate for understanding their success, e.g., the Neural Tangent Kernel approach fails to capture their key feature learning ability, while recent analyses on feature learning are typically problem-specific. This work proposes a unified analysis framework for two-layer networks trained by gradient descent. The framework is centered around the principle of feature learning from gradients, and its effectiveness is demonstrated by applications in several prototypical problems, such as mixtures of Gaussians and parity functions.
The framework also sheds light on interesting network learning phenomena such as feature learning beyond kernels and the lottery ticket hypothesis.

Provable Guarantees for Neural Networks via Gradient Feature Learning
Zhenmei Shi', Junyi Wei", Yingyu Liang
University of Wisconsin, Madison
zhmeishi wisc wisc.edu, jeei3wwiss. htt.i/3wwiic..d..333.. wisc. edu
Abstract
Neural networks have achieved remarkable empirical performance, while the current theoretical analysis not adequate for understanding their success, c.g., the Neural Tangent Kernel approach fails to capture their key feature learning ability, while recent analyses on feature learning are typically problem-specific. This work proposes unified analysis framework for two-layer networks trained by gradient descent. The framework is centered around the principle of feature learning from gradients, and its effectiveness is demonstrated by applications in several prototypical problems such as mixtures of Gaussians and parity functions. The framework also sheds light on interesting network learning phenomena such as feature learning beyond kernels and the lottery ticket hypothesis.
Introduction
Neural network (NN) learning has achieved remarkable empirical success and has been a main driving force for the recent progress in machine learning and artificial intelligence. On the other hand, theoretical understandings significantly lag behind. Traditional analysis approaches are not adequate due the overparameterization of practical networks and the non-convex optimization in the training via gradient descent. One line work (e.g. [9, 31, 38, 60, 71, 123] and many others) shows under proper conditions, heavily overparameterizedt networks data-independent features, i.e., a linear function on the Neural Tangent Kernel (NTK), While making weak assumptions about the data and thus applicable to various settings, this approach requires the network learning to be approximately using fixed data-independent features (i.c., the kernel regime, or fixed feature methods). It thus fails to capture the feature learning ability of networks (i.e., to learn feature mapping for the inputs which allow accurate prediction). which is widely believed to be the key factor their empirical success many applications (c. g.. [54, 77, 117, 119)). To study feature learning in networks, recent line of work (c.g. [5, 6, 14, 33, 52, 72, 76, 116] and others) shows examples where networks provably enjoy advantages over fixed feature methods (including NT'K), under different settings and assumptions. While providing more insights, these studies typically focus on specific problems, and their analyses exploit the specific properties of the problems and appear to be unrelated to each other. Is there common principle for feature learning in networks via gradient descent? there a unified analysis framework that can clarify the principle and also lead to provable error guarantees for prototypical problem settings?
approximately linear models over
In this work, we take a step toward this goal by proposing a gradient feature learning framework for analyzing two-layer network learning by gradient descent. (1) The framework makes essentially no assumption about the data distribution and can be applied to various problems. Furthermore. it is centered around features from gradients, clearly illustrating how gradient descent leads to feature learning in networks and subsequently accurate predictions. competitive with the optimal in family of networks that use the features induced by gradients on the
(2) It leads to error guarantees
"Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

data distribution. Then for a specific problem with structured data distributions, if the optimal in the induced family is small, the framework gives a small error guarantee.
We then apply the framework to several prototypical problems: mixtures of Gaussians, parity functions, linear data, and multiple-index models. These have been used for studying network learning (in particular, for the feature learning ability), but with different and seemingly unrelated analyses. In contrast, straightforward applications of our framework give small error guarantees, where the main effort is to compute the optimal in the induced family. Furthermore, in some cases, such as parities, we can handle more general data distributions than in the existing work.
Finally, we also demonstrate that the framework sheds light on several interesting network learning phenomena or implications such as feature learning beyond the kernel regime, lottery ticket hypothesis simplicity bias, learning over different data distributions, and new perspectives about roadmaps forward. Due to space limitations, we present implications about features beyond the kernel regime and LTH in the main body but defer the other implications in Appendix C with a brief here. (1) For simplicity bias, it is generally believed that the optimization has some implicit regularization effect that restricts learning dynamics to a low capacity subset of the whole hypothesis class, so can lead good generalization [53, 90]. Our framework provides an explanation that the learning first learns simpler functions and then more sophisticated ones. (2) For learning over different data distributions, we provide data-dependent non-vacuous guarantees, as our framework can be viewed as using the optimal prudient-induced NN to measure or quantify the "complexity" of the problem. For easier problems, this quantity is smaller, and our framework can give a better error bound to derive guarantees. (3) For new perspectives about roadmaps forward, our framework suggests the strong representation power of NN actually the key to successful learning, while traditional ones suggest strong representation power leads to vacuous generalization bounds [19, 33]. Thus, we suggest a different analysis road. Traditional analysis typically first reasons about the optimal based on the whole function class then analyzes how NN learns proper features and reaches the optimal. In contrast, our framework defines feature family first, and then reasons about the optimal based on it.
2 Related Work
Neural Networks Learning Analysis. Recently there has been an increasing interest in the analysis of network learning. One line of work connects the sufficiently over- wer-ppaaamterizee neural network to linear methods around its initialization like NTK (e.g. [9, 11, 20, 21, 31, 38, 49, 60, 62, 69, 71, 78, 82, 91, 93, 95, 1 114, 121, 122] and more), so that the neural network training is a convex problem. The key idea is that it suffices to consider the first-order Tyler expansion of the neural network around the origin when the initialization is large enough. However, NTK lies in the lazy training (kernel) regime that excludes feature learning [29, 50, 68, 113]. Many studies (e.g. [2, 5, 8, 12, 14, 22, 26, 33, 37, 51, 52, 57, 58, 70,7 72, 73, 76, 99, 112, 115, 116] and more) show that neural networks take advantage over NTK empirically and theoretically. Another line of work is the mean-field (MF) analysis of neural networks (e.g. [27, 28, 36, 79, 80, 100, 106] and more). The insight is to see the training dynamics of a sufficiently large-width neural network as PDE. It uses a smaller initialization than the NTK so that the parameters may move away from the initialization. However, the MF does not provide explicit convergence rates and requires an unrealistically large width of the neural network. One more line of work is neural networks max- margin analysis (e.g. [30, 47, 48, 56, 61, 63, 74, 75, 83, 85, 86, 107, 109] and more). They need a strong assumption that the convergence starts from weights having perfect training accuracy, while feature learning happens in the early stage of training. To explain the success of neural networks beyond the limitation mentioned above, some work introduces the low intrinsic dimension of data distributions [17, 18, 23, 24, 25, 44, 67, 104, 108, 1 124]. Another recent line of work is that a trained network can exactly recover the ground truth or optimal solution or teacher network [3, 4, 10, 39, 84, 87, 94, 96, 120]. but they have strong assumptions
data distribution or model structure, c.g., Gaussian marginals. [1, 40, 55, 1 110, 111] show that training dynamics of neural networks have multiple phases, e.g., feature learning at the beginning, and then dynamics in convex optimization which requires proxy convexity [43] or PL condition [65] or special data structure.
Feature Learning Based on Gradient Analysis. A recent line of work is studying how features emerge from the gradient. [7, 46] consider linear separable data and show that the first few gradient steps can learn good features, and the later steps learn a good network on neurons with these features. [33, 45, 105] have similar conclusions on non-linear data (e.g., parity functions), while in their problems one feature is sufficient for accurate prediction (i.c., single-index data model).
2

---

[32] considers multiple-indee index with low-degree polynomials as labeling functions and shows that one-step gradient update can learn multiple features that lead accurate prediction. [13, 81] studies one gradient step feature improvements at different learning rates. [97] proposes Recursive Feature Machines show the mechanism of recursively feature learning but without giving final loss guarantee. These studies consider specific problems and exploit properties of the data to analyze the gradient delicately, while our work provides a general framework applicable to different problems.
3 Gradient Feature Learning Framework
Problem Setup. We denote [r]] := {1, 2, .......7 n2 and Ō(+), ((),, (2()) omit the log term inside. Let X C Rd denote the input space, J/ R the label space. Let D be arbitrary data distribution over X' x V. Denote the class of two-layer networks with m neurons as:
Fd.m = (.a...b) ((a.w.)))) = a [o(WT x b)]
Σ a [a(oo((; bb)] (I)
xE(m]
where (2) max(z,0) 0) is the ReLU activation function, E IR is the second layer weight, W E Rddxm is the first layer weight, w, is the i-th column of W (i.e., the weight for the i-th neuron), and b E Rmmh is the bias for the neurons. For technical simplicity, we only train a, W but not b. Let superscript (t) denote the time step, e.g., f(at), w().b) denote the network time step f. Denote E = (a. W, b), =(0 = (9, w(!), b). The goal of neural network learning is to minimize the expected risk, i.e., CD(f)=
(xxxx (J); where L(x.x) (f) (x)(()) is the loss an example (x. 3) for some loss function E(-), e.g., the hinge loss {(z) = max 1-1-2) 2}, and the logistic loss {(z) = log(1 + exp(-2)] We also consider & regularization. The regularized loss with regularization coefficient/ À is (f) == (f( + [| + ||al Given a training set with n from D, the empirical risk and its regularized version are:
i.i.d. samples 2 = ((0)
[z( 1 Σ f wa)
== 2 + ||alB))
([W))
(2)
[E(x]
Then the training process is summarized in Algorithm 1.
Algorithm 1 Network Training via Gradient Descent
Initialize (g) W(O),b) b) for t= 1 to do Sample Z(t- D
a! a 1) 17 (t)
2ic-1 (fze-1), w(1) w(- (0) ( Vw 220-1 W (fe(-n)
end for
In the whole paper, we need some natural assumptions about the data and the loss.
Assumption 3.1. We assume [[[[||||2 S. Bzi, ||] < Br2. ||x||2 S B2 and for any label
we have lyl < 1. We assume the loss function ((-) is I-Lipschitz convex decreasing function, normalized E(0) (0) @(1), and E(co) = O.
Remark 3.2. The above are natural assumptions. Møst input distributions have the bounded norms required, and the typical binary classification y = {=1} satisfies the requirement. Also, the most popular loss functions satisfy the assumption, e.g., the hinge loss and logistic loss.
3.1 Warm Up: Simple Setting with Frozen First Layer
To illustrate some high-level intuition, we first consider a simple setting where the first layer is frozen after one gradient update, i.e., no updates W for 2 2 in Algorithm
The first idea our framework is to provide guarantees compared to the optimal in a family of networks. Here let consider networks with specific weights for the first layer:
Definition 3.3. For some fixed W E Rdxm b Rd, and parameter Baz- consider the following family of networks FW.,, B. and the optimal approximation network loss in this family:
Fw.b = {f(,..b) Famml ||l|| < Ba2),
OPTw.b.B. JeFw.b.922
min
LD (3)
3

---

The second idea is to compare to networks using features from gradient descent. As an illustrative example, we now provide guarantees compared to networks with first layer weights W (i.e., the weights after the first gradient step):
Theorem 3.4 (Simple Setting). Assume Lz (a),))))))) is L-smooth a. Let
= O. for all E (2.3....,I7) Training by Algorithm I with no updates for the first layer after the first gradient step, w.h.p., there exists t E [T) such that
nf(0== = ,, A()
CD( I OPT + L(Ila 13++BB)
B2, W()) B2 Ibb
T
Intuitively, the theorem shows that the weight w()) after one-step gradient gives good set of neurons in the sense that there exists classifier on top of these neurons with low loss, then the network will learn to approximate this good classifier and achieve low loss. The proof is based on standard convex optimization and the Rademacher complexity (details in Appendix D.I).
Such approach, while simple, has been used to obtain interesting results network learning in existing work, which shows that W can indeed give good neurons due to the structure of the special problems considered (e.g., parities on uniform inputs [15], or polynomials on a subspace [32]). However, it is unclear whether such intuition still yield useful guarantees for other problems. So, for our purpose of building a general framework covering more prototypical problems, the challenge is what features from gradient descent should be considered so that the family of networks for comparison can achieve a low loss on other problems. The other challenge is that we would like consider the typical case where the first layer weights are not frozen. In the following, will introduce the core concept of Gradient Features to address the first challenge, and stipulate proper geometric properties of Gradient Features for the second challenge.
3.2 Core Concepts in the Gradient Feature Learning Framework
Now, we will introduce the core concept in our framework, Gradient Features, and use it to build the family of net- works to derive guarantees. As mentioned, we consider the setting where the first layer is not frozen. After the network learns good features, to ensure the updates later gradient steps of the first layer are still benign for feature learning, we need some geometric conditions about the gradient features, which are measured by parameters in the definition of Gradient Features. The conditions are general enough, so that, shown in Section 4, many prototypical problems satisfy them and the induced family of networks enjoys low loss, leading useful guarantees. We begin by considering what features can be learned via gradients. Note that the gradient w.r.t. W; is
Gradient Festure being under Mixture Gaussians data
DLD(J) a, (x,)) [E' (C(y (xf(x)) Jo ((w,,x) bb))] Ow:
a, (x.,) [L" (x))xxXXX Inspired by this, we define the following notion: Definition 3.5 (Simplified Gradient Vector). For any E R4, b E R. a Simplified Gradient Vector is
> bi]]-
An illustration of Gradient Fea- ture, i.c., Definition 3.7 with random initial- ization (Gaussian), under Mixture of three Gaussian clusters in 3-dimension data space with blue/green/orange color. The Gradient Feature stays in three cones. where each cen- ter of the cone aligns with the corresponding Gaussian cluster center.
Figure 1:
G(w, b) = E(x.))~Dlyxxll T x b]].
(4)
Remark 3.6. Note that the definition of G(w, b) ignores the term l' (x)) in the gradient, where is the model function. In the early stage of training (or the first gradient step), (((-) is approximately constant, i.e., (x)](x)) ss (0) due to the symmetric initialization (see Equation (8)).
Definition 3.7 (Gradient Feature). For a unit vector D E IRd with |D = 1, and a a E (0, 1), direction neighborhood (cone) CD,y is defined as:
CD := (w (w, (w, 1/ ///l//////////// (1-771 7))
(5)

---

Let w E RR, E R be random variables drawn from some distribution W, B. A Gradient Feature set with parameters p. y, Bg is defined as:
Bp (W, B) = {D)s))
Pr w.b
(r][ CD.,. |G(w,, b)ll2 BG 8 -=b//b]]] 2 p).
(6)
Remark 3.8. When clear from context, write it as Bp Gradient features (see Figure 1 for illustration) are simply normalized vectors D that are given (approximately) by the simplified gradient vectors. (Similarly, the normalized scalar is given by the bias b.) To be a usefiul gradient feature we require the direction to be "hit" by sufficiently large simplified gradient vectors with sufficient large probability, so as to be distinguished from noise and remain useful throughout the gradient steps. Later we will use the gradient features when W, B are the initialization distributions.
To make use of the gradient features, we consider the following family of networks using these features and with bounded norms, and will provide guarantees compared to the best in this family: Definition 3.9 (Gradient Feature Induced Networks). The Grudient Feature Induced Networks are:
S = (f(..W.)) Famm Vi E [m], | S Bai: ||ll| S Baz, bw,,bl//bll) S, |b, Bb), where S is some Gradient Feature set and BF (Bat, BB2, Bb) are some parameters.
Remark 3.10. In above definition, the weight and bias of a neuron are simply the scalings of some item in the feature S (for simplicity the scaling of w; is absorbed into the scaling of a, and bı). Definition 3.11 (Optimal Approximation via Gradient Features). The optimal approximation network and loss using Gradient Feature Induced Networks Farr.Br.s defined as:
f" argmin CD(f),
Or. ==
min
CD(f).
fEFa.r.Bp.
fEFan.g.
3.3 Provable Guarantee via Gradient Feature Learning
To obtain the guarantees, we first specify the symmetric initialization. It is convenient for the analysis and is typical in existing analysis (c.g., [7, 32, 33, 105]), though some other initialization can also work. Formally, we train a two-layer network with 4m neurons, (aa, E Fd,am- We initialize (0) (0) from Gaussians and b; from constant for i E {l....,m m). and initialize the parameters for E [m I,..., ..43 accordingly to get a zero output initial network. Specifically:
a
for i E {1,.... m) a
W; (0)
N0000011 o. b, b,
for iric E {m + 1, 2m3
i-m b, = -bi-m.
(8)
for i E (2m + 1,.... 4m}
a,
(0)
2m
(0)
2m b; = Di-2m:
where 22 ou, 5 > are hyper- aapeer-aramttrss.. After initialization, a, W are updated as in Algorithm 1. We are now ready to present our main result in the framework.
Theorem 3.12 (Main Result), Assume Assumption 3.1. For any e,8 (0, 1), if m S ed and
1
Bb
1
1
log (7) (7))²
=Ω
Det
rB.1 BBaBBal Bal
V
Be
Vo
p
T =Ω (-]
(ViBozBoBu
+ mbb)
Vlg m v BBB
(mp)
Bu (mp)
n
m²p
Br Br
B
B²
B
Tm
+
+
-
+
log nl
2r2 B2, 12211 BG
B,
Buz
B
B2,
[((O)]=
then with initialization (8) and proper apper-parameter values, we have with probability N 1 - over the initialization and training samples, there exists E [T] in Algorithm with:
Pr[sign(Sz (x)) + v) < LD (JZN)
VBaa logn Bcle
S OPT S.Pddr.Br.S.. dr + +r r B 2y + O
+
5

---

Intuitively. the theorem shows when data distribution admits a small approximation error by some "ground-truth network with neurons using gradient features from Sp..Bc (i.e., a small optimal approximate loss OPT B.r,Bp.S,.s..e the gradient descent training can successfully learn good neural networks with sufficiently many m neurons.
Now we discuss the requirements and the error guarantee. Viewing boundedness parameters Bar Bar 1
pe a polynomial overpa- rameterization compared to the "ground-truth" network. The proof shows that such an overparam- eterization is needed such that some neurons can capture the gradient features given by gradient descent. This consistent with existing analysis about overparameterization network learning, and also consistent with existing empirical observations.
etc. as constants, then the number m of neurons learned is roughly é
The error bound consists of three terms. The last term e can be made arbitrarily small, while the other two depend on the concrete data distribution. Specifically, with larger r and y, the second term increases. While the first term (the optimal approximation loss) decreases, since larger r means a larger "ground-trutte So. there is a trade-off between these two terms. When we later apply the framework to concrete problems (e.g.. mixtures of Gaussians, parity functions), we will show that depending on the specific data distribution, we can choose the proper values for r to make the error small. This then leads to error guarantees for the concrete problems and demonstrates the unifying power of the framework. Please refer to Appendix D.3 for more discussion about our problem setup and our core concept, e.g.. parameter choice, early stopping, the role of 8, activation functions, and so on.
network family, and a larger 7 means a larger Gradient Feature set Sp.
Proof Sketch. The intuition in the proof of Theorem 3.12 is closely related to the notion of Gradient Features. First, the gradient descent will produce gradients that approximate the features in Ba Then, the gradient descent update gives a good set of neurons, such that there exists an accurate classifier using these neurons with loss comparable to the optimal approximation loss. Finally, the training will learn to approximate the accurate classifier, resulting in the desired error guarantee. The complete proof is in Appendix D) (the population version in Appendix D.2 and the empirical version in Appendix D.4), including the proper values for hyper-parameter such as r/(() in Theorem D.17. Below, we briefly sketch the key ideas and omit the technical details.
We first show that large subset of neurons has gradients at the first step as good features. (The claim can be extended to multiple steps; for simplicity, we follow existing work (c.g.. [33, 105]) and present only the first step.) Let V, denote the gradient of the i-th neuron Vw, ( ). Denote the subset of neurons with nice gradients approximating feature (D,8) as:
G(D,s).Nice : f [2m] : s b./lb.), (V., D) (1 - 7) IV || V. 2 a
(0) Ba]. (9)
Lemma 3.13 (Feature Emergence). For any size subset {(D,, S1),.., ... [D,,,)) S Sp.yBa. with probability at least - re" -m(m)), for all j E [r]. we have (G 2 4
mp
This is because V, (0)aa (0) B(x.v) Tyo' (( (w:(0) (0) .x) b, x (())aa (0) G(w (0) b,). Now
consider sj = +1 (the case 1 is similar). Since w: is initialized by Gaussians, by V,'s connection
to Gradient Features, we can see that for all i E [m]. Pr [iE G(D, Nion 2 2. The lemma
follows from concentration via a large enough m, i.e., sufficient overparameterrvzaton. The gradients allow obtaining a set of neurons approximating the "ground-truth" network with comparable loss:
Lemma 3.14 (Existence of Good Networks). For any 5 E (0, 1). with proper
values, with probability at least 1
8, there is ā such that ||āloo O (r (mm) and (1..w.),)(x
4m ā,o
((w!) xx - -b:) satisfies
(w! (1)
2Br VmpBe
Coll(aw(u..)) S
+ B2i
(V5+ + V
Given the good set of neurons, we finally show that the remaining gradient steps can learn an accurate classifier. Intuitively, with small step sizes r/19), the weights of the first layer w; do not change too much (stay in a neighborhood) while the second layer weights grow, and thus the learning is similar to convex learning using the good set of neurons. Technically, we adopt the online convex optimization analysis (Theorem D.5) in [33] to get the final loss guarantec in Theorem 3..

---

4
Applications in Special Cases
In this section will apply the gradient feature learning framework some specific problems, corresponding to concrete data distributions D. We primarily focus on prototypical problems for analyzing feature learning in networks. We will present here the results for mixtures of Gaussians and parity functions, and include the complete proofs and some other results in Appendix E.
4.1 Mixtures Gaussians
Mixtures of Gaussians are among the most fundamental and widely used statistical models. Recently, has been used to study neural network learning, in particular, the effect of gradient descent for feature learning of two-layer neural networks and the advantage over fixed feature methods [46, 99].
Data Distributions. We follow notations from [99]. The data are from mixture of r high- dimensional Gaussians, and each Gaussian is assigned to one of two possible labels in y - {+1}. Let Syy) S [r] denote the of indices of Gaussians associated with the label y. The data distribution is then: q(x y) (x)) y). g(x y) = EjEs(y) (x)N((x). where Ng (x) is a multivariate normal distribution with mean My, covariance Σ,. and py are chosen such that (xx, y) is correctly normalized. We will make some assumptions about the Gaussians, for which we first introduce some notations.
Dj
Gg := my:=M//V M//Vd, Byn min M|a sc(r]
Bya == max | [ل5 Pn == min Pi- jeir
|
jE(r]
Assumption 4.1. & ST S be a parameter that will control our final error guarantee. Assume Equiprobable labels: ((-1) g(+1)= 1/2.
For all j E [r]. Σj oj Fjlaxd: Let оB =- mäxjelr] Jj and 0B ;= max B
S 2d,
pn 20 Ω (1/d VTan+ 2ogg //g Br B22 S d.
The Gaussians are well-separated: for all + E [r], we have - -1 (D, where
0 0 < min 2r
log d
B,2
Remark 4.2. The first two assumptions are for simplicity; they can be relaxed. We can generalize our analysis to the mixture of Gaussians with unbalanced label probabilities and general covariances. The third assumption is to make sure that each Gaussian has a good amount of probability mass to be learned. The remaining assumptions are to make sure that the Gaussians are well rell-separated and can be distinguished by the learning algorithm.
We are now ready to apply the framework to these data distributions, for which we only need to compute the Gradient Feature set and the corresponding optimal approximation loss.
Lemma 4.3 (Mixtures of Gaussians: Gradient Features).
(Dj, +1) Sp for all j E [r]. where
Byd
BG PBB V
oBt
V log do 0 To 2/B2,)
Y
1..5
(u.97
Let (*(x) EE=1 VT log
Min
[o (D,,x)
2 log B-)- whose hinge loss is most d
10.3r- Vilgad
Given the values gradient feature parameters p, 7, BG and the optimal approximation loss OPT Bp.... the framework immediately leads to the following guarantee:
Theorem 4.4 (Mixtures of Gaussians: Main Result). Assume Assumption 4.1. For any ε, E (0, 1), when Algorithm I uses hinge loss with
polo poly 1 ((Tog++//1p) 2/B,,,, 1 -
se,, )
T= poly (m), p (m)
PB
and proper hyper-paramnttes, then with probability at least 1 - 8, there exists E [T] such that
/2r
Pr[sign(ffm (x)) + yl 60.4- 0.8 +
7

---

The theorem shows that gradient descent can learn to a small error via learning the gradient features, given proper ppprr-paraettes. In particular. we need sufficient overparameterization (a sufficiently large number m of neurons). When ØB 2 /B,, 1 is constant which is the prototypical interesting case, and we choose a constant т, then m is polynomial in the key parameters t, 2 d, d, pæ and the error bound inverse polynomial in d. The complete proof is given in Appendix E.2.
[46] studies (almost) linear separable cases while our setting includes non- non-linear separable cases, c.g., XOR. [99] mainly studies neural network classification 4 Gaussian clusters with XOR structured labels, while our setting is much more general, e.g., our cluster number can extend up to 2d.
4.1.1 Mixtures of Gaussians: Beyond the Kernel Regime
As discussed in the introduction, it is important for the analysis to go beyond fixed feature methods such as NTK (i.e., the kernel regime), so as to capture the feature learning ability which is believed to be the key factor for the empirical success. We first review the fixed feature methods. Following [33], suppose is a data-independent feature mapping of dimension N with bounded features, i.e., L : X' + 1- 1] For B O, the family of linear models on V with bounded norm B is H. = (h(x) h(x) = (V (x), w), ||wl|2 I B}. This can capture linear models on fixed finite- dimensional feature maps. c.g., NTK, and also infinite dimensional feature maps, e.g., kernels like RBF, that can be approximated by feature maps of polynomial dimensions [64, 98, 105].
Our framework indeed goes beyond fixed features and shows features from gradients are more powerful than features from random initialization, e-g., NTK. Our framework can show the advantage of network learning over kernel methods under the setting of [99] (4 Gaussian clusters with XOR structured labels). For large enough d, our framework only needs roughly 52 (log d)) neurons and Q (llog d)²) samples to achieve arbitrary small constant error (see Theorem E.18 when ØB = 1), while fixed feature methods need (dd)) features and (dd2) samples to achieve nontrivial errors (as proved in [99]). Moreover, [99] uses ODE to simulate the optimization process for the 2-layer networks learning XOR-shaped Gaussian mixture with (111) neurons and gives convincing evidence that R(d) samples is enough to learn it, yet they do not give a rigorous convergence guarantee for this problem. We successfully derive a convergence guarantee and we require a much smaller sample size Ω ((log d) For the proof (detailed in Appendix E.3), only need to calculate the p, 7, BG of the data distribution carefully and then inject these numbers into Theorem 3.12.
4.2 Parity Functions
Parity functions are canonical family of learning problems in computational learning theory, usually for showing theoretical computational barriers [103]. The typical sparse parties over d-dim binary inputs o E [=1] are ILKE Ø; where A S [d] is subset of dimensions. Recent studies have shown that when the distribution of inputs & has structures rather than uniform, neural networks can perform feature learning and finally learn parity functions with a small eπror, while methods without feature learning, e.g. NTK, cannot achieve as good results [33, 76, 105]. Thus, this has been a prototypical setting for studying feature learning phenomena in networks. Here we consider a generalization of this problem and show that our framework can show successful learning via gradient descent.
Data Distributions. Suppose M E RdxDD is an unknown dictionary with D columns that can be regarded as patterns. For simplicity, assume d = D and M is orthonormal. Let é E Rd be a hidden representation vector. Let A C [D] be a subset of size rk corresponding to the class relevant patterns and r is an odd number. Then the input is generated by MØ, and some function on @A generates the label. WLOG, let A = j e [r]. A₁ = { 1)k- + 1....,jkk. Then the input x and the class label y are given by:
A = {rk + I,... dd. Also, we split A such that for all
x = Mo. y - 9" (oA) = sign Σ XOR())))).
(10)
JE(n]
where 9" is the ground-truth labeling function mapping from R8* of ó with indices in A, and (OR()A,) IIIE Φi the parity function. We still need to specify the distribution X' of d. which determines the structure of the input distribution:
= y {+1}. A is the sub-vector
X == (1 2rpx)Xe + Σ Pa Xz,-)
(11)
5E[r]
8

---

+1,
w.p. Po
For all corresponding OΛL in X, we have 4I E A independently: 01
1, w.p. Po
0, w.p. 1 2po
where Po controls the signal noise ratio: if Po is large, then there are many nonzero entries in A which are noise interfering with the learning of the ground-true -ruth labeling function on A. For corresponding Ωs. any E [r]. we have
In Kj UA, =
+1 and @A only have zero elements. 1]" and 0,,4, only have zero elements.
In X, - 111,
In Ny, we have Φ, draw from {+1,- 1 uniformly.
In short, we have parity functions each corresponding block of k dimensions; X,,+ and x,,- stands for the component providing strong signal for the j-th parity; Xy corresponds to uniform distribution unrelated to any parity and providing weak learning signal; A the noise part. The label depends on the sum of the parity functions.
Assumption 4.5. Let 8 S T S be a parameter that will control our final error guarantee. Assume is an odd number and: k 2 S2(T logd),
d + log d),
Po =0=
rk 2-Tk
PA 4.
Remark 4.6. We set up the problem to be more general than the parity function learning in existing work. If = 1, the labeling function reduces to the traditional k-sparse parties of d bits. The provide enough large signals for learning. Note that when k T6 = L.Po = 1.2 our analysis also holds, which shows our framework beyond the kernel regime (discuss in detail Section 4.2.1).
assumptions require k, d. and P. to be sufficiently large
To apply our framework, again we only need to compute the Gradient Feature set and the correspond- Exe,, M.
ing optimal loss. We first define the Gradient Features: For all j E [r], let Dj = "TEii M.
Lemma 4.7 (Parity Functions: Gradient Features). We have (Dz, (D,3 -1) E Sp for all j E [r]. where
p=
( ( Arrlogd d doldd
7 = -
BG = VipA -
VT
(12)
dt
With features Sp.y.Bor let (etff(()) = Ej=1 E vTllo ((D)) 2im
2a (D)). - -i/p))) + (1D),x)
2/- (
whose hinge loss is 0.
Above, we show that Dj is the "indicator function" for the subset A₁ that we can build the optimal neural network based on such directions. Given the values on gradient feature parameters and the optimal approximation loss, the framework immediately leads the following guarantee:
Theorem 4.8 (Parity Functions: Main Result). Assume Assumption 4.5. For any e, E (0,1),
when Algorithm uses hinge loss with
= poly ()
1 d0(r) k,,, se) T= (m), n= = poly (m) =
1
and proper hyper-purameters, then with probability at least 1 - 8, there exists E [T] such that
3r +e.
(r(sign(/z(oxx)) # v) < d(7-3)/2
The theorem shows that gradient descent can learn to a small error in this problem. We also need sufficient overparameterization When r is a constant (c.g., r = 1 in existing work), and we choose constant 7, m polynomial in 1 x: d, k; pu and the error bound is inverse polynomial d. The proof in Appendix E,4. Our setting is more general than that [33, 76] which corresponds to Po = 3. [105] study single index learning, where one feature direction enough for a two-layer network to recover the label, while our setting considers r directions D D,...., D,, so the network needs to learn multiple directions to get small error.
M 1,r = 1,P4

---

4.2.1 Parity Functions: Beyond the Kernel Regime
Again, we show that our framework indeed goes beyond fixed features under parity functions. Our problem setting in Section 4.2 is general enough to include the problem setting in [33]. Their lower bound for fixed feature methods directly applies to case and leads to the following:
Proposition 4.9. There exists a data distribution in the parity learning setting in Section 4.2 with M 1, = 1, = 3 To Po
such that all h E HB have hinge-loss at least 2 2²'2"
/NB
This means to get hnverse-polynoning tially large size, i.c. either the number of features N or the norm B needs to be exponential in In contrast, Theorem 4.8 shows our framework guarantees a small loss with a polynomially large model, runtime, and sample complexity. Clearly, our framework is beyond the fixed feature methods.
small loss, fixed feature models need to have an exponen-
Parities on Uniform Inputs. When = P. O, our problem setting will degenerate to the classic sparse parity function on uniform input distribution. This has also been used for analyzing network learning [16]. For this case, our framework can get a k:20(k) log(k) network width bound and (d)) sample complexity bound, matching those in [16]. This then again confirms the advantage of network learning over kernel methods that requires dd(k) dimensions as shown in [16]. See the full statement in Theorem E.31, details in Appendix E.5, and alternative analysis in Appendix E.6.
Further Implications and Conclusion
Our general framework sheds light several interesting phenomena in NN learning observed in practice. Feature learning beyond the kernel regime has been discussed in Section 4.1.1 and Section 4.2.1. Here discuss the LTH and defer more implications such as simplicity bias, learning over different data distributions, and new perspectives about roadmaps forward in Appendix C.
Lottery Ticket Hypothesis (LTH). Another interesting phenomenon is the LTH [41]: randomly- initialized networks contain subnetworks that when trained in isolation reach test accuracy comparable to the original network in similar number of iterations. Later studies (e.g., [42]) show that LTH is more stable when subnetworks are found in the network after a few gradient steps.
Our framework provides an explanation for two-layer networks: the lottery ticket subnetwork contains exactly those neurons whose gradient feature approximates the weights of the "ground-truth" network f"; they may not exist at initialization but can be found after the first gradient step. More precisely, Lemma 14 shows that after the first gradient step, there is a sparse second-layer weight with ||āloo = O (rymp) such that using this weight on the hidden neurons gives a network with a small loss. Let U be the support of a. Equivalently, there is small-loss subnetwork f with only neurons in U and with second-layer weight ăy on these neurons. Following the same proof of Theorem 3.12:
Proposition 5.1. In the same selling of Theorem 3.12 but only considering the subnetwork supported on U after the first gradient step, with the same requirements on m and T, with proper hyper- parameter values, we have the same guarantee: with probability 2 5. there is E [T] with
Pr(sign(f(..) # #v S OPTd.,.Br.Sp... +rBaa Bri 27 O
(B.a (BB.2logm
log
+
BavT
This essentially formally proves LTH for two-layer networks, showing (a) the existence of the winning lottery subnetwork and (b) that gradient descent the subnetwork can learn to similar loss in similar runtime as on the whole network. In particular, (b) is novel and not analyzed in existing work.
We provide our work's broader impacts and limitations (e.g., statement of recovering existing results and some failure cases beyond our framework) in Appendix and Appendix respectively.
Conclusion. We propose general framework for analyzing two-layer neural network learning by gradient descent and show that can lead to provable guarantees for several prototypical problem settings for analyzing network learning. In particular, our framework goes beyond fixed feature methods, e.g., NTK. sheds light on several interesting phenomena NN learning, e-g., the lottery ticket hypothesis and simplicity bias. Future directions include: (1) How to extend the framework to deeper networks? (2) While the current framework focuses on the gradient features in the early gradient steps, whether feature learning also happens in later steps and if so how to formalize that?
10

---

Acknowledgements
The work is partially supported by Air Foree Grant FA9550-18-1-011 the National Science Founda- tion (NSF) Grants 2008559-I1S, 2023239- DMS, and CCF-2046710.
References
[1] 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. PMLR, 2022.
[2] Emmanuel Abbe, Samy Bengio, Elisabetta Cornacchia, Jon Kleinberg, Aryo Lotfi, Maithra Raghu, and Chiyuan Zhang. Learning reason with neural networks: Generalization, unseen data and boolean measures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[3] Shunta Akiyama and Taiji Suzuki. learnability via gradient method for two-layer relu neural networks in teacher-student setting. In International Conference on Machine Learning, pages 152-162. PMLR, 2021.
[4] Shunta Akiyama and Taiji Suzuki. Excess risk of two-layer reLU neural networks teacher- student settings and its superiority kernel methods. In The Eleventh International Conference on Learning Representations, 2023.
[5] Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? In Advances in Neural Information Processing Systems, 2019.
[6] Zeyuan Allen-Zhi and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413, 2020.
[7] Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pages 977-988. IEEE, 2022.
[8] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overpa- rameterized neural networks, going beyond two layers. In Advances in neural information processing systems, 2019.
[9] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameetrizaaion.. In International Conference on Machine Learning, 2019.
[10] Sanjeev Arora, Naday Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2018.
[11] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Leurning, pages 322-332. PMLR, 2019.
[12] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv: 1904. 2019.
[13] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. arXiv preprint 2rXiv:2205.01445, 2022.
[14] Yu Bai and Jason D Lec. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations, 2019,
[15] Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprint arXiv: 087999 2022.

---

[16] Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham M Kakade, Cyril Zhang, et al. Hidden progress in deep learning: Sgd learns parities near the computational limit. In Advances in Neural Information Processing Systems, 2022.
[17] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 2020.
[18] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 2022.
[19] Avrim Blum and Ronald L. Rivest. Training a 3-node neural network is np-complete. In Advances in neural information processing systems, pages 494-501, 1989.
[20] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information Processing Systems, 2019.
[21] Yuan Cao, Zhiying Fang, Yue Wu, Muag-Xunnnnaa Zhou, and Quanquan Gu. Towards understand- ing the spectral bias of deep learning. 2020.
[22] Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[23] Niladri S Chatterji, Philip M Long, and Peter L Bartlett. When does gradient descent with logistic loss find interpolating two-layer networks? Journal of Machine Learning Research, pages 1-48, 2021.
[24] Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Efficient approximation of deep relu networks for functions on low dimensional manifolds. Advances neural information processing systems, 32:8174-8184, 2019.
[25] Minshuo Chen, Haoming Jiang. Wenjing Liao, and Tuo Zhao. Nonparametric regression low-dimensional manifolds using deep relu networks: Function approximation and statistical recovery. arXiv preprini arXiv:1908) 01842,884222 2019.
[26] Minshuo Chen, Yu Bai, Jason D Lee, Tuo Zhao, Huan Wang, Caiming Xiong, and Richard Socher. Towards understanding hierarchical learning: Benefits of neural representations. arXiv preprint arXiv:2006. / 13336, 2020.
[27] Zhengdao Chen, Eric Ganden-Eijnder and Joan Bruna. On feature learning in neural networks with global convergence guarantees. In International Conference on Learning Representations, 2022.
[28] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
[29] Lenaic Chizat and Francis Bach. A note lazy training in supervised differentiable program- ming. arXiv preprint arXiv:1812.07955 2018.
[30] Lenaic 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. PMLR, 2020.
[31] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training differentiable program- ming. In Advances in Neural Information Processing Systems, 2019.
[32] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. In Conference on Learning Theory. PMLR, 2022.
[33] Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems, 33, 2020.
[34] Amit Daniely and Gal Vardi. Hardness of learning neural networks with natural weights. Advances in Neural Information Processing Systems, 33:930-940, 2020.
12

---

[35] Amit Daniely. Nathan Srebro, and Gal Vardi. Efficiently learning neural networks: What assumptions may suffice? arXiv preprint arXiv:2302.07426, 2023.
[36] Zhiyan Ding, Shi Chen, Qin Li, and Stephen J Wright. Overparameterizatioon of deep resnet: zero loss and mean-field analysis. The Journal of Machine Learning Research, 2022.
[37] Xialiang Dou and Tengyuan Liang. Training neural networks as learning data-adaptive kernels: Provable representation and approximation benefits. Journal of the American Statistical Association, 2020.
[38] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. In International Conference Machine Learning, 2019.
[39] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations, 8.
[40] Yu Feng and Yuhai Tu. Phases of learning dynamics in artificial neural networks: in the absence or presence of mislabeled data. Machine Learning: Science and Technology, 2021.
[41] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2018.
[42] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint 10X3v:1903.1111 2019.
[43] Spencer Frei and Quanquan Gu. Proxy convexity: A unified framework for the analysis of neural networks trained by gradient descent. Advances in Neural Information Processing Systems, 34, 2021.
[44] Spencer Frei, Yuan Cao, and Quanquan Gu. Provable generalization of sgd-trained neural net- works of any width in the presence of adversarial label noise. arXiv preprint 2021.
[45] Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202. 07622 2022.
[46] Spencer Frei, Gal Vardi, Peter L Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu networks trained on high-dimensional data. arXiv preprint arXiv:2210.07082, 2022.
[47] Spencer Frei, Gal Vardi, Peter L Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. arXiv preprint 2333:2233300162. 2023.
[48] Spencer Frei, Gal Vardi, Peter L Bartlett, and Nathan Srebro. The double-edged sword of implicit bias: Generalization vs. robustness in relu networks. arXiv preprint 0:X455 2023.
[49] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, page 113301, 2020.
[50] Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physics Reports, 924:1-18. 2021.
[51] Behrooz Ghorbani, Song Mei. Theodor Misiakiewicz. and Andrea Montanari. Limitations of lazy training of two-layers neural networks. arXiv preprint arXiv: 0906699098999 2019.
[52] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems, 2020.
13

---

[53] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien, Implicit regularization of discrete gradient dynamics linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[54] Ross Girshick, Jeff Donahue, Trevor Darrell, and Jitendra Malik. Rich feature hierarchies for accurate object detection and semantic segmentation. In Computer Vision and Pattern Recognition, 2014.
[55] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
[56] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages 1832 1841. PMLR, 2018.
[57] Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference Learning Representations, 2019.
(58) Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. In International conference on machine learning, pages 4542-4551. PMLR, 2020.
[59] Arthur Jacot. Implicit bias of large depth networks: notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations, 2023.
[60] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, 2018.
[61] Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory, pages 1772-1798. PMLR, 2019.
[62] Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent achieve arbitrarily small test error with shallow relu networks. In International Conference on Learning Representations, 2019.
[63] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment deep learning. Advances in Neural Information Processing Systems, 33:17176-17186 2020.
[64] Pritish Kamath, Omar Montasser, and Nathan Srebro. Approximate is good enough: Proba- bilistic variants of dimensional and margin complexity. In Conference on Learning Theory, 2020.
[65] Hamed Karimi, Julie Nutini, and Mark Schmidr. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewizz condition. In Joint European Conference on Machine Learning and Knowledge Discovery Databuses, pages 795-811. Springer, 2016.
[66] Jonas Moritz Kohler and Aurelien Lucchi. Sub- Sub-sampled cubic regularization for non-convex optimization. In International Conference on Machine Learning. PMLR, 2017.
[67] Guy Kornowski, Gilad Yehudai, and Ohad Shamir. From tempered to benign overfitting in relu neural networks. arXiv preprint arXiv: 15141, 2023.
[68] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks gaussian processes. In International Conference on Learning Representations, 2018.
[69] Jachoon Lec, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 2019.
[70] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. Advances in Neural Information Processing Systems, 33:15156-15172. 2020.
14

---

[71] Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems, 2018.
[72] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Conference on Learning Theory, 2020.
[73] Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks infinite-width limit. Journal of Machine Learning Research, 2021.
[74] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2019.
[75] Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent on two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems,
2021.
[76] Eran Malach, Pritish Kamath, Emmanuel Abbe, and Nathan Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. arXiv preprint arXiv:2103.01210, 2021.
[77] Christopher D Manning, Kevin Clark, John Hewitt, Urvashi Khandelwal, and Omer Levy. Emergent linguistic structure in artificial neural networks trained by self-supervision. Proceed- ings of the National Academy of Sciences, pages 50046-30054 2020.
[78] Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
[79] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 2018.
[80] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-fre free bounds and kernel limit. In Conference on Learning Theory, pages 2388-2464 PMLR, 2019.
[81] Behrad Moniri. Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of non- linear feature learning with one gradient step two-layer neural networks. arXiv preprint arXiv:2310.078)) 2023.
[82] Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. The Annals of Statistics, 2022.
[83] Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale training accuracy. Advances in Neural Information Processing Systems, 33, 2020.
[84] Alireza Mousavi-Hossein Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv: 2022.
[85] Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry, Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning, pages 4683-4692. PMLR, 2019.
[86] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.
[87] Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3051-3055 PMLR, 2019.
[88] Vaishnavh Nagarajan and Zico Kolter. Uniform convergence may be unable to explain generalization in deep learning. Advances in Neural Information Processing Systems, 2019.
15

---

[89] Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L Edelman, Fred Zhang, and Boaz Barak. Sgd on neural networks learns functions of increasing complexity. arXiv preprint arXiv: /905. /604, 2019.
[90] Behnam Neyshabur. Implicit regularization in deep learning. arXiv preprint
2017.
[91] Roman Novak, Lechao Xiao, Jachoon Lee, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pen- nington, and Jascha Sohl-Dickstein. Bayesian convolutional neural networks with many channels are gaussian processes. In International Conference Learning Representations, 2019.
[92] Ryan O'Donnell. Analysis of boolean functions. Cambridge University Press, 2014.
[93] Samet Oymak and Mahdi Soltanolkotabi. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In International Conference on Machine Learning, pages 4951-4960. PMLR, 2019.
[94] Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparamettriiaatoo:: Global convergence guarantees for training shallow neural networks. IEEE Journal Selected Areas in Information Theory, pages 84-105, 2020.
[95] Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi. Generalization guaran- tees for neural networks via harnessing the low-rank structure of the jacobian. arXiv preprint arXiv:1906.0533 2019.
[96] Vardan Papyan, XY Han, and David Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, pages 24652-24663, 2020.
[97] Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Mechanism of feature learning in deep fully connected networks and kernel machines that recursively learn features, 2023.
[98] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems, 2008
[99] 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 Leurning, pages 8936-8947. PMLR, 2021.
[100] Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. In The Eleventh International Conference on Learning Representations, 2023.
[101] Itay Safran, Ronen Eldan, and Ohad Shamir. Depth separations in neural networks: what is actually being separated? In Conference on Learning Theory, pages 2664-2666. PMLR, 2019,
[102] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateck Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In NeurIPS, 2020.
[103] Shai Shalev-Shwartz, Ohad Shamir, and Shaked Shammah. Failures of gradient-based deep learning. In International Conference on Machine Learning, pages 3067-3075. PMLR, 2017.
[104] Zhenmei Shi, Yifei Ming, Ying Fan, Frederic Sala, and Yingyu Liang. Domain generaliza- tion with nuclear norm regularization. In NeurIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications, 2022.
[105] Zhenmei Shi, Junyi Wei, and Yingyu Liang. A theoretical analysis feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations, 2022.
[106] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: central limit theorem. Stochastic Processes and their Applications, pages 1820-1852, 2020.
16

---

[107] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, pages 2822-2878, 2018.
[108] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34:23831-2388 2021.
[109] Matus Telgarsky. Feature selection with gradient descent two-layer networks in low-rotation regimes. arXiv preprini arXiv:2208. 02789, 2022.
[110] Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks. arXiv preprint arXiv:2202.0029: 2022.
[111] Yifei Wang, Jonathan Lacotte, and Mert Pilanci. The hidden convex optimization landscape of two-layer relu neural networks: an exact characterization of the optimal solutions. arXiv e-prints, pages arXiv-2006, 2020.
[112] Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
[113] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, 2020.
[114] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian pro- cess behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.0476/ 2019.
[115] Greg Yang and Edward Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522, 2020.
[116] Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems, 2019.
[117] Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks, In European Conference on Computer Vision, 2014.
[118] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
[119] Chiyuan Zhang, Samy Bengio, and Yoram Singer. Are all layers created equal? arXiv preprint arXiv:1902.01996,20 2019.
[120] Mo Zhou, Rong Ge, and Chi Jin. A local convergence theory for mildly wer-paramettrizzd two-layer neural network. In COLT, 2021.
[121] Difan Zou and Quanquan Gu. An improved analysis of training over-parameterized deep neural networks, Advances in neural information processing systems, 32, 2019.
[122] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes ever-parameterized deep relu networks. arXiv preprint arXiv: 1.08888, 2018.
[123] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep ReLU networks. Machine Learning, 109:467-499 2020.
[124] Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. The benefits of mixup for feature learning. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 43423-43479, PMLR, 2023.
17

---

Appendix
Contents
1
Introduction
2 Related Work
3 Gradient Feature Learning Framework
3.1 Warm Up: A Simple Setting with Frozen First Layer
3.2 Core Concepts in the Gradient Feature Learning Framework
3.3
Provable Guarantee via Gradient Feature Learning
Applications in Special Cases
4.1
Mixtures of Gaussians
4.1.1
Mixtures of Gaussians: Beyond the Kernel Regime
4.2
Parity Functions
................................
4.2.1
Parity Functions: Beyond the Kernel Regime
10
5 Further Implications and Conelusion
10
A Broader Impacts
19
Limitations
19
C More Further Implications
20
Gradient Feature Learning Framework
D.1
Simplified Gradient Feature Learning Framework
D.2
Gradient Feature Learning Framework under Expected Risk
22
D.2.1 Feature Learning
23
D.2.2
Good Network Exists
24
D.2.3
Learning an Accurate Classifier
26
D.3
More Discussion abut Setting
30
D.4
Gradient Feature Learning Framework under Empirical Risk with Sample Complexity
E
Applications in Special Cases
38
E.1
Linear Data
38
E.2 Mixture of Gaussians
40
E.2.1
Problem Setup
40
E.2.2
Mixture of Gaussians: Feature Learning
41
E.2.3
Mixture of Gaussians: Final Guarantee
46
E.3
Mixture of Gaussians XOR
48
E.3.1
Mixture of Gaussians XOR: Feature Learning
48
18

---

E.3.2
Mixture of Gaussians XOR: Final Guarantee
52
E,4
Parity Functions
52
................................
E.4.1
Problem Setup
................................
52
E,4.2
Parity Functions: Feature Learning
................................
54
E.4.3
Parity Functions: Final Guarantee
58
E.5
Uniform Parity Functions
59
E.5.1
Uniform Parity Functions: Final Guarantee
60
E.6
Uniform Parity Functions: Alternative Analysis
61
E.6.1
Modified General Feature Leamingg Framework for Uniform Parity Functions
61
E.6.2
Feature Learning of Uniform Parity Functions
66
E.7
Multiple Index Model with Low Degree Polynomial
68
E.7.1
Problem Setup
68
E.7.2
Multiple Index Model: Final Guarantee
68
F Auxiliary Lemmas
70
Appendix A discusses the potential socictal impact our work. Appendix B describes the limitations of our work. In Appendix C, present framework implications about simplicity bias. The complete proof of our main results is given in Appendix D. We present the case study of linear data in Appendix E.I, mixtures of Gaussians in Appendix E.2 and Appendix E.3, parity functions in Appendix E.4, Appendix E.5 and Appendix E.6, and multiple- indeiple-index models in Appendix E.7. We put the auxiliary lemmas in Appendix F.
A Broader Impacts
Our paper is purely theoretical in nature, and thus we do not anticipate an immediate negative ethical impact. We provide unified theoretical framework that be applied to different theoretical problems. We propose the two key ideas of gradient feature and gradient feature-induced induced neural networks not only show their ability unify several current works but also open new direction of thinking with respect to the learning process. These notations have the potential to be extended to multi-layer gradient features and multi-step learning, and this work is only first step.
On the other hand, this work may lead to a better understanding and inspire the development of improved network learning methods, which may have positive impact on the theoretical machine- learning community. It may also beneficial to engineering-inlliee machine-learning researchers.
B Limitations
Recover Existing Results. bounds in existing work.
The framework may may not recover the width sample complexity
1. The framework can give matching bounds the existing work in some cases, like parities over uniform inputs (Appendix E.5).
2. In some other cases, it gives polynomial error bounds not the same as those in the existing work (e.g., for parities over structured inputs). This is because our work is analyzing general cases, and thus may not give better than or the same bounds as those in special cases, since special cases have more properties that can be exploited to get potentially better bounds. On the other hand, our bounds can already show the advantage over kernel methods (e.g.. Proposition 4.9).
We would like emphasize that our contribution providing an analysis framework that can (1) formalize the unifying principles of learning features from gradients in network training, and (2) give
19

---

polynomial error bounds for prototypical problems. Our focus is not to recover the guarantees in existing work.
Failure Cases. There are some failure cases that gradient feature learning framework cannot cover:
1. In [101], they constructed function that is easy to approximate using a 3-layer network but not approximable by any 2-layer network. Since the function is not approximable by any 2-layer network, it cannot be approximated by the gradient-induced networks as well, so OPT will be large. As a result, the final error will be large.
2. In uniform parity data distribution, considering an odd number of features rather than even, i.e., k is odd number in Assumption E.28, we show that our gradient feature set is empty even when p in Equation (6) is exponentially small, thus the OPT is a positive constant since the gradient induced network can only be constants. Meanwhile, the neural network won't 't be able to learn this data distribution because its gradient is always O through the training, and the final error equals OPT.
The first case corresponds to the approximation hardness of 2-layer networks, while the second case gives learning hardness example. The above two cases show that if there is approximation or learning hardness, our gradient feature learning framework may be vacuous because the optimal model in the gradient feature class has large risk, then the ground-truth mapping from inputs to labels is not learnable by gradient descent. These analyses are consistent with previous works [15, 101].
C More Further Implications
Our general framework also sheds some light on several interesting phenomena in neural network (NN) learning observed in practice. Feature learning beyond the kernel regime has been discussed in Section 4.1.1 and Section 4.2. The lottery ticket hypothesis (LTH) has been discussed in Section 5. Below we discuss other implications.
Implicit Regularization/Snnpiitt Bias. It is now well known that practical NN are overparame- terized and traditional uniform convergence bounds cannot adequately explain their generalization performance [59, 88, 118]. It is generally believed that the optimization has some implicit regulariza- tion effect that restricts learning dynamics to subset of the whole hypothesis class, which is not of high capacity so can lead to good generalization [53, 90]. Furthermore, learning dynamics tend to first learn simple functions and then learn more and more sophisticated ones (referred as simplicity bias) [89, 102]. However, it remains elusive to formalize such simplicity bias.
Our framework provides a candidate explanation: the learning dynamics first learn to approximate the best network in a smaller family of gradient feature induced networks Fd.r.By.S and then learn to approximate the best in a larger family. Consider the number of neurons for illustration. Letri & r2. and let 7i and T2 be their corresponding runtime bounds for T in the main Theorem 3.12. Clearly, Ti < T2. Then, at time Ti, the theorem guarantees the learning dynamics learn to approximate the best in the family Fd.ri,B..S BB.. with 71 neurons, but not for the larger family Br.., Later, at time T2, the learning dynamics learn to approximate the best in the larger family Fd.ra,Bp.S. That is, the learning first learns simpler functions and then more sophisticated ones where the simplicity bias is measured by the size of the family of gradient feature-induced networks. The implicit regularization is then restricting to networks approximating smaller families of gradient feature-induce networks. Furthermore, we can also conclude that for SGD-optimized NN, its actual representation power is from the subset of NN based on gradient features, instead of the whole set of NN. This view helps explain the simplicity bias/implicit regularization phenomenon of NN learning in practice.
Learning over Different Data Distributions. Our framework articulates the following key princi- ples (pointed out for specific problems in existing work but not articulated more generally):
Role of gradient: the gradient leads to the emergence of good features, which is useful for the learning of upper layers in later stages.
From features to solutions: learned features in early steps will not be distorted, if not improved, in later stages. The training dynamic for upper layers will eventually learn good combination of hidden neurons based on gradient features, giving a good solution.
20

---

Then, more interesting insights are obtained from the generality of the framework. To build general framework, the meaningful error guarantees should be data-dependent, since NN learning on general data distributions is hard and data-independent guarantees will be vacuous [34, 35]. Comparing the optimal in a family of "ground-truth" functions (inspired by agnostic learning in learning theory) is a useful method to obtain the data-dependent bound. We further construct the "ground-truth" functions using properties of the training dynamics, i.c. gradient features. This greatly facilitates the analysis of the training dynamics and is the key to obtaining the final guarantees. On the other hand, the framework can also be viewed as using the optimal by gradient-induced NN to measure or quantify the "complexity" of the problem. For easier problems, this quantity is smaller, and our framework can give a better error bound. So this provides a united way to derive guarantees for specifie problems.
New Perspectives about Roadmaps Forward. We argue new perspective about the connection between the strong representation power and the successful learning of NN. Traditionally, the strong representation power of NN is the key reason for hardness results of NN learning: NN has strong representation power and can encode hard learning questions, so they are hard to learn. See the proof in SQ bound from [33] NP-hardness from [19]. The strong representation power also causes trouble for the statistical aspect: it leads to vacuous generalization bounds when traditional uniform convergence tools used.
Our framework suggests a perspective in sharp contrast: the strong representation power of NN with gradient features is actually the key to successful learning. More concretely, the optimal error of the gradient feature-induced NN being small (i.c., strong representation power for a given data distribution) can lead to a small guarantee, which is the key to successful learning. The above new perspective suggests a different analysis road than traditional ones. Traditional analysis typically first reasons about the optimal based on the whole function class, i.e. the ground truth, then analyze how NN learns proper features and reaches the optimal. In contrast, our framework defines feature family first, and then reasons about the optimal based on it.
Our framework provides the foundation for future work on analyzing gradient- based NN learning, which may inspire future directions including but not limited to (1) defining a new feature family for 2-layer NN rather than gradient feature, (2) considering deep NN and introducing new gradient features (c.g., gradient feature notion for upper layers), (3) defining different gradient feature family at different training stages (c.g.. gradient feature notion for later stages). In particular, the challenges in the later-stage analysis are: (a) the weights in the later stage will not be as normal as the initialization, and we need new tools to analyze their properties; (b) to show that the later-stage features eventually lead to good solution, may need new analysis tools for the non-convex optimization due to the changes in the first layer weights.
D Gradient Feature Learning Framework
We first prove Simplified Gradient Feature Learning Framework in Appendix D.I, which only con- siders one-step gradient feature learning. Then, we prove our Gradient Feature Learning Framework, e.g., no freezing of the first layer. In Appendix D.2, we consider population loss simplify the proof. Then, we provide more discussion about our problem setup and our core concept in Appendix D.3. Finally, we prove our Gradient Feature Learning Framework under empirical loss considering sample complexity in Appendix D.4.
D.1 Simplified Gradient Feature Learning Framework
Algorithm 2 Training by Algorithm with no updates for the first layer after the first gradient step Initialize [a w(.))) (ato,oo),
E Fa,mi Sample 2 ~ D
Get (1) (1) b) by one gradient step update and fix W(),, b
for 1 - 2 to T do (1) a (!-1) VV 2Z(())))))
end for
Theorem 3.4 (Simple Setting). Assume Ez (Jo. ((n,was)))) is L-smooth to Let 11 I. (t) 0. for all 32. Training by Algorithu 1 with no updates for the first layer after the
21

---

first gradient step, w.h.p., there exists t E [T] such that ((0,(, & OPT wi) ba
( L(l 1( Baa) + 1 (I 17 Bg- +
Proof of Theorem 3.4. Recall that
B.. = (f(.. E Fa, ||al|2 S Baz),
W.
min
Lp(f). (13)
fEFw.b,B.z
We denote f = argmin JeFw...B.a Cf(f) and f = argmin BeF denote their second layer weights respectively. Then, we have
(z(.). We use a" and ä*
(ol((alt) =LD(f(a)) (().)) - w(),,))
(14)
+ Cz(f ww, (z(f(a,.oo.b)) wwo.)))
(15)
+ Σz(F(a ( wwi))))) Ez(f(a,wwon
(16)
+ Ez(sia w())))))) Coll(ar.w)),p)
(17)
+ CD(f(a <Dlf,a.....,,.))
(18)
Ez( $ CD Co(Siano wwa) - (2 ww)b)) b))
(19)
+ (2(()))) - Ez(J(ar,
(00)
(21)
(f( Colf(ar,wo,,))
(22)
+ OPTww),bBB.
(23)
Fixing W(1), b and optimizing a only is convex optimization problem. Note that n S Z. where Ez is L-smooth a. Thus with gradient descent, we have
T
|| (1)
T
Σ
Ez (f(ao, wo, - Ez (J(ar ww),)))
2TT
(24)
Then our theorem gets proved by Lemma F.9 and generalization bounds based on Rademacher complexity.
D.2 Gradient Feature Learning Framework under Expected Risk
We consider the following training process under population loss to simplify the proof. We prove our Gradient Feature Learning Framework under empirical loss considering sample complexity in Appendix D.4.
Algorithm 3 Network Training via Gradient Descent
Initialize (a) W()), b) as in Equation (8)
for t= to T do !!) = (!-1) (V (0)
(Jz(-1))
( = ((- - (+) Cb (fgu-1))
end for
Given an input distribution, we can get a Gradient Feature set Sp and (*(x) En ato((w3. x) -0 -b²), where fo E Fa.r,Br.sp.... ty is Gradient Feature Induced networks defined in Definition 3.1 Considering training by Algorithm 3, we have the following results.
22

---

Theorem D.1 (Gradient Feature Learning Framework under Expected Risk). Assume Assumption 3.1. For any e, (0,1), m I ed and
1
BB1 Byl
Bb
v8 + 1 (log
(7))²
=Ω
(25)
p
c
Be
P
1
VrrBaz
(V BG (mp) ))
V Vlog
1
T T=
mb)
(26)
(mp)
then with proper hyper-parameter values, we have with probability 2 6, there exists E [T] in Algorithm 3 with
Pr|sign(fzm) (x)) + yl S CD (fEw) Bp + rr..B v 27 te,
(27)
See the full statement and proof in Theorem D.9. Below, we show some lemmas used in the analysis of population loss.
D.2.1 Feature Learning
We first show that a large subset of neurons has gradients at the first step as good features. Definition D.2 (Nice Gradients Set. Equivalent to Equation (9)). We define
N), = E [m]:
D) (1 7) w
(1)
N ((0)a (0) BG
11,NNe := {i [2m] [m]] (w/) (1) D (1 7)
(0)a (0)
Ba
where y, Bg is the same in the Definition 3.7.
Lemma D.3 (Feature Emergence. Full Statement of Lemma 3.13). Let x() -f..- For any r size subset (1).,.)).... (Dr, } } S Sp.,Ba. with probability at least 1 - 2re cmmp where > 0 is universal constant, we have that for all j E (P).s)GG)
mp 1
Proof of Lemma D.3. By symmetric initialization and Lemma F.I, we have for all i E [2m]
w!" 77 (0)a (0) ((x.v)) yyo' [( (w:))) xx) b;
(28)
17 (0)a G(w!)) (,))
(29)
For all E [r], as (Dj, 83) Sp by Lemma F33,
(1) if 8y = +1, for all i E [m]. we have
Pr i E G(D NiNe
(30)
Dj
[ | 2 ((()a))))))) BG
Prr
(1 (11,)),
(31)
w/ 112
(
(1 7), ||m > n7 (1) (0)aa) (1) D.)
(0)
b;
=Pr
BG, [b,] 8j
(32)
w!"
2 Pr
G(w (0)
b,) E CD [[G(w b,)ll2 2 Bc.
b;
Sj,
(0) (G(W (0) bb)), b,), > 0
b
p
Z5 2'
(33)
(2) if 8y = 11, for all i e [2m] \ [m], similarly we have
Pr [i E Niie 7.
p
(34)
2
By concentration inequality, (Chernoff's s inequality under small deviations), we have
Pr (GD,,
mp < 2e-cen cmp -mpp
(35)
We complete the proof by union bound.
23

---

D.2.2 Good Network Exists
Then, the gradients allow for obtaining set of neurons approximating the "ground- truthh network with comparable loss.
Lemma D.4 (Existence of Good Networks. Full Statement of Lemma 3.14). Let A .payy For any
B. E (0, Ba), let a. = θ -P(0) Bg and 8 = 2re- Then, with probability at least - ò
over the initialization, there exists å, 's such that f(a,wc),))()) = E ā,o ( ((w/",,,) b,) satisfies
BA B, /mpBaB.
Co(f(a.wo).b)) S rBa
BBIV Be -P.T...Bp.Sp.... (36)
and ||āllo = -0 (v((p))) ||ā||2 00
B,BB,
[āll====0
=0
B. B, 6(mp))
b(mp))
Proof of Lemma D.4. Recall f*(x) Lj=1 a30((W). bj), where f" E Bp is
bj
defined in Definition 3.11 and let 8j = [by]- By Lemma D.3, with probability at least 61, 81 = 2re cmp, for all j E [r]. we have [G(m3 (P/W..)))). N.)) 2 mp have (0)) ((w (0) b.) Di only depend w and b,, which is independent of a
4 Then for all i E G(w; Nioe C [2m], we
(0) Given
Definition 3.7, we have
bj E b.)ll2
(e(o)n (1) Ba Bd
B₁
-1(00)
5 (0))
(1) Bx1
(37)
b
b
We split [r] into = (j E [r] bi1 B.J.
bi E [r] : bj S BB. and I'+ = {j E [r]:
by 2 B.). Let ta = B. Then we know that for all j E E _T_, for all i E G(w,,e (w)
B.
we have
by
Pr (0) ~N
a (0) ((0)) (G (0)
bi)ll
S Ea
(38)
b S + ca
Pr n. ~N(0,a2)
1- ca S -a (00) b,)ll2
(39)
b
|G((w (0) bi)
= Pr gov(0.1)
1 1- Ea S g
$1+6 Ea
(40)
BoB.
BGB.
BG
Pr go.1(0.11
(1- (a-
S (c)
|[G(w:") b:)ll2(bū)
|G((w b.))
BB Be
=0
(41)
[[G(w (0)
bi)l22(b71
<a Ba Be
20
(42)
B Bo
=Ω
(43)
/mp
Thus, with probability 2 (v++p) over a;
(0)
we have
lbj|
i BoB.
a (0) (00) (0)
ba) 112
b
1 S Eas
a
(0)
=
(()))
(44)
Similarly, for j E P. for all i E G(wz, with probability ?
over a (0)
we have
mmp
a "I( bi)lla 1
B.
1
5 En.
a,
(0)
= 0
(1()))
b
BB
(45)
24

---

For all j E [r], let A, S (. be the set of i's such that condition Equation (44) or Equation (45) are satisfied. By Chernoff bound and union bound, with probability at least 1-d2, 62 re" re-vmp, for all j E [r] we have |A5l 2 mmpm)
We have for Vj E ET+UT_, Vi E Ag,
bb]
wi
(1)
xx) - xw)
(46)
b
< a (0) ((0))"))G((? (0) bi)lla
|bj] w
(1)
w
(1)
+
w
(1)
wj
|
(47)
b tw/"1e
l! Iw/2112 l12
l| Tw/"1e ll2
S(Ea + /2))|x|22.
(48)
Similarly, for Vj E
Be
w x) - (wj: x) S (ea + /27)||
(49)
b
If Aj. -T- set ā;
aj
[bi| [n,
i E Az. E I, set ā;
aj
B |A 1
otherwise set a;
0,
we have ||āllo = O (r(mp))). ||āl||2 =0
B. 0(mp))
||ā|| 0
Ba B.
d(mpp)
Finally, we have
(50)
=Co(f(a.woo.)) - CD((')+++++))))))) CCol) + LD(f*)
(51)
<B(x)) [[), - [x(x]] + LD(
(52)
Σ 0 ( (m/ x b) + Σ ã (( (w!") (m.["')) :x) =8) Σ a300 [w] x) - -b²)
2m
(((x.))
1=
+ (D(
(53)
bbj 0
<B(x).)
a
((w!!)
( ()) xx)
(x) - b) - < x) -
x(w -b})
(54)
A
jET EEN;
+ (xxx))
Σ
a
1
|b
(w!)) xx) +b) - o [(W x) - -b3)
(55)
|
JET iEA,
+ (x..)
Σ a iEA,
1
Bg
(ww!")
(1) (,/",)) xx)=b) b) =((3] - b,}) + CD(
(56)
|A
<B(x.u)
a
1
bbj
(w?
(1) .x) - -(w3,
(57)
A
JET- HEA,
+ +xxx))
Σ Σe.
aj
1
bb
w
(1) (1)
.x) x
(58)
|
JET REA;
+ (
Σ a JET iEA;
1
B
w
(1) x) + B. - (w)x x)
+ LD(
(59)
[
b
Sr|la" 1l (Ea 227) /27) Ex,)) x||2 + | a B. + LD(
(60)
SrBeI BaI + 22 + B. B. + OP Dr.TaBBriiS....
(61)
We finish the proof by union bound and d > 6, + 82.
25

---

D.2.3 Learning an Accurate Classifier
We will use the following theorem from existing work to prove that gradient descent learns good classifier (Theorem D.9). Theorem D.I is simply a direct corollary of Theorem D.9.
Theorem D.5 (Theorem 13 in [33]). Fix some 1,. and let f1..... fr be some sequence of convex functions. Fix some 8,, and assume we update 0,41 = 0r f.(.). Then for every 8 the following holds:
T T f(()))
T
T
S T
Σ f((0))
110 + 161112 + 2nT
Σ
(0t)))) + "T
T t=1
To apply the theorem we first present few lemmas bounding the change in the network during steps. Lemma D.6 (Bound of [ ). Assume the same conditions as in Lemma D.4, and 2 log m, with probability least 1
m² over the initialization, ||a 0
bylog m
and for
=f Bc
all i [4m], we have [|w; (0) ||2 = O (owvd) Finally, |a = 0 (- (
for all E [Am], [| 12 = O
Myyoogm... Big B.
Proof of Lemma D.6. By Lemma F.4, we have || = 0 at least 1 2mz by property of maximum ii.d Gaussians. For any E [4m], by Lemma F.5 and 2 log m, we have
by Eglg m BB
with probability
(1)
Pr
d + 2 4d 8log(m) 0o
(62)
02
m4
Thus, by union bound, with probability least O @ow (owvd).
2m for all i E [4m]], we have ||w! ||2
For all E [4m], we have
la (1) = -1))00)00))) ((x,)) [o (( (w:")
(0) xx)x bi)
(63)
S- (-)())(0)(( l + +b)
(64)
<0 ( (m)) vd+
(65)
- 11) ( (0)
E(x.os [oo' [( (x/ b, bb]x]]
(66)
b bVlogmma1
<0
(67)
Bg
Lemma D.7 (Bound of =00). Assume the same conditions as in Lemma D.6, and let = ()) for all
LE (2,3,..., T}, 0 B. S o(1), and 0 > = (1) for all E (2.3,..., T), for all E [4m].
we have
la. (0) 50 |a. (1 || (2 + B
+ nb
(68)
Iw/2 w
(1) 2|2 50
(11 (2 + 2 |w! (1)
(1)
||2 Lnj BB1b)
(69)
26

---

Proof of Lemma D.7. For all i E [4m], by Lemma D.6.
|a (1)
[1 - n)Aja 1111 - (x.))) (' ('U [o ((w, (mm xx) bi }]
(70)
(1 - n) 1) (1
+ 17 E [x,)) [[o ((w!- ( x) - bb:)
(71)
8 (1-1)
+ 1p( Bzil l| w/!""1//1
|/2 + 6)
(72)
1)
+ n1 Il w! (t-1) - w 11) 12 + 17 BB1 /w w/''22 + n/
(73)
a 1)
+ B1 |/! ((-1) w/ 12 + 7
(74)
where we denote Z, = Bn l!w,"" 1 112 + b. Then we give a bound of the first layer's weights change,
|w (t) w; 1 112
(75)
(1 n))) (t-1) na;
(xxx)) B ['(USEu )(s))yoo [<w/-- xx)x -by] x w; w/'l| (76)
(1-1)
<||w!" (t-1) w! (
1|2 + nB a 11
(77)
Combine two bounds, we can get
2-2
(0) Sla
(t(
+ nZz + (1)Ba1)²
1
a
(1)
(78)
h-1
1-2
a
Σ
|a
10)
(1 - 12)
Σ
lay
+ +ni...
(79)
1=1
1=1
Let h(1) = () h(2))
2 2 1 + n2, and h(t + 2) = 2/1(t + 111 (1 (n1Bz1)))h((( + n/ for
n E N4, by Lemma F.8, we have
Z₁
htt)
(1 - nBBB)) (t-1) + c1 + ,nBr1)(t
(80)
nB
C1
1
|a
(1)
+
2,
|a
(1)
+
(81)
2
BB2 #1
nB
C2 2
1
|a
(1 +
2,
+
|a (1) + 2
(82)
BB #1
nBa1
Thus, by |G1 S C2. and 0 < B11 S o(1), we have
la
(8)
<h(1) hil - 1)
(83)
=- 01 (1 - -n (1-2) + nn22102 + nBB1) (1-2)
(84)
C2 + n2
(85)
B0(2n)
(86)
Similarly, by binomial approximation, we also have
w (t)
w, (1)
112
Bn -1 1)
(87)
(-Zz
Z;
-1
+ - 7B ((-2) + cz((1 + B)x) (C-2)
(88)
nB2
Z:
SnBB0
+ - -(t - (2))))Bn)) + c2 + (t - 22)
(89)
BB
7
SnBBO
+a + c2 + (c2 -c1)tt/ c1)tn
(00)
nB2 m1
|a.") (1)
la
(1)
+ 17 tn
SnBBO
(91)
BBB1
00 (mlaa(
(1)
+ 11 (.).
(92)
27

---

We finish the proof by plugging Z,, 22 into the bound.
Lemma D.8 (Bound of Loss Gap and Gradient). Assume the same conditions as in Lemma D.7, for all t E [T], we have
(CDl((aw..b) - LD( SB
2 |a/o
max [E[4m]
||w w ( 1l2
(93)
for all [T] for all i [4m], we have
@CD(f)))
BB (11w
w 1l2 llw b.
(94)
0a!")
Proof of Lemma D.8. It follows from that
(D(a(a.wo..) CD( (a w()
(95)
< ((x.))/(aa ((t)),)))) (x) f(a bx)
(96)
0 [( (w?") xx) b.) - -o (o((w()) ,x x)
S E(x,)
max
b;
(97)
=E[4m]
$ Bz1 ||
max E [m
[w
(1)
112
(98)
Also, we have
@LD(JEw)
((xa)) (xf=)(())) [o ((w/") x) b, b1]
(99)
0a,a)
<B ||w (1) 112 +
(100)
SBa (I (1) wi 122 + ||, (1) 12) +b.
w
(101)
We are now ready to prove the main theorem. Theorem D.9 (Online Convex Optimization. Full Statement of Theorem D.I). Consider training by Algorithm 3, and any E (0, 1). Assume 2 log m. Set
Ow > 0, > 0, 17 = n, A = ffr all E 12.3....,77 T)
(1 = min (1b))
b(mp))
1(1)
Oa =
(0(()) vd d +b)
1115²
(e(0)) BaVBCe,
Let 0 < TTBBI S o(1), Ω ( + (log (f))²) initialization, there exists E [T] such that
With probability at least 1 - å over the
2Ba1
Bb
CD (fzu) Br. r
+ B 2
(102)
(mp) 1
Bc
Vlog VlogmBu Bmg Bur (mp)*
+1 +0 Bag
+1 +n mb) 0
/ BB
(yTb2(mp))
Furthermore, for any (0, 1). set
VrBai ) m = Ω ( pel
=
Bb rB Bra
+
+ p
(log
(;))²² (103)
rr
88
Be
v8
17 0
T +
(104)
Bo + mb) (mp))
Vlog mB (mp)+ BB,Bo +
nB (mp))
we have there exists E [T] with
Prisign(fgo) (xx # yl S LD (few) SOPTo,r.B..o Spo..g BB..... 22 +e.
(105)
28

---

(log (3))² we have 2re- vnn + m² I 8. For any B, (0, Ba), when oa θ -(((0) B. by Theorem D.5, Lemma D.4, Lemma D.8, with probability least 1 over the initialization, we have
Proof of Theorem D.9. By m ? J
CD (fzw)
(106)
T
T (CD(f(a,wo..) Lo( + (a()(a
(107)
|
OLD (SEmT)
+ 2n
+ |l2mm 4mm) max re(4m]
(108)
oa!"
B Bo
SPPTdn.B.S..... + rBa
+ Bz1 / 27 + Be
(109)
VmpBa Be
+ Ba Ba|lā|/a 2
|alo
max
(?)
(T) ( ll2
(110)
=E[4m]
|ālé 4m 4m Bm 1|
b
+ 2xyT
00 + n)
max
|w
w 1|2 + max ||
(2
1|2 +
(111)
ic(4m)
if[4m]
B,
By Lemma D.4, Lemma D.6, Lemma D.7, when sy = =
min min(0On),0(n))
we have
(R. vd+hh)
BozBo
|ā =0 (r(mp)²).
|ā112 |ā12= =0
((mp))
(112)
|| 1lo -0 (=)
(113)
= miin(O(n).0(n)))
(114)
liVlogB.y
max xE =E[4m]
||
(1)
2 =0
(115)
BcB.
max
|w
w w,(')) =0
TriBalla()1 1000 + B2 se [4m]
max
|| l| (
+ Tr? BaD
(116)
se [4mm
=0
Try² B2
max sE[4m]
|
12
(117)
B
Set B. = (mp)) 7 V
B,1
Be Be we have da =
(-r00) Ba B.
7mmp)
which satisfy the requirements.
Then.
Σ LD (fsw)
(118)
T
2Bm1
Bb
<OPT Br rBal
+ Ba BaI/27
(119)
(mp) i V Be
+ (V B²
Bal
+ mmBz1
o
(bylogmBu
+
+0
Ba BZ (1T6((mp))
BGB.
Br1
2B
Bb
SOPTarBB.... + r
+ Bad 2
(120)
(mp)
V
Be
Vlog m mB BlogmBml (mp)+
B2
+1 (VrBaaB.TTB2, +bb) 0
+1) +0
(121)
VB,Ba
(17T)2(mp)
29

---

Furthermore, for any E (0, 1). set
B Bo BB m = Ω
Bb BG
8 ))
-
r rBal Bal
+
+
log
(122)
VrBal
pet
v0
P
7 0
T= θ
(123)
Br B+1
+ mb)
ml Bua (mp)) BB
+ +11
nBar (mp)
(mp)*
we have
Σ LD (Jzw) SOPT d.r, Bp. + +rBal
T
2BBl
Bb
+ BaY 227 B2I
+ 6
(124)
(mp)
Bg
But Ba2
+ +
(125)
((mp))
Bo + + Bzi /
(126)
We finish the proof as the 0-1 classification error is bounded by the loss function, e.g.. [[sign(f(x)) # vl f(vf(x)) where E(0) =
((O)
D.3 More Discussion abut Setting
Range of Owe In practice, the value of on cannot be arbitrary, because its choice will have an effect on the Gradient Feature set Sp.,BB- On the other hand, d 2 log is natural assumption, otherwise, the two-layer neural networks may fall in the NTK regime.
Parameter Choice. We use À = 1/1 in the first step so that the neural network will totally forget its initialization, leading to the feature emergence here. This common setting for analysis convenience in previous work, e.g., [32, 33, 105]. initialization and large step size for the first few steps). as long as after the gradient update, the gradient dominates the neuron weights. We use À 0 afterward as the regularization effect is weak in our analysis. We can extend our analysis to À being a small value.
We can extend this to other choices (e.g.. small
Early Stopping. Our analysis divides network learning into two stages: the feature learning stage, and then classifier learning over the good features. The feature learning stage is simplified one gradient step for the convenience of analysis, while in practice feature learning can happen in multiple steps. The current framework focuses on gradient features in the early gradient steps, while feature learning can also happen in later steps, in particular for more complicated data. It is an interesting direction to extend the analysis to longer training horizon.
Role of s. The s encodes the sign of the bias term, which is important. Recall that we do not update the bias term for simplicity. Let's consider simple toy example. Assume we have fi (z) [ + 1). fz) ag w 1) and fs(x) aza (w] 2), where o is ReLU activation function which is a homogeneous function.
1. The sign of the bias term is important. We a2 a2a(w,) (w: x - 1) for any a1, W1, 42, W2. This means that fi fx) and f2(x) are intrinsically different and have different active patterns. Thus, we need to handle the sign of the bias term carefully.
that we always have [jo(Fā + *
2. The scaling of the bias is absorbed. On the other hand, we can see that ( x + 1) = (3 x + 2) when 01 = 20,, 2w1 Wig It means that the scale of the bias term less important, which can be absorbed into other terms.
Thus, we only need to handle bias with different signs carefully.
Gradient Feature Distribution. We may define gradient feature distribution rather than gradient feature set. However, we find that the technical tools used in this continuous setting are pretty different from the discrete version.
30

---

Activation Functions. We can change the ReLU activation function to sublinear activation function, e.g. leaky ReLU, sigmoid, to get similar conclusion. First, need to introduce corresponding gradient feature set, and then we can make it by following the same analysis pipeline. For simplicity, we present ReL. only.
D.4 Gradient Feature Learning Framework under Empirical Risk with Sample Complexity
In this section, we consider training with empirical risk. Intuitively, the proof is straightforward from the proof for population loss. We can simply replace the population loss with the empirical loss, which will introduce an error term in the gradient analysis. We use concentration inequality to control the error term and show that the
term depends pnverse-polynomiilly on the sample size
Definition D.10 (Empirical Simplified Gradient Vector). Recall Z = {(x(1, "y)) ( for any E Rd b R, an Empirical Simplified Gradient Vector defined as
G(w, = Iy (9 [TTTT x(() xx())) > b]].
(127)
[E(m]
Definition D.11 (Empirical Gradient Feature). Recall 2 {(x(() (e))) be random variables drawn from some distribution W, B. An Empirical Gradient Feature set with parameters p, 7 BG is defined
Hegj) ler R. b E R
Bp.r.Bo (W.B) = (D.s) Pr G(w,b) b) CD., and [IG(w,b)lk BG and
b 16
p
When clear from contexl, write it as Sp.y.Be
Considering training by Algorithm 1, we have the following results.
Theorem 3.12 (Main Result). Assume Assumption 3.1. For any e, 8 E (0,1), ifm S and
4
Bb Bc
(log (2))
m =Ω
rBa Bz1
+
pe
vo
-Ω
VrBozBoBul
+ mb)
oog m BB
(mp)
B, (mp)
n
m² m² Bb
(mp)
Bra
B²
1
Ba2 |(((0)2
Tm
logn
2 2r2 Ba
B,Bg
Bz2
B%
B2,
then with initialization (8) and proper hyper- aaper-parrmeter values, we have with probability 2 1 6 over the initialization and training samples, there exists E [T] in Algorithm I with:
Pr[sign(fz( (x)) # y] S CD (fg()
BB2 log
+ + +r Bay 27 +
+
BGl (0))
See the full statement and proof in Theorem D.17. Below, we show some lemmas used in the analysis under empirical loss.
Lemma D.12 (Empirical Gradient Concentration Bound). When logn
with probability at B.a'
B.
least 1 - O (+) over training samples, for all i [4m], we have
(āz(J=)
(LD(=z)
0 |ai| B, log
(128)
dw;
Ow;
n
00 (fz)
@LD(fz)
[willz Bad log
50
(129)
da:
Oa:
n
||a|lo|la ||ll|||| oo (maxie
|| B. + b) + +1) v Vlog
Ez (fz) Lo (fz) 50
(130)
n
31

---

Proof of Lemma D.12. First, we define,
z 2 f=(x") do' ((w,,"") b,) x(1)
(131)
- -E(x,y) 1 (uJ=(x)) (o' ((w,, x) b:)] x]
(132)
As (2)
S 1, gl 1, I I, we have (1) zero-mean random vector with || 112 S 2B as
|| [g] - B,2. Then by Vector Bernstein Inequality, Lemma 18 in [66]. for 0 < 2 < Ba2
well as
have
Pr (1
@fz(f=) Ow,
@CD(f=)
2 |a,/
Pr
22
(133)
Ow;
n IE(n]
1
S exxp
-n 8
7
(134)
Thus, = n B, log ni, with probability at least 1 - (1)
have
(fz(
@CD(fz)
| [aily Bx2 log n
5
(135)
Ow;
Ow;
n
2
On the other hand, by Bernstein Inequality, for 0 we have
OCz(íz)
@Cp(f=)
Pr
> = wsll2
(136)
Øa;
0a;
C (1"" (x((0))) (or ((w,x()) - -b:))
Pr
(137)
[E(n]
[()))) (( -billl > =||wall2
(138)
12
<2exp
(139)
B₂2 + Bzz²/
Thus, when log
B²
let n- / BB2 log n. with probability at least - 0 (H), we have
Baa
9c (fz)
@LD(FE)
||w, 1 Brz logn
<0
(140)
oa;
Oa;
n
Finally, we have
Ez (Jz) - (D)
(141)
(1y0a ( x (CWTx(D-b)] -bi) - B(x.N)~D [e (yaT [o(WTx x x-W]x-b)]]] b)])]
(142)
1=1
By Assumption 3.1, we have aT (o(W x() b)]) E( [ (ya) (o b)])] is w. +b) +2 By Hoeffding's
zero-mean random variable, with bound 2||a|/o /aalolooo (max;e(4m) inequality, for all z > O, we have
z2n
Pr (
(fe) (fz) CD (f=) 2 2)) 2exp
(-(lallo [max w Wil|aB. + b) 1)2
Thus, with probability at least (t), we have
Ez (fe) Co (fz) so (Iallo/lallo
||sll2 + +8) + Vlgg
(143)
n2
32

---

The gradients allow for obtaining set of neurons approximating the "ground-truie -rutth" network with comparable loss.
Lemma D.13 (Existence of Good Networks under Empirical Risk). Suppose
log 72
B T(()) Let ( qlay. For any B. E (0, Bb). let oa = = PI
and 8 2re - Then, with probability at least - over the initialization and training samples,
there exists à, 's such that (aawao,, - 4mm ā.o (( (w!))) b:) satisfies
(D(((a,wo).b))
(144)
2B2 Bb
Ba 20 v
BB log Bale(0))n
SrBal
+ Bz1 Bay 2n
+ Be + OPT O....Tr.,,B... Be.. (145)
/mpBoB.
||ā||2 00 =0
and ||āllo = O (f(mp)²))
Ba b(mp))
||ā|l00 =0
Ba 6(mp))
Proof of Lemma D.13. Denote = and B 0 BBg log
Note that by symmet- ric initialization, we have (3)f00 (x)) [e'(0] for any X E X, so that, by Lemma D.12, we have ||ã(w) [0) b.) - G(w! (0) b,) I POOT with probability at least 1 p Thus, by
nz
union bound, we can see that Bp S S p-pott Tor Be ool Consequently, we Exactly follow the proof in Or Spr.i
have OPT drr,Br,3, y-r mii pal 1e/101
Lemma D.4 by replacing Bp to S
p--rrt De p Ba To
Then, we finish the proof by
S 2 [OO S (1 - (1-1/V2)Bc
We will use Theorem D.5 prove that gradient descent learns a good classifier (Theorem D.17). Theorem 3.12 is simply direct corollary of Theorem D.17. To apply the theorem we first present few lemmas bounding the change the network during steps.
Lemma D.14 (Bound of =(0) _() under Empirical Risk), Assume the same conditions as in
Lemma D).3 and d 2 log m, with probability at least 1 8 m O (m) over the initialization
and training samples, ||a (0) a0 O
by/og [(0
) and for all E [4m]], we have (0)))
am (o Finully, ||a(") = (n) (B))
+ (2) awvdB log ) and for all
n
i [4m]. 12 0
bulo Bg
byTog m. logn
Proof of Lemma D.14. The proof exactly follows the proof of Lemma D.6 with Lemma D.12.
Lemma D.15 (Bound of (t)) under Empirical Risk). Assume the same conditions in Lemma D. 14, and let r) (t()) for all E (2,3.... 2......,TT, T}. 0 < TyBB I o(1), and 0 = X X for all over training samples. for all E [4m], for
t t {2, ... With probability at least 1 - O ( Tm n
all E (2.3....,T) we have
T),
b
(1) 50
|a. |a ||w
+
+ 1
(146)
Bun + vB.2 log n2
Iw(")
w!
(2) 50
tn
Br log Bu1
(1) + tn?
Buy +
v BBz logn
2
l|
(
l12
n
n
v B log
+ trp
BBl +
6
(147)
n2
33

---

Proof of Lemma D.15. that, we have
The proof exactly follows the proof of Lemma D.7 with Lemma D.12. Note
(4) la
a
[t-1)
(TBBB1//WW (
(1-111122 112 1 +b)
| (t-1)
22 Buz log
(148)
n
+ 17 Bu1 /B22log logn
(t-l)
|| (w("-!)= w/" 112 + 1Z
a
(149)
n²
where we denote Z, = Bx1 + /Bza log
l/w; (2) 12 + b. Similarly, we have
n²
B2 log n n
[ wi (1) < s|lw (111)
w (1) 12 + Bx1 +
aa 1111)1
(150)
We finish the proof by following the same arguments in the proof of Lemma D.7 and union bound.
O
Lemma D.16 (Bound of Loss Gap and Gradient under Empirical Risk), Assume the same conditions as in Lemma D.15. With probability at least - O (2), for all t E [T', we have
Ezw (fl,wwo.b)) CD(f(a,wa),))
(151)
(max:e l|w!") (maxsell 12 + b) + +1) Vlogn
00
(152)
nt
+ Bu1 a 22 aallo
max || (1) [ám]
w (1)
(153)
With probability at least - O (L) for all E [T], i [4m] we have
ac dāzu (fEw)
|| (w(""l/v 1/2 VllVBBe log
||
< BBi (1) w1 1111/ [|w (l)))))) +6+
(154)
da!"
n
Proof of Lemma D.16. By Lemma D.8 and Lemma D.12, with probability least 1 O (2) for all
I EE [T], we have
Ezm (fa, - Colf(a
(155)
$ Ezo (Ka,wo)))) Lo(l(a,wo)..)))) bi) + LLD ((,w,,))))) LDll (a(()a
(156)
||ā||o/|a|loo (maxig|am)
( ||w (1) 122B - 5) + +1) Vlogn n
so
(157)
+ B ||ll2
|allo
max sE(mm]
||
w
(1)
112.
(158)
By Lemma D.8 and Lemma D.12, with probability at least 1 - 0 (²), for all t E [T]. e [4m] we have
z.. (JEw)
|| |Iw!" (1 2/ Bug logn
< (ll
w. + | +b)
(159)
da(") (1)
n
We now ready prove the main theorem.
Theorem D.17 (Online Convex Optimization under Empirical Risk. Full Statement of Theorem 3.12).
Consider training by Algorithm 1, and any å (0, 1). Assume d > log m. Set
0n > 0, > 0, (1 7) A 0.for E 3..3..... .T..
(1) = min(O(n), ((b))}
(B00 20) 1
b(mpp)
= 0
(Tj
(1) BaB
34

---

Let
T1
o(1,
m
2 (
+ (log (5))²) and
log n
1 1
With probability at least 6 over the initial- [T] such that
B
ization and training samples, there exists
LD (fg.)
(160)
2V2Bml (mp)
Bb
v VB,2 logn
<OPT d.r Sp.. + rBal
V
Be
Bul
2770
Bale (a]((()))n3)
(161)
Vlog m mB (mp)
B8
+ 11 (ViBaz +mb)
+1
+ 0
(162)
v BB,Bo
1/TB2 (7pp)
Vlogn
(r
B.
bylogm(mp)
0
(163)
nă
BB
Bz1
(
Diylogmin m(mp)
+ Boi b
Br +b
+2
(164)
BBBG
log
O mn
bylog m()))
TT) B21
v B22
(165)
n
VB,Ba
Furthermore, for any E (0, set
B BB
Bb
1
1
log (7))
b =
VrBui )
= Ω
rBal rBaaBBst Bul
+
pes
BG
v8
p
1) =
+ mù) (
T
(mp))
JB B. Ba1
+1 VByBa
(mp)
mn mmppB By (log m)²
(mp) Br2 log
B
1
1
B,d
Tm
=Ω
logn n
222 222 Ba
BBBG
B
p
B
B21
[e((0)2
have there exists E [T] with
(x(sign(/=w))xx) + v) SLD (Jzw)
(166)
VBa log Ba (0) (n
a Sp. + rBaal Bay
27 + +
+ (167)
Proof of Theorem D.17. We follow the proof in Theorem D.9. By Ω ( (log (t))²
V8
and n
log n
Ω
B
B²
+
(0)=
B.g
we have 2re m + 0 [mm S 8. For any
B. (0, Bb), when da = θ "(())")) Ba B, by Theorem D.5, Lemma D.12, Lemma D.13,
35

---

Lemma D.16, with probability at least 1 - 8 over the initialization and training samples, we have
T Lp (fs@)
(168)
==1
T
ICD (fsw) -2g (fww)) 1
1
Ez (fee
(169)
T
T
Σ
LD (f=w) Ezc (fzo) 1+
Σ
Eze (fa.)o))))) - CD(f(a,
(170)
T
aal2
θEzw (fz(e)
+ +CDls(,wo.b)
+
+ (2 (1) 12 + 47m)
max (E(TT)
(171)
2nyL"
Da! (1)
<Bai|áll2
a
max Em
w (T)
w;
(1)
2
(172)
||@lo||ā|| ml (T) bc) (maxie(4m] ww (T) 112,, +b) + +2) Vlogn
+0
(173)
n
2B Bo m.
v BB log
+ Or + rB
+ BBI B
+0 20
+ B.
(174)
BG (0)]mz
||
+ 44mBa1((aa" +
(175)
2nyT
max
l/w (T) -w
w/(")22 112 max
|w 12 +
(1)
+ O
max [4m]
|| (1)
2 Brz log
16(4m]
=E|dm]
Ba1
Bunt
we have 0a - 0
86mp) 1
Set B.
Ba
2Bg
(=r(om)" B.i BB
which satisfy the require-
(mp)
ments. By Lemma D.13, Lemma D.14, Lemma D.15,
Ω
B2 (0))) )), when 17(1)
Bug
log n
min) we have
||l|o =0 (1(mp))).
|āل|2 - ( (( |ā|| =0
[(mp))
b(mp))
Ba Br
(176)
/dB log
|a (1) 1o00 (10 Vd 8) (1)
(177)
n2
= =min(O(n).0(T))
(178)
[Ta 50 ||a"8)|0
+ max || (1) 112 + =E(mm]
+ mb
(179)
Bu1 + BBa log
50
max
(0 ||2 + Ba
(180)
(c(4m]
36

---

bllogmBar
bllog my BB2 log
max
l|w;
(1)
112 =0
(181)
+E[4m]
BaB.
(e') BB B,
bllogmBe1
(182)
BeBe
/oog (mp))
=0
(183)
BB
v Bxz log
max [!w
w (1)
112 =0
Try
B
a;
(1)
(184)
:6[4m]
n
BB2 log B21 +
TT/
Iw (1)
(185)
n²
/Ba2 log
+ TT2
By1 +
(186)
nz
=0
Tr22
max =E[4m]
(2)
1|2 +
(187)
Ba
Then, following the proof Theorem D.9, we have
LD (Szw)
(188)
222BBm
By
VB,2 log
<OPT Br Br. Spring + rBa
(mp)
V BG
+ Buy
27 +0
BG (0)
(189)
Vlog AlogmBz1 (mp)²
B3
+1 (VrBa2 mb) 0
+1 +0
(190)
VB,BG
r/Ti2 (mp))
(|||||||||| mllaa [1| (7)
(max |aa l|w 112 + b) + +2) Vlog n
+0
(191)
n
| 112 B log n
mn ma
+0
(192)
n²
2v2Bml
Bb
v BB logn
<PP d.r. Br Sprino + rBal
(mp)
V
Be
Bud 27 +
Ba (()(e())
(193)
Vlog m mB (
B
+ (VrBaz + mbb)
+1
+ O
(194)
VB,Ba
n)Tb2 (mp)²
Vogn
(( (rBat Bb
bmpl BB
o
(195)
n
B
(( Bllp)
Bir +b
+2
(196)
BB
Vlog n nz
(Vllgg m(mp)+
O mm
B21b
VBz
(197)
VB,Bo
37

---

Furthermore, for any E (0, 1). set
BB Buz B
Bs Bc
1
1
6
= Ω
Ball V
log (8)
Vr
pea
V8
p
c
1) -
T = θ
VrB..B.Bal Bat
+ mb]
Vlg B,, (mp) BB
(BB)(mpp)
(mp)
n
m pB a B. Bb (logg
(mp) $ B22 log m
B
1
1
1
Tm
=Ω
+ +
+
+
logn
2²T² 2r2B2 BG
BG
BBBG
B
p
B
B
[e((0)2
8
and note that BG S Brg S Bz and v Brn s Bg naturally, we have
1
LD (J2(o)
(198)
T
2/22m1
Bb
BBz log
<PT d. Sp
+rBa
+ B 27 +0
(199)
(mp)
Bg
Bale((0)
Bu BagBB. BB
Vlogn
mB Ba2v /B, (Bpp)))) log m
+
(200)
2
62 (mpp
n2
rB.1 v /B
Vlogno
Cv By2 log m (mp)
(201)
n
VB,Ba
v BB2 log +c. Ball(())n
<OPT dr,,Brr.BBB....... Br.Sp + + rrBuBBiy 27 +0
(202)
We finish the proof as the 0-1 classification error is bounded by the loss function, c.g.. [[sign(((x)) # vl < f(yf(x)) where E(0) 1.
((0)
E Applications in Special Cases
We present the study of linear data in Appendix E.1, mixtures of Gaussians in Appendix E.2 and Appendix E.3, parity functions in Appendix E.4, Appendix E.5 and Appendix E.6, and multiple-index models in Appendix E.7.
In special case applications, we consider binary classification with hinge loss, e.g., {(z) = max(1 2, 0}. Let X' Rdd be the input space, and y {=1} be the label space.
Remark E.1 (Hinge Loss and Logistic Loss). Both hinge loss and logistic loss can be used in special cases and general cases. For convenience, we use hinge loss in special cases, where we can directly get the ground-trun truth NN close form of the optimal solution which has zero loss. For logistic loss, there no zero-loss solution. We can still show that the OPT value has an exponentially small upper bound at the cost of more computation.
E.1 Linear Data
Data Distributions. Suppose two labels are cquiprobable, i.e., Ely -1] Ely +1 1
The input data are linearly separable and there is ground truth direction w*, where || = such that (w".x) > O. We also assume w the complementary space of the ground truth, i.c. the components of input data being orthogonal with the ground truth independent of the label y. We define the input data signal level as p== Ely (w*, x)) > 0 and the margin as 3 == min(x.)) (w" x) >
0, where P is the projection operator
We call this data distribution Diinear.
38

---

Lemma E.2 (Linear Data: Gradient Feature Set). Letb d' Be1oT. where T is any number large
enough to satisfy d=/2-
> 52
yBa1
For Disncor setting, we have -1) E Sp,,Bo where
1 p 2'
vBa2 pdt
/B22
7
BG = p-
(203)
d=/2-
Proof of Lemma E.2. By data distribution, we have
B(x.n) [yx] = pw'.
(204)
Define Ssure: : {i E [m] I[; 10) 122 S 2v 2/ }. For all i E [m]] we have
Pr(i E Ssure] ? 1 1l2 S 2 do 2
(205)
2
For all E Ssures by Markov's inequality and considering neuron + m. have
Pr x
(
(w (wifm=+ x) - bitm <0
0
= x
Pr
w
+ b, <
(206)
< Pr
W
1/2 |x|x|12 N bi
(207)
x
d= Bu1
Pr
|x||2 2
(208)
2
<
(209)
d d=
For all E Ssures by Hölder's inequality, we have
E (x.g)
y/ 1 0'
K w wwit x - bi+m
(0)
x
(210)
1 o' K (w!") : xx + +b²] x
(0)
E (x.,)
(211)
[[[[ 1
(1 (1-0') o'' [<m:0) .x) + +b:]] )
(0)
(212)
VBz2
<0
(213)
/2-2-
We have
G(( (0) i+ms b b,-m): wr
(G(w (0) -b;), w'
(214)
[[[ iimo berm) l12
[[((w (0) -b:) 12
<1 p - θ
vBag
/r/2-3 3
(215)
p θ
vBeg
/=/2-
VBz
=
= 7.
(216)
pd
We finish the proof by
by-m [b:-ml
1.
Lemma E.3 (Linear Data: Existence of Good Networks). Assume the same conditions as in
Lemma E.2. Define
1 f*(x)
1 xx)
3 (
B [o(--"xxx) x)).
(217)
For Diinear setting, we have fz Fy. where r = BF (Ba1, Baz, Bb)
1 7 0
BB
Ba = p- 0
B
We also have
B P B
pd x
/122
OPT d. Br Bp.Spo.B 0.
39

---

Proof of Lemma E.3. By Lemma E.2 and Lemma F.3, we have f E F..r.Br.Spnn
We also have
Br. - D (f*)
(218)
= ((xx))))innco (fx)))
(219)
0.
(220)
Theorem E,4 (Linear Data: Main Result), For Diincar setting, for any d E (0,1) and for any 6 E (0, 1) when
poly ( se.
1 1111)
T= T= poly (m, Ba)) n= ==poly poly m. B 8' 8''
(221)
trained by Algorithm with hinge loss, with probability at least over the initialization, with proper saperrpprrameter, there exists t E [T] such that
Pr[sign[Szww (x)) + v) Sc.
(222)
Proof of Theorem E.4. Let = d" Br where T is number large enough to satisfy 22/2-k
Q
B2,2 and O
B (
s 2
By Lemma E.3, we have f E Fd,r, Sp.
where
r = 2, Bp = (Bal, Baz, Bp) =
2 .
- - -
BG = = pd
B
D
B
We also have d,r Sy.ni.
R n Injecting above parameters into Theo-
Adjust ow such that b d" d' = O
VrB.i
rem 3.12, we have with probability at least over the initialization, with proper hyper-parameters, there exists E [T] such that
B B
B Br1 (logn)d + t (223)
Pr|sign(fevo) (x)) # vl 50
82
+0
B /
E.2
Mixture of Gaussians
We recap the problem setup in Section 4.1 for readers' convenience.
E.2.1 Problem Setup
Data Distributions.
We follow the notations from [99]. The data are from mixture of high- dimensional Gaussians, and each Gaussian assigned to one of two possible labels in y = {+1}. Let S(y) S [r] denote the set of indices of the Gaussians associated with the label y. The data distribution is then:
q(x, v) = (x)))(xly), g(xly) = Σ (x))
(224)
JES(x)
where (x,x)) is multivariate normal distribution with mean (4 and covariance Ej, and py are chosen such that g(x, y) is correctly normalized.
We call this data distribution Pmiiature
We will make some assumptions about the Gaussians, for which we first introduce some notations. For all E [r]. let v()) E (+1,-11 be the label for (x)..
D, Hj 12sll2''
Bo 1//v. B, min مäbl jEl
B,2 max |,b pa = min Pj- jElr
JEl
Assumption F.5 (Mixture Gaussians. Recap of Assumption 4.1). Les 8 S r S d be parameter that will control our final error guarantee. Assume
40

---

Equiprobable labels: q(-1) g(+1) = 11/2.
For all j E [r] Ejj = oj =0laad Let OB = max oj and Ø = max Buz)
S 2d,
pa 2 Ω
ran+? log
Bpt Bp S d.
The Gaussians are well-separated: for all E [r], we have -1 (Di, Dj) where
0.05 min 2r
log d
B2
Below, we define sufficient condition that randomly initialized weights will fall in nice gradients set after the first gradient step update.
Definition E.6 (Mixture of Gaussians: Subset of Nice Gradients Set). Recall w; (0) is the weight for the i-th neuron at initialization. For all j E [r]. let Sp,, S [m]] be those neurons that satisfy
wi
Mi
2
Csure,1b,
I 2b for all j' # j.J [r].
@(Vdow).
E.2.2 Mixture of Gaussians: Feature Learning
We show the important Lemma E,7 first and defer other Lemmas after
Lemma E.7 (Mixture of Gaussians: Gradient Feature Set. Part statement of Lemma 4.3), Ler rrdlooo log do BB+. where Cb is large enough universal constant. For Dmiuuture setting, we have (D,++) E Sp for all j E [r]. where
Csure,1 3. Csure,2 2. b CoVTd
Ba
-
VT Volog d B4
( 2/(282,))
7 10.9--13
(225)
BG PBBBIVV -
OB+
(226)
N.9r
Proof of Lemma E.7. For all j [r]. by Lemma E.10, for all i Sure
(o(w.")
(b.), .b))..))
(227)
[[G(W(".b) (0) .b))ll
K (G(w (0) ([( b.),
<1
(228)
(G(W b.). DJ)
+ maxxp -0-00111+12-11
(G(w ,b)))))))))))))))))) bi).
(229)
(C(w!)) ,b))
G(w, (0) b,), D.)
+ max =0.11D 11==1
<1
(230)
B
+oo+ m
B,Baiad (100 +
Bns0
-0 O (
--
TB
(231)
pj BBB Vd - 0B+
(0.9r
1.9 7,
(232)
where the last incquality follows BaI 2 S
JB+
T log
41

---

Thus, we have G(w; (1)
b.) .b.)
CD and (G(w.") bi). DJ) = [[G(w; O b.lll2 - Br1. b
[b,]
+1. Thus, by Lemma E.8, we have
G(w,b) CD2,, and [|G(w, b)] 2 Ba and b
Pr
+1
(233)
w.b
[b]
2 Pr i E Spre
(234)
2p.
(235)
Thus, (Dj:+)) E Ba. We finish the proof.
Below Lemmas used in the proof of Lemma E.7. In Lemma E.8, calculate p used in Sp P.r,Ba Bi- Lemma E.8 (Mixture of Gaussians: Geometry at Initialization. Lemma B.2 in [7]). Assume the same conditions as in Lemma E.7, recall for all E [m), (()) ~ NO, oo Taxd), over the random initialization, have for all i [m],j j E [r].
B,1
Pr [i SD 2 O
(236)
VT log do ( 2
Vog
Proof of Lemma E.8. Recall for all / E [r], ju vI/Vd.
WLOG, let j т. For all I E [r- 1]. We define Z1 E Ir 1] (D), 2 0 and
Zg = {1e 1] -1 < (D,, D,) < -0}. WLOG, let Z [r1]. Z2 = (ri + 1,...,72 r2). where
0 m S S We define the following events
C = (w! M
- Cswre.bb) & {1< w (0)
11) < Csurezb.).
(237)
We define space A span (HI...., (r,) and fa = PA-MA, where PA is the projection operator on the complementary space of A. For l E Zzz we also define ju 14 1l 1
(r).))) and the event
à = (( (0) ia) - Csure & = [[( i1) S Osure
(238)
For E Zz, we have sr = jm phy, where p 2 0. So (w, = (w, - p(W, S (w,
when (w) 2 O. As a result, we have
š { w (..,.,)) 11) 2 bi SGn((w w/
.14) (w/"
Csure.1bi.
(239)
By Assumption 4.1, we have
1 < 1 r0 < 0 <
/ 1.
(240)
|
We also have,
Pr
w
(0)
H4)
2 Csure, bi, St, St....,... S-
(241)
Pr
(0)
14
> Csure, bis S1, Sl,....,rr]
(242)
Pr
(0)
> Csure, by, S1,
S Sritt........ Sr2
(243)
Pr
(0)
> Csure, bi, EE.... Ene E feith 8
(244)
Pr
(0)
144)
2 Csure, by Gi,... Sra & Sr1t1.......................... 2,
Pr G (..... & ..........
42

---

For the first condition in Definition E.6, we have,
(0) Mr) 2 Csure, by E1 ............ Sr. Sr S Sra]
Pr
Pr
(245)
Pr
pr Mr - (tr) 2 Csure,1 bi (1....,,r.]
(246)
2 Pr
Pr - My Pt) 2 Csure,
w (0)
Hr - Pr)) 2 0f2....,..
(247)
Pr
14 + Mr pr) 2 b
: , - Mr) ) 2 0, ,,.,,...,cr.
2n]]
(248)
W Mr - Br) 2 0 0(,,....,Gr 6m]]
Pr
(249)
(0) Prr + Mr - -pr) 2 Csure, bolrer.bb/(( (w/( 14, Pr)) 2 0, O.
Pr
(250)
(0) [rr 2 Csure, by| ( (w/0) =H4, -4 R..) 2 20. S...,....... Gm.
Pr
(251)
Pr
(w!)) PTr)
2 Csure, bi
(252)
2
1|rr
20 log do B+ ( 2
(253)
where the last equality following that ft, is orthogonal with Pi....: Mry and the property of the standard Gaussian vector, and the last inequality follows Lemma F66
For the second condition in Definition E.6, by Lemma F.6, have,
P1 = Pr [G]] 1 - 0
|
(254)
VT log do dl (Czron+ 2 2/H|| 1))
P2 = P [3][a]] 2 Pr [ 2 1
1
(255)
Vt log do d (Ctroz (801
(256)
Pr-1 Pr (5r Ei,..., Sr. {r2] > Pr [ixz]] > Pr [5m)]
(257)
11/2-1112 Volog B+ (CTTon,
2 1-0 A
(258)
On the other hand, if X is (k) random variable. Then we have
Pr(X 2 k + 2 kx + 2r) S =0
(259)
Therefore, by assumption B, 2 Ω JB+ v Flogd
we have
Pr w 2 d + 2/ (9CTTo842/(2133) + 2) dlog d
(260)
+ 2 rBB
2 222B + 2) log d
(261)
50
(262)
d2 (9C33000 (0 2/(28;))
43

---

Recall Bμ1 = je(re [ Buz = maxje(() 18.l12. Thus, by union bound, we have
Pr Pr[i [iEE Sup,
(263)
> he/r)) - 0
(264)
d2 d 2
(1-
Bpl
TB
20
(265)
logdo (12B31)
VTlog do (
(266)
d2 (1CCTTBB"
Bp1 Volog do B+ d 100 2 ((2B))
20
(267)
In Lemma E.9, we compute the activation pattern for the neurons in SurrSuee-
Lemma E.9 (Mixture of Gaussians: Activation Pattern), Assume the same conditions as in Lemma E.7, for all j E [r], SD,,Suee we have
(I) When x Ng 0 Taxa), the activation probability satisfies,
Pr [(w. .x) -b; 2 -b,20] 20 d=
(268)
x
(2) For all +' j.j' E [r], when Ny E.r)) the activation probability satisfies,
(w/))) (0) x) b, 2 S 0
Pr x-Nylogry
(269)
d
Proof of Lemma E.9. In the proof, we need b = Cov o where Co is large enough universal constant. For the first statement, when x (uj of Idxx)), by Csure,1 2 3. have
BB.
Pr
(w: (0) x)
b, 2 0
Pr Ni0.
(w! (0) xx)
(1 Csure,1) (270)
2 b, :xxxx/
(0)
Pr
xx)
(271)
x~N(0.03
b;
Pr
W x
S
(272)
x~N(0,,,
b,²
21 exp
-exp
(273)
Oido
21 01-0 d"
(274)
where the third inequality follows the Chernoff bound and symmetricity of the Gaussian vector.
For the second statement, we prove similarly by 0 < Csure.2 S 2.
Then, Lemma E.10 gives gradients of neurons in S Suree It shows that these gradients are highly aligned with Dj.
Lemma E.10 (Mixture of Gaussians: Feature Emergence). Assume the same conditions as in
Lemma E.7, for all j E [r]. i SD,,Sures we have
(E(x.n) (vo'
(( (w/")) (0) xx) - b.) x U(o D2)
(275)
2p, Bun BuI vdd
1- dF )) By20
0B O
(276)
d=
d0.97
44

---

For any unit vector Dj which is orthogonal with D,, we have
(0) xx b.) x D3) =B ( (EEo0)) [oo' ((w,")
1
1
(d-
+
+OB+0
(u.95
(277)
Proof of Lemma 10. For all j E [r], i E Sures we have
Bxxx)) (uo' ((w/") .x) - -b:) x
(278)
Σ Pillx~N((x) (aal ((m; 10) xx) - bi) x]
(279)
(E(1]
P1 1100000 (o' ((w:"0 x + m1) - -b:) (x (x+ mi)
=
(280)
Thus, by Lemma F.7 and Lemma E.9,
(Exx0) Ioo' ((w,") xx) bi) x D3)
(281)
Py Dy.Ex-N(0..ss /o/ ((w,") (1) x + p2) ) -b.) (x+ + pby DD]
(282)
Σ D1 FU (d' ((w,", x + mi b.) (x + pc) DDj]
00
+
(283)
(E(r]
2pzu, Dj (1-0(#))
Σ MilTT D2 (
1
(284)
et)
(E(T),1
PPj |Ex.ViOss (0' ((w.
, +1² - -b:) x D,,]]
(0)
(285)
E Pi [Ecur (o' ((w))) (0)
x + m1) ) b.) x
xTD]]] Dj
(286)
1E(T),yS
2p;BuiVd (1-0 1
(1- B
1
(287)
d=-
Pj (x.N(0,o3)1) [[1 (1-0' (1-0(( (0)
.x+ (4)) )
- -b:) 1 x D.
(288)
-E pi |R [o' (ww/)) x m1)
(d((
(0)
b.
(289)
-p, vā (1-0 (2))) Byz0
1
(290)
d=
PPj BE (11 (1-0' ((w:0)
(<w,0x++11 x ms)
- -b:) x D3]
(291)
Σ pt |EN-NOoo (1r/))
(c((
(0) x m1) b, x D3
(292)
(E(r),
(1-0(2)) -B,20
( 1 -OB+0 0
2p;Baiva 11- or
(293)
d
p..r
45

---

For any unit vector Dj which is orthogonal with Dj, similarly, we have
( (Eixmn) (vo' ((m/"),) - bi) x DJ)
(294)
Ex-Aiio,sn11 (o' ((w:)) x + u2s) bi) D
Spp
(295)
Σ Pt
Bxm (o' ((w!") (0) x mi) bi) (x + mi) T DJ
+
(296)
1E|r)/*)
Pj |Ex (0' ((w!"),x
(o' xx+ n1)) b.) x D
< B2
(297)
d=-
+
Σ
pt
(o' ((w!") , mx) bi) x D2j
(298)
1E|r)/*i
1
1
SBp20
+
+OB
0
(299)
d--
(0.9T
E.2.3 Mixture of Gaussians: Final Guarantee
Lemma E.11 (Mixture of Gaussians: Existence of Good Networks. Part statement of Lemma 4.3).
Assume the same conditions as in Lemma E.7. Define
f(() Σ
yGn V log do
o -
- -2 dog
(300)
j=1
For D Dinssture setting, we have for Farr,r where BF (Bal, Baz, Bb)
Vr
B, 2 dogl P O Vologdon+ d (
(Flogdozs VT log dog
=
-oot-13, BG = PB B d O
7al (grot)
and BB = (B + BB) Vdd BBg = (Bp2 Bd.)
We also have dr Mo I 3
dr
0 1 Vlodd
Proof of Lemma E. We can check B.1 (Bp2 + +O++) =Bpz + OB+)dd by direct calculation. By Lemma E.7, we have f Fdr..r..Sp...e
For any jE [r]. B,d 2 Ω 0B+ / /zledd /
> 40 Alord we have
Pr x-N, laxa)
(Dy,x) 2 log do 2 V log dogd
Tog
(301)
Pr
(Dj,x) + +1 2 V log d
(302)
x~N,((0, laxa)
Pr
(D,,x) + VdB,i1 2 do 2 v V do
(303)
(0, laxa)
2
Pr
(D,,x) 2 - -V log do
(304)
Nv,,00, laxa)
(305)
where the last inequality follows Chernoff bound.
46

---

For any / i j, E [r]. by 0 -
ast
Flogd
we have
B,
d
Pr
(D(,x) - 2 20
(306)
XnNg(Hs, laxda
(Di,x) +0B 2 2 0]
Pr
(307)
x~,3500, laxa)
Pr
(Dixx) 2 Flo d+
(308)
(U,,(U,, Laxd)
1
(309)
d
Thus, have
Pr (x) > 1]
(310)
.x.y)/DD intture
2
Pr
(Dj,x) - (Dj,x)-2V45
dog 1 2 v log
(311)
jEIr
Σ Pj
Σ
Pr
(DIx)) 2/ T log do <0]
(312)
jE(r)
[A)JE(()
2
21
(313)
d"
We also have
B(x,u) (aau))/Dmittt
[x) [1]1)[(x)]]
(314)
v / log do B+ B
(Dj,x) -2 log n < V do
Pj
Pr
x~Ng(Ms;;14xaa)
V log do
jE(r)
+ Σ jE(r]
Pj
Σ
Ex-,s(es,
0f ( - -2 (o Bog >0]
(D1,x) - 2V log do VT log do
dr
1
+
Σ
Pj
Σ
Ex-N5(0.,
o' [D(,x) > V Bo
(Di, x)
V log do B+
V log do
jE(r]
ItjLEtr
di
1
/ log d
1
Σ
Pj
0 [Dr:x) > V Vilogd (Dixx)]
(315)
JE(r)
1
4
(316)
dr
0.9r VT logd'
where the second last inequality follows Lemma F.7 and r S 2d. Thus, we have
OPT d,y, By Bp
< ((x.y)
[x()]((x)) Hyy (x) 1] (x))
(317)
= -E(x,))~Daaee
(318)
(x,y) (.x.y)-Pmmsturr
[x) 1)l//((xx)]
(xxx)) [llf((x) S 1]]
3 < dr
4
+ 1 V
(319)
Theorem 4.4 (Mixtures of Gaussians: Main Result). Assume Assumption 4.1. For any c, E (0.1).
when Algorithm 1 uses hinge loss with
m -poyy 8' (ropt+ 2/Bb,),
1
< ed T = poly (m), p (m)
Pg
47

---

and proper then with probability at least 1 8, there exists E
[T] such that
V2r - +
Pr[sign(/E:oxx) # yl
Proof of Theorem 4.4. Let = CoVTrloog w where Cb is a large enough universal constant.
By Lemma E.11, we have f BrinBa.S where BF
(Ba1, Baz, Bb)
Vologdogg VT log (VTlogdaB+) P
B₁
dog a ( (
7
20.3--13, BG pB Vd 0 0 (0.35)
and Ba (B oB oB-)Va, Ba (Bμ2 Pd.
We also have drr Bp. Mp S dr + d0.0--1 'Flogd'
B Bu
Adjust on such that CoVTd o
Injecting above parameters
Vr
into Theorem 3.12, we have with probability at least 1 - ā over the initialization, with proper hyper-parameters, there exists E [T] such that
Pr[sign(Szmn) (x)) + ul
(320)
3 dr
/2rB,e
B (rBaiBe Bp1 B22 (log n)
+0
e/2
d0.9r-1 V
log
d0097- 155/2 T log do
VBn
2r d0.4r +
<
(321)
E.3 Mixture of Gaussians XOR
We consider a special Mixture of Gaussians distribution studied in [99]. Consider the same data distribution in Appendix E.2.1 and Definition E.6 with the following assumptions.
Assumption E.12 (Mixture of Gaussians in [99]). Assume four Gaussians cluster with XOR-lik pattem, for any J
= 4 and p1 = P3 = P1 = 4.
#1 = - U2 #3 = - 14 and ||11|22 = || = ulll2 = uall2 = Vd and (M1 us) =
For all j E [4], وi = σ laxd and S OB S Flo Fooglog log
M(1) = F(2) and y(3) = 4(4) -1
We denote this data distribution as D Dmizture-zor setting.
E.3.1 Mixture of Gaussians XOR: Feature Learning
Lemma E.13 (Mixture of Gaussians in [99]: Gradient Feature Set). Let Csure,1 = 6 Csure,2 = loglo b = Vrd log d and large enough. For Dmisture-zan selling, we have (Dj. +1) E Spr.Bo for all [4], where
dB 7
p= 0
(322)
18r0
Vd
V Vol log do (log d)
Vd (1 (1-11-0)
(log d) d) ))
Bg 4
-
0B0
(323)
(log 00.088
48

---

Proof of Lemma E. For all j E [r]. by Lemma E.16, for all E Sp, Sure,
(((w.") (0) b.), D;)
(324)
|[([w]) ,b,)ll2
(G(w;)) bi), Di)
21- ( b.),
(325)
(.), D;) 2 + maxp D. -0 G(w! b.). DJ)
(G(w)) (0) bi), D3)
(326)
(o(m:") (l(w.",.)).)) + maxp] D -0
112=1 ((w (0) b,), Dj
31
(327)
1+
0
Flx 017
t vi
1-0
-nn
(lox 57
(log 2jo..10
σBO
(log dju.D18F
(328)
Vd 1 -
000
1 (log d50.0187
(log d) st
<
7
(329)
vd
Thus, we have G(w! (0) b:) E CD2, and (G( b.). Di) < ||G(w;) (0)
b:)lla s B b
[b.]
+1. Thus, by Lemma E. E.14, we have
G(w,b) CD and ||G((, 2 BG and b
Pr w.b
|b] +1
(330)
Pr
i E Sp
(331)
pp.
(332)
Thus, (D)++1) Sp., We finish the proof.
Lemma E.14 (Mixture Gaussians in [99]: Geometry at Initialization). Assume the same conditions as in Lemma F. Recall for all E [m]. w; NO, og Taxa), over the random initialization, we have for all E [m],j (4).
Pr i Su 20
(333)
1 V log log do (logd)
25
Proof of Lemma E. 14. WLOG, let = 1. By Assumption E. for the first condition in Defini- tion E.6, we have,
Pr (w. (0) M1 Csrre, by 20
(334)
20007 25
V log log do (log d)
where the the last inequality follows Lemma F.6.
For the second condition in Definition E.6, by Lemma F.6, we have,
Pr
[ (w:") (0)
(22)
< Cs 21 2/T OB σ
(335)
On the other hand, is a (k) random variable. Then we have
Pr(X 2 k + 2 kx + 2r) S =
(336)
49

---

Therefore, we have
18ToB
18TT
Pr
(0)
d
2
2
dlog log d +2
+ +2)
log logd
(337)
2
25
25
1
<0
(338)
1 25
(log d)2 (log d)
Thus, by union bound, we have
Pr i E Spy, 20
(339)
18 25
V log log do (log d)
Lemma E.15 (Mixture of Gaussians in [99]: Activation Pattern), Assume the same conditions as in
Lemma E.13, for all j [4], E SDr,S we have
(1) When x ~ Nz(15, TBIxxa). the activation probability satisfies,
(w
b, 2 0 2 (log d) 50
Pr
(340)
x-Ng(05:0
(2) For all # j.3' E [4], when x ~ Nyrlyy (Blax)) the activation probability satisfies,
K ww x - b, 2 ol s so
( (log d)
(1og
Pr x-Nglugs
(341)
Proof of Lemma 1. the proof, we need VTd log d B. For the first statement, when x ~ N (04, T Ixxa), by Csure,1 2 g. we have
Pr
K (w(,x)
-b, 2 20
Pr N10,0nlaxa)
w
2 211 1b
(342)
x-Ni(15,08lxx)
b;
Pr
w
xx)
z 5
(343)
(w:0)xx
b;
=1
Pr
S
(344)
x-N10.7p
5
b 2
21 exp
(345)
50dl 2
21 (log d) 50
(346)
where the third inequality follows the Chernoff bound and symmetricity of the Gaussian vector.
For the second statement, we prove similarly by 0 Csure,2
/2
/ log log
Then, Lemma E.16 gives gradients of neurons in Sure It shows that these gradients highly aligned with Dj.
Lemma E.16 (Mixture of Gaussians in [99]: Feature Emergence). Assume the same conditions as in Lemma E.13, for all j E [4], i E Sa,,Sares we have
( (Bix.v) [80' (( (ww!0) x) b:
x (355)
(347)
2 4 Vd 1 1-0
(log d) BBO
(348)
(log d)o.0188)
For any unit vector D which is orthogonal with D,, we have
(IEROW [VO' ((w,0) xx - b.) x D Sa00
(log djo.0187
(349)
50

---

Proof of Lemma E.1 For j E [4], i E Suyre we have
B(x,y)
(yO' ((m!"0,))
(0) xx)
- b,:) -b;) x
(350)
BxmNi(x) (yo' (nw:" xx) - -b:) x
(351)
4
|E[4]
Uin (x. (o' ((m:0) x+ mi) -b)) (x + mi)
M())
(352)
[E|4]
Thus, by Lemma F.7 and Lemma E.15,
(ool ((w!") (0) .x/' -b:) x . (G) ()))
B(xx))
(353)
4 Exwsiooondaxa (o' ((w,"), x + Mi) - b.) (x + uz) DD]
(354)
Σ (UNIIS (o' ((w, x m1) -b.) (x + un) Dj)
+
(355)
1E|S),**
1
(1-0(105)) (Iogg))5t) )) Σ
1 1310 (13)
2 3
Dj
(356)
d5
RE[4),43
4 ExcsiOo (o' ((w!") : + 11)) - -b:) x D2}
1
(357)
E
1 4
(Ba-WOOo oo' ((w,"), x mi) bi) x D21
(358)
(E(4),
Vd 1- 1
(logd) d) 50 ))
1
-
1
E Ex. [ (1-0' -o' (ooo(()) ((m:")++ 11) -b;) - -1) x D.
x
Σ Bx--(000011)
oo' ((w. ( x + 14) bi) x T DJ11
(359)
1E(
1 Vd (1-0
(log )n ((lgg)) ))- -4
1
- | [[1
1
[11-0'( -0'
x m1) - -b))) x DD]]]
4
Σ 4 /Bx.W/...o)) (o' ((w, x + m1 - b,) x TD,]]
(360)
1E|4),
1 Vd 11 - (1-0(
(log (logd) /t ((1gg))))
1
)) - (
005
(log dj0.0187)
(361)
For any unit vector Dj- which is orthogonal with D,, similarly, we have
E (xxx) loo' ((m."" .x)) b.) x D Exx-1000e (o' ((w:", x mi) b.) x DJ:]]
(362)
(363)
4
1 Z | (o' ((w,"), x + (1) bb,) (x + 11) DD!
Σ
(364)
1E(4),
Bx- (o' (w:"'xx+i)) b.) x Dj
(365)
+
Σ
I Be.NCooo fo' ((w,0),, x + m1) -b.) x xTD
(366)
1E(4))73
<aB0
(log dj0.0187)
(367)
51

---

E.3.2 Mixture of Gaussians XOR: Final Guarantee
Lemma E.17 (Mixture of Gaussians in [99]: Existence of Good Networks). Assume the same conditions as Lemma and let r and when 0 < 7 2 o log
Define
a (1D)))) - 2/ 2V log Flog doB) Flog log d
(x) =
(368)
For Dmizture-zan zor sefting, we have S" E Jd.r where BF (Ba), Baz, Bb)
VFlog VFlogdon 2 (lologd - (
oR va
= 4, BG = 1/2 vd and
an-(log d)
BA1 = (1+o Bar (1 + (B) We also have OPT By
I dt + d0.9 d0.95 / logd
Proof of Lemma E.17. We finish the proof by following the proof of Lemma E.11
Theorem E.18 (Mixture of Gaussians in [99]: Main Result). For Dmizture nor setling with Assump- tion 12, when is large enough, for any E (0, 1) and for any 6 E (0,1) when
m =Ω BBogo 8 ((D08
(
2
1+0B
1
S ed
(369)
ed
v8
T =poly(o 1/e, logd),
(370)
m (1 + oz) á max {aB (log d
Tm 8
+ o (log d) o +
(371)
trained by Algorithm I with hinge loss, with probability at least over the initialization and training samples, with proper eyper-parameters there exists E [T'] such that
Pr[sign(fEco (x)) # v) so ((1+0))
( di
(logn)# ni
)) + )
(372)
Proof of Theorem 1.18. Let b Valog log do
By Lemma E.17, let T and have frr Fd.r.p,,S where BF (Bal, Baz; Bb)
when 7 0
o7 log
V log n 2 log doy)) p= S
ag va"
Be / and
V log don'
(B-(logd)"
B, = (1+an)vd, Ba2 +on))d. We also have OPT d.r. Bp
OPTd,r.Be
d p.. Voll
Adjust ow such that vd log log do e
B B.ZB
Injecting above parameters into 8 over the initialization, with proper hyper-
VrB.1
Theorem 3.12, we have with probability least parameters, there exists / E [T] such that
Pr[sign(fsw (x)) + y] 50 ((
(
(log n)
)) +
(373)
dt
n
E.4
Parity Functions
We recap the problem setup in Section 4.2 for readers' convenience.
E.4.1
Problem Setup
Data Distributions.
Suppose M E RdxD is an unknown dictionary with D columns that can be regarded as patterns. For simplicity, assume d and M orthonormal. Let Φ E Rd be a hidden representation vector. Let A S [D] be subset of size rk corresponding to the class relevant patterns
52

---

and r is an odd number. Then the input is generated by Mő, and some function on O generates the label. WLOG, let {I..... rk), frk I..... d]. Also, we split A such that for all j E [r]. A3 ((j 1)k + ,....,jk. j.k.. Then the input x and the class label are given by:
x Má, y= ( sign
(OR(DA))
(374)
where 9" is the ground-truth labeling function mapping from R* to y {+1}, O^ is the sub-vector of & with indices in A, and ( = I,E Φ1 is the parity function.
We still need to specify the distribution of Φ, which determines the structure of the input distribution:
X = 2rpa) +
Σ Pa(X,, Xz-))
(375)
jE(r]
For all corresponding A in X, have vl A+, independently:
+,,
w.p. Po
Φi =
w.p. Po
0,
w.p. 1 - 2p.
where Po controls the signal noise ratio: if po is large, then there are many nonzero entries in A which are noise interfering with the learning of the ground-truth labeling function on A.
For corresponding @A. any [r], we have
In 1j.t [+1, +1 +1..... + +1]T and @AYA, only have zero elements. Xy. 0A, = [-1, l.... - 1]T and AAYA, only have zero elements. In Xv. have Φ draw from {+1, - 1)rh uniformly. We call this data distribution D Pparity-
Assumption E.19 (Parity Functions. Recap of Assumption 4.5). Let 8 S S be a parameter that will control our final error guarantee. Assume l is an odd number and:
1 M Q(( log d), d rk (ogrrogdd) p. po=0 p r (drrk)
PA
(376)
d'
Remark F.20. The assumptions require k. d, and PA to be sufficiently large so as to provide enough large signals for learning. When O ( means that the signal noise ratio is constant: the expected norm of & and that of A are comparable.
To apply our framework, again we only need compute the parameters in the Gradient Feature set and the corresponding optimal approximation loss. To this end, we first define the gradient features: For all j E [r]. let
EIEA, M
D, EIEA, Mill
(377)
Remark E.21. Our data distribution is symmetric, which means for any Idd
-y (A) and - M(-6),
P(d)=P((e) Φ).
E (x.))) [yx] 0.
Below, we define sufficient condition that randomly initialized weights will fall in nice gradients set after the first gradient step update.
Definition E.22 (Parity Functions: Subset of Nice Gradients Set). Recall w (0) is the weight for the i-th neuron at initialization. For all j E [r], let Sure S [m]] be those neurons that satisfy
(0) )
Csure.11 b,
2
vk
53

---

w (0)
C v/
by for all j' + j.j'' e [r].
PAw
I ((Tkko
PA-W!
S <olvd rkka
where PA. PA- the projection operator the space MA and M
E.4.2 Parity Functions: Feature Learning
We show the important Lemma E.23 first and defer other Lemmas after it.
Lemma E.23 (Parity Functions: Gradient Feature Set. Part statement of Lemma 4.7). Let Csure.1
2' Csure.2 = 1. b = CoVTrk log da where Co is a large enough universal constant. For Dparity setting, we have (Dz, +1), (D), -1) Spr,Bo for all j E [r], where
Vi
-
Vrrlog d d(903rr/8) - dF-2 BG = VipA - O
(378)
dt
Proof of Lemma E.23. Note that for all a E [d], we have M/
x x= Φ1. For allall j E [r]] by Lemma E.26,
for all i E Sures when 2
G(w, (U) b.), b.).D;)) (1-))(1v-) (0) b.)ll2
(379)
ElEr, Mi
G(w!"), bi),
(1- (1- (0) bb))ll
(380)
VTb
2VKPA - 0
vh h
d"
-
1
dr-2
-
kp + Σ
0
dc
1
(381)
FE|d
Vk
(VKPA + 0
1
2VKPA - 0
1
(382)
d"
dr-2
d=
VipA
V
1
0
(383)
d
dr
d
>0.
(384)
Thus, we have G(w: (0) b.) E CD,,, and VipA -
0
vk [G(w (0) b.)lla S VkpA+O
d
S
dr )
bil +1. Thus, by Lemma E.24, we have
Pr w.b
G(w,b) E CDS and ||G(w, 2 Bg and
+1
(385)
b
2 Pr [ E Spp,
(386)
2p.
(387)
Thus, (1,.+1) E Bp. Since Ex [yx] 0. by Lemma F.2 and considering e [2m] \ [m].
we have (D): 1) Sp We finish the proof.
Below are Lemmas used in the proof of Lemma E.23. In Lemma E.24, we calculate used in Be.
Lemma E.24 (Parity Functions: Geometry at Initialization. Lemma B.2 in [7]), Assume the same conditions as in Lemma E.23, recall for all i E [m], w!") o0 over the random initialization, we have for all i E [m], E [r].
Sure Pr i E Sure
2
(388)
Tr log d. ((003T1/8))
54

---

Proof of Lemma E.24. For every i E [m], j.J' E [r], # j', by Lemma F.6.
P1 = =P ( (w! D2)
Osure.1 bi
Vk
=
Tr log d (x0C3rr/8)
(389)
P2 = Pr ( ( Dyy
Csure.
bi
(390)
Vk
Vrrlogd d d(C3rr/8)
On the other hand, if is (k) random variable, by Lemma F.5, we have
Pr(X 2 k + 2vka + 2r) Se
(391)
Therefore, by assumption rk: 2 ( d), r 12(T log d), we have
(0) 2 rk +2 2 2yrklog 2 logd
Pr
Paw Paws
(392)
0
(393)
d ( (907)r/8)
Pr
PA W
2 (d - r1k) +2 (9C + 2)(d - rk)log d + T9 2) log
0
(394)
d2 ((903rr/8)
Thus, by union bound, and .,,,.., D, being orthogonal with each other, we have
Pr [i Sure] 2p1 (1 p21 13 d(903rr/8)
(395)
1
(1-
=
(396)
VTrlog d d(903Tr/8)
vTr Vrrlogd d d(C3rr/B)
d2 ((9C377/8))
(397)
1
(398)
Vrr d. a(902rr/8)
In Lemma E.25, we compute the activation pattern for the neurons in Sure Lemma E.25 (Parity Functions: Activation Pattern). Assume the same conditions in Lemma E.23,
for all j E [r] i Sure, we have (I) When X, we have
12 (r(ko)))
Pr x
Σ IEA-
Mirr)
21
S exp
(399)
(2) When x Xy, we have
Pr X
w (0)
2t
< exx
(400)
Orkoo
IEA
(3) When x Xv, the activation probability satisfies,
Pr Xo
w; (0) MMOi)
b, 20 -b,20
50 d
(401)
[E
55

---

(4) When x ~ Rz.t, the activation probability satisfies,
(
Pr
(0) Miri) bi 20 21 21-0 53
(402)
LE/d
(5) For all j' + J.J' E [r]. s E {t, }. when x ~ Ry or x ~ Zz.-, the activation probability satisfies,
Pr
(0) Mirr) - b; 2 0
S
(403)
d
EE
Proof of Lemma E.25. For the first statement, when x X, note that
(w, .M,)
(0)
Φi is mean-zero
sub-Gaussian random variable with sub- nub-Gaussion norm e ((ww:)) Mi)
VPb)
Pr
w (0)
MiOI)
2t
Pr x-x
w (0)
Mi
Φi
zt
(404)
x
IEAL
8
S exp
(405)
(w!",,M)) Mi) 2
EIE O
Do
$ exp
(406)
+ ((d - rk)oo
exxp
(407)
Orkko
where the inequality follows general Hoeffding's inequality. For the second statement, when x ~ Ny, by Hoeffding's inequality,
Pr
w
Midr) 26
= Pr x
W
M,
2
(408)
X-
IEA
IEA
12
exp
(409)
ElEA
,M)))
t2
Sexp
(410)
Orko
In the proof the third to the statement, we need CLVTrk log daw, where Cp a large enough universal constant.
For the third statement, when x Xv. by union bound and previous statements,
Pr
( M1 - b, 2
(411)
LE [d]
<
Pr
!0!
M1Oi
b;
+ Pr xo
(
M
>
bi
(412)
2
2
LEA
LEA
<0
(413)
d
56

---

For the forth statement, when x ~ Xz.+, by Osure,1 2 3 and previous statements,
Pr
(0)
MiOi) - b; 2 0
(414)
leld]
Pr
Σ (0) Mioi) + Σ IEA;
( .Mii) +
Σ IEA
w;
M 2 bi
(415)
+
IEA)A,
Σ (0) Miri) 2 (1 - (1-Csure,,l I IEAL
2 Pr x
(416)
2 Pr x
Σ
( .M 2
b; 2
(417)
IEAL
21-0
(418)
For the last statement, we prove similarly by < Csure.2 S 1
Then, Lemma E.26 gives gradients of neurons in Surrr It shows that these gradients are highly aligned with Dy. Lemma E.26 (Parity Functions: Feature Emergence). Assume the same conditions as in Lemma E.23, for all j E [r], i Sp,,Sures we have the following holds:
(1) For all I E Aj. we have
PA
1 dt
S E(x,0) (yo' ((m!)) xx) - bi Φi
+
1
(419)
d
(2) For all l E Ay, any j' j.J E [r], we have
E(x.0) yya'
b. क1 2 d=
(420)
(3) For all l E A+, we have
E(x,0) yyo'
-b) bb
00
(
(421)
Proof of Lemma E.26. For all / E [d]. we have
(oo' ((w!") ,,)x) b; Φ1
E(xx))
(422)
PPA
Σ (Exwxi. [o'(( ((w!!) (0) .x) b.:) 01 -Exxx, [o' ((w,0) (x/",x))
x)
b. de
(423)
LE(r)
+ (1 2rpa)Exwxt) Jyo' (( ((w:")) (0) xx -b,) du).
(424)
For the first statement, for all E Aj. by Lemma E.25 (3) and (4), we have
E ((xxv) Iuo' ((w;0) (w/"))
.x/ b. ob
(425)
(xxxX,+ 0 ((mx))
(w:") 10) .x) - bi)] + Ex 0 (wr)0) .x) - -b.)))
(426)
+ (1 - 22 "oo (( (w!()
(0)
x
- -b:) 12
(427)
1
2P4 1 1-0 0
0
(428)
d
d
2PA - 0
1
(429)
d
57

---

and we also have
E (x.9) Lool ((w: (0) xx/ b,) 01
(430)
PP (x [o' (w::), x) -b:)] - Bx d ((w! xx - b,) b))])
(431)
+ (1 - -2rp Exwry yoo' (<w!)) .x) .x) - b; 01
(432)
SP + O
(433)
d
Similarly, for the second statement, for all I E Ay, any j' + j.5 E [r], by Lemma E.25 (3) and (5),
we have
B(x.)) yyo' ((w!"),x ix) bi) 0x
(434)
PA (Exmndzy (o' ((w, xx) b.) + Ex [o' ((w!" - b.) bi))) +
1
dF
(435)
dt
For the third statement, for all E A by Lemma E.25 (3), (4), (5), we have
B(xx))
(y0' ((w!!))
(0) x - by Φ1
(436)
Σ Ex (o' (((( ((w/)) (0) x) - b.) 01] - -x ((w/ .x) - -b,) 01
1
SPA
+
dr
te(ri
Ex (o'( (((1!!0
,,) xx) - b.) 01 - Ex-N, - [o' ((w, - b.) 2b + +
1
SpA
d
Ex [11- (( w (0) xx - b.)) 04)
-Bxuity- [11-0' (( (wx/)) .x) - b.)) 61
SPA
+ PA Ex~xx [01] - Bx [0il] +0 dt
(437)
[ (1-0' (1---((w//) ((v/"),, ,x) -b:)) 01]
-xxuts [[1 (1-0' (1-0'((w)) ((w/0) .x) -b:) -b:) Φ1
-PA
Ex
1
0
(438)
dt
00
(439)
dt
where the second inequality follows 2rpA < and the third inequality follows the triangle incquality.
E.4.3 Parity Functions: Final Guarantee
Lemma E,27 (Parity Functions: Existence of Good Networks. Part statement of Lemma 4.7), Assume the same conditions as in Lemma E.23. Define
fx)x) V
(440)
[ (D),
2i k 1 - 11 20 (D)).x) -22
2i -k) 20-k) +0 (D)),x) )
2i - k + 1} )
Vi
Vi
vi
For Dparity setting, we have 8rr E Fu.3r(4411,B,.Ssc... where BF (Bal, Baz, Bb)
=
(2VT. 2//h,2/rk(k + 1), 1/2 p= p=
vrr an , JT-Z, BG = ViTpA - vT dt)
and Bal B.2 d. We also have
Br 0.
58

---

Proof of Lemma E.27. We can check B,1 = Vd, B22 - d by direct calculation. By Lemma E.23,
we have f E F.
We note that
(D))
2i - ---11 - 20 (D).
2i - k:
(Dj. (
2i k ++11
+
(441)
V
V
bump function for (Dj, x) at 2i-k VR
We can check that yf (x) 2 1. Thus, we have
OPT Br Spo.
Dppprrity (f*)
(442)
= (xxx)) (F*)
(443)
0.
(444)
Theorem 4.8 (Parity Functions: Main Result). Algorithm I uses hinge loss with
Assume Assumption 4.5. For any c, E (0, 1), when
= poly ( dolr) ,,, PA
1 se², poly (m), n = poly (m)
and proper hyper-parameters, then with probability at least 8, there exists E [T'] such that
3r 3r/h
Pr(sign(fzo + S 31/2 te.
Proof of Theorem 4.8. Let 5 = CoVTrk log do where C. is large enough universal con- stant. By Lemma E.27, we have f E Fagzrkk+ Bppr.. where BF = (Ba1, Buz, B,) =
(2Vh. 2. rk(k 1),
k+1 vk
O
#. Be VipA 0 A
VTr dd (20grr/s)
and Bu1 = Vd, Bu2 = d. We also have OPT Bp 0.
Injecting above parameters into Theorem 3.12. we have with probability at least 1 - 8 over the initialization, with proper hyper- parameters, there exists i. (T] such that
Adjust Tw such that b C.VTrk log
θ
B BB VrB.:
2/ +0 Pr(sign([=wxx)) (x)) # v] 3)/2
3rv (rBaiBa (1BBBB Ban (logn)+) + d(r- te.
/Bant
F.5 Uniform Parity Functions
We consider the sparse parity problem over the uniform data distribution studied in [15]. We use the properties of the problem to prove the key lemma (i.e., the existence of good networks) in our framework and then derive the final guarantee from our theorem of the simple setting (Theorem 3.4). We provide Theorem E.31 as (1) use it as warm-up and (2) follow the original analysis [15] to give comparison. We will provide Theorem E.40 alternative version that trains both layers.
Consider the same data distribution in Appendix E.4.1 and Definition E.22 with the following assumptions.
Assumption E.28 (Uniform Parity Functions). We We follow the data distribution in Appendix E.4.1 J. Let r = 1,PA = O,Po 2 M Taxa and 2 2k22, and is an even number.
We denote this data distribution as D uné u form setting.
To apply framework, again only need to compute the parameters in the Gradient Feature set and the corresponding optimal approximation loss. To this end, we first define the gradient features:
let
Eien M
D
(445)
EIEA M1l12
We follow the initialization and training dynamic in [15].
59

---

Initialization and Loss.
We use hinge loss and we have unbiased initialization, for all i E [m],
(0) Unif((#1]), 0) Unif( b;
1//kk)
(446)
Training Process. We use the following one-step training algorithm for this specific data distribu- tion.
Algorithm 4 Network Training via Gradient Descent [15]. Special case of Algorithm 2
Initialize (a) W b) as in Equation (8) and Equation (446); Sample Z Darritt marity-unifore
W(1) W (2 (())
(1) aa (V
(0))
for - to T do
a (1-1) 1/ V (J
end for
Use the notation in Section 5.3 of [92], for every S E [n], s.t. |S| = k, define
2 (d-1) d
EA == Maj(S) = (-1)
(447)
d-1
Lemma E.29 (Uniform Parity Functions: Existence of Good Networks. Rephrase of Lemma 5 in [15]). For every e, E (0, 1/2), denoting 16 2ddl/2ll2 let r7(1) k|Ex-11' A(1) n,
m 2 k- 2 (k)())/), 2 3 log(4dm/ and d 2 Ω (K log(kd/<c)), w.p. at least 1 - 25 over the initialization and the training samples, there exists IR" with |āl] I 8k and ||||| S 8kvk: such that ((a,W)),p) b)) satisfies
CDpartoy (f(a Sc.
(448)
Additionally, holds that |o (W(1))))) x- b)llec 1d
Remark E.30. In [15], they update the bias term in the first gradient step. However, if we check the proof carefully, can see that the fixed bias still goes through all their analysis.
E.5.1 Uniform Parity Functions: Final Guarantee
Considering training by Algorithm 4, we have the following results.
Theorem E.31 (Uniform Parity Functions: Main Result), Fix E (0,1/2) and let m 2 2 (k. 2 log(k/()), n 2 Ω (= log(kd/c) (og(m//c)) + k2mdd d 2 Ω (R* log(kd/c)). 1
62
k|Ea-1l' A(1) n00), and (0))) = O (#m)), for all t (2,3.... T). If 2 Ω (2mdd then training by Algorithm with hinge loss, w.h.p. over the initialization and the training samples, there exists E [T] such that
Let (1)
(()) + y) S LD Dppatto--un/rrmaaa W S.
(449)
Proof of Theorem E.31. By Lemma E.29, w.h.p., we have for properly chosen hyper-parameters,
OPT S partty
(f (fa,w()),)) < 3
(450)
60

---

We compute the L-smooth constant of Ég (fo., a
VaLz (am,)))))))) Valz ( ((a,))))))))))
(451)
(e'
- (uflaw..o.,))33 ((sx))))[((W)))) (W x x-b)
(452)
2
[(a,, - ( x-
(453)
71
xE2
|ai
22lla
( x - b)
(454)
xEZ
By the Lemma E.29, we have Ilo(W(1))
x- x-b)lla S d. Thus, we have,
ooow()xx x- b)ll" /
L L=0
(455)
XE2
< m).
(456)
This means that we can let 17 0 (M/)) and we will get our convergence result. Note that we have
(1)) O and ||āl|2 0 (kv/). So, if we choose T 2 Ω
tat there exists / E [T] such that
Ez (J(a w..,)) - Lz ((a.wo).))) S 0 (
(1a -al3
S c/3.
T
| B2+ b
We also have
Then our theorem gets proved by Theorem 3.4.
E.6 Uniform Parity Functions: Alternative Analysis
It is also possible unify [15] into general Gradient Feature Learning Framework by mildly modifying the framework in Theorem 3.12. In order to do that, we first need to use a different metric in the definition of gradient features.
E.6.1 Modified General Feature Learning Framework for Uniform Parity Functions
Definition E.32 (Gradient Feature with Infinity Norm). For a unit vector
Rdd with
[|D|2 1, and a Yoo E (0, direction neighborhood (cone) Drrw is defined as: CD
[w] - DD| <Tra Let E Rd, E IR be random variables drawn from some distribu-
tion W. B. A Gradient Feature set with parameters p, Yoor Be, Bei is defined as:
Pr Ba. Bai (W,B) = { (D,s) Pr w.b
800
G(w, E CD Bei 2 [[G((,, (((w,))ll2 BG.
2DJ
b
When clear from context, write it as S
BB.
Definition E.33 (Optimal Approximation via Gradient Features with Infinity Norm), The Optimal Approximation network and loss using gradient feature induced networks By. Sp Be defined as:
are
argmin feF Bp... P
Lo(f),
(457)
B
OP Bc. feTa.r. Meros By 861
min
LD(f).
(458)
We consider the data distribution in Appendix E.4.1 with Assumption E.28, i.e., D Dparity-unde tontyy-uunformm
in Appendix E.5. Note that with this dataset, we have ||x|| < Broc = 1. We use the following unbiased initialization:
for i (1, m.. a (0)
N(O, 2.)
(0)
{+1]d b = b 1,
for i E (m + 1,..,, 2m)
a;
(0)
(0) (0)
i-m: b, -b;-m:
(0)
for i E {2m + I,... 4m}
a
2m
w.
w (0)
2m b; = bb-2m
(459)
61

---

Let V, denote the gradient of the i-th neuron Vw. (ffzo)) gradients approximating feature (D,s) s) as:
Denote the subset of neurons with nice
(DD. = {i E [2m)
b;
V V
(0) Bei 2
(0) ||ell2 2 a Be }.
Toc:
a
b,
Lemma E.34 (Existence of Good Networks. Modified Version of Lemma 3.14 Under Uniform Parity Setting). Let ( शोड. For any Be E (0, Bb), let oa = θ ( and 8 = vmp + 8.. Then, with probability at least 5 over the initialization, there exists à, á such that (()))))))))) Σim å,o (( (ww"), x)
b.) satisfies
B₂ BrI
Bai B,
Co(f(a,wo))))) S rBol
+ v 2log( + Be) + O..d,,, S.
MmpBa.
p. Bg.
||āl|== =0 0 and ||āloo - 1 (f(pp)))
|lālloc ( B.
=0(B)L)))
Br
b(mp)
((mp))
Proof of Lemma E.34. Recall f*(x)
Σj= afo((w3.x)
bj), where f
b
Fdr,,r rpice Re BGi is defined in Definition E.33 and let sj = bbj By Lemma D.3, with
probability at least 1 - 5,, 8, 2remomp, for all j E [r]. we have |G (w)og))Noo
2 mp Then
4
for all E G Gwy Nice C [2m], we have G(w,
(0) b,) of only depend on w; and b,
which is independent of a Given Definition E.32, we have
bj T E ((())) (1) Bml 6
Bb
BB
-((0)) G( b.)lla
(0))
BR
(460)
b
We split [r] into I' = E [r]
bil < B.J.- = Li E
[r] bj < -B. and I'+ = {- E [r]
bj 2 B.). Let Ea
Bai Bs
Then know that for all j E T UT for all E (w, (87) Nice"
B.
we have
billl bb
(0) ((()) (0)
Pr a ~N(0
a
S Ea
(461)
bb
Pr
Ea S -a; -a""0 (0 (0)
bi)lla
S 11 + Ea
(462)
10) 2)
b
||G(w (0)
b2)
Pr g>>(00.))
1 6a S 90
11+6a
(463)
BeB.
BGB.
BG B.
Pr g0,11
(1 - (1-
S (1+
(0)
(I((w bo)laabb
6a Be
-0
(464)
([G(w(
(0)
b:)lla/bā1
BaBaB.
20
(465)
BG
=Ω
(466)
/mp
Thus, with probability Ω
over , we have
b 1 S Eai
a
b
a
0
-e((0)n())
(467)
Bc
62

---

Similarly, for j E r, for all E (wrj.o)).Nice with probability ? (pT))) over a (0)
we have
a (00 (0)
b.llk
B
S cas a;
(0)
b (0) Ba B.
(468)
For all j [r], let Aj (w; is),NN)e be the of i's such that condition Equation (467) Equation (468) are satisfied. By Chernoff bound and union bound, with probability at least 1 82, 82 re r for all [r] we have [Ay] 2 mmpp) We have for vjj r+ -I-, Vi e Ay.
bb
w .x) - -(wj. x)
(1)
(469)
b
[bÿ] w(
w((
(1)
0)) (1 bi)ll
x
+
w x
||w] 1 112
[w 1
S(a xx|/2 + 22oogd)dd.o.
(470)
With probability 1 + by Hoeffding's inequality. Similarly, for Vj E TT Vi E Aj.
Be
b w (1)
x -ww,, x) S <al1x112 (2log(d)d7
(471)
-b; [b,] if A T--UT set berā; =a 1.
if E Ay. T, T, set à; = =a i
B.
otherwise set à;
we have ||āllo = (r(mp))) 1|āl/22 =
Bu
||āl| =
B Ba 6(mp))
6(mp))
63

---

Finally, we have
(D(((a.wo).))) -Coll(a, - CD(f") LD(l)
(472)
(473)
(x))))) (a) - -x(x]] + (D(
(474)
2m
Σ a (( (w!"'xx xx) -b) + Σ à (x) +b] Σ aj w
(x.))
wj - bU)
i=1
m=+1
+ (f*)
(475)
(1x)).))
Σ aj
1
b a (( [ xx -b) <o((w) x) 133)--b) -b3)
(476)
AAll
JET+ iEA,
+ ((x,)))
Σ
aj |ay1
1
|b3|
((w!") xx) + +b) o(( x) -b²)
(477)
b
jET iEN,
+ (x,x)
aj jEriEE,
1
Bg
( ((w!") -/)x)))) xxx=) -8) o ((w3,x)) - -bi) + (f()
(1)
(478)
A,
b
<E
aj
1
by
(1) x) - (w]xx))
(479)
[Ag
JET i
+ ((xw)
Σ
a
1
[bj]
(!!) (1) x) (w/. x)
(480)
A
JET- iEA,
+ [(x.v)
ΣΣ iET
aj
1
Be
wi xx + Be - (wwj xx)
+ fD(l))
(481)
AA
b
<r||a" blo (x.B(x.)lll)))
log(d) + ||||a" Be + CD(f")
(482)
BrBa (ca BuI + (2log(d) + [ Bm B. + OPT OPT Br. "o: Res"
(483)
We finish the proof by union bound and 5 2 61 +
dz-
Lemma E.35 (Empirical Gradient Concentration Bound for Single Coordinate). For i [m], when (log(d)) with probability at least 1 O (exp (-n$) over training samples, have
oīz(fe) @LD(fe) Owis
<0 ( |a. |a.|Booc
Vj E [d].
(484)
Owis
n
Proof of Lemma E.35. First, we define,
2 D1 [o' [d'((w,x")) b:) x0
2 - -E(x,) 1e" (uf=(x)))) Jo' ((w,, x) b:)] x,)
(485)
(486)
As (e(=)) < 1, 5 1, (=)( < 1, we have Z is zero-mean random variable with
( 2.
<
2
as well as E
(
< 4B22 Then by Bernstein Inequality, for 0 < < 2B we have
[2
aīz(fz)
(LD(fz)
Pr
Ow,s
Owas
2 |a//z
= Pr
2
(
2
(487)
IEin
2
S exp
-n
(488)
8B 8B
64

---

Thus, for some [m]] when n 2 (log(d))o with probability at least 0 0(exp (-nà) union bound over j [d], have, for Vj [d].
)
from
Oāz(fz) @CD(fE) Ow.. Ows,j
a Buoo
(489)
n
Lemma E.36 (Existence of Good Networks
under Empirical Setting).
Risk.
Modi- n
fied version of Lemma D.13 Under Uniform
Parity
Suppose
Q ((( (BB
3
B__ VBn + log Bole(oy)
Bany
B +
Bux Boile(00]
(log(d))o) Let
2
n For any
and 8 2re V + 1 Then, with proba- bility at least 1 8 over the initialization and training samples, there exists āg ''s such that
Be E (0, Bb), let da O
,x) Eim â;o (( (w?)) xx b.) satisfies
CD (D(F(a
(490)
(2 Ben Bb
B. + 2 /ologd (To +0 (Ble(())nn3) )) + Be + +PT O.. por Bai
Brse
SrBal
and ||āllo (allo=0(((mm) (r(mp)))) ||āll2 = =0 0 BraBy
||āllo =0
Ba 8(mp) 2
b(mp)
Proof of Lemma E.36. Denote p = 0 exp θ (-nt) and B = Bacc n
Note that by symmetric initialization, we have (y (x)) (0) for any E X, so that, by Lemma E.35, have
((( (0) b:); G(w;
(0) bi)s
with probability least p. Thus, by union bound, Poo
we can see that 5:0
C 500
Consequently, we have
Bo
Bal
Ba 1o OPPT.
Bai to
OPT Br.
S p. Bc1
Exactly follow the
p-P:
Ba "c T:oo Mot ाक
proof in Lemma D.4 by replacing S0c Be,
S00
Then,
p-p Bg
Bol
BB [(
Ban + 16
8 we finish the proof by p S 2' [e(o]] $ (1 1/v2) (((O)
I 1)
Theorem E.37 (Online Convex Optimization under Empirical Risk. Modified version of Theo- rem D.17 Under Uniform Parity Setting ). Consider training by Algorithm 1, and any 8 E (0, 1). Assume d log m, ( Set
nn 0, 10 7, 1 0 all 1 E (2, ...,....77)
(1) = θ min (mb))
101 1
0a =
(e((0)(()) BBBs B
b(mp)) &
[BOl V db
(11)
Tn
S
o(1),
m
2 a + (log (3))² and n
((P
((
Bu + log
Tm pd
(1 +
1. B
3
With probability at least over the
Bea
Bea
(0)
65

---

initialization and training samples, there exists t E [T] such that
CD (fgw)
(491)
2/2vBBB Bm Bei
Bb
B2s0
<PPT Br.n + rBal
+ /2log(d)d Yoo +0
(mp)
BG
(cle(0)
VlogmBai (mp)
+1 (VrBa2B,TTDL, + +m) o
+1 +0
(492)
VB
1/T(2(mpp)
1 0
BBaiBB Bs
(//log(mp))
+
(493)
nz
VB,Bo
B21
Bollog m(mp)
+ +TTrB
B2 +
+
(494)
VB.B
by (mp)
+
O
mnj
+ Tup
VBz2
(495)
BB
73
Furthermore, for any e E (0, set
-
B Bo B
m 8
rBal v Bzi Bal
Bb
+
1
+
log
())
r
Vr
pet
Ba
vo
1
T=
B1r (mp)
.rBaaB.Bss + +mb)
Vlgg MogmB.i(mp) BB.Bo
+
(mp))
"m,,B", Bo(mp) log m
Bx
Tm pd
B200
n =Ω
+
+
+og
(1 +
+
erBai
VBaa
BG
Ba
[C(O]]
have there exists E [T] with
Pr(sign(fzo) (x) # yl S LD (Szw)
(496)
(
(cle( )) +c.
B200
<OPT d.r, Sr Rci + r V llo(dd)
7o0 +0
(497)
Proof of Theorem E.37. Proof of the theorem and parameter choices remain the same as Theo-
rem D.17 except for setting B. BB.,
Bay
V Ba
Ba
and apply Lemma E.36.
(mp)
E.6.2 Feature Learning of Uniform Parity Functions denote
Gi.j B(x.v) yyo' (w x)
((
b;
Xj
(498)
n-1
& = (-1)
2
2
12 (n-1))
(499)
n-1
1
2
l:
1-
Lemma E.38 (Uniform Parity Functions: Gradient Feature Learning. Corollary of Lemma 3 in [15]).
Assume that n 2 2(k + 1)2. Then, the following holds:
if A, then
Gi.j Ek-
(0) 1EAV))
II
(wil)
(500)
66

---

i & A, then
Di. E
II (wi//)
(501)
(EAu(j)
Lemma E.39 (Uniform Parity Functions: Existence of Good Networks (Alternative)). Assume the same condition as in Lemma E.38. Define
EIEA M
D
(502)
1EEE Mrl2
and
(x)
Σ 1)'v.s
(503)
L=0
2i k - -i=k-1) 20 (D,x) 2i - +o (D)x)
2i k +
(D,x)
V
Vk
Vk
For 'D pority-unkform setting, we have f" E 14. 2.700 Ba.
where BB =
BG.BGi
(Ba1, Baz, Bb) = 2 1)),
k+1 v
). p= θ (2-TT) Yoc = =0
d-b
vk
Ba = O ((B)))
=
(d k) and Bu1 Vd, Brg d. We also have OPT a.TT,33(+1 +1) Bp, Bq. 0.
Proof of Lemma E.39. Fix index i, with probability P1 @(2 we will have w (0) sign(a ()) sign(Ex-1). for Vj. For w!" that satisfy these conditions, we will
have:
sign( a
(0)
gi. [6-1]) Vj E
(504)
sign(a. a Ii. |Ex+ils vj $ A.
(505)
Then by Lemma 4 in [15], we have
sign(a, (0) (G(((() B) D
I max
(506)
[G(w (0) B)]l
k c d-
1
V
(d k) k)
k/2 + d
Σ Vk
(507)
- k
and
[sign((a (0) (ww (0) b)ll V 2 (d k) (cd0()))
(508)
vk From here, we can see that we set Yoo d-k BG BGi VklEk-1 (d E) +1
p = Pt, we will have (D,+1), -1) E S.c P.7no Be, Ben Ba by our symmetric initialization. As result, we have f" E 1). 700 Ba. Finally, it is easy to verify that f*(x) = XOR(X thus O.
D-7ee OPT BS Bc.
Theorem E.40 (Uniform Parity Functions: Main Result (Alternative)). For D wni form setting, for any E (0, 1) satisfying ó < O( and for any E (0, 1) when
m = poly log 200h) =
.T= (10(1)) n -n= (
(509)
trained by Algorithm with hinge loss, with probability at least over the initialization, with proper hyper-parameters, there exists t [T] such that
12 /dlo((d)) d)
Prisign(fgon (x)) + vl
+ +c.
(510)
d
Proof of Theorem E.40. Plug the values of parameters into Theorem E.37 and directly get the result.
D
67

---

E.7
Multiple Index Model with Low Degree Polynomial
E.7.1 Problem Setup
The multiple-index data problem has been used for studying network learning [18, 32). We consider proving guarantees for the setting in [32], following our framework. We use the properties of the problem prove the key lemma (i.e., the existence of good networks) in our framework and then derive the final guarantee from our theorem of the simple setting (Theorem 3.4).
Data Distributions. We draw input from the distribution Dx NO Taxa), and we assume the target function is 9" 9*(x 12 + where g" is degree polynomial normalized so that Ex-Dx 19" (x)² = 1.
Assumption E.41. There exists linearly independent vectors u1...., Ur such that g (x) = g((x, u1),.... ..., (x, U,)). H == Ex~DxIV VV'g g" (xx) has rank r, where H is Hessian matrix.
Definition E.42. Denote the normalized condition number of H by
||A||
(511)
Initialization and Loss. For Vi E [m]. we use the following initialization:
(0) (-1,1), w (0) .o (0.z and b; = 0.
1
(512)
For this regression problem, we use mean square loss:
LD (fz) = Ex-Dx [(f=(x) - (²()))
(513)
Training Process. We the following one-step training algorithm for this specific data distribu- tion.
Algorithm 5 Network Training via Gradient Descent [32]. Special case of Algorithm 2
Initialize (0) ( b) as in Equation (8) and Equation (512); Sample 2 ~ Dxx
Ex 9 (x), 3 = ExEz 9" (x(x
= (x) p 8-x
( = (
(V (fz + (()(())
Re-initialize b, ~ ~
for == 2 to T do
(*) (1 (1) 22
end for
Lemma E.43 (Multiple Index Model with Low Degree Polynomial: Existence of Good Networks. Rephrase of Lemma 25 in [32]). Assume 2 drk? log( (nmd)) 2 Caxr3/2 and 2 rT 2T (Ci log(nedd))644 for sufficiently large constants Ca. Ct, and let (1) V [C, log(ndd and poly(m.dj) there exists à E TD such that (a.w)))))) satisfies
d
(1) J/... Then with probability
Lox (Fra.wo),))) so
+ r- xT (Cr log(nmd)) 6T+1)
(514)
n
m
and
|| 000 r (C (ogggddd)))) 6r
(515)
E.7.2 Multiple Index Model: Final Guarantee
Considering training by Algorithm 5, we have the following results.
68

---

Theorem E.44 (Multiple Index Model with Low Degree Polynomial: Main Result), Assume n 2 a (dTn' (C' log(nmd))** 7+1 + -m), d 2 CaRr and m 2 Ω (Er" 20 (C, Log(mmd)) for d
sufficiently large constants Ca. Ct. Let 17 V (Ci log( and A(I) = n and rj 1)(1)
@(m for all (2,3....,T). For any E (0, if 20 Ωt
m²
then with properly set such thut
purameters and Algorithm 5, with high probability that there exists & [T]
((lt)) w)). S.
(516)
Proof of Theorem E.44. By Lemma E.43, we have for properly chosen hyper-parametersy
_
r #x²T(CA log(nmd)) 6r+1)
OPT S LDx (fl, 0o
(517)
(518)
3
We compute the I-smooth constant of La Vla..wo... t
V VaCz (Jon) vaīz ((o2)
(519)
2 (Jia (w),) 9" -w,,, o x
(520)
n
2
2 a oW x -b - a a (w x x-
x b)) T( x- -bb b)
(521)
n
XEZ
2 |aa a2lla [w x- b)112 27
(522)
x
By the proof of Lemma 25 in [32], have for Vi E [4m], with probability at least 1 (m) ((w,,) < 1, with some large polynomial poly(m, d). As a result, we have
W()T
S m
1
S O(m).
(523)
poly(m, d)
XE2
Thus, have,
x x-b) b) 2
L L=
(W(
(524)
17
XEZ
T11) x -b1l2-
00
(525)
XEZ
12 ((1)
50
|
(526)
XEZ
(m).
(527)
This means that we can let = θ (m-1) and we will get our convergence result. We can bound
||a (2) l|2 and ||ā||2 by (|2 0 (vm) and ||||| O (r'* (,'a² (C, loR(nmd)) O(e). So. if
choose T > Ω there exists (T] such that Ez (S(al) Lz (fa,.ar,))) s 1 (1) 1ap--al 2ll?
e/3.
We also have (l(l)))
B2
5/4 Then theorem gets proved by Theorem 3.4.
69

---

Discussion. We would like to unify [32], wheih are very closely related to our framework: their analysis for multiple index data follows the same principle and analysis approach as our general frame- work, although it does not completely fit into our Theorem 3.12 1.12 due to some technical differences. We can cover with our Theorem 3.4.
Our work and [32] share the same principle and analysis approach. [32] shows that the first layer learns good features by one gradient step update, which can approximate the true labels by low- degree polynomial function. Then, a classifier (the second layer) is trained on top of the learned first layer which leads to the final guarantees. This is consistent with our framework: we first show that the first layer learns good features by one gradient step update, which can approximate the true labels, and then show good classifier can be learned on the first layer.
Our work and [32] have technical differences. First, in the second stage, [32] fix the first layer and only update the top layer which is convex optimization. Our framework allows updates in the first layer and uses online convex learning techniques for the analysis. Second, they consider the square loss (this is used calculate Hermite coefficients explicitly for gradients, which are useful in the low-degree polynomial function approximation). While in our online convex learning analysis, we need boundedness of the derivative of the loss show that the first layer weights' changes are bounded in the second stage. Given the above two technicalities, we analyze their training algorithm (Algorithm 2) which fixes the first layer weights and fits into our Theorem 3.4.
F Auxiliary Lemmas
In this section, we present some Lemmas used frequently,
Lemma F.1 (Lemmas on Gradients).
(f)
(f(x
(w
(528)
Ow1
Ow,
Owam
0 (f)))
a,l (Uf(((()))) (o' ((w;, x) 13--by -bb)] x,
(529)
Ow:
[LD([E))
@LD(JE) @LD(F=)
(530)
Ow1
Ow;
Owam
@LD(fa) (x,(),)) ['(sF((x) jo' (w:;x) - -b:) x].
(531)
Ow:
OLD(fz) B(x,)) ((((ff(x)) /o (w,,x) - -b,)ll-
(532)
da,
Proof. These can be verified by direct calculation.
Lemma F2 (Property of Symmetric Initialization). For any x E IRd, we have fe (x) = 0. For
all i E [2m], we have w(1) - -w(1] W42m When input data is symmetric, Le, E(x,v)xx)) 0, for all
i [m], we have w!! w 14m
(1)
Proof of Lemma F2. By symmetric initialization, we have fgw (x) -. For all i [2m],, we have
w! (1)
()) E (xxx)
y3o' (w: (1) .x) b, x
(533)
())) (0)
B.a.... (xx
yoo (w.t2 x)
bi+2m x
(534)
2m
= -
(535)
70

---

When E (xx,) yx] O, for all i E [m], we have
-))) E (x..)
y0
w
(0)
.x)
b;
(536)
=0)aa (0) E (x,2)
yo
w
i+m
(0)
Di+m x
(537)
E (x.
yoo
1+m
(0)
Ditm
x yx
(538)
=0)
E (x:9)
[-yo'
{
w (0)
i+m:
D&+m
(539)
=w
(I) i+m
(540)
Lemma F3 (Property of Direction Neighborhood). If E CD.,. we have pw for any A 0. We also have $ CD.,. Also, if (D,s) 8) E Spy.Ba. we have s) E Spr.Ba.
Proof. These can be verified by direct calculation.
Lemma F44 (Maximum Gaussian Tail Bound). Mn is the maximum of n ii.d. standard normal Gaussian. Then
Pr Mn 22ogn
se==
(541)
2ogn
Proof.
These can be verified by direct calculation.
Lemma E.5 (Qhi-squared Tail Bound). X a x2 (k) random variable. Then, V2 E we have
Pr(X 2 + 2 + 22) S =*
(542)
Proof. These can be verified by direct calculation.
Lemma F66 (Gaussian Tail Bound). If is standard Gaussian and 2 >> 0, we have
1
2/2
Pr g~A(0011)
19> 2]
1/
(543)
2 +1
v 2m
Proof. These can be verified by direct calculation.
Lemma E7 (Gaussian Tail Expectation Bound). Ifg is standard Gaussian and z IR, have
E Eg- > z]g] go,N(
[g 2]
(544)
Proof of Lemma F.7. For any p (0, 1), we have 2en (2p
dx
2p 0.9
(545)
V2TT
where /2erf-1(2p 1) is the quantile function of the standard Gaussian. We finish the proof by replacing p to be Prg-Ni(..)l9 > z].
Lemma F88. If function g satisfy hnnn + 2) = 2h(n 1) (1 p²)n(n) + f N4 where p.3 > 0, then h(n)
+ c1 p)" + c2(1 + p)" where C1, C2 only depends on h(1) and h(2).
Proof. These can be verified by direct calculation.
Lemma E9 (Rademacher Complexity Bounds. Rephrase of Lemma 48 in [32]). For fixed W, b, let F = (fa.w.)) ||a|| < Baz)- Then,
B2
W 2 B3
[b|
R(F)= S
(546)
7
71


Reviews:
[{"review_id": "LeQGy9Ot7v", "reviewer": "Reviewer_zxbj", "summary": "The paper defines the concept of \"gradient features\" which capture the features that the network can learn after one step of gradient descent. The paper then instantiates this framework to prove optimization and generalization guarantees for various statistical learning problems.", "questions": "- What is the role of the $s$ in the gradient feature? It seems to encode the sign of $b$ but I don't immediately see why this is important.\n\n- For simple $k$-parity, my understanding is that the learned gradient feature $D_1$ is $1_{A}$, the indicator function for the subset $A$. This would therefore compute parity by computing the parity of $\\sum_{i \\in A} x_i$. Is this correct?", "limitations": "The authors adequately addressed the limitations of this paper.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "- The paper develops a general framework which formalizes the \"one step\" feature learning trick which has recently become popular. This provides an easy to use framework which makes it easy to derive sharp sample complexity guarantees for a number of well defined statistical learning problems.\n\n- The paper instantiates this framework in a number of settings (mixtures of Gaussians, parity, multi-index models) and re-derives a number of sample complexity results.", "weaknesses": "- It appears that this framework can only handle features that can be learned directly at initialization (i.e. through the \"one-step trick\"). In particular, it seems unable to handle multi-step feature learning (e.g. the merged staircase property in [1]).\n\n[1] Abbe et al. (2022) \"The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing thorough suggestions! For the \"one-step trick\" question, please refer to the global response above. Below we address the other comments.\n\n### More about multi-step feature learning\nWe would like to mention that the early-stage analysis (“one-step trick”) is an important and necessary foundation for multi-step analysis. In our analysis, we show that when training time is at a certain range (denoted as $T$), during online convex optimization, the first layer weights will stay in a small regime while the second layer weights will converge to a good classifier based on these gradient features. However, the NN may improve its performance if we continue training NN beyond $T$. Our gradient features are defined based on initialization, i.e., $f^{(0)}$. It is natural to think about whether we can define a new gradient feature set upon $f^{(T)}$ and whether we can get a better guarantee based on the new gradient feature set. If the answer is yes, then we could provide a framework with multi-stage/step feature learning. \n\nIn [1], they introduce a very sophisticated data distribution whose properties are exploited to address the above challenges smartly. The analysis of multi-step feature learning in a general framework is still open and will be our future work, as we mentioned in our Conclusion section “While the current framework focuses on the gradient features in the early gradient steps, whether feature learning also happens in later steps and if so how to formalize that?”. \n\n### Role of $s$\nYes, the $s$ encodes the sign of the bias term, which is important. Recall that we do not update the bias term for simplicity. Let’s consider a simple toy example. Assume we have $f_1(x) = a_1 ReLU(w_1^\\top x + 1)$, $f_2(x) = a_2 ReLU(w_2^\\top x - 1)$ and $f_3(x) = a_3 ReLU(w_3^\\top x + 2)$. \n- The sign of the bias term is important. We can see that we always have $a_1 ReLU(w_1^\\top x + 1) \\neq a_2 ReLU(w_2^\\top x - 1)$ for any $a_1, w_1, a_2, w_2$. It means that $f_1(x)$ and $f_2(x)$ are intrinsically different and have different active patterns. Thus, we need to handle the sign of the bias term carefully. \n- The scaling of the bias is absorbed. On the other hand, we can see that $a_1 ReLU(w_1^\\top x + 1) = a_3 ReLU(w_3^\\top x + 2)$ when $a_1 = 2 a_3, 2w_1 = w_3$. It means the scale of the bias term is less important, which can be absorbed into other terms. \n\nThus, we only need to handle bias with different signs carefully. \n\n### $k$-parity\nYes, it is correct that $D_1$ is $1_A$, the indicator function for the subset $A$ and we build the optimal neural network based on such directions. \n\n[1] Abbe, E., Adsera, E. B., & Misiakiewicz, T. (2022). The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks.\n"}, {"review_id": "2uX1NZ6eLm", "reviewer": "Reviewer_3J3m", "summary": "This paper introduces a general framework for studying feature learning in two-layer NNs. This framework covers feature learning in different examples such as linear classification, mixture of Gaussians and parity functions (it also gives some intuition about the learned features). The neural network under study has one hidden layer and is initialized in a symmetric manner. The training algorithm used is mini-batch GD with fresh samples. In the first iteration, features are learned in some hidden neurons. During the next iterations both layers get updated, however the change in the hidden layer is controlled. \n\n", "questions": "- Q1. Is it possible to extend the results to the continuous setting and $\\ell_2$ loss function (at the beginning $\\ell_2$ also acts like $-y\\hat{y}$)\n- Q2. What are the general insights given by the framework (other than the unification)?\n- Q3. Is it possible to use other activation functions?\n- Q4. Why in Theorems 33 and 46 the algorithm that trains both layers is not used?\n\nSuggestions:\n- It would be really great if some intuition about the gradient feature learning framework can be provided (maybe even in the abstract). \n- The parity formulation in section 4.2 was difficult to understand for me; maybe some changes would help. ", "limitations": "This is a theoretical work and there is not negative societal impact. I think the limitations have been discussed adequately (nonetheless see above). ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper does not freeze the first layer's weights, instead it trains all parameters together and controls the change in the first layer's weights with small learning rate. \n- The literature review is quite extensive; the paper can generally help people interested in feature learning to become familiar with this direction of research. The paper's framework has also been applied to several different learning problems.  \n- I personally liked the connection to the lottery ticket hypothesis. I'd suggest moving (more of) it to the main. ", "weaknesses": "- There are a few limitations which are quite common in the current literature of deep learning theory and have been discussed in the paper: e.g., the feature learning is actually done in the first iteration of gradient descent. See below for some questions regarding the limitations of the framework. \n- Personally I wonder if there are new insights given by this framework or not? (Also see below)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing thorough suggestions! For the limitation of the first gradient descent learning and Q2 General insights, please refer to the global response above. Below we address the other comments.\n\n### Q1 Continous setting and square loss\nIt is possible to extend to a continuous setting and this is one of the works in progress. We can define a gradient feature distribution rather than a gradient feature set. However, we find the technical tools used in the continuous setting are pretty different from the discrete version. \n\nFor the $\\ell\\_2$ loss, it works well in the early-stage where feature learning happens. However, in the later stage analysis (online convex learning part), we need to show that the hidden layer weights stay in a small neighborhood while the top layer weights get updated to a good solution. For the squared loss, it may still be possible to do this control. The current argument easily bounds the change for logistic/hinge losses, but it cannot be directly adopted in the case of squared loss (due to its faster growth). We conjecture a more careful step-by-step inductive argument being needed for that case. \n\n\n### Q3 Activation functions \nYeah, we can change the ReLU activation function to a sub-linear activation function, e.g., leaky ReLU, sigmoid, to get a similar conclusion. First, we need to introduce a corresponding gradient feature set and then we can make it by following the same analysis pipeline. For simplicity, we present ReLU only. \n\n\n### Q4 Theorems 33 and 46\nFor Theorem 33, we provided Theorem 42 as an alternative version that trains both layers. We provide Theorem 33 because (1) use it as a warm-up and (2) follow the original analysis in [1] to give a comparison. \n\nFor Theorem 46: this is because we would like to **unify** previous work. [2] are very closely related to our framework: their analysis for multiple index data follows the same principle and analysis approach as our general framework, although it does not completely fit into our Theorem 3.12 due to some technical differences. We can cover it with our Theorem 3.4. \n- The same principle and analysis approach: [2] shows that the first layer learns good features by one gradient step update, which can approximate the true labels by a low-degree polynomial function. Then, a classifier (the second layer) is trained on top of the learned first layer which leads to the final guarantees. This is consistent with our framework: we first show that the first layer learns good features by one gradient step update, which can approximate the true labels, and then show a good classifier can be learned on the first layer.\n- Technical differences: First, in the second stage, [2] fix the first layer and only update the top layer which is a convex optimization. Our framework allows updates in the first layer and uses online convex learning techniques for the analysis. Second, they consider the square loss (this is used to calculate Hermite coefficients explicitly for gradients, which are useful in the low-degree polynomial function approximation). While in our online convex learning analysis, we need boundedness of the derivative of the loss to show that the first layer weights’ changes are bounded in the second stage. Given the above two technicalities, we analyze their training algorithm (Algorithm 2) which fixes the first layer weights, which currently do not fit directly into our Theorem 3.12 but can fit into Theorem 3.4.\n\n### Parity formulation\nWe will polish the formulation to make the setting more clear. We provide a high-level intuition here. We have $r$ parity functions each corresponding to a block of $k$ dimensions;  $\\mathcal{X}\\_{j,+}$ and $\\mathcal{X}\\_{j,-}$ stands for the component providing a strong signal for the $j$-th parity; $\\mathcal{X}\\_U$ corresponds to uniform distribution unrelated to any parity and providing weak learning signal; $A^\\perp$ is the noise part. The label depends on the sum of the $r$ parity functions. \n\n[1] Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., & Zhang, C. (2022). Hidden progress in deep learning: Sgd learns parities near the computational limit.\n\n[2] Damian, A., Lee, J., & Soltanolkotabi, M. (2022). Neural networks can learn representations with gradient descent.\n"}, {"review_id": "7ijVvHwVMK", "reviewer": "Reviewer_H9VP", "summary": "This paper proposes a general framework of feature learning for two-layer neural networks trained with gradient descent. This framework covers a variety of classification losses and data distributions. Specifically, the authors establish that the loss of neural networks trained with gradient descent is comparable to that of networks with first layers in the direction of a simplified gradient and optimal second layers. The framework is then specialized to multiple examples, including Gaussian mixture classification and learning parity functions.", "questions": "* Is it possible to characterize the reason behind the suboptimality in the $\\tilde{\\mathcal{O}}(d^{1.5})$ sample complexity of this work in comparison with the $O(d)$ sample complexity of [92] to learn a mixture of 4 Gaussians with XOR-like structure?\n\n* A number of figures can be added to improve the presentation of the paper, e.g. for the definition of gradient features and introducing the Gaussian mixtures and the input distribution for learning parity functions.\n\n* Including the choices of step size in Theorem 3.12 can provide additional insights. Specifically, it seems like the theorem is still operating in the regime of a few gradient steps for the first layer, similar to prior work. A more explicit discussion of the training of the first layer would clarify the similarities and distinctions with prior work.", "limitations": "The authors have mostly discussed the limitations of their work. It would help the readers better understand the limitations if there is also a discussion on how realistic the training assumption for the first layer is (i.e. only a few steps or many steps), and also the potential suboptimality in the rate of learning mixtures of Gaussians in comparison with prior work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The presented framework is rather general and unifies the approaches of several recent works, while covering distributions that were not previously covered in the literature.\n* The studied problem is exciting and significant for the community as the existence of such a framework is necessary given the ad-hoc nature of current feature learning analyses.\n* The related literature is covered extensively.\n", "weaknesses": "* The general guarantees provided by the framework (Theorem 3.12) are limited to the performance of the optimal network with first-layer directions approximately aligned with the gradient. Such a statement does not directly imply feature learning for general problems/distributions, as a key part of feature learning seems to be showing that the gradient directions are indeed useful. To that end, the guarantees are more or less similar to prior work and suboptimal in some cases (e.g. Gaussian mixtures).\n\n* The sample complexities of Theorems 4.4 and 4.8 are only stated as polynomials in problem parameters, thus while outperforming kernel methods, provide limited insight in comparison with prior works. Furthermore, the $\\tilde{O}(d^{1.5})$ sample complexity to learn an XOR label from a mixture of 4 Gaussians in Section 4.1.1 seems to be suboptimal, as [92] shows a sample complexity of $O(d)$ for learning the same problem.\n\n* The presentation of the paper can be improved, with specific examples given below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing thorough suggestions! For the theorem still operating in the regime of a few gradient steps for the first layer, please refer to the global response above. Below we address the other comments. In short, we can **improve** our results from $\\tilde{O}(d^{1.5})$ to be $\\tilde{O}(d)$ in sample size under the setting of a mixture of 4 Gaussians with an XOR-like structure. \n\n### Theorem 3.12 are limited to the performance of the optimal network with first-layer directions approximately aligned with the gradient\n\nWe agree that our framework cannot directly imply guarantees for general problems/distributions. We want to emphasize that our key contributions are the concept of gradient features and the idea of using networks with gradient features as baselines to quantify the learning errors. We view the current work as a first step in exploiting the full power of these ideas. Even when we use the framework with the current analysis, we can already unify the prototypical questions/analyses in existing work and also obtain interesting insights. This shows the great potential of the framework. See more discussion in the global response above.\n\n\n### Suboptimal results in the mixture of Gaussians\n\n[92] use ODE to simulate the optimization process for the 2NN learning XOR-shaped Gaussian mixture and give convincing evidence that $O(d)$ number of samples is enough to learn the XOR-shaped Gaussian mixture, yet they did not give a rigorous convergence guarantee for this problem. We successfully derived a convergence guarantee while we required a slightly larger sample size $\\tilde{O}(d^{1.5})$.\n\nMoreover, we can improve our results from $\\tilde{O}(d^{1.5})$ to be $\\tilde{O}(d)$ in sample size. Note that in Lemma D.17 Equation (153) and (154), it still holds if we choose $z = {\\log n \\over n^{-{1 \\over 2}}}$. Then, all $n^{- {1\\over 3}}$ in sample size will change to $ {\\log n \\over n^{-{1 \\over 2}}}$. The probability term will change from $O(\\exp(- n^{1\\over 3}))$ to $O({1 \\over n^2})$. Then, we can still get the final guarantee with $\\tilde{O}(d)$ in a mixture of 4 Gaussians with an XOR-like structure. In our original submission, we do not provide the tightest bound, we will update it in our paper. \n\nOn the other hand, as we mentioned in the limitation section, our framework may or may not recover the width or sample complexity bounds in existing work. This is because our work is analyzing general cases, and thus may not give better than or the same bounds as those in special cases since special cases have more properties that can be exploited to get potentially better bounds.\n\nWe proposed the two key ideas of gradient feature and gradient feature-induced neural networks not only to show their ability to unify several current works but also to open a new direction of thinking with respect to the learning process. These notations have the potential to be extended to multi-layer gradient features and multi-step learning, and this work is only our first step. \n\n### Figures\nGreat suggestion! We made a plot about Gradient Feature under Mixture of Gaussians data in the  **[additional rebuttal pdf](https://openreview.net/attachment?id=BMPAso6Sns&name=pdf)**. We will plot more figures in the revision.  \n"}, {"review_id": "gRUN9M3KJl", "reviewer": "Reviewer_Sj27", "summary": "This paper proposes a general framework for analyzing feature learning in two-layer ReLU neural networks. The idea is to consider the class of two-layer ReLU networks with “gradient features”, i.e. features aligned with the gradients of the loss induced by the distributions of the data and initial model parameters. The main result (Theorem 3.12) is that gradient descent on the two-layer ReLU network achieves generalization error close to the that of the optimal model in this class under weak assumptions. Instantiations of this result are provided for the special cases of the data being generated by a mixture of gaussians and parity functions. In these settings, the optimal loss among models in the gradient feature class is computed, with final generalization rates strictly improving over kernel methods and matching or surpassing the best known results in the feature learning literature.", "questions": "Minor note: the restrictions on $\\tau$ should appear in the statements of Theorems 4.4 and 4.8.", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper analyzes a highly relevant topic.\n\n2. The idea is simple, yet powerful and novel, to my knowledge. The proposed notion of gradient features can likely be used to show feature learning guarantees for models beyond two-layer ReLUs. At present, it allows for competitive feature learning guarantees with very general boundedness assumptions on the data distribution and smoothness of the loss function.\n\n3. The two instantiations of the main result are very helpful to concretize the significance of the main result. Both are very well-studied settings, and the provided guarantees strictly improve upon kernel methods in both cases, and match the rates for learning parities albeit with more general assumptions.\n\n4. The presentation is clear and rigorous.\n\n5. The related works are well-covered.\n", "weaknesses": "1. On its own, it is difficult to gauge the significance of the main result, since in general we do not know whether gradient features can lead to a good model. Indeed, it would be helpful to have more discussion of when gradient features can lead to a good model beyond the two special cases.\n\n2. Related to above, a result showing that if the optimal model in the gradient feature class has large risk, then the ground-truth mapping from inputs to labels is not learnable by gradient descent would strengthen the paper.\n\n3. Like other feature learning studies, the analysis leverages feature learning happening in the early steps of gradient descent, while subsequent updates are shown only to not corrupt the learned features, although it is unclear how well this aligns with practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing thorough suggestions! For feature learning happening in the early steps of gradient descent, please refer to the global response above. Below we address the other comments. In short, we provide some failure cases that our framework cannot cover. \n\n### More discussion\nIn the Appendix, we also provide a linear data model and multiple-index data models. We admit that for general data distributions without detailed information, we may not know whether the gradient feature is good. In our case study, we can build a “ground-truth” network on gradient features. However, for arbitrary data distribution or labeling functions, the “ground-truth” networks may not exist. See examples in the Failure case below. \n\nOn the other hand, one cannot hope for non-vacuous bounds for general problems, given the various hardness results of network learning on general problems, e.g., [1,2]. We agree that given a general problem, we may not have an easy way to compute the “complexity’’ quantity of the problem to get guarantees. This by itself is an interesting and challenging question: given a general problem, is it possible to determine if network learning can have non-vacuous guarantees? We conjecture a negative answer, but this is beyond the scope of the current work and left for future study. See more discussion in the global response above.\n\n### Failure case \nThere are two failure cases we can think of currently:\n- In  [3], they constructed a function that is easy to approximate using a 3-layer network but not approximable by any 2-layer network. Since the function is not approximable by any 2-layer network, it cannot be approximated by the gradient-induced networks as well, so OPT will be large. As a result, the final error will be large. \n- In uniform parity data distribution, considering an odd number of features rather than even, i.e., $k$ is an odd number in Assumption E.30, we can show that our gradient feature set is empty even when $p$ in Equation (5) is exponentially small, thus the OPT is a positive constant since the gradient induced network can only be constants. Meanwhile, the neural network won’t be able to learn this data distribution because its gradient is always 0 through the training, and the final error equals OPT.\n- The above two cases give examples that if the optimal model in the gradient feature class has a large risk, then the ground-truth mapping from inputs to labels is not learnable by gradient descent.\n\n[1] Daniely, A., & Vardi, G. (2020). Hardness of learning neural networks with natural weights.\n\n[2] Daniely, A., Srebro, N., & Vardi, G. (2023). Efficiently Learning Neural Networks: What Assumptions May Suffice?\n\n[3] Safran, I., Eldan, R., & Shamir, O. (2019). Depth separations in neural networks: what is actually being separated?"}]

Author rebuttal:
We thank all reviewers for their constructive and valuable feedback.

We are glad that all reviewers unanimously agree that our theoretical analysis is novel, exciting, powerful, and significant. Reviewers find that our paper provides an easy-to-use, general, and unified framework (zxbj, H9VP, Sj27) that makes it easy to derive sharp sample complexity guarantees strictly improve upon kernel methods (zxbj, 3J3m, Sj27). Reviewers believe that our framework is necessary given the ad-hoc nature of current feature learning analyses (3J3m, H9VP). We are encouraged that reviewers agree that the paper can generally help people interested in feature learning to become familiar with this direction of research (3J3m, H9VP, Sj27).

Here we address the early-stage feature learning problem raised by all reviewers and provide more insights about our framework. We will address the other comments in individual responses to each reviewer.

We agree that early-stage feature learning is a limitation of many latest feature learning analysis papers (zxbj), and we emphasize that our framework provides new insights including some for potentially going beyond the early-stage feature learning. The **challenges in the later-stage analysis** are: (1) the weights in the later stage will not be as normal as the initialization, and we need new tools to analyze their properties; (2) to show that the later-stage features eventually lead to a good solution, we may need new analysis tools for the nonconvex optimization due to the changes in the first layer weights.

On the other hand, when building the general framework, we indeed (1) pin down the key principle behind learning over different data distributions and (2) get new insights.

1. Our framework articulates the following key principles (pointed out for specific problems in existing work but not articulated more generally):
- **Role of gradient**: The gradient leads to the emergence of good features, which is useful for the learning of upper layers in later stages.
- **From features to solutions**: Learned features in early steps will not be distorted, if not improved, in later stages. The training dynamic for upper layers will eventually learn a good combination of hidden neurons based on gradient features, giving a good solution.
2. Some other interesting insights are obtained from the generality of the framework:
- To build a general framework, the meaningful error guarantees should be data-dependent, since NN learning on general data distributions is hard and data-independent guarantees will be vacuous. Comparing the optimal in a family of “ground-truth’’ functions (inspired by agnostic learning in learning theory) is a useful method to obtain the data-dependent bound. We further construct the “ground-truth’’ functions using properties of the training dynamics, i.e., gradient features. This greatly facilitates the analysis of the training dynamics and is the key to obtaining the final guarantees.
- The framework can also be viewed as **using the optimal by gradient-induced NN to measure the “complexity’’** of the problem. For easier problems, this quantity is smaller, and our framework can give a better error bound. So this provides a united way to derive guarantees for specific problems.
- It is important to validate the effectiveness of a general framework by applying it to prototypical problems. Such applications can also help clarify what’s crucial or less crucial in existing analyses for specific problems.
- For an SGD-optimized NN, its **actual representation power** is from the subset of NN based on gradient features, instead of the whole set of NN. This view helps explain the simplicity bias/implicit regularization phenomenon of NN learning in practice.
- Our framework goes **beyond NTK** as we use features from gradients rather than features from random initialization. It means features from gradients are more powerful.
3. More broadly, our framework may give new perspectives about roadmaps forward.
- We argue a new perspective about the connection between the strong representation power and the successful learning of NN. **Traditionally, the strong representation power of NN is the key reason for hardness results of NN learning**: NN has strong representation power and can encode hard learning questions, so they are hard to learn. See the proof in SQ bound from [1] or NP-hardness from [2]. The strong representation power also causes trouble for the statistical aspect: it leads to vacuous generalization bounds when traditional uniform convergence tools are used.
- **Our framework suggests a new perspective in sharp contrast: the strong representation power of NN with gradient features is actually the key to successful learning**. More concretely, the optimal error of the gradient feature-induced NN being small (i.e., strong representation power for a given data distribution) can lead to a small guarantee, which is the key to successful learning.
- The above new perspective suggests a different analysis road than traditional ones. Traditional analysis typically first reasons about the optimal based on the whole function class, i.e. the ground truth, then analyze how NN learns proper features and reaches the optimal. In contrast, our framework defines feature family first, and then reasons about the optimal based on it.
- Our framework provides the foundation for future work on analyzing gradient-based NN learning, which may inspire future directions including but not limited to (1) defining a new feature family for 2-layer NN rather than gradient feature, (2) considering deep NN and introducing new gradient features (e.g., gradient feature notion for upper layers), (3) defining different gradient feature family at different training stages (e.g., gradient feature notion for later stages).

[1] Daniely, A., & Malach, E. (2020). Learning parities with neural networks.

[2] Blum, A., & Rivest, R. (1988). Training a 3-node neural network is NP-complete.

Meta-review: The paper unifies and extends several prior work on feature learning using the "one-step" large batch gradient descent trick via the concept of "gradient features". The key idea is to quantify the useful features that one large step of gradient descent can recover from initialization. On the technical front, the analysis follows similar analyses in the literature with the added extension of training both layers in the second phase of learning. The paper also connects the observations to LTH and simplicity bias among others which are interesting.

The main limitations of the current work is the one-step analysis compared to the multi-step analysis that have been performed for certain settings in prior work (staircase/multi-index), and the lack of clarity on what new insights this framework can have beyond the few examples discussed in the paper (which were somewhat understood before). Despite these limitations, the reviewers and I unanimously agree that this paper provides a useful unification which future work can rely on to build more complex analyses. Therefore, I recommend accepting the paper. I encourage the authors to incorporate the improved bounds for the mixture of Gaussians case, the plot from the rebuttal, and other relevant discussion points that came up in the exchanges.
Decision: Accept (poster)