Stacey: Promoting Stochastic Steepest Descent via Accelerated $\ell_p$-Smooth Nonconvex Optimization

We thank the reviewer for their comments and suggestions.

Error Bar and Standard Deviation No error bars or standard deviation in tables

We ran 3 random seeds to obtain the error bar.

CIFAR:

ImageNet:

LLM:

Baseline Settings 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.

We 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.

Performance Boost 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

We 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.

We thank the reviewer for their thoughtful comments and suggestions, and we respond to their questions below.

Missing Reference

Thanks for the suggestion. We will provide a citation and include it as a baseline in the updated manuscript.

Hyperparameters

Though 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.

Questions and Comments on Experiments

1. Lack of error bars/variation

For space reasons, we kindly point the reviewer to our rebuttal for Reviewer ddGc for the tables of results with error bars.

2. Smaller scale classical optimisation problems are missing

While 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.

3. ImageNet result epochs

Thanks for catching this. This is a typo from an outdated table in the appendix, which we will update accordingly.

4. Tables of some results seem to be missing final

If 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).

5. The number of iterations shown vary

The 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.

6. x) The transformer pertaining experiments details

LLama-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.

7. xi) Missing results for Stochastic ℓp Descent.

Stochastic 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.

8. (Q4) Stacey(p,2) vs. Stacey(p,p)?

Basing 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.

9. (Q6) Stacey for p=2?

This 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.

Choice of $p$

(Q2) automatically work out "p"?

One 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.

(Q3) adjusting Stacey p choice per layer?

This 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).

Theoretical Properties of Stacey

The 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.

Providing 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.

We thank the reviewer for their thoughtful comments and suggestions, and we respond to their questions below.

Detailed computational complexity in terms of runtime and memory consumption compared to SGD, Adam, and Lion

Runtime

Let $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.

SGD

Key steps:

1. $m_i \leftarrow \beta m_i + (1-\beta)\nabla_i \quad (\text{2 multiplications, 1 addition})$.
2. $x_i \leftarrow x_i - \alpha m_i \quad (\text{1 multiplication, 1 addition})$.

Approximate ops per parameter: ~5–6 scalar ops (fewer if no momentum is used).

Adam

Key steps:

1. $m_i \leftarrow \beta_1 m_i + (1-\beta_1)\nabla_i \quad (\text{2 multiplications, 1 addition})$.
2. $v_i \leftarrow \beta_2 v_i + (1-\beta_2)\nabla_i^2 \quad (\text{3 multiplications, 1 addition})$.
3. $x_i \leftarrow x_i - \alpha\frac{m_i}{\sqrt{v_i} + \epsilon} \quad (\text{1 square root, 1 division 1 multiplication, 1 addition})$.

Approximate operations per parameter: ~9–12 scalar ops (slightly more if you count the bias-correction steps separately).

Lion

Key steps:

1. $m_i \leftarrow \beta m_i + (1-\beta)\mathrm{sign}(\nabla_i) \quad (\text{2 multiplications, 1 addition, 1 sign operation})$.
2. $x_i \leftarrow x_i - \alpha\,\mathrm{sign}(m_i) \quad (\text{1 sign, 1 multiplication, 1 addition})$.

Approximate operations per parameter: ~6–7 scalar ops (including sign as an operation).

Stacey

Key steps:

1. Update the “momentum-like” buffer $m_i$ (coordinate-wise re-scaling).
2. Update the dual vector $z_i$ (coordinate-wise multiplications/additions).
3. Combine $m_i$ and $z_i$ to get the final parameter update.

Representative breakdown:

1. Update $m_i$: ~3–5 ops.
2. Update $z_i$: ~3–4 ops (coordinate-wise re-scaling plus addition).
3. Final parameter update: ~2–3 ops (linear combination and one addition/subtraction).

Approximate operations per parameter: ~9–12 scalar ops.

Memory Footprint

SGD:Momentum buffer (if used). Total auxiliary overhead: $d$

Adam: First moment and second moment. Total auxiliary overhead: $2d$

Lion: Momentum-like buffer. Total auxiliary overhead: $d$

Stacey: Momentum-like buffer and dual vector. Total auxiliary overhead: $2d$

Thus, 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).

Choice of $p$: intuition, ablation, practical heuristics

Although 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.

Regarding 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.

Second-order curvature-aware optimizers: citations and discussion

We 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.

Extending 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.