SGD is well-known for large-scale optimization. In my mind, there are two (and only two) fundamental improvements since the original introduction of SGD: (1) variance reduction, and (2) acceleration. The following picture predicts —in theory— the performance difference between the naive SGD, the variance-reduced SGD, and the accelerated SGD:
Here, is the number of stochastic iterations and a full-pass of the dataset corresponds to iterations. Some simple conclusions from the picture includes:
- As long as is more than a few multiples of , then the naive SGD should never be used. The red curve is always slower than variance reduction / acceleration.
- Accelerated SGD is no slower than variance-reduced SGD in all parameter regimes. The blue and purple curves touch exactly at the same point when is a few copies of . Ignoring the log factor, if , then accelerated SGD in theory is times faster than variance-reduced SGD.
Consider again the famous composite convex minimization problem
Assume for simplicity that each is -smooth, and the function is not strongly convex. (These assumptions can be safely removed.) Then, the convergence rates of the vanilla SGD, variance-reduced SGD (best result in this paper), and accelerated SGD (best result in this paper) are respectively:
Here, is a variance upper bound on . The relationship between and depends on the structure of ; in general, can be up to .
The Main Algorithmic Idea
Recall that the theory of acceleration is first introduced by Nesterov and studied in full-gradient and coordinate-gradient settings. The main idea is to use momentum, sometimes referred to as Nesterov momentum.
Informally speaking, instead of moving in the negative-gradient direction , one can move to for some momentum parameter . This reduces the convergence rate from to , which is quadratically faster. (The actual algorithm is a big more involved.)
Unfortunately, Nesterov momentum does not work when stochastic gradients are present. There exist datasets (see the end of the post) for which momentum fails to give the convergence. The underlying reason is because the error incurred by stochastic gradients somehow get blown up by Nesterov momentum.
In this recent paper of mine, the mystery of acceleration has finally been unraveled for stochastic-gradient descent. At a high-level, on top of Nesterov momentum, I introduced a brand new Katyusha momentum, serving as a negative momentum to cancel out the error accumulated by Nesterov momentum.
Informally, in each of the first iterations, we compute where is the variance-reduced stochastic gradient (see this post). Then, define the true next point . As illustrated by the picture below, we are essentially putting a magnet around the starting point , and retract to it back to with a rate .Then, once every iterations, we move the magnet to . See the picture belowThe black vectors are negative stochastic-gradient directions, and the green vectors are Nesterov momentum directions. I called these orange vectors the Katyusha momentum, which —seemingly making negative progress— in fact stabilize the algorithm. This novel idea ensures that the accelerated convergence rate is obtained firmly. This is the first known convergence rate among stochastic methods.
PS: the actual algorithm is a bit more involved than the pictures. See my paper for details.
An Experimental Illustration
The following simple experiment (on dataset RCV1) illustrates the comparison between variance-reduced SGD (red), my Katyusha method (black), and the naive “accelerated” method that does not use Katyusha momentum (dashed, tuned with the best parameter). It is clear from this picture that, without Katyusha momentum, the accelerated convergence rate may not be achievable.
What is Next?
My Katyusha method remains to be the first and only method that converges in the rate among stochastic methods for convex composite minimization.
When I first wrote the paper, I predicted that the Katyusha momentum should work in a bigger agenda and everyone working in large-scale optimization should give Katyusha a hug. I’m glad to see this week that a group of researchers have already applied Katyusha momentum to ADMM, and making accelerated stochastic convergence possible there.
How about non-convex optimization? Well, the naive Nesterov-accelerated SGD has been applied to deep learning, and works reasonably well. I bet that Katyusha-accelerated SGD works even better for deep learning but I don’t have the time to perform numerical experiments. In general, there is no theory of acceleration for non-convex optimization.