Loss Functions and Optimizers (SGD, Adam)

You are building a classifier and you want it to make as few mistakes as possible. So the obvious thing to optimise is the number of mistakes. Count them, and tell the optimiser to make that count go down.

Try it and nothing happens. Not "it trains badly" — literally nothing happens. The weights do not move at all.

The reason is worth sitting with, because it explains why the field's loss functions look the way they do. Nudge a weight by 0.0001. Almost certainly, no prediction flips from right to wrong or wrong to right, so the mistake count is unchanged and the gradient is exactly zero. Nudge it enough and one prediction flips, and the count jumps by a whole unit — an infinite slope at a single point. The error count as a function of the weights is a staircase: flat everywhere, with vertical cliffs. Gradient descent on a staircase has nothing to descend.

A loss function is not a measure of how good the model is. It is a differentiable stand-in for how good the model is, chosen so that gradient descent has a slope to follow.

Accuracy is the thing you care about. Cross-entropy is the thing you optimise. They are correlated but not identical, and confusing them is the source of a good deal of confusion later on.

Regression losses, and the outlier that hijacks your model

For predicting a continuous number, the two standard candidates are squared error and absolute error:

They look interchangeable. They are not, and one bad data point shows why. Suppose you predict 3.0 for five houses whose true prices (in hundreds of thousands) are 3.0, 3.1, 2.9, 3.0 and 30.0. That last one is a data-entry error — someone typed an extra zero.

Under MSE, the single corrupted example supplies 99.3% of the gradient. The model will spend its capacity chasing a typo. Under MAE, it supplies 33% — still more than its share of one-fifth, but not a takeover.

That is the trade-off in one table. MSE punishes large errors disproportionately, which is right when large errors really are disproportionately bad and your data is clean. MAE treats all errors in proportion, which is right when outliers are noise rather than signal. MAE's weakness is its constant gradient: as the model gets close to correct, the gradient stays at $\pm 1$ instead of shrinking, so the final convergence is jittery.

Huber loss takes both: quadratic near zero (smooth convergence), linear far away (outlier resistance).

1import torch.nn as nn23nn.MSELoss()                        # squared error4nn.L1Loss()                         # absolute error5nn.HuberLoss(delta=1.0)             # quadratic inside delta, linear outside

One practical note that catches people out: if your target is house price in pounds, MSE values will be in the order of $10^{10}$ and the gradients will be enormous. Either scale the target (predict price in units of £100,000, or predict $\log(\text{price})$) or expect to use a much smaller learning rate. The loss value being large is not itself a problem; the gradients being large is.

Classification losses, and why squared error is wrong here

Cross-entropy measures the distance between two probability distributions — what the model believes and what is actually true:

For single-label classification only one $y_c$ is 1, so this collapses to $-\log(\text{probability assigned to the correct class})$. Read that as: how surprised was the model by the right answer?

That $\ln$ curve is the point: the penalty grows without bound as the model becomes confidently wrong. A model that says "99% certain" and is wrong should be punished far more than one that says "51% certain" and is wrong, and cross-entropy encodes exactly that.

Now the reason not to use MSE on classification. Suppose the true label is 1 and a sigmoid output currently says 0.01 — as wrong as it is possible to be. What gradient reaches the pre-activation $z$?

MSE's gradient is 50 times smaller — and it is smallest exactly when the model is most wrong, because $\sigma'$ has collapsed in the saturated region. A model using MSE on a sigmoid output learns most slowly from its worst mistakes. Cross-entropy's $\sigma'$ term cancels algebraically against the $\log$, leaving the beautifully simple $\hat{y} - y$: the gradient is just the error, and it is largest when the model is most wrong.

Picking the right one

For heavy imbalance — fraud at 0.1% of transactions, say — plain cross-entropy lets the 99.9% majority dominate the average, and the model learns to predict "not fraud" always. Class weights fix this bluntly by multiplying each class's loss by a factor. Focal loss, $-(1-\hat{y})^{\gamma}\log \hat{y}$, does it more surgically: the $(1-\hat{y})^\gamma$ factor shrinks the contribution of examples the model already gets right, so training concentrates on the hard cases.

Gradient descent, and the one hyperparameter that matters most

Given a loss and its gradient, the update rule is one line:

Everything hinges on $\eta$, the learning rate. Take the simplest possible loss, $L(w) = w^2$, whose gradient is $2w$, starting at $w = 1$. The update becomes $w \leftarrow w(1 - 2\eta)$, so behaviour depends entirely on the multiplier $(1-2\eta)$:

Two orders of magnitude in $\eta$ separate "too slow to be useful" from "explodes". This is why learning rate is the first hyperparameter to tune and usually the only one that matters much.

How much data per step

The noise in mini-batch gradients is not purely a cost. It acts as a mild regulariser, letting the optimiser rattle out of narrow, sharp minima that tend to generalise poorly. Very large batches produce cleaner gradients and often worse test accuracy for exactly this reason.

Momentum: remembering which way you were going

