Stochastic Training is Not Necessary for Generalization

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is the backbone of optimization for neural networks , going back at least as far as LeCun et al . ( 1998a ) , and SGD is the de-facto tool for optimizing the parameters of modern neural networks ( Krizhevsky et al. , 2012 ; He et al. , 2015a ; Brown et al. , 2020 ) . A central reason for the success of stochastic gradient descent is its efficiency in the face of large datasets – a noisy estimate of the loss function gradient is generally sufficient to improve the parameters of a neural network and can be computed much faster than a full gradient over the entire training set . At the same time , folk wisdom dictates that small-batch SGD is not only faster but also has a unique bias towards good loss function minima that can not be replicated with full batch gradient descent . Some even believe that stochastic sampling is the fundamental force behind the success of neural networks . These popular beliefs are linked to various properties of SGD , such as its gradient noise , fast escape from saddle points , and its uncanny ability to avoid sub-optimal local minima ( Hendrik , 2017 ; LeCun , 2018 ) . These properties are also attributed in varying degrees to all mini-batched first-order optimizers , such as Adam ( Kingma & Ba , 2015 ) and others ( Schmidt et al. , 2020 ) . But why does stochastic mini-batching really aid generalization ? In this work , we set out to isolate mechanisms which underlie the benefits of SGD and use these mechanisms to replicate the empirical benefits of SGD without stochasticity . In this way , we provide a counterexample to the hypothesis that stochastic mini-batching , which leads to noisy estimates of the gradient of the loss function , is fundamental for the strong generalization success of over-parameterized neural networks . We show that a standard ResNet-18 can be trained with batch size 50K ( the entire training dataset ) and still achieve 95.67 % ( ±0.08 ) validation accuracy on CIFAR-10 , which is comparable to the same network trained with a strong SGD baseline , provided data augmentation is used for both methods . We then extend these findings to train without ( random ) data augmentations , for an entirely nonstochastic full-batch training routine with exact computation of the full loss gradient , while still achieving over 95 % accuracy . Because existing training routines are heavily optimized for smallbatch SGD , the success of our experiments requires us to eschew standard training parameters in favor of more training steps , aggressive gradient clipping , and explicit regularization terms . The existence of this example raises questions about the role of stochastic mini-batching , and by extension gradient noise , in generalization . In particular , it shows that the practical effects of such gradient noise can be captured by explicit , non-stochastic , regularization . This shows that deep learning succeeds even in the absence of mini-batched training . A number of authors have studied relatively large batch training , often finding trade-offs between batch size and model performance ( Yamazaki et al. , 2019 ; Mikami et al. , 2019 ; You et al. , 2020 ) . However , the goal of these studies has been first and foremost to accelerate training speed ( Goyal et al. , 2018 ; Jia et al. , 2018 ) , with maintaining accuracy as a secondary goal . In this study , we seek to achieve high performance on full-batch training at all costs . Our focus is not on fast runtimes or ultra-efficient parallelism , but rather on the implications of our experiments for deep learning theory . In fact , the extremely high cost of each full-batch update makes GD far less efficient than a conventional SGD training loop . We begin our discussion by reviewing the literature on SGD and describing various studies that have sought to explain various successes of deep learning through the lens of stochastic sampling . Then , we explain the hyper-parameters needed to achieve strong results in the full-batch setting and present benchmark results using a range of settings , both with and without data augmentation . 2 PERSPECTIVES ON GENERALIZATION VIA SGD . The widespread success of SGD in practical neural network implementations has inspired theorists to investigate the gradient noise created by stochastic sampling as a potential source of observed generalization phenomena in neural networks . This section will cover some of the recent literature concerning hypothesized effects of stochastic mini-batch gradient descent ( SGD ) . We explicitly focus on generalization effects of SGD in this work . Other possible sources of generalization for neural networks have been proposed that do not lean on stochastic sampling , for example generalization results that only require overparametrization ( Neyshabur et al. , 2018 ; Advani et al. , 2020 ) , large width ( Golubeva et al. , 2021 ) , and well-behaved initialization schemes ( Wu et al. , 2017 ; Mehta et al. , 2020 ) . We will not discuss these here . Furthermore , because we wish to isolate the effect of stochastic sampling in our experiments , we fix an architecture and network hyperparameters in our studies , acknowledging that they were likely chosen because of their synergy with SGD . Notation : We denote the optimization objective for training a neural network by L ( x , θ ) , where θ represents network parameters , and x is a single data sample . Over a dataset X of N data points , { xi } Ni=1 , the neural network training problem is the minimization of L ( θ ) = 1 N ∑ x∈X L ( x , θ ) . ( 1 ) This objective can be optimized via first-order optimization , of which the simplest form is descent in the direction of the negative gradient with respect to parameters θ on a batch B of data points and with step size τk : θk+1 = θk − τk 1 |B| ∑ x∈B ∇L ( x , θk ) . ( 2 ) Now , full-batch gradient descent corresponds to descent on the full dataset B = X , stochastic gradient descent corresponds to sampling a single random data point B = { x } ∼ X ( with or without replacement ) , and mini-batch stochastic gradient descent corresponds to sampling S data points B = { xj } Sj=1 , xj ∼ X at once . When sampling without replacement , the set is commonly reset after all elements are depleted . Although stochastic gradient descent has been used intermittently in applications of pattern recognition as far back as the 90 ’ s , its advantages were debated as late as Wilson & Martinez ( 2003 ) , who in support of SGD discuss its efficiency benefits ( which would become much more prominent in the following years due to increasing dataset sizes ) , in addition to earlier ideas that stochastic training can escape from local minima , and its relationship to Brownian motion and “ quasi-annealing ” , both of which are also discussed in practical guides such as LeCun et al . ( 1998b ) . SGD and critical points While early results from an optimization perspective were concerned with showing the effectiveness and convergence properties of SGD ( Bottou , 2010 ) , later ideas focused on the generalization benefits of stochastic training via navigating the optimization landscape , finding global minima , and avoiding bad local minima and saddlepoints . Ge et al . ( 2015 ) show that stochastic descent is advantageous compared to full-batch gradient descent ( GD ) in its ability to escape saddle points . Although the same conditions actually also allow vanilla gradient descent to avoid saddle-points ( Lee et al. , 2016 ) , full-batch descent is slowed down significantly by the existence of saddle points compared to stochastically perturbed variants ( Du et al. , 2017 ) . Random perturbations also appear necessary to facilitate escape from saddle points in Jin et al . ( 2019 ) . It is also noted by some authors that higher-order optimization , which can alleviate these issues , does perform better in the large-batch regimes ( Martens & Grosse , 2020 ; Yadav , 2020 ; Anil et al. , 2021 ) . Related works further study a critical mini-batch size ( Ma et al. , 2018 ; Jain et al. , 2018 ) after which SGD behaves similarly to full-batch gradient descent ( GD ) and converges slowly . It is unclear though whether the analysis of sub-optimal critical points can explain the benefits of SGD , given that modern neural networks can generally be trained to reach global minima even with deterministic algorithms ( for wide enough networks ( Du et al. , 2019 ) ) . It has been postulated that “ good ” minima that generalize well share geometric properties that make it likely for SGD to find them ( Huang et al. , 2020 ) . Flatness and Noise Shapes One such geometric property of a global minimizer is its flatness ( Hochreiter & Schmidhuber , 1997 ) . Empirically , Keskar et al . ( 2016 ) discuss the advantages of small-batch stochastic gradient descent and propose that finding flat basins is a benefit of smallbatch SGD : Large-batch training converges to models with both lower generalization and sharper minimizers . Although flatness is difficult to measure ( Dinh et al. , 2017 ) , flatness based measures appear to be the most promising tool for predicting generalization in Jiang et al . ( 2019 ) . The analysis of such stochastic effects is often facilitated by considering the stochastic differential equation that arises for small enough step sizes τ from Eq . ( 2 ) under the assumption that the gradient noise is effectively a Gaussian random variable : dθt = −∇L ( θt ) dt+ √ τΣt dWt , ( 3 ) where Σt represents the covariance of gradient noise at time t , and Wt is a Brownian motion modeling it . The magnitude of Σt is inversely proportional to mini-batch size ( Jastrzębski et al. , 2018 ) , and it is also connected to the flatness of minima reached by SGD in Dai & Zhu ( 2018 ) and Jastrzębski et al . ( 2018 ) if Σt is isotropic . Analysis therein as well as in Le ( 2018 ) provides evidence that the step size should increase linearly with the batch size to keep the magnitude of noise fixed . However , the anisotropy of Σt is strong enough to generate behavior that qualitatively differs from Brownian motion around critical points ( Chaudhari & Soatto , 2018 ; Simsekli et al. , 2019 ) and isotropic diffusion is insufficient to explain generalization benefits in Saxe et al . ( 2019 ) . The shape of Σt is thus further discussed in Zhu et al . ( 2019 ) where anisotropic noise induced by SGD is found to be beneficial to reach flat minima in contrast to isotropic noise , Zhou et al . ( 2020 ) where it is contrasted with noise induced by Adam ( Kingma & Ba , 2015 ) , and HaoChen et al . ( 2020 ) who discuss that such parameter-dependent noise , also induced by label noise , biases SGD towards well-generalizing minima . Empirical studies in Wen et al . ( 2020 ) ; Wu et al . ( 2020 ) and Li et al . ( 2021 ) show that large-batch training can be improved by adding the right kind of anisotropic noise . Notably , in all of these works , the noise introduced by SGD is in the end both unbiased and ( mostly ) Gaussian , and its disappearance in full-batch gradient descent should remove its beneficial effects . However , Eq . ( 3 ) only approximates SGD to first-order , while for non-vanishing step sizes τ , Li et al . ( 2017 ) find that a second-order approximation , dθt = −∇ ( L ( θt ) + τ 4 ||∇L ( θ ) ||2 ) dt+ √ τΣt dWt , ( 4 ) does include an implicit bias proportional to the step size . Later studies such as Li et al . ( 2020 ) discuss the importance of large initial learning rates , which are also not well modeled by first-order SDE analysis but have a noticeable impact on generalization . An explicit , non-stochastic bias ? Several of these theoretical investigations into the nature of generalization via SGD rely on earlier intuitions that this generalization effect would not be capturable by explicit regularization Arora et al . ( 2019a ) , who write that “ standard regularizers may not be rich enough to fully encompass the implicit regularization brought forth by gradient-based optimization ” and further rule out norm-based regularizers rigorously . Similar statements have already been shown for the generalization effects of overparametrization in Arora et al . ( 2018 ) who show that no regularizer exists that could replicate the effects of overparametrization in deep linear networks . Yet , Barrett & Dherin ( 2020 ) ; Smith et al . ( 2020b ) find that the implicit regularization induced by GD and SGD can be analyzed via backward-error analysis and a scalar regularizer can be derived . The implicit generalization of mini-batched gradient descent with batches B ∈ B can be ( up to third-order terms and sampling without replacement ) described explicitly by the modified loss function L ( θ ) + τ 4|B| ∑ B∈B ∣∣∣∣∣ ∣∣∣∣∣ 1|B|∑ x∈B ∇L ( x , θ ) ∣∣∣∣∣ ∣∣∣∣∣ 2 , ( 5 ) which simplifies for gradient descent to L ( θ ) + τ 4 ||∇L ( θ ) ||2 , ( 6 ) as found in Barrett & Dherin ( 2020 ) . Training with this regularizer can induce the generalization benefits of larger learning rates , even if optimized with small learning rates , and induce benefits in generalization behavior for small batch sizes when training moderately larger batch sizes . However , Smith et al . ( 2020b ) “ expect this phenomenon to break down for very large batch sizes ” . Related are discussions in Roberts ( 2018 ) and Poggio & Cooper ( 2020 ) , who show a setting in which SGD can be shown to converge to a critical point where ∇L ( xi , θ ) = 0 holds separately for each data point x , a condition which implies that the regularizer of Eq . ( 5 ) is zero . Large-batch training in practice In response to Keskar et al . ( 2016 ) , Hoffer et al . ( 2017 ) show that the adverse effects of ( moderately ) large batch training can be mitigated by improved hyperparameters – tuning learning rates , optimization steps , and batch normalization behavior . A resulting line of work suggests hyperparameter improvements that successively allow larger batch sizes , ( You et al. , 2017 ) with reduced trade-offs in generalization . Yet , parity in generalization between small and large batch training has proven elusive in many applications , even after extensive hyperparameter studies in De et al . ( 2017 ) ; Golmant et al . ( 2018 ) ; Masters & Luschi ( 2018 ) and Smith et al . ( 2020a ) . Golmant et al . ( 2018 ) go on to discuss that this is not only a problem of generalization in their experiments but also one of optimization during training , as they find that the number of iterations it takes to even reach low training loss increases significantly after the critical batch size is surpassed . Conversely , Shallue et al . ( 2019 ) find that training in a large-batch regime is often still possible , but this is dependent on finding an appropriate learning rate that is not predicted by simple scaling rules , and it also depends on choosing appropriate hyperparameters and momentum that may differ from their small-batch counterparts . This reduction of possible learning rates that converge reliably is also discussed in Masters & Luschi ( 2018 ) , but a significant gap in generalization is observed in Smith et al . ( 2020a ) even after grid-searching for an optimal learning rate . Empirical studies continue to optimize hyperparameters for large-batch training with reasonable sacrifices in generalization performance , including learning rate scaling and warmup ( Goyal et al. , 2018 ; You et al. , 2019a ) , adaptive optimizers ( You et al. , 2017 ; 2019b ) , omitting weight regularization on scales and biases ( Jia et al. , 2018 ) , adaptive momentum ( Mikami et al. , 2019 ) , second-order optimization ( Osawa et al. , 2019 ) , and label smoothing Yamazaki et al . ( 2019 ) . Yet , You et al . ( 2020 ) find that full-batch gradient descent can not be tuned to reach the performance of SGD , even when optimizing for long periods , indicating a fundamental “ limit of batch size ” . The difficulty of achieving good generalization with large batches has been linked to instability of training . As discussed in Cohen et al . ( 2020 ) ; Gilmer et al . ( 2021 ) , training with GD progressively increases the sharpness of the objective function until training destabilizes in a sudden loss spike . Surprisingly however , the algorithm does not diverge , but quickly recovers and continues to decrease non-monotonically , while sharpness remains close to a stability threshold . This phenomenon of non-monotone , but effective training close to a stability threshold is also found in Lewkowycz et al . ( 2020 ) .

Summary: The paper examines the role of SGD noise in the generalization performance of neural networks. In particular, the paper examines if the generalization performance of SGD-trained models can be replicated with explicitly regularized models trained with full-batch gradient descent. The authors show that this is indeed possible for a number of different architectures (ResNet, DenseNet, VGG) trained on CIFAR-10 dataset. Based on these observations, the authors conclude that any theory that relies exclusively on stochasticity of training to explain generalization in neural networks is unlikely to capture the true phenomena responsible for the success of deep learning.