Title: Provable Guarantees for Neural Networks via Gradient Feature Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: neural networks

gradient descent

feature learning

provable guarantees

theoretical analysis
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

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.

1 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 to the overparameterization of practical networks and the non-convex optimization in the training via gradient descent. One line of work (e.g. [9, 31, 38, 60, 71, 123] and many others) shows under proper conditions, heavily overparameterized networks are approximately linear models over 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.e., the kernel regime, or fixed feature methods). It thus fails to capture the feature learning ability of networks (i.e., to learn a feature mapping for the inputs which allow accurate prediction), which is widely believed to be the key factor to their empirical success in many applications (e.g., [54, 77, 117, 119]). To study feature learning in networks, a recent line of work (e.g. [5, 6, 14, 33, 52, 72, 76, 116] and others) shows examples where networks provably enjoy advantages over fixed feature methods (including NTK), 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 a common principle for feature learning in networks via gradient descent? Is there a unified analysis framework that can clarify the principle and also lead to provable error guarantees for prototypical problem settings?

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. (2) It leads to error guarantees competitive with the optimal in a family of networks that use the features induced by gradients on the

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 (LTH), 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 to 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 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 to derive guarantees. (3) For new perspectives about roadmaps forward, our framework suggests the strong representation power of NN is 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-parameterized 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, 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, 6, 8, 12, 14, 22, 26, 33, 37, 51, 52, 57, 58, 70, 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 a 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, 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 on data distribution or model structure, e.g., Gaussian marginals. [1, 40, 55, 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.e., single-index data model).

[32] considers multiple-index with low-degree polynomials as labeling functions and shows that a one-step gradient update can learn multiple features that lead to accurate prediction. [13, 81] studies one gradient step feature improvements at different learning rates. [97] proposes Recursive Feature Machines to show the mechanism of recursively feature learning but without giving a 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 [ n ] := { 1 , 2 , . . . , n } and ˜ O ( · ) , ˜ Θ( · ) , ˜ Ω( · ) to omit the log term inside. Let X ⊆ R d denote the input space, Y ⊆ R the label space. Let D be an arbitrary data distribution over X × Y . Denote the class of two-layer networks with m neurons as:

where σ ( z ) = max( z, 0) is the ReLU activation function, a ∈ R m is the second layer weight, W ∈ R d × m is the first layer weight, w i is the i -th column of W (i.e., the weight for the i -th neuron), and b ∈ R m 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 ( a ( t ) , W ( t ) , b ) denote the network at time step t . Denote Ξ := ( a , W , b ) , Ξ ( t ) := ( a ( t ) , W ( t ) , b ) . The goal of neural network learning is to minimize the expected risk, i.e., L D ( f ) := E ( x ,y ) ∼D L ( x ,y ) ( f ) , where L ( x ,y ) ( f ) = ℓ ( yf ( x )) is the loss on an example ( x , y ) for some loss function ℓ ( · ) , e.g., the hinge loss ℓ ( z ) = max { 0 , 1 -z } , and the logistic loss ℓ ( z ) = log[1 + exp( -z )] . We also consider ℓ 2 regularization. The regularized loss with regularization coefficient λ is L λ D ( f ) := L D ( f ) + λ 2 ( ∥ W ∥ 2 F + ∥ a ∥ 2 2 ) . Given a training set with n i.i.d. samples Z = { ( x ( l ) , y ( l ) ) } l ∈ [ n ] from D , the empirical risk and its regularized version are:

Then the training process is summarized in Algorithm 1.

Algorithm 1 Network Training via Gradient Descent

In the whole paper, we need some natural assumptions about the data and the loss.

Assumption 3.1. We assume E [ ∥ x ∥ 2 ] ≤ B x 1 , E [ ∥ x ∥ 2 2 ] ≤ B x 2 , ∥ x ∥ 2 ≤ B x and for any label y , we have | y | ≤ 1 . We assume the loss function ℓ ( · ) is a 1-Lipschitz convex decreasing function, normalized ℓ (0) = 1 , | ℓ ′ (0) | = Θ(1) , and ℓ ( ∞ ) = 0 .

Remark 3.2. The above are natural assumptions. Most 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: A 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 to W for t ≥ 2 in Algorithm 1.

The first idea of our framework is to provide guarantees compared to the optimal in a family of networks. Here let us consider networks with specific weights for the first layer:

Definition 3.3. For some fixed W ∈ R d × m , b ∈ R d , and a parameter B a 2 , consider the following family of networks F W , b ,B a 2 , and the optimal approximation network loss in this family:

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 (1) (i.e., the weights after the first gradient step):

Intuitively, the theorem shows that if the weight W (1) after a one-step gradient gives a good set of neurons in the sense that there exists a 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.1).

Such an approach, while simple, has been used to obtain interesting results on network learning in existing work, which shows that W (1) 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 can 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 to consider the typical case where the first layer weights are not frozen. In the following, we 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 networks 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 in 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, as shown in Section 4, many prototypical problems satisfy them and the induced family of networks enjoys low loss, leading to useful guarantees. We begin by considering what features can be learned via gradients. Note that the gradient w.r.t. w i is

Inspired by this, we define the following notion:

Definition 3.5 (Simplified Gradient Vector) . For any w ∈ R d , b ∈ R , a Simplified Gradient Vector is

Remark 3.6. Note that the definition of G ( w , b ) ignores the term ℓ ′ ( yf ( x )) in the gradient, where f is the model

Gradient Feature being cones under Mixture of Gaussians data

20

10

0

10

20

30

40

40

30

20

10

0

10

20

30

20

10

0

10

20

30

40

function. In the early stage of training (or the first gradient step), ℓ ′ ( · ) is approximately a constant, i.e., ℓ ′ ( yf ( x )) ≈ ℓ ′ (0) due to the symmetric initialization (see Equation (8) ).

Definition 3.7 (Gradient Feature) . For a unit vector D ∈ R d with ∥ D ∥ 2 = 1 , and a γ ∈ (0 , 1) , a direction neighborhood (cone) C D,γ is defined as:

Let w ∈ R d , b ∈ R be random variables drawn from some distribution W , B . A Gradient Feature set with parameters p, γ, B G is defined as:

Remark 3.8. When clear from context, write it as S p,γ,B G . 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 s is given by the bias b .) To be a useful 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 Gradient Feature Induced Networks are: F d,m,B F ,S := { f ( a , W , b ) ∈ F d,m ∣ ∣ ∀ i ∈ [ m ] , | a i | ≤ B a 1 , ∥ a ∥ 2 ≤ B a 2 , ( w i , b i / | b i | ) ∈ S, | b i | ≤ B b } , where S is some Gradient Feature set and B F := ( B a 1 , B a 2 , B b ) 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 set S (for simplicity the scaling of w i is absorbed into the scaling of a i and b i ). Definition 3.11 (Optimal Approximation via Gradient Features) . The optimal approximation network and loss using Gradient Feature Induced Networks F d,r,B F ,S are defined as:

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 (e.g., [7, 32, 33, 105]), though some other initialization can also work. Formally, we train a two-layer network with 4 m neurons, f ( a , W , b ) ∈ F d, 4 m . We initialize a (0) i , w (0) i from Gaussians and b i from a constant for i ∈ { 1 , . . . , m } , and initialize the parameters for i ∈ { m +1 , . . . , 4 m } accordingly to get a zero output initial network. Specifically:

where σ 2 a , σ 2 w , ˜ b > 0 are hyper-parameters. 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 ϵ, δ ∈ (0 , 1) , if m ≤ e d and

then with initialization (8) and proper hyper-parameter values, we have with probability ≥ 1 -δ over the initialization and training samples, there exists t ∈ [ T ] in Algorithm 1 with:

̸

Intuitively, the theorem shows when a data distribution admits a small approximation error by some 'ground-truth' network with r neurons using gradient features from S p,γ,B G (i.e., a small optimal approximate loss OPT d,r,B F ,S p,γ,B G ), 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 B a 1 , B x 1 etc. as constants, then the number m of neurons learned is roughly ˜ Θ ( r 4 pϵ 4 ) , a polynomial overparameterization compared to the 'ground-truth' network. The proof shows that such an overparameterization is needed such that some neurons can capture the gradient features given by gradient descent. This is consistent with existing analysis about overparameterization network learning, and also consistent with existing empirical observations.

The error bound consists of three terms. The last term ϵ can be made arbitrarily small, while the other two depend on the concrete data distribution. Specifically, with larger r and γ , the second term increases. While the first term (the optimal approximation loss) decreases, since a larger r means a larger 'ground-truth' network family, and a larger γ means a larger Gradient Feature set S p,γ,B G . 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 s , activation functions, and so on.

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 S p,γ,B G . 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-parameters such as η ( t ) in Theorem D.17. Below, we briefly sketch the key ideas and omit the technical details.

We first show that a 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 (e.g., [33, 105]) and present only the first step.) Let ∇ i denote the gradient of the i -th neuron ∇ w i L D ( f Ξ (0) ) . Denote the subset of neurons with nice gradients approximating feature ( D,s ) as:

Lemma 3.13 (Feature Emergence) . For any r size subset { ( D 1 , s 1 ) , . . . , ( D r , s r ) } ⊆ S p,γ,B G , with probability at least 1 -re -Θ( mp ) , for all j ∈ [ r ] , we have | G ( D j ,s j ) ,Nice | ≥ mp 4 .

This is because ∇ i = ℓ ′ (0) a (0) i E ( x ,y ) [ yσ ′ [〈 w (0) i , x 〉 -b i ] x ] = ℓ ′ (0) a (0) i G ( w (0) i , b i ) . Now consider s j = +1 (the case -1 is similar). Since w i is initialized by Gaussians, by ∇ i 's connection to Gradient Features, we can see that for all i ∈ [ m ] , Pr [ i ∈ G ( D j , +1) ,Nice ] ≥ p 2 . The lemma follows from concentration via a large enough m , i.e., sufficient overparameterization. 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 δ ∈ (0 , 1) , with proper hyper-parameter values, with probability at least 1 -δ , there is ˜ a such that ∥ ˜ a ∥ 0 = O ( r √ mp ) and f (˜ a , W (1) , b ) ( x ) = ∑ 4 m i =1 ˜ a i σ (〈 w (1) i , x 〉 -b i ) satisfies

Given the good set of neurons, we finally show that the remaining gradient steps can learn an accurate classifier. Intuitively, with small step sizes η ( t ) , the weights of the first layer w i 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 guarantee in Theorem 3.12.

4 Applications in Special Cases

In this section we will apply the gradient feature learning framework to 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 of Gaussians

Mixtures of Gaussians are among the most fundamental and widely used statistical models. Recently, it 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 a mixture of r highdimensional Gaussians, and each Gaussian is assigned to one of two possible labels in Y = {± 1 } . Let S ( y ) ⊆ [ r ] denote the set of indices of Gaussians associated with the label y . The data distribution is then: q ( x , y ) = q ( y ) q ( x | y ) , q ( x | y ) = ∑ j ∈S ( y ) p j N j ( x ) , where N j ( x ) is a multivariate normal distribution with mean µ j , covariance Σ j , and p j are chosen such that q ( x , y ) is correctly normalized. We will make some assumptions about the Gaussians, for which we first introduce some notations.

Assumption 4.1. Let 8 ≤ τ ≤ d be a parameter that will control our final error guarantee. Assume

Equiprobable labels: q ( -1) = q (+1) = 1 / 2 .

For all j ∈ [ r ] , Σ j = σ j I d × d . Let σ B := max j ∈ [ r ] σ j and σ B + := max { σ B , B µ 2 } .

r ≤ 2 d, p B ≥ 1 2 d , Ω ( 1 /d + √ τσ B + 2 log d/d ) ≤ B µ 1 ≤ B µ 2 ≤ d .

̸

The Gaussians are well-separated: for all i = j ∈ [ r ] , we have -1 ≤ ⟨ D i , D j ⟩ ≤ θ , where 0 ≤ θ ≤ min { 1 2 r , σ B + B µ 2 √ τ log d d } .

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-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) . ( D j , +1) ∈ S p,γ,B G for all j ∈ [ r ] , where

Given the values on gradient feature parameters p, γ, B G and the optimal approximation loss OPT d,r,B F ,S p,γ,B G , the framework immediately leads to the following guarantee:

Theorem 4.4 (Mixtures of Gaussians: Main Result) . Assume Assumption 4.1. For any ϵ, δ ∈ (0 , 1) , when Algorithm 1 uses hinge loss with

and proper hyper-parameters, then with probability at least 1 -δ , there exists t ∈ [ T ] such that

̸

