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:  This paper investigates the generalization benefits of using momentum when training neural networks with gradient descent. In contrast with existing literature, which studies momentum mostly empirically and in the stochastic setting, the authors develop a theoretical explanation for why generalization improves when optimization is performed using full-batch gradient descent with momentum (GD+M) than without it (GD). Their analysis focuses on a specific structure of the data and the learning problem, and the authors argue that similar assumptions apply to the real-world datasets used in practice. The authors prove two main results: one about GD learning large-margin data and overfitting to small-margin data and the other about GD+M successfully learning all the data thanks to historical gradients. Interestingly, the authors provide a counterexample where momentum does not aid generalization (or even worsen it), which offers good food for thought.