Minibatch and Momentum Model-based Methods for Stochastic Weakly Convex Optimization
Authors: Qi Deng
Wenzhi Gao
Venue: NeurIPS 2021 Poster (2021)
Keywords: model-based optimization
weakly convex optimization
stochastic optimization
algorithm stability
momentum methods

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 model- based 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 model-based method based on Nesterov’s momentum is presented, for which we establish an optimal complexity bound for reaching optimality.


Minibatch and Momentum Model-based Methods for Stochastic Weakly Convex Optimization
Qi Deng'
Wenzhi Gao?
School Information Management and Engineering Shanghai University of Finance and Economics
q "gidengssufe.ee edu.
"gwzd163. shufe. edu.
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 stochastic weakly convex optimization. First, we propose new minibatch model- based methods by involving a of samples to approximate the model function 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 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 extrapolation terms.
While mainly focusing on weakly convex optimization, we also extend our work convex optimization. We apply the minibatch and extrapolated model-based methods to stochastic convex optimization, for which we provide new complexity bound and promising linear speedup batch size. Moreover, an accelerated model- based method based on Nesterov's 's momentum is presented, for which we establish an optimal complexity bound for reaching optimality.
Introduction
In this paper, we are interested in the following stochastic optimization problem:
min X = E(=)=EE==(ffr,,
(1)
where {)) stands for the loss function, sample ę follows certain distribution E, and X' is closed convex set. We assume that f(-, is weakly convex, namely, the sum of f(x, {) and quadratic function =|2 is convex (A > 0). This type of non-smooth non-convex functions can be found in variety of machine learning applications, such as phase retrieval, robust PCA and low rank decomposition [8]. To solve problem ([). we consider the stochastic moodel-based method (SMOD, 11). which comprises 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 (NcurIPS 2021).

---

