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: The paper presents several minibatch and momentum extensions of stochastic model-based optimization methods (SMOD) for learning with non-smooth and weakly convex loss functions. For SMOD with minibatch size $m$ and iteration count $K$, the authors showed via a uniform stability argument that the algorithm converges to stationary point at the rate of $O(\frac{1}{\sqrt{mK}})$. Similar convergence rates have been established for an extrapolated variant of SMOD with or without minibatching. The actual performance of the proposed methods is evaluated in a set of experiments on synthetic and real-data robust phase retrieval.