Bandwidth-based Step-Sizes for Non-Convex  Stochastic Optimization

√ 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 . 1 INTRODUCTION . 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 . 1.1 CONTRIBUTIONS . 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 . 1.2 RELATED WORK . 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 . 2 PROBLEM SET-UP . 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 .

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