measure have been established for the SMOD family. In addition, empirical results 19 [5] indicate that SMOD exhibits remarkable robustness to hyper- parameeer tuning and often outperforms SGD.
Despite much recent progress, our understanding of model-based methods for weakly convex opti- mization 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 sequential in nature. It is unclear whether minibatching, which is immensely used in training learning models, can improve the perfor- mance of SMOD when the problem is non-smooth. Particularly, the current best complexity bound ( from ( which is regardless of batch size, is unsatisfactory. Were this bound tight, a sequen- tial 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, is crucial to know whether minibatching can improve the complexity bound of the SMOD family 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 and 2) non-smooth optimization with 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.
23 10 19
Contributions. Our work is motivated by the aforementioned challenge to make SMOD more practi- cally efficient. We summarize the contributions as follows. First, we extend SMOD to the minibatch setting and develop sharper rates of convergence stationarity. Leveraging the tool of algorithm stability (16) 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 brox-linear are inherently parallclizable: 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-conver (particularly, weakly convex) optimization. Morcover, 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 mo- mentum term. We develop a unified Lyapunov analysis show that worst-case complexity of O(1/st) 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 l function (Soft-Thie (oottThresholling) to enhance solution sparsity. We summarize the complexity results in Table
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 e-optimal solution, our proposed method exhibits an OL1/ + 1/(m=²)) complexity bound in the worst case. Moreover, we develop new minibatch SMOD based on Nesterov's momentum, achieving the ((1/E1/2 + 1/(m²?)) optimal complexity bound. Note that a similar complexity result, explicitly relying on the smoothness assumption, has been shown in recent study 77. Compared to this work, our analysis makes weaker assumptions, showing that smoothness is not a must-have for many -bdel- algorithms, such SPL and SPP, 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. [8] 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

---

Table I: Complexity of SMOD to reach E 1//ffll S (M: minibatch; E: Extrapolation, m: batch
size)
Algorithms
Problem
Current Best
Ours ((1/44
M + SGD
f: non-smooth
O
M Prox. SGD
1 + wi Essmoohh
(1/ (1/4 +
( 1/62)
M SPL/SPP
f: non-smooth
O
(11((m=* 1/E2)
E SGD
f: non-smooth
0(1/c) 22
(11/c4)
E+ + Prox. SCD
f = + w;; tsmoooh
((1/E4)
E + SPL/SPP
f: non-smooth
((1/E4)
M+ + SGD
f: non-smooth
O(1/E*)
25
O(1/E*)
M+ E + Prox. SGD
f = L + wi t:smooth
((1/(mE*) + 1//22) (11((me*) + +1/22)
M E SPL/SPP
f: non-smooth
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 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 using stability analysis was proposed by Wang et al. (30), albeit with different motivation from distributed stochastic optimization. Robustness and fast convergence of model-based methods have been shown on various statistical learning problems B 13
16 5 Drusvyatskiy and Paquette [13] give 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 certain finite-sum and stochastic problems by using variance-reduction. Momentum and accelerated methods for convex stochastic optimization be referred from [4 26
develop the convergence rate of stochastic momentum method for smooth non-convex optimization.
The study 10 2 3T
2 Background
Throughout the paper, we use ||- to denote the Euclidean norm and (..-) -) to denote the Buclidean
inner product. We assume that /(x) is bounded below. i.e., min, [(x) > -c. The subdifferential
0z()) of function f (z) is the of vectors V E I2d that satisfy: fly) 2 f(x) + (v,y- yu,y-ar x} + o[[2= vIl), as y - xz Any such vector in 0f is called subgradient and is denoted by f'(x) E ()))) for simplicity. We say that a point is stationary if 0 E 0f(x)) + Nx)x)), where the normal cone Nx (z:) defined as Nx(x) A {d: : (d, - x) S O, Vy E X). For a set S, define the set distance to
by: ||S| e [ O|I, x E S}. is natural to use the quantity [[āf(z) + Nx(x)]l- to measure the stationarity of point
-oreau-envelope. The y-Moreau-ene -enveloeee of f is defined by fa( mingex (f(v) + 2 2 ) and the proximal mapping associated with f()) is defined by prox,, (x) - argmin (f(y) zu || 31 Assume that f(x is X-weakly convex, then for μ < X-1, the Moreau envelope fx(() differentiable and gradient is [x) = 1 - - prox (z)).
The SMOD family iteratively computes the proximal map associated with model function fat Ex): x A++ argmin Jo (zt, Ex) + Tkk || 11]
(2)
xEX
2
where } are i.i.d. samples. Typical algorithms and the accompanied models are described below.
Stochastic (Proximal) Gradient Descent: consider the composite function f(x, f(r, E) + (a) where f(x, {) is a data-driven and weakly-convex loss term and w(x) is a convex regularizer such as fi-penalty, SGD applies the model function:
fg [z, E = fly. {) + (L( v) + w(x).
Stochastic Prox-linear (SPL): consider the composition function f(x,E) = M(((,,) where h((-E) is convex continuous and C(x,, E) is a continuously differentiable map. We perform partial
3

---

linearization to obtain the model
fz) h(C(v, + (VC(y. - 4)).
(4)
Stochastic Proximal Point (SPP): compute
2 with full stochastic function:
f, (z.
(5)
Throughout the paper, assume that far E) is continuous and j--weklyy convex, and that the model
function fx(:..) satisfies the following assumptions 9
A1: For any E, the model function (u, {) is X-weakly convex in (A 2 0).
A2: Tightness condition: fa <) = f(x, E), Vx E &, E~E.
A3: One-sided quadratic approximation: (v: {) [(.:)) S f/11=- g1?,, Vx, y X,E~=
A4: Lipschitz continuity: There exists L > 0 that fz( - (v. & L || yll, for any I. 2 E X, { ~~E.
Remark Assumption A2 is quite standard and will be used only in the convergence proof. Com- bining Al and A3 we immediately have that {xr 5) is (A + ---weally convex. Thus, it suffices to assume that j + Assumptions |2 A4jean can be slightly relaxed by replacing the uniform bound with bound on expectation over ç, leading to only a minor adjustment to the analysis.
Denote è - prox;/p (22) argmin,, {f(x) flyy x ==²²} for some p > μ. Davis and Drusvyatskiy 9 revealed striking feature of Moreau envelope to characterize stationarity:
| - x|| - [VV1 and [ + Nx S (1/p1/(())1-
Namely, point with small gradient norm f1)a stays in the proximity of nearly-stationary point d. With this observation, they show the first complexity result of SMOD for non-smooth non- convex optimization: min mini EllV (x*)(>
(({R) Note that this rate 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.
SMOD with minibatches
In this section, we present a minibatch SMOD method which takes small batch of i.i.d. samples to estimate the model function. The overall procedure is detailed in Algorithm[] Within each iteration, Algorithm]] 1 forms stochastic model function fa (., B&) ma
mig fax Ek,s) parameterized at x* by sampling over my i.i.d. samples B₁ E En Then performs proximal update to get the next iterate 2**+1 We will illustrate the main convergence results of Algorithm I the proof details in Appendix sections. But first, let us present an additional assumption.
and leave all
A5: Two-sided quadratic bound: for any x. y E X, & ~ E, (v. E) f(v: E) T1| yll2.
Remark 2. Assumption A5 is vital for improved convergence analysis. While it is slightly stronger than A3 A5 is indeed satisfied by the SMOD family in most contexts: 1) For SPP SPP[[A5] is trivially satisfied
by taking (z <) = (v, E). 2) For SPL, minimize composition function (x. hCCCe (x)) where h(-) is a e -ipschitzz convex function and Ce(+) is a -ipsccit smooth map. In view of
A5 is verified with I (y, <) (y,, < Ci (v - Ce (z) - VC&
S16 |x
3) For SGD, A5 is satisfied if E(-,E) () is c3 Lipschitz smooth for some c2 > 0, as (3. f fly. (e(g.))
(z. {) v((x, <) (y z) Sx || y We note that A5 is not satisfied by SGD when the loss {(., {) is also non non-smooth. Unfortunately, there seems to be little hope to accelerate SGD in such case since the convergence rate of SGD already matches the rate of deterministic subgradient method.
We present an improved complexity analysis SMOD by leveraging the framework of algorithm stability 6 27 In stark contrast to its standard application in characterizing the algorithm gener- alization performance, stability analysis is applied to determine how the variation of minibatch affects the estimation of the model function in each algorithm iteration.

---

Algorithm Stochastic Model-based Method with Minibatches (SMOD)
Input: x² : for k 1 to K do Sample a minibatch Bg (E..1,.... mm and update wh++1 by solving
fe (x, Ck.i) + 7k || x4 2012
min zEA
(6)
ma
2
end for
Notations. Let = (ć1,
{m} be a batch of i..d samples and B(y) = B ( by
replacing & with an i.i.d. copy f, and B' = {&1 .,...., J. Let h(-,E) E) be stochastic model function, and denote hyy, B) = m Ei h(y Er). The stochastic proximal mapping associated with h(, B) is defined by prox ph (z, B) argmingex {h(y, B) + l > for some > 0. We denote t A prox (x. B) for brevity. We say that the stochastic proximal mapping prox, ph is e-stable if, for any x E X, we have
EB, ( { har cD] Se
where is an index chosen from (1.2....,m) m} uniformly at random.
The next lemma exploits the stability of proximal mapping associated with the model function.
Lemma 3.1. Let fe(., B) be a stochastic model function under the assumptions A7 A4 For y e
(X, oo), vectors z and y, the proximal mapping prox,,, proxx,///. j./n/yy B) = argmin {fa (z.
? is e-stable with m(7-A)"
21.
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
EBa [ k41 B&) Ex (x (x
[))ox] S Skk: EK m²(7x-X)
21
(8)
where dk is the o-algebra generating [B.)sisk-1.
Note that since =+ is dependent Bk, (x*+11 B&)) is not unbiased estimator of Be[Sa+ (x*+1_ { However, the stability argument identifies that the expected approximation error is a decreasing function of batch size my. This observation is the key to the sharp analysis of minibatch stochastic algorithms. With all the tools at hands, we obtain the key descent property in the following theorem.
Theorem 3.3. Suppose that p > X TT, T, 7k 2 T. A5 and all the assumptions in Lemma 3.1 hold. Let [.] abbreviates EB. [ok] and Ea be given by (8) then we have
(p-A-7) p(78 -
22 7) [VSi/(a"")"" f1 Ex [f [a
PE
2 + P 2A
Next, we specify the rate of convergence to stationarity using a constant stepsize policy.
Theorem 3.4. Under the ussumptions of Theorem 3 let A S1/0(x*) min, J (x), my = m, and 7h = ? = max(p T, tu where rj auym
/K and ou E (0, oc). Let k be an index chosen in
(1. K} uniformly, then we have
(2p - A)4
A (A
E [ 11
+
20mpI.
(10)
P P-
K
Cku
vmk
Remark 3. The performance of SMOD depends on dg and batch size m. 10p implies that when batch size is fixed, the best rate is obtained at og = 2p Since both and L are unknown,
5

---

Algorithm 2 Stochastic Extrapolated Model-Based Method (SEMOD)
Input: x0, x², 3.
for k = 1 to K do Sample data 5 and update:
/² = x² + B(x - 2
(II)
x *+1 argmin (2)(xx + -112=- s/1
(12)
#EN
end for
hyper-
hhper-paramen parameter tuning over an is required to obtain good empirical performance. For the simplicity
of theoretical analysis, let us take aq añ. Hence, to obtain an iterate whose Moreau envelop has
expected gradient norm smaller than E, the total iteration count is To = max { ), O LPA me" }. For
small batch size m (i.e. m = )(1/=?)) the second term in max(,) dominates the bound Te, yielding total complexity of O 124 by a factor of m.
Note that this complexity bound is better than the O (2 bound 9
me
Remark 4. Theorem 3.4 implies that SCD can he accelerated by minibatching on the smooth composite problems 33 but leaves out the more general problems where f(x, {) non-smooth and weakly convex. In the latter case, showing any improved rate of minibatch SGD substantially more challenging. Without additional knowledge, the O (.A)) complexity of SGD already matches the best result for deterministic subgradient method (c.f. (9]. It remains unknown whether such ((1/s*) bound is tight or not, and a possible direction to obtain sharper complexity bound is by exploiting the non-smooth structure information such 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 fIr {) = h(C((z, {)) where h(a) |a). For SPL, it is easy to transform the corresponding subproblem to an O(m²) -dimenniooall quadratic program (QP) in the dual space (e.g. )2). The dual QP can be efficiently solved in parallel, for example, by fast interior point solver. For SPP, we show that the subproblem can be solved by deterministic prox-linear method at rapid linear convergence rate. Note that the SPP subproblem is especially well-conditioned because our stepsize policy ensures large strongly convex parameter 7 - A. We refer to the appendix for more technical details.
4
SMOD with momentum
We present model-based method by incorporating additional extrapolation term, and record this stochastic extrapolated model- based method in Algorithm(2] Each iteration of Algorithm[2 consists of two steps, first, extrapolation step performed get an auxiliary update y²* Then a random sample &. is collected and the proximal mapping, associated with the model function J Ex),, is computed at y to obtain the new point = For ease of exposition, we take constant values of stepsize and extrapolation term
Note that A can be interpreted gradient descent step with a broader class of proximal mappings. To see this intuition, we combine and
extension of the momentum SGD by replacing the
12 to get
argmin {1a (x, 7 Nx - xn) + 2 11x x =11]
(13)
=EX
2
If we choose the linear model (3) i.e.,
flat Et) + (a* ( { 2h), and assume
X' = Rdd then the update
13) has the following form:
r. = t²
f . 3( k-1 Inh
(14)
Define v A ( - x² and apply it to 14 then Algorithm reduces to the heavy-ball method
x++ = f* gh)) + 300
(15)
x k+1 = t kt1
(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 ((3I) for our problem.
To overcome this issue, we present unified convergence analysis which can deal with all the model functions and is amenable to both constrained and composite problems.
Our theoretical analysis of Aigerithm2Irelies on different potential function from the one in the previous section. Let define the auxiliary variable
3 B
(x x R-1).
(17)
The following lemma proves some approximate descent property by adopting the potential function + [1:p(=*)
2(70-18) 1| 112 and measuring the quantity of VI ///p\/(=*))
Lemma 4.1. Assume that 2 2(T + X) and 3 E [0, 1). Let 0 = 1 B. Then we have
(p 1 X0)
PL2
Ex [[1ppp( (=++))]]] + (n02 - -.320-1)(( (7 X02) 2p(70 - X0) |V/1/p )1? S 11/p(=")
+ p(7 + 1880022)
(1120
-x--111 Ex 1|1 - 13)
2(90 20)
p(7 - 032003)
-
Ex Ex(ll
(18)
4(7-\) A)
Invoking Lemma4.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 À + p -- for some " then
if we choose x²2 = x", and set = 7000 VV K
2p
(10 +
pL
(19)
P X
K
VK
where k- is an index chosen in f1.2..... K} uniformly at random.
Remark 5. Despite the fact that convergence established for all 7o >> O, we can see that the optimal
To would be % = Vo which gives the bound E| [ )]1] 2p (089g 2A + 2L In practice, we can set "Yo to a suboptimal value and obtain a possibly loose upper-bound.
pA
K
0K
Remark 6. Since - is an extrapolated solution, it may not be feasible. It is desirable to show optimality guarantee at iterates x². Note that using Lemma 4.1 and the parameters in Theorem 4.2
it is easy to show that E[[l=* x* 1 = x Based
17 we have ||=* x²*
330--2 2E[||x" z.kk 111' = O ). Using Lipschitz smoothness of Moreau envelop, we can show
B[[I (1/p(a* [[S] converges at the same O (x)) rate is shown in Theorem 4.2
Combining momentum and minibatching, we develop a minibatch version of Algorithm 2 that takes batch of samples B₁ 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 |Bx| m and 7 = O(v K m ), then fi;l/f = = + 1
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 rate of
view of this result, the deterministic part of our rate is consistent with the best rate for the heavy-ball method. For example, see 11 17 Moreover, the stochastic part of the
Vm
O [
rate is improved from the O 1 rate of Theorem 4.4 9 by factor of vm. VK
7

---

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 heavy-ball type momentum, it would be interesting to consider Nesterov's acceleration. To this end, we present model-basee method with Nesterov type momentum. Thanks to the stability argument, we obtain the following improved rate of convergence: O k² + V We note that 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.
6 Experiments
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 a, Rd and nonnegative scalars b; IR, the goal of phase retrieval is to recover the true signal x² from the measurement by = [(a;, x') > Due to the potential corruption in the dataset, we consider the following penalized formulation
1 Σ
min IR
(a., x)² -bl
bil
(20)
where we impose fossoss to promote robustness and stability (cf.
15 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 Rnxd 9s ~ (O, 1) and D E Rd is diagonal matrix whose diagonal entries are evenly distributed in [1/x,1], Here x 2 1 plays the role of condition number (large k makes problem hard). The measurements are generated by b, x=) &&C; (1 < <i S n) with G (N, 25), 6, Bernoulli(Ptat), where Ptail E [0, 1] controls the fraction of corrupted observations on expectation.
2) Real data. We consider zipcode, a dataset of 16 x 16 handwritten digits collected from (21] Following the setup in [5]. let E 256x2256 be a normalized Hadamard matrix such that hij E ( - ] H = HT and H = H-1 Then we generate k = 3 diagonal sign matrices S S2, S3 such that each diagonal element of Sy is uniformly sampled from {-1, Last we set A [HS,, 2S,,, HSS.'
(3x22566) x56 As for the true signal and measurements, each image is represented by data matrix X E R 16 16 and gets vectorized to x= = vec(X). To simulate the case of corruption, we set measurements b @pmn (Ar*), where &pen (-) denotes element-wiss squaring and setting a fraction Plan of entries to 0 on expectation.
In the first experiment, illustrate that SMOD methods enjoy linear speedup in the size of minibatches and exhibit strong robustmess to the stepsize policy. We conduct comparison on SPL and SCD and describe the detailed experiment setup as follows.
1) Dataset generation. We generate four testing cases: the synthetic datasets with (s, Pla)) (10, 0.2), and (10, 0.3); zipcode with digit images of id 2 and 24;
2) Initial point. We set the initial point x (= 20) (0, Ia) for synthetic data and x! x (0, for zipcode;
3) Stopping criterion. We set the stopping criterion to be f(x²) S 1.5f, where f = f(*)) ) is the corrupted objective evaluated at the true signal a::
3) Stepsize. We set the parameter 7 og /K/m where is the batch size; For synthetic dataset, we test 10 evenly spaced ag values in range [10-1, logarithmic scale, and for zipcode dataset we set such range of og to [101,
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, (2,32,44);
6) Sub-problems The solution to the proximal sub-problems is left in the appendix.

---

For each algorithm, speedup from minibatching is quantified as TT /TT where Tw is the total number of iterations for reaching the desired accuracy, with batch size m and the best initial stepsize ag among values specified above. Specially, if an algorithm fails to reach desired accuracy after running out of 400 epochs, we set iteration number to the maximum.
320
120
328
3DD
GFL
SPL
02
00
0
10
30
00
30
40
10
40
Figure 1: Speedup over minibatch sizes. The left two are for synthetic datasets K = 10, Pfail E (0.2,0.3) Digit datasets: digit image (id:24) with pratt E {0.2, 0.3)-
18
3-600
560
6P.
SPL
30
30
16'
39a
x30"
-307
Q-500
L
1 -5500
30 KG05
3PL
-1-SPP
10-
10
5-
Figure 2: From left to right: synthetic datasets with E 32) and zipcode image (id=24) with m E {8, 323. x-axis: initial stepsize og- y-axis (first row): speedup over the sequential version: / (0o) where Th (Oo) stands for the number of iterations when using batch size m. and initial stepsize ag. y-axis (second row): Total number of iterations.
Figure D plots the speedup of each algorithm over different values of batch size according to the average of 20 independent runs. can be seen that SPL exhibits a linear acceleration over the batch size, which confirms 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 ag in Figure (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 initial stepsizes in Figure 2 (2nd row) for batch size = 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 ( 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 (x, Pra)) (10,0.2) and

---

(10, 0.3); zipcode with digit images of id 2 and Peait E {0.2, 0.3). We set E [10-? 10" = 0.6 for synthetic data, and ag [10", 10'], 0.9 for zipcode dataset. The of settings are the same as in minibatch with m 1.
Figure plots the number of epochs e-accuracy over initial stepsize op- It can be seen that with properly selected momentum parameters (SEGD, SEPL, SEPP) all suggest improved convergence when stepsize relatively small.
-0400
500
6 9P..
3 KP
SE50 1EP1
SEGO
360
SEG0
200
200
300
SFL
3PP
110
SEPL
10a
Figure 3: From left to right: synthetic datasets with K. 10, Pran E {0.2, 0.3), 8 = 0.6 and zipcode image (id=2) with Prait E [0.2, 0.3), = 0.9. x-axis: initial stepsize ag. y-axis: number of epochs on reaching desired accuracy
SOD SPP
500
SED 0
100
SEPL
300
520%
300
SPF
SEPL
300
02 SECO
SEFP
X SEPP
100
100
10" 107
16'
16²
Figure 4: From left to right: synthetic datasets with x = 10, Phil = 0.2, = 0.6, m E 32) and Phil
zipcode image (id=24) with Prait 0.3, 3 0.9, E (1, 32). x-axis: initial stepsize oo. y-axis: number of epochs for reaching desired accuracy
In the last experiment, attempt 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 can be seen that minibatch SMOD, when combined with momentum, exhibits even better convergence performance and robustness.
7 Discussion
On broad class of non-smooth non-convex (particularly, weakly convex) problems, we make stochastic model-based methods more efficient by leveraging minibatching and momentum techniques that are well-known only for SGD. Applying algorithm stability for optimization analysis a key step to achieving improved convergence rate over the batch size. This perspective appears to be interesting for stochastic optimization in much broader context. Although some progress is made, we are unable to show whether minibatches can accelerate SGD when the objective does not have smooth component. Note that the complexity of SGD already matches the best bound of full subgradient method. would be interesting to know whether this bound for SGD is tight 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 [2
8 Acknowledgement and disclosure of funding
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).
10

---

References
[1] H. Asi and C. Duchi. The importance of better models in stochastic optimization. Proceedings of the National Academy of Sciences, 116(46):22244 22930, 2019.
[2] H. Asi, K. Chadha, G. Cheng, and J. C. Duchi. Minibatch stochastic approximate proximal point methods. Advances in Neural Information Processing Systems, 33, 2020.
[3] R. Bassily, V. Feldman, C. Guzmán, and K. Talwar. Stability of stochastic gradient descent on nonsmooth convex losses. Advances in Neural Information Processing Systems, 33, 2020.
[4] L. Berrada, A. Zisserman, and M. P. Kumar. Deep frank-wolfe for neural network optimization. In ICLR 2019: : 7th International Conference on Learning Representations, 2019.
[5] A. Botev, H. Ritter, and D. Barber. Practical gauss-newton optimisation for deep learning. In International Conference on Machine Learning, pages 557-565. PMLR, 2017.
[6] O. Bousquet and A. Elissefff. Stability and generalization. Journal of machine learning research, 29Mr))499-552, 2002.
[7] K. Chadha, G. Cheng, and J. C. Duchi. Accelerated, optimal, and parallel: Some results on model-baaed stochastic optimization, arXiv preprini W.Niv:2111.0269000 2021.
[8] V. Charisopoulos, Y. Chen, D. Davis, M. Díaz, L. Ding, and D. Drusvyatskiy. Low-rank matrix recovery with composite optimization: good conditioning and rapid convergence, arXiv preprint arXiv: /904. 10020, 2019.
[9] D. Davis and D. Drusvyatskiy. Stochastic model-based minimization of weakly convex functions. Siam Journal on Optimization, 29(1):207-239, 2019.
[10] A. Defazio. Understanding the role of momentum in non-convex optimization: Practical insights
from Iyapunov analysis. arXiv preprint 20Xiv:200000044000 2020.
[11] O. Dekel, R. Baadd-Bachrcch... O. Shamir. and L. Xiao. Optimal distributed online prediction using bithhe. Journal of Machine Learning Research, 13(1), 2012.
[12] Diakonikolas and M. I. Jordan. Generalized momentum- based methods: a hamiltonian
perspective. SIAM Journal on Optimization, 31(1):915-944, 2021.
[13] D. Drusvyatskiy and C. Paquette. Efficiency of minimizing compositions of convex functions and smooth maps. Mathematical Programming, pages 1-56, 2018.
[14] C. Duchi and F. Ruan. Stochastic methods for composite and weakly convex optimization problems. SIAM Journal on Optimization, 28(4):3229-3259, 2018
[15] J. C. Duchi and F. Ruan. Solving (most) of a set of quadratic equalities: Composite optimization for robust phase retrieval. Information and Inference: Journal of the IMA, 8(3):471-529, 2019.
[16] T. Frerix, T. Möllenhoff, M. Moeller, and D. Cremers. Proximal backpropagation. In Interna- tional Conference on Learning Representations, 2018
[17] E. Ghadimi, H. R. Feyzmahdavian, and M. Johansson. Global convergence of the heavy-ball method for convex optimization. In 2015 European control conference (ECC), pages 310-315. IEEE, 2015.
[18] S. Ghadimi and G. Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal Optimization, 23(4):2341-233 2013. ISSN 1052-6234.
[19] I. Gitman, H. Lang, P. Zhang, and L. Xiao. Understanding the role of momentum in stochastic gradient methods. In Advances in Neural Information Processing Systems, volume 32, pages 6633-9643, 2019.
[20] M. Hardt, B. Recht, and Y. Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, pages 1225-1234. PMLR, 2016.
[21] Hastic, R. Tibshirani, Friedman, T. Hastie, J. Friedman, and R. Tibshirani. The elements of statistical learning, volume 2. Springer, 2009.
[22] G. Kornowski and O. Shamir. Oracle complexity in nonsmooth nonconvex optimization. arXiv preprint arXiv:2104.0676, 2021.
[23] Y. Liu, Y. Gao, and Yin. An improved analysis stochastic gradient descent with momentum. arXiv preprint arXi::2007.0798, 2020.

---

[24] N. Loizou and P. Richtárik. Momentum and stochastic momentum for stochastic gradient, newton, proximal point and subspace descent methods. Computational Optimization and Appli- cations, 77(3):653-710, 2020.
[25] V. Mai and M. Johansson. Convergence of a stochastic gradient method with momentum for non-smooth non-convex optimization. In Proceedings of the 37th International Conference on Machine Learning, pages 6630-6639, 2020.
[26] O. Sebbouh, R. M. Gower, and A. Defazio. On the convergence of the stochastic heavy ball method. arXiv preprint arXiv:2006. 07867, 2020.
[27] S. Shalev-Shwartz, O. Shamir, N. Srebro, and K. Sridharan. Learnability, stability and uniform convergence. The Journal of Machine Learning Research, 11 11:2635-2670 2010.
[28] I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning, In International Conference on Machine Learning, pages 1139-
[29] M. Takáč, P. Richtárik, and N. Srebro. Distributed mini-batch sdca. arXiv preprint arXiv:1507.0880 2015.
[30] J. Wang, W. Wang, and N. Srebro. Memory and communication efficient distributed stochastic optimization with minibatch prox. Conference on Learning Theory, pages 1882-1919, PMLR, 2017.
[3]] Y. Yan, T. Yang, Z. Li, Q. Lin, and Y. Yang. A unified analysis of stochastic momentum methods for deep learning. In Proceedings of the Eventy-Seventh International Joint Conference on Artificial Intelligence, pages 2955-2961, 2018.
[32] J. Zhang and L. Xiao. Stochastic variance-reduced prox-linear algorithms for nonconvex composite optimization. Mathematical Programming, pages 1-43, 2021.
[33] J. Zhang. H. Lin, S. Jegelka, A. Jadbabaie, and S. Sra. Complexity of finding stationary points of nonsmooth nonconvex functions. arXiv preprint arXiv:2000.041:30, 2020.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See remarks 3, 4 in Section 3 and Discussion section
(c) Did you discuss any potential negative societal impacts of your work? [No] The paper addresses theoretical questions on algorithm complexity, which, to the best of our knowledge, has no negative social impact
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See assump- tions in Section and 3
(b) Did you include complete proofs of all theoretical results? [Yes] Proof is left in the appendix
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either the supplemental material or as a URL)? [Yes] Code is supplied in the supplemental materials
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section/]for details of experiments
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [No] From the experiments the error bars are relatively thin and the results are presented by taking average.
12

---

(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [No] The main goal of experiments is to demonstrate our theoretical foundings, thereby only showing the iteration complexity of algorithms.
4. If you using existing assets (e-g., code, data, models) curating/releasing assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] zipcode dataset is referenced from [21].
(b) Did you mention the license of the assets? [No] The dataset used is published on an open site without license.
(c) Did you include any new assets either the supplemental material or as URL? [No] The experiments do not involve new datasets.
(d) Did you discuss whether and how consent
obtained from people whose data you rre're using/curating? [No] An open dataset is used.
(c) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [No] The dataset has been open for years and only involves zipcode digits.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [No] No crowdsourcing human object is involved.
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [No] No crowdsourcing or human object is involved.
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [No] No crowdsourcing or human object is involved.
13


Reviews:
[{"review_id": "wGGk737Js6s", "reviewer": "Reviewer_h9LY", "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.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is very well-written and easy to read.\nThe authors propose to use the concept in algorithm stability for the analysis of convergence, which is novel to me since the algorithm stability is initially proposed for the analysis of generalization. They obtain their main convergence results based on bounding the effect of one-step stochastic proximal mapping, i.e., the stability of each iteration. This aligns with the intuition of using SMOD over SGD, i.e., for better robustness under non-smooth weakly-convex setting. In their analysis, using minibatch in SMOD is able to provide linear speed up since the batch size is shown in the stability (Lemma 3.1) and can help optimize the choice of $\\eta$ and the stepsize $\\gamma$ in Theorem 3.4. \nThey also use the same way to analyze the method of SMOD with heavyball momentum.  Moreover, in the Appendix, they show results on both heavyball and Nesterov momentum under convex setting.\nTheir experiments are standard among the literatures of SMOD. With the results shown in the experiments, their theoretical results are verified.\n\nMy questions:\n1. For Theorem 3.4, in the proof, we can see $\\gamma$ cannot be too large since we require $\\gamma \\leq \\rho + \\tau + \\lambda + \\eta$. As the main reason to choose SMOD is for the wide range of the choice of step size, to better show the linear speedup of minibatch SMOD, can the author change the condition on $\\gamma$ in Theorem 3.4 to a range of $\\gamma$ instead of just $\\gamma=\\max(\\rho+\\tau, \\lambda+\\eta)$?  This is also related to the experiments in this paper:  the initial stepsize $\\alpha_0$ is in a range (e.g., Figure 2) and it is unknown why this range is useful for verifying the results of Theorem 3.4.\n\n2. In Line 57, for the analysis of SMOD with momentum, the authors mention they develop a unified Lyapunov analysis. Where is the Lyapunov function actually used for analysis? It may be implicit shown in the proof but I did not find it.\n\n3. For the experiments, in Figure 1, the authors only compared SPL and SGD, is there any reason that SPP is not shown in this Figure?  Also in Figure 1, the difference between SPL and SGD is relatively small for $\\alpha_0\\in[10^{-1}, 10^0]$, so why in Figure 3, the range in x-axis is chosen as $[10^{-2}, 10^0]$? We can see that for the first subfigure in Figure 4, there are some interesting results around $\\alpha_0=10^1$. This extreme case ($\\alpha_0\\geq 10^1$) may help us understand the momentum minibatching SMOD method better (compared with Theorem 4.3).\n\n\n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "16"}, "extra_text": {}}, {"review_id": "eMEBb9B29UJ", "reviewer": "Reviewer_5Yxy", "summary": "This paper studies minibatch and momentum algorithm for non-smooth non-convex optimizations. A linear speedup property of minibatch is achieved in non-smooth non-convex optimization.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Most interestingly the usage of minibatch is demonstrated to have a linear speed up property in nonsmooth setting, which is rare in the area. However, there are several limitations about the problem setting being considered. Although this paper claims to target on non-smooth optimization, the problems considered are several special classes of problems. \n\n1. The SSP assumes exact solution of sub-problem at each iteration, which is a very strong assumption of power of the subroutine. Taking the complexity of subroutines into consideration, it is very likely that the complexity will be much higher. \n\n2. The Prox. SGD requires smoothness of \\ell(.), which setting has been well studied in existing works including in [8]. \n\n3. What is more, SPL requires smoothness of inner function C(.) (refer to line 133 of the submission), this should be clearly pointed out in early sections and explicitly in assumptions. Because the fact that the inner function where variance arise is smooth makes the algorithms quite restricted to this special class of problem. \n\n4. I also noticed a paper that studies a similar setting as the SPL and gets better sample complexity of O(1/\\epsilon^3), while this paper gets O(1/\\epsilon^4) sample complexity. See [Zhang, Junyu, and Lin Xiao. \"Stochastic variance-reduced prox-linear algorithms for nonconvex composite optimization.\" arXiv preprint arXiv:2004.04357 (2020)]. Since this reference uses variance techniques, it is understandable that they get better rates, but it would be helpful to mention and compare with it. \n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "-_xrGG3mE_s", "reviewer": "Reviewer_Cw8j", "summary": "The paper studies the convergence of stochastic model-based methods with mini-batching and proves theoretical speedup for quadratically accurate models such as the proximal point and prox-linear ones.\nA momentum extension is also studied. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe keys to their analysis are algorithm stability and the new Assumption A5. Assumption A5 brings us close to a situation as we would have been if Lispchitz continuous gradient was assumed in the analysis of SGD. In particular, we are granted some kind of *decent lemma*. \n\nIndeed, a similar condition has been used to unify the analysis of several model-based methods [Damek, 2020]. More concretely, the accelerated rate O(1/k^2) can be obtained for the proximal point and prox-linear methods on *non-smooth* deterministic problems. So I would say that under A5, the linear speedup is rather natural; the proof may not be trivial though.\n\nThe author(s) mentioned the restricted convexity (RSC) assumption when compared to [Asi et al., 2020], can you comment on the difference with a somewhat more recent version of that paper [ACCD21]? \nIt seems that the RSC requires some lower quadratic growth (QG) of the original function f(x, \\xi), while A5 requires QG on both sides but with respect to the model f_y(x, xi). Which one is more conservative? Can you give an example? \n\nIn eq. (7), both the samples are \\xi'??\n\nAs the author(s) mentioned that removing A6+A7 does not affect the main convergence results, can you drop them?\n\n\nReferences: \n\n[Dam20] Damek Davis, Lecture Notes for ORIE 6300: Mathematical Programming I, 2020\n\n[ACCD21] Chadha, Karan, Gary Cheng, and John C. Duchi. \"Accelerated, Optimal, and Parallel: Some Results on Model-Based Stochastic Optimization.\" arXiv preprint arXiv:2101.02696 (2021). ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The bound concerns only the number of proximal steps and not the total complexity. The author did not tell us how to carry out the mini-batch update in general? Without an efficient procedure, it would be hard to compete with gradient-based methods. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper analyzes two extensions of the recently studied (stochastic) model based optimization methods: i) minibatch variant, ii) momentum variant. The authors focus on a class of nonsmooth nonconvex functions which become convex after sufficiently strong Tikhonov regularization. The analysis approach is based on leveraging a recent result by Davis and Drusvyatskiy characterizing the degree of stationarity of a point via the gradient of its Moreau envelope, and on utilizing tools from algorithmic stability literature. The authors recover several existing and several new convergence results.

The reviewers assessed the paper favorably. There was a reasonable discussion between the reviewers and the authors. I have read the paper myself and agree that the work is of sufficient quality to be published in NeurIPS. I have noticed a number of grammatical and stylistic issues and would recommend the authors to carefully proofread the paper.
Decision: Accept (Poster)