SMG: A Shuffling Gradient-Based Method with Momentum

Most training tasks in supervised learning are boiled down to solving the following finite-sum minimization:

where $f(\cdot;i):\mathbb{R}^{d}\to\mathbb{R}$ is a given smooth and possibly nonconvex function for $i\in[n]:=\left\{1,\cdots,n\right\}$.

Problem (1) looks simple, but covers various convex and nonconvex applications in machine learning and statistical learning, including, but not limited to, logistic regression, multi-kernel learning, conditional random fields, and neural networks. Especially, (1) covers the empirical risk minimization as a special case. Solution methods for approximately solving (1) have been widely studied in the literature under different sets of assumptions. The most common approach is perhaps stochastic gradient-type (SGD) methods (Robbins & Monro, 1951; Ghadimi & Lan, 2013; Bottou et al., 2018; Nguyen et al., 2018, 2019) and their variants.

Motivation. While SGD and its variants rely on randomized sampling strategies with replacement, gradient-based methods using without-replacement strategies are often easier and faster to implement. Moreover, practical evidence (Bottou, 2009) has shown that they usually produce a faster decrease of the training loss. Randomized shuffling strategies (also viewed as sampling without replacement) allow the algorithm to use exactly one function component $f(\cdot;i)$ at each epoch compared to SGD, which has only statistical convergence guarantees (e.g., in expectation or with high probability). However, very often, the analysis of shuffling methods is more challenging than SGD due to the lack of statistical independence.

In the deterministic case, single permutation (also called shuffle once, or single shuffling) and incremental gradient methods can be considered as special cases of the shuffling gradient-based methods we study in this paper. One special shuffling strategy is randomized reshuffling, which is broadly used in practice, where we use a different random permutation at each epoch. Alternatively, in recent years, it has been shown that many gradient-based methods with momentum update can notably boost the convergence speed both in theory and practice (Nesterov, 2004; Dozat, 2016; Wang et al., 2020). These methods have been widely used in both convex and nonconvex settings, especially, in deep learning community. Remarkably, Nesterov’s accelerated method (Nesterov, 1983) has made a revolution in large-scale convex optimization in the last two decades, and has been largely exploited in nonconvex problems. The developments we have discussed here motivate us to raise the following research question:

Can we combine both shuffling strategy and momentum scheme to develop new provable gradient-based algorithms for handling (1)?

In this paper, we answer this question affirmatively by proposing a novel algorithm called Shuffling Momentum Gradient (SMG). We establish its convergence guarantees for different shuffling strategies, and in particular, randomized reshuffling strategy. We also investigate different variants of our method.

Our contribution. To this end, our contributions in this paper can be summarized as follows.

• (a)

  We develop a novel shuffling gradient-based method with momentum (Algorithm 1 in Section 2) for approximating a stationary point of the nonconvex minimization problem (1). Our algorithm covers any shuffling strategy ranging from deterministic to randomized, including incremental, single shuffling, and randomized reshuffling variants.

• (b)

  We establish the convergence of our method in the nonconvex setting and achieve the state-of-the-art $\operatorname{\mathcal{O}}\left(1/{T^{2/3}}\right)$ convergence rate under standard assumptions (i.e. the $L$-smoothness and bounded variance conditions), where $T$ is the number of epochs. For randomized reshuffling strategy, we can improve our convergence rate up to $\operatorname{\mathcal{O}}\left(1/({n^{1/3}T^{2/3})}\right)$.

• (c)

  We study different strategies for selecting learning rates (LR), including constant, diminishing, exponential, and cosine scheduled learning rates. In all cases, we prove the same convergence rate of the corresponding variants without any additional assumption.

• (d)

  When a single shuffling strategy is used, we show that a momentum strategy can be incorporated directly at each iteration of the shuffling gradient method to obtain a different variant as presented in Algorithm 2. We analyze the convergence of this algorithm and achieve the same $\operatorname{\mathcal{O}}(1/T^{2/3})$ epoch-wise convergence rate, but under a bounded gradient assumption instead of the bounded variance as for the SMG algorithm.

