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: The paper analyzes joint scalings of the parameter initialization and the learning rate, with respect to the limit of infinite width, in the context of two-layer neural networks with stochastic gradient descent and binary logistic loss. It proposes some “dynamically stable” conditions and identifies a range of scalings that satisfy these conditions. This covers the neural tangent kernel (NTK) and mean field (MF) scalings, as well as others. The paper then proposes to add an extra correction function to the initialization of the MF limit and argues experimentally that this correction can give a proxy for standard neural networks.