O N F LAT M INIMA , L ARGE M ARGINS AND G ENERALIZABILITY
Anonymous authors
Paper under double-blind review
A BSTRACT
The intuitive connection to robustness and convincing empirical evidence have made the ﬂatness of the loss surface an attractive measure of generalizability for neural networks.
Yet it suffers from various problems such as computational difﬁculties, reparametrization issues, and a growing concern that it may only be an epiphenomenon of optimization methods.
We provide empirical evidence that under the cross-entropy loss once a neural network reaches a non-trivial training error, the ﬂatness correlates (via Pearson Correlation Coefﬁcient) well to the classiﬁcation margins, which allows us to better reason about the concerns surrounding ﬂatness.
Our results lead to the practical recommendation that when assessing generalizability one should consider a margin-based measure instead, as it is computationally more efﬁcient, provides further insight, and is highly correlated to ﬂatness.
We also use our insight to replace the misleading folklore that smallbatch methods generalize better because they are able to escape sharp minima.
Instead we argue that large-batch methods did not have enough time to maximize margins and hence generalize worse.
I NTRODUCTION
Understanding under which conditions a neural network will generalize from seen to unseen data is crucial, as it motivates design choices and principles which can greatly improve performance.
Complexity or generalization measures are used to quantify the properties of a neural network which lead to good generalization.
Currently however, established complexity measures such as VC-Dimension (Vapnik, 1998) or Rademacher Complexity (Bartlett & Mendelson, 2002) do not correlate with the generalizability of neural networks (e.g. see Zhang et al.
Hence many recommendations, such as reducing model complexity, early stopping, or adding explicit regularization are also not applicable or necessary anymore.
Therefore, there is an ongoing effort to devise new complexity measures that may guide recommendations on how to obtain models that generalize well.
A popular approach is to consider the ﬂatness of the loss surface around a neural network.
Hochreiter & Schmidhuber (1997) used the minimum description length (MDL) argument of Hinton & Van Camp (1993) to claim that the ﬂatness of a minimum can also be used as a generalization measure.
Motivated by this new measure Hochreiter & Schmidhuber (1997), and more recently Chaudhari et al.
(2019), developed algorithms with explicit regularization intended to converge to ﬂat solutions.
Keskar et al.
(2016) then presented empirical evidence that ﬂatness relates to improved generalizability and used it to explain the behavior of stochastic gradient descent (SGD) with large and small-batch sizes.
Other works since have empirically corroborated that ﬂatter minima generalize better (e.g. Jiang et al.
; Li et al.
; Bosman et al.
There are however various issues that are still unresolved, which makes using ﬂatness for constructing practical deep learning recommendations difﬁcult.
For one, ﬂatness is computationally expensive to compute.
The most common way to compute the ﬂatness is via the Hessian, which grows quadratically in the number of parameters; this becomes too large when used with modern networks containing millions of parameters.
It is also not clear to what extent ﬂatness is a true measure of generalizability, capable of discerning which neural network will or will not generalize.
Dinh et al.
(2017) showed that reparametrizations affect ﬂatness and a ﬂat model can be made arbitrarily sharp without changing any of its generalization properties.
In addition Probably Approximately Correct (PAC-Bayes) bounds that bound the generalizability in terms of the ﬂatness are also either affected by rescaling, impossible to evaluate or loose (Neyshabur et al., 2017; Arora et al., 2018; Petzka et al., 2020).
While there have been solutions attempting to prevent issues around reparametrization (Liang et al., 2019; Tsuzuku et al., 2019), it remains to establish whether ﬂatness is an epiphenomenon of stochastic gradient descent or other complexity measures as Achille et al.
and Jastrzebski et al.
(2018) are suggesting.
This motivates investigating possible correlations to more well-understood measures of generalization that may help alleviate issues surrounding ﬂat minima, while allowing ﬂat minima to be used when appropriate.
In this paper we will demonstrate a correlation to classiﬁcation margins, which are a well-understood generalization measure.
Margins represent the linearized distance to the decision boundaries of the classiﬁcation region (Elsayed et al., 2018).
An immediate consequence of such a relationship is that to assess generalizability, we could now simply use a computationally cheap and more robust margin based complexity measure.
Our contributions will demonstrate further practical implications of the relationship between margins and ﬂatness which open doors to valuable future work such as a better understanding of why and when a model generalizes and more principled algorithm design.
• We prove that under certain conditions ﬂatness and margins are strongly correlated.
We do so by deriving the Hessian trace for the afﬁne classiﬁer.
Based on its form, we derive an expression in terms of classiﬁcation margins which we show correlates well with the Hessian trace, with increasing training accuracy for various neural network architectures.
By being able relate the two complexity measures, we are now able to provide various practical recommendations, and offer different perspectives on phenomena that may not be explainable without such a view.
These are shown in the following contributions
• We use our insight to replace the misleading folklore that, unlike large-batch methods, small-batch methods are able to escape sharp minima (Keskar et al., 2016).
We instead employ a margin perspective and use our empirical results along with recent results by Banburski et al.
and Hoffer et al.
(2017) to argue that a large batch method was unable to train long enough to maximize the margins.
With our explanation, we help reframe the small and large-batch discussion and build further intuition.
• We show that once a neural network is able to correctly predict the label of every element in the training set it can be made arbitrarily ﬂat by scaling the last layer.
We are motivated by the relationship to margins which suffer from the same issue.
We highlight this scaling issue because, in some instances, it may still be beneﬁcial for algorithm design to be guided by convergence to ﬂat regions.
Hence, we need to account for scaling issues which make it difﬁcult to use ﬂatness to assess whether a network generalizes better than another.
Other works have made connections between ﬂatness and well-behaved classiﬁcation margins via visualizations (see Huang et al.
; Wang et al.
, but they have not demonstrated a quantiﬁable relationship.
Further work has used both the classiﬁcation margins and ﬂatness to construct PAC-Bayes bounds (Neyshabur et al., 2017; Arora et al., 2018), and have related ﬂatness to increased robustness (Petzka et al., 2020; Borovykh et al., 2019) however they did not show when and to what extent these quantities are related.
We structure the paper as follows.
In Section 2, we discuss both our notation and our motivation choosing the cross-entropy loss and the Hessian trace as the ﬂatness measure and provide further background on the classiﬁcation margins.
In Section 3, we present our contribution showing a strong correlation between the margins and ﬂatness by deriving.
In Section 4, we combine recent results based on classiﬁcation margins to offer a different perspective on the misleading folklore on why larger-batch methods generalize worse.
In Section 5, we highlight that networks can be made arbitrarily ﬂat.
Lastly, we offer our thoughts and future work in the Section 6.
P ROBLEM S ETTING
We ﬁrst deﬁne the basic notation that we use for a classiﬁcation task.
We let X represent the input space and Y = { 1 , ..., C } the output space where C are the number of possible classes.
The network architecture is given by φ : Θ × X → R |Y| where Θ is the corresponding parameter space.
We measure the performance of a parameter vector by deﬁning some loss function (cid:96) : R C × Y → R .
If we have have a joint probability distribution D relating input and output space then we would like to minimize the expected loss L D ( θ ) =
E ( x,y ) ∼D
[ (cid:96) ( φ ( θ, x ) , y )] .
Since we usually only have access to some ﬁnite dataset D , we denote the empirical loss by ˜ L D ( θ ) =
| D | (cid:80)
| D | i =1 (cid:96) ( φ ( θ, x i ) , y i ) .
If L D and ˜ L D are close, then we would say a model generalizes well, as we were able to train on a ﬁnite dataset and extrapolate to the true distribution.
We will use the cross-entropy loss which is given by (cid:96) ( φ ( θ, x ) , y )
= − log( S y ( φ ( θ, x ))) where the softmax function S : R C → R C is given by S ( a ) i = [Equation]
(see Goodfellow et al [Equation]
The choice of the cross-entropy function as the loss function has a signiﬁcant impact on how the ﬂatness measure behaves.
Unlike the multiclass mean squared error (MMSE), exponential type losses such as the cross-entropy loss on neural networks have been shown to include implicit regularization which leads to margin maximizing solutions for neural networks (Banburski et al., 2019).
Also, various properties for ﬂat minima which have been proven for the MMSE loss by Mulayoff & Michaeli are not applicable to the cross-entropy loss, further highlighting the fundamental differences between the loss functions.
While the MMSE loss has shown some promise for many classiﬁcation tasks (Hui & Belkin, 2020) the cross-entropy loss is still the loss which is most used and was primarily used for the empirical evidence around ﬂat minima (Keskar et al., 2016; Chaudhari et al., 2019), which motivates our choice.
The qualitative description of a ﬂat region was given by Hochreiter & Schmidhuber (1997) as “a large connected region in parameter space where the error remains approximately constant".
We measure the ﬂatness by the trace of the Hessian of the loss with respect to the parameters (in short the Hessian trace) denoted by Tr ( H θ ( ˜ L D ( θ ))
(Dinh et al., 2017).
Since the Hessian is symmetric, the Hessian trace is equivalent to the sum of its eigenvalues which for a ﬁxed parameter space is proportional to the expected increase of the second order approximation of the loss around a ﬁxed minimum θ in a random direction θ (cid:48) with θ (cid:48) ∼ N ( θ, I ) .
Since we apply ﬂatness arguments only close to minima, we assume that all eigenvalues are positive and that the Hessian trace is a good measure of ﬂatness Sagun et al.
Even though the Hessian is only an approximation of ﬂatness, the Hessian is often preferred as it allows us to reason about various directions in parameter space via its eigenvectors and eigenvalues (see Sagun et al. ; Chaudhari et al. and alleviates the issue of inﬁnitely long but sharp ridges making a minimum inﬁnitely ﬂat (Dinh et al., 2017; Freeman & Bruna, 2016).
The Hessian has also been linked to feature robustness via its use in the second order approximation of the loss (e.g. Petzka et al.
; Borovykh et al.
and is a promising quantity to relate to the margins.
As we are working with non-linear functions it is intractable to compute exact distances to the decision boundary, therefore we use a measure which is related to the linearized distance as described in Elsayed et al.
Under this view, larger margins are better because the data is further from the decision boundary.
Speciﬁcally, we deﬁne the margins as in Neyshabur et al.
(2017): for some vector v ∈ R C and label y
we let the margin of v be γ ( v, y ) =
| v y − max j (cid:54) =
y v j | .
Since we use the margin in different contexts we deﬁne the output margins γ ( φ ( θ, x ) , y ) and the margins of the model output after the softmax layer γ ( S ( φ ( θ, x )) , y ) .
Due to the intuition of margins relating to the regularity of the classiﬁcation regions, they have been proven and shown to be a good generalization measure for linear networks (Langford & Shawe-Taylor, 2003) and later for neural networks (see Bartlett et al. ; Jiang et al. (2018; 2019)) when correctly adjusted.
Due to results by Banburski et al.
(2019) and Soudry et al.
, Poggio et al.
(2019) claimed that a large part of the mystery around generalizability has been solved, since standard optimization methods are maximizing the margin instead of memorizing data.
T HE M ARGIN AND H ESSIAN T RACE R ELATIONSHIP 3. T HE A FFINE C ROSS -E NTROPY H ESSIAN T RACE
Generally, it is difﬁcult to derive a closed form solution of the Hessian trace due to the non-linear nature of neural networks.
To gain insight into what may determine the ﬂatness or sharpness of a solution we consider an afﬁne prediction function for which we derive the following simple and insightful expression for the Hessian trace:
Proposition 3.1 (Afﬁne Cross-Entropy Hessian Trace (ACEHT)) .
Assume an afﬁne predictor given by φ (( θ, b ) , x ) = θx + b where ( θ, b ) ∈ R C × d × R C = Θ .
Then the trace of the Hessian under the cross-entropy loss assuming our predictor function is:
[Equation]
The derivation is in Appendix C. We immediately observe that the trace of the Hessian is a product of both the size of the input and 1 − η ( S ( φ ( θ, x ))) where η ( S ( φ ( θ, x )))
= (cid:80) Cj =1 S 2 j ( φ (Θ , x )) , where we can view 1 − η ( S ( φ ( θ, x )))
as a conﬁdence measure.
In the visualization provided in Figure 1 we clearly see that 1 − η ( S ( φ ( θ, x ))) is only zero when the predictor predicts one class with probability 1 , regardless of whether it is the correct class or not.
When the model is least conﬁdent, namely when every entry is predicted with probability 1 /C , then 1 − η ( S ( φ ( θ, x ))) is also highest.
Hence, in the afﬁne case with a cross-entropy loss the Hessian trace can be seen as an indication of the model conﬁdence in its prediction.
This conﬁdence interpretation is also connected to classiﬁcation margins by observing that S y ≥ γ ( S ( φ ( θ, x )) , y ) and hence (
[Equation]
−
[Equation]
x )) , y ) .
Therefore, if the margins are large then the region will also be ﬂat.
The intuition for this is that the error in the upper bound becomes smaller as S y becomes larger, i.e. when the model predicts correctly and conﬁdently.
We will also provide evidence for a converse, i.e. a ﬂat minimum has large margins, in the following experimental sections.
Finally, we note that without the expression in Proposition 3.1 we would not have been able to derive the upper bound 1 − γ 2 ( S ( φ ( θ, x )) , y ) without guesswork.
[Figure]
[Figure]
[Figure]
3. E XTENSION TO THE N ON -L INEAR C ASE
Now we will attempt to extend the derivation of the previous section to the non-linear case.
This is a challenging undertaking so we will resort to numerical evidence.
To extend the results from the afﬁne case we will consider both the ACEHT and the upper bound ACEHT ( S ( φ ( θ, x ))))
≤ | x | (1 − γ 2 ( S ( φ ( θ, x )) , y )) to which we refer as the "margin bound".
We will compare both quantities to the empirically derived Hessian trace.
To compute the empirical Hessian trace we use the PyHessian package (Yao et al., 2019) which implements Hutchinson’s method (Bai et al., 1996; Avron & Toledo, 2011).
To compare the quantities we will compare them in terms of their distributions over the data.
Speciﬁcally, let ( X, Y ) ∼ D and ﬁx
θ then we compute the Pearson Correlation Coefﬁcient (rvalue)
(Lee Rodgers & Nicewander, 1988) between the random variables
Tr ( H ( (cid:96) ( φ ( θ, X )) , Y )) and ACEHT ( S ( φ ( θ, X ))) and similarly for the margin bound.
The choice of the r-value is natural because in the afﬁne case the ACEHT and the Hessian trace are equivalent, therefore a linear relationship should be expected.
Our method is also more general than just comparing some statistic, such as the average (which is generally used for ﬂatness measures), of the above random variables.
For example, while the smallest margin over the dataset is commonly used a generalization measure (Bartlett et al., 2017; Jiang et al., 2019; Neyshabur et al., 2017), Jiang et al.
(2018) showed that higher moments of the distribution are a much better predictor for generalizability as we will also see in Section 4.
Figure 2 is an examples of such a ﬁt for an afﬁne predictor.
While the high r-value of 0.97 conﬁrms our analytic results, we also observe that the ﬁt is not perfect, as would be expected due to the exact relationship.
The inaccuracies are due to the numerical methods used and become more pronounced the higher the Hessian trace is.
To avoid outliers heavily impacting the linear regression model in the non-linear case, we will use the SciPy function LocalOutlierFactor (Breunig et al., 2000) to remove outliers before ﬁtting the line.
With this we prevent hand picking points to skew the results and will also stabilize our results.
3.2. E MPIRICAL E VIDENCE
We present our results using the convolutional neural network LeNet on the MNIST dataset as they are representative of what we have observed on other architectures, hyperparameters, and datsets (see Appendix B). Our results use stochastic gradient descent with a ﬁxed learning rate and batch size to achieve an appropriate performance on the classiﬁcation task.
Because of the computational difﬁculty of computing the empirical Hessian trace for every single element in the input data, we consider 1,000 randomly selected datapoints from the training-set.
To highlight the computational difﬁculty of using even very optimized numerical tools, such as PyHessian, we note that it takes us roughly 1,5 hours to compute the Hessian trace for the whole MNIST dataset while it only takesseconds for the margins.
In Figure 3 we present the plots related to the correlation of the empirical Hessian trace to the ACEHT and margin bound over the randomly sampled datapoints.
Figures 3a and 3b show that for most of training, the correlation is between 0.8 and 1.
Combining Figures 3a and 3c it can be seen that the r-value increases with the model training accuracy.
Furthermore, the datapoint which are incorrectly predicted do not show a correlation.
With that we conﬁrm the intuition that indeed, ﬂatter solution are more robust and have larger margins.
While we have found ﬂatness and margins to be highly correlated in scenarios in which others have identiﬁed ﬂatness to be a good generalization measure (Jiang et al., 2019; Keskar et al., 2016; Chaudhari et al., 2019), it may just be that this is also an epiphenomenon of stochastic gradient descent or some other process and there may be situations in which the relationship does not hold.
However, our general advice to consider margins more is not impacted by this.
In the scenario where generalizability and ﬂatness have been linked, we have also shown that margins and ﬂatness are correlated, hence it is advantageous to use margins instead due to computational reasons or for more complete intuition.
The only situation in which it is more likely that margins and ﬂatness are not correlated is when ﬂatness has not yet been linked to generalizability.
In such a situation it may also be better to use the better understood margin measure instead of using a ﬂatness measure to assess generalizability.
In the next section we will consider the ﬁrst case, where we examine a general scenario in which ﬂatness has been used to reason about generalizability and offer a more insightful margin perspective. [Figure]
P ERSPECTIVE ON L ARGE AND S MALL -B ATCH METHODS
We now show how our results lead to a better understanding of phenomena which have been misleadingly attributed to ﬂat minima.
To do so, we consider the experiments which rekindled the debate around ﬂat minima by Keskar et al.
(2016), where ﬂatness was used to explain why smallbatch methods tend to generalize better than large-batch methods.
The idea was that small-batch methods converge to ﬂatter minima due to them being able to "escape" sharp minima more easily.
However, it has been shown that the minima of both methods appear to be in the same attractive basin (Sagun et al., 2017; Freeman & Bruna, 2016; Draxler et al., 2018), meaning that small-batch methods do not seem to escape any attractive basin but are merely in a different area of the same attractive basin.
While the results gave credence to ﬂatter minima generalizing better, ﬂatter minima do not seem to provide the full picture for why large-batch methods tend to do worse and we believe that an explanation in terms of the margins is more illuminating.
4. E XPERIMENT S ETUP
We will replicate the experiment by Keskar et al.
for a fully connected network with batchnormalized layers on the MNIST dataset as described in Appendix A.
We chose the large-batch size to be 4096 and the small-batch size to be 256.
To have a fair comparison, we use the same seed and take 10,000 gradient steps for both methods, instead of basing the stopping time on epochs.
We also used stochastic gradient descent without Momentum.
With our setup we observe a similar phenomenon as Keskar et al.
(2016) in Table 1.
The small and large-batch method both attain the same training accuracy and comparable training loss.
However, the small-batch method is at a considerably ﬂatter minimum and generalizes better than the large-batch method.
We will now show that instead of considering the ﬂatness, it would be more insightful to consider margins to explain the difference in generalizability.
While the upper bound of ACEHT is in terms of the softmax margins, we consider the output margins in this section.
The reason is that most margin based generalization measures use the output margins.
Another more practical reason is that towards the end of training, the softmax margins are all very close to 1 making it difﬁcult to visualize and observe the distribution.
We also do not use a normalized version of the margins (such as Bartlett et al. ; Jiang et al.
Our reasoning is that because we use the same architecture, the same dataset, and train in a similar manner the margin distributions will be comparable.
[Table]
4. A M ARGIN P ERSPECTIVE ON L ARGE AND S MALL -B ATCH S IZES
In Figure 4 we see that the output margins and the Hessian trace are correlate as expected from Section 3.
We can also roughly see that the small-batch method has fewer low margins than the largebatch methods.
To emphasize this difference we consider Figure 4c where we plot the histogram and box-plot of the output margin distribution for both the large-batch and small-batch method.
We also display the skewness of each, which is the third moment centered around the mean.
The box-plots and the skewness conﬁrm that the small-batch method is dominated by large margins indicating better generalizability (as discussed in Bartlett et al.
; Jiang et al.
The idea with a left-skewed margin distribution is that the tail with low margin datapoints is mostly compromised of outliers and will not massively affect the robustness to input perturbations.
This soft-margin SVM perspective is in contrast to hard-margin SVMs where the margin is deﬁned to be the minimum of all the distances to the decision boundary (Shalev-Shwartz & Ben-David, 2014).
If a hard-margin view was adopted, then the small-batch method would be predicted to generalize worse, because it has the smallest margin as we see in Figure 4c.
However, the distribution of the small-batch method is also more left skewed, which would point to this minimum being an outlier rather than being indicative of generalizability.
We now want to explain why the small-batch method generalizes well.
As observed in Jastrz˛ebski et al.
(2017) a smaller batch-size is similar to a larger learning-rate, hence at every step the process will advance further than a large batch-method would.
It has already been noted by Hoffer et al.
(2017) that training longer leads the large-batch method to generalize just as well as the small-batch method because it had time to "catch up", even though the decrease in training loss may be barely noticeable.
We have also seen that SGD converges to margin maximizing solutions by Banburski et al.
Therefore, a method that is able to train or advance further, will also be closer to a margin maximizing solution.
We therefore expect that large-batch methods not having had enough time to maximize margins is the driving force behind the large vs small-batch phenomenon. [Figure]
B ECOMING F LATTER WITH I NCREASING M ARGINS
Reparametrization problems such as shown by Dinh et al.
(2017) are neither a new phenomenon nor should they necessarily discourage the design of algorithms which attempt to ﬁnd ﬂat minima.
Rather they inform on what aspects of a generalization measure need to be adjusted to allow them to be used in a practical setting.
For SVMs, the problem of scaling the hyperplane normal to increase margins of correctly classiﬁed points is solved by scaling the normal to make it a unit vector, transforming the functional margin into the geometric margin (Shalev-Shwartz & Ben-David, 2014).
In the case of neural networks, it is also known that scaling the last layer leads to an increase in the margins for data which has been correctly predicted (Neyshabur et al., 2017).
This scaling issues has been successfully addressed (see Bartlett et al. ; Elsayed et al. ; Jiang et al.
Due to the relationship to the classiﬁcation margins it is natural to ask if ﬂatness suffers from a similar problem.
We conﬁrm this with the following Proposition:
Proposition 5.1.
For a given neural network φ let T α : Θ → Θ be such that for all x ∈ X and θ ∈ Θ we have φ ( T α ( θ ) , x ) =
αφ ( θ, x ) .
Now assume that
θ (cid:48) ∈ Θ and a datapoint ( x (cid:48) , y (cid:48) ) for which argmax [Equation]
x
[Equation] [Equation] = y [Equation]
then
[Equation] [Equation]
The proof is in the Appendix D. From the Proposition we immediately derive the following Corollary:
Corollary 5.2.
Assume that φ and θ predict every datapoint in a set D correctly then
[Equation]
Due to the Corollary, if a network has achieved full training accuracy, then the network is equivalent under the T α transformation to an arbitrarily ﬂat network.
We note that there exists such a T α transform for most networks.
Scaling the last layer is one simple instance of such a transform.
Another is that for fully connected and convolutional networks with ReLU non-linearities we observe that by the non-negative homogeneity scaling each layer also results in a valid T α transformation.
The crucial property of the T α map is that it does not change the relative order of the model outputs and therefore, given two networks which have achieved full training accuracy we can not determine which network should generalize better based solely on the ﬂatness of the local-geometry.
We note that Banburski et al.
(2019) mentioned such an issue but they did not discuss it in the context of ﬂat minima and their arguments relied on further structure which we believe is less illuminating than our presentation and proofs.
C ONCLUSIONS
In this paper, we have related ﬂatness to the classiﬁcation margins in a principled manner, in contrast to other works that have made a more intuitive or less quantiﬁable connection (Huang et al., 2019;
Wang et al., 2018; Neyshabur et al., 2017; Petzka et al., 2020).
Our results lead to the immediate practical recommendation of using margins instead of the computationally expensive ﬂatness to assess generalizability.
We also use our results to replace the misleading notion that small-batch methods generalize better because they "escape" sharp minima, instead arguing that small-batch methods have more time to maximize margins.
We were also motivated by the ﬂatness and margin relationship to highlight that neural networks can be made arbitrarily ﬂat.
This implies that the generalizability of two networks can not be distinguished based on ﬂatness and hence needs to be addressed to make ﬂatness a viable generalization measure.
Based on our results, future work may assess whether ﬂatness is an epiphenomenon of the optimization methods, because now recent work on margins (e.g. Banburski et al.
; Soudry et al.
(2018)) can be applied to reason about ﬂatness.
Furthermore, by relating properties of the parameter space (ﬂatness) to properties of the input space (margin) there is now an opportunity to further explore results such as by Sagun et al.
(2017), where they found that the Hessian, with respect to the parameters of a neural network upon convergence, has as many positive eigenvalues as the number of classes in the dataset used.
Overall, our results enable more principled discussion on how ﬂatness may contribute to generalizability.
A A PPENDIX : N ETWORK A RCHITECTURE AND D ATASETS A. N ETWORK A RCHITECTURE
We implement the convolutional neural network LeNet-5 as described in LeCun et al.
Our fully connected neural network with batch normalized layers (FCNBN) is inspired by Keskar et al.
It has a 784-dimensional (MNIST) or 1024-dimensional (CIFAR10) input layer followed by three batch-normalized Ioffe & Szegedy (2015) layers with ReLU non-linearities and a 10dimensional output layer.
A. D ATASETS
[Table]
B A PPENDIX : F LATNESS AND M ARGIN C ORRELATION
Here we present further evidence of the ﬂatness and margin correlation discussed in Section 3.
Like in Section 3 we have used appropriate learning rates and batch sizes to get a reasonable performance for the task, and have observed our results to hold for different hyperparameters.
One instance where we demonstrate two different batch-sizes is for the Fully Connected Network with Batch Normalization on MNIST (Section B.1) where we present results for a batch size of 256 and 4096.
We again only consider 1,000 randomly selected datapoints from the training-set due to the computational difﬁcult of computing the Hessian trace.
If the network achieves full training accuracy and there are no incorrectly classiﬁed datapoints, we set the r-value to zero.
Overall, we observe the same results as in Section 3 and a correlation between 0.8 and 1.
As before, the correlation increases with increasing training accuracy for correctly predicted datapoints.
B. F ULLY C ONNECTED N ETWORK WITH B ATCH N ORMALIZATION ON MNIST B.1.
[Table]
[Figure]
[Figure]
[Figure] [Figure] [Figure] [Figure]
B.1. B
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Table]
[Figure]
[Table]
[Table]
B. L E N ET ON CIFAR
[Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Figure]
[Table]
[Figure]
[Table]
[Table]
[Table]
B
F
ULLY C ONNECTED N ETWORK WITH B ATCH N ORMALIZATION
ON CIFAR
C A PPENDIX : D ERIVATIVES OF THE C ROSS -E NTROPY L OSS C. G ENERAL F ORM
For the general form we consider the cross-entropy loss for a predictor function which is scaled by some scalar α .
Speciﬁcally, we assume an arbitrary input-output pair ( x, y ) ∈
X × Y and will compute the partial derivatives with respect to the parameters θ of the predictor function αφ ( θ, x ) .
Since the equations can become very long we will declutter the notation by letting S = S ( αφ ( θ, x )) , φ = φ ( θ, x ) and for two d-dimensional vectors x,
y ∈ R
d we write (cid:104) x, y (cid:105) =
(cid:80) di
=1 x i y i .
We also denote elementwise multiplication by (cid:12) and let Φ be a matrix such that (Φ) ij = φ j .
Lemma C.1.
The ﬁrst partial derivative of the cross-entropy loss with respect to an element θ i is:
[Equation]
Proof.
[Equation]
With some manipulation we compute ∂ θ i S y :
[Equation]
Combining Equations 3 and 4 we obtain Lemma C
[Equation] [Equation]
Lemma C.2.
The second partial derivative of the cross-entropy loss with respect to elements θ s and θ t is:
[Equation]
Proof.
Differentiating the ﬁrst order derivative given by Lemma C.1 we obtain by the multi-variable chain rule:
[Equation]
To compute ∂ θ t S in Equation 5 we use Equation 4 and obtain:
[Equation]
which after some simpliﬁcation reduces to:
[Equation]
Combining Equations 5 and 6 we obtain Lemma C.2:
[Equation]
C. A FFINE C ROSS -E NTROPY H ESSIAN T RACE
We now present the proof of Proposition 3.1:
Proof.
Throughout the proof we make use of Lemma C.2 and let α = 1 .
We also notice that any second derivative with respect to φ (( θ, b ) , x ) is zero since φ is an afﬁne classiﬁer.
We ﬁrst consider the derivatives with respect to elements of θ where we use θ i,j to denote the element in the i th row and j th column of the matrix θ .
Notice that ∂ θ i,j φ = x j e i which we write as x ij .
The second order derivatives are given by:
[Equation]
when computing the trace we only compute the elements on the diagonal and hence we get:
[Equation]
Now we consider derivatives with respect to elements of b and notice that ∂ b i φ = e i .
For the second derivative we then get:
[Equation]
Finally, summing up the diagonal of the total Hessian we get:
[Equation]
where we used the fact that (cid:80) i S i = 1 .
D A PPENDIX : S CALING P ROOF
To prove Proposition 5.1 we ﬁrst prove the following lemma:
Lemma D.1.
Assume that the argmax of φ is the correct class y and is unique then for k ∈ N , k[-≥ 
-]
≥and i (cid:54) =
y :
[Equation]
Proof.
Let y be such that φ y = max k ∈{ 1 ,...,C } φ k .
For i (cid:54) =
y we have:
[Equation]
where the last line follows from applying L’Hopital’s rule k times.
Since we assumed that y is the only y ∈ { 1 , ..., N } such that φ y = max k ∈{ 1 ,...,N φ k , we have that φ k < φ y for all k (cid:54) =
y .
Hence, as α → ∞ we have e α ( φ k − φ y )
→ 0 .
Therefore:
[Equation]
α →∞
[Equation]
We are now ready to prove Proposition 5.1:
Proof.
We ﬁrst show that the term − α 2 (cid:104)
∂
θ i φ, S (cid:12) ( ∂ θ t φ − (cid:104) ∂ θ t Φ , S (cid:105) ) (cid:105) ) always goes to zero.
Expanding we get:
[Equation]

 We now show that each term in the sum goes to zero.
Consider l (cid:54) =
y : | α
[Equation]
by the Triangle Inequality and
[Equation]
S
[Equation]
≤ α 2 S
[Equation]
C
We let M = (cid:80) Ck =1 | ∂ θ i φ l ( ∂ θ t φ l − ∂ θ t φ k S k ) | and note that M < ∞ for all 0 < α < ∞ since 0 ≤ S k ≤ 1 , ∂ θ i φ l , ∂ θ t φ l , ∂ θ t φ k are constants, and it is a ﬁnite sum.
By Lemma D.1 as α → ∞ we have α 2 S l C → 0 and hence
α 2 (cid:80) Ck =1 ∂ θ i φ l S l ( ∂ θ t
φ l − ∂ θ t φ k S k ) → 0 .
We now consider l = y :
[Equation]
[Equation]
s since S
[Equation]
and using similar arguments and Lemma D.1 follows that this term is zero in the limit.
It is also obvious that α 2 (cid:80) Ck =1 ,
k (cid:54) =
y | ∂ θ t φ k S k | goes to zero.
We are left with showing that α ( ∂ θ t ∂ θ i φ y − (cid:104) ∂ θ t ∂ θ i φ, S (cid:105) goes to zero, this is only guaranteed when y is the true label).
We will use the same method as above:
[Equation]
and the result follows using again Lemma D.1.