S TABILITY
OF S TOCHASTIC G RADIENT
M ETHOD
WITH
M OMENTUM
FOR
S TRONGLY
C ONVEX
L OSS
F UNCTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
While momentum-based methods, in conjunction with the stochastic gradient descent, are widely used when training machine learning models, there is little theoretical understanding on the generalization error of such methods.
In practice, the momentum parameter is often chosen in a heuristic fashion with little theoretical guidance.
In this work we use the framework of algorithmic stability to provide an upper-bound on the generalization error for the class of strongly convex loss functions, under mild technical assumptions.
Our bound decays to zero inversely with the size of the training set, and increases as the momentum parameter is increased.
We also develop an upper-bound on the expected true risk, in terms of the number of training steps, the size of the training set, and the momentum parameter.
I NTRODUCTION
A fundamental issue for any machine learning algorithm is its ability to generalize from the training dataset to the test data.
A classical framework used to study the generalization error in machine learning is PAC learning (Vapnik and Chervonenkis, 1971; Valiant, 1984).
However, the associated bounds using this approach can be conservative.
Recently, the notion of uniform stability, introduced in the seminal work of Bousquet and Elisseeff (Bousquet and Elisseef, 2002), is leveraged to analyze the generalization error of the stochastic gradient method (SGM) (Hardt et al., 2016).
The result in (Hardt et al., 2016) is a substantial step forward, since SGM is widely used in many practical systems.
This method is scalable, robust, and widely adopted in a broad range of problems.
To accelerate the convergence of SGM, a momentum term is often added in the iterative update of the stochastic gradient (Goodfellow et al., 2016).
This approach has a long history, with proven beneﬁts in various settings.
The heavy-ball momentum method was ﬁrst introduced by Polyak (Polyak, 1964), where a weighted version of the previous update is added to the current gradient update.
Polyak motivated his method by its resemblance to a heavy ball moving in a potential well deﬁned by the objective function.
Momentum methods have been used to accelerate the backpropagation algorithm when training neural networks (Rumelhart et al., 1986).
Intuitively, adding momentum accelerates convergence by circumventing sharp curvatures and long ravines of the sublevel sets of the objective function (Wilson et al., 2018).
For example, Ochs et al. has presented an illustrative example to show that the momentum can potentially avoid local minima (Ochs et al., 2015).
Nesterov has proposed an accelerated gradient method, which converges as O (1 /k 2 ) where k is the number of iterations (Nesterov, 1983).
However, the Netstrov momentum does not seem to improve the rate of convergence for stochastic gradient (Goodfellow et al., 2016, Section 8.3.3).
In this work, we focus on the heavy-ball momentum.
Although momentum methods are well known to improve the convergence in SGM, their effect on the generalization error is not well understood.
In this work, we ﬁrst build upon the framework in (Hardt et al., 2016) to obtain a bound on the generalization error of SGM with momentum (SGMM) for the case of strongly convex loss functions.
Our bound is independent of the number of training iterations and decreases inversely with the size of the training set.
Secondly, we develop an upper-bound on the optimization error, which quantiﬁes the gap between the empirical risk of
SGMM and the global optimum.
Our bound can be made arbitrarily small by choosing sufﬁciently many iterations and a sufﬁciently small learning rate.
Finally, we establish an upper-bound on the expected true risk of SGMM as a function of various problem parameters.
We note that the class of strongly convex loss functions appears in several important machine learning problems, including linear and logistic regression with a weight decay regularization term.
Other related works: convergence analysis of ﬁrst order methods with momentum is studied in (Nesterov, 1983;
Ochs et al., 2014; Su et al., 2014; Ghadimi et al., 2015; Lessard et al., 2016;
Yang et al., 2016; Loizou and Richt´arik, 2018; Gadat et al., 2016).
Most of these works consider the deterministic setting for gradient update.
Only a few works have analyzed the stochastic setting (Yang et al., 2016; Loizou and Richt´arik, 2018; Gadat et al., 2016).
Our convergence analysis results are not directly comparable with these works due to their different assumptions regarding the properties of loss functions.
In particular, we analyze the convergence of SGMM for a smooth and strongly convex loss function as in (Hardt et al., 2016), which is new.
First-order methods with noisy gradient are studied in (Kidambi et al., 2018) and references therein.
In (Kidambi et al., 2018), the authors show that there exists linear regression problems for which
SGM outperforms SGMM in terms of convergence.
Our main focus in this work is on the generalization, and hence true risk, of SGMM.
We are aware of only one similar work in this regard, which provides stability bounds for quadratic loss functions (Chen et al., 2018).
In this paper, we obtain stability bounds for the general case of strongly convex loss functions.
In addition, unlike (Chen et al., 2018), our results show that machine learning models can be trained for multiple epochs of SGMM with bounded generalization errors.
Notation: We use E [ · ] to denote the expectation and (cid:107) · (cid:107) to represent the Euclidean norm of a vector.
We use lower-case bold font to denote vectors.
We use sans-serif font to denote random quantities.
Sets and scalars are represented by calligraphic and standard fonts, respectively.
G ENERALIZATION ERROR AND STABILITY
We consider a general supervised learning problem, where S =
{ z 1 , · · · , z n } denotes the set of samples of size n drawn i.i.d.
from some space Z with an unknown distribution D .
We assume a learning model described by parameter vector w .
Let f ( w ; z ) denote the loss of the model described by parameter w on example z ∈ Z .
Our ultimate goal is to minimize the true or population risk:
[Equation]
Since the distribution D is unknown, we replace the objective by the empirical risk, i.e.,
[Equation]
We assume w =
A
( S ) for a potentially randomized algorithm A ( · ) .
In order to ﬁnd an upper-bound on the true risk, we consider the generalization error, which is the expected difference of empirical and true risk:
[Equation]
Finally, to upper bound (cid:15) g , we consider uniform stability:
Deﬁnition 1 Let S and S (cid:48) denote two data sets from space Z n such that S and S (cid:48) differ in at most one example.
Algorithm A is (cid:15) s -uniformly stable if for all data sets S , S (cid:48) , we have
[Equation]
It is shown in (Hardt et al., 2016) that uniform stability implies generalization in expectation:
Theorem 1 (Hardt et al., 2016)
If A is an (cid:15) s -uniformly stable algorithm, then the generalization error of A is upper-bounded by (cid:15) s .
Theorem 1 shows that it is enough to control the uniform stability of an algorithm to upper bound the generalization error.
2. A SSUMPTIONS ON THE LOSS FUNCTION
In our analysis, we will assume that the loss function satisﬁes the following properties.
Deﬁnition
A function f : Ψ → R is L -Lipschitz if for all u , v ∈ Ψ we have | f ( u ) − f ( v ) | ≤
L (cid:107) u − v (cid:107) .
Deﬁnition
A function f : Ψ → R is β -smooth if for all u , v ∈ Ψ we have (cid:107)∇ f ( u ) − ∇ f ( v ) (cid:107) ≤ β (cid:107) u − v (cid:107) .
Deﬁnition 4 A function f : Ψ → R is γ -strongly convex if for all u , v ∈ Ψ we have f ( u ) ≥
[Equation]
We assume that the parameter space Ω is a convex set.
Furthermore, for the loss function to be
L -Lipschitz and and strongly convex, we further assume that Ω is compact.
Since Ω is compact, the
SGMM update requires projection.
2. S TOCHASTIC GRADIENT METHOD WITH MOMENTUM
The update rule for projected SGMM is given by:
[Equation]
(5) where P denotes the Euclidean projection onto Ω , α > 0 is the learning rate 1 , µ > 0 is the momentum parameter, i t is a randomly selected index, and f ( w t ; z i t ) is the loss evaluated on sample z i t .
In
SGMM, we run the update (5) iteratively for T steps and let w T denote the ﬁnal output.
Note that there are two typical approaches to select i t .
The ﬁrst approach is to select i t ∈ { 1 , · · · , n } uniformly at random at each iteration.
The second approach is to permutate { 1 , · · · , n } randomly once and then select the examples repeatedly in a cyclic manner.
Our results are valid for both approaches.
The key quantity of interest in this paper is the generalization error for SGMM given by:
[Equation]
since the randomness in A arises from the choice of i 0 , · · · , i T − 1 .
M AIN RESULTS
In the following, we assume that the loss function f ( · ; z ) is β -smooth, L -Lipschitz, and γ -strongly convex for all z .
Theorem 2 (Stability bound) Suppose that the SGMM update (5) is executed for T steps with constant learning rate α and momentum µ .
Provided that αβγβ + γ − 1 2 ≤ µ < [Equation] and α ≤ [Equation]
SGMM satisﬁes (cid:15) s -uniform stability where
[Equation]
The result in Theorem 2 implies that the stability bound decreases inversely with the size of the training set.
It increases as the momentum parameter µ increases.
These properties are also veriﬁed in our experimental evaluation.
Theorem 3 (Convergence bound) Suppose that the SGMM update (5) is executed for T steps with constant learning rate α and momentum µ .
Then we have
[Equation]
where ˆ w T denotes the average of T steps of the algorithm, i.e., ˆ w T = 1 T +1 (cid:80) Tt =
w t , R S ( w ) =
n (cid:80) ni
=1 f ( w ; z i ) , w ∗ S = arg min
w R S ( w ) , W 0 = E S ,A [ R S ( w 0 ) − R S ( w T )] , W 1 = E S ,A [ (cid:107) w 0 − w ∗ S (cid:107) 2 ] , W 2 = E S ,A
[ (cid:107) ˆ w T − w ∗ S (cid:107) 2 ] , and W 3 = 1 T +
(cid:80)
Tt =0 E S ,A
[ (cid:107) w t − w t − 1 (cid:107) 2 ] .
Theorem 3 bounds the optimization error, i.e., the expected difference between the empirical risk achieved by SGMM and the global minimum.
Upon setting µ = 0 and γ = 0 in (7), we can recover the classical bound on optimization error for SGM (Nemirovski and Yudin., 1983), (Hardt et al., 2016, Theorem 5.2).
The ﬁrst two terms in (7) vanish as T increases.
The terms with negative sign improve the convergence due to the strongly convexity.
The last term depends on the learning rate, α , the momentum parameter µ , and the Lipschitz constant L .
This term can be controlled by selecting α sufﬁciently small.
Proposition 1 (Upper-bound on true risk) Suppose that the SGMM update (5) is executed for T steps with constant learning rate α and momentum µ , satisfying the conditions in Theorem 2 and µ ( β + γ ) (cid:28) αβγ .
Then, setting α = 1 − µL (cid:113) W 1 T , we have:
[Equation]
2(1 − µ ) + 2 L 2 ( β + γ ) nβγC (8) where C ∆ = 1 − 3 µL ( β + γ ) √ T (1 − µ ) βγ √ W
and ˆ w T as well as the constants W 0 , · · · , W 3 are deﬁned in Theorem 3.
Proposition 1 provides a bound on the expected true risk of SGMM in terms of the global minimum of the empirical risk.
The bound in (8) is obtained by combining Theorem 2 and Theorem 3 and minimizing the expression over α .
The choice of α simpliﬁes considerably when µ is sufﬁciently small, as stated in Proposition 1.
Due to the page constraint, the proof of this result is provided in the supplementary material.
Note that the ﬁrst two terms in (8) vanish as T increases.
The last term in (8) vanishes as the number of samples n increases.
P ROOF OF T HEOREM 2 (S TABILITY BOUND )
Following (Hardt et al., 2016), we track the divergence of two different iterative sequences of update rules with the same starting point.
However, our analysis is more involved as the presence of momentum term requires a more careful bound on the iterative expressions.
To keep the notation uncluttered, we ﬁrst consider SGMM without projection and defer the discussion of projection to the end of this proof.
Let S
=
{ z
, · · · , z n } and S (cid:48) = { z (cid:48) 1 , · · · , z (cid:48) n } be two samples of size n that differ in at most one example.
Let w T and w (cid:48)
T denote the outputs of SGMM on S and S (cid:48) , respectively.
We consider the updates w t +1 =
G t ( w t ) + µ ( w t − w t − 1 ) and w (cid:48) t +1 =
G (cid:48) t ( w (cid:48) t ) + µ ( w (cid:48) t − w (cid:48) t − 1 ) with G t ( w t ) =
w t − α ∇ w f ( w t ; z i t ) and G (cid:48) t ( w (cid:48) t ) =
w (cid:48)
t − α ∇ w f ( w (cid:48) t ; z (cid:48) i t ) , respectively, for t = 1 , · · · , T .
We denote δ t ∆ = (cid:107) w t − w (cid:48)
t (cid:107) .
Suppose w 0 = w (cid:48) 0 , i.e., δ 0 = 0 .
We ﬁrst establish an upper-bound on E A [ δ t +1 ] in terms of E A [ δ t ] and E A
[ δ t − 1 ] in the following lemma, whose proof is provided in the supplementary document.
Lemma 1 Provided that α ≤ 2 β + γ , an upper-bound on E A
[ δ t +1 ] is given by
[Equation]
Using the result of Lemma 1, in the following, we develop an upper bound on E A
[ δ T ] .
Let us consider the recursion
[Equation]
with ˜ δ
= δ 0 = 0 .
Upon inspecting (10) it is clear that
[Equation]
as we simply drop the remainder of positive terms.
Substituting (11) into (10), we have
[Equation] [Equation]
where the second inequality holds due to µ ≥ αβγβ + γ − 12 .
Noting that E A [˜ δ t ] ≥ E A
[ δ t ] for all t including T , we have
[Equation]
where the second expression holds since 0 ≤ µ < αβγ 3( β + γ ) is assumed.
Applying the L -Lipschitz property on f ( · , z ) , it follows that
[Equation]
Since this bound holds for all S , S (cid:48) and z , we obtain an upper-bound on the uniform stability and the proof is complete.
Our stability bound in Theorem 2 holds for the projected SGMM update (5) because Euclidean projection does not increase the distance between projected points (the argument is essentially analogous to (Hardt et al., 2016, Lemma 4.6)).
In particular, note that Lemma 1 holds for the projected SGMM.
P ROOF OF T HEOREM 3 (C ONVERGENCE BOUND )
Again, we ﬁrst consider SGMM without projection and discuss the extension to projection at the end of this proof.
Our proof is inspired by the convergence analysis in (Yang et al., 2016; Ghadimi et al., 2015) for a convex loss function with bounded variance and time-decaying learning rate.
Different from these works, we analyze the convergence of SGMM for a smooth and strongly convex loss function with constant learning rate.
To facilitate the convergence analysis, we deﬁne:
[Equation]
with p 0 = 0 .
Substituting into the SGMM update [Equation] the parameter recursion is given by
[Equation]
[Equation]
It follows that
[Equation]
Substituting p t (15) into (17)
[Equation]
the recursion (16) can be written as
[Equation] [Equation]
Upon taking the expectation with respect to i t in (18) we have
[Equation]
where we use the fact that (cid:107)∇ w f ( w t ; z i t ) (cid:107) ≤ L , due to L -Lipschitz, and that E i t [ ∇ w f ( w t ; z i t )]
= ∇ w R S ( w t ) .
Furthermore, since R S ( · ) is a γ -strongly convex function, for all w t and w t − 1 , we have
[Equation]
Substituting (20) in (19), we have
[Equation]
Taking expectation over i 0 , · · · , i t for a given S , summing (21) for t = 0 , · · · , T , and rearranging terms, we have 2 α
[Equation]
Since (cid:107) · (cid:107) is a convex function, for all w T and w , we have
[Equation]
Furthermore, due to convexity of R S ( · ) , we have
[Equation]
Taking expectation over S , applying inequalities (23) and (24) into (22), and substituting w = w ∗ S , we obtain (7) and the proof is complete.
Our convergence bound in Theorem 3 can be extended to projected SGMM (5).
Let use denote y t +1 ∆ =
w t + µ ( w t − w t − 1 ) − α ∇ w f ( w t ; z i t ) .
Then, for any feasible w ∈ Ω , (17) holds for y t +1 , i
[Equation]
e., (cid:107)
[Equation]
+ µ
[Equation]
Note that the LHS of (25) can be written as
[Equation]
We note that µ w t + (1 − µ ) w ∈ Ω for any w ∈ Ω and w t ∈ Ω since Ω is convex.
Now in projected SGMM, we have
[Equation]
[Equation]
since projection a point onto Ω moves it closer to any point in Ω .
This shows inequality (19) holds, and the convergence results do not change.
[Figure] [Figure]
E XPERIMENTAL EVALUATION
In this section
[Figure]
we validate the insights obtained in our theoretical results in experimental evaluation.
Our main goal is to study how adding momentum affects the convergence and generalization of SGM.
We study the performance of SGMM when applied to the notMINIST dataset.
Please note that similar results are provided for the MNIST dataset in the supplementary document.
We train a logistic regression model with the weight decay regularization using SGMM for binary classiﬁcation on the two-class notMNIST dataset that contains the images from letter classes “C” and “J”, which leads to a smooth and strongly convex loss function.
We set the learning rate α = 0 .
The weight decay coefﬁcient and the minibatch size are set to 0.001 and 10, respectively.
We use 100 SGMM realizations to evaluate the average performance.
We compare the training and generalization performance of SGM without momentum with that of SGMM under µ = 0 .
5 and µ = 0 .
9 , which are common momentum values used in practice (Goodfellow et al., 2016, Section 8.3.2).
The generalization error (with respect to cross entropy) and training error versus the number of training samples, n , under SGMM with ﬁxed T = 1000 iterations are shown in Figures 1a and 1b, respectively, for µ = 0 , 0 . 5 , 0 .
In Figures 2a and 2b, we plot the generalization error (with respect to classiﬁcation accuracy) and the training accuracy as a function of the number of training samples for the same dataset.
First, we observe that the generalization error (with respect to both cross entropy and classiﬁcation accuracy) decreases as n increases for all values of µ , which is suggested by our stability upper-bound in Theorem 2.
In addition, for sufﬁciently large n , we observe that the generalization error increases with µ , consistent with Theorem 2.
On the other hand, the training error increases as n increases, which is expected.
We can observe that adding momentum reduces training error as it improves the convergence rate.
The training accuracy also improves by adding momentum as illustrated in Fig. 2b.
[Figure] [Figure]
In order to study the optimization error of SGMM
[Figure]
we show the training error and test error versus the number of epochs, under SGMM trained with n = 500 samples in Figures 3a and 3b, respectively.
We plot the classiﬁcation accuracy for training and test datasets in Figures 4a and 4b, respectively.
We observe that the training error decreases as the number of epochs increases for all values of µ , which is consistent with the convergence analysis in Theorem 3.
Furthermore, as expected, we see that adding momentum improves the training error and accuracy.
However, as the number of epochs increases, we note that the beneﬁt of momentum on the test error and accuracy becomes negligible.
This happens because adding momentum also results in a higher generalization error thus penalizing the gain in training error.
C ONCLUSIONS
We study the generalization error and convergence of SGMM for the class of strongly convex loss functions, under mild technical conditions.
We establish an upper-bound on the generalization error, which decreases with the size of the training set, and increases as the momentum parameter is increased.
Secondly, we analyze the convergence of SGMM during training, by establishing an upper-bound on the gap between the empirical risk of SGMM and the global minimum.
Our proposed bound reduces to a classical bound on the optimization error of SGM (Nemirovski and Yudin., 1983) for convex functions, when the momentum parameter is set to zero.
Finally, we establish an upper-bound on the expected difference between the true risk of SGMM and the global minimum of the empirical risk, and illustrate how it scales with the number of training steps and the size of the training set.
Although our results are established for the case when the learning rate is constant, they can be easily extended to the case when the learning rate decreases with the number of iterations.
We also present experimental evaluations on the notMNIST dataset and show that the numerical plots are consistent with our theoretical bounds on the generalization error and the convergence gap.