12.5.4 — Making Training Work

The training loop of 12.5.3 is five lines. Getting a deep network to actually learn with it is another matter. Cartwise's first attempt at a deeper photo checker, twenty layers instead of two, did one of three things depending on the run: the loss sat flat and never moved, or it shot to nan within a hundred steps, or it fell beautifully on the training photos while the validation score got worse. None of these were bugs in the loop.

They were the three classic failures of deep learning, and each has a well-understood cause:

1. The signal dies or explodes as it passes through many layers, so nothing learns or everything breaks. The cures are careful starting weights, normalisation, residual connections and gradient clipping.
2. The network memorises the training data. The cures are dropout, weight decay, early stopping and augmentation.
3. Something is simply wrong, and the loss gives you only vague clues. The cure is a debugging routine.

All of the cures are small, most are a single line of code, and together they are the difference between deep learning working and not.

1. Starting weights: neither too small nor too large

Before any training, every weight has to start somewhere. Two tempting choices are both wrong, and seeing why explains what the right choice must do.

All zeros fails completely. If every weight in a layer starts at zero, every neuron in that layer computes the same output, receives exactly the same gradient in backpropagation, and makes exactly the same update. They stay identical forever. A layer of 128 neurons behaves like one neuron copied 128 times. Starting values must be random, to break this symmetry.

Random, but how big? Think about what one layer does to the size of the numbers passing through it. A neuron adds up n inputs, each multiplied by a weight. If the inputs have a typical size (variance) of v and the weights have variance \sigma^2, the sum of n such products has variance about n \sigma^2 v. So each layer multiplies the variance of the signal by roughly n \sigma^2. A ReLU then zeroes about half the values, which halves it again. Per layer, the signal's variance is multiplied by about:

\text{factor} = \tfrac{1}{2}\, n\, \sigma^2

Now take a network with 512 neurons per layer and try three weight scales:

Weight standard deviation σ	Factor per layer	After 10 layers
0.01	\tfrac{1}{2} \times 512 \times 0.0001 \approx 0.026	about 10^{-16}: the signal has vanished
0.1	\tfrac{1}{2} \times 512 \times 0.01 = 2.56	about 12,000: the signal has exploded
\sqrt{2/512} \approx 0.0625	exactly 1	unchanged

The lesson is in the last row. To keep the signal the same size through any number of layers, choose the weights' variance so that the factor is exactly 1: \sigma^2 = 2/n. This is He initialisation, published by Kaiming He and colleagues in 2015 for ReLU networks. The earlier version for tanh and sigmoid networks, Xavier (or Glorot) initialisation, by Xavier Glorot and Yoshua Bengio in 2010, uses \sigma^2 = 1/n because those activations do not zero half their inputs. The same argument applies in reverse to gradients flowing backwards, which is why it matters for learning too.

You rarely set this yourself: PyTorch's layers start with a sensible scheme already. But it explains why deep networks were so hard to train before 2010, and why a network you have built by hand from unusual pieces can mysteriously refuse to learn.

2. Normalisation: re-centre the signal at every layer

Good starting weights keep the signal well behaved at the start of training. But as the weights change during training, the size and centre of each layer's outputs drift, and every later layer has to keep adapting to inputs whose scale keeps moving. Normalisation layers fix the scale at every layer, all the time, by applying the standardisation of 12.2.4 inside the network: subtract the mean, divide by the standard deviation.

Layer normalisation

Layer normalisation, from Jimmy Ba, Jamie Kiros and Geoffrey Hinton in 2016, normalises across the values of one example. Take a layer's output for one photo, \mathbf{x} = (2, 4, 6, 8):

1. The mean is (2 + 4 + 6 + 8)/4 = 5.
2. The variance is \big((-3)^2 + (-1)^2 + 1^2 + 3^2\big)/4 = 20/4 = 5, so the standard deviation is \sqrt{5} \approx 2.236.
3. Each value becomes (x - 5)/2.236: that gives (-1.342,\ -0.447,\ 0.447,\ 1.342).
4. Finally, multiply by a learned scale \gamma (gamma) and add a learned shift \beta (beta), one pair per position. They start at 1 and 0, so at first nothing changes, but training can adjust them if a different scale works better.

