Title: STABILITY OF STOCHASTIC GRADIENT METHOD WITH MOMENTUM FOR STRONGLY CONVEX LOSS FUNCTIONS Anonymous authors Paper under double-blind review

Abstract: 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.

1. INTRODUCTION:

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 benefits in various settings. The heavy-ball momentum method was first 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 defined 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 first 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 quantifies the gap between the empirical risk of SGMM and the global optimum. Our bound can be made arbitrarily small by choosing sufficiently many iterations and a sufficiently 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 first 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árik, 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árik, 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 • 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.

2. GENERALIZATION 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:

R(w):

∆ = E z∼D f (w; z).

(1)

Since the distribution D is unknown, we replace the objective by the empirical risk, i.e., R S (w)

EQUATION

We assume w = A(S) for a potentially randomized algorithm A(•). In order to find 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 g , we consider uniform stability:

Definition 1 Let S and S denote two data sets from space Z n such that S and S differ in at most one example. Algorithm A is s -uniformly stable if for all data sets S, S , 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 s -uniformly stable algorithm, then the generalization error of A is upper-bounded by s .

Theorem 1 shows that it is enough to control the uniform stability of an algorithm to upper bound the generalization error.

2.1. ASSUMPTIONS ON THE LOSS FUNCTION:

In our analysis, we will assume that the loss function satisfies the following properties.

Definition 2 A function f : Ψ → R is L-Lipschitz if for all u, v ∈ Ψ we have |f (u) -f (v)| ≤ L u -v . Definition 3 A function f : Ψ → R is β-smooth if for all u, v ∈ Ψ we have ∇f (u) -∇f (v) ≤ β u -v . Definition 4 A function f : Ψ → R is γ-strongly convex if for all u, v ∈ Ψ we have f (u) ≥ f (v) + ∇f (v) T (u -v) + γ 2 u -v 2 .

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.2. STOCHASTIC GRADIENT METHOD WITH MOMENTUM:

The update rule for projected SGMM is given by:

EQUATION

where P denotes the Euclidean projection onto Ω, α > 0 is the learning rate1 , µ > 0 is the momentum parameter, i t is a randomly selected index, and f (w t ; z it ) is the loss evaluated on sample z it . In SGMM, we run the update (5) iteratively for T steps and let w T denote the final output. Note that there are two typical approaches to select i t . The first 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:

g = E S,A [R(w T ) -R S (w T )] = E S,i0,••• ,i T -1 [R(w T ) -R S (w T )] since the randomness in A arises from the choice of i 0 , • • • , i T -1 .

3. MAIN 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 ≤ µ < αβγ 3(β+γ) and α ≤ 2 β+γ , SGMM satisfies 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 verified 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

E S,A [R S ( ŵT ) -R S (w * S )] ≤ µW 0 (1 -µ)T + (1 -µ)W 1 2αT - γW 2 2 - µγW 3 2(1 -µ) + αL 2 2(1 -µ) (7)

where ŵT denotes the average of T steps of the algorithm, i.e., ŵT = 1

T +1 T t=0 w t , R S (w) = 1 n n i=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 [ w 0 - w * S 2 ], W 2 = E S,A [ ŵT -w * S 2 ], and W 3 = 1 T +1 T t=0 E S,A [ w t -w t-1 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 first 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 α sufficiently 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

µ(β + γ) αβγ. Then, setting α = 1-µ L W1

T , we have:

EQUATION

)

where

C ∆ = 1-3µL(β+γ) √ T (1-µ)βγ √

W1 and ŵT as well as the constants W 0 , • • • , W 3 are defined 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 α simplifies considerably when µ is sufficiently 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 first two terms in (8) vanish as T increases. The last term in (8) vanishes as the number of samples n increases.

