A Unified Analysis of Stochastic Momentum Methods for Deep Learning

1 Introduction

Momentum methods have a long history dating back to 1960's. [Polyak, 1964] proposed a heavy-ball (HB) method that uses the previous two iterates when computing the next one. The original motivation of momentum methods is to speed up the convergence for convex optimization. For a twice continuously differential strongly convex and smooth objective function, Polyak's analysis yields an accelerated linear convergence rate over the standard gradient method. In 1983, [Nesterov, 1983] proposed an accelerated gradient (NAG) method, which is also deemed as a momentum method and achieves the optimal O(1/t2 ) convergence rate for convex smooth optimization 1 , which has a clear advantage over standard gradient method with O(1/t) convergence for the same problem. It was later on shown to have an accelerated linear convergence rate for smooth and strongly convex optimization problems [Nesterov, 2004]. Both the HB method and the NAG method use a momentum term in updating the solution, i.e., the difference between current iterate and the previous iterate. Therefore, both methods have been referred to as momentum methods in literature.

1 t is the number of iterations.

Due to recently increasing interests in deep learning, the stochastic variants of HB and NAG methods have been employed broadly in optimizing deep neural networks [Krizhevsky et al., 2012; Hinton et al., 2012]. [Sutskever et al., 2013] are probably the first to study SNAG and compare it with SHB for optimizing deep neural networks. Although they have some interesting findings of these two methods in deep learning (e.g., a distinct improvement in performance of SNAG is usually observed in their experiments), their analysis and argument are mostly restricted to convex problems. Moreover, some questions remained unanswered, e.g., (i) Do SHB and SNAG enjoy faster convergence rates than SG for deep learning (non-convex optimization problems) as in the convex and deterministic setting? (ii) If not, what is the advantage of these two methods over SG? (iii) Why does SNAG often yield improvement over SHB?

In this paper, we propose and analyze a unified framework for stochastic momentum methods and stochastic gradient method aiming at bringing answers and more insights to above questions. We summarize our results and contributions as follows:

• We propose a unified stochastic momentum framework parameterized by a free scalar parameter. The framework reduces to SHB, SNAG and SG by setting three different values for the free parameter.
• We present a unified convergence analysis of the gradient's norm of the training objective of these stochastic methods for non-convex optimization, revealing the same rate of convergence for three variants.
• We analyze the generalization error of the unified framework through the uniform stability approach. The result exhibits a clear advantage of stochastic momentum methods, i.e., adding a momentum helps generalization.
• Our empirical results for learning deep neural networks

complete the unified view and analysis by showing that (i) there is no clear advantage of SHB and SNAG over SG in convergence speed of the training error; (ii) the advantage of SHB and SNAG lies at better generalization due to more stability; (iii) SNAG usually achieves the best tradeoff between speed of convergence in training error and stability of testing error among the three stochastic methods.

2 More RelatedWork

There is much analysis on the momentum methods for deterministic optimization. Nesterov pioneered the work of accelerated gradient methods for smooth convex optimization [Nesterov, 2004]. The convergence analysis of HB has been recently extended to smooth functions for both convex [Ghadimi et al., 2014; Ochs et al., 2015] and non-convex deterministic optimization [Ochs et al., 2014; Ochs, 2016]. As the rising popularity of deep neural networks, the stochastic variants of HB and NAG have been employed widely for training neural networks and leading to tremendous success for many problems in computer vision and speech recognition [Krizhevsky et al., 2012; Hinton et al., 2012; Sutskever et al., 2013]. However, their stochastic variants in non-convex optimization are under-explored.

It is worth mentioning that two recent works have established the convergence results of the SG method [Ghadimi and Lan, 2013] and the stochastic version of a different variant of accelerated gradient method for non-convex optimization [Ghadimi and Lan, 2016]. However, the variant of accelerated gradient method in [Ghadimi and Lan, 2016] is hard to be explained in the framework of momentum methods and is not widely employed for optimizing deep neural networks. Moreover, their analysis is not applicable to the SHB method. Hence, from a theoretical standpoint, it is still interesting to analyze the stochastic variants of the Nesterov's accelerated gradient method and the HB method for stochastic non-convex optimization, which are extensively employed for learning deep neural networks. Our unified analysis shows that they enjoy the same order of convergence rate as the SG method, which conincides with the results in [Ghadimi and Lan, 2013; Ghadimi and Lan, 2016].

