PAGE: A Simple and Optimal Probabilistic Gradient Estimator for Nonconvex Optimization

Abstract

In this paper, we propose a novel stochastic gradient estimator -- ProbAbilistic Gradient Estimator (PAGE) -- for nonconvex optimization. PAGE is easy to implement as it is designed via a small adjustment to vanilla SGD: in each iteration, PAGE uses the vanilla minibatch SGD update with probability $p_t$ or reuses the previous gradient with a small adjustment, at a much lower computational cost, with probability $1-p_t$. We give a simple formula for the optimal choice of $p_t$. Moreover, we prove the first tight lower bound $\Omega (n+\frac{\sqrt{n}}{{\epsilon }^2})$ for nonconvex finite-sum problems, which also leads to a tight lower bound $\Omega (b+\frac{\sqrt{b}}{{\epsilon }^2})$ for nonconvex online problems, where $b:=\min \{\frac{{\sigma }^2}{{\epsilon }^2},n\}$. Then, we show that PAGE obtains the optimal convergence results $O(n+\frac{\sqrt{n}}{{\epsilon }^2})$ (finite-sum) and $O(b+\frac{\sqrt{b}}{{\epsilon }^2})$ (online) matching our lower bounds for both nonconvex finite-sum and online problems. Besides, we also show that for nonconvex functions satisfying the Polyak-Łojasiewicz (PL) condition, PAGE can automatically switch to a faster linear convergence rate $O(\cdot \log \frac{1}{\epsilon })$. Finally, we conduct several deep learning experiments (e.g., LeNet, VGG, ResNet) on real datasets in PyTorch showing that PAGE not only converges much faster than SGD in training but also achieves the higher test accuracy, validating the optimal theoretical results and confirming the practical superiority of PAGE.

25 pages; accepted by ICML 2021 (long talk)

References in corpus (11)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives
• EF21: A New, Simpler, Theoretically Better, and Practically Faster Error Feedback
• A Unified Analysis of Stochastic Gradient Methods for Nonconvex Federated Optimization
• MARINA: Faster Non-Convex Distributed Learning with Compression
• ZeroSARAH: Efficient Nonconvex Finite-Sum Optimization with Zero Full Gradient Computation
• L-SVRG and L-Katyusha with Arbitrary Sampling
• Stabilized SVRG: Simple Variance Reduction for Nonconvex Optimization
• Lower Bounds for Smooth Nonconvex Finite-Sum Optimization
• ANITA: An Optimal Loopless Accelerated Variance-Reduced Gradient Method

Cited by in corpus (10)

• EF21: A New, Simpler, Theoretically Better, and Practically Faster Error Feedback
• MARINA: Faster Non-Convex Distributed Learning with Compression
• The Complexity of Nonconvex-Strongly-Concave Minimax Optimization
• Variance Reduction on General Adaptive Stochastic Mirror Descent
• Permutation Compressors for Provably Faster Distributed Nonconvex Optimization
• An Expectation-Maximization Perspective on Federated Learning
• Random-reshuffled SARAH does not need a full gradient computations
• Faster Perturbed Stochastic Gradient Methods for Finding Local Minima
• Accelerated Stochastic ExtraGradient: Mixing Hessian and Gradient Similarity to Reduce Communication in Distributed and Federated Learning
• Geom-SPIDER-EM: Faster Variance Reduced Stochastic Expectation Maximization for Nonconvex Finite-Sum Optimization