The theorem shows that gradient descent can learn to a small error via learning the gradient features, given proper hyper-parameters. In particular, we need sufficient overparameterization (a sufficiently large number m of neurons). When σ B + 2 /B 2 µ 1 is a constant which is the prototypical interesting case, and we choose a constant τ , then m is polynomial in the key parameters 1 δ , 1 ϵ , d, r, 1 p B , and the error bound is inverse polynomial in d . The complete proof is given in Appendix E.2.

[46] studies (almost) linear separable cases while our setting includes non-linear separable cases, e.g., XOR. [99] mainly studies neural network classification on 4 Gaussian clusters with XOR structured labels, while our setting is much more general, e.g., our cluster number can extend up to 2 d .

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., Ψ : X → [ -1 , 1] N . For B > 0 , the family of linear models on Ψ with bounded norm B is H B = { h (˜ x ) : h (˜ x ) = ⟨ Ψ(˜ x ) , w ⟩ , ∥ w ∥ 2 ≤ B } . This can capture linear models on fixed finitedimensional feature maps, e.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 Ω(log d ) neurons and Ω ( (log d ) 2 ) samples to achieve arbitrary small constant error (see Theorem E.18 when σ B = 1 ), while fixed feature methods need Ω( d 2 ) features and Ω( d 2 ) 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 Ω(1) neurons and gives convincing evidence that Ω( 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 ) 2 ) . For the proof (detailed in Appendix E.3), we only need to calculate the p, γ, B G of the data distribution carefully and then inject these numbers into Theorem 3.12.

4.2 Parity Functions

