Dynamically Stable Infinite-Width Limits of Neural Classifiers

1 INTRODUCTION . For a couple of decades neural networks have proved to be useful in a variety of applications . However , their theoretical understanding is still lacking . Several recent works have tried to simplify the object of study by approximating a training dynamics of a finite-width neural network with its limit counterpart in the limit of a large number of hidden units ; we refer it as an `` infinite-width '' limit . The exact type of the limit training dynamics depends on how hyperparameters of the training dynamics scale with width . In particular , two different types of limit models have been already extensively discussed in the literature : an NTK model ( Jacot et al. , 2018 ) and a mean-field limit model ( Mei et al. , 2018 ; 2019 ; Rotskoff & Vanden-Eijnden , 2019 ; Sirignano & Spiliopoulos , 2020 ; Chizat & Bach , 2018 ; Yarotsky , 2018 ) . A recent work ( Golikov , 2020 ) attempted to provide a link between these two different types of limit models by building a framework for choosing a scaling of hyperparameters that lead to a `` well-defined '' limit model . Our work is the next step in this direction . We study infinite-width limits for networks with a single hidden layer trained to minimize cross-entropy loss with gradient descent . Our contributions are following . 1 . We develop a framework for reasoning about scaling of hyperparameters , which allows one to infer scaling parameters that allow for a dynamically stable model evolution in the limit of infinite width . This framework allows us to derive both mean-field and NTK limits that have been extensively studied in the literature , as well as the `` intermediate limit '' introduced in Golikov ( 2020 ) . 2 . Our framework demonstrates that there are only 13 distinct stable model evolution equations in the limit of infinite width that can be induced by scaling hyperparameters of a finite-width model . Each distinct limit model corresponds to a region ( two- , one- , or zero-dimensional ) of a green band of the Figure 1 , left . 3 . We consider a list of properties that are statisfied by the evolution of finite-width models , but not generally are for its infinite-width limits . We demonstrate that mean-field and NTK limit models , as well as `` sym-default '' limit model which was not discussed in the literature previously , are special in the sense that they satisfy most of these properties among all limit models induced by hyperparameter scalings . We propose a model modification that allows for all of these properties in the limit of infinite width and call the corresponding limit `` initialization-corrected mean-field limit ( IC-MF ) '' . 4 . We discuss the ability of limit models to approximate the training dynamics of finite-width ones . We show that our proposed IC-MF limiting model is the best among all other possible limit models . While our present analysis is restricted to networks with a single hidden layer , we discuss a high-level plan for generalizing it to deep nets , as well as an expected outcome of this research program , in App . H . 2 TRAINING A ONE HIDDEN LAYER NET WITH SGD . Here we consider training a one hidden layer net fd with d hidden units with SGD . We assume the hyperparameters , namely , initialization variances and learning rates , are scaled as power-laws of d. Each scaling induces a limit model f∞ = limd→∞ fd . We present a notion of dynamical stability , which states that the change of logits after a single gradient step is comparable to logits themselves . We derive a necessary condition for dynamical stability in terms of the power-law exponents of hyperparameters . We then present a list of conditions that divide the class of scalings into 13 subclasses ; each subclass corresponds to a unique distinct limit model . Consider a one hidden layer network : f ( x ; a , W ) = aTφ ( WTx ) = d∑ r=1 arφ ( w T r x ) , ( 1 ) where x ∈ Rdx , W = [ w1 , . . . , wd ] ∈ Rdx×d , and a = [ a1 , . . . , ad ] T ∈ Rd . We assume a nonlinearity to be real analytic and asymptotically linear : φ ( z ) = Θz→∞ ( z ) . Such a nonlinearity can be , e.g . `` leaky softplus '' : φ ( z ) = ln ( 1 + ez ) − α ln ( 1 + e−z ) for α > 0 . This is a technical assumption introduced to simplify proofs . Note that we have used traditional leaky ReLUs in our experiments : see App . F for details . We assume the loss function ` ( y , z ) to be the standard binary cross-entropy loss : ` ( y , z ) = ln ( 1 + e−yz ) , where labels y ∈ { −1 , 1 } . The data distribution loss is defined as L ( a , W ) = E x , y∼D ` ( y , f ( x ; a , W ) ) . We assume that the data distribution D does not depend on width d. Weights are initialized with isotropic gaussians with zero means : w ( 0 ) r ∼ N ( 0 , σ2wI ) , a ( 0 ) r ∼ N ( 0 , σ2a ) ∀r = 1 . . . d. The evolution of weights is driven by the stochastic gradient descent ( SGD ) : ∆θ ( k ) = θ ( k+1 ) − θ ( k ) = −ηθ ∂ ` ( y ( k ) θ , f ( x ( k ) θ ; a , W ) ) ∂θ , ( x ( k ) θ , y ( k ) θ ) ∼ D , ( 2 ) where θ is either a or W . We assume that gradients for a and W are estimated using independent data samples ( x ( k ) a , y ( k ) a ) and ( x ( k ) w , y ( k ) w ) . While this assumption is indeed non-standard , we note that corresponding stochastic gradients still give unbiased estimates for true gradients . Moreover , we have used either full-batch GD or standard mini-batch SGD in our experiments : see App . F for details . Define : â ( k ) r = a ( k ) r σa , ŵ ( k ) r = w ( k ) r σw , η̂a = ηa σ2a , η̂w = ηw σ2w . ( 3 ) Then the dynamics transforms to : ∆θ̂ ( k ) r = η̂θ ∂ ` ( y ( k ) θ , f ( x ( k ) θ ; σaâ , σwŴ ) ) ∂θ̂r , ( x ( k ) θ , y ( k ) θ ) ∼ D , ( 4 ) while scaled initial conditions become : â ( 0 ) r ∼ N ( 0 , 1 ) , ŵ ( 0 ) r ∼ N ( 0 , I ) ∀r = 1 . . . d. By expanding gradients , we get the following : ∆â ( k ) r = −η̂aσa∇ ( k ) fd ` ( x ( k ) a , y ( k ) a ) φ ( σwŵ ( k ) , T r x ( k ) a ) , â ( 0 ) r ∼ N ( 0 , 1 ) , ( 5 ) ∆ŵ ( k ) r = −η̂wσaσw∇ ( k ) fd ` ( x ( k ) w , y ( k ) w ) â ( k ) r φ ′ ( . . . ) x ( k ) w , ŵ ( 0 ) r ∼ N ( 0 , I ) , ( 6 ) ∇ ( k ) fd ` ( x , y ) = ∂ ` ( y , z ) ∂z ∣∣∣∣ z=f ( k ) d ( x ) = −y 1 + exp ( f ( k ) d ( x ) y ) , f ( k ) d ( x ) = σa d∑ r=1 â ( k ) r φ ( σwŵ ( k ) , T r x ) . Without loss of generality assume σw = 1 ( we can rescale inputs x otherwise ) . We shall omit a subscript of σa from now on . Assume hyperparameters that drive the dynamics obey power-law dependence on d : σ ( d ) = σ∗ × ( d/d∗ ) qσ , η̂a ( d ) = η̂∗a × ( d/d∗ ) q̃a , η̂w ( d ) = η̂∗w × ( d/d∗ ) q̃w . ( 7 ) Given this , a network of width d∗ has hyperparameters σ∗ and η̂∗a∨w . Here and then we write `` a∨w '' meaning `` a or w '' . This assumption is quite natural : for He initialization ( He et al. , 2015 ) commonly used in practice σ ∝ d−1/2 , while we keep learning rates in the original parameterization constant while changing width by default : ηa∨w = const , which implies η̂a ∝ d and η̂w ∝ d0 . On the other hand , NTK scaling ( Jacot et al. , 2018 ; Lee et al. , 2019 ) requires scaled learning rates to be constants : η̂a∨w ∝ d0 . Scaling exponents ( qσ , q̃a , q̃w ) together with proportionality factors ( d∗ , σ∗ , η̂∗a , η̂ ∗ w ) define a limit model f ( k ) ∞ ( x ) = limd→∞ f ( k ) d ( x ) . We call a model `` dynamically stable in the limit of large width '' if it satisfies the following condition which we state formally in Appendix A : Condition 1 ( informal version of Condition 4 in Appendix A ) . Let ∆f ( k ) d ( x ) = f ( k+1 ) d ( x ) −f ( k ) d ( x ) . ∃kbalance ∈ N : ∀k ≥ kbalance ∆f ( k ) d fkbalanced stays finite for large d. Roughly speaking , this condition states that the change of logits after a single step is comparable to logits themselves . This means that the model learns . Note that this condition is weaker than the one used in Golikov ( 2020 ) , because it allows logits to vanish or diverge with width . Such situations are fine , because only logit signs matter for the binary classification . For simplicity assume q̃a = q̃w = q̃ . We prove the following in Appendix B.1 : Proposition 1 . Suppose q̃a = q̃w = q̃ and D is a continuous distribution . Then Condition 1 requires qσ + q̃ ∈ [ −1/2 , 0 ] to hold . This statement gives a necessary condition for growth rates of σ and η̂ to lead to a well-defined limit model evolution . This condition corresponds to a band in ( qσ , q̃ ) -plane : see Figure 1 , left . We refer it as a `` band of dynamical stability '' . Each point of this band corresponds to a dynamically stable limit model evolution . We present several conditions that separate the dynamical stability band into regions . We then show that each region corresponds to a single limit model evolution . We start with defining tangent kernels . Since φ is smooth , we have : ∆f ( k ) d ( x ) = f ( k+1 ) d ( x ) − f ( k ) d ( x ) = d∑ r=1 ∂fd ( x ) ∂θ̂r ∣∣∣∣ θ̂r=θ̂ ( k ) r ∆θ̂ ( k ) r +Oη̂∗a∨w→0 ( η̂ ∗ aη̂ ∗ w + η̂ ∗,2 w ) = = −η̂∗a∇ ( k ) fd ` ( x ( k ) a , y ( k ) a ) K ( k ) a , d ( x , x ( k ) a ) − η̂∗w∇ ( k ) fd ` ( x ( k ) w , y ( k ) w ) K ( k ) w , d ( x , x ( k ) w ) +O ( η̂ ∗ aη̂ ∗ w+ η̂ ∗,2 w ) , ( 8 ) where we have defined kernels : K ( k ) a , d ( x , x ′ ) = ( d/d∗ ) q̃aσ2 d∑ r=1 φ ( ŵ ( k ) , Tr x ) φ ( ŵ ( k ) , T r x ′ ) , ( 9 ) K ( k ) w , d ( x , x ′ ) = ( d/d∗ ) q̃wσ2 d∑ r=1 |â ( k ) r |2φ′ ( ŵ ( k ) , Tr x ) φ′ ( ŵ ( k ) , Tr x′ ) xTx′ . ( 10 ) Here we deviate from the traditional definition of tangent kernels ( e.g . from Jacot et al . ( 2018 ) ) in embedding learning rate growth factors into kernels . This is done for avoiding 0 ×∞ ambiguity when η̂a∨w grows with d while σ vanishes so that `` a learning rate times a kernel '' stays finite . This is the case for the mean-field scaling : η̂a∨w ∝ d , while σ ∝ d−1 . While for the NTK scaling kernels stop evolving with k in the limit of large d , this is not the case generally . Indeed , for the mean-field scaling mentioned above we have : K ( k ) a , d ( x , x ′ ) = σ∗,2 ( d/d∗ ) −1 d∑ r=1 φ ( ŵ ( k ) , Tr x ) φ ( ŵ ( k ) , T r x ′ ) . ( 11 ) Similarly to the NTK case , the kernel above converges due to the Law of Large Numbers , however in contrast to the NTK case the weights evolve in the limit : ŵ ( k ) r 9 ŵ ( 0 ) r . This is due to the fact that weight increments are proportional to η̂wσ which is ∝ d0 for the mean-field scaling but ∝ d−1/2 for the NTK one . For this reason , similarly to model increments ∆f ( k ) d we define kernel increments : ∆K ( k ) a∨w , d ( x , x ′ ) = K ( k+1 ) a∨w , d ( x , x ′ ) −K ( k ) a∨w , d ( x , x ′ ) . ( 12 ) Condition 2 ( informal version of Condition 5 in Appendix A ) . Following conditions separate the band of dynamical stability ( Figure 1 , left ) : 1. f ( 0 ) d stays finite for large d. 2 . K ( 0 ) a∨w , d stays finite for large d. 3 . K ( 0 ) a∨w , d/f ( 0 ) d stays finite for large d. 4 . ∆K ( 0 ) a∨w , d/K ( 0 ) a∨w , d stays finite for large d. We prove the following in Appendix B.2 : Proposition 2 ( Separating conditions ) . Given Condition 1 , Condition 2 reads as , point by point : 1 . A limit model at initialization is finite : qσ + 1/2 = 0 . 2 . Tangent kernels at initialization are finite : 2qσ + q̃ + 1 = 0 . 3 . Tangent kernels and a limit model are of the same order at initialization : qσ + q̃ + 1/2 = 0 . 4 . Tangent kernels start to evolve : qσ + q̃ = 0 . We have also checked this Proposition numerically for limit models discussed below : see Figure 1 , right . Each condition corresponds to a straight line in the ( qσ , q̃ ) -plane : see Figure 1 , left . These four lines divide the well-definiteness band into 13 regions : three are two-dimensional , seven are onedimensional , and three are zero-dimensional . In Appendix C we show that each region corresponds to a single distinct limit model evolution ; we also list corresponding evolution equations . Note that a segment ( a one-dimensional region ) that corresponds to the Condition 2-2 exactly coincides with a family of `` intermediate scalings '' introduced in Golikov ( 2020 ) .

Summary: This paper proposed a general framework to derive different stable limiting behaviors of the dynamics of two-layers neural networks, under different parameterization of the hyper-parameters. For certain choices of hyper-parameters, this recovers the mean-field limit and the NTK limit. This paper also proposed certain properties of the limiting dynamics and showed that using these properties as the classification criteria, there are only a finite number of distinct models in the limit. This paper also proposed a novel initialization-corrected mean-field limit that satisfies all properties.