Towards understanding how momentum improves generalization in deep learning

1 INTRODUCTION . It is commonly accepted that adding momentum to an optimization algorithm is required to optimally train a large-scale deep network . Most of the modern architectures maintain during the training process a heavy momentum close to 1 ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; He et al. , 2016 ; Zagoruyko & Komodakis , 2016 ) . Indeed , it has been empirically observed that architectures trained with momentum outperform those which are trained without ( Sutskever et al. , 2013 ) . Several papers have attempted to explain this phenomenon . From the optimization perspective , Defazio ( 2020 ) assert that momentum yields faster convergence of the training loss since , at the early stages , it cancels out the noise from the stochastic gradients . On the other hand , Leclerc & Madry ( 2020 ) empirically observes that momentum yields faster training convergence only when the learning rate is small . While these works shed light on how momentum acts on neural network training , they fail to capture the generalization improvement induced by momentum ( Sutskever et al. , 2013 ) . Besides , the noise reduction property of momentum advocated by Defazio ( 2020 ) seems to even contradict the observation that , in deep learning , having a large noise in the training improves generalization ( Li et al. , 2019 ; HaoChen et al. , 2020 ) . To the best of our knowledge , there is no existing work which theoretically explains how momentum improves generalization in deep learning . Therefore , this paper aims to close this gap and addresses the following question : Is the higher generalization induced by momentum tied to the stochastic noise of the gradient ? If not , what is the underlying mechanism of momentum improving generalization in deep learning ? In this paper , we empirically verify that the generalization improvement induced by momentum is not tied to the stochasticity of the gradient . Indeed , as reported in Figure 1 , momentum improves generalization more significantly for full batch GD than for SGD in CIFAR object recognition tasks . Motivated by this empirical observation and the fact that the stochastic noise influences generalization , we theoretically study how gradient descent with momentum ( GD+M ) can generalize better than vanilla gradient descent ( GD ) . We therefore only focus on the contribution of momentum of the true gradient on generalization . The question we address concerns algorithmic regularization which characterizes the generalization of an optimization algorithm when multiple global solutions exist in over-parameterized deep learning model Soudry et al . ( 2018 ) ; Lyu & Li ( 2019 ) ; Ji & Telgarsky ( 2019 ) ; Chizat & Bach ( 2020 ) ; Gunasekar et al . ( 2018 ) ; Arora et al . ( 2019 ) . This regularization arises in deep learning mainly due to the nonconvexity of the objective function . Indeed , this latter can create multiple global minima scattered in the space that vastly differ in terms of generalization . Algorithmic regularization is induced by and depends on many factors such as learning rate and batch size ( Goyal et al. , 2017 ; Hoffer et al. , 2017 ; Keskar et al. , 2016 ; Smith et al. , 2018 ) , initialization Allen-Zhu & Li ( 2020 ) , adaptive step-size ( Kingma & Ba , 2014 ; Neyshabur et al. , 2015 ; Wilson et al. , 2017 ) , batch normalization ( Arora et al. , 2018 ; Hoffer et al. , 2019 ; Ioffe & Szegedy , 2015 ) and dropout ( Srivastava et al. , 2014 ; Wei et al. , 2020 ) . However , none of these works theoretically analyzes the regularization induced by momentum . We therefore start our investigation by raising the following question : Does momentum unconditionally improve generalization in deep learning ? This question could be positively answered given the success of momentum for learning distinct architectures such as ResNets ( He et al. , 2016 ) or BERT ( Devlin et al. , 2018 ) . However , we here empirically give a negative answer through the following synthetic example in deep learning . We consider a binary classification problem where data-points are generated from a standard normal distribution and labels are outputs of teacher networks . Starting from the same initialization , we train different over-parametrized student networks using GD and GD+M . Based on Table 1 , whether the target function is simple ( linear ) or complex ( neural network ) , momentum does not improve generalization even when using a non-linear neural network as learner . The same observation holds for SGD/SGD+M as shown in the Appendix . Therefore , momentum does not always lead to a higher generalization in deep learning . Instead , such benefit seems to heavily depend on both the structure of the data and the learning problem . On which data set does momentum help generalization ? In this paper , in order to determine the underlying mechanism produced by momentum to improve generalization , we design a binary classification problem with a simple data structure where training a two-layer ( overparameterized ) convolutional network with momentum provably improves generalization in deep learning . It is built upon a data distribution that relies on the concepts of feature and margin . Informally , each example in this distribution is a 1D image having P patches . One of the patches ( the signal patch ) contains a feature we want to learn and all the others are Gaussian random noise with small variance . Mathematically , one can think of a feature as a vectorw∗ ∈ Rd.We assume that our training examples are divided into large margin data where the signal is αw∗ with α constant and small margin data where the signal is βw∗ with β 1 . Intuitively , the second type of data is inherently noisier as the margin is small and therefore , a classifier would struggle more to generalize on this type of data . We underline that all the examples share the same feature but differ in the intensity of the signal . We consider a training dataset of size N with the following split for µ̂ 1 : ( 1− µ̂ ) N datapoints are with large margin , µ̂N datapoints are with small margin data . ( D ) Figure 2 sketches equation ( D ) in a 2D setting . We emphasize that datasets having similar features and different margins are common in the real-world . Examples include object-recognition datasets such as CIFAR ( Krizhevsky et al. , 2009 ) or Imagenet ( Deng et al. , 2009 ) ( for example , the “ wheel feature ” of a car can be strong or weak depending on the orientation of the car ) . More specifically , we believe that the dataset ( D ) can be viewed as a simplified model of these object-recognition datasets . In this context , the following informal theorems characterize the generalization of the GD and GD+M models . They dramatically simplify Theorem 3.1 and Theorem 3.2 but highlight the intuitions behind our results . Theorem 1.1 ( Informal , GD+M ) . There exists a dataset of the form ( D ) with size N such that a two-layer ( over-parameterized ) convolutional network trained with GD+M : 1. initially only learns large margin data from the ( 1− µ̂ ) N examples . 2. has large historical gradients that contain the feature w∗ present in small margin data . 3. keeps learning the feature in the small margin data using its momentum historical gradients . The model thus reaches zero training error and perfectly classify large and small margin data at test . Theorem 1.2 ( Informal , GD ) . There exists a dataset of the form ( D ) with sizeN such that a two-layer ( over-parameterized ) convolutional network trained with GD : 1. initially only learns large margin data from the from the ( 1− µ̂ ) N examples . 2. has small gradient after learning these data . 3. memorizes the remaining small margin data from the µ̂N examples using the noises . The model thus reaches zero training and manages to classify the large margin data at test . However , it fails to classify the small margin data because of the memorization step during training . Why does GD+M generalize better than GD ? Since the large margin data are dominant , GD focus in priority on these examples to decrease its training loss . However , after fitting this data , it significantly lowers its gradient . The gradient is thus not large enough for learning the small margin data . Similarly , GD+M fits the large margin data and subsequently gets a small gradient . However , contrary to GD , GD+M has large historical gradients in his momentum gradient . These gradients helped to learn the feature in the large margin data . They also help to learn small margin data since all the examples share the same feature . GD+M therefore uses his momentum to learn the small margin data . We name this process historical feature amplification and believe that it is key to understand why momentum improves generalization . Empirical justification . We also provide an empirical justification that such phenomenon does happen in a real-world setting as reported in Figure 4 . In this experiment , we create small margin data in the CIFAR-10 dataset by respectively lowering the resolution of 10 % of the training and test images , adding Gaussian noise of variance 0.005 and randomly shuffling the RGB channels . Figure 4 shows that even though both algortihms reach zero training error and 100 % training accuracy , GD+M gets higher generalization than GD on this decimated dataset . Above all , at test , GD+M performs as well on small and large margin data while GD does relatively worse on small margin data.Indeed , the relative accuracy drop for GD+M is 80.36/83.32 = 0.97 while for GD is equal to 65.14/73.69 = 0.88 . Our paper is organized as follows . In Section 2 , we formally define the data distribution equation ( D ) , the model and algorithms we use to learn it . Lastly , Section 3 presents our main theorems and provide a proof sketch in Section 4 and Section 5 . Additional experiments can be found in the Appendix . MORE RELATED WORK Momentum in convex setting . GD+M ( a.k.a . heavy ball or Polyak momentum ) consists in calculating the exponentially weighted average of the past gradients and using it to update the weights . For convex functions near a strict twice-differentiable minimum , GD+M is optimal regarding local convergence rate Polyak ( 1963 ; 1964 ) ; Nemirovskij & Yudin ( 1983 ) ; Nesterov ( 2003 ) . However , it may fail to converge globally for general strongly convex twice-differentiable functions Lessard et al . ( 2015 ) and is no longer optimal for the class of smooth convex functions . In the stochastic setting , GD+M is more sensitive to noise in the gradients ; that is , to preserve their improved convergence rates , significantly less noise is required d ’ Aspremont ( 2008 ) ; Schmidt et al . ( 2011 ) ; Devolder et al . ( 2014 ) ; Kidambi et al . ( 2018 ) . Finally , other momentum methods are extensively used for convex functions such as Nesterov ’ s accelerated gradient Nesterov ( 1983 ) . Our paper focuses on the use of GD+M and contrary to the aforementioned papers , our setting is non-convex and we mainly focus on the generalization of the model learned by GD and GD+M when both methods converge to global optimal . We underline that contrary to the non-convex world , generalization is typically disentangled with optimization for ( strictly ) convex functions . Non-convex optimization with momentum . A long line of work consists in understanding the convergence speed of momentum methods when optimizing non-convex functions . Mai & Johansson ( 2020 ) ; Liu et al . ( 2020 ) ; Cutkosky & Mehta ( 2020 ) ; Defazio ( 2020 ) show that SGD+M reaches a stationary point as fast as SGD under diverse assumptions . Besides , Leclerc & Madry ( 2020 ) empirically shows that momentum accelerates neural network training for small learning rates and slows it down otherwise . Our paper differs from these works as we work in the batch setting and theoretically investigate the generalization benefits brought by momentum ( and not the training ones ) . Generalization with momentum . Momentum-based methods such as SGD+M , RMSProp ( Tieleman & Hinton , 2012 ) and Adam ( Kingma & Ba , 2014 ) are standard in deep learning training since the seminal work of Sutskever et al . ( 2013 ) . Although its well accepted that Momentum improve generalization in deep learning , only a few works formally investigate the role of momentum in generalization . Leclerc & Madry ( 2020 ) empirically reports that momentum yields higher generalization when using a large learning rate . However , they assert that this benefit can be obtained by applying an even larger learning rate on vanilla SGD . We suspect that this observation is due to batch normalization ( BN ) which is known to dramatically bias the algorithm ’ s generalization ( Lyu & Li , 2019 ) . In Appendix , we report that BN reduces the generalization gain of momentum comparing to without BN . To our knowledge , our work is first that theoretically investigate the generalization of momentum in deep learning .

Summary: The authors provide a new perspective on the why momentum is useful for generalization in neural networks.  They provide motivating intuition around why momentum is not always useful (including a great toy example), as well as empirical experiments on CIFAR-10.  Finally, the authors prove a variety of theorems on a synthetic problem where momentum provably results in greater generalization.