Year: 2021

Abstract:
Many popular learning-rate schedules for deep neural networks combine a decaying trend with local perturbations that attempt to escape saddle points and bad local minima. We derive convergence guarantees for bandwidth-based step-sizes, a general class of learning-rates that are allowed to vary in a banded region. This framework includes many popular cyclic and non-monotonic step-sizes for which no theoretical guarantees were previously known. We provide worst-case guarantees for SGD on smooth non-convex problems under several bandwidth-based step sizes, including stagewise 1/ √ t and the popular step-decay (“constant and then drop by a constant”), which is also shown to be optimal. Moreover, we show that its momentum variant converges as fast as SGD with the bandwidth-based step-decay step-size. Finally, we propose novel step-size schemes in the bandwidth-based family and verify their efficiency on several deep neural network training tasks.

text: √ t and the popular step-decay (“constant and then
drop by a constant”), which is also shown to be optimal. Moreover, we show that its momentum variant converges as fast as SGD with the bandwidth-based step-decay step-size. Finally, we propose novel step-size schemes in the bandwidth-based family and verify their efficiency on several deep neural network training tasks.

heading: 1 INTRODUCTION
text: Stochastic gradient methods including stochastic gradient descent (SGD) (Robbins and Monro, 1951) and its accelerated variants (e.g., SGD with momentum (Polyak, 1964; Sutskever et al., 2013)) have become the algorithmic workhorse in much of machine learning. The step-size (learning rate) is the most important hyper-parameter for controlling the speed at which gradient-based methods converge to stationarity. For problems with multiple local minima, the step-size also affects which local optimum the optimization process converges to. It therefore needs to be both well-designed and well-tuned to make SGD and its variants effective in practice.
In the deep learning literature, cyclical step-sizes (Loshchilov and Hutter, 2017; Smith, 2017) and non-monotonic schedules (Keskar and Saon, 2015; An et al., 2017; Seong et al., 2018; Loizou et al., 2021) have attracted strong recent interest, with significant benefits for non-convex problems with poor local minima or saddle points (Seong et al., 2018). Popular cyclical schedules include the cosine step-size (cosine with restart) (Loshchilov and Hutter, 2017) and the triangular policy (Smith, 2017), which have become the default choices in some deep learning libraries, e.g., PyTorch and TensorFlow (cf. lr scheduler.CyclicLR and CosineAnnealingLR). However, non-monotonic policies are much more complex to analyze than decaying ones, and theoretical results for these non-monotonic policies are scarce. This motivates us to focus on a bandwidth step-size framework, in which
mδ(t) ≤ ηt ≤Mδ(t) (1) for some boundary function δ(t) and positive constants m and M . This framework allows for nonmonotonic step-sizes and covers most of the situations discussed above. In particular, it includes the cosine (Loshchilov and Hutter, 2017), triangular (Smith, 2017), sine wave (An et al., 2017) step-sizes as special cases. The framework provides a uniform convergence rate guarantee for all step-size policies which remain in the band (1). This gives a lot of freedom to design novel step-sizes schedules with improved practical performance without loosing track of their theoretical convergence guarantee.
The generic bandwidth framework has recently been proposed by Wang and Yuan (2021), but they only analyzed strongly convex problems. We believe that more significant potential lies in the non-convex regime. For non-convex problems, non-monotonic step-sizes have distinct advantages, helping iterates to escape local minima and producing final iterates of high quality. In the paper, we
demonstrate this point on both a simple toy example and on large-scale neural network training tasks. Our main contribution is a sequence of non-asymptotic convergence results for the bandwidth stepsize on non-convex optimization problems, based on the popular “constant and then drop” step-size schedules (Krizhevsky et al., 2012; He et al., 2016; Hazan and Kale, 2014; Ge et al., 2019; Wang et al., 2021). This allows non-monotonic variations both within each (inner) stage and between stages.

heading: 1.1 CONTRIBUTIONS
text: Inspired by the strong potential of non-monotonic step-size schedules demonstrated above, we extend the bandwidth-based step-size framework to “constant and then drop” (multi-stage) profiles, where the bands stay constant throughout each stage and drops between stages. We provide convergence guarantees for both SGD and its momentum variant (SGDM) on non-convex problems. Specifically,
• We establish worst-case theoretical guarantees for SGD with bandwidth step-size on smooth nonconvex problems. We (i) derive an optimal rate for SGD under a bandwidth step-size with δ(t) = 1/ √ t; (ii) and achieve optimal and near-optimal rates for step-decay (constant
and then drop by a constant), improving the results by Wang et al. (2021).
• We provide worst-case theoretical guarantees for SGDM with bandwidth-based step-decay step-size in the smooth non-convex setting. To the best of our knowledge, these are the first results that provide optimal (Theorem 4.3) and near-optimal (Theorem 4.2) results for momentum with step-decay step-sizes. Moreover, our results significantly improve the convergence results from Liu et al. (2020) (see Remark 4.4).
• Our analysis results also provide state-of-the-art theoretical guarantees for cosine (Loshchilov and Hutter, 2017) and triangular (Smith, 2017) step-sizes if their boundary functions are within our bands. Especially, we improve the result of Li et al. (2021) for cosine step-size and achieve a state-of-art rate (see Remark 3.4). Moreover, our results first provide the convergence guarantees for triangular step-size (Smith, 2017).
• We propose novel, possibly non-monotonic, step-size schedules (e.g., step-decay with linearmode and cosine-mode) based on the bandwidth-based framework and demonstrate their efficacy on several large-scale neural network training tasks.

heading: 1.2 RELATED WORK
text: This subsection reviews the theoretical development of the SGD algorithm and its momentum variant in the smooth non-convex setting, with a special focus on different step-size policies.
SGD for nonconvex problems The first non-asymptotic convergence of SGD to a stationary point of a general smooth non-convex function was established in Ghadimi and Lan (2013). The authors proved that a constant step-size O(1/ √ T ) attains a convergence rate of O(1/ √ T ), where T is the iteration budget. To the best of our knowledge, this rate is not improvable and was proven to be tight up to a constant without additional assumptions (Drori and Shamir, 2020). For the 1/ √ t decay step-size, an O(lnT/ √ T ) rate can be easily obtained from (Ghadimi and Lan, 2013). This rate can be improved to the optimal by selecting a random iterate using weights proportional to the inverse of the step-size (Wang et al., 2021). The sampling rule in (Wang et al., 2021) depends on the step-size and is easily applicable to different step-size policies. Thus, in this paper, we choose a similar sampling rule as (Wang et al., 2021) to favor the later iterates when selecting the output for SGD and its momentum variant.
Step-decay step-sizes Recently, the theoretical performance of step-decay or stagewise strategies has attracted an increasing attention due to their excellent practical performance (Yuan et al., 2019; Ge et al., 2019; Chen et al., 2019; Li et al., 2021; Wang et al., 2021). For a class of least-squares problems, Ge et al. (2019) established a near-optimalO(lnT/T ) rate for the step-decay step-size (cut by 2 every T/ log2(T ) iterates) and showed that step-decay can perform better than the polynomial decay step-size. Stochastic optimization methods with stagewise step-sizes decaying as 1/t were analyzed in Chen et al. (2019). A near-optimal rate for the continuous version of step-decay, called exp-decay, as well as for cosine decay step-sizes under the Polyak-Lójasiewicz (PL) condition and a general smooth assumption were established in Li et al. (2021). However, in the smooth case, to achieve such results for exponential and cosine decay step-sizes, the initial step-size is required to
be bounded by O(1/ √ T ). This is obviously impractical when the number of iterations T is large. Near-optimal rates (up to lnT ) of SGD with step-decay step-size in several general settings including strongly convex, convex and smooth (non-convex) problems were proved in Wang et al. (2021). They also removed the restriction on the initial step-size for exponential decay step-sizes. Empirical evidences have been given in (Wang and Yuan, 2021) that bandwidth-based strategies can improve the performance of the step-decay step-size on some large scale neural network tasks. However, no theoretical guarantees for non-convex problems were given.
SGD with momentum on nonconvex problems The momentum variant of SGD (SGDM) has been widely used in deep neural networks (Krizhevsky et al., 2012; Sutskever et al., 2013; He et al., 2016; Zagoruyko and Komodakis, 2016). Due to its practical success on neural networks, its theoretical performance is now attracting a lot of interest, especially for nonconvex problems (Yan et al., 2018; Gadat et al., 2018; Chen et al., 2019; Gitman et al., 2019; Mai and Johansson, 2020; Liu et al., 2020; Defazio, 2020). Under the assumption of bounded gradients, Yan et al. (2018) proposed a unified analysis framework for stochastic momentum methods and proved an optimal O(1/ √ T ) rate under constant step-sizes. A similar result for the Nesterov-accelerated variant was established in Ghadimi and Lan (2016). However, studies related to the multi-stage performance of SGD with momentum is lacking and far from being complete. Reference Chen et al. (2019) considers a momentum method with a stagewise step-size, but the method is a proximal point algorithm with extra averaging between stages, and not the widely used momentum SGD considered here. More recently, Liu et al. (2020) established the convergence for multi-stage SGDM and provided empirical evidence to show that multi-stage SGDM is faster. However, their results require an inverse relationship between stage length and step-size which limits the initial stage length or step-size. A detailed comparison with (Liu et al., 2020) will be given in Section 4 (see Remark 4.4).
Organization: The rest of this paper is organized as follows. Notations and basic definitions are introduced in Section 2. Our novel theoretical results for SGD and its momentum variant (SGDM) under bandwidth-based step-sizes are introduced in Sections 3 and 4, respectively. Numerical experiments are presented and reported in Section 5. Finally, Section 6 concludes the paper.

heading: 2 PROBLEM SET-UP
text: We study the following, possibly non-convex, stochastic optimization problem
min x∈Rd
f(x) = Eξ∼Ξ[f(x; ξ)] (2)
where ξ is a random variable drawn from some (unknown) probability distribution Ξ and f(x; ξ) is the instantaneous loss function over the variable x ∈ Rd. We consider stochastic gradient methods that generate iterates xt according to
xt+1 = xt − ηtdt (3)
where ηt is the step-size and dt the search direction (e.g., dt = ∇f(xt; ξ) for SGD). We assume that there are constants m > 0 and M ≥ m, and two functions n(t) and δ(t): R→ R such that such that
ηt = n(t)δ(t), ∀ t ≥ 1, where n(t) ∈ [m,M ] and δ(t) is monotonically decreasing function satisfying δ(1) = 1. Note that even though the boundary function δ(t) is monotonic, the step-size itself is not restricted to be. Throughout the paper, we make the following assumptions: Assumption 1. The loss function f satisfies ‖∇f(x)−∇f(y)‖ ≤ L‖x−y‖ for every x, y ∈ dom (f). Assumption 2. For any input vector x, the stochastic gradient oracle O returns a vector g such that (a) E[‖g −∇f(x)‖2] ≤ ρ ‖∇f(x)‖2 + σ where ρ ≥ 0 and σ ≥ 0; (b) E[‖g‖2] ≤ G2.

heading: 3 NON-ASYMPTOTIC CONVERGENCE OF SGD WITH BANDWIDTH-BASED STEP-SIZE
text: In this section, we provide the first non-asymptotic convergence guarantees for SGD with bandwidthbased step-sizes on smooth non-convex problems. The results consider a general family of bandwidthbased step-sizes which includes the classical multi-stage SGD as a special case.
Algorithm 1 SGD with Bandwidth-based Step-Size
1: Input: initial point x11, # iterations T , # stages N , stage length {St} N t=1 such that ∑N t=1 St = T ,
the sequences {δ(t)}Nt=1 and { {n(t, i)}Sti=1 }N t=1 ∈ [m,M ] with 0 < m ≤M
2: for t = 1 : N do 3: for i = 1 : St do 4: Query a stochastic gradient oracleO at xti to get a vector gti such that E[gti | F ti ] = ∇f(xti)1 5: Update step-size ηti = n(t, i)δ(t) 6: xti+1 = x t i − ηtigti 7: end for 8: xt+11 = x t St+1
9: end for 10: Return: x̂T is uniformly chosen from { xt ∗ 1 , x t∗ 2 , · · · , xt ∗
St∗
} , where the integer t∗ is chosen from
{1, 2, · · · , N} with probability Pt = δ−1(t)/( ∑N l=1 δ −1(l))
Algorithm 1 details our bandwidth-based version of the popular “constant and then drop” policy for SGD. Here, the boundary function δ(t) is adjusted in an outer stage, and the length of each stage St is allowed to vary. Similar to Wang et al. (2021), the output distribution depends on the inverse of δ(t), hence puts more weight on the final iterates. By considering specific combinations of δ(t) and St, this framework allows us to analyze several important multi-stage SGD algorithms, including those with constant, polynomial-decay and step-decay step-sizes. For example, we consider the step-decay step-size by letting n(t, i) = m and δ(t) = 1/αt−1 where m denotes its initial step-size and α > 1. Many interesting results on polynomial-decay step-size (e.g., δ(t) = 1/ √ t, we called it 1/ √ t-band) are given in Appendix A.

