Three Mechanisms of Weight Decay Regularization

1 INTRODUCTION . Weight decay has long been a standard trick to improve the generalization performance of neural networks ( Krogh & Hertz , 1992 ; Bos & Chug , 1996 ) by encouraging the weights to be small in magnitude . It is widely interpreted as a form of L2 regularization because it can be derived from the gradient of the L2 norm of the weights in the gradient descent setting . However , several findings cast doubt on this interpretation : • Weight decay has sometimes been observed to improve training accuracy , not just generalization performance ( e.g . Krizhevsky et al . ( 2012 ) ) . • Loshchilov & Hutter ( 2017 ) found that when using Adam ( Kingma & Ba , 2014 ) as the optimizer , literally applying weight decay ( i.e . scaling the weights by a factor less than 1 in each iteration ) enabled far better generalization than adding an L2 regularizer to the training objective . • Weight decay is widely used in networks with Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) . In principle , weight decay regularization should have no effect in this case , since one can scale the weights by a small factor without changing the network ’ s predictions . Hence , it does not meaningfully constrain the network ’ s capacity . The effect of weight decay remains poorly understood , and we lack clear guidelines for which tasks and architectures it is likely to help or hurt . A better understanding of the role of weight decay would help us design more efficient and robust neural network architectures . In order to better understand the effect of weight decay , we experimented with both weight decay and L2 regularization applied to image classifiers using three different optimization algorithms : SGD , Adam , and Kronecker-Factored Approximate Curvature ( K-FAC ) ( Martens & Grosse , 2015 ) . Consistent with the observations of Loshchilov & Hutter ( 2017 ) , we found that weight decay consistently outperformed L2 regularization in cases where they differ . Weight decay gave an especially strong performance boost to the K-FAC optimizer , and closed most of the generalization gaps between first- and second-order optimizers , as well as between small and large batches . We then investigated the reasons for weight decay ’ s performance boost . Surprisingly , we identified three distinct mechanisms by which weight decay has a regularizing effect , depending on the particular algorithm and architecture : 1 . In our experiments with first-order optimization methods ( SGD and Adam ) on networks with BN , we found that it acts by way of the effective learning rate . Specifically , weight decay reduces the scale of the weights , increasing the effective learning rate , thereby increasing the regularization effect of gradient noise ( Neelakantan et al. , 2015 ; Keskar et al. , 2016 ) . As evidence , we found that almost all of the regularization effect of weight decay was due to applying it to layers with BN ( for which weight decay is meaningless ) . Furthermore , when we computed the effective learning rate for the network with weight decay , and applied the same effective learning rate to a network without weight decay , this captured the full regularization effect . 2 . We show that when K-FAC is applied to a linear network using the Gauss-Newton metric ( K-FAC-G ) , weight decay is equivalent to regularizing the squared Frobenius norm of the input-output Jacobian ( which was shown by Novak et al . ( 2018 ) to improve generalization ) . Empirically , we found that even for ( nonlinear ) classification networks , the Gauss-Newton norm ( which K-FAC with weight decay is implicitly regularizing ) is highly correlated with the Jacobian norm , and that K-FAC with weight decay significantly reduces the Jacobian norm . 3 . Because the idealized , undamped version of K-FAC is invariant to affine reparameterizations , the implicit learning rate effect described above should not apply . However , in practice the approximate curvature matrix is damped by adding a multiple of the identity matrix , and this damping is not scale-invariant . We show that without weight decay , the weights grow large , causing the effective damping term to increase . If the effective damping term grows large enough to dominate the curvature term , it effectively turns K-FAC into a first-order optimizer . Weight decay keeps the effective damping term small , enabling K-FAC to retain its second-order properties , and hence improving generalization . Hence , we have identified three distinct mechanisms by which weight decay improves generalization , depending on the optimization algorithm and network architecture . Our results underscore the subtlety and complexity of neural network training : the final performance numbers obscure a variety of complex interactions between phenomena . While more analysis and experimentation is needed to understand how broadly each of our three mechanisms applies ( and to find additional mechanisms ! ) , our work provides a starting point for understanding practical regularization effects in neural network training . 2 PRELIMINARIES . Supervised learning . Given a training set S consisting of training pairs { x , y } , and a neural network fθ ( x ) with parameters θ ( including weights and biases ) , our goal is to minimize the emprical risk expressed as an average of a loss ` over the training set : L ( θ ) ≡ 1N ∑ ( x , y ) ∼S ` ( y , fθ ( x ) ) . Stochastic Gradient Descent . To minimize the empirical risk L ( θ ) , stochastic gradient descent ( SGD ) is used extensively in deep learning community . Typically , gradient descent methods can be derived from the framework of steepest descent with respect to standard Euclidean metric in parameter space . Specifically , gradient descent minimizes the following surrogate objective in each iteration : h ( θ ) = ∆θ > ∇θL ( θ ) + 1/ηD ( θ , θ + ∆θ ) , ( 1 ) where the distance ( or dissimilarity ) function D ( θ , θ + ∆θ ) is chosen as 12‖∆θ‖ 2 2 . In this case , solving equation 1 yields ∆θ = −η∇θL ( θ ) , where η is the learning rate . Natural gradient . Though popular , gradient descent methods often struggle to navigate “ valleys ” in the loss surface with ill-conditioned curvature ( Martens , 2010 ) . Natural gradient descent , as a variant of second-order methods ( Martens , 2014 ) , is able to make more progress per iteration by taking into account the curvature information . One way to motivate natural gradient descent is to show that it can be derived by adapting steepest descent formulation , much like gradient descnet , except using an alternative local distance . The distance function which leads to natural gradient is the KL divergence on the model ’ s predictive distribution DKL ( pθ ‖ pθ+∆θ ) ≈ 12∆θ > F∆θ , where F ( θ ) is the Fisher information matrix1 ( Amari , 1998 ) : F = E [ ∇θ log p ( y|x , θ ) ∇θ log p ( y|x , θ ) > ] . ( 2 ) Applying this distance function to equation 1 , we have θt+1 ← θt − ηF−1∇θL ( θ ) . Gauss-Newton algorithm . Another sensible distance function in equation 1 is the L2 distance on the output ( logits ) of the neural network , i.e . 12‖fθ+∆θ − fθ‖ 2 2 . This leads to the classical Gauss-Newton algorithm which updates the parameters by θt+1 ← θt − ηG−1∇θL ( θ ) , where the Gauss-Newton ( GN ) matrix is defined as G = E [ J > θ Jθ ] , ( 3 ) and Jθ is the Jacobian of fθ ( x ) w.r.t θ . The Gauss-Newton algorithm , much like natural gradient descent , is also invariant to the specific parameterization of neural network function fθ . Two curvature matrices . It has been shown that the GN matrix is equivalent to the Fisher matrix in the case of regression task with squared error loss ( Heskes , 2000 ) . However , they are not identical for the case of classification , where cross-entropy loss is commonly used . Nevertheless , Martens ( 2014 ) showed that the Fisher matrix is equivalent to generalized GN matrix when model prediction p ( y|x , θ ) corresponds to exponential family model with natural parameters given by fθ ( x ) , where the generalized GN matrix is given by G = E [ J > θ H ` Jθ ] , ( 4 ) and H ` is the Hessian of ` ( y , z ) w.r.t z , evaluated at z = fθ ( x ) . In regression with squared error loss , the Hessian H ` happens to be identity matrix . Preconditioned gradient descent . Given the fact that both natural gradient descent and GaussNewton algorithm precondition the gradient with an extra curvature matrix C ( θ ) ( including the Fisher matrix and GN matrix ) , we also term them preconditioned gradient descent for convenience . K-FAC . As modern neural networks may contain millions of parameters , computing and storing the exact curvature matrix and its inverse is impractical . Kronecker-factored approximate curvature ( K-FAC ) ( Martens & Grosse , 2015 ) uses a Kronecker-factored approximation to the curvature matrix to perform efficient approximate natural gradient updates . As shown by Luk & Grosse ( 2018 ) , K-FAC can be applied to general pullback metric , including Fisher metric and the Gauss-Newton metric . For more details , we refer reader to Appendix F or Martens & Grosse ( 2015 ) . Batch Normalization . Broadly speaking , Batch Normalization ( BN ) is a mechanism that aims to stabilize the distribution ( over a mini-batch ) of inputs to a given network layer during training . This is achieved by augmenting the network with additional layers that subtract the mean µ and divide by the standard deviation σ . Typically , the normalized inputs are also scaled and shifted based on trainable parameters γ and β : BN ( x ) = x− µ σ · γ + β . ( 5 ) For clarity , we ignore the parameters γ and β , which do not impact the performance in practice . This is not surprising , since with ReLU activations , only the γ of the last layer affects network ’ s outputs which can be merged with the softmax layer weights ( as also pointed out by van Laarhoven ( 2017 ) ) . 3 THE EFFECTIVENESS OF WEIGHT DECAY . Our goal is to understand weight decay regularization in the context of training deep neural networks . Towards this , we first discuss the relationship between L2 regularization and weight decay in different optimizers . 1The underlying distribution in equation 2 has been left ambiguous . Throughout the experiments , we sample the targets from the model ’ s predictions , as done in Martens & Grosse ( 2015 ) ; Zhang et al . ( 2017 ) Gradient descent with weight decay is defined by the following update rule : θt+1 ← ( 1− ηβ ) θt − η∇L ( θt ) , where β defines the rate of the weight decay per step and η is the learning rate . In this case , weight decay is equivalent to L2 regularization . However , the two differ when the gradient update is preconditioned by a matrix C−1 , as in Adam or K-FAC . The preconditioned gradient descent update with L2 regularization is given by θt+1 ← ( I− ηβC−1 ) θt − ηC−1∇θL ( θt ) , ( 6 ) whereas the weight decay update is given by θt+1 ← ( 1− ηβ ) θt − ηC−1∇θL ( θt ) . ( 7 ) The difference between these updates is whether the preconditioner is applied to θt . The latter update can be interpreted as the preconditioned gradient descent update on a regularized objective where the regularizer is the squared C-norm ‖θ‖2C = θ > Cθ . If C is adapted based on statistics collected during training , as in Adam or K-FAC , this interpretation holds only approximately because gradient descent on ‖θ‖2C would require differentiating through C. However , this approximate regularization term can still yield insight into the behavior of weight decay . ( As we discuss later , this observation informs some , but not all , of the empirical phenomena we have observed . ) Though the difference between the two updates may appear subtle , we find that it makes a substantial difference in terms of generalization performance . Initial Experiments . We now present some empirical findings about the effectiveness of weight decay which the rest of the paper is devoted to explaining . Our experiments were carried out on two different datasets : CIFAR-10 and CIFAR-100 ( Krizhevsky & Hinton , 2009 ) with varied batch sizes . We test VGG16 ( Simonyan & Zisserman , 2014 ) and ResNet32 ( He et al. , 2016 ) on both CIFAR-10 and CIFAR-100 ( for more details , see Appendix A ) . In particular , we investigate three different optimization algorithms : SGD , Adam and K-FAC . We consider two versions of K-FAC , which use the Gauss-Newton matrix ( K-FAC-G ) and Fisher information matrix ( K-FAC-F ) . Figure 1 shows the comparison between weight decay , L2 regularization and the baseline . We also compare weight decay to the baseline on more settings and report the final test accuracies in Table 1 . Finally , the results for large-batch training are summarized in Table 3 . Based on these results , we make the following observations regarding weight decay : 1 . In all experiments , weight decay regularization consistently improved the performance and was more effective than L2 regularization in cases where they differ ( See Figure 1 ) . 2 . Weight decay closed most of the generalization gaps between first- and second-order optimizers , as well as between small and large batches ( See Table 1 and Table 3 ) . 3 . Weight decay significantly improved performance even for BN networks ( See Table 1 ) , where it does not meaningfully constrain the networks ’ capacity . 4 . Finally , we notice that weight decay gave an especially strong performance boost to the K-FAC optimizer when BN was disabled ( see the first and fourth rows in Table 1 ) . In the following section , we seek to explain these phenomena . With further testing , we find that weight decay can work in unexpected ways , especially in the presence of BN .

Summary: This paper identifies and investigates three mechanisms of weight decay regularization. The authors consider weight decay for DNN architectures with/without BN and different types of optimization algorithms (SGD, Adam, and two versions of KFAC). The paper unravels insights on weight decay regularization effects, which cannot be explained only by traditional L2 regularization approach. This understanding is of high importance for the further development of regulations techniques for deep learning.