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 presents an analysis of generalization error of SGD with multiple passes for strongly convex objectives using the framework of algorithmic stability [Bousquet and Elisseef, 2002] and its recent use to analyze generalization error of SGD based methods [Hardt, Recht and Singer 2016]. The problem considered by this work is interesting and raises the possibility of understanding generalization related questions of SGD style methods when augmented with momentum, which is common practice in Deep Learning [Sutskever et al. 2013]. That said, there are some concerns about the results as presented in this paper, which I will elaborate below: - Consider the stability bound admitted by theorem 2: The special case (similar to theorem 3.9 of Hardt et al 2016) when the learning rate alpha = 1/beta (which is the typical learning rate that theory advocates), and setting kappa = beta/gamma where kappa is the condition number of the problem, leads to the following bound on momentum allowed by theorem 2, which is: (something non-positive) <= mu < 1/(3*kappa). This is basically the regime where momentum does not make any difference towards accelerating optimization. Referring to the standard value of momentum for strongly convex functions, we see that the momentum is set as mu = (sqrt(kappa)-1)/(sqrt(kappa)+1) [Nesterov, 1983], or, mu = ( (sqrt(kappa)-1)/(sqrt(kappa)+1) )^2 [Polyak,1964]. Upon simplification of this standard momentum values, we see that mu .approx 1 - 1/sqrt(kappa) which grows close to one as kappa grows large. On the other hand, the momentum values admitted by the paper for their bound is super tiny (which gets to zero as the condition number kappa grows large). This essentially implies there is not much about momentum that is captured by the bound of theorem 2 since there is no characterization of the provided bound for theoretically advocated and practically used parameters for momentum. - In proposition 1, there is no quantitative description of what "sufficiently small" mu (momentum parameter) is - this statement is imprecise. As mentioned in the previous point, sufficiently small mu really is not descriptive of momentum parameters employed in practice (mu in practice typically is >= 0.9). For strongly convex objectives, this should be closer to 1- (1/sqrt(kappa)). Sufficiently small mu parameter essentially does not yield quantitatively different behaviors compared to standard SGD. In summary, while this paper attempts to make progress on an interesting question, but falls short and doesn*t really capture the behavior of these methods that is even mildly reflective of practice (even in terms of the parameter regimes admitted by the bounds proven in the theorems). - This paper does not perform a thorough literature survey of published results. Furthermore, this paper does not present a precise treatment of assumptions (and implications) amongst other works cited in the paper (see for e.g. [4] below). [1] Polyak (1987) presents (generalization) behavior of Heavy Ball momentum with noisy (inexact) gradients. [2] Several efforts in Signal Processing literature do consider the similar setting as one considered by this paper, which is that of Heavy Ball (called accelerated LMS) method with noisy gradients: refer to Proakis (1974), Roy and Shynk (1990), Sharma et al. (1998). [3] Kidambi et al (2018) estimate the "optimization" power (which is a part of characterization of generalization error [Bach and Moulines 2011], since this dominates at the start of optimization) of HB method with Stochastic Gradients and prove that HB+stochastic gradients does not offer any speedup over vanilla SGD. [4] Loizou and Richtarik provide an analysis of stochastic heavy ball with super large batch sizes (so they end up showing accelerated rates) under similar assumptions as considered by this paper, such as assuming the function is smooth and strongly convex. However, the paper dismisses the work of Loizou and Richtarik to be working with a different set of assumptions - this is not really true.