Exploring Curvature Noise in Large-Batch Stochastic Optimization

1 INTRODUCTION . Modern datasets and network architectures in deep learning are becoming increasingly larger which results in longer training time . An ongoing challenge for both researchers and practitioners is how to scale up deep learning while keeping training time manageable . Using stochastic gradient descent ( SGD ) with large batch-sizes offers a potential avenue to address scalability issues . Increasing the batch-size used during training improves data parallelization ( Goyal et al. , 2017 ; You et al. , 2018 ) ; it ensures multiple processors ( GPUs ) have sufficiently useful workload at each iteration hence reducing communication-to-cost ratio . Unfortunately , a severe limitation to employing large-batch training in practice is the so-called “ generalization gap ” . While large-batch yields considerable advantages over small-batch on training loss and error per parameter update , it has been verified empirically in LeCun et al . ( 2012 ) ; Keskar et al . ( 2016 ) that we have the opposite effect for test loss and error . To fully realize the benefits of using large-batches in the distributed synchronous setup , it is necessary to engineer large-batch training in a way such that this generalization gap can be closed without sacrificing too much the training performance . This is precisely the central objective of our paper . In this paper , we propose to add diagonal Fisher curvature noise to large-batch gradient updates . We discuss the motivations behind our approach . Under the typical log-likelihood loss assumption , the difference of large-batch and small-batch gradients can be modeled as a Fisher noise . We can expect that adding this noise directly to large-batch gradients gives small-batch performance . While this remedies the generalization issues associated with large-batch training , the resulting convergence performance will be undesirable . To attain our end goal of designing a method which enjoys both desirable convergence and generalization performance , we reduce the amplitude of the noise by changing the covariance structure from full Fisher to diagonal Fisher . We find that in practice , this improves convergence considerably while maintaining good generalization . We give a theoretical analysis of our proposed method in Section 3 . This is done over the convex quadratic setting which often serves as an excellent “ proxy ” for neural network optimization ( Martens , 2010 ) . We additionally provide numerical experiments on standard deep-learning tasks in Section 4 to demonstrate the efficacy of our proposed method . There are two primary takeaways from our empirical analysis . First , we show adding diagonal Fisher noise to large-batch ( a ) SGD step ( b ) Isotropic Gaussian ( c ) Full Fisher ( d ) Diagonal Fisher gradients does not hinder much the training performance ; while at the same time giving a validation accuracy comparable to small-batch training , thereby correcting the “ generalization gap ” . Second , this validation accuracy was attained in the same number of epochs used for small-batch training . This indicates that our method is data-efficient and does not require any lengthening of the optimization process . 2 PRELIMINARIES . 2.1 BACKGROUND . Excess risk decomposition . Let S = { ( x1 , y1 ) , . . . , ( xN , yN ) } be a training set of N input-target samples drawn i.i.d . from an unknown joint probability distribution D. The family of classifiers of interest to us are neural network outputs f ( xi , θ ) , where θ ∈ Θ ⊂ Rd are parameters of the network and Θ here denotes the parameter space of the network . Let ` ( f ( xi , θ ) , yi ) be the loss function measuring the disagreement between outputs f ( xi , θ ) and targets yi . For convenience , we use the notation ` i ( θk ) to denote ` ( f ( xi , θ ) , yi ) . The expected risk and empirical risk functions are defined to be L ( θ ) ≡ E ( x , y ) ∼D [ ` ( f ( x , θ ) , y ) ] , L ( θ ) ≡ 1 N N∑ i=1 ` i ( θ ) . The standard technique to analyze the interplay between optimization and generalization in statistical learning theory is through excess risk decomposition . After k iterations , the excess risk is defined as L̂ ( θk ) ≡ L ( θk ) − inf θ∈Θ L ( θ ) . From Bottou & Bousquet ( 2008 ) ; Chen et al . ( 2018 ) , the expected excess risk is upper-bounded by ES [ L̂ ( θk ) ] ≤ ES [ L ( θk ) − L ( θk ) ] ︸ ︷︷ ︸ Egen +ES [ L ( θk ) − L ( θ∗ ) ] ︸ ︷︷ ︸ Eopt , ( 1 ) where θ∗ = argminθL ( θ ) here is the empirical risk minimizer . The terms Egen and Eopt are the expected generalization and optimization errors respectively . In the machine learning literature , optimization algorithms are often studied from one perspective : either optimization or generalization . The excess risk decomposition suggests that both aspects should be analyzed together ( Chen et al. , 2018 ; Bottou & Bousquet , 2008 ) since the goal of a good optimization-generalization algorithm in machine learning is to minimize the excess risk in the least amount of iterations . Related Work . In the context of large-scale learning , stochastic algorithms are very popular compared to full-batch methods due to lower computational overhead ( Bottou et al. , 2018 ; Bottou , 1991 ) . Despite a decay in optimization performance ( slower rate of convergence in the convex case ( Moulines & Bach , 2011 ; Bottou et al. , 2018 ) ) , stochastic algorithms possess good generalization properties due to the inherent noise of their gradients . There are vast bodies of research literature devoted to understanding the connection between inherent noise and generalization ; for example , scaling the learning rate or batch-size ( Smith & Le , 2017 ; Goyal et al. , 2017 ; Hoffer et al. , 2017 ) or studying the covariance structure of mini-batch gradients ( Jastrzebski et al. , 2017 ; Xing et al. , 2018 ; Zhu et al. , 2018 ; Li et al. , 2015 ) . In the non-convex setting , the inherent noise allows SGD to efficiently escape from saddle points or shallow local minima which tends to give poor generalization ( Zhu et al. , 2018 ; Daneshmand et al. , 2018 ; Jin et al. , 2017 ) . The trade-off between optimization and generalization was also observed for large-batch versus small-batch training in deep learning : both large-batch and small-batch training can reach similar training loss after sufficient number of iterations . Since large-batch gradients have smaller variance , large-batch typically requires less iterations to reach an optimum . However , large-batch training usually has worse test performance compared to small-batch ( Hoffer et al. , 2017 ; Masters & Luschi , 2018 ; Keskar et al. , 2016 ) . It was recently observed that generalization and optimization should not be decoupled in deep learning ( Neyshabur et al. , 2014 ; 2017a ; b ) . The choice of the batch-size has a direct impact on the trade-off between optimization and generalization . In this paper , we aim to design a novel algorithm for large-batch training by increasing the variance of its gradients such that the resulting generalization performance matches small-batch gradient descent . Moreover , we require that this generalization performance be achieved within a number of iterations comparable to original large-batch . 2.2 MOTIVATION . We tackle the large-batch generalization problem through gradient noise injection . Let BL denote large-batch and ML = |BL| denote size of the large-batch . Consider the following modification of large-batch SGD updates θk+1 = θk − αk∇LML ( θk ) + αkD ( θk ) ξk+1 , ξk+1 ∼ N ( 0 , Id ) . ( 2 ) where αk is the learning rate , ∇LML ( θk ) = 1ML ∑ML i=1 ` i ( θk ) is the large-batch gradient and N ( 0 , Id ) is the multivariate Gaussian distribution with mean zero and identity convariance . We can interpret this modification as injecting a Gaussian noise with mean zero and covariance D ( θk ) D ( θk ) > to the gradients . We now focus our attention on finding a suitable matrix D ( θk ) such that the algorithm in Eqn . 2 has desirable convergence and generalization performance . Intrinsic SGD noise . Let B ⊂ S be a mini-batch drawn uniformly and without replacement from S and M = |B| be the size of this chosen mini-batch . We can write the SGD update rule here as θk+1 = θk − αk∇LM ( θk ) = θk − αk∇L ( θk ) + αk ( ∇L ( θk ) −∇LM ( θk ) ︸ ︷︷ ︸ δk ) where ∇L ( θk ) = 1N ∑N i=1∇ ` i ( θk ) is the full-batch gradient . The difference δk = ∇L ( θk ) − ∇LM ( θk ) is the intrinsic noise stemming from mini-batch gradients . The covariance of δk is Cov ( δk , δk ) = ( N −M NM ) · 1 N N∑ i=1 ( ∇ ` i ( θk ) −∇L ( θk ) ) ( ∇ ` i ( θk ) −∇L ( θk ) ) > , ( 3 ) This result can be found in Hu et al . ( 2017 ) ; Hoffer et al . ( 2017 ) . Suppose that the loss function here is the negative log-likelihood , ` i ( θk ) = − log p ( yi|xi , θk ) , where p ( y|x , θ ) is the density function for the model ’ s predictive distribution . Fisher information matrix . Given an input-target pair ( x , y ) , the Fisher information matrix is defined as the expectation of the outer product of log-likelihood gradients , EPx , Py|x [ ∇ log p ( y|x , θ ) ∇ log p ( y|x , θ ) > ] . ( 4 ) The expectation here is taken with respect to the data distribution Px for inputs x and the model ’ s predictive distribution Py|x for targets y . The Fisher matrix can be realized as the second-order Taylor expansion of the KL-divergence on the space of predictive distributions , hence defining a Riemannian metric on this space ( Amari , 1998 ; Martens , 2014 ) . The Fisher matrix for neural networks captures the local curvature information of the parameter space . We give a detailed description of Fisher matrices for feed-forward and convolutional network architectures in Appendix C. In practical situations , we often sample y from the empirical training distribution rather than the predictive distribution . Therefore , we have the empirical Fisher ( Martens , 2014 ) 1 N ∑ ( xi , yi ) ∈S ∇ log p ( yi|xi , θ ) ∇ log p ( yi|xi , θ ) > . For the rest of this paper , unless otherwise specified , all mentions of “ Fisher matrix ” or F ( θ ) refers to the empirical Fisher . Naive method to close “ generalization gap ” . Having modeled the covariance of the mini-batch SGD noise as a Fisher matrix , we now consider a naive approach to close “ generalization gap ” in large-batch training . Let BS denote small-batch and MS = |BS | denote the size of small-batch . Let us choose D ( θk ) in Eqn . 2 to be D ( θk ) = √ ML −MS MLMS √ F ( θk ) , ( 5 ) where √ F ( θk ) is the square-root of Fisher , √ F ( θk ) √ F ( θk ) > = F ( θk ) . The motivation here is that the difference between large-batch and small-batch gradients can be approximated as a Gaussian noise with mean zero and covariance ML−MSMLMS F ( θk ) ( Jastrzebski et al. , 2017 ; Xing et al. , 2018 ; Zhu et al. , 2018 ; Li et al. , 2015 ) . However , there is an immediate issue with this naive approach : if the Gaussian approximation is reasonable , then this algorithm should have similar behavior as smallbatch training which implies poor convergence performance . Indeed , as shown in the 2D convex example in Fig . 1c , adding full Fisher noise recovers SGD behavior . Furthermore , on the CIFAR-10 image classification task trained using ResNet44 , we see in Fig . 3c that the training performance is almost identical to small-batch during the entire optimization process . Thus , this choice of D ( θk ) does not satisfy our requirement of maintaining desirable convergence .

Summary: It has previously been observed that training deep networks using large batch-sizes leads to a larger generalization gap compared to the gap when training with a relatively small batch-size. This paper proposes to add noise sampled from diagonal "empirical" Fisher matrix to the large batch gradient as a method for closing the generalization gap. The authors motivate the use of empirical Fisher for sampling noise by arguing that the covariance of gradients from small batch-sizes can be seen as approximately equal to a scaled version of the Fisher matrix. It is then pointed out that using the Fisher matrix directly to sample noise could in principle close the generalization gap but would lead to slow converegence similar to SGD with a small batch-size. The authors then claim that the convergence speed is better when noise is sampled from the diagonal Fisher matrix instead of the full Fisher matrix. This claim is proven in theory for a convex quadratic loss surface and experiments are conducted to empirically verify this claim both in the quadratic setting are for realistic deep networks. Finally an efficient method for sampling noise from the diagonal empirical Fisher matrix is proposed.