4. PROOF OF THEOREM 2 (STABILITY 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 first consider SGMM without projection and defer the discussion of projection to the end of this proof. Let S = {z 1 , • • • , z n } and S = {z 1 , • • • , z n } be two samples of size n that differ in at most one example. Let w T and w T denote the outputs of SGMM on S and S , respectively. We consider the updates w t+1 = G t (w t ) + µ(w t -w t-1 ) and w t+1 = G t (w t ) + µ(w t -w t-1 ) with G t (w t ) = w t -α∇ w f (w t ; z it ) and G t (w t ) = w t -α∇ w f (w t ; z it ), respectively, for t = 1, • • • , T . We denote δ t ∆ = w t -w t . Suppose w 0 = w 0 , i.e., δ 0 = 0. We first 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 = 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

)

where the second inequality holds due to µ ≥ αβγ β+γ -1 2 .

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 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.

5. PROOF OF THEOREM 3 (CONVERGENCE BOUND):

Again, we first 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 define:

EQUATION

with p 0 = 0. Substituting into the SGMM update, the parameter recursion is given by

EQUATION

It follows that

EQUATION

)

Substituting p t (15) into ( 17), the recursion ( 16) can be written as

w t+1 + p t+1 -w 2 = w t + p t -w 2 + α 1 -µ 2 ∇ w f (w t ; z it ) 2 - 2αµ (1 -µ) 2 (w t -w t-1 ) T ∇ w f (w t ; z it ) - 2α 1 -µ (w t -w) T ∇ w f (w t ; z it ).

(18) Upon taking the expectation with respect to i t in (18) we have

EQUATION

where we use the fact that ∇ w f (w t ; z it ) ≤ L, due to L-Lipschitz, and that E it [∇ w f (w t ; z it )] = ∇ 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

EQUATION

)

Since • 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 it ). Then, for any feasible w ∈ Ω, (17) holds for y t+1 , i.e.,

EQUATION

)

Note that the LHS of (25) can be written as

y t+1 + µ 1 -µ (y t+1 -w t ) -w 2 = 1 1 -µ y t+1 -µw t + (1 -µ)w .

We note that µw t + (1 -µ)w ∈ Ω for any w ∈ Ω and w t ∈ Ω since Ω is convex. Now in projected SGMM, we have w t+1 -µw t + (1 -µ)w 2 = P(y t+1 ) -µw t + (1 -µ)w 2 ≤ y t+1 -µw t + (1 -µ)w 2 (26) since projection a point onto Ω moves it closer to any point in Ω. This shows inequality (19) holds, and the convergence results do not change. 

6. EXPERIMENTAL EVALUATION:

In this section, 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 classification 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.01. The weight decay coefficient 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 fixed T = 1000 iterations are shown in Figures 1a and 1b , respectively, for µ = 0, 0.5, 0.9. In Figures 2a and 2b , we plot the generalization error (with respect to classification 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 classification accuracy) decreases as n increases for all values of µ, which is suggested by our stability upper-bound in Theorem 2. In addition, for sufficiently 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 . In order to study the optimization error of SGMM, 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 classification accuracy for training and test datasets in Figures 4a and 4b , respectively. We observe that the training error decreases as the number 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 benefit 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.

:

In the following, we assume α is constant over time. In practice, time-decaying α is used in some applications. Please note that our results hold for time-decaying α.

Figures and Tables:

Figure 1: Generalization performance (cross entropy) of logistic regression for notMNIST dataset with T = 1000 iterations and minibatch size 10.

Figure 2: Generalization performance (classification accuracy) of logistic regression for notMNIST dataset with T = 1000 iterations and minibatch size 10.

Figure 3: Training and test error of logistic regression (cross entropy loss) for notMNIST dataset with n = 500 training samples and minibatch size 10.

Figure 4: Training and test accuracy of logistic regression for notMNIST dataset with n = 500 training samples and minibatch size 10.

Peer review (iclr2019):
 This paper studies the algorithmic stability of SGD with momentum and provides an upper-bound on true risk through convergence analysis. This bound clarifies dependencies of convergence speed on the size of dataset and the momentum parameter. The presentation is easy to follow and technically sounds good. SGD with momentum is heavily used for learning linear models and deep neural networks, hence to analyze its convergence behavior is quite important. This paper achieves this goal well by extending a previous result on vanilla SGD in a straightforward manner, although it is not technically difficult.