Momentum on a lopsided bowl
Medium14 min11 hidden tests
A team is tuning an optimizer on a toy loss shaped like a long, narrow bowl: L(w) = 0.5 * sum(c * w ** 2), where each coordinate has its own curvature c, so the gradient is c * w. Plain gradient descent bounces across the steep direction and crawls along the flat one. Momentum keeps a velocity v, a running sum of past gradients, and steps along it. Write momentum_descent(curvature, w0, lr, beta, steps) that runs the updates from w0 and returns the final weights.
Drag the start, tune beta
v = beta · v + c · w, w = w − lr · vAfter 30 steps from (−4.2, 1.6): plain descent is at (−0.248, 0.002) with loss 0.0308; momentum with beta 0.60 is at (−0.003, 0.000) with loss 3.4e−6. With lr × 20 above 1, plain descent overshoots the steep side (curvature 20) and bounces across it, while it crawls along the flat side (curvature 1). Momentum keeps a velocity that builds up along the flat side; with beta near 1 it overshoots and swings instead.
Plain loss 0.0308, momentum loss 3.4e−6.
Examples
The second step is longer because the velocity carries over
momentum_descent([1.0], [1.0], 0.1, 0.9, 2)
returns [0.72]
With beta = 0, momentum is plain gradient descent
momentum_descent([2.0], [3.0], 0.1, 0.0, 2)
returns [1.92]
Requirements
curvatureandw0are lists or 1-D arrays of the same length,lris a positive float,betais at least 0 and below 1, andstepsis 0 or more.vstarts at zero. Each step computes the gradientc * wat the current weights, then setsv = beta * v + grad, thenw = w - lr * v.- Return the weights after
stepssteps as a 1-D float array. Withsteps = 0, returnw0as floats.
The idea behind it
Plain gradient descent steps along the current gradient: w = w - lr * grad. On a loss shaped like a narrow valley, the gradient points mostly across the valley, so the weights bounce side to side and make slow progress along it. Momentum adds memory. It keeps a velocity v = beta * v + grad, a sum of past gradients where older ones fade by a factor of beta each step, and moves along v. Gradients that keep pointing the same way, along the valley floor, add up and speed the weights along. Gradients that flip sign, across the valley, partly cancel. With beta = 0.9, a steady gradient builds a step up to 10 times longer than plain descent. PyTorch's SGD with momentum uses this exact form. Some texts scale the gradient by 1 - beta, which makes v an average of past gradients. That only changes the effective learning rate, but code and tuning must agree on one.