Plain gradient descent has a characteristic pathology. In a valley that is steep across and shallow along — which describes most loss landscapes — the gradient points mostly across the valley, so the optimiser bounces from wall to wall while creeping slowly towards the minimum.

Momentum accumulates a running average of past gradients:

The side-to-side components alternate in sign and cancel out in the average. The along-the-valley component has a consistent sign and accumulates. With $\beta = 0.9$ the effective step in a consistent direction grows to roughly $1/(1-\beta) = 10$ times the plain gradient step. That is where the speed-up comes from.

Adam: a learning rate for every parameter

One global learning rate is a compromise. Some parameters — the weights on a rare feature — see gradients only occasionally and need large steps; others see large gradients constantly and need small ones. Adam maintains per-parameter statistics and scales each step accordingly.

1import numpy as np23def adam_step(w, grad, state, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):4    state["t"] += 15    t = state["t"]67    # 1st moment: running mean of the gradient (direction, i.e. momentum)8    state["m"] = b1 * state["m"] + (1 - b1) * grad9    # 2nd moment: running mean of the SQUARED gradient (magnitude)10    state["v"] = b2 * state["v"] + (1 - b2) * grad ** 21112    # Bias correction: both averages start at 0 and are far too small early on13    m_hat = state["m"] / (1 - b1 ** t)14    v_hat = state["v"] / (1 - b2 ** t)1516    # Parameters with historically large gradients get smaller steps17    return w - lr * m_hat / (np.sqrt(v_hat) + eps)

The bias correction deserves a concrete look, because it is the part people skip and then wonder why their from-scratch Adam trains oddly. At step 1, with $m$ initialised to zero and $\beta_1 = 0.9$:

The estimate of the average gradient is ten times too small, purely because the average started at zero. Dividing by $(1 - \beta_1^1) = 0.1$ restores it to $g$. By step 100, $\beta_1^{100} \approx 0.00003$ and the correction is negligible — it only matters for the first hundred or so steps, which are precisely the steps that set the trajectory of the whole run.

AdamW, and a weight-decay bug that lasted years

Adding L2 regularisation by tacking $\lambda w$ onto the gradient works fine with plain SGD. With Adam it does not, because that penalty term then gets divided by $\sqrt{\hat{v}}$ along with everything else — so parameters with large gradients receive less regularisation, which is the opposite of the intent. AdamW decouples the two: it applies the decay directly to the weights, outside the adaptive scaling.

1import torch23# Wrong for Adam: weight_decay is folded into the gradient4opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-2)56# Right: decoupled decay7opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)

If you are using weight decay at all, use AdamW. It is a one-word change and it reliably improves generalisation.

Which optimiser, and what to set the learning rate to

A defensible default: AdamW at $10^{-3}$, betas left alone, weight decay 0.01. Get a model training, get a baseline number, and only then consider whether a tuned SGD-with-momentum schedule is worth the effort. It sometimes reaches a slightly better final score on vision tasks, but it takes real tuning to get there and Adam gets you 95% of the way with none.

Schedules

A constant learning rate is a compromise between "large enough to make progress early" and "small enough to settle at the end". Schedules stop compromising:

1from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau23# Smoothly anneal from the initial LR down to ~0 over the whole run.4# Currently the most reliable general-purpose choice.5sched = CosineAnnealingLR(opt, T_max=num_epochs)67# Or react to the data: cut the LR by 10x whenever validation loss stalls.8sched = ReduceLROnPlateau(opt, mode="min", factor=0.1, patience=5)910for epoch in range(num_epochs):11    train_one_epoch(model, opt)12    val = validate(model)13    sched.step()          # for CosineAnnealingLR14    # sched.step(val)     # for ReduceLROnPlateau -- it needs the metric

A characteristic sight in real training curves is a sudden drop in loss exactly when the schedule cuts the learning rate. That is not the model suddenly understanding something; it is the optimiser finally being able to settle into a minimum it had been orbiting.

Reading a training run through this lens

When training goes wrong, the loss curve tells you which of these components to suspect, and the diagnosis is usually fast.

Loss is nan within a few steps. Almost always the learning rate, occasionally a $\log(0)$ from applying softmax twice. Divide the learning rate by ten and see if it survives. If it does, you found it.

Loss falls then plateaus far above zero. Either the model lacks capacity, or the learning rate is now too large for the region you are in. Add a schedule before you add layers — it is cheaper and it is more often the answer.

Loss oscillates without trending down. Learning rate too high, or batch size too small so the gradient estimate is dominated by noise. Try halving the learning rate first.

Loss barely moves from its initial value. Learning rate too small by orders of magnitude, or inputs not normalised, or — the embarrassing one — the optimiser is holding a different set of tensors from the ones used in the forward pass: it was built for an earlier copy of the model, or the layers live in a plain Python list and were never registered.

Training loss looks perfect and the model is useless in production. The loss was never the thing you cared about. Check the metric you actually care about — precision at a fixed recall, say — alongside the loss on every validation pass. A model can drive cross-entropy down beautifully by becoming very confident about the majority class while being worthless on the minority class you actually built it for.