heading: 3.1 CONVERGENCE UNDER BANDWIDTH STEP-DECAY STEP-SIZE
text: Another important step-size is Step-Decay (“constant and then drop by a constant”), which is popular and widely used in practice, e.g. for neural network training (Krizhevsky et al., 2012; He et al., 2016). In this subsection, we analyze bandwidth step-sizes that include step-decay as a special case.
For Step-Decay, the stage length St is typically a hyper-parameter selected by experience. We first analyze a bandwidth version of the algorithm analyzed in [Theorem 3.2](Wang et al., 2021), namely Algorithm 1 with N = b(logα T )/2c outer loops where α > 1, each with a constant length of St = d2T/ logα T e. The logarithmic dependence of N on T leads to a small number of stages in practice, and was demonstrated to perform well in deep neural network tasks (Wang et al., 2021). Theorem 3.1. Under Assumptions 1 and 2(a), and assume that there exists a constant ∆0 > 0 such that E[f(xt1)− f∗] ≤ ∆0 for each t ≥ 1 where f∗ = min f(x), if we run Algorithm 1 with T > α2, ηti ≤ 1/((ρ + 1)L), N = b(logα T )/2c, St = d2T/ logα T e, and δ(t) = 1/αt−1 for 1 ≤ t ≤ N , where α > 1 then
E[‖∇f(x̂T )‖2] ≤ (
∆0 2αm + αM2Lσ 2m
) (α− 1)
lnα · lnT√ T − α .
Theorem 3.1 establishes a near-optimal (up to lnT ) rate for the step-decay bandwidth scheme which matches the result achieved at its boundaries i.e., ηti = mδ(t) or η t i = Mδ(t) (Wang et al., 2021). As the next theorem shows, this guarantee can be improved by appropriate tuning of the stage length St. Remark 3.2. (Justification of uniformly bounds on the function values) In Theorem 3.1, we require that the expectation of the function value at each outer iterate E[f(xt1)] is uniformly upper bounded. As shown by Shi et al. (2020), the function values at the iterates of SGD can be controlled (bounded) by the initial state provided the step-size is bounded by 1/L. So the assumption is fair if the initial state is settled. Nevertheless, this assumption (or its stronger version that the objective function is bounded) is commonly used or implied in optimization (Hazan et al., 2015; Xu et al., 2019b; 2020; 2019a) and statistic machine learning (Vapnik, 1998; Cortes et al., 2019) , and it has never been violated in our numerical experiments.
1We use F ti to denote σ-algebra formed by all the random information before current iterate xti and xti ∈ F ti .
Theorem 3.3. Under Assumptions 1 and 2(a), and assume that there exists a constant ∆0 > 0 such that E[f(xt1)− f∗] ≤ ∆0 for each t ≥ 1, if we run Algorithm 1 with T > α2, ηti ≤ 1/((ρ+ 1)L), S0 = √ T , St = dS0α(t−1)e and δ(t) = 1/αt−1 where α > 1, then
E[‖∇f(x̂T )‖2] ≤ α+ 1
α− 1 ( 2∆0 m + M2Lσ m ) 1√ T +O ( 1 T ) .
Optimal rate for step-decay step-size The theorem shows that if the stage length St increases exponentially, and the length of the first stage is set appropriately, then we can achieve an optimal O(1/ √ T ) rate for the bandwidth step-decay step-size in the non-convex case. If M = m, which means that the bandwidth scheme degenerates to the step-decay type step-size, Theorem 3.3 removes the logarithmic term present in the results of Wang et al. (2021). To the best of our knowledge, this is the first result that demonstrates that vanilla SGD with step-decay step-sizes can achieve the optimal rate for general non-convex problems. The numerical performance of the two step-size schedules in Theorems 3.1 and 3.3 are reported in Figure 4.
Benefits of Theorems 3.5 vs the references of Hazan and Kale (2014); Yuan et al. (2019) Another commonly used step-decay scheme in theory which halves the step-size after each stage and then doubles the length of each stage (e.g., (Hazan and Kale, 2014; Yuan et al., 2019)). In Hazan and Kale (2014), which considers strongly convex problem, the initial stage is very short, S1 = 4, while the analysis in Yuan et al. (2019) for PL functions use an inverse relation between stage length and step-size, which means that a longer initial stage length requires a smaller stepsize. In contrast to these references, Theorem 3.3 considers a step-decay with a long first stage, S1 = d √ T e, which allows us to benefit from a large constant step-size for more iterations.
Remark 3.4. (Guarantees for cyclical step-sizes) In (Loshchilov and Hutter, 2017), the authors decay the step-size with cosine annealing and use ηtmin < η t max to control the range of the step-size. If mδ(t) ≤ ηtmin, ηtmax ≤Mδ(t), then our results provide convergence guarantees for their step-size. To achieve a near-optimal rate, Li et al. (2021, Theorem 4) need to use an initial step-size that is smaller than O(1/ √ T ) which is obviously impractical. In contrast, we allow the cosine step-size to start from a relatively large step-size and then gradually decay (see Theorem 3.1) and also improve the convergence rate to be optimal (Theorem 3.3).

heading: A triangular cyclical step-size is proposed by Smith (2017) which is varied around the two boundaries that drop by a constant after a few iterations. Our analysis first provides theoretical guarantees (e.g.,
text: Theorems 3.1 and 3.3) also for this step-size. The details are shown in Appendix D.2.
The bandwidth step-sizes we consider above are independent on the random information. In recent years, some non-monotonic step-sizes have been proposed that are dependent on the current random information, e.g., the trust-region-ish algorithm (Curtis et al., 2019) and stochastic Polayk stepsize (Loizou et al., 2021). We provide some interesting results for these step-sizes (see Lemma B.2 in Appendix B).

heading: 4 NON-ASYMPTOTIC CONVERGENCE OF SGDM UNDER BANDWIDTH-BASED STEP-SIZE
text: In this section, we establish the first non-asymptotic convergence properties of SGD with momentum (SGDM) under the bandwidth-based step-size on smooth nonconvex problems.
In this scheme, the inner iterations in Step 6 of Algorithm 1 are essentially replaced by
vti+1 = βv t i + (1− β)gti (4)
xti+1 = x t i − ηtivti+1 (5)
for β ∈ (0, 1). We refer to Algorithm 2 in Appendix C for a more detailed description. As in many studies of momentum-methods (e.g. Ghadimi et al. (2015); Yan et al. (2018); Liu et al. (2020); Mai and Johansson (2020)), we establish an iterate relationship of the form E[W t+1] ≤ E[W t] − c0ηE[et] + c1η2, where W t is a Lyapunov function, et is a performance measure (here, et = ‖∇f(·)‖2), η is the step-size and c0 and c1 are constants. However, due to the time-dependent and possibly non-monotonic bandwidth-based step-size, we cannot use the Lyapunov functions suggested in Yan et al. (2018); Liu et al. (2020) but rely on the following non-trivial construction: Lemma 4.1. Suppose that Assumption 1 and Assumption 2(b) hold. Let zti = (1−β)−1(xti−βxti−1) and assume that there exists a constant ∆0 such that E[f(xti) − f∗] ≤ ∆0 for t, i ≥ 1 and the step-size in each stage is monotonically decreasing. Define the function W ti+1
W ti+1 = f(zti+1)− f∗ ηti + r ∥∥xti+1 − xti∥∥2 ηti + 2r[f(xti+1)− f∗],
where r = βL2(1−β2)(1−β)2 . Then, if η t i ≤ 1/L, for any t and i ≥ 2, we have E[W ti+1 | F ti ] ≤W ti +A1 ( 1
ηti − 1 ηti−1
) − ∥∥∇f(xti)∥∥2 + ηti ·B1G2. (6)
where A1 = β∆01−β + ∆z + rG2 L2 , B1 = r(1− β)(2− β) + L 2(1−β)2 , and ∆z = ∆0 1−β + βG2 2(1−β)2L .
Note that even though the step-size is assumed to be monotonically decreasing in each stage, it may be increased between stages, leading to a globally non-monotonic step-size. The proposed bandwidthbased step-sizes (e.g., step-decay with linear or cosine modes) in the numerical experiments and the cosine annealing policy proposed in (Loshchilov and Hutter, 2017) all satisfy this condition. Note that, unlike (Mai and Johansson, 2020; Liu et al., 2020), the momentum parameter β does not rely on the step-size, but can be chosen freely in the interval (0, 1). In particular, our analysis supports the common choice of β = 0.9 used as default in many deep learning libraries (Krizhevsky et al., 2012; He et al., 2016). Similar to Remark 3.2, the function value of the iterates for momentum can also be controlled (bounded) by the initial state given ηti ≤ 1/L; see (Shi, 2021). Therefore, we believe our assumptions are reasonable.
If we restrict the analysis to a single stage, N = 1, the lemma allows to recover the optimal rate for SGDM under the step-size ηti = η0/ √ T (Yan et al., 2018; Mai and Johansson, 2020; Liu et al., 2020;
Defazio, 2020) and to prove, for the first time, an optimal O(1/ √ T ) rate for SGDM under the 1/ √ i stepsize. These results are formalized in Appendix D.1.

