Stochastic Heavy Ball with Polyak Step Size and Armijo Line Search: A General Convergence Analysis

Jiawei Zhang; Qitan Shi; Yuantao Gu

cs.LGSep 29, 2026

Stochastic Heavy Ball with Polyak Step Size and Armijo Line Search: A General Convergence Analysis

Authors: Jiawei Zhang, Qitan Shi, Yuantao Gu

Organizations: Tsinghua University

Abstract

Polyak step size (PS) and Armijo line search (ALS) have received increasing attention in stochastic optimization, with encouraging empirical performance and theoretical guarantees. However, their convergence theory for stochastic heavy ball (SHB) methods remains limited. In this work, we develop a unified convergence analysis for SHB equipped with PS and ALS. To this end, we introduce a modified Armijo rule that closely parallels the Polyak step size, together with a decoupling analysis that isolates the historical dependence induced by momentum. For SHB with standard PS and ALS, we establish expected convergence for strongly convex, convex, and non-convex objectives without interpolation or restrictive conditions on the momentum parameter. Under interpolation or strong growth, we further strengthen the results to almost sure rates and last-iterate convergence. Moreover, for general settings beyond interpolation, we prove almost sure convergence to the exact optimum or to stationarity for SHB with diminishing variants of PS and ALS. These results provide a more comprehensive theoretical view of Polyak step size and Armijo line search for stochastic heavy ball methods.

Figures & tables

Objective	Noise condition	Guarantee	Theorem
Standard PS / ALS: convergence in expectation
Strongly convex	$\sigma^{2}<\infty$	$\mathbb{E}\\|{\bm{\theta}}_{k}-{\bm{\theta}}^{\star}\\|^{2}=\mathcal{O}(r^{k}+\sigma^{2})$	1
Convex	$\sigma^{2}<\infty$	$\mathbb{E}[F(\bar{\bm{\theta}}_{k})-F({\bm{\theta}}^{\star})]=\mathcal{O}(k^{-1}+\sigma^{2})$	2
Non-convex	Weak growth	$\min_{0\leq m<k}\mathbb{E}\\|\nabla F({\bm{\theta}}_{m})\\|^{2}=\mathcal{O}(k^{-1}+\delta)$	3
Standard PS / ALS: almost sure convergence
Convex	Interpolation	$F(\bar{\bm{\theta}}_{k})-F({\bm{\theta}}^{\star})=\mathcal{O}(k^{-1}),\qquad\text{a.s.}$	4

Table 1: Summary of the main results in this paper.

\bar{\bm{\theta}}_{k}

and

\bar{\bm{\theta}}_{k}^{\rm w}

denote the uniform and weighted averages of the iterates, respectively.

S_{k}

denotes the cumulative upper bound of the effective step sizes. Under interpolation or strong growth, we further establish last-iterate convergence. For the diminishing variants, the corresponding convergence rates can be sharpened from big-

\mathcal{O}

to little-

o

bounds.

Figure 1: Convergence of SHB with PS and ALS on convex (top) and non-convex (bottom) objectives. Panels (a)–(b) and (e)–(f) show the standard, non-diminishing variants, while (c)–(d) and (g)–(h) additionally include the diminishing variants. For the standard step sizes, we report the empirical counterparts of the expected quantities in Theorems 2 and 3 : the objective gap at the uniformly averaged iterate for convex objectives, and the running minimum of the squared full-gradient norm for non-convex objectives. For the diminishing step sizes, which admit almost-sure convergence guarantees, we report the pointwise maximum as a proxy across runs of the corresponding pathwise quantities: the weighted-average objective gap for convex objectives and the best-iterate squared gradient norm for non-convex objectives, following Theorems 6 and 7 , respectively.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Full logistic training loss on eight LIBSVM datasets.

Figure 3: ResNet-18 training on CIFAR-10 (top) and CIFAR-100 (bottom). The first two columns compare the PS methods, and the last two compare the LS methods. We report training cross-entropy and test accuracy versus epochs.

Stochastic Optimization

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

See all 209 papers

Non-Acceleration of the Heavy-Ball Method in 1D
Oct 18, 2026Konstantinos Fotopoulos, Michael Helcig, Karl Deck +1Momentum Methods Accelerated Gradient Methods
Department of CS, ETH Zürich, Switzerland
Adaptive Bregman Proximal Stochastic Gradient with a Stabilized Barzilai--Borwein Step Size
Aug 12, 2026Chenhan Jin, Shengze Xu, Binghui Xie +4Nonconvex Stochastic Optimization Adaptive Gradient Methods
1The Chinese University of Hong Kong · University of Utah · 3Beijing Normal-Hong Kong Baptist University +1
In-Expectation Convergence of Stochastic Gradient Methods under Heavy-Tailed Noise
May 30, 2026Zijian LiuStochastic Optimization Nonconvex Stochastic Optimization
Stern School of Business, New York University

Oct 18, 2026J/K move · Enter open · S save

Non-Acceleration of the Heavy-Ball Method in 1D

Konstantinos Fotopoulos, Michael Helcig, Karl Deck, Jannis Alsbach

Department of CS, ETH Zürich, Switzerland

Momentum Methods Accelerated Gradient Methods

Open paper AskPDF arXiv

Despite having an asymptotically optimal rate of convergence on convex quadratics, the acceleration of Polyak's heavy-ball method and a characterization of its dynamics within more general settings have remained long standing open problems. Recently, the work of Goujaud et al. [Math. Program. (2025)] made progress regarding acceleration in the context of smooth, strongly convex minimization: for a small enough inverse condition number, either the method does not accelerate on quadratics, or there exists a smooth strongly convex function such that heavy-ball cycles over the roots of unity. The constructed functions have two variables, thus leaving the question of acceleration over univariate functions open. In this work, we use a Fourier-based analysis to i) provide a proof for the non-acceleration of the heavy-ball method over the class of univariate smooth strongly convex functions. We also ii) provide a strict enlargement of the previous root-of-unity cycling region, with insights regarding the "fractional" sets of Goujaud et al. [arXiv:2502.19916], and iii) give a higher dimensional cycling instance which is not covered by root-of-unity cycling.

Abstract

Stochastic Optimization

Explore similar work

Non-Acceleration of the Heavy-Ball Method in 1D

Adaptive Bregman Proximal Stochastic Gradient with a Stabilized Barzilai--Borwein Step Size

In-Expectation Convergence of Stochastic Gradient Methods under Heavy-Tailed Noise