C ATASTROPHIC F ISHER E XPLOSION : E ARLY P HASE F ISHER M ATRIX I MPACTS G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
The early phase of training has been shown to be important in two ways for deep neural networks.
First, the degree of regularization in this phase signiﬁcantly impacts the ﬁnal generalization.
Second, it is accompanied by a rapid change in the local loss curvature inﬂuenced by regularization choices.
Connecting these two ﬁndings, we show that stochastic gradient descent (SGD) implicitly penalizes the trace of the Fisher Information Matrix (FIM) from the beginning of training.
We argue it is an implicit regularizer in SGD by showing that explicitly penalizing the trace of the FIM can signiﬁcantly improve generalization.
We further show that the early value of the trace of the FIM correlates strongly with the ﬁnal generalization.
We highlight that in the absence of implicit or explicit regularization, the trace of the
FIM can increase to a large value early in training, to which we refer as catastrophic
Fisher explosion.
Finally, to gain insight into the regularization effect of penalizing the trace of the FIM, we show that it limits memorization by reducing the learning speed of examples with noisy labels more than that of the clean examples, and 2) trajectories with a low initial trace of the FIM end in ﬂat minima, which are commonly associated with good generalization.
I NTRODUCTION
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.
Keskar et al.
; Sagun et al.
; 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 ﬁnal 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.
; Jastrz˛ebski et al.
; Golatkar et al.
; Lewkowycz et al.
; Leclerc & Madry (2020) independently suggest that rapid changes in the local curvature of the loss surface in the early phase critically affects the ﬁnal generalization.
Closely related to our work, Lewkowycz et al.
(2020); Jastrz˛ebski 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 classiﬁcation tasks that the value of Tr( F ) early in training correlates with the ﬁnal 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) signiﬁcantly 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
[Figure] [Figure]
( F i ) is recorded during the early phase of training (2-7 epochs, see the main text for details).
The plots show that early Tr( F ) is predictive of ﬁnal generalization.
Analogous results illustrating the inﬂuence of batch size are shown in Appendix A. 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 ﬂat minimum, as characterized by the trace of the Hessian.
I MPLICIT AND EXPLICIT REGULARIZATION OF THE FIM
Fisher Information Matrix Consider a probabilistic classiﬁcation model p θ ( y | x ) , where θ denotes its parameters.
Let (cid:96) ( x , y ; θ ) be the cross-entropy loss function calculated for input x and label y .
Let g ( x , y ; θ ) =
∂∂θ (cid:96) ( x , y ; θ ) denote the gradient computed for an example ( x , y ) .
The central object that we study is the Fisher Information Matrix F deﬁned as
[Equation]
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 (cid:96) ( 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 classiﬁcations 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; Jastrz˛ebski et al., 2019; Gur-Ari et al., 2018; Lewkowycz et al., 2020; Leclerc & Madry, 2020).
In particular,
Lewkowycz et al.
(2020); Jastrz˛ebski 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 Jastrz˛ebski et al.
Our main contribution is to propose and investigate a speciﬁc mechanism by which using a large learning rate or a small batch size implicitly inﬂuences ﬁnal generalization.
Our ﬁrst 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 beneﬁt of studying the FIM is that it can be directly and efﬁciently 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 ) =
E x ∼X , ˆ y ∼ p
θ ( y | x ) (cid:2) (cid:107) ∂∂θ (cid:96) ( x , ˆ y )
(cid:107) 22 (cid:3) .
To regularize Tr( F ) , we add the following term to the loss function:
[Equation]
: B ; θ ) = 1 B
B (cid:88) i = (cid:96) ( x i , y i ; θ ) + α (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 B
B (cid:88) i = g
( x i , ˆ y i ) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) , (2) where ( x 1: B , y 1: B ) is a mini-batch, ˆ y i is sampled from p θ ( y | x i ) , and α is a hyperparameter.
We refer to this regularizer as Fisher penalty.
The formulation is based on the empirical observation that (cid:13)(cid:13)(cid:13) 1 B (cid:80) Bi
=1 g ( x i , ˆ y i ) (cid:13)(cid:13)(cid:13) 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 everyoptimization 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 overﬁtting, which coincides with a sharp increase in Tr( F ) in the early phase of training.
We also show in Appendix D that these effects cannot be explained by the difference in learning speed between runs with smaller and learning rates.
We call this phenomenon the catastrophic Fisher explosion.
E ARLY - 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.
; Zhang et al.
(2019) trained on the ImageNet dataset (Deng et al., 2009),
(2) ResNetinitialized 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 deﬁne Tr( F i ) as Tr( F ) during the initial phase of training.
The early-phase Tr( F ) is measured when the training loss crosses a task-speciﬁc threshold (cid:15) .
For ImageNet, we use learning rates 0.001,
[Table]
0.01, 0.1, and (cid:15) = 3 .
For CIFAR-10, we use learning rates 0.007, 0.01, 0.05, and (cid:15)
= 1 . 2 .
For
CIFAR-100, we use learning rates 0.001, 0.005, 0.01, and (cid:15)
= 3 .
In all cases, training loss reaches (cid:15) 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( F i ) 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 ﬁnd that Tr( F i ) correlates well with the ﬁnal 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.
F ISHER P ENALTY
To better understand the signiﬁcance of the identiﬁed correlation between Tr( F i ) and generalization, we now run experiments in which we directly penalize Tr( F ) .
We focus our attention on the identiﬁed 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)
(depthand 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( F i ) and ﬁnal 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( F i ) .
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 underﬁtting.
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 (cid:107) g x (cid:107) =
∂∂ x (cid:96) ( x , y ) , which we denote by GP x
(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 (cid:107) g r (cid:107) =
∂∂ x (cid:96) ( x , ˆ y ) where ˆ y is sampled from a uniform distribution over the label set (GP r ).
We are not aware of any prior work using GP or GP r in supervised training, with the exception of Alizadeh et al.
(2020) where the authors penalized (cid:96) 1 norm of gradients to compress the network towards the end of training.
We tune the hyperparameters on the validation set.
More speciﬁcally for α , we test 10 different values spaced uniformly between 10 − 1 × v to 10 1 × v on a logarithmic scale with v ∈ R
+ .
For TinyImageNet
[Table]
[Figure]
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 overﬁtting.
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 (cid:107) g r (cid:107) .
We will come back to this observation in the next section.
GP and GP x reduce the early value of Tr( F ) (see Table 4 in the Appendix).
They, however, generally perform worse than Tr( F ) or GP r 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 GP r .
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 ﬁnal 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.
[Figure]
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, ﬁnal generalization is signiﬁcantly worse, and the generalization gap arising from a suboptimal learning rate is not closed.
4. F ISHER P ENALTY R EDUCES M EMORIZATION
It is not self-evident how regularizing Tr( F ) inﬂuences generalization.
In this section, we provide evidence that regularizing Tr( F ) slows down learning on data with noisy labels.
To study this, we replace labels of the examples in the CIFAR-100 dataset (25% or 50% of the training set) with labels sampled uniformly.
While label noise in real datasets is not uniform, methods that perform well with uniform label noise generally are more robust to label noise in real datasets (Jiang et al., 2020a).
We also know that datasets such as CIFAR-100 contain many labeling errors (Song et al., 2020).
As such, examining how Tr( F ) reduces memorization of synthetic label noise provides an insight into how it improves generalization in our prior experiments.
We expect FP to reduce memorization.
When the predictive distribution p θ ( y | x ) and the true label distribution p ∗ ( y | x ) are both uniform, Tr( F ) of the speciﬁc example x is equivalent to the squared loss gradient norm of the sample example.
The proposed Fisher penalty thus minimizes the contribution of the loss gradient from the training examples whose labels were sampled uniformly.
In other words, the Fisher penalty implicitly suppresses learning noisy examples, under the assumption that clean examples’ label distributions are not uniform.
To study whether the above happens in practice, we compare FP to GP x , GP r , and mixup (Zhang et al., 2018).
While mixup is not the state-of-the-art approach to learning with noisy labels, it is competitive among approaches that do not require additional data nor multiple stages of training.
In particular, it is a component in several state-of-the-art approaches (Li et al., 2020; Song et al., 2020).
For gradient norm based regularizers, we evaluate 6 different hyperparameter values spaced uniformly on a logarithmic scale, and for mixup we evaluate β ∈ { 0 . 2 , 0 . 4 , 0 . 8 , 1 . 6 , 3 . 2 , 6 . 4 } .
We experiment with the Wide ResNet and VGG-11 models.
We describe remaining experimental details in the Appendix G.3.
Results We begin by studying the learning dynamics on data with noisy labels through the lens of training accuracy and mini-batch gradient norm.
We show the results for VGG-11 and ResNetin Figure 5 and Figure 9 in the Appendix.
We observe that FP limits the ability of the model to memorize data more strongly than it limits its ability to learn from clean data.
We can further conﬁrm our interpretation of the effect Tr( F ) has on training by studying the gradient norms.
As visible in
Figure 5, the gradient norm on examples with noisy labels is larger than on clean examples, and the ratio is closer to 1 when large regularization is applied.
We report test accuracy (at the best validation point) in Table 3.
We observe that Tr( F ) reduces memorization competitively to mixup.
Furthermore, FP performs similarly to GP r , which agrees with our interpretation of why FP limits learning on examples with noisy labels.
[Figure]
[Table]
[Figure]
E ARLY Tr( F ) INFLUENCES FINAL CURVATURE
To provide further insight why it is important to regularize Tr( F ) during the early phase of training, we establish a connection between the early phase of training and the wide minima hypothesis (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017) which states that ﬂat minima typically correspond to better generalization.
Here, we use Tr( H ) as a measure of ﬂatness.
Experimental setting We investigate how likely it is for an optimization trajectory to end up in a wide minimum in two scenarios.
1) When optimization exhibits small Tr( F ) early on. 2) When optimization exhibits large Tr( F ) early on.
Speciﬁcally, we train two separate ResNet-26 models for 20 epochs using high and low regularization conﬁgurations.
At epoch 20 we record Tr( F ) for each model.
We then use these two models as initialization for 8 separate models each, and continue training using the low regularization conﬁguration with different random seeds.
The motivation behind this experiment is to show that the degree of regularization in the early phase biases the model towards minima with certain ﬂatness ( Tr( H ) ) even though no further high regularization conﬁgurations are used during the rest of the training.
For all these runs, we record the best test accuracy along the optimization trajectory along with Tr( H ) at the point corresponding to the best test accuracy.
We describe the remaining experimental details in Appendix G.4.
Results We present the result in Figure 6 for the CIFAR-100 datasets, and for CIFAR-10 in
Appendix A.4.
A training run that shows a lower Tr( F ) during the early phase is more likely to end up in a wider minimum as opposed to one that reaches large Tr( F ) during the early phase.
This happens despite that the late phases of both sets of models use the low regularization conﬁguration.
The latter runs have a high variance in the best test accuracy and always end up in sharper minima.
In Appendix G.4 we also show evolution of Tr( H ) throughout training, which suggests that this behavior can be attributed to curvature stabilization happening early during training.
R ELATED W ORK
SGD’s implicit regularization effect has been argued to be a critical component of the empirical success of DNNs (Neyshabur, 2017; Zhang et al., 2016).
Much of it is attributed to the choice of hyperparameters (Keskar et al., 2017; Smith & Le, 2018;
Jastrzebski et al., 2017), the low complexity bias induced by gradient descent (Xu, 2018; Jacot et al., 2018;
Hu et al., 2020) or the cross-entropy loss function (Poggio et al., 2018; Soudry et al., 2018).
However, a more mechanistic understanding of how SGD implicitly regularizes DNNs remains a largely unsolved problem.
Prior work on replicating SGD’s implicit regularization focused mainly on the loss curvature at the ﬁnal minimum (Hochreiter & Schmidhuber, 1997).
Chaudhari et al.
(2019) propose a Langevin dynamics based algorithm for ﬁnding update directions that point towards wide minima.
Wen et al.
(2018) propose to ﬁnd wide minima by averaging gradients at the neighborhood of the current parameter state.
In contrast, we shift the focus to the FIM and the early phase of training.
This new perspective allows us to more directly test our theory by explicitly penalizing Tr( F ) .
Penalizing Tr( F ) is related to regularizing the input gradient norm, which was shown to be an effective regularizer for deep neural networks (Drucker & Le Cun, 1992; Varga et al., 2018).
Chatterjee (2020);
Fort et al.
(2020) show that SGD avoids memorization by extracting commonalities between examples due to following gradient descent directions shared between examples.
Our work is complementary.
We argue that SGD implicitly penalizes Tr( F ) , which also reduces memorization.
Concurrently,
Barrett & Dherin (2020) show that SGD implicitly penalizes the gradient norm for large learning rates and propose GP as an explicit regularizer.
Similarly, we found that SGD implicitly regularizes
Tr( F ) , which is the squared gradient norm under labels sampled from p θ ( y | x ) .
In contrast to them, we connected the implicit regularization effect of SGD to large curvature in the early phase.
We also found GP to be a generally less effective regularizer than FP.
C ONCLUSION
Inspired by recent ﬁndings of rapid changes to the local curvature of the loss surface that happen in the early phase of training (Achille et al., 2019; Jastrz˛ebski et al., 2019; Lewkowycz et al., 2020), we investigated more closely the connection between the loss geometry in the early phase of training of neural networks and generalization.
We proposed and investigated a hypothesis that SGD inﬂuences generalization by implicitly penalizing the trace of the Fisher Information Matrix ( Tr( F ) ) from the very beginning of training.
We show that (1) the value of early Tr( F ) correlates with ﬁnal generalization, and (2) explicitly regularizing Tr( F ) can substantially improve generalization.
To gain further insight into the mechanism by which penalizing Tr( F ) improves generalization, we investigated training on noisy data.
We found that penalizing Tr( F ) reduces memorization by penalizing examples with noisy labels more strongly than clean ones, which seems to happen because it penalizes more strongly their gradient norm.
This sheds new light onto implicit regularization effects in SGD, and suggests the utility of penalizing Tr( F ) as an explicit regularizer.
An interesting topic for the future is to put our ﬁndings in the context of transfer and continual learning.
We hypothesize that catastrophic Fisher explosion (the initial growth of Tr( F ) to a large value) can negatively impact not only generalization, but also transferability of the model.
A PPPENDIX A A DDITIONAL RESULTS A. E ARLY PHASE Tr( F ) CORRELATES WITH FINAL GENERALIZATION
In this section, we present the additional experimental results for Section 3.
The experiments with varying batch size for CIFAR-100 and CIFAR-10 are shown in Figure 7.
The conclusions are the same as discussed in the main text in Section 3.
[Figure]
A. F ISHER P ENALTY
We ﬁrst show additional metrics for experiments summarized in Table 1.
In Table 6 we show the ﬁnal training accuracy.
Table 4 conﬁrms that generally all gradient norm regularizers reduce the maximum value of Tr( F ) (we measure Tr( F ) starting from after one epoch of training because Tr( F ) explodes in networks with batch normalization layers at initialization).
Finally, Table 5 conﬁrms that the regularizers incurred a relatively small additional computational cost.
Figure 8 is a counterpart of Figure 3 for the other two models on the CIFAR-10 and the CIFARdatasets.
[Figure]
V a li d a t i o n
A cc u r a c y ( % )
[Figure]
[Table] [Table]
[Table]
A. F ISHER P ENALTY R EDUCES M EMORIZATION
In this section, we describe additional experimental results for Section 4.1.
Figure 9 is the same as
Figure 5, but for ResNet-50.
[Figure]
A. E ARLY Tr( F ) INFLUENCES FINAL CURVATURE
In this section, we present additional experimental results for Section 5.
The experiment on CIFARis shown in Figure 10.
The conclusions are the same as discussed in the main text in Section 5.
[Figure]
test accuracy along the entire trajectory for the corresponding runs (averaged across seeds).
The plots show that ﬂatter optimization trajectories become biased towards ﬂatter minima early during training, at a coarse scale of hyper-parameter values (red vs blue).
Next, to understand why smaller Tr( F ) during early phase is more likely to end up in a wider ﬁnal minimum, we track Tr( H ) during the entire coarse of training and show that it stabilizes early during training.
In this experiment, we create two sets of hyper-parameters: coarse-grained and ﬁne-grained.
For CIFAR-10, we use batch size S ∈ A ∪ B , where A = { 480 , 500 , 520 } and B = { 80 , 100 , 120 } .
For all batch size conﬁgurations, a learning rate of 0.02 is used.
Overloading the symbols A and
B for CIFAR-100, we use learning rate η ∈ A ∪ B ,
where
A = { 0 . 0008 , 0 . 001 , 0 . 0012 }
and
B = { 0 . 008 , 0 . 01 , 0 . 012 } .
For all learning rate conﬁgurations, a batch size of 100 is used.
In both cases, the elements within each set ( A and B ) vary on a ﬁne-grained scale, while the elements across the two sets vary on a coarse-grained scale.
The remaining details and additional experiments can be found in Appendix G.4.
The experiments are shown in Figure 11.
Notice that after initialization (index 0 on x-axis), the ﬁrst value is computed at epoch 10 (at which point previous experiments show that entanglement starts to hold with late phase).
We make three observations in this experiment.
First, the relative ordering of Tr( H ) values for runs between sets A vs B stay the same after the ﬁrst 10 epochs.
Second, the degree of entanglement is higher between any two epochs when looking at runs across sets A and B , while it is weaker when looking at runs within any one the sets.
Finally, test accuracies for set B runs are always higher than those of set A runs, but this trend is not strong for runs within any one set.
Note that the minimum loss values are roughly at a similar scale for each dataset and
B C OMPUTATION OF Tr( H )
We computed Tr( H ) in our experiments using the Hutchinson’s estimator Hutchinson (1990),
[Equation]
where I is the identity matrix, z is a multi-variate standard Gaussian random variable, and z i ’s are i.i.d.
instances of z .
The larger the value of M , the more accurate the approximation is.
We used M = 30 .
To make the above computation efﬁcient, note that the gradient ∂(cid:96)∂θ only needs to be computed once and it can be re-used in the summation over the M samples.
C A PPROXIMATIONS IN F ISHER P ENALTY
In this section, we describe the approximations made in Fisher Penalty in detail.
Recall, that Tr( F ) can be expressed as
[Equation]
[Equation]
In the preliminary experiments, we found empirically that we can use the norm of the expected gradient rather than the expected norm of the gradient, which is a more direct expression of Tr( F ) :
[Equation]
where N and M are the minibatch size and the number of samples from p θ ( y | x n ) , respectively.
This greatly improves the computational efﬁciency.
With N = B and M = 1 , we end up with the following learning objective function:
[Equation]
We found empirically that (cid:13) (cid:13) (cid:13) 1 B (cid:80) Bi
=1 g ( x i , ˆ y i ) (cid:13) (cid:13) (cid:13) 2 , which we denote by Tr( F B ) , and Tr( F ) correlate well during training.
To demonstrate this, we train SimpleCNN on the CIFAR-10 dataset withdifferent learning rates (from 10 − 3 to 10 − 1 ).
The outcome is shown in Figure 12.
We see that for most of the training, with the exception of the ﬁnal phase, Tr( F B ) and Tr( F ) correlate extremely well.
Equally importantly, we ﬁnd that using a large learning affects both Tr( F B ) and Tr( F ) , which further suggests the two are closely connected.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
Figure 13.), achieving 79.7% test accuracy.
We also update the gradient of Tr( F B ) only every 10 optimization steps.
We found empirically it does not affect generalization performance nor the ability to regularize Tr( F ) in our setting.
However, we acknowledge that it is plausible that this choice would have to be reconsidered in training with very large learning rates or with larger models.
Figure 13 compares learning curves of training with FP recomputed every optimization step, or every 10 optimization steps.
For each, we tune the hyperparameter α , checking 10 values equally spaced between 10 − 2 and 10 0 on a logarithmic scale.
We observe that for the optimal value of α , both validation accuracy and Tr( F ) are similar between the two runs.
Both experiments achieve approximately 80% test accuracy.
Finally, to ensure that using the approximation in Equation 2 does not negatively affect how Fisher
Penalty improves generalization or reduces the value of Tr( F ) , we experiment with a variant of
Fisher Penalty without the approximation.
Please recall that we always measure Tr( F ) (i.e. we do not use approximations in computing Tr( F ) that is reported in the plots), regardless of what variant of penalty is used in regularizing the training.
Speciﬁcally, we augment the loss function with the norm of the gradient computed on the ﬁrst example in the mini-batch as follows
[Equation]
We apply this penalty in each optimization step.
We tune the hyperparameter α , checking 10 values equally spaced between 10 − 4 and 10 − 2 on a logarithmic scale.
Figure 14 summarizes the results.
We observe that the best value of α yields 79.7% test accuracy, compared to 80.02% test accuracy yielded by the Fisher Penalty.
The effect on Tr( F ) is also very similar.
We observe that the best run corresponds to maximum value of Tr( F ) of 24.16, compared to that of 21.38 achieved by Fisher Penalty.
These results suggest that the approximation used in Fisher
Penalty only improves the generalization and ﬂattening effects of Fisher Penalty.
D A CLOSER LOOK AT THE SURPRISING EFFECT OF LEARNING RATE ON THE LOSS GEOMETRY IN THE EARLY PHASE OF TRAINING
It is intuitive to hypothesize that the catastrophic Fisher explosion (the initial growth of the value of Tr( F ) ) occurs during training with a large learning rate, but is overlooked due to not sufﬁciently ﬁne-grained computation of Tr( F ) .
In this section we show evidence against this hypothesis based on the literature mentioned in the main text.
We also run additional experiments in which we compute the value of Tr( F ) at each iteration.
The surprising effect of the learning rate on the geometry of the loss surface (e.g. the value of Tr( F ) ) was demonstrated in prior works (Jastrz˛ebski et al., 2019; Golatkar et al., 2019; Lewkowycz et al., 2020; Leclerc & Madry, 2020).
In particular, Jastrzebski et al.
; Lewkowycz et al.
(2020) show that training with large learning rate rapidly escapes regions of high curvature, where curvature is understood as the spectral norm of the Hessian evaluated at the current point of the loss surface.
Perhaps the most direct experimental data against this hypothesis can be found in Anonymous (2021) in Figure 1, where training with Gradient Descent ﬁnds regions of the loss surface with large curvature for small learning rate rapidly in the early phase of training.
We also run the following experiment to provide further evidence against the hypothesis.
We train
SimpleCNN on the CIFAR-10 dataset using two different learning rates, while computing the value of Tr( F ) for every mini-batch.
We use 128 random samples in each iteration to estimate Tr( F ) .
We ﬁnd that training with a large learning rate never (even for a single optimization step) enters a region with the value of Tr( F ) as large as reached during training with a small learning rate.
Figureshows the experimental data.
We also found similar to hold when varying the batch size, see Section E, which further shows that the observed effects cannot be explained by the difference in learning speed incurred by using a small learning rate.
To summarize, both the published evidence of Jastrzebski et al.
; Lewkowycz et al.
Anonymous (2021), as well as our additional experiments are inconsistent with the hypothesis that the results in this paper can be explained by differences in training speed between experiments using large and small learning rates. [Figure] [Figure] = 0.
= 0. E C ATASTROPHIC F ISHER E XPLOSION HOLDS IN TRAINING WITH LARGE BATCH - SIZE In this section [Figure]
we show evidence that the conclusions transfer to large batch size training.
Namely, we show that (1) catastrophic Fisher explosion also occurs in large batch size training
[Figure]
and (2) Fisher
Penalty can improve generalization and close the generalization gap due to using a large batch size (Keskar et al., 2017)
[Figure]
We ﬁrst train SimpleCNN on the CIFAR-10 dataset using three different batch sizes, while computing the value of Tr( F ) for every mini-batch.
We use 128 random samples in each iteration to estimate
Tr( F ) .
Figure 16 summarizes the experiment.
We observe that training with a large batch size enters a region of the loss surface with a substantially larger value of Tr( F ) than the small batch size.
Next, we run a variant of one of the experiments in Table 1.
Instead of using a suboptimal (smaller) learning rate, we use a suboptimal (larger) batch size.
Speciﬁcally, we train SimpleCNN on the
CIFAR-10 dataset (without augmentation) with a 10x larger batch size while keeping learning rate the same.
Using a larger batch size results in 3 .
24% lower test accuracy ( 76 . 94 % compared to 73 .
7% test accuracy, c.f. with Table 1).
We next experiment with Fisher Penalty.
We apply the penalty in each optimization step and use the ﬁrst 128 examples when computing the penalty.
We also use a 2x lower learning rate, which stabilizes training but does not improve generalization on its own (training with this learning rate reaches 73 .
59% test accuracy).
Figure 17 shows Tr( F ) and validation accuracy during training for different values of the penalty.
We observe that Fisher Penalty improves test accuracy from 73 .
59% to 78 .
Applying Fisher Penalty also effectively reduces the peak value of Tr( F ) /
Taken together, the results suggest that Catastrophic Fisher explosion holds in large batch size training; using a small batch size improves generalization by a similar mechanism as using a large batch size, which can be introduced explicitly in the form of Fisher Penalty
[Figure]
F Tr( H ) AND Tr( F ) CORRELATE STRONGLY
We demonstrate a strong correlation between Tr( H ) and Tr( F ) for DenseNet, ResNet-56 and
SimpleCNN in Figure 18.
We calculate Tr( F ) using a mini-batch.
We see that Tr( F ) has a smaller magnitude (because we use the mini-batch gradient which has lower variance), but correlates extremely well with Tr( H ) .
[Figure]
G A DDITIONAL E XPERIMENTAL D ETAILS G. E ARLY PHASE Tr( F ) CORRELATES WITH FINAL GENERALIZATION
Here, we describe additional details for experiments in Section
[Figure] In the experiments with batch size, for CIFAR-10, we use batch sizes 100, 500 and 700, and (cid:15)
For CIFAR-100, we use batch sizes 100, 300 and 700, and (cid:15) = 3 . 5 .
These thresholds are crossed between 2 and 7 epochs across different hyperparameter settings.
The remaining details for CIFAR100 and CIFAR-10 are the same as described in main text.
The optimization details for these datasets are as follows.
ImageNet : No data augmentation was used in order to allow training loss to converge to small values.
We use a batch size of 256.
Training is done
using SGD with momentum set to 0.9, weight decay set to 1 e − 4 , and with
base learning rate as per the aforementioned details. Learning rate is dropped by a factor of 0.1 after 29 epochs and training is ended at around 50 epochs at which most runs converge to small loss values.
No batch normalization is used and weight are initialized using Fixup Zhang et al.
For each hyperparameter setting, we run two experiments with different random seeds due to the computational overhead.
We compute Tr( F ) using 2500 samples (similarly to ? ).
CIFAR-10 :
We used random ﬂipping as data augmentation.
In the experiments with variation in learning rates, we use a batch size of
In the experiments with variation in batch size, we use a learning rate of 0.02. Training is
done using SGD with momentum set to 0.9, weight decay set to 1 e − 5 , and with base learning rate as per the aforementioned details.
Learning rate is dropped by a factor of 0.5 at epochs 60, 120, and 170, and training is ended at 200 epochs at which most runs converge to small loss values.
No batch normalization is used and weight are initialized using Arpit et al.
For each hyperparameter setting , we run experiments
with different random seeds.
We compute Tr( F ) using 5000 samples.
CIFAR-100 : No data augmentation was used for CIFAR-100 to allow training loss to converge to small values.
We used random ﬂipping as data augmentation for CIFAR-10.
In the experiments with variation in learning rates, we use a batch size of 100.
In the experiments with variation in batch size, we use a learning rates of 0.02.
Training is done using SGD with momentum set to 0.9, weight decay set to 1 e − 5 , and with base learning rate as per the aforementioned details.
Learning rate is dropped by a factor of 0.5 at epochs 60, 120, and 170, and training is ended at 200 epochs at which most runs converge to small loss values.
No batch normalization is used and weight are initialized using Arpit et al.
For each hyperparameter setting, we run 32 experiments with different random seeds.
We compute Tr( F ) using 5000 samples.
G. F ISHER P ENALTY
Here, we describe the remaining details for the experiments in Section 4.
We ﬁrst describe how we tune hyperparameters in these experiments.
In the remainder of this section, we describe each setting used in detail .
Tuning hyperparameters In all experiments, we refer to the optimal learning rate η ∗ as the learning rate optimized using grid search.
In most experiments we check 5 different learning rate values uniformly spaced on a logarithmic scale, usually between 10 − 2 and 10 0 .
In some experiments we adapt the range to ensure that the range includes the optimal learning rate.
We tune the learning rate only once for each conﬁguration (i.e. we do not repeat it for different random seeds).
In the ﬁrst setting, for most experiments involving gradient norm regularizers, we use 10 × smaller learning rate than η ∗ .
For TinyImageNet, we use 30 × smaller learning rate than η ∗ .
To pick the regularization coefﬁcient α , we evaluate 10 different values uniformly spaced on a logarithmic scale between 10 − 1 × v to 10 1 × v with v ∈ R + .
We choose the best performing α according to best validation accuracy.
We pick the value of v manually with the aim that the optimal α is included in this range.
We generally found that v = 0 .
01 works well for GP, GP r , and FP.
For GP x we found in some experiments that it is necessary to pick larger values of v .
Measuring Tr( F ) We measure Tr( F ) using the number of examples equal to the batch size used in training.
For experiments with Batch Normalization layers, we use Batch Normalization in evaluation mode due to the practical reason that computing Tr( F ) uses batch size of 1, and hence Tr( F ) is not deﬁned for a network with Batch Normalization layers in training mode.
DenseNet on the CIFAR-100 dataset We use the DenseNet (L=40, k=12) conﬁguration from
Huang et al.
We largely follow the experimental setting in Huang et al.
We use as the validation
We train the model using SGD with momentum of 0.9, a batch size of 128, and weight decay of 0.0001.
Following Huang et al.
(2017), we train forepochs and decay the learning rate by a factor of 0.1 after epochs 150 and 225.
To reduce variance, in testing we update Batch statistics using from the set.
Wide ResNet on the CIFAR-100 dataset We train Wide ResNet (depth 44 and width 3, without
Batch Normalization layers).
We largely follow experimental setting in He et al.
(2015).We use the standard data augmentation and data normalization for CIFAR-100.
We hold out randomexamples as the validation set.
We train the model using SGD with momentum of 0.9, a batch size of 128, weight decay of 0.0010.
Following He et al.
, we train for 300 epochs and decay the learning rate by a factor of 0.1 after epochs 150 and 225.
We remove Batch Normalization layers.
To ensure stable training we use the SkipInit initialization (De & Smith, 2020).
VGG-11 on the CIFAR-100 dataset We adapt the VGG-11 model (Simonyan & Zisserman, 2015) to CIFAR-100.
We do not use dropout nor Batch Normalization layers.
We hold out randomexamples as the validation set.
We use the standard data augmentation (where noted) and data normalization for CIFAR-100.
We train the model using SGD with momentum of 0.9, a batch size of 128, and weight decay of 0.0001.
We train the model for 300 epochs, and decay the learning rate by a factor of 0.1 after every 40 epochs starting from epoch 80.
SimpleCNN on the CIFAR-10 dataset We also run experiments on the CNN example architecture from the Keras example repository (Chollet & others, 2015) 1 , which we change slightly.
Speciﬁcally, we remove dropout and reduce the size of the ﬁnal fully-connected layer to 128.
We train it forepochs and decay the learning rate by a factor of 0.1 after the epochs 150 and 225.
We train the model using SGD with momentum of 0.9, a batch size of 128.
Wide ResNet on the TinyImageNet dataset We train Wide ResNet (depth 44 and width 3, with
Batch Normalization layers) on TinyImageNet Le & Yang (2015).
TinyImageNet consists of subset of 100,000 examples from ImageNet that we downsized to 32 × 32 pixels.
We train the model using
SGD with momentum of 0.9, a batch size of 128, and weight decay of 0.0001.
We train forepochs and decay the learning rate by a factor of 0.1 after epochs 150 and 225.
We train the model using SGD with momentum of 0.9, a batch size of 128.
We do not use validation in TinyImageNet due to its larger size.
To reduce variance, in testing we update Batch Normalization statistics using 100 batches from the training set.
G. F ISHER P ENALTY R EDUCES M EMORIZATION
Here, we describe additional experimental details for Section 4.1.
We use two conﬁgurations described in Section G.2: VGG-11 trained on CIFAR-100 dataset, and Wide ResNe trained on the CIFARdataset.
We tune the regularization coefﬁcient α in the range { 0 . 01 , 0 . 1 , 0 . 31 , 10 } , with the exception of GP x for which we use the range { 10 , 30 , 100 , 300 , 1000 } .
We tuned mixup coefﬁcient in the range { 0 . 4 , 0 . 8 , 1 . 6 , 3 . 2 , 6 . 4 } .
We removed weight decay in these experiments.
G. E ARLY Tr( F ) INFLUENCES FINAL CURVATURE
CIFAR-10 :
We used random ﬂipping as data augmentation for CIFAR-10.
We use a learning rate of 0.02 for all experiments.
Training is done using SGD with momentum 0.9, weight decay 1 e − 5 , and with batch size as shown in ﬁgures.
Learning rate is drop by a factor of 0.5 at 80, 150, andepochs, and training is ended at 250 epochs.
No batch normalization is used and weight are initialized using Arpit et al.
For each batch size, we run 32 experiments with different random seeds.
We compute Tr( F ) using 5000 samples.
CIFAR-100 : No data augmentation is used.
We use a batch size of 100 for all experiments.
Training is done using SGD with momentum 0.9, weight decay 1 e − 5 , and with base learning rate as shown in ﬁgures.
Learning rate is drop by a factor of 0.5 at 80, 150, and 200 epochs, and training is ended at 250 epochs.
No batch normalization is used and weight are initialized using Arpit et al. For
Accessible each learning rate, we run 32 experiments with different random seeds.
We compute Tr( F ) using 5000 samples.