LARGE BATCH SIZE TRAINING OF NEURAL NETWORKS WITH ADVERSARIAL TRAINING AND SECOND-ORDER INFORMATION

1 INTRODUCTION . Finding the right NN architecture for a particular application requires extensive hyper-parameter tuning and architecture search , often times on a very large dataset . The delays associated with training NNs is often the main bottleneck in the design process . One of the ways to address this issue to use large distributed processor clusters ; however , to efficiently utilize each processor , the portion of the batch associated with each processor ( sometimes called the mini-batch ) must grow correspondingly . In the ideal case , the hope is to decrease the computational time proportional to the increase in batch size , without any drop in generalization quality . However , large batch training has a number of well known draw backs . These include degradation of accuracy , poor generalization , and poor robustness to adversarial perturbations ( Keskar et al. , 2016 ; Yao et al. , 2018 ) . In order to address these drawbacks , many solutions have been proposed ( Goyal et al. , 2017 ; You et al. , 2017 ; Devarakonda et al. , 2017 ; Smith et al. , 2017 ; Jia et al. , 2018 ) . However , these methods either work only for particular models on a particular dataset , or they require massive hyperparameter tuning , which is often times not discussed in the presentation of results . Note that while extensive hyper-parameter turning may result in good result tables , it is antithetical to the original motivation of using large batch sizes to reduce training time . One solution to reduce the brittleness of SGD to hyper-parameter tuning is to use second-order methods . Full Newton method with line search is parameter-free , and it does not require a learning rate . This is achieved by using a second-order Taylor series approximation to the loss function , instead of a first-order one as in SGD , to obtain curvature information . Schaul et al . ( 2013 ) ; Xu et al . ( 2017a ; b ) show that Newton/quasi-Newton methods outperform SGD for training NNs . However , their re- sults only consider simple fully connected NNs and auto-encoders . A problem with second-order methods is that they can exacerbate the large batch problem , as by construction they have a higher tendency to get attracted to local minima as compared to SGD . For these reasons , early attempts at using second-order methods for training convolutional NNs have so far not been successful . Ideally , if we could find a regularization scheme to avoid local/bad minima during training , this could resolve many of these issues . In the seminal works of El Ghaoui & Lebret ( 1997 ) ; Xu et al . ( 2009 ) , a very interesting connection was made between robust optimization and regularization . It was shown that the solution to a robust optimization problem for least squares is the same as the solution of a Tikhonov regularized problem ( El Ghaoui & Lebret , 1997 ) . This was also extended to the Lasso problem in Xu et al . ( 2009 ) . Adversarial learning/training methods , which are a special case of robust optimization methods , are usually described as a min-max optimization procedure to make the model more robust . Recent studies with NNs have empirically found that robust optimization usually converges to points in the optimization landscape that are flatter and are more robust to adversarial perturbation ( Yao et al. , 2018 ) . Inspired by these results , we explore whether second order information regularized by robust optimization can be used to do large batch size training of NNs . We show that both classes of methods have properties that can be exploited in the context of large batch training to help reduce the brittleness of SGD with large batch size training , thereby leading to significantly improved results . MAIN CONTRIBUTIONS In more detail , we propose an adaptive batch size method based on curvature information extracted from the Hessian , combined with a robust optimization method . The latter helps regularize against sharp minima , especially during early stages of training . We show that this combination leads to superior testing performance , as compared to the proposed methods for large batch size training . Furthermore , in addition to achieving better testing performance , we show that the total number of SGD updates of our method is significantly lower than state-of-the-art methods for large batch size training . We achieve these results without any additional hyper-parameter tuning of our algorithm ( which would , of course , have helped us to tailor our solution to these experiments ) . Here is a more detailed itemization of the main contributions of this work : • We propose an Adaptive Batch Size method for SGD training that is based on second order information , computed by backpropagating the Hessian operator . Our method automatically changes the batch size and learning rate based on Hessian information . We state and prove a result that this method is convergent for a convex problem . More importantly , we empirically test the algorithm for important non-convex problems in deep learning and show that it achieves equal or better test performance , as compared to small batch SGD ( We refer to this method as ABS ) . • We propose a regularization method using robust training by solving a min-max optimization problem . We combine the second order adaptive batch size method with recent results of Yao et al . ( 2018 ) , which show that robust training can be used to regularize against sharp minima . We show that this combination of Hessian-based adaptive batch size and robust optimization achieves significantly better test performance with little computational overhead ( we refer to this Adaptive Batch Size Adversarial method as ABSA ) . • We test the proposed strategies extensively on a wide range of datasets ( Cifar-10/100 , SVHN , TinyImageNet , and ImageNet ) , using different NNs , including residual networks . Importantly , we use the same hyper-parameters for all of the experiments , and we do not perform any kind of tuning of our hyper-parameters to tailor our results . The empirical results show the clear benefit of our proposed method , as compared to the state-of-the-art . The proposed algorithm achieves equal or better test accuracy ( up to 1 % ) and requires significantly fewer SGD updates ( up to 5× ) . • We empirically show that we can use a block approximation of the Hessian operator ( i.e . the Hessian of the last fewer layers ) to reduce the computational overhead of backpropagating the second order information . This approximation is especially effective for deep NNs . While a number of recent works have discussed adaptive batch size or increasing batch size during training ( Devarakonda et al. , 2017 ; Smith et al. , 2017 ; Friedlander & Schmidt , 2012 ; Balles et al. , 2016 ) , to the best of our knowledge this is the first paper to introduce Hessian information and adversarial training in adaptive batch size training , with extensive testing on many datasets . LIMITATIONS We believe that it is important for every work to state its limitations ( in general , but in particular in this area ) . We were particularly careful to perform extensive experiments and repeated all the reported tests multiple times . We test the algorithm on models ranging from a few layers to hundreds of layers , including residual networks as well as smaller networks such as SqueezeNext . An important limitation is that second order methods have additional overhead for backpropagating the Hessian . Currently , most of the existing frameworks do not support ( memory ) efficient backpropagation of the Hessian ( thus providing a structural bias against these powerful methods ) . However , the complexity of each Hessian matvec is the same as a gradient computation ( Martens , 2010 ) . Our method requires Hessian spectrum , which typically needs ten Hessian matvecs ( for power method iterations to reach a tolerance of 1e-2 ) . Thus , the benefits that we show in terms of testing accuracy and reduced number of updates do come at a cost ( see Table 3 for details ) . We measure this additional overhead and report it in terms of wall clock time . Furthermore , we ( empirically ) show that this power iteration needs to be done only at the end of every epoch , thus significantly reducing the additional overhead . Another limitation is that our theory only holds for convex problems ( under certain smoothness assumptions ) . Proving convergence for non-convex setting requires more involved analysis . Recently , Ward et al . ( 2018 ) has provided interesting theoretical guarantees for AdaGrad ( Duchi et al. , 2011 ) in the non-convex setting . Exploring a similar direction for our method is of interest for future work . Another point is that adaptive batch size , prevents one from utilizing all of the processes , as compared to using large batch throughout the training . However , a large data center can handle and accommodate a growing number of requests for processor resources , which could alleviate this . 2 RELATED WORK . Optimization methods based on SGD are currently the most effective techniques for training NNs , and this is commonly attributed to SGD ’ s ability to escape saddle-points and “ bad ” local minima ( Dauphin et al. , 2014 ) . The sequential nature of weight updates in synchronous SGD limits possibilities for parallel computing . In recent years , there has been considerable effort on breaking this sequential nature , through asynchronous methods ( Zhang et al. , 2015 ) or symbolic execution techniques ( Maleki et al. , 2017 ) . A main problem with asynchronous methods is reproducibility , which , in this case , depends on the number of processes used ( Zheng et al. , 2016 ; Agarwal & Duchi , 2011 ) . Due to this issue , recently there have been attempts to increase parallelization opportunities in synchronous SGD by using large batch size training . With large batches , it is possible to distribute more efficiently the computations to parallel compute nodes ( Gholami et al. , 2018a ) , thus reducing the total training time . However , large batch training often leads to sub-optimal test performance ( Keskar et al. , 2016 ; Yao et al. , 2018 ) . This has been attributed to the observation that large batch size training tends to get attracted to local minima or sharp curvature directions , which are not robust to ( possible ) mismatch between training and testing curves ( Keskar et al. , 2016 ) . A full understanding of this , however , remains elusive . There have been several solutions proposed for alleviating the problem with large batch size training . The first notable work here is Goyal et al . ( 2017 ) , where it was shown that by scaling the learning rate , it is possible to achieve the same testing accuracy for large batches . In particular , ResNet50 model was tested on ImageNet dataset , and it was shown that the baseline accuracy could be recovered up to a batch size of 8192 . However , this approach does not generalize to other networks such as AlexNet ( You et al. , 2017 ) , or other tasks such as NLP . In You et al . ( 2017 ) , an adaptive learning rate method ( called LARS ) was proposed which allowed scaling training to a much larger batch size of 32K with more hyper-parameter tuning . Another notable work is Smith et al . ( 2017 ) ( and also Devarakonda et al . ( 2017 ) ) , which proposed a hybrid increase of batch size and learning rate to accelerate training . In this approach , one would select a strategy to “ anneal ” the batch size during the training . This is based on the idea that large batches contain less “ noise , ” and that could be used much the same way as reducing learning rate during training . More recent work Jia et al . ( 2018 ) ; Puri et al . ( 2018 ) proposed mix-precision method to further explore the limit of large batch training . A recent study has shown that anisotropic noise injection could also help in escaping sharp minima ( Zhu et al. , 2018 ) . The authors showed that the noise from SGD could be viewed as anisotropic , with the Hessian as its covariance matrix . Injecting random noise using the Hessian as covariance was proposed as a method to avoid sharp minima . Another recent work by Yao et al . ( 2018 ) has shown that adversarial training ( or robust optimization ) could be used to “ regularize ” against these sharp minima , with preliminary results showing superior testing performance as compared to other methods . The link between robust optimization and regularization is a very interesting observation that has been theoretically proved in the case of Ridge regression ( El Ghaoui & Lebret , 1997 ) , and Lasso ( Bertsimas et al. , 2011 ) . Shaham et al . ( 2015 ) ; Shrivastava et al . ( 2017 ) used adversarial training and showed that the model training using robust optimization is often times more robust to perturbations , as compared to normal SGD training . Similar observations have been made by others ( Szegedy et al. , 2013 ; Goodfellow et al. , 2014 ) .

Summary: This paper studies the large batch size training of neural networks, and incorporates adversarial training and second-order information to improve the efficiency and effectiveness of the proposed algorithm. In particular, the authors use second-order information to automatically generate the step size and batch size in each iteration, and apply adversarial training as a regularization method to improve the test performance. Finally, the authors demonstrate their algorithm and compare it with the baseline algorithms on a wide range of datasets. This paper is clearly written and has the following strength: