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

Abstract:
Stochastic model-based methods have received increasing attention lately due to their appealing robustness to the stepsize selection and provable efficiency guarantee. We make two important extensions for improving model-based methods on stochastic weakly convex optimization. First, we propose new minibatch modelbased methods by involving a set of samples to approximate the model function in each iteration. For the first time, we show that stochastic algorithms achieve linear speedup over the batch size even for non-smooth and non-convex (particularly, weakly convex) problems. To this end, we develop a novel sensitivity analysis of the proximal mapping involved in each algorithm iteration. Our analysis appears to be of independent interests in more general settings. Second, motivated by the success of momentum stochastic gradient descent, we propose a new stochastic extrapolated model-based method, greatly extending the classic Polyak momentum technique to a wider class of stochastic algorithms for weakly convex optimization. The rate of convergence to some natural stationarity condition is established over a fairly flexible range of extrapolation terms. While mainly focusing on weakly convex optimization, we also extend our work to convex optimization. We apply the minibatch and extrapolated model-based methods to stochastic convex optimization, for which we provide a new complexity bound and promising linear speedup in batch size. Moreover, an accelerated modelbased method based on Nesterov’s momentum is presented, for which we establish an optimal complexity bound for reaching optimality.

heading: 1 Introduction
text: 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.

heading: 2 Background
text: 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.