On the other hand, there exist few studies on analyzing the statistical properties (e.g., the generalization error) of the model learned by the SG method or stochastic momentum methods for minimizing the empirical risk. Conventional studies on the SG method in terms of statistical property focus on one pass learning, i.e., the training examples are passed once [Cesa-Bianchi et al., 2004]. Recently, there emerge several works that aim to establish the statistical properties of the multiple pass SG methods in machine learning [Lin and Rosasco, 2016; Hardt et al., 2016]. The latter work is closely related to the present work, which established the generalization error of the SG method with multiple pass for both convex and non-convex learning problems by analyzing the uniform stability. Nevertheless, it remains an open problem from a theoretical standpoint how the momentum term helps improve the generalization, though it has been observed to yield better performance in practice for deep learning [Sutskever et al., 2013]. Our unified analysis of the uniform stability of the SG method and stochastic momentum methods explicitly exhibit the advantage of the stochastic momentum methods in terms of the generalization error, hence providing the theoretical support for the common wisdom.

In the remainder of the paper, we first review the HB and NAG method, and present their stochastic variants. Then we present a unified view of these momentum methods. Next, we present the convergence and generalization analysis for stochastic momentum methods. In addition, we present empirical results for comparing different methods for optimizing deep neural networks. Finally, we conclude this work.

3 MomentumMethods And Their Stochastic Variants

3.1 Notations and Setup

Let us consider a general setting of learning with deep learning as a special case. Given a set of training examples S = (q1, . . . , qn) sampled from an unknown distribution D, the goal of learning is to find a model x that minimizes the population risk, i.e.,

where ℓ is a loss function, ℓ(x, q) denotes the loss of the model x on the example q and Ω denotes the hypothesis class of the model. Since we cannot compute F (x) due to unknown distribution D, one usually learns a model by minimizing the empirical risk, i.e.,

Two concerns usually present in the above empirical risk minimization approach. First, how fast the optimization algorithm solves Problem (2). This is usually measured by the speed of convergence to the optimal solution. However, it is NP-hard to find the global optimal solution for a general non-convex optimization problem [Hillar and Lim, 2013]. As with many previous works [Ghadimi and Lan, 2013; Ghadimi and Lan, 2016; Reddi et al., 2016; Zhu and Hazan, 2016], we study the convergence rate of an iterative algorithm to the critical point, i.e., a point x∗ such that ∇f(x∗) = 0.

Second, how the model learned by solving Problem (2) generalizes to different data. It is usually measured by the population risk F(̂x) defined in (1). Since the model ̂x is learned from the random samples q1, . . . , qn with randomness in the optimization algorithm itself, the expected population risk E[F(̂x)] is also used for the analysis with the expectation taking over the randomness in the samples and the algorithm itself. One way to assess the expected population risk is the generalization error, i.e., the difference between the population risk and the empirical risk,

We use ∇h(x) to denote the gradient of a smooth function. A function is smooth iff there exists L > 0 such that

where ‖ · ‖ denotes the Euclidean norm. Note that the above inequality does not imply convexity. Through the paper, we assume that ℓ(x, q) a G-Lipschitz continuous and L-smooth non-convex function in x, and assume that Ω = Rd . It follows that f(x) is G-Lipschitz continuous and L-smoothness.

3.2 Stochastic MomentumMethods

We denote by Gk = G(xk; ξk) a stochastic gradient of f(x) at xk depending on a random variable ξk such that E[G(xk; ξk)] = ∇f(xk). In the context of the empirical risk minimization (2), G(xk; ξk) = ∇ℓ(xk; qik ), where ik is a random index sampled from {1, . . . , n}.

There are two variants of momentum methods for solving (2), i.e., HB and NAG. HB was originally proposed for optimizing a smooth and strongly convex objective function. Based on HB, the update of stochastic HB (SHB) is given below for k = 0, . . . ,

with x-1 = x0, where β ∈ [0, 1) is the momentum constant and α is the step size. Equivalently, the above update can be implemented by the following two steps for k = 0, . . . ,:

Based on NAG [Nesterov, 2004], the update of stochastic NAG (SNAG) consists of the two steps below for k = 0, . . . ,:

with y0 = x0. By introducing vk = yk -yk-1 with v0 = 0, the above update can be equivalently written as

Finally, the traditional view of SG can be written as

By comparing (8) to (6), one might argue that the difference between HB and NAG lies at the point for evaluating the gradient [Sutskever et al., 2013]. We will present a different unified view of the three methods that allows us to analyze them in a unified framework. The convergence of HB and NAG has been established for convex optimization [Polyak, 1964; Nesterov, 1983; Nesterov, 2004; Ghadimi et al., 2014; Ochs et al., 2015].