heading: 4.1 CONVERGENCE OF SGDM FOR BANDWIDTH STEP-DECAY STEP-SIZE
text: We now show the convergence complexity of SGDM with the bandwidth step-decay step-size. Here step-decay means that the bandwidth limits are divided by a constant after some iterations.
We first consider the total number of iterations T to be given, the stage length St to be constant, and the number of stages N as a hyper-parameter. Theorem 4.2. Assume the same setting as Lemma 4.1. If given the total number of iterations T ≥ 1, N ≥ 1, St = S = dT/Ne, δ(t) = 1/αt−1 for each 1 ≤ t ≤ N and α > 1, then
E[‖∇f(x̂T )‖2] ≤ W 11 ·N TαN−1 + (αC0 + C2) · N T + (∆z + C1) m · Nα N T +MB1G 2 · N αN−1 (7)
where C0 = r(G 2 L + 2∆0), C1 = A1 + ∆z + ∆0 1−β , and C2 = C0 + A2G 2, A2 = 1 + β2(1−β)2 , and W 11 , A1, B1, ∆z , and r are defined in Lemma 4.1. Furthermore, if N = b(logα T )/2c and St = d2T/ logα T e for each 1 ≤ t ≤ N where α > 1, we have
E[‖∇f(x̂T )‖2] ≤ α2W 11 2 lnα lnT T 3/2 + (αC0 + C2) 2 lnα lnT T + (∆z + C1) 2m lnα lnT√ T + α2MB1G 2 2 lnα lnT√ T . (8)
When N = 1, m ≤ ηti ≤ M and the bound (7) reduces to E[‖∇f(x̂T )‖ 2 ] ≤ O( 1T + 1 mT + M).
If, in particular, m and M are of order O(1/ √ T ), then we can derive the optimal convergence for constant bandwidth step-sizes, comparable to the literature for constant step-sizes (Yan et al., 2018; Liu et al., 2020; Mai and Johansson, 2020; Defazio, 2020).
It is not easy to explicitly minimize the right-hand-side of (7) with respect to N . However, N = b(logα T )/2c attempts to balance the last two terms and appears to be a good choice in practice. The theorem (see (8) establishes an O(lnT/ √ T ) rate under step-decay bandwidth step-size. If M = m, which means that the step-size follows the boundary functions, we get a near-optimal (up to lnT ) rate for stochastic momentum with a step-decay step-size on non-convex problems. We believe that this is the first near-optimal rate for stochastic momentum with step-decay step-size. The next result shows how an exponentially increasing stage-length allow to sharpen this guarantee even further. Theorem 4.3. Suppose the same setting as Lemma 4.1. Consider Algorithm 2, if the functions δ(t) = 1/αt−1 with α > 1, St = dS0αt−1e with S0 = √ T , we have
E[‖∇f(x̂T )‖2] ≤ O ( W 11 T 3/2 + C0 T + ∆z m √ T + C1 m √ T + MB1G 2 √ T ) .
The stage length St in Theorem 4.3 increases exponentially from S1 = d √ T e over N = blogα((α−
1) √ T + 1)c stages, resulting in an O(1/ √ T ) optimal rate for SGDM under the bandwidth-based step-decay scheme. This removes the lnT term of Theorem 8. To the best of our knowledge, this work is the first that is able to achieve an optimal rate for stochastic momentum with step-decay step-size in a general non-convex setting. Remark 4.4. (Better convergence than Liu et al. (2020)) We notice that reference Liu et al. (2020) analyzes multi-stage momentum and obtains the bound
E[‖∇f(x̃)‖2] ≤ O
( f(x1)− f∗ N + σL ∑N t=1 η t N ) . (9)
Here, x̃ is a uniformly sampled iterate (unlike our results, which favour later iterates) and N is the number of stages. The result uses a time-varying momentum parameter, whose value is determined by the step-size ηt, and also assumes an inverse relationship between the step-size and stage-length, i.e. that ηtSt is constant. Hence, N is of O(logα T ) and the convergence guarantee in (9) is of O(1/ logα T ), which is far worse than the rate of Theorem 4.2 and the optimal rate of Theorem 4.3.

heading: 5 NUMERICAL EXPERIMENTS
text: In this section, we design and evaluate several specific step-size policies that belong to the bandwidthbased family. We consider SGD with and without momentum, and compare their performances on neural network training tasks on the CIFAR10 and CIFAR100 datasets.

heading: 5.1 BASELINES AND PARAMETER SELECTION FOR THE BANDWIDTH STEP-SIZES
text: The bandwidth framework allows for a unified and streamlined (worst-case) analysis of all step-size policies that lie in the corresponding band. Within this family, the band gives a lot of freedom in crating innovative step-size policies with additional advantages. In particular, we will design a number of step-size policies that add periodic perturbations to a baseline step-size, attempting to both escape bad local minima and to improve the local convergence properties.
The step-decay bandwidth step-sizes divide the total number of iterations into a small number of stages, in which the boundary functions are constant. The width of the band are determined by the constants m and M . We will explore step-sizes that add a decreasing perturbation within each stage, starting at the upper band at the beginning of the stage, ending at the lower bound at the end of the stage, and decaying as 1/i, 1/ √ i, linearly or according to a cosine function. As baseline, we consider the step-decay step-size that follows the lower boundary function mδ(t). To use the same maximum value for the bandwidth step-sizes, we do not add any perturbation in the first stage; cf. Figure 5. For the 1/ √ t-band, on the other hand, stages correspond to epochs and perturbing the step-size within a stage would be too frequent and lead to bias. Rather, we choose to add similar perturbations
as for the step-decay band, but adjust the perturbation between stages. In our experiments, the two step-size policies perform roughly the same number of periods of perturbations over the training set. As baseline, we consider the step-size ηti = m/ √ t. In all experiments, the hyper-parameters (e.g., m and M ) have been determined using grid search, see Section E.2 for details.

heading: 5.2 NON-MONOTONIC SCHEDULE HELPS TO ESCAPE LOCAL MINIMA
text: To demonstrate the potential benefits of bandwidth-based non-monotonic step-size schedules, we consider the toy example (see Section E.3 for details and further results) from (Shi et al., 2020), which is non-convex and has four local minima2; see Figure 2. We then compare the final iterates of SGD with constant step-sizes (both large and small), step-decay, and a bandwidth-based step-decay step size which we call linear-mode (illustrated in Figure 5). As shown in Figure 2, a large constant step-size more easily escapes the bad local minima to approach the global minimum at (0.7,−0.7) than a small constant step-size. However, with a large constant step-size, the final iterates are scattered and end up far from the global minimum, which also has been observed in Figure 5 of (Shi et al., 2020). Therefore, we have to reduce the step-size at some points to reduce the error. This is exactly the intuition of step-decay step-size. As shown in Figure 2, the scatter plots of SGD with step-decay (red) and step-decay with linear-mode (green) are more concentrated around the global minimum than the constant step-sizes.
To quantify the ability of different step-sizes to escape the local minima, Table 1 reports the percentage of the final iterates under the different step-size policies that are close to each minima. We can see that the ability of the step-decay policy (named baseline) to escape the local minima is slightly worse than the large constant step-size, but Figure 2 shows that the variance of the near-optimal iterates is reduced significantly. In a similar way, we can see that linear-mode not only improves the ability to escape the local minima, but also produces final iterates that are more concentrated around the global optimum. Hence, it appears (at least in this example) that non-monotonic step-size schedules allow SGD to escape local minima and produce final iterates of high quality.

heading: 5.3 NUMERICAL RESULTS ON CIFAR10 AND CIFAR100
text: To illustrate the practical performance of the bandwidth-based step-sizes, we choose the wellknown CIFAR10 and CIFAR100 (Krizhevsky, 2009) image classification datasets. We consider the benchmark experiments of CIFAR10 on ResNet-18 (He et al., 2016) and CIFAR100 on a 28×10 wide residual network (WRN-28-10) (Zagoruyko and Komodakis, 2016), respectively. All the experiments are repeated 5 times to eliminate the influence of randomness.
We begin by evaluating our step-sizes for SGD. The left column of Figure 3 present the results of the 1/ √ t-band step-sizes on the two datasets. As shown in Figure 5, these stepsizes are all non-monotonic. The sudden increase in the step-size leads to a corresponding cliff-like reduction in accuracy followed by a recovery phase that consistently ends up at a better performance than in the previous stage. The three 1/ √ t-band step-sizes achieve significant improvements compared to their baseline (ηti = m/ √ t), in terms of both test loss and test accuracy. Moreover, the linear-mode performs the best compared to other polynomial decaying modes. Then, the results of SGD with step-decay band (described in Section 5.1 or see Figure 5 in Appendix) on CIFAR10 and CIFAR100
2Notation: 1 − 3 denote the local minima at top left, top right and bottom left, respectively; and 4 denotes the global minimum at bottom right.
are given in Figure 3 (middle column), respectively. At the final stage, the bandwidth step-sizes improve both test loss (see Figures 6 and 7 in Appendix) and test accuracy compared to the baseline. In particular, the cosine-mode performs the best on this problem. In the second stage, baseline methods have a sharp boost. Our guess is that the noise accumulates quickly under a relatively large constant step-size. But this phenomenon is only temporary. When we drop the step-size in the third stage, the performance improves.
Next, we evaluate the performance of step-decay bandwidth step-sizes on SGDM. The results are reported in Figure 3 (right column). The first observation from Figure 3 (right column) is that the step-decay bandwidth step-sizes also work well for SGDM, and that again, the cosine-mode performs better than the others. Another interesting observation is that the performance of vanilla SGD with cosine-mode (red) in Figure 3 is comparable to (even better than) SGDM with the baseline step-decay step-size (black) in Figure 3. A similar conclusion can also be made on CIFAR100.

heading: 6 CONCLUSION
text: We have studied a general family of bandwidth step-sizes for non-convex optimization. The family specifies a globally decaying band in which the actual step-size is allowed to vary, and includes both stage-wise and continuously decaying step-size policies as special cases. We have derived convergence rate guarantees for SGD and SGDM under all step-size policies in two important classes of bandwidth step-sizes (1/ √ t and step-decay), some of which are optimal. Our results provide theoretical guarantees for several popular “cyclical” step-sizes (Loshchilov and Hutter, 2017; Smith, 2017), as long as they are tuned to lie within our bands. We have also designed a number of novel step-sizes that add periodic perturbations to the global trend in order to escape bad local minima and to improve the local convergence properties. These step-sizes were shown to have superior practical performance in neural network training tasks on the CIFAR data set.
In the analysis of SGDM, we assume that the stochastic gradient is bounded (see Assumption 2(b)). It is interesting to see how to relax this assumption in some special cases, for example, when the step-size is constant throughout each stage. It would also be interesting to see if the bandwidth framework could be specialized to a more narrow class of step-sizes, for which we can provide even stronger convergence rates.

heading: A CONVERGENCE FOR BANDWIDTH POLYNOMIAL-DECAY STEP-SIZES
text: Our first more specific result considers the 1/ √ t bandwidth step-size with fixed stage length.
Theorem A.1. Under Assumptions 1 and 2(a), and assume that there exists a constant ∆0 > 0 such that E[f(xt1)− f∗] ≤ ∆0 for each t ≥ 1. If the step-size ηti ≤ 1/((ρ+ 1)L), the stage length St = S ≥ 1, and the boundary function δ(t) = 1/ √ t for each 1 ≤ t ≤ N , we have
E[‖∇f(x̂T )‖2] ≤ 3∆0 m · 1√ ST +
3M2Lσ 2m · √ S T . (10)
Theorem A.1 shows how multi-stage SGD with polynomial-decay bandwidth step-sizes converges to a stationary point. In the extreme case that S = 1, the step-size reduces to m/ √ t ≤ ηt1 ≤M/ √ t and our result is comparable to the non-asymptotic optimal rate derived for m = M = η0 in (Wang et al., 2021, Theorem 3.5). Multi-stage vs traditional 1/ √ t step-size (S = 1) In general, during the initial iterations when the first term of (10) dominates the error bound, the multi-stage technique can accelerate the convergence by a larger step-size and longer inner-loop S. However, a large S will make the error bound worse when the noise term begins to dominate the bound. The next theorem analyzes an algorithm with a decreasing stage length.
Theorem A.2. Under Assumptions 1 and 2(a),and assume that there exists a constant ∆0 > 0 such that E[f(xt1)− f∗] ≤ ∆0 for each t ≥ 1. If the step-size ηti ≤ 1/((ρ+ 1)L), the stage length St = dS0/ √ te with S0 = √ T , and δ(t) = 1/ √ t for each 1 ≤ t ≤ N , we have
E[‖∇f(x̂T )‖2] ≤ ( 2∆0 + M2Lσ
3− 2 √ 2 ) · 1 m √ T .
Schedule of Theorem A.2 vs Chen et al. (2019) The theorem establishes an optimal rate for multistage SGD with 1/ √ t bandwidth step-size. Note that Chen et al. (2019) also analyzes a stagewise algorithm with varying stage length, but their step-size decays as 1/t and stage length increases with t. An important novelty with our result is that it uses a long initial stage, S1 = √ T while a large
stage length in Chen et al. (2019) requires a small initial step-size (of O(1/ √ T )). Figure 4 illustrates the performance of different step-size policies: 1) 1/ √ t with St = 1; 2) 1/ √ t with St = dn/be
where n is total sample size and b is the batch size; 3) 1/ √ t with time-decreasing St = dS0/ √ te and
S0 = √ T ; 4) and 1/t step-size with St = S0t and S0 = 12 from Chen et al. (2019). We can see that the step-size policies proposed in Theorem A.2 are more stable and perform the best.
For completeness, we also compare the performance of step-size schedules proposed by Theorems 3.1 and 3.3 in Figure 4 (right). Although step-decay with time-increasing stage length has a superior theoretical convergence guarantee, constant stage length performs better in this particular example.

