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 (6)

• On the Convergence of Adam and Beyond
• SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimator
• AdaGrad stepsizes: Sharp convergence over nonconvex landscapes
• Stochastic Recursive Gradient Algorithm for Nonconvex Optimization
• Accelerating Stochastic Gradient Descent For Least Squares Regression
• Distributed Stochastic Optimization via Adaptive SGD

Cited by in corpus (9)

• 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
• Sample complexity of variance-reduced policy gradient: weaker assumptions and lower bounds
• 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