Momentum-Based Variance Reduction in Non-Convex SGD

By Ashok Cutkosky, Francesco Orabona. 2019

553 citations recorded by the source corpus.

Variance reduction has emerged in recent years as a strong competitor to stochastic gradient descent in non-convex problems, providing the first algorithms to improve upon the converge rate of stochastic gradient descent for finding first-order critical points. However, variance reduction techniques typically require carefully tuned learning rates and willingness to use excessively large "mega-batches" in order to achieve their improved results. We present a new algorithm, STORM, that does not require any batches and makes use of adaptive learning rates, enabling simpler implementation and less hyperparameter tuning. Our technique for removing the batches uses a variant of momentum to achieve variance reduction in non-convex optimization. On smooth losses $F$, STORM finds a point $\boldsymbol{x}$ with $\mathbb{E}[\|\nabla F(\boldsymbol{x})\|]\le O(1/\sqrt{T}+\sigma^{1/3}/T^{1/3})$ in $T$ iterations with $\sigma^2$ variance in the gradients, matching the optimal rate but without requiring knowledge of $\sigma$.

Subfields: Artificial Intelligence, Computational Mechanics, Statistics and Probability

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mechanics

In Statistics and Probability

Citations in the corpus

Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning (2020, methodology): “• We prove that incorporating momentum based variance reduction [4, 33] into each local client update reduces client drift and obtains the optimal statistical rates.”
Projection-Free Adaptive Gradients for Large-Scale Optimization (2020, methodology): “…variance reduction based on the works of Johnson and Zhang (2013) and Fang et al. (2018) respectively, ORGFW uses a form of momentum inspired by Cutkosky and Orabona (2019), and CSFW takes advantage of the additive separability of the objective function in the data samples, when applicable,…”
Lower bounds for non-convex stochastic optimization (2019, methodology): “Recent work [21, 49] offers further refinements of these algorithms that also obtain the −3 rate.”
Hybrid Variance-Reduced SGD Algorithms For Minimax Problems with Nonconvex-Linear Function (2020, background): “When B(1) t ≡ B(2) t and B̂(1) t ≡ B̂(2) t , our estimators reduce the STORM estimators studied in [7] as a special case.”
Momentum Improves Normalized SGD (2020, methodology): “Nevertheless, more recent work in this area has produced algorithms whose updates are very similar to the (unnormalized) momentum update (Cutkosky & Orabona, 2019; Tran-Dinh et al., 2019), suggesting that momentum may operate by reducing variance in SGD. Finally, another technique that also…”
AUC Maximization in the Era of Big Data and AI: A Survey (2022, methodology): “They use the STORM technique [31] to compute ut+1, vt+1, i.”
Learning from History for Byzantine Robust Optimization (2020, methodology): “We now describe how to modify the momentum method with a small correction term to improve its convergence rate (Cutkosky & Orabona, 2019).”
Conservative Stochastic Optimization With Expectation Constraints (2020, methodology): “Lemma 3 Under Assumptions 1-5, the iterates generated by Algorithm 2 satisfy where The proof of Lemma 3 is given in Appendix B and follows along similar lines as Lemma 1 of [35].”