heading: B PROOFS OF LEMMA AND THEOREMS IN SECTION 3
text: Lemma B.1. Suppose that Assumption 1 and Assumption 2(a) hold. If we run the Algorithm 1 with T > 1 and ηti ≤ 1/((ρ+ 1)L), we have
E[‖∇f(x̂T )‖2] ≤ 1 N∑ t=1 Stδ −1(t)
( N∑ t=1 2(E[f(xt1)]− E[f(xt+11 )]) mδ(t)2 + M2Lσ m T ) . (11)
Proof. (of Lemma B.1) The L-smoothness of f (see Assumption 1), i.e., ‖∇f(x)−∇f(y)‖ ≤ L ‖x− y‖ for all x, y ∈ dom(f) implies that
f(x) + 〈∇f(x), y − x〉 − L 2 ‖x− y‖2 ≤ f(y) ≤ f(x) + 〈∇f(x), y − x〉+ L 2 ‖x− y‖2 . (12)
Applying the L-smoothness property of f and recalling Algorithm 1 at current iterate xti, we have
f(xti+1) ≤ f(xti) + 〈 ∇f(xti), xti+1 − xti 〉 + L
2 ∥∥xti+1 − xti∥∥2 ≤ f(xti)− ηti 〈 ∇f(xti), gti 〉 + (ηti) 2L
2 ∥∥gti∥∥2 . Taking conditional expectation of F ti on the above inequality and due to the unbiased estimator gti such that E[gti | F ti ] = ∇f(xti), we obtain that
E[f(xti+1)|F ti ] ≤ f(xti)− ηti ∥∥∇f(xti)∥∥2 + (ηti)2L2 E[∥∥gti∥∥2 |F ti ]. (13)
By Assumption 2(a) that E[‖gti −∇f(xti)‖ 2 | F ti ] ≤ ρ ‖∇f(xti)‖ 2 + σ, we have
E[ ∥∥gti∥∥2 | F ti ] = E[∥∥gti −∇f(xti) +∇f(xti)∥∥2]
= E[ ∥∥gti −∇f(xti)∥∥2] + E[∥∥∇f(xti)∥∥2] ≤ (ρ+ 1) ∥∥∇f(xti)∥∥2 + σ. (14)
Then incorporating the above inequality into (13) gives E[f(xti+1)|F ti ] ≤ f(xti) + ( −ηti + (ηti) 2L(ρ+ 1)
2 )∥∥∇f(xti)∥∥2 + (ηti)2Lσ2 . (15) If step-size ηti ≤ 1/((ρ+ 1)L), we have −ηti + (ηti) 2L(ρ+1) 2 ≤ −η t i/2. For any t ≥ 1, the inequality (15) can be estimated as:
ηti 2 ∥∥∇f(xti)∥∥2 ≤ f(xti)− E[f(xti+1)|F ti ] + (ηti)2Lσ2 . (16) Applying the assumption of step-size that ηti = n(t, i)δ(t) with m ≤ n(t, i) ≤ M for all t ∈ {1, 2, · · · , N} to (16) gives
mδ(t)
2 ∥∥∇f(xti)∥∥2 ≤ f(xti)− E[f(xti+1)|F ti ] + M2Lσ2 δ(t)2, (17) and dividing mδ(t)2/2 into the both sides, we have
δ−1(t) ∥∥∇f(xti)∥∥2 ≤ 2 (f(xti)− E[f(xti+1)|F ti ])mδ(t)2 + M2Lσm . (18)
Recalling the output x̂T of Algorithm 1 and then taking expectation, we have
E[‖∇f(x̂T )‖2] = 1∑N
t=1 Stδ −1(t) N∑ t=1 δ−1(t) · St∑ i=1 E [∥∥∇f(xti)∥∥2] . (19)
Applying (18) recursively from i = 1 to St and using the fact that xt+11 = x t St+1 , the sum of (19) for t ≥ 1 can be estimated as
N∑ t=1 δ−1(t) St∑ i=1 E [∥∥∇f(xti)∥∥2] ≤ N∑ t=1 2(E[f(xt1)]− E[f(xt+11 )]) mδ(t)2 + M2Lσ m T. (20)
Then plugging (20) into (19), we get
E[‖∇f(x̂T )‖2] ≤ 1∑N
t=1 Stδ −1(t) ( N∑ t=1 2(E[f(xt1)]− E[f(xt+11 )]) mδ(t)2 + M2Lσ m T ) as desired.
Proof. (of Theorem A.1) In this case, the step size ηti satisfies that m/ √ t ≤ ηti ≤ M/ √ t, where
δ(t) = 1/ √ t. By the definition of δ(t), we have
N∑ t=1 1 δ(t) ≥ ∫ N t=0 1 δ(t) dt = 2N 3 2 3 . (21)
Applying the assumption that E[f(xt1)− f∗] ≤ ∆0 where f∗ = minx f(x) and the definition of δ(t), we have
N∑ t=1 2(E[f(xt1)]− E[f(xtS+1)]) mδ(t)2 ≤ 2 m N∑ t=1 t(E[f(xt1)− f∗]− E[f(xt+11 )− f∗])
≤ 2 m
( ∆0 +
N∑ t=2 E[f(xt1)− f∗]
) ≤ 2N∆0
m . (22)
Under Assumptions 1 and 2(a) and ηti < 1/((ρ+ 1)L), thus Lemma B.1 holds. Incorporating these inequalities into Lemma B.1 gives
E[‖∇f(x̂T )‖2] ≤ 1 S ∑N t=1 δ −1(t) [ N∑ t=1 2(E[f(xt1)]− E[f(xt+11 )]) mδ(t)2 + SM2Lσ m T ]
≤ 3 2SN3/2 [ 2N∆0 m + M2Lσ m T ] ≤ 3∆0
m · 1√ ST +
3M2Lσ 2m · √ S T .
Then the proof is finished.
Proof. (of Theorem A.2) In this theorem, we consider SGD with the 1/ √ t bandwidth step-size, i.e., m√ t ≤ ηti ≤ M√t for 1 ≤ t ≤ N , where the stage length St = dS0/ √ te with S0 = √ T , and the boundary function δ(t) = 1/ √ t. By the relationship that ∑N t=1 St = T and x ≤ dxe ≤ x+ 1 for all x, we get that
(3− 2 √ 2)T ≤ N ≤ ( √ T
2 + 1)2. (23)
Under Assumptions 1 and 2(a) and ηti < 1/((ρ+ 1)L), thus Lemma B.1 holds. Following the same process as Theorem A.1, the inequality (22) also holds, that is
N∑ t=1 2(E[f(xt1)]− E[f(xtS+1)]) mδ(t)2 ≤ 2N∆0 m . (24)
Then incorporating the above inequalities to Lemma B.1, we have
E[‖∇f(x̂T )‖2] ≤ 1∑N
t=1 Stδ −1(t)
[ 2N∆0 m + M2Lσ m T ] ≤ 1∑N
t=1dS0/ √ te · √ t
[ 2N∆0 m + M2Lσ m · T ]
≤ 1 S0N [ 2N∆0 m + M2Lσ m · T ]
≤ 2∆0 mS0 + M2Lσ (3− 2 √ 2)mS0 = 1 m
( 2∆0 + M2Lσ
3− 2 √ 2
) · 1√
T ,
which concludes the proof.
Proof. (of Theorem 3.1) In this case, the step-size ηti exponentially decays every S iterations. By the definition of ηti , we have δ(t) = 1/α
t−1 for all 1 ≤ t ≤ N . The sum of 1/δ(t) can be estimated as
N∑ t=1 1 δ(t) = N∑ t=1 αt−1 = αN − 1 α− 1 .
Recalling the assumption that E[f(xt1)− f∗] ≤ ∆0 for all t ≥ 1, we have N∑ t=1 2(E[f(xt1)]− E[f(xt+11 )]) mδ(t)2 ≤ 2 m N∑ t=1 α2(t−1)(E[f(xt1)− f∗]− E[f(xt+11 )− f∗])
≤ 2 m
( (α2 − 1)
N∑ t=2 α2(t−2)E[f(xt1)− f∗] + ∆0
)
≤ 2α 2(N−1)∆0 m .
Under Assumptions 1 and 2(a) and ηti < 1/((ρ+ 1)L), thus Lemma B.1 holds. Then applying the result of Lemma B.1 and incorporating the above inequalities into Lemma B.1 gives
E[‖∇f(x̂T )‖2] ≤ 1 S ∑N t=1 δ −1 1 (t) ( N∑ t=1 2(E[f(xt1)]− E[f(xtS+1)]) mδ(t)2 + SM2Lσ m · N∑ t=1 δ2(t) 2 δ(t)2 )
≤ α− 1 S(αN − 1)
( 2α2(N−1)∆0
m + M2Lσ m SN
) .
Substituting the specific values of N = b(logα T )/2c, S = d2T/ logα T e into the above inequality, we have
E[‖∇f(x̂T )‖2] ≤ (α− 1)
S(αN − 1)
( ∆0α 2(N−1)
m + M2Lσ m SN ) ≤ (
∆0 2αm + αM2Lσ 2m
) (α− 1) logα T√
T − α .
then transform the base logα to the natural logarithm ln, we can get the desired result.
Proof. (of Theorem 3.3) If S0 = √ T and St = dS0αt−1e, by ∑N t=1 St = T , the stage length N is
blogα((α− 1) √ T + 1)c. In this case, we have δ(t) = 1/αt−1 for each t ∈ [N ]. Under Assumptions 1 and 2(a) and ηti < 1/((ρ + 1)L), thus Lemma B.1 holds. Before applying the result of Lemma B.1, we first give the following estimations:
N∑ t=1 Stδ −1(t) ≥ N∑ t=1 S0α t−1 · αt−1 = √ T · 1− α 2N 1− α2 = (α− 1)T 32 + 2T α+ 1
(25)
and also
N∑ t=1 2(E[f(xt1)]− E[f(xt+11 )]) mδ(t)2 ≤ 2 m
( (α2 − 1)
N∑ t=2 α2(t−2)E[f(xt1)− f∗] + ∆0
)
≤ 2α 2(N−1)∆0 m = 2((α− 1)
√ T + 1)2∆0
α2m .
Applying these results into Lemma B.1, we have
E[‖∇f(x̂T )‖2] ≤ α+ 1
(α− 1)T 32 + 2T
( 2((α− 1) √ T + 1)2∆0
α2m + M2Lσ m · T
)
≤ α+ 1 α− 1 ( 2∆0 m + M2Lσ m ) 1√ T +O ( 1 T ) .
Thus the proof is complete.
In Lemma B.1, we provide a unified analysis framework for bandwidth step-sizes which are independent on the current random information. Recently, there are some interesting non-monotonic stepsizes, e.g., the trust-region-ish algorithm (Curtis et al., 2019) and stochastic Polyak step-sizes (Loizou et al., 2021), which can also be regarded to be in a band, but those are related to the current stochastic information. We provide a unified framework for these kind of step-sizes below.
Lemma B.2. Under the same conditions as in Lemma B.1, we assume that the step-size ηti satisfies mδ(t) ≤ ηti ≤ Mδ(t) and is dependent on the current random information. If m ≤ 4L(ρ+1) and m ≤M ≤ M̃ := −ρ+ √ ρ2+2(ρ+1)(2+ρ)mL
L(ρ+1) , we have
E[‖∇f(x̃T )‖2] ≤ 1
ψ0ΣT
( ∆0
δ(N)2 + (M −m)σΣT +
LM2σ
2 T
) (26)
where ΣT = ∑N t=1 ∑St i=1 δ(t) −1 and ψ0 = (2 + ρ)m− ρM − LM 2 2 (ρ+ 1).
Proof. We consider the step-size ηti is depended on the current random information and mδ(t) ≤ ηti ≤Mδ(t). By the L-smoothness of f , that is ‖∇f(x)−∇f(y)‖ ≤ L ‖x− y‖, implies that
f(y) ≤ f(x) + 〈∇f(x), y − x〉+ L 2 ‖x− y‖2 (27)
Let x = xti and y = x t i+1, then
f(xti+1) ≤ f(xti)− 〈 ∇f(xti), ηtigti 〉 + L
2 ∥∥ηtigti∥∥2 (28) Next we turn to estimate the product term −〈∇f(xti), ηtigti〉.
− 〈 ∇f(xti), ηtigti 〉 = ηti (∥∥gti −∇f(xti)∥∥2 − ∥∥gti∥∥2 − ∥∥∇f(xti)∥∥2) ≤Mδ(t) ∥∥gti −∇f(xti)∥∥2 −mδ(t)∥∥gti∥∥2 −mδ(t)∥∥∇f(xti)∥∥2 . (29)
Then taking conditional expectation on F ti to (28) and applying Assumption 2(a) and (14), we have E[f(xti+1) | F ti ] ≤ f(xti)−mδ(t) ∥∥∇f(xti)∥∥2 −mδ(t)E[∥∥gti∥∥2 | F ti ] +Mδ(t)E[∥∥gti −∇f(xti)∥∥2 | F ti ]
+ LM2δ(t)2 2 E[ ∥∥gti∥∥2 | F ti ]
≤ f(xti)−mδ(t) ∥∥∇f(xti)∥∥2 −mδ(t)(E[∥∥gti −∇f(xti)∥∥2] + ∥∥∇f(xti)∥∥2)
+Mδ(t)E[ ∥∥gti −∇f(xti)∥∥2 | F ti ] + LM2δ(t)22 (E[∥∥gti −∇f(xti)∥∥2] + ∥∥∇f(xti)∥∥2)
≤ f(xti)− ( 2mδ(t)− LM 2δ(t)2
2
)∥∥∇f(xti)∥∥2 + (Mδ(t)−mδ(t) + LM2δ(t)22 ) E[ ∥∥gti −∇f(xti)∥∥2]
≤ f(xti)− ( 2m− LM 2δ(t)
2 − ρ
( M −m+ LM 2δ(t)
2
)) δ(t) ∥∥∇f(xti)∥∥2 + ( M −m+ LM 2δ(t)
2
) δ(t)σ
where by M ≥ m we have Mδ(t) − mδ(t) + LM 2δ(t)2 2 > 0. Let ψ := 2m − LM2δ(t) 2 − ρ ( M −m+ LM
2δ(t) 2
) . We know that δ(t) ≤ 1, then
ψ ≥ (2 + ρ)m− ρM − LM 2
2 (ρ+ 1). (30)
Let ψ0 = (2 + ρ)m− ρM − LM 2
2 (ρ+ 1). By solving the quadratic expression of M in ψ0, if
m ≤M < M̃ := −ρ+
√ ρ2 + 2(ρ+ 1)(2 + ρ)mL
L(ρ+ 1) (31)
then we have ψ ≥ ψ0 > 0. To guarantee that m ≤M , we require that ψ0(m) > 0, then m ≤ 4L(ρ+1) . Recalling the output x̃ which is selected from {xti} for all i ∈ [St] and t ∈ [N ] with probability Pt ∝ 1/δ(t), we have
E[‖∇f(x̃)‖2] = 1∑N t=1 Stδ(t) −1 N∑ t=1 δ(t)−1 St∑ i=1 E[ ∥∥∇f(xti)∥∥2]
≤ 1 ψ0ΣT
(∑N t=1 ( E[f(xt1)]− E[f(xtSt+1)] ) δ(t)2 + (M −m)σ N∑ t=1 St∑ i=1 δ(t)−1 + LM2σ 2 T ) (a) ≤ 1 ψ0ΣT ( ∆0 δ(N)2 + (M −m)σΣT + LM2σ 2 T ) (32)
where ΣT = ∑N t=1 Stδ(t)
−1 and (a) follows from the assumption that E[f(xt1)− f∗] ≤ ∆0 where f∗ = minx f(x) and xt+11 = x t St+1 .
The above lemma immediately results in the following convergence results for bandwidth step-sizes which depend on the current random information.
• If N = 1 and St = T , that is m ≤ ηti ≤ M and ηti is related to the current random information, we have
E[‖∇f(x̃)‖2] ≤ ∆0 ψ0T + (M −m) + LM2/2 ψ0 σ (33)
In this case, the boundary function δ(t) = 1, so the assumption on function value can be replaced by f(x10)− f∗ is bounded. We can achieve an O(1/T ) convergence rate to reach a neighborhood of the stationary point. This error bound is comparable to the results of (Loizou et al., 2021, Theorem 3.8) and (Curtis et al., 2019, Theorem 3.5).
Algorithm 2 SGDM with Bandwidth-based Step-Size
1: Input: initial point x11 ∈ Rd, v11 = 0, # iterations T , # stages N , stage length {St} N t=1 such that ∑N t=1 St = T , momentum parameter β ∈ (0, 1), the sequences {δ(t)} N t=1 and{
{n(t, i)}Sti=1 }N t=1 ∈ [m,M ] with 0 < m ≤M
2: for t = 1 : N do 3: for i = 1 : St do 4: Query a stochastic gradient oracle O at xti to get a vector gti such that E[gti | F ti ] = ∇f(xti) 5: Update step-size ηti = n(t, i)δ(t), which belongs to the interval [mδ(t),Mδ(t)] 6: vti+1 = βv t i + (1− β)gti 7: xti+1 = x t i − ηtivti+1 8: end for 9: vt+11 = v t St+1 and xt+11 = x t St+1
10: end for 11: Return: x̂T is uniformly chosen from { xt ∗ 1 , x t∗ 2 , · · · , xt ∗
St∗
} , where the integer t∗ is randomly
chosen from {1, 2, · · · , N} with probability Pt = δ−1(t)/( ∑N l=1 δ −1(l))
• If N > 1, St = S and δ(t) = 1/αt−1, we have
E[‖∇f(x̃)‖2] ≤ α− 1 ψ0S(αN − 1)
( ∆0α 2(N−1) + LM2σ
2 T
) +
(M −m)σ ψ0
= 1
ψ0
( NαN−1
T + LM2σ 2 (α− 1)N (αN − 1)
) +
(M −m)σ ψ0
(34)
Compared to the case that N = 1, we observe that increasing N > 1 can improve the error term LM
2/2 ψ0 σ of (33). If M −m ≤ LM 2 2 , i.e., m ≥M − LM2 2 , then LM2/2 ψ0
σ turns out to dominate the error bound of (33). If we increase N appropriately, then the error term of (33) can be improved.

