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- yo,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 0, 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
-oveau-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 (v. {) 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 + (., 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. 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² :Tk 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 [ kt1 B&) Ex (x (x
[))ox] S Skk: EK mu(3x-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 (1. K} uniformly, then we have
/K and ou E (0, oc). Let k be an index chosen in
(2p A)4
A (A
E [ 11
+
20mpI.
(10)
P P-
K
Cku
vmk
Remark 3. The performance of SMOD depends on do 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)(1 + -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 ((m²) -dimensiooall 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 (x),,, 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 ///p1/(=*))
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, 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 {1,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 % = VVLoL 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 0 (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 = = + 1R)
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 SEPF
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 (.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)33229-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.
[31] 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": "uo3NrZlQjN4", "reviewer": "Reviewer_37oF", "summary": "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.", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "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": "Originality: The convergence analysis seems non-trivial. \n\nQuality: The assumptions are too strong. For example, it seems that Assumption 5 is never used in the literature. How to justify it?\n\nClarity: The paper is well organized, and the writing is clear to me. \n\nSignificance: The “improved” convergence result of SMOD requires additional assumptions (e.g., A5, A6, A7), which makes the contribution of this paper limited. With Assumption 5, having such convergence result is not surprising. The Assumption 5 makes the results no general enough and is not applicable to standard stochastic subgradient methods for solving non-smooth problems. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No. The new assumption A5 is never used in the lierature. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "rusdv8qrCE", "reviewer": "Reviewer_Kosw", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "The paper is technically sound and the main results make sense to me. I am particularly impressed with the clear and clean writing style, and how all steps in the proofs are carefully explained. However, I still have some reservations on the novelty and significance of results as specified below:\n\n- Originality: While interesting, the idea of extending SMOD-type methods to minibatch optimization using stability argument is definitely not new. For example, for distributed learning with convex and Lipschitz losses, the benefit of minibatching has already been revealed for stochastic proximal point (SPP) method through the lens of uniform algorithmic stability theory [1]. More precisely, up to minor differences, Theorem 3.2 is identical the Lemma 2 in [1]. Therefore, the only novel point I can spot in this part is using the generalization bounds in Theorem 3.2 to analyze the convergence of SMOD for weakly convex losses, as summarized in Theorem 3.3 & 3.4. The technical part of these theorems is more or less straightforward given the techniques developed in [2] (i.e., reference [8] in the paper). Overall, the degree of novelty of the minibatch extension is fairly low with respect to the prior work in this line. The proposed momentum extension is relatively more impressive, and somewhat novel as far as I know about. \n \n- Significance of results: My main concern regarding the proposed extrapolated variant of SMOD is the strength of results in Theorem 4.2 & 4.3, which seem showing no significant improvement over the vanilla version. Since the method reduces to a heavy-ball acceleration strategy in the SGD case, it is not surprising that it does not substantially improve the rate of convergence for generic convex or non-convex losses. On the other side, it is well known that heavy-ball is beneficial for accelerating batch GD for quadratic programs. Then a natural question is if it is possible for the extrapolated SMOD method to achieve acceleration at least for quadratic losses? Anyway, a careful discussion on the advantages of the momentum extension is highly desirable to justify the usage of method in theory. \n\nMinor comments: \n\n- Clarity: Since Assumption A6 & A7 are of very limited interest unless $M$ and $D$ are extremely small, it is suggested to carrying out the analysis without introducing these void assumptions so that the core results can be presented in a neater way. \n\n- Experiment: As the comparison is made among stochastic optimization methods, it is desirable to provide error bars in the plots to illustrate the variance of computation.\n\n[1] Wang, Wang, and Srebro. Memory and communication efficient distributed stochastic optimization with minibatch prox. COLT, 2017.\n\n[2] Davis and Drusvyatskiy. Stochastic model-based minimization of weakly convex functions.  SIOPT, 2019.\n\n\n\n=== Post rebuttal update ===\n\nThe major concerns raised in my initial review on the momentum extension have been much clarified in the author response. It's good to know, based on the augmented results provided during author discussion phase, that the robustness of SMOD to the choices of step-size can be much preserved by its momentum variant. As for the minibatching acceleration analysis, while  the results for weakly-convex problems are somwhat interesting, IMO the technical contribution of this part is too incremental for NeurIPS given the prior work of Wang et al. (2017) and Davis & Drusvyatskiy (2019).  It is suggested to tune down the claim of this part accordingly.  \n\n Overall, I would still like to change my rating one grid up to 6, apprecaiting the novelty and strength of the momentum extension/analysis of SMOD.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes,  the authors have partially addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "11"}, "extra_text": {}}, {"review_id": "f1oeYq3wFlH", "reviewer": "Reviewer_cuoE", "summary": "This contribution introduces and analyzes mini batch variants of stochastic model-based optimization methods when the loss function is possibly nonconvex and non smooth. Their algorithms are accompanied with a theoretical analysis and numerical experiments.", "questions": "", "limitations": "", "rating": 3, "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": "Overall, I don’t understand what makes those methods different from classical stochastic optimization method. All of those baselines are model based in the sense that the updated parameters throughout the iterations, depend on the loss function and the model employed.\nHence, Algorithm 1 boils down to any proximal and regularized gradient step, which has been around for a long time. \nIt would be greatly appreciated if the authors can clarify what makes their method different from what is in the literature.\n\n- Numerical experiments: The authors only compare, in Figure 1 and 2, their method with plain SGD.\nGiven the sheer size of optimization methods, one would expect much more baselines for fair comparison.\n\nMinor typos:\n- In the checklist: « Did you include the license to the code and datasets? [Yes] See Section ??. » what is the section number?\n- Using the same notation for the stochastic objective function and the unregularized one is confusing.\n- Algorithm 1 writing requires more details: How is the stepsize initialized, how are the m_k decided, what is the output of the method.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No they have not.", "review_text": "", "extra_scores": {"ethical_concerns": "N.A.", "needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "FwXWpQDtgp", "reviewer": "Reviewer_kmZz", "summary": "The paper studied stochastic model based methods (SMOD), and prove complexity bounds for minibatch SMOD and momentum SMOD.\n", "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": "Strength: the paper proves complexity bounds for SMOD with minibatch and momentum, and show that minibatch provides speedup, which is new to the literature. \n\n\nThe complexity result, such as Theorem 3.4, indeed implies  linear speedup, however only for m being not too large. Suppose $ m \\geq 1/\\epsilon^2 $, then the complexity would be $ T_\\epsilon = O(\\Delta/\\epsilon^2) $ which is independent of batch size. As a result, if would be interesting to see in Figure 1, what the plots will be like for even large $m$. In addition, what is the “desired accuracy’’ chosen in the numerical experiments, will the value of accuracy affects the linear speedup of minibatch?\n\nMinor: \nLine 12, “…possibly improve…’’ please remove or rewrite this statements since it’s meaningless. Even in the deterministic setting, momentum is not guaranteed to provides acceleration. However, it might leads to smaller objective function values. It would be good if the authors can check if this is true for the considered problems here. \n\nLine 149, “B = ”, equal sign missing. \n\nLine 205, “Despite such a relation’’, a missing\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See above for limitations and no negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: 
The paper proposes an analyses of model based methods with mini-batching and momentum. The topic is relevant to the Neurips community. But the reviewers were unclear as to how this fits in and compares to existing work, in particular due to Assumption 5. On that note, since most of the models involve a linearization (excluding the proximal point method) Assumption 5 is essentially smoothness. This makes it a bit confusing given the emphasis of the work on the non-smooth setting. Reviewer Kosw also raised questions as to the technical novelty of the mini-batch analysis based on stability analysis, given closely related prior work. Given the above, I will not recommend the paper be published.

The technical contribution of the paper may be sound and novel, but it appears that it wasn't clear or accessible enough for most of the reviewers. Thus in revising the paper, it may simply be a case of re-writing to improve clarity and the positioning of the paper in the literature. For example, Theorem 3.4 is hard to parse, given the large number of new constant definitions, where the parameter of most interest the mini-batch size (m) is hidden away in several new constants, making it hard to appreciate the speed-up offered by increasing the mini-batch size. 

Some further references on SGD + minibatch analysis (though in the smooth case) include [1] in the non-convex setting and [2] in the convex setting. On understanding SGD with momentum in the nonconvex setting, recently (and arguably concurrent work with yours) SGD with momentum was analysed in [3].  

[1] Better Theory for SGD in the Nonconvex World, Ahmed Khaled, Peter Richtárik, 2020

[2]  SGD: General Analysis and Improved Rates, Robert M. Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, Peter Richtárik; ICML 2019.

[3] Aaron Defazio, Momentum via Primal Averaging: Theoretical Insights and Learning Rate Schedules for Non-Convex Optimization, arXiv:2010.00406, 2021
Decision: Accept (Poster)