Getting It to Train, and Getting It to Generalise

Start by looking, instead of guessing. Push a batch of inputs through and print the standard deviation of the activations at each layer, and then the standard deviation of the weight gradients:

def profile(model, x):
    h = x
    for layer in model:
        h = layer(h)
        if isinstance(layer, (nn.Tanh, nn.ReLU)):
            print(f"activation std: {h.std().item():.4f}")
    model(x).sum().backward()
    for p in model.parameters():
        if p.dim() == 2:
            print(f"gradient std: {p.grad.std().item():.2e}")

Three initialisations, same architecture, six layers of $\tanh$:

The first row is the network above, and it is not learning slowly — it has no signal left. By layer four the activation standard deviation has underflowed to zero in four decimal places. Every input produces the same output, the output is a constant, and the gradient of a constant is nothing. The weights were initialised small "to be safe", and small was fatal.

The second row is the opposite failure and it is worth understanding because it is counterintuitive. The activations look healthy — around 0.96 — but that is $\tanh$ saturated, pinned near its limit, exactly the regime Chapter 5 measured as losing a factor of almost ten thousand in gradient. And yet the gradients are enormous: 1940 at the first layer. Both things are true at once. Each backward step multiplies by $W^\top$, and with 128 inputs at unit variance that factor has a gain of about $\sqrt{128} \approx 11$, which overwhelms the shrinkage from the saturated $\tanh$. The gradients grow geometrically on the way back. This is the exploding gradient, and it produces loss values of nan within a few steps in any real training run.

The third row is what you want: activations roughly constant in scale across depth, gradients roughly constant in scale across depth. Nothing dies, nothing explodes.

Initialising well fixes the scale at step zero. It does not keep it fixed: the weights move, and by step five thousand the careful variance argument no longer applies.

Normalisation layers enforce the scale continuously. Given a vector of activations, subtract a mean, divide by a standard deviation, then apply a learned scale $\gamma$ and shift $\beta$ so the layer can undo the normalisation if that turns out to be what it wants:

The only real question is what you average over. Batch normalisation³ takes $\mu$ and $\sigma$ across the batch dimension, one statistic per feature. Layer normalisation⁴ takes them across the features, one statistic per example.

That choice looks minor and decides almost everything downstream:

This is why LayerNorm is the one you will meet again in Chapter 9 unchanged: the transformer block uses it, and it uses it for the reasons in the right-hand column, not because it works better in the abstract.

Four candidate fixes for the dead network: Xavier initialisation, LayerNorm, residual connections, and Adam instead of SGD. The temptation is to apply all four and move on. Do that and you will never know which one mattered, and the next time it happens you will have no method — only a ritual.

So apply them one at a time. Same seed, same data, same architecture, 800 steps:

Read that table the way you would read it at 2 a.m. and the conclusion is: nothing works alone, everything works together, therefore deep learning is alchemy. That conclusion is wrong, and finding out why is the most useful thing in this chapter.

Give each run six times the budget — 5000 steps instead of 800 — and it changes completely:

Now the picture is sharp, and it is a diagnosis rather than a ritual.

Initialisation alone fixes it. Normalisation alone fixes it. Each addresses the actual disease — the forward signal collapsing to zero — and either one is sufficient. At 800 steps they merely looked like partial credit, because they had solved the problem and were still climbing out.

Residual connections and Adam do not fix it, at any budget. Not because they are bad, but because they treat a different disease. A residual connection gives the gradient a path around a blocking layer; that is worth a great deal when the gradient is the problem, and worth nothing when the forward signal is already zero, because a shortcut around a dead layer still carries a dead value. Adam rescales each parameter's step by its own gradient history; that helps when gradients have wildly different magnitudes, and cannot resurrect a network whose output does not depend on its input.

And "nothing" is still exactly 0.6931 after five thousand steps. Not 0.6929. It is not slow; it is dead, and that distinction is visible in a way it was not before, because you have the row that says a fix works to compare against.

From here on this course uses PyTorch. That should be earned rather than announced, so here is exactly what it does that you already know how to do.

