Overview

Training a neural network is an optimization problem. We have a loss function \(J(\theta)\) that measures how wrong the model is, and we look for the parameters \(\theta^\star\) that make it as small as possible:

For linear regression this has a closed-form solution, and it is worth seeing why. With a linear model and squared error, \(J\) is a quadratic function of \(\theta\), so its gradient is linear in \(\theta\) — and asking where the gradient vanishes is then just a system of linear equations, which we know how to solve exactly:

Put a non-linear activation in the middle and that collapses: \(\nabla J = 0\) becomes a system of non-linear equations, with no formula to isolate \(\theta\). And even when a formula exists it asks for a matrix inversion — around \(O(n^3)\) work for \(n\) parameters, hopeless when \(n\) is in the millions. So we do the only thing left: start somewhere and walk downhill.

The gradient points uphill

For a function of several variables, the gradient is the vector of partial derivatives:

Each entry answers a local question: if I increase this one parameter a little, how much does the loss change? Put together, \(\nabla J\) points in the direction of steepest ascent, and its magnitude says how steep the slope is. Since we want to go down, we move against it — hence the minus sign in the update rule:

where:

• \(\theta\) are the model parameters,
• \(\eta\) is the learning rate, a hyperparameter controlling the step size,
• \(\nabla J(\theta)\) is the gradient of the loss with respect to the parameters.

A gradient-based method is any algorithm that finds minima of a function assuming the gradient is cheap to compute. It requires the function to be continuous and differentiable almost everywhere — ReLU has a kink at zero and that is fine¹. Everything in this chapter is a variation on the update rule above: the same skeleton, with smarter ways of choosing the direction and the step size.

The learning rate: the one hyperparameter you must feel

Take the simplest possible loss, \(J(\theta) = \theta^2\), whose gradient is \(2\theta\). The update becomes

The whole behaviour of gradient descent is inside that factor \((1 - 2\eta)\). Play with the learning rate below and watch four qualitatively different worlds appear.

The transition points in that simulator are not arbitrary. For a quadratic whose curvature (second derivative) is \(L\) — the letter is for Lipschitz³ — gradient descent converges only if

Here \(L = 2\), so the boundary sits exactly at \(\eta = 1\) — try it. Below \(1/L = 0.5\) the approach is monotone; between \(1/L\) and \(2/L\) it zig-zags across the minimum; above \(2/L\) it explodes. Real loss surfaces are not quadratic, but near a minimum they look quadratic, and this bound is the reason a learning rate that worked yesterday makes today's model diverge after you changed the architecture: you changed \(L\).

How much data per step? Batch, Stochastic and Mini-Batch

The loss over a dataset of \(N\) samples is an average of per-sample losses:

Computing \(\nabla J\) exactly means a full pass over the data for one parameter update. The three classic variants differ only in how many samples they look at before stepping.

for epoch in range(num_epochs):
    gradients = compute_gradients(X, y, model)      # one pass over everything
    model.parameters -= learning_rate * gradients   # one single update

for epoch in range(num_epochs):
    for i in random_permutation(len(X)):            # shuffle every epoch
        gradients = compute_gradients(X[i], y[i], model)
        model.parameters -= learning_rate * gradients

for epoch in range(num_epochs):
    for batch in get_mini_batches(X, y, batch_size):
        gradients = compute_gradients(batch.X, batch.y, model)
        model.parameters -= learning_rate * gradients

The key quantitative fact: averaging \(B\) independent samples divides the standard deviation of the gradient estimate by \(\sqrt{B}\).

That square root is why batch size has diminishing returns: going from 32 to 128 costs 4× the compute per step and only halves the noise. And the noise never fully disappears — with a constant learning rate, SGD does not converge to the minimum, it converges to a cloud around the minimum whose radius grows with \(\eta\) and shrinks with \(\sqrt{B}\). Watch it happen:

Why plain gradient descent struggles

If the loss surface were a nice round bowl, a single learning rate would do and this chapter would end here. Real surfaces are not round, and three specific pathologies explain every optimizer that follows.

Momentum: give the ball some weight

Instead of stepping in the direction of the current gradient, accumulate an exponential moving average of past gradients and step in that direction. Two state variables now, \(V\) and \(\theta\):

\(V\) is a running average of the gradients: at each step we dampen the old value by \(\beta\) (between 0 and 1) and mix in the new gradient. We then move \(\theta\) along the new momentum \(V\)¹.

Why it fixes ravines. Across the valley the gradient flips sign every step, so consecutive terms cancel in the average and the oscillation is damped. Along the valley the sign never changes, so nothing cancels and the average settles on the gradient itself. One rule, two opposite effects on the two directions — that is the whole trick.

How much history is being averaged is set by

the effective number of past gradients in the moving average: about 10 steps at \(\beta = 0.9\), about 100 at \(\beta = 0.99\). Read it as a memory length, not a speed-up — in the form written above, once \(V\) has settled on a constant gradient the step is \(\eta V = \eta g\), exactly SGD's step. (In fact it arrives late: over a long run momentum trails plain SGD by \(\beta/(1-\beta)\) steps' worth of progress, 9 of them at \(\beta = 0.9\).)

