Papers › (Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum

(Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum

12 Jan 2024arXiv:2401.06738archive 2025-07-28

Anh Dang, Reza Babanezhad, Sharan Vaswani

Stochastic heavy ball momentum (SHB) is commonly used to train machine learning models, and often provides empirical improvements over stochastic gradient descent. By primarily focusing on strongly-convex quadratics, we aim to better understand the theoretical advantage of SHB and subsequently improve the method. For strongly-convex quadratics, Kidambi et al. (2018) show that SHB (with a mini-batch of size $1$) cannot attain accelerated convergence, and hence has no theoretical benefit over SGD. They conjecture that the practical gain of SHB is a by-product of using larger mini-batches. We first substantiate this claim by showing that SHB can attain an accelerated rate when the mini-batch size is larger than a threshold b^* that depends on the condition number κ. Specifically, we prove that with the same step-size and momentum parameters as in the deterministic setting, SHB with a sufficiently large mini-batch size results in an O(exp(-T/(√(κ))) + σ) convergence when measuring the distance to the optimal solution in the ℓ₂ norm, where T is the number of iterations and σ² is the variance in the stochastic gradients. We prove a lower-bound which demonstrates that a κ dependence in b^* is necessary. To ensure convergence to the minimizer, we design a noise-adaptive multi-stage algorithm that results in an O(exp(-T/(√(κ))) + σ/(√(T))) rate when measuring the distance to the optimal solution in the ℓ₂ norm. We also consider the general smooth, strongly-convex setting and propose the first noise-adaptive SHB variant that converges to the minimizer at an O(exp(-T/κ) + σ²/T) rate when measuring the distance to the optimal solution in the squared ℓ₂ norm. We empirically demonstrate the effectiveness of the proposed algorithms.

PaperPDFCode

Code

anh-dang/accelerated_noise_adaptive_shb officialmentioned in paper report

Repository list and official/mentioned flags are the archive's, frozen 2025-07-28. Reachability, where shown, is from one Syntology probe window (2026-09-16 to 2026-09-18); repositories not probed show nothing. GitHub stars are not tracked.

Code Syntology ran Syntology

Not run by Syntology. Nothing on this page verifies that the listed code works.

Results from the paper archive 2025-07-28

No leaderboard rows for this paper in the archive.

Methods

SGD

Report a problem or propose a change · a person checks every report against the paper or source before anything changes; decisions are listed on /corrections