Catastrophic Fisher Explosion: Early Phase Fisher Matrix Impacts Generalization

1 INTRODUCTION . Implicit regularization in gradient-based training of deep neural networks ( DNNs ) remains relatively poorly understood , despite being considered a critical component in their empirical success ( Neyshabur , 2017 ; Zhang et al. , 2016 ; Jiang et al. , 2020b ) . Recent work suggests that the early phase of training of DNNs might hold the key to understanding these implicit regularization effects . Golatkar et al . ( 2019 ) ; Keskar et al . ( 2017 ) ; Sagun et al . ( 2018 ) ; Achille et al . ( 2019 ) show that by introducing regularization later , a drop in performance due to lack of regularization in this phase is hard to recover from , while on the other hand , removing regularization after the early phase has a relatively small effect on the final performance . Other works show that the early phase of training also has a dramatic effect on the trajectory in terms of properties such as the local curvature of the loss surface or the gradient norm ( Jastrzebski et al. , 2020 ; Frankle et al. , 2020 ) . In particular , Achille et al . ( 2019 ) ; Jastrzębski et al . ( 2019 ) ; Golatkar et al . ( 2019 ) ; Lewkowycz et al . ( 2020 ) ; Leclerc & Madry ( 2020 ) independently suggest that rapid changes in the local curvature of the loss surface in the early phase critically affects the final generalization . Closely related to our work , Lewkowycz et al . ( 2020 ) ; Jastrzębski et al . ( 2019 ) show that using a large learning rate has a dramatic effect on the early optimization trajectory in terms of the loss curvature . These observations lead to a question : what is the mechanism by which regularization in the early phase impacts the optimization trajectory and generalization ? We investigate this question mainly through the lens of the Fisher Information Matrix ( FIM ) , a matrix that can be seen as approximating the local curvature of the loss surface in DNNs ( Martens , 2020 ; Thomas et al. , 2020 ) . Our main contribution is to show that the implicit regularization effect of using a large learning rate or a small batch size can be modeled as an implicit penalization of the trace of the FIM ( Tr ( F ) ) from the very beginning of training . We demonstrate on image classification tasks that the value of Tr ( F ) early in training correlates with the final generalization performance across settings with different learning rates or batch sizes . We then show evidence that explicitly regularizing Tr ( F ) ( which we call Fisher penalty ) significantly improves generalization in training with a sub-optimal learning rate . On the other hand , growth of Tr ( F ) early in training , which may occur in practice when using a relatively small learning rate , coincides with poor generalization . We call this phenomenon the catastrophic Fisher explosion . Figure 1 illustrates this effect on the TinyImageNet dataset ( Le & Yang , 2015 ) . Our second contribution is an analysis of why implicitly or explicitly regularizing Tr ( F ) impacts generalization . We reveal two effects of implicit or explicit regularization of Tr ( F ) : ( 1 ) penalizing Tr ( F ) discourages memorizing noisy labels , ( 2 ) small Tr ( F ) in the early phase of training biases optimization towards a flat minimum , as characterized by the trace of the Hessian . 2 IMPLICIT AND EXPLICIT REGULARIZATION OF THE FIM . Fisher Information Matrix Consider a probabilistic classification model pθ ( y|x ) , where θ denotes its parameters . Let ` ( x , y ; θ ) be the cross-entropy loss function calculated for input x and label y . Let g ( x , y ; θ ) = ∂∂θ ` ( x , y ; θ ) denote the gradient computed for an example ( x , y ) . The central object that we study is the Fisher Information Matrix F defined as F ( θ ) = Ex∼X , ŷ∼pθ ( y|x ) [ g ( x , ŷ ) g ( x , ŷ ) T ] , ( 1 ) where the expectation is often approximated using the empirical distribution X̂ induced by the training set . We denote its trace by Tr ( F ) . Later , we also look into the Hessian H ( θ ) = ∂ 2 ∂θ2 ` ( x , y ; θ ) . We denote its trace by Tr ( H ) . The FIM can be seen as an approximation to the Hessian ( Martens , 2020 ) . In particular , as p ( y|x ; θ ) → p̂ ( y|x ) , where p̂ ( y|x ) is the empirical label distribution , the FIM converges to the Hessian . Thomas et al . ( 2020 ) showed on image classifications tasks that Tr ( H ) ≈ Tr ( F ) along the optimization trajectory , which we also evidence in Appendix F. Fisher Penalty Several studies have presented evidence that the early phase has a drastic effect on the trajectory in terms of the local curvature of the loss surface ( Achille et al. , 2019 ; Jastrzębski et al. , 2019 ; Gur-Ari et al. , 2018 ; Lewkowycz et al. , 2020 ; Leclerc & Madry , 2020 ) . In particular , Lewkowycz et al . ( 2020 ) ; Jastrzębski et al . ( 2019 ) show that using a large learning rate in stochastic gradient descent biases training towards low curvature regions of the loss surface very early in training . For example , using a large learning rate in SGD was shown to result in a rapid decay of Tr ( H ) along the optimization trajectory Jastrzębski et al . ( 2019 ) . Our main contribution is to propose and investigate a specific mechanism by which using a large learning rate or a small batch size implicitly influences final generalization . Our first insight is to shift the focus from studying the Hessian , to studying properties of the FIM . Concretely , we hypothesize that using a large learning rate or a small batch size improves generalization by implicitly penalizing Tr ( F ) from the very beginning of training . The benefit of studying the FIM is that it can be directly and efficiently manipulated during training . In order to study the effect of implicit regularization of Tr ( F ) , we introduce a regularizer , which we refer to as Fisher penalty , explicitly penalizing Tr ( F ) . We derive this regularizer in the following way . First , we note that Tr ( F ) can be written as Tr ( F ) = Ex∼X , ŷ∼pθ ( y|x ) [ ‖ ∂∂θ ` ( x , ŷ ) ‖ 2 2 ] . To regularize Tr ( F ) , we add the following term to the loss function : ` ′ ( x1 : B , y1 : B ; θ ) = 1 B B∑ i=1 ` ( xi , yi ; θ ) + α ∥∥∥∥∥ 1B B∑ i=1 g ( xi , ŷi ) ∥∥∥∥∥ 2 , ( 2 ) where ( x1 : B , y1 : B ) is a mini-batch , ŷi is sampled from pθ ( y|xi ) , and α is a hyperparameter . We refer to this regularizer as Fisher penalty . The formulation is based on the empirical observation that ∥∥∥ 1B ∑Bi=1 g ( xi , ŷi ) ∥∥∥2 and Tr ( F ) correlate well during training . Crucially , this allows us to reduce the added computational cost of Fisher penalty to that of a single additional backpropagation call ( Drucker & Le Cun , 1992 ) . Finally , we compute the gradient of the second term only every 10 optimization steps , and in a given iteration use the most recently computed gradient . We discuss these approximations in detail in Appendix C. Catastrophic Fisher Explosion To illustrate the concepts mentioned in this section , we train a Wide ResNet model ( depth 44 , width 3 ) ( Zagoruyko & Komodakis , 2016 ) on the TinyImageNet dataset with SGD and two different learning rates . We illustrate in Figure 1 that the small learning rate leads to dramatic overfitting , which coincides with a sharp increase in Tr ( F ) in the early phase of training . We also show in Appendix D that these effects can not be explained by the difference in learning speed between runs with smaller and learning rates . We call this phenomenon the catastrophic Fisher explosion . 3 EARLY-PHASE Tr ( F ) CORRELATES WITH FINAL GENERALIZATION Using a large learning rate ( η ) or a small batch size ( S ) in SGD steers optimization to a lower curvature region of the loss surface . However , it remains a hotly debated topic whether this explains their strong regularization effect ( Dinh et al. , 2017 ; He et al. , 2019 ; Maddox et al. , 2020 ; Tsuzuku et al. , 2019 ; Yoshida & Miyato , 2017 ) . We begin by studying the connection between Tr ( F ) and generalization in experiments across which we vary η or S in SGD . Experimental setup We run our experiments in two settings : ( 1 ) ResNet-18 with Fixup He et al . ( 2015 ) ; Zhang et al . ( 2019 ) trained on the ImageNet dataset ( Deng et al. , 2009 ) , ( 2 ) ResNet-26 initialized with Arpit et al . ( 2019 ) trained on the CIFAR-10 and CIFAR-100 datasets ( Krizhevsky , 2009 ) . We train each architecture using SGD , with various values of η , S , and random seed . We define Tr ( Fi ) as Tr ( F ) during the initial phase of training . The early-phase Tr ( F ) is measured when the training loss crosses a task-specific threshold . For ImageNet , we use learning rates 0.001 , 0.01 , 0.1 , and = 3.5 . For CIFAR-10 , we use learning rates 0.007 , 0.01 , 0.05 , and = 1.2 . For CIFAR-100 , we use learning rates 0.001 , 0.005 , 0.01 , and = 3.5 . In all cases , training loss reaches between 2 and 7 epochs across different hyper-parameter settings . We repeat similar experiments for different batch sizes in Appendix A.1 . The remaining training details can be found in Appendix G.1 . Results Figure 2 shows the correlation between Tr ( Fi ) and test accuracy across runs with different learning rates . We show results for CIFAR-10 and CIFAR-100 when varying the batch size in Figure 7 in the Appendix . We find that Tr ( Fi ) correlates well with the final generalization in our setting , which provides initial evidence for the importance of Tr ( F ) . It also serves as a stepping stone towards developing a more granular understanding of the role of implicit regularization of Tr ( F ) in the following sections . 4 FISHER PENALTY . To better understand the significance of the identified correlation between Tr ( Fi ) and generalization , we now run experiments in which we directly penalize Tr ( F ) . We focus our attention on the identified effect of high learning rate on Tr ( F ) . Experimental setting We use a similar setting as in the previous section , but we include larger models . We run experiments using Wide ResNet ( Zagoruyko & Komodakis , 2016 ) ( depth 44 and width 3 , with or without BN layers ) , SimpleCNN ( without BN layers ) , DenseNet ( L=40 , K=12 ) ( Huang et al. , 2017 ) and VGG-11 ( Simonyan & Zisserman , 2015 ) . We train these models on either the CIFAR-10 or the CIFAR-100 datasets . Due to larger computational cost , we replace ImageNet with the TinyImageNet dataset ( Le & Yang , 2015 ) in these experiments . To investigate if the correlation of Tr ( Fi ) and final generalization holds more generally , we apply Fisher penalty in two settings . First , we use a learning rate 10-30x smaller than the optimal one , which both incur up to 9 % degradation in test accuracy and results in large value of Tr ( Fi ) . We also remove data augmentation from the CIFAR-10 and the CIFAR-100 datasets to ensure that training with small learning rate does not result in underfitting . In the second setting , we add Fisher penalty in training with an optimized learning rate using grid search ( η∗ ) and train with data augmentation . Fisher penalty penalizes the gradient norm computed using labels sampled from pθ ( y|x ) . We hypothesize that a similar , but weaker , effect can be introduced by other gradient norm regularizers . First , we compare FP to penalizing the input gradient norm ‖gx‖ = ∂∂x ` ( x , y ) , which we denote by GPx ( Varga et al. , 2018 ; Rifai et al. , 2011 ; Drucker & Le Cun , 1992 ) . We also experiment with penalizing the vanilla mini-batch gradient Gulrajani et al . ( 2017 ) , which we denote by GP . Finally , we experiment with penalizing the mini-batch gradient computed with random labels ‖gr‖ = ∂∂x ` ( x , ŷ ) where ŷ is sampled from a uniform distribution over the label set ( GPr ) . We are not aware of any prior work using GP or GPr in supervised training , with the exception of Alizadeh et al . ( 2020 ) where the authors penalized ` 1 norm of gradients to compress the network towards the end of training . We tune the hyperparameters on the validation set . More specifically for α , we test 10 different values spaced uniformly between 10−1×v to 101×v on a logarithmic scale with v ∈ R+ . For TinyImageNet we test 5 alternatives instead . To pick the optimal learning rate , we evaluate 5 values spaced equally on a logarithmic scale . We include the remaining experimental details in the Appendix G.2 . Fisher Penalty improves generalization Table 1 summarizes the results of the main experiment . First , we observe that a suboptimal learning rate ( 10-30x lower than the optimal ) leads to dramatic overfitting . We observe a degradation of up to 9 % in test accuracy , while achieving perfect training accuracy ( see Table 6 in the Appendix ) . Fisher penalty closes the gap in test accuracy between the small and optimal learning rate , and even achieves better performance than the optimal learning rate . A similar performance was observed when minimizing ‖gr‖ . We will come back to this observation in the next section . GP and GPx reduce the early value of Tr ( F ) ( see Table 4 in the Appendix ) . They , however , generally perform worse than Tr ( F ) or GPr and do not fully close the gap between small and optimal learning rate . We hypothesize they improve generalization by a similar but less direct mechanism than Tr ( F ) and GPr . In the second experimental setting , we apply FP to a network trained with the optimal learning rate η∗ . According to Table 2 , Fisher Penalty improves generalization in 4 out of 5 settings . The gap between the baseline and FP is small in 3 out of 5 settings ( below 1 % ) , which is natural given that we already regularize training implicitly by using the optimal η and data augmentation . Geometry and generalization in the early phase of training Here , we investigate the temporal aspect of Fisher Penalty on CIFAR-10 and CIFAR-100 . In particular , we study whether early penalization of Tr ( F ) matters for final generalization . First , we observe that all gradient-norm regularizers reduce the early value of Tr ( F ) closer to Tr ( F ) achieved when trained with the optimal learning rate η∗ . We show this effect with Wide ResNet and VGG-11 on CIFAR-100 in Figure 3 , and for other experimental settings in the Appendix . We also tabulate the maximum achieved values of Tr ( F ) over the optimization trajectory in Appendix A.2 . 0 50 100 Fisher Penalty onset 46 48 50 52 54 Va lid at io n Ac cu ra cy ( % ) ( a ) Wide ResNet CIFAR100 ( w/o aug. ) 0 20 40 60 Fisher Penalty onset 35.0 37.5 40.0 42.5 45.0 47.5 Va lid at io n Ac cu ra cy ( % ) ( b ) VGG-11 CIFAR-100 ( w/o aug. ) 0 50 100 Fisher Penalty onset 72 74 76 78 80 Va lid at io n Ac cu ra cy ( % ) ( c ) Simple CNN CIFAR10 ( w/o aug. ) 0 50 100 Fisher Penalty onset 60 62 64 66 Va lid at io n Ac cu ra cy ( % ) ( d ) DenseNet CIFAR-100 ( w/o aug. ) Figure 4 : Each subplot summarizes an experiment in which we apply Fisher Penalty starting from a certain epoch ( x axis ) and measure the final test accuracy ( y axis ) . Fisher Penalty has to be applied from the beginning of training to close the generalization gap to the optimal learning rate ( c.f . the red horizontal line to the blue horizontal line ) . To test the importance of explicitly penalizing Tr ( F ) early in training , we start applying it after a certain number of epoch E ∈ { 1 , 2 , 4 , 8 , 16 , 32 , 64 , 128 } . We use the best hyperparameter set from the previous experiments . Figure 4 summarizes the results . For both datasets , we observe a consistent pattern . When FP is applied starting from a later epoch , final generalization is significantly worse , and the generalization gap arising from a suboptimal learning rate is not closed .

Summary: This paper empirically investigates the effect of the trace of the Fisher Information Matrix (FIM) early in training has on the generalization of SGD. Authors demonstrate that the effect of optimally chosen learning rate and batch size for SGD can be modeled as an implicit penalty on the trace of FIM. They argue that explicitly penalizing the trace of FIM discourages memorizing noisy labels, thus leading to better generalization. Furthermore, they experimentally show that the early low value of the trace of FIM may bias the optimization towards a flat optimum which has been observed to correlate well with good generalization.