Momentum-Based Variance Reduction in Non-Convex SGD

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 $\mathbf{x}$ with $\mathbb{E}[\Vert \nabla F(\mathbf{x})\Vert ]\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$.

Added Ack

References in corpus (2)

• On the Convergence of Adam and Beyond
• Stochastic Recursive Gradient Algorithm for Nonconvex Optimization

Cited by in corpus (8)

• A Unified Convergence Analysis for Shuffling-Type Gradient Methods
• Distributed Momentum for Byzantine-resilient Learning
• Single-Timescale Stochastic Nonconvex-Concave Optimization for Smooth Nonlinear TD Learning
• Stochastic Recursive Momentum Method for Non-Convex Compositional Optimization
• Almost Tune-Free Variance Reduction
• Enhanced Bilevel Optimization via Bregman Distance
• CDMA: A Practical Cross-Device Federated Learning Algorithm for General Minimax Problems
• On the Convergence Rate of Off-Policy Policy Optimization Methods with Density-Ratio Correction