4 AUnified View of Stochastic Momentum Methods

In this section, we present a unified view of the two (stochastic) momentum methods and (stochastic) gradient methods. We first present the unified framework and then show that HB, NAG and the gradient method are special cases of the unified framework. Denote by G(xk) either a gradient or a stochastic gradient of f(x) at xk.

Let α > 0, β ∈ [0, 1), and s ≥ 0. The updates of the stochastic unified momentum (SUM) method are given by

for k ≥ 0 with ys0 = x0. It is notable that in the update of xk+1, a momentum term is constructed based on the auxiliary sequence {ysk }, whose update is parameterized by s. The following proposition indicates that SUM reduces to the concerned three special cases by setting different values to s.

Proposition 1. SUM (10) reduces to the three variants SG (9), SHB (5) and SNAG (7) by setting s = 1 1-β , s = 0 and s = 1, respectively. Particularly, the update of SG is xk+1 = xk - α 1-β G(xk, ξk), where the step size is α 1-β .

From the above result, we can see that SHB, SNAG and SG are three variants of SUM. Moreover, the SUM view of SG implies that SG can have a larger "effective" step size (i.e., α/(1 - β)) before the gradient G(xk) than that of SHB and SNAG. We note that this is a very important observation about SG since setting a smaller effective step size for SG (e.g., the same as that in SNAG) will yield much worse convergence of training error as observed in experiments.

To facilitate the unified analysis of the stochastic momentummethods, we note that (10) implies the recursions in (13) and (14) given in the following lemma.

Lemma 1. Let pk be

and

Then for any k ≥ 0, we have

Remark: We note that a similar recursion in (13) with s = 0 and s = 1 has been observed and employed to [Ghadimi et al., 2014] for deterministic convex optimization. However, the recursion in (14) for vk (i.e., pk) is a key to our convergence analysis for non-convex optimization and importantly the generalization to any s allows us to analyze SHB, SNAG and SG in a unified framework.

Finally, we present a lemma stating the cumulative effect of updates for each iterate, which will be useful for our generalization error analysis.

Lemma 2. Given the update in (10), for any k ≥ 0 we have

Remark: The above cumulative update reduce to the following three cases for SHB (s = 0), SNAG (s = 1) and SG (s = 1/(1 - β)).

From the above cumulative update, we can see that SHB and SNAG have smaller step size for each stochastic gradient. This is the main reason that SHB and SNAG are more stable than SG, and hence yield a solution with better generalization performance. We will present a more formal analysis of generalization error later.

5 Convergence Analysis of SUM

In this section, we present the convergence results for the empirical risk minimization (2) of the SUM methods. As mentioned before, for deep learning problems the loss function ℓ(x, q) is a non-convex function, which makes finding the global optimal solution an NP-hard problem. Instead, as in many previous works we will present the convergence rates of SUM in terms of the norm of the gradient. We will present the main results first and then sketch the analysis. Detailed proofs are deferred to the supplement due to limit of space.

5.1 Main results

Theorem 1. (Convergence of SUM) Suppose f(x) is a nonconvex and L-smooth function, E[‖G(x; ξ)-∇f(x)‖2 ] ≤ σ2 and ‖∇f(x)‖ ≤ G for any x. Let update (10) run for t iterations with G(xk; ξk). By setting α = min{ 1-β 2L , C √t+1 } we have

Remark: We would like to make several remarks. (i) The assumption on the magnitude of the gradient and the variance of stochastic gradient can be simply replaced by the magnitude of the stochastic gradient, which are standard assumptions in the previous analysis of stochastic gradient method [Ghadimi and Lan, 2013]. (ii) This is the first time that the convergence rate of SHB for non-convex optimization is established. A similar convergence rate of SG and a different stochastic variant of accelerated gradient method has been established in [Ghadimi and Lan, 2013] and [Ghadimi and Lan, 2016], respectively under similar assumptions. (iii) The unified convergence makes it clear that the difference of the convergence bounds for different variants of SUM lies at the term ((1 - β)s - 1)2 , which is equal to β2 , β4 and 0 for SHB, SNAG and SG, respectively. (iv) The step size α of different variants of SUM used in the analysis of Theorem 1 is the same value.

The above result shows that the convergence upper bound of the three methods are of the same order, i.e., O(1/√t) for the gradient's square norm. In addition, when the momentum term β is large, the effect of different values of s in the term would Lβ2 ((1 - β)s - 1)2 (G2 + σ2 ) becomes marginal in contrast to the term Lσ2 (1 - β)2 in the convergence bound. This reveals that SHB and SNAG have no advantage over SG in terms of empirical risk minimization.

