Gradient descent, momentum and Adam from scratch

Once you know which way to nudge a number, you need a rule for how far to step and when to stop. That rule is the optimiser, and this section builds the three you meet in every training script — plain gradient descent, momentum, and Adam — one on top of the other, in a few lines each. Writing them yourself is what turns the learning rate from a knob you copy off a blog post into something you can reason about when a run refuses to converge.

Step downhill#

import numpy as np
rng = np.random.default_rng(0)
x = rng.uniform(0, 10, 200)
y = 3 * x + 2 + rng.normal(0, 1, 200)
X = np.c_[x, np.ones_like(x)]

w = np.zeros(2)
lr = 0.01
for step in range(2000):
    grad = 2 / len(y) * X.T @ (X @ w - y)     # gradient of mean squared error
    w -= lr * grad
print(w)                                        # [2.97471834 2.04697713]

Two things are worth naming in those seven lines. The gradient expression is the derivative of the mean squared error, and it has a readable shape: X @ w minus y is the vector of errors, and multiplying by X transposed asks how much each feature was implicated in them. And the loop has no stopping rule; it runs a fixed 2,000 steps. That is deliberate, because the honest way to decide when to stop is to watch the loss, which the next few blocks do.

def mse(w):
    return ((X @ w - y) ** 2).mean()

def gradient(w):
    return 2 / len(y) * X.T @ (X @ w - y)

best = np.linalg.lstsq(X, y, rcond=None)[0]      # the exact answer, for comparison
print(best, mse(best))       # [2.97469853 2.04711649] 1.042141861405242

The learning rate#

Try lr = 0.0001: it crawls. Try lr = 0.05: it overshoots, the loss grows and the weights explode. There is a range that works, and it depends on the curvature of the loss, which is why the learning rate is the most important setting in all of deep learning.

for lr in (0.0001, 0.001, 0.01, 0.03):
    w = np.zeros(2)
    for step in range(2000):
        w = w - lr * gradient(w)
        if not np.isfinite(w).all():
            break
    ok = np.isfinite(w).all()
    print(lr, np.round(w, 4) if ok else "diverged", round(mse(w), 4) if ok else "inf")

# 0.0001  [3.1808 0.5967]   1.5423    still moving; the intercept is nowhere near 2
# 0.001   [3.0638 1.4204]   1.1355    closer, still short after 2000 steps
# 0.01    [2.9747 2.047 ]   1.0421    at the optimum
# 0.03    [-3.02308336e+253 -4.29660063e+252]  inf   diverged

Look carefully at the 0.0001 row before moving on, because it is the failure that wastes the most time in practice. The slope, 3.18, is almost right; the intercept, 0.60, is a third of the way to 2. The loss curve for that run is smooth and decreasing and looks entirely healthy. Nothing in it says the rate is wrong. The only way to know is to have something to compare against: a closed-form answer, a longer run, or a larger rate.

Mini-batches#

Computing the gradient over the whole dataset every step is slow. Stochastic gradient descent uses a small random batch instead. Each step is noisier and far cheaper, and the noise often helps escape poor regions.

w = np.zeros(2)
lr = 0.01
index = np.arange(len(y))
shuffler = np.random.default_rng(1)

for epoch in range(50):
    shuffler.shuffle(index)
    for start in range(0, len(index), 32):          # batches of 32
        batch = index[start:start + 32]
        g = 2 / len(batch) * X[batch].T @ (X[batch] @ w - y[batch])
        w = w - lr * g

print(np.round(w, 4), round(mse(w), 4))     # [2.9494 1.7267] 1.2568

Momentum#

w = np.zeros(2)
v = np.zeros_like(w)
lr, steps = 0.01, 2000

for step in range(steps):
    g = gradient(w)
    v = 0.9 * v + g          # a running direction, like a ball rolling downhill
    w -= lr * v
print(np.round(w, 4), round(mse(w), 4))     # [2.9747 2.0471] 1.0421

In a long narrow valley, plain gradient descent zigzags across. Momentum averages recent gradients, so the zigzags cancel and progress along the valley accumulates.

The same answer is not the point; the number of steps is. Counting how many updates each method needs to get within one per cent of the best achievable loss on this data gives plain descent 435, momentum 68, and Adam at a learning rate of 0.5 just 81 — and Adam at 0.1 needs 307, which is the reminder that these methods still have a learning rate and it still has to be roughly right.

Adam#

w = np.zeros(2)
m = np.zeros_like(w); v = np.zeros_like(w)
b1, b2, eps = 0.9, 0.999, 1e-8
lr, steps = 0.5, 2000

for t in range(1, steps + 1):
    g = gradient(w)
    m = b1 * m + (1 - b1) * g           # average gradient (momentum)
    v = b2 * v + (1 - b2) * g ** 2      # average squared gradient (scale)
    m_hat = m / (1 - b1 ** t)           # correct the start-up bias towards zero
    v_hat = v / (1 - b2 ** t)
    w -= lr * m_hat / (np.sqrt(v_hat) + eps)
print(np.round(w, 4), round(mse(w), 4))     # [2.9747 2.0471] 1.0421

Adam combines momentum with a per-parameter step size: parameters with consistently large gradients take smaller steps, rarely updated ones take larger ones. It works well with little tuning, which is why it, and its weight decay variant AdamW, are the default for transformers.

The bias correction is the line people skip and it is the reason Adam works from step one. Both running averages start at zero, so after the first update m is only a tenth of the gradient it just saw and v is a thousandth of the squared gradient. Dividing by 1 minus the decay to the power of t undoes exactly that: at t equal to 1 it divides m by 0.1 and v by 0.001, recovering the true values, and by a few hundred steps the correction has faded to nothing. Without it the first few steps are far too small in a way that does real damage on short runs.

Where Adam earns its keep#

On well-scaled data Adam is a convenience. Change one thing — put the first feature in units a thousand times larger, which is exactly what happens when one column is in metres and another in millimetres — and it becomes the difference between training and not training.

X2 = np.c_[x * 1000, np.ones_like(x)]          # the same data, one column rescaled

def grad2(w):
    return 2 / len(y) * X2.T @ (X2 @ w - y)

def mse2(w):
    return ((X2 @ w - y) ** 2).mean()

print(np.linalg.lstsq(X2, y, rcond=None)[0], mse2(np.linalg.lstsq(X2, y, rcond=None)[0]))
# [0.0029747 2.04711649] 1.0421418614052422   -- the same fit, slope divided by 1000

for lr in (1e-8, 1e-7):
    w = np.zeros(2)
    for t in range(5000):
        w = w - lr * grad2(w)
        if not np.isfinite(w).all():
            break
    ok = np.isfinite(w).all()
    print(lr, np.round(w, 6) if ok else "diverged", round(mse2(w), 4) if ok else "inf")
# 1e-08  [3.264e-03 4.900e-05]  2.0383    the intercept has barely moved from zero
# 1e-07  diverged               inf

w = np.zeros(2)
m = np.zeros_like(w); v = np.zeros_like(w)
for t in range(1, 5001):
    g = grad2(w)
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g ** 2
    w -= 0.01 * (m / (1 - b1 ** t)) / (np.sqrt(v / (1 - b2 ** t)) + eps)
print(np.round(w, 6), round(mse2(w), 4))    # [0.002975 2.045165] 1.0421

That is the whole argument for adaptive optimisers in one experiment, and it is also an argument for standardising features when you can. Adam rescues a badly conditioned problem; it does not make conditioning stop mattering.