Hyperparameter Tuning and Implicit Regularization in Minibatch SGD

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is the most popular optimization algorithm in deep learning , but it remains poorly understood . A number of papers propose simple scaling rules that predict how changing the learning rate and batch size will influence the final performance of popular network architectures ( Hoffer et al. , 2017 ; Goyal et al. , 2017 ; Smith et al. , 2017 ; Jastrzębski et al. , 2017 ) . Some of these scaling rules are contradictory , and Shallue et al . ( 2018 ) argue that none of these simple prescriptions work reliably across multiple architectures . Some papers claim SGD with Momentum significantly outperforms SGD without Momentum ( Sutskever et al. , 2013 ) , but others observe little difference between both algorithms in practice ( Kidambi et al. , 2018 ; Zhang et al. , 2019 ) . We hope to clarify this debate . We argue that minibatch stochastic gradient descent exhibits two regimes with different behaviours : a noise dominated regime and a curvature dominated regime ( Ma et al. , 2017b ; McCandlish et al. , 2018 ; Liu & Belkin , 2018 ) . The noise dominated regime typically arises for small or moderate batch sizes , while the curvature dominated regime typically arises when the batch size is large . The curvature dominated regime may also arise if the epoch budget is small or the loss is poorly conditioned ( McCandlish et al. , 2018 ) . Our extensive experiments demonstrate that , 1 . In the noise dominated regime , the final training loss and test accuracy are independent of batch size under a constant epoch budget , and the optimal learning rate increases as the batch size rises . In the curvature dominated regime , the optimal learning rate is independent of batch size , and the training loss and test accuracy degrade with increasing batch size . The critical learning rate which separates the two regimes varies between architectures . 2 . If specific assumptions are satisfied , then the optimal learning rate is proportional to batch size in the noise dominated regime . These assumptions hold for most tasks . However we observe a square root scaling rule when performing language modelling with an LSTM . This is not surprising , since consecutive gradients in a language model are not independent . 3 . SGD with Momentum and learning rate warmup do not outperform vanilla SGD in the noise dominated regime , but they can outperform vanilla SGD in the curvature dominated regime . There is also an active debate regarding the role of stochastic gradients in promoting generalization . It has been suspected for a long time that stochastic gradients sometimes generalize better than full batch gradient descent ( Heskes & Kappen , 1993 ; LeCun et al. , 2012 ) . This topic was revived by Keskar et al . ( 2016 ) , who showed that the test accuracy often falls if one holds the learning rate constant and increases the batch size , even if one continues training until the training loss ceases to fall . Many authors have studied this effect ( Jastrzębski et al. , 2017 ; Smith & Le , 2017 ; Chaudhari & Soatto , 2018 ) , but to our knowledge no paper has demonstrated a clear generalization gap between small and large batch training under a constant step budget on a challenging benchmark while simultaneously tuning the learning rate . This phenomenon has also been questioned by a number of authors . Shallue et al . ( 2018 ) argued that one can reduce the generalization gap between small and large batch sizes if one introduces additional regularization ( we note that this is consistent with the claim that stochastic gradients can enhance generalization ) . Zhang et al . ( 2019 ) suggested that a noisy quadratic model is sufficient to describe the performance of neural networks on both the training set and the test set . In this work , we verify that small or moderately large batch sizes substantially outperform very large batches on the test set in some cases , even when compared under a constant step budget . However the batch size at which the test accuracy begins to degrade can be larger than previously thought . We find that the test accuracy of a 16-4 Wide-ResNet ( Zagoruyko & Komodakis , 2016 ) trained on CIFAR-10 for 9725 updates falls from 94.7 % at a batch size of 4096 to 92.8 % at a batch size of 16384 . When performing language modelling with an LSTM on the Penn TreeBank dataset for 16560 updates ( Zaremba et al. , 2014 ) , the test perplexity rises from 81.7 to 92.2 when the batch size rises from 64 to 256 . We observe no degradation in the final training loss as the batch size rises in either model . These surprising results motivated us to study how the optimal learning rate depends on the epoch budget for a fixed batch size . As expected , the optimal test accuracy is maximized for a finite epoch budget , consistent with the well known phenomenon of early stopping ( Prechelt , 1998 ) . Meanwhile the training loss falls monotonically as the epoch budget increases , consistent with classical optimization theory . More surprisingly , the learning rate that maximizes the final test accuracy decays very slowly as the epoch budget increases , while the learning rate that minimizes the training loss decays rapidly.1 These results provide further evidence that the noise in stochastic gradients can enhance generalization in some cases , and they suggest novel hyper-parameter tuning strategies that may reduce the cost of identifying the optimal learning rate and optimal epoch budget . We describe the noise dominated and curvature dominated regimes of SGD with and without Momentum in section 2 . We focus on the analogy between SGD and stochastic differential equations ( Gardiner et al. , 1985 ; Welling & Teh , 2011 ; Mandt et al. , 2017 ; Li et al. , 2017 ) , but our primary contributions are empirical and many of our conclusions can be derived from different assumptions ( Ma et al. , 2017b ; Zhang et al. , 2019 ) . In section 3 , we provide an empirical study of the relationship between the optimal learning rate and the batch size under a constant epoch budget , which verifies the existence of the two regimes in practice . In section 4 , we study the relationship between the optimal learning rate and the batch size under a constant step budget , which confirms that stochastic gradients can introduce implicit regularization enhancing the test set accuracy . Finally in section 5 , we fix the batch size and consider the relationship between the optimal learning rate and the epoch budget . 2 THE TWO REGIMES OF SGD . Full batch gradient descent is in the “ curvature dominated ” regime , where the optimal learning rate is determined by the curvature of the loss function . The full batch gradient descent update on the ith step is given by ωi+1 = ωi − dCdω ∣∣ ω=ωi , where the loss C ( ω ) = 1N ∑N j=1 C ( ω , xj ) is a function of the parameters ω and the training inputs { xj } Nj=1 , and denotes the learning rate . In order to minimize the loss as quickly as possible , we will set the learning rate at the start of training as large as 1These conclusions might change if the number of updates is very small , in which case the optimal learning rate schedules on the training test and the test set may agree . we can while avoiding divergences or instabilities . To build our intuition for this , we approximate the loss by a strictly convex quadratic , C ( ω ) ≈ 12ω THω . For simplicity , we assume the minimum lies at ω = 0 . Substituting this approximation into the parameter update , we conclude ωi+1 = ωi − ωTi H. In the eigenbasis of H , ωi+1 = ωi ( I − Λ ) , where I denotes the identity matrix and Λ denotes a diagonal matrix comprising the eigenvalues of H . The iterates will converge so long as the learning rate < crit , where crit = 2/λmax is the critical learning rate above which training diverges , and λmax is the largest Hessian eigenvalue . We call this inequality the curvature constraint , and the optimal initial learning rate with full batch gradients will be just below crit . For clarity , although the critical learning rate will perform poorly for high curvature directions of the loss , we can introduce learning rate decay to minimize the loss along these directions later in training ( Ge et al. , 2019 ) . Of course , in realistic loss landscapes this critical learning rate might also change during training . Acceleration methods such as Heavy-Ball Momentum ( referred to as Momentum from here on ) ( Polyak , 1964 ) were designed to enable faster convergence on poorly conditioned loss landscapes . Momentum works by taking an exponential moving average of previous gradients , ωi+1 = ωi − ∑i j=0m i−j dC dω ∣∣ w=wi , where m denotes the momentum coefficient . Gradients in high curvature directions , which often switch sign between updates , partially cancel out . This enables Momentum to take larger steps in low curvature directions while remaining stable in high curvature directions . On quadratic losses for example , Momentum increases the critical learning rate : crit ≤ 2 ( 1 +m ) /λmax ( Goh , 2017 ) , and can minimize the training loss in fewer steps than full batch gradient descent . In practice we do not compute a full batch gradient , we estimate the gradient over a minibatch ( Robbins & Monro , 1951 ; Bottou , 2010 ) . This introduces noise into our parameter updates , and this noise will play a crucial role in the training dynamics in some cases . However when the batch size is large , and the number of training epochs is finite , the noise in the parameter updates is low , and so typically most of training is governed by the curvature of the loss landscape ( similar to full batch gradient descent ) . We call this large batch training regime curvature dominated . When the batch size is in the curvature dominated regime , we expect the optimal initial learning rate to be determined by the critical learning rate crit , and for SGD with Momentum to outperform SGD without Momentum . On the other hand , when the batch size is small , typically most of the training process is governed by the noise in the parameter updates , and we call this small batch training regime noise dominated . In order to build a model of the training dynamics in the noise dominated regime , we must make some assumptions . Following previous work ( Mandt et al. , 2017 ; Li et al. , 2017 ; Smith & Le , 2017 ; Jastrzębski et al. , 2017 ; Simsekli et al. , 2019 ) , we assume the gradients of individual examples are independent samples from an underlying distribution , and that this distribution is not heavy tailed . When the training set size N B , the batch size B 1 , and B N , we can apply the central limit theorem to model the noise in a gradient update by a Gaussian noise source δ , whose covariance is inversely proportional to the batch size . We therefore approximate the SGD update by , ( ωi+1 − ωi ) ≈ − ( dC dω ∣∣∣ ω=ωi + δi√ B ) , ( 1 ) To interpret this update , we introduce the temperature T = /B to obtain , ( ωi+1 − ωi ) ≈ − dC dω ∣∣∣ ω=ωi + √ Tδi . ( 2 ) Equation 2 describes the discretization of a stochastic differential equation with step size and temperature T ( Gardiner et al. , 1985 ) , and we expect the dynamics of SGD to follow this underlying stochastic differential equation , so long as the learning rate crit and the assumptions above are satisfied . When equation 2 holds and the learning rate crit , any two training runs with the same temperature and the same epoch budget should achieve similar performance on both the training set and the test set . Consequently , we usually expect the learning rate to scale linearly with the batch size in the noise dominated regime , and this was observed in many empirical studies ( Krizhevsky , 2014 ; Goyal et al. , 2017 ; Smith et al. , 2017 ; Jastrzębski et al. , 2017 ; McCandlish et al. , 2018 ) . For completeness , we derive this linear scaling rule in appendix B , and we demonstrate that the linear scaling rule can be derived without assuming that the batch size B 1 . However the remaining assumptions above are required , and they are not always satisfied . This linear scaling rule is therefore not valid in all cases ( Shallue et al. , 2018 ; Simsekli et al. , 2019 ) . Empirically however , we have found that all batch sizes in the noise dominated regime achieve similar test accuracies and training losses under a constant epoch budget , even when the optimal learning rate does not obey linear scaling . Further observations on the two regimes : Many previous works have established that SGD with and without Momentum are equivalent in the small learning rate limit when m is fixed ( Orr & Leen , 1994 ; Qian , 1999 ; Yuan et al. , 2016 ) . In this limit , the speed of convergence of SGD with Momentum is governed by the effective learning rate eff = / ( 1 − m ) , and the temperature T = eff/B ( Mandt et al. , 2017 ; Smith & Le , 2017 ) . We therefore expect SGD with and without Momentum to achieve the same final training losses and test accuracies in the noise dominated regime ( where the optimal learning rate is smaller than crit ) . Supporting this claim , Shallue et al . ( 2018 ) confirmed in a recent empirical study that SGD with and without Momentum achieve similar test accuracies in the small batch limit , while SGD with Momentum outperforms vanilla SGD in the large batch limit . In recent years , a number of authors have exploited large batch training and parallel computation to minimize the wallclock time of training deep networks ( De et al. , 2016 ; Smith et al. , 2017 ; Akiba et al. , 2017 ; You et al. , 2018 ; McCandlish et al. , 2018 ; Shallue et al. , 2018 ) . Goyal et al . ( 2017 ) succeeded in training ResNet-50 to over 76 % accuracy in under one hour , and since then this has fallen to just a few minutes ( Mikami et al. , 2018 ; Ying et al. , 2018 ) . Goyal et al . ( 2017 ) also introduced learning rate warmup , and found that it enabled stable training with larger batch sizes ; a result which we confirm in this work . This procedure has a straightforward interpretation within the two regimes . If the critical learning rate increases early in training , then learning rate warmup will increase the largest stable learning rate , which in turn enables efficient training with larger minibatches . Learning rate schedules : In the noise dominated regime , the learning rate increases as the batch size rises , and therefore as we increase the batch size we will eventually invalidate the assumption crit and enter the curvature dominated regime . There may be a transition phase between the two regimes ( Liu & Belkin , 2018 ) , although our experiments suggest this transition can be surprisingly sharp in practice . We note that , with the optimal learning rate schedule , many batch sizes might exhibit both the curvature dominated regime ( typically early in training ) and the noise dominated regime ( typically towards the end of training ) . For example , on simple quadratic loss landscapes at any batch size , the optimal learning rate schedule for minimizing the training loss begins with an initial learning rate close to crit , followed by learning rate decay ( Zhang et al. , 2019 ) . However in practice it is not possible to identify the optimal learning rate schedule within a realistic computation budget . Practitioners prefer simple learning rate schedules , often parameterized by an initial learning rate and a few sharp drops ( He et al. , 2016 ) . These schedules are easy to tune and are also thought to generalize well . For these popular schedules , the optimal initial learning rate would be determined by whether most of the training process is noise dominated or curvature dominated . Furthermore , Smith & Le ( 2017 ) suggested that there may be an optimal temperature at early times that promotes generalization . If this speculation is correct , then the optimal learning rate schedule to maximize the test acccuracy would select the noise dominated regime throughout the whole of training when the batch size is small . Of course , if the loss is extremely poorly conditioned , the critical learning rate may already be optimal when the batch size is 1 , although we have never seen this in practice2 . Implicit regularization : In the noise dominated regime , the temperature T defines the influence of gradient noise on the dynamics . In the noisy quadratic model , our goal during training is to minimize the effect this noise has on the final parameters ( Polyak & Juditsky , 1992 ; Zhang et al. , 2019 ) . However recently many authors have argued that minibatch noise can be beneficial , helping us to select final parameters which perform well on the test set ( Keskar et al. , 2016 ; Jastrzębski et al. , 2017 ; Chaudhari & Soatto , 2018 ) . In this alternative perspective , there may be an “ optimal temperature ” early in training , which drives the parameters towards regions that generalize well ( Smith et al. , 2017 ; Park et al. , 2019 ) . The noisy quadratic model predicts that increasing the batch size may increase the final training loss under a constant epoch budget , but that it should not increase the final training loss under a constant step budget . Note that this model does not make explicit predictions about the test accuracy . We note that Jin et al . ( 2019 ) argued noise enables SGD to escape saddle points in non-convex landscapes , which could enable small batch sizes to achieve lower training losses under both constant epoch and constant step budgets . The optimal temperature perspective predicts that increasing the batch size may increase the final training loss under a constant epoch budget , but it does not predict what will happen to the training loss under a constant step budget . Crucially , it predicts that beyond some threshold batch size ( which may be very large ) , increasing the batch size will decrease the test accuracy under both constant epoch and constant step budgets3 . 2We assume the epoch budget is large enough to fit the training data . Training may occur in the curvature dominated regime at all batch sizes when the epoch budget is very small ( McCandlish et al. , 2018 ) . 3We emphasize that both perspectives agree that learning rate decay should be used at the end of training .

Summary: This paper studies the properties of SGD as a function of batch size and learning rate. Authors argue that SGD has two regimes:  a noise dominated regime (small batch size) and curvature dominated regime (large batch size). Authors conduct through numerical experiments highlighting how learning rate changes as a function of batch size (initially linear growth and then saturates). The critical contribution of this work appears to be the observation that large batch size can be worse than small under same number of steps demonstrating implicit regularization of small batch size.