heading: C PROOFS OF LEMMA AND THEOREMS IN SECTION 4
text: We recall the momentum scheme of Algorithm 2 below
vti+1 = βv t i + (1− β)gti (35)
xti+1 = x t i − ηtivti+1 (36)
where β ∈ (0, 1). Before giving the proofs, we introduce an extra variable zti = xti 1−β − β 1−βx t i−1, then
zti − xti = β
1− β (xti − xti−1), (37)
zti+1 − xti = 1
1− β (xti+1 − xti). (38)
However, the bandwidth-based step-size ηti in our analysis is time dependent and also possibly non-monotonic, so the commonly used equalities xti+1 = x t i − ηgti + β(xti − xti−1) (Yan et al., 2018) or zti+1 = z t i − ηtgti (see lemma 3 of Liu et al. (2020)) do not hold in our analysis. This significantly increases the level of difficulty of the analysis. The results of Lemma 4.1 is based on a sequence of lemmas introduced below. Lemma C.1. Suppose that the objective function f satisfies Assumption 1. At each stage t, the step-size ηti is monotonically decreasing. Then, for i ≥ 2, we have
E[f(zti+1) | F ti ]− f(zti) + β
1− β
( 1− η t i
ηti−1
)( f(xti)− f(xti−1) ) ≤ −ηti ∥∥∇f(xti)∥∥2 + βL2(1− β)2 ∥∥xti − xti−1∥∥2 + L2(1− β)2E[∥∥xti+1 − xti∥∥2 | F ti ].
Proof. Using the L-smoothness of f (Assumption 1) and taking conditional expectation gives
E[f(zti+1) | F ti ]− f(xti) ≤ E[ 〈 ∇f(xti), zti+1 − xti 〉 | F ti ] + L 2 E[ ∥∥zti+1 − xti∥∥2 | F ti ]
≤ 1 1− β
E[ 〈 ∇f(xti), xti+1 − xti 〉 | F ti ] +
L 2(1− β)2 E[ ∥∥xti+1 − xti∥∥2 | F ti ]
≤ 1 1− β
E 〈 ∇f(xti),−ηti((1− β)gti + βvti) 〉 | F ti ] +
L 2(1− β)2 E[ ∥∥xti+1 − xti∥∥2 | F ti ]
≤ −ηti ∥∥∇f(xti)∥∥2 − ηtiβ1− β 〈∇f(xti), vti〉+ L2(1− β)2E[∥∥xti+1 − xti∥∥2 | F ti ]
= −ηti ∥∥∇f(xti)∥∥2 − βηti(1− β)ηti−1 〈∇f(xti), xti−1 − xti〉+ L2(1− β)2E[∥∥xti+1 − xti∥∥2 | F ti ].
Re-using the L-smoothness property of f at zti and x t i gives
f(zti) ≥ f(xti) + 〈 ∇f(xti), zti − xti 〉 − L
2 ∥∥zti − xti∥∥2 = f(xti) + β 1− β 〈 ∇f(xti), xti − xti−1 〉 − Lβ 2 2(1− β)2 ∥∥xti − xti−1∥∥2 . (39)
Then, combining the two inequalities above, we have
E[f(zti+1) | F ti ] ≤ f(zti)− ηti ∥∥∇f(xti)∥∥2 + β1− β ( 1− η t i
ηti−1
)〈 ∇f(xti), xti−1 − xti 〉 + L
2(1− β)2 E[ ∥∥xti+1 − xti∥∥2 | F ti ] + Lβ22(1− β)2 ∥∥xti − xti−1∥∥2 . (40)
The step-size ηti for each stage is monotonically decreasing, i.e. η t i ≤ ηti−1 for i ≥ 2, so(
1− η t i
ηti−1 ) ≥ 0. By the L-smoothness of f , the inner product of (40) can be estimated as〈
∇f(xti), xti−1 − xti 〉 ≤ f(xti−1)− f(xti) + L
2 ∥∥xti − xti−1∥∥2 . (41) Applying (41) into (40), we find
E[f(zti+1) | F ti ]
≤ f(zti)− ηti ∥∥∇f(xti)∥∥2 + β1− β ( 1− η t i
ηti−1
)( f(xti−1)− f(xti) + L
2 ∥∥xti − xti−1∥∥2) + L
2(1− β)2 E[ ∥∥xti+1 − xti∥∥2 | F ti ] + Lβ22(1− β)2 ∥∥xti − xti−1∥∥2 . (42)
Finally, we re-write the above inequality as
E[f(zti+1) | F ti ]− f(zti) + β
1− β
( 1− η t i
ηti−1
)( f(xti)− f(xti−1) ) ≤ −ηti ∥∥∇f(xti)∥∥2 + ( β2L2(1− β)2 + βL2(1− β) )∥∥xti − xti−1∥∥2 + L2(1− β)2E[∥∥xti+1 − xti∥∥2 | F ti ]
≤ −ηti ∥∥∇f(xti)∥∥2 + βL2(1− β)2 ∥∥xti − xti−1∥∥2 + L2(1− β)2E[∥∥xti+1 − xti∥∥2 | F ti ].
The proof is complete.
Lemma C.2. Suppose that the objective function satisfies Assumption 1 and the step-size ηti is monotonically decreasing with ηti ≤ 1L at each stage, then
E [∥∥xti+1 − xti∥∥2 | F ti ]− β2 ∥∥xti − xti−1∥∥2 + 2(1− β)ηti (E[f(xti+1) | F ti ]− f(xti)) ≤ −2(ηti)2(1− β)2 ∥∥∇f(xti)∥∥2 + 2(ηti)2(1− β)2E [∥∥gti∥∥2 | F ti ]+ (ηti)3β(1− β)L∥∥vti∥∥2 .
Proof. First, due to the L-smoothness of the objective function f , we have f(xti+1) ≤ f(xti) + 〈 ∇f(xti), xti+1 − xti 〉 + L
2 ∥∥xti+1 − xti∥∥2 ≤ f(xti) + 〈 ∇f(xti),−ηti((1− β)gti + βvti) 〉 + (ηti) 2L
2 ∥∥(1− β)gti + βvti∥∥2 (a)
≤ f(xti)− (1− β)ηti 〈 ∇f(xti), gti 〉 − βηti 〈 ∇f(xti), vti 〉 + (ηti) 2L
2
( (1− β) ∥∥gti∥∥2 + β ∥∥vti∥∥2) where inequality (a) follows the Cauchy-Schwarz inequality that ‖(1− β)gti + βvti‖
2 ≤ (1 − β) ‖gti‖ 2 + β ‖vti‖ 2 . Then taking conditional expectation on both sides and due to that gti is an unbiased estimator of∇f(xti), i.e., E[gti | F ti ] = ∇f(xti), we have
E[f(xti+1) | F ti ] ≤ f(xti)− ηti(1− β) ∥∥∇f(xti)∥∥2 − βηti 〈∇f(xti), vti〉
+ (ηti) 2(1− β)L 2 E[ ∥∥gti∥∥2 | F ti ] + (ηti)2βL2 ∥∥vti∥∥2 . (43)
We recall the definition of vti+1 and incorporate (35) into (36), then
E[ ∥∥xti+1 − xti∥∥2 | F ti ]
= E[ ∥∥ηti(βvti + (1− β)gti)∥∥2 | F ti ] = (ηti)2E[∥∥βvti + (1− β)gti∥∥2 | F ti ]
(a) = (ηti) 2 ( β2 ∥∥vti∥∥2 + (1− β)2E[∥∥gti∥∥2 | F ti ] + 2β(1− β) 〈vti ,∇f(xti)〉)
(b) = ( ηti ηti−1 )2 β2 ∥∥xti − xti−1∥∥2 + (ηti)2(1− β)((1− β)E[∥∥gti∥∥2] + 2β 〈vti ,∇f(xti)〉) (44)
(c) ≤ β2 ∥∥xti − xti−1∥∥2 + (ηti)2(1− β)((1− β)E[∥∥gti∥∥2] + 2β 〈vti ,∇f(xti)〉) , (45)
where (a) uses the fact that E[gti | F ti ] = ∇f(xti); (b) follows the procedure that xti = xti−1−ηti−1vti ; (c) applies the fact that the step-size per stage is monotonically decreasing, i.e., ηti ≤ ηti−1 for i ≥ 2. Then multiplying 2ηti(1− β) into (43) and combining (45), we get that
E [∥∥xti+1 − xti∥∥2 | F ti ]− β2 ∥∥xti − xti−1∥∥2 + 2ηti(1− β) (E[f(xti+1) | F ti ]− f(xti)) ≤ −2(ηti)2(1− β)2
∥∥∇f(xti)∥∥2 + (ηti) 2(1− β)2(Lηti + 1)E [∥∥gti∥∥2 | F ti ]+ (ηti)3β(1− β)L∥∥vti∥∥2
≤ −2(ηti)2(1− β)2 ∥∥∇f(xti)∥∥2 + 2(ηti)2(1− β)2E [∥∥gti∥∥2 | F ti ]+ (ηti)3β(1− β)L∥∥vti∥∥2
where the last inequality follows from the fact that ηti ≤ 1/L.
Proof. (of Lemma 4.1) First we apply the result of Lemma C.1 and divided by ηti to the both side, we have
E[f(zti+1) | F ti ]− f(zti) ηti + β 1− β
( f(xti)− f(xti−1) ) ηti − β 1− β ( f(xti)− f(xti−1) ) ηti−1
≤ − ∥∥∇f(xti)∥∥2 + ( βL2(1− β)2ηti )∥∥xti − xti−1∥∥2 + ηtiL2(1− β)2E[∥∥vti+1∥∥2 | F ti ]. (46) Then we recall the result of Lemma C.2
E[ ∥∥xti+1 − xti∥∥2 | F ti ]− ∥∥xti − xti−1∥∥2 + 2ηti (E[f(xti+1)]− f(xti)) ≤ −(1− β2)
∥∥xti − xti−1∥∥2 − 2(ηti)2(1− β)2 ∥∥∇f(xti)∥∥2 + 2(ηti) 2(1− β)2E [∥∥gti∥∥2 | F ti ]+ (ηti)3β(1− β)L∥∥vti∥∥2 ,
multiplying a constant r = βL2(1−β2)(1−β)2 > 0 and dividing η t i to the both side, and then incorporating it into (46), we have
E[f(zti+1) | F ti ]− f(zti) ηti + β 1− β
( f(xti)− f(xti−1) ) ηti − β 1− β ( f(xti)− f(xti−1) ) ηti−1
+ r
ηti
( E[ ∥∥xti+1 − xti∥∥2 | F ti ]− ∥∥xti − xti−1∥∥2)+ 2r (E[f(xti+1) | F ti ]− f(xti))
≤ − ∥∥∇f(xti)∥∥2 + 2r(ηti)(1− β)2E[∥∥gti∥∥2 | F ti ] + r(ηti)2β(1− β)L∥∥vti∥∥2
+ ηtiL 2(1− β)2 E[ ∥∥vti+1∥∥2 | F ti ]. (47)
We define a function W ti+1 as follows:
W ti+1 = f(zti+1)− f∗ ηti + r ∥∥xti+1 − xti∥∥2 ηti + 2r[f(xti+1)− f∗].
Because of ηti ≤ ηti−1 at each stage, we have −1/ηti ≤ −1/ηti−1 (i ≥ 2), then
W ti+1 ≤ f(zti+1)− f∗
ηti +
β 1− β (f(xti)− f∗) ηti − β 1− β (f(xti)− f∗) ηti−1
+ r
ηti ∥∥xti+1 − xti∥∥2 + 2r (f(xti+1)− f∗) . (48) Taking conditional expectation on W ti+1 and applying (47) to the above inequality, we have
E[W ti+1 | F ti ] ≤W ti + (f(zti)− f∗) ( 1
ηti − 1 ηti−1
) + β ( f(xti−1)− f∗ ) 1− β ( 1 ηti − 1 ηti−1 ) + r
∥∥xti − xti−1∥∥2( 1ηti − 1ηti−1 ) − ∥∥∇f(xti)∥∥2 + 2r(ηti)(1− β)2E[∥∥gti∥∥2 | F ti ]
+ r(ηti) 2β(1− β)L ∥∥vti∥∥2 + ηtiL2(1− β)2E[∥∥vti+1∥∥2 | F ti ]. (49) We recall that v11 = 0, due to the assumption that E[‖gti‖ 2 ] ≤ G2, and vti+1 is a convex combination of gti and v t i , then by induction if E[‖vti‖ 2 ] ≤ G2, then
E[ ∥∥vti+1∥∥2] = E ∥∥βvti + (1− β)gti∥∥2 ≤ β ∥∥vti∥∥2 + (1− β)E[∥∥gti∥∥2] ≤ G2. (50)
Therefore, we have E[‖vti‖ 2 ] is bounded by G2. Then we apply E[‖gti‖ 2 ] ≤ G2, E[‖vti‖ 2 ] ≤ G2 and ηti ≤ 1/L into (49)
E[W ti+1 | F ti ]−W ti ≤ (f(zti)− f∗) ( 1
ηti − 1 ηti−1
) + β ( f(xti−1)− f∗ ) 1− β ( 1 ηti − 1 ηti−1 ) + r
∥∥xti − xti−1∥∥2( 1ηti − 1ηti−1 ) − ∥∥∇f(xti)∥∥2
+ ηtiG 2
( r(1− β)(2− β) + L
2(1− β)2
) . (51)
The step-size is decreasing at each stage, then ηti ≤ ηti−1 (i ≥ 2), thus 1ηti − 1 ηti−1 ≥ 0. Due to the fact that vti is bounded (see (50)), i.e., E[‖vti‖ 2 ] ≤ G2, we have
E[ ∥∥xti − xti−1∥∥2 = (ηti−1)2E[∥∥vti∥∥2]] ≤ (ηti−1)2G2 ≤ G2L2 . (52)
Recalling the definition of zti , and applying the assumption that E[f(xti)− f∗] ≤ ∆0 for each t, i ≥ 1 and f is L-smooth on its domain, and ηti ≤ 1/L gives
f(zti) ≤ f(xti) + 〈 ∇f(xti), zti − xti 〉 + L
2 ∥∥zti − xti∥∥2 ≤ f(xti) + β 1− β 〈 ∇f(xti), xti − xti−1 〉 + Lβ2 2(1− β)2 ∥∥xti − xti−1∥∥2
(a) ≤ f(xti) + β
1− β
( f(xti)− f(xti−1) + L
2 ∥∥xti − xti−1∥∥2)+ Lβ22(1− β)2 ∥∥xti − xti−1∥∥2 ≤ f(xti) + β 1− β ( f(xti)− f(xti−1) ) + Lβ 2(1− β)2 ∥∥xti − xti−1∥∥2
≤ 1 1− β f(xti)− β 1− β f(xti−1) +
Lβ 2(1− β)2 ∥∥xti − xti−1∥∥2 (53)
where the inequality (a) dues to the fact that f(xti−1) ≤ f(xti) + 〈 ∇f(xti), xti−1 − xti 〉 + L 2
∥∥xti−1 − xti∥∥2. Then we have E[f(zti)− f∗] ≤ 1
1− β ∆0 +
Lβ
2(1− β)2 (ηti−1) 2G2 ≤ ∆0 1− β + βG2 2(1− β)2L . (54)
Let ∆z = ∆01−β+ βG2 2(1−β)2L . Finally, applying (52) and (54), the bounded assumption on E[f(x t i)−f∗], and ηti ≤ 1/L, we have
E[W ti+1 | F ti ] ≤W ti + ( β∆0 1− β + ∆z + rG2 L2 )( 1 ηti − 1 ηti−1 ) − ∥∥∇f(xti)∥∥2
+ ηti
( r(1− β)(2− β) + L
2(1− β)2
) G2
= W ti +A1
( 1
ηti − 1 ηti−1
) − ∥∥∇f(xti)∥∥2 + ηtiB1G2 (55)
where A1 = β∆01−β + ∆z + rG2 L2 , B1 = r(1− β)(2− β) + L 2(1−β)2 and ∆z = ∆0 1−β + βG2 2(1−β)2L .
The bandwidth step-size highly rises the difficulty of the analysis for momentum, especially when the step-size has an increase between the stages, i.e. ηt−1St−1 := η t 0 < η t 1. Before giving the results, we consider two situations:
• ηt0 > η t 1. We can apply Lemma 4.1 from i = 1 to St. Recalling the definition of W t i+1, we
have W t+11 = W t St+1 , then
St∑ i=1 E[ ∥∥∇f(xti)∥∥2] ≤ (E[W t1 ]− E[W t+11 ])+A1( 1ηtSt − 1ηt0 ) +B1G 2 St∑ i=1 ηti . (56)
• Otherwise if ηt0 ≤ ηt1, the results of Lemma 4.1 only hold from i = 2 to St. Then
St∑ i=1 E[ ∥∥∇f(xti)∥∥2] ≤ (E[W t2 ]− E[W t1 ] + E[W t1 ]− E[W t+11 ])+ E[∥∥∇f(xt1)∥∥2]
+A1
( 1
ηtSt − 1 ηt1
) +B1G 2 St∑ i=2 ηti . (57)
For the bandwidth step-size, the initial step-size of stage t, ηt1, is possibly larger than the ending step-size of the previous stage, ηt0. Thus, we can not use the simpler condition (56), but have to rely on (57) in our derivations below.
Lemma C.3. Suppose the same setting as Lemma 4.1, we have
E[‖∇f(x̂T )‖2] ≤ 1∑ δ
( W 11 +
C0 δ(N + 1) + ∆z mδ(N)δ(N + 1) + C1 m N∑ t=1 1 δ(t)2 + C2 N∑ t=1 1 δ(t) +B1G 2MT ) where ∑ δ = ∑N t=1 Stδ(t) −1, C0 = r(G 2 L + 2∆0), C1 = A1 + ∆z + ∆0 1−β , C2 = C0 +A2G 2, and A1, B1, r, and ∆z are defined in Lemma 4.1.
Proof. Applying the result of Lemma 4.1 from i = 2 to St, the step-size ηti ∈ [mδ(t),Mδ(t)], and W t+11 = W t St+1 , we have
St∑ i=1 E[ ∥∥∇f(xti)∥∥2] ≤ (E[W t2 ]− E[W t1 ] + E[W t1 ]− E[W t+11 ])+ E[∥∥∇f(xt1)∥∥2]
+A1
( 1
ηtSt − 1 ηt1
) +B1G 2M(St − 1)δ(t). (58)
Recalling the output of Algorithm 2, we have
E[‖∇f(x̂T )‖2] = 1∑N
t=1 Stδ(t) −1 N∑ t=1 δ(t)−1 St∑ i=1 E[ ∥∥∇f(xti)∥∥2]. (59)
Then we divide δ(t) into the both side of (58), apply (58) from t = 1 to N and let ∑ δ =∑N
t=1 Stδ(t) −1
E[‖∇f(x̂T )‖2] ≤ 1∑ δ ( N∑ t=1 E[W t2 ]− E[W t1 ] + E[W t1 ]− E[W t+11 ] δ(t) +A1 N∑ t=1 1 δ(t) ( 1 ηtSt − 1 ηt0 ))
+ 1∑ δ ( N∑ t=1 E[‖∇f(xt1)‖ 2 ] δ(t) +B1G 2MT ) . (60)
First, we estimate ∑N t=1 E[‖∇f(xt1)‖2] δ(t) . From Lemma C.2, let i = 1, then incorporating the inequalities (43) and (44), we have
2(ηt1) 2(1− β)2E[ ∥∥∇f(xt1)∥∥2] ≤ (ηt1βηt0 )2 ∥∥xt1 − xt0∥∥2 − E[∥∥xt2 − xt1∥∥2] + 2(ηt1(1− β))2E[∥∥gt1∥∥2]
+ (ηt1) 3β(1− β)L ∥∥vt1∥∥2 − 2ηt1(1− β) (f(xt1)− E[f(xt2)]]) . Then dividing 2(ηt1) 2(1− β)2 to the both side and applying the fact that ‖xt1 − xt0‖ = (ηt0)2 ‖vt1‖ 2, E[‖gti‖ 2 ] ≤ G2 and E[‖vti‖ 2 ] ≤ G2 for any i, t, f(xt2)− f∗ ≤ ∆0 and ηt1 ≥ mδ(t), we have
E[ ∥∥∇f(xt1)∥∥2] ≤ G2(1 + β2(1− β)2 ) + ∆0 (1− β)ηt1 ≤ A2G2 + ∆0 (1− β)mδ(t) (61)
where A2 = ( 1 + β2(1−β)2 ) . Then
N∑ t=1 E[‖∇f(xt1)‖ 2 ] δ(t) ≤ N∑ t=1 A2G 2 δ(t) + ∆0 (1− β)m N∑ t=1 1 δ2(t) . (62)
Next we turn to estimate ∑N t=1 E[W t2−W t 1 ] δ(t) . Recalling the definition of W t i+1, we have W t i ≥ 0. Applying the inequalities (52), (54) and the assumption that f(xti)− f∗ ≤ ∆0 for any i, t, we have
E[W t2 ]− E[W t1 ] ≤ E[W t2 ] := E[f(zt2)− f∗] ηt1 + rE[‖xt2 − xt1‖ 2 ] ηt1 + 2rE[f(xt2)− f∗]
≤ ∆z mδ(t) + r
( G2
L + 2∆0
) , (63)
dividing δ(t) and applying (63) from t = 1 to N , we have
N∑ t=1 E[W t2 ]− E[W t1 ] δ(t) ≤ ∆z m N∑ t=1 1 δ(t)2 + r ( G2 L + 2∆0 ) N∑ t=1 1 δ(t) . (64)
Then we consider
N∑ t=1 E[W t1 ]− E[W t+11 ] δ(t) = N∑ t=1 ( E[W t1 ] δ(t) − E[W t+1 1 ] δ1(t+ 1) ) + N∑ t=1 ( 1 δ1(t+ 1) − 1 δ(t) ) E[W t+11 ]
≤ W 1 1
δ1(1) + N∑ t=1 ( 1 δ1(t+ 1) − 1 δ(t) ) E[W t+11 ]. (65)
Recalling the definition of E[W t+11 ],
E[W t+11 ] = E[f(zt+11 )− f∗] ηtSt + rE[ ∥∥xt+11 − xt+10 ∥∥2] ηtSt + 2rE[f(xt+11 )− f∗],
and applying the assumption that E[f(xti)− f∗] ≤ ∆0 and E[f(zti)− f∗] ≤ ∆z , and η t+1 0 = η t St , E[ ∥∥xt+11 − xt+10 ∥∥2] = (ηtSt)2E[∥∥vt+11 ∥∥2] ≤ (ηtSt)2G2, we have
E[W t+11 ] ≤ ∆z ηtSt + rηtStG 2 + 2r∆0
(a) ≤ ∆z mδ(t) + r
( G2
L + 2∆0
) (66)
where (a) follows from ηtSt ≥ mδ(t) and η t St ≤ 1/L. Applying (66) into (65), we have
N∑ t=1 E[W t1 ]− E[W t+11 ] δ(t) ≤ W 1 1 δ(1) + r ( G2 L + 2∆0 ) δ(N + 1) + ∆z m N∑ t=1 ( 1 δ(t+ 1)δ(t) − 1 δ(t)2 )
≤ W 1 1 δ(1) + r ( G2 L + 2∆0 ) δ(N + 1) + ∆z m N∑ t=1 ( 1 δ(t+ 1)δ(t) − 1 δ(t)δ(t− 1) )
≤ W 1 1 δ(1) + r ( G2 L + 2∆0 ) δ(N + 1) + ∆z mδ(N)δ(N + 1) , (67)
where the second inequality follows that δ(t) is decreasing, so δ(t) ≤ δ(t − 1), then −1/δ(t) ≤ −1/δ(t− 1). Finally, due to that ηtSt ∈ [mδ(t),Mδ(t)], we have
N∑ t=1 1 δ(t) ( 1 ηtSt − 1 ηt0 ) ≤ N∑ t=1
1
mδ(t)2 . (68)
Incorporate the inequalities (62), (64), (67) and (68) into (60), we have
E[‖∇f(x̂T )‖2] ≤ 1∑ δ W 11 δ(1) + r ( G2 L + 2∆0 ) δ(N + 1) + ∆z mδ(N)δ(N + 1) +A1 N∑ t=1
1
mδ(t)2  +
1∑ δ ( ∆z m N∑ t=1 1 δ(t)2 + r ( G2 L + 2∆0 ) N∑ t=1 1 δ(t) +B1G 2MT )
+ 1∑ δ ( N∑ t=1 A2G 2 δ(t) + ∆0 (1− β)m N∑ t=1 1 δ21(t) ) . (69)
The above result can be re-written as (recall δ(1) = 1)
E[‖∇f(x̂T )‖2]
≤ 1∑ δ
W 11 + r ( G2 L + 2∆0 ) δ(N + 1) + ∆z mδ(N)δ(N + 1) + (A1 + ∆z + ∆0 1−β ) m N∑ t=1 1 δ(t)2  +
1∑ δ
(( r ( G2
L + 2∆0
) +A2G 2 ) N∑ t=1 1 δ(t) +B1G 2MT )
≤ 1∑ δ
( W 11 +
C0 δ(N + 1) + + ∆z mδ(N)δ(N + 1) + C1 m N∑ t=1 1 δ(t)2 + C2 N∑ t=1 1 δ(t) +B1G 2MT ) where C0 = r(G 2 L + 2∆0), C1 = A1 + ∆z + ∆0 1−β , and C2 = C0 +A2G 2.
Proof. (of Theorem 4.2) In this case, given the total number iteration T ≥ 1, the number of stages N ≥ 1, St = S = dT/Ne, δ(t) = 1/αt−1 for each 1 ≤ t ≤ N , then the boundary function at the final stage δ(N) = 1/αN−1 and δ(N + 1) = 1/αN . Applying the specific value of δ(t) and N gives
N∑ t=1 1 δ(t) = αN − 1 α− 1
(70)
N∑ t=1 1 δ(t)2 = α2N − 1 α2 − 1 . (71)
Then by St = dT/Ne and (70), we easily get∑ δ = N∑ t=1 St/δ(t) ≥ T (αN − 1) N(α− 1) . (72)
We then plug the above results into Lemma C.3,
E[‖∇f(x̂T )‖2]
≤ N(α− 1) T (αN − 1)
( W 11 + C0α N + ∆zα 2N−1
m + C1(α 2N − 1) m(α2 − 1) + C2(α N − 1) α− 1 +MB1G 2T ) ≤W 11 N
TαN−1 + (αC0 + C2)
N T + (∆z + C1) m
NαN
T +MB1G
2 N
αN−1 . (73)
where C0 = r(G 2 L + 2∆0), C1 = A1 + ∆z + ∆0 1−β , and C2 = C0 +A2G 2, A2 = 1 + β2(1−β)2 , and W 11 , A1, B1, ∆z , and r are defined in Lemma 4.1.
Especially, we consider the number of outer-stage N = b(logα T )/2c, the stage length St = d2T/ logα T e, and the boundary functions δ(t) = 1/αt−1 for all t ∈ {1, 2, · · · , N}. Let N = b(logα T )/2c, we have
E[‖∇f(x̂T )‖2] ≤ α2W 11 2 lnα lnT T 3/2 + (αC0 + C2) 2 lnα lnT T + (∆z + C1) 2m lnα lnT√ T + α2MB1G 2 2 lnα lnT√ T .
(74) Therefore, we complete the proof.
Proof. (of Theorem 4.3) In this theorem, we consider the boundary functions δ(t) = 1/αt−1, and the stage length St = dS0αt−1e with S0 = √ T , then we have the number of stages N = blogα((α − 1) √ T + 1)c and δ(N) = αN−1. Next we estimate ∑ δ = ∑N t=1 Stδ
−1(t) ≥ (α − 1)2T 3/2 + 2(α− 1)T . Then applying these results into Lemma C.3, we have
E[‖∇f(x̂T )‖2] ≤ O ( W 11 T 3/2 + C0 T + ∆z m √ T + C1 m √ T + MB1G 2 √ T ) . (75)

