On the Dynamics and Convergence of Weight Normalization for Training Neural Networks

1 INTRODUCTION . Dynamic normalization in neural networks is a re-parametrization procedure between the layers that improves stability during training and leads to faster convergence . This approach was popularized with the introduction of Batch Normalization ( BatchNorm ) in [ 20 ] and has led to a plethora of additional normalization methods , notably including Layer Normalization ( LayerNorm ) [ 6 ] and Weight Normalization ( WeightNorm ) [ 28 ] . WeightNorm was proposed as a method that emulates BatchNorm and benefits from similar stability and convergence properties . Moreover , WeightNorm has the advantage of not requiring a batch setting , therefore considerably reducing the computational overhead that is imposed by BatchNorm [ 16 ] . WeightNorm is widely used in training of neural networks and is the focus of this work . Today , normalization methods are ubiquitous in the training of neural nets since in practice they significantly improve the convergence speed and stability in training . Despite the impressive empirical results and massive popularity of dynamic normalization methods , explaining their utility and proving that they converge when training with non-smooth , non-convex loss functions has remained an unsolved problem . In this paper we provide sufficient conditions on the data , initialization , and over-parameterization for dynamically normalized ReLU networks to converge to a global minimum of the loss function , and rigorously illustrate the utility of normalization methods . Consider the class of 2-layer ReLU neural networks f : Rd → R parameterized by ( W , c ) ∈ Rm×d × Rm as f ( x ; W , c ) = 1√ m m∑ k=1 ckσ ( w > k x ) . ( 1.1 ) Here the activation function is the ReLU , σ ( s ) = max { s , 0 } [ 26 ] , m denotes the width of the second layer , and f is normalized accordingly by a factor √ m. We investigate gradient descent training with WeightNorm for ( 1.1 ) , which re-parameterizes the network in terms of ( V , g , c ) ∈ Rm×d×Rm×Rm as f ( x ; V , g , c ) = 1√ m m∑ k=1 ckσ ( gk · v > k x ‖vk‖2 ) . ( 1.2 ) This gives a similar parameterization to [ 14 ] that study convergence of gradient optimization of convolutional filters on Gaussian data . We consider the regression task , optimizing with respect to the L2 loss with random parameter initialization and focus on the over-parametrized regime , meaning that m > n , where n is the number of training samples . The neural network function class ( 1.1 ) has been studied in many papers including [ 3 , 15 , 33 , 36 ] along with other similar over-parameterized architectures [ 1 , 14 , 23 ] . An exuberant series of recent works prove that feed-forward ReLU networks converge to zero training error when trained with gradient descent from random initialization . Nonetheless , to the best of our knowledge , there are no proofs that ReLU networks trained with dynamic normalization on general data converge to a global minimum . This is in part because normalization methods completely change the optimization landscape during training . Here we show that neural networks of the form given above converge at linear rate when trained with gradient descent and WeightNorm . The analysis is based on the over-parameterization of the networks , which allows for guaranteed descent while the gradient is non-zero . For regression training , a group of papers studied the trajectory of the networks ’ predictions and showed that they evolve via a “ neural tangent kernel ” ( NTK ) as introduced by Jacot et al . [ 21 ] . The latter paper studies neural network convergence in the continuous limit of infinite width overparameterization , while the works of [ 3 , 15 , 27 , 33 , 36 ] analyze the finite width setting . For finitewidth over-parameterized networks , the training evolution also exhibits a kernel that takes the form of a Gram matrix . In these works , the convergence rate is dictated by the least eigenvalue of the kernel . We build on this fact , and also on the general ideas of the proof of [ 15 ] and the refined work of [ 3 ] . Compared with un-normalized training , we prove that normalized networks follow a modified kernel evolution that features a “ length-direction ” decomposition of the NTK . This leads to two convergence regimes in WeightNorm training and explains the utility of WeightNorm from the perspective of the NTK . In the settings considered , WeightNorm significantly reduces the amount of over-parameterization needed for provable convergence , as compared with un-normalized settings . The decomposition of the NTK also connects to observations of [ 12 ] that discuss “ lazy training ” which refers to a training regime where the weights of the network stay close to their initialization ( see Section 6 ) . Further , we present a more careful analysis that leads to improved over-parameterization bounds as compared with [ 15 ] . In this work we rigorously analyze the dynamics of weight normalization training and its convergence from the perspective of the neural tangent kernel . We discover WeightNorm training has two regimes with distinct behaviors . The main contributions of this work are : • We prove the first general convergence result for dynamically normalized 2-layer ReLU networks trained with gradient descent . Our formulation does not assume the existence of a teacher network and has mild assumptions on the training data . • We explain the utility of normalization methods via a decomposition of the neural tangent kernel . In the analysis we highlight two distinct convergence regimes and give a concrete example of “ lazy training ” for finite-step gradient descent . • It is shown that finite-step gradient descent converges for all weight magnitudes at initialization and we significantly reduce the amount of over-parameterization required for provable convergence as compared with un-normalized training . The paper is organized as follows . In Section 2 we provide background on WeightNorm and derive key evolution dynamics of training in Section 3 . We present and discuss our main results , alongside with the idea of the proof , in Section 4 . We discuss related work in Section 5 , and offer a discussion of our results and their implications to dynamic normalization training and “ lazy training ” in Section 6 . Proofs are presented in the Appendix . 2 WEIGHTNORM . Here we give an overview of the WeightNorm procedure and review some known properties of normalization methods . Notation We use lowercase , lowercase boldface , and uppercase boldface letters to denote scalars , vectors and matrices resp . We denote the Rademacher distribution as U { 1 , −1 } and write N ( µ , Σ ) for a Gaussian with mean µ and covariance Σ . Training points are denoted by x1 . . .xn ∈ Rd and parameters of the first layer by vk ∈ Rd . We use σ ( x ) : = max { x , 0 } , and write ‖ · ‖2 , ‖ · ‖F for the spectral and Frobenius norms for matrices . λmin ( A ) is used to denote the minimum eigenvalue of a matrix A and 〈· , ·〉 denotes the Euclidean inner product . For a vector v denote the ` 2 vector norm as ‖v‖2 and for a positive definite matrix S define the induced vector norm ‖v‖S : = √ v > Sv . The projections of x onto u and u⊥ are defined as xu : = uu > x ‖u‖22 , xu ⊥ : = x ( I− uu > ‖u‖22 ) . Denote the indicator function of event A as 1A and for a weight vector at time t , vk ( t ) , and data point xi we denote 1ik ( t ) : = 1 { vk ( t ) > xi≥ 0 } . WeightNorm procedure For a single neuron σ ( w > x ) , WeightNorm re-parametrizes the weight w ∈ Rd in terms of v ∈ Rd , g ∈ R as w ( v , g ) = g · v ‖v‖2 , σ ( g · v > x ‖v‖2 ) . ( 2.1 ) This decouples the magnitude and direction of each weight vector ( referred as the “ length-direction ” decomposition ) . In comparison , for BatchNorm each output w > x is normalized according to the average statistics in a batch . We can draw the following analogy between WeightNorm and BatchNorm if the inputs xi are centered ( Ex = 0 ) and the covariance matrix is known ( Exx > = S ) . In this case , batch training with BatchNorm amounts to σ ( γ · w > x√ Ex ( w > xx > w ) ) = σ ( γ · w > x√ w > Sw ) = σ ( γ · w > x ‖w‖S ) . ( 2.2 ) From this prospective , WeightNorm is a special case of ( 2.2 ) with S = I [ 22 , 28 ] . Properties of WeightNorm We start by giving an overview of known properties of WeightNorm that will be used to derive the gradient flow dynamics of WeightNorm training . For re-parameterization ( 2.1 ) of a network function f that is initially parameterized with a weight w , the gradient ∇wf relates to the gradients∇vf , ∇gf by the identities ∇vf = g ‖v‖2 ( ∇wf ) v ⊥ , ∇gf = ( ∇wf ) v. This implies that∇vf · v = 0 for each input x and parameter v. For gradient flow , this orthogonality results in ‖v ( 0 ) ‖2 = ‖v ( t ) ‖2 for all t. For gradient descent ( with step size η ) the discretization in conjunction with orthogonality leads to increasing parameter magnitudes during training [ 4 , 19 , 28 ] , as illustrated in Figure 1 , ‖v ( s+ 1 ) ‖22 = ‖v ( s ) ‖22 + η2‖∇vf‖22 ≥ ‖v ( s ) ‖22 . ( 2.3 ) Problem Setup We analyze ( 1.1 ) with WeightNorm training ( 1.2 ) , so that f ( x ; V , c , g ) = 1√ m m∑ k=1 ckσ ( gk · v > k x ‖vk‖2 ) . We take an initialization in the spirit of [ 17 , 28 ] : vk ( 0 ) ∼ N ( 0 , α2I ) , ck ∼ U { −1 , 1 } , and gk ( 0 ) = ‖vk ( 0 ) ‖2/α . ( 2.4 ) Where α is the variance of vk at initialization . The initialization of gk ( 0 ) is therefore taken to be independent of α . We remark that the initialization ( 2.4 ) gives the same initial output distribution as in methods that study the un-normalized network class ( 1.1 ) . The parameters of the network are optimized using the training data { ( x1 , y1 ) , . . . , ( xn , yn ) } with respect to the square loss L ( f ) = 1 2 n∑ i=1 ( f ( xi ) − yi ) 2 = 1 2 ‖f − y‖22 , ( 2.5 ) where f = ( f1 , f2 , . . . , fn ) > = ( f ( x1 ) , f ( x2 ) , . . . , f ( xn ) ) > and y = ( y1 , y2 , . . . , yn ) > .

Summary: This paper presents a general proof of the convergence of two-layer ReLU networks with weight normalization trained with gradient descent. Weight normalization re-parameterizes the weights to decouple the directions and lengths of kernels. Depending on the lengths of kernels the training process can be divided into two regimes, corresponding to updates of lengths and directions, respectively. One of the regimes naturally corresponds to lazy training where the directions remain stable. And there are transitions from one regime to the other when the lengths gradually change during the training process.