Neural Networks Efficiently Learn Low-Dimensional Representations with SGD
Abstract
We study the problem of training a two-layer neural network (NN) of arbitrary width using stochastic gradient descent (SGD) where the input x ∈ R d is Gaussian and
the target y ∈ R follows a multiple-index model, i.e., y = g
( ⟨ u 1 , x ⟩ , . . . , ⟨ u k , x ⟩ ) with a noisy link function g .
We prove that the first-layer weights of the NN converge to the k -dimensional principal subspace spanned by the vectors u 1 , . . .
, u k of the true model, when online SGD with weight decay is used for training.
This phenomenon has several important consequences when k ≪ d .
First, by employing uniform convergence on this smaller subspace, we establish a generalization error bound of O ( (cid:112)
kd/T ) after T iterations of SGD, which is independent of the width of the NN.
We further demonstrate that, SGD-trained ReLU NNs can learn a single-index target of the form y =
f ( ⟨ u , x ⟩ ) + ϵ by recovering the principal direction, with a sample complexity linear in d (up to log factors), where f is a monotonic function with at most polynomial growth, and ϵ is the noise.
This is in contrast to the known d Ω( p ) sample requirement to learn any degree p polynomial in the kernel regime, and it shows that NNs trained with SGD can outperform the neural tangent kernel at initialization.
1. Introduction
The task of learning an unknown statistical (teacher) model using data is fundamental in many areas of learning theory.
There has been a considerable amount of research dedicated to this task, especially when the trained (student) model is a neural network (NN), providing precise and nonasymptotic guarantees in various settings [1, 2, 7, 8, 23, 34, 62, 66, 74, 75].
As evident from these works, explaining the remarkable learning capabilities of NNs requires arguments beyond the classical learning theory [72].
The connection among NNs and kernel methods has been particularly useful towards this expedition [22, 38].
In particular, a two-layer NN with randomly initialized and untrained weights is an example of a random features model [58], and regression on the second layer captures several interesting phenomena that NNs exhibit in practice [45, 47], e.g. cusp in the learning curve.
However, NNs also inherit favorable characteristics from the optimization procedure [3, 31, 44, 59, 69], which cannot be captured by associating NNs with regression on random features.
Indeed, recent works have established a separation between NNs and kernel methods, relying on the emergence of representation learning as a consequence of gradient-based training [1, 8, 9, 23], which often exhibits a natural bias towards low-complexity models.
A theme that has emerged repeatedly in modern learning theory is the implicit regularization effect provided by the training dynamics [51].
Specifically, [64] has inspired an abundance of recent works focusing on the implicit bias of gradient descent favoring, in some sense, low-complexity models, e.g. by achieving min-norm and/or max-margin solutions despite the lack of any explicit regularization [21, 33, 36, 39, 43, 56].
However, these works mainly consider linear models or unrealistically wide NNs, and the notion of reduced complexity as well as its implications on generalization varies.
A concrete example in this domain is compressiblity and its connection to generalization [5, 65].
Indeed, when a trained NN can be compressed into a smaller NN with similar prediction behavior, the resulting models exhibit similar generalization performance, while the latter is classically linked to better generalization.
In this paper, we demonstrate the emergence of low-complexity structures during the training procedure.
More specifically, we consider training a two-layer student NN with arbitrary width m where the input x ∈ R d is Gaussian and the target y ∈ R follows a multiple-index teacher model, i.e. y = g ( ⟨ u 1 , x ⟩ , . . .
, ⟨ u k , x ⟩ ; ϵ ) with a link function g and a noise ϵ independent of the input.
In this setting, we prove that the first-layer weights trained by online stochastic gradient descent (SGD) with weight decay converge to the k -dimensional subspace span( u 1 , , . . .
, u k ) , which we refer to as the principal subspace .
Our primary focus is the case where the target values depend only on a few important directions along the input, i.e. k ≪ d , which induces a low-dimensional structure on the SGD-trained first-layer weights, whose impact on generalization is profound.
We summarize our contributions as follows
• We show in Theorem 2 that NNs learn low-dimensional representations by proving that the iterates of online SGD on the first layer of a two-layer NN with width m converge to √ mε neighborhood of the principal subspace after O ( d/ε 2 ) iterations, with high probability.
The error tolerance of √ mε is sufficient to guarantee that the risk of SGD iterates and that of its orthogonal projection to the principal subspace are within O ( ε ) distance.
• We demonstrate the impact of learning low-dimensional representations with two applications.
– For a single-index target y = f ( ⟨ u , x ⟩ ) + ϵ with certain link functions f , we prove in Theorem 4 that ReLU networks of width m can learn this target with an SGD-based procedure (Algorithm 1) with number of samples T and an excess risk estimate of ˜ O ( (cid:112) d/T + 1 /m ) , with high probability (see the illustration in Figure 1).
In particular, the sample complexity is (almost) linear in d , even when f is a monotonic polynomial of any (fixed) odd degree p .
– Based on a uniform convergence argument on the principal subspace, we prove in Theoremthat T iterations of SGD will produce a model with generalization error of O ( (cid:112) kd/T ) , with high probability.
Remarkably, this rate is independent of the width m of the NN, even in the case k ≍ d where the target is any function of the input, and not necessarily low-dimensional.
The rest of the paper is organized as follows.
We discuss the notation and the related work in the remainder of this section.
We describe the problem formulation and preliminaries in Section 2.
Our main result on SGD is presented in Section 3.
We discuss two implications of our main theorem in
Section 4, where we provide results on learnability and generalization gap in Sections 4.1 and 4.
respectively.
We finally conclude with a brief discussion in Section 5.
Notation.
For a loss function ℓ : R 2 → R , let ∂ i ℓ and ∂ 2 ij ℓ denote its partial derivatives with respect to i th and j th inputs for i, j ∈ { 1 , 2 } .
For quantities a and b , a ≲ b implies a ≤ Cb for an absolute constant C , and a ≍ b implies both a ≳ b and a ≲ b .
Finally, Unif ( A ) denotes the uniform distribution over a set A and N (0 , I d ) denotes the d -dimensional isotropic Gaussian distribution.
1.1. Related work
Training dynamics of NNs.
Several works have demonstrated learnability in a special case of teacher-student setting where the teacher model is similar to the student NN being trained [16, 42, 73–75].
This setting has also been studied through the lens of loss landscape [61] and optimization over measures [2].
We stress that our results work under misspecification and hold for generic teacher models that are not necessarily NNs with similar architecture to the student.
Additional related works on different scaling regimes of wide networks is provided in Appendix A.
Feature learning with multiple-index teacher models.
The task of learning a target of an unknown low-dimensional function of the input is fundamental in statistics [41].
Several recent works in learning theory literature have also focused on this problem, with an aim to demonstrate NNs can learn useful feature representations, outperforming kernel methods [12, 32].
In particular, [1] studies the necessary and sufficient conditions for learning with sample complexity linear in d with inputs on the hypercube, in the mean-field limit.
Closer to our setting are the recent works [8, 9, 23] which demonstrate a clear separation between NNs and kernel methods, leveraging the effect of representation learning.
However, their analysis considers a single (full) gradient step on the first-layer weights followed by training the second-layer parameters.
In contrast, in our learnability result, we consider training both layers with SGD, which induces essentially different learning dynamics.
Generalization bounds for SGD.
A popular algorithm-dependent approach for studying generalization is through algorithmic stability [14, 15, 28], which has been used to study the generalization behavior of gradient-based methods in various settings [11, 27, 37, 40].
Other approaches include studying the low-dimensional structure of the trajectory [55, 63] or the invariant measure of continuous-time approximations of SGD [17], and employing information-theoretic tools [50].
Among these, [10] also shows that SGD yields compressible networks; however, they assume the mean-field approximation holds and the SGD iterates converge to a heavy-tailed distribution.
2. Preliminaries: Neural Networks and the Principal Subspace
For an input x ∈ R d , we consider training a two-layer neural network (NN) with m neurons ˆ y ( x ; W , a , b ) = m (cid:88) i = a i σ ( ⟨ w i , x ⟩ + b i ) , (2.1) where σ is the activation function, { w i } 1 ≤ i ≤ m are the first-layer weights collected in the rows of the matrix W ∈ R m × d , b ∈ R m is the bias, and a ∈ R m is the second-layer weights.
We assume x ∼ N (0 , I d ) and the target is generated from a multiple-index (teacher) model given by y = g ( ⟨ u 1 , x ⟩ , . . . , ⟨ u k , x ⟩ ; ϵ ) , (
[Equation]
2) for a weakly differentiable link function g : R k +
→ R and a noise ϵ independent of x .
The Gaussian input is a rather standard assumption in the literature, especially in recent works that consider the student-teacher setup; see e.g. [23, 61, 75].
The multiple-index teacher model (2.2) can encode a broad class of input-output relations through the non-linear link function, including a multi-layer fully-connected NN with arbitrary depth and width and weakly differentiable activations.
While our results remain valid regardless of how k and d compare, they are most insightful when k ≪ d ; thus, we specifically consider this regime when interpreting the results.
We also collect the teacher weights { u i } 1 ≤ i ≤ k in the rows of the matrix U ∈ R k × d and use y =
g
( Ux ; ϵ ) for simplicity.
For a given loss function ℓ (ˆ y, y ) , we consider the population and the empirical risks
R ( W , a , b ) =
E [ ℓ (ˆ y ( x ; W , a , b ) , y )] and
[Equation]
where the expectation is over the data distribution.
Similarly, for some τ ≥ 1 , the truncated loss is defined as ℓ τ (ˆ y, y ) :
=
ℓ (ˆ y, y ) ∧ τ with the corresponding risks R τ and ˆ R τ .
We define the L 2regularized population risk with a penalty parameter λ ≥ 0 as
[Equation]
To minimize (2.3), we use stochastic gradient descent (SGD) over the first-layer weights, where we are interested in the convergence of iterates to the principal subspace defined by the teacher weights
[Equation]
Notice that the principal subspace satisfies S ( U ) ⊆
R m × d , and its dimension is mk as opposed to the ambient dimension of md , with any matrix in this subspace having rank at most k .
For any vector v ∈ R d , we let v ∥ denote the orthogonal projection of v onto span( u 1 , . . .
, u k ) and v ⊥ := v − v ∥ .
Similarly, for a matrix W ∈ R m × d , we define W ∥ and W ⊥ by applying the projection to each row.
The smoothness properties of the activation σ play an important role in our analysis.
As such, we consider two scenarios, with different requirements on the loss function.
Assumption 1.A (Smooth activation)
The activation function σ satisfies | σ ( z ) | , | σ ′ ( z ) | , | σ ′′ ( z ) | ≤ 1 for all z ∈ R , the loss is ℓ (ˆ y, y )
=
12 (ˆ y − y ) 2 for simplicity, and y satisfies | y | ≤ K almost surely.
Assumption 1.B (ReLU activation)
The activation function σ is σ ( z ) =
max( z, 0) for z ∈ R .
The loss satisfies 0 ≤ ∂ 21 ℓ (ˆ y, y ) ≤ 1 , | ∂ 1 ℓ (ˆ y, y ) | ≤ 1 , and | ∂ 212 ℓ (ˆ y, y ) | ≤ 1 .
Commonly used activations such as sigmoid and tanh satisfy Assumption 1.A.
For ReLU activation in Assumption 1.B, we choose σ ′ ( z ) =
1 ( z ≥ 0) as its weak derivative.
We highlight that
Assumption 1.B is satisfied by common Lipschitz and convex loss functions such as the Huber loss
[Equation]
as well as the logistic loss ℓ L (ˆ y, y ) :
= log(1 + e − ˆ yy ) , up to appropriate scaling constants.
3. Convergence of Stochastic Gradient Descent
We now consider stochastic gradient descent (SGD) in the online setting where at each iteration t , we have access to a new data point ( x ( t ) , y ( t ) ) drawn independently of the previous samples from the same distribution.
We update the first-layer weights W t with a time varying step size η t and weight decay, according to the update rule
[Equation]
The above algorithm can be used to minimize the population risk (2.3) in practice [57], even in certain non-convex landscapes [70].
We will use the following initialization for SGD
Assumption 2 (Initialization) For all 1 ≤ i, j ≤ m , we initialize the NN weights and biases with √ dW 0 ij iid ∼ N (0 , 1) , ma 0 j iid ∼ Unif ([ − 1 , 1]) , and b 0 j iid ∼ Unif ( {− 1 , 1 } ) .
This initialization is standard in the mean-field regime as it allows feature learning. However, we only use it to simplify the exposition.
Indeed, we can initialize W and a with any scheme that guarantees ∥ W ∥ F ≲ √ m and ∥ a ∥ ∞ ≲ m − 1 with high probability.
Further, initialization of b mostly matters in the analysis of ReLU activation.
Next, we show that the population gradient admits a certain decomposition which plays a central role in our analysis.
For smooth activations, the below result is a remarkable consequence of Stein’s lemma, which provides a certain alignment between the true statistical model (teacher) and the model being trained (student), which has profound impact on the learning dynamics.
We generalize this result for ReLU through a sequence of smooth approximations (see Appendix C.1 for details).
Lemma 1 Under Assumption 1.A or 1.B , the gradient of the population risk can be written as
[Equation]
for some H ( W ) ∈ R m × d and
D ( W ) ∈ R k × d (with explicit forms provided in Appendix C.1).
The following result, proved in Appendix D via a recursion on the moment generating function of ∥ W t ⊥ ∥ F , demonstrates the algorithmic implications of Lemma 1, and shows that the iterates converges to the principal subspace.
Theorem 2 Consider running T SGD iterations (3.1) when the activation and loss either satisfy
Assumption 1.A or Assumption 1.B , and the initialization satisfies Assumption 2.
Let ζ : =
E [ | y | ]+ under Assumption 1.A and ζ := 2 (cid:112)
/eπ under Assumption 1.B. Choose the decreasing step size η t = m 2( t + t ∗ )
+1 γ ( t + t ∗ +1) 2 , ˜ λ ≥ γ + ζ , and t ∗ ≍ ˜ λγ for any γ > 0 .
Then, for λ = ˜ λm , with probability at least 1 − δ ,
[Equation]
whenever m ≳ log(1 /δ ) and T ≳
[Equation]
Remark
The above result provides a convergence rate of O ( (cid:112) d/T ) for the normalized quantity ∥ W T ⊥ ∥ F / √ m .
We justify this exposition by noting that since the risk is locally Lipschitz with a constant O (1 / √ m ) , the condition ∥ W T ⊥ ∥ F ≲ √ mε is sufficient to guarantee that the risk of W T and that of its orthogonal projection to the principal subspace W T ∥ are within O ( ε ) distance.
The above result states that, with a number of samples linear in the input dimension d , SGD is able to learn (approximately) low-dimensional weights, exhibiting an implicit bias towards lowcomplexity models.
The interplay between two forces is in effect here.
The most important one is the linear relationship between the first-layer weights and the input in both student and teacher models together with the input distribution.
The alignment described in Lemma 1 yields sparsified weights in a basis defined by the teacher network, effectively reducing the input dimension from d to k .
The second force is the explicit L 2 -regularization.
We emphasize that L 2 -regularization does not play the main role in this sparsification; even though it may provide shrinkage to zero, Lpenalty will in general produce non-sparse solutions.
However, it is still required as the explicit Lregularization ensures that SGD avoids critical points outside of the principal subspace.
We refer to
Section 5 for further discussion.
Notice that while Theorem 2 does not have any implications on the convergence behavior of the orthogonal projection W ∥ , in the next section, we show that the implied low-dimensional structure is sufficient to provide guarantees on the generalization error and learnability of SGD.
Moreover, we showcase by a non-convex example in Appendix F that the regularization in this Theorem does not in general imply (strong) convexity in the population landscape.
4. Implications of Low-Dimensionality 4.1. Learning Single-Index Targets
An essential characteristic of NNs is their ability to learn useful representations, which allows them to adapt to the underlying misspecified statistical model.
Although this fundamental property has been the guiding principle in all empirical studies, it was mathematically proven only recently for gradient-based training [1, 8, 9, 23, 29]; see also a survey of prior works in [46].
Our results in the previous section are in the same spirit, establishing the convergence of SGD to the principal subspace which is indeed a span of useful directions associated with the target function being learned.
As such, we leverage the learned low-dimensional representations to demonstrate that SGD is capable of learning a target function of the form y = f
( ⟨ u , x ⟩ ) + ϵ with a number of samples linear in d (up to logarithmic factors).
For simplicity, we work with the Huber loss below; however, our analysis can accommodate any Lipschitz and convex loss at the expense of a more detailed analysis.
Algorithm
Input: a 0 , b 0 ∈ R m , W 0 ∈
R m × d , { ( x ( t ) , y ( t ) ) } 0 ≤ t ≤ T − 1 , ( η t ) t ≥ 0 , ( η ′ t ) t ≥ 0 , λ , λ ′ , ∆ .
1: for t = 0 , ..., T − 1 do 2:
W t +1 =
(1 − η t λ ) W t − η t ∇ W ℓ (ˆ y ( x ( t ) ; W t , a 0 , b 0 ) , y ( t ) ) .
3: end for 4: Let b j ∼ Unif ( − ∆ , ∆) for 1 ≤ j ≤ m .
5: for t = 0 , ..., T ′ − 1 do 6:
Sample i t ∼
T − 1 } .
7: a t +1 = (1 − η ′ t λ ′ )
a t − η ′ t ∇ a ℓ (ˆ y ( x ( i t ) ; W T , a t , b ) , y ( i t ) ) 8: end for 9: return ( W T , a T ′ , b ) .
The following Theorem states the learning guarantee of Algorithm 1, and is proved in Appendix E.2 via recovering the principal direction according to Theorem 2, along with the generalization gap of Theorem 5, and a universal approximation argument similar to [23].
We highlight that the recent works [8, 9, 23] perform only one gradient step on the first layer weights, whereas in
Algorithm 1, we train the entire NN with SGD.
Theorem 4 Suppose that the data is from a single-index model y =
f
( ⟨ u , x ⟩ )+ ϵ with a monotone differentiable f and 1 -sub-Gaussian noise ϵ , and Assumption 1.B holds.
Further, let ∥ u ∥
=
| f (0) | < 1 , and consider the Huber loss (2.4) for simplicity.
Consider running Algorithm 1 with the initialization 0 < a 0 j = a ≲ 1 /m , 0 < b 0 j = b ≲ 1 , and w
j = w 0 ∼ N (0 , 1 d I d ) for all j with the hyper-parameters λ = ˜ λm
= γm + 2 ab (cid:113) 2 eπ for any γ ≍ 1 , η t = m 2( t ∗ + t )+1 γ ( t ∗ + t +1) 2 with t ∗ ≍ γ − 1 , η ′ t = 2 t +1 λ
′ ( t +1) 2 , and ∆ ≍ (cid:112) log( T/δ ) .
Then, for T ≳
( d + log( 1 δ ))
∨ ( ˜ λγd log( mδ )) , some λ ′ >
(see (E.7) ), and sufficiently large T ′ (see (E.8) ), with probability at least 1 − δ ,
[Equation]
where ∆ ∗ := ∆ sup | z | ≲ ∆ | f ′′ ( z ) | which is poly(log( T/δ )) when f ′′ has at most polynomial growth.
Notice that T determines the sample complexity of running Algorithm 1 as it only requires T data samples.
As such, the above result implies that a ReLU NN trained with SGD can learn any monotone polynomial with a sample complexity linear in the input dimension d , up to logarithmic factors.
In comparison, [23] considers training the first-layer weights with one gradient descent step with a carefully chosen weight decay, and obtains a sample complexity of d 2 to learn any unknown degree p polynomial multiple-index target.
Note that learning any degree p polynomial using rotationally invariant kernels requires d Ω( p ) samples for a variety of input distributions including isotropic Gaussian [24]; thus, our result shows that SGD is able to efficiently learn a target function where kernel methods cannot.
It is worth emphasizing that in Theorem 4, the width of the network m grows with poly log( T/δ ) , which is in contrast to the neural tangent kernel (NTK) regime where m is required to grow with poly( T ) .
On a separate note, [18] proposes a method that can train NNs to learn a polynomial of a few directions with a number of samples linear in the input dimension d ; yet, the proposed algorithm is not a simple variant of SGD and requires a non-trivial warm-start initialization.
We refer to [26] for a review of classical literature on learning single and multiple-index models.
Specifically, they provide a procedure with sample complexity at least linear in d to recover the principal direction of the single-index model with Gaussian inputs.
4.2. Generalization Gap
For a given learning algorithm, the gap between its empirical and population risks is termed as the generalization gap (not to be confused with excess risk), and establishing convergence estimates for this quantity is a fundamental problem in learning theory.
Classical results rely on uniform convergence over the feasible domain containing the weights; thus, they apply to any learning algorithm including SGD [52].
However, these bounds often diverge with the width of the NN, yielding vacuous estimates in the overparameterized regime [72].
To alleviate this, recent works considered establishing estimates for a specific learning algorithm; see e.g. [37, 55, 64, 71].
Here, we are interested in deriving an estimate for the generalization gap over the SGD-trained first-layer weights, which holds uniformly over the second layer weights and biases.
More specifically, we study, after T iterations of SGD (3.1) initialized with ( W 0 , a 0 , b 0 ) , the following quantity
E ( W T ):= sup S R τ ( W T , a , b ) − ˆ R τ ( W T , a , b ) with S :
= { a , b ∈ R m : ∥ a ∥ 2 ≤ r a √ m , ∥ b ∥ ∞ ≤ r b } , where the scaling ensures ˆ y =
O (1)
when ∥ w j ∥ 2 ≍ 1 , which is the setting considered in
Theorem 4.
We state the following bound on E ( W T ) ; the proof is provided in Appendix E.1, and it is based on a covering argument over the smaller dimensional principal subspace implied by
Theorem 2.
Theorem
Consider the setting of Theorem 2.
For any δ > 0 , if T ≳ ( d + log(1 /δ )) ∨ ( ˜ λγd log( m/δ )) , then with probability at least 1 − δ ,
[Equation]
The above bound is independent of the width m of the NN, and only grows with the dimension of the input space d and that of the principal subspace k ; thus, producing non-vacuous estimates in the overparametrized regime where m is large.
Further, the bound is stable in the number of SGD iterations T , that is, it converges to zero as T → ∞ .
We remark that generalization bounds that rely on algorithmic stability are optimal for strongly convex objectives [37]; yet, they lead to unstable diverging bounds in non-convex settings as T → ∞ .
As such, these techniques often require early stopping , which is clearly not needed in our result.
5. Conclusion
We studied the dynamics of SGD with weight decay on two-layer NNs, and proved that under a multiple-index teacher model, the first-layer weights converge to the principal subspace, i.e. the span of the weights of the teacher.
This phenomenon is of particular interest when the target depends on the input along a few important directions.
In this setting, we proved novel generalization bounds for SGD via uniform convergence on the lowdimensional principal subspace.
Further, we proved that two-layer ReLU networks can learn a single-index target with a monotone link that has at most polynomial growth, using online SGD, with a number of samples almost linear in d .
Thus, as an implication of low-dimensionality, we established a separation between kernel methods and trained NNs where the former suffers from the curse of dimensionality.
Two principal forces are responsible for the emergence of the low-dimensional structure.
The main one is the linear interaction between the Gaussian input and the first-layer weights in both student and teacher models.
The secondary one is the weight decay which allows SGD to avoid critical points outside of principal subspace.
Figure 2 shows the convergence behavior in absence of weight decay.
Understanding more precisely the range of λ that implies convergence to the principal subspace, as well as investigating the possibility of learning multiple-index models using this convergence, are left as important directions for future studies.
Appendix A. Further Related Work
Two scaling regimes of neural networks have seen a surge of recent interest.
In the regime of lazy training [22], the parameters hardly move from initialization and the NN does not learn useful features, behaving like a kernel method [4, 6, 25, 38, 54].
However, many works have shown that deep learning is more powerful than kernel models [30, 32, 69], establishing a clear separation between them; thus, several important characteristics of NNs cannot be captured in the lazy training regime [31].
In the other scaling regime, gradient descent on infinitely wide NNs reduces to Wasserstein gradient flow, which is known as the mean-field regime where feature learning is possible [19, 20, 48, 53, 60].
However, these results mostly hold for infinite or very wide NNs, and quantitative guarantees are difficult to obtain in this regime.
The setting we consider in the current paper is different from both of these regimes, e.g. we may allow for NNs of arbitrary width which are not necessarily excessively overparameterized; yet, we still use the mean-field scaling when initializing the weights of the student NN.
On a separate note, we point out that our learnability result of Thoerem 4 is consistent with [13]; they establish a sharp sample complexity of ˜ O
( d 1 ∨ ( I− 2) ) to learn a target with online SGD using the same activation f in the student network, where I is the information exponent ( I = 1 in the above case due to the monotonocity of f ).
Despite assuming the link function f is known, we highlight that their setting covers I ≥ 1 , whereas Theorem 4 is a proof concept to demonstrate the learnability implications of convergence to the principal subspace, even when f is unknown.
Appendix B. Additional Notations
For vectors v and u , we use ⟨ v , u ⟩ and v ◦ u to denote their Euclidean inner product and the element-wise product, and we use ∥ v ∥ p and diag( v ) to denote the L p -norm and the diagonal matrix whose diagonal entries are v .
For matrices V and W , we use ⟨ V , W ⟩ F , ∥ V ∥ F , and ∥ V ∥to denote the Frobenius inner product, Frobenius norm, and the operator norm, respectively.
For an activation function σ : R → R , σ ′ and σ ′′ denote its first and second (weak) derivatives, which are applied element-wise for vector inputs.
The symbol ∇ · is reserved for the gradient operator, and we frequently use ∇ ℓ to denote ∇ W ℓ when it is clear from the context.
In the appendix, we will prove the statements of the main text in a more general formulation.
In particular, for smooth activations, we assume sup | σ ′ | ≤ β 1 and sup | σ ′′ | ≤ β 2 for some β 1 , β 2 ∈
R + , and we denote sup | σ | ≤ β 0 , β 0 ∈ (0 , ∞ ] .
We will consider the following general case for the bias vector b ∈ R m : b j iid ∼ D b , such that | b j | ≥ b ∗ > 0 , for some b ∗ > 0 .
This setting clearly covers the case of b j = ± 1 from the initialization of Assumption 2.
Throughout the appendix, C will denote a generic positive absolute constant (e.g. 10), whose value may change from line to line.
We use the shorthand notation σ a , b ( W x ) to denote a ◦ σ ( W x + b ) , and similarly for σ ′ a , b ( W x ) and σ ′′ a , b ( W x ) .
We use the notations vec( A ) ∈ R mn for the vectorized representation of a matrix
A ∈ R m × n , and A ⊗ B for the Kronecker product of two matrices A ∈ R m × n and B ∈ R p × q ; we recall that the Kronecker product is an mp × nq block matrix comprised of m × n blocks of shape p × q , where block ( i, j ) is given by A ij B .
Appendix C. Population Gradient Decomposition C.1. Proof of Lemma
In what follows
[Equation]
∇ ⊤ is the Jacobian matrix and ∇ is the transpose of Jacobian for vector valued functions, which is the same as gradient for real-valued functions
[Equation]
When σ is twice differentiable (Assumption 1.A)
[Equation]
standard matrix calculations yield
[Equation]
) where (a) follows from the dominated convergence theorem and (b) follows from the Stein’s lemma, and ∇ g ϵ is the weak derivative of g ϵ w.r.t.
its inputs.
Recall that Stein’s lemma (Gaussian integration by parts) states that for x ∼ N (0 , I d ) and weakly differentiable f : R d → R , we have E [ ∇ f ( x )] =
E [ x f ( x )] .
Combining the above calculations with the gradient of the regularization term, with
[Equation]
where ∇ g ϵ is the weak derivative of g ϵ w.r.t.
its inputs, and
[Equation]
the proof is complete for smooth activations.
For ReLU activations and ℓ satisfying Assumption 1.B, we introduce the following smooth approximation
[Equation]
Then we have
[Equation]
As σ ′′ τ ≥ 0 , the critical step is to show lim ι →∞ E [ σ ′′ ι ( ⟨ w , x ⟩ + b )] < ∞ , uniformly for all w .
Let z = ⟨ w , x ⟩ + b .
Then z ∼ N ( b, ∥ w ∥ 22 ) , and
[Equation]
where (a) follows from the Gaussian tail bound 1 − Φ( x ) ≤ e − x 2 / 2 √ 2 πx , where Φ is the standard Gaussian
CDF; (b) holds for large enough ι ; and (c) holds by considering supremum over ∥ w ∥ 2 .
Thus
E [ σ ′′ ι ( ⟨ w j , x ⟩ + b j )] ≤ 2 | b j | (cid:113) 2 eπ
and consequently,
[Equation]
where b ∗ =
min 1 ≤ j ≤ m | b j | .
Moreover, as σ ′ ι ( W x + b ) converges a.s. (i.e. except when ⟨ w j , x ⟩ + b j = 0 for some j ) to σ ′ ( W x + b ) , by the dominated convergence theorem,
[Equation]
We can immediately observe from the dominated convergence theorem that D ι ( W ) → D ( W ) as ι → ∞ with D ( W ) given in (C.2).
Moreover, we let H ( W ) =
lim
ι →∞ H ι ( W ) , and observe that
[Equation]
This finishes the proof of Lemma 1.
In the case of smooth activations (Assumption 1.A), the following bounds will be useful.
Lemma 6 Let R ( W ) :
= E [ ℓ (ˆ y ( x ; W , a , b ) , y )] be the unregularized population risk.
Under
Assumption 1.A we have
[Equation]
Proof Assumption 1.A requires ℓ (ˆ y, y )
=
12 (ˆ y − y ) 2 .
Hence by definition of H ,
[Equation]
The first term is positive semi-definite and it can be easily bounded:
[Equation]
for an arbitrary vector v ∈ R m .
For the second term, we have
[Equation]
and E [ | ˆ y ( x ; W ) − y | ] ≤ (cid:112) 2 R ( W ) by Jensen’s inequality.
Appendix D. Proofs of Section
We begin by characterizing the tail behavior of the stochastic gradient noise in the SGD updates (3.1) through the following lemma.
Lemma
For any fixed W ∈ R m × k , let
[Equation]
denote the zero-mean stochastic noise in the gradient of the loss function ℓ when ( x , y ) ∼ P , and recall that
[Equation]
Suppose sup ˆ y,y | ∂ 1 ℓ (ˆ y, y ) | ≤ κ .
Then for any V ∈ R m × d , the zero-mean random variable ⟨ V , Γ ⟩ F is Cβ 1 κ ∥
a ∥ 2 ∥ V ∥ F -sub-Gaussian.
Proof We use the shorthand notation ∇ ℓ := ∇ W ℓ ( W x , y ) and ∇ R :=
∇ W R ( W ) .
We compute the following
[Equation]
where (a) and (b) follow from the Minkowski and Jensen inequalities respectively.
Furthermore, we have
[Equation]
where the last inequality follows from Gaussianity of V x and Lemma 25.
Hence
[Equation]
Invoking Lemma 21 implies sub-Gaussianity of ⟨ V , ∇ ℓ − ∇ R ⟩ F and completes the proof.
We proceed by presenting a lemma which constitutes the main part of the proof of Theoremvia establishing a recursive bound on the moment generating function (MGF) of ∥ W t ⊥ ∥ 2 F , which will in turn be used to prove high probability statements for ∥ W t ⊥ ∥ 2 F .
Lemma 8 Consider running the iterates of SGD (3.1) , under either Assumption 1.A or 1.B , with step size sequence { η t } t ≥ 0 that is either constant η t = η or decreasing η t = m 2( t ∗ + t )
+1 γ ( t ∗ + t +1) 2 (cf.
[35, Theorem 3.2]) .
Let κ :=
sup | ∂ 1 ℓ (ˆ y, y ) | , κ := β 1 ∥ a ∥ 2 κ , ζ := β 2 ( E [ | y | ] + β 0 ∥ a ∥ 1 ) under
Assumption 1.A and ζ := 2 (cid:112) 2 / ( eπ ) under Assumption 1.B, and finally ˜ ϱ := λ + β 1 ∥ a ∥ 22 + ζ ∥ a ∥ ∞ .
Suppose η 0 ≲ ˜ ϱ − 1 .
Let F t denote the sigma algebra generated by { W j } t j =0 , and let { A t } t ≥be a sequence of decreasing events (i.e. A t +1 ⊆
A t ), such that A t ∈ F t and on A t we have
H ( W t )
+ λ
I m ⪰ γm I m .
Then, for every t ≥ 0 , with probability at least P ( A t ) − δ ,
[Equation]
Proof Let F t denote the sigma algebra generated by { W j } tj =0 .
Recall from Lemma 7 that we define
[Equation]
with
[Equation]
Then for the SGD updates we have
[Equation]
By projecting the iterates onto the orthogonal complement of the principal subspace
[Equation]
Let M t :
= I m − η t ( H ( W t )
+ λ
I m ) .
Then
[Equation]
by observing that 1 A t +1 ≤ 1 A t , for any 0 ≤ s ≲ γ mη t κwe have
[Equation]
where the last inequality follows from the Cauchy-Schwartz inequality for conditional expectation.
Moreover, it is straightforward to observe that ∥∇ ℓ (ˆ y ( x ( t ) ; W t ) , y ( t )
[Equation]
∥
[Equation]
≤ κ 2 ∥ x ∥ 22 , hence
[Equation]
Note that by Jensen’s inequality
[Equation]
Consequently
[Equation]
where the second inequality follows from Lemma 26 for 4 sη 2 t κ 2 ≤ 1 / 4 .
Since s ≲ γ mη t κ 2 , in order to satisfy the condition of Lemma 26 we need to ensure η t γ/m ≲ , which is guaranteed by our η t ˜ ϱ ≲ 1 assumption for a suitably small absolute constant, as γ/m ≤ λ ≤ ˜ ϱ .
Next, we bound the last term in (D.2).
Let V := − 4 sη t M t W t ⊥ .
Then by Lemma 7 we have
[Equation]
Putting things back together in (D.2) and using the tower property of expectation, we have
[Equation]
Next, we bound ∥ M t ∥ 2 .
By definition of A t , we can already ensure H ( W t ) ⪰ γm
I m in (D.3).
Recall the definition of H ( W t )
[Equation]
Notice that 0 ≤ ∂ 21 ℓ (ˆ y ( x ; W ) , y ) ≤ 1 under either Assumption 1.A or Assumption 1.B. Moreover we have, | ∂ 1 ℓ (ˆ y, y ) | ≤ κ .
Thus,
[Equation]
Therefore, 0 ⪯
[Equation]
As a result ∥ M t ∥ 2 ≤ 1 − η t γm .
Combined with (D.3) we have
E (cid:2)
A t +1 exp (cid:0) s ∥ W t +1 ⊥ ∥ 2 F (cid:1) | F
[Equation]
≤ E
[Equation]
1 A t exp (cid:0) s
[Equation]
1 + Csη 2 t
[Equation] [Equation]
1 − η t γm ) 2 ∥ W t ⊥ ∥ 2 F + Csη 2 t dκ 2 (cid:1) | F 0 (cid:3) ≤
[Equation]
where the second inequality holds by the fact that Csη 2 t κ 2 ≤ η t γ/m , which in turn holds when a small enough absolute constant is chosen in 0 ≤ s ≲ γ mη t κ 2 .
Also notice that for decreasing step size,
[Equation]
(and the above holds trivially for constant step size), thus when s ≤ C 1 γ η t κ 2 for some absolute constant
C 1 , we have s (1 − η t γ ) ≤ C 1 γ η t − 1 κ 2 with the same absolute constant.
Hence we are allowed to expand the recursion (D.4), which implies
[Equation]
for all 0 ≤ s ≲ γ mη t − 1 κ 2 .
Moreover, direct calculation implies that with both constant and decreasing step sizes of Lemma 8, we have (cid:80)
[Equation] [Equation]
=
η 2 i
[Equation]
−
[Equation]
≤
[Equation]
(with C = 1 for constant step size).
Thus, for all 0 ≤ s ≲
[Equation]
Finally, we can apply a Chernoff bound to obtain
[Equation]
Cmη t κ 2 d γ
[Equation]