Parity functions are a 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 ϕ ∈ {± 1 } d are ∏ i ∈ A ϕ i where A ⊆ [ d ] is a 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 error, 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 ∈ R d × D is an unknown dictionary with D columns that can be regarded as patterns. For simplicity, assume d = D and M is orthonormal. Let ϕ ∈ R d be a hidden representation vector. Let A ⊆ [ 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 = { 1 , . . . , rk } , A ⊥ = { rk +1 , . . . , d } . Also, we split A such that for all j ∈ [ r ] , A j = { ( j -1) k +1 , . . . , jk } . Then the input x and the class label y are given by:

where g ∗ is the ground-truth labeling function mapping from R rk to Y = {± 1 } , ϕ A is the sub-vector of ϕ with indices in A , and XOR ( ϕ A j ) = ∏ l ∈ A j ϕ l is the parity function. We still need to specify the distribution X of ϕ , which determines the structure of the input distribution:

where p o controls the signal noise ratio: if p o 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 j ∈ [ r ] , we have

In X j, + , ϕ A j = [+1 , +1 , . . . , +1] ⊤ and ϕ A \ A j only have zero elements.

In X j, -, ϕ A j = [ -1 , -1 , . . . , -1] ⊤ and ϕ A \ A j only have zero elements.

In X U , we have ϕ A draw from { +1 , -1 } rk uniformly.

In short, we have r parity functions each corresponding to a block of k dimensions; X j, + and X j, -stands for the component providing a strong signal for the j -th parity; X U corresponds to uniform distribution unrelated to any parity and providing weak learning signal; A ⊥ is the noise part. The label depends on the sum of the r parity functions.

Assumption 4.5. Let 8 ≤ τ ≤ d be a parameter that will control our final error guarantee. Assume k is an odd number and: k ≥ Ω( τ log d ) , d ≥ rk +Ω( τr log d ) , p o = O ( rk d -rk ) , p A ≥ 1 d .

Remark 4.6. We set up the problem to be more general than the parity function learning in existing work. If r = 1 , the labeling function reduces to the traditional k -sparse parties of d bits. The assumptions require k, d , and p A to be sufficiently large so as to provide enough large signals for learning. Note that when k = d 16 , r = 1 , p o = 1 2 , our analysis also holds, which shows our framework is beyond the kernel regime (discuss in detail in Section 4.2.1).

To apply our framework, again we only need to compute the Gradient Feature set and the corresponding optimal loss. We first define the Gradient Features: For all j ∈ [ r ] , let D j = ∑ l ∈ Aj M l ∥ ∑ l ∈ Aj M l ∥ 2 .

Lemma 4.7 (Parity Functions: Gradient Features) . We have ( D j , +1) , ( D j , -1) ∈ S p,γ,B G for all j ∈ [ r ] , where

With gradient features from S p,γ,B G , let f ∗ ( x ) = ∑ r j =1 ∑ k i =0 ( -1) i +1 √ k [ σ ( ⟨ D j , x ⟩ -2 i -k -1 √ k ) -2 σ ( ⟨ D j , x ⟩ -2 i -k √ k ) + σ ( ⟨ D j , x ⟩ -2 i -k +1 √ k )] whose hinge loss is 0.

Above, we show that D j is the 'indicator function' for the subset A j so 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 to the following guarantee:

Theorem 4.8 (Parity Functions: Main Result) . Assume Assumption 4.5. For any ϵ, δ ∈ (0 , 1) , when Algorithm 1 uses hinge loss with

and proper hyper-parameters, then with probability at least 1 -δ , there exists t ∈ [ T ] such that

̸

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 (e.g., r = 1 in existing work), and we choose a constant τ , m is polynomial in 1 δ , 1 ϵ , d, k, 1 p A , and the error bound is inverse polynomial in d . The proof is in Appendix E.4. Our setting is more general than that in [33, 76] which corresponds to M = I, r = 1 , p A = 1 4 , p o = 1 2 . [105] study single index learning, where one feature direction is enough for a two-layer network to recover the label, while our setting considers r directions D 1 , . . . , D r , so the network needs to learn multiple directions to get a small error.

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 our 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 = I, r = 1 , p A = 1 4 , k = d 16 , p o = 1 2 , such that all h ∈ H B have hinge-loss at least 1 2 -√ NB 2 k √ 2 .

This means to get an inverse-polynomially small loss, fixed feature models need to have an exponentially large size, i.e., either the number of features N or the norm B needs to be exponential in k . 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.

Parities on Uniform Inputs. When r = 1 , p A = 0 , our problem setting will degenerate to the classic sparse parity function on a uniform input distribution. This has also been used for analyzing network learning [16]. For this case, our framework can get a k 2 O ( k ) log( k ) network width bound and a O ( d k ) sample complexity bound, matching those in [16]. This then again confirms the advantage of network learning over kernel methods that requires d Ω( 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.

5 Further Implications and Conclusion

Our general framework sheds light on 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 we 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]: randomlyinitialized networks contain subnetworks that when trained in isolation reach test accuracy comparable to the original network in a 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 3.14 shows that after the first gradient step, there is a sparse second-layer weight ˜ a with ∥ ˜ a ∥ 0 = O ( r √ mp ) , 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 a small-loss subnetwork f U Ξ with only neurons in U and with second-layer weight ˜ a U on these neurons. Following the same proof of Theorem 3.12:

Proposition 5.1. In the same setting 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 hyperparameter values, we have the same guarantee: with probability ≥ 1 -δ , there is t ∈ [ T ] with

̸

This essentially formally proves LTH for two-layer networks, showing (a) the existence of the winning lottery subnetwork and (b) that gradient descent on 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 A and Appendix B respectively.

Conclusion. We propose a general framework for analyzing two-layer neural network learning by gradient descent and show that it 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. It sheds light on several interesting phenomena in 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?

References

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory . PMLR, 2022.

Emmanuel Abbe, Samy Bengio, Elisabetta Cornacchia, Jon Kleinberg, Aryo Lotfi, Maithra Raghu, and Chiyuan Zhang. Learning to 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.

Shunta Akiyama and Taiji Suzuki. On 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.

Shunta Akiyama and Taiji Suzuki. Excess risk of two-layer reLU neural networks in teacherstudent settings and its superiority to kernel methods. In The Eleventh International Conference on Learning Representations , 2023.

Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? In Advances in Neural Information Processing Systems , 2019.

Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413 , 2020.

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.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in neural information processing systems , 2019.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , 2019.

Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations , 2018.

Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning , pages 322-332. PMLR, 2019.

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.11955 , 2019.

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 arXiv:2205.01445 , 2022.

Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations , 2019.

Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprint arXiv:2207.08799 , 2022.

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.

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.

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.

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.

Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information Processing Systems , 2019.

Yuan Cao, Zhiying Fang, Yue Wu, Ding-Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning, 2020.

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.

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.

Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Efficient approximation of deep relu networks for functions on low dimensional manifolds. Advances in neural information processing systems , 32:8174-8184, 2019.

Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Nonparametric regression on low-dimensional manifolds using deep relu networks: Function approximation and statistical recovery. arXiv preprint arXiv:1908.01842 , 2019.

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.13436 , 2020.

Zhengdao Chen, Eric Vanden-Eijnden, and Joan Bruna. On feature learning in neural networks with global convergence guarantees. In International Conference on Learning Representations , 2022.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. Advances in neural information processing systems , 31, 2018.

Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv:1812.07956 , 2018.

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.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems , 2019.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory . PMLR, 2022.

Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems , 33, 2020.

Amit Daniely and Gal Vardi. Hardness of learning neural networks with natural weights. Advances in Neural Information Processing Systems , 33:930-940, 2020.

Amit Daniely, Nathan Srebro, and Gal Vardi. Efficiently learning neural networks: What assumptions may suffice? arXiv preprint arXiv:2302.07426 , 2023.

Zhiyan Ding, Shi Chen, Qin Li, and Stephen J Wright. Overparameterization of deep resnet: zero loss and mean-field analysis. The Journal of Machine Learning Research , 2022.

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.

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning , 2019.

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2018.

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.

Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations , 2018.

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint arXiv:1903.01611 , 2019.

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.

Spencer Frei, Yuan Cao, and Quanquan Gu. Provable generalization of sgd-trained neural networks of any width in the presence of adversarial label noise. arXiv preprint arXiv:2101.01152 , 2021.

Spencer Frei, Niladri S Chatterji, and Peter L Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202.07626 , 2022.

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.

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 arXiv:2303.01462 , 2023.

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 arXiv:2303.01456 , 2023.

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.

Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Landscape and training regimes in deep learning. Physics Reports , 924:1-18, 2021.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural networks. arXiv preprint arXiv:1906.08899 , 2019.

Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems , 2020.

Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems , 32, 2019.

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.

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.

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.

Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations , 2019.

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.

Arthur Jacot. Implicit bias of large depth networks: a notion of rank for nonlinear functions. In The Eleventh International Conference on Learning Representations , 2023.

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.

Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory , pages 1772-1798. PMLR, 2019.

Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. In International Conference on Learning Representations , 2019.

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems , 33:17176-17186, 2020.

Pritish Kamath, Omar Montasser, and Nathan Srebro. Approximate is good enough: Probabilistic variants of dimensional and margin complexity. In Conference on Learning Theory , 2020.

Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases , pages 795-811. Springer, 2016.

Jonas Moritz Kohler and Aurelien Lucchi. Sub-sampled cubic regularization for non-convex optimization. In International Conference on Machine Learning . PMLR, 2017.

Guy Kornowski, Gilad Yehudai, and Ohad Shamir. From tempered to benign overfitting in relu neural networks. arXiv preprint arXiv:2305.15141 , 2023.

Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. In International Conference on Learning Representations , 2018.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems , 2019.

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.

Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems , 2018.

Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Conference on Learning Theory , 2020.

Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at infinite-width limit. Journal of Machine Learning Research , 2021.

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations , 2019.

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 , 34:12978-12991, 2021.

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.

Christopher D Manning, Kevin Clark, John Hewitt, Urvashi Khandelwal, and Omer Levy. Emergent linguistic structure in artificial neural networks trained by self-supervision. Proceedings of the National Academy of Sciences , pages 30046-30054, 2020.

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.

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.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , pages 2388-2464. PMLR, 2019.

Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of nonlinear feature learning with one gradient step in two-layer neural networks. arXiv preprint arXiv:2310.07891 , 2023.

Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. The Annals of Statistics , 2022.

Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Advances in Neural Information Processing Systems , 33, 2020.

Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with sgd. arXiv preprint arXiv:2209.14863 , 2022.

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.

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.

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-3059. PMLR, 2019.

Vaishnavh Nagarajan and J Zico Kolter. Uniform convergence may be unable to explain generalization in deep learning. Advances in Neural Information Processing Systems , 2019.

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:1905.11604 , 2019.

Behnam Neyshabur. Implicit regularization in deep learning. arXiv preprint arXiv:1709.01953 , 2017.

Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian convolutional neural networks with many channels are gaussian processes. In International Conference on Learning Representations , 2019.

Ryan O'Donnell. Analysis of boolean functions . Cambridge University Press, 2014.

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.

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , pages 84-105, 2020.

Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi. Generalization guarantees for neural networks via harnessing the low-rank structure of the jacobian. arXiv preprint arXiv:1906.05392 , 2019.

Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences , pages 24652-24663, 2020.

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.

Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems , 2008.

Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborov. Classifying highdimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning , pages 8936-8947. PMLR, 2021.

Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. In The Eleventh International Conference on Learning Representations , 2023.

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.

Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In NeurIPS , 2020.

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.

Zhenmei Shi, Yifei Ming, Ying Fan, Frederic Sala, and Yingyu Liang. Domain generalization with nuclear norm regularization. In NeurIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications , 2022.

Zhenmei Shi, Junyi Wei, and Yingyu Liang. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations , 2022.

Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications , pages 1820-1852, 2020.

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.

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-23843, 2021.

Matus Telgarsky. Feature selection with gradient descent on two-layer networks in low-rotation regimes. arXiv preprint arXiv:2208.02789 , 2022.

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.00293 , 2022.

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.

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.

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.

Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760 , 2019.

Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522 , 2020.

Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems , 2019.

Matthew D Zeiler and Rob Fergus. Visualizing and understanding convolutional networks. In European Conference on Computer Vision , 2014.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations , 2017.

Chiyuan Zhang, Samy Bengio, and Yoram Singer. Are all layers created equal? arXiv preprint arXiv:1902.01996 , 2019.

Mo Zhou, Rong Ge, and Chi Jin. A local convergence theory for mildly over-parameterized two-layer neural network. In COLT , 2021.

Difan Zou and Quanquan Gu. An improved analysis of training over-parameterized deep neural networks. Advances in neural information processing systems , 32, 2019.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint arXiv:1811.08888 , 2018.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep ReLU networks. Machine Learning , 109:467-492, 2020.

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.

Contents

Appendix

Appendix A discusses the potential societal impact of our work. Appendix B describes the limitations of our work. In Appendix C, we present our 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.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. 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 a unified theoretical framework that can be applied to different theoretical problems. We propose 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.

On the other hand, this work may lead to a better understanding and inspire the development of improved network learning methods, which may have a positive impact on the theoretical machinelearning community. It may also be beneficial to engineering-inclined machine-learning researchers.

B Limitations

Recover Existing Results. The framework may or may not recover the width or sample complexity bounds in existing work.

The framework can give matching bounds as the existing work in some cases, like parities over uniform inputs (Appendix E.5).

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 to emphasize that our contribution is providing an analysis framework that can (1) formalize the unifying principles of learning features from gradients in network training, and (2) give

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:

In [101], 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.

In uniform parity data distribution, considering an odd number of features rather than even, i.e., k is an odd number in Assumption E.28, we can 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 be able to learn this data distribution because its gradient is always 0 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 a learning hardness example. The above two cases show that if there is an approximation or learning hardness, our gradient feature learning framework may be vacuous because 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. 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.1. The lottery ticket hypothesis (LTH) has been discussed in Section 5. Below we discuss other implications.

Implicit Regularization/Simplicity Bias. It is now well known that practical NN are overparameterized and traditional uniform convergence bounds cannot adequately explain their generalization performance [59, 88, 118]. It is generally believed that the optimization has some implicit regularization effect that restricts learning dynamics to a 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 to 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 F d,r,B F ,S and then learn to approximate the best in a larger family. Consider the number of neurons r for illustration. Let r 1 ≪ r 2 , and let T 1 and T 2 be their corresponding runtime bounds for T in the main Theorem 3.12. Clearly, T 1 ≪ T 2 . Then, at time T 1 , the theorem guarantees the learning dynamics learn to approximate the best in the family F d,r 1 ,B F ,S with r 1 neurons, but not for the larger family F d,r 2 ,B F ,S . Later, at time T 2 , the learning dynamics learn to approximate the best in the larger family F d,r 2 ,B F ,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-induced networks. Furthermore, we can also conclude that 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.

Learning over Different Data Distributions. 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.

Then, more 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 [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.e., 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 specific problems.

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 [33] or 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 are 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.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). 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 a good solution, we 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 a Simplified Gradient Feature Learning Framework in Appendix D.1, which only considers 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 to 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 1 with no updates for the first layer after the first gradient step

Theorem 3.4 (Simple Setting) . Assume ˜ L Z ( f ( a , W (1) , b ) ) is L -smooth to a . Let η ( t ) = 1 L , λ ( t ) = 0 , for all t ∈ { 2 , 3 , . . . , T } . Training by Algorithm 1 with no updates for the first layer after the

Proof of Theorem 3.4. Recall that

We denote f ∗ = argmin f ∈F W , b ,B a 2 L D ( f ) and ˜ f ∗ = argmin f ∈F W , b ,B a 2 ˜ L Z ( f ) . We use a ∗ and ˜ a ∗ to denote their second layer weights respectively. Then, we have

Fixing W (1) , b and optimizing a only is a convex optimization problem. Note that η ≤ 1 L , where ˜ L Z is L -smooth to a . Thus with gradient descent, we have

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

Given an input distribution, we can get a Gradient Feature set S p,γ,B G and f ∗ ( x ) = ∑ r j =1 a ∗ j σ ( 〈 w ∗ j , x 〉 -b ∗ j ) , where f ∗ ∈ F d,r,B F ,S p,γ,B G is a Gradient Feature Induced networks defined in Definition 3.11. Considering training by Algorithm 3, we have the following results.

Theorem D.1 (Gradient Feature Learning Framework under Expected Risk) . Assume Assumption 3.1. For any ϵ, δ ∈ (0 , 1) , if m ≤ e d and

then with proper hyper-parameter values, we have with probability ≥ 1 -δ , there exists t ∈ [ T ] in Algorithm 3 with

̸

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

where γ, B G is the same in the Definition 3.7.

Lemma D.3 (Feature Emergence. Full Statement of Lemma 3.13) . Let λ (1) = 1 η (1) . For any r size subset { ( D 1 , s 1 ) , . . . , ( D r , s r ) } ⊆ S p,γ,B G , with probability at least 1 -2 re -cmp where c > 0 is a universal constant, we have that for all j ∈ [ r ] , | G ( D j ,s j ) ,Nice | ≥ mp 4 .

Proof of Lemma D.3. By symmetric initialization and Lemma F.1, we have for all i ∈ [2 m ]

For all j ∈ [ r ] , as ( D j , s j ) ∈ S p,γ,B G , by Lemma F.3, (1) if s j = +1 , for all i ∈ [ m ] , we have

(2) if s j = -1 , for all i ∈ [2 m ] \ [ m ] , similarly we have

By concentration inequality, (Chernoff's inequality under small deviations), we have

We complete the proof by union bound.

D.2.2 Good Network Exists

Then, the gradients allow for obtaining a set of neurons approximating the 'ground-truth' network with comparable loss.

Lemma D.4 (Existence of Good Networks. Full Statement of Lemma 3.14) . Let λ (1) = 1 η (1) . For any B ϵ ∈ (0 , B b ) , let σ a = Θ ( ˜ b -ℓ ′ (0) η (1) B G B ϵ ) and δ = 2 re - √ mp . Then, with probability at least 1 -δ over the initialization, there exists ˜ a i 's such that f (˜ a , W (1) , b ) ( x ) = ∑ 4 m i =1 ˜ a i σ (〈 w (1) i , x 〉 -b i ) satisfies

Proof of Lemma D.4. Recall f ∗ ( x ) = ∑ r j =1 a ∗ j σ ( 〈 w ∗ j , x 〉 -b ∗ j ) , where f ∗ ∈ F d,r,B F ,S p,γ,B G is defined in Definition 3.11 and let s ∗ j = b ∗ j | b ∗ j | . By Lemma D.3, with probability at least 1 -δ 1 , δ 1 = 2 re -cmp , for all j ∈ [ r ] , we have | G ( w ∗ j ,s ∗ j ) ,Nice | ≥ mp 4 . Then for all i ∈ G ( w ∗ j ,s ∗ j ) ,Nice ⊆ [2 m ] , we have -ℓ ′ (0) η (1) G ( w (0) i , b i ) b ∗ j ˜ b only depend on w (0) i and b i , which is independent of a (0) i . Given Definition 3.7, we have

We split [ r ] into Γ = { j ∈ [ r ] : | b ∗ j | < B ϵ } , Γ -= { j ∈ [ r ] : b ∗ j ≤ -B ϵ } and Γ + = { j ∈ [ r ] : b ∗ j ≥ B ϵ } . Let ϵ a = B x 1 B b √ mpB G B ϵ . Then we know that for all j ∈ Γ + ∪ Γ -, for all i ∈ G ( w ∗ j ,s ∗ j ) ,Nice , we have

Thus, with probability Ω ( 1 √ mp ) over a (0) i , we have

Similarly, for j ∈ Γ , for all i ∈ G ( w ∗ j ,s ∗ j ) ,Nice , with probability Ω ( 1 √ mp ) over a (0) i , we have

For all j ∈ [ r ] , let Λ j ⊆ G ( w ∗ j ,s ∗ j ) ,Nice 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 -δ 2 , δ 2 = re - √ mp , for all j ∈ [ r ] we have | Λ j | ≥ Ω( √ mp ) .

We have for ∀ j ∈ Γ + ∪ Γ -, ∀ i ∈ Λ j ,

Similarly, for ∀ j ∈ Γ , ∀ i ∈ Λ j ,

If i ∈ Λ j , j ∈ Γ + ∪ Γ -, set ˜ a i = a ∗ j | b ∗ j | | Λ j | ˜ b , if i ∈ Λ j , j ∈ Γ , set ˜ a i = a ∗ j B ϵ | Λ j | ˜ b , otherwise set ˜ a i = 0 , we have ∥ ˜ a ∥ 0 = O ( r ( mp ) 1 2 ) , ∥ ˜ a ∥ 2 = O ( B a 2 B b ˜ b ( mp ) 1 4 ) , ∥ ˜ a ∥ ∞ = O ( B a 1 B b ˜ b ( mp ) 1 2 ) .

Finally, we have

We finish the proof by union bound and δ ≥ δ 1 + δ 2 .

D.2.3 Learning an Accurate Classifier

We will use the following theorem from existing work to prove that gradient descent learns a good classifier (Theorem D.9). Theorem D.1 is simply a direct corollary of Theorem D.9.

Theorem D.5 (Theorem 13 in [33]) . Fix some η , and let f 1 , . . . , f T be some sequence of convex functions. Fix some θ 1 , and assume we update θ t +1 = θ t -η ∇ f t ( θ t ) . Then for every θ ∗ the following holds:

To apply the theorem we first present a few lemmas bounding the change in the network during steps.

Lemma D.6 (Bound of Ξ (0) , Ξ (1) ) . Assume the same conditions as in Lemma D.4, and d ≥ log m , with probability at least 1 -δ -1 m 2 over the initialization, ∥ a (0) ∥ ∞ = O ( ˜ b √ log m -ℓ ′ (0) η (1) B G B ϵ ) , and for all i ∈ [4 m ] , we have ∥ w (0) i ∥ 2 = O ( σ w √ d ) . Finally, ∥ a (1) ∥ ∞ = O ( -η (1) ℓ ′ (0)( B x 1 σ w √ d + ˜ b ) ) , and for all i ∈ [4 m ] , ∥ w (1) i ∥ 2 = O ( ˜ b √ log mB x 1 B G B ϵ ) .

Proof of Lemma D.6. By Lemma F.4, we have ∥ a (0) ∥ ∞ = O ( ˜ b √ log m -ℓ ′ (0) η (1) B G B ϵ ) with probability at least 1 -1 2 m 2 by property of maximum i.i.d Gaussians. For any i ∈ [4 m ] , by Lemma F.5 and d ≥ log m , we have

Thus, by union bound, with probability at least 1 -1 2 m 2 , for all i ∈ [4 m ] , we have ∥ w (0) i ∥ 2 = O ( σ w √ d ) .

For all i ∈ [4 m ] , we have

Lemma D.7 (Bound of Ξ ( t ) ) . Assume the same conditions as in Lemma D.6, and let η = η ( t ) for all t ∈ { 2 , 3 , . . . , T } , 0 < TηB x 1 ≤ o (1) , and 0 = λ = λ ( t ) for all t ∈ { 2 , 3 , . . . , T } , for all i ∈ [4 m ] , we have

Proof of Lemma D.7. For all i ∈ [4 m ] , by Lemma D.6,

where we denote Z i = B x 1 ∥ w (1) i ∥ 2 + ˜ b . Then we give a bound of the first layer's weights change,

Combine two bounds, we can get

Let h (1) = | a (1) i | , h (2) = 2 | a (1) i | + ηZ i and h ( t +2) = 2 h ( t +1) -(1 -( ηB x 1 ) 2 ) h ( t ) + ηZ i for n ∈ N + , by Lemma F.8, we have

Thus, by | c 1 | ≤ c 2 , and 0 < TηB x 1 ≤ o (1) , we have

Similarly, by binomial approximation, we also have

We finish the proof by plugging Z i , c 2 into the bound.

Lemma D.8 (Bound of Loss Gap and Gradient) . Assume the same conditions as in Lemma D.7, for all t ∈ [ T ] , we have

and for all t ∈ [ T ] , for all i ∈ [4 m ] , we have

Proof of Lemma D.8. It follows from that

Also, we have

We are now ready to prove the main theorem.

Theorem D.9 (Online Convex Optimization. Full Statement of Theorem D.1) . Consider training by Algorithm 3, and any δ ∈ (0 , 1) . Assume d ≥ log m . Set

Let 0 < TηB x 1 ≤ o (1) , m = Ω ( 1 √ δ + 1 p ( log ( r δ )) 2 ) . With probability at least 1 -δ over the initialization, there exists t ∈ [ T ] such that

Furthermore, for any ϵ ∈ (0 , 1) , set

we have there exists t ∈ [ T ] with

̸

Proof of Theorem D.9. By m = Ω ( 1 √ δ + 1 p ( log ( r δ )) 2 ) we have 2 re - √ mp + 1 m 2 ≤ δ . For any B ϵ ∈ (0 , B b ) , when σ a = Θ ( ˜ b -ℓ ′ (0) η (1) B G B ϵ ) , by Theorem D.5, Lemma D.4, Lemma D.8, with probability at least 1 -δ over the initialization, we have

By Lemma D.4, Lemma D.6, Lemma D.7, when η (1) = Θ ( min { O ( η ) ,O ( η ˜ b ) } -ℓ ′ (0)( B x 1 σ w √ d +˜ b ) ) , we have

Set B ϵ = B x 1 ( mp ) 1 4 √ B b B G , we have σ a = Θ ( ˜ b ( mp ) 1 4 -ℓ ′ (0) η (1) B x 1 √ B G B b ) which satisfy the requirements. Then,

Furthermore, for any ϵ ∈ (0 , 1) , set

we have

We finish the proof as the 0-1 classification error is bounded by the loss function, e.g., I [ sign ( f ( x )) = y ] ≤ ℓ ( yf ( x )) ℓ (0) , where ℓ (0) = 1 .

̸

D.3 More Discussion abut Setting

Range of σ w . In practice, the value of σ w cannot be arbitrary, because its choice will have an effect on the Gradient Feature set S p,γ,B G . On the other hand, d ≥ log m is a natural assumption, otherwise, the two-layer neural networks may fall in the NTK regime.

Parameter Choice. We use λ = 1 /η in the first step so that the neural network will totally forget its initialization, leading to the feature emergence here. This is a common setting for analysis convenience in previous work, e.g., [32, 33, 105]. We can extend this to other choices (e.g., small 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.

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 to one gradient step for the convenience of analysis, while in practice feature learning can happen in multiple steps. The current framework focuses on the 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 a 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 a simple toy example. Assume we have f 1 ( x ) = a 1 σ ( w ⊤ 1 x +1) , f 2 ( x ) = a 2 σ ( w ⊤ 2 x -1) and f 3 ( x ) = a 3 σ ( w ⊤ 3 x +2) , where σ is ReLU activation function which is a homogeneous function.

̸

The sign of the bias term is important. We can see that we always have a 1 σ ( w ⊤ 1 x +1) = a 2 σ ( w ⊤ 2 x -1) for any a 1 , w 1 , a 2 , w 2 . This 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.

The scaling of the bias is absorbed. On the other hand, we can see that a 1 σ ( w ⊤ 1 x +1) = a 3 σ ( w ⊤ 3 x +2) when a 1 = 2 a 3 , 2 w 1 = w 3 . It means that the scale of the bias term is less important, which can be absorbed into other terms.

Thus, we only need to handle bias with different signs carefully.

Gradient Feature Distribution. Wemaydefine a gradient feature distribution rather than a gradient feature set. However, we find that the technical tools used in this continuous setting are pretty different from the discrete version.

Activation Functions. We can change the ReLU activation function to a sublinear 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.

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 error term depends inverse-polynomially on the sample size n .

Definition D.10 (Empirical Simplified Gradient Vector) . Recall Z = { ( x ( l ) , y ( l ) ) } l ∈ [ n ] , for any w ∈ R d , b ∈ R , an Empirical Simplified Gradient Vector is defined as

Definition D.11 (Empirical Gradient Feature) . Recall Z = { ( x ( l ) , y ( l ) ) } l ∈ [ n ] , let w ∈ R d , b ∈ R be random variables drawn from some distribution W , B . An Empirical Gradient Feature set with parameters p, γ, B G is defined as:

When clear from context, write it as ˜ S p,γ,B G .

Considering training by Algorithm 1, we have the following results.

Theorem 3.12 (Main Result) . Assume Assumption 3.1. For any ϵ, δ ∈ (0 , 1) , if m ≤ e d and

then with initialization (8) and proper hyper-parameter values, we have with probability ≥ 1 -δ over the initialization and training samples, there exists t ∈ [ T ] in Algorithm 1 with:

̸

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 n log n > B 2 x B x 2 , with probability at least 1 -O ( 1 n ) over training samples, for all i ∈ [4 m ] , we have

Proof of Lemma D.12. First, we define,

As | ℓ ′ ( z ) | ≤ 1 , | y | ≤ 1 , | σ ′ ( z ) | ≤ 1 , we have z ( l ) is zero-mean random vector with ∥ ∥ z ( l ) ∥ ∥ 2 ≤ 2 B x as well as E [ ∥ ∥ z ( l ) ∥ ∥ 2 2 ] ≤ B x 2 . Then by Vector Bernstein Inequality, Lemma 18 in [66], for 0 < z < B x 2 B x we have

Thus, let z = n -1 2 √ B x 2 log n , with probability at least 1 -O ( 1 n ) , we have

Thus, when n log n > B x B x 2 , let z = n -1 2 B x 2 log n , with probability at least 1 -O ( 1 n ) , we have

Finally, we have

By Assumption 3.1, we have ℓ ( y ( l ) a ⊤ [ σ ( W ⊤ x ( l ) -b ) ]) -E ( x ,y ) ∼D [ ℓ ( y a ⊤ [ σ ( W ⊤ x -b ) ])] is a zero-mean random variable, with bound 2 ∥ a ∥ 0 ∥ a ∥ ∞ (max i ∈ [4 m ] ∥ w i ∥ 2 B x + ˜ b )+2 . By Hoeffding's inequality, for all z > 0 , we have

Thus, with probability at least 1 -O ( 1 n ) , we have

The gradients allow for obtaining a set of neurons approximating the 'ground-truth' network with comparable loss.

Lemma D.13 (Existence of Good Networks under Empirical Risk) . Suppose n log n > Ω ( B 2 x B x 2 + 1 p + B x 2 B 2 G | ℓ ′ (0) | 2 ) . Let λ (1) = 1 η (1) . For any B ϵ ∈ (0 , B b ) , let σ a = Θ ( ˜ b -| ℓ ′ (0) | η (1) B G B ϵ ) and δ = 2 re -√ mp 2 . Then, with probability at least 1 -δ over the initialization and training samples, there exists ˜ a i 's such that f (˜ a , W (1) , b ) ( x ) = ∑ 4 m i =1 ˜ a i σ (〈 w (1) i , x 〉 -b i ) satisfies

Proof of Lemma D.13. Denote ρ = O ( 1 n ) and β = O ( √ B x 2 log n n 1 2 ) . Note that by symmetric initialization, we have ℓ ′ ( yf Ξ (0) ( x )) = | ℓ ′ (0) | for any x ∈ X , so that, by Lemma D.12, we have ∥ ∥ ∥ ˜ G ( w (0) i , b i ) -G ( w (0) i , b i ) ∥ ∥ ∥ 2 ≤ β | ℓ ′ (0) | with probability at least 1 -ρ . Thus, by union bound, we can see that S p,γ,B G ⊆ ˜ S p -ρ,γ + β BG | ℓ ′ (0) | ,B G -β | ℓ ′ (0) | . Consequently, we have OPT d,r,B F , ˜ S p -ρ,γ + β BG | ℓ ′ (0) | ,B G -β | ℓ ′ (0) | ≤ OPT d,r,B F ,S p,γ,B G . Exactly follow the proof in Lemma D.4 by replacing S p,γ,B G to ˜ S p -ρ,γ + β BG | ℓ ′ (0) | ,B G -β | ℓ ′ (0) | . Then, we finish the proof by ρ ≤ p 2 , β | ℓ ′ (0) | ≤ (1 -1 / √ 2) B G .

We will use Theorem D.5 to prove that gradient descent learns a good classifier (Theorem D.17). Theorem 3.12 is simply a direct corollary of Theorem D.17. To apply the theorem we first present a few lemmas bounding the change in the network during steps.

Lemma D.14 (Bound of Ξ (0) , Ξ (1) under Empirical Risk) . Assume the same conditions as in Lemma D.13, and d ≥ log m , with probability at least 1 -δ -1 m 2 -O ( m n ) over the initialization and training samples, ∥ a (0) ∥ ∞ = O ( ˜ b √ log m | ℓ ′ (0) | η (1) B G B ϵ ) , and for all i ∈ [4 m ] , we have ∥ w (0) i ∥ 2 = O ( σ w √ d ) . Finally, ∥ a (1) ∥ ∞ = O ( η (1) | ℓ ′ (0) | ( B x 1 σ w √ d + ˜ b ) + η (1) σ w √ dB x 2 log n n 1 2 ) , and for all i ∈ [4 m ] , ∥ w (1) i ∥ 2 = O ( ˜ b √ log mB x 1 B G B ϵ + ˜ b √ log mB x 2 log n | ℓ ′ (0) | B G B ϵ n 1 2 ) .

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 as in Lemma D.14, and let η = η ( t ) for all t ∈ { 2 , 3 , . . . , T } , 0 < TηB x 1 ≤ o (1) , and 0 = λ = λ ( t ) for all t ∈ { 2 , 3 , . . . , T } . With probability at least 1 -O ( Tm n ) over training samples, for all i ∈ [4 m ] , for all t ∈ { 2 , 3 , . . . , T } , we have

Proof of Lemma D.15. The proof exactly follows the proof of Lemma D.7 with Lemma D.12. Note that, we have

where we denote Z i = ( B x 1 + √ B x 2 log n n 1 2 ) ∥ w (1) i ∥ 2 + ˜ b . Similarly, we have √

Wefinish the proof by following the same arguments in the proof of Lemma D.7 and union bound.

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 1 -O ( T n ) , for all t ∈ [ T ] , we have

With probability at least 1 -O ( T n ) , for all t ∈ [ T ] , i ∈ [4 m ] we have

Proof of Lemma D.16. By Lemma D.8 and Lemma D.12, with probability at least 1 -O ( T n ) , for all t ∈ [ T ] , we have

By Lemma D.8 and Lemma D.12, with probability at least 1 -O ( T n ) , for all t ∈ [ T ] , i ∈ [4 m ] we have

We are now ready to 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

Let 0 < TηB x 1 ≤ o (1) , m = Ω ( 1 √ δ + 1 p ( log ( r δ )) 2 ) and n log n > Ω ( B 2 x B x 2 + 1 p + ( 1 B 2 G + 1 B 2 x 1 ) B x 2 | ℓ ′ (0) | 2 + Tm δ ) . With probability at least 1 -δ over the initialization and training samples, there exists t ∈ [ T ] such that

Furthermore, for any ϵ ∈ (0 , 1) , set

we have there exists t ∈ [ T ] with

̸

Proof of Theorem D.17. We follow the proof in Theorem D.9. By m = Ω ( 1 √ δ + 1 p ( log ( r δ )) 2 ) and n log n > Ω ( B 2 x B x 2 + 1 p + B x 2 B 2 G | ℓ ′ (0) | 2 + Tm δ ) , we have 2 re -√ mp 2 + 1 m 2 + O ( Tm n ) ≤ δ . For any B ϵ ∈ (0 , B b ) , when σ a = Θ ( ˜ b -ℓ ′ (0) η (1) B G B ϵ ) , by Theorem D.5, Lemma D.12, Lemma D.13,

Lemma D.16, with probability at least 1 -δ over the initialization and training samples, we have

Set B ϵ = B x 1 ( mp ) 1 4 √ 2 B b B G , we have σ a = Θ ( ˜ b ( mp ) 1 4 -ℓ ′ (0) η (1) B x 1 √ B G B b ) which satisfy the requirements. By Lemma D.13, Lemma D.14, Lemma D.15, n log n > Ω ( B x 2 B 2 x 1 | ℓ ′ (0) | 2 ) , when η (1) = Θ ( min { O ( η ) ,O ( η ˜ b ) } -ℓ ′ (0)( B x 1 σ w √ d +˜ b ) ) , we have

Then, following the proof in Theorem D.9, we have

Furthermore, for any ϵ ∈ (0 , 1) , set

and note that B G ≤ B x 1 ≤ B x and √ B x 2 ≤ B x naturally, we have

̸

We finish the proof as the 0-1 classification error is bounded by the loss function, e.g., I [ sign ( f ( x )) = y ] ≤ ℓ ( yf ( x )) ℓ (0) , where ℓ (0) = 1 .

E Applications in Special Cases

We present the case 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 -z, 0 } . Let X = R d 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-truth NN close form of the optimal solution which has zero loss. For logistic loss, there is 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 equiprobable, i.e., E [ y = -1] = E [ y = +1] = 1 2 . The input data are linearly separable and there is a ground truth direction w ∗ , where ∥ w ∗ ∥ 2 = 1 , such that y ⟨ w ∗ , x ⟩ > 0 . We also assume E [ yP w ∗⊥ x ] = 0 , where P w ∗⊥ is the projection operator on the complementary space of the ground truth, i.e., the components of input data being orthogonal with the ground truth are independent of the label y . We define the input data signal level as ρ := E [ y ⟨ w ∗ , x ⟩ ] > 0 and the margin as β := min ( x ,y ) y ⟨ w ∗ , x ⟩ > 0 .

We call this data distribution D linear .

Lemma E.2 (Linear Data: Gradient Feature Set) . Let ˜ b = d τ B x 1 σ w , where τ is any number large enough to satisfy d τ/ 2 -1 4 > Ω ( √ B x 2 ρ ) . For D linear setting, we have ( w ∗ , -1) ∈ S p,γ,B G where

Proof of Lemma E.2. By data distribution, we have

Define S Sure : { i ∈ [ m ] : ∥ w (0) i ∥ 2 ≤ 2 √ dσ w } . For all i ∈ [ m ] , we have

For all i ∈ S Sure , by Markov's inequality and considering neuron i + m , we have

For all i ∈ S Sure , by Hölder's inequality, we have

We have

We finish the proof by b i + m | b i + m | = -1 .

Lemma E.3 (Linear Data: Existence of Good Networks) . Assume the same conditions as in Lemma E.2. Define

For D linear setting, we have f ∗ ∈ F d,r,B F ,S p,γ,B G , where r = 2 , B F = ( B a 1 , B a 2 , B b ) = ( 1 β , √ 2 β , 1 B 2 x 1 ) , p = 1 2 , γ = Θ ( √ B x 2 ρd τ/ 2 -1 4 ) , B G = ρ -Θ ( √ B x 2 d τ/ 2 -1 4 ) . We also have OPT d,r,B F ,S p,γ,B G = 0 .

Proof of Lemma E.3. By Lemma E.2 and Lemma F.3, we have f ∗ ∈ F d,r,B F ,S p,γ,B G . We also have

Theorem E.4 (Linear Data: Main Result) . For D linear setting, for any δ ∈ (0 , 1) and for any ϵ ∈ (0 , 1) when

trained by Algorithm 1 with hinge loss, with probability at least 1 -δ over the initialization, with proper hyper-parameters, there exists t ∈ [ T ] such that

̸

Proof of Theorem E.4. Let ˜ b = d τ B x 1 σ w , where τ is a number large enough to satisfy d τ/ 2 -1 4 > Ω ( √ B x 2 ρ ) and O ( B x 1 B 1 4 x 2 β √ ρd τ/ 4 -1 8 ) ≤ ϵ 2 . By Lemma E.3, we have f ∗ ∈ F d,r,B F ,S p,γ,B G , where r = 2 , B F = ( B a 1 , B a 2 , B b ) = ( 1 β , √ 2 β , 1 B 2 x 1 ) , p = 1 2 , γ = Θ ( √ B x 2 ρd τ/ 2 -1 4 ) , B G = ρ -Θ ( √ B x 2 d τ/ 2 -1 4 ) We also have OPT d,r,B F ,S p,γ,B G = 0 .

Adjust σ w such that ˜ b = d τ B x 1 σ w = Θ ( B 1 4 G B a 2 B 3 4 b √ rB a 1 ) . Injecting above parameters into Theorem 3.12, we have with probability at least 1 -δ over the initialization, with proper hyper-parameters, there exists t ∈ [ T ] such that

̸

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 a mixture of r highdimensional Gaussians, and each Gaussian is assigned to one of two possible labels in Y = {± 1 } . Let S ( y ) ⊆ [ r ] denote the set of indices of the Gaussians associated with the label y . The data distribution is then:

where N j ( x ) is a multivariate normal distribution with mean µ j and covariance Σ j , and p j are chosen such that q ( x , y ) is correctly normalized.

We call this data distribution D mixture .

We will make some assumptions about the Gaussians, for which we first introduce some notations. For all j ∈ [ r ] , let y ( j ) ∈ { +1 , -1 } be the label for N j ( x ) .

Assumption E.5 (Mixture of Gaussians. Recap of Assumption 4.1) . Let 8 ≤ τ ≤ d be a parameter that will control our final error guarantee. Assume

Equiprobable labels: q ( -1) = q (+1) = 1 / 2 .

For all j ∈ [ r ] , Σ j = σ j I d × d . Let σ B := max j ∈ [ r ] σ j and σ B + := max { σ B , B µ 2 } .

̸

The Gaussians are well-separated: for all i = j ∈ [ r ] , we have -1 ≤ ⟨ D i , D j ⟩ ≤ θ , where 0 ≤ θ ≤ min { 1 2 r , σ B + B µ 2 √ τ log d d } .

Below, we define a 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) i is the weight for the i -th neuron at initialization. For all j ∈ [ r ] , let S D j ,Sure ⊆ [ m ] be those neurons that satisfy

̸

E.2.2 Mixture of Gaussians: Feature Learning

We show the important Lemma E.7 first and defer other Lemmas after it.

Lemma E.7 (Mixture of Gaussians: Gradient Feature Set. Part statement of Lemma 4.3) . Let C Sure, 1 = 3 2 , C Sure, 2 = 1 2 , ˜ b = C b √ τd log dσ w σ B + , where C b is a large enough universal constant. For D mixture setting, we have ( D j , +1) ∈ S p,γ,B G for all j ∈ [ r ] , where

Proof of Lemma E.7. For all j ∈ [ r ] , by Lemma E.10, for all i ∈ S D j ,Sure ,

where the last inequality follows B µ 1 ≥ Ω ( σ B + √ τ log d d ) .

Thus, we have G ( w (0) i , b i ) ∈ C D j ,γ and ∣ ∣ ∣ 〈 G ( w (0) i , b i ) , D j 〉∣ ∣ ∣ ≤ ∥ G ( w (0) i , b i ) ∥ 2 ≤ B x 1 , b i | b i | = +1 . Thus, by Lemma E.8, we have

Thus, ( D j , +1) ∈ S p,γ,B G . We finish the proof.

Below are Lemmas used in the proof of Lemma E.7. In Lemma E.8, we calculate p used in S p,γ,B G . 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 i ∈ [ m ] , w (0) i ∼ N (0 , σ 2 w I d × d ) , over the random initialization, we have for all i ∈ [ m ] , j ∈ [ r ] ,

Proof of Lemma E.8. Recall for all l ∈ [ r ] , ˜ µ l = µ l / √ d .

WLOG, let j = r . For all l ∈ [ r -1] . We define Z 1 = { l ∈ [ r -1] : ⟨ D l , D r ⟩ ≥ -θ } and Z 2 = { l ∈ [ r -1] : -1 < ⟨ D l , D r ⟩ < -θ } . WLOG, let Z 1 = [ r 1 ] , Z 2 = { r 1 +1 , . . . , r 2 } , where 0 ≤ r 1 ≤ r 2 ≤ r -1 . We define the following events

We define space A = span ( µ 1 , . . . , µ r 1 ) and ˆ µ r = P A ⊥ µ r , where P A ⊥ is the projection operator on the complementary space of A . For l ∈ Z 2 , we also define ˙ µ l = µ l -⟨ µ l ,µ r ⟩ µ r ∥ µ r ∥ 2 2 , and the event

For l ∈ Z 2 , we have µ l = ˙ µ l -ρµ r , where ρ ≥ 0 . So ⟨ w , µ l ⟩ = ⟨ w , ˙ µ l ⟩ -ρ ⟨ w , µ r ⟩ ≤ ⟨ w , ˙ µ l ⟩ when ⟨ w , µ r ⟩ ≥ 0 . As a result, we have

We also have,

For the first condition in Definition E.6, we have,

where the last equality following that ˆ µ r is orthogonal with µ 1 , . . . , µ r 1 and the property of the standard Gaussian vector, and the last inequality follows Lemma F.6.

For the second condition in Definition E.6, by Lemma F.6, we have,

On the other hand, if X is a χ 2 ( k ) random variable. Then we have

Therefore, by assumption B µ 1 ≥ Ω ( σ B + √ τ log d d ) , we have

Recall B µ 1 = min j ∈ [ r ] ∥ ˜ µ j ∥ 2 , B µ 2 = max j ∈ [ r ] ∥ ˜ µ j ∥ 2 . Thus, by union bound, we have

In Lemma E.9, we compute the activation pattern for the neurons in S D j ,Sure .

LemmaE.9 (Mixture of Gaussians: Activation Pattern) . Assume the same conditions as in Lemma E.7, for all j ∈ [ r ] , i ∈ S D j ,Sure , we have

(1) When x ∼ N j ( µ j , σ j I d × d ) , the activation probability satisfies,

̸

(2) For all j ′ = j, j ′ ∈ [ r ] , when x ∼ N j ′ ( µ j ′ , Σ j ′ ) , the activation probability satisfies,

Proof of Lemma E.9. In the proof, we need ˜ b = C b √ τd log dσ w σ B + , where C b is a large enough universal constant. For the first statement, when x ∼ N j ( µ j , σ j I d × d ) , by C Sure, 1 ≥ 3 2 , we have

where the third inequality follows the Chernoff bound and symmetricity of the Gaussian vector.

For the second statement, we prove similarly by 0 < C Sure, 2 ≤ 1 2 .

Then, Lemma E.10 gives gradients of neurons in S D j ,Sure . It shows that these gradients are highly aligned with D j .

Lemma E.10 (Mixture of Gaussians: Feature Emergence) . Assume the same conditions as in Lemma E.7, for all j ∈ [ r ] , i ∈ S D j ,Sure , we have

For any unit vector D ⊥ j which is orthogonal with D j , we have

Proof of Lemma E.10. For all j ∈ [ r ] , i ∈ S D j ,Sure , we have

Thus, by Lemma F.7 and Lemma E.9,

̸

̸

̸

̸

̸

For any unit vector D ⊥ j which is orthogonal with D j , similarly, we have

̸

̸

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

For D mixture setting, we have f ∗ ∈ F d,r,B F ,S p,γ,B G , where B F = ( B a 1 , B a 2 , B b ) = ( 1 √ τ log dσ B + , √ r √ τ log dσ B + , 2 √ τ log dσ B + ) , p = Θ ( B µ 1 √ τ log dσ B + · d ( 9 C 2 b τσ B + 2 / (2 B 2 µ 1 ) ) ) , γ = 1 d 0 . 9 τ -1 . 5 , B G = p B B µ 1 √ d -O ( σ B + d 0 . 9 τ ) and B x 1 = ( B µ 2 + σ B + ) √ d, B x 2 = ( B µ 2 + σ B + ) 2 d . We also have OPT d,r,B F ,S p,γ,B G ≤ 3 d τ + 4 d 0 . 9 τ -1 √ τ log d .

Proof of Lemma E.11. We can check B x 1 = ( B µ 2 + σ B + ) √ d, B x 2 = ( B µ 2 + σ B + ) 2 d by direct calculation. By Lemma E.7, we have f ∗ ∈ F d,r,B F ,S p,γ,B G .

For any j ∈ [ r ] , by B µ 1 ≥ Ω ( σ B + √ τ log d d ) ≥ 4 σ B + √ τ log d d , we have

where the last inequality follows Chernoff bound.

̸

Thus, we have

We also have

̸

where the second last inequality follows Lemma F.7 and r ≤ 2 d . Thus, we have

Theorem 4.4 (Mixtures of Gaussians: Main Result) . Assume Assumption 4.1. For any ϵ, δ ∈ (0 , 1) , when Algorithm 1 uses hinge loss with

̸

̸

̸

and proper hyper-parameters, then with probability at least 1 -δ , there exists t ∈ [ T ] such that

̸

Proof of Theorem 4.4. Let ˜ b = C b √ τd log dσ w σ B + , where C b is a large enough universal constant. By Lemma E.11, we have f ∗ ∈ F d,r,B F ,S p,γ,B G , where B F = ( B a 1 , B a 2 , B b ) = ( 1 √ τ log dσ B + , √ r √ τ log dσ B + , 2 √ τ log dσ B + ) , p = Θ ( B µ 1 √ τ log dσ B + · d ( 9 C 2 b τσ B + 2 / (2 B 2 µ 1 ) ) ) , γ = 1 d 0 . 9 τ -1 . 5 , B G = p B B µ 1 √ d -O ( σ B + d 0 . 9 τ ) and B x 1 = ( B µ 2 + σ B + ) √ d, B x 2 = ( B µ 2 + σ B + ) 2 d . We also have OPT d,r,B F ,S p,γ,B G ≤ 3 d τ + 4 d 0 . 9 τ -1 √ τ log d .

Adjust σ w such that ˜ b = C b √ τd log dσ w σ B + = Θ ( B 1 4 G B a 2 B 3 4 b √ rB a 1 ) . Injecting above parameters into Theorem 3.12, we have with probability at least 1 -δ over the initialization, with proper hyper-parameters, there exists t ∈ [ T ] such that

̸

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-like pattern, for any τ > 0 ,

For all j ∈ [4] , Σ j = σ B I d × d and 1 ≤ σ B ≤ √ d τ log log d .

y (1) = y (2) = 1 and y (3) = y (4) = -1 .

We denote this data distribution as D mixture -xor setting.

E.3.1 Mixture of Gaussians - XOR: Feature Learning

Lemma E.13 (Mixture of Gaussians in [99]: Gradient Feature Set) . Let C Sure, 1 = 6 5 , C Sure, 2 = √ 2 √ τ log log d , ˜ b = √ τd log log dσ w σ B and d is large enough. For D mixture -xor setting, we have ( D j , +1) ∈ S p,γ,B G for all j ∈ [4] , where

Proof of Lemma E.13. For all j ∈ [ r ] , by Lemma E.16, for all i ∈ S D j ,Sure ,

Thus, we have G ( w (0) i , b i ) ∈ C D j ,γ and ∣ ∣ ∣ 〈 G ( w (0) i , b i ) , D j 〉∣ ∣ ∣ ≤ ∥ G ( w (0) i , b i ) ∥ 2 ≤ B x 1 , b i | b i | = +1 . Thus, by Lemma E.14, we have

Thus, ( D j , +1) ∈ S p,γ,B G . We finish the proof.

Lemma E.14 (Mixture of Gaussians in [99]: Geometry at Initialization) . Assume the same conditions as in Lemma E.13. Recall for all i ∈ [ m ] , w (0) i ∼ N (0 , σ 2 w I d × d ) , over the random initialization, we have for all i ∈ [ m ] , j ∈ [4] ,

Proof of Lemma E.14. WLOG, let j = 1 . By Assumption E.12, for the first condition in Definition E.6, we have,

where the the last inequality follows Lemma F.6.

For the second condition in Definition E.6, by Lemma F.6, we have,

On the other hand, if X is a χ 2 ( k ) random variable. Then we have

Therefore, we have

Thus, by union bound, we have

Lemma E.15 (Mixture of Gaussians in [99]: Activation Pattern) . Assume the same conditions as in Lemma E.13, for all j ∈ [4] , i ∈ S D j ,Sure , we have

(1) When x ∼ N j ( µ j , σ B I d × d ) , the activation probability satisfies,

̸

(2) For all j ′ = j, j ′ ∈ [4] , when x ∼ N j ′ ( µ j ′ , σ B I d × d ) , the activation probability satisfies,

Proof of Lemma E.15. In the proof, we need ˜ b = √ τd log log dσ w σ B . For the first statement, when x ∼ N j ( µ j , σ B I d × d ) , by C Sure, 1 ≥ 6 5 , we have

where the third inequality follows the Chernoff bound and symmetricity of the Gaussian vector.

For the second statement, we prove similarly by 0 < C Sure, 2 ≤ √ 2 √ τ log log d .

Then, Lemma E.16 gives gradients of neurons in S D j ,Sure . It shows that these gradients are highly aligned with D j .

Lemma E.16 (Mixture of Gaussians in [99]: Feature Emergence) . Assume the same conditions as in Lemma E.13, for all j ∈ [4] , i ∈ S D j ,Sure , we have

For any unit vector D ⊥ j which is orthogonal with D j , we have

Proof of Lemma E.16. For all j ∈ [4] , i ∈ S D j ,Sure , we have

Thus, by Lemma F.7 and Lemma E.15,

̸

̸

For any unit vector D ⊥ j which is orthogonal with D j , similarly, we have

̸

̸

̸

̸

̸

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 in Lemma E.13 and let τ = 1 and when 0 < ˜ τ ≤ O ( d σ 2 B log d ) . Define

For D mixture -xor setting, we have f ∗ ∈ F d,r,B F ,S p,γ,B G , where B F = ( B a 1 , B a 2 , B b ) = ( 1 √ ˜ τ log dσ B , 2 √ ˜ τ log dσ B , 2 √ ˜ τ log dσ B ) , p = Ω ( 1 σ B · (log d ) σ 2 B ) , γ = σ B √ d , r = 4 , B G = 1 5 √ d and B x 1 = (1 + σ B ) √ d, B x 2 = (1 + σ B ) 2 d . We also have OPT d,r,B F ,S p,γ,B G ≤ 3 d ˜ τ + 4 d 0 . 9˜ τ -1 √ ˜ τ log d .

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 D mixture -xor setting with Assumption E.12, when d is large enough, for any δ ∈ (0 , 1) and for any ϵ ∈ (0 , 1) when

trained by Algorithm 1 with hinge loss, with probability at least 1 -δ over the initialization and training samples, with proper hyper-parameters, there exists t ∈ [ T ] such that

̸

Proof of Theorem E.18. Let ˜ b = √ d log log dσ w σ B . By Lemma E.17, let τ = 1 and when ˜ τ = O ( d σ 2 B log d ) , we have f ∗ ∈ F d,r,B F ,S p,γ,B G , where B F = ( B a 1 , B a 2 , B b ) = ( 1 √ ˜ τ log dσ B , 2 √ ˜ τ log dσ B , 2 √ ˜ τ log dσ B ) , p = Ω ( 1 σ B · (log d ) σ 2 B ) , γ = σ B √ d , r = 4 , B G = 1 5 √ d and B x 1 = (1 + σ B ) √ d, B x 2 = (1 + σ B ) 2 d . We also have OPT d,r,B F ,S p,γ,B G ≤ 3 d ˜ τ + 4 d 0 . 9˜ τ -1 √ ˜ τ log d .

Adjust σ w such that ˜ b = √ d log log dσ w σ B = Θ ( B 1 4 G B a 2 B 3 4 b √ rB a 1 ) . Injecting above parameters into Theorem 3.12, we have with probability at least 1 -δ over the initialization, with proper hyperparameters, there exists t ∈ [ T ] such that

̸

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 ∈ R d × D is an unknown dictionary with D columns that can be regarded as patterns. For simplicity, assume d = D and M is orthonormal. Let ϕ ∈ R d be a hidden representation vector. Let A ⊆ [ 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 = { 1 , . . . , rk } , A ⊥ = { rk +1 , . . . , d } . Also, we split A such that for all j ∈ [ r ] , A j = { ( j -1) k +1 , . . . , jk } . Then the input x and the class label y are given by:

where g ∗ is the ground-truth labeling function mapping from R rk to Y = {± 1 } , ϕ A is the sub-vector of ϕ with indices in A , and XOR ( ϕ A j ) = ∏ l ∈ A j ϕ l is the parity function.

We still need to specify the distribution of ϕ , which determines the structure of the input distribution:

For all corresponding ϕ A ⊥ in X , we have ∀ l ∈ A ⊥ , independently:

where p o controls the signal noise ratio: if p o 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 j ∈ [ r ] , we have

In X j, + , ϕ A j = [+1 , +1 , . . . , +1] ⊤ and ϕ A \ A j only have zero elements.

In X j, -, ϕ A j = [ -1 , -1 , . . . , -1] ⊤ and ϕ A \ A j only have zero elements.

In X U , we have ϕ A draw from { +1 , -1 } rk uniformly.

We call this data distribution D parity .

Assumption E.19 (Parity Functions. Recap of Assumption 4.5) . Let 8 ≤ τ ≤ d be a parameter that will control our final error guarantee. Assume k is an odd number and:

Remark E.20. The assumptions require k, d , and p A to be sufficiently large so as to provide enough large signals for learning. When p o = Θ( rk d -rk ) means that the signal noise ratio is constant: the expected norm of ϕ A and that of ϕ A ⊥ are comparable.

To apply our framework, again we 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: For all j ∈ [ r ] , let

Remark E.21. Our data distribution is symmetric, which means for any ϕ ∈ R d :

-y = g ∗ ( -ϕ A ) and -x = M ( -ϕ ) ,

P ( ϕ ) = P ( -ϕ ) ,

E ( x ,y ) [ y x ] = 0 .

Below, we define a 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) i is the weight for the i -th neuron at initialization. For all j ∈ [ r ] , let S D j ,Sure ⊆ [ m ] be those neurons that satisfy

̸

where P A , P A ⊥ are the projection operator on the space M A and M A ⊥ .

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 C Sure, 1 = 3 2 , C Sure, 2 = 1 2 , ˜ b = C b √ τrk log dσ w , where C b is a large enough universal constant. For D parity setting, we have ( D j , +1) , ( D j , -1) ∈ S p,γ,B G for all j ∈ [ r ] , where

Proof of Lemma E.23. Note that for all l ∈ [ d ] , we have M ⊤ l x = ϕ l . For all j ∈ [ r ] , by Lemma E.26, for all i ∈ S D j ,Sure , when γ = 1 d τ -2 ,

Thus, we have G ( w (0) i , b i ) ∈ C D j ,γ and √ kp A -O ( √ k d τ ) ≤ ∥ G ( w (0) i , b i ) ∥ 2 ≤ √ kp A + O ( 1 d τ -1 2 ) , b i | b i | = +1 . Thus, by Lemma E.24, we have

Thus, ( D j , +1) ∈ S p,γ,B G . Since E ( x ,y ) [ y x ] = 0 , by Lemma F.2 and considering i ∈ [2 m ] \ [ m ] , we have ( D j , -1) ∈ S p,γ,B G . We finish the proof.

Below are Lemmas used in the proof of Lemma E.23. In Lemma E.24, we calculate p used in S p,γ,B G .

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 ∈ [ m ] , w (0) i ∼ N (0 , σ 2 w I d × d ) , over the random initialization, we have for all i ∈ [ m ] , j ∈ [ r ] ,

̸

Proof of Lemma E.24. For every i ∈ [ m ] , j, j ′ ∈ [ r ] , j = j ′ , by Lemma F.6,

On the other hand, if X is a χ 2 ( k ) random variable, by Lemma F.5, we have

Therefore, by assumption rk ≥ Ω( τr log d ) , d -rk ≥ Ω( τr log d ) , we have

Thus, by union bound, and D 1 , . . . , D r being orthogonal with each other, we have

In Lemma E.25, we compute the activation pattern for the neurons in S D j ,Sure .

Lemma E.25 (Parity Functions: Activation Pattern) . Assume the same conditions as in Lemma E.23, for all j ∈ [ r ] , i ∈ S D j ,Sure , we have

(1) When x ∼ X , we have

(2) When x ∼ X U , we have

(3) When x ∼ X U , the activation probability satisfies,

(4) When x ∼ X j, + , the activation probability satisfies,

̸

(5) For all j ′ = j, j ′ ∈ [ r ] , s ∈ { + , -} , when x ∼ X j ′ ,s , or x ∼ X j, -, the activation probability satisfies,

Proof of Lemma E.25. For the first statement, when x ∼ X , note that 〈 w (0) i , M l 〉 ϕ l is a mean-zero sub-Gaussian random variable with sub-Gaussion norm Θ (∣ ∣ ∣ 〈 w (0) i , M l 〉∣ ∣ ∣ √ p o ) .

where the inequality follows general Hoeffding's inequality.

For the second statement, when x ∼ X U , by Hoeffding's inequality,

In the proof of the third to the last statement, we need ˜ b = C b √ τrk log dσ w , where C b is a large enough universal constant.

For the third statement, when x ∼ X U , by union bound and previous statements,

For the forth statement, when x ∼ X j, + , by C Sure, 1 ≥ 3 2 and previous statements,

For the last statement, we prove similarly by 0 < C Sure, 2 ≤ 1 .

2

Then, Lemma E.26 gives gradients of neurons in S D j ,Sure . It shows that these gradients are highly aligned with D j .

Lemma E.26 (Parity Functions: Feature Emergence) . Assume the same conditions as in Lemma E.23, for all j ∈ [ r ] , i ∈ S D j ,Sure , we have the following holds:

(1) For all l ∈ A j , we have

̸

(2) For all l ∈ A j ′ , any j ′ = j, j ′ ∈ [ r ] , we have

(3) For all l ∈ A ⊥ , we have

Proof of Lemma E.26. For all l ∈ [ d ] , we have

For the first statement, for all l ∈ A j , by Lemma E.25 (3) and (4), we have

and we also have

̸

Similarly, for the second statement, for all l ∈ A j ′ , any j ′ = j, j ′ ∈ [ r ] , by Lemma E.25 (3) and (5), we have

For the third statement, for all l ∈ A ⊥ , by Lemma E.25 (3), (4), (5), we have

where the second inequality follows 2 rp A ≤ 1 and the third inequality follows the triangle inequality.

E.4.3 Parity Functions: Final Guarantee

LemmaE.27 (Parity Functions: Existence of Good Networks. Part statement of Lemma 4.7) . Assume the same conditions as in Lemma E.23. Define

Proof of Lemma E.27. We can check B x 1 = √ d, B x 2 = d by direct calculation. By Lemma E.23, we have f ∗ ∈ F d, 3 r ( k +1) ,B F ,S p,γ,B G . We note that

is a bump function for ⟨ D j , x ⟩ at 2 i -k √ k . We can check that yf ∗ ( x ) ≥ 1 . Thus, we have

Theorem 4.8 (Parity Functions: Main Result) . Assume Assumption 4.5. For any ϵ, δ ∈ (0 , 1) , when Algorithm 1 uses hinge loss with

and proper hyper-parameters, then with probability at least 1 -δ , there exists t ∈ [ T ] such that

̸

Proof of Theorem 4.8. Let ˜ b = C b √ τrk log dσ w , where C b is a large enough universal constant. By Lemma E.27, we have f ∗ ∈ F d, 3 r ( k +1) ,B F ,S , where B F = ( B a 1 , B a 2 , B b ) =

p,γ,B G √ ( ) √ √ )

