Minibatch and Momentum Model-based Methods for Stochastic Weakly Convex Optimization

1 Introduction . In this paper , we are interested in the following stochastic optimization problem : min x ∈ X f ( x ) = Eξ∼Ξ [ f ( x , ξ ) ] ( 1 ) where f ( · , ξ ) stands for the loss function , sample ξ follows certain distribution Ξ , and X is a closed convex set . We assume that f ( · , ξ ) is weakly convex , namely , the sum of f ( x , ξ ) and a quadratic function λ2 ‖x‖ 2 is convex ( λ > 0 ) . This type of non-smooth non-convex functions can be found in a variety of machine learning applications , such as phase retrieval , robust PCA and low rank decomposition [ 8 ] . To solve problem ( 1 ) , we consider the stochastic model-based method ( SMOD , [ 14 , 9 , 1 ] ) , which comprises a large class of stochastic algorithms ( including stochastic ( sub ) gradient descent , proximal point , among others ) . Recent work [ 14 , 9 ] show that SMOD exhibits promising convergence property : both asymptotic convergence and rates of convergence to certain stationarity 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . measure have been established for the SMOD family . In addition , empirical results [ 9 , 15 ] indicate that SMOD exhibits remarkable robustness to hyper-parameter tuning and often outperforms SGD . Despite much recent progress , our understanding of model-based methods for weakly convex optimization is still quite limited . Particularly , it is still unknown whether SMOD is competitive against modern SGD used in practice . We highlight some important remaining questions . First , despite the appealing robustness and stable convergence , the SMOD family is sequential in nature . It is unclear whether minibatching , which is immensely used in training learning models , can improve the performance of SMOD when the problem is non-smooth . Particularly , the current best complexity bound O ( L 2 ε4 ) from [ 9 ] , which is regardless of batch size , is unsatisfactory . Were this bound tight , a sequential algorithm ( using one sample per iteration ) would be optimal : it offers the highest processing speed per iteration as well as the best iteration complexity . Therefore , it is crucial to know whether minibatching can improve the complexity bound of the SMOD family or the current bound is tight . Second , in modern applications , momentum technique has been playing a vital role in large-scale non-convex optimization ( see [ 31 , 28 ] ) . In spite of its effectiveness , to the best of our knowledge , momentum technique has been provably efficient only in 1 ) unconstrained smooth optimization [ 23 , 10 , 19 ] and 2 ) non-smooth optimization with a simple constraint [ 25 ] , which constitute only a portion of the interesting applications . From the practical aspect , it is peculiarly desirable to know whether momentum technique is applicable beyond in SGD and whether it can benefit the SMOD algorithm family in the non-smooth and non-convex setting . Contributions . Our work is motivated by the aforementioned challenge to make SMOD more practically efficient . We summarize the contributions as follows . First , we extend SMOD to the minibatch setting and develop sharper rates of convergence to stationarity . Leveraging the tool of algorithm stability ( [ 6 , 27 , 20 ] ) , we provide a nearly complete recipe on when minibatching would be helpful even in presence of non-smoothness . Our theory implies that stochastic proximal point and stochastic prox-linear are inherently parallelizable : both algorithms achieve linear speedup over the minibatch size . To the best of our knowledge , this is the first time that these minibatch stochastic algorithms are proven to exhibit such an acceleration even for non-smooth and non-convex ( particularly , weakly convex ) optimization . Moreover , our theory recovers the complexity of minibatch ( proximal ) SGD in [ 9 ] , showing that ( proximal ) SGD enjoys the same linear speedup by minibatching for smooth composite problems with non-smooth regularizers or with constrained domain . Second , we present new extrapolated model-based methods by incorporating a Polyak-type momentum term . We develop a unified Lyapunov analysis to show that a worst-case complexity of O ( 1/ε4 ) holds for all momentum SMOD algorithms . To the best of our knowledge , these are the first complexity results of momentum stochastic prox-linear and stochastic proximal point for non-smooth non-convex optimization . Since our analysis offers complexity guarantees for momentum SGD and its proximal variant , our work appears to be more general than a recent study [ 25 ] , which only proves the convergence of momentum projected SGD . Proximal SGD is more advantageous in composite optimization , where the non-smooth term is often involved via its proximal operator rather than the subgradient . For example , in the Lasso problem , it is often favorable to invoke the proximal operator of ` 1 function ( Soft-Thresholding ) to enhance solution sparsity . We summarize the complexity results in Table 1 . Third , we develop new convergence results of SMOD for convex optimization , showing that minibatch extrapolated SMOD achieves a promising linear speedup over the batch size under some mild condition . Specifically , to obtain some ε-optimal solution , our proposed method exhibits an O ( 1/ε+ 1/ ( mε2 ) ) complexity bound in the worst case . Moreover , we develop a new minibatch SMOD based on Nesterov ’ s momentum , achieving the O ( 1/ε1/2 + 1/ ( mε2 ) ) optimal complexity bound . Note that a similar complexity result , explicitly relying on the smoothness assumption , has been shown in a recent study [ 7 ] . Compared to this work , our analysis makes weaker assumptions , showing that smoothness is not a must-have for many model-based algorithms , such as SPL and SPP , to get sharper complexity bound . Other related work . For smooth and composite optimization , it is well known that SGD can be linearly accelerated by minibatching ( c.f . [ 11 , 18 , 29 ] ) . Minibatch model-based methods have been studied primarily in the convex setting . Asi et al . [ 2 ] investigates the speedups of minibatch stochastic model-based methods in the convex smooth , restricted strongly convex and convex interpolation settings , respectively . Since their assumptions differ from ours , the technique does not readily apply to the non-convex setting . Chadha et al . [ 7 ] studies the accelerated minibatch model-based methods for convex smooth and convex interpolated problems . The interpolation setting , where the model can perfectly fit the data , is not considered in our paper . Algorithm stability [ 6 , 27 ] —an important technique for analyzing the generalization performance of stochastic algorithms [ 20 , 3 ] , is the key tool to obtain some of our convergence results . In contrast to the traditional work , our paper employs the stability argument to obtain sharper optimization convergence rates ( with respect to the batch size ) . See Section 3 . As noted by an anonymous reviewer , a similar idea of using stability analysis was proposed by Wang et al . [ 30 ] , albeit with a different motivation from distributed stochastic optimization . Robustness and fast convergence of model-based methods have been shown on various statistical learning problems [ 8 , 15 , 1 , 4 , 16 , 5 ] . Drusvyatskiy and Paquette [ 13 ] give a complete complexity analysis of the accelerated proximal-linear methods for deterministic optimization . Zhang and Xiao [ 32 ] further improve the convergence rates of prox-linear methods on certain finite-sum and stochastic problems by using variance-reduction . Momentum and accelerated methods for convex stochastic optimization can be referred from [ 24 , 26 ] . The study [ 10 , 23 , 31 ] develop the convergence rate of stochastic momentum method for smooth non-convex optimization . 2 Background . Throughout the paper , we use ‖ · ‖ to denote the Euclidean norm and 〈· , ·〉 to denote the Euclidean inner product . We assume that f ( x ) is bounded below . i.e. , minx f ( x ) > −∞ . The subdifferential ∂f ( x ) of function f ( x ) is the set of vectors v ∈ Rd that satisfy : f ( y ) ≥ f ( x ) + 〈v , y− x〉+ o ( ‖x− y‖ ) , as y → x . Any such vector in ∂f ( x ) is called a subgradient and is denoted by f ′ ( x ) ∈ ∂f ( x ) for simplicity . We say that a point x is stationary if 0 ∈ ∂f ( x ) + NX ( x ) , where the normal cone NX ( x ) is defined as NX ( x ) , { d : 〈d , y− x〉 ≤ 0 , ∀y ∈ X } . For a set S , define the set distance to 0 by : ‖S‖− , inf { ‖x− 0‖ , x ∈ S } . It is natural to use the quantity ‖∂f ( x ) +NX ( x ) ‖− to measure the stationarity of point x. Moreau-envelope . The µ-Moreau-envelope of f is defined by fµ ( x ) , miny∈X { f ( y ) + 12µ‖x− y‖2 } and the proximal mapping associated with f ( · ) is defined by proxµf ( x ) , argminy∈X { f ( y ) + 1 2µ‖x− y‖ 2 } . Assume that f ( x ) is λ-weakly convex , then for µ < λ−1 , the Moreau envelope fµ ( · ) is differentiable and its gradient is∇fµ ( x ) = µ−1 ( x− proxµf ( x ) ) . The SMOD family iteratively computes the proximal map associated with a model function fxk ( · , ξk ) : xk+1 = argmin x∈X { fxk ( x , ξk ) + γk 2 ‖x− xk‖2 } , ( 2 ) where { ξk } are i.i.d . samples . Typical algorithms and the accompanied models are described below . Stochastic ( Proximal ) Gradient Descent : consider the composite function f ( x , ξ ) = ` ( x , ξ ) +ω ( x ) where ` ( x , ξ ) is a data-driven and weakly-convex loss term and ω ( x ) is a convex regularizer such as ` 1-penalty . SGD applies the model function : fy ( x , ξ ) = ` ( y , ξ ) + 〈 ` ′ ( y , ξ ) , x− y 〉 + ω ( x ) . ( 3 ) Stochastic Prox-linear ( SPL ) : consider the composition function f ( x , ξ ) = h ( C ( x , ξ ) ) where h ( · , ξ ) is convex continuous and C ( x , ξ ) is a continuously differentiable map . We perform partial linearization to obtain the model fy ( x , ξ ) = h ( C ( y , ξ ) + 〈∇C ( y , ξ ) , x− y〉 ) . ( 4 ) Stochastic Proximal Point ( SPP ) : compute ( 2 ) with full stochastic function : fy ( x , ξ ) = f ( x , ξ ) . ( 5 ) Throughout the paper , we assume that f ( x , ξ ) is continuous and µ-weakly convex , and that the model function fx ( · , · ) satisfies the following assumptions [ 9 ] . A1 : For any ξ ∼ Ξ , the model function fx ( y , ξ ) is λ-weakly convex in y ( λ ≥ 0 ) . A2 : Tightness condition : fx ( x , ξ ) = f ( x , ξ ) , ∀x ∈ X , ξ ∼ Ξ. A3 : One-sided quadratic approximation : fx ( y , ξ ) − f ( y , ξ ) ≤ τ2‖x− y‖ 2 , ∀x , y ∈ X , ξ ∼ Ξ. A4 : Lipschitz continuity : There exists L > 0 that fx ( z , ξ ) − fx ( y , ξ ) ≤ L ‖z − y‖ , for any x , y , z ∈ X , ξ ∼ Ξ . Remark 1 . Assumption A2 is quite standard and will be used only in the convergence proof . Combining A1 and A3 , we immediately have that f ( x , ξ ) is ( λ+ τ ) -weakly convex . Thus , it suffices to assume that µ < τ + λ. Assumptions A2-A4 can be slightly relaxed by replacing the uniform bound with a bound on expectation over ξ , leading to only a minor adjustment to the analysis . Denote x̂ , proxf/ρ ( x ) = argminy { f ( y ) + ρ2‖y − x‖ 2 } for some ρ > µ. Davis and Drusvyatskiy [ 9 ] revealed a striking feature of Moreau envelope to characterize stationarity : ‖x̂− x‖ = ρ−1‖∇f1/ρ ( x ) ‖ , and ‖∂f ( x̂ ) +NX ( x̂ ) ‖− ≤ ‖∇f1/ρ ( x ) ‖ . Namely , a point x with small gradient norm ‖∇f1/ρ ( x ) ‖ stays in the proximity of a nearly-stationary point x̂ . With this observation , they show the first complexity result of SMOD for non-smooth nonconvex optimization : min1≤k≤K E [ ‖∇f1/ρ ( xk ) ‖ ] 2 ≤ O ( L√K ) . Note that this rate is regardless of the size of minibatches since it does not explicitly use any information of the samples other than the Lispchitzness of the model function . Due to this limitation , it remains unclear whether minibatching can further improve the convergence rate of SMOD . 3 SMOD with minibatches In this section , we present a minibatch SMOD method which takes a small batch of i.i.d . samples to estimate the model function . The overall procedure is detailed in Algorithm 1 . Within each iteration , Algorithm 1 forms a stochastic model function fxk ( · , Bk ) = 1mk ∑mk i=1 fxk ( x , ξk , i ) parameterized at xk by sampling over mk i.i.d . samples Bk = ξk,1 , . . . , ξk , mk . Then it performs proximal update to get the next iterate xk+1 . We will illustrate the main convergence results of Algorithm 1 and leave all the proof details in Appendix sections . But first , let us present an additional assumption . A5 : Two-sided quadratic bound : for any x , y ∈ X , ξ ∼ Ξ , ∣∣fx ( y , ξ ) − f ( y , ξ ) ∣∣ ≤ τ2‖x− y‖2 . Remark 2 . Assumption A5 is vital for our improved convergence analysis . While it is slightly stronger than A3 , A5 is indeed satisfied by the SMOD family in most contexts : 1 ) For SPP , A5 is trivially satisfied by taking fx ( y , ξ ) = f ( y , ξ ) . 2 ) For SPL , we minimize a composition function f ( x , ξ ) = h ( Cξ ( x ) ) where h ( · ) is a c1-Lipschitz convex function and Cξ ( · ) is a c2-Lipschitz smooth map . In view of ( 4 ) , A5 is verified with |fx ( y , ξ ) − f ( y , ξ ) | ≤ c1 ∥∥Cξ ( y ) −Cξ ( x ) −∇Cξ ( x ) T ( y− x ) ∥∥ ≤ c1c22 ‖x− y‖2 . 3 ) For SGD , A5 is satisfied if ` ( · , ξ ) is c3-Lipschitz smooth for some c3 > 0 , as |fx ( y , ξ ) − f ( y , ξ ) | ≤ | ` ( y , ξ ) − ` ( x , ξ ) −∇ ` ( x , ξ ) T ( y − x ) | ≤ c32 ‖x− y‖ 2 . We note that A5 is not satisfied by SGD when the loss ` ( · , ξ ) is also non-smooth . Unfortunately , there seems to be little hope to accelerate SGD in such a case since the convergence rate of SGD already matches the rate of deterministic subgradient method . We present an improved complexity analysis of SMOD by leveraging the framework of algorithm stability [ 6 , 27 ] . In stark contrast to its standard application in characterizing the algorithm generalization performance , stability analysis is applied to determine how the variation of a minibatch affects the estimation of the model function in each algorithm iteration . Algorithm 1 Stochastic Model-based Method with Minibatches ( SMOD ) Input : x1 , γk ; for k = 1 to K do Sample a minibatch Bk = { ξk,1 , . . . , ξk , mk } and update xk+1 by solving min x∈X { 1 mk mk∑ i=1 fxk ( x , ξk , i ) + γk 2 ∥∥x− xk∥∥2 } ( 6 ) end for Notations . Let B = { ξ1 , ξ2 , . . . , ξm } be a batch of i.i.d . samples and B ( i ) = B \ { ξi } ∪ { ξ′i } by replacing ξi with an i.i.d . copy ξ′i , and B ′ = { ξ′1 , ξ′2 , . . . , ξ′m } . Let h ( · , ξ ) be a stochastic model function , and denote h ( y , B ) = 1m ∑m i=1 h ( y , ξi ) . The stochastic proximal mapping associated with h ( · , B ) is defined by proxρh ( x , B ) , argminy∈X { h ( y , B ) + 12ρ‖y − x‖ 2 } for some ρ > 0 . We denote x+B , proxρh ( x , B ) for brevity . We say that the stochastic proximal mapping proxρh is ε-stable if , for any x ∈ X , we have∣∣EB , B′ , i [ h ( x+B ( i ) , ξ′i ) − h ( x+B , ξ′i ) ] ∣∣ ≤ ε , ( 7 ) where i is an index chosen from { 1 , 2 , . . . , m } uniformly at random . The next lemma exploits the stability of proximal mapping associated with the model function . Lemma 3.1 . Let fz ( · , B ) be a stochastic model function under the assumptions A1-A4 . For γ ∈ ( λ , ∞ ) , vectors z and y , the proximal mapping proxfz/γ ( y , B ) = argminx∈X { fz ( x , B ) + γ 2 ‖x− y‖2 } is ε-stable with ε = 2L 2 m ( γ−λ ) . Applying Lemma 3.1 , we obtain the error bound for approximating the full model function in the next theorem . Theorem 3.2 . Under all the assumptions of Lemma 3.1 , we have∣∣EBk [ fxk ( xk+1 , Bk ) − Eξfxk ( xk+1 , ξ ) |σk ] ∣∣ ≤ εk , εk = 2L2mk ( γk−λ ) . ( 8 ) where σk is the σ-algebra generating { Bi } 1≤i≤k−1 . Note that since xk+1 is dependent on Bk , fxk ( xk+1 , Bk ) is not an unbiased estimator of Eξ [ fxk ( xk+1 , ξ ) ] . However , the stability argument identifies that the expected approximation error is a decreasing function of batch size mk . This observation is the key to the sharp analysis of minibatch stochastic algorithms . With all the tools at our hands , we obtain the key descent property in the following theorem . Theorem 3.3 . Suppose that ρ > λ+ τ , γk ≥ ρ+ τ , A5 and all the assumptions in Lemma 3.1 hold . Let Ek [ · ] abbreviates EBk [ · |σk ] and εk be given by ( 8 ) , then we have ( ρ− λ− τ ) ρ ( γk + ρ− 2λ− τ ) ‖∇f1/ρ ( xk ) ‖2 ≤ f1/ρ ( xk ) − Ek [ f1/ρ ( x k+1 ) ] + ρεk γk + ρ− 2λ− τ . ( 9 ) Next , we specify the rate of convergence to stationarity using a constant stepsize policy . Theorem 3.4 . Under the assumptions of Theorem 3.3 , let ∆ = f1/ρ ( x1 ) − minx f ( x ) , mk = m , and γk = γ = max { ρ+ τ , λ+ η } where η = √ K α0 √ m and α0 ∈ ( 0 , ∞ ) . Let k∗ be an index chosen in { 1 , 2 , . . . , K } uniformly , then we have E [ ‖∇f1/ρ ( xk ∗ ) ‖2 ] ≤ ρ ρ− λ− τ [ ( 2ρ− λ ) ∆ K + ( ∆ α0 + 2α0ρL 2 ) 1√ mK ] . ( 10 ) Remark 3 . The performance of SMOD depends on α0 and batch size m. ( 10 ) implies that when batch size is fixed , the best rate is obtained at α∗0 = √ ∆ 2ρ 1 L . Since both ∆ and L are unknown , Algorithm 2 Stochastic Extrapolated Model-Based Method ( SEMOD ) Input : x0 , x1 , β , γ ; for k = 1 to K do Sample data ξk and update : yk = xk + β ( xk − xk−1 ) ( 11 ) xk+1 = argmin x∈X { fxk ( x , ξ k ) + γ 2 ‖x− yk‖2 } ( 12 ) end for hyper-parameter tuning over α0 is required to obtain good empirical performance . For the simplicity of theoretical analysis , let us take α0 = α∗0 . Hence , to obtain an iterate whose Moreau envelop has expected gradient norm smaller than ε , the total iteration count is Tε = max { O ( ∆ε2 ) , O ( L2∆ mε4 ) } . For small batch size m ( i.e . m = o ( 1/ε2 ) ) , the second term in max ( , ) dominates the bound Tε , yielding a total complexity of O ( L 2∆ mε4 ) . Note that this complexity bound is better than the O ( L2∆ ε4 ) bound [ 9 ] by a factor of m. Remark 4 . Theorem 3.4 implies that SGD can be accelerated by minibatching on the smooth composite problems ( 3 ) but leaves out the more general problems where ` ( x , ξ ) is non-smooth and weakly convex . In the latter case , showing any improved rate of minibatch SGD is substantially more challenging . Without additional knowledge , the O ( L 2∆ ε4 ) complexity of SGD already matches the best result for deterministic subgradient method ( c.f . [ 9 ] ) . It remains unknown whether such O ( 1/ε4 ) bound is tight or not , and a possible direction to obtain sharper complexity bound is by exploiting the non-smooth structure information such as sharpness . Solving the subproblems . SGD is embarrassingly parallelizable by simply averaging the stochastic subgradients . We highlight how to solve the proximal subproblems for SPL and SPP . Consider the composition function f ( x , ξ ) = h ( C ( x , ξ ) ) where h ( a ) = |a| . For SPL , it is easy to transform the corresponding subproblem to an O ( mk ) -dimensional quadratic program ( QP ) in the dual space ( e.g . [ 2 ] ) . The dual QP can be efficiently solved in parallel , for example , by a fast interior point solver . For SPP , we show that the subproblem can be solved by a deterministic prox-linear method at a rapid linear convergence rate . Note that the SPP subproblem is especially well-conditioned because our stepsize policy ensures a large strongly convex parameter γ − λ . We refer to the appendix for more technical details . 4 SMOD with momentum We present a new model-based method by incorporating an additional extrapolation term , and we record this stochastic extrapolated model-based method in Algorithm 2 . Each iteration of Algorithm 2 consists of two steps , first , an extrapolation step is performed to get an auxiliary update yk . Then a random sample ξk is collected and the proximal mapping , associated with the model function fxk ( · , ξk ) , is computed at yk to obtain the new point xk+1 . For ease of exposition , we take constant values of stepsize and extrapolation term . Note that Algorithm 2 can be interpreted as an extension of the momentum SGD by replacing the gradient descent step with a broader class of proximal mappings . To see this intuition , we combine ( 11 ) and ( 12 ) to get xk+1 = argmin x∈X { fxk ( x , ξ k ) + γβ〈xk−1 − xk , x− xk〉+ γ 2 ‖x− xk‖2 } , ( 13 ) If we choose the linear model ( 3 ) , i.e. , fxk ( x , ξk ) = f ( xk , ξk ) + 〈f ′ ( xk , ξk ) , x− xk〉 , and assume X = Rd , then the update ( 13 ) has the following form : xk+1 = xk − γ−1f ′ ( xk , ξk ) − β ( xk−1 − xk ) . ( 14 ) Define vk , γ ( xk−1 − xk ) and apply it to ( 14 ) , then Algorithm 2 reduces to the heavy-ball method vk+1 = f ′ ( xk , ξk ) + βvk , ( 15 ) xk+1 = xk − γ−1vk+1 . ( 16 ) Despite such relation , the gradient averaging view ( 15 ) only applies to SGD for unconstrained optimization , which limits the use of standard analysis of heavy-ball method ( [ 31 ] ) for our problem . To overcome this issue , we present a unified convergence analysis which can deal with all the model functions and is amenable to both constrained and composite problems . Our theoretical analysis of Algorithm 2 relies on a different potential function from the one in the previous section . Let us define the auxiliary variable zk , xk + β 1− β ( xk − xk−1 ) . ( 17 ) The following lemma proves some approximate descent property by adopting the potential function f1/ρ ( z k ) + ρ ( γβ+ρβ 2θ−2 ) 2 ( γθ−λθ ) ‖x k − xk−1‖2 and measuring the quantity of ‖∇f1/ρ ( zk ) ‖ . Lemma 4.1 . Assume that ρ ≥ 2 ( τ + λ ) and β ∈ [ 0 , 1 ) . Let θ = 1− β . Then we have ( ρ− λθ ) 2ρ ( γθ − λθ ) ‖∇f1/ρ ( zk ) ‖2 ≤ f1/ρ ( zk ) − Ek [ f1/ρ ( z k+1 ) ] + ρL2 ( γθ2 − ρβ2θ−1 ) ( γθ2 − λθ2 ) + ρ ( γβ + ρβ2θ−2 ) 2 ( γθ − λθ ) ( ‖xk − xk−1‖2 − Ek [ ‖xk+1 − xk‖2 ] ) − ρ ( γ − ρβ 2θ−3 ) 4 ( γ − λ ) Ek [ ‖xk+1 − xk‖2 ] . ( 18 ) Invoking Lemma 4.1 and specifying the stepsize policy , we obtain the main convergence result of Algorithm 2 in the following theorem . Theorem 4.2 . Under assumptions of Lemma 4.1 , if we choose x1 = x0 , and set γ = γ0θ−1 √ K + λ+ ρβ2θ−3 for some γ0 > 0 , then E [ ‖∇f1/ρ ( zk ∗ ) ‖2 ] ≤ 2ρ ρ− λ [ ρβ2θ−2∆ K + ( γ0∆ + ρL2 θγ0 ) 1√ K ] ( 19 ) where k∗ is an index chosen in { 1 , 2 , . . . , K } uniformly at random . Remark 5 . Despite the fact that convergence is established for all γ0 > 0 , we can see that the optimal γ0 would be γ0 = √ ρ ∆θL , which gives the bound E [ ‖∇f1/ρ ( z k∗ ) ‖2 ] ≤ 2ρρ−λ ( ρβ2θ−2∆ K +2L √ ρ∆ θK ) . In practice , we can set γ0 to a suboptimal value and obtain a possibly loose upper-bound . Remark 6 . Since zk is an extrapolated solution , it may not be feasible . It is desirable to show optimality guarantee at iterates xk . Note that using Lemma 4.1 and the parameters in Theorem 4.2 , it is easy to show that E [ ‖xk∗ − xk∗−1‖2 ] = O ( 1K ) . Based on ( 17 ) we have ‖z k∗ − xk∗‖2 = β2θ−2E [ ‖xk∗ − xk∗−1‖2 ] = O ( 1K ) . Using Lipschitz smoothness of Moreau envelop , we can show E [ ‖∇f1/ρ ( xk ∗ ) ‖2 ] converges at the same O ( 1√ K ) rate as is shown in Theorem 4.2 . Combining momentum and minibatching , we develop a minibatch version of Algorithm 2 that takes a batch of samples Bk in each iteration . The convergence analysis of this minibatch SEMOD is more involving . We leave the details in the Appendix but informally state the main result below . Theorem 4.3 ( Informal ) . In the minibatch SEMOD , suppose that A5 holds , the batch size |Bk| = m and γ = O ( √ K m ) , then E [ ‖∇f1/ρ ( z k∗ ) ‖2 ] = O ( 1 K + √ 1 mK ) . 5 SMOD for convex optimization Besides the study on non-convex optimization , we also apply model-based methods to stochastic convex optimization . Due to the space limit , we highlight main theoretical results but defer all the technical details to the Appendix section . We show that if certain assumption adapted from A5 for the convex setting holds , the function gap of minibatching SEMOD will converge at a rate of O ( 1 K + 1√ mK ) . In view of this result , the deterministic part of our rate is consistent with the best O ( 1K ) rate for the heavy-ball method . For example , see [ 12 , 17 ] . Moreover , the stochastic part of the rate is improved from the O ( 1√ K ) rate of Theorem 4.4 [ 9 ] by a factor of √ m. An important question arises naturally : Can we further improve the convergence rate of model-based methods for stochastic convex optimization ? Due to the widely known limitation of heavy-ball type momentum , it would be interesting to consider Nesterov ’ s acceleration . To this end , we present a model-based method with Nesterov type momentum . Thanks to the stability argument , we obtain the following improved rate of convergence : O ( 1 K2 + 1√ mK ) . We note that a similar convergence rate for minibatching model-based methods is obtained in a recent paper [ 7 ] . However , their result requires the assumption that the stochastic function is Lipschitz smooth while our assumption is much weaker .

Summary: This paper proves the effectiveness of using minibatch and momentum in stochastic model-based methods (SMOD) under the non-smooth weakly-convex setting. The authors utilize the concept in algorithm stability framework to show that minibatch SMOD has linear speedup over the batch size. They also give the first complexity results of momentum stochastic proximal-linear (SPL) and proximal point (SPP) algorithms for non-smooth weakly-convex optimization.