heading: D SUPPLEMENTARY CONVERGENCE RESULTS
text: D.1 CONVERGENCE OF SGDM WITH CONSTANT AND 1/ √ t BANDWIDTH DECAYING
STEP-SIZE
We focus on a single stage that N = 1. In this case, we first consider the constant step-size ηti = η0/ √ T . Recalling the result of Lemma 4.1 and letting ηti = η0/ √ T for each i ≥ 1, we have
1
T T∑ i=1 E[ ∥∥∇f(x1i )∥∥2] ≤ (W 11T + B1G2η0√T ) where B1 = r(1− β)(2− β) + L2(1−β)2 and r = βL 2(1−β2)(1−β)2 . Next we turn to analyze the 1/ √ t bandwidth step-size η1i ∈ [m/ √ i,M/ √ i] (which is also monotonic decreasing). Recalling the result of Lemma 4.1
E[W ti+1 | F ti ] ≤W ti +A1 ( 1
ηti − 1 ηti−1
) − ∥∥∇f(xti)∥∥2 + ηtiB1G2 (76)
and applying the result from i = 1 to T , we have
1
T T∑ i=1 E[ ∥∥∇f(x1i )∥∥2] ≤ 1T (( W 11 − E[W 1T+1] ) +A1 ( 1 η1T − 1 η10 ) +B1G 2 T∑ i=1 η1i ) .
Then applying the step-size η1i ∈ [m/ √ i,M/ √ i] gives
1
T T∑ i=1 E[ ∥∥∇f(x1i )∥∥2] ≤ (W 11T + A1m√T + 2MB1G2√T ) . (77)
Thus, we can achieve an O(1/ √ T ) optimal rate for SGDM with 1/ √ t bandwidth step-size on nonconvex problems. When M = m, then the step-size reduces to η1i = m/ √ i, we also provide the
convergence guarantee for the commonly used 1/ √ i decaying step-size.
D.2 CONVERGENCE GUARANTEES FOR CYCLICAL STEP-SIZES
In Loshchilov and Hutter (2017), the authors proposed a cosine annealing step-size
ηt = ηtmin + 1
2
( ηtmax − ηtmin ) (1 + cos(Tcur · π/Tt)). (78)
where ηtmin and η t max are ranges of the step-size, and Tcur accounts for how many epochs since the beginning of the current stage and Tt accounts for the current stage length (epoch). At each stage t, the step-size is monotonically decaying within the range ηtmin and η t max. In this paper, we propose a general bandwidth framework for step-size which can cover this situation as long as mδ(t) ≤ ηtmin and ηtmax ≤Mδ(t). If the ranges {ηtmin, ηtmax} and the stage length are chosen as for example 1/ √ t in Theorems A.1 and A.2 or step-decay in Theorems 3.1 and 3.3, the theoretical convergence of SGD under the cosine annealing step-size is guaranteed by our analysis in Section 3. Moreover, because the cosine annealing is monotonic at each stage, so the convergence of SGD with momentum under the cosine annealing policy is also guaranteed by the analysis of Section 4 as long as the ranges ηtmin, η t max are within our bands. To the best of our knowledge, Li et al. (2021) provides a convergence guarantee for cosine step-size. However, to achieve a near-optimal rate for the general smooth (non-convex) problems, the initial step-size is required to be bounded by O(1/ √ T ) which is obviously impractical when the total number of iteration T is large (also discussed in related work). In our framework, the cosine step-size is allowed to start from a larger step-size and gradually decay. Besides, our results (e.g., Theorems 3.3 and 4.3) provide state-of-the-art convergence guarantees for cosine step-size which remove the log T term of Li et al. (2021).
Another interesting example is triangular cyclical step-size proposed by Smith (2017), which sets minimum ηtmin and maximum η t max boundaries and the learning rate cyclically varies (linearly
increasing then linearly decreasing) in these bounds. In each stage, the step-size is non-monotonic. In their paper, the author also consider a variant which cuts ηtmin and η t max in half after each stage. This is exactly the step-decay boundary we discussed. Our analysis in Section 3 can provide the convergence guarantees for such kinds of step-sizes. However, such cyclical step-size is not monotonic in each stage, so our analysis for SGDM in Section 4 is not suitable for this situation.

