Momentum-Based Variance Reduction in Non-Convex SGD

TL;DR

This paper introduces STORM, a momentum-based variance reduction algorithm for non-convex stochastic gradient descent that eliminates the need for batch sizes and hyperparameter tuning, achieving optimal convergence rates.

Contribution

The paper proposes STORM, a novel variance reduction method using momentum that simplifies implementation and removes the need for batch size tuning in non-convex optimization.

Findings

Achieves optimal convergence rate without batch sizes or variance knowledge.

Uses adaptive learning rates for simpler implementation.

Matches the best known theoretical convergence bounds.

Abstract

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