Stacey: Promoting Stochastic Steepest Descent via Accelerated $\ell_p$-Smooth Nonconvex Optimization
Authors: Xinyu Luo
Site Bai
Bolian Li
Petros Drineas
Ruqi Zhang
Brian Bullins
Venue: ICML 2025 poster (2025)
Keywords: Non-convex Optimization
Non-Euclidean Acceleration
Stochastic Steepest Descent
Primary area: optimization->stochastic
TL;DR: We propose Stacey, an accelerated $\ell_p$ steepest descent algorithm that leverages non-Euclidean geometry to achieve faster convergence and higher accuracy.

Abstract: While popular optimization methods such as SGD, AdamW, and Lion depend on steepest descent updates in either $\ell_2$ or $\ell_\infty$ norms, there remains a critical gap in handling the non-Euclidean structure observed in modern deep networks training. In this work, we address this need by introducing a new accelerated $\ell_p$ steepest descent algorithm, called Stacey, which uses interpolated primal-dual iterate sequences to effectively navigate non-Euclidean smooth optimization tasks. In addition to providing novel theoretical guarantees for the foundations of our algorithm, we empirically compare our approach against these popular methods on tasks including image classification and language model (LLM) pretraining, demonstrating both faster convergence and higher final accuracy. We further evaluate different values of $p$ across various models and datasets, underscoring the importance and efficiency of non-Euclidean approaches over standard Euclidean methods. Code can be found at https://github.com/xinyuluo8561/Stacey.