heading: E ADDITIONAL DETAILS OF THE EXPERIMENTS ON BANDWIDTH STEP-SIZES
text: In this section, we provide additional details about the numerical experiments in Section 5.
E.1 HOW TO DESIGN THE BANDWIDTH STEP-SIZES AND SELECT PARAMETERS
To better understand the bandwidth step-sizes tested in the numerical experiments, we visualize the step-size ηti (y-axis is log(η t i)) vs the number of epochs in Figure 5. We first consider the popular “step-decay” policy as the baseline. During the first stage, the bandwidth step-size follows the lower bound. From the second stage and on, we let the initial step-size in each stage to be equal to the upper bound and the last step-size in the stage to reach the lower bound. Our numerical experience has shown that the best performance is obtained when the initial step-size of each stage is larger than the final step-size of the previous stage, which means that the step-size experiences a sudden increase before it decreases again. For the step-decay band, we consider four decay modes: 1/i, 1/ √ i, linearly, and according to a cosine function (Loshchilov and Hutter, 2017) and update the step size each epoch. If the training size is n and sample size per iteration is b, then one epoch is n/b iterations. We also adopt the polynomial 1/ √ t step-size as the boundary function, named 1/ √ t-band, and update the step-size every epoch. We add similar perturbation as for the step-decay band, but we do not apply the perturbation per stage. Otherwise, the perturbation is too frequent and just increases the variance of the iterates. We tune the frequency of the perturbations (denoted as N0) to undergo a similar number of cycles as the step-decay perturbations. In the first cycle, the bandwidth step-sizes agree with their lower bound, just as for step-decay. From the second cycle, we begin to add the decreasing perturbations, e.g., 1/ √ i, 1/i and linearly. As discussed above, these perturbations are only adjusted between stages. Several different 1/ √ t bandwidth step-sizes are shown in Figure 5.
We perform a grid search for the initial step-size η0 ∈ {0.01, 0.05, 0.1, 0.5, 1, 5} of the baseline stepsizes. For step-decay step-size (baseline), we select the decay factor α from {1.5, 2, 3, 4, · · · , 12} and set the number of stages N = blogα T/2c according to Theorems 3.1 and 4.2. We choose the lower bound parameter m = η0 to agree with the baseline. The bandwidth s = M/m = αθ where θ ∈ {0.5, 0.8, 1, 1.2, 1.3, 1.5, 1.8}. If θ > 1, it means that the starting step-size at the current stage is larger than the ending step-size from the previous stage. For the 1/ √ t-band step-sizes, we choose ηti = η0/(1 + a √ t) as the baseline , and select the best a > 0 to make the final step-size reach the interval {0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1}. Moreover, we select the lower bound parameter m = η0 to make sure that the lower bound agrees with the baseline. For the upper bound parameter M , we do a grid search for s = M/m ∈ {2, 3, 4, 5, 6}. The number of perturbation cycles N0 for the 1/ √ t-band step-sizes is chosen from {1, 2, 3, 4}. All the hyper-parameters are selected to work best according to their performance on the test dataset.
E.2 EXPERIMENTS DETAILS ON CIFAR10 AND CIFAR100 DATASETS
In this subsection, we will give the implementation details for the experiments on CIFAR10 and CIFAR100. The benchmark datasets CIFAR10 and CIFAR100 (Krizhevsky, 2009) both consist of 60000 colour images (50000 training images and the rest 10000 images for testing). The maximum epochs called for the two datasets is 180 and the batch size is 128. All the experiments on CIFAR datasets are implemented in Python 3.7.4 and run on 2 x Nvidia Tesla V100 SXM2 GPUs with 32GB RAM. All experiments are repeated 5 times to eliminate the effect of the randomness. The performance of different algorithms is evaluated in terms of their loss function value and classification accuracy on the training and test datasets. All the results for training loss, test loss, and test accuracy are reported in Figures 6, 7, and 3.
For CIFAR 10, we train an 18-layer Resident Network model (He et al., 2016) called ResNet-18. We first test the vanlia SGD with a weight decay of 0.0005. The initial step-size η0 = 1 and a = 1.41618 for 1/ √ t step-size (baseline). For 1/ √ t-band step-sizes, we set: s = 2 and N0 = 3 for 1/ √ i mode; s = 3 and N0 = 3 for 1/i mode; and s = 4 and N0 = 3 for linear mode. For the step-decay step-size (baseline) and also the bandwidth step-sizes, the initial step-size η0 = 0.5 and decay factor α = 6. For step-decay band step-sizes, the parameter θ is 1.3 for the 1/ √ i, 1/i and linear modes, and is 1.2 for cosine mode. We also implement the SGD with momentum (SGDM) algorithm, with the momentum parameter of 0.9 and a weight decay of 0.0005. For the step-decay step-size, the initial step-size η0 is 0.05 and the decay factor α is 6. We choose the same initial step-size and decay factor for the step-decay bandwidth step-sizes. The best θ is 1.3 for the four decay modes.
In a similar way, we also detail our parameter selection for the experiments on CIFAR100. On this data set, we train a 28× 10 wide residual network (WRN-28-10) (Zagoruyko and Komodakis, 2016). We first implement vanilla SGD with a weight decay of 0.0005. For the baseline of the 1/ √ t band, we set η0 = 0.5 and a = 3.65224. For the 1/ √ t-band step-sizes we use: s = 2 and N0 = 3 for
1/ √ i-mode; s = 4 and N0 = 3 for 1/i-mode; s = 4 and N0 = 2 for linear-mode. For the step-decay band, we choose η0 = 0.5 and α = 6. The parameter θ is set to 1.3 for 1/i mode and 1.2 for the other modes. Then we also apply the step-decay band step-sizes on the SGD with momentum (SGDM) algorithm, where the momentum parameter is 0.9 and the weight decay parameter is 0.0005. We set the initial step-size η0 = 0.1 and α = 6 for the baseline and other step-decay band step-sizes; θ = 1.2 for 1/ √ i and linear modes; and θ = 1.3 for 1/i and cosine modes.
E.3 EXPERIMENTS RESULTS OF THE TOY EXAMPLE ON BANDWIDTH STEP-SIZE
In this subsection, we describe the toy example and also report additional results using other bandwidth step-sizes.
The loss-function is the two-dimensional (x, y ∈ R are the variables) non-convex function: f(x, y) = ( (x+ 0.7)2 + 0.1 ) (x− 0.7)2 + (y + 0.7)2 ( (y − 0.7)2 + 0.1 ) which has four local minima (denoted by 1 to 4 3), one of which is global ( 4 ). We execute 10000 algorithm runs with an initial point (−0.9, 0.9). The total number of iterations is set to T = 3000. The gradient noise is drawn from the standard normal distribution. The setting of the experiments follows (Shi et al., 2020). The step-sizes we tested in this part are similar to Section E.1. The difference is that here we update the step-size per iterate instead of per epoch, as we did on the CIFAR datasets. We report the percentage (%) of the final iterate close to each local minima in Table 2. Note that the results for constant (large and small) step-size, step-decay (baseline), and step-decay with linear have already been presented in the introduction of the main document. The large constant step-size is 0.1 and the small constant step-size is 0.05. As we can see, 1/ √ t-band with 1/ √ i, 1/i and linear modes more likely to escape the bad local minima and find the global solution than their baseline. Except the result of step-decay with linear (shown in Figure 2 and Table 1), we also find that other bandwidth step-sizes achieve good performance and work better than the baseline.
In the toy experiment, we also test the performance of step-decay (baseline) and step-decay with linear on different initial points. We evenly select 100 initial points x0 from the region [−1, 1]× [−1, 1]. At each initial point, we repeat the same process as previous initial point (−0.9, 0.9) and do 10000 runs. We record the percentage of final iterate close to global minima 4 at each initial point in Figure 8. The x-axis denotes the distance of the initial point x0 and the global minima x∗. We can see that compared to its baseline, step-decay with linear is less sensitive to the selection of the initial point and a relatively large θ > 1 works better if the initial point is far from the global minima.
3Notation: 1 denotes the local minima at (−0.7, 0.7); 2 denotes the local minima at (0.7, 0.7) ; 3 denotes the local minima at (−0.7,−0.7); 4 denotes the global minima at (0.7,−0.7).