p,γ,B G

Adjust σ w such that ˜ b = C b √ τrk log dσ w = Θ ( B 1 4 G B a 2 B 3 4 b √ rB a 1 ) . Injecting above parameters into Theorem 3.12, we have with probability at least 1 -δ over the initialization, with proper hyperparameters, there exists t ∈ [ T ] such that

̸

E.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 a warm-up and (2) follow the original analysis in [15] to give a comparison. We will provide Theorem E.40 as an 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 follow the data distribution in Appendix E.4.1. Let r = 1 , p A = 0 , p o = 1 2 , M = I d × d and d ≥ 2 k 2 , and k is an even number.

We denote this data distribution as D parity -uniform setting.

To apply our framework, again we 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

We follow the initialization and training dynamic in [15].

Initialization and Loss. We use hinge loss and we have unbiased initialization, for all i ∈ [ m ] ,

Training Process. We use the following one-step training algorithm for this specific data distribution.

Algorithm 4 Network Training via Gradient Descent [15]. Special case of Algorithm 2

Initialize ( a (0) , W (0) , b ) as in Equation (8) and Equation (446); Sample Z ∼ D n parity -uniform W (1) = W (0) -η (1) ( ∇ W ˜ L Z ( f Ξ (0) ) + λ (1) W (0) ) a (1) = a (0) -η (1) ( ∇ a ˜ L Z ( f Ξ (0) ) + λ (1) a (0) ) for t = 2 to T do a ( t ) = a ( t -1) -η ( t ) ∇ a ˜ L Z ( f Ξ ( t -1) ) end for

