O N FEATURE LEARNING IN NEURAL NETWORKS WITH GLOBAL CONVERGENCE GUARANTEES
Anonymous authors
Paper under double-blind review
A BSTRACT
We study the gradient flow optimization of over-parameterized neural networks (NNs) in a setup that allows feature learning while admitting non-asymptotic global convergence guarantees.
First, we prove that for wide shallow NNs under the mean-field (MF) scaling and with a general class of activation functions, when the input dimension is at least the size of the training set, the training loss converges to zero at a linear rate under gradient flow.
Building upon this analysis, we study a model of wide multi-layer NNs with random and untrained weights in earlier layers, and also prove a linear-rate convergence of the training loss to zero regardless of the input dimension.
We also show empirically that, unlike in the Neural Tangent Kernel (NTK) regime, our multi-layer model exhibits feature learning and can achieve better generalization performance than its NTK counterpart
I NTRODUCTION
The training of neural networks (NNs) can often succeed despite being a non-convex optimization problem and involving a large number of parameters.
To understand this mystery, one recent direction is to take advantage of the over-parametrization - that the NNs have large widths - by studying their infinite-width limits under certain scalings of the parameters.
For example, under the NTK or “ 1 / √ width " scaling , the gradient descent (GD) dynamics of the training of a wide NN can be approximated by the linearized dynamics around initialization, and in the infinite-width limit, we obtain a kernel method with the Neural Tangent Kernel (NTK) (37).
Theoretical guarantees for optimization and generalization can be obtained for wide NNs under the NTK scaling (19; 5).
Nonetheless, it has been pointed out that the NTK analysis replies on a lack of feature learning / representation learning (16; 72), which is a crucial ingredient to the success of deep learning, and is therefore not adequate for explaining the success of NNs (27; 41).
Alternatively, for shallow (i.e. one-hidden-layer) NNs, if we initialize the parameters according to a scaling of “ 1 / width ”, we can derive the mean-field (MF) limit of shallow NNs.
In this limit, feature learning occurs, and the training dynamics can be described by the Wasserstein gradient flow of a probability measure on the space of the parameters, which converges to a global minimizer of the loss function under certain conditions (59; 50; 63; 14).
Generalization guarantees have also been proved for learning with shallow NNs under the MF scaling via a corresponding function space (6; 48).
Currently, however, three limitations exist to the mean-field analysis of NNs.
First, the global convergence guarantees for shallow NNs only hold in the infinite-width limit, and the deviation between finite-width NNs and their infinite-width limits has only been studied asymptotically (12).
Second, a convergence rate is yet to be established except under special assumptions or with modifications to the GD algorithm (13; 34; 43; 38).
Second, while several works have attempted to extend the MF formulation to deep (i.e. multi-layer) NNs (4; 62; 52; 23; 57), there has been less of a concensus on what the right model should be than in the shallow case.
More generally, we still lack a tractable model for the optimization of multi-layer NNs that, unlike NTK, allows feature learning.
The aim of this work is to obtain theoretical guarantees of fast global convergence of feature learning in both shallow and multi-layer NNs, and we do so by examining the optimization properties of the model described in the next section, which generalizes shallow NNs.
1. S UMMARY OF MAIN CONTRIBUTIONS
We consider the scenario of training an NN model to fit a set of n training data with input dimension d under least-squares regression, where the model parameters are optimized by gradient flow.
With most choices of the activation function, we prove that:
For a shallow NN, if the hidden layer is sufficiently wide and the input data are linearly independent (requiring n ≤ d ), then with high probability, the training loss converges to zero at a linear rate.
For a multi-layer NN where we only training the second-to-last layer, if the hidden-layer widths are sufficiently large, then with high probability, the training loss converges to zero at a linear rate.
Unlike for shallow NNs, here we no longer need the requirement that n ≤ d , demonstrating a benefit of jointly having depth and width.
To our knowledge, we are presenting the first theoretical guarantee of linear-rate convergence of feature learning in multi-layer NNs.
We also run numerical experiments that complement the theory.
1. R ELATED WORKS
Over-parameterized NNs, NTK and lazy training.
Many recent works have studied the optimization landscape of NNs and the benefits of over-parameterization (24; 67; 60; 66; 64; 54; 65; 42; 3; 40; 76; 15; 75; 17; 39; 9; 32; 20; 31; 61; 10).
One influential idea is the Neural Tangent Kernel (NTK) (37), which characterizes the behavior of GD on NNs by its linearization around initialization, and becomes exact if the width tends to infinity under a particular scaling of the parameters (e.g. replacing 1 /m with 1 / √ m in (1)).
In particular, when the network width is polynomially large in the number of training set size, the training loss converges to global minima at a linear rate under
GD (19; 5; 55).
Nonetheless, in the NTK limit, due to a relatively large scaling of the parameters at initialization, the hidden-layer features do not move significantly (37; 19; 16).
For this reason, the NTK scaling has been called the lazy-training regime, as opposed to a feature-learning or rich regime (16; 72; 26).
Several works have investigated the differences between the two regimes both in theory (28; 29; 69; 47) and in practice (27; 41).
In addition, several works have generalized the
NTK analysis by considering higher-order Taylor approximations of the GD dynamics or finite-width corrections to the NTK (2; 35; 7; 33).
Mean-field theory of NNs.
An alternative path has been taken to study shallow NNs in the mean-field scaling (as in (1)), where the infinite-width limit is analogous to the thermodynamic or hydrodynamic limit of interacting particle systems (59; 50; 63; 14; 49; 70).
Thanks to the interchangeability of the parameters, the neural network is equivalently characterized by a probability measure on the space of its parameters, and the training can then be described by a Wasserstein gradient flow followed by this probability measure, which, in the infinite-width limit, converges to global mimima under mild conditions.
Regarding convergence rate, ref.
(71) proves that if we train a shallow NN to fit a Lipschitz target function under population loss, the convergence rate cannot beat the curse of dimensionality.
In contrast, we will study the setting of empirical risk minimization, where there are finitely many training data.
Ref.
(34) shows that mean field Langevin dynamics on shallow NNs can converge exponentially to global minimizers in over-regularized scenarios, but we focus on gradient flow without entropic regularization.
Besides the question of optimization, shallow NNs under this scaling represent functions in the Barron space (48) or variation-norm function space (6), which provide theoretical guarantees on generalization as well as fluctuation in training (12).
Several works have proposed different mean-field limits of wide multi-layer NNs and proved convergence guarantees (4; 62; 51; 52; 23; 57; 22), but questions remain.
First, due to the presence of different symmetries in a multi-layer network compared to a shallow network (57), the limiting object at the infinite-width limit is often quite complicated.
Second, it has been pointed out that under the MF scaling of a multi-layer network, an i.i.d. initialization of the weights would lead to a collapse of the diversity of neurons in the middle layers, diminishing the effect of having large widths (23).
In addition, while another line of work develops MF models of residual models (8; 46; 21; 36), we are interested in multi-layer NN models with a large width in every layer.
Feature learning in deep NNs.
Ref.
(1) demonstrates the importance of hierarchical learning by proving the existence of concept classes than can be learned efficiently by a deep NN with quadratic activations but not by non-hierarchical models.
Ref.
(11) studies the optimization landscape and generalization properties of a hierarchical model that is similar to ours in spirit, where an untrained embedding of the input is passed into a trainable shallow model, and prove an improvement in sample complexity in learning polynomials by having neural network outputs as the embedding.
However, the trainable models they consider are not shallow NNs but their linearized and quadratic-Taylor approximations, and furthermore the convergence rate of the training is not known.
Ref.
(73) shows the existence of an infinite-width limit of deep NNs under a new scaling, which exhibits feature learning, but the convergence of the training dynamics is not understood.
Our pseudo-multi-layer NN models adopt the same scaling, and our focus is on proving non-asymptotic convergence guarantees.
P ROBLEM SETUP 2. M ODEL
We summarize some of our notations in Appendix A. Let Ω ⊆ R d denote the input space, and let x = [ x 1 , ..., x d ] ⊺ ∈ Ω denote a generic input data vector.
A shallow NN model under the MF scaling can be written as:
[Equation]
where m is the width, W ∈ R m × d and c
= [ c 1 , ..., c m ] ∈ R m are the first- and second-layer weight parameters of the model, and σ : R → R is the activation function.
For simplicity of presentation, we neglect the bias terms.
In this paper, we study a more general type of models with the following form:
[Equation]
where W ∈ R m × D and c =
[ c 1 , ..., c m ] ∈ R m are parameters of the model, and ϕ 1 , ..., ϕ D are a set of functions from Ω to R that we call the embedding .
Each of h 1 , ..., h m is a function from Ω to R , and we will refer to them as the (hidden-layer) feature map or activations .
For simplicity, we write Φ( x ) =
[ ϕ i ( x ) , ...
ϕ D ( x )] ⊺ ∈ R D .
We consider two types of the embedding, Φ :
Fixed embedding D is fixed and Φ is deterministic.
In the simplest example, we set D = d and ϕ j ( x ) = x j , ∀ j ∈ [ D ] , and recover the shallow NN model in (1).
More generally, our definition includes cases where Φ is a deterministic transformation of an input vector in Ω into an embedding vector in R D .
This can be understood as input pre-processing or feature engineering.
High-dimensional random embedding D is large and Φ is random.
For instance, we can sample z j i.i.d.
in R d and set ϕ j ( x ) = σ (cid:16) 1 √ d z ⊺ j x (cid:17) , equivalent to setting ϕ 1 , ..., ϕ m as the hidden-layer activations of a shallow NN with randomly-initialized first-layer weights.
Then, the model becomes
[Equation]
Thus, we obtain a 3 -layer feed-forward NN whose first-layer weights are random and fixed, and we call it a pseudo- 3 -layer (pseudo- 3 L) NN .
Note that the scaling in this model is different from both the NTK scaling ( 1 / √ m instead of 1 /m in (4)) and the MF scaling for multi-layer NNs adopted in (4; 62; 51; 57; 23) ( 1 /D instead of 1 / √ D in (5)).
We show in Appendix B that when σ is homogeneous, this scaling is consistent with the Xavier initialization of neural network parameters up to a reparameterization (30; 56).
We also show in Appendix C that a special case of this scaling is equivalent to the maximum-update scaling scheme proposed in (73).
Numerical experiments that compare different scalings are described in Section 4.
2. T RAINING WITH GRADIENT FLOW
Consider the scenario of supervised least-squares regression, where we are given a set of n training data points together with their target values, ( x 1 , y 1 ) , ..., ( x n , y n ) ∈
Ω × R .
We fit our models by minimizing the empirical least-squares loss:
[Equation]
To do so, we first initialize the parameters randomly by sampling each c i and W ij i.i.d.
from probability measures π c ∈ P ( R ) and π w ∈ P
( R D ) , respectively, and then perform GD on W .
For simplicity of analysis, we leave c untrained , and further assume that
Assumption 1.
π c = 12 δ ˆ c
( dc ) + 12 δ − ˆ c ( dc ) for some ˆ c > 0 independent from m , which is the law of a scaled Rademacher random variable.
If σ is Lipschitz, it is differentiable almost everywhere, and we write σ ′ ( x ) for the value of its derivative when σ is differentiable at x and 0 otherwise.
When σ is differentiable at h i ( x ) , there is
[Equation]
[Equation]
Thus, the gradient of the loss function with respect to W ij is given by
[Equation]
Thus, we can perform GD updates on W according to the following rule (consistent with standard
GD for Xavier-initialzed NNs, as discussed in Appendix B): ∀ i ∈ [ m ] and ∀ j ∈
[ D ] ,
[Equation]
where δ > 0 is the step size.
In the limit of infinitesimal step size ( δ → 0 ), the evolution of the parameters during training is described by the gradient flow equation: if we use the superscript t ≥to denote time elaposed during training, the time-derivative of the parameters is given by
[Equation]
where f t denotes the output function and h t 1 , ..., h tm denote the hidden-layer feature maps determined by the parameters at time t .
Then, induced by the evolution of W t , each h t i evolves according to ∀
[Equation]
where we define a kernel function G :
Ω × Ω → R as
[Equation]
[Equation]
Accordingly, the output function f t satisfies
[Equation]
Thus, the loss function L t :=
L [ f t ] evolves according to
[Equation]
where we define the symmetric Gram matrix G ∈ R n × n with entries G ab = G ba =
G ( x a , x b ) , and use λ min ( G ) to denote its least eigenvalue.
We will also use G min = min a ∈ [ n ] G aa and
G max = max a ∈ [ n ] G aa to denote the minimum and maximum diagonal entries of G , respectively.
Since G is positive semidefinite, we see that ˙
L t ≤ 0 , which means that the loss value is indeed non-increasing along the gradient flow trajectory.
Feature learning Compared to the NTK scaling of neural networks, the crucial difference is the 1 /m instead of 1 / √ m factor in (2).
It is known that under the NTK scaling, due to the 1 / √ m factor, the movement of the feature maps, h 1 , ..., h m , is only of order O (1 / √ m ) before the function value changes by an amount of order Ω(1) .
While this greatly simplifies the convergence analysis, it also implies a lack of learning of the hidden representations.
By contrast, with the 1 /m factor in (2), if σ is Lipschitz with Lipschitz constant L σ , there is | f t 2 ( x ) − f [Equation] 1 ( x ) | ≤ L σ [Equation]
x ) | , ∀ t 1 , t 2 ≥ 0 .
Therefore, regardless of m and D ,
[Equation]
which implies that the average movement of the feature maps is at least on the same order as the change in function value, and thus the NTK undergoes nontrivial movement.
In Appendix C, we further justify the occurrence of feature learning using the framework developed in (73).
C ONVERGENCE ANALYSIS 3. M ODELS WITH A FIXED EMBEDDING
Based on (14), to prove that the training loss converges to zero, we need a tighter upper bound on the value of ˙ L t than 0.
Indeed, if G is positive definite, which depends on Φ and the training data, we can achieve this as follows.
As a warm-up, if we use activation functions whose derivative has a magnitude uniformly bounded from below by a constant K σ ′ > 0 , such as linear, cubic or (smoothed) Leaky ReLU activations, we can directly establish a Polyak-Lojasiewicz (PL) condition (58; 45),
[Equation]
(16) and hence L t ≤ L 0 e − 2ˆ c 2 λ min ( G )( K σ ′ )
2 t/n , implying that the training loss decays to 0 at a linear rate.
For more general choices of the activation function, a challenge is to guarantee that, heuristically speaking, for each a ∈ [ n ] , σ ′ (cid:0) h
i ( x a ) (cid:1) does not become near zero for too many i ∈ [ m ] before the loss vanishes.
To facilitate a fine-grained analysis, we need the following mild assumption on σ :
Assumption 2.
σ is Lipschitz with Lipschitz constant L σ , and there exists an open interval
I =
( I l , I r ) ⊆ R on which σ is differentiable and | σ ′ | is lower-bounded by some K σ ′ > 0 .
Intuitively, I is an active region of σ , within which the derivative has a magnitude bounded away from zero.
This assumption is satisfied by the majority of activation functions in practice, including smooth ones such as tanh and sigmoid as well as non-smooth ones such as ReLU.
Then, under the following initialization scheme, we prove a general result for models with a fixed embedding.
Assumption 3.
π w is the law of N (0 , Id D ) , i.e., the product distribution of D independent standard
Gaussian random variables.
Theorem 1 (Fixed embedding) .
Suppose that Assumptions 1, 2 and 3 are satisfied, and λ min ( G ) > 0 .
Then ∃ ˆ c 0 , r and C > 0 such that ∀ δ > 0 , if ˆ c ≥ n 2 ˆ c 0 /λ min ( G ) and m ≥ C
1 + ˆ c 2 / ∥ y ∥ 2 ) log ( n/δ ) , then with probability at least− δ , it holds that ∀ t ≥ 0 ,
[Equation]
The values of ˆ c 0 , r and C depend on I, G min , G max , L σ and K σ ′ (but not m , n , d , D , δ , λ min ( G ) ).
The result is proved in Appendix E, and below we briefly describe the intuition.
For each a ∈ [ n ] , among all neurons i ∈ [ m ] , we can keep track of the proportion of neurons for which h ti ( x a ) ∈
I .
In particular, we show that if the proportion is large enough at initialization (shown by Lemma 3 in
Appendix E.2 under Assumption 3), then it cannot drop dramatically without a simultaneous decrease of the loss value, as long as the c i ’s are not too small in magnitude.
This property of the dynamics is formalized in the following lemma:
Lemma 1. Consider the dynamics of L t and (cid:8) h ti ( x a ) (cid:9) i ∈
[ m ] ,a ∈ [ n ] governed by (11) and (14) .
Assume that λ min ( G ) > 0 , and ∀ i ∈ [ m ] , | c i | =
ˆ c > 0 .
Under Assumption 2, define η t = min a ∈ [ n ] (cid:40)m m (cid:88) i
{+= 
+}
h ti ( x a ) ∈
I (cid:41) , ∀ t ≥ 0 and ˜ η 0 = min a ∈ [ n ] (cid:40)m m (cid:88) i
{+= 1 h 
+}
i ( x a ) ∈ ( 2 Il + Ir,
Il +2 Ir)
(cid:41) (18)
Then ∀ t ≥ 0 , there is
[Equation]
[Equation] where κ ( λ 1 , λ 2 ) = 9 λ 2 ( I r − I l ) 2 λ 1 K σ ′ .
3.1. E XAMPLE : S HALLOW NEURAL NETWORKS WHEN n ≤ d
In the case of shallow NNs, G = G (0) ∈ R n × n , where
[Equation]
Thus, G (0) is positive definite if and only the training data set (cid:8) x 1 , ..., x n } ⊆
R d consists of linearlyindependent vectors, which is possible (and expected if the training data are sampled independently from some non-degenerate distribution)
when n ≤ d .
In that case, Theorem 1 implies
Corollary 1 (Shallow NN with n ≤ d ) .
Suppose that Assumptions 1, 2 and 3 are satisfied.
If the training data are linearly-independent vectors, then under the gradient flow (10) on the first-layer weights of the shallow NN, the training loss converges to zero at a linear rate.
While the assumption that n ≤ d is restrictive, we note that existing convergence rate guarantees for the GD-type training of shallow NNs in the MF scaling need strong additional assumptions (38), modifications to the GD algorithm (34; 13; 53), or restrictions to certain special tasks (43).
3. M ODELS WITH A HIGH - DIMENSIONAL RANDOM EMBEDDING
A clear limitation of Corollary 1 is that it is only applicable when n ≤ d , since otherwise the Gram matrix G (0) cannot be positive definite.
With this in mind, we consider the type of models where D is large and Φ is a high-dimensional random embedding.
While the Gram matrix G in this case is also random, we are interested in its limit as D tends to infinity.
Inspired by the analysis above, we desire the following condition on the random embedding:
Condition 1 (Concentration of G around a positive definite matrix) .
There exists a (deterministic) positive definite matrix ¯ G ∈ R n × n with least eigenvalue λ min ( ¯ G ) > 0 such that ∀ δ, u > 0 , ∃ D min ( δ, u ) > 0 such that if D ≥ D min ( δ, u ) , then P ( ∥ G − ¯ G ∥ 2 > u ) < δ .
Since the sampling of c and W is independent from the randomness of Φ , Condition 1 is sufficient for us to apply Lemma 1 and obtain the following global convergence guarantee that extends Theoremto models with a high-dimensional random embedding.
The proof is given in Appendix F.
Theorem 2 (High-dimensional random embedding) .
Suppose Assumptions 1, 2 and 3 and Conditionhold.
Then ∃ ˆ c 0 , r and C > 0 such that ∀ δ > 0 , if ˆ c ≥ n 2 ˆ c 0 /λ min ( ¯ G ) , m ≥ C (1+ˆ c 2 / ∥
y ∥ 2 ) log ( n/δ ) and D ≥ D min ( 12 δ, 12 λ min ( ¯ G )) , then with probability at least 1 − δ , it holds that ∀ t ≥ 0 ,
[Equation]
The values of ˆ c 0 , r and C depend on I, ¯ G min , ¯ G max , L σ and K σ ′ (but not m , n , d , D , δ , λ min ( ¯ G ) ).
3.2. E XAMPLE : P SEUDO - THREE - LAYER NEURAL NETWORKS
Consider the pseudo- 3 L NN model defined in (4).
In this case, the Gram matrix is G (1) , defined by
[Equation]
[Equation]
If z 1 , ..., z m are sampled i.i.d.
from a probability measure π z on R d and fixed during training, then the limiting Gram matrix, denoted by ¯ G (1) ∈ R n × n , is given by
[Equation]
Thus, for the convergence result, the assumption we need on the limiting Gram matrix is
Assumption 4.
The matrix ¯ G (1) , which depends on the choice of σ and the training set, is positive definite with least eigenvalue λ min (cid:0) ¯ G (1) (cid:1) > 0 and ( ¯ G (
[Equation] This assumption also plays an important role in the NTK analysis, and it is satisfied if, for example, π z = N (0 , I d ) , no two data points are parallel and σ is either the ReLU function (19) or analytic and not a polynomial (18).
When Assumption 4 is satisfied, as long as σ is Lipschitz and π z is sub-Gaussian (including Gaussian, uniform and Rademacher distributions, for example), we can use standard concentration techniques to verify Condition 1.
Thus, Theorem 2 implies that
Theorem 3 (Pseudo- 3 L NN) .
Suppose Assumptions 1, 2, 3 and 4 hold, and π z is sub-Gaussian.
Then ∃ ˆ c 0 , r , C 1 and C 2 > 0 such that ∀ δ > 0 , if ˆ c ≥ n 2 ˆ c 0 /λ min ( ¯ G (1) ) , m ≥ C 1 (1 + ˆ c 2 / ∥ y ∥ 2 ) log ( n/δ ) and D ≥ C 2 n 2 log( n/δ ) /λ
min ( ¯ G (1) ) 2 , then with probability at least 1 − δ , it holds that ∀ t ≥ 0 ,
[Equation]
The values of ˆ c 0 , r , C 1 and C 2 depend on I, ¯ G (1)min , ¯ G (1)max , K σ ′ as well as the sub-Gaussian norm of µ z (but not m , n , d , D , δ , λ min ( ¯ G (1) ) ).
The proof is given in Appendix G. Compared to Corollary 1 for shallow NNs, a highlight of Theoremis that the requirement of n ≤ d is no longer needed.
This demonstrates an advantage of the highdimensional random embedding, realized as the activations of the first hidden layer in the pseudo- 3 L
NN, which illustrates a benefit of jointly having depth and width in NNs.
Compared to the NTK analysis (18), our result assumes the same level of over-parameterization, but crucially allows feature training to occur, as discussed in Section 2.2.
3.2. G ENERALIZATION TO DEEPER MODELS
By setting Φ to be the activations of the second-to-last hidden-layer of a multi-layer NN [Equation] we can obtain generalizations of the pseudo- 3 L NN to deeper architectures.
For example, in the feed-forward case, we can obtain the following pseudo-
L -layer NN :
[Equation]
[Table]
where ¯ W (1)
[Table]
[Table]
¯ W ( L − 3) ∈ R D × D and z [Table]
..., z D ∈ R d are sampled randomly and fixed. This model can be written in the form of (4) and (5) with ϕ j ( x ) = σ (cid:0) h ( L − 2)
j ( x ) (cid:1) .
The corresponding Gram matrix is recursively defined and also appears in the NTK analysis (18).
In particular, the results in (18) imply that if σ is analytic and not a polynomial [Table] then Condition 1 holds
[Table]
and hence similar global convergence results can be obtained as corollaries of Theorem 2.
N UMERICAL EXPERIMENTS
Additional results and details of the experiments are provided in Appendix H
[Table]
4. E XPERIMENT 1: C ONVERGENCE OF TRAINING WHEN FITTING RANDOM DATA
We train a shallow NN to fit a randomly labeled data set { ( x 1 , y 1 ) , ..., ( x n , y n ) } with d = 20 , where we sample each x a i.i.d.
from N (0 , I 20 ) , with I 20 being the 20 × 20 identity matrix, and each y a i.i.d.
uniformly on [ − 12 , 12 ] and independently from the x a ’s. We see from Figure 1 that the convergence happens at a nearly linear rate when n = 20 and 40 , and the rate decreases as n becomes larger.
This is coherent with our theoretical result (Corollary 1) and also echoes the negative result by (71) on the convergence rate of optimizing a shallow NN using population loss, which implies a worsening of the convergence rate as the number of data points increases to infinity.
4. E XPERIMENT 2: B ENEFIT OF INPUT EMBEDDING
We consider a model defined by (2) and (3) with Φ( x ) = vec ( xx ⊺ )
∈ R
d 2 with d = 30 , and call this model a shallow NN augmented with quadratic embedding .
We compare this model against a plain shallow NN (without extra embedding), both with m = 8192 , on a regression task where the target y is given by another shallow NN augmented with quadratic embedding with m = 5 and varying n ’s.
We see from Figure 2 that the augmented shallow NN indeed achieves lower test error given the same number of training samples, demonstrating the benefit of receiving a good embedding.
4. E XPERIMENT 3: F EATURE LEARNING V . S . LAZY TRAINING
We consider the pseudo- 3 L NN model defined in (4) and (5) with D set to be m (i.e. both hidden layers having the same width, m ), and compare it with the same model under NTK and MF scalings as described in Table 1.
We adopt the data set used in (69) (more details in Appendix H.3), and train the models by unregularized least-squares regression for varying n ’s and m ’s.
First, we see from the top-left plot in Figure 4 that, consistently across different m , the training loss converges at a linear rate for the model under our scaling, which is coherent with Theorem 3.
Second, we see from the second row that feature learning occurs in the model under our scaling but very little in the model under NTK scaling, as expected (16).
Note also that in the MF scaling, the feature maps h 1 ( x ) , ..., h m ( x ) concentrate near 0 at initialization due to the small scaling, but gains diversity during training.
Third, we see from Figure 3 that our scaling yields the smallest test errors out of the three, and as n grows the test error also decreases faster in the MF scaling than in the NTK scaling, which indicates the advantage of feature learning compared to lazy training.
C ONCLUSIONS AND LIMITATIONS
We consider a general type of models that includes shallow and pseudo-multi-layer NNs, which exhibits feature learning when trained via gradient flow, and prove non-asymptotic global convergence [Figure] [Figure] [Figure]
[Figure] guarantees that accommodates a general class of activation functions. For a randomly-initialized shallow NN, we prove that by performing gradient flow on the input-layer weights, the training loss converges to zero in a linear rate if the number of training data does not exceed the input dimension.
For a randomly-initialized multi-layer NN, we prove that by performing gradient flow on the weights in the second-to-last layer, the same result holds except there is no more requirement on the input dimension.
We also perform numerical experiments to demonstrate the advantage of feature learning in our pseudo-multi-layer NNs relative to their counterparts under the NTK scaling.
Our work focuses on only the optimization and not the approximation or generalization of NNs.
In addition, as the current theoretical results on global convergence neglect the bias terms and assume that the last-layer weights are untrained, and a more general version is left for future work.
A A DDITIONAL NOTATIONS
• For a positive integer n , we let [ n ] denote the set { 1 , ..., n } .
• We use i, j (as subscripts) to index the neurons in the hidden layers, a, b (as subscripts or superscripts) to index different training data points, t (as a superscript) to denote the training time / time parameter in gradient flow.
• We write (cid:80) a for 1 n (cid:80) na =1 .
• We use bold letters (e.g. x , z , c , y ) to denote vectors.
• We use W and { W ij } i ∈ [ m ] ,j ∈
[ D ] interchangeably to refer to the same set of parameters.
B C ONSISTENCY OF THE SCALING AND GD UPDATE RULE WITH X AVIER INITIALIZATION
Consider a three-layer network defined by
[Equation]
[Equation] with weight parameters (cid:8) θ (1) jk (cid:9) j,k ∈ [ m ] , (cid:8) θ (2) ij (cid:9) i,j ∈ [ m ] and (cid:8) θ (3) i (cid:9) i ∈ [ m ] are initialized according to
Xavier initialization, which means that we sample each θ (1) jk i.i.d.
from N (0 , 1 m + d ) , each θ (2) ij i.i.d.
from N (0 , 12 m ) , and each θ (3) i i.i.d.
from N (0 , 1 m +1 ) .
If m ≫ d , both N (0 , 1 m + d ) and N (0 , 1 m +1 ) can be approximated by N (0 , 1 m ) .
Then, up to this approximation, by redefining c i = √ mθ (3) i ,
W ij = √ mθ (2) ij and z jk = √ mθ (1) jk , we can write
[Equation]
and note that c i , W ij and z jk are all initialized i.i.d.
of order O (1) .
In addition, if σ is homogeneous, this is then equivalent to (4) and (5) when D = m .
Moreover, there is ∂f∂W ij ( x ) = 1 √
m ∂f ∂θ (2) ij ( x ) .
Then, since performing GD on θ (2) ij with step size δ means updating θ (2) ij according to
[Equation]
(29) this is equivalent to updating W ij according to
[Equation]
[Equation]
which justifies the m factor on the right-hand-side of (9).
C R ELATIONSHIP TO THE ABC - PARAMETERIZATION (73) AND FEATURE LEARNING
Consider the pseudo- L -layer NN model defined in Section 3.2.2 in the case where D = m >> d . In the language of abc-parameterization defined in (73), our parameterization corresponds to having [Equation]
Furthermore, as we explain in Appendix B, the appropriate learning rate scales linearly with m (as in (9)), which corresponds to having
[Equation]
Meanwhile, the maximum-update ( µ P) scaling (73) is characterized by setting [Equation] Recall the symmetry of abc-parameterization derived in (73), which states that one gets a different but equivalent abc-parameterization by setting [Equation]
Since our parameterization can be obtained from the maximum-update parameterization by applying the transformation above with θ = 12 , they are equivalent in the function space.
In particular, for our parameterization, the r parameter defined in (73) can be computed as
[Equation]
Hence, according to (73), our parameterization exhibits feature learning.
D P ROOF OF L EMMA
Since we assume that G is positive definite and | c i | =
ˆ c > 0 , ∀ i ∈ [ m ] , we can derive from (14) that
[Equation]
[Equation]
f
[Equation]
[Equation]
[Equation]
[Equation]
(cid:1)
[Equation]
[ n ]
[Equation]
i
[Equation]
[Equation]
ti
[Equation]
) ∈
I (cid:41) (37)
Since I is an open interval, ∃ ξ > 0 such that we can find a subinterval I 0 ⊆
I such that the distance between I 0 and the boundaries of I (if I is bounded on either side) is no less than ξ , i.e.,
[Equation]
[Equation]
In particular, we can choose ξ =
( I r − I l ) and I
= ( I l + ξ, I r − ξ ) .
Then there is
[Equation]
and so min
[Equation]
Thus, we have ˙
[Equation]
Meanwhile, since
[Equation] [Equation]
there is m
[Equation] [Equation]
Therefore
[Equation]
[Equation]
Since ∀ ξ ∈ R
[Equation]
there is
[Equation]
45) we derive that, ∀ a ∈ [ n ] ,
[Equation]
where we set C 1 = λ max ( G ) ( λ min ( G ))
− 12 ξ − 1 > 0 for simplicity.
As a consequence,
[Equation]
Define
[Equation]
Note that at t = 0 , there is ˜ η t = min a ∈ [ n ] (cid:110) 1 m (cid:80) mi = [Equation]
[-=1 1 h 
-]x [Equation]
) ∈ [Equation]
0 (cid:111) .
Then, on one hand, we know from (40) and (47) that ∀ t ≥ 0 , [Equation] [Equation]
[Equation]
Hence, (41) implies that
[Equation]
On the other hand, by the definition of ˜ η t ,
[Equation]
where we set C 2 = 2 − 12 C 1 ( λ min [Equation]
) − [Equation]
for simplicity.
Therefore [Equation] when η t > 0 ,
[Equation]
[Equation]
which implies that
[Equation]
and so ∀ t ≥ 0 ,
[Equation]
E P ROOF OF T HEOREM
To apply Lemma 1, we need two additional lemmas to be proved in Appendix E.1 and E.2.
The first one guarantees that the loss value at initialization, L 0 , is upper-bounded with high probability:
Lemma 2.
∀ δ > 0 , if m ≥ Ω (cid:0)
[Equation]
c
[Equation]
log
[Equation] nδ − [Equation] G max (cid:1) , then with probability at least 1 − δ , there is
[Equation]
[Equation] The second one proves that ˜ η 0 is lower-bounded with high probability, which heuristically says that there is indeed a nontrivial proportion of neurons in the central part of the active region of σ , for every a ∈ [ n ] :
Lemma 3.
∀ δ > 0 , if m ≥
[Equation]
2 , then with probability at least 1 − δ , there is
[Equation] [Equation]
where K ( I, λ 1 , λ 2 ) =
[Equation]
√
[Equation]
πλ 2 ( I
[Equation]
−
[Equation]
−
[Equation] {| [Equation] | , |
[Equation]
(cid:111) is a positive number that depends on
I , λ 1 and λ 2 .
With these two concentration lemmas, we deduce that, ∀ δ > 0 , if m ≥ Ω (cid:0) (1 + ˆ c 2 ) log (cid:0) 2 nδ (cid:1)(cid:1) , then with probability at least δ , there is ∀ t ≥ 0 ,
[Equation] [Equation]
where K ( I, G min , G max ) is defined as in Lemma 3.
Therefore, if our choice of ˆ c satisfies
[Equation]
then there is ∀ t ≥ 0 ,
[Equation]
in which case (50) gives
[Equation] [Equation]
which will allow us to finally conclude that
[Equation]
[Equation]
Note that (60) establishes a PL condition.
Several other convergence analyses of NNs have also relied on variants of the PL condition (25; 44; 74).
E. P ROOF OF L EMMA
Proof.
Since at initialization, (cid:8) c
i (cid:9) i ∈ [ m ] and (cid:8) W 0 ij (cid:9) i ∈ [ m ] ,j ∈
[ D ] are both sampled i.i.d.
and (cid:8) c i (cid:9) i ∈ [ m ] have mean zero, we know that ∀ a ∈ [ n ] , f
( x a ) = 1 m (cid:80) mi =1 c i σ (cid:0) h ti ( x a ) (cid:1) is the sample mean of i.i.d.
random variables with zero-mean.
Moreover, since (cid:8) W 0 ij (cid:9) i ∈ [ m ] ,j ∈
[ D ] are sampled from N (0 , 1) , we know that ∀ i ∈ [ m ] , the random variable c
i σ (cid:0) h ti ( x a ) (cid:1) is sub-Gaussian, with sub-Gaussian norm (68)
[Equation]
) where M SG > 0 is some absolute constant.
Thus, by Hoeffding’s inequality (68), ∀ a ∈ [ n ] , ∀ r > 0 ,
[Equation]
where K is some absolute constant
[Equation]
Hence, by union bound,
[Equation]
Thus, ∀ δ > 0 , if
[Equation]
then with probability at least 1 − δ , there is
[Equation]
E. P ROOF OF L EMMA
Since each W 0 ij are sampled i.i.d.
from N (0 , 1) , we know that ∀ a ∈ [ n ] , independently for each i ∈ [ m ] , h
i ( x a ) follows a Gaussian distribution with mean 0 and variance G aa .
Therefore, m (cid:88)
[Equation]
Hence, by Hoeffding’s inequality, ∀ a ∈ [ n ] , ∀ r > 0 ,
P (cid:32)m m (cid:88) i =
[Equation]
(cid:33)
[Equation] (cid:8) [Equation] (cid:9) (68) ∀ a ∈ [ n ] , choosing r = 12 π ( I 0 ; G aa ) , we then get
P (cid:32)
[Equation]
and so by union bound
[Equation]
Since ∀ b ∈ [ n ] , there is G min ≤ G bb ≤ G max ,
[Equation]
Letting K ( I, λ 1 , λ 2 ) = 1 6 √ 2 πλ 2 ( I r − I l ) exp (cid:110) −
[Equation]
{| I l | , | I r | [Equation] [Equation]
[Equation]
(cid:111) > 0 , we can then write
[Equation]
Thus, ∀ δ > 0 , if m ≥
[Equation]
[Equation] ( I [Equation] G min ,G max
[Equation]
) 2 , then with probability at least 1 − δ , it holds that ˜ η 0 > K
( I, G min , G max ) > 0 .
F P ROOF OF T HEOREM
By Condition 1, we know that ∀ δ > 0 , if D ≥ D min ( 12 δ, 12 λ min ( ¯ G )) , then with probability at least 1 − 12 δ , there is ∥ G − ¯ G ∥ 2 ≤ 12 λ min ( ¯ G ) , and hence λ min ( G ) ≥ 12 λ min ( ¯ G ) , G min ≥ 12 ¯ G min , λ max ( G ) ≤ λ max ( ¯ G ) + 12 λ min ( ¯ G ) ≤ 2 λ max ( ¯ G ) , and G max ≤ 2 ¯ G max .
We then perform the following analysis conditioned on the event that ∥ G − ¯ G ∥ 2 ≤ 12 λ min ( ¯ G ) .
Since the sampling of (cid:8) c i (cid:9) i ∈ [ m ] and (cid:8) W 0 ij (cid:9) i ∈ [ m ] ,j ∈
[ D ] is independent from the realization of G , we know from Lemma 3 that if m ≥ [Equation] K
[Equation]
G [Equation] ) 2 , then with probability at least 1 − 14 δ , there is [Equation]
[Equation]
From Lemma 2, we also know that if m ≥ Ω (cid:0) ˆ c 2 log (cid:0) nδ (cid:1) λ max (
[Equation]
G ) (cid:1)
[Equation]
ˆ c 2 log (cid:0) nδ (cid:1) λ max ( ¯ G ) (cid:1) , then with probability at least 1 − 14 δ , there is L 0 ≤ ∥ y ∥ 2 .
Therefore, in total, we know that with probability at least− δ , the following conditions all hold:
[Equation]
in which case, by applying Lemma 1 with G = G (1) , we get
[Equation]
[Equation] where K 1 ( λ 1 , λ 2 ) = 92 λ − 1 1 λ 2 K − 1 σ ′
( I r − I l ) > 0 .
Thus, by the definition of K 1 ( · , · )
[Equation]
we know that
[Equation]
and so ∀ t ≥ 0 ,
[Equation]
Therefore, if our choice of ˆ c satisfies
[Equation]
then there is ∀ t ≥ 0 ,
[Equation]
Hence, (50) implies that ∀ t ≥ 0 ,
[Equation]
and therefore L
[Equation]
G P ROOF OF T HEOREM
In view of Theorem 2, it is sufficient to verify that Condition 1 holds for D min ( δ, u ) =
Ω (cid:0) n 2 u − 2 log( nδ − 1 ) (cid:1) , which is given by the following lemma:
Lemma 4.
∀ δ ≥ 0 , if D ≥ Ω (cid:0) n 2 u − 2 log( nδ − 1 ) (cid:1) , then with probability at least 1 − δ , [Equation] G.
[Equation]
P ROOF OF L EMMA
Let Z be a random vector on R d with law given by π z , and then we can write ¯ G (1) ab =
E [ σ ( x ⊺ a Z ) σ ( x ⊺ b Z )] for a, b ∈ [ n ] .
By assumption, Z is sub-gaussian with sub-gaussian norm
[Equation]
[Equation]
Thus, ∀ a ∈ [ n ] , we have
[Equation] [Equation]
Hence [Equation] by Lemma 2.7.7 in (68)
[Equation]
we know that ∀ a, b ∈ [ n ] , σ ( x ⊺ a Z ) σ ( x ⊺ b Z ) is a sub-exponential random variable with sub-exponential norm [Equation]
[Equation]
Then, by Bernstein’s inequality (Theorem 2.8.1 in (68)), since each z j is sampled i [Equation] i [Equation]
d. from π z , we have that ∀ a, b ∈ [ n ] and ∀ u > 0 ,
[Equation]
where K > 0 is some absolute constant.
In other words, for any δ ′ > 0 , if
[Equation]
[Figure]
then we have (cid:12)(cid:12)(cid:12) G (1) ab − ¯ G (1) ab (cid:12)(cid:12)(cid:12) ≥ u with probability at least 1 − δ .
If we choose u = u ′ n and δ ′ =
δn 2 , then we get, if
[Equation]
then
P
[Equation]
Hence, with probability at least 1 − δ , we have
[Equation]
H F URTHER DETAILS OF THE NUMERICAL EXPERIMENTS
In our models, (cid:8) c i (cid:9) i ∈ [ m ] are sampled i.i.d.
from the Rademacher distribution µ c = 12 δ 1 + 12 δ − 1 , (cid:8) z j (cid:9) j ∈
[ D ] are sampled i.i.d.
from N (0 , I d ) , and (cid:8) W ij (cid:9) i ∈ [ m ] ,j ∈
[ D ] are initialized by sampling i.i.d.
from N (0 , 1) .
In the model under NTK scaling, we additionally symmetrize the model at initialization according to the strategy used in (16) to ensure that the function value at initialization does not blow up.
We choose train the models using 50000 steps of (full-batch) GD with step size δ = 1 .
When a test error is computed, we use a test set of size 500 generated by sampling i.i.d.
from the same distribution as the training set.
The experiments are run with NVIDIA GPUs (1080ti and titanrtx).
H. E XPERIMENT
We choose σ to be tanh .
For each choice of n , we run the experiment with 5 different random seeds, and Figure 1 plots the evolution of the training loss during GD averaged over the 5 runs with m = 8192 .
Figure 5 is the same plot as Figure 1 except for with m = 4096 .
We see that the two two plots agree well.
H. E XPERIMENT
We choose σ to be ReLU.
For each choice of n and each of the two models, we experiment withdifferent random seeds, and Figure 2 plots the test error at the 50000 GD step averaged over theruns ± its standard deviation.
[Figure]
In Figure 6, we plot the evolution of the training loss and test error over GD for the two models with m = 2048 and 8192 and different choices of n , averaged over 5 runs with different random seeds.
We see in particular that the difference between the two choices of m is negligible, suggesting that it is unlikely to obtain performance improvements with further over-parameterization.
H. E XPERIMENT
We choose σ to be ReLU and input dimension d = 50 .
We use a training set of size n = 600 for the results reported in Figure 4.
The data set is inspired by (69): We sample both the training and the test set i.i.d.
from the distribution ( x , y ) ∼ D on R d +1 , under which the joint distribution of ( x 1 , x 2 , y ) is
[Equation]
and x 3 , ..., x d each follow the uniform distribution in [ − 1 , 1] , independently from each other as well as x 1 , x 2 and y .
Figures 7 and 8 are the same plot as Figure 4 except for setting n = 400 and 800 , respectively.
We see that as n increases, test error improves for all three models, and the features learned by the pseudo- 3 L
NN (in our scaling) become “sharper” visually.
[Figure] [Figure]