Below, we present a result with different step sizes α for different variants of SUM in the analysis, which sheds more insights of different methods.

Theorem 2. (Convergence of SUM) Suppose f(x) is a nonconvex and L-smooth function, E[‖G(x; ξ) - ∇f(x)‖2 ] ≤ σ2 and ‖∇f(x)‖ ≤ G for any x. Let update (10) run for t iterations with G(xk; ξk). By setting α = min{ 1-β 2L[1+((1-β)s-1)2 ] , C √t+1 } we have

Remark: The above result allows us to possibly set a larger initial value of α for SG (where s = 1/(1 - β)) and SNAG (where s = 1) than that for SHB (where s = 0). Our empirical studies for deep learning also confirms this point.

5.2 Generalization Error Analysis of SUM

In this section, we provide a unified analysis for the generalization error of the solution returned by SUM after a finite number of iterations. By employing the unified analysis, we are able to analyze the effect of the scalar s on the generalization error. Our analysis is inspired by [Hardt et al., 2016], which leverages the uniform stability of a randomized algorithm [Bousquet and Elisseeff, 2002] to bound the generalization error of multiple pass SG method. To this end, we first introduce the uniform stability and its connection with generalization error.

Let A : S → Rd denote a randomized algorithm that generates a model A(S) from the set of training samples S of size n. The uniform stability measures that how likely the prediction of the learned model on any sample q would change if one example in S is changed to a different data. In particular, let S ′ denote a set of training examples that differ from S in one example. The algorithm A is said to be ϵ-uniform stable, if

The following proposition states that the generalization error of A(S) is bounded by the uniform stability of A.

Proposition 2. (Theorem 2.2 in [Hardt et al., 2016]) For a randomized algorithm A : S → R,

The above proposition allows us to use the uniform stability of a randomized algorithm as a proxy of the generalization error. Below, we will show that SHB and SNAG are more uniform stable than SG, which exhibits that SHB and SNAG has potentially smaller generalization error than SG.

To proceed, we assume that loss function is G-Lipschitz continuous, then |ℓ(A(S), q) - ℓ(A(S ′ ), q)| ≤ G‖A(S) - A(S ′ )‖ and ϵ(A, n) ≤ supS,S ′ E[‖A(S) -A(S ′ )‖]. To analyze the uniform stability of SUM, we will assume that there are two instances of SUM starting from the same initial solution, with one running on S and the other one running on S ′ , where S and S ′ differs only at one example. Let xt = xt(S) denote the t-th iterate of the first instance and x′t = xt(S ′ ) denote the t-th iterate of the second instance. Below, we establish a result showing how ∆t = E[‖xt -x′t ‖] grows based on the unified framework in Lemma 2.

Proposition 3. Assume that ‖∇ℓ(x, q)‖2 ≤ G for any x and q and ℓ(x, q) is L-smooth w.r.t x. For two data sets S, S ′ that differs at one example, let xt and x′t denote the t-th iterates of running SUM for the empirical risk minimization on S and S′ , we have

Remark: From the above result, we can easily analyze how the value of s affects the growth of ∆t that implies the growth of the generalization error of xt. The values of ηtk for the three variants (i.e., SHB, SNAG and SG) are given by ηtk (SHB) = 1 1-β - βt-k+1 1-β , ηtk (SNAG) = 1 1-β - βt-k+2 1-β and ηtk (SG) = 1 1-β , respectively. It is obvious that ηtk (SHB) < ηtk (SNAG) < ηtk (SG). As a result, ∆t of SG grows faster than that of SNAG, and then followed by ∆t of SHB. Since the generalization error of xt is bounded by ∆t up to a constant, we can conclude that by running the same number of iterations, the generalization error of the model returned by SHB and SNAG is potentially smaller than that of SG.

More Discussion. So far, we have analyzed the convergence rate for optimizing the empirical risk and the generalization error of the learned models of different variants of SUM, which provide answers to the questions raised at the beginning except the last one (why is SNAG observed to yield improvement on the prediction performance over SHB by some studies [Sutskever et al., 2013]). Next, we show that how our analysis can shed lights on this question. In fact, the population risk of xk that is usually assessed in practice by the testing error can be decomposed into three parts, consisting of the optimization error, the generalization error, and an optimization independent term, i.e.,

