Papers
Topics
Authors
Recent
Gemini 2.5 Flash
Gemini 2.5 Flash
120 tokens/sec
GPT-4o
7 tokens/sec
Gemini 2.5 Pro Pro
46 tokens/sec
o3 Pro
4 tokens/sec
GPT-4.1 Pro
38 tokens/sec
DeepSeek R1 via Azure Pro
28 tokens/sec
2000 character limit reached

Non-convex Stochastic Composite Optimization with Polyak Momentum (2403.02967v4)

Published 5 Mar 2024 in math.OC and cs.LG

Abstract: 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.

Definition Search Book Streamline Icon: https://streamlinehq.com
References (45)
  1. Zeyuan Allen-Zhu. How to make the gradients small stochastically: Even faster convex and nonconvex sgd. Advances in Neural Information Processing Systems, 31, 2018.
  2. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1-2):165–214, 2023.
  3. Principled analyses and design of first-order methods with inexact proximal operators. Mathematical Programming, 201(1):185–230, 2023.
  4. Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning, 8(3-4):231–357, 2015.
  5. Efficient algorithms for minimizing compositions of convex functions and random functions and its applications in network revenue management. arXiv preprint arXiv:2205.01774, 2022.
  6. Proximal splitting methods in signal processing, 2010.
  7. Momentum improves normalized sgd. In International conference on machine learning, pages 2260–2268. PMLR, 2020.
  8. Momentum-based variance reduction in non-convex sgd. Advances in neural information processing systems, 32, 2019.
  9. Stochastic model-based minimization of weakly convex functions. SIAM Journal on Optimization, 29(1):207–239, 2019.
  10. Finite-time analysis of distributed td (0) with linear function approximation on multi-agent reinforcement learning. In International Conference on Machine Learning, pages 1626–1635. PMLR, 2019.
  11. Spider: Near-optimal non-convex optimization via stochastic path-integrated differential estimator. Advances in neural information processing systems, 31, 2018.
  12. Momentum provably improves error feedback! In Advances in Neural Information Processing Systems, 2023.
  13. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization i: A generic algorithmic framework. SIAM Journal on Optimization, 22(4):1469–1492, 2012.
  14. Accelerated gradient methods for nonconvex nonlinear and stochastic programming, 2013.
  15. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization, ii: shrinking procedures and optimal algorithms. SIAM Journal on Optimization, 23(4):2061–2089, 2013.
  16. Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization, 2013.
  17. Statistically preconditioned accelerated gradient method for distributed optimization. In International conference on machine learning, pages 4203–4227. PMLR, 2020.
  18. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, pages 9965–10040. PMLR, 2022.
  19. Federated reinforcement learning with environment heterogeneity. In International Conference on Artificial Intelligence and Statistics, pages 18–37. PMLR, 2022.
  20. On large-batch training for deep learning: Generalization gap and sharp minima. In Proceedings of International Conference on Learning Representations, 2017.
  21. Federated reinforcement learning: Linear speedup under markovian sampling. In International Conference on Machine Learning, pages 10997–11057. PMLR, 2022.
  22. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
  23. Cifar-10 dataset, 2014.
  24. Efficient backprop. Neural Networks: Tricks of the Trade, Springer-Verlag, 2012.
  25. On the last iterate convergence of momentum methods. In International Conference on Algorithmic Learning Theory, pages 699–717. PMLR, 2022.
  26. Sparse convolutional neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2015.
  27. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems, 33:18261–18271, 2020.
  28. D. Russell Luke. Proximal Methods for Image Processing, pages 165–202. Springer International Publishing, Cham, 2020.
  29. Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally! In International Conference on Machine Learning, pages 15750–15769. PMLR, 2022.
  30. Problem complexity and method efficiency in optimization. Wiley-Interscience, 1983.
  31. Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013.
  32. Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathematics and mathematical physics, 4(5):1–17, 1964.
  33. Making gradient descent optimal for strongly convex stochastic optimization. In Proceedings of the 29th International Coference on International Conference on Machine Learning, pages 1571–1578, 2012.
  34. The future of digital health with federated learning. NPJ digital medicine, 3(1):119, 2020.
  35. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In Conference on Learning Theory, pages 3935–3971. PMLR, 2021.
  36. Sgd: The role of implicit regularization, batch-size and multiple-epochs. Advances In Neural Information Processing Systems, 34:27422–27433, 2021.
  37. A hybrid stochastic optimization framework for composite nonconvex optimization. Mathematical Programming, 191(2):1005–1071, 2022.
  38. Quickly finding a benign region via heavy ball momentum in non-convex optimization. arXiv preprint arXiv:2010.01449, 2020.
  39. Spiderboost and momentum: Faster variance reduction algorithms. Advances in Neural Information Processing Systems, 32, 2019.
  40. The general inefficiency of batch training for gradient descent learning. Neural Networks, 2003.
  41. Two losses are better than one: Faster optimization using a cheaper proxy, 2023.
  42. Momentum-based variance-reduced proximal stochastic gradient method for composite nonconvex stochastic optimization, 2022.
  43. On the linear speedup analysis of communication efficient momentum sgd for distributed non-convex optimization. In International Conference on Machine Learning, pages 7184–7193. PMLR, 2019.
  44. Complexity of finding stationary points of nonconvex nonsmooth functions. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 11173–11182. PMLR, 13–18 Jul 2020.
  45. Junyu Zhang. Stochastic bergman proximal gradient method revisited: Kernel conditioning and painless variance reduction. arXiv preprint arXiv:2401.03155, 2024.
Citations (2)
View on Semantic Scholar

Summary

We haven't generated a summary for this paper yet.

Ai Sparkles Streamline Icon: https://streamlinehq.com Summarize Now
X Twitter Logo Streamline Icon: https://streamlinehq.com

Tweets