mlsanityData to 14 Oct 2026

Non-convex Stochastic Composite Optimization with Polyak Momentum

arXiv 2403.02967 ↗First posted 5 Mar 2024math.OC cs.LGcited by 9

Not on the map. Why?

Abstract (arXiv)

The stochastic proximal gradient method is a powerful generalization of the widely used stochastic gradient descent (SGD) method and has found numerous applications in Machine Learning. However, it is notoriously known that this method fails to converge in non-convex settings where the stochastic noise is significant (i.e. when only small or bounded batch sizes are used). In this paper, we focus on the stochastic proximal gradient method with Polyak momentum. We prove this method attains an optimal convergence rate for non-convex composite optimization problems, regardless of batch size. Additionally, we rigorously analyze the variance reduction effect of the Polyak momentum in the composite optimization setting and we show the method also converges when the proximal step can only be solved inexactly. Finally, we provide numerical experiments to validate our theoretical results.

Key facts. Cited by 9 papers in the map (data through 14 Oct 2026). Builds on 29 of its 46 references that are in the map. First posted on arXiv 5 Mar 2024.

Builds on

29 of 46 references.

References · 29

  1. 2015Deep Residual Learning for Image Recognitioncited by 52,096
  2. 2016On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minimacited by 1,352
  3. 2020The Future of Digital Health with Federated Learningcited by 633
  4. 2019Lower Bounds for Non-Convex Stochastic Optimizationcited by 307
  5. 2019Momentum-Based Variance Reduction in Non-Convex SGDcited by 306
  6. 2018SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimatorcited by 287
  7. 2019On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimizationcited by 180
  8. 2018Stochastic model-based minimization of weakly convex functionscited by 159
  9. 2013Mini-batch Stochastic Approximation Methods for Nonconvex Stochastic Composite Optimizationcited by 143
  10. 2020Momentum Improves Normalized SGDcited by 139
  11. 2022ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!cited by 134
  12. 2020An Improved Analysis of Stochastic Gradient Descent with Momentumcited by 119
17 more
  1. 2013Accelerated Gradient Methods for Nonconvex Nonlinear and Stochastic Programmingcited by 75
  2. 2022Federated Reinforcement Learning with Environment Heterogeneitycited by 63
  3. 2018How To Make the Gradients Small Stochastically: Even Faster Convex and Nonconvex SGDcited by 57
  4. 2022Towards understanding how momentum improves generalization in deep learningcited by 40
  5. 2023Momentum Provably Improves Error Feedback!cited by 31
  6. 2020Statistically Preconditioned Accelerated Gradient Method for Distributed Optimizationcited by 30
  7. 2019A Hybrid Stochastic Optimization Framework for Stochastic Composite Nonconvex Optimizationcited by 27
  8. 2022Federated Stochastic Approximation under Markov Noise and Heterogeneity: Applications in Reinforcement Learningcited by 27
  9. 2021SGD: The Role of Implicit Regularization, Batch-size and Multiple-epochscited by 22
  10. 2020Momentum-based variance-reduced proximal stochastic gradient method for composite nonconvex stochastic optimizationcited by 14
  11. 2021On the Last Iterate Convergence of Momentum Methodscited by 10
  12. 2023Two Losses Are Better Than One: Faster Optimization Using a Cheaper Proxycited by 7
  13. 2019Auto-Sizing the Transformer Network: Improving Speed, Efficiency, and Performance for Low-Resource Machine Translationcited by 6
  14. 2024Stochastic Bregman Proximal Gradient Method Revisited: Kernel Conditioning and Painless Variance Reductioncited by 6
  15. 2020Principled Analyses and Design of First-Order Methods with Inexact Proximal Operatorscited by 5
  16. 2020Quickly Finding a Benign Region via Heavy Ball Momentum in Non-Convex Optimizationcited by 5
  17. 2022Efficient Algorithms for A Class of Stochastic Hidden Convex Optimization and Its Applications in Network Revenue Managementcited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 2; 2025: 6; 2026: 1.

9 papers.

  1. 2024Variance-reduced first-order methods for deterministically constrained stochastic nonconvex optimization with strong convergence guaranteescited by 11
  2. 2025Complexity of normalized stochastic first-order methods with momentum under heavy-tailed noisecited by 11
  3. 2025Composite Optimization with Error Feedback: the Dual Averaging Approachcited by 2
  4. 2025Stochastic Difference-of-Convex Optimization with Momentumcited by 2
  5. 2025Non-Convex Federated Optimization under Cost-Aware Client Selectioncited by 1
  6. 2024A stochastic first-order method with multi-extrapolated momentum for highly smooth unconstrained optimizationcited by 0
  7. 2025Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentumcited by 0
  8. 2025Bias-Variance Trade-off for Clipped Stochastic First-Order Methods: From Bounded Variance to Infinite Meancited by 0
  9. 2026MoSSP: A Momentum-Based Single-Loop Stochastic Penalty Method for Nonconvex Constrained DC-Regularized Optimizationcited by 0