Research map / Optimization / Nonconvex and smooth optimization
Non-convex Stochastic Composite Optimization with Polyak Momentum
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.
Shows stochastic proximal gradient with Polyak momentum converges to a stationary point of non-convex composite objectives at the optimal O(eps^-2) rate without mega-batches, and extends to inexact proximal steps.
Problem
Stochastic proximal gradient methods for non-convex composite optimization fail to converge to stationary points beyond the gradient noise variance unless large (error-dependent) batches are used.
Approach
Replace the stochastic gradient in the proximal step with a Polyak momentum-aggregated gradient; the momentum error ||m_k - grad f(x_k)||^2 shrinks at the same rate as the gradient norm, removing the need for mega-batches, with a Lyapunov analysis and a lower bound showing vanilla SPG cannot beat the noise floor.
Key results
- Vanilla stochastic proximal gradient with batch size 1 stalls at error around 22, 672, and 14381 for sigma=5, 25, 125 on a synthetic quadratic, while the momentum method reaches the tolerance in all cases.
- Theory: with momentum, O(eps^-2) stochastic gradient calls suffice without mega-batches (rate O(L Phi0 sigma^2/eps^2 + L Phi0/eps)), matching the optimal non-convex rate.
- On CIFAR-10 with ResNet-18 and l_inf,1 regularization, the momentum method outperforms the vanilla method in training loss and test accuracy and gives smoother test accuracy curves.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add Polyak momentum to stochastic proximal gradient | Stochastic proximal gradient · Synthetic quadratic problem | error | Helped |
| Use batch size 16 instead of batch size 1 | Vanilla stochastic proximal gradient · Synthetic quadratic problem | convergence to tolerance | Mixed |
| Use batch size 64 instead of batch size 1 | Vanilla stochastic proximal gradient · Synthetic quadratic problem | convergence to tolerance | Mixed |
| Add Polyak momentum to stochastic proximal gradient | Stochastic proximal gradient with ℓ∞,1 regularization · CIFAR-10, ResNet-18 | training loss | Helped |
| Add Polyak momentum to stochastic proximal gradient | Stochastic proximal gradient with ℓ∞,1 regularization · CIFAR-10, ResNet-18 | test accuracy | Helped |
| Add Polyak momentum to stochastic proximal gradient | Stochastic proximal gradient with statistical preconditioning · CIFAR-10 | training loss | Helped |
| Add Polyak momentum to stochastic proximal gradient | Stochastic proximal gradient with statistical preconditioning · CIFAR-10 | test accuracy | Helped |
Surprising findings
- The vanilla stochastic proximal gradient method cannot converge beyond the gradient noise variance even on a simple quadratic with a large quadratic regularizer, for any fixed stepsize choice.
Limitations
- Without an initial mini-batch, the convergence bound has an extra O(sigma^4/eps^2 + sigma^2/eps) term.
- Analysis assumes smooth psi in the main text; non-differentiable convex psi handled only in an appendix sketch.
- Experiments are small-scale (synthetic quadratic and a single CIFAR-10 ResNet-18 setup).
Experimental setup
Datasets: CIFAR-10 synthetic quadratic problem
Benchmarks: CIFAR-10
Metrics: squared gradient norm / error on synthetic quadratic training loss test accuracy
Baselines: vanilla stochastic proximal gradient method Prox-SpiderBoost PStorm
Builds on: stochastic proximal gradient method heavy-ball momentum SPIDER STORM
Closest papers by meaning
- Convergence of a Stochastic Gradient Method with Momentum for Non-Smooth Non-Convex Optimization (ICML 2020 · similarity 0.93)
- Minibatch and Momentum Model-based Methods for Stochastic Weakly Convex Optimization (NeurIPS 2021 · similarity 0.90)
- Dynamics of SGD with Stochastic Polyak Stepsizes: Truly Adaptive Variants and Convergence to Exact Solution (NeurIPS 2022 · similarity 0.90)
- Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance (ICLR 2025 · similarity 0.90)
- Proximal Stochastic Methods for Nonsmooth Nonconvex Finite-Sum Optimization (NeurIPS 2016 · similarity 0.89)
- New Perspectives on the Polyak Stepsize: Surrogate Functions and Negative Results (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.