Optimization for Data Science · Lecture 5 demo
Momentum, acceleration, and restarts
Four short experiments, meant to be opened one at a time. Every method uses one gradient per iteration. Nothing runs until you press Run; click a landscape to move the starting point.
1Gradient descent zigzags
On an ill-conditioned quadratic, gradient descent bounces between the steep walls of a narrow valley while it creeps along the valley floor. The step size cannot fix both directions at once.
Parameter space
Contours of f. Click to choose the starting point.
Objective
f(wk)/f(w0), logarithmic scale
- Run with the default step. The steep coordinate is multiplied by 1 − ηL < 0 at every step, so it flips sign: that is the zigzag. The flat coordinate barely shrinks.
- Press η = 1/L. The zigzag disappears after one step, but progress along the valley is no faster.
- Increase κ. The number of iterations grows like κ, whatever the step size.
Mathematical setup
f(w) = ½(L⟨u1,w⟩² + μ⟨u2,w⟩²), L = 1, μ = 1/κ, w⋆ = 0, wk+1 = wk − η∇f(wk).
The valley floor u2 is tilted by 25°. Each eigen-coordinate evolves on its own: the coordinate with curvature λ is multiplied by 1 − ηλ at every step. The dashed line is q2k with q = max{|1 − ηL|, |1 − ημ|}. The best fixed step, 2/(L+μ), makes both factors equal to (κ−1)/(κ+1).
2Momentum helps on quadratics
Heavy ball adds a fraction of the previous step, wk+1 = wk − η∇f(wk) + β(wk − wk−1). Along the valley the steps point the same way and accumulate. Polyak's choice also takes long steps (η close to 4/L), so the first few iterations overshoot across the valley before the momentum damps the oscillation.
Parameter space
Same quadratic and start for both methods. Click to move the start.
Objective
f(wk)/f(w0), logarithmic scale
- Run with Polyak's choice. Watch the first few steps overshoot across the valley (f even increases briefly), then settle. Heavy ball still needs about √κ times fewer iterations than GD.
- Increase κ to 1000. GD's count grows like κ, heavy ball's like √κ.
- Move β away from Polyak's value. Too little momentum and progress along the valley slows; too much and the path overshoots and spirals.
Mathematical setup
Same quadratic as in Section 1. Polyak's choice is η = 4/(√L+√μ)² and β = ρ², with ρ = (√κ−1)/(√κ+1), the rate from the Lecture 4 lower bound. On a quadratic each eigen-coordinate follows ck+1 = (1+β−ηλ)ck − βck−1; with Polyak's choice every root has modulus ρ. The dashed lines are r2k, where r is the largest root modulus over λ ∈ [μ, L]: (κ−1)/(κ+1) for GD and ρ for heavy ball with Polyak's choice. Heavy ball's actual curve can lag its dashed line by a polynomial factor (double roots at λ = μ and λ = L).
3…but momentum can fail
This function is 1-strongly convex and 25-smooth: it has curvature 25 everywhere except on [1, 2], where its curvature is 1. Heavy ball uses Polyak's parameters for κ = 25, which are optimal on every quadratic in this class. From many starting points it never converges.
The function and the iterates
Shaded: the interval [1, 2] with curvature 1.
Distance to the minimizer
|wk − w⋆|, logarithmic scale; exact zeros drawn at the bottom
- Run from w0 = 3.3. After a few iterations heavy ball settles into a cycle of three points and stays there.
- Drag the starting point. The strip shows every start in [−5, 5] from which heavy ball fails.
- Gradient descent with step 1/L converges from every start. Here it even lands on w⋆ within two steps.
This is not a matter of tuning. Goujaud, Taylor, and Dieuleveut proved that for every choice of η and β, heavy ball is either not accelerated even on quadratics, or fails to converge on some smooth strongly convex function.
Mathematical setup
f′(w) = 25w for w < 1, w + 24 for 1 ≤ w < 2, 25w − 24 for w ≥ 2; f(0) = 0.
Heavy ball uses η = 4/(√25+1)² = 1/9 and β = ρ² = 4/9. The example is from Lessard, Recht, and Packard (SIAM J. Optim., 2016, Section 4.6).
4Restarted accelerated gradient descent
Nesterov's method evaluates the gradient after the momentum step, and its momentum grows like k/(k+3). Restarting it every N iterations throws its momentum away and starts again from the current point. When heavy ball converges it can be faster still; the difference is the guarantee.
Parameter space
Objective gap
(f(wk) − f⋆)/(f(w0) − f⋆), logarithmic scale; ticks mark restarts
- Run from the default start. Heavy ball is trapped in a cycle, GD crawls, AGD converges with ripples, and restarted AGD converges linearly.
- Click a start whose w₁ lies outside the red strip. There heavy ball converges, and it is the fastest. Above the red strip it never converges.
- Change N, or switch to the adaptive rule, and compare the restarted curve with AGD.
- Switch to logistic regression. Its curvature is only local, and AGD's iterates stay in a ball around the solution.
Mathematical setup
AGD: wk+1 = yk − ∇f(yk)/L, yk+1 = wk+1 + βk(wk+1 − wk), βk = (λk+1 − 1)/λk+2, λk+1 = (1 + √(1+4λk²))/2.
A restart sets y = w and λ back to its initial value, so the momentum starts again from zero. GD uses η = 1/L. Heavy ball uses Polyak's choice for the stated μ and L. In the first problem, f(w) = φ(w1) + w2²/8, where φ is Section 3's function: L = 25 and μ = 1/4. Because f is separable, heavy ball fails exactly when w1 starts in the red strip. In logistic regression (80 points in ℝ², no intercept), L = λmax(XᵀX)/(4n), and μ is the smallest Hessian eigenvalue at the solution: f is strongly convex near w⋆ but not globally.