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.

arXiv comment25 pages; accepted by ICML 2021 (long talk)

References in corpus (22)

• 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
• SPIDER: Near-Optimal Non-Convex Optimization via Stochastic Path Integrated Differential Estimator
• A Universal Catalyst for First-Order Optimization
• A Simple Proximal Stochastic Gradient Method for Nonsmooth Nonconvex Optimization
• Acceleration for Compressed Gradient Descent in Distributed and Federated Optimization
• Stochastic Nested Variance Reduction for Nonconvex Optimization
• Better Theory for SGD in the Nonconvex World
• EF21: A New, Simpler, Theoretically Better, and Practically Faster Error Feedback
• An optimal randomized incremental gradient method
• A unified variance-reduced accelerated gradient method for convex optimization
• Nonconvex Variance Reduced Optimization with Arbitrary Sampling
• A Unified Analysis of Stochastic Gradient Methods for Nonconvex Federated Optimization
• MARINA: Faster Non-Convex Distributed Learning with Compression
• L-SVRG and L-Katyusha with Arbitrary Sampling
• ZeroSARAH: Efficient Nonconvex Finite-Sum Optimization with Zero Full Gradient Computation
• 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
• A General Analysis Framework of Lower Complexity Bounds for Finite-Sum Optimization
• Improving SAGA via a Probabilistic Interpolation with Gradient Descent

Cited by in corpus (13)

• EF21: A New, Simpler, Theoretically Better, and Practically Faster Error Feedback
• Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers
• MARINA: Faster Non-Convex Distributed Learning with Compression
• A Second look at Exponential and Cosine Step Sizes: Simplicity, Adaptivity, and Performance
• The Complexity of Nonconvex-Strongly-Concave Minimax Optimization
• Sample complexity of variance-reduced policy gradient: weaker assumptions and lower bounds
• Permutation Compressors for Provably Faster Distributed Nonconvex Optimization
• Variance Reduction on General Adaptive Stochastic Mirror Descent
• Random-reshuffled SARAH does not need a full gradient computations
• An Expectation-Maximization Perspective on Federated Learning
• Geom-SPIDER-EM: Faster Variance Reduced Stochastic Expectation Maximization for Nonconvex Finite-Sum Optimization
• Accelerated Stochastic ExtraGradient: Mixing Hessian and Gradient Similarity to Reduce Communication in Distributed and Federated Learning
• Faster Perturbed Stochastic Gradient Methods for Finding Local Minima