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: The paper proposes to understand the effect of weight decay and design schedules for weight decay in Adam based on the theory. The first contribution is to understand the dynamics of weight decay in a neural network by assuming the loss function is quadratic. The theoretical analysis shows that weight decay biases stationary points but accelerates convergence. Second, the paper proposes to increase the weight decay when a large batch size is used based on the convergence acceleration property of WD. Finally, AdamS is proposed to improve the stability of Adam and some empirical results show its effectiveness using various neural architectures on small-scale datasets (not good on ImageNet).