What momentum buys, then, is not a bigger step but a better direction: what alternates cancels, what persists survives. It also keeps moving on stale gradients — which is why it coasts through small bumps, and why it can sail past a minimum instead of braking at it.

Nesterov: look before you leap

Nesterov Accelerated Gradient⁹¹⁰ evaluates the gradient after the momentum step, not before:

The idea is a correction term: since we already know inertia is going to carry us to \(\theta_t - \eta\beta V_t\), we may as well measure the slope there. If the surface has started to climb again, the gradient at the look-ahead point brakes earlier than it would otherwise. In practice it gives a small but consistent improvement¹¹, and it is one flag away: torch.optim.SGD(..., momentum=0.9, nesterov=True).

Adaptive methods: one learning rate per parameter

Momentum fixes the direction. It does not fix the fact that a single \(\eta\) has to serve every parameter in the network — the embedding of a rare word, whose gradient is almost always zero, and a bias in the first layer, whose gradient is large at every step.

The shared idea of every adaptive method: keep a per-parameter estimate of the typical gradient magnitude, and divide by it. Parameters with consistently large gradients take smaller steps; parameters with tiny gradients take larger ones. The update becomes roughly scale-free.

AdaGrad — the first idea, and its flaw

Accumulate the sum of all squared gradients seen so far and divide by its square root¹²:

This works beautifully for sparse features and it is where the whole family starts. But \(G_t\) is a sum that never decreases, so the effective learning rate \(\eta/\sqrt{G_t}\) decays monotonically toward zero. On a long deep-learning run, AdaGrad stops learning long before it has converged.

RMSProp — replace the sum by a moving average

The fix¹³ is one character wide: turn the sum into an exponential moving average, so old gradients are forgotten.

where:

• \(V_t\) is the moving average of the squared gradients (typically \(\beta = 0.9\); PyTorch calls it alpha and defaults to 0.99),
• \(\epsilon \approx 10^{-8}\) is a small constant for numerical stability, preventing division by zero.

Now \(\sqrt{V_t}\) is the root mean square of the recent gradients — hence the name — and it tracks the current terrain instead of the whole history. Note that \(g/\sqrt{\overline{g^2}}\) is dimensionless and roughly \(\pm 1\): RMSProp is close to "take a step of size \(\eta\) in the sign of the gradient", which is precisely what makes it immune to a badly scaled loss.

Adam — momentum and RMSProp in one

Adaptive Moment Estimation keeps both moving averages: the first moment (the mean of the gradients — momentum) and the second moment (the uncentered variance — RMSProp).

1. Compute the gradient:

   \[g_t = \nabla J(\theta_t)\]

2. First moment (direction, with inertia):

   \[m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t\]

3. Second moment (scale):

   \[v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2\]

4. Bias correction:

   \[\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}\]

5. Update:

   \[\theta_{t+1} = \theta_t - \eta \, \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}\]

where \(\beta_1, \beta_2\) control the decay of the two averages and \(\epsilon\) guards the division. The defaults are \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\)¹⁴¹⁵, and they are remarkably robust — Adam is the closest thing we have to an optimizer that works out of the box.

Adam is the default for training deep networks: it combines momentum's smoothing with per-parameter adaptivity, and it needs far less learning-rate tuning than SGD. The price is memory — two extra tensors the size of the model — and a well-documented tendency to generalize slightly worse than well-tuned SGD+momentum on convolutional vision tasks.¹⁶

AdamW — decoupled weight decay

AdamW¹⁷ separates weight decay from the adaptive update, which is mathematically the correct thing to do:

where \(\lambda\) is the weight-decay coefficient applied directly to the weights, not to the gradient.

In the original Adam, L2 regularization is implemented by adding \(\lambda\theta\) to the gradient before the adaptive scaling — so it gets divided by \(\sqrt{\hat{v}_t}\) too. The consequence is perverse: parameters with large gradients end up being regularized less, which is the opposite of what you want. AdamW applies the decay outside the adaptive term, restoring the intended behaviour. It is the standard for training Transformers and LLMs (BERT, GPT, LLaMA), usually with \(\lambda \approx 0.01\)–\(0.1\).

Simulator: the optimizer race

Everything above, side by side. Pick a landscape, set the learning rate, and watch five optimizers start from the same point. The bottom panel shows the loss on a logarithmic scale — the shape of those curves is what you actually see in a training log.

Learning rate schedules

The noise simulator showed why a constant learning rate cannot be optimal: big steps explore fast but leave you orbiting the minimum; small steps settle precisely but take forever to get there. So do both — start big, end small.

Two extra ingredients matter in practice. Warmup ramps \(\eta\) up from ~0 over the first few hundred/thousand steps: at initialization the gradients are large and Adam's second moment estimate is still garbage, so full-size steps early can wreck the model — this is essential for Transformers¹⁸. And cosine decay brings \(\eta\) smoothly to zero by the end of training, which empirically beats step decay on almost everything.¹⁹

