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: Authors present a theoretical understanding for the "transition to linearity" phenomenon of wide neural networks with linear output layer under gradient descent. Presented line of thought develops from the perspective where one considers a neural network as a recursively-built assembly of sub-models which correspond to individual neurons from the preceding layers. Authors prove that the transition to linearity phenomenon when the number of sub-models converge to infinity (i.e., the network becomes wider) is a result of assembling weak sub-models that do not dominate the assembly from network initialization.