An optimiser is a rule for turning gradients into parameter updates. Plain gradient descent uses the gradient. Momentum uses a running average of it, which smooths out the noise and builds speed along directions that stay consistent:

v = beta * v + p.grad          
p -= lr * v                    

Adam⁵ keeps two running averages — of the gradient and of the gradient squared — and divides one by the square root of the other, so each parameter gets a step scaled to its own recent gradient magnitude:

m = b1 * m + (1 - b1) * g          # mean of the gradient          
v = b2 * v + (1 - b2) * g * g      # mean of the squared gradient  
m_hat = m / (1 - b1 ** t)          # bias correction: both averages start at zero
v_hat = v / (1 - b2 ** t)
p -= lr * m_hat / (v_hat.sqrt() + eps)   

Ten lines. Run both against torch.optim on the same problem for 50 steps:

SGD+momentum   by hand [2.7781870365142822, -1.0304985046386719]
               torch   [2.7781870365142822, -1.0304983854293823]   max |diff| = 1.19e-07
Adam           by hand [0.4893140196800232, -0.46317872405052185]
               torch   [0.48931416869163513, -0.46317875385284424]   max |diff| = 1.49e-07

Identical to float32 precision. torch.optim.Adam is those five lines, plus decades of care about edge cases and a C++ kernel. That is the trade you are making from here on: not magic for understanding, but speed for lines you have already written.

The usual explanation of Adam is "adaptive per-parameter learning rates", which is a description rather than a reason. The reason is geometry, and it can be measured.

Take a loss whose curvature differs between directions: steep in one, shallow in another. SGD has one global learning rate, so it must pick a value small enough to be stable in the steepest direction — and that value is then far too small for the shallow one, where progress crawls. This is what causes the classic picture of gradient descent zig-zagging down a narrow valley.

Two curvature ratios, three optimisers, 300 steps, and each optimiser given the best learning rate from a sweep so nobody is handicapped:

At a ratio of ten, everything works and there is nothing to discuss. At a thousand, plain SGD cannot reach the answer at any learning rate tried — its best result is still an error of 1.93 — and it diverges outright at half the rates. Adam lands exactly on the target and diverges at none of them.

That last column is the practical reason Adam is the default. It is not that Adam finds better solutions; on well-conditioned problems tuned SGD often matches or beats it. It is that Adam is far less sensitive to the learning rate you picked, and real networks have curvature ratios much worse than a thousand across their millions of parameters.

Two more pieces belong here and both are one line. Gradient clipping rescales the gradient vector whenever its norm exceeds a threshold, which turns the "loss suddenly jumps to a huge value" row of the diagnostic table into a non-event. And learning rate schedules: a short warmup from near-zero over the first few hundred steps, because Adam's variance estimates are garbage until they have seen some gradients and a full-size step taken on garbage can wreck an initialisation; then cosine decay toward zero, because ending a run with the same step size you started with means jittering around the minimum instead of settling into it.

Everything so far was about getting the loss down. Now the harder half, because the loss going down is not the goal — it is a proxy for the goal, and the proxy fails in a specific and famous way.

Twelve points from a smooth function with a little noise. Fit polynomials of increasing degree:

Degree 11 through 12 points passes through every single one exactly — train error zero to six decimal places — and is eight times worse than degree 5 on data it has not seen. Ask degree 3 and degree 11 to predict at $x = 3.25$, just outside the training range:

degree  3: predicts   -1.053   (truth -0.012)
degree 11: predicts  +61.224   (truth -0.012)

Sixty-one, where the answer is approximately zero. The model did not learn the function; it learned the twelve points, and between them it does whatever the arithmetic demands.

This is overfitting, and its opposite — degree 1, which cannot represent the curve at all and is bad everywhere — is underfitting. The classical account splits a model's expected error into three parts: bias, the error from the model being too rigid to represent the truth; variance, the error from the model being so flexible that it chases the noise in this particular sample; and irreducible noise, which nothing fixes. Simple models are biased, flexible models are high-variance, and the classical prescription is to find the sweet spot in the middle — degree 5 in the table above.

The standard tools all attack the variance term:

