Abstract:
Stochastic gradient descent (SGD) with momentum is widely used for training modern deep learning architectures. While it is well understood that using momentum can lead to faster convergence rate in various settings, it has also been observed that momentum yields higher generalization. Prior work argue that momentum stabilizes the SGD noise during training and this leads to higher generalization. In this paper, we take the opposite view to this result and first empirically show that gradient descent with momentum (GD+M) significantly improves generalization comparing to gradient descent (GD) in many deep learning tasks. From this observation, we formally study how momentum improves generalization in deep learning. We devise a binary classification setting where a two-layer (over-parameterized) convolutional neural network trained with GD+M provably generalizes better than the same network trained with vanilla GD, when both algorithms start from the same random initialization. The key insight in our analysis is that momentum is beneficial in datasets where the examples share some features but differ in their margin. Contrary to the GD model that memorizes the small margin data, GD+M can still learn the features in these data thanks to its historical gradients. We also empirically verify this learning process of momentum in real-world settings.

heading: 1 INTRODUCTION
text: 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.

heading: 2 SETTING AND ALGORITHMS
text: In this section, we first introduce a formal definition of the data distribution equation (D) and the neural network model we use to learn it. We finally present the GD and GD+M algorithms. General notations. For a matrix W ∈ Rm×d, we denote by wr its r-th row. For a function f : Rm×d → R, we denote by ∇wrf(W ) the gradient of f with respect to wr and ∇f(W ) the gradient with respect to W. For an optimization algorithm updating a vector w, w(t) represents its iterate at time t. We use Id for the d × d identity matrix and 1m the all-ones vector of dimension m. Finally, we use the asymptotic complexity notations when defining the different constants in the paper. We use Õ, Θ̃, Ω̃ to hide logarithmic dependency on d. Data distribution. We define our data distribution D as follows.
Each sample from D consists in an input data X and a label y that are generated as: 1. The label y is uniformly sampled from {−1, 1}. 2. Each data-point X = (X[1], . . . , X[P ]) consists in P patches where each X[j] ∈ Rd. 3. Signal patch: for one patch P (X) ∈ [P ], we have X[P (X)] = cw∗, where c ∈ R, w∗ ∈ Rd and ‖w∗‖2 = 1. (D)
4. The distribution of c satisfies that
c = { αy with probability 1− µ βy with probability µ .
5. Noisy patches: for all the other patches j ∈ [P ]\{P (X)},X[j] ∼ N (0, (I−w∗w∗>)σ2Id).
We precise that we sample the noisy patches in the orthogonal complement of w∗ to have a simpler analysis. To present the simplest result, we assume that the values in equation (D) satisfy α = d0.49, β = 1
polylog(d) √ d α, σ = 1√ d and P ∈ [2,polylog(d)].
Using this model, we generate a training dataset Z = {(Xi, yi)}i∈[N ] where Xi = (Xi[j])j∈[P ]. We focus on the case where µ = 1/poly(d) and N = Θ ( 1 µ ) . We let Z to be partitioned in two sets Z1 and Z2 such that Z1 gathers the large margin data while Z2 the small margin ones. Lastly, we define µ̂ = |Z2|N the fraction of small margin data. Learner model. We use a two-layer convolutional neural network with cubic activation to learn the training dataset Z . This model is the simplest non-linear network since a quadratic activation
would only output positive labels and mismatch our labeling function. The first layer weights are W ∈ Rm×d and the second layer is fixed to 1m. Given a input data X , the output of the model is
fW (X) = m∑ r=1 P∑ j=1 〈wr, X[j]〉3. (CNN)
The number of neurons is set as m = polylog(d) to ensure that (CNN) is mildly over-parametrized. Training objective. We fit the training dataset Z using (CNN) and solve the logistic regression problem
min W∈Rm×d
1
N N∑ i=1 log (1 + exp (−yifW (Xi))) + λ 2 ‖W‖22 := L̂(W ). (P)
(P) sheds light on our choice of cubic activation in (CNN). Indeed, it is the smallest polynomial degree that makes the training objective (P) non-convex and compatible with our dataset. Linear or quadratic activations would respectively make the problem convex or all the labels positive. Here, we pick λ ∈ [ 0, 1poly(d)N ] .
Importance of non-convexity. When λ > 0, if the loss 1N ∑N i=1 log (1 + exp (−yifW (Xi))) is convex, then there is a unique global optimal solution, so the choice of optimization algorithm does not matter. In our case, due to the non-convexity of the training objective, GD + M converges to a different (approximate) global optimal comparing to GD, with better generalization properties.
Test error. We assess the quality of a predictor Ŵ using the classical 0-1 loss used in binary classification. Given a sample (X, y), the individual test (classification) error is defined as L (X, y) = 1{f
Ŵ (X)y < 0}. While L measures the error of fW on an individual data-point,
we are interested in the test error that measures the average loss over data points generated from (D) and defined as
L (f Ŵ ) := E(X,y)∼D[L (fŴ (X), y)]. (TE)
Algorithms. We solve the training problem equation (P) using GD and GD+M. GD is defined by
W (t+1) = W (t) − η∇L̂(W (t)), for t ≥ 0, (GD)
where η > 0 is the learning rate. On the other hand, GD+M is defined by the update rule{ g(t+1) = γg(t) + (1− γ)∇L̂(W (t)) W (t+1) = W (t) − ηg(t+1) , for t ≥ 0. (GD+M)
where γ ∈ (0, 1) is momentum factor. We now detail how to set parameters in (GD) and (GD+M). Parametrization 2.1. When running GD and GD+M on equation (P), the number of iterations is T ∈ [ poly(d)N/(η), dO(log d)/(η) ] . For both algorithms, the weights w(0)1 , . . . , w (0) m are initialized using independent samples from a normal distribution N (0, σ20Id) where σ20 = polylog(d)
d . The learning rate is set as:
1. GD: the learning rate may take any reasonable value η ∈ (0, Õ(1)].
2. GD+M: the learning rate is a large learning rate: η = Θ̃(1).1
Lastly, the momentum factor in GD+M is set to be γ = 1− polylog(d)d .
Our Parametrization 2.1 matches with the parameters used in practice as the weights are generally initialized from Gaussian with small variance and momentum is set close to 1 (Sutskever et al., 2013).

heading: 3 MAIN RESULTS
text: We now formally state our main theorems regarding the generalization of models trained using equation (GD) and equation (GD+M) on the training set Z generated by equation (D). As announced in the introduction, we show that the GD+M model incurs a generalization error that is dramatically smaller than the GD model. Before introducing the main result, we define some notations:
1This is consistent with the empirical observation that only momentum with large learning rate improves generalization (Sutskever et al., 2013)
Main objects. Let r ∈ [m], i ∈ [N ], j ∈ P\{P (Xi)}, γ ∈ (0, 1) and t ≥ 0. We are mainly interested in w(t)r , the r-th weight of the network,∇wr L̂(W (t)) the gradient of the training loss w.r.t. wr, g (t) r the momentum gradient defined by g (t+1) r = γg (t) r + (1− γ)∇wr L̂(W (t)). The analysis lies on the projection of these objects on the feature w∗ and on noisy patches Xi[j]. We introduce the following notations for the component of the learned weights along feature and noise directions:
– Projection on w∗: c(t)r = 〈w(t)r , w∗〉.
– Projection on Xi[j] : Ξ (t) i,j,r = 〈w (t) r , Xi[j]〉.
– Total noise: Ξ(t)i = ∑m r=1 ∑ j∈[P ]\{P (Xi)}〈w (t) r , Xi[j]〉3.
– Maximum signal: let rmax = argmaxr∈[m]c (t) r , c(t) = c (t) rmax
Theorem 3.1. Assume that we run GD on (P) for T iterations with parameters set as in Parametrization 2.1. With probability at least 1− o(1), the weights learned by GD
1. partially learn the feature: for all r ∈ [m], |c(T )r | ≤ Õ(1/α).
2. memorize from small margin data: for all i ∈ Z2, Ξ(t)i ≥ Ω̃(1). Consequently, the training error is smaller than O(µ/poly(d)) and the test error is at least Ω̃(µ).
Intuitively, the training process of the GD model is described as follows. Since the large margin data are dominant in Z , the gradient points mainly in the direction of the feature w∗. Therefore, GD eventually learns the feature in Z1 (Lemma 4.1) and the gradients from Z1 quickly become small. Afterwards, the gradient is dominated by the gradients from Z2 (Lemma 4.2). Because Z2 has small margin, the full gradient is now directed by the noisy patches. It implies that GD memorizes noise in Z2 (Lemma 4.4). Since these gradients also control the amount of remaining feature to be learned (Lemma 4.3), we conclude that the GD model partially learns the feature and introduces a huge noise component in the learned weights. We provide a proof sketch of Theorem 3.1 in Section 4. Theorem 3.2. Assume that we run GD+M on equation (P) for T iterations with parameters set as in Parametrization 2.1. With probability at least 1− o(1), the weights learned by GD+M
1. (at least for one of them) is highly correlated with the feature: c(T ) > Ω̃(1/β).
2. are barely correlated with noise: for all r ∈ [m], for all i ∈ [N ] and j ∈ [P ]. |Ξ(T )i,j,r| ≤ Õ(σ0). Consequently, the training loss and the test error are at most O(µ/poly(d)).
Intuitively, the GD+M model follows this training process. Similarly to GD, it first fits the Z1 (Lemma 5.1). Contrary to GD, the momentum gradient is still highly correlated with w∗ after this step (Lemma 5.2). Indeed, the key difference is that momentum accumulates historical gradients. Since these gradients were accumulated when learning large margin data, the direction of momentum gradient is highly biased towards w∗. Therefore, the GD+M model amplifies the feature from these historical gradients to learn the feature in small margin data (Lemma 5.3). Subsequently, the gradient becomes small (Lemma 5.4) and the weights are no longer updated. Therefore, the GD+M model manages to ignore the noisy patches (Lemma 5.5) and learns the feature from both Z1 and Z2. We provide a proof sketch of Theorem 3.2 in Section 5.
To state the proof, we further decompose the gradients along signal and noise directions.
– Projection on w∗: G (t)r = 〈∇wr L̂(Wt), w∗〉 and G (t) r = 〈g(t)r , w∗〉.
– Projection on Xi[j] : G (t) i,j,r = 〈∇wr L̂(W (t)), Xi[j]〉, G (t) i,j,r = 〈g (t) r , Xi[j]〉.
– Maximum signal: let rmax = argmaxr∈[m]c (t) r , c(t) = c (t) rmax and G(t) = G (t) rmax .
Signal and noise iterates. Our analysis is build upon a decomposition of the updates equation (GD) and equation (GD+M) on w∗ and Xi[j]. These decompositions are respectively defined as follows:
c(t+1)r = c (t) r − ηG (t)r (GD-S) Ξ (t+1) i,j,r = Ξ (t) i,j,r − ηG (t) i,j,r (GD-N)
{ G(t+1)r = γG(t)r + (1− γ)G (t)r c (t+1) r = c (t) r − ηG(t+1)r (GDM-S)
{ G
(t+1) i,j,r = γG (t) i,j,r + (1− γ)G (t) i,j,r
Ξ (t+1) i,j,r = Ξ (t) i,j,r −G (t+1) i,j,r
(GDM-N) We detail how to use these dynamics to analyze GD+M and GD in Section 4 and Section 5. Our analysis heavily depends on the gradients of the training loss which involve sigmoid(x) = (1 + e−x)−1 . We define the derivative of a data-point i as `(t)i = sigmoid(−yifW (t)(Xi)), the derivatives ν
(t) k = 1 N ∑ i∈Zk ` (t) i for k ∈ {1, 2} and the full derivative ν(t) = ν (t) 1 + ν (t) 2 .

heading: 4 ANALYSIS OF GD
text: In this section, we provide a proof sketch for Theorem 3.1 that reflects the behavior of GD with λ = 0. A more detailed proof (extending to λ > 0) can be found in the Appendix.
Step 1: Learning Z1. At the beginning of the learning process, the gradient is mostly dominated by the gradients coming from the Z1 samples. Since these data have large margin, the gradient is thus highly correlated with w∗ and c(t)r increases as shown in the following Lemma. Lemma 4.1. For all r ∈ [m] and t ≥ 0, equation (GD-S) is simplified as:
c(t+1)r ≥ c(t)r + Θ̃(η)α3(c(t)r )2 · sigmoid(− ∑t s=1 α 3(c (t) s )3).
Consequently, after T0 = Θ̃ (
1 ηα3σ0
) iterations, for all t ∈ [T0, T ], we have c(t) ≥ Ω̃(1/α).
Intuitively, the increment in the update in Lemma 4.1 is non-zero when the sigmoid is not too small which is equivalent to c(t) ≤ Õ(1/α). Therefore, c(t) keeps increasing until reaching this threshold. After this step, the Z1 data have small gradient and therefore, GD has learned these data.
Lemma 4.2. Let T0 = Θ̃ (
1 ηα3σ0
) . After t ∈ [T0, T ] iterations, the Z1 derivative is bounded as
ν (t) 1 ≤ Õ
( 1
η(t−T0+1)α
) + Õ ( β3
α
) ν
(t) 2 . The full derivative is ν
(t) ≤ Õ (
1 η(t−T0+1)α +
( 1 + β 3
α
) ν
(t) 2
) .
By our choice of parameter, Lemma 4.2 indicates that the full gradient is dominated by the gradients from Z2 data after T0 = Ω̃ ( 1 µ̂ηα ) . Consequently, ν(t)2 also rules the amount of feature learnt by GD.
Lemma 4.3. Let T0 = Θ̃ (
1 ηα3σ0
) . For t ∈ [T0, T ], equation (GD-S) becomes c(t+1) ≤ Õ(1/α) +
Õ(ηβ3/α) ∑t τ=T0 ν (τ) 2 .
Lemma 4.3 implies that quantifying the decrease rate of ν(t)2 provides an estimate on the quantity of feature learnt by the model. We remark that ν(t)2 = sigmoid(β 3 ∑m s=1(c (t) s )3 + Ξ (t) i ) for some i ∈ Z2. We thus need to determine whether the feature or the noise terms dominates in the sigmoid.
Step 2: Memorizing Z2. We now show that the total correlation between the weights and the noise in Z2 data increases until being large.
Lemma 4.4. Let t ≥ 0 and i ∈ Z2. Assume that Ξ(t)i ≤ Õ(1). Then, equation (GD-N) can be simplified as:
yiΞ (t+1) i,j,r ≥ yiΞ (0) i,j,r +
Θ̃(ησ2d)
N
t∑ τ=0 (Ξ (τ) i,j,r) 2 − Õ
( Pσ2 √ d
α
) .
Let T1 = Õ (
N σ0σ √ dσ2d
) . Therefore, Ξ(t)i ≥ Ω̃(1), for t ∈ [T1, T ] and thus GD memorizes.
By Lemma 4.4, in the gradient of Z2 data, the noise term dominates the feature term (which scales as Õ(β3)). Consequently, the algorithm memorizes the Z2 data which implies a fast decay of ν(t)2 .
Lemma 4.5. Let T1 = Õ (
N σ0σ √ dσ2d
) . For t ∈ [T1, T ], we have ∑t τ=0 ν (τ) 2 ≤ Õ ( 1 ησ0 ) .
Combining Lemma 4.5 and Lemma 4.3, we prove that GD partially learns the feature. Lemma 4.6. For t ≤ T , the signal component satisfies c(t) ≤ Õ(1/α).
Lemma 4.4 and Lemma 4.6 respectively yield the first two items in Theorem 3.1. Bounds on the training and test errors are respectively obtained by plugging these results in (P) and (TE).

heading: 5 ANALYSIS OF GD+M
text: In this section, we provide a proof sketch for Theorem 3.2 that reflects the behavior of GD+M with λ = 0. A more detailed proof (also extending to λ > 0) can be found in the Appendix.
Step 1: Learning Z1. Similarly to GD, by our initialization choice, the early gradients and so, momentum gradients are large. They are also spanned by the feature w∗ and therefore, the GD+M model also increases its correlation with w∗. Lemma 5.1. For all r ∈ [m] and t ≥ 0, as long as c(t) ≤ Õ(1/α), the momentum update equation (GDM-S) is simplified as:
−G(t+1)r = −γG(t)r + (1− γ)Θ(α3)(c(t)r )2 Consequently, after T0 = Θ̃ (
1 σ0α2 + 11−γ
) iterations, for all t ∈ [T0, T ], we have c(t) ≥ Ω̃(1/α).
Step 2: LearningZ2. Contrary to GD, GD+M has a large momentum that contains w∗ after Step 1. Lemma 5.2. Let T0 = Θ̃ ( 1
σ0α3 + 11−γ
) . For t ∈ [T0, T ], we have G(t) ≥ Ω̃( √ 1− γ/α).
Lemma 5.2 hints an important distinction between GD and GD+M: while the current gradient along w∗ is small at time T0, the momentum gradient stores historical gradients that are spanned by w∗. It amplifies the feature present in previous gradients to learn the feature from small margin data.
Lemma 5.3. Let T0 = Θ̃ (
1 σ0α3 + 11−γ
) . After T1 = T0 + Θ̃ ( 1
1−γ
) iterations, for t ∈ [T1, T ], we
have c(t) ≥ Ω̃ (
1√ 1−γα
) . Our choice of parameter in Section 2, this implies c(t) ≥ Ω̃(1/β).
Lemma 5.3 states that at least one of the weights that is highly correlated with the feature compared to GD where c(t) = Õ(1). This result implies that ν(t) converges fast.
Lemma 5.4. Let T0 = Θ̃ (
1 ησ0α3 + 11−γ
) . After T1 = T0 + Θ̃ ( 1
1−γ
) iterations, for t ∈ [T1, T ],
ν(t) ≤ Õ (
1 η(t−T1+1)β
) .
With this fast convergence, Lemma 5.4 implies that the correlation of the weights with the noisy patches does not have enough time to increase and thus, remains small. Lemma 5.5. Let i ∈ [N ], j ∈ [P ]\{P (Xi)} and r ∈ [m]. For t ≥ 0, equation (GDM-N) can be rewritten as |G(t+1)i,j,r | ≤ γ|G (t) i,j,r| + (1 − γ)Õ(σ20σ4d2)ν(t). As a consequence, after t ∈ [T1, T ] iterations, we thus have |Ξ(t)i,j,r| ≤ Õ(σ0σ √ d).
Lemma 5.3 and Lemma 5.5 respectively yield the two first items in Theorem 3.2.

heading: 6 DISCUSSION
text: Our work is a first step towards understanding the algorithmic regularization of momentum and leaves room for improvements. We constructed a data distribution where historical feature amplification may explain the generalization improvement of momentum. However, it would be interesting to understand whether this phenomenon is the only reason or whether there are other mechanisms explaining momentum’s benefits.An interesting setting for this question is NLP where momentum is used to train large models as BERT (Devlin et al., 2018). Lastly, our analysis is in the batch setting to isolate the generalization induced by momentum. It would be interesting to understand how the stochastic noise and the momentum together contribute to the generalization of a neural network.

Review summaries:
 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.

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.

Through some experiments, this paper claims that momentum does not always lead to a higher generalization in deep learning and such benefit seems to heavily depend on both the structure of the data and the learning problem. Then, the authors considered a certain data structure (large and small margin data) and learning problem (binary classification problem with 2-layer CNN model and logistic loss). Under this special case, the authors shown that both GD and GDM reach zero training error and perfectly classify large margin data, but GD fails to classify the small margin data while GDM can still perfect classify small margin data due to the historical gradients caused by the momentum.