Use the notation in Section 5.3 of [92], for every S ∈ [ n ] , s.t. | S | = k , we define

Lemma E.29 (Uniform Parity Functions: Existence of Good Networks. Rephrase of Lemma 5 in [15]) . For every ϵ, δ ∈ (0 , 1 / 2) , denoting τ = | ξ k -1 | 16 k √ 2 d log(32 k 3 d/ϵ ) , let η (1) = 1 k | ξ k -1 | , λ (1) = 1 η (1) , m ≥ k · 2 k log( k/δ ) , n ≥ 2 τ 2 log(4 dm/δ ) and d ≥ Ω ( k 4 log( kd/ϵ ) ) , w.p. at least 1 -2 δ over the initialization and the training samples, there exists ˜ a ∈ R m with ∥ ˜ a ∥ ∞ ≤ 8 k and ∥ ˜ a ∥ 2 ≤ 8 k √ k such that f (˜ a , W (1) , b ) satisfies

Additionally, it holds that ∥ σ ( W (1) ⊤ x -b ) ∥ ∞ ≤ d +1 .

Remark E.30. In [15], they update the bias term in the first gradient step. However, if we check the proof carefully, we 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 ϵ ∈ (0 , 1 / 2) and let m ≥ Ω ( k · 2 k log( k/ϵ ) ) , n ≥ Ω ( k 7 / 6 d ( d k -1 ) log( kd/ϵ ) log( dm/ϵ ) + k 3 md 2 ϵ 2 ) , d ≥ Ω ( k 4 log( kd/ϵ ) ) . Let η (1) = 1 k | ξ k -1 | , λ (1) = 1 η (1) , and η = η ( t ) = Θ ( 1 d 2 m ) , for all t ∈ { 2 , 3 , . . . , T } . If T ≥ Ω ( k 3 md 2 ϵ ) , then training by Algorithm 4 with hinge loss, w.h.p. over the initialization and the training samples, there exists t ∈ [ T ] such that

