Momentum-based variance-reduced proximal stochastic gradient method for composite nonconvex stochastic optimization
This work addresses the need for efficient stochastic optimization in online and large-scale machine learning by reducing sample requirements, though it is incremental as it builds on existing variance reduction techniques.
The paper tackles the problem of slow convergence in stochastic gradient methods for nonconvex nonsmooth optimization by proposing PStorm, a momentum-based variance-reduced method that achieves optimal complexity O(ε^{-3}) using only one or O(1) samples per update, enabling real-time online learning and better generalization in neural network training.
Stochastic gradient methods (SGMs) have been extensively used for solving stochastic problems or large-scale machine learning problems. Recent works employ various techniques to improve the convergence rate of SGMs for both convex and nonconvex cases. Most of them require a large number of samples in some or all iterations of the improved SGMs. In this paper, we propose a new SGM, named PStorm, for solving nonconvex nonsmooth stochastic problems. With a momentum-based variance reduction technique, PStorm can achieve the optimal complexity result $O(\varepsilon^{-3})$ to produce a stochastic $\varepsilon$-stationary solution, if a mean-squared smoothness condition holds. Different from existing optimal methods, PStorm can achieve the ${O}(\varepsilon^{-3})$ result by using only one or $O(1)$ samples in every update. With this property, PStorm can be applied to online learning problems that favor real-time decisions based on one or $O(1)$ new observations. In addition, for large-scale machine learning problems, PStorm can generalize better by small-batch training than other optimal methods that require large-batch training and the vanilla SGM, as we demonstrate on training a sparse fully-connected neural network and a sparse convolutional neural network.