Source-linked AI summary
Momentum-Based Variance Reduction in Non-Convex SGD
Ashok Cutkosky, Francesco Orabona
TL;DR
Non-convex variance-reduction methods improve SGD’s convergence but typically require carefully tuned learning rates and large checkpoint batches. The paper introduces STORM, which uses momentum-based recursive variance reduction with adaptive learning rates and no batches. STORM achieves the optimal convergence guarantee while adapting to gradient noise without knowing its level.
Problem
Existing non-convex variance-reduction methods improve convergence but rely on carefully tuned learning rates and giant checkpoint batches.
Method
STORM uses a momentum variant that recursively reduces gradient variance while adapting its stepsize and momentum term from past gradients, without checkpoint gradients or batches.
Results
STORM achieves the optimal convergence rate O(1/T^1/3) and adapts to the noise level without requiring knowledge of it.
Takeaways & Limitations
Removing checkpoint gradients and using adaptive learning rates makes STORM easier to tune while retaining the optimal convergence guarantee.
Takeaways & Limitations
The analysis assumes unbiased stochastic gradients, bounded gradient noise, and smooth sample losses; its adaptive analysis also assumes G-Lipschitz sample losses.
Abstract
from arXiv · showhide
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 $\boldsymbol{x}$ with $\mathbb{E}[\|\nabla F(\boldsymbol{x})\|]\le O(1/\sqrt{T}+σ^{1/3}/T^{1/3})$ in $T$ iterations with $σ^2$ variance in the gradients, matching the optimal rate but without requiring knowledge of $σ$.
1 Introduction
The paper studies finding critical points in stochastic, non-convex optimization, where SGD requires carefully tuned learning rates and achieves an O(1/T^1/4) guarantee. It introduces STORM, which uses momentum-based variance reduction, adaptive learning rates, and no batches to achieve O(1/T^1/3).
- Non-convex optimization relaxes finding a true minimum to finding a critical point using only stochastic gradients.Finding a true minimum may be NP-hard.
- SGD achieves E[∥∇F(x_t)∥] ≤ O(1/T^1/4) when its learning rates are properly tuned.The guarantee applies to a randomly selected iterate.
- Variance reduction improves non-convex convergence rates from O(1/T^1/4) through O(1/T^3/10) to O(1/T^1/3), but has seen limited practical success.The cited practical concerns include non-adaptive learning rates and giant checkpoint batches.
- STORM uses a momentum variant for variance reduction, requires no batches or checkpoint gradients, and automatically adapts its learning rate to gradient variance.The method is presented as STOchastic Recursive Momentum.
- STORM achieves the optimal convergence rate O(1/T^1/3) without requiring knowledge of the noise level.
2 Related Work
Related variance-reduction methods improve convergence but generally rely on checkpoint gradients, large sample batches, and carefully selected parameters. Adaptive learning rates offer data-dependent tuning, though prior adaptive variance reduction was limited to convex losses.
- Earlier non-convex variance-reduction methods require gradients computed at checkpoint points.Checkpoint sample counts may be as large as O(T) and typically no smaller than O(T^2/3).
- Some checkpoint-free exceptions do not improve over plain SGD, while a contemporaneous method matches STORM’s rate but still requires one checkpoint gradient.
- Adaptive learning-rate schemes choose η_t from data to reduce manual tuning and can improve non-convex SGD’s guarantee to O(1/T + (σ^2/T)^1/4).
- Prior adaptive variance reduction known to the authors applied only to convex losses.
3 Notation and Assumptions
The analysis assumes unbiased stochastic gradients from independent samples, bounded gradient noise, and smooth sample losses. The adaptive analysis additionally assumes Lipschitz sample losses, though that assumption can be removed at the expense of adaptivity to σ.
- Vectors are denoted with bold letters, and a · b denotes the inner product between vectors a and b.
- The stochastic oracle uses independent ξ_t with E[f(x, ξ_t) | x] = F(x), and each update accesses two gradients on the same ξ_t at different points.
- Gradient noise is bounded by E[∥∇f(x, ξ_t) − ∇F(x)∥^2] ≤ σ^2.
- Sample losses are assumed differentiable and L-smooth, while the adaptive analysis also assumes G-Lipschitzness.The Lipschitz assumption can be lifted at the expense of adaptivity to σ.
4 Momentum and Variance Reduction
The paper modifies momentum by adding a same-sample gradient difference, producing a recursive direction estimate that can reduce variance without checkpoint gradients. Its error contracts while smoothness controls the added difference term.
- Momentum and Variance Reduction: Classic momentum uses an exponential average of past observed gradients instead of the current gradient.The cited description uses a small momentum parameter, such as a = 0.1.
- Momentum and Variance Reduction: Noise can nullify momentum’s theoretical gain in stochastic gradients, motivating a variant designed specifically for variance reduction.
- Momentum and Variance Reduction: The proposed direction adds (1 − a)(∇f(x_t, ξ_t) − ∇f(x_{t−1}, ξ_t)) to the momentum update using two gradients from the same sample.
- Momentum and Variance Reduction: Unlike standard checkpoint-based variance reduction, the update does not require gradients evaluated at checkpoint points.
- STORM: The algorithm initializes with one stochastic gradient, updates recursively, and returns a uniformly selected iterate in theory or the final iterate in practice.
- Momentum and Variance Reduction: The direction error contracts recursively, while its stochastic and smoothness terms are controlled by choosing sufficiently small a and η.The smoothness-controlled difference is of order O(ηd_{t−1}).
5 Storm: STOchastic Recursive Momentum
Storm achieves variance reduction without checkpoint gradients or batches by combining recursive momentum with adaptive learning rates. Its analysis uses a time-varying Lyapunov potential and establishes stochastic convergence at the optimal O(1/T^1/3) rate without requiring prior knowledge of the gradient-noise level.
- Algorithm: Storm uses past gradients to adapt both the stepsize and momentum term, achieving variance reduction without checkpoint gradients or batches.The method’s momentum variant is designed to replace the giant batches used to construct low-noise checkpoint gradients.
- Convergence guarantee: O(1/T^1/3) is Storm’s stochastic convergence rate, matching the optimal rate previously obtained by checkpoint-based variance-reduction methods.The result applies in expectation to the norm of the gradients in the stochastic case.
- Convergence guarantee: O(ln T / sqrt(T)) is the convergence rate when there is no gradient noise.Storm achieves both the noiseless and stochastic rates automatically, without knowing the noise level or tuning stepsizes.
- Assumptions: The main theorem assumes each stochastic loss f(x, ξ) is G-Lipschitz, but an oracle-tuned σ substitution can remove this requirement while retaining checkpoint-free updates.The alternative is described in Appendix B rather than in the primary theorem statement.
- Analysis: A time-varying potential Φ_t = F(x_t) + z_t||ϵ_t||^2 tracks both objective value and update error, enabling the checkpoint-free analysis.The weighting satisfies z_t ∝ η_t^-1, and the proof bounds successive potential differences before summing them over iterations.
- Analysis: The variance-estimate recurrence suggests ||ϵ_t||^2 = O(1/t^2/3 + ||∇F(x_t)||^2), making the gradient estimate accurate when the true gradient is sufficiently large.The resulting equilibrium scale for ||∇F(x_t)||^2 is O(1/T^2/3).
6 Empirical Validation
Storm was evaluated on CIFAR-10 image recognition with a ResNet and compared against AdaGrad and Adam using one tuned parameter per algorithm. Storm was somewhat faster in training loss and accuracy, but only marginally better than AdaGrad on test accuracy.
- Experimental setup: CIFAR-10 image recognition was tested with a ResNet implementation, comparing Storm against AdaGrad and Adam.Learning rates for AdaGrad and Adam were swept, while Storm used fixed w = k = 0.1 and swept c, leaving one parameter to tune for each algorithm.
- Results: Storm appeared somewhat faster than the baselines in training loss and training accuracy measured by iterations.The reported convergence proof applies only to training loss because the experiments make multiple passes over the dataset.
- Results: Storm was only marginally better than AdaGrad on test accuracy.The authors leave open whether regularization could trade Storm’s stronger training-loss performance for better test performance.
7 Conclusion
The paper concludes that Storm combines variance reduction with adaptive learning rates without checkpoint gradients, making it easier to tune while retaining an optimal convergence guarantee. Its update also resembles practical SGD with momentum, suggesting a possible connection between momentum and variance reduction.
- Conclusion: Storm removes checkpoint gradients and uses adaptive learning rates, avoiding checkpoint-size and checkpoint-frequency choices.The algorithm adapts to the problem’s noise level without requiring prior knowledge of that parameter.
- Conclusion: Storm obtains the optimal convergence guarantee for stochastic, smooth, non-convex critical-point finding.The paper also reports fewer iterations than baseline algorithms when optimizing CIFAR-10 with a ResNet.
- Momentum connection: Storm’s update formula is strikingly similar to the standard SGD-with-momentum heuristic used in practice.The algorithm is not precisely the same as SGD with momentum, but the authors view the similarity as intuitive evidence that momentum may perform variance reduction.
A Extra Lemmas
The appendix restates and proves supporting lemmas using convexity, smoothness, independence, Young’s inequality, and norm bounds. These lemmas establish technical relationships used in the main convergence analysis.
- A Extra Lemmas: The appendix restates and proves lemmas supporting the paper’s analysis.It includes a standard convexity consequence and a technical observation important for proving Lemma 2.
- Lemma 1: The proof of Lemma 1 uses smoothness of F and the algorithm’s update definition.Young’s inequality, a squared-norm bound, and η_t ≤ 1/4L are used in successive inequalities.
- Technical observation: The technical observation relies on ϵ_t−1 and x_t being independent of ξ_t under the stated hypothesis.Expectations are taken first over ξ_t and then over earlier samples ξ_1, ..., ξ_t−1.
- Lemma 4: Lemma 4 bounds a sequence using concavity of the logarithm and summation over t = 1, ..., T.The proof starts with a logarithmic concavity inequality and sums both sides to obtain the stated bound.
A.1 Proof of Lemma 2
The proof of Lemma 2 expands the estimator error and bounds its squared norm through prior lemmas, norm inequalities, and smoothness assumptions.
- A.1 Proof of Lemma 2: Lemma 2 is deferred to the appendix and restated there for reference.The proof begins by observing two relations before expanding the estimator error.
- Error expansion: The estimator error is written as ϵ_t = ∇f(x_t, ξ_t) + (1 − a_t)(d_t−1 − ∇f(x_t−1, ξ_t)) − ∇F(x_t).This expression follows from the definition of ϵ_t and the notation in Algorithm 1.
- Error bound: The resulting bound uses Lemma 3, squared-norm inequalities, conditions (5) and (6), and Lipschitzness and smoothness of f.Each successive inequality applies one of these ingredients to control the estimator error.
B Non-adaptive Bound Without Lipschitz Assumption
Algorithm 2 removes the G-Lipschitz requirement by replacing G-dependent quantities with the oracle-tuned variance parameter σ, while retaining STORM’s convergence structure. Its deterministic adaptive schedule supports a batch-free recursive-momentum update and a bound presented as nearly identical to Theorem 1.
- Algorithm: The update uses η_t = k/(w+σ_t^2)^(1/3), recursively combines two stochastic gradients, and returns a uniformly selected iterate.The algorithm samples ξ_t+1 and updates d_t+1 using the correction ∇f(x_t+1,ξ_t+1)−∇f(x_t,ξ_t+1).
- Removing the Lipschitz requirement: Algorithm 2 replaces G and G_t with the oracle-tuned value σ, removing the G-Lipschitz requirement.This sacrifices adaptivity to σ.
- Guarantee: Theorem 2’s convergence guarantee is nearly identical to Theorem 1, but exchanges adaptivity to σ for removal of the G-Lipschitz assumption.The theorem is stated for k = bσ under the assumptions in Section 3.
- Proof strategy: The proof analyzes the potential Φ_t = F(x_t) + 1/(32L^2η_{t−1})||ε_t||^2 and bounds its one-step change.The analysis uses η_t ≤ 1/(4L), controls the adaptive-step terms, and then sums the resulting inequalities over t.
- Proof strategy: Because η_t is deterministic and independent of ∇F(x_t), the proof can avoid Cauchy–Schwarz in its remaining steps.This simplifies the non-adaptive analysis relative to the adaptive result.