T OWARDS UNDERSTANDING HOW MOMENTUM IMPROVES GENERALIZATION IN DEEP LEARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst empirically show that gradient descent with momentum (GD+M) signiﬁcantly 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 classiﬁcation 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 beneﬁcial 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.
I NTRODUCTION
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 signiﬁcantly for full batch GD than for SGD in CIFAR object recognition tasks.
Motivated by this empirical observation and the fact that the stochastic noise inﬂuences 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
[Figure]
model Soudry et al.
; Lyu & Li (2019); Ji & Telgarsky (2019); Chizat & Bach (2020); Gunasekar et al.
Arora et al.
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 classiﬁcation 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 beneﬁt 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 classiﬁcation 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 vector w ∗ ∈ R d .
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 β (cid:28) 1 .
Intuitively, the second type of data is inherently noisier as the margin is small and therefore, a classiﬁer would struggle more to generalize on this type of data.
We
[Table]
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 ˆ µ (cid:28) 1 :
[Equation]
) N datapoints are with large margin , ˆ µN datapoints are with small margin data .
[Figure]
, the following the GD GD
dramatically Theorem Theorem highlight 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:
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.
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 size N such that a two-layer (over-parameterized) convolutional network trained with GD:
initially only learns large margin data from the from the (1 − ˆ µ )
N examples.
2. has small gradient after learning these data.
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 ﬁtting this data, it signiﬁcantly lowers its gradient.
The gradient is thus not large enough for learning the small margin data.
Similarly, GD+M ﬁts the large margin data and subsequently gets a small gradient. However, [Figure]
[Figure] 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 ampliﬁcation and believe that it is key to understand why momentum improves generalization [Figure]
Empirical justiﬁcation.
We also provide an empirical justiﬁcation 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 shufﬂing 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 .
32 = 0 .
97 while for GD is equal to 65 .
69 = 0 . 88 .
Our paper is organized as follows.
In Section 2, we formally deﬁne 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.
M ORE RELATED W ORK
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.
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, signiﬁcantly less noise is required d’Aspremont (2008);
Schmidt et al.
(2011); Devolder et al.
(2014); Kidambi et al.
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 beneﬁts 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.
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 beneﬁt 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 ﬁrst that theoretically investigate the generalization of momentum in deep learning.
S ETTING AND ALGORITHMS
In this section, we ﬁrst introduce a formal deﬁnition of the data distribution equation (D) and the neural network model we use to learn it.
We ﬁnally present the GD and GD+M algorithms.
General notations.
For a matrix W ∈ R m × d , we denote by w r its r -th row.
For a function f : R m × d → R , we denote by ∇ w r f ( W ) the gradient of f with respect to w r 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 I d for the d × d identity matrix and 1 m the all-ones vector of dimension m. Finally, we use the asymptotic complexity notations when deﬁning the different constants in the paper.
We use ˜ O, ˜Θ , ˜Ω to hide logarithmic dependency on d .
Data distribution.
We deﬁne our data distribution D as follows.
Each sample from D consists in an input data X and a label y that are generated as:
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 ] ∈ R d . 3.
Signal patch: for one patch P ( X ) ∈
[ P ] , we have X [ P ( X )]
= cw ∗ , where c ∈ R , w ∗ ∈ R d and (cid:107) w ∗ (cid:107) 2 = 1 . [Equation] D
[Equation]
The distribution of c satisﬁes that
[Equation]
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 ∗(cid:62) )
σ 2 I d ) .
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 α = d 0 . 49 , β = 1 polylog(d) √ d α , σ = 1 √ d and P ∈ [2 , polylog(d)] .
Using this model, we generate a training dataset Z = { ( X i , y i ) } i ∈
[ N ] where X i =
( X i [ j ]) j ∈
[ P ] .
We focus on the case where µ = 1 / poly( d ) and N = Θ (cid:16) 1 µ (cid:17) .
We let Z to be partitioned in two sets Zand Z 2 such that Z 1 gathers the large margin data while Z 2 the small margin ones.
Lastly, we deﬁne ˆ µ = |Z 2 | 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 ﬁrst layer weights are
W ∈ R m × d and the second layer is ﬁxed to 1 m .
Given a input data X , the output of the model is
[Equation]
The number of neurons is set as m = polylog( d ) to ensure that (CNN) is mildly over-parametrized.
Training objective.
We ﬁt the training dataset Z using (CNN) and solve the logistic regression problem
[Equation]
(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 λ ∈ (cid:104) 0 ,
[Equation] (cid:105) . Importance of non-convexity. When λ > 0 , if the loss 1 N (cid:80) Ni =1 log (1 + exp ( − y i f W ( X i ))) 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 (cid:99) W using the classical 0-1 loss used in binary classiﬁcation.
Given a sample ( X, y ) , the individual test (classiﬁcation) error is deﬁned as L ( X, y )
=
1 { f (cid:99) W
( X ) y < 0 } .
While L measures the error of f W on an individual data-point, we are interested in the test error that measures the average loss over data points generated from ( D ) and deﬁned as
[Equation]
Algorithms.
We solve the training problem equation (P) using GD and GD+M. GD is deﬁned by
[Equation]
where η > 0 is the learning rate.
On the other hand, GD+M is deﬁned by the update rule
[Equation]
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 ∈ (cid:2) poly(d) N/ ( η ) , d O (log d ) / ( η ) (cid:3) .
For both algorithms, the weights w (0)1 , . . .
, w (0) m are initialized using independent samples from a normal distribution N (0 , σ 20 I d ) where σ 20 = polylog( d ) d .
The learning rate is set as:
GD: the learning rate may take any reasonable value η ∈ (0 , ˜ O
GD+M: the learning rate is a large learning rate: η = ˜Θ(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).
M AIN RESULTS
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 deﬁne some notations:
Main objects.
Let r ∈
[ m ] , i ∈
[ N ] , j ∈ P \{ P ( X i ) } , γ ∈ (0 , 1) and t ≥ 0 .
We are mainly interested in w ( t ) r , the r -th weight of the network, ∇ w r (cid:98) L ( W ( t ) )
the gradient of the training loss w.r.t.
w r , g ( t ) r the momentum gradient deﬁned by g ( t +1)
r
[Equation]
L
[Equation]
W ( t ) ) .
The analysis lies on the projection of these objects on the feature w ∗ and on noisy patches
X i
[ j ] .
We introduce the following notations for the component of the learned weights along feature and noise directions:
– Projection on w ∗ : c ( t ) r = (cid:104) w ( t ) r , w ∗ (cid:105) .
– Projection on X i [ j ] : Ξ ( t ) i,j,r = (cid:104) w ( t ) r , X i
[ j ] (cid:105) .
– Total noise: Ξ ( t ) i = (cid:80)
mr =1 (cid:80) j ∈
[ P ] \{ P ( X i ) } (cid:104) w ( t ) r , X i
[ j ] (cid:105) 3 .
– Maximum signal: let r max = argmax r ∈ [ m ] c ( t ) r , c ( t ) =
c ( t ) r max
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 /α ) .
memorize from small margin data: for all i ∈
Z 2 , Ξ ( 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 Z 1 (Lemma 4.1) and the gradients from Z 1 quickly become small.
Afterwards, the gradient is dominated by the gradients from Z 2 (Lemma 4.2).
Because Z 2 has small margin, the full gradient is now directed by the noisy patches.
It implies that GD memorizes noise in
Z 2 (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 ﬁrst ﬁts the Z(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 ampliﬁes 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 Z 1 and Z 2 .
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 = (cid:104)∇ w r (cid:98) L ( W t ) , w ∗ (cid:105) and G ( t ) r = (cid:104) g ( t ) r , w ∗ (cid:105) .
– Projection on X i [ j ] : G ( t ) i,j,r = (cid:104)∇ w r (cid:98) L ( W ( t ) ) , X i [ j ] (cid:105) , G ( t ) i,j,r = (cid:104) g ( t ) r , X i [ j ] (cid:105) .
– Maximum signal: let r max = argmax r ∈ [ m ] c ( t ) r , c ( t ) =
c ( t ) r max and G ( t ) =
G ( t ) r max .
Signal and noise iterates.
Our analysis is build upon a decomposition of the updates equation (GD) and equation (GD+M) on w ∗ and X i [ j ] .
These decompositions are respectively deﬁned as follows:
[Equation]
[Equation]
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 deﬁne the derivative of a data-point i as (cid:96) ( t ) i = sigmoid( − y i f W ( t ) ( X i )) , the derivatives ν ( t ) k
= 1 N (cid:80) i ∈Z k (cid:96) ( t ) i for k ∈ { 1 , 2 } and the full derivative ν ( t )
=
ν ( t )
+ ν ( t ) 2 .
A NALYSIS OF GD
In this section, we provide a proof sketch for Theorem 3.1 that reﬂects the behavior of GD with λ = 0 .
A more detailed proof (extending to λ > 0 ) can be found in the Appendix.
Step 1: Learning Z 1 .
At the beginning of the learning process, the gradient is mostly dominated by the gradients coming from the Z 1 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 simpliﬁed as:
[Equation]
− (cid:80)
[Equation] Consequently, after T 0 = ˜Θ (cid:16) [Equation] iterations, for all t ∈ [ T 0 , 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 Z 1 data have small gradient and therefore, GD has learned these data.
Lemma 4.2.
Let T 0 = ˜Θ (cid:16) [Equation] [Equation] 0 (cid:17) .
After t ∈ [ T 0 , T ] iterations, the Z 1 derivative is bounded as ν [Equation]
derivative
[Equation]
By our choice of parameter, Lemma 4.2 indicates that the full gradient is dominated by the gradients from Z 2 data after T 0 = ˜Ω (cid:16) 1ˆ µηα (cid:17) .
Consequently, ν ( t ) 2 also rules the amount of feature learnt by GD.
Lemma 4.3.
Let T 0 = ˜Θ (cid:16) [Equation]
For t ∈ [ T 0 , T ] , equation (GD-S) becomes c ( t +1) ≤ ˜ O
(1 /α ) + ˜ O
( ηβ 3 /α ) (cid:80) tτ
=
T
ν ( τ ) 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 (cid:80) ms =1 ( c ( t ) s ) 3 + Ξ ( t ) i ) for some i ∈
Z 2 .
We thus need to determine whether the feature or the noise terms dominates in the sigmoid.
Step 2: Memorizing Z 2 .
We now show that the total correlation between the weights and the noise in Z 2 data increases until being large.
Lemma 4.4.
Let t ≥
and i ∈
Z 2 .
Assume that Ξ ( t ) i ≤ ˜ O (1) .
Then, equation (GD-N) can be simpliﬁed as:
[Equation]
Let T 1 = ˜ O (cid:16)
N σ 0 σ √ dσ 2 d (cid:17) .
Therefore, Ξ ( t ) i ≥ ˜Ω(1) , for t ∈ [ T 1 , T ] and thus GD memorizes.
By Lemma 4.4, in the gradient of Z 2 data, the noise term dominates the feature term (which scales as ˜ O ( β 3 ) ).
Consequently, the algorithm memorizes the Z 2 data which implies a fast decay of ν ( t ) 2 .
Lemma 4.5.
Let T 1 = ˜ O (cid:16)
[Equation]
For t ∈ [ T 1 , T ] , we have (cid:80) tτ
[Equation]
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 satisﬁes c ( t ) ≤ ˜ O (1 /α ) .
Lemma 4.4 and Lemma 4.6 respectively yield the ﬁrst two items in Theorem 3.1.
Bounds on the training and test errors are respectively obtained by plugging these results in (P) and (TE).
A NALYSIS OF GD+M
In this section, we provide a proof sketch for Theorem 3.2 that reﬂects the behavior of GD+M with λ = 0 .
A more detailed proof (also extending to λ > 0 ) can be found in the Appendix.
Step 1: Learning Z 1 .
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 simpliﬁed as:
[Equation]
Consequently, after T 0 = ˜Θ (cid:16) 1 σ
α 2 + [Equation] γ (cid:17) iterations, for all t ∈
[ T 0 , T ] , we have c ( t ) ≥ ˜Ω(1 /α ) .
Step 2: Learning Z 2 .
Contrary to GD, GD+M has a large momentum that contains w ∗ after Step 1.
Lemma 5.2.
Let T 0 = ˜Θ (cid:16) 1 σ
α 3 + 11 − γ (cid:17) .
For t ∈ [ T 0 , 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 T 0 , the momentum gradient stores historical gradients that are spanned by w ∗ .
It ampliﬁes the feature present in previous gradients to learn the feature from small margin data.
Lemma 5.3.
Let T 0 = ˜Θ (cid:16) 1 σ
α 3 + 11 − γ (cid:17) .
After T 1 = T 0 + ˜Θ (cid:16) 11 − γ (cid:17) iterations, for t ∈ [ T 1 , T ] , we have c ( t ) ≥ ˜Ω (cid:16) 1 √ 1 − γα (cid:17) .
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 T 0 = ˜Θ (cid:16) [-1 σ 
-][Equation]
[Equation] (cid:17) .
After T 1 = T 0 + ˜Θ (cid:16) 11 − γ (cid:17) iterations, for t ∈ [ T 1 , T ] , ν [Equation] t ) ≤ ˜ O [Equation]
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 ( X i ) } 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 σ 4 d 2 ) ν
( t ) .
As a consequence, after t ∈ [ T 1 , T ] iterations, we thus have | Ξ ( t ) i,j,r | ≤ ˜ O
( σ 0 σ √ d ) .
Lemma 5.3 and Lemma 5.5 respectively yield the two ﬁrst items in Theorem 3.2.
D ISCUSSION
Our work is a ﬁrst step towards understanding the algorithmic regularization of momentum and leaves room for improvements.
We constructed a data distribution where historical feature ampliﬁcation 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 beneﬁts.
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.