NAG-GS: A Semi-Implicit Stochastic Optimizer with Large-Step Stability

Abstract

Stochastic first-order methods are widely used for large-scale optimization and machine learning, where their performance often depends on the learning rate. We study NAG-GS, a stochastic momentum method based on a semi-implicit Gauss--Seidel discretization of a Nesterov-type flow. For strongly convex quadratics, we determine the exact range of stable step sizes for fixed damping and the step size that minimizes the worst-case spectral radius. When the damping equals the strong-convexity parameter, the method recovers the accelerated asymptotic factor $(\sqrt{\kappa }-1)/(\sqrt{\kappa }+1)$, where $\kappa$ is the condition number. We also prove that the quadratic iteration converges linearly with dynamic damping when its limiting fixed-damping iteration is stable. For general strongly convex, possibly nonsmooth objectives, we establish an $O(1/k)$ bound on the expected objective gap of the last iterate, using a conditionally unbiased stochastic first-order oracle with bounded second moment. The result holds with dynamic damping from any ${\gamma }_0>0$. We show that NAG-GS can also be written as a Heavy-Ball iteration with explicitly determined step sizes and momentum coefficients. We test the stochastic result on nonsmooth hinge-loss classification problems, and further evaluate NAG-GS on quadratic problems, logistic regression, convolutional neural networks, RoBERTa, and Vision Transformers. The experiments show that NAG-GS can remain effective at learning rates where standard optimizers become unstable or lose accuracy.

arXiv comment30 pages, 13 figures. Revised title and exposition; stochastic convergence with dynamic damping, Heavy-Ball equivalence, and nonsmooth numerical experiments