Research map / Optimization / Optimizers
Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes three Polyak-type adaptive step-sizes (MomSPSmax, MomDecSPS, MomAdaSPS) for stochastic heavy ball momentum, derived via the iterate moving average viewpoint, with convergence guarantees for convex smooth problems with and without interpolation.
Problem
Stochastic heavy ball (SGD with momentum) requires costly tuning of step-size and momentum; naive stochastic Polyak step-size combined with momentum can diverge, and prior adaptive momentum methods lacked exact convergence without interpolation.
Approach
Use the iterate moving average (IMA) reformulation of SHB to derive momentum-corrected Polyak step-sizes (multiplying by 1-beta), yielding a neighborhood-convergent variant and two decreasing variants that converge to the exact minimizer for any beta in [0,1) without knowing problem parameters.
Key results
- MomSPSmax with momentum outperforms SPSmax without momentum on CIFAR-10 and CIFAR-100 with ResNet34, and has competitive generalization against other optimizers.
- Naive SPSmax with momentum diverges for larger beta, whereas the proposed MomSPSmax converges for all tested beta on logistic regression.
- Proposed decreasing variants MomDecSPS and MomAdaSPS are proven to converge to the exact minimizer for any beta in [0,1) (O(1/sqrt(T)) and O(1/T + sigma/sqrt(T)) respectively).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SPSmax directly with heavy-ball momentum instead of scaling the step-size by 1−β | SHB · Synthetic logistic regression | convergence stability | Hurt |
| Scale SPSmax by the correcting factor 1−β (MomSPSmax) | SHB · Synthetic logistic regression | convergence stability | Helped |
| Use MomDecSPS with momentum rather than its no-momentum counterpart | SHB · LibSVM convex classification datasets | training loss and test accuracy | Helped |
| Use MomAdaSPS with momentum rather than its no-momentum counterpart | SHB · LibSVM convex classification datasets | training loss and test accuracy | Helped |
| Use MomSPSmax rather than SPSmax without momentum | SHB · CIFAR-10, ResNet34 | performance | Helped |
| Use MomSPSmax rather than SPSmax without momentum | SHB · CIFAR-100, ResNet34 | performance | Helped |
| Use MomSPSmax rather than vanilla Polyak step-size | Deterministic heavy-ball method · Least squares | relative error | Helped |
| Move the MomDecSPS correcting factor 1−β outside the minimum (ablation) | MomDecSPS · Experiments reported in the paper | practical performance | Hurt |
| Move the MomAdaSPS correcting factor 1−β outside the minimum (ablation) | MomAdaSPS · Experiments reported in the paper | practical performance | Hurt |
| Set momentum β=0 instead of using momentum (ablation) | MomSPSmax · SHB optimization | convergence rate: O(1/T) | Mixed |
| Set momentum β=0 instead of using momentum (ablation) | MomDecSPS · SHB optimization | convergence rate: O(1/√T) | No clear effect |
| Set momentum β=0 instead of using momentum (ablation) | MomAdaSPS · SHB optimization | convergence rate: O(1/T) | No clear effect |
Surprising findings
- Naively combining SPSmax with heavy ball momentum leads to divergence even on simple convex smooth logistic regression problems when beta is large.
- Theory suggests beta=0 is the best choice for the constants C1 and C2 in the neighborhood convergence bound, since they increase with beta.
Limitations
- MomSPSmax converges only to a neighborhood of the solution in the non-interpolated setting.
- Restriction on beta required for the neighborhood theorem; decreasing-step-size theorems assume bounded iterates.
- Constant gamma_b must be set/smoothed by the user, with sensitivity to its choice noted as a problem-dependent tuning issue.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 LIBSVM datasets
Benchmarks: CIFAR-10 CIFAR-100 LIBSVM
Metrics: training loss test accuracy relative error
Baselines: SGD SHB with constant step-size ADAM ALR-SMAG SPSmax AdaGrad-Norm AGM AHB Gradient Descent with Polyak step-size
Builds on: stochastic Polyak step-size Iterate Moving Average AdaSPS DecSPS Polyak heavy ball
Closest papers by meaning
- The Role of Momentum Parameters in the Optimal Convergence of Adaptive Polyak's Heavy-ball Methods (ICLR 2021 · similarity 0.92)
- Generalized Polyak Step Size for First Order Optimization with Momentum (ICML 2023 · similarity 0.91)
- Non-convex Stochastic Composite Optimization with Polyak Momentum (ICML 2024 · similarity 0.90)
- Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise (ICLR 2024 · similarity 0.90)
- Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration (NeurIPS 2026 · similarity 0.88)
- Accelerated Over-Relaxation Heavy-Ball Method: Achieving Global Accelerated Convergence with Broad Generalization (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.