Transition to Linearity of Wide Neural Networks is an Emerging Property of Assembling Weak Models

1 INTRODUCTION . Success of gradient descent methods for optimizing neural networks , which generally correspond to highly non-convex loss functions , has long been a challenge to theoretical analysis . A series of recent works including Du et al . ( 2018 ; 2019 ) ; Allen-Zhu et al . ( 2019 ) ; Zou et al . ( 2018 ) ; Oymak & Soltanolkotabi ( 2020 ) showed that convergence can indeed be shown for certain types of wide networks . Remarkably , Jacot et al . ( 2018 ) demonstrated that , when the network width goes to infinity , the Neural Tangent Kernel ( NTK ) of the network becomes constant during training with gradient flow , a continuous time limit of gradient descent . Based on that Lee et al . ( 2019 ) showed that the training dynamics of gradient flow in the space of parameters for the infinite width network is equivalent to that of a linear model . As discussed in ( Liu et al. , 2020a ) , the constancy of the Neural Tangent Kernel stems from the fact that wide neural networks with linear output layer “ transition to linearity ” as their width increases to infinity . The network becomes progressively more linear inO ( 1 ) -neighborhoods around the network initialization , as the network width grows . Specifically , consider the neural network as a function f ( θ ) of its parameters θ and write the Taylor expansion with the Lagrange remainder term : f ( θ ) = f ( θ0 ) +∇f ( θ0 ) T ( θ − θ0 ) + 1 2 ( θ − θ0 ) TH ( ξ ) ( θ − θ0 ) , ( 1 ) where H is the Hessian , the second derivative matrix of f , and ξ is a point between θ and θ0 . The “ transition to linearity ” is the phenomenon when the quadratic term 12 ( θ− θ0 ) TH ( ξ ) ( θ− θ0 ) tends to zero in an O ( 1 ) ball around θ0 as the network width increases to infinity . In particular , transition to linearity provides a method for showing that the square loss function of wide neural networks satisfy the Polyak-Lojasiewicz ( PL ) inequality which guarantees convergence of ( Stochastic ) Gradient Descent to a global minimum ( Liu et al. , 2020b ) . While the existing analyses demonstrate transition to linearity mathematically , the underlying structure of that phenomenon does not appear to be fully clear . What structural properties account for it ? Are they specific to neural networks , or also apply to more general models ? In this paper , we aim to explain this phenomenon from a more structural point of view and to reveal some underlying mechanisms . We provide a new perspective in which a neural network , as well as each of its hidden layer neurons before activation , can be viewed as an “ assembly model ” , built up by linearly assembling a set of sub-models , corresponding to the neurons from the previous layer . Specifically , a ( l + 1 ) -layer pre-activation neuron α̃ ( l+1 ) = 1√ m m∑ i=1 wiα ( l ) i , as a linear combination of the l-th layer neurons α ( l ) i , is considered as an assembly model , whereas the l-layer neurons α ( l ) i are considered as sub-models . Furthermore , each pre-activation α̃ ( l ) i is also an assembly model constructed from ( l − 1 ) -th layer neurons α ( l−1 ) i . In this sense , the neural network is considered as a multi-level assembly model . To show O ( 1 ) -neighborhood linearity , we prove that the quadratic term in the Taylor expansion Eq . ( 1 ) vanishes in O ( 1 ) -neighborhoods of network initialization , as a consequence of assembling sufficiently many sub-models . To illustrate the idea of assembling , we start with a simple case : assembling independent sub-models in Section 2 . The key finding is that , as long as the assembly model is not dominated by one or a few sub-models , the quadratic term in the Taylor expansion Eq . ( 1 ) becomes small in O ( 1 ) neighborhoods of the parameter space when the number of sub-models m is tends to infinity . This means that as m increases to infinity , the assembly model becomes a linear function of parameters . Since we put almost no requirements on the form of sub-models , it is the assembling process that results in the linearity of the assembly model . This case includes two-layer neural networks as examples . For deep neural networks ( Section 3 ) , we consider a neural network as a multi-level assembly model each neuron is considered as an assembly model constructed iteratively from all the neurons from the previous layer . We then follow an inductive argument : using near-linearity of previous-layer pre-activation neurons to show the linearity of next-layer pre-activation neurons . Our key finding is that , when the network width is large , the neurons within the same layer become essentially independent to each other in the sense that their gradient directions are orthogonal in the parameter space . This orthogonality allows for the existence of a new coordinate system such that these neuron gradients are parallel to the new coordinate axes . Within the new coordinate system , one can apply the argument of assembling independent sub-models , and obtain the O ( 1 ) -neighborhood linearity of the pre-activation neurons , as well as the output of the neural network . We further point out that this assembling viewpoint and the O ( 1 ) -neighborhood linearity can be extended to more general neural network architectures , e.g. , DenseNet . We end this section by commenting on a few closely related concepts and point out some of the differences . Boosting . Boosting ( Schapire , 1990 ; Freund , 1995 ) , which is a popular method that combines multiple “ weak ” models to produce a powerful ensembe model , has a similar form with the assembly model , Eq . ( 4 ) , see , e.g. , Friedman et al . ( 2000 ) . However , we note a few key differences between the two . First , in boosting , each “ weak ” model is trained separately on the dataset and the coefficients ( i.e. , vi in Eq . ( 4 ) ) of the weak models are determined by the model performance . In training an assembly model , the “ weak ” sub-models ( i.e. , hidden layer neurons ) , are never directly evaluated on the training dataset , and the coefficients vi are considered as parameters of the assembly model and trained by gradient descent . Second , in boosting , different data samples may have different sample weights . In assembling , the data samples always have a uniform weight . Bagging and Reservoir computing . Bagging ( Breiman , 1996 ) and Reservoir computing ( Jaeger , 2001 ) are two other ways of combining multiple sub-models to build a single model . In bagging , each sub-model is individually trained and the ensemble model is simply the average or the max of the sub-models . In reservoir computing , each sub-model is fixed , and only the coefficients of the linear combination are trainable . Notation . We use bold lowercase letters , e.g. , v , to denote vectors , capital letters , e.g. , W , to denote matrices . We use ‖ · ‖ to denote the Euclidean norm of vectors and spectral norm ( i.e. , matrix 2-norm ) for matrices . We denote the set { 1 , 2 , · · · , n } as [ n ] . We use ∇wf to denote the partial derivative of f with respect to w. When w represents all of the parameters of f , we omit the subscript , i.e. , ∇f . 2 ASSEMBLING INDEPENDENT MODELS . In this section , we consider assembling a sufficiently large number m of independent sub-models and show that the resulting assembly model is ( approximately ) linear in O ( 1 ) -neighborhood of the model parameter space . Ingredients : sub-models . Let ’ s consider a set of m ( sub- ) models { gi } mi=1 , where each model gi , with a set of model parameters wi ∈ Di ⊂ Pi = Rpi , takes an input x ∈ Dx ⊂ Rd and outputs a scalar prediction gi ( wi ; x ) . Here Pi is the parameter space of sub-model gi and Dx is the domain of the input . By independence , we mean that any two distinct models gi and gj share no common model parameters : Pi ∩ Pj = { 0 } , ∀i 6= j ∈ [ m ] . ( 2 ) In this sense , the change of gi ’ s output , as a result of change of it parameters wi , does not affect the outputs of the other models gj where j 6= i . In addition , we require that there is no dominating sub-models over the others . Specifically , we assume the following in this section : Assumption 1 ( No dominating sub-models ) . There exists a constant c ∈ ( 0 , 1 ) independent of m such that , for any parameter setting { wi : wi ∈ Di } mi=1 and input x ∈ Dx , median [ |g1 ( w1 ; x ) | , · · · , |gm ( wm ; x ) | ] max [ |g1 ( w1 ; x ) | , · · · , |gm ( wm ; x ) | ] ≥ c. ( 3 ) Further more , we assume that max [ |g1 ( w1 ; x ) | , · · · , |gm ( wm ; x ) | ] ∈ [ a , b ] ⊂ R for some constants a > 0 and b > 0 . Remark 1 . This assumption guarantees that most of the outputs of the sub-models are at the same order , typically O ( 1 ) . It makes sure that the resulting assembly model is not dominated by one or a minor portion of the sub-models . This is typically seen for the neurons of neural networks . We further make the following technical assumption , which is common in the literature . Assumption 2 ( Twice differentiablity and smoothness ) . Each sub-model gi is twice differentiable with respect to the parameters wi , and is β-smooth in the parameter space : there exists a positive number β such that , for any i ∈ [ m ] , ‖∇gi ( wi ; x ) −∇gi ( w′i ; x ) ‖ ≤ β‖wi −w′i‖ . This assumption makes sure that for each gi , the gradient is well-defined and its Hessian has a bounded spectral norm . Assembly model . Based on these sub-models { gi } mi=1 , we construct an assembly model ( or supermodel ) f , using linear combination as follows : f ( θ ; x ) : = 1 s ( m ) m∑ i=1 vigi ( wi ; x ) , ( 4 ) where vi is the weight of the sub-model gi , and 1/s ( m ) , as a function of m , is the scaling factor . Here , we denote θ as the set of parameters of the assembly model , θ : = ( w1 , . . . , wm ) . The total number of parameters that f has is p = ∑m i=1 pi . Typical choices for the weights vi are setting vi = 1 for all i ∈ [ m ] , or random i.i.d . drawing vi from some zero-mean probability distribution . Remark 2 . For the ease of analysis and simplicity of notation , we assumed that the outputs of assembly model f and sub-models gi are scalars . It is not difficult to see that our analysis below also applies to the scenarios of finite dimensional model outputs . The scaling factor 1/s ( m ) . The presence of the scaling factor 1/s ( m ) is to keep the output of assembly model f at the order O ( 1 ) w.r.t . m. In general , we expect that s ( m ) grows with m. In particular , when vi are chosen i.i.d . from a probability distribution with mean zero , e.g. , { −1 , 1 } , the sum in Eq . ( 4 ) is expected to be of the order √ m , under the Assumption 1 . In this case , the scaling factor 1/s ( m ) = O ( 1/ √ m ) . Now , we will show the the assembly model f has a small quadratic term in its Taylor expansion , when the size of the set of sub-models i.e . m , is sufficiently large . Theorem 1 . Consider the assembly model f constructed in Eq . ( 4 ) with each vi either set to be 1 or randomly drawn from { −1 , 1 } , and suppose Assumption 1 and 2 hold . Given a positive number R > 0 and a parameter setting θ0 ∈ Rp , for any θ ∈ Rp such that ‖θ− θ0‖ ≤ R , the absolute value of the quadratic term in Taylor expansion Eq . ( 1 ) is bounded by : ∣∣∣∣12 ( θ − θ0 ) TH ( ξ ) ( θ − θ0 ) ∣∣∣∣ ≤ βR22s ( m ) . ( 5 ) Proof . In what follows , we don ’ t explicitly write out the dependence on θ and x in the Hessian H : = ∂ 2f ∂θ2 , for simplicity of notation . We further denote Hgi : = ∂2gi ∂w2i as the Hessian of gi . We decompose the assembly model Hessian H as a linear combination of Hgi . By the definition of the assembly model in Eq . ( 4 ) , an arbitrary entry Hjk of the Hessian can be written as : Hjk = 1 s ( m ) ∑m i=1 vi ∂2gi ∂θj∂θk . Here θj , θk are two individual parameters of the assembly model f . Note that ∂2gi ∂θj∂θk is non-zero , only if both θj and θk are parameters of sub-model gi , i.e. , θj , θk ∈ Pi . Hence , the assembly model Hessian can be decomposed as a linear combination of sub-model Hessians : H = 1s ( m ) ∑m i=1 viHgi . Therefore , the quadratic term becomes 1 s ( m ) m∑ i=1 vi [ 1 2 ( θ − θ0 ) THgi ( ξ ) ( θ − θ0 ) ] = 1 s ( m ) m∑ i=1 vi [ 1 2 ( wi −wi,0 ) THgi ( ξ ) ( wi −wi,0 ) ] and its absolute value is bounded by∣∣∣1 2 ( θ − θ0 ) TH ( ξ ) ( θ − θ0 ) ∣∣∣ ≤ 1 2s ( m ) m∑ i=1 |vi| · ‖Hgi‖ · ‖wi −wi,0‖2 ≤ β 2s ( m ) m∑ i=1 ‖wi −wi,0‖2 In the second inequality , we used that |vi| = 1 and that gi is β-smooth . Because of the independence of the sub-models as seen in Eq . ( 2 ) , the summation in the above equation becomes ‖θ−θ0‖2 , which is bounded by R2 , as stated in the theorem condition . Therefore , we conclude the theorem . It is important to note that β is a constant and 1/s ( m ) = O ( 1/ √ m ) . Then we have the following corollary in the limiting case . Corollary 1 . Consider the assembly model f under the same setting as in Theorem 1 . If the number of sub-models m increases to infinity , then for all parameter θ0 and input x , the quadratic term∣∣∣1 2 ( θ − θ0 ) TH ( ξ ) ( θ − θ0 ) ∣∣∣→ 0 , ( 6 ) as long as θ is within an O ( 1 ) -neighborhood of θ0 , i.e. , ‖θ − θ0‖2 ≤ R for some constant R > 0 . Example : Two-layer neural networks . A good example of this kind of assembly model is the two-layer neural network . A two-layer neural network is mathematically defined as : f ( W , v ; x ) = 1√ m m∑ i=1 uiσ ( w T i x ) , ( 7 ) wherem is the number of hidden layer neurons , σ ( · ) is the activation function , x ∈ Rd is the network input , and W ∈ Rm×d and u ∈ Rm are the parameters for the first and second layer , respectively . Here , we can view the i-th hidden neuron and all the parameters wi and ui that connect to it as the i-th sub-model : gi = uiσ ( wTi x ) . We see that these sub-models do not share parameters , and each sub-model has d+ 1 parameters . In addition , the weights of the sub-models are all 1 . By Corollary 1 , the two-layer neural network becomes a linear model in any O ( 1 ) -neighborhoods , as the network widthm increases to infinity . This is consistent with the previous observation that a two-layer neural network transitions to linearity Liu et al . ( 2020a ) . Gaussian distributed weights vi . Another common way to set the weights vi of the sub-models is to independently draw each vi fromN ( 0 , 1 ) . In this case , vi is unbounded , but with high probability the quadratic term is still O ( log ( m ) / √ m ) . Please see the detailed analysis in Appendix A . Hierarchy of assembly models . In principle , we can consider the assembly model f , together with multiple similar models independent to each other , as sub-models , and construct a higher-level assembly model . Repeating this procedures , we can have a hierarchy of assembly models . Our analysis above also applies to this case and each assembly model is also O ( 1 ) -neighborhood linear , when the number of its sub-models is sufficiently large . However , one drawback of this hierarchy is that the total number of parameters of the highest level assembly model increase exponentially with the number of levels . We will see shortly that wide neural networks , as a hierarchy of assembly models , allows overlapping between sub-models , but still keeping the O ( 1 ) -neighborhood linearity .

Summary: This paper aims to explain the phenomenon that neural networks with infinite width tend to be linear in the neighborhood of initial optimization points, from an assemble model perspective. The key point is to prove that the quadratic term from Taylor expansion disappears when the width approaches infinity. The authors first proved in the simple two-layer case and then extend the results to deep neural networks with L layers.