SGD with Momentum

Plain gradient descent forgets everything between steps: it moves according to the slope it feels right now. Momentum gives the optimizer a memory. After this lesson your program carries a velocity from step to step and moves a parameter with it.

Why plain SGD zig-zags

Loss surfaces in real networks are often long narrow valleys: very steep across the valley, almost flat along it. The gradient points mostly across the valley wall, so plain SGD bounces from one wall to the other and creeps forward slowly. A learning rate small enough to stop the bouncing makes the creep along the valley painfully slow.

The idea: a ball with mass

Think of the parameter as a ball rolling on the loss surface. Plain SGD is a ball with no mass that changes direction instantly. Momentum gives it inertia: the current gradient nudges an existing velocity instead of replacing it. Back-and-forth components cancel across steps, while a component that keeps pointing the same way builds up speed.

The update rule

With velocity v (starting at 0), momentum coefficient beta (typically 0.9), learning rate lr and gradient g:

v      <-  beta * v + g
theta  <-  theta - lr * v

The new gradient is added to the decayed old velocity, and the parameter moves by the whole velocity. Two sanity checks: with beta = 0 the velocity is just g, so this is exactly plain SGD; and with a constant gradient the velocity settles at g / (1 - beta), so the effective step is 1 / (1 - beta) times bigger, 10 times for beta = 0.9.

A worked trace

Start at theta = 4, lr = 0.05, beta = 0.8, and feed gradients 1, 1, -1:

step 1   g = 1    v = 0.8*0   + 1 = 1.0     theta = 4     - 0.05*1.0  = 3.95
step 2   g = 1    v = 0.8*1.0 + 1 = 1.8     theta = 3.95  - 0.05*1.8  = 3.86
step 3   g = -1   v = 0.8*1.8 - 1 = 0.44    theta = 3.86  - 0.05*0.44 = 3.838

At step 3 the gradient reversed sign, yet theta kept moving down because the velocity is still positive. That overshoot is the price of inertia, and the reason too large a beta can oscillate around the minimum.

State and its pitfalls

Momentum is stateful: the optimizer must remember v between calls, one velocity per parameter. Two details follow.

• Starting a new run must reset the velocity to 0. Leftover velocity from an earlier run silently bends the first steps of the next.
• Changing the learning rate mid-training affects only future steps. The velocity already accumulated is not rescaled.

Common mistakes

• Updating theta with the raw gradient instead of the velocity.
• Using the new velocity in one place and the old one in another within a single step.
• Using the damped form v = beta*v + (1-beta)*g. That is a different, equally valid optimizer variant, but it changes every number. This lesson uses the classical form above.
• Forgetting to zero v on INIT.

Your exercise

The input is a script of INIT, LR, MOM, STEP and PRINT commands. The starter keeps theta, v, lr and beta, and applies your function at every STEP. Implement momentum_step, which takes the current theta and v and returns their new values.