Hessian-based Analysis of Large Batch Training and Robustness to Adversaries

Introduction

During the training of a Neural Network (NN), we are given a set of input data ${\bf x}$ with the corresponding labels $y$ drawn from an unknown distribution $\mathcal{P}$. In practice, we only observe a set of discrete examples drawn from $\mathcal{P}$, and train the NN to learn this unknown distribution. This is typically a non-convex optimization problem, in which the choice of hyper-parameters would highly affect the convergence properties. In particular, it has been observed that using large batch size for training often results in convergence to points with poor convergence properties. The main motivation for using large batch is the increased opportunities for data parallelism which can be used to reduce training time . Recently, there have been several works that have proposed different methods to avoid the performance loss with large batch . However, these methods do not work for all networks and datasets. This has motivated us to revisit the original problem and study how the optimization with large batch size affects the convergence behavior.

We first start by analyzing how the Hessian spectrum and gradient change during training for small batch and compare it to large batch size and then draw connection with robust training.

In particular, we aim to answer the following questions: Q1

How is the training for large batch size different than small batch size? Equivalently, what is the difference between the local geometry of the neighborhood that the model converges when large batch size is used as compared to small batch?

A1 We backpropagate the second-derivative and compute its spectrum during training. The results show that despite the arguments regarding prevalence of saddle-points plaguing optimization , that is actually not the problem with large batch size training, even when batch size is increased to the gradient descent limit. In , an approximate numerical method was used to approximate the maximal eigenvalue at a point. Here, by directly computing the spectrum of the true Hessian, we show that large batch size progressively gets trapped in areas with noticeably larger spectrum (and not just the dominant eigenvalue). For details please see §2, especially Figs. 1, 2 and 4.

Q2 What is the connection between robust optimization and large batch size training? Equivalently, how does the batch size affect the robustness of the model to adversarial perturbation?

A2 We show that robust optimization is antithetical to large batch training, in the sense that it favors areas with small spectrum (aka flat minimas). We show that points converged with large batch size are significantly more prone to adversarial attacks as compared to a model trained with small batch size. Furthermore, we show that robust training progressively favors the opposite, leading to points with flat spectrum and robust to adversarial perturbation. We provide empirical and theoretical proof that the inner loop of the robust optimization, where we find the worst case, is a saddle-free optimization problem almost everywhere. Details are discussed in §3, especially Table 1, 7 and Figs. 4, 6.

Limitations: We believe it is critical for every paper to clearly state limitations. In this work, we have made an effort to avoid reporting just the best results, and repeated all the experiments at least three times and found all the findings to be consistent. Furthermore, we performed the tests on multiple datasets and multiple models, including a residual network, to avoid getting results that may be specific to a particular test. The main limitation is that we do not propose a solution for large batch training. We do offer analytical insights into the relationship between large batch and robust training, but we do not fully resolve the problem of large batch training. There have been several approaches to increasing batch size proposed so far , but they only work for particular cases and require extensive hyper-parameter tuning. We are performing an in-depth follow up study to use the results of this paper to better guide large batch size training.

Related Work. Deep neural networks have achieved good performance for a wide range of applications. The diversity of the different problems that a DNN can be used for, has been related to their efficiency in function approximation . However the work of showed that not only the network can perform well on a real dataset, but it can also memorize randomly labeled data very well. Moreover, the performance of the network is highly dependent on the hyper-parameters used for training. In particular, recent studies have shown that Neural Networks can easily be fooled by imperceptible perturbations to input data . Moreover, multiple studies have found that large batch size training suffers from poor generalization capability .

