# Non-convex Stochastic Composite Optimization with Polyak Momentum

arXiv: https://arxiv.org/abs/2403.02967 · HTML: https://mlsanity.com/paper/2403.02967/

## 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 (most cited first, 29 of 29)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [The Future of Digital Health with Federated Learning](https://mlsanity.com/paper/2003.08119/) (2020)
- [Lower Bounds for Non-Convex Stochastic Optimization](https://mlsanity.com/paper/1912.02365/) (2019)
- [Momentum-Based Variance Reduction in Non-Convex SGD](https://mlsanity.com/paper/1905.10018/) (2019)
- [SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimator](https://mlsanity.com/paper/1807.01695/) (2018)
- [On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization](https://mlsanity.com/paper/1905.03817/) (2019)
- [Stochastic model-based minimization of weakly convex functions](https://mlsanity.com/paper/1803.06523/) (2018)
- [Mini-batch Stochastic Approximation Methods for Nonconvex Stochastic Composite Optimization](https://mlsanity.com/paper/1308.6594/) (2013)
- [Momentum Improves Normalized SGD](https://mlsanity.com/paper/2002.03305/) (2020)
- [ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!](https://mlsanity.com/paper/2202.09357/) (2022)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [Accelerated Gradient Methods for Nonconvex Nonlinear and Stochastic Programming](https://mlsanity.com/paper/1310.3787/) (2013)
- [Federated Reinforcement Learning with Environment Heterogeneity](https://mlsanity.com/paper/2204.02634/) (2022)
- [How To Make the Gradients Small Stochastically: Even Faster Convex and Nonconvex SGD](https://mlsanity.com/paper/1801.02982/) (2018)
- [Towards understanding how momentum improves generalization in deep learning](https://mlsanity.com/paper/2207.05931/) (2022)
- [Momentum Provably Improves Error Feedback!](https://mlsanity.com/paper/2305.15155/) (2023)
- [Statistically Preconditioned Accelerated Gradient Method for Distributed Optimization](https://mlsanity.com/paper/2002.10726/) (2020)
- [A Hybrid Stochastic Optimization Framework for Stochastic Composite Nonconvex Optimization](https://mlsanity.com/paper/1907.03793/) (2019)
- [Federated Stochastic Approximation under Markov Noise and Heterogeneity: Applications in Reinforcement Learning](https://mlsanity.com/paper/2206.10185/) (2022)
- [SGD: The Role of Implicit Regularization, Batch-size and Multiple-epochs](https://mlsanity.com/paper/2107.05074/) (2021)
- [Momentum-based variance-reduced proximal stochastic gradient method for composite nonconvex stochastic optimization](https://mlsanity.com/paper/2006.00425/) (2020)
- [On the Last Iterate Convergence of Momentum Methods](https://mlsanity.com/paper/2102.07002/) (2021)
- [Two Losses Are Better Than One: Faster Optimization Using a Cheaper Proxy](https://mlsanity.com/paper/2302.03542/) (2023)
- [Auto-Sizing the Transformer Network: Improving Speed, Efficiency, and Performance for Low-Resource Machine Translation](https://mlsanity.com/paper/1910.06717/) (2019)
- [Stochastic Bregman Proximal Gradient Method Revisited: Kernel Conditioning and Painless Variance Reduction](https://mlsanity.com/paper/2401.03155/) (2024)
- [Principled Analyses and Design of First-Order Methods with Inexact Proximal Operators](https://mlsanity.com/paper/2006.06041/) (2020)
- [Quickly Finding a Benign Region via Heavy Ball Momentum in Non-Convex Optimization](https://mlsanity.com/paper/2010.01449/) (2020)
- [Efficient Algorithms for A Class of Stochastic Hidden Convex Optimization and Its Applications in Network Revenue Management](https://mlsanity.com/paper/2205.01774/) (2022)

## Built on by (most cited first, 9 of 9)

- [Variance-reduced first-order methods for deterministically constrained stochastic nonconvex optimization with strong convergence guarantees](https://mlsanity.com/paper/2409.09906/) (2024)
- [Complexity of normalized stochastic first-order methods with momentum under heavy-tailed noise](https://mlsanity.com/paper/2506.11214/) (2025)
- [Composite Optimization with Error Feedback: the Dual Averaging Approach](https://mlsanity.com/paper/2510.03507/) (2025)
- [Stochastic Difference-of-Convex Optimization with Momentum](https://mlsanity.com/paper/2510.17503/) (2025)
- [Non-Convex Federated Optimization under Cost-Aware Client Selection](https://mlsanity.com/paper/2512.05327/) (2025)
- [A stochastic first-order method with multi-extrapolated momentum for highly smooth unconstrained optimization](https://mlsanity.com/paper/2412.14488/) (2024)
- [Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum](https://mlsanity.com/paper/2504.12742/) (2025)
- [Bias-Variance Trade-off for Clipped Stochastic First-Order Methods: From Bounded Variance to Infinite Mean](https://mlsanity.com/paper/2512.14686/) (2025)
- [MoSSP: A Momentum-Based Single-Loop Stochastic Penalty Method for Nonconvex Constrained DC-Regularized Optimization](https://mlsanity.com/paper/2605.29635/) (2026)

## 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.