Understanding and Scheduling Weight Decay

1 INTRODUCTION . Weight decay is a popular and even necessary regularization technique for training deep neural networks that generalize well ( Krogh and Hertz , 1992 ) . People commonly use L2 regularization as “ weight decay ” for training of deep neural networks and interpret it as a Gaussian prior over the model weights ( David , 1992 ; Graves , 2011 ) . This is true for vanilla Stochastic Gradient Descent ( SGD ) . However , Loshchilov and Hutter ( 2018 ) revealed that , when the learning rate is adaptive , the commonly used L2 regularization is not identical to the vanilla weight decay proposed by Hanson and Pratt ( 1989 ) : θt = ( 1− λ′ ) θt−1 − η ∂L ( θt−1 ) ∂θ , ( 1 ) where λ′ is the weight decay hyperparameter , η is the learning rate , and L ( θt−1 ) is the loss function over one minibatch . We denote the loss function over the whole training dataset as L ( θ ) . A conventional belief is that , from the Bayesian perspective , weight decay indicates a Gaussian prior over the model parameters , which regularize the complexity of the learned model ( Graves , 2011 ) . However , we argue that this Bayesian interpretation can not explain the empirical success of weight decay in deep learning . On one hand , Wenzel et al . ( 2020 ) demonstrated that the conventional Bayesian posterior usually does not work well in deep learning . This challenged the meaning of weight decay as a proper prior in the Bayesian interpretation . On the other hand , Lewkowycz and GurAri ( 2020 ) and our experimental results ( Figure 2 ) both suggested that weight decay is surprisingly unnecessary for long-time training . This observation critically challenged the Bayesian interpretation of weight decay in deep learning . The Bayesian interpretation believes that the prior knowledge over model parameters truly exist . Thus , the prior knowledge should not depend on the training time of deep neural networks , which contradicts with the experimental results in Figure 2 . Recently researchers started to explore other explanations for weight decay . Zhang et al . ( 2019b ) revealed three different roles of weight decay . Some papers ( Van Laarhoven , 2017 ; Zhang et al. , 2019b ) argued that weigh decay increases the effective learning rate but has no regularization effect when combined with scale-invariant normalization layers , such as batch normalization ( Ioffe and Szegedy , 2015 ) . Li et al . ( 2020 ) defined a new “ intrinsic learning rate ” parameter that is the product of the vanilla learning rate η and the weight decay factor λ and argued that the convergence time depends on this intrinsic learning rate . However , people still expect to further explore the theoretical mechanism of weight decay , because weight decay sometimes shows mysterious behaviors beyond the existing theoretical understanding . As we discussed above , the Bayesian interpretation has been critically challenged by recent empirical findings ( Lewkowycz and Gur-Ari , 2020 ) . The interpretation of increasing the effective learning rate only holds for scale-invariant normalization layers , which is a narrow scope . Contributions . In this paper , we try to theoretically understand the novel behaviors of weight decay and explore the potential of weight decay . Our paper mainly has three contributions . First , we propose a novel theoretical interpretation of weight decay from the perspective of learning dynamics . We study how the stability of the stationary points and the convergence time depend on weight decay , which corresponds to two main effects of weight decay . Second , we propose a novel weight-decay linear scaling rule ( See Rule 2 in Section 3 ) for large-batch training . We discover that the existing linear scaling rule ( Krizhevsky , 2014 ; Goyal et al. , 2017 ) for large-batch training that proportionally increases the learning rate may sometimes be harmful to large-batch training due to bad convergence . Instead , increasing weight decay is a more suitable solution in this case . Third , to the best of our knowledge , we are the first to design an effective learning-rate-aware scheduler for weight decay which may significantly improve test performance . Our work fills the gap in this direction . We note that a trivial weight decay scheduling method recently proposed by Lewkowycz and Gur-Ari ( 2020 ) has limited practical value because it only works well with constant learning rate , as the original paper claimed . Inspired by the stability of the stationary points , we show that weight decay should be coupled with the effective learning rate . We call the proposed method Stable Weight Decay ( SWD ) . Adam ( Kingma and Ba , 2015 ) with SWD ( AdamS ) is displayed in Algorithm 2 . The results in Table 1 fully supports that SWD can significantly improve Adam . Structure of this paper . In Section 2 , we analyze the dynamics of weight decay . In Section 3 , we study how weight decay can improve large-batch training . In Section 4 , we propose SWD . In Section 5 , we empirically study SWD . In Section 6 , we conclude our work . 2 UNDERSTANDING THE DYNAMICS OF WEIGHT DECAY . In this section , we provide a way to understand weight decay from a viewpoint of learning dynamics . Weight decay should be coupled with the learning rate scheduler . The vanilla weight decay described by Hanson and Pratt ( 1989 ) is given by Equation ( 1 ) . A more popular implementation for vanilla SGD in modern deep learning libraries , such as PyTorch ( Paszke et al. , 2019 ) and TensorFlow ( Abadi et al. , 2016 ) , is given by θt = ( 1− ηtλ ) θt−1 − ηt ∂L ( θt−1 ) ∂θ , ( 2 ) where ηt is the learning rate at the t-th step and weight decay is coupled with the learning rate scheduler . While previous work did not studied why weight decay should be coupled with the learning rate scheduler , this trick has been adopted by most deep learning libraries . We first take vanilla SGD as the studied example , where no Momentum is involved . It is easy to see that vanilla SGD with L2 regularization λ2 ‖θ‖ 2 is also given by Equation ( 2 ) . Suppose the learning rate is fixed in the whole training procedure . Then Equation ( 2 ) will be identical to Equation ( 1 ) if we simply choose λ′ = ηλ . However , learning rate decay is quite important for training of deep neural networks . So Equation ( 2 ) is not identical to Equation ( 1 ) in practice . Which implementation is better ? In Figure 1 , we empirically verified that the popular Equation- ( 2 ) based weight decay indeed outperforms the vanilla implementation in Equation ( 1 ) . We argue that Equation- ( 2 ) -based weight decay is theoretically better than Equation- ( 1 ) -based weight decay in terms of the stability of the stationary points . We define the stability as Definition 1 . Definition 1 ( The stability of the stationary point ) . The stationary point is stable if a stationary point is fixed during training . Otherwise , we define it as an unstable stationary point . The training objective at step-t of Equation ( 1 ) is L ( θ ) + λ ′ ηt ‖θ‖2 , while that of Equation ( 2 ) is L ( θ ) + λ‖θ‖2 . Thus , the stationary points of the regularized loss given by Equation ( 2 ) are stable during training , while that given by Equation ( 1 ) is unstable due to the dynamic training objective . Weight decay matters to the stability . As the stability of the stationary points critically depends on the weight decay scheduler , we define stable weight decay accordingly in Definition 2 . Definition 2 ( Stable Weight Decay ) . The weight decay is stable if the stationary points are stable during training . Otherwise , we define it as unstable weight decay . We propose Theorem 1 and prove that even deterministic Gradient Descent ( GD ) with vanilla weight decay may not converge to any non-zero stationary point due the instability of the stationary points . Theorem 1 ( Non-convergence of GD with unstable weight decay ) . Suppose learning dynamics is governed by GD with vanilla weight decay ( Equation ( 1 ) ) and the learning rate ηt ∈ ( 0 , +∞ ) holds . If ∃δ that satisfies 0 < δ ≤ |ηt − ηt+1| for any t > t0 , then the learning dynamics can not converge to any non-zero stationary point satisfying the condition lim t→+∞ [ ‖∇L ( θt ) ‖2 + ‖∇L ( θt+1 ) ‖2 ] = 0 . Remark . We leave the proof in Appendix A.1 . Even if the gradient is zero at the t-th step , the gradient at the t+ 1-th step will not be zero . Unstable weight decay causes changing the objective and hence the solution does not converge to any non-zero stationary point . Thus , GD with unstable weight decay has no theoretical convergence guarantee like GD ( Wolfe , 1969 ) . This explains why weight decay should be coupled with the learning rate to keep the stationary points stable during training . We conjecture that the stability of the stationary points should be a desirable property in deep learning . Dynamics of weight decay . Lewkowycz and Gur-Ari ( 2020 ) empirically discovered that , if we let the training time be approximately inverse to weight decay , the test error can monotonically decrease as weight decay decreases , which is beyond the existing theoretical understanding . We will show that it is easy to understand this empirical observation by analyzing the dynamics of weight decay . We first state Assumption 1 . Assumption 1 which holds near minima are common and useful for analyzing the solution or the convergence behavior in related papers ( Mandt et al. , 2017 ; Neyshabur et al. , 2017 ; Xie et al. , 2021a ; Li et al. , 2017 ; Zhang et al. , 2019a ; Zhou et al. , 2020 ) . It is known that stochastic optimization can escape multiple bad minima ( Xie et al. , 2021b ; Zhu et al. , 2019 ; Kleinberg et al. , 2018 ) and finally converge to a good minimum during training . Thus , Assumption 1 can be repeatedly applied near each minimum during training . Assumption 1 . The loss function around a minimum θ ? can be written as L ( θ ) = L ( θ ? ) + 1 2 ( θ − θ ? ) > H ( θ − θ ? ) , where H is the Hessian at θ ? . We propose Theorem 2 and demonstrate how the learned model parameters θ may depend on weight decay . We leave the proof in Appendix A.2 . Theorem 2 ( Dynamics of weight decay ) . Suppose Assumption 1 holds , and learning dynamics is governed by SGD with weight decay ( Equation ( 2 ) ) . If the learning rate scheduler is constant , the expected θt after t iterations satisfies the following : E [ θt − θ ? ] = [ I − η ( H + λI ) ] t [ θ0 −H ( H + λI ) −1θ ? ] − λ ( H + λI ) −1θ ? , ( 3 ) where I is the identity matrix . If the learning rate scheduler is given by the sequence { η1 , η2 , . . . , ηt } , the expected θt after t iterations satisfies the following : E [ θt − θ ? ] = t∏ k=1 [ I − ηk ( H + λI ) ] [ θ0 −H ( H + λI ) −1θ ? ] − λ ( H + λI ) −1θ ? . ( 4 ) Theorem 2 is a useful tool to understand convergence behaviors which may repeatedly happen during searching minima of deep learning . While Theorem 2 depends on Assumption 1 , we will show that it is remarkably consistent with novel empirical behaviors of weight decay and can also reveal novel insights . Theorem 2 indicates that weight decay mainly has two effects : 1 ) biasing the expected learned solution and 2 ) accelerating the convergence . Effect 1 . Weight decay biases the stationary points . Theorem 2 provides us two key insights about weight decay . The first insight is Corollary 1 derived from Theorem 2 by choosing the long-time limit . Corollary 1 indicates that large weight decay may not learn accurate minima well . This insight is consistent with the regularization-based understanding of weight decay ( Krogh and Hertz , 1992 ) . Corollary 1 . Suppose the conditions of Theorem 2 hold and Hi + λ > 0 , where Hi is the i-th eigenvalue of H . Then , in the limit of long training time that t→ +∞ , the displacement from the true minimum θ ? to the expected learned solution along the i-th direction is given by µi = −λ ( Hi + λ ) −1θ ? , ( 5 ) where the bias µ = limt→∞ E [ θt − θ ? ] . Figure 2 suggests that , given long enough training time inverse to weight decay , decreasing the weight decay hyperparameter may consistently improve test performance , which is also verified by Lewkowycz and Gur-Ari ( 2020 ) . This observation is contradicted with the conventional regularizationbased understanding ( Krogh and Hertz , 1992 ) , but fully supports Corollary 1 that weight decay is bad , if we already have good convergence . Effect 2 . Weight decay accelerates convergence . However , we discover that the long-time posterior can not explain the convergence behaviors . Figure 2 also suggests that , we need to increase the number of training iterations by 1λ for maintaining good convergence , which is also discovered by Lewkowycz and Gur-Ari ( 2020 ) and Li et al . ( 2020 ) . To explain why the convergence time depends on 1λ , we need to explore the dynamics of weight decay before approaching the convergence . The second insight is Corollary 2 that estimates the number of iterations for good convergence . Corollary 2 . Suppose the conditions of Theorem 2 hold and δi = Hi + λ > 0 , where Hi is the i-th eigenvalue of H . Then , the convergence time in the i-th direction is tconvergence , i = O ( ( Hi + λ ) −1η−1 ) . ( 6 ) It is well-known that the Hessian in deep learning has a very small number of large eigenvalues and a large number of nearly zero eigenvalues ( Sagun et al. , 2016 ; 2017 ) . Thus , we have Hi + λ ≈ λ for most directions . It means that the convergence time critically depend on the dynamics in the subspace corresponding to those nearly zero eigenvalues : tconvergence = O ( λ−1η−1 ) . ( 7 ) Equation ( 7 ) reveals two facts . First , Equation ( 7 ) supports the statement that the initial learning rate and the weight decay hyperparameter λ are linearly coupled for maintaining performance ( Loshchilov and Hutter , 2018 ; Van Laarhoven , 2017 ; Zhang et al. , 2019b ; Hoffer et al. , 2018 ) . Second , Equation ( 7 ) demonstrates that the number of iterations ( not the number of epochs or training time ) for good convergence depends onO ( λ−1η−1 ) . This not only explains the relation in Figure 2 , but also reveals that we should increase the number of iterations instead of the number of epochs by 1λ . We note that the second fact is also reported by Lewkowycz and Gur-Ari ( 2020 ) , while Lewkowycz and Gur-Ari ( 2020 ) did not present formal theoretical analysis like Theorem 2 . In our theoretical analysis , weight decay mainly has two effects . The first effect is to bias the expected solution in the long-time limit . The first effect is negative for learning accurate minima . The second effect is to accelerate the convergence in the subspace corresponding to nearly zero eigenvalues of the Hessian . The second effect is positive for accelerating training . Tuning weight decay needs to balance the tradeoff between the two effects .

Summary: In this paper, the authors propose a new interpretation on the effect of weight decay in learning dynamics of neural network. Based on their interpretation, they design a weight-decay linear scaling rule for large-batch training, and a learning-rate-aware scheduler for weight decay in common settings. They verify the effectiveness of the proposed method on CIFAR, Imagenet experiments.