Here we focus on the latter two aspects of training neural networks. presented results showing that large batches converge to a “sharper minima”. It was argued that even if the sharp minima has the same training loss as the flat one, small discrepancies between the test data and the training data can easily lead to poor generalization performance . The fact that “flat minimas” generalize well goes back to the earlier work of . The authors relate flat minima to the theory of minimum description length , and proposed an optimization method to actually favor flat minimas. There have been several similar attempts to change the optimization algorithm to find “better” regions . For instance, proposed entropy-SGD, which uses Langevin dynamics to augment the loss functional to favor flat regions of the “energy landscape”. The notion of flat/sharpness does not have a precise definition. A detailed comparison of different metrics is discussed in , where the authors show that sharp minimas can also generalize well. The authors also argued that the sharpness can be arbitrarily changed by reparametrization of the weights. However, this won’t happen when considering the same model and just changing the training hyper-parameters which is the case here. In , the authors proposed that the training can be viewed as a stochastic differential equation, and argued that the optimum batch size is proportional to the training size and the learning rate.

As our results show, there is an interleaved connection by studying when NNs do not work well. found that they can easily fool a NN with very good generalization by slightly perturbing the inputs. The perturbation magnitude is most of the time imperceptible to human eye, but can completely change the networks prediction. They introduced an effective adversarial attack algorithm known as Fast Gradient Sign Method (FGSM). They related the vulnerability of the Neural Network to linear classifiers and showed that RBF models, despite achieving much smaller generalization performance, are considerably more robust to FGSM attacks. The FGSM method was then extended in to an iterative FGSM, which performs multiple gradient ascend steps to compute the adversarial perturbation. Adversarial attack based on iterative FGSM was found to be stronger than the original one step FGSM. Various defenses have been proposed to resist adversarial attacks . We will later show that there is an interleaved connection between robustness of the model and the large batch size problem.

The structure of this paper is as follows: We present the results by first analyzing how the spectrum changes during training, and test the generalization performances of the model for different batch sizes in §2. In §3, we discuss details of how adversarial attack/training is performed. In particular, we provide theoretical proof that finding adversarial perturbation is a saddle-free problem under certain conditions, and test the robustness of the model for different batch sizes. Also, we present results showing how robust training affects the spectrum with empirical studies. Finally, in §4 we provide concluding remarks.

Large Batch, Generalization Gap and Hessian Spectrum

Setup: The architecture for the networks used is reported in Table 6. In the text, we refer to each architecture by the abbreviation used in this table. Unless otherwise specified, each of the batch sizes are trained until a training loss of 0.001 or better is achieved. Different batches are trained under the same conditions, and no weight decay or dropout is used.

We first focus on large batch size training versus small batch and report the results for large batch training for C1 network on CIFAR-10 dataset, and M1 network on MNIST are shown in Table 1, and Table 7, respectively. As one can see, after a certain point increasing batch size results in performance degradation on the test dataset. This is in line with results in the literature .

As discussed before, one popular argument about large batch size’s poor generalization accuracy has been that large batches tend to get attracted to “sharp” minimas of the training loss. In an approximate metric was used to measure curvature of the loss function for a given model parameter. Here, we directly compute the Hessian spectrum, using power iteration by back-propagating the matvec of the Hessian . Unless otherwise noted, we continue the power iterations until a relative error of $1\text{\sc{e}-}{4}$ reached for each eigenvalue. The Hessian spectrum for different batches is shown in Fig. 1. Moreover, the value of the dominant eigenvalue, denoted by $\lambda_{1}^{\theta}$, is reported in Table 1, and Table 2, respectively (Additional result for MNIST tested using LeNet-5 is given in appendix. Please see Table 7). From Fig. 1, we can clearly see that for all the experiments, large batches have a noticeably larger Hessian spectrum both in the dominant eigenvalue as well as the rest of the 19 eigenvalues. However, note that curvature is a very local measure. It would be more informative to study how the loss functional behaves in a neighborhood around the point that the model has converged. To visually demonstrate this, we have plotted how the total loss changes when the model parameters are perturbed along the dominant eigenvector as shown in Fig. 2, and Fig. 7 for C1 and M1 models, respectively. We can clearly see that the large batch size models have been attracted to areas with higher curvature for both the test and training losses.

This is reflected in the visual figures. We have also added a 3D plot, where we perturb the parameters of C1 model along both the first and second eigenvectors as shown in Fig. 3. The visual results are in line with the numbers shown for the Hessian spectrum (see $\lambda_{1}^{\theta}$) in Table 1, and Table 7. For instance, note the value of $\lambda_{1}^{\theta}$ for the training and test loss for $B=256,\ 2048$ in Table 1 and compare the corresponding results in Fig. 3.

