Research map / Optimization / Nonconvex and smooth optimization
Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise
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.
Gives a non-asymptotic last-iterate convergence bound for stochastic heavy-ball with step-decay schedules on quadratic objectives under anisotropic gradient noise, showing O~(sqrt(kappa)) bias acceleration over SGD in large-batch settings.
Problem
Whether stochastic heavy-ball momentum with decaying learning rates provably accelerates SGD (roughly sqrt(kappa) in the bias term) on quadratic objectives under anisotropic gradient noise, especially in large-batch settings.
Approach
A matrix-product bound showing that the product of step-decay-dependent 2x2 heavy-ball transition matrices is upper bounded by the power of a single matrix, using sign structure of columns, so SHB retains fast bias decay with large learning rates and variance control with small ones.
Key results
- On ridge regression (a4a), SHB + step decay reaches excess loss 0.13e-2 at M=512 versus 2.44e-2 for SGD + step decay and 2.10e-2 for SGD constant.
- On CIFAR-10 with batch 128x16 for 100 epochs, SHB with step decay improves ResNet18 test accuracy to 89.78% versus 83.82% for SGD.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SHB with step decay instead of SHB with a constant learning rate (ablation) | SHB on a4a ridge regression · a4a, batch size 512 | training loss: 0.86$\pm$0.55 -> 0.13$\pm$0.03 | Helped |
| Use SHB with step decay instead of SHB with a constant learning rate (ablation) | SHB on a4a ridge regression · a4a, batch size 128 | training loss: 0.55$\pm$0.26 -> 0.01$\pm$0.00 | Helped |
| Use SHB with step decay instead of SHB with a constant learning rate (ablation) | SHB on a4a ridge regression · a4a, batch size 32 | training loss: 1.03$\pm$0.35 -> 0.03$\pm$0.02 | Helped |
| Use SHB with step decay instead of SHB with a constant learning rate (ablation) | SHB on a4a ridge regression · a4a, batch size 8 | training loss: 0.97$\pm$0.58 -> 0.06$\pm$0.05 | Mixed |
| Use SHB instead of SGD with step decay | a4a ridge regression · a4a, batch size 512 | training loss: 2.44$\pm$0.45 -> 0.13$\pm$0.03 | Helped |
| Use SHB instead of SGD | CIFAR-10, ResNet18 · batch size 128, 10 epochs | accuracy (%): 81.19$\pm$0.93 -> 85.16$\pm$2.30 | Helped |
| Use SHB instead of SGD | CIFAR-10, DenseNet121 · batch size 128, 10 epochs | accuracy (%): 88.58$\pm$0.23 -> 88.63$\pm$0.27 | No clear effect |
| Use SHB instead of SGD | CIFAR-10, MobileNetV2 · batch size 128, 10 epochs | accuracy (%): 82.90$\pm$0.37 -> 86.23$\pm$0.23 | Helped |
| Use SHB instead of SGD | CIFAR-10, ResNet18 · batch size 128×16, 100 epochs | accuracy (%): 83.82$\pm$0.42 -> 89.78$\pm$0.23 | Helped |
| Use SHB instead of SGD | CIFAR-10, DenseNet121 · batch size 128×16, 100 epochs | accuracy (%): 89.28$\pm$0.23 -> 92.46$\pm$0.15 | Helped |
| Use SHB instead of SGD | CIFAR-10, MobileNetV2 · batch size 128×16, 100 epochs | accuracy (%): 84.37$\pm$0.77 -> 89.57$\pm$0.18 | Helped |
Surprising findings
- SHB with constant learning rate performs no better than SGD at small batch sizes, so the benefit depends on the step-decay schedule and large batch.
Limitations
- Analysis restricted to quadratic objectives.
- Theorem requires specific stage-length, decay-factor and total-iteration conditions.
- Empirical CIFAR-10 experiments use minimal momentum tuning and are not a strict test of the theory.
Experimental setup
Datasets: a4a CIFAR-10
Benchmarks: a4a CIFAR-10
Metrics: training loss excess risk f(w)-f(w*) test accuracy
Baselines: SGD + constant learning rate SGD + step decay SHB + constant learning rate
Builds on: stochastic heavy ball bias-variance decomposition for SGD anisotropic gradient noise assumption
Closest papers by meaning
- Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration (NeurIPS 2026 · similarity 0.91)
- Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance (ICLR 2025 · similarity 0.90)
- Accelerated Over-Relaxation Heavy-Ball Method: Achieving Global Accelerated Convergence with Broad Generalization (ICLR 2025 · similarity 0.90)
- SGD with memory: fundamental properties and stochastic acceleration (ICLR 2025 · similarity 0.89)
- Gradient correlation is a key ingredient to accelerate SGD with momentum (ICLR 2025 · similarity 0.88)
- Nonlinearly Preconditioned Gradient Methods: Momentum and Stochastic Analysis (NeurIPS 2025 · similarity 0.88)
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.