Gradient Descent, in depth: the update rule from first principles, a worked example by hand, learning rates, batches, momentum and debugging training

The idea from first principles

Training defines a loss function L(theta): a single number that measures how badly a model with parameters theta fits the data. Mean squared error for regression and cross-entropy for classification are the usual choices; the cross-entropy derivation shows where the latter comes from. The goal is to find parameters with a low loss.

The gradient, written g = grad L(theta), is the vector of partial derivatives, one per parameter. Each entry says how fast the loss changes if you nudge that one parameter. Taken together, the gradient points in the direction in which the loss increases fastest. A first-order Taylor expansion makes the consequence precise: for a small step d, L(theta + d) is approximately L(theta) + g . d. Choosing d = -lr * g makes the change approximately -lr * |g|^2, which is negative whenever the gradient is not zero. So the update

theta_next = theta - lr * grad L(theta)

reduces the loss, as long as the learning rate lr is small enough that the linear approximation holds. Everything interesting about gradient descent lives in that qualifier. Too small and training takes forever; too large and the step overshoots into a region where the approximation was wrong and the loss goes up.

A worked example by hand

Fit a one-parameter model y = w * x to three points, (1, 2), (2, 4) and (3, 6). The right answer is obviously w = 2. The mean squared error is L(w) = (1/3) * sum of (w * x - y)^2. Since y = 2x here, each term is x^2 (w - 2)^2, so the derivative is

dL/dw = (2/3) * sum(x^2) * (w - 2) = (2/3) * 14 * (w - 2) = (28/3) * (w - 2)

Start at w = 0 with a learning rate of 0.05. The gradient is (28/3) * (-2) = -18.67, so the update moves w to 0 - 0.05 * (-18.67) = 0.933. At w = 0.933 the gradient is (28/3) * (-1.067) = -9.96, and the next w is 1.431. Then 1.697, then 1.838. The error w - 2 goes -2, -1.067, -0.569, -0.303, -0.162: each step multiplies it by the same factor.

That factor is 1 - lr * 28/3 = 1 - 0.467 = 0.533. In general the error is multiplied by (1 - lr * h) each step, where h = 28/3 is the second derivative, the curvature of the loss. Three regimes follow. If lr * h is between 0 and 1, the error shrinks steadily. If it is between 1 and 2, the error flips sign each step but still shrinks: the iterate overshoots and bounces in. If it is above 2, the factor's magnitude exceeds one and the error grows without bound. Here that threshold is lr = 2 / h = 6/28, about 0.214. Try lr = 0.25: the factor is -1.33, and w goes 0, 4.67, -1.56, 6.74, growing more violently each step. A learning rate just 17 percent above the threshold turns convergence into divergence.

The best learning rate for this problem is lr = 1/h, about 0.107, which reaches the answer in one step. The general lesson is the one to remember: the curvature of the loss sets the largest stable learning rate, and you usually do not know the curvature in advance, which is why learning rates are tuned.

Many parameters: curvature in every direction

With many parameters the loss has a curvature in every direction, captured by the eigenvalues of the Hessian matrix of second derivatives. The one-dimensional analysis applies along each eigenvector separately. Stability requires lr below 2 divided by the largest eigenvalue, while progress along the flattest direction shrinks the error by only (1 - lr * smallest eigenvalue) per step. The ratio of the two, the condition number, therefore controls how slow gradient descent is. A long, narrow valley is a badly conditioned problem: the steep walls force a small step, the gentle floor needs a big one, and the iterate zigzags.

This is the practical reason for feature scaling. If one input is measured in metres and another in millimetres, their weights see curvatures that differ by a factor of a million, and plain gradient descent crawls. Standardising inputs to zero mean and unit variance, and normalisation layers inside deep networks, make the surface rounder so one learning rate suits every direction. Adaptive optimizers such as Adam go further by scaling each parameter's step by a running estimate of its gradient magnitude, which is a cheap, diagonal correction for conditioning.

A NumPy implementation of full-batch gradient descent for linear regression makes the loop concrete. Because the inputs are already standardised, a learning rate of 0.1 converges in a few dozen steps:

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(1000, 3))
true_w, true_b = np.array([2.0, -1.0, 0.5]), 0.3
y = X @ true_w + true_b + rng.normal(scale=0.1, size=1000)

w, b, lr = np.zeros(3), 0.0, 0.1
for step in range(200):
    err = X @ w + b - y                  # forward: residuals
    loss = np.mean(err ** 2)             # mean squared error
    grad_w = 2 * X.T @ err / len(y)      # dL/dw
    grad_b = 2 * err.mean()              # dL/db
    w -= lr * grad_w                     # update
    b -= lr * grad_b
    if step % 50 == 0:
        print(step, round(loss, 4), np.linalg.norm(grad_w).round(4))
print(w.round(3), round(b, 3))           # close to [2, -1, 0.5] and 0.3

Batch, mini-batch and stochastic