heading: 6 Experiments
text: In this section, we examine the empirical performance of our proposed methods through experiments on the problem of robust phase retrieval. (Additional experiments on blind deconvolution are given in Appendix section). Given a set of vectors ai ∈ Rd and nonnegative scalars bi ∈ R+, the goal of phase retrieval is to recover the true signal x∗ from the measurement bi = |〈ai, x∗〉|2. Due to the potential corruption in the dataset, we consider the following penalized formulation
min x ∈ Rd
1
n n∑ i=1 ∣∣〈ai, x〉2 − bi∣∣ (20) where we impose `1-loss to promote robustness and stability (cf. [15, 9, 25]).
Data Preparation. We conduct experiments on both synthetic and real datasets.
1) Synthetic data. Synthetic data is generated following the setup in [25]. We set n = 300, d = 100 and select x∗ from unit sphere uniformly at random. Moreover, we generate A = QD where Q ∈ Rn×d, qij ∼ N (0, 1) and D ∈ Rd is a diagonal matrix whose diagonal entries are evenly distributed in [1/κ, 1]. Here κ ≥ 1 plays the role of condition number (large κ makes problem hard). The measurements are generated by bi = 〈ai, x∗〉2 + δiζi (1 ≤ i ≤ n) with ζi ∼ N (0, 25), δi ∼ Bernoulli(pfail), where pfail ∈ [0, 1] controls the fraction of corrupted observations on expectation. 2) Real data. We consider zipcode, a dataset of 16 × 16 handwritten digits collected from [21]. Following the setup in [15], let H ∈ R256×256 be a normalized Hadamard matrix such that hij ∈{
1 16 ,− 1 16 } , H = HT and H = H−1. Then we generate k = 3 diagonal sign matrices S1, S2, S3 such that each diagonal element of Sk is uniformly sampled from {−1, 1}. Last we set A = [HS1, HS2, HS3]
T ∈ R(3×256)×256. As for the true signal and measurements, each image is represented by a data matrix X ∈ R16×16 and gets vectorized to x∗ = vec(X). To simulate the case of corruption, we set measurements b = φpfail(Ax
∗), where φpfail(·) denotes element-wise squaring and setting a fraction pfail of entries to 0 on expectation.
In the first experiment, we illustrate that SMOD methods enjoy linear speedup in the size of minibatches and exhibit strong robustness to the stepsize policy. We conduct comparison on SPL and SGD and describe the detailed experiment setup as follows.
1) Dataset generation. We generate four testing cases: the synthetic datasets with (κ, pfail) = (10, 0.2), and (10, 0.3); zipcode with digit images of id 2 and 24;
2) Initial point. We set the initial point x1(= x0) ∼ N (0, Id) for synthetic data and x1 = x∗ + N (0, Id) for zipcode;
3) Stopping criterion. We set the stopping criterion to be f(xk) ≤ 1.5f̂ , where f̂ = f(x∗) is the corrupted objective evaluated at the true signal x∗;
3) Stepsize. We set the parameter γ = α−10 √ K/m where m is the batch size; For synthetic dataset, we test 10 evenly spaced α0 values in range [10−1, 102] on logarithmic scale, and for zipcode dataset we set such range of α0 to [101, 103];
4) Maximum iteration. We set the maximum number of epochs to be 200 and 400 respectively for minibatch and momentum related tests;
5) Batch size. We take minibatch size m from the range {1, 4, 8, 16, 32, 64}; 6) Sub-problems The solution to the proximal sub-problems is left in the appendix.
For each algorithm, speedup from minibatching is quantified as T ∗1 /T ∗ m where T ∗ m is the total number of iterations for reaching the desired accuracy, with batch size m and the best initial stepsize α0 among values specified above. Specially, if an algorithm fails to reach desired accuracy after running out of 400 epochs, we set its iteration number to the maximum.
Figure 1 plots the speedup of each algorithm over different values of batch size according to the average of 20 independent runs. It can be seen that SPL exhibits a linear acceleration over the batch size, which confirms our theoretical analysis. Moreover, we find SGD admits considerable acceleration using minibatches, and sometimes the speedup performance matches that of SPL and SPP. This observation seems to suggest the effectiveness of minibatch SGD in practice, despite the lack of theoretical support.
Next, we investigate the sensitivity of minibatch acceleration to the choice of initial stepsizes. We plot the algorithm speedup over the initial stepsize α0 in Figure 2 (1st row). It can be readily seen that SGD, SPL and SPP all achieve considerable minibatch acceleration when choosing the initial stepsize properly. However, SPL and SPP enjoy a much wider range of initial stepsizes for good speedup performance, and hence, lays more robust performance than SGD. To further illustrate the robustness of SPL and SPP, we compare the efficiency of both algorithms in the minibatch setting. In contrast to the previous comparison on the relative scale, we directly compare the iteration complexity of the two algorithms. We plot the total iteration number over the choice of initial stepsizes in Figure 2 (2nd row) for batch size m = 8 and 32. We observe that minibatch SPL(SPP)s exhibits promising performance for a wide range of stepsize policies, while minibatch SGD quickly diverges for large stepsizes. Overall, our experiment complements the recent work [9], which shows that SPL (SPP) is more robust than SGD in the sequential setting.
Our second experiment investigates the performance of the proposed momentum methods. We compare three model-based methods (SGD, SPL, SPP) and extrapolated model-based methods (SEGD, SEPL, SEPP). We generate four testing cases: the synthetic datasets with (κ, pfail) = (10, 0.2) and
(10, 0.3); zipcode with digit images of id 2 and pfail ∈ {0.2, 0.3}. We set α0 ∈ [10−2, 100], β = 0.6 for synthetic data, and set α0 ∈ [100, 101], β = 0.9 for zipcode dataset. The rest of settings are the same as in minibatch with m = 1.
Figure 3 plots the number of epochs to ε-accuracy over initial stepsize a0. It can be seen that with properly selected momentum parameters (SEGD, SEPL, SEPP) all suggest improved convergence when stepsize is relatively small.
In the last experiment, we attempt to exploit the performance of the compared algorithms when minibatching and momentum are applied simultaneously. The parameter setting is the same as that of the second experiment, except that we choose m ∈ {1, 32}. Results are plotted in Figure 4 and it can be seen that minibatch SMOD, when combined with momentum, exhibits even better convergence performance and robustness.

heading: 7 Discussion
text: On a broad class of non-smooth non-convex (particularly, weakly convex) problems, we make stochastic model-based methods more efficient by leveraging minibatching and momentum—two techniques that are well-known only for SGD. Applying algorithm stability for optimization analysis is a key step to achieving improved convergence rate over the batch size. This perspective appears to be interesting for stochastic optimization in a much broader context. Although some progress is made, we are unable to show whether minibatches can accelerate SGD when the objective does not have a smooth component. Note that the complexity of SGD already matches the best bound of full subgradient method. It would be interesting to know whether this bound for SGD is tight or improvable. It would also be interesting to study the lower bound of SGD (and other stochastic algorithms) in the non-smooth setting. Some interesting recent results can be referred from [22, 33].

heading: 8 Acknowledgement and disclosure of funding
text: The authors are grateful to the Area Chairs and the anonymous reviewers for their constructive suggestions. QD was partially supported by National Natural Science Foundation of China (Grant 11831002, 72150001).

Review summaries:
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.

The paper proposed stochastic model-based method and its momentum version for solving non-smooth non-convex optimization problems. The convergence results of the proposed methods are established. The authors also provide some experimental results for verifying the effectiveness of the proposed methods.

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.