Our $\operatorname{\mathcal{O}}(1/T^{2/3})$ convergence rate is the best known so far for shuffling gradient-type methods in nonconvex optimization (Nguyen et al., 2020; Mishchenko et al., 2020). However, like (Mishchenko et al., 2020), our SMG method only requires a generalized bounded variance assumption (Assumption 1(c)), which is weaker and more standard than the bounded component gradient assumption used in existing works. Algorithm 2 uses the same set of assumptions as in (Nguyen et al., 2020) to achieve the same rate, but has a momentum update. For the randomized reshuffling strategy, our $\operatorname{\mathcal{O}}\left(1/({n^{1/3}T^{2/3})}\right)$ convergence rate also matches the rate of the without-momentum algorithm in (Mishchenko et al., 2020). It leads to the total of iterations $nT=\operatorname{\mathcal{O}}(\sqrt{n}\varepsilon^{-3})$.

We emphasize that, in many existing momentum variants, the momentum $m_{i}^{(t)}$ is updated recursively at each iteration as $m_{i+1}^{(t)}:=\beta m_{i}^{(t)}+(1-\beta)g_{i}^{(t)}$ for a given weight $\beta\in(0,1)$. This update shows that the momentum $m_{i+1}^{(t)}$ incorporates all the past gradient terms $g_{i}^{(t)},g_{i-1}^{(t)},g_{i-1}^{(t)},\dots,g_{0}^{(t)}$ with exponential decay weights $1,\beta,\beta^{2},\dots,\beta^{i}$, respectively. However, in shuffling methods, the convergence guarantee is often obtained in epoch. Based on this observation, we modify the classical momentum update in the shuffling method as shown in Algorithm 1. More specifically, the momentum term $m_{0}^{(t)}$ is fixed at the beginning of each epoch, and an auxiliary sequence $\{v_{i}^{(t)}\}$ is introduced to keep track of the gradient average to update the momentum term in the next epoch. This modification makes Algorithm 1 fundamentally different from existing momentum-based methods. This new algorithm still achieves $\operatorname{\mathcal{O}}(1/T^{2/3})$ epoch-wise convergence rate under standard assumptions. To the best of our knowledge, our work is the first analyzing convergence rate guarantees of shuffling-type gradient methods with momentum under standard assumptions.

Besides Algorithm 1, we also exploit recent advanced strategies for selecting learning rates, including exponential and cosine scheduled learning rates. These two strategies have shown state-of-the-art performance in practice (Smith, 2017; Loshchilov & Hutter, 2017; Li et al., 2020). Therefore, it is worth incorporating them in shuffling methods.

Related work. Let us briefly review the most related works to our methods studied in this paper.

Shuffling gradient-based methods. Shuffling gradient-type methods for solving (1) have been widely studied in the literature in recent years (Bottou, 2009; Gürbüzbalaban et al., 2019; Shamir, 2016; Haochen & Sra, 2019; Nguyen et al., 2020) for both convex and nonconvex settings. It was empirically investigated in a short note (Bottou, 2009) and also discussed in (Bottou, 2012). These methods have also been implemented in several software packages such as TensorFlow and PyTorch, broadly used in machine learning (Abadi et al., 2015; Paszke et al., 2019).

In the strongly convex case, shuffling methods have been extensively studied in (Ahn et al., 2020; Gürbüzbalaban et al., 2019; Haochen & Sra, 2019; Safran & Shamir, 2020; Nagaraj et al., 2019; Rajput et al., 2020; Nguyen et al., 2020; Mishchenko et al., 2020) under different assumptions. The best known convergence rate in this case is $\operatorname{\mathcal{O}}(1/(nT)^{2}+1/(nT^{3}))$, which matches the lower bound rate studied in (Safran & Shamir, 2020) up to some constant factor. Most results in the convex case are for the incremental gradient variant, which are studied in (Nedic & Bertsekas, 2001; Nedić & Bertsekas, 2001). Convergence results of shuffling methods on the general convex case are investigated in (Shamir, 2016; Mishchenko et al., 2020), where (Mishchenko et al., 2020) provides a unified approach to cover different settings. The authors in (Ying et al., 2017) combine a randomized shuffling and a variance reduction technique (e.g., SAGA (Defazio et al., 2014) and SVRG (Johnson & Zhang, 2013)) to develop a new variant. They show a linear convergence rate for strongly convex problems but using an energy function, which is unclear how to convert it into known convergence criteria.

