[Paper Review] Variance Reduction for Faster Non-Convex Optimization

Zeyuan Allen-Zhu, Elad Hazan|arXiv (Cornell University)|Mar 17, 2016

Stochastic Gradient Optimization TechniquesComputer Science25 references126 citations

TL;DR

The paper introduces a variance-reduced stochastic method for non-convex optimization that achieves an O(n^{2/3} / ε) iteration rate to reach an ε-stationary point, outperforming gradient descent and SGD without extra assumptions.

ABSTRACT

We consider the fundamental problem in non-convex optimization of efficiently reaching a stationary point. In contrast to the convex case, in the long history of this basic problem, the only known theoretical results on first-order non-convex optimization remain to be full gradient descent that converges in $O(1/\varepsilon)$ iterations for smooth objectives, and stochastic gradient descent that converges in $O(1/\varepsilon^2)$ iterations for objectives that are sum of smooth functions. We provide the first improvement in this line of research. Our result is based on the variance reduction trick recently introduced to convex optimization, as well as a brand new analysis of variance reduction that is suitable for non-convex optimization. For objectives that are sum of smooth functions, our first-order minibatch stochastic method converges with an $O(1/\varepsilon)$ rate, and is faster than full gradient descent by $Ω(n^{1/3})$. We demonstrate the effectiveness of our methods on empirical risk minimizations with non-convex loss functions and training neural nets.

Motivation & Objective

• Motivate efficient finding of stationary points for non-convex objectives.
• Improve upon GD/SGD with a variance-reduction approach suitable for non-convex loss.
• Develop an SVRG-based algorithm with provable O(n^{2/3}L(f(x0)−f(x*)) / ε) convergence rate.
• Extend variance-reduction techniques to non-convex settings and analyze variance upper bounds.
• Demonstrate empirical effectiveness on ERM with non-convex loss and neural networks.

Proposed method

• Adopt a variance-reduced gradient estimator analogous to SVRG for non-convex objectives.
• Use an epoch structure with a snapshot point xs0 and inner iterations to form e∇k = ∇fi(xsk) − ∇fi(xs0) + ∇f(xs0).
• Set inner-iteration length m = n and step size η = Θ(1/(n^{2/3}L)).
• Partition each epoch into subepochs to telescope the variance bound, leveraging a mirror-descent analysis.
• Show that the gradient estimator variance is bounded by O(∥xsk − xs0∥^2) and relate this to objective decrease.
• Provide algorithm variants and discuss mini-batch, non-uniform smoothness, and extensions to sum-of-non-convex objectives.

Experimental results

Research questions

• RQ1 Can variance-reduction techniques yield faster convergence to ε-stationary points in non-convex optimization than GD/SGD?
• RQ2 What are the appropriate variance bounds and analysis techniques for non-convex objectives to enable such acceleration?
• RQ3 How can SVRG be adapted (snapshot selection, epoch/subepoch structure) to non-convex losses?
• RQ4 Do these methods extend to ERM with non-convex losses and neural networks in practice?

Key findings

• The proposed non-convex SVRG variant achieves an ε-stationary point in O(n^{2/3}L(f(x0)−f(x*)) / ε) iterations.
• Each iteration of SVRG is as fast as SGD and n times faster than full gradient descent, yielding a theoretical Ω(n^{1/3}) speedup over GD.
• The variance bound is established as O(∥xsk − xs0∥^2), enabling epoch/subepoch analysis to guarantee progress.
• With m = n and η = Θ(1/(n^{2/3}L)), the algorithm outputs x with E[∥∇f(x)∥^2] ≤ O(L(f(xφ)−min f) / (S n^{1/3})).
• Experiments on ERM with non-convex losses and neural nets show SVRG can match or exceed SGD in training speed, especially for smaller ε and non-convex losses.