̸

Proof of Theorem E.31. By Lemma E.29, w.h.p., we have for properly chosen hyper-parameters,

We compute the L -smooth constant of ˜ L Z ( f ( a , W (1) , b ) ) to a .

By the Lemma E.29, we have ∥ σ ( W (1) ⊤ x -b ) ∥ ∞ ≤ d +1 . Thus, we have,

This means that we can let η = Θ ( 1 d 2 m ) and we will get our convergence result. Note that we have a (1) = 0 and ∥ ˜ a ∥ 2 = O ( k √ k ) . So, if we choose T ≥ Ω ( k 3 ϵη ) , there exists t ∈ [ T ] such that ˜ L Z ( f ( a ( t ) , W (1) , b ) ) -˜ L Z ( f (˜ a , W (1) , b ) ) ≤ O ( L ∥ a (1) -˜ a ∥ 2 2 T ) ≤ ϵ/ 3 .

E.6 Uniform Parity Functions: Alternative Analysis

It is also possible to unify [15] into our 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

d

Definition E.32 (Gradient Feature with Infinity Norm) . For a unit vector D ∈ R with ∥ D ∥ 2 = 1 , and a γ ∞ ∈ (0 , 1) , a direction neighborhood (cone) C ∞ D,γ ∞ is defined as: C ∞ D,γ ∞ := { w ∣ ∣ ∣ ∥ ∥ ∥ w ∥ w ∥ -D ∥ ∥ ∥ ∞ < γ ∞ } . Let w ∈ R d , b ∈ R be random variables drawn from some distribution W , B . A Gradient Feature set with parameters p, γ ∞ , B G , B G 1 is defined as:

