math.OCSep 29, 2022

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

Authors: Valentin Leplat, Daniil Merkulov, Aleksandr Katrutsa, Daniel Bershatsky, Olga Tsymboi, Ivan Oseledets

Organizations: Skolkovo Institute of Science and Technology, Moscow, Russia · Moscow Institute of Physics and Technology, Moscow, Russia · AIRI, Moscow, Russia · Sber AI Lab, Moscow, Russia

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 (κ−1)/(κ+1)(\sqrtκ-1)/(\sqrtκ+1), where κκ 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)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 γ0>0γ_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.

Figures & tables

ProblemDampingStep sizeResult
Strongly convex quadratic, deterministicfixed γ>0\gamma>0constantexact Schur-stability interval; optimal fixed step for γ≥μ\gamma\geq\mu
Strongly convex quadratic, deterministicdynamic γk\gamma_{k}constantlinear convergence when the limiting γ=μ\gamma=\mu system is Schur stable
Strongly convex quadratic, stochasticfixed γ>0\gamma>0constantstationary covariance inside the Schur-stable region
General strongly convex, possibly nonsmoothdynamic γk\gamma_{k}decreasingO(1/k)O(1/k) rate in expectation under a conditionally unbiased oracle
Nonconvex learning problemsfixed, tunedfixed/tuned experimentallyempirical study only
Table 1: Parameter regimes considered in the analysis.
Figure 1: Number of iterations needed to reach the target tolerance as a function of the learning rate. The value 101010^{10} indicates divergence within the test protocol.
MethodGap after 10 epochsGap after 40 epochsTest accuracy after 40 epochs
NAG-GS1.32×10−6±6.54×10−71.32\times 10^{-6}\pm 6.54\times 10^{-7}5.71×10−8±2.99×10−85.71\times 10^{-8}\pm 2.99\times 10^{-8}85.10±0.03\mathbf{85.10\pm 0.03}
SSGD1.02×10−6±4.37×10−71.02\times 10^{-6}\pm 4.37\times 10^{-7}1.42×10−7±6.96×10−81.42\times 10^{-7}\pm 6.96\times 10^{-8}85.07±0.0285.07\pm 0.02
Pegasos4.00×10−7±1.79×10−74.00\times 10^{-7}\pm 1.79\times 10^{-7}5.37×10−8±8.70×10−9\mathbf{5.37\times 10^{-8}\pm 8.70\times 10^{-9}}85.09±0.0285.09\pm 0.02
AdaGrad5.07×10−4±1.12×10−45.07\times 10^{-4}\pm 1.12\times 10^{-4}2.61×10−4±9.67×10−52.61\times 10^{-4}\pm 9.67\times 10^{-5}85.01±0.7185.01\pm 0.71
Table 2: Binary MNIST classification with the strongly convex hinge-loss objective ( 37 ). The values are averaged over five random seeds. The objective gap is measured with respect to the high-accuracy reference solution.
Figure 2: Nonsmooth strongly convex binary MNIST experiment. The curves show the mean over five random seeds. NAG-GS uses γ=μ\gamma=\mu and the decreasing step size from Theorem 4.2 .
Figure 3: NAG-GS and matched HB ( γ=μ\gamma=\mu , αk=2/(k+2)\alpha_{k}=2/(k+2) ), each evaluating its own subgradient on the same batches. Lines show means; shading shows the min–max range over five paired seeds. Annotations give the maximum paired objective difference across all seeds and recorded iterates.
Figure 4: Dynamic update ( 6 ) with αk=2/(k+2)\alpha_{k}=2/(k+2) : means over five paired seeds. Shading shows the seed range for γ0=μ\gamma_{0}=\mu . Final spread compares the largest and smallest final mean gaps across the six initializations, relative to the smallest. Gray segments indicate a 1/k1/k slope, not a bound.
Learning rateEpochsOptimizerBest test accuracy
10−310^{-3}40NAG-GS88.47±0.0588.47\pm 0.05
SGD-M92.01±0.0492.01\pm 0.04
AdamW92.71±0.01\mathbf{92.71\pm 0.01}
AC-SA92.67±0.0592.67\pm 0.05
10−210^{-2}20NAG-GS90.99±0.0290.99\pm 0.02
SGD-M92.50±0.04\mathbf{92.50\pm 0.04}
Table 3: Multiclass logistic regression on MNIST. Accuracies are averaged over three random seeds.
Figure 5: Best MNIST test accuracy as a function of batch size for several learning rates. The learning-rate legend in the NAG-GS panel applies to all four plots.
Learning rateOptimizerBest test accuracy
10−310^{-3}NAG-GS66.15±0.8166.15\pm 0.81
SGD-MW78.84±0.2178.84\pm 0.21
10−210^{-2}NAG-GS78.60±0.4078.60\pm 0.40
SGD-MW81.58±0.2381.58\pm 0.23
0.10.1NAG-GS81.50±0.1481.50\pm 0.14
SGD-MW–
Table 4: Best VGG-11 test accuracy on CIFAR-10 over 50 epochs, averaged over three random seeds. A dash indicates divergence.
Figure 6: VGG-11 on CIFAR-10. NAG-GS with α=0.15\alpha=0.15 reaches high test accuracy faster in the first epochs than SGD-MW with α=0.01\alpha=0.01 . Curves are averaged over three random seeds.
Figure 7: Best VGG-11 test accuracy as a function of batch size and learning rate.
Figure 8: NAG-GS and SGD-MW on ResNet-20/CIFAR-10.
OptimizerCoLAMNLIMRPCQNLIQQPRTESST2STS-BWNLI
AdamW61.6087.5688.2492.6291.6978.3494.9590.6856.34
NAG-GS61.6087.2490.6992.5991.0177.9794.5090.2156.34
Table 5: RoBERTa fine-tuning on GLUE.
StageNAG-GSAdamW
After 1 epoch0.84190.8269
After 25 epochs0.86060.8324
Table 6: Vision Transformer fine-tuning on Food-101.
Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Distance between the empirical mean and the minimizer for several step sizes in the stochastic quadratic tests.
Figure 10: Initial and stationary samples in the case μ>L/2\mu>L/2 .
Figure 11: Initial and stationary samples in the case μ<L/2\mu<L/2 . The spread increases when the step approaches the stability boundary.
Figure 12: Empirical distribution comparisons in two scalar nonconvex tests.
Figure 13: Classification-error landscape for ResNet-20/CIFAR-10 under the NAG-GS hyperparameter search.
OptimizerMean, sVariance, s 2Rel. MeanRel. Variance
SGD0.4580.0081.01.0
NAG-GS1.6480.0453.65.5
SGD-M3.3740.0427.45.2
SGD-MW3.5120.0377.74.7
AdamW5.2080.10211.412.6
Adam7.9190.16917.320.8
Appendix
Table 7: Measured duration of one optimizer update in the ResNet-20/CIFAR-10 implementation.
More in this topic

Stochastic Optimization

Stochastic optimization studies how to find good decisions or model parameters when objectives, data, or feedback are random or noisy.

See all 210 papers

Explore similar work

CardsList

  1. Ordinary Nonconvex SGD under Distance-Dependent Moments: Finite-Horizon Stationarity and Nagaev Bounds

    Sep 24, 2026Wei Biao WuNonconvex Stochastic OptimizationStochastic Optimization Convergence

  2. A Unified Lyapunov-IQC Framework for Uniform Stability of Smooth Quadratic First-Order Accelerated Optimizers

    May 8, 2026Don Li, Dacian DaescuLyapunov StabilityAlgorithmic Stability

    Department of Mathematics & Statistics, Portland State University, Portland, OR 97201, USA

  3. Convergence of Large-Stepsize SGD for the Cross-entropy Loss via Stochastic Lyapunov Stability

    Jun 29, 2026Konstantinos Emmanouilidis, Lachlan MacDonald, Salma Tarmoun +1Edge of StabilityStochastic Optimization Convergence

    Department of CIS, University of Pennsylvania · IDEAS, University of Pennsylvania · AMCS, University of Pennsylvania +1