Gradient Descent Variants: SGD, Momentum and Adam

The backpropagation subtopic ended with the basic update rule: new_weight = old_weight - learning_rate * gradient. Applied repeatedly over many examples, this is gradient descent — walking downhill on the loss 'landscape' toward a minimum, like a hiker feeling which way the ground slopes and taking a step in the steepest-downhill direction. Plain gradient descent (often called SGD, Stochastic Gradient Descent, when you compute the gradient from a small random batch of examples rather than the whole dataset each step) works, but it has two well-known annoyances that better optimizers fix.

The first annoyance: SGD can be jittery and slow, especially in 'ravine' shaped loss landscapes — steep in one direction, gently sloped in another — where SGD zig-zags back and forth across the steep direction while making frustratingly slow progress along the gentle one. Momentum fixes this with an analogy straight from physics: instead of only looking at the CURRENT gradient, it keeps a running 'velocity' — an exponentially-weighted average of past gradients — and moves in that accumulated direction. Like a heavy ball rolling downhill, it builds up speed in directions that consistently point downhill (the gentle, consistent slope) while the zig-zagging in the steep direction partially cancels out (since it flips sign every step, canceling in the average). This means faster, smoother progress toward the minimum.

The second annoyance: SGD uses the SAME learning rate for every single weight, even though different weights often need very different step sizes — some weights' gradients are consistently large, others are tiny and infrequent. Adam (Adaptive Moment Estimation) fixes this by tracking, for EACH weight individually, both a momentum-style running average of the gradient (the 'first moment,' capturing direction) AND a running average of the SQUARED gradient (the 'second moment,' capturing typical magnitude). It then divides the effective step by that magnitude estimate — weights with typically large gradients get smaller effective steps (avoiding overshooting), and weights with typically small/rare gradients get proportionally larger steps (so they still learn at a reasonable pace). This combination of momentum-like direction smoothing plus per-weight adaptive step sizing is why Adam converges quickly with little manual tuning, and is the default optimizer choice for most deep learning today, including virtually all LLM training.

# Comparing plain SGD, SGD+Momentum, and a simplified Adam update rule
# on the same 1D toy loss: L(w) = (w - 3)^2  (minimum at w=3).

def gradient(w: float) -> float:
    return 2 * (w - 3)   # dL/dw for L(w) = (w-3)^2

def sgd_step(w, lr=0.1):
    return w - lr * gradient(w)

def momentum_step(w, velocity, lr=0.1, beta=0.9):
    g = gradient(w)
    velocity = beta * velocity + (1 - beta) * g   # running average of gradients
    return w - lr * velocity, velocity

def adam_step(w, m, v, t, lr=0.3, beta1=0.9, beta2=0.999, eps=1e-8):
    g = gradient(w)
    m = beta1 * m + (1 - beta1) * g          # 1st moment: direction
    v = beta2 * v + (1 - beta2) * (g ** 2)   # 2nd moment: magnitude
    m_hat = m / (1 - beta1 ** t)              # bias correction (early steps)
    v_hat = v / (1 - beta2 ** t)
    w = w - lr * m_hat / (v_hat ** 0.5 + eps)  # step scaled by gradient magnitude
    return w, m, v

w_sgd, w_mom, w_adam = 0.0, 0.0, 0.0
velocity, m, v = 0.0, 0.0, 0.0
for t in range(1, 11):
    w_sgd = sgd_step(w_sgd)
    w_mom, velocity = momentum_step(w_mom, velocity)
    w_adam, m, v = adam_step(w_adam, m, v, t)
print(f"after 10 steps -> SGD: {w_sgd:.3f}  Momentum: {w_mom:.3f}  Adam: {w_adam:.3f}")
print("(target: 3.000 -- notice how each optimizer approaches it differently)")