General convergence rates of stochastic heavy ball with Polyak step size and Armijo line search

Polyak step size (PS) and Armijo line search (ALS) have received increasing attention in stochastic optimization in recent years, with encouraging empirical performance and theoretical guarantees. However, their convergence properties for stochastic heavy ball (SHB) methods remain relatively limited. In this work, we bridge this gap by establishing general convergence guarantees for SHB equipped with PS and ALS. To enable a unified analysis, we first introduce a slightly modified Armijo rule that more closely parallels the Polyak step size. Building on a decoupling technique for the SHB iterates, we establish convergence rates for both step-size rules under substantially more general assumptions. Specifically, without assuming interpolation or imposing restrictive conditions on the momentum parameter, we establish convergence guarantees in expectation for SHB with PS or ALS on both convex and non-convex objectives. Under interpolation, we further provide almost sure last-iterate convergence for both settings. These results complement and extend the existing theoretical understanding of adaptive step-size rules for stochastic heavy ball methods.