A recent argument has been that saddle-points in high dimension plague optimization for neural networks . We have computed the dominant eigenvalue of the Hessian along with the total gradient during training and report it in Fig. 4. As we can see, large batch size progressively gets attracted to areas with larger spectrum, but it clearly does not get stuck in saddle points .

Large Batch, Adversarial Attack and Robust training

We first give a brief overview of adversarial attack and robust training and then present results connecting these with large batch size training.

Here we focus on white-box adversarial attack, and in particular the optimization-based approach both for the attack and defense. Suppose ${\mathcal{M}}(\theta)$ is a learning model (the neural network architecture), and $({\bf x},y)$ are the input data and the corresponding labels. The loss functional of the network with parameter $\theta$ on $({\bf x},y)$ is denoted by $\mathcal{J}(\theta,{\bf x},y)$. For adversarial attack, we seek a perturbation $\Delta{\bf x}$ (with a bounded $L_{\infty}$ or $L_{2}$ norm) such that it maximizes $\mathcal{J}(\theta,{\bf x},y)$:

where $\mathcal{U}$ is an admissibility set for acceptable perturbation (typically restricting the magnitude of the perturbation). A typical choice for this set is $\mathcal{U}={\bf B}({\bf x},\epsilon)$, a ball of radius $\epsilon$ centered at ${\bf x}$. A popular method for approximately computing $\Delta{\bf x}$, is Fast Gradient Sign Method , where the gradient of the loss functional is computed w.r.t. inputs, and the perturbation is set to:

This is not the only attack method possible. Other approaches include an iterative FGSM method (FGSM-10) or using other norms such as $L_{2}$ norm instead of $L_{\infty}$ (We denote the $L_{2}$ method by $L_{2}Grad$ in our results). Here we also use a second-order attack, where we use the Hessian w.r.t. input to precondition the gradient direction with second order information; please see Table 5 in Appendix for details.

One method to defend against such adversarial attacks, is to perform robust training :

Solving this min-max optimization problem at each iteration requires first finding the worst adversarial perturbation that maximizes the loss, and then updating the model parameters $\theta$ for those cases. Since adversarial examples have to be generated at every iteration, it would not be feasible to find the exact perturbation that maximizes the objective function. Instead, a popular method is to perform a single or multiple gradient ascents to approximately compute $\Delta{\bf x}$. After computing $\Delta{\bf x}$ at each iteration, a typical optimization step (variant of SGD) is performed to update $\theta$.

Next we show that solving the maximization part is actually a saddle-free problem almost everywhere. This property assures us that the Hessian w.r.t input does not have negative eigenvalue which allows us to use CG for performing Newton solver for our 2nd order adversarial perturbation tests in §3.4. This results might also be helpful for finding better optimization strategies for GANS.

2 Adversarial perturbation: A saddle-free problem

Recall that our loss functional is $\mathcal{J}(\theta;{\bf x},y)$. We make following assumptions for the model to help show our theoretical result,

We assume the model’s activation functions are strictly ReLu activation, and all layers are either convolution or fully connected. Here, Batch Normalization layers are accepted. Note that even though the ReLu activation has discontinuity at origin, i.e. $x=0$, ReLu function is twice differentiable almost everywhere.

The following theorem shows that the problem of finding an adversarial perturbation that maximized $\mathcal{J}$, is a saddle-free optimization problem, with a Positive-Semi-Definite (PSD) Hessian w.r.t. input almost everywhere. For details on the proof please see Appendix. A.1.

With Assumption. 1, for a DNN, its loss functional \mathcal{J}(\mbox{\boldmath\theta\unboldmath},{\bf x},y) is a saddle-free function w.r.t. input ${\bf x}$ almost everywhere, i.e.

From the proof of Theorem 1, we could immediately get the following proposition of DNNs:

3 Large Batch Training and Robustness