STACEY: Promoting Stochastic Steepest Descent via Accelerated
4p-Smooth Nonconvex Optimization
Xinyu Luo" 1
Cedar Site Bai Bolian Li
Petros Drineas'
Ruqi Zhang
Brian Bullins'
Abstract
more recent methods like signSGD (Bernstein al., 2018) and Lion (Chen et al., 2023) focus on using the sign of the (stochastic) gradient.
While popular optimization methods such as SGD, AdamW, and Lion depend steepest de- scent updates in either & or for norms, there re- mains a critical gap in handling the non-Fuclidean structure observed in modern deep networks train- ing. In this work, we address this need by intro- ducing new uccelerated l steepest descent al- gorithm, called STACEY, which uses interpolated primal-dual iterate sequences effectively nav- igate Euclidedd smooth optimization tasks. In addition to providing novel theoretical guar- antees for the foundations of our algorithm, we empirically compare our approach against these popular methods on tasks including image classi- fication and language model (LLM) pretraining. demonstrating both faster convergence and higher final accuracy. We further evaluate different val- ues of across various models and datasets, un- derscoring the importance and efficiency of non- Euclidean approaches over standard Euclidean methods. Code can be found at https:// /// github. com/ /it.yuuul@8////ttcr....
Although these algorithms have shown impressive empirical performance (sometimes exceeding that of standard adap- tive methods), their theoretical analyses typically rely on Euclidean (i.e., E2) or -based assumptions. Specifically, crucial to guarantees of finding e-approximate stationary points (Carmon et al., 2017; Ghadimi & Lan, 2013; Jin et al., 2017) two related choices: (i) the norm used define stationarity, and (ii) the corresponding notion of smoothness. Classical analyses in deep learning often adopt Euclidean smoothness (Ghadimi & Lan, 2013), while signSGD relies on le based assumptions (Bernstein et al., 2018; Balles et al., 2020).
Yet, there is mounting evidence both theoretical and empir- moal-suggesting that a more flexible Ep perspective cap- ture the geometric structure of complex deep network objec- tives far better than either p 2 or co alone (Adolphs et al., 2019; Cohen et al., 2021; Ghorbani et al., 2019; Jiang et al., 2024; Li et al., 2020a; Papyan, 2018). For instance, depending on the shape of the loss surface and the distri- bution of gradients across coordinates, certain E, norms with < p < oc may lead to faster descent improved generalization. This leaves open significant gap: How can we develop and analyze optimizers in alternative non- Euclidean regimes, namely those with general Ep norms where p E (2, o.)?
1. Introduction
Stochastic first-order methods have proven essential for effi ciently training modern deep learning models. Beyond the basic stochastic gradient descent (SGD) algorithm (Robbins & Monro, 1951) and its momentum- based variants (Nes- terov, 1983; Polyak, 1964), a variety of adaptive methods have been developed, such AdaGrad (Duchi al., 201 la), Adam (Kingma & Ba, 2015), and AdamW (Loshchilov & Hutter, 2019), which incorporate second-moment gradient information to provide per-coordinate scaling. Meanwhile,
To address this question, we propose a novel approach that we term STACEY (Stochastic Steepest Descent with Acceleration). Our development builds on insights from both f,p stespee descent and non-Euclidean acceleration techniques (Allen-Zin & Orecchia, 2017; Diakonikolas & Guzmán, 2024; Nemirovskii & Nesterov, 1985; Nesterov, 2005), combining primal-dual (aataalldala(iaakooiia (Diakonikolas & Orecchia, 2019) iterates with an interpolation scheme designed specif- ically for by-based smoothness. While the notion of ac- celeration is well understood in the classical (Euclidean) setting (Nesterov, 1983), extending it to arbitrary &p norms introduces a fundamental tradc-off: although we may at- tain improved geometry dependence (and thus potentially faster practical convergence in certain regimes), the theoret-
'Equal contribution 'Department of Computer Science, Pur- due University, Indiana, USA. Correspondence to: Xinyu Luo <luo466 purdue edu>,, Cedar Site Bai <bail23 puux, Bolian Li <li4468 @ edud..
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated Ep-Smooth Nonconvex Optimization
ical "acceleration exponent" necessarily decreases from 2 toward 1 as p grows large (Guzmán & Nemirovski, 2015; Nemirovskii & Nesterov, 1985). Nonetheless, by situating STACEY within this continuum of non Euclidean optimizers, we can reap meaningful benefits over purely Euclidean (e.g-- SGD) and purely sign-based (e.g., SignSGD) methods on modern, large-scale tasks.
an &g majorization-based smoothness condition," 1 showing that in expectation, | [+)]] can be driven below a pre- scribed threshold. Similarly, Balles et al. (2020) investi- gated the geometric underpinnings of sign-based updates, highlighting how they relate to fo steepest descent. Re- cent work on Lion (Chen et al., 2023) and its generalization Lion-K (Chen et al., 2024) further underscores the empirical benefits of sign-driven coordinates in large-scale tasks.
Our Contributions:
However, sign-based approaches (p 20) represent just one extreme of non-Euclidean geometry. The other well- studied example is the classical by-based regime (c.g.. vanilla SGD) (Ghadimi & Lan, 2013; Robbins & Monro, 1951). where standard notions of Euclidean smoothness and approximate stationarity (+)lll S E underpin core theo- retical results. Interpolating between these extremes (l2 and ( by considering fp norms for 2 p co has remained comparatively underexplored in the stochastic, non-convex setting. One challenge is that, unlike in the Euclidean case, the coordinate-scaling in an lp steepest-descent update is not merely a straightforward unbiased estimator of the full- batch direction, making a standard "SGD-sttle" analysis more involved.
Accelerated f,-based method. Drawing inspiration from primal-dual interpolation techniques in the con- vex setting (Allen-Zhu & Orecchia, 2017; Nesterov, 2005; Diakonikolas & Guzmán, 2024), we design STACEY, an accelerated E, descent algorithm specif- ically tailored to -un-Euclidean smooth optimization (Section 4).
General Ep convergence guarantees for non-convex problems. We first establish stochastic fp steepest descent guarantees E [+)] S e at a rate of *), under standard prriance-bounder and f-- smoothness assumptions, where we let p¹ == p-1 (Section 4.1). Our results strictly generalize previous guarantees for signSGD (p = o.).
_-
Methods for curvature-awarc optimization. Another line of work exploits local geometry by incorporating second-order information, such as the Hessian or Fisher information matrix, and develops techniques for their ef- ficient approximation. K-FAC (Martens & Grosse, 2015) approximates the Fisher information matrix using layer-wise fronecker-factordd preconditioners for efficient second- order updates. Shampoo (Gupta et al., 2018; 8; Morwani et al., 2025; Vyas et al., 2025) similarly employs per-dimension Kronecker- factored preconditioners to approximate the gra- dient's mecond-moment matrix, enabling scalable curvature- aware optimization for tensor-structured parameters. Sophia (Liu al., 2024) further improves scalability by approxi- mating the diagonal of the Hessian using second-order mo- mentum. In contrast, our method, STACEY, is a first-order approach that leverages non-Euclidean geometry, rather than local curvature, through a differing Ep norm.
Practical performance on large-scale tasks. We com- pare STACEY against popular optimizers such as SGD, Adam, AdamW, and Lion tasks ranging from im- age classification to pretraining large language models (Section 5). Our experiments show that STACEY can converge faster and achieve higher accuracy than these baselines, particularly when the geometry of the objec- tive departs significantly from the Euclidean setting.
Flexible norm choices. We further evaluate different values of p E (2, 2o) across various model architec- tures and datasets, illustrating the potential advantages of tailoring the choice of norm to the problem geome- try.
Taken together, our results highlight the importance of non- Euclidean perspectives for contemporary machine learning tasks, offering both theoretical insight and practical improve- ment over classical (E2-based) and sign-based (bae-based) optimizers.
Why ba-ba methods help for large models. A key motivation for exploring ---norms with p E oo) stems from recent studies on the Hessian spectrum of large neural networks (Ghorbani et al., 2019; Papyan, 2018). In particu- lar, Ghorbani et al. (2019) provide evidence that the Hessian eigenvalue density can be highly non-uniform, leading to large curvature in certain subspaces while others remain comparatively flat. Under standard ba-baded (Euclidean) as- sumptions, these directions of high curvature can inflate the global smoothness parameter L2, potentially slowing con-
2. Related Work
Methods for non-Euclidean geometries. A significant line of research has studied sign-based methods, which can be viewed as (stochastic) steepest descent under the lo norm. For instance, Bernstein et al. (2018) introduced signSGD and analyzed its convergence properties through
We provide a comparison of &2 majorization and & smooth- ness conditions Appendix B.
2

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E Ep-Smooth Nonconvex Optimization
vergence or complicating optimization. By transitioning to fp-smoothness for p > 2, one can sometimes leverage rea- sonable Lipschitz constant Lp. as high-curvature directions may not always dominate in the same way.
Importantly, we emphasize the appearance here of ||-1|2 not only in measuring smoothness, but also for the | x 11 term.
Unfortunately, the standard analysis of AGD does not read- ily adapt to alternative notions of smoothness, as the de- sign of the algorithm is, in a sense, specific to Euclidean settings; we refer the reader to the work of (Allen-Zhu & Orecchia, 2017) for further discussion of this basic incom- patibility. Nevertheless, several works (Diakonikolas & Guzmán, 2024; Nemirovskii & Nesterov, 1985; Nesterov, 2005; Song et al., 2021) Orecchia (2017) with optimal rates in the Euclidean set- ting (Nesterov, 2018; Bai & Bullins, 2025), have provided techniques for accelerating in non-Fuclidean settings. In particular, the approach of Nemirovskii & Nesterov (1985), for convex fxx) that is smoohh with respect to I1-Ilp- leads to guarantees of the form
Formally, a function f is f--smooth if, for all y E IR.
- S Lp |0 1 sllp:
where p² I is the dual value. Thus, the choice of p shifts how curvature in different coordinates or sub- spaces affects f.. Because large-scale models often exhibit anisotropic Hessians (Cohen et al., 2021; Li et al., 2020a), an Ep analysis can better mirror the true geometric struc- ture of the objective. This observation aligns with analyses in (Balles et al., 2020), where sign-based methods (i.e., (x) can exploit flat directions effectively; by continuity, lp norms for p E (2, 0o) may interpolate between purely Euclidean and purely sign-driven behaviors.
that of Allen- &
Two main factors motivate the study of general lp np-norms (2 < p < 00) in large-scale training:
Tppl1500 - x"11
f(xr)
SO f(rr)- ff
(1)
T
pi2
1. Hessian Geometry and Tail Behavior. Large neu- ral networks often exhibit Hessians whose eigenval- ues and singular vectors follow nontrivial (sometimes heavy-tailed) distributions (Ghorbani et al., 2019; Pa- pyan, 2018). By choosing p to better accommodate outlier directions or to exploit more uniform curvature across coordinates, one can leverage better effective ey-smoothness constant Lep-
(See also, e.g., Theorem 2 in (Diakonikolas & Guzmán, 2024).) Moreover, these rates are likewise known to be tight (Guzmán & Nemirovski, 2015).
Looking closely at these convergence guarantees, we may first note that, for p = 2, the rate in equation recovers that of Nesterov (1983). On the other hand, for p - oc, while || - 13 can, at best, be as small as d | - x 11. we also have that limp-oo T- benefit of acceleration disappears altogether-ant and in fact this (limiting) rate essentially matches that of unaccelerated for steepest descent (Kelner et al., 2014). Consequently, these observations reveal the opportunity afforded by (non- Euclidean) b.-based accelerated methods for other values of p < xc, resulting from this trade-off between the depen- dence on the problem geometry and the rate of acceleration.
21. = TT-1in in which case the
2. Balancing Sparse and Dense Updates. Methods at = oc (sign-based) produce coordinate-wise updates of the same magnitude, while ba-based approaches "spread out" updates proportionally to gradient mag- nitudes. In high dimensions, intermediate lp steps can yield a better trade-off between these extremes, potentially improving both speed of descent and gener- alization (Cohen et al., 2021; Li et al., 2020a).
3. Preliminaries and Assumptions
Trade-off's for non-Euclidean acceleration. Alongside this matter of defining (and parameterizing) smoothness, there is second lens through which we observe the poten- tial for general p, namely that of acceleration (Allen-Zhu & Orecchia, 2017; Bai & Bullins, 2024a; Nemirovskii & Nesterov, 1985; Nesterov, 1983; 2005). As we further dis- cuss in Section 4. there is a fundamental trade-off (for con- vex settings) between the rate of acceleration and the norm used to measure the initial distance to the optimal solution. Concretely, it is well known that, for convex fzz that is L-smooth with respect to 1|-1,,, the classic accelerated gradi- ent descent (AGD) method of (Nesterov, 1983) converges at the rate f(XT) - f(x*) ) < O Ellxo-= and this rate is indeed tight (Nesterov, 2018; Nemirovskij & Yudin, 1983).
Throughout we let ||-Il and |I-I|+ denote a general norm and its dual, respectively. In addition, we specify II- to denote the standard lp norm (1 S p S oc) and || Il-IIp- := (1-llp/ip-1) to denote its dual norm. For symmetric M E Raxd s.t. M > O, we further let ||-I|as denote the standard matrix norm, i.c., ||
Vz Mæ for x e RR For a vector o Rd, we use superscript, i.e., v(0) to denote the 7th coordinate of v, and we let diag(e) denote the diagonal matrix such that diag(r)i. = v(0. We use subscript, e.g.. O,, to denote a vector in the tth iteration. For brevity, we use Jr for the true gradient (f(),) and ÿ. for the stochastic gradient g(,)). We use sgn (.) to denote the sign function and TtJ to denote the indicator function.
72
3

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E, fp-Smooth Nonconvex Optimization
Algorithm 1 (TACEY Optimizer
Algorithm 2 (p.p) Optimizer
input p, B1, 82, 0; T, 1/, A,
input p, 3,, 52, a, T, 1): 6, A
initialize 0b, 20: mo +- 0
initialize 9o. 20, ma < 0
1: while fu+1 not converged do
1: while 0(+1 not converged do
2: ft 9 s.t. [Elg] = (0.)
2: ft < s.t. Eløj = (0.)
3: Ct-1 f 8im++ + (1 - B1)g
3: c++1 + Bim, + (1 - 31)9e
4:
s (x) = (sf(2), ... [xx) where
4: s xx)) = (.), ((()))
where
x1())
x1
(x)) = [z(00 5-1 +
(de
of (x) [x(+)
Vie
p-
p-1 +
5: Dt+1 + 0: - 8//86 (c++1)
5: V++1 < 02 1718 (c(+1)
6: 21+1 = 2 ac(+1
6: 2/41
34 E [d]
7: 0.+1 = T + (1 - -7) - 70
|//0 u1) (1) -an +
8: M = 82 + (1 - -B2)Øt
9: end while
7: Oc+1 T2++1 (1 T) 102
10:
return 0z+1
8: M++1 = 82ms (1 B2)9
9: end while
10: return 8[+1
We may also consider the following equivalent definition of bp smoothness.
Algorithm 3 Stochastic lp Descent
Assumption (Smoothness in l, norm). Let f: R.d + R be L-smooth w.t.. ll- for p 2 2. Then, for all z, y E Rd,
input p, 7, S.00
I: for t 0 to T - do
2: s(x) = [s (z). (z(2 where
[f()) - f(x) (f(x) (v- x)| S lly - x 11
sx(x) x(1)
2
Vi E [d]
4. Accelerating Stochastic Steepest Descent
3: 0c+1 0, ns (gx)
gr s.t B = (0,)
Inspired by previous techniques in non-Euclidean acceler- ation (Allen-Zhu & Orecchia, 2017; Nesterov, 2005) as well as their successes, e.g., (Bullins, 2020; Jambulapati
4: end for
5: return 0r
et al., 2019; Sherman. 2017; Sidford & Tian, 2018) we introduce a practical acceleration scheme called STACEY (Algorithm 1), which is specifically designed for by-based methods. Central to our approach is its reliance on two se- quences of stochastic steps: 1) one sequence based on the standard lp steepest descent direction (line 5), which we show is theoretically well-grounded in the stochastic non- convex setting; 2) another sequence- whose combination with the first ultimately leads to acceleration- on a gradient descent direction (line 6), whose details we will further discuss.
141 -
which is derived from
o
the closed form of
{ (n9i, A 0-0)) + 110 11
0.+1 = arg min
When = 0o, Algorithm 3 reduces exactly to signSGD (Bernstein et al., 2018).
For p > 2, we show in Theorem I that stochastic 4p descent converges in expectation to an c-approximate stationary point with respect the dual norm at rate of O(E 4), thereby generalizing the previous guarantees for signSGD (p xo). In addition, we provide here proof sketch, deferring the complete proof to Appendix A. Curiously, as we will see, moving from the Eg setting (or even from the f setting) introduces certain technical considerations that need to be addressed non-trivially. As standard in stochastic and
4.1. Analyzing Stochastic Ep Descent
In this section, we present the stochastic é, descent al- gorithm, which serves as the fundamental framework of our approach, and establish its convergence guarantees. As shown in Algorithm 3, its update step takes the un- scaled form² of its counterpart in the deterministic setting
This is in line with signSGD (Bersstein al., 2018) compared to the scaled form in (Balles et al., 2020). In addition, we adopt
the unscaled version for clearer convergence analysis and more practical implementation.
4

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E --Smooth Nonconvex Optimization
non-Euclidean settings (Ghadimi & Lan, 2013; Bernstein al., 2018), we rely on the following assumptions.
we have
L
G
Assumption 2 (Unbiased Estimate). The stochastic gradi- ent ÿ is an unbiased estimate of the true gradient g. That is, 1
R
Σ
ge
lpp-
t=0
I.
(oo ( + 2p -- - ||قا||
Assumption 3 (Bounded Variance). For some data &, the variance of each coordinate of the stochastic gradient is bounded, i.e. Vi [d]. E g S of.
No
p p-1
Lee, Algorithm 3 takes N E O (c=)) gradient queries to reach an e-approximate stationary point.
Corollary 1. By Assumption 3. Blllg gaاle S o² where for a= | 22. 5 = [o...... _oa]
Proof Sketch. Starting with Assumption and the descent step in Algorithm 3,
Corollary 2. If the stochastic gradient is an n-sample mini- batch estimate, then'r Vi E [d]. Bllg() - [ of
[(+++) S s(0.) 1 (gt, s(ge)) + + (ye, s(g:) - -((r)
Assumption 4 (Bounded gradient). For G > 0, p 2 2. and p" where" 1 + p- = 1, |r1 G.
Lap
+ 2 ((()))
Corollary 3. By Assumption 4, we know that
(a) l|91اp- || I E |11 S G with Jensen's inequality.
where A = nll In conventional first-order anal- ysis, the inner product term B is supposed to cancel out after taking expectation. In contrast, the closed- form stochastic fp descent update is coordinate wise re-scaled, which makes the descent step biased, that is, E[s(ā)] # s(f())). In the literature on biased gradient descent (Stich & Ajallocian, 2020; Demidovich et al., 2023), the bias terms simply accumulate as constants and do not decay with the iterations. Thus, this term requires novel techniques to guarantee convergence. (1((
(b) Vield, |0 S G and |
We briefly justify the necessity of Assumption 4, which arises from additional technical challenges. Specifically, the coordinate-wis re-scaled update introduces bias un- der standard assumptions, preventing the direct application of conventional expectation and variance analyses as we later elaborate in detail. Notably, similar assumptions are also made when analyzing problems with complex struc- tures, such as stochastic compositional (Wang al., 2017), composite (Wang et al., 2024; Duchi al., 201 1b), and federated optimization (Li et al., 2020b; Yuan et al., 2021; Bai & Bullins, 2024b). Now we introduce the convergence result for lp steepest descent in the stochastic non-convex setting.
Noticing that sz(z)
B = 11
sgn
91
19/
sgn
10
=9 Σ
9i
(1)
(195 | pー +
13t 1=b
I (sen (v( H (o6t)
i=1
+
gt 1=1
(
19
p'a
10t
11 "Psau (9t") =sga ) (z!") ()
Theorem 1 (Main). Running Algorithm 3 on some (possibly non-convex) function that satisfies Assumptions 1 to 4 yields
Denote the first term as B₁ and the second B2. The 1g. (9) 1 p- + (g( | term in Bi can be bounded by 2G with Corollary 3, after which we take expectation, turning the indicator into a probability, and Lemma
T T
5o ff
L
E
<
1/T
2
0=0
in Appendix A.I shows E [Bi] Markov's inequality.
< FnG [alli using Vnn
1 T-1
2p-1 Go |لل|
Tht
B₂ requires more sophisticated handling since we cannot push the expectation through due to the data dependence of the term
where f(0) and f+= = f(0),, T4 is the batch size in iteration t and L, o, and G are constants from Assumption 1, 3, 4. Further letting the batch size T4 = T. the number of gradient call N T for iterations. With 1) =
1g (0)
1 م| nor does
P sgn (9) = sgn 9 give us much information. We instead take the oeroth-order Taylor expansion so
sgn
5

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E, Ep-Smooth Nonconvex Optimization
that 4 E [d] (()) between Di
(8) and Dt such that
(1)
(primal) steepest descent perspective (line 4 Algorithm 1), while the other functions in dual capacity (line 5 Algo- rithm 1). An apparent distinction is that momentum, as separate functionality, can be applied on top of the accel- eration scheme in STACEY (p.2). as demonstrated in lines 3 and 7 of Algorithm 1, for both the steepest descent and the (Euclidean) mirror descent.
) # (sl" ((0)
lo
p-d
+ p
In addition, we have
(6( p-1 - 99 p-1
A Buclidean- two-sequence interpolation was adopted by Schedule-Free SGD/AdamW (Defazio et al., 2024), which removes explicit learning-rate schedules while re- taining strong performance. In the realm of non-Euclidean methods, we contrast our algorithm with Lion-K. (Chen et al., 2024; Bernstein et al., 2018). While at first glance it may seem that these methods may simply be a rewriting of each other (based on the choice of parameters), a closer inspection on the very first step reveals that such is not the case:
1
sgn(((1) (((0))
3 (9.))
p-
(1(0) = sgn
Furthermore, given sgn
it is either
9i
(6)
|<00
<
Di or (8)
(1)
(8)
Ap-
pendix A.1 Lemma 3 shows that E either case.
[B2]
in
Term C is usually turned into mean-squared error that co- incides with variance in an unbiased setting, which the bounded variance assumption can directly handle. This is not the case for our setting. It is worth noting that the analysis of signSGD (Bernstein et al., 2018), a special case of the l, setting with p oc, was able to push through due to its update being in the very form of the sign of the gradient, which is in itself bounded by the constant 1. Our update, in contrast, is more complicated with the absolute value of the coordinates of the gradient in the denomina- tor, which is only lower bounded by O, or some e > 0 at best. Therefore, we directly apply Assumption and
Lion-KC 01 = 11==-//((( ((1 - -3))0))
STACEY (p:2): 01 = -- - T)ns ((L - 91)30)
-a f 31)
where KC(-) = I1-Ipp and s (.) is defined in Algorithm 1. The key difference of STACEY (p.2) lies in the convex combi- nation of a steepest descent step and a gradient descent step, whereas Lion-KC is composed of only the steepest descent step. They coincide only when T = 0 for STACEY i.e., completely getting rid of the "'couplin"", which then defeats the purpose of our acceleration. In addition, there is no choice of parameters for Lion-K to recover linear coupling. As a result, they are not iterate-equivalent, which further highlights the fundamental difference between "momentum" and "acceleration", distinction which, crucially, does not appear in the case of standard (Euclidean) AGD, i.c., when both steepest and mirror descent steps are with respect to Euclidean norms.
C= -- 11 T < = Moving term A to the left hand side, telescoping across iterations, and dividing both sides by 1/T completes the proof.
2
4.2. f, acceleration
We would note that for smooth convex optimization, (deter- ministic) gradient descent can be accelerated to achieve a rate of (()TT)). However, for stochastic first-order meth- ods, has been shown that a) in convex settings, SGD can- not improve upon the standard ((1/VT) rate when noise parameter is large enough (Agarwal et al., 2009), and b) in first-order smooth non-conver settings, SGD cannot be accelerated (in theory) without additional assumptions (in terms of gradient norm minimization), due to known lower bounds (Arjevani et al., 2023). Nevertheless, standard practical implementations of SGD are frequently designed to introduce some notion of acceleration with momentum (e.g., (Bernstein et al., 2018: Sutskever et al., 2013)), "push- ing" the converging sequence further along the direction of previous gradients.
Further inspired by the fact that (p.2) breaks the symmetry (in primal and dual trajectories) by coupling an lp steepest descent step with an based mirror descent step, we present the natural variant STACKY((.P) (Algorithm 2), for which we group Lp steepest descent with a mirror descent step having 1 11-113 (whose pd -ordr uniform convexity is useful for non-Fuclidean acceleration (Adil et al., 2024; Contreras et al., 2024; Song et al., 2021)) as its distance generating function. The closed-form mirror descent update is presented in line 5 of the algorithm.
5. Experiments
In this section, we present empirical evidence that the STACEY optimizer outperforms other optimizers in both convergence speed and accuracy. We evaluate STACEY's 's ef- fectiveness on image classification (Section 5.1), and LLM
In contrast, we take the view of acceleration not as a "push- ing" (in the Euclidean sense), but rather as a (dynamic) interpolation of two iterate sequences: one acting from a
6

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated fp-Smooth Nonconvex Optimization
Table Image classification CIFAR the 50th, 100th, and 200th epochs. demonstrating both superior accuracy and faster convergence.
STACEY consistently outperforms other optimizers,
Training NLL
Testing ACC (%)
Optimizer
@50 epach
epoch
@200 epoch
@50 epoch
100 epoch
@ @200 epoch
SGD Momentum
0.0567 000117
0.0441
0.0014
0.0352
0.0012
91.15 = 0.30
92.02 0.24
92.76
Adam
0.0401 0.0077
0.0182 1 0.0017
0.0083
0.0010
91.69 1 10188
92.13 1 1.16
92.66 0.36
A
0.0590 0.0010
0.0278
i: 0.0000
0.0195
0.0015
90.59
1 0.37
91.47 1: 0.42
92.12 0.00
Lion (Chen al., 2023)
0.1006 0.0571
0.2226
1.140
0.0245
0.0043
89.38 2.02
89.19 1. 1,88
92.15 0.3.2
STACEY (p.p)
0.0423
0.0009
0.0118
0.0014
0.0021
0.0011
91.88
1:211
92.79 0.16
93.79 1.0.8
STACEY (p.2)
0.0614 0.0031
0.0131 + 0.0027
0.0014 0.0005
90.83 1 0.32
92.70 1 0.28
93.54 1.0.0
Table Image classification on ImageNet the 20th, 40th, and 60th epochs. STACEY demonstrates superior test accuracy and faster convergence compared to other optimizers.
Training NLL
Testing Top-1 ACC (%)
Optimizer
@0 epoch
@44 epoch
@60 epoch
@ epoch
@740 epoch
@60 epoch
SGD w/ Momentum
2.0731 1:
1.7926 1: 0,9822 +: 0.0017
1.4993 0.00003
56.34 1
63.54 0.09
68.81 0.54
AdamW
1.3337 +: 0.0008
0.7395 1: 0.0029
66.12 0.53
68.47
0.14
69.31
1.005
Lion (Chen et al., 2023)
1.3529 +: 0.0007
1.0948 +: 0.0126
0.8605 +: 0.0045
67.66 1.003
68.43
0.10
69.62 2.111
STACEY (Pa))
1,4680 = 0.0150
1.1636
0.0159
1.0324 + 0.0100
66,93 0.10
69.15
0.15
69.87
0014
STACEY (p.2)
1.8376 0.0134
1.3781
+: 0.0187
1.1983 0.0120
60.89 0.12
0.16 66.34
67.56 0.15
125
4.
Thinigg
Frech
Egech
Tocation
(a) Training loss of STACEY (p.p) (b) Testing ACC of STACEY (p.p)
(a) Training loss of STACEY (b) Testing ACC of STACEY
b>
Tesus Recuuao (55)
1aa
Training
500
140
8..0
Epech
Fpnch
(c) Training loss of STACEY (n.2) (d) Testing ACC of STACEY
(c) Training loss of STACEY (p.2) (d) Testing ACC of STACEY
Figure I. Learning curves of CIFAR classification with varying no-no..
Figure 2. Learning curves ImageNet classification the first epochs with varying ny-norm
In real-world large datasets, such as training from scratch on ImageNet (Deng et al., 2009) and LLM (LLaMA (Tou- vron et al., 2023)) pretraining on C4 dataset, we further demonstrate the necessity of utilizing different fp-norms for specific tasks. For example, in the CIFAR (Krizhevsky, 2009) image classification, an Ey-norm for p close to 2 de- livers the best performance (Section 5.1), consistent with the effectiveness of Euclidean suceidean-based optimizers. In contrast, an no-morm with p around 3 proves more effective LLM
pretraining (Section 5.2). The hyperparameter choices and tuning are summarized in Appendix C.
In all experiments, STACEY optimizer by comparing it against other optimiz- ers as baselines including SGD (with momentum) (Nes- terov, 1983; Polyak, 1964), Adam (Kingma & Ba, 2015), W (Loshchilov & Hutter, 2019), and Lion (Chen et al., 2023).
underscore the efficiency of the

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E Ep-Smooth Nonconvex Optimization
Table Training and testing loss of L.LM pre-training at series of steps. The proposed STACEY optimizer consistently achieves lower loss than baselines at all steps.
Training Loss
Testing Loss
Optimizer
65% step
10K steps
@020k steps
@330 steps
step
e0kk steps
@20k steps
@30 steps
SGD w/ Momentum
6,6704 + 0.01 29
6,5205 = 0.0088
6,4206 = 0.0055
6,3920 :: 0,004
6.6558
0.0131
6.1150=0.0000
6,4173 0.0038
6,3909 #: 0.0008
Adam
6,4548
0,0028
6,3647
0.0037
6,2851 0.0030
6.2485
0.0028
6.4493 0.0011
6.3646 :: 0.0035
6.2820 0.0037
6.2480 0.0022
A
5.6655
0.0095
5,5172
0.0081
5,4401 1: 0.0091
5,4268
1: 0.0096
5.6510 1: 0.0099
0.0099
5,5171
1: 0.0080
5.4350
1: 0.0088
5.4240 1: 0,0009
(Chen al., 2023)
6.8722
0.0656
6,8190
0.0549
6,8021
0.0451
6.7794
0.0425
6.8624 0.0587
6,8220
0.0500
6.7954
6.7733-10.00411
STACEY
5,4016
0.0107
4,9938
0.0209
4.6492
0.0112
4.4962
0.0123
5.3616
1 0.0068
4,9655
0.0169
4.6372 1 0.0116
4,4879 1: 0.0122
STACEY(p.2)
6,2492
0.0060
6,0038
0.0319
5.7210
0.0363
5,5841
0.0379
6.2312 0.006
5,9867 1 0.0313
5.7022 170022+0.0377
0.0055
yields the best performance for the CIFAR dataset when using the ResNet18 architecture.
5
Traisns Loo
Tertnng Lons
15
12
Training from scratch on ImageNet. We train ResNet50 (He et al., 2016) with a batch size 256 on ImageNet (Deng al., 2009) for 60 epochs. 3 The learning rate schedule is cosine decay with 10K steps of warm-up, and the mix- mix-precision training is used to reduce the memory footprint. The learning curves are shown in Table 2.
5
(a) Training loss of STACEY(p.p)) (p.p) (b) Testing loss of STACEY (p.p)
5.2. Pretraining Large Language Models (LLMs)
fi
Teaing Lose
We pretrain 11ama-100m (Touvron et al., 2023) on the C4 subser' using various optimizers with cosine scheduler. The training and testing loss results, as presented in Ta- ble 3, show the advantage of STACEY over alternative algo- rithms. We additionally compare in Fig. 3 the performance of STACEY across different choices of p. whereby we ob- serve the best performance when p 3, which contrasts with the best results being observed when p = 2 in the CI- FAR image classification tasks, as discussed in Section 5.1.
63
Iication
(c) Training loss of STACEY (x.2) (d) Testing loss STACEY (p,2)
Figure 3. I Learning curves of LLM pretraining the first 30K iterations with varying -n-nrrm.
pretraining (Section 5.2). These results highlight the impor- tance of developing Eucldiedd optimizers and adjusting the choice of fp-norm to enhance performance across dif- ferent tasks, and we would note this choice may further benefit from, e.g., parameter-free approaches (Jacobsen & Cutkosky, 2022).
5.3. Discussion
As we observe throughout the experiments, STACEY demon- strates superior performance over SGD, which showcases its ability to adapt to a broader range of non-Fuclidean geome- tries. This adaptability verifies STACEY's 's convergence for general lp- -p-norms, making better choice for optimization tasks that present complex geometries and extend beyond the conventional Euclidean frameworks.
5.1. Image Classification
We demonstrate improved accuracy and faster convergence of the STACEY optimizer across image classification tasks of varying scales, consistent with our algorithm's 's design for acceleration.
Compared with Adam (Kingma & Ba, 2015) and AdamW (Loshchilov & Hutter, 2019), the results of STACEY suggests that the introduced acceleration technique well- aligned with the principles of non-Euclidean optimization. In addition, they highlight how STACEY's acceleration mechanism, which is designed for wider range of non- Euclidean structure, can yield better performance than tradi-
Training from scratch on CIFAR. We train ResNet18 (He et al., 2016) dataset (Krizhevsky, 2009) for 200 epochs, with the results presented in Table 1. We report training NLL and testing accuracy at the 50th, 100th, and 200th epochs. The proposed STACEY optimizer consistently outperforms all compared optimizers. As shown in Fig. p-norm of 2
the CIFAR
Due to computational resource limitations, the batch sizes used this paper are smaller than those in Lion's original pa- per (Chen et al., 2024). 'nttpa:/ wttpgggiinaacc.cc//aa.aooo datablations. 64 -wwsss/tessss.ss.....o subsets.
8

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E Ep-Smooth Nonconvex Optimization
tional adaptive methods.
partially supported by NSF AF 2209509 and NSF CDSE 2152687.
Furthermore, STACEY's improved performance over Lion (Chen et al., 2023) highlights the effectiveness of in- terpolating primal and dual sequences as an acceleration strategy, in contrast to simply incorporating momentum. The primal-dual interpolation ensures a more balanced and stable progression towards optimality, leveraging informa- tion from both primal and dual sequences. This strategy allows STACEY achieve faster convergence, even in chal- lenging settings and complex tasks like large-scale image classification and pretraining LLMs.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Algorithmic efficiency. We observe that STACEY has a 2d memory overhead, as needs to store both momemtum and a dual vector. This matches the memory overhead of Adam, which requires storing two moment vectors, and the per-iteration cost, in terms of basic arithmetic operations, is also comparable to that of Adam. Whereas methods such as SGD with momentum and Lion require only a single momentum vector, we would note that the overhead of the additional dual variable in STACEY is precisely what enables its b based acceleration.
Adil, D., Bullins, B., Jambulapati, A., and Sidford, A. Con- optimization with p-norm oracles. arXiv preprint arXiv:2410.24158, 2024. (Cited on page 6.)
Adolphs, L., Kohler, J., and Lucchi, A. Ellipsoidal trust region methods and the marginal value of hessian in- formation for neural network training. arXiv preprint arXiv: (ate::9955002/11 2019. (Cited on page 1.)
Agarwal, A., Wainwright, M. J., Bartlett, P., and Raviku- mar, P. Information-theoreti lower bounds on the oracle complexity of convex optimization. Advances in Neural Information Processing Systems, 22, 2009. (Cited on page 6.)
6. Conclusion
In this paper, we have presented a new approach to stochas- tic non-convex optimization by leveraging non-Euclicnn Euclidedn tp geometry. We first established that stochastic Ep steep- est descent converges at a rate of O(e expectation to stationary point under fy-smoothness assumptions, thus strictly generalizing previous analyses for signSGD (p x.). Building on these foundations, we introduced STACEY, an accelerated algorithm that combines stochas- tic lp descent with primal-dual interpolation techniques to effeetively navigate non-Euclides optimization landscapes.
Allen-Zhu, 7. and Orecchia, I.. Linear coupling: An ulti- mate unification of gradient and mirror descent. In 8th Innovations in Theoretical Computer Science Conference. Schloss Bagstuhl-Leebnii--Zeee (Cited on pages 1, 2, 3, and 4.)
für Informatik, 2017.
Arjevani, Y., Carmon, Y., Duchi, J. C., Foster, D. J., Srebro, N., and Woodworth, B. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199 (1):165-214, 2023. (Cited on page 6.)
Our results highlight how acceleration in lp spaces can yield improved geometry-dependenl performance compared to Euclidean and -bsed updates. In extensive experi- ments on large-scale image classification and language mod- eling, STACEY consistently achieved faster convergence and higher accuracy than popular optimizers such as SGD, AdamW, W. and Lion. Moreover, we demonstrated the versa- tility of choosing different p E (2, oo) to tailor the descent geometry diverse model architectures and datasets. Over- all, our contributions underscore both the theoretical and practical benefits of pursuing -on-Euclddean perspectives for addressing the complexities of modern machine learning
Bai, S. and Bullins, B. Faster acceleration for steepest descent. arXiv preprint 2xXiv:2409.19200 2024a. (Cited page 3.)
Bai, S. and Bullins, B. Local composite saddle point op- timization. In International Conference on Learning Representations, 2024b. (Cited on page 5.)
Bai, S. and Bullins, B. Tight lower bounds under asymmet- ric high-order Hölder smoothness and uniform convexity. In International Conference on Learning Representations, 2025. (Cited on page 3.)
tasks.
Balles, L.,, Pedregosa, F., and Roux, L. The geometry of sign gradient descent. arXiv preprint arXiv:2002.08056, 2020. (Cited on pages 2, 3, 4, and 17.)
Acknowledgements
We thank Jincheng Zhou for helpful discussions related to the experiments implementation. Petros Drineas was
Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anand-
kumar, A. signsgd: Compressed optimisation for non-

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E fp-Smooth Nonconvex Optimization
convex problems. In International Conference on Ma- chine Learning, pp. 560-569. PMLR, 2018. (Cited pages 2, 4, 5, 6, and 16.)
Diakonikolas, and Orecchia, L. The approximate duality gap technique: A unified theory of first-order methods. SIAM Journal on Optimization, 29(1):600-689 2019. (Cited on page 1.)
Bullins, B. Highly smooth minimization of non-smooth problems. In Conference on Learning Theory, pp. 988- 1030. PMLR, 2020. (Cited on page 4.)
Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 201 la. (Cited on page 1.)
Carmon, Y., Duchi, J. C., Hinder, O., and Sidford, A. "Con- Vex until proven guilty": dimension-free acceleration of gradient descent on non-convex functions. In Interna- tional Conference on Machine Learning, pp. 654-663. PMLR. 2017. (Cited on page 1.)
Duchi, J. C., Agarwal, A., and Wainwright, M. Dual aver- aging for distributed optimization: Convergence analysis and network scaling. IEEE Transactions on Automatic control, 57(3):592 5265996666666 201 1b. (Cited on page 5.)
Chen, I.,, Liu, B., Liang, K., and Liu, Q. Lion secretly solves constrained optimization: As æyapunov predicts. In International Conference on Learning Representations, 2024. (Cited pages 2, 6, and 8.)
Ghadimi, S. and Lan, G. Stochastic first-and zeroth -oderr methods for nonconvex stochastic programming. SIAM journal on optimization, 23(4))2341-2368, 2013. (Cited on pages 1, 2, and 5.)
Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Pham, H., Dong, X., Luong, T., Hsieh, C.-J,, Lu, Y., et al. Sym- bolic discovery of optimization algorithms. Advances in Neural Information Processing Systems, 36, 2023. (Cited on pages 1, 2, 7, 8, 9, and 17.)
Ghorbani, B., Krishnan, S., and Xiao, Y. An investiga- tion into neural net optimization via Hessian eigenvalue density. In International Conference on Machine Learn- ing. pp. 21022222222222 PMLR, 2019. (Cited on pages 1, 2, and 3.)
Gupta, V., Koren, T., and Singer, Y. Shampoo: Pre- conditioned stochastic tensor optimization. In Interna- tional Conference on Machine Learning, pp. 1842-1850. PMLR, 2018. (Cited on page 2.)
Cohen, J., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. Gradient descent on neural networks typically oc- curs at the edge of stability. In International Conference on Learning Representations, 2021. (Cited on pages and 3.)
Guzmán, C. and Nemirovski, A. On lower complexity bounds for large-scale smooth convex optimization. Jour- nal of Complexity, 31(1):1-14, 2015. (Cited on pages 2 and 3.)
Contreras, J. P., Guzmán, C., and Martinez-Rubio, D. Non- euclidean high-order smooth convex optimization. arXiv preprint arXiv:2411 1.0987, 2024. (Cited on page 6.)
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016. (Cited on page 8.)
Defazio, A., Yang. X., Khaled, A. Mishchenko, K., Mehta, H., and Cutkosky, A. The road less scheduled. Ad- vances in Neural Information Processing Systems, 37: 9974-10007, 2024. (Cited on page 6.)
Jacobsen, A. and Cutkosky, A. Parameter- free mirror de- scent. Conference on Learning Theory, pp. 4160-4211 PMLR, 2022. (Cited on page 8.)
Demidovich, Y., Malinovsky, G., Sokolov, I., and Richtárik, P. A guide through the zoo of biased SGD. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https: /pppr//iie//eeeeee/eeeeem / Doruumid==Ctv4Nyyall
Jambulapati, A., Sidford, A., and Tian, K. A direct Ö(1/c)
iteration parallel algorithm for optimal transport. Ad- vances in Neural Information Processing Systems, 32, 2019. (Cited on page 4.)
(Cited on page 5.)
Deng, J., Dong, W., Socher, R., Li, L.JJ,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pat- tern Recognition, pp. 248 248-255. leee, 2009. (Cited on pages 7, 8, and 17.)
Jiang, K., Malik, D., and Li, Y. How does adaptive optimiza- tion impact local neural network geometry? Advances in Neural Information Processing Systems, 36, 2024. (Cited on page 1.)
Diakonikolas. J. and Guzmán, C. Complementary com- posite minimization, small gradients in general norms, and applications. Mathematical Programming, pp. 1-45, 2024. (Cited on pages 1, 2, and 3.)
Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., and Jordan, M. How to escape saddle points efficiently. In Interna- tional Conference on Machine Learning, pp. 1724-1733 1732. PMLR, 2017. (Cited on page 1.)
10

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E, fp-Smooth Nonconvex Optimization
Karimi, H., Nutini, J., and Schmidt, M. Linear conver- gence of gradient and proximal-gradient methods under the Polyak-Lojasiewiez condition. In Machine Learn- ing and Knowledge Discovery in Databases: European Conference, ECML. PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16, pp. 795- 811. Springer, 2016. (Cited on page 16.)
Nemirovskij, A. S. and Yudin, D. B. Problem Com- plexity and Method Efficiency in Optimization. A Wiley-Interscience publication. Wiley, 1983. ISBN 9780471103455. (Cited on page 3.)
Nesterov, Y. method for solving the convex programming problem with convergence rate O(1/k2). In Dokl. Akad. Nauk SSSR, volume 269, pp. 543, 1983, (Cited pages 1, 3, and 7.)
Kelner, J. A., Lee, Y. T., Orecchia, L.,, and Sidford, A. An almost-linear-timg algorithm for approximate max flow in undirected graphs, and its multicommodity generaliza- tions. In Proceedings of the Twenty-Fifth Annual ACM- SIAM Symposium on Discrete Algoritums, pp. 217-226.
Nesterov, Y. Smooth minimization of non-smooth functions. Mahematical programming, 103:127-152, 2005. (Cited on pages 2, 3, and 4.)
SIAM, 2014. (Cited on page 3.)
Nesterov, Y. Lectures on convex optimization, volume 137.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations. 2015. (Cited on pages 7, 8, and 17.)
Springer, 2018. (Cited on page 3.)
Papyan, V. The full spectrum of deepnet Hessians scale: Dynamics with sgd training and sample size. arXiv preprint 2, and 3.)
Krizhevsky, A. Learning multiple layers of features from tiny images. Master's thesis. University of Tront, 2009. (Cited on pages 7, 8, and 17.)
2018. (Cited on pages 1,
Polyak, B. T. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964. (Cited on pages 1 and 7.)
Li, X., Gu, Q. Zhou, Y., Chen, T., and Banerjee, A. Hessian based analysis of sgd for deep nets: Dynamics and gener- alization. In Proceedings of the 2020 SIAM International Conference on Data Mining, pp. 190-198. SIAM, 2020a. (Cited on pages and 3.)
Robbins, H. and Monro, S. A stochastic approximation method. The Annals of Mathematical Statistics, pp. 400- 407, 1951. (Cited on pages and 2.)
Li, X., Huang, K., Yang. W., Wang, S., and Zhang, Z. On the convergence of fedavg on non-iid data. In International Conference on Learning Representations, 2020b. (Cited on page 5.)
Sherman, J. Area-convexity, &oo regularization, and undi- rected multicommodity flow. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Comput- ing. pp. 452 460, 2017. (Cited on page 4.)
Liu, H., Li, Z., Hall, D. L. W., Liang, P., and Ma, T. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In International Conference on Learn- ing Representations, 2024. (Cited on page 2.)
Sidford, A. and Tian, K. Coordinate methods for accelerat- ing lo regression and faster approximate maximum flow. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science, pp. 922-933. IEEE, 2018. (Cited on page 4.)
Loshehilov, and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019. (Cited on pages 1, 7, 8, and 17.)
Martens. J. and Grosse, R. Optimizing neural networks with froneccer-factored approximate curvature. In Interna- tional conference on machine learning, pp. 2408-2417. PMLR, 2015. (Cited on page 2.)
Song, C., Jiang, Y., and Ma. Y. Unified acceleration of high- order algorithms under general holder continuity. SIAM Journal on Optimization, 31(3):1777-1826 2021. (Cited on pages 3 and 6.)
Morwani, D., Shapira, I., Vyas, N., eran malach, Kakade, S. M., and Janson, L. A new perspective on shampoo's preconditioner. In The Thirteenth In- ternational Conference on Learning Representations, 2025. URL https: http://ppp...ee..neeeeem wwt.:nettooornoeee id=c62I3Cp8c6, (Cited on page 2.)
Stich, S. U. and Ajalloeian, A. Analysis of sgd with biased gradient estimators. arXiv preprini S0t::220080005... 2020. (Cited on page 5.)
Sutskever, I., Martens, J., Dahl, G., and Hinton, G. On the importance of initialization and momentum in deep learn- ing. In International Conference on Machine Learning, pp. 1139-1147. PMLR, 2013. (Cited on page 6.)
Nemirovskii, A. S. and Nesterov, Y. E. Optimal methods of smooth convex minimization. USSR Computational Mathematics and Mathematical Physics, 25(2):21-30, 1985. (Cited on pages 1, 2, and 3.)
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E.,

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E fp-Smooth Nonconvex Optimization
Azhar, F., al. Llama: Open and efficient foundation language models. arXiv preprint 2 2023. (Cited on pages 7 and 8.)
Vyas, N., Morwani, D., Zhao, R., Shapira, I., Brand- fonbrener, D., Janson, L., and Kakade, S. M. SOAP: Improving and stabilizing shampoo using adam for language modeling. In The Thirteenth In- ternational Conference on Learning Representations, 2025. URL https::/ http://ppprreeeeeeeeeeeeeeee n forum? ID=IxXXE.P (Cited on page 2.)
Wang, M., Fang, E. X., and Liu, H. Stochastic composi- tional gradient descent: algorithms for minimizing com- positions of expected-value functions. Mathematical Pro- gramming, 161:419 449, 2017. (Cited on page 5.)
Wang, Y., Chen, S., Jiang, W., Yang. W., Wan, Y., and Zhang, L. Online composite optimization between stochastic and adversarial environments. In The Thirty- eighth Anmual Conference on Neural Information Pro- cessing Systems, 2024. URL https: / pppprevee....eeeeeoeeee/. net/ forum? dd=
(Cited on page 5.)
Yuan, H., Zaheer, M., and Reddi, S. Federated composite optimization. In International Conference on Machine Learning, pp. 12253-12266 PMLR, 2021. (Cited on page 5.)
12

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated fp-Smooth Nonconvex Optimization
A. Proofs
A.1. Complete Proof for Theorem
Theorem
Running Algorithm 3 on some (possibly non-convex) f that satisfies Assumptions I to 4 yields
So fo-5 T/T
L-1
1 T
2p
G |dllz
E
Σ
1pp
T
2
where fo fo= = Further letting the batch size n2 T. the number of gradient call N T for T iterations. With rj Lic we
f(0)) and f f(0)), ne the batch size iteration and L, G, and G are constants from Assumption 3, 4.
have
E
19
1pp
1 N
G.G (5-
2p G |ō|| +
T
p-1
ie., Algorithm 3 takes N E (e gradient queries to reach an e-approximate stationary point.
Proof.
Starting with Assumption 1 and the descent step in Algorithm 3,
T,
f(0++1) s f(t,) + (ge, 8c+1 - 0.) + |0(+1 - lell
2
L
= =((0) + 1 (gt, - -g()) + ((()))
2
Ln²
= f(0)) - n (9t, s(g))) + 1 (ge. s(gt) s(o))
+
2
B
Now we analyze these terms one by one.
91
A =
(
i=
19t
Σ
9
(8)
p
|1
=
For term B,
d
B
1
91
(1)
gt
19/
1=1
8g1
9t
1gg
sgn
g
10
(1)
$ 17
91
(1)
sgn
gt
(2)
19.
(3)
sgn
gr.
(2)
p-
77
91
(4)
l0!
ssnn (ol
sgg (o!
11
17
91
(8)
19t
(8)
15t
Jann (at =sgon (7(99
1=1
B4
13

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated fp-Smooth Nonconvex Optimization
Bi is bounded in expectation by 2n ||
in Lemma 2 and B2 is bounded in expectation by #GF | in Lemma 3.
1p-1)/0
Lap
C
2
3
Lai²
(8)
P-
9
2
Lup
119..11
2
Laj²
G
Therefore,
7(2p 182p-1)GF+ 1)G 10111 GG
y
|1ge [bp s(0,) s(0++1) +
(p /
2
By telescoping through t 0,... T L, get
1
f(0o) (106)-((0) f(0r)
1
T 1
(2p 1) 1-
L
EE
1gell
llp
T
1/
T
1) Tm n
2
1=
Lemma 2.
E
1
d
9t
(2)
(19.
(9)
- +
II
(9(
Fugn (3!
2n
|lll
1=1
Proof. By Corollary (b).
d
(
(196 (0)
(+)
E
1J
9t
)08ss (0.0) #agn (3.
i=1
d
<
21/ p-
EI
91
sggu 9.0))
#egn [s
9:
(1)
P
sgn (6
*
sgn gp(
d
S G
(2)
(3)
91
d
211 <
9:
(4)
1=1
9t
d
(0) 27
<
221
E
9t
1== 03
<
Gm -- ||611
where for the last three inequalities we used Markov's inequality, Jensen's 's inequality, and Assumption 3.
14

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E 5--Smotth Nonconvex Optimization
Lemma 3. E
17
1
96
(6)
(8)
106
(8)
sgn (o{") - (xf*)
nG ||ll
(p-
Proof. Denoting
E
sgn
gi (
sgn
g ]
as 1= [-], and IP
sgm
9i (0)
sgn
91
(6)
as
IP:
|-
1
9t
9
(1)
p-
sEn (3( =gn [5
n
10. 10 -
10
TP
nBB
96
198 (0)
10!2
P|
PP
n
0 --
sgn (1))
|500 ママ
g
(
9. ))
10.
p-
P P[-]
p-
= =
sgu(((1))
|510))
P|-1
1 p
E
|,00
06
P|=],
p
(1) at 91 with Lagrange remainder,
(6)
in which the second equality holds by taking the zeroth order Taylor expansion of 9i
and (1) is in the range from Pr to
(8) (()
Given sgn Di = sgn g by the definition of ((0) in the Lagrange remainder,
we must have either ge
9:
(0)
or 9t 2 5 (0) 2 9r Now we analyze these two cases respectively. We
write out the derivations separately for clarity and simplicity, alternatively one can merge these two cases with the law of total expectation.
(1)
If
ge
(i)
[())
(() 9t
then
d
17
B
9i
(
[(9
gi
96
(0)
P|-|
P
"
E
(1)
((0
P
96
(3)
96
(1)
P|-]
P
p
n
EE
(())
P-T
91
(6)
91
IP
p
11
9t
(())
91
(
[P-]]
p
G
E
9
(4)
9
P|-|
P
p
G
E
96
p
PP-
G
E
9t
(6)
(Jensen's 's)
p
GT
(Assumption 3)
p
VTtz
G/ ||-||
(p-1) 1)
15

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated 6 fp-Smotth Nonconvex Optimization
(2)
If
2
(0))
then
"
(1)
E
91
9t
p
d
11
E =
9t
(0)
1t
(1)
9t
(8)
(0)
IP
p
(p-
[)
E
9t
(
(0)
36
(9)
P[-]
17 d
22
9t
E
9t
Schwarz)
7 Σ
E
96
(4)
212 p2
o?
(Assumption 3)
T4
11
(
of
9
(Jensen's 's)
T2e
17
Var
+
E
gt
(8)
2
p-1
(Variance Definition)
Y1g
11
Σ
(1)
2
E g6
2(2-m)
9t
d
17
9r
(8)
p-1
O2
(Assumption 2)
v TTe
1=1
G |
1) m) 74
Combining these two cases together (e.g., by the law of total expectation) completes the proof.
B. le Majorization and lp Smoothness
An assumption of interest, studied by Bernstein et al. (2018) (as well as Karimi et al. (2016)), is that of & majorization (with respect to L. [I,,..., La)), meaning that for all x. y Rdd
d - 100
( f(x) - ()()) V - x)
2 1=1
We may equivalently express this condition as 1-smoothness w.r.t. 11.1|z. where L = diag(T.) T.), i.e., for all x. E Red
[Vv(()) - (z) S ly =
Interestingly, we may observe that, for any 1 < S oc and letting p² = = we have
(V(())
V=(a))la/// < Vf(x) vf(x)) S -11 2111 S 11211//110 1lv x1|20-
1 1/2
where the first inequality holds by reverse Hölder's inequality, i.c., for u, u u Rd, E (n(0 0000002 - (where we choose q = 3p-1). and the last inequality holds by Hölder's inequality.
16

---

STACEY: Promoting Stochastic Steepest Deseent via Accelerated E Ep-Smooth Nonconvex Optimization
Rearranging, have (y) V/(=))]20//(pp I || Ilv =1lep- and it follows that, for p 2, majorization implies |ل 10__--smoothnss, w.r.t. I|-1pp- Thus, while this condition is sufficient to entail fp smoothness (as previously noted by (Balles al., 2020) in the case of p oo), nevertheless prefer work directly with Ep smoothness assumptions, we believe they provide more natural pairing for the methods we consider.
C. Hyperparameter Choices
summarize the hyperparameters used our experiments Tables 4, 5, and 6. These hyperparameters are determined through a grid search. Specifically, we perform search to identify appropriate values for the fn-norm, learning rate 7, a, and weight decay X. This process involves initial rough comparison across range of magnitudes, followed by more precise grid search to determine the optimal values.
For fair comparison, all experimental settings, apart from the listed hyperparameters, follow the original papers of AdamW (Loshchilov & Hutter, 2019) and Lion (Chen et al., 2023), and are kept consistent across all optimizers. For example, data augmentations for ImageNet (Deng et al., 2009) and CIFAR (Krizhevsky, 2009) all include random cropping and random horizontal flipping.
Table 4. CIFAR byper-parameters.
Model
Optimizer
Batch
p
Learning
Schedule
À
ResNet-18
Momentum
cosine decay
0.0002
ResNet-18 18
Adam (Kingma & Ba, 2015)
128
0.001
cosine decay
0.999
0.0005
1e-8
R8sNet-18
AdamW (Loshchilov & Hutter, 2019)
128
0.01
cosine decay
0.9
0.999
0.0005
le-8
RosNet-18
(Chen al., 2023)
0.001
decay
0.01
R8
STACRY (p.r)
128
2
0,1
cosine decay
0.1
0.9
0.99
0.01
0.001
1o-12
ResNet-
STACEY
128
2
0.1
cosine decay
0.1
0.9
0.99
0.01
0.001
1e:12
Table 5. ImageNet hyper-parameters
Model
Optimizer
Batch Size
Learning Rale
Schedule
9g
ResNet-50
256
decay
ResNei--0
AdamW (Loshchilov & Hutter, 2019)
256
0.002
cosine decay
0.9
0.999
0.005
le.4
al., 2023)
decay
0.99
ResNet-50
STACEY(poe)
256
0.01
decay 0.001
0.9
0.999
0.001
0.001
ResNet-50
STACEY (p.2)
256
288
0.01
cosine decay
0.001
0.9
0.999
0.001
0.001
10-8
Table 6. Hyper parameeasssssrrs for LLM pretraining.
Model
Optimizer
Size
F.
Learning
Schedule
11ame 100m
SGD w/ Momentunt
16
0.01
cosine decay
0.0005
11ama- m0
Adam (Kingma & Ba, 2015)
0.0001
cosine decay
0.9
0.999
0.01
1e.8
11ama- 10Cm
AdamW (1 (Loshchilov & Hutter, 2019)
16
0.0001
cosine decay
0.999
0.05
1e-8
11ama 100m
Lion (Chen al., 2023)
16
0.05
cosine decay
0.9
0.999
0.01
1lama 10Cm
(T.CEY
3
0.01
cosine docay
0.1
0.99
0.01
0.001
1c-8
11ama- 10
STACEY
16
2.8
0.01
cosine decay
0.1
0.99
0.0005
0.001
1c-8
17


Reviews:
[{"review_id": "oMAPeLJwVo", "reviewer": "Reviewer_NAH2", "summary": "This paper introduces **STACEY**, a novel optimization algorithm designed to accelerate stochastic steepest descent via ℓp-smooth nonconvex optimization. The key contributions of this work include:\n- The development of **STACEY**, which incorporates primal-dual iterate interpolation to improve convergence rates for non-Euclidean smooth optimization problems.\n- A theoretical framework that generalizes both SGD (when \\( p = 2 \\)) and signSGD (when \\( p = \\infty \\)), with a **convergence guarantee of \\( O(\\epsilon^{-4}) \\)** under standard assumptions.\n- Empirical results demonstrating **superior convergence speed and final accuracy** compared to existing optimization methods, including SGD, Adam, AdamW, and Lion, on large-scale deep learning tasks such as image classification (CIFAR, ImageNet) and large language model (LLM) pretraining.\n- A study on how different values of \\( p \\) affect performance, showing that non-Euclidean norms can be more effective in certain settings than traditional ℓ2-based methods.\n\nThe paper is **well-written, well-structured, and presents both strong theoretical contributions and compelling empirical results**.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We thank the reviewer for their thoughtful comments and suggestions, and we respond to their questions below.\n\n> Detailed **computational complexity in terms of runtime and memory consumption** compared to SGD, Adam, and Lion\n\n**Runtime**\n\nLet $d$ be the number of parameters, and let each “basic operation” refer to simple scalar arithmetic (e.g., an addition, multiplication, or sign). We will ignore lower-level details (e.g., hardware vectorization) and focus on how many scalar operations are performed per parameter per iteration. \n\n**SGD**\n\nKey steps:  \n\n1) $m_i \\leftarrow \\beta m_i + (1-\\beta)\\nabla_i \\quad (\\text{2 multiplications, 1 addition})$.\n2) $x_i \\leftarrow x_i - \\alpha m_i \\quad (\\text{1 multiplication, 1 addition})$.\n\nApproximate ops per parameter: ~5–6 scalar ops (fewer if no momentum is used).\n\n**Adam**\n\nKey steps: \n\n1) $m_i \\leftarrow \\beta_1 m_i + (1-\\beta_1)\\nabla_i\n\\quad (\\text{2 multiplications, 1 addition})$.\n2) $v_i \\leftarrow \\beta_2 v_i + (1-\\beta_2)\\nabla_i^2\n\\quad (\\text{3 multiplications, 1 addition})$.\n3) $x_i \\leftarrow x_i - \\alpha\\frac{m_i}{\\sqrt{v_i} + \\epsilon}\n\\quad (\\text{1 square root, 1 division 1 multiplication, 1 addition})$.\n\nApproximate operations per parameter: ~9–12 scalar ops (slightly more if you count the bias-correction steps separately).\n\n**Lion**\n\nKey steps:\n\n1) $m_i \\leftarrow \\beta m_i + (1-\\beta)\\mathrm{sign}(\\nabla_i)\n\\quad (\\text{2 multiplications, 1 addition, 1 sign operation})$.\n2) $x_i \\leftarrow x_i - \\alpha\\,\\mathrm{sign}(m_i) \\quad (\\text{1 sign, 1 multiplication, 1 addition})$.\n\nApproximate operations per parameter: ~6–7 scalar ops (including sign as an operation).\n\n**Stacey**\n\nKey steps:\n\n1) Update the “momentum-like” buffer $m_i$ (coordinate-wise re-scaling).\n2) Update the dual vector $z_i$ (coordinate-wise multiplications/additions).\n3) Combine $m_i$ and $z_i$ to get the final parameter update.\n\nRepresentative breakdown:\n1) Update $m_i$: ~3–5 ops.\n2) Update $z_i$: ~3–4 ops (coordinate-wise re-scaling plus addition).\n3) Final parameter update: ~2–3 ops (linear combination and one addition/subtraction).\n\nApproximate operations per parameter: ~9–12 scalar ops.\n\n**Memory Footprint**\n\n**SGD**:Momentum buffer (if used). Total auxiliary overhead: $d$\n\n**Adam**: First moment and second moment. Total auxiliary overhead: $2d$\n\n**Lion**: Momentum-like buffer. Total auxiliary overhead: $d$\n\n**Stacey**: Momentum-like buffer and dual vector. Total auxiliary overhead: $2d$\n\nThus, Stacey’s memory requirement is similar to Adam, and slightly more than other single-pass gradient methods (though we note that its overhead compared to e.g. SGD, Lion comes precisely from the additional dual vector).\n\n> **Choice of $p$: intuition, ablation, practical heuristics**\n\nAlthough it can be challenging to provide intuition for acceleration (even in the convex $p=2$, i.e. Nesterov's, case), at a high level the analysis carefully balances the primal and dual updates, leveraging the uniform convexity of $\\|\\cdot\\|_p^p$ and an alternative smoothness-derived upper bound, as in Definition 2 and Lemma 1 in [Diakonikolas & Guzmán 2024]. We will include an ablation study of the influence of $p$ on the acceleration rate in the updated manuscript, to help provide additional insight for these theoretical results.\n\nRegarding heuristics for choosing $p$, one idea would be to use occasional Hessian information to determine how to adjust $p$ over time, based on the spectrum density [Ghorbani et al. 2019]. Additionally, there has been much work on developing automated means of hyperparameter tuning, including e.g. parameter-free methods [Jacobson & Cutkosky 2022], from which we may hope to draw inspiration.\n\n> **Second-order curvature-aware optimizers: citations and discussion** \n\nWe thank the reviewer for pointing out these helpful references. We will cite them and incorporate the following discussion in the revised version. Shampoo [Gupta et. al., 2018], K-FAC [Martens & Grosse, 2015] and their follow-ups are indeed representative works of curvature-aware optimization methods. One notable difference is that these works are second-order methods that exploit the structure of the Hessian or Fisher information and focus on techniques for their efficient approximation, whereas our method, Stacey, is a first-order approach that explores non-Euclidean geometry though a differing $\\ell_p$ norm. We will include comparisons with these curvature-aware optimizers in the updated version of our manuscript.\n\nExtending Stacey to second-order methods is, in our view, a promising direction that aligns with our own considerations for future research as well. It is natural to investigate the non-Euclidean counterpart of a preconditioned gradient step, just as we have done in the first-order setting. Such a method holds the potential to benefit from both the curvature awareness provided by the Hessian information and the geometric advantages of operating under a differing $\\ell_p$ norm.", "claims_and_evidence": "The paper makes several claims, all of which are generally well-supported:\n- **Claim:** STACEY improves convergence rates over traditional SGD and adaptive optimizers.  \n  **Evidence:** Empirical evaluations on CIFAR, ImageNet, and LLM pretraining demonstrate improved speed and final accuracy.\n- **Claim:** The proposed method generalizes previous optimization approaches by considering a broader class of ℓp-norms.  \n  **Evidence:** Theoretical analysis rigorously proves this generalization.\n- **Claim:** STACEY benefits from the flexibility of choosing different \\( p \\) values.  \n  **Evidence:** Experiments with different values of \\( p \\) show that problem-specific norm choices can yield better performance.\n\nOne area where **the claims could be strengthened** is in providing a more detailed computational complexity comparison to confirm the practical efficiency of STACEY in large-scale settings.\n\n---", "methods_and_evaluation_criteria": "The methods and evaluation criteria are **appropriate and well-justified**:\n- **Optimization Benchmarks:** The paper evaluates STACEY on well-established benchmarks, including CIFAR-10, ImageNet, and LLM pretraining tasks.  \n- **Comparisons:** The algorithm is compared against widely used optimizers (SGD, Adam, AdamW, Lion), which are the **correct baselines** for this type of work.  \n- **Evaluation Metrics:** The paper reports **training loss, test accuracy, and convergence speed**, which are standard and relevant evaluation criteria for optimization algorithms in deep learning.\n\nHowever, **an additional runtime or computational cost comparison** would be useful to fully assess the trade-offs of using STACEY in practice.\n\n---", "theoretical_claims": "I reviewed the theoretical claims and found them **mostly sound**:\n- The **convergence proof for stochastic ℓp steepest descent** follows standard assumptions and **logically extends previous results**.\n- The **use of primal-dual interpolation** is well-motivated, and the acceleration rate follows existing literature on non-Euclidean acceleration.\n- The **generalization to different ℓp norms** appears correct and aligns with prior research on optimization under non-Euclidean norms.\n\nOne possible **area for clarification** is the effect of different values of \\( p \\) on the acceleration exponent. **An intuitive explanation of how acceleration scales with \\( p \\) would improve clarity.**\n\n---", "experimental_designs_or_analyses": "The **experimental design is well-structured**, with meaningful comparisons. I verified the validity of:\n- **Image classification experiments** on CIFAR and ImageNet.\n- **LLM pretraining experiments** on the C4 dataset.\n- **Ablation studies** exploring the role of \\( p \\).\n\nPotential improvement:\n- A **detailed analysis of computational efficiency (runtime per iteration, memory overhead, etc.)** would make the comparisons more complete.\n- It would be useful to explore **STACEY’s performance in additional non-Euclidean problems**, such as adversarial training or reinforcement learning.\n\n---", "relation_to_broader_scientific_literature": "This paper is **well-grounded in prior research** on stochastic optimization and non-Euclidean geometry in machine learning:\n- It builds upon **SGD (Robbins & Monro, 1951), signSGD (Bernstein et al., 2018), and AdamW (Loshchilov & Hutter, 2019)**.\n- It connects with recent studies on **non-Euclidean optimization (Diakonikolas & Guzmán, 2024)** and **adaptive gradient methods**.\n- The idea of **primal-dual interpolation** is influenced by **work on non-Euclidean acceleration (Allen-Zhu & Orecchia, 2017; Nemirovskii & Nesterov, 1985)**.\n\nThis paper **extends these ideas in a meaningful way**, demonstrating both theoretical and empirical improvements.\n\n---", "essential_references_not_discussed": "The paper **covers most essential references**, but **a comparison with curvature-aware optimizers (e.g., Shampoo, K-FAC)** would be useful. These methods also attempt to handle **non-Euclidean optimization challenges**, making them relevant to the discussion.\n\nI recommend citing **Gupta et al. (2018) on Shampoo** and **Martens & Grosse (2015) on K-FAC** to highlight how STACEY differs from these approaches.\n\n---", "other_strengths_and_weaknesses": "### **Strengths**\n- **Strong theoretical foundation** with **generalized convergence guarantees**.\n- **Comprehensive empirical validation** showing consistent improvements over baselines.\n- **Clear writing and structured explanations**.\n\n### **Weaknesses**\n- **No computational cost analysis** (memory and runtime comparisons are missing).\n- **Hyperparameter selection for \\( p \\) is unclear** (no systematic guidance provided).\n- **Missing comparisons with curvature-aware optimizers** (e.g., Shampoo, K-FAC).\n\n---", "other_comments_or_suggestions": "- **Clarify the intuition behind acceleration for different values of \\( p \\)**.\n- **Include an ablation study on the effect of different values of \\( p \\) on generalization performance**.\n- **Discuss whether STACEY could be extended to second-order methods or mixed-order approaches**.\n\n---", "questions_for_authors": "1. **How does the per-iteration computational cost of STACEY compare to SGD, Adam, and Lion in terms of runtime and memory consumption?**  \n2. **Could you provide a practical heuristic or automated procedure for selecting \\( p \\) in different tasks?**  \n3. **How does STACEY perform when compared to second-order methods like Shampoo or K-FAC? Would these methods benefit from a similar ℓp-based approach?**  \n\n---", "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": "", "extra_scores": {}}, {"review_id": "2ENhW8kci2", "reviewer": "Reviewer_d7Gi", "summary": "This paper introduces Stacey an optimisation algorithm targeted at training deep neural networks (DNNs). Stacey generalises SignSGD and conventional SGD, in a p-norm sense, where SGD uses the 2-norm to measure distance and SignSGD uses the inf-norm. On top of this Stacey include a acceleration scheme to aid the speed of convergence. The paper offers a theoretical result, specifically a convergence rate for an non-accelerated version of Stacey on smooth stochastic problems with bounded gradient and variance.  Additionally there is an empirical evaluation of Stacey on standard DNN benchmarks against popular deep learning algorithms.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We thank the reviewer for their thoughtful comments and suggestions, and we respond to their questions below.\n\n> **Missing Reference**\n\nThanks for the suggestion. We will provide a citation and include it as a baseline in the updated manuscript.\n\n> **Hyperparameters**\n\nThough we may tune $\\tau$ and $\\alpha$, we found that, similar to SGD and Adam, the best choice comes from a small set, i.e., $\\tau \\in [0.001]$ and $\\alpha \\in [0.1, 0.01, 0.001]$, and so we believe this provides a fair comparison with other methods. We did run comparisons where a limited number of hyperparameters were tuned, while the rest were left at default values, as such defaults let us reduce the scope of the search.\n\n> **Questions and Comments on Experiments**\n> 1. Lack of error bars/variation\n\nFor space reasons, we kindly point the reviewer to our rebuttal for Reviewer ddGc for the tables of results with error bars.\n\n> 2. Smaller scale classical optimisation problems are missing \n\nWhile our method was designed with large-scale models in mind, we will include smaller scale experiments, and we would kindly ask the reviewer for suggested problems.\n\n> 3. ImageNet result epochs\n\nThanks for catching this. This is a typo from an outdated table in the appendix, which we will update accordingly.\n\n> 4. Tables of some results seem to be missing final\n\nIf we understand correctly, we will include tables of the PPL/transformer pretraining results, with error bars/variance included (as in the tables provided in the rebuttal).\n\n> 5. The number of iterations shown vary\n\nThe iteration counts in Figures 2 and 4 differ because they represent two distinct experimental setups: one for ImageNet classification, the other for LLM pretraining, each with its own training configuration and iteration schedule.\n\n> 6.  x) The transformer pertaining experiments details\n\nLLama-100 is adopted from the github repo of \"GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection\" [Zhao et al. 2024]. The $p=3.3$ line's ppl increases because it is within a range that our algorithm cannot converge.\n\n\n> 7. xi) Missing results for Stochastic ℓp Descent.\n\nStochastic lp descent is a special case of our algorithm; thus we have presented the better empirical versions of our algorithm (with acceleration), though we will include ablation studies in the updated manuscript.\n\n> 8. (Q4) Stacey(p,2) vs. Stacey(p,p)?\n\nBasing our intuition on the theory for the convex case, coupling a non-Eulidean primal with a Euclidean dual update (Theorem 1 in [Bai & Bullins 2024]), vs. non-Euclidean primal and dual updates (e.g. Theorem 2 in [Diakonikolas & Guzmán 2024]), leads to different trade offs between the problem geometry, initial iterate, and acceleration exponent, such that neither is uniformly better than the other.\n\n> 9. (Q6) Stacey for p=2?\n\nThis occurs due to numerical differences from handling general $p$. However, we acknowledge this could cause confusion, so we will provide the algorithm for the simplified case of Stacey(2,2), along with a discussion of this distinction.\n\n> **Choice of $p$**\n\n> (Q2) automatically work out \"p\"?\n\nOne idea would be to use occasional Hessian information to adjust $p$ over time, based on the spectrum density [Ghorbani et al. 2019]. There has also been much work on developing automated means of hyperparameter tuning, including e.g. parameter-free methods [Jacobson & Cutkosky 2022], from which we hope to draw inspiration.\n\n> (Q3) adjusting Stacey p choice per layer?\n\nThis is a wonderful suggestion. We may even benefit from a per-layer basis if we wished to leverage layer-wise Hessian information as a diagnostic tool (being far more compute-friendly).\n\n> **Theoretical Properties of Stacey**\n\nThe acceleration framework of Stacey builds on linear coupling [Allen-Zhu & Orecchia, 2017], which is optimal in the first-order smooth and convex setting, and HASD [Bai & Bullins, 2024], which achieves faster convergence in the $\\ell_p$ smooth non-Euclidean setting. Beyond the deterministic and convex regime, the core foundation of Stacey ---- non-convex stochastic $\\ell_p$ steepest descent ---- achieves a first-order oracle complexity of $\\mathcal{O}(\\epsilon^{-4})$, which we discuss as tight in Section 4.1 and 4.2. \n\nProviding a further theoretical characterization of Stacey’s acceleration remains challenging. Notably, even for widely used algorithms like Adam, existing theory typically shows only convergence or recovers first-order optimal regret, with limited formal evidence of superiority over other accelerated first-order methods. We suspect capturing the acceleration of Stacey in theory would require more refined and tailored assumptions. While we leave this as an open direction for future work, we highlight that Stacey generalizes both SignSGD and Lion ---- as discussed in the third paragraph of Section 4.2 ---- which suggests it offers greater flexibility and is more likely to admit improved theoretical properties.", "claims_and_evidence": "The theoretical claims seem well supported. \n\nThe empirical claims are supported however I do have some doubts about the fairness of the empirical evaluation when comparing to other methods, however it is difficult to be totally objective here, given the natural differences between optimisers. Inevitably the wider community will need to judge how Stacey performs in practice compared to existing methods.", "methods_and_evaluation_criteria": "The benchmarks look appropriate, it would have been nice to see some smaller networks and more classical optimisation problems considered, given these don't take much compute.", "theoretical_claims": "I did not thoroughly check the proofs due to time. The Theoretical Claims presented are not for Stacey but a far simpler non-accelerated algorithm.", "experimental_designs_or_analyses": "The experimental design looks okay, though there are definitely some flaws here, see weakness section.", "relation_to_broader_scientific_literature": "The Relation To Broader Scientific Literature is well motivated.", "essential_references_not_discussed": "This recent paper is missing, and would be great to see included as a baseline, specially given its lower number of hyperparameters. \n\nDefazio A, Yang X, Khaled A, Mishchenko K, Mehta H, Cutkosky A. The road less scheduled. Advances in Neural Information Processing Systems. 2024 Dec 16;37:9974-10007.", "other_strengths_and_weaknesses": "*Strengths*\n1) The paper is well written.\n2) The idea behind the paper seems neat and some related theoretical results are provided.\n3) The benchmarks considered look appropriate, of course it would be nice to see more, including some more classical non-deep learning optimisation problems.\n4) This seems a promising direction for further research to build on top of.\n\n*Weaknesses*\n1) The empirical experimental section has some flaws:\ni) Some recent baselines are missing from the experimental section, specifically the one introduced last year in \"The Road Less Scheduled\" \nii) The amount of hyperparameters tuned for Stacey seems to far exceed the other methods, making the comparison hard to judge.\niii) It is not clear how robust Stacey is to its hyperparameters, nor how much hyperparameter tuning is required to get good results. \niv) There is lack of error bars or mention of variation between runs\nv) Smaller scale classical optimisation problems are missing\nvi) For the ImageNet results It is not clear why the results are presented at epoch 60 rather than the typical epoch 90 given at least some of the experiment were run for 90 epochs, as report in the appendix\nvii) Tables of some results seem to be missing final, such as PPL on transformer pretraining.\nviii) The number of iterations shown vary between plots for no explained reason.\nix) Some of the results for Stacey(2,2) seem to be different between plots (figure 5&6) & (figure 3 vs 11&12)\nx) The transformer pertaining experiments seem to be missing some important details, model architecture what is LLama-100?, why test ppl increases, unexplained differences in performance between plots see point ix.\nxi) Missing results for Stochastic ℓp Descent.\n\n2) The totality of the above critics makes me wonder how much the results are being presented in a way to make them seem best. I would suggest spending a little more time in the appendix making it clear why specific choices were made. Without the empirical results the paper doesn't offer enough of a contribution in my opinion so it is essential a reader is not wondering why some seemly odd choice have been made in the way the experiments have been conducted and presented. My score is assuming greater clarity is given on the experiments in the final version of the paper.\n\n3) The theoretical results are not presented for Stacey but a far simpler non-accelerated algorithm.", "other_comments_or_suggestions": "It would be great to detail the grids searched over in terms of hyperparameters, not just the final values.\n\nSome idea of the robustness of stacey to its hyperparameters would really help sell its practicality.\n\nPlease explain why the number of iteration at which results (both tables and plots) are shown seem to vary so much.", "questions_for_authors": "1) Stacey has *a lot* of hyperparameters, many more than typically tuned for SGD and Adam, do you think your comparison is fair given this fact? Did you run a comparison where the same number of limited hyperparameters (say 1 or 2 only) are adjusted for all algorithms and the rest are left at their default values? Much of Adam (&AdamW) success is due to their robustness to it's hyperparamters, algorithms that need extensive hyperparameter tuning are unlikely to be used in practice. The experiments present in the paper do not make it clear to me how useful Stacey is as a practical algorithm. \n\n2) The p norm used in Stacey is a hyperparameter, do you think it would be possible to adept Stacey to automatically work out the best setting of \"p\" for a given problem during the optimisation process? This would would help reduce the amount of hyperparamters required.\n\n3) Further to the above, I understand the p norm used to measure distance in Stacey is fixed for all parameters, do you think it would be possible to extend stacey so it is adaptive adjusting the p-norm used, say per layer? \n\n4) For some tasks Stacey(p,2) does better and some Stacey(p,p) does better what do you think might be causing this discrepancy? \n\n5) Do you think Stacey with acceleration enjoys any theoretically properties? Did you make any progress to this end?\n\n6) Is Stacey(p,p) equivalent to Stacey(p,2) and Stacey(2,p), If p = 2. If so, why do they seem to behave differently in your experiments?", "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": "", "extra_scores": {}}, {"review_id": "dFQhQaw0L7", "reviewer": "Reviewer_ddGc", "summary": "The paper uses different mixed Lp norms to run SGD", "questions": "", "limitations": "", "rating": 1, "confidence": null, "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": "We thank the reviewer for their comments and suggestions.\n\n>**Error Bar and Standard Deviation**\n> No error bars or standard deviation in tables\n\nWe ran 3 random seeds to obtain the error bar.\n\nCIFAR:\n\n| **Optimizer** | **Train NLL @50** | **Train NLL @100** | **Train NLL @200** | **Test ACC @50** | **Test ACC @100** | **Test ACC @200** |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|\n| SGD w/ Momentum | 0.0567 ± 0.0017 | 0.0441 ± 0.0014 | 0.0352 ± 0.0012 | 91.15 ± 0.30 | 92.02 ± 0.24 | 92.76 ± 0.13 |\n| Adam | 0.0401 ± 0.0017 | 0.0182 ± 0.0017 | 0.0083 ± 0.0010 | 91.69 ± 0.18 | 92.13 ± 0.16 | 92.66 ± 0.36 |\n| AdamW | 0.0590 ± 0.0010 | 0.0278 ± 0.0009 | 0.0195 ± 0.0015 | 90.59 ± 0.37 | 91.47 ± 0.42 | 92.12 ± 0.07 |\n| Lion | 0.1006 ± 0.0571 | 0.2226 ± 0.1410 | 0.0245 ± 0.0043 | 89.38 ± 2.02 | 89.19 ± 1.88 | 92.15 ± 0.32 |\n| Stacey(p,p) | 0.0423 ± 0.0009 | 0.0118 ± 0.0014 | 0.0021 ± 0.0011 | 91.88 ± 0.21 | 92.79 ± 0.16 | 93.79 ± 0.38 |\n| Stacey(p,2) | 0.0614 ± 0.0031 | 0.0131 ± 0.0027 | 0.0014 ± 0.0005 | 90.83 ± 0.32 | 92.70 ± 0.28 | 93.54 ± 0.06 |\n\n\nImageNet: \n\n| **Optimizer** | **Train NLL @20** | **Train NLL @40** | **Train NLL @60** | **Test Top-1 ACC @20** | **Test Top-1 ACC @40** | **Test Top-1 ACC @60** |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|\n| Stacey(p,p) | 1.4680 ± 0.0150 | 1.1636 ± 0.0159 | 1.0324 ± 0.0100 | 66.93 ± 0.10 | 69.15 ± 0.15 | 69.87 ± 0.14 |\n\n\nLLM:\n\n| **Optimizer** | **Train loss @5K** | **Train loss @10K** | **Train loss @20K** | **Train loss @30K** | **Test loss @5K** | **Test loss @10K** | **Test loss @20K** | **Test loss @30K** |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n| SGD w/ Momentum | 6.6704 ± 0.0129 | 6.5205 ± 0.0088 | 6.4206 ± 0.0055 | 6.3920 ± 0.0048 | 6.6558 ± 0.0131 | 6.5150 ± 0.0085 | 6.4173 ± 0.0038 | 6.3909 ± 0.0038 |\n| Adam | 6.4548 ± 0.0028 | 6.3647 ± 0.0037 | 6.2851 ± 0.0030 | 6.2485 ± 0.0028 | 6.4493 ± 0.0017 | 6.3646 ± 0.0035 | 6.2820 ± 0.0037 | 6.2480 ± 0.0028 |\n| AdamW | 5.6655 ± 0.0095 | 5.5172 ± 0.0081 | 5.4401 ± 0.0091 | 5.4268 ± 0.0096 | 5.6510 ± 0.0099 | 5.5171 ± 0.0080 | 5.4350 ± 0.0088 | 5.4240 ± 0.0093 |\n| Lion | 6.8722 ± 0.0656 | 6.8190 ± 0.0549 | 6.8021 ± 0.0451 | 6.7794 ± 0.0425 | 6.8624 ± 0.0587 | 6.8220 ± 0.0500 | 6.7954 ± 0.0438 | 6.7733 ± 0.0413 |\n| Stacey(p,p) | 5.4016 ± 0.0107 | 4.9938 ± 0.0209 | 4.6492 ± 0.0112 | 4.4962 ± 0.0123 | 5.3616 ± 0.0068 | 4.9655 ± 0.0169 | 4.6372 ± 0.0116 | 4.4879 ± 0.0132 |\n| Stacey(p,2) | 6.2492 ± 0.0060 | 6.0038 ± 0.0319 | 5.7210 ± 0.0363 | 5.5841 ± 0.0379 | 6.2312 ± 0.0065 | 5.9867 ± 0.0313 | 5.7062 ± 0.0375 | 5.5755 ± 0.0375 |\n\n\n\n>**Baseline Settings**\n> Learning rates for the baseline Adam is very off for LLMs; its set to 1e-4 where it should really be set to 1e-3 or higher. Also, the epsilon for Adam is also off. Setting epsilon to 1e-8 basically makes Adam act as SGD + M. On LLMs the epsilon should be set closer to 1e-17.\n\nWe would kindly ask the reviewer to provide additional details regarding which LLMs they are referring to, as we did not observe a notable difference using the suggested parameters for our setting.\n| **Optimizer** | **lr** | **eps** | **Test loss @15K** | **Test loss @30K** |\n|:---:|:---:|:---:|:---:|:---:|\n| Adam | 1e-3 | 1e-17 | 6.4120 | 6.2341 |\n| Adam (Our settings) | 1e-4 | 1e-8 | 6.3102 | 6.2485 |\n\n>**Performance Boost**\n>I downloaded the code and ran it on a few benchmarks. I did not see a boost in modded-nanoGPT, which is a properly tuned baseline. I also did not see a boost vs Adam in some toy problems like xor \n\nWe would kindly ask the reviewer to elaborate on the details of their evaluation, as otherwise, we are unable to provide a proper response or explanation.", "claims_and_evidence": "The main takeaway is that there are different Lp norms boost performance of optimization for different problems. For CNN's for example they find L2 to work best but for LLMs they find another L3 to work better for example. Unfortunately there are no error bars or standard deviation in tables. Also the learning rates for the baseline Adam is very off for LLMs, its set to 1e-4 where it should really be set to 1e-3 or higher. Also the epsilon for Adam is also off. Setting epsilon to 1e-8 basically makes Adam act as SGD + M. On LLMs the epsilon should be set closer to 1e-17.", "methods_and_evaluation_criteria": "I believe the baselines were not set correctly.", "theoretical_claims": "The paper does a good job of contextualizing their method in modern optimization.\n\nThe paper gives some extensions to simpler claims from but good overall. \n\n[1] Guillaume Garrigos, Robert M. Gower, Handbook of Convergence Theorems for (Stochastic) Gradient Methods\n\n[2] Ahmed Khaled, Peter Richtárik, Better Theory for SGD in the Nonconvex World", "experimental_designs_or_analyses": "please see other sections.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The paper lacks variance bars or standard deviations as well as weakly tuned baselines.", "other_comments_or_suggestions": "Fix baselines.", "questions_for_authors": "NA", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces Stacey, an accelerated l_p steepest descent algorithm leveraging primal-dual interpolation for non-Euclidean optimization. Theoretical guarantees and empirical evaluations on image classification and LLM pretraining suggest improved convergence speed, final accuracy, and benefits from adaptive p-values, advocating for non-Euclidean methods over conventional approaches.

The reviews are highly mixed: Reviewers d7Gi  and Reviwer  NAH2 strongly support the work, emphasizing its novel theoretical insights into optimizer design, which could inspire future research. However, Reviewer ddGc raises significant concerns. The primary criticism centers on the experimental validity, arguing that Stacey’s performance may not surpass Adam, let alone state-of-the-art optimizers (e.g., Shampoo, Muon, Scion). The AC also notes Stacey’s hyperparameter complexity, which may limit practical adoption.

While the empirical claims require further validation, the theoretical contributions—particularly the framework for non-Euclidean optimization—align with ICML’s theory-oriented focus. Balancing these factors, the AC leans toward acceptance, as the work provides  inspirations for future research into simpler, effective optimizers for realistic deep learning scenarios.
Decision: Accept (poster)