where x∗ is the optimal solution to the empirical risk minimization problem. An informal analysis follows: our Theorem 2 implies that SNAG converges potentially faster than SHB in terms of the optimization error, while Proposition 3 implies that SHB has potentially smaller generalization error. If the optimization error of SNAG decreases faster than the generalization error increases comparing with SHB, then SNAG could yield a solution with a smaller population risk. However, a rigorous analysis of the optimization error is complicated by the non-convexity of the problem. In next section, we will present empirical results to corroborate and complete our theoretical analysis.

6 Empirical Studies

In this section, we present empirical results on the nonconvex optimization of deep neural networks. We train a deep convolutional neural network (CNN) for classification on two benchmark datasets, i.e., CIFAR-10 and CIFAR-100. Both datasets contain 50, 000 training images of size 32 ∗ 32 from 10 classes (CIFAR-10) or 100 classes (CIFAR-100) and 10, 000 testing images of the same size. The employed CNN consists of 3 convolutional layers and 2 fully-connected layers. Each convolutional layer is followed by a max pooling layer. The output of the last fully-connected layer is fed into a 10-class or 100-class softmax loss function. We emphasize that we do not intend to obtain the state-of-the-art prediction performance by trying different network structures and different engineering tricks, but instead focus our attention on verifying the theoretical analysis. We compare the three variants of SUM, i.e., SHB, SNAG, and SG, which corresponds to s = 0, s = 1 and s = 1/(1 - β) in (10). We fix the momentum constant β = 0.9 and the regularization parameter of weights to 0.0005. We use a mini-batch of size 128 to compute a stochastic gradient at each iteration. All three methods use the same initialization. We follow the procedure in [Krizhevsky et al., 2012] to set the step size α, i.e., initially giving a relatively large step size and and decreasing the step size by 10 times after certain number of iterations when observing the performance on testing data saturates.

Results on CIFAR-10. We first present the convergence results of different methods with the best initial step size. In particular, for the initial step size, we search in a range ({0.001, 0.002, 0.005, 0.01, 0.02, 0.05}) for different methods and select the best one that yields the fastest convergence in training error. In particular, for SHB the best initial step size is 0.001 and that for SNAG and SG is 0.01. In fact, a larger initial step size (e.g, 0.002) for SHB gives a divergent result. The training and testing error of different methods versus the number of iterations is plotted in Figure 1. This result is consistent with our convergence result in Theorem 2.

Next, we plot the performance of different methods with the same initial step size 0.001 in Figure 2. We report the training error, the testing error and their absolute difference in Figure 2(a), 2(b) and 2(c), respectively. We use the absolute difference between the training and testing error as an estimate of the generalization error. We can see that the convergence of training error of the three methods are very close, which is consistent with our theoretical result in Theorem 1. Moreover, the behavior of the absolute difference between the training and testing error is also consistent with the theoretical result in Proposition 3, i.e, SG has a larger generalization error than SNAG and SHB.

Results on CIFAR-100. We plot the training and testing error and their absolute difference of the three methods with the same initial step sizes (0.002) in Figure 2(d), 2(e) and 2(f), respectively. We observe similar trends in the training error and the generalization error, i.e., the three methods have similar performance (convergence speed) in training error but exhibit different degree of generalization error. The testing error curve shows that SNAG achieves the best prediction performance on the testing data.

Finally, we present a comparison of SUM with different values of s including variants besides SHB, SNAG and SG. In particular, we compare SUM with s ∈ {0, 0.5, 1, 2, 10}

and the same initial step size 0.002. Note that s = 0 corresponds to SHB, s = 1 corresponds to SNAG, s = 10 corresponds to SG since β = 0.9 and s = 2 corresponds to a new variant. The results on the CIFAR-100 data are shown in Figure 3. From the results, we can observe that the convergence of training error for different variants perform similarly. For the generalization error, we observe a clear trend from s = 10 to s = 0 in that the generalization error decreases.

7 Conclusion

We have developed a unified framework of stochastic momentum methods that subsumes SHB, SNAG and SG as special cases, which have been widely adopted in training deep neural network. We also analyzed convergence of the training for non-convex optimization and generalization error for learning of the unified stochastic momentum methods. The unified framework and analysis bring more insights about differences between different methods and help explain experimental results for optimization deep neural networks. In particular, the momentum term helps improve the generalization performance but not helps speed up the training process.

Acknowledgements

This work is partially supported by NSF-1545995, the Data to Decisions Cooperative Research Centre and ARC DP180100106. Most work of Y. Yan was done when he was visiting the University of Iowa.