Here, we test the robustness of the models trained with different batches to an adversarial attack. We use Fast Gradient Sign Method for all the experiments (we did not see any difference with FGSM-10 attack). The adversarial performance is measured by the fraction of correctly classified adversarial inputs. We report the performance for both the training and test datasets for different values of $\epsilon=0.02,0.01$ ($\epsilon$ is the metric for the adversarial perturbation magnitude in $L_{\infty}$ norm). The performance results for C1, and C2 models on CIFAR-10, CR model on CIFAR-100, are reported in the last two columns of Tables 1,and 2 (MNIST results are given in appendix, Table 7). The interesting observation is that for all the cases, large batches are considerably more prone to adversarial attacks as compared to small batches. This means that not only the model design affects the robustness of the model, but also the hyper-parameters used during optimization, and in particular the properties of the point that the model has converged to.

From this result, there seems to be a strong correlation between the spectrum of the Hessian w.r.t. $\theta$ and how robust the model is. However, we want to emphasize that in general there is no correlation between the Hessian w.r.t. weights and the robustness of the model w.r.t. the input. For instance, consider a two variable function \mathcal{J}(\mbox{\boldmath\theta\unboldmath},{\bf x}) (we treat $\theta$ and ${\bf x}$ as two single variables), for which the Hessian spectrum of $\theta$ has no correlation to robustness of $\mathcal{J}$ w.r.t. ${\bf x}$. This can be easily demonstrated for a least squares problem, $L=\|\theta{\bf x}-{\bf y}\|_{2}^{2}$. It is not hard to see the Hessian of $\theta$ and ${\bf x}$ are, ${\bf x}{\bf x}^{T}$ and $\theta\theta^{T}$, respectively. Therefore, in general we cannot link the Hessian spectrum w.r.t. weights to robustness of the network. However, the numerical results for all the neural networks show that models that have higher Hessian spectrum w.r.t. $\theta$ are also more prone to adversarial attacks. A potential explanation for this would be to look at how the gradient and Hessian w.r.t. input (i.e. ${\bf x}$) would change for different batch sizes. We have computed the dominant eigenvalue of this Hessian using power iteration for each individual input sample for both training and testing datasets. Furthermore, we have computed the norm of the gradient w.r.t. ${\bf x}$ for these datasets as well. These two metrics are reported in $\lambda_{1}^{\bf x}$, and $\|\nabla_{x}\mathcal{J}\|$; see Table 1 for details. The results on all of our experiments show that these two metrics actually do not correlate with the adversarial accuracy. For instance, consider C1 model with $B=512$. It has both smaller gradient and smaller Hessian eigenvalue w.r.t. ${\bf x}$ as compared to $B=32$, but it performs acidly worse under adversarial attack. One possible reason for this could be that the decision boundaries for large batches are less stable, such that with small adversarial perturbation the model gets fooled.

4 Adversarial Training and Hessian Spectrum

In this part, we study how the Hessian spectrum and the landscape of the loss functional change after adversarial training is performed. Here, we fix the batch size (and all other optimization hyper-parameters) and use five different adversarial training methods as described in §3.1.

For the sake of clarity let us denote $\mathcal{D}$ to be the test dataset which can be the original clean test dataset or one created by using an adversarial method. For instance, we denote $\mathcal{D}_{FGSM}$ to be the adversarial dataset generated by FGSM, and $\mathcal{D}_{clean}$ to be the original clean test dataset.

Setup: For the MNIST experiments, we train a standard LeNet on MNIST dataset (using M1 network). For the original training, we set the learning rate to 0.01 and momentum to 0.9, and decay the learning rate by half after every 5 epochs, for a total of 100 epochs. Then we perform an additional five epochs of adversarial training with a learning rate of $0.01$. The perturbation magnitude, $\epsilon$, is set to $0.1$ for $L_{\infty}$ attack and $2.8$ for $L_{2}$ attack. We also present results for C3 model on CIFAR-10, using the same hyper-parameters, except that the training is performed for 100 epochs. Afterwards, adversarial training is performed for a subsequent $10$ epochs with a learning rate of $0.01$ and momentum of $0.9$ (the learning rate is decayed by half after five epochs). Furthermore, the adversarial perturbation magnitude is set to $\epsilon=0.02$ for $L_{\infty}$ attack and $1.2$ for $L_{2}$ attack.