S p,γ ∞ ,B G ,B G 1 .

Definition E.33 (Optimal Approximation via Gradient Features with Infinity Norm) . The Optimal Approximation network and loss using gradient feature induced networks F d,r,B F ,S ∞ p,γ ∞ ,B G ,B G 1 are defined as:

We consider the data distribution in Appendix E.4.1 with Assumption E.28, i.e., D parity -uniform in Appendix E.5. Note that with this dataset, we have ∥ x ∥ ∞ ≤ B x ∞ = 1 . We use the following unbiased initialization:

Let ∇ i denote the gradient of the i -th neuron ∇ w i L D ( f Ξ (0) ) . Denote the subset of neurons with nice gradients approximating feature ( D,s ) as:

Lemma E.34 (Existence of Good Networks. Modified Version of Lemma 3.14 Under Uniform Parity Setting) . Let λ (1) = 1 η (1) . For any B ϵ ∈ (0 , B b ) , let σ a = Θ ( ˜ b -ℓ ′ (0) η (1) B G B ϵ ) and δ = 2 re - √ mp + 1 d 2 . Then, with probability at least 1 -δ over the initialization, there exists ˜ a i 's such that f (˜ a , W (1) , b ) ( x ) = ∑ 4 m i =1 ˜ a i σ (〈 w (1) i , x 〉 -b i ) satisfies

Proof of Lemma E.34. Recall f ∗ ( x ) = ∑ r j =1 a ∗ j σ ( 〈 w ∗ j , x 〉 -b ∗ j ) , where f ∗ ∈ F d,r,B F ,S ∞ p,γ ∞ ,B G ,B G 1 is defined in Definition E.33 and let s ∗ j = b ∗ j | b ∗ j | . By Lemma D.3, with probability at least 1 -δ 1 , δ 1 = 2 re -cmp , for all j ∈ [ r ] , we have | G ∞ ( w ∗ j ,s ∗ j ) ,Nice | ≥ mp 4 . Then for all i ∈ G ∞ ( w ∗ j ,s ∗ j ) ,Nice ⊆ [2 m ] , we have -ℓ ′ (0) η (1) G ( w (0) i , b i ) b ∗ j ˜ b only depend on w (0) i and b i , which is independent of a (0) i . Given Definition E.32, we have

We split [ r ] into Γ = { j ∈ [ r ] : | b ∗ j | < B ϵ } , Γ -= { j ∈ [ r ] : b ∗ j ≤ -B ϵ } and Γ + = { j ∈ [ r ] : b ∗ j ≥ B ϵ } . Let ϵ a = B G 1 B b √ mpB G B ϵ . Then we know that for all j ∈ Γ + ∪ Γ -, for all i ∈ G ∞ ( w ∗ j ,s ∗ j ) ,Nice , we have

Thus, with probability Ω ( 1 √ mp ) over a (0) i , we have

Similarly, for j ∈ Γ , for all i ∈ G ∞ ( w ∗ j ,s ∗ j ) ,Nice , with probability Ω ( 1 √ mp ) over a (0) i , we have

For all j ∈ [ r ] , let Λ j ⊆ G ∞ ( w ∗ j ,s ∗ j ) ,Nice be the set of i 's such that condition Equation (467) or Equation (468) are satisfied. By Chernoff bound and union bound, with probability at least 1 -δ 2 , δ 2 = re - √ mp , for all j ∈ [ r ] we have | Λ j | ≥ Ω( √ mp ) . We have for ∀ j ∈ Γ + ∪ Γ -, ∀ i ∈ Λ j ,

With probability 1 -1 d 2 by Hoeffding's inequality. Similarly, for ∀ j ∈ Γ , ∀ i ∈ Λ j ,

If i ∈ Λ j , j ∈ Γ + ∪ Γ -, set ˜ a i = a ∗ j | b ∗ j | | Λ j | ˜ b , if i ∈ Λ j , j ∈ Γ , set ˜ a i = a ∗ j B ϵ | Λ j | ˜ b , otherwise set ˜ a i = 0 , we have ∥ ˜ a ∥ 0 = O ( r ( mp ) 1 2 ) , ∥ ˜ a ∥ 2 = O ( B a 2 B b ˜ b ( mp ) 1 4 ) , ∥ ˜ a ∥ ∞ = O ( B a 1 B b ˜ b ( mp ) 1 2 ) .

Finally, we have

We finish the proof by union bound and δ ≥ δ 1 + δ 2 + 1 d 2 .

Lemma E.35 (Empirical Gradient Concentration Bound for Single Coordinate) . For i ∈ [ m ] , when n ≥ (log( d )) 6 , with probability at least 1 -O ( exp ( -n 1 3 )) over training samples, we have

Proof of Lemma E.35. First, we define,

∣ ∣

As | ℓ ′ ( z ) | ≤ 1 , | y | ≤ 1 , | σ ′ ( z ) | ≤ 1 , we have z ( l ) i,j is zero-mean random variable with ∣ ∣ z ( l ) i,j ∣ ∣ ≤ 2 B x ∞ as well as E [ ∣ ∣ ∣ z ( l ) i,j ∣ ∣ ∣ 2 2 ] ≤ 4 B 2 x ∞ . Then by Bernstein Inequality, for 0 < z < 2 B x ∞ , we have

Thus, for some i ∈ [ m ] , when n ≥ (log( d )) 6 , with probability at least 1 -O ( expΘ ( -n 1 3 )) , from a union bound over j ∈ [ d ] , we have, for ∀ j ∈ [ d ] ,

