Part II · ML foundations · Chapter 13

Training deep networks well

Plain gradient descent works on small networks. Deep ones need a handful of extra techniques: careful starting weights, smarter optimisers, normalisation layers, dropout, residual connections and a learning-rate schedule. Each was invented to remove a specific bottleneck.

Initialisation

Bottleneck: with bad starting weights, signals either explode or fade away as they pass through many layers, before any learning has happened.

A neuron with 100 inputs adds up 100 products. If the inputs and weights each have a standard deviation of 1, the sum's standard deviation is about √100 = 10 (variances add, chapter 8). A tanh or sigmoid fed numbers of size 10 sits on its flat part, where the slope is nearly zero. Ten layers later the signal is nonsense. If instead the weights have standard deviation 1/√100 = 0.1, the sum's standard deviation is about 1, and it stays about 1 layer after layer.

• Xavier (Glorot) initialisation: draw weights uniformly between ±√(6 / (n_in + n_out)). For a 100 → 100 layer that is ±0.173. Suited to tanh and sigmoid. The trainer in chapter 12 uses it.
• He initialisation: standard deviation √(2 / n_in), here 0.141. The extra factor of 2 makes up for ReLU zeroing half its inputs.
• Never start all weights at the same value: every neuron in a layer would compute the same thing and receive the same update forever. Random starts break the symmetry.

Optimisers

Bottleneck: a single learning rate is too big for steep directions and too small for shallow ones (chapter 6), and full-batch gradients are too expensive.

Mini-batch SGD

Estimate the gradient on a random batch of, say, 32 or 512 examples instead of the whole training set (chapter 6). Cheaper steps, with noise that averages out.

Momentum

Keep a running "velocity" that adds up past gradients with a decay factor β, usually 0.9, and step along the velocity:

Worked with a constant gradient of 1: the velocity goes 1, 1.9, 2.71, 3.44, 4.10, and settles at 1 / (1 − 0.9) = 10. Along a long gentle slope, steps grow up to ten times larger. Across a narrow valley, gradients flip sign each step and cancel in the velocity, so the zig-zag is damped. A heavy ball rolling down the valley, rather than a ball with no mass.

Adam and AdamW

Adam keeps two running averages per parameter: of the gradient (m, like momentum) and of the squared gradient (s). It divides one by the square root of the other, so every parameter takes steps of roughly the same size whatever the scale of its gradient.

Worked first step, gradient 0.5, η = 0.001:

1. m = 0.1 × 0.5 = 0.05; s = 0.001 × 0.25 = 0.00025.
2. Corrected: m̂ = 0.05 / (1 − 0.9) = 0.5; ŝ = 0.00025 / (1 − 0.999) = 0.25.
3. Step: 0.001 × 0.5 / √0.25 = 0.001 × 0.5 / 0.5 = 0.001.

Try it with a gradient of 50 instead: the step is still 0.001. That scale-independence is why Adam is the default for neural networks. AdamW is Adam with weight decay applied directly to the weights rather than through the gradient (chapter 11); almost all language models are trained with it.

Learning-rate schedules

• Warmup: start with a tiny learning rate and raise it over the first few hundred or thousand steps. Adam's averages are unreliable at the start, and big early steps can wreck a large model.
• Cosine decay: then lower it smoothly: η_t = η_max × ½(1 + cos(πt/T)). At 0%, 25%, 50%, 75% and 100% of training the multiplier is 1, 0.854, 0.5, 0.146 and 0. Big steps early to make progress, small steps late to settle into the bottom of the valley.

Gradient clipping

One bad batch can produce a huge gradient and undo hours of training. Clipping caps the gradient's length: if ‖g‖ is larger than a limit c, rescale it to length c. Worked: g = (3, 4) has length 5; clipped to 1 it becomes (0.6, 0.8), same direction, smaller step.

Normalisation layers

Bottleneck: as earlier layers change during training, the numbers each later layer receives drift in scale, so every layer is chasing a moving target.

• Batch normalisation normalises each feature across the examples in a batch (to mean 0, standard deviation 1), then rescales with two learned numbers γ and β. It made deep image networks much easier to train, but depends on batch statistics, which is awkward for text of varying length and for small batches.
• Layer normalisation normalises across the features of one example instead. For the vector (2, 4, 6) it gives (−1.225, 0, 1.225), exactly the z-scores of chapter 8; then with γ = 2 and β = 0.5 the output is (−1.949, 0.5, 2.949). It needs no batch, so transformers use it, or its cheaper cousin RMSNorm, which skips subtracting the mean.

Dropout

Bottleneck: big networks overfit, with neurons that only work in fragile combinations.

Dropout switches off a random fraction p of a layer's outputs at each training step and scales the survivors by 1/(1 − p) so the expected total is unchanged. Worked with p = 0.5: outputs (2, 4, 6, 8) with the random mask (1, 0, 1, 0) become (4, 0, 12, 0). At test time nothing is dropped and nothing is scaled. Each neuron learns to be useful without relying on any particular partner, a cheap way to train many overlapping networks at once.

Residual connections

Bottleneck: in very deep networks the gradient reaching the early layers is a product of many small factors (chapter 5) and vanishes.

A residual connection adds a layer's input to its output: y = x + F(x). The layer only has to learn a correction to its input. By the chain rule the slope of y with respect to x is 1 + F′(x): even if the layer's own slope is tiny, the 1 passes the gradient straight through. A stack of residual layers has a gradient "highway" from the loss back to the first layer. Residual networks of more than 100 layers trained easily once this was added, and every transformer block uses two residual connections (chapter 18).

Lower precision

Training in 16-bit numbers (bf16) instead of 32-bit halves memory and roughly doubles speed on modern GPUs. bf16 keeps fp32's range but fewer digits of precision; chapter 39 shows the bit layouts and why that trade suits training.

If you only call trained models

If you only call trained models, none of this runs in your code, but it still matters in two practical places. When you compare embedding or language models, their model cards describe training in exactly these terms. And the fine-tuning projects in chapter 41 train with this chapter's Adam and a fixed learning rate (the Adam class in labs/finetune-distil/ftcommon.py), while the real-model libraries they point to take AdamW, warmup and weight decay as training arguments, with these meanings.