On feature learning in shallow and multi-layer neural networks with global convergence guarantees

1 INTRODUCTION . 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 focus on the over-parameterized regime , where 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 equipped 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 form of lazy training that excludes the learning of features or representations [ 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 ] . Third , while several works have proposed to extend the MF formulation to deep ( i.e . multi-layer ) NNs [ 4 , 62 , 52 , 23 , 57 ] , there is less concensus on what the right model should be than for the shallow case . More generally , we still lack a model for the optimization of shallow and multi-layer NNs that goes beyond lazy training while admitting fast global convergence . The aim of this work is to obtain theoretical guarantees of fast global convergence of feature learning in both shallow and a type of multi-layer NNs . 1.1 SUMMARY 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 : 1 . 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 . 2 . For a multi-layer NN where we only train 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.2 RELATED 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 meanfield 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 can not 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 quadraticTaylor approximations , and furthermore the convergence rate of the training is not known . Ref . [ 73 ] proposes an alternative scaling under which there exists an infinite-width limit of deep NNs that exhibits feature learning , but properties of its training dynamics is still not well-understood . Our pseudo-multi-layer NN models adopt the same scaling , and our focus is on proving non-asymptotic convergence guarantees . 2 PROBLEM SETUP . 2.1 MODEL . We summarize our notations in Appendix A . Let Ω ⊆ Rd denote the input space , and let x = [ x1 , ... , xd ] ᵀ ∈ Ω denote a generic input data vector . A shallow NN model under the MF scaling can be written as : f ( x ) = 1 m m∑ i=1 ciσ ( 1√ d d∑ j=1 Wijxj ) , ( 1 ) where m is the width , W ∈ Rm×d and c = [ c1 , ... , cm ] ∈ Rm 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 : f ( x ) = 1 m m∑ i=1 ciσ ( hi ( x ) ) , ( 2 ) ∀i ∈ [ m ] : hi ( x ) = 1√ D D∑ j=1 Wijφj ( x ) , ( 3 ) where W ∈ Rm×D and c = [ c1 , ... , cm ] ∈ Rm are parameters of the model , and φ1 , ... , φD are a set of functions from Ω to R that we call the embedding . Each of h1 , ... , hm 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 ) ] ᵀ ∈ RD . 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 ) = xj , ∀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 RD . 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 zj i.i.d . in Rd and set φj ( x ) = σ ( 1√ d zᵀj x ) , equivalent to setting φ1 , ... , φm as the hidden-layer activations of a shallow NN with randomly-initialized first-layer weights . Then , the model becomes f ( x ) = 1 m m∑ i=1 ciσ ( hi ( x ) ) , ( 4 ) ∀i ∈ [ m ] : hi ( x ) = 1√ D D∑ j=1 Wijσ ( 1√ d zᵀj x ) . ( 5 ) Thus , we obtain a 3-layer feed-forward NN whose first-layer weights are random and fixed , and we call it a partially-trained 3-layer ( p-3L ) 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 .

Summary: This paper studies the optimization of a pseudo-three-layer network. The first (output) layer is scaled by a law of large number (LNN) scaling, while the second layer is scaled by a central limit theorem (CLT)-scaling. The inputs are fixed random feature embeddings. Only the weights of the second layer are learnable. They proved that when the second layer is sufficiently over-parameterized (wrt the number of samples), GD flow converges to a zero-loss solution exponentially fast. In particular, the authors claim that, unlike the NTK regime, their setting exhibits feature learnings.