[Equation]

By choosing ε =
[Equation] [Equation]
γ
[Equation]
and the largest possible s ≲ γ mη t κ 2 , we obtain
[Equation]
Taking another expectation to remove conditioning on initialization completes the proof.
The proof of Theorem 2 for decreasing step size follows by a direct computation of the quantities in Lemma 8 and is presented below.
We remark that proving the same result with a constant step size is essentially similar.
D.1. Proof of Theorem
This part is directly implied by Lemma 8.
The following argument holds on an event where ∥ W ∥ F ≲ √ m , which happens with probability at least 1 − O
( δ ) .
In order to see this connection, we will first present an improved statement over Lemma 6 for the case of smooth activations.
Recall the definition of H ( W ) for the squared error loss ℓ (ˆ y, y ) =
(ˆ y − y ) 2,
[Equation]
Notice that under Assumption 1.A we have | ˆ y | ≤ β 0 ∥ a ∥ 1 .
Then basic matrix algebra similar to that of Lemma 6 along with the triangle inequality shows
[Equation]
( β 0 ∥ a ∥ 1 + E
[ | y | ]) (cid:1) I m .
Therefore, with λ ≥ γ/m + β 2 ∥ a ∥ ∞ ( ∥ a ∥ 1 β 0 + E
[ | y | ]) , we have H ( W )
+ λ
I m ⪰ γ/m I m for all W .
In addition, | ∂ 1 ℓ (ˆ y, y ) | ≤ β 0 ∥ a ∥ 1 + K by the triangle inequality.
Thus we can invoke
Lemma 8 with η t = mγ (cid:16) 1 − ( t ∗ + t ) 2 ( t ∗ + t +1) 2 (cid:17) , κ = β
∥
a ∥ 1 + K , ζ = β 2 κ and 1 A t = 1 .
Recall that in the statement of the theorem, β 0 =
β 1 = β 2 = 1 , K ≲ 1 , ∥ a ∥ ∞ ≤ 1 /m , ∥ a ∥ 2 ≤ 1 / √ m , and ∥ a ∥ 1 ≤ 1 , hence ˜ ϱ
≍ λ , and with t ∗ ≍ ˜ λγ we can guarantee η t λ ≲ 1 .
As the step size condition of
Lemma 8 is satisfied, the desired result follows.
Similarly, for ReLU we have | ∂ 1 ℓ (ˆ y, y ) | ≤ 1 by Assumption 1.B, and for λ ≥ γ/m + 2 ∥ a ∥ ∞ b ∗ (cid:113)eπ , we have H ( W )
+ λ
I m ⪰ γ/m I m .
Hence this time, we can invoke Lemma 8 with the same decreasing η t and 1 A t = 1 , κ = 1 , and ζ = 2 b ∗ (cid:113) 2 eπ (recall b ∗ = 1 in the statement of the theorem).
Appendix E. Proofs of SectionE.1. Proof of Theorem
As our arguments are based on the Rademacher complexity of a two-layer neural network, we require the knowledge of the norm of W t .
We prove a high probability bound for this norm in the following lemma.
Lemma 9 Under Assumption 1.A or 1.B with either decreasing or constant step size as in Theorem 2, let κ = sup ˆ
y,y
| ∂ 1 ℓ (ˆ y, y ) | < ∞ and κ ∞ := β 1 ∥ a ∥ ∞ κ .
Then for any t ≥ 1 , with probability at least 1 − m exp (cid:16) − γtd 2 mλ (cid:17) we have for all 1 ≤ j ≤ m
[Equation]
Proof First, we prove that for any t ≥ 0 and 0 ≤ s ≤√ d κ ∞ η t , we have
[Equation]
The base case of t = 0 is trivial, and for the induction step we have
[Equation]
where (a) holds since ∥ x
∥ 2 is a 1-Lipschitz function of a standard Gaussian random vector, thus it is sub-Gaussian with parameter 1 (Lemma 22) and additionally E [ ∥ x ∥ 2 ] ≤ √ d , and (b) holds by the induction hypothesis (notice that for decreasing step size s (1 − η t λ ) ≤ 2 √ d κ ∞ η t − 1 by (D.5)).
Next, we apply the following Chernoff bound,
[Equation]
which holds for any 0 ≤ s ≤ 2 √ d κ ∞ η t .
Choosing the largest s possible and noting that η t ≤ 4 mγt yields an exp (cid:16) − γtd 2 mλ (cid:17) upper bound on the conditional probability, which followed by taking expectation removes the randomness of conditioning on w 0 j .
Finally applying a union bound gives us the desired bound.
In addition, we would like to approximate R τ ( W T ) and ˆ R τ ( W T ) with R τ ( W T ∥ ) and ˆ R τ ( W T ∥ ) respectively.
As a result, we will investigate the Lipschitzness of the population and empirical risk in the next lemma.
Lemma
Under either Assumption 1.A or 1.B , the truncated risk W (cid:55)→ R τ ( W ) is √ 2 τβ 1 ∥ a ∥ 2Lipschitz.
Moreover, for T ≥ d + log(1 /δ )
with probability at least 1 − δ over the stochasticity of { x ( t ) } 0 ≤ t ≤ T − 1 , the truncated empirical risk W (cid:55)→
ˆ R τ ( W ) is Cτβ 1 ∥ a ∥ 2 -Lipschitz for some absolute constant C .
Proof We begin by the simple observation that ˆ y
(cid:55)→ ℓ (ˆ y, y )
∧ τ is √ 2 τ -Lipschitz when ℓ (ˆ y, y )
= 1 / 2(ˆ y − y )
and 1 -Lipschitz when | ∂ 1 ℓ (ˆ y, y ) | ≤ 1 .
As τ ≥ 1 , we can consider both of them as √ 2 τ
Lipschitz.
Thus by Jensen’s inequality
[Equation]
where (a) follows from the Cauchy-Schwartz inequality.
Note that Equation (E.3) also holds for | ˆ R τ ( W ) − ˆ R τ ( W ′ )
| when expectation is over the empirical distribution given by the training samples, meaning
[Equation]
By Example 6.3 of [68], with probability at least 1 − δ we have
[Equation]
which completes the proof.
Lemma 11 Suppose either Assumptions 1.A or 1.B hold.
Denote the loss with ℓ (ˆ y, y ) =
ℓ (ˆ y − y ) , ˜ S
=
[Equation]
and
G =
[Equation]
for ˜ x ∈ R k and y ∈ R .
Let R ( G ) denote the Rademacher complexity of the function class
G (see
Lemma 11 for definition).
Then with ˜ x ∼ N (0 , U ) we have
[Equation]
where T is the number of samples.
Proof Let F = { (˜ x , y )
(cid:55)→ f a , ˜ W
(˜ x , y ) : ( ˜ W , a , b ) ∈ ˜ S } for f a , ˜ W
(˜ x , y ) =
ˆ y (˜ x ; ˜ W , a , b ) − y .
Define g
z ) :=
ℓ ( z ) ∧ τ , and notice G = { (˜ x , y ) (cid:55)→ g ( f a , ˜ W
(˜ x , y )) :
f a , ˜ W ∈ F} , and that g is a √ 2 τ -Lipschitz (thus √ 2 τ -Lipschitz as well, for τ > 1 ) function.
Then by Talagrand’s contraction principle we have R ( G ) ≤ √ 2 τ R ( F ) .
Moreover, let { ξ i } 1 ≤ i ≤ n be a sequence of i.i.d.
Rademacher random variables.
Then similar to the Rademacher bound of [23]
[Equation]
where (a) holds by H¨older’s inequality and the fact that ∥ a ∥ 1 ≤ √ m ∥ a ∥ 2 ≤ r a , and (b) follows from the fact that σ is β 1 Lipschitz, thus another application of Talagrand’s contraction principle.
Proof [ Proof of Theorem 5 ] Let E 1 denote the event of Lemma 10.
We begin with the following decomposition for generalization error which holds on E 1 ,
[Equation]
where the upper bound follows from Lemma 10.
Consequently, sup a
[Equation]
We begin by upper bounding the first term.
From Theorem 2, on an event E 2 we have with probability at least 1 − O ( δ )
[Equation]
Next, we bound the second term in (E.5).
For each W , define ˜ W :
= U † W ∥ , where U † is the
Moore–Penrose pseudo-inverse of U .
Then, since we have the representation W ∥ =
MU for some M ∈ R m × k ,
[Equation]
Thus, W x = ˜ W ˜ x and ℓ (ˆ y ( x ; W , a , b ) , y ) =
ℓ (ˆ y (˜ x ; ˜ W , a , b ) , y ) for ˜ x
= Ux , when W is in the principal subspace,
i
[Equation]
e. W = W ∥ .
Let E 3 denote the event of Lemma
[Equation]
on which
[Equation]
and consequently
[Equation]
for any 1 ≤ j ≤ m .
Define r ˜ w T as the RHS bound above.
Then on Esup
[Equation]
where we recall
[Equation]
Additionally define
[Equation]
Then Lemma 24 and Lemma 11 yield
[Equation]
Besides, as the loss is bounded by τ , by McDiarmid’s inequality, on an event E 4 which happens with probability at least 1 − O ( δ ) we have
[Equation]
and consequently on ∩ 4 i =1 E i sup
[Equation]
Finally, observe that ∥ a ∥ 1 ≤ √ m ∥ a ∥ 2 ≤ r a and ∥ U † ∥ 2 = ∥ U ∥ − 1≤ √ k ∥ U ∥ − 1 F , thus with probability at least 1 − o ( δ ) ,
[Equation]
We remark that in the setting of Theorem 2 which is adapted in Theorem 5, ∥ a ∥ ∞ ≲ m − 1 , thus κ ∞ ≲ m − 1 .
Finally, we observe that r w ≤ √ 2 m with probability at least 1 − O ( δ ) over initialization, which completes the proof.
E.2. Proof of Theorem
Note that due to the special symmetry in the initialization of Algorithm 1, while training the first layer, all neurons have an identical value, i.e. w tj = w t for all j , and that the stochastic gradient with respect to any neuron can be denote by ∇ ℓ = a∂ 1 ℓ (ˆ y, y ) σ ′ ( ⟨ w , x ⟩ + b ) x .
Furthermore, ∇ w j R λ ( W ) will also be identical for all j , which due to the population gradient formula (3.2), we denote by
[Equation]
where h ( w ) =
(cid:80) mj
=1 H ij ( W ) and d ( w ) =
a E (cid:2) ∂ 212 ℓ (ˆ y, y ) σ ′ ( ⟨ w , x ⟩ + b ) f ′ ( ⟨ u , x ⟩ ) (cid:3) .
Additionally, via the arguments in the proof of Lemma 1, it is not difficult to observe γ/m ≤ h ( w ) + λ ≲ m − 1 .
Furthermore, similar to the arguments of Lemma 7, ⟨∇ ℓ, v ⟩ is Ca ∥ v ∥ 2 -sub-Gaussian for any v ∈ R d .
Next, we will derive a lower bound for (cid:10) w t , u (cid:11) to argue that useful features are learned, which first requires obtaining a sharper upper bound on ∥ w t ∥ 2 than that of Lemma 9.
This improvement is possible due to considering the special case of w tj = w t here.
Lemma 12 Suppose t ≥ d .
Then,
[Equation]
with probability at least 1 − exp( − C ( t ∗ + t )) .
In particular, using the union bound, we have sup
[Equation]
with probability at least 1 − exp( − C ( t ∗ + t 0 ))
− exp( − Cd ) .
Proof Let e t := ∇ w ℓ − ∇ w R λ .
Then we have
[Equation]
Recall that (cid:10) e t , v (cid:11) is Ca ∥ v ∥ 2 -sub-Gaussian, and F t is the sigma algebra generated by { w j } 0 ≤ j ≤ t .
Let ω t := w t − η t ∇ w R λ .
Then, for any 0 ≤ s ≲ γη t a 2 ,
E (cid:2) exp (cid:0) s ∥ w t +1 ∥ 22 (cid:1) | F 0 (cid:3)
=
E (cid:2) exp (cid:0) s ∥ ω t ∥ 22 − 2 sη t (cid:10) ω t , e t (cid:11) + sη 2 t
∥ e t ∥ 22 (cid:1) | F 0 (cid:3) ≤
[Equation]
(cid:104) exp (cid:0)
[Equation]
∥ ω
[Equation]
∥ 22 (cid:1) E (cid:2) exp (cid:0) − 4 sη
[Equation]
(cid:10) ω t , e t (cid:11)(cid:1) | F t (cid:3) 12 E (cid:2) exp (cid:0) 2 sη 2 t ∥
e t ∥ 22 (cid:1) | F t (cid:3) 12 | F 0 (cid:105) .
By sub-Gaussianity of (cid:10) ω t , e t (cid:11) we have E (cid:2) exp( − 4 sη t (cid:10) ω t , e t (cid:11) ) | F t (cid:3) ≤ exp( Cs 2 η 2 t a 2 ∥ ω t ∥ 22 ) .
Moreover, as ∥∇ ℓ ∥ 2 ≤ | a |∥ x ∥ 2 , by Jensen’s inequality
[Equation]
Thus E (cid:2) exp(2 sη 2 t ∥ e t ∥ 22 ) | F t (cid:3) ≤ exp( Csη 2 t a 2 d ) for s ≲ 1 η 2 t a 2 (which holds by s ≲ γη t a 2 , see the proof of Lemma 8 for more details), i.e. we have
[Equation]
Recall that by our choice of η t , 0 ≤ (1 − η t ( h ( w t ) + λ )) ≤ 1 − η t γm (cf.
proof of Lemma 8), and ω t
= (1 − η t ( h ( w t ) + λ ))
w t − η t d ( w t ) u .
As ∥ u ∥ 2 = 1 and | d ( w t ) | ≲ | a | , we have
[Equation]
where (a) holds by Young’s inequality and (b) holds for η t γ ≲ 1 with a sufficiently small absolute constant.
Therefore, for s ≲
[Equation]
Expanding the recursion yields,
[Equation]
Finally, we apply a Chernoff bound with the maximum choice of s ≲ γη t a 2 , and combine it with the fact that ∥ w 0 ∥ 2 ≲ 1 with probability at least 1 − exp( − Cd ) .
Lemma 13 Suppose mab < 1 − | f (0) | .
Then, we have | (cid:10) w t , u (cid:11) | ≳ 1 with probability at least 1 − 2 exp( − Ct ) − exp( − Cd ) .
Proof We will only prove for the case where f is increasing as the case for decreasing f is similar.
We begin by proving an upper bound for d ( w ) when ∥ w ∥ 2 ≲ 1 .
By the triangle inequality,
[Equation]
Furthermore, | ˆ y | ≤ ma ( |⟨ w , x ⟩| + b ) .
Thus, for |⟨ w , x ⟩| ≤
[Equation]
∧
[Equation]
|⟨ u , x ⟩| ≲ 1 and | ϵ | ≲ 1 for sufficiently small absolute constants, we have | ˆ y − y | ≤ 1 hence ∂ 212 ℓ (ˆ y, y )
= − 1 .
Then we have,
[Equation]
y, y )
σ ′ ( ⟨ w , x ⟩ + β ) f ′ ( ⟨ u , x ⟩ ) (cid:3) ≲ a E (cid:2) 1 ( | ϵ | ≲ 1) 1 ( |⟨ w , x ⟩| ≲ 1) 1 ( |⟨ u , x ⟩| ≲ 1) f ′ ( ⟨ u , x ⟩ ) (cid:3)
= a E [ 1 ( | ϵ | ≲ 1)]
E (cid:2) 1 ( |⟨ w , x ⟩| ≲ 1) 1 ( |⟨ u , x ⟩| ≲ 1) f ′ ( ⟨ u , x ⟩ ) (cid:3)
≲ − a. where the last line is obtained by considering supremum over ∥ w ∥ 2 ≲ 1 .
Let A t = { sup t 0 ≤ t ′ ≤ t ∥ w t ′ ∥ 2 ≲ 1 } .
Then,
[Equation]
where (a) follows from the sub-Gausianity of the stochastic noise in the gradient, (b) follows since (cid:10) ∇R λ ( w t ) , u (cid:11) =
d ( w t ) by definition, and (c) holds for s ≲ η t − 1 with a sufficiently small absolute constant.
Notice that by the condition on t ∗ inherited from Theorem 2, 1 − η t ( h ( w t ) + λ )) > 0 , and Since s (1 − η t ( h ( w t ) + λ )) ≤ s (1 − η t γ m ) , we can expand the recursion,
E
[Equation]
where in the second inequality we used a ≍ m − 1 and γ ≍ 1 .
Applying the Chernoff bound implies that (cid:10) w t , u (cid:11) ≳ 1 with probability at least 1 − P (cid:0)
A Ct (cid:1) − exp( − Ct ) ≤ 1 − exp( − C ( t ∗ + t 0 ))
− exp( − Cd ) − exp( − Ct ) .
Finally the result follows by letting t
=
t 2 .
We have proven that | (cid:10) w t , u (cid:11) | ≳ 1 while ∥ w t ⊥ ∥ 2 → 0 .
This fact shows that the features learned in the first layer are useful.
What remains to be shown is an approximation result, such that for a carefully constructed second layer, the network can approximate polynomials of the desired type.
This type of approximation using random biases has been adopted from [23].
We first present an approximation result using infinite neurons.
Lemma 14 Let 0 < | α | ≤ r and b ∼ Unif ( − 2 r ∆ , 2 r ∆) .
For any smooth f : R → R , let ˜
f
α : R →
R be a smooth function such that ˜ f
α ( z ) =
f ( z ) for | z | ≤ r ∆ | α | and ˜ f
α ( − 2 r ∆ α ) =
˜ f
′ α ( − 2 r ∆ α ) = 0 .
Then, for | z | ≤ ∆ we have
[Equation]
Proof Using integration by parts, we have
[Equation]
Now, by a concentration argument, we state an approximation result with finitely many neurons.
Lemma 15 Let r ∗ ≤ | α j | ≤ r and b j ∼ Unif ( − 2 r ∆ , 2 r ∆) .
Then there exists a ( α j , b j ) such that for any fixed z ∈ [ − ∆ , ∆] , with probability at least 1 − δ over the choice of ( b j ) , we have
[Equation]
Moreover, ∥ a ∥ 2 ≲ r ∆ √ m sup | z |≤ r ∆ r ∗ | f ′′ ( z ) | .
Proof Let ˜ f
α ( z ) be a candidate in Lemma 14, which can be obtained by e.g. extending f with suitable polynomials (notice that ˜ f
α only needs to be twice differentiable on its domain).
Now choose a j = 4 r ∆ m ˜ f ′′ α j ( − b j α j ) .
Then Lemma 14 ensures that
[Equation]
As ˜ f
α ( z ) =
f ( z ) for all | z | ≤ r ∆ | α j | , we have ˜ f ′′ α j ( − b j α j ) =
f ′′ α j ( − b j α j ) .
It immediately follows that ∥ a ∥ 2 ≤ Cr ∆ √ m sup | z |≤ r ∆ r ∗ | f ′′ ( z ) | and | a j σ ( αz + b j ) | ≤ C ( r ∆) 2 m sup | z |≤ r ∆ r ∗ | f ′′ ( z ) | .
Applying the
Hoeffding’s inequality finishes the proof.
In the following lemma, we will briefly record useful properties of W T which will be of help for invoking the above approximation results and providing guarantees when the second layer is optimized by SGD.
Through the rest of the proof, we will add the mild assumption that d ≳ log(1 /δ ) .
Otherwise, we need to add e − Cd to the probability of failure in Theorem 4.
Lemma 16 Suppose T ≳
d + log(1 /δ ) .
Then with probability at least 1 − δ over the choice of ( b j ) 1 ≤ j ≤ m and { ( x ( t ) , y ( t ) ) } T − 1 t =0 , the following statements hold:
∥ w Tj ∥ 2 ≍ | (cid:68) w Tj , u (cid:69) | ≍ 1 for all 1 ≤ j ≤ m .
∥ 1 T (cid:80) T − 1 t =0 x ( t ) ( x ( t ) )
⊤ ∥ 2 ≲ 1 . 3.
∥ W T ⊥ ∥ F ≲ (cid:113) m
( d +log(1 /δ )) T . 4.
| (cid:10) u , x ( t ) (cid:11) | ≲ ∆ for all 1 ≤ j ≤ m and 0 ≤ t ≤ T − 1 . 5.
∥ W T x ( t ) ∥ 2 ≲ √ m ( √ d + ∆) for all 0 ≤ t ≤ T − 1 .
Proof We will show that each of the events hold with probability (w.p.)
at least 1 − O ( δ ) .
Recall from Lemma 12 that ∥ w Tj ∥ 2 ≲ 1 for all j w.p. ≥ 1 −O ( δ ) , which implies the same for (cid:68) w Tj , u (cid:69) .
On the other hand, from Lemma 13, | (cid:68) w Tj , u (cid:69) | ≳ 1 for all j w.p. ≥ 1 − O ( δ ) .
Combining these events implies that | (cid:68) w Tj , u (cid:69) | ≍ 1 .
The fact that ∥ 1 T (cid:80) T − 1 t =0 x ( t ) ( x ( t ) )
⊤ ∥ 2 ≲ 1 w p
1 −
( δ
T d + log(1 /δ ) follows from the statement of Lemma 23.
Furthermore, ∥ W T ⊥ ∥ F ≲ (cid:113) m
( d +log(1 /δ ))
T w.p. 1 − O
( δ ) follows from Theorem 2.
Note that as (cid:10) u , x ( t ) (cid:11) ∼ N (0 , 1) , by the choice of ∆ , we have (cid:10) u , x ( t ) (cid:11) ≳ 1 w.p. ≥ 1 − O
( δ/T ) , thus | (cid:10) u , x ( t ) (cid:11) | ≲ ∆ w.p. ≥ 1 − O ( δ ) by a union bound.
Finally, we have
[Equation]
The first term is already bounded by √ m ∆ with probability at least 1 − O ( δ ) .
Moreover, recall that ∥ x ( t ) ∥ 2 − E (cid:2)
∥
x
( t ) ∥ 2 (cid:3) is 1-sub-Gaussian, thus by the union bound ∥ x ( t ) ∥ 2 −√ d ≲ (cid:112) log( T/δ ) ≲ ∆ for all 0 ≤ t ≤ T − 1 .
Thus w.p. ≥ 1 −O ( δ ) we have ∥ W T x ( t ) ∥ 2 ≲ √ m ( √ d +∆) which completes the proof.
From this point onwards, we will denote the Huber loss with ℓ H (ˆ y, y ) =
ℓ (ˆ y − y ) .
Notice that ℓ H is 1-Lischitz.
Lemma 17 Recall
[Equation]
the empirical risk of W T given by Algorithm 1.
Let ∆ ≍ (cid:113) log( Tδ ) , ∆ ∗ := ∆ sup | z | ≲ ∆ | f ′′ ( z ) | , and b j i.i.d.
∼ Unif ( − ∆ , ∆) .
Then, with probability at least 1 − δ (over the randomness of ( b j ) 1 ≤ j ≤ m and { x ( t ) , y ( t ) } T − 1 t =0 hence W T ), for T ≳
d + log(1 /δ ) , there exists a ∗ with ∥ a ∗ ∥ 2 ≲ ∆ ∗ √ m such that
[Equation]
Proof We will condition the following discussion on the event of Lemma 16.
Let α j = (cid:68) w Tj , u (cid:69) , and let a ∗ be constructed according to Lemma 15.
By the Lipschitzness of the Huber loss, for an inividual sample ( x , y ) we have
[Equation]
Moreover, by the Cauchy-Schwartz inequality
[Equation]
Additionally, since ∥ 1 T (cid:80) T − 1 t =0 x ( t ) ( x ( t ) )
⊤ ∥ 2 ≲ 1 , by Jensen’s inequality,
[Equation]
On the other hand
[Equation]
let z ( t ) :
= (cid:10) u , x ( t ) (cid:11) ≤ ∆ .
Then, we can apply Lemma 15, which states that with probability 1 − O ( δ ) over the choice of ( b j ) 1 ≤ j ≤ m ,
[Equation]
Combining the events above
[Equation]
we have with probability at least 1 − δ ,
[Equation]
The final step to apply a concentration for (cid:80) T − 1 t =
ℓ H ( ϵ ( t ) ) .
Note that as ℓ H ( ϵ ) ≤ | ϵ | , if | ϵ | is Csub-Gaussian, then ℓ H ( ϵ ) − E
[ ℓ H ( ϵ )] is also sub-Gaussian with an absolute constant as parameter (can be verified e.g. by Lemma 21).
Then a sub-Gaussian concentration bound implies that
[Equation]
, which finishes the proof.
Let E S [ · ] denote expectation w.r.t.
the random sampling of SGD used to train a , hence conditioned on { x ( t ) , y ( t ) } T − 1 t =0 .
Also, define the stochastic noise in the gradient w.r.t.
a as
[Equation]
Notice that E S (cid:2) e ta (cid:3) = 0 .
Lemma
On the event of Lemma 16 and with ( b j ) i.i.d.
∼ Unif ( − ∆ , ∆) , consider the mapping a (cid:55)→ ˆ R λ ′ ( a ) .
Then, ∇ 2 a ˆ R λ ′ ( a ) ≾ m ∆ 2 + λ ′ , and ∥ e ta ∥ 2 ≲ √ mr ( √ d + ∆) .
Proof For ∇ 2 a ˆ R ( a ) , and any v ∈ R m with ∥ v ∥ 2 = 1 , we have the following computation:
[Equation]
t
[Equation]
where (a) holds since ∥ 1 T (cid:80) T − 1 t =
x ( t ) x ( t ) ⊤ ∥ 2 ≲ 1 .
Thus ∇ 2 a ˆ R λ ′ ( a ) ≾ m ∆ 2 + λ ′ .
On the other hand, as ∥ W T x ( t ) ∥ 2 ≲ √ m ( √ d + ∆) for all 0 ≤ t ≤ T − 1 , we have
[Equation]
Now we can analyze the SGD run on the second layer a to give a high probability statement for ˆ R λ ′ ( a T ) .
As ˆ R λ ′ ( a ) is a smooth and strongly convex function of a , we will state the following well-known elementary convergence result of SGD for smooth and strongly convex functions with bounded noise, which we present in a high-probability framework suitable for our analysis.
Lemma 19 Let R : R m → R be a µ -strongly convex function satisfying µ I m ⪯ ∇ 2 a R ( a ) ⪯ L I m .
Suppose we run the SGD iterates a t +1 =
a t − η t g t with E (cid:2) g t | a t (cid:3) = ∇ a R ( a t ) and
∥ g t ∥ 2 ≤ G .
Choose η t = 2 t +1 µ ( t +1) 2 .
Then with probability at least 1 − δ
[Equation]
µT
[Equation]
where R ∗ = arg min a R ( a ) .
Proof Let e t =
g t − ∇ a R ( a t ) denote the stochastic noise.
By the smoothness property of R , we have
[Equation]
Notice that by Jensen’s inequality, ∥∇ a R ( a t ) ∥ 2 ≤ G , thus ∥ e t ∥ 2 ≤ 2 G and the zero-mean random variable ∥ (cid:10) ∇ a R ( a t ) , e t (cid:11) ∥ 2 is 2 G ∥∇ a R ( a t ) ∥ 2 -sub-Gaussian conditioned on a t .
Now, we can establish the following recursive bound on the MGF of R t :=
R ( a t ) − R ∗ .
For 0 ≤ s ≤ 1 4 η t G 2 we have
[Equation]
where (a) follows since R ( a ) is strongly convex thus satisfies the Polyak-Łojasiewicz inequality 2 λ
[Equation]
R ( a ) − R ∗ ) ≤ ∥∇
[Equation]
R ( a ) ∥ 22 .
As s (1 − η t µ ) ≤ 1 4 η t − 1 G 2 (cf.
(D.5))
, we can expand the recursion and have
[Equation]
Finally, applying a Chernoff bound using s = (4 η t − 1 G 2 ) − 1 concludes the proof.
We are finally in a position to complete the proof of Theorem 4.
Proof [ Proof of Theorem 4 ] We will consider the event of Lemma 16 on which from Lemmawe know with probabilility at least 1 − δ over the dataset and ( b j ) 1 ≤ j ≤ m we have
[Equation]
Notice that a (cid:55)→ ˆ R ( W , a , b ) is a convex function.
Thus by strong duality, there exists λ ′ > 0 such that the value of the above constrained minimization problem is equal to the value of the following regularized minimization problem,
[Equation]
Explicitly, this λ ′ can be chosen such that the unique solution to
[Equation]
has ∥ a ∗ ∥ 2 ≲ ∆ ∗ √ m .
Notice that this a ∗ is the unique solution to arg min a ˆ R λ ′ ( W T , a , b ) .
Moreover, from Lemma 19 we have ˆ
[Equation]
and by strong convexity
[Equation]
Thus, with sufficiently large T ′ such that
[Equation]
) we have ∥ a T ′ ∥ 2 ≲ ∆ ∗ √ m and
[Equation]
Finally, we invoke Theorem 5, to close the generalization gap and get
[Equation]
Appendix F. Example of Non-Convex R λ ( W )
Here, we outline an examples for which R λ ( W ) is non-convex on a neighborhood around W =while λ = ˜
λm satisfies the condition in Theorem 2.
For simplicity of exposition, we fix a = 1 m m where 1 m is the vector of all ones.
It is easy to observe that the results hold with high probability when a follows the initialization of Assumption 2 as well.
This example is constructed for the smooth activation case of Theorem 2.
We require σ ( ± 1) = − β
= − 1 (which automatically implies σ ′ ( ± 1) = 0 as σ attains its minimum) and σ ′′ ( ± 1) =
β 2 = 1 .
Then simplifying ∇ 2 R λ ( 0 ) yields
[Equation]
Thus, we need to show that ˜ λ − 1 m I d − E [ y xx ⊤ ] m is not PSD.
Let y = 1 2 (1 + tanh( ⟨ w , x ⟩ 2 − ∥ w ∥ 22 )) and ˜ λ
= 1 + E
[ y ] + γ (notice that y ≥ 0 thus ˜ λ indeed satisfies the assumption in Theorem 2) with
[Equation]
Then we have w ⊤
[Equation]
∥ w ∥ 22 − 12 E
[Equation]
w ,
[Equation]
22 ) tanh
[Equation] [Equation]
w
[Equation] [Equation]
m < 0 .
Therefore, we have shown that R λ ( W ) is not convex on a neighborhood around zero.
Appendix G. Auxiliary Lemmas
In order to be explicit, we state the following definitions and lemmas that will be used in the proof of Theorem 5.
We only state the next definitions and lemmas and refer the reader to [68] and [67] for proof and more details.
Definition 20 [68, Definitions 2.2 and 2.7] A real-valued random variable z is said to be ν
-subGaussian
if for all s ∈ R we have E [exp( sz )] ≤ exp( s E [ z ] + s 2 ν 2 2 ) , and is said to be u -subexponential if for all | s | ≤ 1 ν we have E [exp( sz )] ≤ E (cid:104) exp( s E
[ z ] + s 2 ν 2 2 ) (cid:105) .
Lemma 21 [67, Propositions 2.5.2 and 2.7.1] Suppose z is a zero-mean random variable and
E [ | z | p ] 1 p ≤ L √ p for all p ≥ 1 .
Then z is cL -sub-Gaussian for an absolute constant c > 0 , i.e.
E [exp( sz )] ≤ exp( s 2 c 2 L 2) for all s ∈ R .
Similarly, suppose E [ | z | p ] 1 p ≤ Lp .
Then z is cL -subexponential for an absolute constant c > 0 .
Lemma 22 [68, Theorem 2.26] Let x ∼ N (0 , I d ) .
If f : R d → R is L -Lipschitz, then f ( x ) is sub-Gaussian with parameter L .
Lemma 23 [68, Example 6.3] Let { x ( i ) } 1 ≤ i ≤ n be a sequence of i.i.d.
standard Gaussian random vectors x i ∼ N (0 , I d ) .
It holds with probability at least 1 − δ that
[Equation]
where C is an absolute constant.
The next lemma is the well-known symmetrization argument that upper bounds the expected value of an empirical process with Rademacher complexity.
Lemma 24 [49, Theorem 3.3] Let F be a class functions f : R p
→ R for some p > 0 .
For a number of samples T and a probability distribution P on R p , define the Rademacher complexity of
F as
[Equation]
where { x ( t ) } T − 1 t =0 i.i.d.
∼ P and { ξ t } T − 1 t =0 are independent Rademacher random variables (i.e. ±equiprobably).
Then the following holds,
[Equation]
Furthermore, we have the following fact for standard normal random vectors.
Lemma 25 Let x ∼ N (0 , I d ) .
There exists an absolute constant C > 0 such that for any V ∈
R m × k and p ≥ 1 we have
[Equation]
Proof First of all, ∥ V x ∥ 2 is a ∥ V ∥ 2 -Lipschitz function of x , thus Lemma 22 applies and ∥ V x ∥is sub-Gaussian.
Furthermore, by applying Lemma 21 to ∥ V x ∥ 2 − E [ ∥ V x ∥ 2 ] and Minkowski’s inequality, we have
[Equation]
where the last inequality follows from Jensen’s inequality.
Lemma 26 Let x ∼ N (0 , I d ) .
Then E (cid:2) exp( c ∥ x ∥ 22 ) (cid:3) ≤ exp(2 cd ) for c ≤ 1 / 4 .
Proof Gaussian integration yields E (cid:2) exp( cx 2 i ) (cid:3) = 1 √ 1 − 2 c .
Furthermore, for c ≤ 14 we have 1 √ 1 − 2 c ≤ exp(2 c ) .