Source-linked AI summary
PAGE: A Simple and Optimal Probabilistic Gradient Estimator for Nonconvex Optimization
Zhize Li, Hongyan Bao, Xiangliang Zhang, Peter Richtárik
TL;DR
The paper asks how to obtain efficient and provably optimal stochastic-gradient methods for nonconvex finite-sum and online optimization. It proposes PAGE, which probabilistically switches between minibatch SGD and a low-cost adjusted reuse of the previous gradient. PAGE matches tight lower bounds in both settings, can achieve faster linear rates under the PL condition, and outperforms SGD in the reported deep-learning experiments.
Problem
The paper addresses gaps in tight lower bounds and simple optimal methods for nonconvex finite-sum and online optimization.
Method
PAGE probabilistically uses a minibatch SGD update or a low-cost adjusted reuse of the previous gradient.
Results
PAGE matches tight lower bounds for finite-sum and online problems, achieves faster linear convergence under the PL condition, and reports faster training with higher test accuracy than SGD.
Takeaways & Limitations
PAGE provides a simple, easy-to-implement estimator with optimal convergence guarantees and reported practical superiority in deep-learning experiments.
Takeaways & Limitations
The convergence target is an ϵ-approximate solution satisfying E[∥∇f(x̂)∥] ≤ ϵ, and some analyses require assumptions such as bounded variance or the PL condition.
Abstract
from arXiv · showhide
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 $Ω(n+\frac{\sqrt{n}}{ε^2})$ for nonconvex finite-sum problems, which also leads to a tight lower bound $Ω(b+\frac{\sqrt{b}}{ε^2})$ for nonconvex online problems, where $b:= \min\{\frac{σ^2}{ε^2}, n\}$. Then, we show that PAGE obtains the optimal convergence results $O(n+\frac{\sqrt{n}}{ε^2})$ (finite-sum) and $O(b+\frac{\sqrt{b}}{ε^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}ε)$. 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.
1. Introduction
The paper studies stochastic-gradient complexity for differentiable, possibly nonconvex finite-sum and online optimization problems. It situates PAGE among efforts to improve the dependence of nonconvex methods on dataset size and accuracy.
- Nonconvex optimization appears in robust regression, matrix recovery, sparse recovery, and supervised learning.
- The problem: The paper considers differentiable, possibly nonconvex objectives with finite-sum and online forms.The online form applies when the number of samples is very large or infinite.
- Gradient complexity: Gradient complexity counts stochastic gradient computations needed to find an ϵ-approximate solution satisfying E[∥∇f(x̂)∥] ≤ ϵ.
- GD has complexity O(n/ϵ^2), whereas bounded-variance SGD has complexity O(1/ϵ^4) in the nonconvex regime.SGD typically uses a constant minibatch instead of GD’s full batch.
- Variance-reduced methods such as SVRG, SAGA, and SARAH were developed to improve dependence on n and/or ϵ, but much prior analysis focused on convex problems.
2. Our Contributions
The paper addresses gaps in optimality guarantees and implementation complexity for nonconvex finite-sum and online optimization. PAGE combines probabilistic switching with tight lower-bound analysis, including faster rates under the PL condition.
- Motivation: Prior optimal nonconvex finite-sum methods lacked tight lower bounds in the online regime and often used complicated double-loop designs.The existing lower bound applied only in the small-data regime n ≤ O(1/ϵ^4).
- PAGE estimator: PAGE uses minibatch SGD with probability p_t, or reuses the previous gradient with a small adjustment and lower cost with probability 1 − p_t.The estimator is designed as a small adjustment to vanilla minibatch SGD.
- Optimality: PAGE achieves convergence results matching tight lower bounds for both nonconvex finite-sum and online problems.
- PL condition: PAGE can automatically switch to a faster linear convergence rate under the PL condition, despite global nonconvexity.
- Experiments: Deep-learning experiments with LeNet, VGG, and ResNet report faster training convergence and higher test accuracy than SGD.The experiments use real datasets in PyTorch.
- Parameter choice: The choice p_t ≡ b′/(b+b′) is sufficient for PAGE to obtain optimal convergence rates.
3. Notation and Assumptions
The paper establishes notation and standard assumptions for analyzing stochastic-gradient methods in nonconvex optimization. It also introduces the PL condition as the basis for faster linear convergence results.
- Notation: The notation defines [n], vector and matrix norms, inner products, asymptotic notation, and the initial objective gap Δ_0.
- Assumptions: Convergence analysis uses standard assumptions appropriate to the setting, including bounded stochastic-gradient variance and average L-smoothness.
- PL condition: Under the Polyak-Łojasiewicz condition, the paper proves faster linear convergence rates for globally nonconvex functions.
4. General Convergence Results
PAGE is analyzed for nonconvex finite-sum and online problems, with convergence guarantees, matching lower bounds, and reductions to GD or SGD under specific parameter choices.
- Finite-sum problems: PAGE’s finite-sum convergence theorem targets an ϵ-approximate solution under the stated smoothness assumptions and minibatch parameters.
- Finite-sum problems: PAGE recovers GD when p_t ≡ 1 and b = n, including the corresponding finite-sum convergence guarantee.
- Finite-sum problems: PAGE’s optimal finite-sum parameter setting uses a fixed probability p_t ≡ b′/(b+b′), followed by a matching lower bound.
- Finite-sum problems: The finite-sum lower bound is Ω(n + Δ_0L√n/ϵ^2) stochastic gradient computations for linear-span first-order algorithms.
- Online problems: In the online setting, PAGE is analyzed under bounded variance and obtains an optimal parameterized convergence result with fixed p_t.
- Algorithmic reductions: PAGE reduces to vanilla minibatch SGD when p_t ≡ 1, and its expected per-iteration gradient cost is pb + (1−p)b′.
- Online problems: The online lower bound is Ω(b + Δ_0L√b/ϵ^2), with b = min{σ^2/ϵ^2, n}, and matches PAGE’s online convergence result.
5. Better Convergence under PL Condition
Under the PL condition, PAGE achieves faster linear convergence and can switch to that rate in regions where an otherwise globally nonconvex objective satisfies PL locally.
- PL convergence: The PL condition enables faster linear convergence for PAGE than the sublinear nonconvex rate.
- Finite-sum PL results: PAGE’s finite-sum PL theorem provides an ϵ-solution guarantee with fixed probability p_t ≡ p and secondary minibatch size b′ < b.
- Finite-sum PL results: The corresponding finite-sum gradient complexity is #grad = b + T(pb + (1−p)b′).
- Local PL behavior: With the stated parameter choice, PAGE can automatically switch to a faster linear rate O(· log 1/ϵ) where the PL condition holds locally.
- Online PL results: The online PL theorem likewise uses a minibatch, a secondary minibatch, and fixed probability to establish an ϵ-solution guarantee.
6. Experiments
Experiments compare PAGE with vanilla minibatch SGD across LeNet, VGG, and ResNet models on MNIST and CIFAR-10. PAGE converges faster in training and achieves higher test accuracy, with a larger gap at larger minibatch sizes.
- Experimental setup: PAGE and vanilla SGD are compared on LeNet, VGG, and ResNet models using MNIST and CIFAR-10 datasets in PyTorch.Experiments run on NVIDIA Tesla V100 GPUs.
- Results: PAGE converges much faster than SGD in training and achieves higher test accuracy across the reported experiments.The experiments use training loss and test accuracy plotted against gradient computations.
- Results: The performance gap between PAGE and SGD becomes larger as minibatch size increases.The comparison covers minibatch sizes b = 64, 256, and 512 across MNIST and CIFAR-10 experiments.
- Additional experiments: Additional experiments show faster PAGE convergence for both training and test loss, plus higher test accuracy and faster training-accuracy convergence.These comparisons use ResNet18 on MNIST and ResNet34 on CIFAR-10.
7. Conclusion
The paper concludes that PAGE provides simple, optimal convergence guarantees for nonconvex finite-sum and online optimization. It also supports faster PL-condition convergence and practical deep-learning performance, with an analysis designed for adoption and generalization.
- Contributions: PAGE achieves optimal convergence results matching tight lower bounds for nonconvex finite-sum and online problems.The conclusion identifies PAGE as simple and optimal in both settings.
- Contributions: Under the PL condition, PAGE automatically switches to a faster linear convergence rate.This applies to nonconvex functions satisfying the PL condition.
- Experiments: PyTorch experiments with LeNet, VGG, and ResNet confirm PAGE’s practical superiority in deep-learning experiments.The conclusion presents these experiments as support for the theoretical results.
- Broader relevance: The convergence analysis is described as simple and clean, facilitating adoption and generalization to other work.The paper connects this analysis to subsequent communication-efficient distributed-learning breakthroughs.
A.1. Proof of Main Theorem 1
The finite-sum proof establishes PAGE’s convergence and gradient-complexity bounds by combining smoothness-based descent with a variance bound for its gradient estimator. The resulting theorem targets an ϵ-approximate solution under specified minibatch and probability choices.
- Main theorem: Theorem 1 bounds PAGE’s iterations for finding an ϵ-approximate solution in the nonconvex finite-sum problem.The theorem assumes average L-smoothness and specifies stepsize, minibatch, secondary minibatch, and probability parameters.
- Gradient complexity: The corresponding gradient complexity equals the initial minibatch cost plus T times PAGE’s expected per-iteration gradient cost.The initial b gradients come from computing g0.
- Proof strategy: The proof uses the PAGE estimator’s variance relation together with the descent inequality for L-smooth functions.The argument combines the estimator bound with a stepsize-dependent inequality before summing over iterations.
- Proof strategy: The proof selects a random iterate and applies Jensen’s inequality to derive the expected gradient-norm guarantee.The output is randomly chosen from the iterates generated by PAGE.
A.2. Proofs of Corollaries 1 and 2
The corollaries recover gradient descent as a special case and establish PAGE’s optimal finite-sum rate, while the online analysis extends the framework under bounded variance. Lower-bound arguments show the finite-sum rate is unavoidable for linear-span first-order algorithms.
- Corollary 1: Setting p_t ≡ 1 reduces PAGE to gradient descent with full minibatch b = n.The resulting gradient complexity is the full-batch cost per iteration.
- Corollary 2: Choosing p = b′/(b+b′) yields an expected PAGE estimator cost of 2bb′/(b+b′) stochastic gradients per iteration.This choice underlies the optimal finite-sum gradient-complexity result.
- Lower bound: The finite-sum lower bound is Ω(n + Δ0L√n/ϵ^2) stochastic gradient computations for linear-span first-order algorithms.The construction combines an Ω(n) term with the previously established Ω(Δ0L√n/ϵ^2) term.
- Online setting: The online theorem assumes bounded variance and uses b = min{⌈2σ^2/ϵ^2⌉, n} with a smaller secondary minibatch b′.It bounds iterations and gradient complexity for an ϵ-approximate solution in the online problem.
- Proof strategy: The finite-sum and online proofs apply the same smoothness-based descent framework to PAGE’s estimator variance.The online proof invokes the bounded-variance and average-smoothness assumptions.
B.2. Proofs of Corollaries 3, 4 and 5
The corollaries recover SGD as a special case and establish PAGE's optimal online complexity using a probability that balances the two minibatch sizes. The lower bound matches this complexity for online nonconvex optimization.
- Corollary 3: p_t ≡ 1 reduces PAGE to vanilla minibatch SGD, with stepsize η ≤ 1/L and iteration bound T = 4∆0L/ϵ^2 + 1.
- Corollary 3: b stochastic gradients per iteration yield SGD gradient complexity #grad = b + Tb.
- Corollary 4: p_t ≡ b′/(b+b′) gives PAGE an expected 2bb′/(b+b′) stochastic gradients per iteration.
- Corollary 4: PAGE attains the optimal nonconvex online complexity under the stated stepsize and minibatch conditions.
- Corollary 5: Any linear-span first-order algorithm requires Ω(b + ∆0L√b/ϵ^2) stochastic gradient computations for an ϵ-approximate online solution.
C. Missing Proofs for Nonconvex Finite-Sum Problems under PL Condition
Under the PL condition, the finite-sum analysis defines a potential combining suboptimality and estimator error, then derives convergence and gradient-complexity bounds for PAGE. A specific parameter choice yields the finite-sum PL corollary.
- Main theorem: Theorem 5 bounds PAGE iterations and stochastic gradient computations for finding an ϵ-solution in the nonconvex finite-sum PL setting.
- Proof strategy: The PL analysis uses Φ_t = f(x_t) − f* + β||g_t − ∇f(x_t)||^2 to combine objective error with gradient-estimator error.
- Corollary 6: Corollary 6 sets b = n, b′ ≤ √b, and p_t ≡ b′/(b+b′) under a PL-compatible stepsize.
- Corollary 6: With this parameterization, PAGE has an explicit iteration bound and gradient-complexity bound for reaching an ϵ-solution.
- Corollary 6: The expected per-iteration estimator cost is 2bb′/(b+b′) stochastic gradients.
D. Missing Proofs for Nonconvex Online Problems under PL Condition
The online PL analysis applies the same potential-based proof strategy under bounded variance and selects minibatch parameters using σ^2, μ, ϵ, and n. It provides iteration and gradient-complexity bounds for PAGE.
- Assumptions: The online PL setting assumes bounded variance in addition to smoothness and the PL condition.
- Main theorem: Theorem 6 derives PAGE iteration and gradient-complexity bounds for finding an ϵ-solution in the online PL problem.
- Proof strategy: The proof defines Φ_t = f(x_t) − f* + β||g_t − ∇f(x_t)||^2 and combines the PL descent relation with estimator-error control.
- Corollary 7: The online PL corollary gives explicit iteration and gradient-complexity bounds under this parameter setting.