Lemma E.36 (Existence of Good Networks under Empirical Risk. Modified version of Lemma D.13 Under Uniform Parity Setting) . Suppose n > Ω ( ( B x √ B x 2 +log 1 p + B x ∞ B G | ℓ ′ (0) | + B x ∞ B G 1 | ℓ ′ (0) | ) 3 +(log( d )) 6 ) . Let λ (1) = 1 η (1) . For any B ϵ ∈ (0 , B b ) , let σ a = Θ ( ˜ b -| ℓ ′ (0) | η (1) B G B ϵ ) and δ = 2 re -√ mp 2 + 1 d 2 . Then, with probability at least 1 -δ over the initialization and training samples, there exists ˜ a i 's such that f (˜ a , W (1) , b ) ( x ) = ∑ 4 m i =1 ˜ a i σ (〈 w (1) i , x 〉 -b i ) satisfies

Proof of Lemma E.36. Denote ρ = O ( expΘ ( -n 1 3 )) and β = O ( B x ∞ n 1 3 ) . Note that by symmetric initialization, we have ℓ ′ ( yf Ξ (0) ( x )) = | ℓ ′ (0) | for any x ∈ X , so that, by Lemma E.35, we have ∣ ∣ ∣ ˜ G ( w (0) i , b i ) j -G ( w (0) i , b i ) j ∣ ∣ ∣ ≤ β | ℓ ′ (0) | with probability at least 1 -ρ . Thus, by union bound, we can see that S ∞ p,γ ∞ ,B G ,B G 1 ⊆ ˜ S ∞ p -ρ,γ ∞ + β BG | ℓ ′ (0) | ,B G -β | ℓ ′ (0) | ,B G 1 + β | ℓ ′ (0) | . Consequently, we have OPT d,r,B F , ˜ S ∞ p -ρ,γ ∞ + β BG | ℓ ′ (0) | ,B G -β | ℓ ′ (0) | ,B G 1 + β | ℓ ′ (0) | ≤ OPT d,r,B F ,S ∞ p,γ ∞ ,B G ,B G 1 . Exactly follow the proof in Lemma D.4 by replacing S ∞ p,γ ∞ ,B G ,B G 1 to ˜ S ∞ p -ρ,γ ∞ + β BG | ℓ ′ (0) | ,B G -β | ℓ ′ (0) | ,B G 1 + β | ℓ ′ (0) | . Then, we finish the proof by ρ ≤ p 2 , β | ℓ ′ (0) | ≤ (1 -1 / √ 2) B G , β | ℓ ′ (0) | ≤ ( √ 2 -1) B G 1 .

Theorem E.37 (Online Convex Optimization under Empirical Risk. Modified version of Theorem D.17 Under Uniform Parity Setting ) . Consider training by Algorithm 1, and any δ ∈ (0 , 1) . Assume d ≥ log m,δ ≤ O ( 1 d 2 ) . Set

Let 0 < TηB x 1 ≤ o (1) , m = Ω ( 1 √ δ + 1 p ( log ( r δ )) 2 ) and n > Ω ( ( B x √ B x 2 +log Tm pδ +(1 + 1 B G + 1 B G 1 ) B x ∞ | ℓ ′ (0) | ) 3 ) . With probability at least 1 -δ over the

initialization and training samples, there exists t ∈ [ T ] such that

Furthermore, for any ϵ ∈ (0 , 1) , set

we have there exists t ∈ [ T ] with

̸

Proof of Theorem E.37. Proof of the theorem and parameter choices remain the same as Theorem D.17 except for setting B ϵ = √ B x 1 B G 1 ( mp ) 1 4 √ B b B G and apply Lemma E.36.

E.6.2 Feature Learning of Uniform Parity Functions

We denote

Lemma E.38 (Uniform Parity Functions: Gradient Feature Learning. Corollary of Lemma 3 in [15]) . Assume that n ≥ 2( k +1) 2 . Then, the following holds:

If j ∈ A , then

If i / ∈ A , then and

For D parity -uniform setting, we have f ∗ ∈ F d, 3( k +1) ,B F ,S ∞ p,γ ∞ ,B G ,B G 1 where B F = ( B a 1 , B a 2 , B b ) = ( 2 √ k, 2 √ ( k ( k +1)) , k +1 √ k ) , p = Θ ( 1 2 k -1 ) , γ ∞ = O ( √ k d -k ) , B G = Θ( B G 1 ) = Θ( d -k ) and B x 1 = √ d , B x 2 = d . We also have OPT d, 3( k +1) ,B F ,S ∞ p,γ ∞ ,B G ,B G 1 = 0 .

Proof of Lemma E.39. Fix index i , with probability p 1 = Θ(2 -k ) , we will have w (0) i,j = sign ( a (0) i ) · sign ( ξ k -1 ) , for ∀ j . For w (0) i that satisfy these conditions, we will have:

Then by Lemma 4 in [15], we have

and

From here, we can see that if we set γ ∞ = k d -k , B G = B G 1 = √ k | ξ k -1 | 2 +( d -k ) | ξ k +1 | 2 , p = p 1 , we will have ( D, +1) , ( D, -1) ∈ S ∞ p,γ ∞ ,B G ,B G 1 by our symmetric initialization. As a result, we have f ∗ ∈ F d, 3( k +1) ,B F ,S ∞ p,γ ∞ ,B G ,B G 1 . Finally, it is easy to verify that f ∗ ( x ) = XOR ( x A ) , thus OPT d, 3( k +1) ,B F ,S ∞ p,γ ∞ ,B G ,B G 1 = 0 .

Theorem E.40 (Uniform Parity Functions: Main Result (Alternative)) . For D parity -uniform setting, for any δ ∈ (0 , 1) satisfying δ ≤ O ( 1 d 2 ) and for any ϵ ∈ (0 , 1) when

trained by Algorithm 1 with hinge loss, with probability at least 1 -δ over the initialization, with proper hyper-parameters, there exists t ∈ [ T ] such that

̸

Proof of Theorem E.40. Plug the values of parameters into Theorem E.37 and directly get the result.

Lemma E.39 (Uniform Parity Functions: Existence of Good Networks (Alternative)) . Assume the same condition as in Lemma E.38. Define

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 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).

Data Distributions. We draw input from the distribution D X = N (0 , I d × d ) , and we assume the target function is g ∗ ( x ) : R d - → R , where g ∗ is a degree τ polynomial normalized so that E x ∼D X [ g ∗ ( x ) 2 ] = 1 .

Assumption E.41. There exists linearly independent vectors u 1 , . . . , u r such that g ∗ ( x ) = g ( ⟨ x , u 1 ⟩ , . . . , ⟨ x , u r ⟩ ) . H := E x ∼D X [ ∇ 2 g ∗ ( x )] has rank r , where H is a Hessian matrix.

Definition E.42. Denote the normalized condition number of H by

Initialization and Loss. For ∀ i ∈ [ m ] , we use the following initialization:

For this regression problem, we use mean square loss:

Training Process. We use the following one-step training algorithm for this specific data distribution.

Algorithm 5 Network Training via Gradient Descent [32]. Special case of Algorithm 2

end for

Lemma E.43 (Multiple Index Model with Low Degree Polynomial: Existence of Good Networks. Rephrase of Lemma 25 in [32]) . Assume n ≥ d 2 rκ 2 ( C l log( nmd )) τ +1 , d ≥ C d κr 3 / 2 , and m ≥ r τ κ 2 τ ( C l log( nmd )) 6 τ +1 for sufficiently large constants C d , C l , and let η (1) = √ d ( C l log( nmd )) 3 and λ (1) = 1 η (1) . Then with probability 1 -1 poly ( m,d ) , there exists ˜ a ∈ R m such that f (˜ a , W (1) , b ) satisfies

and

E.7.2 Multiple Index Model: Final Guarantee

Considering training by Algorithm 5, we have the following results.

Theorem E.44 (Multiple Index Model with Low Degree Polynomial: Main Result) . Assume n ≥ Ω ( d 2 rκ 2 ( C l log( nmd )) τ +1 + m ) , d ≥ C d κr 3 / 2 , and m ≥ Ω ( 1 ϵ r τ κ 2 τ ( C l log( nmd )) 6 τ +1 ) for sufficiently large constants C d , C l . Let η (1) = √ d ( C l log( nmd )) 3 and λ (1) = 1 η (1) , and η = η ( t ) = Θ( m -1 ) , for all t ∈ { 2 , 3 , . . . , T } . For any ϵ ∈ (0 , 1) , if T ≥ Ω ( m 2 ϵ ) , then with properly set parameters and Algorithm 5, with high probability that there exists t ∈ [ T ] such that

Proof of Theorem E.44. By Lemma E.43, we have for properly chosen hyper-parameters,

We compute the L -smooth constant of ˜ L Z ( f ( a , W (1) , b ) ) to a .

By the proof of Lemma 25 in [32], we have for ∀ i ∈ [4 m ] , with probability at least 1 -1 poly ( m,d ) , |⟨ w i , x ⟩| ≤ 1 , with some large polynomial poly ( m,d ) . As a result, we have

This means that we can let η = Θ ( m -1 ) and we will get our convergence result. We can bound ∥ a (1) ∥ 2 and ∥ ˜ a ∥ 2 by ∥ a (1) ∥ 2 = O ( √ m ) and ∥ ˜ a ∥ 2 = O ( r τ κ 2 τ ( C l log( nmd )) 6 τ m ) = O ( ϵ ) . So, if we choose T ≥ Ω ( m ϵη ) , there exists t ∈ [ T ] such that ˜ L Z ( f ( a ( t ) , W (1) , b ) ) -˜ L Z ( f (˜ a , W (1) , b ) ) ≤ O ( L ∥ a (1) -˜ a ∥ 2 2 T ) ≤ ϵ/ 3 .

We also have √ ∥ ˜ a ∥ 2 2 ( ∥ W (1) ∥ 2 F B 2 x + ∥ b ∥ 2 2 ) n ≤ ϵ 3 . Then our theorem gets proved by Theorem 3.4.

Thus, we have,

Discussion. We would like to unify [32], whcih 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.

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 a lowdegree 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.

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 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 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) .

Proof. These can be verified by direct calculation.

Lemma F.2 (Property of Symmetric Initialization) . For any x ∈ R d , we have f Ξ (0) ( x ) = 0 . For all i ∈ [2 m ] , we have w (1) i = -w (1) i +2 m . When input data is symmetric, i.e, E ( x ,y ) [ y x ] = 0 , for all i ∈ [ m ] , we have w (1) i = w (1) i + m .

Proof of Lemma F.2. By symmetric initialization, we have f Ξ (0) ( x ) = 0 . For all i ∈ [2 m ] , we have

When E ( x ,y ) [ y x ] = 0 , for all i ∈ [ m ] , we have

̸

Lemma F.3 (Property of Direction Neighborhood) . If w ∈ C D,γ , we have ρ w ∈ C D,γ for any ρ = 0 . We also have 0 / ∈ C D,γ . Also, if ( D,s ) ∈ S p,γ,B G , we have ( -D,s ) ∈ S p,γ,B G .

Proof. These can be verified by direct calculation.

Lemma F.4 (Maximum Gaussian Tail Bound) . M n is the maximum of n i.i.d. standard normal Gaussian. Then

Proof. These can be verified by direct calculation.

Lemma F.5 (Chi-squared Tail Bound) . If X is a χ 2 ( k ) random variable. Then, ∀ z ∈ R , we have

Proof. These can be verified by direct calculation.

Lemma F.6 (Gaussian Tail Bound) . If g is standard Gaussian and z > 0 , we have

Proof. These can be verified by direct calculation.

Lemma F.7 (Gaussian Tail Expectation Bound) . If g is standard Gaussian and z ∈ R , we have

Proof of Lemma F.7. For any p ∈ (0 , 1) , we have

where √ 2 erf -1 (2 p -1) is the quantile function of the standard Gaussian. We finish the proof by replacing p to be Pr g ∼N (0 , 1) [ g > z ] .

Lemma F.8. If a function g satisfy h ( n +2) = 2 h ( n +1) -(1 -ρ 2 ) h ( n ) + β for n ∈ N + where ρ, β > 0 , then h ( n ) = -β ρ 2 + c 1 (1 -ρ ) n + c 2 (1 + ρ ) n , where c 1 , c 2 only depends on h (1) and h (2) .

Proof. These can be verified by direct calculation.

Lemma F.9 (Rademacher Complexity Bounds. Rephrase of Lemma 48 in [32]) . For fixed W , b , let F = { f ( a , W , b ) : ∥ a ∥ ≤ B a 2 } . Then,

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nMinor note: the restrictions on $\\tau$ should appear in the statements of Theorems 4.4 and 4.8.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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)\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThis is a theoretical work and there is not negative societal impact. I think the limitations have been discussed adequately (nonetheless see above).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

comments: 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.

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.

- 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]).

[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"

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.

- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-10