The results are shown in Table 3, 4. We can see that after adversarial training the model becomes more robust to these attacks. Note that the accuracy of different adversarial attacks varies, which is expected since the various strengths of different attack method. In addition, all adversarial training methods improve the robustness on adversarial dataset, though they lose some accuracy on $\mathcal{D}_{clean}$, which is consistent with the observations in . As an example, consider the second row of Table 3 which shows the results when FGSM is used for robust training. The performance of this model when tested against the $L_{2}GRAD$ attack method is 63.46% as opposed to 14.32% of the original model ($\mathcal{M}_{ORI}$). The rest of the rows show the results for different algorithms. In §A.4, we give further discussion about the performances of first order and second order attacks.

The main question here is how the landscape of the loss functional is changed after these robust optimizations are performed? We first show a 1-D parametric interpolation between the original model parameters $\theta$ and that of the robustified models, as shown in Fig. 5 (see Fig. 11 for all cases) and 10. Notice the robust models are at a point that has smaller curvature as compared to the original model. To exactly quantify this, we compute the spectrum of the Hessian as shown in Fig. 6, and 12. Besides the full Hessian spectrum, we also report the spectrum of sub-sampled Hessian. The latter is computed by randomly selecting a subset of the training dataset. We denote the size of this subset as $B_{H}$ to avoid confusion with the training batch size. In particular, we report results for $B_{H}=1$ and $B_{H}=320$. There are several important observations here. First, notice that the spectrum of the robust models is noticeably smaller than the original model. This means that the min-max problem of Eq. 3 favors areas with lower curvature. Second, note that even though the total Hessian shows that we have converged to a point with positive curvature (at least based on the top 20 eigenvalues), but that is not necessarily the case when we look at individual samples (i.e. $B_{H}=1$). For a randomly selected batch of $B_{H}=1$, we see that we have actually converged to a point that has both positive and negative curvatures, with a non-zero gradient (meaning it is not a saddle point). To the best of our knowledge this is a new finding, but one that is expected as SGD optimizes the expected loss instead of individual ones.

Now going back to Fig. 4, we show how the spectrum changes during training when we use robust optimization. We can clearly see that with robust optimization the solver is pushed to areas with smaller spectrum. This is a very interesting finding and shows the possibility of using robust training as a systematic means to bias the solver to avoid sharp minimas. A preliminary result is shown in Table 8, where we can see that robust optimization performs better for large batch size training as opposed to the baseline, or when we use the method proposed by . However, we emphasize that the goal is to perform analysis to better understand the problems with large batch size training. More extensive tests are needed before one could claim that robust optimization performs better than other methods.

Conclusion

In this work, we studied NNs through the lens of the Hessian operator. In particular, we studied large batch size training and its connection with stability of the model in the presence of white-box adversarial attacks. By computing the Hessian spectrum we provided several points of evidence that show that large batch size training tends to get attracted to areas with higher Hessian spectrum. We reported the eigenvalues of the Hessian w.r.t. whole dataset, and plotted the landscape of the loss when perturbed along the dominant eigenvector. Visual results were in line with the numerical values for the spectrum. Our empirical results show that adversarial attacks/training and large batches are closely related. We provided several empirical results on multiple datasets that show large batch size training is more prone to adversarial attacks (more results are provided in the supplementary material). This means that not only is the model design important, but also that the optimization hyper-parameters can drastically affect a network’s robustness. Furthermore, we observed that robust training is antithetical to large batch size training, in the sense that it favors areas with noticeably smaller Hessian spectrum w.r.t. $\theta$.

The results show that the robustness of the model does not (at least directly) correlate with the Hessian w.r.t. ${\bf x}$. We also found that this Hessian is actually a PSD matrix, meaning that the problem of finding the adversarial perturbation is actually a saddle-free problem almost everywhere for cases that satisfy this criterion 1. Furthermore, we showed that even though the model may converge to an area with positive curvature when considering all of the training dataset (i.e. total loss), if we look at individual samples then the Hessian can actually have significant negative eigenvalues. From an optimization viewpoint, this is due to the fact that SGD optimizes the expected loss and not the individual per sample loss.

