A T HEORETICAL A NALYSIS ON F EATURE L EARNING IN N EURAL N ETWORKS : E MERGENCE FROM I NPUTS AND A DVANTAGE OVER F IXED F EATURES
A BSTRACT
An important characteristic of neural networks is their ability to learn representations of the input data with effective features for prediction, which is believed to be a key factor to their superior empirical performance.
To better understand the source and benefit of feature learning in neural networks, we consider learning problems motivated by practical data, where the labels are determined by a set of class relevant patterns and the inputs are generated from these along with some background patterns.
We prove that neural networks trained by gradient descent can succeed on these problems.
The success relies on the emergence and improvement of effective features, which are learned among exponentially many candidates efficiently by exploiting the data (in particular, the structure of the input distribution).
In contrast, no linear models on data-independent features of polynomial sizes can learn to as good errors.
Furthermore, if the specific input structure is removed, then no polynomial algorithm in the Statistical Query model can learn even weakly.
These results provide theoretical evidence showing that feature learning in neural networks depends strongly on the input structure and leads to the superior performance.
Our preliminary experimental results on synthetic and real data also provide positive support.
I NTRODUCTION
Various empirical studies have shown that an important characteristic of neural networks is their feature learning ability, i e., to learn a feature mapping for the inputs which allow accurate prediction (e.g., Zeiler & Fergus (2014);
Girshick et al.
; Zhang et al.
; Manning et al.
This is widely believed to be a key factor to their remarkable success in many applications, in particular, an advantage over traditional machine learning methods.
To understand their success, it is then crucial to understand the source and benefit of feature learning in neural networks.
Empirical observations show that networks can learn neurons that correspond to different semantic patterns in the inputs (e.g., eyes, bird shapes, tires, etc.
in images (Zeiler & Fergus, 2014; Girshick et al., 2014)).
Moreover, recent progress (e.g., Caron et al.
; Chen et al.
(2020b)
; He et al.
; Jing & Tian (2020)) shows that one can even learn a feature mapping using only unlabeled inputs and then learn an accurate predictor (usually a linear function) on it using labeled data.
This further demonstrates the feature learning ability of neural networks and that these input distributions contain important information for learning useful features.
These empirical observations strongly suggest that the structure of the input distribution is crucial for feature learning and feature learning is crucial for the strong performance.
However, it is largely unclear how practical training methods (gradient descent or its variants) learn important patterns from the inputs and whether this is necessary for obtaining the superior performance, since the empirical studies do not exclude the possibility that some other training methods can achieve similar performance without feature learning or with feature learning that does not exploit the input structure.
Rigorous theoretical investigations are thus needed for answering these fundamental questions: How can effective features emerge from inputs in the training dynamics of gradient descent?
Is learning features from inputs necessary for the superior performance?
Compared to the abundant empirical evidence, the theoretical understanding still remains largely open.
One line of work (e.g. Jacot et al.
(2018); Li & Liang (2018);
Du et al.
Allen-Zhu et al.
; Zou et al. Chizat et al.
(2019) and many others) shows in certain regime, sufficiently overparameterized networks are approximately linear models, i.e., a linear function on the Neural
Tangent Kernel (NTK).
This falls into the traditional approach of linear models on fixed features, which also includes random features (Rahimi & Recht, 2008) and other kernel methods (Kamath et al., 2020).
The kernel viewpoint thus does not explain feature learning in networks nor the advantage over fixed features.
A recent line of work (e.g. Daniely & Malach (2020); Bai & Lee (2019); Ghorbani et al.
; Yehudai & Shamir (2019); Allen-Zhu & Li (2019; 2020a); Li et al.
(2020); Malach et al.
(2021) and others) shows examples where networks provably enjoy advantages over fixed features, under different settings and assumptions.
While providing insightful results separating the two approaches, most studies have not investigated if the input structure is crucial for feature learning and thus the advantage.
Also, most studies have not analyzed how gradient descent can learn important input patterns as effective features, or rely on strong assumptions like models or data atypical in practice (e.g., special networks, Gaussian data, etc).
Towards a more thorough understanding, we propose to analyze learning problems motivated by practical data, where the labels are determined by a set of class relevant patterns and the inputs are generated from these along with some background patterns.
We use comparison for our study: (1) by comparing network learning approaches with fixed feature approaches on these problems, we analyze the emergence of effective features and demonstrate feature learning leads to the advantage over fixed features; (2) by comparing these problems to those with the input structure removed, we demonstrate that the input structure is crucial for feature learning and prediction performance.
More precisely, we obtain the following results.
We first prove that two-layer networks trained by gradient descent can efficiently learn to small errors on these problems, and then prove that no linear models on fixed features of polynomial sizes can learn to as good errors.
These two results thus establish the provable advantage of networks and implies that feature learning leads to this advantage.
More importantly, our analysis reveals the dynamics of feature learning: the network first learns a rough approximation of the effective features, then improves them to get a set of good features, and finally learns an accurate classifier on these features.
Notably, the improvement of the effective features in the second stage is needed for obtaining the provable advantage.
The analysis also reveals the emergence and improvement of the effective features are by exploiting the data, and in particular, they rely on the input structure.
To formalize this, we further prove the third result: if the specific input structure is removed and replaced by a uniform distribution, then no polynomial algorithm can even weakly learn in the Statistical Query (SQ) learning model, not to mention the advantage over fixed features.
Since SQ learning includes essentially all known algorithms (in particular, mini-batch stochastic gradient descent used in practice), this implies that feature learning depends strongly on the input structure.
Finally, we perform simulations on synthetic data to verify our results.
We also perform experiments on real data and observe similar phenomena, which show that our analysis provides useful insights for the practical network learning.
Our analysis then provides theoretical support for the following principle: feature learning in neural networks depends strongly on the input structure and leads to the superior performance.
In particular, our results make it explicit that learning features from the input structure is crucial for the superior performance.
This suggests that input-distribution-free analysis (e.g., traditional PAC learning) may not be able to explain the practical success, and advocates an emphasis of the input structure in the analysis.
While these results are for our proposed problem setting and network learning in practice can be more complicated, the insights obtained match existing empirical observations and are supported by our experiments.
The compelling evidence hopefully can attract more attention to further studies on modeling the input structure and analyzing feature learning.
R ELATED W ORK
This section provides an overview while more technical discussions can be found in Appendix A.
Neural Tangent Kernel (NTK) and Linearization of Neural Networks.
One line of work (e.g. Jacot et al.
(2018); Li & Liang (2018);
Matthews et al.
; Lee et al.
; Novak et al.
Yang (2019); Du et al.
Allen-Zhu et al.
; Zou et al. ; Ji & Telgarsky (2019); Cao et al.
Geiger et al.
Chizat et al.
(2019) and more) explains the success of sufficiently over-parameterized neural network by connecting them to linear methods like NTK. Though their approaches are different, they all base on the observation that when the network is sufficiently large, the weights stay close to the initialization during the training, and training is similar to solving a kernel method problem.
This is typically referred to as the NTK regime, or lazy training, or linearization.
However, networks used in practice are usually not large enough to enter this regime, and the weights are frequently observed to traverse away from the initialization.
Furthermore, in this regime, network learning is essentially the traditional approach of linear methods over fixed features, which cannot establish or explain feature learning and the advantage of network learning.
Advantage of Neural Networks over Linear Models on Fixed Features.
Since the superior network learning results via gradient descent are not well explained by the NTK view, a recent line of work has turned to learning settings where neural networks provably have advantage over linear models on fixed features (e.g. Daniely & Malach (2020); Refinetti et al.
; Malach et al.
Dou & Liang (2020); Bai & Lee (2019); Ghorbani et al.
; Allen-Zhu & Li (2019); see the great summary in Malach et al. (2021)).
While formally establishing the advantage, they have not thoroughly answered the two fundamental questions this work focuses on; in particular, most existing work has not studied whether the input structure is a crucial factor for feature learning and thus the advantage, and/or has not considered how the features are learned in more practical training scenarios.
For example, Ghorbani et al.
(2020) show the advantage of networks in approximation power and Dou & Liang (2020) show their statistical advantage, but they do not consider the learning dynamics (i.e., how the training method obtains the good network).
Allen-Zhu & Li (2019) prove the advantage of the networks for PAC learning with labels given by a depth-2 ResNet and Allen-Zhu & Li (2020a) prove for Gaussian inputs with labels given by a multiple-layer network, while neither considers the influence of the input structure on feature learning or the advantage.
Daniely & Malach (2020) prove the advantage of the networks for learning sparse parities on specific input distributions that help gradient descent learn effective features for prediction, and Malach et al.
(2021) consider similar learning problems but with specifically designed differentiable models, while our work analyzes data distributions and models closer to those in practice and also explicitly focuses on whether the input structure is needed for the learning.
There are also other theoretical studies on feature learning in networks (e.g. Yehudai & Ohad (2020);
Zhou et al.
; Diakonikolas et al.
Frei et al.
(2020)), which however do not directly relate feature learning to the input structure or the advantage of network learning.
P ROBLEM S ETUP
To motivate our setup, consider images with various kinds of patterns like lines and rectangles.
Some patterns are relevant for the labels (e.g., rectangles for distinguishing indoor or outdoor images), while the others are not.
If the image contains a sufficient number of the former, then we are confident that the image belongs to a certain class.
Dictionary learning or sparse coding is a classic model of such data (e.g., Olshausen & Field (1997); Vinje & Gallant (2000); Blei et al.
We thus model the patterns as a dictionary, generate a hidden vector indicating the presence of the patterns, and generate the input and label from this vector.
Let X =
R d be the input space, and Y = {± 1 } be the label space.
Suppose M ∈ R d × D is an unknown dictionary with D columns that can be regarded as patterns.
For simplicity, assume M is orthonormal.
Let ˜ ϕ ∈ { 0 , 1 } D be a hidden vector that indicates the presence of each pattern.
Let
A ⊆
[ D ] be a subset of size k corresponding to the class relevant patterns.
Then the input is generated by M ˜ ϕ , and the label can be any binary function on the number of class relevant patterns.
More precisely, let P ⊆
[ k ] .
Given A and P , we first sample ˜ ϕ
from a distribution D ˜ ϕ , and then generate the input ˜ x and the class label y from ˜ ϕ :
[Equation]
Learning with Input Structure.
We allow quite general D ˜ ϕ with the following assumptions:
(A0)
The class probabilities are balanced:
Pr[ (cid:80) i ∈
A ˜ ϕ i ∈ P ] = 1 / 2 .
(A1)
The patterns in A are correlated with the labels with the same correlation: for any i ∈ A , γ = E [ y ˜ ϕ i ] − E
[ y ] E [ ˜ ϕ i ] > 0 .
(A2) Each pattern outside A is identically distributed and independent of all other patterns.
Let p o :=
Pr[˜ ϕ i = 1] and without loss of generality assume p o ≤ 1 / 2 .
Let D ( A, P, D ˜ ϕ ) denote the distribution on (˜ x, y ) for some A, P, and D ˜ ϕ .
Given parameters Ξ =
( d, D, k, γ, p o ) , the family F Ξ of distributions include all D ( A, P, D ˜ ϕ ) with A ⊆
[ D ] , P ⊆
[ k ] , and D ˜ ϕ satisfying the above assumptions.
The labeling function includes some interesting special cases:
Example 1.
Suppose P = { i ∈ [ k ] : i > k/ 2 } for some threshold, i.e., we will set the label y
= + when more than a half of the relevant patterns are presented in the input.
Example 2.
Suppose k is odd, and let P
= { i ∈ [ k ] : i is odd } , i.e., the labels are given by the parity function on ˜ ϕ j ( j ∈ A ) .
This is useful to prove our lower bounds via the properties of parities.
Appendix F presents results for more general settings (e.g., incoherent dictionary, unbalanced classes, etc.).
On the other hand, our problem setup does not include some important data models.
In particular, one would like to model hierarchical representations often observed in practical data and believed to be important for deep learning.
We leave such more general cases for future work.
Learning Without Input Structure.
For comparison, we also consider learning problems without input structure.
The data are generated as above but with different distributions D ˜ ϕ
A1’)
The patterns are uniform over { 1 } D for any i ∈ [ D ]
, Pr [˜ ϕ i = 1] = 1 / independently.
Given parameters Ξ 0 = ( d, D, k ) , the family F Ξ 0 of distributions without input structure is the set of all the distributions with A ⊆
[ D ] , P ⊆
[ k ] and D ˜ ϕ satisfying the above assumptions.
3. N EURAL N ETWORK L EARNING
Networks.
We consider training a two-layer network via gradient descent on the data distribution:
[Equation]
where w i ∈ R d , b i , a i ∈ R , and σ ( z ) = min(1 , max( z, 0)) is the truncated rectified linear unit (ReLU) activation function.
Let θ = { w i , b i , a i } 2 mi =1 denote all the parameters, and let superscript ( t ) denote the time step, e.g., g ( t ) denote the network at time step t with θ ( t )
=
{ w ( t ) i , b ( t ) i , a ( t ) i } .
Loss Function.
Similar to typical practice, we will normalize the data for learning: first compute x = (˜ x − E [˜ x ]) /
˜ σ where ˜ σ 2 = E (cid:80) di
=1 (˜ x i − E
[˜ x i ]) 2 is the variance of the data, and then train on ( x, y ) .
This is equivalent to setting ϕ =
(˜ ϕ − E [˜ ϕ ]) /
˜ σ and generating x = Mϕ .
For (˜ x, y ) from D and the normalized ( x, y ) , we will simply say ( x, y ) ∼ D .
For the training, we consider the hinge-loss ℓ ( y, ˆ y )
=
max { 1 − y ˆ y, 0 } .
We will inject some noise ξ to the neurons for the convenience of the analysis.
(This can be viewed as using a smoothed version of the activation ˜ σ ( z ) =
E ξ σ ( z + ξ ) similar to those in existing studies like Allen-Zhu & Li (2020b);
Malach et al.
See Section 5 for more explanations.)
Formally, the loss is:
[Equation] x [Equation] y )
[Equation] x [Equation] ξ ))] , where g ( x ; ξ ) =
[Equation]
where ξ ∼ N (0 , σ 2 ξ I m × m ) are independent Gaussian noise.
Let L D ( g ) denote the typical hinge-loss without noise.
We also consider ℓ 2 regularization: R ( g ; λ a , λ w ) = (cid:80) 2 mi =1 λ a | a i | 2 + λ w ∥ w i ∥ 22 with regularization coefficients λ a , λ w .
Training Process.
We first perform an unbiased initialization: for every i ∈ [ m ] , initialize w (0) i ∼
N (0 , σ 2 w I d × d ) with σ w = 1 /k , b (0) i ∼ N (cid:0) 0 , σ 2 b (cid:1) with σ b = 1 /k 2 , a (0) i ∼ N (0 , σ 2 a ) with σ a = ˜ σ 2 / ( γk 2 ) , and then set w (0) m
+ i = w (0) i , b (0) m
+ i
=
b (0)
i , a (0) m
+ i = − a (0) i .
We then do gradient updates:
[Equation]
, for
[Equation] [Equation]
[Equation] [Equation]
(4) for some choice of the hyperparameters η ( t ) , λ ( t ) a , λ ( t ) w , σ ( t ) ξ , and T .
M AIN R ESULTS
Provable Guarantee for Neural Networks.
The network learning has the following guarantee:
Theorem 1.
For any δ, ϵ ∈ (0 , 1) , if k = Ω (cid:0) log 2 ( D/ ( δγ )) (cid:1) , p o = Ω( k 2 /D ) , and max { Ω( k 12 /ϵ 3 / 2 ) , D } ≤ m ≤ poly ( D ) , then with properly set hyperparameters, for any D ∈ F Ξ , with probability at least 1 − δ , there exists t ∈
[ T ] such that Pr[sign( g ( t ) ( x ))
̸ = y ] ≤ L D ( g ( t ) ) ≤ ϵ.
The theorem shows that for a wide range of the background pattern probability p o and the number of class relevant patterns k , the network trained by gradient descent can obtain a small classification error.
More importantly, the analysis shows the success comes from feature learning.
In the early stages, the network learns and improves the neuron weights such that on the features (i.e., the neurons’ outputs) there is an accurate classifier; afterwards it learns such a classifier. The next section will provide a detailed discussion on the feature learning.
Lower Bound for Fixed Features.
Empirical observations and Theorem 1 do not exclude the possibility that some methods without feature learning can achieve similar performance.
We thus prove a lower bound for the fixed feature approach, i.e., linear models on data-independent features.
Theorem 2.
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 } .
If 3 < k ≤ D/ 16 and k is odd, then there exists
D ∈ F Ξ such that all h ∈ H B have hinge-loss at least p o (cid:16) 1 − √ 2 NB 2 k (cid:17) .
So using fixed features independent of the data cannot get loss nontrivially smaller than p o unless with exponentially large models.
In contrast, viewing the neurons σ ( ⟨ w i , x ⟩ + b i ) as learned features, network learning can achieve any loss ϵ ∈ (0 , 1) with models of polynomial sizes. We emphasize the lower bound is because the feature map Ψ is independent of the data.
Indeed, there exists a small linear model on a small dimensional feature map allowing 0 loss for each data distribution in our problem set F Ξ (Lemma 5). However, this feature map Ψ ∗ is different for different data distribution in F Ξ , i.e., depends on the data.
On the other hand, the feature map Ψ in the lower bound is data-independent, i.e., fixed before seeing the data. For Ψ to work simultaneously for all distributions in F Ξ , it needs to have exponential dimensions.
Intuitively, it needs a large number of features, so that there are some features to approximate each Ψ ∗ i . There are exponentially many data distributions in F Ξ , and thus exponentially many data-dependent features Ψ ∗ i , which requires Ψ to have an exponentially large dimension. Network learning updates the hidden neurons using the data and can learn to move the features to the right positions to approximate the ground-truth data-dependent features Ψ ∗ , so it does not need an exponentially large dimension feature map.
The theorem directly applies to linear models on fixed finite-dimensional feature maps, e.g., linear models on the input or random feature approaches (Rahimi & Recht, 2008).
It also implies lower bounds to infinite dimensional feature maps (e.g., some kernels) that can be approximated by feature maps of polynomial dimensions.
For example, Claim 1 in Rahimi & Recht (2008) implies that a function f using shift-invariant kernels (e.g., RBF) can be approximated by a model ⟨ Ψ(˜ x ) , w ⟩ with the dimension N and weight norm B bounded by polynomials of the related parameters of f like its RKHS norm and the input dimension.
Then our theorem implies some related parameter of f needs to be exponential in k for f to get nontrivial loss, formalized in Corollary 3.
Kamath et al.
(2020) has more discussions on approximating kernels with finite dimensional maps.
Corollary 3.
For any function f using a shift-invariant kernel K with RKHS norm bounded by L , or f ( x ) = (cid:80)
i
α i K
( z i , x ) for some data points z i and || α || 2 ≤ L . If 3 < k ≤ D/ 16 and k is odd, then there exists D ∈ F Ξ such that f has hinge-loss at least p o (cid:16) 1 − poly ( d,L ) 2 k (cid:17) − 1 poly ( d,L ) . Lower Bound for Without Input Structure.
Existing results do not exclude the possibility that some learning methods without exploiting the input structure can achieve strong performance.
To show the necessity of the input structure, we consider learning F Ξ 0 with input structure removed.
We obtain a lower bound for such learning problems in the classic Statistical Query (SQ) model (Kearns, 1998).
In this model, the algorithm can only receive information about the data through statistical queries.
A statistical query is specified by some polynomially-computable property predicate Q of labeled instances and a tolerance parameter τ ∈ [0 , 1] .
For a query ( Q, τ ) , the algorithm receives a response ˆ P Q ∈
[ P Q − τ, P Q + τ ] , where P Q =
Pr[ Q ( x, y ) is true ] .
Notice that a query can be simulated using the average of roughly O (1 /τ 2 ) random data samples with high probability.
The SQ model captures almost all common learning algorithms (except Gaussian elimination) including the commonly used mini-batch SGD, and thus is suitable for our purpose.
Theorem 4.
For any algorithm in the Statistical Query model that can learn over F Ξ 0 to classification error less than 12 − 1 ( Dk ) 3 , either the number of queries or 1 /τ must be at least 12 (cid:0) Dk (cid:1) 1 / 3 .
The theorem shows that without the input structure, polynomial algorithms in the SQ model cannot get a classification error nontrivially smaller than random guessing.
The comparison to the result for with input structure then shows that the input structure is crucial for network learning, in particular, for achieving the advantage over fixed feature models.
P ROOF S KETCHES
Here we provide the sketch of our analysis, focusing on the key intuition and discussing some interesting implications.
The complete proofs are included in Appendix B-D.
5. P ROVABLE G UARANTEES OF N EURAL N ETWORKS
Overall Intuition.
We first show that there is a two-layer network that can represent the target labeling function, whose neurons can be viewed as the “ground-truth” features to be learned.
We then show that after the first gradient step, the hidden neurons of the trained network become close to the ground-truth: their weights contain large components along the class relevant patterns but small along the background patterns.
We further show that in the second gradient step, these features get improved: the “signal-noise” ratio between the components for class relevant patterns and those for the background ones becomes larger, giving a set of good features.
Finally, we show that the remaining steps learn an accurate classifier on these features.
Existence of A Good Network.
We show that there is a two-layer network that can fit the labels.
Lemma 5.
For any D ∈ F Ξ , there exists a network g ∗ ( x )
= (cid:80) ni
=1 a ∗ i σ ( ⟨ w ∗ i , x ⟩ + b ∗ i ) with y =
g ∗ ( x ) for any ( x, y ) ∼ D .
Furthermore, the number of neurons n = 3( k + 1) , | a ∗ i | ≤ 32 k, 1 / (32 k ) ≤ | b ∗ i | ≤ 1 / 2 , w ∗ i = ˜
σ (cid:80)
j ∈
A M j / (4 k ) , and |⟨ w ∗ i , x ⟩ + b ∗ i | ≤ 1 for any i ∈ [ n ] and ( x, y ) ∼ D .
In particular, the weights of the neurons are proportional to (cid:80)
j ∈
A M j , the sum of the class relevant patterns.
We thus focus on analyzing how the network learns such neuron weights.
Feature Emergence in the First Gradient Step.
The gradient for w i (ignoring the noise) is:
[Equation]
where the last step is due to g ( x ) = 0 by the unbiased initialization.
Let q j = ⟨ M j , w i ⟩ denote the component along the direction of the pattern M j .
Then the component of the gradient on M j is:
[Equation]
The key intuition is that with the randomness of ϕ ℓ (and potentially that of the injected noise ξ ), the random variable under σ ′ is not significantly affected by a small subset of ϕ ℓ q ℓ .
For example, for class relevant patterns j ∈ A , let I [ D ] := σ ′ (cid:104)(cid:80) ℓ ∈
[ D ] ϕ ℓ q ℓ + b i (cid:105) and I − A :
= σ ′ (cid:104)(cid:80) ℓ ̸∈
A ϕ ℓ q ℓ + b i (cid:105) .
We have I
[ D ] ≈ I − A and thus:
[Equation]
since y only depends on ϕ j ( j ∈ A ) .
Then the gradient has a nontrivial component along the pattern.
Similarly, for background patterns j ̸∈ A , the component of the gradient along M j is close to 0 .
Lemma 6 (Informal) .
Assume p o , k as in Theorem 1 and σ (1) ξ < 1 /k , then with high probability ∂∂w i L D ( g (0) ; σ (1) ξ )
= − a (0) i (cid:80)
Dj =1 M j T j where for a small ϵ e :
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ O ( ϵ e / ˜ σ ) with β ∈ [Ω(1) , 1] ; • if j ̸∈ A , then | T j | ≤ O ( σ 2 ϕ ϵ e ˜ σ ) .
By setting λ (1) w = 1 / (2 η (1) ) , we have w (1) i =
[Equation] i [Equation] Dj
[Equation]
M [Equation]
For small p o , e.g., p o =
˜ O
( k 2 /D ) , these neurons can already allow accurate prediction.
However, for such small p o , we cannot show a provable advantage of networks over fixed features.
On the other hand, for larger p o meaning a significant number of background patterns in the input, the approximation error terms T j ( j ̸∈ A ) together can overwhelm the signals T j ( j ∈ A ) and lead to bad prediction, even though each term is small.
Fortunately, we will show that the second gradient step can improve the weights by decreasing the ratio between T j ( j ̸∈ A ) and T j ( j ∈ A ) .
Feature Improvement in the Second Gradient Step.
We note that by setting a small η (1) , after the update we still have yg ( x ; ξ ) < 1 for most ( x, y ) ∼ D and thus the gradient in the second step is:
[Equation]
We can then follow the intuition for the first step again.
For j ∈ A , the component ⟨ M j , [Equation] L
[Equation] g ) ⟩ is roughly proportional to γ ˜ σ E [ I − A,ξ ] where I − A,ξ := σ ′ (cid:104)(cid:80) ℓ ̸∈
A ϕ ℓ q ℓ + b i + ξ i (cid:105) .
While ϕ ℓ q ℓ may not have large enough variance, the injected noise ξ i makes sure that a nontrivial amount of data activate the neuron.
Then I − A,ξ ̸ = 0 , leading to a nontrivial component along M j , similar to the first step.
On the other hand, for j ̸∈ A , the approximation error term T j depends on how well σ ′ (cid:104)(cid:80) ℓ ̸∈ A,ℓ ̸ =
j ϕ ℓ q ℓ + b i + ξ i (cid:105) approximates σ ′ (cid:104)(cid:80) ℓ ∈
[ D ] ϕ ℓ q ℓ + b i + ξ i (cid:105) .
Since the q ℓ ’s (the weight’s component along M ℓ ) in the second step are small compared to those in the first step, we can then get a small error term T j .
So the ratio between T j ( j ̸∈ A ) over T j ( j ∈ A ) improves after the second step, giving better features allowing accurate prediction.
Classifier Learning Stage.
Given the learned features, we are then ready to show the remaining gradient steps can learn accurate classifiers.
Intuitively, with small hyperparameter values ( η ( t ) =
k
Tm 1 / 3 , λ ( t ) a = λ ( t ) w ≤
[Equation]
0 for 2 < t ≤ T = m 4 / 3 ), the first layer’s weights do not change too much and thus the learning is similar to convex learning using the learned features.
Formally, our proof uses the online convex optimization technique in Daniely & Malach (2020).
L OWER B OUNDS
The lower bounds are based on the following observation: our problem setup is general enough to include learning sparse parity functions.
Consider an odd k , and let P
= { i ∈ [ k ] : i is odd } .
Then y is given by Π A ( z ) :
=
(cid:81)
j ∈
A z j for z j = 2˜ ϕ j − 1 , i.e., the parity function on z j ( j ∈ A ) .
Then known results for learning parity functions can be applied to prove our lower bounds.
Lower Bound for Fixed Features.
We show that F Ξ contains learning problems that consist of a mixture of two distributions with weights p o and 1 − p o respectively, where in the first distribution
D (1) A , ˜ x is given by the uniform distribution over ˜ ϕ and the label y is given by the parity function on A .
On such D (1) A , Daniely & Malach (2020) shows that exponentially large models over fixed features is needed to get nontrivial loss.
Intuitively, there are exponentially many labeling functions Π A that are uncorrelated (i.e., “orthogonal” to each other): E [Π A 1 Π A 2 ] = 0 for any A 1 and A 2 .
Note that the best approximation of Π A by a fixed set of features Ψ i ’s is its projection on the linear span of the features.
Then with polynomial-size models, there always exists some Π A far from the linear span.
Remark.
It is instructive to compare to network learning, which finds the effective weights (cid:80)
j ∈
A M j among the exponentially many candidates corresponding to different A ’s. This can be done efficiently by exploiting the data since the gradient is roughly proportional to E { yx }
= (cid:80)
j ∈
A M j .
The network then learns data-dependent features on which polynomial size linear models can achieve small loss.
Lower Bound for Learning without Input Structure.
Clearly, F Ξ 0 contains the distributions D (1) A described above.
The lower bound then follows from classic SQ learning results (Blum et al., 1994).
Remark.
The SQ lower bound analysis does not apply to F Ξ , because in F Ξ the input distribution is related the labeling function.
This allows networks to learn with polynomial time/sample.
While both the labeling function and the input distribution affect the learning, few existing studies explicitly point out the importance of the input structure.
We thus emphasize the input structure is crucial for networks to learn effective features and achieve superior performance.
E XPERIMENTS
Our experiments mainly focus on feature learning and the effect of the input structure.
We first perform simulations on our learning problems to (1) verify our main theorems on the benefit of feature learning and the effect of input structure; (2) verify our analysis of feature learning in networks.
We then check if our insights carry over to real data: (3) whether similar feature learning is presented in real network/data; (4) whether damaging the input structure lowers the performance.
The results are consistent with our analysis and provide positive support for the theory.
Below we present part of the results and include the complete experimental details and results in Appendix E.
Simulation: Verification of the Main Results.
We generate data according to our problem setup, with d = 500 , D = 100 , k = 5 , p o = 1 / 2 , a randomly sampled A , and labels given by the parity function.
We then train a two-layer network with m =
following our learning process, and for comparison, we also use two fixed feature methods (the NTK and random feature methods based on the same network).
Finally, we also use these three methods on the data distribution with the input structure removed (i.e.,
F Ξ 0 in Theorem 4).
Figure 1 shows that the results are consistent with our results.
Network learning gets high test accuracy while the two fixed feature methods get significantly lower accuracy.
Furthermore, when the input structure is removed, all three methods get test accuracy similar to random guessing.
Simulation: Feature Learning in Networks.
We compute the cosine similarities between the weights w i ’s and visualize them by Multidimensional Scaling.
(Recall that our analysis is on the directions of the weights without considering their scaling , and thus it is important to choose cosine similarity rather than say the typical Euclidean distance.)
Figure 2 shows that the results are as predicted by our analysis.
After the first gradient step, some weights begin to cluster around the ground-truth (cid:80)
j ∈
A M j (or − (cid:80)
j ∈
A M j due to the a i in the gradient update which can be positive or negative).
After the second step, the weights get improved and well-aligned with the ground-truth (with cosine similarities > 0 . 99 ).
Furthermore, if a classifier is trained on the features after the first step, the test accuracy is about 52% ; if the same is done after the second step, the test accuracy is about 100% .
This demonstrates while some effective features emerge in the first step, they need to be improved in the second step to get accurate prediction.
Real Data: Feature Learning in Networks.
We perform experiments on MNIST (LeCun et al., 1998; Deng, 2012), CIFAR10 (Krizhevsky, 2012), and SVHN (Netzer et al., 2011).
On MNIST, we train a two-layer network with m = 50 on the subset with labels 0/1 and visualize the neurons’ weights as in the simulation.
Figure 3 shows a similar feature learning phenomenon: effective features emerge after a few steps and then get improved to form two clusters.
Similar results are observed on other datasets.
These suggest the insights obtained in our analysis are also applicable to the real data.
[Figure]
Real Data: The Effect of Input Structure.
Since we cannot directly manipulate the input distribution of real data, we perform controlled experiments by injecting different inputs.
For labeled dataset
L and injected input U , we first train a teacher network fitting L , then use the teacher network to give labels on a mixture of inputs from L and U , and finally train a student network on this new dataset M consisting of the mixed inputs and the teacher network’s labels.
Checking the student’ performance on different parts of M and comparing to those by directly training the student on the original data L can reveal the impact of changing the input structure.
We use MNIST, CIFAR10, or
SVHN as L , and use Gaussian or images in Tiny ImageNet (Le & Yang, 2015) as U .
The networks for MNIST are two-layer with m = 9 , and those for CIFAR10/SVHN are ResNet-18 convolutional neural networks (He et al., 2016).
Figure 4 shows the results on an equal mixture of data and Gaussian.
It presents the test accuracy of the student on the original data part, the Gaussian part, and the whole mixture.
For example, on
CIFAR10, the network learns well over the CIFAR10 part (with accuracy similar to directly training on the original data) but learns slower with worse accuracy on the Gaussian part.
Furthermore, the accuracy on the whole mixture is lower than that of training on the original CIFAR10.
This shows that the input structure indeed has a significant impact on the learning.
While MNIST+Gaussian shows a less significant trend (possibly because the tasks are simpler), the other datasets show similar significant trends as CIFAR10+Gaussian (the results using Tiny ImageNet are in the appendix).
E THICS S TATEMENT
Our paper is mostly theoretical in nature and thus we foresee no immediate negative ethical impact.
We are of the opinion that our theoretical framework may lead to better understanding and inspire development of improved network learning methods, which may have a positive impact in practice.
In addition to the theoretical machine learning community, we perceive that our conceptual message that the input structure is crucial for the network learning’s performance can be beneficial to engineeringinclined machine learning researchers.
R EPRODUCIBILITY S TATEMENT
For theoretical results in the Section 4, a complete proof is provided in the Appendix B-D.
The theoretical results and complete proofs for a setting more general than that in the main text are provided in the Appendix F. For experiments in the Section 6, complete details and experimental results are provided in the Appendix Section
E. The source code with explanations and comments is provided in the supplementary material.  A CKNOWLEDGEMENT
The work is partially supported by Air Force Grant FA9550-18-1-0166, the National Science Foundation (NSF) Grants 2008559-IIS and CCF-2046710.
Appendix
Section A presents more technical discussion on related work.
Section B-D provides the complete proofs for our results in the main text.
Section E provides the complete details and experimental results for our experiments.
Finally, Section F provides the theoretical results and complete proofs for a setting more general than that in the main text, allowing incoherent dictionaries, unbalanced classes, and Gaussian noise in the data. C ONTENTS 1 Introduction 2 Related Work 3 Problem Setup 3.
Neural Network Learning . .
4 Main Results 5 Proof Sketches 5.
Provable Guarantees of Neural Networks . . . . . . . . . . . . . . . . . . . . . . .
Lower Bounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6 Experiments 7 Ethics Statement 8 Reproducibility Statement 9 Acknowledgement A More Technical Discussion on Related Work B Complete Proofs for Provable Guarantees of Neural Networks
B.1 Existence of A Good Network . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2 Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
{+B.
+}
Some Auxiliary Lemmas . .
B.4 Feature Emergence: First Gradient Step . . . . . . . . . . . . . . . . . . . . . . .
B.5 Feature Improvement: Second Gradient Step . . . . . . . . . . . . . . . . . . . . .
B.
.7 Proof
C Lower Bound for Linear Models on Fixed Feature Mappings
D Lower Bound for Learning without Input Structure
E Complete Experimental Results
E.1 Simulation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.1.
Parity Labeling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.1.
Interval Labeling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.2 More Simulation Result in Various Settings . .
E.2.
Varying Input Data Dimension . .
E.2.
Varying Class Imbalance Ratio . . .
E.2.
Varying Sample Size . .
E.3 Experiments on More Data Generation Models . . . . . . . . . . . . . . . . . . .
E.3.
Hidden Representation Labeling . . . . . . . . . . . . . . . . . . . . . . .
E.3.
Two-layer Networks on Mixture of Gaussians . . . . . . . . . . . . . . . .
E.4 Real Data: Feature Learning in Networks . . . . . . . . . . . . . . . . . . . . . .
E.4.
CNNs on Binary Cifar10: Feature Learning in Networks . . . . . . . . . .
{+E.
+}
Real Data: The Effect of Input Structure . . . . . . . . . . . . . . . . . . . . . . .
E.5.
Experimental Methodology . .
E.5.
Experimental Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.5.
Larger Network on MNIST for Checking The Effect of Input Structure . .
E.5.
Empirical Verification of Our Method . . . . . . . . . . . . . . . . . . . .
F Provable Guarantees for Neural Networks in A More General Setting
F.
Problem Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.1.
Neural Network Learning . .
F.
Main Result . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Existence of A Good Network . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Some Auxiliary Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Feature Emergence: First Gradient Step . . . . . . . . . . . . . . . . . . . . . . .
F.
Feature Improvement: Second Gradient Step . . . . . . . . . . . . . . . . . . . . .
F.
Classifier Learning Stage and Main Theorem . .
A M ORE T ECHNICAL D ISCUSSION ON R ELATED W ORK
Advantage of Neural Networks over Linear Models on Fixed Features.
A recent line of work has turned to show learning settings where network learning provably has advantage over linear models on fixed features; see the nice summary in Malach et al.
Here we highlight the results and focuses of the existing related studies and discuss the differences from ours.
Yehudai & Shamir (2019) shows that the random feature method fails to learn even a single ReLU neuron on Gaussian inputs unless its size is exponentially large in dimension.
This points out the limitation of the random feature method (belonging to the fixed feature approach) but does not consider feature learning in networks.
Some studies show that a single ReLU neuron can be learnt by gradient descent (Yehudai & Ohad, 2020; Diakonikolas et al., 2020;
Frei et al., 2020).
The analysis typically involves feature learning.
However, their focus is different: they do not show the advantage over fixed feature methods and do not consider the effect of the input structures.
Zhou et al.
(2021) shows that in a special teacher-student setting, the student network will do exact local convergence in a surprising way that all student neurons will converge to one of the teacher neurons.
The work does not consider the effect of the input structure nor the advantage over fixed features.
Dou & Liang (2020) explains the advantage of network learning by constructing adaptive Reproducing Kernel Hilbert Space (RKHS) indexed by the training process of the neural network, and shows that adaptive RKHS benefits from a smaller function space containing the residue comparing to RKHS.
The work shows the statistical advantage of networks over data-independent kernels, but does not consider the optimization for learning the network.
Ghorbani et al.
(2020) considers data generated from a hidden vector with two subsets of variables, each uniformly distributed in a high-dimensional sphere (with a different radius), while the label is determined by only the first subset of variables.
It shows the existence of good neural networks that can overcome the curse of dimensionality by representing the best low-dimensional hidden structure.
However, it studies the approximation power of neural networks rather than the learning, i.e., it does not show how to learn the good network.
Fang et al.
(2019) argues that in the infinite width limit, a two-layer neural network will learn a nearly optimal feature representation in the distribution sense, thanks to the convexity of the limit problem.
It is unclear how this result helps to understand the feature learning procedure for practical networks, which is usually a non-convex process.
Chen et al.
(2020a) considers a fixed, randomly initialized neural network as a representation function fed into another trainable network which is the quadratic Taylor model of a wide two-layer network.
It shows that learning over the random representation can achieve improved sample complexities compared to learning over the raw data.
However, the representation considered is not learned, which is different from our focus on feature learning.
Allen-Zhu & Li (2020a) considers Gaussian inputs with labels given by a multiple-layer network with quadratic activations and skip connections (with the assumption of information gap on the weights), and studies training a deep network with quadratic activation.
It shows that the trained network can learn proper representations and obtain small errors while no polynomial fixed feature methods can.
On the other hand, it does not focus on the influence of input structure on feature learning: note that its input distribution contains no information about the “ground-truth” features in the target network.
It also points out that the learned features get improved during training: higher-level layers will help lower-level layers to improve by backpropagating correction signals.
Our analysis also shows feature improvement which however is by signals from the input distribution.
Allen-Zhu & Li (2019) considers PAC learning with labels given by a depth-2 ResNet, and studies training an overparameterized depth-2 ResNet (using uniform inputs over Boolean hypercube as an example).
It shows the trained network can obtain small errors while no polynomial kernel methods can obtain as good errors.
Similar to Allen-Zhu & Li (2020a), it does not focus on the influence of input structure on feature learning or the advantage of networks.
Allen-Zhu & Li (2020c) studies how ensemble of deep learning models can improve test accuracy and how the ensemble can be distilled into a single model.
It develops a theory which assumes the data has multi-view structure and shows that the ensemble of independently trained networks can provably improve test accuracy and the ensemble can also be provably distilled into a single model.
The analysis also relies on showing that the data structure can help the ensemble and the distillation.
On the other hand, their focus is on ensembles and is quite different from ours: the analysis is on showing the multi-view input structure allows the ensembles of networks to improve over single ones and ensembles of fixed feature mappings do not have improvement.
While our focus is on supervisedly learning one single network that outperforms the fixed feature approaches.
Daniely & Malach (2020) considers the task of learning sparse parities with two-layer networks, and the analysis suggests that the ability to learn the label-correlated features also seems to be critical towards the success of neural networks, although the authors did not explore much in this direction.
Malach et al.
(2021) also considers similar learning problems but with specifically designed models for the problems.
The learning problems considered in Daniely & Malach (2020); Malach et al.
(2021) have input distributions that leak information about the target labeling function, which is similar to our setting, and their analysis also shows that the first gradient descent can learn a set of good features and later steps can learn an accurate classifier on top.
Our work is inspired by their studies, while there are some important differences.
First, their focuses are different from ours.
Daniely & Malach (2020) focuses on showing neural networks can learn targets (i.e., k -parity functions) that are inherently non-linear.
Our analysis generalizes to more general distributions, including practically motivated ones.
Malach et al.
(2021) focuses on strong separations between learning with gradient descent on differentiable models (including typical neural networks) and learning using the corresponding tangent kernels.
The analysis is on specific differentiable models, while our work is on two-layer neural networks similar to practical ones.
Second, our analysis relies on the feature improvement in the second gradient step.
This is not an artifact of the analysis but comes from our problem setup.
While in Daniely & Malach (2020) the data distribution allows some neurons to be sufficiently good after the first gradient step and needs no feature improvement, our setup is more general where the data distribution may not have a similar strong benign effect and thus needs feature improvement in the second gradient step. Most related to our work is Daniely & Malach (2020).
Therefore, we provide a detailed discussion to highlight the connections and differences.
Our problem setting is more general than that in Daniely & Malach (2020).
To see this, let our dictionary be the identity matrix, the set P to be the odd numbers (i.e., the labeling function is a sparse parity).
{+Furthermore, let the distribution of the hidden representation be an equal mixture of the following two: (a) D 1 : Uniform distribution over the hypercube. (b) D 2 : Irrelevant patterns ˜ ϕ j ( j ̸∈ A ) have appearance probability p 
+}
= 1 / 2 .
And the distribution of relevant patterns ˜ ϕ
j ( j ∈ A ) is: all 0’s with probability 1/2, and all 1’s with probability 1/2.
Then our problem setting reduces to their setting (up to scaling/translation of ˜ ϕ j ’s).
On the other hand, in general our setting allows for more choices for the labeling, the dictionary, and the distributions over ˜ ϕ .
Upper bound: Because of the more general setting, our upper bound proof requires technical novelty .
Recall that in their work, the input distribution is essentially a mixture of D 1 and
D 2 above.
In D 2 , the relevant patterns ˜ ϕ
j ( j ∈ A ) have the specific structure of all 0’s or all 1’s with probability 1/2.
This allows to show that neurons with weight w satisfying (cid:80) j ∈
A w j = 0 will have good gradients: small components from irrelevant patterns (their
Lemma 7) and large components from relevant patterns (their Lemma 8).
However, in our setting, the relevant patterns do not have this specific structure, and thus their proof technique is not applicable (or can be applied only when we have an exponentially large number of hidden neurons so that some hit the good positions at random initialization).
What we showed is that the gradient has some correlation with the good feature direction.
{+So after the first gradient step, the neuron weights are not good yet but are in a better position for further improvement (in particular, their setting corresponds to p 
+}
= 1 / 2 which means large noise in the weights after the first step
; see discussion after our Lemma 6 in Section 5).
Then
the latter gradient steps are able to improve the weights to better “signal-to-noise-ratio”.
In summary, our proof does not rely on their specific input structure or an exponentially large number of hidden neurons for hitting some good positions.
The key is that the good feature will emerge with the help of the input structure, and once in a better position, the neurons’ weights can be improved to the desired quality. 3.
Lower bound:
On the other hand, our lower bound is proved by a reduction to the lower bound results in Daniely & Malach (2020).
They have shown that D 1 above can lead to large errors for fixed feature models of polynomial size.
Our proof is essentially constructing a mixture of D 1 and D 2 with mixture weights p 0 and ( 1 − p 0 ) , and applying their lower bound for D 1 .
See our proof in Appendix C. 4.
Conceptually, our work belongs to the same line of research as Daniely & Malach (2020), to analyze how feature learning leads to the superior performance of networks.
While their analysis also relies on feature learning from good gradients induced by input structure, their focus is more on separating network learning and fixed feature models and has not explicitly explored the impact of input structures (while we agree that such an explicit study will not be difficult in their setting).
More importantly, their input distribution is specific and atypical in practice, which allows a specific type of feature learning (as explained in the above discussion on upper bounds).
Our work thus considers a more general setting that is motivated by practical problems.
Our results then bring theoretical insights closer for explaining the feature learning in practice and provide some positive evidence for the importance of analysis under proper models of the input distributions.
Sparse Coding and Subspace Data Models.
To analyze neural networks’ performance, various data models have been considered.
A practical way to model the underlying structure of data is by assuming that a set of hidden variables exists and the input data is a high dimensional projection of the hidden vector (possibly with noise).
Along this line, the classic sparse coding model has been used in existing works for analyzing networks.
Koehler & Risteski (2018) considers such a data distribution where the label is given by a linear function on the hidden sparse vector, but studies the approximation power of networks and classic polynomial methods rather than the learning.
Allen-Zhu & Li (2020b) considers similar data distributions, but studies the performance of networks under adversarial perturbations.
Another type of related data models assumes that the label is determined by a subset of hidden variables.
Ghorbani et al.
(2020) considers a hidden vector with two subsets of variables, each uniformly distributed in a high-dimensional sphere (with a different radius), while the label is determined by only the first subset of variables.
However, Ghorbani et al.
(2020) studies the approximation power of neural networks rather than the learning.
Compared to these studies, our work assumes the input is given by a dictionary multiplied with a hidden vector (not necessarily sparse) while the label is determined by a subset of the hidden vector, as motivated by pattern recognition applications in practice.
Furthermore, we focus on the learning ability of networks instead of approximation.
B C OMPLETE P ROOFS FOR P ROVABLE G UARANTEES OF N EURAL N ETWORKS
We first make a few remarks about the proof.
Remark.
The analysis can be carried out for more gradient steps following similar intuition, while we analyze two steps for simplicity.
Remark.
Readers may notice that the network can be overparameterized.
With sufficient overparameterization and proper initialization and step sizes, network learning becomes approximately
NTK.
However, here our learning scheme allows going beyond this kernel regime: we use aggressive gradient updates λ ( t )
w = 1 / (2 η ( t ) )
in the first two steps, completely forgetting the old weights to learn effective features.
Using proper initialization and aggressive updates early to escape the kernel regime has been studied in existing work (e.g., Woodworth et al.
; Li et al.
Our result thus adds another concrete example.
Notations.
For a vector v and an index set I , let v
I denote the vector containing the entries of v indexed by I , and v − I denote the vector containing the entries of v with indices outside I .
By initialization, w (0) i for i ∈ [ m ] are i.i.d.
copies of the same random variable w (0) ∼ N (0 , σ 2 w I d × d ) ; similar for a (0) and b (0) .
Let q ℓ := ⟨ w (0) , M ℓ ⟩ , then ⟨ w (0) , x ⟩ = ⟨ ϕ, q ⟩ .
Similarly, define q ( t ) i,ℓ := ⟨ w ( t ) i , M ℓ ⟩ .
Let σ 2 ϕ :
= p o (1 − p o ) / ˜ σ 2 denote the variance of ϕ ℓ for ℓ ̸∈ A .
We also define the following sets to denote typical initialization.
For a fixed δ ∈ (0 , 1) , define G
[Equation]
B. E XISTENCE OF A G OOD N ETWORK
we first show that there exists a network that can fit the data distribution.
Lemma 7.
For some s, a, b ∈ R with a, b ≥ 0 , define a function δ s,a,b : R → R as δ s,a,b
[Equation]
[Equation]
(8) where σ r ( z ) =
max { z, 0 } is the ReLU activation function.
Then δ s,a,b ( z ) =
  when z ≤ s − b, a ( z − s )
+ ab when s − b ≤ z ≤ s, − a ( z − s ) + ab when s ≤ z ≤ s + b, when s + b ≤ z.
That is, δ s,a,b ( z ) linearly interpolates between ( s − b, 0) , ( s, ab ) , ( s + b, 0) when z ∈ [ s − b, s + b ] , and is 0 elsewhere.
Proof of Lemma 7.
This can be simply verified for the four cases of the value of z .
Lemma 8 (Restatement of Lemma 5) .
For any D ∈ F Ξ , there exists a network g ∗ ( x )
= (cid:80) ni
=1 a ∗ i σ ( ⟨ w ∗ i , x ⟩ + b ∗ i ) with y =
g ∗ ( x ) for any ( x, y ) ∼ D .
Furthermore, the number of neurons n = 3( k + 1) , | a ∗ i | ≤ 32 k, 1 / (32 k ) ≤ | b ∗ i | ≤
Proof of Lemma 5.
Let w = ˜
σ (cid:80)
j ∈
A M j
and let µ = (cid:80)
j ∈
A E [˜ ϕ j ] .
We have
[Equation]
Then by Lemma 7,
[Equation]
for any ( x, y ) ∼ D .
Similarly,
[Equation]
= y (18) for any ( x, y ) ∼ D .
Note that the bias terms in g ∗ 1 and g ∗ 2 have distance at least 1 / 4 , then at least one of them satisfies that all its bias terms have absolute value ≥ 1 / 8 .
Pick that one and denote it as g ( x )
= (cid:80)
[Equation] =1 a i σ r ( ⟨ w i , x ⟩ + b i ) .
By the positive homogeneity of σ r , we have
[Equation]
Since for any ( x, y ) ∼ D , |⟨ w i , x ⟩ / (4 k ) + b i / (4 k ) | ≤ 1 , then g ( x ) = n (cid:88) i =
[Equation] [Equation]
b
[Equation]
/ (4 k )) (20) where σ is the truncated ReLU.
Now we can set a ∗ i
= 4 ka i , w ∗ i =
w i / (4 k ) , b ∗ i
= b i / (4 k ) , to get our final g ∗ .
B. I NITIALIZATION
We first show that with high probability, the initial weights are in typical positions.
Lemma 9.
For any δ ∈ (0 , 1) , with probability at least 1 − δ − 2 exp ( − Θ( D − k )) over w (0) , σ [Equation] w ( D [Equation]
k ) /
[Equation]
≤ (cid:88)
[Equation] q [Equation] With probability at least 1 − δ over b (0) ,
[Equation]
With probability at least 1 − δ over a (0) , [Equation]
Proof of Lemma 9.
From q ∼ N (0 , σ 2 w I d × d ) , we have:
• With probability ≥ 1 − δ/ 2 , max ℓ | q ℓ | ≤ (cid:113) 2 σ 2 w log Dδ , and • For any subset S ⊆
[ D ] , with probability ≥ 1 − 2 exp ( − Θ( | S | )) , ∥ q S ∥ 22 ∈ (cid:16) | S | σ 2 w 2 , 3 | S | σ 2 w
(cid:17)
Similar for b (0) and a (0) .
The lemma then follows.
Lemma 10.
We have:
• With probability ≥ 1 − δ − 2 m exp( − Θ( D − k )) over w (0)
i ’s, for all i ∈ [2 m ] , w (0) i ∈
G w ( δ ) .
• With probability ≥ 1 − δ over b (0)
i ’s, for all i ∈ [2 m ] , b (0) i ∈
G b ( δ ) .
• With probability ≥ 1 − δ over a (0) i ’s, for all i ∈ [2 m ] , a (0) i ∈ G a ( δ ) .
Proof of Lemma 10.
This follows from Lemma 9 by union bound.
The following lemma about the typical w (0) i ’s will be useful for later analysis.
Lemma 11.
Fix δ ∈ (0 , 1) .
For any w (0) i ∈
G w ( δ ) , we have
[Equation]
Consequently, when p o = Ω( k 2 /D ) and k = Ω(log 2 ( Dm/δ )) ,
[Equation]
Proof of Lemma 11.
Note that for ℓ ̸∈ A , E [ ϕ ℓ ] = 0 , E [ ϕ 2 ℓ ] =
σ 2 ϕ , and E [ | ϕ ℓ | 3 ] =
Θ( σ 2 ϕ / ˜ σ ) .
Then the statement follows from Berry-Esseen Theorem.
B. S OME A UXILIARY L EMMAS
The expression of the gradients will be used frequently.
Lemma 12.
[Equation]
Proof of Lemma 12.
It follows from straightforward calculation.
We now show that a small subset of the entries in ϕ, q does not affect the probability distribution of ⟨ ϕ, q ⟩ much [Equation]
Lemma 13.
Suppose ν ∼ N (0 , σ 2 ) .
For any B ⊇ A and any b :
[Equation]
Similarly,
[Equation]
Proof of Lemma 13.
Note that for ℓ ̸∈ A , E [ ϕ ℓ ] = 0 , E [ ϕ 2 ℓ ] =
σ 2 ϕ , and E [ | ϕ ℓ | 3 ] =
Θ( σ 2 ϕ / ˜ σ ) .
Let t = |⟨ ϕ B , q B ⟩| .
Then by the Berry-Esseen Theorem,
[Equation]
The second statement follows from a similar argument.
We also have the following auxiliary lemma for later calculations.
Lemma 14.
[Equation]
Proof of Lemma 14.
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
B. F EATURE E MERGENCE : F IRST G RADIENT S TEP
We will show that w.h.p.
over the initialization, after the first gradient step, there are neurons that represent good features.
We begin with analyzing the gradients.
Lemma 15 (Full version of Lemma 6) .
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
G w ( δ ) , b (0) i ∈ G b ( δ ) for all i ∈ [2 m ] .
Let
[Equation]
If p o = Ω( k 2 /D ) , k = Ω(log 2 ( Dm/δ )) , and σ (1) ξ < 1 /k , then
[Equation]
where T j satisfies:
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ O ( ϵ e / ˜ σ ) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if j ̸∈ A , then | T j | ≤ O ( σ 2 ϕ ϵ e ˜ σ ) .
Proof of Lemma 15.
Consider one neuron index i and omit the subscript i in the parameters.
Since the unbiased initialization leads to g (0) ( x ; ξ (1) )
= 0 , we have
[Equation]
[Equation]
First, consider j ∈ A
[Equation]
[Equation]
Let
[Equation] [Equation]
We have
[Equation] [Equation]
Then by Lemma 13,
[Equation]
Note that (cid:80) ℓ ̸∈ A Var ( ϕ ℓ q ℓ ) =
Θ( σ 2 ϕ σ 2 w ( D − k ))
=
Θ( σ 2 w ) , and | ϕ ℓ | ≤ 1˜ σ , max ℓ | q ℓ | ≤ σ w (cid:112) 2 log( Dm/δ ) .
Applying Bernstein’s inequality for bounded distributions, we have:
[Equation]
We also have:
[Equation]
[Equation]
Therefore,
[Equation]
[Equation]
where the last step follows from the assumption on σ w and k .
A similar argument gives:
[Equation]
Then we have
[Equation]
where the last step is from Lemma 14.
Furthermore,
[Equation]
By Lemma 11, the assumption on p o , and (63), we have
[Equation]
[Equation]
This leads to
[Equation]
By Lemma 14, E ϕ A { yϕ j } = γ/ ˜ σ .
Therefore,
[Equation]
Now, consider j ̸∈ A .
Let B denote A ∪ { j } .
[Equation]
Let
[Equation]
Similar as above, we have | E ξ (1) ( I b − I ′ b ) | ≤ O ( ϵ e ) by Lemma 13.
Then by Lemma 14,
[Equation]
Furthermore,
[Equation]
Therefore,
[Equation]
Lemma 16.
Under the same assumptions as in Lemma 15,
[Equation]
where | T b | ≤ O ( ϵ e ) .
Proof of Lemma 16.
Consider one neuron index i and omit the subscript i in the parameters.
Since the unbiased initialization leads to g (0) ( x ; ξ (1) )
= 0 , we have
[Equation]
[Equation]
Similar to the proof in Lemma 6,
[Equation]
Then
[Equation] [Equation]
=
[Equation]
(96) ≤ O ( ϵ e )
E ϕ A {| y |}
[Equation]
≤ O ( ϵ e ) .
[Equation]
Also,
[Equation]
Therefore, | T b | ≤ O ( ϵ e ) .
Lemma 17.
We have
[Equation]
where | T a | ≤ O (max ℓ q (0) i,ℓ ) . So if w (0) i
[Equation] , | T [Equation] |
[Equation]
Proof of Lemma 17.
Consider one neuron index i and omit the subscript i in the parameters.
Since the unbiased initialization leads to g (0) ( x ; ξ (1) )
= 0 , we have
[Equation]
Let ϕ ′
A be an independent copy of ϕ A [Equation] ϕ ′ be the vector obtained by replacing in ϕ the entries ϕ
A with ϕ ′
A
[Equation]
and let x ′ =
Mϕ ′ and its label is y ′ .
Then
[Equation]
Since σ is 1-Lipschitz,
[Equation]
With the bounds on the gradient
[Equation]
we now summarize the results for the weights after the first gradient step.
Lemma 18.
Set λ (1) w = 1 / (2 η (1) ) , λ (1) a = λ (1) b = 0 , σ (1) ξ = 1 /k 3 / 2 .
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
G w ( δ ) , b (0) i ∈ G b ( δ ) for all i ∈ [2 m ] .
If p o = Ω( k 2 /D ) , k = Ω(log 2 ( Dm/δ )) , then for all i ∈ [ m ] , w (1) i = (cid:80) Dℓ =1 q (1) i,ℓ M ℓ satisfying
• if ℓ ∈ A , then | q (1) i,ℓ − η (1) a (0) i βγ/ ˜ σ | ≤ O (cid:18) | η (1) a (0) i | ϵ e ˜ σ (cid:19) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if ℓ ̸∈ A , then | q (1) i,ℓ | ≤ O (cid:16) σ 2 ϕ | η (1) a (0) i | ϵ e ˜ σ (cid:17) ;
and
• b (1) i =
b (0) i + η (1) a (0)
i T b where | T b | =
O ( ϵ e ) ; • a (1) i =
a (0) i + η (1) T a where | T a | =
O ( σ w (cid:112) log( Dm/δ )) .
Proof of Lemma 18.
This follows from Lemma 10 and Lemma 15-17.
B. F EATURE I MPROVEMENT : S ECOND G RADIENT S TEP
We first show that with properly set η (1) , for most x , | g (1) ( x ; σ (2) ξ ) | < 1 and thus yg (1) ( x ; σ (2) ξ ) < 1 .
Lemma 19.
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
G w ( δ ) , b (0) i ∈ G b ( δ ) , a (0) i ∈ G a ( δ ) for all i ∈ [2 m ] .
If p o = Ω( k 2 /D ) , k = Ω(log 2 ( Dm/δ )) , σ a ≤ ˜ σ 2 / ( γk 2 ) , η (1) =
O (cid:16) γ kmσ a (cid:17) , and σ (2) ξ ≤ 1 /k , then with probability ≥ 1 − exp( − Θ( k )) over ( x, y ) , we have yg (1) ( x ; σ (2) ξ ) < 1 .
Furthermore, for any i ∈ [2 m ] , (cid:12)(cid:12)(cid:12) ⟨ w (1) i , x ⟩ (cid:12)(cid:12)(cid:12) = (cid:12)(cid:12)(cid:12) ⟨ q (1) i , ϕ ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , (cid:12)(cid:12)(cid:12) ⟨ ( q (1) i ) − A , ϕ − A ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , and (cid:12)(cid:12)(cid:12) b (1) i − b (1) m
+ i (cid:12)(cid:12)(cid:12) =
O ( | η (1) a (0) i | ϵ e ) .
Proof of Lemma 19.
Note that w (0) i = w (0) m
+ i , b (0) i
= b (0) m
+ i , and
a (0) i = − a (0) m
+ i .
Then the gradient for w i is the negation of that for w m + i , the gradient for b
i is the negation of that for b m + i , and the gradient for a
i is the same as that for a m
+ i .
With probability ≥ 1 − exp( − Θ(max { 2 p o ( D − k ) , k } )) , among all j ̸∈ A , we have that at most 2 p o
( D − k )
+ k of ϕ j are (1 − p o ) /
˜ σ , while the others are − p o / ˜ σ .
For data points with ϕ satisfying this, we have:
[Equation] [Equation]
Then we have
[Equation]
(cid:12)(cid:12)(cid:12) [Equation] (cid:12)(cid:12)(cid:12) [Equation]
[Equation]
(cid:12)(cid:12)(cid:12)(cid:17) .
We have | T a |
=
O ( σ w (cid:112) log( Dm/δ )) , and (cid:12)(cid:12)(cid:12) ⟨ w (1) i , x ⟩ (cid:12)(cid:12)(cid:12) ≤ [Equation] [Equation] [Equation]
[Equation]
[Equation]
[Equation] p [Equation] (131) (cid:12)(cid:12)(cid:12)
[Equation]
(1) i (cid:12)(cid:12)(cid:12) ≤ (cid:12)(cid:12)(cid:12) b (0) i (cid:12)(cid:12)(cid:12) + (cid:12)(cid:12)(cid:12) η (1) a (0) i
[Equation] b (cid:12)(cid:12)(cid:12) (132) ≤ [Equation] log
[Equation] m [Equation] k
[Equation]
[Equation]
i ϵ
[Equation]
(cid:12)(cid:12)(cid:12)
[Equation] [Equation]
Then we have
[Equation] [Equation]
[Equation]
Then (cid:12)(cid:12)(cid:12) yg (1) ( x ; σ (2) ξ ) (cid:12)(cid:12)(cid:12) < 1 .
Finally, the statement on (cid:12)(cid:12)(cid:12) ⟨ ( q (1) i ) − A , ϕ − A ⟩ (cid:12)(cid:12)(cid:12) follows from a similar calculation on (cid:12)(cid:12) (cid:12) ⟨ w (1) i , x ⟩ (cid:12)(cid:12) (cid:12) = (cid:12)(cid:12) (cid:12) ⟨ q (1) i , ϕ ⟩ (cid:12)(cid:12) (cid:12) .
We are now ready to analyze the gradients in the second gradient step.
Lemma 20.
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
[Equation] i
[Equation] ) i ∈ G a ( δ ) for all i ∈ [2 m ] .
Let ϵ e2 :
= O (cid:18) η (1) | a (0)
i | k
( γ + ϵ e ) ˜ σ 2 σ (2) ξ (cid:19) + exp( − Θ( k )) .
If k = Ω(log 2 ( Dm/δ )) and k = O ( D ) , σ a ≤
[Equation]
γk
[Equation]
[Equation]
[Equation] (cid:16) [Equation] and
[Equation] [Equation] = 1 /k
[Equation]
2 , then
[Equation]
where T j satisfies:
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ O ( ϵ e2 / ˜ σ + η (1) /σ
(2) ξ
[Equation]
1) | a (
[Equation]
i | ϵ e / (˜ σσ (2) ξ )) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if j ̸∈ A , then | T j | ≤ 1˜ σ exp( − Θ( k ))
+ O ( σ 2 ϕ ϵ e2 ˜ σ ) .
Proof of Lemma 20.
Consider one neuron index i and omit the subscript i in the parameters.
By
Lemma 19, Pr[ yg (1) ( x ; ξ (2) [Equation] > 1] ≤ exp( − Θ( k )) .
Let I x
= I [ yg (1) ( x ; ξ (2) [Equation] ≤ 1]
[Equation]
T
[Equation] [Equation] [Equation]
Let T j 1 :=
E ( x,y ) ∼D ,
[Equation]
) (cid:8) y I [ ⟨ w (1) , x ⟩ + b (1) + ξ (2) ∈ (0 , 1)]
ϕ j (cid:9) .
We have
[Equation]
[Equation]
[Equation]
[Equation]
So it is sufficient to bound T j 1 .
For simplicity, we use q as a shorthand for q (1) i .
First, consider j ∈ A .
[Equation]
Let
[Equation]
By the property of the Gaussian ξ (2) , that |⟨ ϕ A , q A ⟩| =
O ( η (1) | a (0)
i | k
( γ + ϵ e ) ˜ σ) , and that |⟨ ϕ, q ⟩| = |⟨ w (1) i , x ⟩| = O ( η (1) /γ ) <
O (1 /k ) and |⟨ ϕ − A , q − A ⟩| =
O ( η (1) /γ ) <
O (1 /k ) , we have | I a − I ′ a | ≤ (cid:12)(cid:12)(cid:12) (cid:12) Pr ξ (2) (cid:104) ⟨ ϕ, q ⟩ + b (1) + ξ (2) ≥ 0 (cid:105) − Pr ξ (2) (cid:104) ⟨ ϕ − A , q − A ⟩ + b (1) + ξ (2) ≥ 0 (cid:105)(cid:12)(cid:12)(cid:12) (cid:12) (155) + Pr ξ (2) (cid:104) ⟨ ϕ, q ⟩ + b (1) + ξ (2) ≥ 1 (cid:105) + Pr ξ (2) (cid:104) ⟨ ϕ − A , q − A ⟩ + b (1) + ξ (2) ≥ 1 (cid:105) (156) =
O (cid:32) η
(1) | a (0)
i | k
( γ + ϵ e ) ˜ σ 2 σ (2) ξ (cid:33) + exp( − Θ( k ))
=
O ( ϵ e2 ) .
This leads to
[Equation]
where the last step is from Lemma 14.
Furthermore,
[Equation]
By Lemma
[Equation]
we have |⟨ ϕ − A , q − A ⟩| ≤ O ( η (1) ˜ σ/γ ) .
Also, | b (1) − b (0) | ≤ O ( η (1) | a (0) i | ϵ e ) .
By the property of ξ (2) ,
[Equation]
On the other hand,
[Equation]
and β only depends on b (0) .
By Lemma [Equation] E ϕ A { yϕ j } = γ/ ˜ σ .
Therefore,
[Equation]
Now
[Equation]
consider j ̸∈ A .
Let B denote A ∪ { j } .
[Equation]
[Equation]
Let
[Equation] [Equation]
[Equation]
Similar as above, we have | I b −
I ′ b | ≤ ϵ e2 .
Then by Lemma 14,
[Equation] [Equation]
I
[Equation]
[Equation] [Equation]
[Equation]
Furthermore,
[Equation]
Therefore, | T j 1 | ≤ O ( σ 2 ϕ ϵ e2 ˜ σ ) .
Lemma 21.
Under the same assumptions as in Lemma 20,
[Equation]
where | T b | ≤ exp( − Ω( k ))
+ O ( ϵ e2 ) .
Proof of Lemma 21.
Consider one neuron index i and omit the subscript i in the parameters.
By
Lemma 19, Pr[ yg (1) ( x ; ξ (2) ) > 1] ≤ exp( − Ω( k )) .
Let I x
= I [ yg (1) ( x ; ξ (2) ) ≤ 1] . [Equation] [Equation]
[Equation]
Let T b 1 :=
E ( x,y ) ∼D ,ξ (2) (cid:8) y I [ ⟨ w (1) , x ⟩ + b (1) + ξ (2) ∈ (0 , 1)] (cid:9) .
We have
[Equation] [Equation]
[Equation] [Equation]
So it is sufficient to bound T b
For simplicity, we use q as a shorthand for q (1) i .
[Equation]
Let
[Equation]
Similar as in Lemma
[Equation]
we have | I b − I ′ b | ≤ ϵ e2 .
Then by Lemma 14,
[Equation]
[Equation]
[Equation]
Furthermore,
[Equation]
Therefore, | T b 1 | ≤ O ( ϵ e2 ) and the statement follows.
Lemma 22.
Under the same assumptions as in Lemma 20,
[Equation]
(198) where |
Proof of Lemma 22.
Consider one neuron index i and omit the subscript i in the parameters.
By
Lemma 19, Pr[ yg (1) ( x ; ξ (2) ) > 1] ≤ exp( − Ω( k )) .
Let I x
= I [ yg (1) ( x ; ξ (2) ) ≤ 1] .
[Equation]
Let T a 1 :=
E ( x,y ) ∼D (cid:8)
y E ξ (2) σ ( ⟨ w (1) , x ⟩ + b (1) + ξ (2) ) (cid:9) .
We have
[Equation]
[Equation]
So it is sufficient to bound T a 1 .
For simplicity, we use q as a shorthand for q (1) i .
Let ϕ ′
A be an independent copy of ϕ A , ϕ ′ be the vector obtained by replacing in ϕ the entries ϕ
A with ϕ ′ A , and let x ′ =
Mϕ ′ and its label is y ′ .
Then
[Equation]
where the fourth step follows from that σ is 1-Lipschitz, the third to the last step from Lemma 19, and the second to the last step from Lemma 18.
With the above lemmas about the gradients, we are now ready to show that at the end of the second step, we get a good set of features for accurate prediction.
Lemma 23.
Set
[Equation]
Fix δ ∈ (0 , O (1 /k 3 )) .
If p o = Ω( k 2 /D ) , k = Ω (cid:16) log 2 (cid:16) Dmδγ (cid:17)(cid:17) , and m ≥ max { Ω( k 4 ) , D } , then with probability at least 1 − δ over the initialization, there exist ˜ a i ’s such that ˜ g ( x ) :
= (cid:80) 2 mi =1 ˜ a i σ ( ⟨ w (2) i , x ⟩ + b (2) i ) satisfies L D (˜ g )
= 0 .
Furthermore, ∥ ˜ a ∥
= O ( m/k ) , ∥ ˜ a ∥ ∞
=
O ( k 5 /m ) , and ∥ ˜ a [-(cid:107) 
-]∥ 22 =
O ( k 9 /m ) .
Finally, ∥ a (2) ∥ ∞ =
O (cid:0) 1 km 2 (cid:1) , ∥ w (2) i ∥ 2 = O (˜ σ/k ) , and | b (2) i
[-(cid:107) 
 = O (˜ σ/k ) , and | b (2) i-] | =
O (1 /k 2 ) for all i ∈ [2 m ] .
Proof of Lemma 23.
By Lemma 5, there exists a network g ∗ ( x )
= (cid:80) 3( k +1) ℓ = a ∗ ℓ σ ( ⟨ w ∗ ℓ , x ⟩ + b ∗ ℓ ) satisfying g ∗ ( x ) for all ( x, y ) ∼ D .
Furthermore, | a ∗ i | ≤ 32 k, 1 / (32 k ) ≤ | b ∗ i | ≤ 1 / 2 , w ∗ i = ˜
σ (cid:80)
j ∈
A M j / (4 k ) , and |⟨ w ∗ i , x ⟩ + b ∗ i | ≤ 1 for any i ∈ [ n ] and ( x, y ) ∼ D .
Now we fix an ℓ , and show that with high probability there is a neuron in g (2) that can approximate the ℓ -th neuron in g ∗ .
By Lemma 10, with probability 1 − 2 δ over w (0) i ’s, they are all in G w ( δ ) ; with probability 1 − δ over a (0)
i ’s, they are all in G a ( δ ) ; with probability 1 − δ over b (0)
i ’s, they are all in G b ( δ ) .
Under these events, by Lemma 18, Lemma 20 and 21, for any neuron i ∈ [2 m ] , we have
[Equation]
where
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ ϵ w 1 : = O ( ϵ e2 / ˜ σ + η (1) /σ (2) ξ + η (1) | a (0)
i | ϵ e / (˜ σσ (2) ξ )) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if j ̸∈ A , then | T j | ≤ ϵ w 2 :=
1˜ σ exp( − Θ( k ))
+ O ( σ 2 ϕ ϵ e2 ˜ σ ) .
• | T b | ≤ ϵ b :=
1˜ σ exp( − Θ( k ))
+ O ( ϵ e2 ) .
Given the initialization, with probability Ω(1) over b (0) i , we have
[Equation]
b
[Equation]
Finally, since 4 k | b ∗ ℓ | βγ | b (0) i | ˜ σ 2 ∈ [Ω( k 2 γ/ ˜ σ 2 ) , O ( k 3 γ/ ˜ σ 2 )] and depends only on w (0) i , b (0) i , we have that for ϵ a = Θ(1 /k 2 ) , with probability Ω( ϵ a ) > δ over a (0) i ,
[Equation]
Let n a = ϵ a m/ 4 .
For the given value of m , by (219)-(222) we have with probability ≥ 1 − 5 δ over the initialization, for each ℓ there is a different set of neurons I ℓ ⊆
[ m ] with | I ℓ | = n a and such that for each i ℓ ∈ I ℓ ,
[Equation]
We also have
[Equation]
[Table]
[Equation]
For the given values of parameters
[Equation]
we have
[Equation] [Equation]
Therefore,
[Equation]
[Equation]
We also have by Lemma 18 and 21:
[Equation]
Now, construct ˜ a such that ˜ a i ℓ = 2 a
[Equation]
b
[Equation]
| n a for each ℓ and each i ℓ ∈ I ℓ , and ˜ a i = 0 elsewhere.
Then
[Equation]
=
[Equation]
[Equation]
≤ 1 .
[Equation]
Here the second equation follows from that σ is positive-homogeneous in [0 , 1] , |⟨ w ∗ ℓ , x ⟩ + b ∗ ℓ | ≤ 1 , | b (0) i ℓ | / | b ∗ ℓ | ≤ 1 .
This guarantees y ˜ g
( x ) ≥ 1 .
Changing the scaling of δ leads to the statement.
Finally, the bounds on ˜ a follow from the above calculation.
The bound on ∥ a (2) ∥ 2 follows from
Lemma 22, and those on ∥ w (2) i ∥ 2 and ∥ b (2) i ∥ 2 follow from (219)(220) and the bounds on a (1) i and b (1) i in Lemma 18.
B. C LASSIFIER L EARNING S TAGE
Once we have a set of good features, we are now ready to prove that the later steps will learn an accurate classifier.
The intuition is that the first layer’s weights do not change much and the second layer
’s weights get updated till achieving good accuracy.
In particular, we will employ the online optimization technique from Daniely & Malach (2020).
We begin by showing that the first layer’s weights do not change too much.
Lemma 24.
Assume the same conditions as in Lemma 23.
Suppose for t > 2 , λ ( t ) a = λ ( t ) w = λ , η ( t ) =
η for some λ, η ∈ (0 , 1) , and σ ( t ) ξ = 0 .
Then for any t > 2 and i ∈ [2 m ] ,
[Equation]
[Equation]
Proof of Lemma 24.
First, we bound the size of | a ( t ) i | :
[Equation]
which leads to
[Equation]
where | a (2) i | =
O (cid:0) 1 km 2 (cid:1) .
We are now to bound the change of w ( t ) i and b ( t ) i .
[Equation]
[Table]
[Equation]
[Table]
[Equation]
269) leading to
[Equation]
Note that (cid:12)(cid:12)(cid:12) b (2) i (cid:12)(cid:12)(cid:12)
=
O (1 /k 2 ) .
Lemma 25.
Assume the same conditions as in Lemma 24.
Let g ( t ) ˜ a ( x ) = (cid:80) mi
=1 ˜ a i σ ( ⟨ w ( t ) i , x ⟩ + b ( t ) i ) .
Then
[Equation]
Proof of Lemma 25.
It follows from that
[Equation]
and Lemma 24.
B. P ROOF OF T HEOREM
Based on the above lemmas
[Equation]
following the same argument as in the proof of Theorem 2 in Daniely &
Malach (2020) [Equation] we get our main theorem
[Equation]
Theorem 26 (Full version of Theorem 1) .
Set
[Equation]
[Equation]
, for 2 < t
[Equation]
T. [Equation] For any δ ∈ (0 , 1) , if p o = Ω( k 2 /D ) , k = Ω (cid:16) log 2 (cid:16) Dδγ (cid:17)(cid:17) , max { Ω( k 4 ) , D } ≤ m ≤ poly ( D ) , then we have for any D ∈ F Ξ , with probability at least 1 − δ , there exists t ∈
[ T ] such that
[Equation]
Consequently, for any ϵ ∈ (0 , 1) , if T = m 4 / 3 , and max { Ω( k 12 /ϵ 3 / 2 ) , D } ≤ m ≤ poly ( D ) , then
[Equation]
Proof of Theorem 1.
Consider ˜ L D ( g ( t ) )
=
E [ ℓ ( g ( t ) , y )] + λ ( t ) a ∥ a ( t ) ∥ 22 .
Note that the gradient update using ˜ L D ( g ( t ) ) is the same as the update in our learning algorithm.
Then by Theorem 27, Lemma 23, and Lemma 25,
[Equation]
The statement follows from that 0-1 classification error is bounded by the hinge-loss.
Theorem 27 (Theorem 13 in Daniely & Malach (2020)) .
Fix some η , and let f 1 , . .
, f T be some sequence of convex functions.
Fix some θ 1 , and assume we update θ t {++
+}
=
θ t − η ∇ f t ( θ t ) .
Then for every θ ∗ the following holds:
[Equation]
C L OWER B OUND FOR L INEAR M ODELS ON F IXED F EATURE M APPINGS
Theorem 28 (Restatement of Theorem 2) .
Suppose Ψ is a data-independent feature mapping of dimension N with bounded features, i.e., Ψ : X → [ − 1 , 1] N .
Define for B > 0 :
[Equation]
Then, if 3 < k ≤ D/ 16 and k is odd, then there exists D ∈ F Ξ such that all h ∈ H B have hinge-loss at least p o (cid:16) 1 − √ 2 NB 2 k (cid:17) .
Proof of Theorem 2.
We first show that F Ξ contains some distributions that are essentially sparse parity learning problems, and then we invoke the lower bound result from existing work for such problems.
Consider D defined as follows.
• Let P =
{ i ∈ [ k ] : i is odd } .
That is, if there are odd numbers of 1 ’s in ˜ ϕ A , then y
= +1 .
Let D (0)˜ ϕ be a distribution where all entries ˜
ϕ j are i i
d. with Pr[˜ ϕ j = 0] =
Pr[˜ ϕ j = 1] = 1 / 2 .
Let D (0) be the distribution over (˜ x, y ) induced by D (0)˜ ϕ and the above P .
Let D (1)˜ ϕ be a distribution where all entries ˜ ϕ
j for j ̸∈ A are i i
d. with Pr[˜ ϕ j = 1] =
p o / (2 − 2 p o ) , while Pr[˜ ϕ A =
=
Pr[˜ ϕ
A = (1 , 1 , . . .
= 1 / 2 .
Let D (1) be the distribution over (˜ x, y ) induced by D (1)˜ ϕ and the above P .
Let D mix A = p o D (0) + (1 − p
o ) D (1) .
It can be verified that such distributions are included in F Ξ for γ = Θ(1) .
Assume for contradiction that for all D ∈ F Ξ , there exists h ∗ ∈ H B such that h = ⟨ Ψ , w ∗ ⟩ loss smaller than p o (cid:16) 1 − √ 2 NB 2 k (cid:17) .
Then for all the distributions D mix A defined above, we have
[Equation]
Now let D z be a distribution over z ∈ {− 1 , +1 } D with i.i.d.
entries z j and Pr[ z j = − 1] =
Pr[ z j = +1] = 1 / 2 .
Let f A ( z )
=
(cid:81)
j ∈
A z j be the k -sparse parity functions.
Let Ψ ′
( z ) =
Ψ( M ( z + 1) / 2) .
Then we have h ′ ( z ) = ⟨ Ψ ′ ( z ) , w ∗ ⟩ such that for all A ,
[Equation]
This is contradictory to Theorem 29.
The following theorem is implicit in the proof in Theorem 1 in Daniely & Malach (2020).
Theorem 29.
For a subset A ⊆
[ D ] of size k , let
the distribution D A over ( z, y ) defined as follows:
z is uniform over {± 1 } D and y = (cid:81) i ∈
A z i .
Fix some Ψ : {± 1 } D → [ − 1 , +1] N , and define:
[Equation]
If k is odd and k ≤ D/ 16 , then there exists some A such that
[Equation]
We now prove the corollary.
Corollary 30 (Restatement of Corollary 3) .
For any function f using a shift-invariant kernel K with RKHS norm bounded by L , or f ( x ) = (cid:80)
i
α i K
( z i , x ) for some data points z i and ||
α || 2 ≤ L .
If 3 < k ≤ D/ 16 and k is odd, then there exists D ∈ F Ξ such that f have hinge-loss at least p o (1 − poly ( d,L ) 2 k ) − 1 poly ( d,L ) .
Proof.
By Claim 1 in Rahimi & Recht (2008), for any ν > 0 , there exists N = poly ( d, 1 /ν )
Fourier features Ψ j that can approximate the shift-invariant kernel up to error ν .
For any ϵ > 0 , consider (cid:80) i
α i ⟨ Ψ( z i ) , Φ( x ) ⟩ = ⟨ (cid:80)
i α i Ψ( z i ) , Ψ( x ) ⟩ .
Let w
= (cid:80) i
α i Ψ
( z i ) and let ν
= O ( ϵL ) , then ⟨ Ψ( x ) , w ⟩ approximates f ( x ) upto error ϵ and N = poly ( d, L, 1 /ϵ ) and the norm of w bounded by B = poly ( d, L, 1 /ϵ ) .
The reasoning is the same for f in the RKHS form, replacing sum with integral.
By Theorem 2, ⟨ Ψ( x ) , w ⟩ has hinge-loss at least p 0 (1 − √ 2 NB 2 k ) .
Thus, the function f has loss at least p 0 (1 − poly ( d,L, 1 /ϵ ) 2 k ) − ϵ .
Choose ϵ
= 1 poly ( d,L )
, we get the bound.
D L OWER B OUND FOR L EARNING WITHOUT I NPUT S TRUCTURE
First recall the Statistical Query model (Kearns, 1998).
In this model, the learning algorithm can only receive information about the data through statistical queries.
A statistical query is specified by some property predicate Q of labeled instances, and a tolerance parameter τ ∈ [0 , 1] .
When the algorithm asks a statistical query ( Q, τ ) , it receives a response ˆ P Q ∈
[ P Q − τ, P Q + τ ] , where
P Q =
Pr[ Q ( x, y ) is true ] .
Q is also required to be polynomially computable, i.e., for any ( x, y )
Q ( x, y ) can be computed in polynomial time.
Notice that a statistical query can be simulated by empirical average of a large random sample of data of size roughly O (1 /τ 2 ) to assure the tolerance τ with high probability.
Blum et al.
(1994) introduces the notion of Statistical Query dimension, which is convenient for our purpose.
Definition 31 (Definition 2 in Blum et al.
For concept class C and distribution D , the statistical query dimension SQ-DIM ( C, D ) is the largest number d such that C contains d concepts c 1 , . .
, c d that are nearly pairwise uncorrelated: specifically, for all i ̸
= j ,
[Equation]
[Equation]
Theorem 32 (Theorem 12 in Blum et al.
In order to learn C to error less than 1 / 2 − 1 /din the Statistical Query model, where d = SQ-DIM ( C, D ) , either the number of queries or 1 /τ must be at least 12 d 1 / 3 .
We now use the above tools to prove our lower bound.
Theorem 33 (Restatement of Theorem 4) .
For any algorithm in the Statistical Query model that can learn over F Ξ 0 to error less {+than 1 2 − 
+}
( Dk ) 3 , either the number of queries or 1 /τ must be at least 12 (cid:0) Dk (cid:1) 1 / 3 .
Proof of Theorem 4.
Consider the following concept class and marginal distribution:
Let D be the distribution over ˜ x , given by ˜ x
= M ˜ ϕ and ˜ ϕ j are i.i.d.
with Pr[˜ ϕ j = 0] =
Pr[˜ ϕ j = 1] = 1 / 2 .
• Let C be the class of functions y =
g A (˜ ϕ ) =
I [ (cid:80) j (1 − ˜ ϕ j ) is odd ] for different A ⊆
[ D ] .
The distributions over (˜ x, y ) induced by ( C, D ) are a subset of F Ξ 0 .
It is then sufficient to show that
SQ-DIM ( C, D ) ≥ (cid:0)
Dk (cid:1) .
It is easy to see that C are essentially the sparse parity functions:
if z j = 2˜ ϕ j − 1 , then g A (˜ ϕ ) = (cid:81) j ∈
A z j .
This then implies that the g A ’s are uncorrelated, so SQ-DIM ( C, D ) ≥ (cid:0) Dk (cid:1) .
E C OMPLETE E XPERIMENTAL R ESULTS
Our experiments mainly focus on feature learning and the effect of the input structure.
We first perform simulations on our learning problems to (1) verify our main theorems on the benefit of feature learning and the effect of input structure (2) verify our analysis of feature learning in networks.
We then check if our insights carry over to real data: (3) whether similar feature learning is presented in real network/data; (4) whether damaging the input structure lowers the performance.
The results are consistent with our analysis and provide positive support for the theory.
The experiments were ran 5 times with different random seeds, and the average results (accuracy) are reported.
The standard deviations of the results are smaller than 0.5% and thus we do not present them for clarity.
The hardware specifications are 4 Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz, 16 GB RAM, and one NVIDIA GPU GTX1080.
E. S IMULATION
We train a two-layer network following our learning process.
We use two fixed feature methods: the
NTK (Fang et al., 2021) and random feature (RF) methods based on the same network and random initialization as the network learning.
More precisely, in the NTK method, we randomly initialize the network and take its NTK and learn a classifier on it.
In the RF method, we freeze the first layer of the network, and train the second layer (on the random features given by the frozen neurons).
The training step number is the same as that in network learning.
We also test these three methods on the data distribution with input structure removed (i.e., F Ξ 0 in Theorem 4).
For comparison, we take the representation of our two-layer network at step one/step two, named One Step/Two Step (fix the weight of the 1st layer after the first step/second step to train the weight of the second layer), and train the best classifiers on top of them.
Recall that our analysis is on the directions of the weights without considering their scaling , and thus it is important to choose cosine similarity rather than the typical ℓ 2 distance.
Thus, we use metric Cos
Similarity max { i ∈ [2 m ] } cos( w i , (cid:80)
j ∈
A M j ) in our tables, and use Multidimensional Scaling to plot the weights distribution.
The simulation dataset size is 50000.
During training, the batch size is 1000, while for the first two steps we use the approximate full gradient (batch size is 50000).
Each step is corresponding to one weights update.
E.1. P ARITY L ABELING
Setting.
We generate data according to the parity function data distributions used in our proof of the lower bound for fixed features (Theorem 2), with d = 500 , D = 100 , k = 5 , p o = 1 / 2 , with a randomly sampled A .
More precisely, we consider D defined as follows.
• Let P =
{ i ∈ [ k ] : i is odd } .
That is, if there are odd numbers of 1 ’s in ˜ ϕ A , then y
= +1 .
Let D (0)˜ ϕ be a distribution where all entries ˜ ϕ
j are i.i.d.
with Pr[˜ ϕ j = 0] =
Pr[˜ ϕ j = 1] = 1 / 2 .
Let D (0) be the distribution over (˜ x, y ) induced by D (0)˜ ϕ and the above P .
Let D (1)˜ ϕ be a distribution where all entries ˜ ϕ
j for j ̸∈ A are i.i.d.
with Pr[˜ ϕ j = 1] =
p o / (2 − 2 p o ) , while Pr[˜ ϕ A =
=
Pr[˜ ϕ
A = (1 , 1 , . . .
= 1 / 2 .
Let D (1) be the distribution over (˜ x, y ) induced by D (1)˜ ϕ and the above P .
Let D mix A = p o D (0) + (1 − p
o ) D (1) .
The network and the training follow Section 3, where the network size is m =
300 and the training time T = 600 steps.
Verification of the Main Results.
Figure 5 shows that the results are consistent with our analysis.
Network learning gets high test accuracy while the two fixed feature methods get significantly lower accuracy.
Furthermore, when
[Table]
Feature Learning in Networks.
Figure 6 shows that the results are as predicted by our analysis.
After the first gradient step, some weights begin to cluster around the ground-truth (cid:80)
j ∈
A M j (or − (cid:80)
j ∈
A M j due to we have a i in the gradient update which can be positive or negative).
After the second step the weights get improved and well-aligned with the ground-truth (with cosine similarity > 0 . 99 ).
Table 1 shows the results for different methods.
Recall that the Cos Similarity metric is max { i ∈ [2 m ] } cos( w i , (cid:80)
j ∈
A M j ) , which reports the cosine value of the closest one.
One Step refers to the method where we take the neurons after one gradient step, freeze their weights, and train a classifier on top; similar for Two Step.
One Step gets test accuracy about 52% , while Two Step gets accuracy about 100% .
This demonstrates that while some effective feature emerge in the first step, they need to be improved in the second step for accurate prediction.
NTK, random feature, One Step all failed, while Network and Two Step can achieve 100% test accuracy.
Network w/o structure refers to training the network on data without the input structure.
It overfits the training dataset with 52% test accuracy.
[Figure]
[Table]
Setting.
We also tried interval function, where y = 1 if (cid:80) i ∈
A ˜ ϕ i is in the range [ t 1 , t 2 ] with t 1 =and t 2 = 30 , otherwise y = − 1 .
We use d = 500 , D = 100 , k = 30 .
The ˜ ϕ i ’s are independent
[Table]
and
Pr[˜ ϕ i = 1] = 2 / 3 for any i ∈
A
[Table]
and Pr[˜ ϕ i = 1] = 1 / 2 otherwise.
When the input structure is removed, we set Pr[˜ ϕ i = 1] = 1 / 2 for all i ’s.
The network and training again follows Section 3 with a network size m = 100 and the training time
T = 200 steps.
Verification of the Main Results.
Figure 7 shows that network learning learns the fastest, NTK learns slower but reaches similar test accuracy, while random feature can only reach a decent but lower accuracy.
This is because for such simpler labeling functions, fixed feature methods can still achieve good performance (note that the lower bound does not hold for such a case), while the performance depends on what fixed features to use.
Furthermore, when the input structure is removed, the methods still get similar (or only slightly worse) performance as with input structure.
This shows that when the labeling function is simple, the help of the input structure for learning may not be needed.
In the experiments on real data, we will show that when the input structure is changed, it indeed leads to lower performance which suggests that the labeling function in practice is typically more complicated than this interval labeling setting, and the help of the input structure is significant for learning.
Feature Learning in Networks.
Figure 8 shows the phenomenon of feature learning similar to that in the parity labeling setting.
Table 2 shows the test accuracy of six different methods.
Random feature and One Step failed, while Network, NTK and Two Step succeed showing that interval labeling setting is a simpler case than parity labeling setting.
E. M ORE S IMULATION R ESULT IN V ARIOUS S ETTINGS
We show the robustness of our simulation results by studying the learning behaviors in a variety of settings including different sample size, input data dimension and class imbalance.
We reuse the same setting as the simulation in the main text (details in E.1.1), vary different parameters, and report the accuracy, the cosine similarities between the learned weights, and the visualization of the neuron weights.
E.2. V ARYING I NPUT D ATA D IMENSION
In the simulation experiments in the main text, the input data dimension d is 500.
Here we change the input data dimension to 100 and 2000.
All other configurations follow E.1.1.
Verification of the Main Results.
Figure 9 shows that our claim is robust under different input data dimensions.
The performance of network learning is superior over NTK and random feature approaches on inputs with structure, and on inputs without structure, all three methods fail.
[Figure]
d =
Network
NTK
RF
One Step
Two Step
Network w/o structure
Train Acc 100.
Test Acc 100.
Cos Similarity 1.
NA 0.
d =
[Table]
Feature Learning in Networks.
Figure 10 visualizes the neuron weights.
It shows similar results to that in E.1.1: the weights gets updated to to the effective feature in the first two steps, forming clusters.
Table 3 shows some quantitative results.
In particular, the average cosine similarities between neuron weights and the effective features after two steps are close to 1, showing that they match the effective features.
E.2. V ARYING C LASS I MBALANCE R ATIO
The experiments in the main text has 25000 training samples for each class.
Here we keep the total sample size 50000 but use different class imbalance ratios, which is the class − 1 sample size divide by the total sample size.
Verification of the Main Results.
Figure 11 shows that our claim is robust under different class imbalance ratios.
The results are similar to those for balanced classes, except that NTK becomes less stable.
[Figure]
Feature Learning in Networks.
Figure 12 visualizes the neurons’ weights.
Again, the observation is similar to that for balanced classes.
Table 4 shows some quantitative results which are also similar to those for balanced classes.
[Table]
E.2. V ARYING S AMPLE S IZE
Here we change the sample size 50000 in Section E.1.1 to be 25000 and 10000.
For sample size 25000, we observe similar results.
For sample size 10000, we observe over-fitting (test accuracy much lower than train accuracy).
Therefore, for sample size 10000 we reduces the size of the network (i
[Table]
e., number of hidden neurons) from m = 300 to m = 50 .
Verification of the Main Results.
Figure 13 shows that our claim is robust under different sample sizes.
In particular, the network learning still outperforms the NTK and random feature approaches on structured inputs.
n =
Network
[Table]
Train
[Table]
Test
[Table]
Cos Similarity
[Table]
Feature Learning in Networks.
Figure 14 and Table 5 show that the phenomenon of feature learning for different samples is similar to that in E.1.1.
[Figure]
E. E XPERIMENTS ON M ORE D ATA G ENERATION M ODELS
In this section we consider some additional data distributions and run the simulation experiments, in particular, focusing on the feature learning phenomenon.
Note that our analysis is for the setting where the input distributions have structure revealing some information about the labeling function.
(More precisely, the labeling function is specified by A and P , while the input distribution also depends on them.)
We therefore consider two other data generation mechanisms where the labeling function also has connections to the input distributions.
E.3. H IDDEN R EPRESENTATION L ABELING
Here we consider the following data model: first uniformly at random select ˜ ϕ
A from a set of binary vectors, and assign label 1 to some and -1 to others; sample irrelevant patterns ˜ ϕ
− A uniformly at random; generate the input x = M ˜ ϕ .
We randomly select 50 binary vectors for each label, with d = 500 , D = 250 , k = 50 , p o = 1 / 2 .
This is a generalization of the distribution D (1) , a component in the distribution of our simulation experiments (see the proof of Theorem 2 for details).
Recall the definition of D (1) :
˜ ϕ
A is uniform on only two values [+1 , . . .
, +1] and [0 , . . .
, 0] , and uniform over irrelevant patterns; the value
[+1 , . . .
, +1] corresponds to one class and [0 , . .
. , 0] correspond to another class.
Our data model here generalizes D (1) to more than 2 values.
The visualization is shown in Figure 15.
We can observe similar feature learning phenomena, and the neuron weights are updated to form clusters.
E.3. T WO - LAYER N ETWORKS ON M IXTURE OF G AUSSIANS
To further support our intuition of feature learning, we run experiments on mixture of Gaussians.
Data.
Let X =
R d be the input space, and Y = {± 1 } be the label space.
Suppose M ∈ R
d
× k is an dictionary with k orthonormal columns.
Let ε i , i = 1 , . . .
, k be i.i.d symmetric Bernoulli random variables, and g ∼ N (0 , σ 2 r kd I d ) .
Then we generate the input x and class label y by:
[Equation]
In this case, 2 k Gaussian clusters will be created.
The centers of the Gaussian clusters (cid:80) ki =1 ± M : i lie on the vertices of a hyper cube, and the label of each Gaussian cluster is determined by the parity function on the vertices of the hyper cube.
Note that the labeling function is roughly equivalent to a network: y = (cid:80) ni
=1 a i ReLU ( ⟨ c i , x ⟩ ) where c i ’s are the Gaussian centers, and a i ∝ 1 for Gaussian components with label 1 and a i ∝ − 1 for those with label -1.
Setting.
We then train a two-layer network with m =
800 hidden neurons on data sets generated as above with different chosen k ’s and d ’s.
The training follows typical practice (not the hyperparameters in our analysis).
In this setting, we expect the neural network to learn the effective features: the directions of Gaussian cluster centers.
Result.
We run experiments with different settings.
The parameters are shown in Table 6.
From
Figure 16 we can see that some neurons learn the directions of Gaussian centers, and each Gaussian center is covered by some neurons, which matches our expectation.
[Table]
E. R EAL D ATA : F EATURE L EARNING IN N ETWORKS
We take the subset of MNIST (Deng, 2012) with labels 0/1, CIFAR10 (Krizhevsky, 2012) with labels airplane/automobile and SVHN (Netzer et al., 2011) with labels 0/1, and train a two-layer network [Figure] with m = 50 .
We use traditional weight initialization method (random Gaussian) and training method (SGD with momentum = 0 .
95 without regularization) in this section, for our purpose of investigating the training dynamics in practice.
Then we visualize the neurons’ weights following the same method in the simulation.
Figure 17,
Figure 18 and Figure 19 show a similar feature learning phenomenon: effective features emerge after a few steps, and then get improved to form clusters.
This shows the insights obtained on our learning problems are also applicable to the real data. [Table]
E.4. CNN S ON B INARY C IFAR 10: F EATURE L EARNING IN N ETWORKS
Setting.
We use ResNet( m ) , which is a ResNet-18 convolutional neural network (He et al., 2016) with m filters in the first residual block.
It is obtained by scaling the number of filters in each block proportionally from the standard ResNet-18 network which is ResNet(64) .
We use ResNet(128) and ResNet(256) in this experiment.
{+We train our model on Binary CIFAR
+}
(Krizhevsky, 2012) with labels airplane/automobile for 20 epochs.
The final test accuracy of ResNet(128) is 95.75% and that of ResNet(256) is 93.8%.
Results.
Figure 20 visualizes the filters’ weights of different residual blocks in ResNet(128) at
Epoch 0, 3, and 20, and Figure 21 shows those in ResNet(256) .
They show that feature learning happens in the early stage, and show that there are some clusters of weights (e.g., the red and green points).
These colored points are selected at Epoch 20.
We first visualize the weights at Epoch 20, and then hand pick the points that roughly form two clusters (i.e., the points in the same cluster are close to each other while those in different clusters are far away).
We assign red and green colors to the two clusters at Epoch 20, and then assign these weights with the same color in Epoch 0 and 3.
Finally, we compute the cosine similarities and show that the hand picked points are indeed roughly clusters in the high-dimension.
In particular, we have the following three observations.
[Figure]
First, we can see that the filter weights change significantly during the early stage of the training, indicating feature learning happens in the early stage: the change between Epoch 0 and Epoch 3 is much more significant than that between Epoch 3 and Epoch 20.
Second we can also verify that the feature learning is guided by the gradients: the gradients of a filter in the early gradient steps point to similar directions (and thus the updated filter will learn this direction).
More precisely for a selected filter we average the gradients every 10 gradient steps (so to reduce the variance due to mini-batch) and get v 1 , v 2 and v 3 for the first 30 steps and compute their cosine similarities and those to their average.
Table 7 shows the results.
In general the similarities
[Figure]
We are high indicating they point to similar directions.
(Note that a similarity of 0 .
6 is regarded as very significant as the filters are in a high dimension of 3 × 3 × 1024 = 9216 ).
Third, we also observe some clustering effect of the filter weights, though not as significant as in our simulations.
For example, in the red and green clusters in Figure 20(a) for the first residual block, the average cosine similarity for filter weights in the red cluster is about 0.62 and that for the green is about 0.7, while the cosine similarity between the two clusters’ centers is about -0.72.
This shows significant similarities within the cluster while difference between clusters.
Note that the clustering is less significant than our simulation experiments.
This is because practical data have more patterns (i.e., effective feature directions) to be learned than our synthetic data, and also the practical network is not as overparameterized as in our simulation.
Then filters are likely to learn different patterns (or their mixtures) without forming significant clusters.
The results of
ResNet(256) show more significant clustering than ResNet(128) , which supports our explanation.
On the other hand, we emphasize that the key insight of our analysis is that the gradient guides the learning of effective features in the early stage of training (rather than the clustering), which is verified as discussed above.
E. R EAL D ATA : T HE E FFECT OF I NPUT S TRUCTURE
To study the influence of the input structure, we propose to keep the labeling function unchanged, vary the input distributions, and exam the change of the loss surface and the training dynamics.
We first describe the detailed experimental methodology, which allows us to generate data with similar labeling function but different input distributions.
Then we perform experiments on the generated datasets to investigate the change of the learning due to the change in the input distributions, and present the experimental results.
Finally, we also perform experiments to verify the intuition behind our experimental method.
E.5. E XPERIMENTAL M ETHODOLOGY
We consider the following experimental method.
Given an original dataset L = { ( x i , y i ) } ni =1 (e.g.,
CIFAR10) and an unlabeled dataset U =
{ ˜ x i } mi =1 from a proposed distribution P U (e.g., Gaussians), first extend the labeling function of L to U , giving synthetic labels ˜
y i to ˜ x i .
Then train a neural network on the union of L and the synthetic data L U = { (˜ x i , ˜ y i ) } mi =1 .
By investigating the new training dynamics, in particular the difference on the original part L and the synthetic part L U , we can see the effect of the input structure.
The original dataset should be from real-world data, since one of our goals is to compare them with synthetic data, and identify the properties of real-world data important for the success of learning.
A natural idea is to first learn a powerful network f ( x ) (called the teacher) on L to approximate the true labeling function, then apply f on U to generate synthetic labels, and finally train another network (called the student) on the synthetic data and original data.
However, we found that naïvely implementation of this idea fails miserably: the support of L and U can be typically different, and the powerful network learned over L can have entirely different behavior on U .
Therefore, we need to control the size of the teacher f
so that the labeling on U has similar complexity as that on L .
For our purpose, we can define the complexity of the labeling on L as the minimum size of the teacher achieving an approximation error ϵ for a chosen ϵ , if the ground-truth data distribution of L is known.
However, given only limited data, we cannot faithfully estimate the needed size of the teacher, and need to take into account the variance introduced by the finite data.
Our key idea is to use the U-shaped curve of the bias-variance trade-off and select the size of the teacher at the minimum of the U-shaped curve.
Since recent works (Belkin et al., 2019; Nakkiran et al., 2020) show that neural networks can have a double descent curve for the error v.s. model complexity, we thus plot the double descent curve, and find the minimum in the classical regime (corresponding to the traditional U-shape curve).
Our method is designed based on the following two reasons.
First, on the U-shaped curve, the complexity of the network is still roughly controlled by that of the number of parameters.
The local minimum of the U-shaped curve is a good measurement of the complexity of the data.
If the ground-truth is much more complicated than the teacher, then increasing the teacher’s size leads to a significant decrease in the approximation error (bias) compared to a small increase in the variance, that is, we will be on the left-hand side of the U-shaped.
In contrast, on the right-hand side of the
U-shaped, increasing the teacher’s size leads to a small decrease in the bias compared to a significant increase in the variance.
That is, the complexity of the ground-truth is comparable to or lower than the teacher.
So the local minimum approximates the complexity of the ground-truth labeling function.
Second, the local minimum point is chosen to get the best approximation of the true labels.
This helps to maintain the labeling from the real-world data and thus helps our investigation on the input, since too drastic change in the labeling can affect the training.
We note that the method is not perfect.
First, the teacher at the local minimum of U-shape may not have very high accuracy, especially on more complicated data.
To alleviate this, we also use the teacher to give synthetic labels y
′ i to x i in L , and train the student network on L ′ =
{ ( x i , y ′ i ) } ni =1 .
Though this introduces some differences from the original labels, it is acceptable for our purpose of studying the inputs.
Furthermore, ensuring the consistency of the labels on the original input in L and
U is important in our experiments.
Second, the measurement is an approximation due to variance.
Since only limited labeled data is available, it’s important and necessary to calibrate the measurement w.r.t.
the level of variance on the given dataset.
Method Description.
Algorithm 1 presents the details.
For a fixed network architecture for the teacher f , it first varies the network size and plots the double descent curve.
Then it selects the local minimum in the classic regime of U-shape and trains the teacher with the corresponding size.
In practice, we observed that the teacher might have unbalanced probabilities for different classes on U if its training does not take into account U .
Therefore, we propose the following heuristic regularization using x ∈ U , where λ is a regularization weight, and f ( x ) is the probabilities over classes given by the teacher:
[Equation]
Here, R 1 ( x ) guarantees that each kind of label has the same average probability to be generated, and
R 2 ( x ) pushes the probability away from uniform to avoid the case that the class probabilities for each data point converge to uniform.
Algorithm 1 Learning the teacher network to generate synthetic labels for studying the effect of the input structure
Input: teacher architecture f , labeled dataset L = { ( x i , y i ) } ni =1 , unlabeled dataset U =
{ ˜ x i } mi =1 .
Let i to be the size of f , f i to be the teacher of size i .
for i = 1 to n do
Train f i on L and let l i denote the test loss end for
Plot l i v.s. i , identify the classical regime, and the size i t corresponding to the local minimum in classical regime.
Train f i t on L with a regularizer R ( x ) on U defined in (291).
Output: f i t
E.5. E XPERIMENTAL R ESULTS
Network models.
Here we use one-hidden-layer fully-connected networks with m hidden units and quadratic activation functions.
The network is denoted as FC( m ) .
We use ResNet( m ) , which is a
ResNet-18 convolutional neural network (He et al., 2016) with m filters in the first residual block.
It is obtained by scaling the number of filters in each block proportionally from the standard ResNetnetwork which is ResNet(64) .
Datasets.
We use MNIST (Deng, 2012), CIFAR10 (Krizhevsky, 2012) and SVHN (Netzer et al., 2011) as L , and use Gaussian and images in Tiny ImageNet (Le & Yang, 2015) as U .
We generate the mixture data, where the fraction of the unlabeled data is denoted as α .
Setup.
We first use Algorithm 1 on the labeled data L and the unlabeled data U to get a synthetic labeling function (the teacher network) and then use it to give synthetic labels on a mixture of inputs from L and U .
For MNIST, the teacher network learned is FC(9) , where the number of the hidden units is determined by Algorithm 1.
See empirical verification in Figure 26.
For CIFAR10 and SVHN, the teacher networks are ResNet(5) and ResNet(2) , respectively, as determined by our method.
The student network for MNIST is FC(9) , and those for CIFAR10 and SVHN are ResNet(9) and
ResNet(8) , respectively.
Finally, we train the student networks on these new datasets with perturbed input distributions.
[Figure] [Figure] [Figure] Figure 22 shows the results on an equal mixture of data and Gaussian.
It presents the test accuracy of the student on the original data part, the Gaussian part, and the whole mixture.
For example, for CIFAR10, the test accuracy on the whole mixture is lower than that of training on the original
CIFAR10, showing that the input structure indeed has a significant impact on the learning.
Furthermore, the network learns well over the CIFAR10 part (with accuracy similar to that on the original data) but learns slower with worse accuracy on the Gaussian part.
This suggests that the
CIFAR10 input structure is still helping the network to learn effective features.
While the results on
MNIST+Gaussian do not show a significant trend (possibly because the tasks there are simpler), the results on SVHN+Gaussian show similar significant trends as CIFAR10+Gaussian.
Figure 24 shows the results when we vary the fraction of the Gaussian data α .
We observe that the test accuracy curve on the original part and that on the synthetic part have roughly the same trend for different α as before, further verifying our insights.
Figure 23 shows the results when mixed with Tiny ImageNet data instead of Gaussians.
It shows a similar trend, while the performance on the Tiny ImageNet part is higher than that on the Gaussian part.
This suggests that compared to Gaussians, the Tiny ImageNet data has helpful input structures, though not as helpful as that on the original data for learning the particular labeling. E.5. L ARGER N ETWORK ON MNIST FOR C HECKING T HE E FFECT OF I NPUT S TRUCTURE
Here we perform the experiment on MNIST as in E.5.2, but for a network with m = 50 hidden neurons rather than m = 9 .
Figure 25 shows similar results as those for m = 9 : the learning on the MNIST input part is faster and better than that on the Gaussian input part. The separation between the two is actually more significant than that for m = 9 .
This then also supports our insight about the effect of input structures E.5. E MPIRICAL V ERIFICATION OF O UR M ETHOD
[Figure] We also perform experiments to verify the intuition behind our methodology, i.e., the method gives a synthetic labeling function with roughly the same complexity on the original inputs and the injected inputs.
We first use our method on MNIST and samples (of the same size as MNIST) from a Gaussian to get the teacher FC(9) ; the double descent curve is in Figure 26(a).
Then we train students on the Gaussian data with synthetic labels from the teacher, and plot the double descent curve for the students in Figure 26(b).
The local minimums of the two U-shapes are roughly the same, matching our reasoning.
Then we also train larger teachers and plot the double descent curve for students on
Gaussian data.
Figure 26(c)
Teacher size 50 .
Figure 26(d)
Teacher size 500 .
The local minimum of the U-shape becomes larger when the teacher gets larger, again matching our reasoning.
F P ROVABLE G UARANTEES FOR N EURAL N ETWORKS IN A M ORE G ENERAL S ETTING
This section provides the analysis in a more general setting.
We first describe the learning problems, and then provide the proofs following similar intuitions as for the simpler settings in the main text.
F. P ROBLEM S ETUP
Let X =
R d be the input space, and Y = {± 1 } be the label space.
Suppose M ∈ R d × D is a dictionary with D elements, where each element M j can be regarded as a pattern.
We assume quite general incoherent dictionary:
(D) M is µ -incoherent, i.e., the columns of M are unit vectors, and for any i ̸
= j, |⟨ M i , M j ⟩| ≤ [Equation]
/ √ d.
Note that the setting in the main text corresponds to µ = 0 .
Let ˜ ϕ ∈ { 0 , 1 } D be a hidden vector that indicates the presence of each pattern, and D ˜ ϕ a distribution for ˜ ϕ .
Let A ⊆
[ D ] be a subset of size k corresponding to the class relevant patterns.
Let P ⊆
[ k ] .
We first sample ˜ ϕ
from D ˜ ϕ , and then generate the input ˜ x and the class label y from ˜ ϕ, A, P by:
[Equation]
where the Gaussian noise ζ ∼ N (0 , σ 2 ζ I d × d ) is independent from ˜ ϕ .
Note that the setting in the main text corresponds to σ ζ = 0 .
We allow general D ˜ ϕ with the following assumptions:
(A1)
The patterns in A are correlated with the labels: for any i ∈ A , for v ∈ {± 1 } let γ v = E
[ y ˜ ϕ i | y
= v ] , then γ :=
( γ +1 + γ − 1 ) / 2 > 0 .
(A2)
The patterns outside A are independent of the patterns in A .
Note that we allow imbalanced classes.
Let p min := min(Pr[ y = − 1] , Pr[ y = +1]) .
If the classes are balanced, then the assumption (A1) implies the assumption (A1) in the main text, so the setting here is more general.
(A2) is also more general, in particular, allowing dependence between irrelevant patterns and non-identical distributions for them.
Let D ( A, P, D ˜ ϕ ) denote the distribution on (˜ x, y ) corresponding to some A, P, and D ˜ ϕ .
Given parameters Ξ =
( d, D, k, γ, p o , µ, σ ζ ) , the family F Ξ of distributions for learning is the set of all
D ( A, P, D ˜ ϕ ) with A ⊆
[ D ] , P ⊆
[ k ] , and D ˜ ϕ satisfying the above assumptions.
One special case is the mixture of two Gaussians.
Example.
Suppose M has one single column v , and y
=
+
if ˜ ϕ = 1 and y = − 1 otherwise.
Then the data distribution is simply a mixture of two Gaussians: ˜ x ∼ v 2 + N ( y v 2 , σ 2 ζ I d × d ) .
F.1. N EURAL N ETWORK L EARNING
Again, we will normalize the data for learning: we first compute x = (˜ x − E [˜ x ]) /
˜ σ where ˜ σ 2 := (cid:80) di
=1 (˜ x i − E
[˜ x i ]) 2 = (cid:80)
j ∈
[ D ] Var (˜ ϕ j ) + dσ 2 ζ is the variance of the data, and then train on ( x, y ) .
This is equivalent to setting ϕ =
(˜ ϕ − E [˜ ϕ ]) /
˜ σ and generating x = Mϕ + ζ/σ ζ .
For (˜ x, y ) from D and the normalized ( x, y ) , we will simply say ( x, y ) ∼ D .
The learning will be the same as that in the main text, except the following.
We will use a small σ 2 w = ˜
σ 2 / poly ( Dm ) .
And we will use a weighted loss to handle the imbalanced classes in the first two steps for feature learning, and then use the unweighted loss in the remaining steps.
Formally, the weighted loss is:
[Equation]
[Equation]
where the class weights α v = 1 2Pr[ y = v ] for v ∈ {± 1 } .
F. M AIN R ESULT
In this setting, we have the following theorem:
Theorem 34.
Set
[Equation]
For any δ ∈ (0 , O (1 /k 3 )) , if µ ≤ O ( √ d/D ) , σ ζ ≤
[Equation]
, m ≥ max { Ω( k 4 ) , D, d } , then we have for any D ∈ F Ξ , with probability at least 1 − δ , there exists t ∈
[ T ] such that
[Equation]
Consequently, for any ϵ ∈ (0 , 1) , if T = m 4 / 3 , and m ≥ max { Ω( k 12 /ϵ 3 / 2 ) , D } , then
[Equation]
The rest of the section is devoted to the proof of this theorem.
F. N OTATIONS
Recall some notations that we will use throughout the analysis.
For a vector v and an index set I , let v
I denote the vector containing the entries of v indexed by I , and v − I denote the vector containing the entries of v with indices outside I .
Let ρ :=
M ⊤ M .
Then we have ρ jj = 1 for any j , and | ρ jℓ | ≤ µ/ √ d for any j ̸ = ℓ .
By initialization, w (0) i for i ∈ [ m ] are i.i.d.
copies of the same random variable w (0) ∼ N (0 , σ 2 w I d × d ) ; similar for a (0) and b (0) .
Let σ 2 ϕj :=
p o j (1 − p o j ) / ˜ σ 2 denote the variance of ϕ ℓ for ℓ ̸∈ A , where p o j = Pr[˜ ϕ j = 1] .
Let p o be the value such that with probability 1 − exp( − Ω( k )) , (cid:80) j ̸∈ A ˜ ϕ
j ≤ p o
( D − k ) for some p o ∈ [0 , 1] .
That is, p o is an upper bound on the density of ˜ ϕ j with high probability.
Let q ℓ := ⟨ w (0) , M ℓ ⟩ .
Similarly, define q ( t ) i,ℓ := ⟨ w ( t ) i , M ℓ ⟩ .
We also define the following sets to denote typical initialization.
For a fixed δ ∈ (0 , 1) , define
G w ( δ ) :
= (cid:40)
[Equation]
2 log( m/δ ) } .
F. E XISTENCE OF A G OOD N ETWORK
We first show that there exists a network that can fit the data distribution.
Lemma 35.
Suppose kµ √ d p o D ˜ σ ≤ 116 .
For any D ∈ F Ξ , there exists a network g ∗ ( x )
= (cid:80) ni
=1 a ∗ i σ ( ⟨ w ∗ i , x ⟩ + b ∗ i ) which satisfies
[Equation]
Furthermore, the number of neurons n = 3( k + 1) , | a ∗ i | ≤ 64 k, 1 / (64 k ) ≤ | b ∗ i | ≤ 1 / 4 , w ∗ i = ˜
σ (cid:80)
j ∈
A M j / (8 k ) , and |⟨ w ∗ i , x ⟩ + b ∗ i | ≤ 1 for any i ∈ [ n ] and ( x, y ) ∼ D .
Consequently, if furthermore we have kµ/ √ d < 1 and σ ζ < 1 /k , then
[Equation]
Proof of Lemma 35.
Let w = ˜ σ (cid:80)
j ∈
A M j and let
u
= (cid:80)
j ∈
A E [˜ ϕ j ] .
We have
[Equation]
[Equation] With probability ≥ 1 − exp( − Ω( k )) , among all j ̸∈ A , we have that at most p
o
( D − k ) of ϕ j are (1 − p o ) /
˜ σ , while the others are − p o / ˜ σ , and thus
[Equation]
Furthermore, ⟨ w, ζ ⟩ ∼ N (0 , σ 2 ζ ∥ w ∥ 22 ) and ∥ w ∥ 22 ≤ ˜ σ 2 ( k + k 2 µ/ √ d ) , we have
[Equation]
For good data points with ϕ and ζ satisfying the above, we have | ϵ x | ≤ 1 / 8 .
By Lemma 7,
[Equation]
Then for good data points, we have yg ∗ 1 ( x ) ≥ 1 .
Similarly,
[Equation]
Then for good data points, we have yg ∗ 2 ( x ) ≥ 1 .
Note that the bias terms in g ∗ 1 and g ∗ 2 have distance at least 1 / 4 , then at least one of them satisfies that all its bias terms have absolute value ≥ 1 / 8 .
Pick that one and denote it as g ( x )
= (cid:80) ni
=1 a i σ r ( ⟨ w i , x ⟩ + b i ) .
By the positive homogeneity of σ r , we have
[Equation]
Since for any good data points, |⟨ w i , x ⟩ / (8 k ) + b i / (8 k ) | ≤ 1 , then g ( x ) = n
[Equation]
i =
[Equation]
where σ is the truncated ReLU.
Now we can set a ∗ i = 8 ka i , w ∗ i =
w i / (8 k ) , b ∗ i
= b i / (8 k ) , to get our final g ∗ .
F. I NITIALIZATION
We first show that with high probability, the initial weights are in typical positions.
Lemma 36.
Suppose Dµ/ √ d ≤ 1 / 16 .
For any δ ∈ (0 , 1) , with probability at least 1 − δ − 2 exp ( − Θ( D − k )) over w (0) ,
[Equation]
With probability at least 1 − δ over b (0) ,
[Equation]
With probability at least 1 − δ over a (0) ,
[Equation]
Proof of Lemma 36.
The bound on ∥ w (0) ∥ 22 follows from the property of Gaussians.
Note that q = M ⊤ w (0) ∼ N (0 , σ 2 w ρ ) for the matrix ρ = M ⊤ M .
We have with probability ≥ 1 − δ/ 2 , max ℓ | q ℓ | ≤ (cid:113) 2 σ 2 w log Dδ .
For any subset S ⊆
[ D ] , let ρ S denote the submatrix of ρ containing the rows and columns indexed by S .
Then q S =
M ⊤ w (0) ∼ N (0 , σ 2 w ρ S ) .
By diagonalizing ρ S and then applying Bernstein’s inequality, we have with probability ≥ 1 − 2 exp ( − Θ( | S | / ∥ ρ ∥ 2 ) , ∥ q S ∥ 22 ∈ (cid:16) ( ∥ ρ S ∥ 2 F − | S | 4 ) σ 2 w , ( ∥ ρ S ∥ 2 F + | S | 4 ) σ 2 w (cid:17) .
By Gershgorin circle theorem, we have
[Equation]
Similarly, we have
[Equation]
The bounds on q then follow.
The bounds on b (0) and a (0) follow from the property of Gaussians.
Lemma 37.
Suppose Dµ/ √ d ≤ 1 / 16 .
We have:
• With probability ≥ 1 − δ − 2 m exp( − Θ( D − k )) over w (0)
i ’s, for all i ∈ [2 m ] , w (0) i ∈
G w ( δ ) .
• With probability ≥ 1 − δ over b (0)
i ’s, for all i ∈ [2 m ] , b (0) i ∈
G b ( δ ) .
• With probability ≥ 1 − δ over a (0) i ’s, for all i ∈ [2 m ] , a (0) i ∈ G a ( δ ) .
Proof of Lemma 37.
This follows from Lemma 36 by union bound.
F. S OME A UXILIARY L EMMAS
The expression of the gradients will be used frequently.
Lemma 38.
[Equation]
Proof of Lemma 38.
It follows from straightforward calculation.
We also have the following auxiliary lemma for later calculations.
Lemma 39.
[Equation]
Proof of Lemma
[Equation]
F. F EATURE E MERGENCE : F IRST G RADIENT S TEP
We will show that w.h.p.
over the initialization, after the first gradient step, there are neurons that represent good features.
We begin with analyzing the gradients.
Lemma 40.
Fix δ ∈ (0 , 1) and suppose w (0)
i
[Equation]
G
[Equation]
, b [Equation] 0) i ∈ G b ( δ ) for all i ∈ [2 m ] .
Let
[Equation]
If σ 2 ζ σ 2 w d/ ˜ σ 2 =
O (1 /k ) , p o = Ω( k 2 /D ) , k = Ω(log 2 ( Dmd/δ )) , and σ (1) ξ
= O (1 /k ) , then
[Equation]
[Equation] where T j satisfies:
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ O ( ϵ e / ˜ σ ) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if j ̸∈ A , then | T j | ≤ O ( σ 2 ϕj ϵ e ˜ σ ) ; • | ν j | ≤ O (cid:18) σ ζ √ log( k ) ˜ σ ϵ ν (cid:19)
+ σ ζ d ˜ σ e − Θ( k ) .
Proof of Lemma 40.
Consider one neuron index i and omit the subscript i in the parameters.
Since the unbiased initialization leads to g (0) ( x ; ξ (1) ) =
[Equation] we have
[Equation]
First, consider j ∈ A .
[Equation]
where ι := ⟨ w (0) , ζ/ ˜ σ ⟩ .
Let
[Equation] [Equation]
Note that |⟨ ϕ A , q A ⟩| = O
( kσ w √ 2log( D/δ ) ˜ σ) , and that | ι | = |⟨ w (0) i , ζ/ ˜ σ ⟩| = O ( √ dσ ξ σ w √ 2log( D/δ ) ˜ σ ) , and that |⟨ ϕ, q ⟩| , |⟨ ϕ − A , q − A ⟩| are O
( Dσ w √ 2log( D/δ ) ˜ σ) .
When σ w is sufficiently small, by the property of the Gaussian ξ (1) , we have
[Equation] [Equation]
+ b (0) + ξ (1) ≥ 0 (cid:105) − Pr ξ (1) (cid:104) ⟨ ϕ − A , q − A ⟩ + ι + b (0) + ξ (1) ≥ 0 (cid:105)(cid:12)(cid:12)(cid:12)(cid:12) (353) + Pr ξ (1) (cid:104) ⟨ ϕ, q ⟩ + ι + b (0) + ξ (1) ≥ 1 (cid:105) + Pr ξ (1) (cid:104) ⟨ ϕ − A , q − A ⟩ + ι + b (0) + ξ (1) ≥ 1 (cid:105) (354) = O ( ϵ e ) . In summary, | E ζ,
[Equation]
[Equation] ) | = O ( ϵ
[Equation]
Then we have (cid:12)(cid:12) [Equation] (357) ≤
[Equation]
I
[Equation]
(358) ≤ O (
[Equation]
E ϕ
[Equation]
{| α
[Equation]
yϕ
[Equation]
|} (359) ≤
[Equation]
[Equation]
where the last step is from Lemma 39.
Furthermore,
[Equation]
When σ w is sufficiently small
[Equation]
we have
[Equation]
This leads to
[Equation]
By Lemma [Equation] E ϕ A { α y yϕ j } = γ/ ˜ σ .
Therefore,
[Equation]
Now
[Equation]
consider j ̸∈ A .
Let B denote A ∪ { j } .
[Equation]
Let
[Equation]
Similar as above
[Equation]
we have | E ζ,ξ (1) ( I b − I ′ b ) | ≤ O ( ϵ e ) .
Then by Lemma 39,
[Equation]
Furthermore,
[Equation]
Therefore,
[Equation]
Finally, consider ν j .
[Equation]
where ι j := w (0) j ζ j / ˜ σ and ι − j := ⟨ w (0) , ζ/ ˜ σ ⟩ − ι j .
With probability ≥ 1 − d exp( − Θ( k )) over ζ , for any j , | ζ j | ≤ O ( σ ζ (cid:112) log( k )) .
Let G ζ denote this event.
Let
[Equation] [Equation]
Similar as above, we have | E ζ
[ I j − I ′ j |G ζ ] | ≤ O ( ϵ ν ) .
Then
[Equation]
[Equation]
A
[Equation]
The first term is bounded by
[Equation]
[Equation]
[Equation]
[Equation]
The second term is bounded by
[Equation]
[Equation]
(cid:12)(cid:12)(cid:12) E
[Equation]
ϕ
[Equation]
A (cid:26)
[Equation]
[Equation]
(cid:27)(cid:12) (cid:12)(cid:12)(cid:12)
[Equation]
de
[Equation]
( k ) (396)
≤ E ϕ A
(cid:12)(cid:12)(cid:12) α y y ˜
σ (cid:12)(cid:12)(cid:12) × E ζ {| ζ j || − G ζ } × de − Θ( k ) (397) ≤ σ ζ ˜ σ × de − Θ( k ) (398) ≤ σ ζ d ˜ σ e − Θ( k ) .
Furthermore,
E
[Equation]
E ζ [Equation] [ I ′
[Equation]
] = 0 .
Therefore, | ν j | ≤ O (cid:32) σ ζ (cid:112) log( k ) ˜ σ ϵ ν (cid:33)
+ σ ζ d ˜ σ e − Θ( k ) .
Lemma 41.
Under the same assumptions as in Lemma 40,
[Equation]
where | T b | ≤ O ( ϵ e ) .
Proof of Lemma 41.
Consider one neuron index i and omit the subscript i in the parameters.
Since the unbiased initialization leads to g (0) ( x ; ξ (1) )
= 0 , we have
[Equation]
[Equation] where ι := ⟨ w (0) , ζ/ ˜ σ ⟩ .
Similar to the proof in Lemma 40,
[Equation]
q − A ⟩ + ι + b (0) + ξ (1) ∈ (0 , 1)] (cid:33)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) =
O ( ϵ e ) .
Then (cid:12)(cid:12)(cid:12)(cid:12) T b − E ϕ A ,ζ (cid:26) α y y Pr ϕ − A ,ξ (1) [ ⟨ ϕ − A , q − A ⟩ + ι + b (0) + ξ (1) ∈ (0 , 1)] (cid:27)(cid:12)(cid:12)(cid:12)(cid:12) (409) =
E ϕ A ,ζ (cid:40) | α y y | (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) Pr ϕ − A ,ξ (1) [ ⟨ ϕ, q ⟩ + ι + b (0) + ξ (1) ∈ (0 , 1)] (410) − Pr ϕ − A ,ξ (1) [ ⟨ ϕ − A , q − A ⟩ + ι + b (0) + ξ (1) ∈ (0 , 1)] (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:41) (411) ≤
O ( ϵ e ) E ϕ A {| α y y |} (412) ≤ O ( ϵ e ) .
Also,
E ϕ A ,ζ (cid:26) α y y Pr ϕ − A ,ξ (1) [ ⟨ ϕ − A , q − A ⟩ + ι + b (0) + ξ (1) ∈ (0 , 1)] (cid:27) (414) =
E ϕ A { α y y } Pr ϕ − A ,ζ,ξ (1) [ ⟨ ϕ − A , q − A ⟩ + ι + b (0) + ξ (1) ∈ (0 , 1)] (415) = 0 .
Therefore, | T b | ≤ O ( ϵ e ) .
Lemma 42.
We have
[Equation]
[Equation] where | T a | ≤ O (max ℓ q (0) i,ℓ ) .
So if w (0) i ∈ G ( [Equation] | T a | ≤ [Equation]
log
[Equation] Dm [Equation] Proof of Lemma 42.
Consider one neuron index i and omit the subscript i in the parameters.
Since the unbiased initialization leads to g (0) ( x ; ξ
[Equation]
= 0 , we have
[Equation]
Let ϕ ′
A be an independent copy of ϕ A
[Equation]
ϕ ′ be the vector obtained by replacing in ϕ the entries ϕ A with ϕ ′
A
[Equation]
and let x ′ =
Mϕ ′ + ζ/ ˜ σ and its label is y ′ .
Then
[Equation]
Since σ is 1-Lipschitz,
[Equation]
With the bounds on the gradient, we now summarize the results for the weights after the first gradient step.
Lemma 43.
Set
[Equation]
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
G w ( δ ) , b (0) i ∈ G b ( δ ) for all i ∈ [2 m ] .
If k = Ω(log 2 ( Dm/δ )) , then for all i ∈ [ m ] , w (1) i
[Equation]
satisfying
• if ℓ ∈ A , then | q (1) i,ℓ − η (1) a (0) i βγ/ ˜ σ | ≤ O (cid:18) | η (1) a (0) i | ϵ e ˜ σ (cid:19) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if ℓ ̸∈ A , then | q (1) i,ℓ | ≤ O (cid:16) | η (1) a (0) i | σ 2 ϕℓ ϵ e ˜ σ (cid:17) ;
• | υ j | ≤ O (cid:18) | η (1) a (0) i | (cid:18) σ ζ √ log( k ) ˜ σ
ϵ ν + σ ζ d ˜ σ e − Θ( k ) (cid:19)(cid:19) .
and
• b (1) i =
b (0) i + η (1) a (0)
i T b where | T b | =
O ( ϵ e ) ; • a (1) i =
a (0) i + η (1) T a where | T a | =
O ( σ w (cid:112) log( Dm/δ )) .
Proof of Lemma 43.
This follows from Lemma 37 and Lemma 40-42.
F. F EATURE I MPROVEMENT : S ECOND G RADIENT S TEP
We first show that with properly set η (1) , for most x , | g (1) ( x ; σ (2) ξ ) | < 1 and thus yg (1) ( x ; σ (2) ξ ) < 1 .
Lemma 44.
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
G w ( δ ) , b (0) i ∈ G b ( δ ) , a (0) i ∈ G a ( δ ) for all i ∈ [2 m ] .
If Dµ/ √ d ≤ 1 / 16 , σ ζ ˜ σ = O (1) , σ 2 ζ d/ ˜ σ 2 =
O
(1) , k = Ω(log
( Dm/δ )) , σ a ≤ ˜ σ 2 / ( γk 2 ) , η (1) =
O (cid:16) γ kmσ a ˜ σ (cid:17) , and σ (2) ξ ≤ 1 /k , then with probability ≥ 1 − ( d + D ) exp( − Ω( k )) over ( x, y ) , we have yg (1) ( x ; σ (2) ξ ) < 1 .
Furthermore, for any i ∈ [2 m ] , (cid:12)(cid:12)(cid:12) ⟨ w (1) i , ζ/ ˜ σ ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , (cid:12)(cid:12)(cid:12) ⟨ q (1) i , ϕ ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , and (cid:12)(cid:12)(cid:12) ⟨ ( q (1) i ) − A , ϕ − A ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , and for any j ∈ [ d ] , ℓ ∈ [ D ] , | ζ j | ≤ O ( σ ζ (cid:112) log( k )) and |⟨ ζ, D ℓ ⟩| ≤ O ( σ ζ (cid:112) log( k )) .
Proof of Lemma 44.
Note that w (0)
i
[Equation]
= [Equation]
i .
Then the gradient for w m
+ i is the negation of that for w m
+ i , the gradient for b m
+ i is the negation of that for b m + i , and the gradient for a m
+ i is the same as that for a m
+ i .
[Equation]
[Equation]
i
[Equation]
[Equation]
m (cid:88) i =
[Equation]
a (1) i E ξ (2) σ ( ⟨ w (1) i , x ⟩ + b (1)
i
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
(cid:17)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (433) ≤ (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) m (cid:88) i = (cid:16) a (1) i E ξ (2) σ ( ⟨ w (1) i , x ⟩ [Equation] b (1) i
[Equation]
[Equation]
[Equation] (cid:17)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (434) + (cid:12) (cid:12)(cid:12)(cid:12)(cid:12) m (cid:88) i = (cid:16) − a (1) m + i E ξ (2) σ ( ⟨ w (1) i , x ⟩ + b (1) i + ξ (2) i )
+ a (1) m + i E ξ (2) σ ( ⟨ w (1) m
+ i , x ⟩ + b (1) m + i + ξ (2) i ) (cid:17) (cid:12) (cid:12)(cid:12)(cid:12)(cid:12) . Then we have (cid:12)(cid:12)(cid:12) g (1) ( x ; σ (2) ξ ) (cid:12)(cid:12)(cid:12) ≤ m (cid:88) i = (cid:12)(cid:12)(cid:12) 2 η (1) T a E ξ (2) σ ( ⟨ w (1) i , x ⟩ + b (1) i + ξ (2) i ) (cid:12)(cid:12)(cid:12) (436) + m (cid:88) i = (cid:12)(cid:12)(cid:12) a (1) m
+ i (cid:12)(cid:12)(cid:12) (cid:16)(cid:12)(cid:12)(cid:12) ⟨ w (1) i − w (1) m
+ i , x ⟩ (cid:12)(cid:12)(cid:12) + (cid:12)(cid:12)(cid:12) b (1) i − b (1) m
+ i (cid:12)(cid:12)(cid:12)(cid:17) (437)
≤ m (cid:88) i =
(cid:12)(cid:12)(cid:12) 2 η (1) T a (cid:12)(cid:12)(cid:12) (cid:16)(cid:12)(cid:12)(cid:12) ⟨ w (1) i , x ⟩ + b (1) i (cid:12)(cid:12)(cid:12) + E ξ (2) (cid:12)(cid:12)(cid:12) ξ (2)
i (cid:12)(cid:12)(cid:12)(cid:17) (438) + m (cid:88) i = (cid:12)(cid:12)(cid:12) a (1) m + i (cid:12)(cid:12)(cid:12) (cid:16)(cid:12)(cid:12)(cid:12) ⟨ w (1) i − w (1) m + i , x ⟩ (cid:12)(cid:12)(cid:12) + (cid:12)(cid:12)(cid:12) b (1) i − b (1) m + i (cid:12)(cid:12)(cid:12)(cid:17) .
[Equation]
With probability ≥ 1 − exp( − Ω( k )) , among all j ̸∈ A , we have that at most p
o
( D − k ) of ϕ j are (1 − p
o j ) / ˜ σ , while the others are − p o j / ˜ σ .
With probability ≥ 1 − ( d + D ) exp( − Ω( k )) over ζ , for any j , |
ζ
j | ≤ O ( σ ζ (cid:112) log( k )) and |⟨ ζ, D ℓ ⟩| ≤ O ( σ ζ (cid:112) log( k )) .
For data points with ϕ and ζ satisfying these, we have:
[Equation]
Proof of Claim 1.
[Equation]
D (cid:88) ℓ = q (1) i,ℓ M ℓ , ζ/ ˜ σ ⟩ (cid:12) (cid:12)(cid:12)(cid:12)(cid:12) + (cid:12)(cid:12) (cid:12) (cid:12) (cid:12)(cid:12) ⟨ υ,
D (cid:88)
j = ϕ j M j ⟩ (cid:12)(cid:12) (cid:12) (cid:12) (cid:12)(cid:12) + |⟨ υ, ζ/ ˜ σ ⟩| .
For each term above we bound as follows.
Note that when σ w is sufficiently small, ϵ e =
O ( k log 1 / 2 ( Dm/δ ) / √ D ) .
Let
[Equation] [Equation]
[Equation]
C
[Equation]
p
[Equation]
Then | a
[Equation]
[Equation] [Equation] (448) | [Equation]
[Equation]
Then by the assumption on µ , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ⟨ D (cid:88) ℓ = q (1) i,ℓ M ℓ , D (cid:88) j =
[Equation]
[Equation]
) = (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:88) ℓ ∈ A ⟨ q (1)
i
[Equation]
ℓ
[Equation]
ℓ
[Equation]
ℓ
[Equation]
ℓ
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
ℓ
[Equation]
[Equation]
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (451) ≤ O ( | η (1) a (0) i | ) (cid:18) B 1 C 1 + B 2 C 2 + µ √ d
kB 1 ( C 1 + C 2 ) + DB 2 ( C 1 + C 2 ))
(cid:19) (452) ≤ O ( | η (1) a (0) i | ) (cid:18) B 1 C 1 + B 2 C 2 + kµ √ dB 1 C 2 + B 2 C 1 (cid:19) (453) ≤
O ( η (1) ) (1 /γ + ˜ σ/γ + 1 /γ + 1 /γ ) (454) ≤ O ( η (1) /γ )
(1 + ˜ σ ) .
By the assumption on σ ζ ,
[Equation]
Also note that | ν j | ≤ O (cid:16) σ ζ log 2 ( Dm/δ ) ˜ σ √ D (cid:17) .
Then by the assumption on σ ζ , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ⟨ υ,
D (cid:88)
j = ϕ
j M j ⟩ (cid:12) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (461) ≤
[Equation]
σ ζ log
[Equation]
˜
[Equation]
[Equation]
≤ O ( | η (1) /γ ) . Finally, we have |
[Equation]
We also have:
[Equation]
Then we have
[Equation]
Then (cid:12)(cid:12)(cid:12) yg
[Equation]
< 1 . Finally, the statement on (cid:12)(cid:12)(cid:12) ⟨ ( q (1) i ) − A , ϕ − A ⟩ (cid:12)(cid:12)(cid:12) follows from a similar calculation on (cid:12)(cid:12) (cid:12) ⟨ w (1) i , x ⟩ (cid:12)(cid:12) (cid:12) =
(cid:12)(cid:12) (cid:12) ⟨ q (1) i , ϕ ⟩ (cid:12)(cid:12) (cid:12) .
We are now ready to analyze the gradients in the second gradient step.
Lemma 45.
Fix δ ∈ (0 , 1) and suppose w (0) i ∈
G w ( δ ) , b (0) i ∈ G b ( δ ) , a (0) i ∈ G a ( δ ) for all i ∈ [2 m ] .
Let ϵ e2 :
= O (cid:18) η (1) | a (0)
i | k
( γ + ϵ e ) ˜ σ 2 σ (2) ξ (cid:19) + exp( − Θ( k )) .
If Dµ/ √ d ≤ 1 / 16 , σ ζ ˜ σ = O (1) , σ 2 ζ d/ ˜ σ 2 =
O
(1) , k = Ω(log
( Dm/δ )) , σ a ≤ ˜ σ 2 / ( γk 2 ) , η (1) =
O (cid:16) γ kmσ a ˜ σ (cid:17) , and σ (2) ξ = 1 /k 3 / 2 , then ∂
∂w i L D ( g (1) ; σ (2) ξ ) = − a (1) i   D (cid:88)
j = M j T j + ν 
 (479) where T j satisfies:
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ O ( ϵ e2 / ˜ σ + η (1) /σ (2) ξ + η [Equation]
i | ϵ e / (˜ σσ (2) ξ )) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if j ̸∈ A ,
then | T j | ≤ 1˜ σ exp( − Ω( k ))
+ O ( σ 2 ϕ ˜ σϵ e2 ) ; • | ν j | ≤ O (cid:18) η (1) σ ζ γσ (2) ξ (cid:19) + exp( − Ω( k )) .
Proof of Lemma 45.
Consider one neuron index i and omit the subscript i in the parameters.
By Lemma 44, with probability at least 1 − ( d + D ) exp( − Ω( k ))
= 1 − exp( − Ω( k )) over ( x, y ) , yg (1) ( x ; ξ (2) ) > 1 and furthermore, for any i ∈ [2 m ] , (cid:12)(cid:12)(cid:12) ⟨ w (1) i , ζ/ ˜ σ ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , (cid:12)(cid:12)(cid:12) ⟨ q (1) i , ϕ ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , and (cid:12)(cid:12)(cid:12) ⟨ ( q (1) i ) − A , ϕ − A ⟩ (cid:12)(cid:12)(cid:12) =
O ( η (1) ˜ σ/γ ) , and for any j ∈ [ d ] , ℓ ∈ [ D ] , | ζ j | ≤ O ( σ ζ (cid:112) log( k )) and |⟨ ζ, D ℓ ⟩| ≤ O ( σ ζ (cid:112) log( k )) .
Let I x be the indicator of this event.
[Equation]
[Equation]
Let T j 1 :=
E ( x,y ) ∼D ,ξ (2) (cid:8) α
y y
I [ ⟨ w (1) , x ⟩ + b (1) + ξ (2) ∈ (0 , 1)]
ϕ j (cid:9) .
We have
[Equation]
Similarly, let ν ′ :
= E ( x,y ) ∼D ,ξ (2) (cid:110) α y yζ ˜
σ I [ ⟨ w (1) , x ⟩ + b (1) + ξ (2) ∈ (0 , 1)] (cid:111) .
We have
[Equation]
So it is sufficient to bound T j
and ν ′ .
For simplicity, we use q as a shorthand for q (1) i .
First, consider j ∈ A .
[Equation]
where ι := ⟨ w (1) , ζ/ ˜ σ ⟩ . Let
I a
[Equation]
=
[Equation]
[Equation]
By the property of the Gaussian ξ (2) , that |⟨ ϕ A , q A ⟩| =
O ( η (1) | a (0)
i | k
( γ + ϵ e ) ˜ σ) , and that | ι | = |⟨ w (1) i , ζ/ ˜ σ ⟩| , |⟨ ϕ, q ⟩| , |⟨ ϕ − A , q − A ⟩| are all O ( η (1) ˜ σ/γ ) <
O (1 /k ) , we have
[Equation]
[Equation]
[Equation]
This leads to
[Equation]
where the last step is from Lemma
[Equation]
Furthermore,
[Equation]
By Lemma 19, we have |⟨ ϕ − A , q − A ⟩ + ι | ≤ O ( η (1) ˜ σ/γ ) .
Also, | b (1) − b (0) | ≤ O ( η (1) | a [Equation] 0) i | ϵ e ) .
By the property of ξ (2) ,
[Equation]
On the other hand
[Equation]
and β only depends on b (0) .
By Lemma 39, E ϕ A { α y yϕ j } = γ/ ˜ σ .
Therefore,
[Equation]
Now, consider j ̸∈ A .
Let B denote A ∪ { j } .
[Equation]
Let
[Equation]
Similar as above
[Equation]
we have | I b − I ′ b | ≤ ϵ e2 .
Then by Lemma 39,
[Equation]
Furthermore,
[Equation]
b ]
[Equation]
[Equation]
Therefore, |
[Equation]
[Equation]
Finally, consider ν ′ j .
[Equation]
Let
[Equation]
Since | ι | ≤ O ( η (1) ˜ σ/γ )
[Equation]
we have | I j − I ′
j | ≤ O ( η (1) ˜ σ/ ( γσ (2) ξ )) .
Then
[Equation]
Furthermore,
[Equation]
Therefore,
[Equation]
Lemma 46.
Under the same assumptions as in Lemma 45,
[Equation]
where | T b | ≤ exp( − Ω( k ))
+ O ( ϵ e2 ) .
Proof of Lemma 46.
Consider one neuron index i and omit the subscript i in the parameters.
By
Lemma 44, Pr[ yg (1) ( x ; ξ (2) ) > 1] ≤ exp( − Ω( k )) .
Let I x
= I [ yg (1) ( x ; ξ (2) ) ≤ 1] .
[Equation]
[Equation]
Let T b 1 :=
E ( x,y ) ∼D ,ξ (2) (cid:8) α
y y
I [ ⟨ w (1) , x ⟩ + b (1) + ξ (2) ∈ (0 , 1)] (cid:9) .
We have
[Equation]
[Equation]
So it is sufficient to bound T b
For simplicity, we use q as a shorthand for q (1) i .
[Equation]
Let
[Equation]
Similar as in Lemma 20, we have | I b −
I ′ b | ≤ ϵ e2 .
Then by Lemma 39,
[Equation]
Furthermore,
[Equation]
Therefore, | T b 1 | ≤ O ( ϵ e2 ) and the statement follows.
Lemma 47.
Under the same assumptions as in Lemma 45,
[Equation]
where | T a | = O (cid:16) η (1) ˜ σ γp min (cid:17) + exp( − Ω( k ))
poly (cid:16) dDp min (cid:17) .
Proof of Lemma 47.
Consider one neuron index i and omit the subscript i in the parameters.
By
Lemma 44, Pr[ yg (1) ( x ; ξ (2) ) > 1] ≤ exp( − Θ( k )) .
Let I x
= I [ yg (1) ( x ; ξ ( [Equation] [Equation]
[Equation]
Let T a 1 :=
E ( x,y ) ∼D (cid:8)
α y y E ξ (2) σ ( ⟨ w (1) , x ⟩ + b (1) + ξ (2) ) (cid:9) .
We have
[Equation]
a
[Equation]
a
[Equation]
[Equation]
[Equation]
So it is sufficient to bound T a 1 .
For simplicity, we use q as a shorthand for q (1) i .
Let ϕ ′
A be an independent copy of ϕ A , ϕ ′ be the vector obtained by replacing in ϕ the entries ϕ A with ϕ ′
A
[Equation]
and let x ′ =
Mϕ ′ + ζ and its label is y ′ .
Then
[Equation]
[Equation]
[Equation]
where the fourth step follows from that σ is 1-Lipschitz, and the second to the last line from Lemmaand that (cid:12)(cid:12) ⟨ w (1) , Mϕ ⟩ (cid:12)(cid:12) ≤ ∥ w (1) ∥ ∞ (cid:112) d
∥ Mϕ ∥ 22 .
With the above lemmas about the gradients, we are now ready to show that at the end of the second step, we get a good set of features for accurate prediction.
Lemma 48.
Set
[Equation]
Fix δ ∈ (0 , O (1 /k 3 )) .
If Dµ/ √ d ≤ 1 / 16 , σ ζ ˜ σ = O (1) , σ 2 ζ d/ ˜ σ 2 =
O
(1) , k = Ω (cid:16) log 2 (cid:16) Dmdδγp min (cid:17)(cid:17) , and m ≥ max { Ω( k 4 ) , D, d } , then with probability at least 1 − δ over the initialization, there exist ˜ a i ’s such that ˜ g ( x ) :
= (cid:80) 2 mi =1 ˜ a i σ ( ⟨ w (2) i , x ⟩ + b (2) i ) satisfies L D (˜ g ) ≤ exp( − Ω( k )) .
Furthermore, ∥ ˜ a ∥
= O ( m/k ) , ∥ ˜ a ∥ ∞
=
O ( k 5 /m ) , and ∥ ˜ a [-(cid:107) 
-]∥ 22 =
O ( k 9 /m ) .
Finally, ∥ a (2) ∥ ∞ =
O (cid:0) 1 km 2 (cid:1) , ∥ w (2) i ∥ 2 = O (˜ σ/k ) , and | b (2) i
[-(cid:107) 
 = O (˜ σ/k ) , and | b (2) i-] | =
O (1 /k 2 ) for all i ∈ [2 m ] .
Proof of Lemma 48.
By Lemma 35, there exists a network g ∗ ( x )
= (cid:80) 3( k [Equation] [Equation] ℓ =
[Equation]
) satisfying
[Equation]
Furthermore, the number of neurons n = 3( k + 1) , | a ∗ i | ≤ 64 k, 1 / (64 k ) ≤ | b ∗ i | ≤ 1 / 4 , w ∗ i = ˜
σ (cid:80)
j ∈
A M j / (8 k ) , and |⟨ w ∗ i , x ⟩ + b ∗ i | ≤ 1 for any i ∈ [ n ] and ( x, y ) ∼ D .
Now we fix an ℓ , and show that with high probability there is a neuron in g (2) that can approximate the ℓ -th neuron in g ∗ .
With probability ≥ 1 − exp( − Ω(max { 2 p o
( D − k ) , k } )) , among all j ̸∈ A , we have that at most 2 p
o
( D − k )
+ k of ϕ j are (1 − p o ) /
˜ σ , while the others are − p o / ˜ σ .
With probability ≥ 1 − ( d + D ) exp( − Ω( k )) over ζ , for any j , | ζ j | ≤ O ( σ ζ (cid:112) log( k )) and |⟨ ζ, D ℓ ⟩| ≤ O ( σ ζ (cid:112) log( k )) . Below we consider data points with ϕ and ζ satisfying these.
By Lemma 37, with probability 1 − 2 δ over w (0) i ’s, they are all in G w ( δ ) ; with probability 1 − δ over a (0)
i ’s, they are all in G a ( δ ) ; with probability 1 − δ over b (0)
i ’s, they are all in G b ( δ ) .
Under these events, by Lemma 43, Lemma 45 and 46, for any neuron i ∈ [2 m ] , we have
[Equation]
where
• if j ∈ A , then | T j − βγ/ ˜ σ | ≤ ϵ w 1 : = O ( ϵ e2 / ˜ σ + η (1) /σ (2) ξ + η (1) | a (0)
i | ϵ e / (˜ σσ (2) ξ )) , where β ∈ [Ω(1) , 1] and depends only on w (0) i , b (0) i ; • if j ̸∈ A ,
then | T j | ≤ ϵ w 2 :=
1˜ σ exp( − Ω( k ))
+ O ( σ 2 ϕ ˜ σϵ e2 )
; • | ν j | ≤
ϵ ν := O (cid:18) η (1) σ ζ γσ (2) ξ (cid:19) + exp( − Ω( k )) .
• | T b | ≤ ϵ b := exp( − Ω( k ))
+ O ( ϵ e2 ) .
Given the initialization, with probability Ω(1) over b (0) i , we have
[Equation]
Finally, since 8 k | b ∗ ℓ | βγ | b (0) i | ˜ σ 2 ∈ [Ω( k 2 γ/ ˜ σ 2 ) , O ( k 3 γ/ ˜ σ 2 )] and depends only on w (0) i , b (0) i , we have that for ϵ a = Θ(1 /k 2 ) , with probability Ω( ϵ a ) > δ over a (0) i ,
[Equation]
Let n a = ϵ a m/ 4 .
For the given value of m , by (570)-(573) we have with probability ≥ 1 − 5 δ over the initialization, for each ℓ there is a different set of neurons I ℓ ⊆
[ m ] with | I ℓ | = n a and such that for each i ℓ ∈ I ℓ ,
[Equation]
Now, construct ˜ a such that ˜ a i ℓ = 2 a ∗ ℓ | b ∗ ℓ | | b (0) iℓ | n a for each ℓ and each i ℓ ∈ I ℓ , and ˜ a i = 0 elsewhere.
To show that it gives accurate predictions, we first consider bounding some error terms.
For the given values of parameters, we have
[Equation]
We also have the following useful claims.
[Equation]
Proof of Claim 2.
[Equation]
Claim 3.
[Equation]
Proof of Claim 3.
[Equation]
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
[Equation]
[Table]
is
[Equation]
[Table]
[Equation]
Then by (576)-(580), ϵ ϕ := (cid:32) kϵ
[Equation]
(cid:19) .
We have (cid:12)(cid:12)(cid:12) a (1) i ℓ − a (0) i ℓ (cid:12)(cid:12)(cid:12) =
O ( η (1) σ w (cid:112) log( Dm/δ )) .
So the first term is bounded by
[Equation]
(cid:12)(cid:12)(cid:12)
[Equation]
(606) ≤ O (cid:18) ˜
σ
[Equation]
(cid:19)
[Equation]
O (cid:18)
[Equation]
m (cid:19) .
By Claim 2, the second term is bounded by (cid:12)(cid:12)(cid:12)
[Equation]
1) i ℓ −
[Equation]
(0) i ℓ (cid:12)(cid:12)(cid:12) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
Combining the bounds on the two terms leads to the claim
[Equation]
Claim 4.
[Equation]
Proof of Claim 4.
By Lemma 43 and 46:
[Equation]
We are now ready to show ˜ g is close to 2 g ∗ .
[Equation]
(cid:33)  (cid:12)(cid:12)(cid:12) (cid:12) (cid:12)(cid:12) .
Here the second equation follows from that σ is positive-homogeneous in [0 , 1] , |⟨ w ∗ ℓ , x ⟩ + b ∗ ℓ | ≤ 1 , | b (0) i ℓ | / | b ∗ ℓ | ≤ 1 .
By Claim 3 and 4, the first term is bounded by:
[Equation]
By Claim 2, the second term is bounded by:
[Table]
[Equation]
Then
[Equation]
This guarantees y ˜ g
( x ) ≥ 1 .
Changing the scaling of δ leads to the statement.
Finally, the bounds on ˜ a follow from the above calculation.
The bound on ∥ a (2) ∥ 2 follows from
Lemma 47, and those on ∥ w (2) i ∥ 2 and ∥ b (2) i ∥ 2 follow from (570)(571) and the bounds on a (1) i and b (1) i in Lemma 43.
F.
C LASSIFIER L EARNING S TAGE AND M AIN T HEOREM
Once we have a good set of features in Lemma 48, we can follow exactly the same argument as in
Section B.6 and B.7 for the simplified setting, and arrive at the main theorem for the general setting:
Theorem 49 (Restatement of Theorem 34) .
Set
[Equation]
For any δ ∈ (0 , O (1 /k 3 )) , if µ ≤ O ( √ d/D ) , σ ζ ≤ O (min { 1 / ˜ σ,
[Equation]
, m ≥ max { Ω( k 4 ) , D, d } , then we have for any D ∈ F Ξ , with probability at least 1 − δ , there exists t ∈
[ T ] such that
[Equation]
Consequently, for any ϵ ∈ (0 , 1) , if T = m 4 / 3 , and m ≥ max { Ω( k 12 /ϵ 3 / 2 ) , D } , then
Pr[sign( g (
[Equation]
[Equation]
[-We train our model on Binary CIFAR
-]