Nesterov Lookahead Momentum
Compare a gradient taken at a momentum lookahead with one taken at the current point.
By Riano · Published · Updated · Report a correction
Nesterov Lookahead Momentum
Standard momentum, often called the heavy-ball method, updates the current position using the current velocity and the gradient at the current position. Nesterov momentum improves upon this by evaluating the gradient not at the current position, but at a "lookahead" position where the momentum alone would take us.
Worked Example
Consider the objective function f(x) = 0.5 * a * x2. Let a = 1, giving f(x) = 0.5 * x2 with gradient f'(x) = x. Assume step size (eta) = 0.25, momentum (mu) = 0.5, and initial state x0 = 2, v0 = 0.
- Update 1:
Lookahead: x0 + mu * v0 = 2 + 0.5 * 0 = 2.
Gradient at lookahead: 2.
Velocity update: v1 = mu * v0 - eta * 2 = 0.5 * 0 - 0.25 * 2 = -0.5.
Position update: x1 = x0 + v1 = 2 - 0.5 = 1.5.
Next Lookahead: x1 + mu * v1 = 1.5 + 0.5 * (-0.5) = 1.25. - Update 2:
Gradient at lookahead (1.25): 1.25.
Velocity update: v2 = mu * v1 - eta * 1.25 = 0.5 * (-0.5) - 0.25 * 1.25 = -0.25 - 0.3125 = -0.5625.
Position update: x2 = x1 + v2 = 1.5 - 0.5625 = 0.9375.
Next Lookahead: x2 + mu * v2 = 0.9375 + 0.5 * (-0.5625) = 0.65625.
Experiment
Explore these differences interactively in the Nesterov Lookahead Momentum module.
- Set Curvature to 1, Step size to 0.25, and Momentum to 0.5.
- Observe that at Updates 0, both Nesterov and Heavy ball methods start at position 2.
- Change Updates to 1. Both methods update to 1.5 because the initial velocity is zero, making the lookahead identical to the current position.
- Change Updates to 2. Nesterov momentum uses the gradient at the lookahead (1.25), resulting in a position of 0.9375. The Heavy ball method uses the gradient at the current position (1.5), resulting in a position of 0.875. The trajectories diverge.
Limitations
The visualization is limited to a 1D objective (plotted as a 2D parabola) to clearly trace the path of the position and its momentum lookahead. Real-world loss functions are highly dimensional and non-convex. Furthermore, for extremely large step sizes or high curvature, the discrete updates can overshoot significantly or diverge entirely, and this simulation shows the exact finite numerical results rather than clipping or artificially ensuring convergence.
References
The original formulation of Nesterov momentum provides a more aggressive update by predicting the next position before calculating the gradient. PyTorch implements this parameterization as an option within its SGD optimizer.
- PyTorch SGD Documentation - Explains the practical implementation details of SGD with Nesterov momentum.
Experience it interactively
Adjust parameters, observe in real time, and build deep intuition with Riano’s interactive Nesterov Lookahead Momentum module.
Try Nesterov Lookahead Momentum on Riano →