SGD with Momentum

Plain stochastic gradient descent updates a parameter by stepping directly opposite its gradient: theta -= lr * grad. That works, but it's short-sighted — it reacts only to the current gradient and ignores everything about the trajectory so far. On loss surfaces with narrow ravines (steep in one direction, shallow in another — extremely common in real neural nets), plain SGD oscillates back and forth across the ravine instead of making steady progress along it. Momentum fixes this by giving the optimizer memory.

The physical intuition

Think of theta as a ball rolling on the loss surface. Plain SGD is a ball with no mass — it changes direction instantly every step, following whatever the local gradient says right now. Momentum gives the ball mass and inertia: it keeps rolling in the direction it's been going, and the current gradient only nudges that existing velocity rather than replacing it outright. This smooths out oscillations (they cancel out across steps) and accelerates movement along directions where the gradient has been consistently pointing the same way.

The update rule (classical / Polyak momentum)

v_{t+1} = beta * v_t + g_t
theta    = theta - lr * v_{t+1}

• v is the velocity — an accumulator that persists across steps.
• beta (commonly 0.9) controls how much of the old velocity survives each step. beta = 0 collapses this exactly to plain SGD (v_{t+1} = g_t, so theta -= lr * g_t) — a good sanity check when you implement it.
• Each step, the new gradient g_t is added on top of the decayed previous velocity, and the parameter moves by the full velocity, not just the raw gradient.

Why this needs careful state management

Momentum is stateful across calls — unlike a pure function of (theta, grad), the optimizer must remember v between STEP calls. Two implementation details fall out of that:

• INIT must reset v to zero. If you start a new trajectory (a new parameter, or restarting training) without zeroing velocity, leftover momentum from a previous run will bias the very first steps of the new one — a subtle bug that's easy to miss because it only shows up as "training behaves a bit differently the second time," not a crash.
• LR and MOM changes apply going forward, not retroactively. If a learning-rate schedule changes lr mid-training, the existing accumulated velocity v doesn't get rescaled — only future steps use the new lr.

Worked trace

Starting theta = 10.0, lr = 0.1, beta = 0.9, v = 0:

step 1: grad = 2.0
  v = 0.9*0 + 2.0 = 2.0
  theta = 10.0 - 0.1*2.0 = 9.8

step 2: grad = 2.0
  v = 0.9*2.0 + 2.0 = 3.8
  theta = 9.8 - 0.1*3.8 = 9.42

Notice the effective step size grows across identical gradients (2.0 → 3.8 in the velocity) — that's momentum accelerating in a direction the gradient keeps agreeing with, exactly the behavior that makes it converge faster than plain SGD on realistic loss landscapes, and exactly why it's the default upgrade every real training loop applies before reaching for fancier optimizers like Adam.