In the nonconvex case, (Nguyen et al., 2020) first shows $\operatorname{\mathcal{O}}(1/T^{2/3})$ convergence rate for a general class of shuffling gradient methods under the $L$-smoothness and bounded gradient assumptions on (1). This analysis is then extended in (Mishchenko et al., 2020) to a more relaxed assumption. The authors in (Meng et al., 2019) study different distributed SGD variants with shuffling for strongly convex, general convex, and nonconvex problems. An incremental gradient method for weakly convex problems is investigated in (Li et al., 2019), where the authors show $\operatorname{\mathcal{O}}(1/T^{1/2})$ convergence rate as in standard SGD. To the best of our knowledge, the best known rate of shuffling gradient methods for the nonconvex case under standard assumptions is $\operatorname{\mathcal{O}}(1/T^{2/3})$ as shown in (Nguyen et al., 2020; Mishchenko et al., 2020).

Our Algorithm 1 developed in this paper is a nontrivial momentum variant of the general shuffling method in (Nguyen et al., 2020) but our analysis uses a standard bounded variance assumption instead of bounded gradient one.

Momentum-based methods. Gradient methods with momentum were studied in early works for convex problems such as heavy-ball, inertial, and Nesterov’s accelerated gradient methods (Polyak, 1964; Nesterov, 2004). Nesterov’s accelerated method is the most influent scheme and achieves optimal convergence rate for convex problems. While momentum-based methods are not yet known to improve theoretical convergence rates in the nonconvex setting, they show significantly encouraging performance in practice (Dozat, 2016; Wang et al., 2020), especially in the deep learning community. However, the momentum strategy has not yet been exploited in shuffling methods.

Adaptive learning rate schemes. Gradient-type methods with adaptive learning rates such as AdaGrad (Duchi et al., 2011) and Adam (Kingma & Ba, 2014) have shown state-of-the-art performance in several optimization applications. Recently, many adaptive schemes for learning rates have been proposed such as diminishing (Nguyen et al., 2020), exponential scheduled (Li et al., 2020), and cosine scheduled (Smith, 2017; Loshchilov & Hutter, 2017). In (Li et al., 2020), the authors analyze convergence guarantees for the exponential and cosine learning rates in SGD. These adaptive learning rates have also been empirically studied in the literature, especially in machine learning tasks. Their convergence guarantees have also been investigated accordingly under certain assumptions.

Although the analyses for adaptive learning rates are generally non-trivial, our theoretical results in this paper are flexible enough to cover various different learning rates. However, we only exploit the diminishing, exponential scheduled, and cosine scheduled schemes for our shuffling methods with momentum. We establish that in the last two cases, our algorithm still achieves state-of-the-art $\operatorname{\mathcal{O}}(T^{-2/3})$ (and possibly up to $\operatorname{\mathcal{O}}(n^{-1/3}T^{-2/3})$) epoch-wise rates.

Content. The rest of this paper is organized as follows. Section 2 describes our novel method, Shuffling Momentum Gradient (Algorithm 1). Section 3 investigates its convergence rate under different shuffling-type strategies and different learning rates. Section 4 proposes an algorithm with traditional momentum update for single shuffling strategy. Section 5 presents our numerical simulations. Due to space limit, the convergence analysis of our methods, all technical proofs, and additional experiments are deferred to the Supplementary Document (Supp. Doc.).