import torch, math

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

warmup, total = 500, 10_000
def factor(step):                                   # returns a multiplier on lr
    if step < warmup:
        return (step + 1) / warmup                  # linear ramp up (never exactly 0)
    p = (step - warmup) / (total - warmup)          # progress in [0, 1]
    return 0.5 * (1 + math.cos(math.pi * p))        # cosine down to 0

sched = torch.optim.lr_scheduler.LambdaLR(opt, factor)

for x, y in loader:
    loss = criterion(model(x), y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # see below
    opt.step()
    sched.step()                                    # per step, not per epoch

Choosing an optimizer in practice

Side by side

In short: Batch GD is exact but impractical; SGD buys speed with noise; Momentum smooths that noise and accelerates along consistent directions; RMSProp fixes the scale of each parameter; and Adam/AdamW combines the last two, which is why it is the default — although a well-tuned SGD+momentum still generalizes better on some vision benchmarks.

The right choice depends on the size of the dataset, the compute available and how much tuning you can afford — validate on held-out data rather than trusting a table.

Sensible defaults

Troubleshooting a training run

The optimizers, from scratch

All of them are ten lines. Written side by side, the family resemblance is obvious — every one is θ -= η · (direction) / (scale).

import numpy as np

class SGD:
    def __init__(self, lr=0.01, beta=0.0):
        self.lr, self.beta, self.v = lr, beta, None

    def step(self, theta, grad):
        if self.v is None:
            self.v = np.zeros_like(theta)
        self.v = self.beta * self.v + (1 - self.beta) * grad     # momentum (beta=0 → plain SGD)
        return theta - self.lr * self.v


class RMSProp:
    def __init__(self, lr=0.01, beta=0.9, eps=1e-8):
        self.lr, self.beta, self.eps, self.v = lr, beta, eps, None

    def step(self, theta, grad):
        if self.v is None:
            self.v = np.zeros_like(theta)
        self.v = self.beta * self.v + (1 - self.beta) * grad ** 2   # average squared magnitude
        return theta - self.lr * grad / (np.sqrt(self.v) + self.eps)


class Adam:
    def __init__(self, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.0):
        self.lr, self.b1, self.b2, self.eps, self.wd = lr, b1, b2, eps, weight_decay
        self.m = self.v = None
        self.t = 0

    def step(self, theta, grad):
        if self.m is None:
            self.m = np.zeros_like(theta)
            self.v = np.zeros_like(theta)
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * grad            # 1st moment: direction
        self.v = self.b2 * self.v + (1 - self.b2) * grad ** 2       # 2nd moment: scale
        m_hat = self.m / (1 - self.b1 ** self.t)                    # bias correction
        v_hat = self.v / (1 - self.b2 ** self.t)
        update = m_hat / (np.sqrt(v_hat) + self.eps)
        return theta - self.lr * (update + self.wd * theta)         # AdamW: decay outside the scaling

# The same thing, in PyTorch
torch.optim.SGD(p, lr=0.1, momentum=0.9, nesterov=True, weight_decay=5e-4)
torch.optim.RMSprop(p, lr=1e-3, alpha=0.9)
torch.optim.Adam(p, lr=1e-3, betas=(0.9, 0.999), eps=1e-8)
torch.optim.AdamW(p, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)

Key takeaways

1. Gradient descent is one line; everything else is about choosing the step size well.
2. A quadratic converges only when \(\eta < 2/L\). Divergence is not bad luck, it is arithmetic.
3. Mini-batches trade noise for speed, and the noise falls only as \(1/\sqrt{B}\). A constant learning rate leaves you orbiting the minimum — hence schedules.
4. The real obstacles are ill-conditioning and saddle points, not bad local minima.
5. Momentum averages gradients over time, so what alternates cancels and what persists survives. \(1/(1-\beta)\) is the length of that memory — and, in PyTorch's convention, also the factor by which the step grows.
6. Adaptive methods (AdaGrad → RMSProp → Adam) give every parameter its own step size by dividing by the RMS of recent gradients. Adam adds momentum and bias correction on top.
7. AdamW + warmup + cosine is the modern default. Tune the learning rate first, and almost nothing else.

Additional Resources

Three things that carry on from where this page stops, in the order that makes sense:

1. Who's Adam and What's He Optimizing? — Kundu, S. Walks the same road as this chapter — SGD, momentum, RMSProp, Adam — in animation. The fastest way to revise the whole page:

2. An overview of gradient descent optimization algorithms — Ruder, S. Gathers every method here into one notation, plus the ones left out (Adadelta, Nadam, AMSGrad). Read it once the chapter is closed, to see the full map.

3. Why Momentum Really Works — Goh, G., Distill (2017). An entirely interactive article on what the condition number does to convergence. The natural next step after the momentum section: it proves with sliders what is left as intuition here.

References

The works cited through the text, in order of appearance: