Gradient descent

1. Intuition

Training means finding parameter values that make the loss small. The loss is a function of the parameters: change a weight, and the loss changes. For one parameter you could plot the loss as a curve and read off the bottom. With millions of parameters you cannot plot anything, but you can compute the local slope in every direction at once: the gradient.

The gradient points uphill, in the direction where the loss grows fastest. So you move the other way, by a small amount, and recompute. That's the whole algorithm.

Three things decide how well it works:

• The step size (learning rate). Too small and progress is painfully slow. Too large and you jump over the valley, land higher on the other side, and can bounce out entirely.
• The shape of the valley. A round bowl is easy. A long narrow valley is hard: the steep walls want tiny steps, the gentle floor wants big ones, and one learning rate can't satisfy both.
• How you estimate the slope. Using all the data is accurate but slow. Using a small random batch is noisy but much cheaper per step.

2. Visualization

Try it (predict first, then check):

1. Predict whether η = 0.19 converges on the narrow bowl. Step it and watch the zig-zag.
2. Press Compare all three and step 30 times. Which optimizer is closest to the minimum, and is that what you expected?

What to notice:

• At $\eta = 0.19$ the steep direction overshoots every step: $\theta_2$ flips sign and shrinks slowly, the classic zig-zag.
• At $\eta = 0.21$ each overshoot is larger than the last. The loss grows forever: divergence.
• Momentum builds speed along the valley floor but overshoots and oscillates; Adam takes steps of roughly equal size in each coordinate regardless of steepness.
• On this simple quadratic, plain gradient descent with a well-chosen learning rate is the fastest of the four. Fancier optimizers earn their keep on messier surfaces and when you can't tune the learning rate carefully.

3. The math

Symbols

The update rule

Why subtract? The first-order Taylor approximation says $L(\theta + \Delta) \approx L(\theta) + \nabla L \cdot \Delta$. Choosing $\Delta = -\eta \nabla L$ gives $L(\theta) - \eta \lVert \nabla L \rVert^2$, which is smaller than $L(\theta)$ for a small enough $\eta > 0$.

Batch, stochastic, mini-batch

The training loss is an average over $n$ examples, $L = \frac{1}{n}\sum_i \ell_i$, so its gradient is the average of per-example gradients. In practice we estimate it from a random mini-batch $B$:

• Batch gradient descent: $|B| = n$. Exact gradient, one expensive step per pass over the data.
• Stochastic gradient descent (SGD): $|B| = 1$. Very noisy, very cheap.
• Mini-batch: $|B|$ between about 32 and a few thousand. The default, because GPUs process a batch in parallel and the noise even helps escape sharp, poorly generalizing minima. One pass over the dataset is an epoch.

Momentum and Adam

Momentum keeps a running velocity, so consistent directions accelerate and alternating directions cancel:

Adam keeps a running mean of the gradient ($m$) and of its square ($s$), corrects both for starting at zero, and divides, so each parameter gets its own step size:

Here $g_t$ is the gradient at step $t$, squares are element-wise, and $\epsilon \approx 10^{-8}$ prevents division by zero.

Worked example

Loss $f(\theta_1, \theta_2) = \tfrac12(\theta_1^2 + 10\,\theta_2^2)$. Partial derivatives: $\partial f/\partial\theta_1 = \theta_1$ and $\partial f/\partial\theta_2 = 10\,\theta_2$. Start at $\theta_0 = (-3, 1.6)$, where $f = \tfrac12(9 + 10 \times 2.56) = \tfrac12(34.6) = 17.3$.

Gradient at the start: $\nabla f = (-3,\; 10 \times 1.6) = (-3,\; 16)$.

Learning rate 0.1:

The loss drops to $\tfrac12(7.29 + 0) = 3.645$. The steep coordinate hit zero in one step, because $1 - \eta \times 10 = 0$. From now on $\theta_1$ shrinks by a factor $1 - 0.1 = 0.9$ per step: $-2.43, -2.187, \dots$

Learning rate 0.19: the steep coordinate is multiplied each step by $1 - 0.19 \times 10 = -0.9$: $1.6 \to -1.44 \to 1.296 \to -1.166$. It flips sign every step but shrinks: zig-zag convergence. The loss goes $17.3 \to 13.32 \to 10.34 \to 8.07$.

Learning rate 0.21: the factor is $1 - 2.1 = -1.1$, so $1.6 \to -1.76 \to 1.936 \to -2.13$. The loss goes $17.3 \to 18.30 \to 20.49 \to 23.77$, up every step. Diverged.

The general rule: along a direction with curvature $\lambda$ (here 1 and 10), each step multiplies the error by $1 - \eta\lambda$. Convergence needs $|1 - \eta\lambda| < 1$, so $\eta < 2/\lambda_{\max} = 2/10 = 0.2$. The steepest direction sets the speed limit for every direction.

Momentum ($\eta = 0.02$, $\beta = 0.9$), two steps: $v_1 = (-3, 16)$, so $\theta_1 = (-3, 1.6) - 0.02(-3, 16) = (-2.94, 1.28)$. New gradient $(-2.94, 12.8)$, so $v_2 = 0.9(-3, 16) + (-2.94, 12.8) = (-5.64, 27.2)$ and $\theta_2 = (-2.94 + 0.1128,\; 1.28 - 0.544) = (-2.8272,\; 0.736)$. The velocity in $\theta_1$ grew from 3 to 5.64 because the gradient kept pointing the same way.

Adam ($\eta = 0.1$), first step: $m_1 = 0.1\,g$ and $s_1 = 0.001\,g^2$; after bias correction $\hat m = g$ and $\hat s = g^2$, so the step is $\eta\, g/|g| = 0.1 \times \text{sign}(g)$: $\theta_1 = (-2.9, 1.5)$. Both coordinates move by exactly 0.1, even though one gradient is 5 times larger. That per-parameter normalization is Adam's main idea.

4. Implementation

The full script is code/11-gradient-descent-backprop/optimizers_scratch.py (NumPy only). Each optimizer is a small class with one method, step(params, grads), which is the same interface PyTorch's optimizers expose. The gradient comes from outside; the optimizer only decides how to use it.

Build it from scratch

Step 1: plain gradient descent. No memory between steps. This is the update rule from section 3 and nothing else.

class SGD:
    def __init__(self, lr):
        self.lr = lr

    def step(self, p, g):
        return p - self.lr * g

Predict: on the bowl from the worked example, which lr makes the steep coordinate land exactly on 0 in one step?

Step 2: momentum. Keep a velocity v that adds each new gradient to a decayed copy of the old one. Where gradients keep pointing the same way, v grows; where they flip sign (a zig-zag), they cancel. This uses the course's convention $v_t = \beta v_{t-1} + g_t$.

class Momentum:
    def __init__(self, lr, beta=0.9):
        self.lr, self.beta, self.v = lr, beta, None

    def step(self, p, g):
        self.v = g if self.v is None else self.beta * self.v + g
        return p - self.lr * self.v

Predict: if the gradient were the same constant $g$ forever, what would v approach? (Sum the geometric series.)

Step 3: Adam. Two running averages: m of the gradient (momentum) and s of its square (its typical size). Dividing one by the square root of the other gives every parameter a step of about lr, whatever its gradient scale. Both averages start at 0, so early on they are too small; dividing by $1 - \beta^t$ corrects that.

    def step(self, p, g):
        if self.m is None:
            self.m, self.s = np.zeros_like(p), np.zeros_like(p)
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * g          # average gradient
        self.s = self.b2 * self.s + (1 - self.b2) * g ** 2     # average squared gradient
        m_hat = self.m / (1 - self.b1 ** self.t)               # undo the bias toward 0
        s_hat = self.s / (1 - self.b2 ** self.t)
        return p - self.lr * (m_hat / (np.sqrt(s_hat) + self.eps) + self.wd * p)

Predict: without the two bias-correction lines, is the first step bigger or smaller than lr? The script prints the answer.

Step 4: AdamW. The last line above already has a self.wd * p term. It shrinks weights directly, outside Adam's per-parameter scaling. That is decoupled weight decay, which is what AdamW means. Adding an L2 penalty to the loss instead would push the decay through g, so Adam would rescale it, and parameters with large gradients would barely be regularized.

def AdamW(lr, weight_decay=0.01, **kw):
    return Adam(lr, weight_decay=weight_decay, **kw)

Predict: with zero gradient, lr = 0.1, and weight_decay = 0.1, what is a weight of 1.0 after 10 steps?

Step 5: the training loop. Every optimizer plugs into the same loop: compute the gradient, take a step, record where you are.

def run(opt, grad, start, steps):
    p, path = np.array(start, float), [np.array(start, float)]
    for _ in range(steps):
        p = opt.step(p, grad(p))
        path.append(p.copy())
    return np.array(path)

Predict: on the Rosenbrock valley (a long, curved, flat-bottomed valley), which of the three reaches the minimum at (1, 1) within 3,000 steps?

Use the library

PyTorch splits the same two jobs: autograd fills .grad for every parameter, then the optimizer reads it and updates in place. torch.optim.AdamW is step 4. One convention differs: PyTorch's SGD momentum uses the same $v_t = \beta v_{t-1} + g_t$ as ours, while some textbooks use $(1-\beta) g_t$, which changes the effective learning rate by a factor of 10 at $\beta = 0.9$.

import torch

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
# or torch.optim.SGD(..., momentum=0.9)
for x_batch, y_batch in loader:           # mini-batches
    optimizer.zero_grad()                 # clear old gradients
    loss = torch.nn.functional.mse_loss(model(x_batch), y_batch)
    loss.backward()                       # compute gradients (backpropagation)
    optimizer.step()                      # θ ← θ − η · (update rule)

Run it

uv run code/11-gradient-descent-backprop/optimizers_scratch.py

Key lines of the output:

        GD lr=0.1: step 1 [-2.7  0. ], step 2 [-2.43  0.  ], f after 2 = 2.9525
 momentum lr=0.02: step 1 [-2.94  1.28], step 2 [-2.8272  0.736 ], f after 2 = 6.7050
      Adam lr=0.1: step 1 [-2.9  1.5], step 2 [-2.8001  1.4002], f after 2 = 13.7234
lr=0.21: f after 50 steps = 1.76e+05
raw m/sqrt(s) = 3.1623; corrected = 1.0000
     SGD: end [-1.2487  1.567 ], loss 5.06e+00
momentum: end [0.7811 0.6092], loss 4.80e-02
    Adam: end [1. 1.], loss 3.15e-11
after 10 steps: 0.9044 (= 0.99^10 = 0.9044)

The first three lines are the worked example in section 3. On Rosenbrock, the learning rate that keeps plain SGD stable is so small that it barely moves along the valley floor; momentum gets close, and Adam's per-parameter scaling reaches the minimum. The figure script for the bowl paths is code/11-gradient-descent-backprop/gradient_descent.py.

5. Engineering

Picking a learning rate. Try values on a log scale (1e-1, 3e-2, 1e-2, ...). Pick the largest one that trains smoothly, then reduce it during training with a schedule (step decay, cosine decay). Large models often use a short warmup where the learning rate ramps up from near zero. A learning-rate range test (increase it each batch, plot the loss) finds the cliff quickly.

Read the loss curve.

• Flat from the start: learning rate too small, a bug (no gradient flowing), or dead activations.
• Decreasing then exploding into NaN: learning rate too large, or exploding gradients. Lower $\eta$, add gradient clipping, check for unscaled inputs.
• Very noisy: batch too small or learning rate too high; average over more steps before judging.

Conditioning matters. Unscaled features create narrow valleys (curvature differs wildly by direction), which force a tiny learning rate. Standardizing inputs and using normalization layers makes the surface rounder. Adam partially compensates by giving each parameter its own step size.

Which optimizer? Adam (or AdamW, which handles weight decay correctly) is the default for transformers and most deep networks because it works with little tuning. SGD with momentum is still competitive for many vision models and sometimes generalizes better. For convex problems with few parameters (linear and logistic regression), second-order or quasi-Newton methods such as L-BFGS are often faster.

Cost. Each step costs one forward and one backward pass over the batch, roughly proportional to (batch size × parameters). Adam stores two extra numbers per parameter, so its state alone takes twice the memory of the weights, which matters at LLM scale.

Watch out

Failure modes. Divergence (loss becomes NaN), getting stuck on plateaus or saddle points, oscillating around the minimum with a constant learning rate, and convergence to a poor local minimum on non-convex losses. In deep networks, vanishing or exploding gradients can make some layers learn far faster than others (covered in backpropagation).

Common mistakes

• Judging a learning rate after a handful of noisy steps.
• Changing the batch size without adjusting the learning rate (larger batches usually allow a larger $\eta$).
• Forgetting optimizer.zero_grad(), so gradients accumulate across steps.
• Comparing optimizers each at its default learning rate instead of each at its own tuned one.
• Treating a lower training loss as success without checking validation loss.

6. Knowledge check

Practice exercise. For $f(\theta) = 2\theta^2$, starting at $\theta = 1$: compute three gradient-descent steps with $\eta = 0.1$, then find the largest learning rate that still converges.

Implementation challenge. Implement mini-batch SGD for the linear regression from Module 3 on a synthetic dataset of 10,000 points. Plot the training loss per step for batch sizes 1, 32, and 10,000 with the same number of examples processed, and explain the differences.

Summary

• Gradient descent repeats $\theta \leftarrow \theta - \eta \nabla L$: compute the slope, step downhill.
• Along a direction of curvature $\lambda$, the error is multiplied by $1 - \eta\lambda$ each step, so $\eta < 2/\lambda_{\max}$ for convergence; steep directions set the limit.
• Mini-batches trade gradient accuracy for many cheaper steps and are the standard.
• Momentum accumulates velocity along consistent directions; Adam scales each parameter's step by its gradient history.
• Read the loss curve: flat, exploding, and noisy curves each point to different fixes.

Next: Backpropagation

Related: Linear regression · Training and regularization · Forward pass

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Optimization, Optimizers.

Interview angle