• L2 regularisation (weight decay) adds $\lambda \lVert w \rVert^2$ to the loss, pulling weights toward zero and making the function smoother. In the table above, degree 11's largest coefficient does the damage; penalising size defuses it.
• L1 adds $\lambda \sum |w_i|$ instead. The difference is not cosmetic: L2's gradient is proportional to the weight and so shrinks as the weight does, approaching zero without arriving, while L1's gradient is a constant $\pm\lambda$ that keeps pushing all the way. L1 therefore produces weights that are exactly zero — it selects features. L2 produces small weights. Use L2 when you want smoothness, L1 when you want sparsity.
• Dropout⁷ zeroes a random subset of activations on each training step, so no unit can rely on any particular other unit being present.
• Early stopping watches the validation loss and stops when it turns upward.
• Data augmentation manufactures more training examples from the ones you have, which attacks the problem at its source: overfitting is a shortage of data as much as an excess of parameters.
• Cross-validation splits the data $k$ ways and trains $k$ times, which buys a reliable estimate of test error when you have too little data to spare a held-out set.

Now the fact that breaks the picture.

The bias-variance story says that past the sweet spot, more parameters mean worse generalisation. Modern language models have far more parameters than the classical rules allow for the data they see, and generalise superbly. Both of those statements are true, and reconciling them is the most useful thing in this chapter.

Forty training points, twenty-dimensional inputs, random ReLU features, and the number of features $P$ swept from 2 to 5000 — with the minimum-norm solution chosen whenever there are many that fit:

Read it in three parts. Up to $P/n = 0.5$ the classical story holds exactly: error falls, then starts to rise. At $P = n = 40$ — the interpolation threshold, where the model has exactly enough parameters to pass through every training point — the test error peaks, at 5.81, five times worse than the small model. That peak is the classical warning, and it is real.

Then it descends again. And it keeps descending, past $P = 5n$, past $P = 37n$, all the way to $P = 125n$, where the test error of 0.5664 is better than the best under-parameterised model ever achieved. A model with 5000 parameters fitted to 40 points is the best model in the table.

This is double descent,⁸⁹ and the mechanism is visible in the last column. Once $P > n$ there are infinitely many parameter settings that fit the training data exactly, and which one you get depends on how you choose. The minimum-norm solution picks the smallest, and $\lVert w \rVert$ shows what that means: it peaks at 14.83 right at the threshold — where there is exactly one interpolating solution and you are stuck with it, however extreme — and then falls monotonically as $P$ grows, because more parameters means more interpolating solutions to choose from, which means the smallest available one gets smaller. At $P = 5000$ the norm is 0.18, eighty times smaller than at the threshold.

So the extra parameters are not adding complexity. They are adding choice, and the selection rule spends that choice on simplicity. The regularisation is not in the loss function; it is in the algorithm. Gradient descent from a small initialisation has a documented bias toward small-norm solutions, which is why this behaviour shows up in real networks trained the ordinary way and not only in the linear algebra above.

The practical consequence, which Chapter 10 depends on: "the model has more parameters than data, so it will overfit" is not a valid argument. It was a good rule when models lived to the left of the threshold. Everything interesting now lives far to the right of it, where the rule reverses.

The tools in this chapter are enough to train a network that works on data you can put in a table: rows of numbers, a column of labels.

Language is not that. Before a model can predict the next word, something has to decide what a "word" even is — and the answer is neither letters nor words, but a vocabulary the model learns from the raw bytes of the training data. That decision, made once before training starts, determines how many things the model can say, how much a request costs, and why models that can pass a law exam cannot reliably count the letters in strawberry.

Chapter 7 builds a tokenizer.

For the residual connections used above, He et al., Deep Residual Learning for Image Recognition (arXiv:1512.03385). Andrej Karpathy's Building makemore Part 3: Activations & Gradients, BatchNorm walks through the activation-histogram diagnostic on a real model and is the best hands-on treatment of the first half of this chapter. Yaser Abu-Mostafa's Learning From Data lectures 8 and 11–13 give the classical generalisation theory properly, including the parts this chapter compressed into a paragraph.