Review summaries:
This paper presents a general framework for analyzing SGD with a bandwidth-based step size. This step size scheme uses a monotonically decreasing boundary function along with upper and lower bound constants to cover several non-monotonic step sizes strategies, including epoch-wise step decay, cosine annealing and triangular step sizes that are commonly adopted in training deep neural networks. They provide near-optimal convergence guarantees for smooth, non-convex functions when the boundary function is $1/\sqrt{t}$, where $t$ can be thought of as the epoch counter. The analysis works by dividing the total number of iterations into stages, where the stage length can either be constant or increasing exponentially. By appropriately selecting the number of stages and the stage length, they obtained optimal rate for the stage-wise step decay scheme. Furthermore, near-optimal and optimal rates are also provided for SGD with heavy-ball momentum under the bandwidth step size assumption that covers the stage-wise step decay setting. Lastly, they proposed non-monotonic step size schedules (step-decay with linear/cosine perturbation) and compared their empirical performance on several deep learning benchmarks.

This paper investigates a general class of step-sizes for stochastic gradient algorithms. The step-sizes are chosen based on some boundary function $\delta(t)$ where $\delta(t)$ decays exponentially fast throughout the stages. Then the step-sizes are bounded between two bandwidth $m \delta(t)$ and $M \delta(t)$ where $m < M$. The bandwidth-based scheme allows some flexibilities in the choice of learning rates, although the theoretical results are restricted to some settings where every hyperparameter is carefully chosen. The authors provide the theoretical results for these settings for stochastic gradient descent (SGD) and for momentum version (SGD-M), and show that these schemes attain the standard rate of convergence for this nonconvex stochastic problem. They show the numerical results for these step-sizes using a toy example and classification example for CIFAR dataset.  

This paper studies the bandwidth based step size optimization mainly from theoretical guarantee and experimental validation is also provided. For step-decay SGD, the paper proves the non-asymptotic convergence and an optimal rate is derived. In the SGDM scenario, an optimal convergence result is provided and beats the result from previous literature. Bandwidth schedules are first tested on a toy example showing its power to avoid local minima, then further experiments of SGDM and SGD are performed on CIFAR dataset and achieve a better final accuracy compared with baseline.