The loss over a dataset is an average over examples, so its gradient is an average of per-example gradients. Computing it over all examples, full-batch gradient descent, gives the exact direction but costs a full pass over the data per step. Stochastic gradient descent uses one example per step: very cheap, very noisy. Mini-batch gradient descent, which is what everyone actually runs, averages over a batch of, say, 32 to a few thousand examples. The batch gradient is an unbiased estimate of the full one, and its noise falls as one over the square root of the batch size.

Batch size is partly a hardware decision. GPUs reach good utilisation only when each matrix multiplication is large enough, so batches are made as big as memory and throughput allow, and gradient accumulation sums several small batches before one update when memory is the limit. The noise is not purely a cost: it helps escape saddle points and sharp minima, and very large batches often need more careful tuning to generalise as well. When you multiply the batch size by k, a common starting heuristic, from Goyal and colleagues' 2017 work on large-batch ImageNet training, is to multiply the learning rate by k and add a warmup period; treat it as a starting point to tune from, not a law.

Momentum and its descendants

Momentum keeps a running velocity, an exponentially weighted sum of past gradients, and steps along it:

v = beta * v + g
theta = theta - lr * v          # beta is typically 0.9

Along the floor of a valley successive gradients agree, so velocity builds up and progress accelerates by up to a factor of 1 / (1 - beta), ten for beta = 0.9. Across the valley they alternate in sign and cancel, damping the zigzag. Nesterov momentum evaluates the gradient at the point the velocity is about to carry you to, which reacts to curvature slightly earlier. Adam combines momentum with per-parameter scaling, and AdamW decouples weight decay from that scaling; the SGD vs Adam vs AdamW article derives each update and its memory cost. Plain SGD with momentum remains competitive for many vision models; Adam-family optimizers are the default for transformers.

Learning-rate schedules

A single fixed learning rate is a compromise: large steps make fast early progress, but near a minimum the gradient noise from mini-batches keeps a large-step iterate bouncing around instead of settling. Schedules change the rate during training. Warmup ramps it from near zero over the first few hundred or thousand steps, which protects the early phase when parameters and adaptive optimizer statistics are poorly scaled; it matters most for transformers and large batches. Decay then reduces it, by steps at fixed epochs, by a cosine curve to near zero, or linearly. The one-cycle policy used in the PyTorch loop below combines both.

To find a starting point, run a learning-rate range test: train for a few hundred steps while increasing the rate exponentially, plot loss against rate, and choose a value somewhat below where the loss stops falling and starts to climb. That empirically locates the curvature limit the worked example derived analytically. Systematic search over learning rate, batch size and schedule is covered in the hyperparameter optimisation article.

A production training loop

The PyTorch loop below has the pieces a real run needs: zeroing gradients, which otherwise accumulate across steps, gradient-norm clipping to bound the size of any single update, a check for non-finite loss, and a per-step schedule. It also logs the learning rate and gradient norm alongside the loss, because those are the numbers you will need when something goes wrong.

import torch

model = MyModel().cuda()
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=0.1, total_steps=EPOCHS * len(loader))
loss_fn = torch.nn.CrossEntropyLoss()

for epoch in range(EPOCHS):
    model.train()
    for x, y in loader:
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        opt.zero_grad(set_to_none=True)          # gradients accumulate otherwise
        loss = loss_fn(model(x), y)              # forward
        loss.backward()                          # backward: fills p.grad
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        if not torch.isfinite(loss):
            raise RuntimeError(f"non-finite loss at epoch {epoch}")
        opt.step()                               # update
        sched.step()                             # per-step schedule
    log(epoch=epoch, loss=loss.item(), grad_norm=gnorm.item(), lr=sched.get_last_lr()[0])

In mixed precision training a gradient scaler multiplies the loss before the backward pass and unscales before clipping and stepping, so small gradients do not underflow in 16-bit formats; the mixed precision article explains the mechanics. Around the loop, the training pipeline article covers data loading, checkpointing and evaluation.

Diagnostics worth logging

• Training and validation loss per step or epoch. Smooth descent is healthy; a flat line from step one points at a learning rate far too small, frozen parameters or a bug in the data; a rising validation loss with falling training loss is overfitting, not an optimisation problem.
• Global gradient norm. Sudden spikes precede loss spikes; a norm stuck near zero means vanishing gradients or a disconnected graph. If clipping fires on nearly every step, the threshold or the learning rate is wrong.
• Update-to-weight ratio. The norm of lr * update divided by the norm of the weights, per layer. A widely used rule of thumb puts healthy values around one in a thousand; far above means steps are too big for that layer, far below means it is barely learning.
• The learning rate actually applied. Schedulers are easy to misconfigure, for example stepping per epoch when the schedule was defined per step.

Failure modes

What to do next

1. Reproduce the worked example in a few lines of Python with lr = 0.05, 0.107, 0.2 and 0.25, and watch the three regimes.
2. Run the NumPy loop, then remove the input standardisation by scaling one feature by 1,000 and see how far the learning rate must drop.
3. In your own training loop, log loss, global gradient norm, the applied learning rate and per-layer update-to-weight ratios.
4. Run a learning-rate range test before your next real training run and pick the rate from the plot.
5. Add warmup plus cosine or one-cycle decay, and gradient clipping, as defaults.
6. Read the optimizer derivations next, so you know what momentum and Adam change about the basic update.