References

Appendix A Appendix

In this section, we give the proof of Theorem 1. The first thing we want to point out is that, although we prove the Hessians of these NNs are positive semi-definite almost everywhere, these NNs are not convex w.r.t. inputs, i.e., ${\bf x}$. The discontinuity of ReLU is the cause. (For instance, consider a combination of two step functions in 1-D, e.g. $f(x)=1_{x\geq 1}+1_{x\geq 2}$ is not a convex function but has second derivative almost everywhere.) However, this has an important implication, that the problem is saddle-free.

Before we go to the proof of Theorem 1, let us prove the following lemma for cross-entropy loss with soft-max layer.

Let $s_{d}=\sum_{j=1}^{d}e^{{\bf s}_{j}}$, ${\bf p}_{i}=\frac{e^{{\bf s}_{i}}}{s_{d}}$, and then it follows that

Then, the second order derivative of $L$ w.r.t. ${\bf s}_{i}{\bf s}_{j}$ is

Assume the input of the soft-max layer is ${\bf s}$ and the cross-entropy is $L({\bf s},y)$. Based on Chain Rule, it follows that

From Assumption. 1 we know that all the layers before the soft-max are either linear or ReLU, which indicates $\dfrac{\partial^{2}{\bf s}}{\partial{\bf x}^{2}}={\bf 0}$ (a tensor) almost everywhere. Therefore, applying chain rule again for the above equation,

It is easy to see \dfrac{\partial^{2}\mathcal{J}(\mbox{\boldmath\theta\unboldmath},{\bf x},y)}{\partial{\bf x}^{2}}\succeq 0 almost everywhere since $\dfrac{\partial^{2}L}{\partial{\bf s}^{2}}\succeq 0$ from Lemma 3.

From above we could see that the Hessian of NNs w.r.t. ${\bf x}$ is at most a rank $c$ (the number of class) matrix, since the rank of the Hessian matrix

is dominated by the term $\frac{\partial^{2}L}{\partial{\bf s}^{2}}$, which is at most rank $c$.

A.2 Attacks Mentioned in Paper

In this section, we show the details about the attacks used in our paper. Please see Table 5 for details.

A.3 Models Mentioned in Paper

In this section, we give the details about the NNs used in our paper. For clarification, We omit the ReLu activation here. However, in practice, we implement ReLu regularity. Also, for all convolution layers, we add padding to make sure there is no dimension reduction. We denote Conv(a,a,b) as a convolution layer having b channels with a by a filters, MP(a,a) as a a by a max-pooling layer, FN(a) as a fully-connect layer with a output and SM(a) is the soft-max layer with a output. For our Conv(5,5,b) (Conv(3,3,b))layers, the stride is 2 (1). See Table 6 for details of all models used in this paper.

A.4 Discussion on Second Order Method

Although second order adversarial attack looks well for MNIST (see Table 3), but for most our experiments on CIFAR-10 (see Table 4), the second order methods are weaker than variations of the gradient based methods. Also, notice that the robust models trained by second order method are also more prone to attack on CIFAR-10, particularly $\mathcal{M}_{FHSM}$ and $\mathcal{M}_{L_{2}HESS}$. We give two potential explanation here.

First note that the Hessian w.r.t. input is a low rank matrix. In fact, as mentioned above, the rank of the input Hessian for CIFAR-10 is at most ten; see Proposition 2, the matrix itself is $3K\times 3K$. Even though we use inexact Newton method along with Conjugate Gradient solver, but this low rank nature creates numerical problems. Designing preconditioners for second-order attack is part of our future work. The second point is that, as we saw in the previous section the input Hessian does not directly correlate with how robust the network is. In fact, the most effective attack method would be to perturb the input towards the decision boundary, instead of just maximizing the loss.

A.5 More Numerical Result for §2 and 3

In this section, we provide more numerical results for §2 and 3. All conclusions from the numerical results are consistency with those in §2 and 3.