Spherical Motion Dynamics: Learning Dynamics of Normalized Neural Network using SGD and Weight Decay

1 Introduction . Normalization techniques ( e.g . Batch Normalization ( Ioffe & Szegedy , 2015 ) or its variants ) are one of the most commonly adopted techniques for training deep neural networks ( DNN ) . A typical normalization can be formulated as following : consider a single unit in a neural network , the input is X , the weight of linear layer is w ( bias is included in w ) , then its output is y ( X ; w ; γ ; β ) = g ( Xw − µ ( Xw ) σ ( wX ) γ + β ) , ( 1 ) where g is a nonlinear activation function like ReLU or sigmoid , µ , σ are mean and standard deviation computed across specific dimension of Xw ( like Batch Normalization ( Ioffe & Szegedy , 2015 ) , Layer Normalization Ba et al . ( 2016 ) , Group Normalization ( Wu & He , 2018 ) , etc. ) . β , γ are learnable 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . parameters to remedy for the limited range of normalized feature map . Aside from normalizing feature map , Salimans & Kingma ( 2016 ) normalizes weight by l2 norm instead : y ( X ; w ; γ ; β ) = g ( X w ||w||2 γ + β ) , ( 2 ) where || · ||2 denotes l2 norm of a vector . Though formulated in different manners , all normalization techniques mentioned above share an interesting property : scale-invariance Definition 1 ( Scale-invariance ) . Given loss function L ( w ) , w is scale-invariant w.r.t . L if and only if ∀k ∈ R+ , we have L ( w ) = L ( kw ) . By definition of scale-invariance , we can directly derive the following properties of scale-invariant weights in Lemma 1 Lemma 1 . If w is scale-invariant with respect to L ( w ) , then for all k > 0 , we have : 〈wt , ∂L ∂w ∣∣∣ w=wt 〉 = 0 ( 3 ) ∂L ∂w ∣∣∣ w=kwt = 1 k · ∂L ∂w ∣∣∣ w=wt . ( 4 ) Proof is in appendix . Lemma 1 is also discussed in Hoffer et al . ( 2018 ) ; van Laarhoven ( 2017 ) ; Li & Arora ( 2020 ) ; Li et al . ( 2020 ) , it makes the learning dynamics of normalized neural network exhibit an interesting phenomenon when using Stochastic Gradient Descent ( SGD ) with Weight Decay ( WD ) : a typical SGD update rule with WD is wt+1 = wt − η ( ∂L ∂w ∣∣∣ w=wt + λwt ) = ( 1− ηλ ) wt − η ∂L ∂w ∣∣∣ w=wt , ( 5 ) where η denotes learning rate , λ denotes WD factor . Then dynamics of wt is like a physical process – a satellite ’ s motion around the earth ( see illustration in Fig.1 ) : according to Eq . ( 3 ) , −η∂L/∂w ∣∣ w=wt ( green line in Fig.1 ) is always perpendicular to wt , providing “ centrifugal effect ” to make ||wt+1||2 larger than ||wt||2 ; while −ηλwt ( red line in Fig.1 ) is always in the opposite direction of wt , providing “ centripetal effect ” to make ||wt+1||2 smaller than ||wt||2 . Due to this “ tug of war ” between “ centrifugal effect ” and “ centripetal effect ” , the norm of weight will change accordingly . Though the norm of scale-invariant weight doesn ’ t influence the loss at all ( by definition ) , but it will affect the optimizing trajectory by changing the scale of gradient ( Eq.4 ) . Therefore , to decouple the relation between weight norm and gradient norm , a common way ( van Laarhoven , 2017 ; Hoffer et al. , 2017 ; Chiley et al. , 2019 ) is to regard the learning dynamics of scale-invariant weight as a manifold learning process restricted on a unit sphere , using “ effective learning rate '' , defined as η/||wt||22 . Akin to this reason , in this paper the learning dynamics of scale-invariant weight is viewed as a motion on a unit sphere , whose motion speed is determined by weight norm ( in original weight space ) , gradient norm , and hyper-parameters . Therefore , we formally name it as Spherical Motion Dynamcis ( SMD ) and focus on discussing its unique properties . Concept of “ Equilibrium ” Since “ tug-of-war ” in SMD determines the relative sizes of ||wt+1||2 and ||wt||2 , a question naturally arises : what will happen if ||wt+1||2 = ||wt||2 ? van Laarhoven ( 2017 ) discuss this question first ; Chiley et al . ( 2019 ) named this state where ||wt+1||2 = ||wt||2 in SMD as “ equilibrium ” , and discuss its properties ; Li & Arora ( 2020 ) derives a lemma about equilibrium in SGD with Momentum ( SGDM ) ; Li et al . ( 2020 ) ; Kunin et al . ( 2021 ) establish the equilibrium in continuous approximation model . However , most early literature ( van Laarhoven , 2017 ; Chiley et al. , 2019 ; Li & Arora , 2020 ; Kunin et al. , 2021 ) do not discuss a fundamental question : “ Does equilibrium really occur in practice ? ” van Laarhoven ( 2017 ) intuitively explains that equilibrium is caused by convergence of optimization . But there exists a contradiction between the interpretation of van Laarhoven ( 2017 ) and traditional view of optimization : if equilibrium ( ||wt||2 = ||wt+1||2 ) is caused by convergence of optimization , then gradient of loss ∂L/∂w ∣∣ w=wt should be 0 , which makes the balance of “ centrifugal effect ” and “ centripetal effect ” impossible to reach ( since centripetal effect comes from ∂L/∂w ∣∣ w=wt ) . Therefore , van Laarhoven ( 2017 ) ; Chiley et al . ( 2019 ) ; Li & Arora ( 2020 ) ; Kunin et al . ( 2021 ) all essentially regard equilibrium as an assumption , and do not justify its existence in neither empirical nor theoretical aspects . “ Equilibrium ” was not even a phenomenon observed in practice , but only a concept until recently . Recent work ( Li et al. , 2020 ) successfully exhibits the existence of equilibrium by formulating SGD in Eq . ( 5 ) via a Stochastic Differential Equation ( SDE ) in the continuous time limit . They theoretically prove equilibrium can be reached in SDE settings via convergence of ||w||t , and the convergence time is O ( 1/ ( λη ) ) ( λ , η denote WD factor and learning rate respectively ) . However , due to the gap between discrete formulation of SGD and continuous formulation of SDE , theoretical results derived from SDE model can only provide intuitive understanding on empirical observations , some of which are even incorrect ( will be discussed latter ) ; Besides , SDE can hardly take SGD with momentum ( Polyak , 1964 ) into account , which has become default setting in nearly all kinds of deep learning tasks . In summary , a thorough understanding on cause of “ equilibrium ” and its impact to learning dynamics of normalized neural network is still needed . In this paper , we comprehensively reveal Spherical Motion Dynamics ( SMD ) , i.e . the learning dynamics of normalized neural network using SGD ( M ) and weight decay ( WD ) . Our analysis on SMD is directly established on discrete settings . We interpret why equilibrium can be reached in SMD in both theoretical and empirical aspects , and show how SMD affects the optimization trajectory of neural network . Specifically , our contributions are • We introduce the assumptions which can lead to equilibrium in SMD , and justify their reasonableness by sufficient experiments . We also prove under given assumptions , equilibrium can be reached as weight norm approach to its theoretical value in a linear rate regime . Our theorem show equilibrium is a dynamic state in SMD , convergence of weight norm is neither the precondition nor the consequence of equilibrium . This conclusion almost refutes the assumptions of all previous related work ( van Laarhoven , 2017 ; Chiley et al. , 2019 ; Li & Arora , 2020 ; Kunin et al. , 2021 ) , and the conclusion of Li et al . ( 2020 ) . • We define a novel index , angular update , to measure the change of normalized neural network within a single iteration . We also derive its theoretical value in equilibrium . Our results show that angular update is better than norm of weight to indicate if equilibrium has been reached reached in SMD . Our empirical results further show angular update is an important index to reflect the effect of SMD and equilibrium ; • We verify our theorems on different computer vision tasks ( including one of most challenging datasets ImageNet ( Russakovsky et al. , 2015 ) and MSCOCO ( Lin et al. , 2014 ) ) with various networks structures . Experiments show the theoretical value of angular update and weight norm agree well with empirical observation . We also show how SMD influence the optimization trajectory of normalized neural network by controlling angular update . We believe SMD is one of the key reason why learning dynamics of normalized neural network is not consistent with traditional optimization theory ( Li et al. , 2020 ) . We think it is of great potential to take SMD and its equilibrium state into account while studying leaning dynamics of modern normalized neural network or designing novel efficient training strategy . 2 Theoretical results . In this section , we theoretically formulate Spherical Motion Dynamics ( SMD ) in discrete SGD/SGDM settings , and provide a precise description on “ equilibrium ” phenomenon ; Then we prove equilibrium can be reached in SMD under specific assumptions ; Finally a new index , whose theoretical value in equilibrium can be derived , is proposed to indicate the state of SMD . A new definition need to be given at first . Eq . ( 4 ) implies though norm of scale-invariant weights does not affect the output of neural network , it can influence norm of gradients , thus we define unit gradient in order to decouple weight norm and gradient norm . Definition 2 ( Unit Gradient ) . If wt 6= 0 , w̃ = w/||w||2 , the unit gradient of ∂L/∂w|w=wt is ∂L/∂w|w=w̃t . According to the definition of unit gradient , the unit gradient norm is independent of weight norm . Specifically , by setting k as 1/||wt||2 in Eq . ( 4 ) , the relation among weight norm , gradient and unit gradient is ∂L ∂w ∣∣∣ w=wt = 1 ||wt|| · ∂L ∂w ∣∣∣ w=w̃t . ( 6 ) Now , we can depict equilibrium of SGD and SGDM in theorem 1 , 2 respectively . Theorem 1 . ( Equilibrium in SGD ) Assume the loss function is L ( X ; w ) with scale-invariant weight w , denote gt = ∂L∂w ∣∣ Xt , wt , g̃t = gt · ||wt||2 . Consider the update rule of SGD with weight decay , wt+1 = wt − η · ( gt + λwt ) ( 7 ) where λ , η ∈ ( 0 , 1 ) . If the following assumptions hold : 1 ) λη 1 ( o ( λη ) can be omitted ) ; 2 ) Let Lt = E [ ||g̃t||22|wt ] . ∃V ∈ R+ , ∀t ∈ N+ , E [ ( ||g̃t||22 − Lt ) 2|wt ] ≤ V ; 3 ) ∀t ∈ N+ , Lt satisfies |Lt+1 − Lt| < 4 √ V ( λη ) 3/2 ; 4 ) ∃l ∈ R+ , ∀t ∈ N+ , ||g̃t||22 > l , l > 2 [ 2λη 1−2λη ] 2Lt . Then ∃B > 0 , ∀t ∈ N+ , w∗t = 4 √ Lt−1η/ ( 2λ ) , we have E [ ||wt||22 − ( w∗t ) 2 ] 2 ≤ ( 1− 2λη ) tB + 2V η2 l . ( 8 ) Remark 1 . The theoretical value of weight norm w∗t in Theorem 1 is consistent with the magnitude of weight norm ( O ( 4 √ η/λ ) ) in equilibrium in van Laarhoven ( 2017 ) , though van Laarhoven ( 2017 ) assumes the equilibrium has been reached in advance , hence van Laarhoven ( 2017 ) can not provide the approaching rate and scale of bias/variance . The vanishing term ( ( 1 − 2λη ) tB ) in Eq . ( 8 ) is consistent with the mixing time O ( 1/ ( λη ) ) presented in Li et al . ( 2020 ) . The proof can be seen in appendix . Assumption 1 is consistent with commonly used settings in practice ( Goyal et al. , 2017 ; He et al. , 2017 ; Ma et al. , 2018 ) ; Assumptions 2 , 3 , 4 all concern unit gradient : unit gradient norm should change smoothly ( assumption 3 ) with bounded variance ( assumption 2 ) ; besides , unit gradient norm should have a lower bound ( assumption 4 ) . We will see these assumptions can easily hold in practice in section 4.1 . Remark 2 . Note in theorem 1 , we assume ||g̃t||22 has a uniform lowere bound l , while E [ ( ||g̃t||22 − Lt ) 2|wt ] has a uniform upper bound V . But these are too strong assumptions and may not hold in practice . We only use them for ease of demonstration and proof . In fact the lower bound of ||g̃t||22 and upper bound of E [ ( ||g̃t||22 − Lt ) 2|wt ] can be replaced with smoothly varying functions lt and Vt respectively to remove unnecessary assumptions on uniformness . The “ trend ” to reach equilibirum is actually a local property , which can be inferred from the proof . True meaning of “ equilibrium ” : a dynamic state of SMD Recall as we demonstrate in introduction , the concept of equilibrium is originally established on the assumption that weight norm is steady ( ||wt||2 = ||wt+1||2 ) . But the assumption ( ||wt||2 = ||wt+1||2 ) is unrealistic due to the complex dynamics of training process and the variance of stochastic gradients . Now theorem 1 provides a realistic meaning of equilibrium in SGD settings : equilibrium is just a dynamic state of SMD , meaning ||wt||22 oscillates around the theoretical value ( w∗t ) 2 determined by hyperparameters and unit gradient norm . Its variance is bounded by 2V η2/l , which is relatively small comparing with ( w∗t ) 4 because 2V η2 l / ( w∗t ) 4 = 4V λη Lt−1l = O ( λη ) 1 . ( 9 ) Aside form the stochastic behavior of ||wt||22 , the “ dynamic state ” also reflects in the variation of the theoretical value ( w∗t ) 2 . Because ( w∗t ) 2 is determined by Lt , which is allowed to change smoothly across the whole training process in assumption 2 ( See more discussion in appendix ) . In summary , the sign of equilibrium is neither the convergence of weight norm ||wt||22 ( van Laarhoven , 2017 ; Chiley et al. , 2019 ) nor the convergence of ||wt||2 in expectation ( Li et al. , 2020 ) . The real sign of equilibrium is whether E||wt||22 is close to its theoretical value ( w∗t ) 2 . Theorem 1 also shows the dynamic equilibrium can be reached in a linear rate regime when vanishing term is larger than constant term in Eq. ( 8 ) . The approaching rate is only determined by predefined parameters λ , η . Moreover , based on the proof of theorem 1 , the cause of equilibrium is independent of optimization process at all , which implies the possibility that equilibrium can be reached long before the convergence of loss function , which refutes the conjections that “ equilibirum ” a equilibirum distribution of function space . Now we extend theorem 1 to momentum case . SGDM is more complex than SGD since momentum is not always perpendicular to the weight , hence we need to modify assumptions . Theorem 2 . ( Equilibrium in SGDM ) Considering the update rule of SGDM ( heavy ball method ( Polyak , 1964 ) ) : vt = αvt−1 + gt + λwt ( 10 ) wt+1 = wt − ηvt ( 11 ) where λ , η ∈ ( 0 , 1 ) , α ∈ ( 12 , 1 ) . If following assumptions hold : 5 ) λη 1 , λη < ( 1− √ α ) 2 ; 6 ) Define ht = ||gt||22 + 2α〈vt−1 , gt〉 , h̃t = ht · ||wt||22 , Lt = E [ h̃t|wt ] . ∃V ∈ R+ , ∀t ∈ N+ , E [ ( h̃t − Lt ) 2|wt ] ≤ V ; 7 ) ∀t ∈ N+ , Lt satisfies |Lt+1 − Lt| < 4 √ V ( λη ) 3/2 ; 8 ) ∃l ∈ R+ , ∀t ∈ N+ , h̃t > l > 2 [ 6λη ( 1−α ) 3 ( 1+α ) −8λη ( 1−α ) ] 2Lt ; then ∃B , C > 0 , C only depends on α , w∗t = 4 √ Lt−1η/ ( λ ( 1− α ) ( 2− λη/ ( 1 + α ) ) ) , we have E [ ||wt||22 − ( w∗t ) 2 ] 2 ≤ ( 1− 2λη 1− α ) tB + V η2 l C , ( 12 ) Remark 3 . So far , no other work rigorously prove equilibrium can be reached in SGDM . The most relevant work ( Li et al. , 2020 ) only provides a conjecture on convergence rate of weight norm in SGDM . By regarding SGDM as SGD with larger learning rate , they guess that the mixing time to reach equilibrium in SGDM case should be O ( 1/ ( λη ) ) , same order as mixing time in SGD case . Their conjecture can not provide further insight on difference between SGD and SGDM . While our results ( vanishing terms in Eq . ( 8 ) , ( 12 ) respectively ) clearly reflect the difference : the approaching rate of SGDM should be 1/ ( 1− α ) times larger than rate of SGD with same ηλ . α is usually set as 0.9 in practice , hence SGDM can reach equilibrium condition much faster than SGD . Proof can be seen in appendix . Like assumption 1 , assumption 5 also holds for commonly used hyperparameter settings ; Assumption 6 , 7 , 8 concerns not unit gradient norm ||g̃t||22 but an adjusted value h̃t which dominates the expectation and variance of ||wt||22 . We empirically find the expectation of 〈vt−1 , gt〉 is very close to 0 , therefore the behavior of h̃t is similar to that of ||g̃t||22 ( see Figure 3 ( d ) ) . We leave theoretical analysis on h̃t as future work . The experiments on justification of assumptions 6 , 7 , 8 can be seen in Figure 3 . Comparing with Eq . ( 8 ) and Eq . ( 12 ) , we can infer with same η , λ , SGDM can reach equilibrium state much faster than SGD , but it may have a larger variance , our experiments also verify our claim ( see Figure 3 ( b ) , 3 ( e ) ) . We have derived the theoretical value of weight norm in equilibrium , it allows us to check if equilibrium has been reached in practice . But the theoretical value of weight norm still relies on the expectation of unit gradient norm , which is not easy to compute in practice . Besides , weight norm is of little value for studying normalized models since their weight is scale-invariant . Hence we introduce an new and meaningful index , angular update , to reflect the effect of SMD and equilibrium . Definition 3 ( Angular Update ) . Let wt denote a scale-invariant weight from a neural network at iteration t , then angular update ∆t is defined as ∆t = ] ( wt , wt+1 ) = arccos ( 〈wt , wt+1〉 ||wt|| · ||wt+1|| ) , ( 13 ) where ] ( · , · ) denotes the angle between two vectors , 〈· , ·〉 denotes the inner product . Angular update has a concrete geometric meaning ( see illustration in Figure 1 ) : it is exactly the geodesic distance between w̃t and w̃t+1 on Sp−1 , where w̃t = wt/||wt||2 , w̃t+1 = wt+1/||wt+1||2 . Comparing with the Euclidean distance ||wt+1−wt||2 , angular update ∆t better reflects the effective update of the scale-invariant weight wt on its intrinsic domain Sp−1 . Angular update is determined by the relative sizes of gradient norm and weight norm , while the relative sizes of gradient/weight norm are influenced by SMD , hence angular update is strongly affected by SMD . The following theorem exhibits the behavior of angular update when equilibrium is reached in SMD . Theorem 3 . ( Theoretical value of Angular Update ) In SGD ( SGDM ) case , if assumptions in theorem 1 ( 2 ) hold , η2 1 , t is sufficiently large so that vanishing terms in Eq . ( 8 ) , ( 12 ) can be omitted , then with probability at least 1− 3 √ V Ltl we have |∆t − √ 2λη 1 + α | < O ( 3 √ V Ltl ) . ( 14 ) In SGD case , α = 0 . Remark 4 . Results of SGD and SGDM case are summarized in Eq . ( 14 ) in order to highlight the connection between SGD and SGDM . Theoretical value of angular update in Theorem 3 is partially consistent with previous works ( Chiley et al. , 2019 ; Li & Arora , 2020 ; Li et al. , 2020 ; Kunin et al. , 2021 ) , detailed discussion is buried in appendix . Note bias term in right side of Eq . ( 14 ) is of O ( 3 √ V/Ltl ) , which is too large comparing with its empirical value ( see Figure 3 ( c ) , 3 ( f ) ) , we leave it as a future work to improve the bound in Eq. ( 14 ) . Proof is in appendix . According to theorem 3 , the theoretical value of angular update in equilibrium only depends on hyper-parameters : learning rate η , WD factor λ , and momentum factor α . Hence comparing with behavior of weight norm , angular update provides an easier way to check whether equilibrium is reached . Since equilibrium can be reached in a linear rate regime as theorem 1 , 2 demonstrate , theorem 3 implies update efficiency of scale-invariant weights within a single step eventually will be determined only by predefined hyperparameters , regardless other attributes of the weights ( shape , size , position in network structure , or effects from other weights ) .

Summary: This paper investigates the time-evolution of weights in normalized neural networks (e.g. one in which every layer has batch-norm applied). It proposes theorems regarding convergence to equilibrium, in which the norm of the weights approaches a fixed value. Compared to prior work, which established the fixed value, this work additionally gives the approach rate, and the scale of the variance due to stochastic gradients. The paper introduces the angular update parameter, and shows this approaches a fixed value. The paper emphasizes that equilibrium is not reached immediately, but is approached dynamically. It also asserts that equilibrium is not about convergence of stochastic oscillations to zero, but instead convergence to a fixed value, up to oscillations. The theoretical results are then investigated empirically, showing agreement with theoretical results for the weight norm and angular update, and the dynamical approach to equilibrium.