\text{LayerNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

Read it as "gamma times the standardised x, plus beta", where \mu (mu) is the mean of this example's values, \sigma^2 their variance, and the tiny \epsilon prevents dividing by zero.

import torch, torch.nn as nn

ln = nn.LayerNorm(4)                               # γ starts at 1, β at 0
print(ln(torch.tensor([2., 4., 6., 8.])).round(decimals=3))
# → tensor([-1.3420, -0.4470,  0.4470,  1.3420], grad_fn=<RoundBackward1>)

The printout matches the hand calculation. (The trailing grad_fn just says PyTorch is tracking it for backpropagation.)

Batch normalisation, and why transformers do not use it

Batch normalisation, introduced by Sergey Ioffe and Christian Szegedy in 2015, came first and normalises in the other direction: for each feature, across all the examples in the batch. It was a large step forward for image networks and is still common in them. But it has drawbacks that come straight from its definition. It ties the examples in a batch together, so the output for one photo depends on which other photos happen to share its batch. It behaves badly with small batches, where the batch's mean and spread are noisy estimates. And at inference, when there may be a single example and no batch at all, it must switch to running averages saved during training, so the model behaves differently in training and in production.

Layer normalisation has none of these problems: each example is normalised on its own, so batch size does not matter and training and inference behave identically. That, plus the fact that it handles sentences of different lengths naturally, is why every transformer uses layer normalisation, or its simplified relative RMSNorm, which skips subtracting the mean and only divides by the root of the mean square, and is used in models like Llama.

3. Residual connections: a shortcut for the gradient

Even with good initialisation and normalisation, networks deeper than about twenty layers used to train worse than shallower ones, which made no sense: a deeper network could, in principle, just copy the shallower one and pass the result through its extra layers unchanged. In practice, the extra layers could not easily learn to "do nothing", and gradients still weakened through the long chain.

In 2015, Kaiming He and colleagues at Microsoft Research fixed this with an almost trivial change. Instead of a block of layers computing a new output f(\mathbf{x}) that replaces its input, it computes a change and adds it to the input:

\text{output} = \mathbf{x} + f(\mathbf{x})

A residual block. The layers compute only a change, f(x), which is added to the untouched input arriving along the green shortcut. Going backwards, the gradient can return along the shortcut without being multiplied by any layer's slope, so it cannot vanish.

Why does this help so much? Look at the slope. The derivative of \mathbf{x} + f(\mathbf{x}) with respect to \mathbf{x} is 1 + f'(\mathbf{x}). In backpropagation from 12.5.2, the gradient gets multiplied by this at every block. However small f' becomes, that "1" lets the gradient travel back along the shortcut, the green path in the figure, unchanged. A hundred residual blocks cannot shrink the gradient to nothing, because there is always a clean path straight from the loss to the earliest layers.

There is a second, more intuitive reason. If a block's best contribution is "do nothing", it only has to learn f(\mathbf{x}) = 0, which is easy: push its weights towards zero. Without the shortcut it would have to learn to reproduce its input exactly, which is hard. Residual networks, called ResNets, trained well at 152 layers, won the 2015 ImageNet competition, and residual connections are now in essentially every deep architecture. A transformer is a stack of residual blocks: every attention layer and every feed-forward layer in an LLM has a shortcut around it, as 12.7.4 shows.

4. Gradient clipping: a ceiling on each step

Sometimes a single batch produces a gradient hundreds of times larger than usual: an unusual example, or a sharp cliff in the loss landscape. One update with that gradient can throw the weights so far that training never recovers, and the loss turns into nan. Gradient clipping puts a ceiling on the size of each step. The usual form measures the length (the norm, from 12.2.1) of the whole gradient vector, and if it exceeds a limit, scales the vector down to that limit while keeping its direction.

Suppose the gradient is (30, 40), with length \sqrt{30^2 + 40^2} = 50, and the limit is 1.0. Clipping divides by 50, giving (0.6, 0.8): the same direction, length 1. Normal-sized gradients below the limit are left untouched.

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # (1)
opt.step()

Line (1) goes between backward() and step(), so that the step uses the clipped gradient. It costs almost nothing and saves many training runs, which is why nearly every large-model training script includes it.

5. Regularisation for networks

A network big enough to learn photos is also big enough to memorise every training photo, the overfitting of 12.3.3. Four tools are standard.

Dropout. During training, randomly switch off each neuron with some probability p, typically 0.1 to 0.5, choosing afresh for every batch. Dropout was introduced by Geoffrey Hinton's group, with the full paper by Nitish Srivastava and colleagues in 2014. The effect is that no neuron can count on any particular other neuron being present, so the network cannot build fragile chains of neurons that only work together; it has to spread each piece of knowledge across many neurons, which generalises better. You can also think of it as training a huge number of slightly different networks that share weights, and averaging them at the end.

There is one detail that matters. If half the neurons are switched off during training but all are on at inference, the next layer would suddenly receive about twice as much total signal. So PyTorch scales the surviving values up during training by 1/(1 - p). With p = 0.5, activations (2, 4, 6, 8) and a random mask that keeps the first and third, the output is (4, 0, 12, 0): each survivor is doubled, so the expected total stays the same. At inference, with model.eval(), dropout does nothing at all, which is one reason that call matters.

Weight decay. Shrink every weight slightly at every step, the L2 idea of 12.3.3, applied correctly through AdamW as 12.5.3 explained.

Early stopping. Keep a copy of the model whenever the validation loss reaches a new best, and stop when it has not improved for several epochs in a row; the number of epochs you wait is called the patience. Then use the saved best copy. This picks the bottom of the validation U-curve from 12.3.3 automatically.

best, patience, bad_epochs = float("inf"), 3, 0
for epoch in range(100):
    train_one_epoch(model)
    val_loss, _ = evaluate(model, val_dl)
    if val_loss < best:                                       # (1) new best
        best, bad_epochs = val_loss, 0
        torch.save(model.state_dict(), "best.pt")             # (2) keep it
    else:
        bad_epochs += 1
        if bad_epochs >= patience:                            # (3) give up
            break
model.load_state_dict(torch.load("best.pt"))                  # (4) use the best, not the last

Line (1) checks for a new best validation loss. Line (2) saves the model's weights; state_dict() is PyTorch's dictionary of every weight tensor by name. Line (3) stops after three epochs without improvement. Line (4) restores the best weights, because the model at the moment you stopped is already a few epochs past its best.

Augmentation and label smoothing. Data augmentation from 12.4.3 is often the strongest regulariser of all for images. Label smoothing trains classifiers towards a target of, say, 0.9 for the right class instead of exactly 1.0, sharing the remaining 0.1 among the others, which stops the network from becoming absurdly overconfident and tends to improve calibration from 12.3.2.

6. Debugging: a routine that finds most problems

When a network will not train, the loss curve is almost the only clue, and it is a vague one. Experienced engineers follow the same routine every time, in this order.

Check the loss at the very first step. Before any training, a classifier with random weights should be guessing roughly evenly, so its cross-entropy should be about \ln(\text{number of classes}): \ln 2 \approx 0.693 for yes-or-no, \ln 4 \approx 1.386 for the four support teams. If the starting loss is 15, something is badly wrong before training even begins: perhaps the inputs were not standardised, or a softmax is being applied twice.

Overfit one small batch. Take ten or twenty training examples and train on only those, for a few hundred steps. Any working network, however badly tuned, should drive the loss on ten examples almost to zero, because it can simply memorise them. If it cannot, the problem is a bug, not a hyperparameter: labels misaligned with inputs, a shape being broadcast wrongly, gradients not flowing to some layer, an optimiser created before the model was moved to the GPU. This one check finds more problems than any other, and it takes two minutes.

Then read the curves. With the full data, plot training loss and validation loss together. The pattern tells you which way to go:

Four loss curves and what each one means. Blue is training loss, red is validation loss. Reading them together, rather than either alone, is what tells you which fix to apply.

• Both high and flat. Underfitting or a bug. If the one-batch test passed, the model may be too small, or the learning rate far too low.
• Both falling. Working. Keep training.
• Training falls, validation turns up. Overfitting. Add dropout or weight decay, stop earlier, augment, or get more data.
• The loss shoots up or becomes nan. Lower the learning rate, add gradient clipping, add warmup, and look for a log(0) or a division by zero.

Finally, watch the metric you actually care about, not just the loss. Loss is what the optimiser minimises; recall at review capacity, or whatever the problem card from 12.1.3 says, is what decides whether the model ships. The two usually move together, but not always, and only one of them is the goal.

Recall

• Zero starting weights keep every neuron identical forever. Each layer multiplies signal variance by about \tfrac12 n\sigma^2 with ReLU, so He initialisation sets \sigma^2 = 2/n to keep it at 1 (Xavier uses 1/n for tanh).
• Normalisation re-standardises inside the network. Layer normalisation works per example and behaves identically in training and inference, so transformers use it (or RMSNorm); batch normalisation works across the batch and suits image networks.
• Residual connections compute \mathbf{x} + f(\mathbf{x}): the slope 1 + f' gives gradients a path that cannot vanish, and "do nothing" becomes easy to learn. Every transformer is a stack of them.
• Gradient clipping rescales an oversized gradient to a maximum length, keeping its direction.
• Regularise with dropout (survivors scaled by 1/(1-p), off at inference), weight decay, early stopping with patience, augmentation and label smoothing. Debug in order: check the first loss (about \ln C), overfit ten examples, read training and validation curves together, track the real metric.

Self-test: With 1,024 neurons per layer and ReLU, what weight standard deviation keeps the signal steady? Why does a residual block find "do nothing" easier than a plain block? Your network cannot get the loss on ten examples below 0.6. What does that tell you, and what does it rule out?

Next: 12.5.5 Convolutional Networks — so far every layer connected everything to everything; for photos that wastes millions of weights, and building the shape of images into the network is what made computers see.