In this section, we describe our new shuffling gradient algorithm with momentum in Algorithm 1. We also compare our algorithm with existing methods and discuss its per-iteration complexity and possible modifications, e.g., mini-batch.

Clearly, if $\beta=0$, then Algorithm 1 reduces to the standard shuffling gradient method as in (Nguyen et al., 2020; Mishchenko et al., 2020). Since each inner iteration of our method uses one component $f(\cdot,i)$, we use $\eta_{i}^{(t)}=\frac{\eta_{t}}{n}$ in Algorithm 1 to make it consistent with one full gradient, which consists of $n$ gradient components. This form looks different from the learning rates used in previous work for SGD (Shamir, 2016; Mishchenko et al., 2020), however, it does not necessary make our learning rate smaller. In the same order of training samples $n$, our learning rate matches the one in (Mishchenko et al., 2020) as well as the state-of-the-art complexity results. In fact, the detailed learning rate $\eta_{t}$ used in Algorithm 1 will be specified in Section 3.

Comparison. Unlike existing momentum methods where $m_{i}^{(t)}$ in (2) is updated recursively as $m_{i+1}^{(t)}:=\beta m_{i}^{(t)}+(1-\beta)g_{i}^{(t)}$ for $\beta\in(0,1)$, we instead fix the first term $m_{0}^{(i)}$ in (2) at each epoch. It is only updated at the end of each epoch by averaging all the gradient components $\{g_{i}^{(t)}\}_{i=0}^{n-1}$ evaluated in such an epoch. To avoid storing these gradients, we introduce an auxiliary variable $v_{i}^{(t)}$ to keep track of the gradient average. Here, we fix the weight $\beta$ for the sake of our analysis, but it is possible to make $\beta$ adaptive.

Our new method is based on the following observations. First, while SGD generally uses an unbiased estimator for the full gradient, shuffling gradient methods often do not have such a nice property, making them difficult to analyze. Due to this fact, updating the momentum at each inner iteration does not seem preferable since it could make the estimator deviate from the true gradient. Therefore, we consider updating the momentum after each epoch. Second, unlike the traditional momentum with exponential decay weights, our momentum term $m_{0}^{(t)}$ is an equal-weighted average of all the past gradients in an epoch, but evaluated at different points $w_{i}^{(t-1)}$ in the inner loop. Based on these observations, the momentum term should aid the full gradient $\nabla{F}$ instead of its component $\nabla{f}(\cdot;i)$, leading to the update at the end of each epoch.

It is also worth noting that our algorithm is fundamentally different from variance reduction methods. For instance, SVRG and SARAH variants require evaluating full gradient at each snapshot point while our method always uses single component gradient. Hence, our algorithm does not require a full gradient evaluation at each outer iteration, and our momentum term does not require full gradient of $f$.

When the learning rate $\eta_{i}^{(t)}$ is fixed at each epoch as $\eta_{i}^{(t)}:=\frac{\eta_{t}}{n}$, where $\eta_{t}>0$, we can derive from (2) that

where $e_{t}:=\tilde{w}_{t-1}-\tilde{w}_{t-2}+\beta\eta_{t-1}\tilde{m}_{t-2}$. Here, $e_{t}$ plays a role as a momentum or an inertial term, but it is different from the usual momentum term. However, we still name Algorithm 1 the Shuffling Momentum Gradient since it is inspired by momentum methods.

Per-iteration complexity. The per-iteration complexity of Algorithm 1 is almost the same as in standard shuffling gradient schemes, see, e.g., (Shamir, 2016). It needs to store two additional vectors $m_{i+1}^{(t)}$ and $v_{i}^{(t)}$, and performs two vector additions and three scalar-vector multiplications. Such additional costs are very mild.

Shuffling strategies. Our convergence guarantee in Section 3 holds for any permutation $\pi^{(t)}$ of $\{1,2,\cdots,n\}$, including deterministic and randomized ones. Therefore, our method covers any shuffling strategy, including incremental, single shuffling, and randomized reshuffling variants as special cases. We highlight that our convergence result for the randomized reshuffling variant is significantly better than the general ones as we will explain in detail in Section 3.

Mini-batch. Algorithm 1 also works with mini-batches, and our theory can be slightly adapted to establish the same convergence rate for mini-batch variants. However, it remains unclear to us how mini-batches can improve the convergence rate guarantee of Algorithm 1 in the general case where the shuffling strategy is not specified.

We analyze the convergence of Algorithm 1 under standard assumptions, which is organized as follows.

In this section, we modify the single-shuffling gradient method by directly incorporating a momentum term at each iteration. We prove for the first time that this variant still achieves state-of-the-art convergence rate guarantee under the smoothness and bounded gradient (i.e. Assumption 2) assumptions as in existing shuffling methods. This variant, though somewhat special, also covers an incremental gradient method with momentum as a special case.

Our new momentum algorithm using single shuffling strategy for solving (1) is presented in Algorithm 2.

Besides fixing the permutation $\pi$, Algorithm 2 is different from Algorithm 1 at the point of updating $m_{i}^{(t)}$. Here, $m_{i+1}^{(t)}$ is updated from $m_{i}^{(t)}$ instead of $m_{0}^{(t)}$ as in Algorithm 1. In addition, we update $\tilde{m}_{t}:=m_{n}^{(t)}$ instead of the epoch gradient average. Note that we can write the main update of Algorithm 2 as

which exactly reduces to existing momentum updates.

Incremental gradient method with momentum. If we choose $\pi:=[n]$, then we obtain the well-known incremental gradient variant, but with momentum. Hence, Algorithm 2 is still new compared to the standard incremental gradient algorithm (Bertsekas, 2011). To prove convergence of Algorithm 2, we replace Assumption 1(c) by the following:

Now, we state the convergence of Algorithm 2 in the following theorem as our third main result.

Compared to (9) in Theorem 2, (10) strongly depends on the weight $\beta$ as $\frac{4\beta^{n}G^{2}}{1-\beta^{n}}$ appears on the right-hand side of (10). Hence, $\beta$ must be chosen in a specific form to obtain desired convergence rate.

Theorem 3 provides a key bound to derive concrete convergence rates in the following two corollaries.

With a constant learning rate as in Corollary 7, the convergence rate of Algorithm 2 is

which depends on $L[F(\tilde{w}_{0})-F_{*}]$, $\|\nabla{F}(\tilde{w}_{0})\|^{2}$, and $G^{2}$, and is slightly different from Corollary 1.

In order to examine our algorithms, we present two numerical experiments for different nonconvex problems and compare them with some state-of-the-art SGD-type and shuffling gradient methods.

We have developed two new shuffling gradient algorithms with momentum for solving nonconvex finite-sum minimization problems. Our Algorithm 1 is novel and can work with any shuffling strategy, while Algorithm 2 is similar to existing momentum methods using single shuffling strategy. Our methods achieve the state-of-the-art $\operatorname{\mathcal{O}}(1/T^{2/3})$ convergence rate under standard assumptions using different learning rates and shuffling strategies. When a randomized reshuffling strategy is exploited for Algorithm 1, we can further improve our rates by a fraction $n^{1/3}$ of the data size $n$, matching the best known results in the non-momentum shuffling method. Our numerical results also show encouraging performance of the new methods.

We believe that our analysis framework can be extended to study non-asymptotic rates for some recent adaptive SGD schemes such as Adam (Kingma & Ba, 2014) and AdaGrad (Duchi et al., 2011) as well as variance-reduced methods such as SARAH (Nguyen et al., 2017) under shuffling strategies. It is also interesting to extend our framework to other problems such as minimax and federated learning. We will further investigate these opportunities in the future.

The authors would like to thank the reviewers for their useful comments and suggestions which helped to improve the exposition in this paper. The work of Q. Tran-Dinh has partly been supported by the Office of Naval Research under Grant No. ONR-N00014-20-1-2088 (2020-2023).