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

TL;DR

Stacey introduces an accelerated $\, ext{l}_p$-steepest descent algorithm tailored for non-Euclidean optimization in deep learning, offering theoretical guarantees and empirical improvements over traditional methods.

Contribution

The paper presents a novel accelerated $\, ext{l}_p$-steepest descent algorithm with primal-dual iterates, addressing non-Euclidean structures in deep network training.

Findings

Faster convergence compared to SGD, AdamW, and Lion.

Higher final accuracy in image classification and LLM pretraining.

Effectiveness of non-Euclidean approaches demonstrated across datasets.

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 .