12.4.2 — Architectures, and What Makes Training Work

A fully connected layer treats every input independently of every other. For a 224×224 colour image that is 150,528 inputs, and one hidden layer of 1,000 neurons needs 150 million weights — for one layer, before anything useful happens.

Worse, it has no idea that neighbouring pixels are related, or that a cat in the top-left is the same cat in the bottom-right. Every architecture below is a way of building a known property of the data into the model, so the model does not have to learn it from scratch with billions of parameters.

1. Convolutional networks: locality and weight sharing

A convolution slides a small window — a kernel, typically 3×3 — across the image, computing a dot product at each position.

kernel               patch of image        output
[ 1  0 -1 ]          [ 10  10  40 ]
[ 1  0 -1 ]    ·     [ 12  11  45 ]   =   (10+12+9) − (40+45+42) = −96
[ 1  0 -1 ]          [  9  10  42 ]

That kernel detects a vertical edge — it is large where the left side is bright and the right side dark. A convolutional layer learns dozens of such kernels, and each produces a feature map showing where its pattern occurs.

Two ideas do all the work.

Parameter sharing. One 3×3 kernel over 3 colour channels is 27 weights, used at every position in the image. Compare with 150 million. And it encodes something true: an edge is an edge wherever it appears.

Locality. A pixel relates to its neighbours, not to one across the picture. Early layers see small patches; because each later layer looks at the outputs of the previous one, the receptive field grows with depth — so early layers detect edges, middle layers detect textures and parts, and deep layers detect objects. That hierarchy is learned, not designed, and inspecting it was one of the results that made the field take deep learning seriously.

Pooling downsamples — max-pooling takes the largest value in each 2×2 block — shrinking the map, widening the receptive field, and giving a little tolerance to position. Many modern architectures use strided convolutions instead, which do the same thing with learnable weights.

A typical stack: convolution → activation → convolution → activation → pool, repeated with more channels and smaller spatial size each time, ending in a classifier.

Convolutional networks remain the sensible default for images at ordinary data scales. Vision transformers (Chapter 12.5.1's mechanism applied to image patches) beat them given very large datasets, precisely because they have fewer built-in assumptions and must learn locality from data.

2. Recurrent networks, and why they lost

Sequences need memory: the meaning of a word depends on what came before.

An RNN keeps a hidden state and updates it one step at a time:

h_t = \tanh(W_h h_{t-1} + W_x x_t + b)

The same weights are applied at every step — parameter sharing again, now across time.

Training uses backpropagation through time: unroll the loop and apply Chapter 12.4.1's method. And that is where it breaks. A 100-token sequence is a 100-layer chain, so the gradient is a product of 100 factors — vanishing or exploding, per the same argument. A plain RNN cannot learn a dependency more than a few steps back.

LSTM fixes it with gates. A separate cell state runs through the sequence with only small linear interactions, and three learned gates control it: a forget gate deciding what to discard, an input gate deciding what to add, and an output gate deciding what to expose. The cell state is a highway with an almost-uninterrupted gradient path, which is what allows learning across hundreds of steps. GRU is a simpler two-gate variant that performs comparably.

They worked, and they lost to transformers for one reason: h_t depends on h_{t-1}, so the sequence must be processed in order. No parallelism across a sequence. When the winning strategy became "train on far more data", an architecture that cannot use a whole GPU on one sequence could not compete. Chapter 12.5.1's opening is exactly this problem being removed.

3. The three things that made deep networks trainable

Residual connections. ResNet (2015) added a shortcut: \text{output} = f(x) + x. Trivial, and decisive — it took usable depth from about 20 layers to over 100.

Why it works, in gradient terms: the derivative of f(x) + x with respect to x is f'(x) + 1. That +1 gives the gradient a path backwards that cannot vanish, no matter how small f' becomes. Every deep architecture since, including every transformer, uses residual connections.

Normalisation. As values pass through layers, their scale drifts, and layers must keep adapting to their inputs' changing distribution.

Batch normalisation normalises each feature across the examples in a batch. Effective for vision, and it couples examples in a batch together, which makes small batches unstable and inference behave differently from training (it must use running averages).

Layer normalisation normalises across the features of one example. It is batch-independent, so batch size does not affect it and training and inference behave identically. That is why transformers use layer norm — it also works with variable-length sequences, which batch norm handles badly.

Where to put it matters more than it sounds. Post-norm (normalise after the residual addition) was the original transformer design and needs careful warmup; pre-norm (normalise before the sublayer) is far more stable and is what large models use.

Better initialisation. Sizing initial weights to the layer's fan-in — He initialisation for ReLU, Xavier for tanh — keeps activations from shrinking or exploding as they propagate. Chapter 12.4.1 covers why zero initialisation fails outright.

4. Regularisation: fighting overfitting

Dropout randomly zeroes a fraction of activations during training (0.1–0.5), and turns off at inference. The effect is that no neuron can rely on any particular other neuron being present, so the network cannot build fragile co-adapted paths. It is also loosely an ensemble of many sub-networks sharing weights.

Weight decay penalises large weights, the same idea as ridge regression (Chapter 12.3). AdamW exists because naive weight decay interacts badly with adaptive optimisers.

Early stopping — watch validation loss and stop when it starts rising while training loss still falls. That divergence is the definition of overfitting, visible on a chart.

Data augmentation — generate variations: crop, flip, rotate, adjust colour for images; synonym replacement or back-translation for text. The most effective regulariser is more data, and augmentation is the cheapest way to approximate it. Be careful that the augmentation preserves the label: flipping a photo horizontally is fine, flipping a photo of text is not.

Label smoothing — train toward 0.9 instead of 1.0 for the correct class, which stops the model becoming absurdly overconfident and improves calibration (Chapter 12.3).

5. Making training fit on the hardware

Batch size interacts with everything. Larger batches give more stable gradients and better hardware use; very large batches can generalise slightly worse and need a scaled-up learning rate. The practical constraint is memory, and the usual limit is activations rather than weights.

Mixed precision stores weights and computes in 16-bit rather than 32-bit. Roughly half the memory and substantially faster on hardware with tensor cores. bfloat16 is preferred over float16 because it keeps float32's exponent range and sacrifices mantissa bits (Chapter 1.4), so it does not overflow or underflow during training — which is the failure that made early 16-bit training need loss scaling.

Gradient accumulation simulates a large batch on small hardware: run several mini-batches, sum the gradients, and update once. Same result as a big batch, at the same total time and a fraction of the memory.

Gradient checkpointing discards intermediate activations and recomputes them during the backward pass. Roughly 30% slower, and it can cut memory dramatically — the standard trade when a model nearly fits.

Distributing across machines comes in three shapes, and the names are worth knowing:

• Data parallel — every device holds the whole model and a slice of the batch; gradients are averaged. Simple, and it requires the model to fit on one device.
• Tensor parallel — one layer's matrices are split across devices. High communication, used within a machine.
• Pipeline parallel — different layers on different devices, with micro-batches flowing through. Introduces idle time (the "bubble") that micro-batching reduces.

Large models use all three at once, and frameworks such as DeepSpeed and FSDP also shard the optimiser state, which is often larger than the model itself — Adam keeps two extra values per parameter, so optimiser state alone is roughly twice the weights.

Checkpoint frequently. A multi-day run without checkpoints is a multi-day run you will repeat.

6. Transfer learning

Almost nobody trains from scratch, and this is why.

A model pre-trained on a huge general dataset has already learned useful low-level structure — edges and textures for images, syntax and word relationships for text. Adapting it to your task needs orders of magnitude less data and compute.

Three levels of adaptation:

Feature extraction — freeze everything, train only a new final layer. Fast, works with hundreds of examples, and cannot adapt the internal representations.

Fine-tuning — unfreeze some or all layers and continue training at a much lower learning rate (typically 10–100× lower). Higher rates destroy the pre-trained knowledge, an effect called catastrophic forgetting.

Parameter-efficient fine-tuning — train a small number of added parameters and leave the base model frozen. This is how large models are adapted now, and Chapter 12.8 covers LoRA and its family.

7. Generative architectures, briefly

Autoencoders compress to a small latent representation and reconstruct. Useful for anomaly detection — a high reconstruction error means "unlike anything seen in training".

GANs train a generator and a discriminator against each other. They produced remarkable images and are notoriously unstable to train, including mode collapse, where the generator finds one output that fools the discriminator and produces only that.

Diffusion models are what displaced them for images. Train a model to remove a small amount of noise from a noisy image; then start from pure noise and denoise repeatedly. The training objective is far more stable than an adversarial game, which is the main reason they won, and the iterative process is why generation takes many steps and is slower than a GAN's single pass.

8. Reading a training run

Watch training loss and validation loss together, because the gap between them is the diagnosis:

• Both high and flat — underfitting, a bug, or a learning rate far too low. Try overfitting ten examples first (Chapter 12.4.1).
• Both falling — working. Keep going.
• Training falls, validation rises — overfitting. Stop early, regularise, or get more data.
• Loss becomes NaN — exploding gradients or a numerical problem. Clip gradients, lower the learning rate, check for a division by zero or a log(0).
• Loss spikes then recovers — usually a bad batch; frequent spikes mean the learning rate is too high.

Track a metric you actually care about alongside the loss. Loss is what you optimise; accuracy, recall or a business measure is what you ship. They can move in opposite directions, and only one of them decides whether the model is useful.

Recall

• Architectures encode a known property of the data: convolutions encode locality and weight sharing (27 weights reused everywhere instead of 150 million), and the receptive field grows with depth, producing a learned edges → textures → objects hierarchy.
• RNNs lost because h_t depends on h_{t-1} — no parallelism across a sequence. LSTM gates gave the cell state an almost-uninterrupted gradient path, which is what allowed long dependencies at all.
• Residual connections make the derivative f'(x) + 1, giving gradients a path that cannot vanish — the change that took usable depth past 100 layers, and used by every transformer.
• Layer norm over batch norm for transformers: batch-independent, identical at training and inference, and fine with variable-length sequences. Pre-norm is far more stable than post-norm.
• Regularisation: dropout stops fragile co-adaptation, weight decay shrinks weights, early stopping triggers when validation loss rises while training loss falls, and augmentation approximates more data — the strongest regulariser of all.
• bfloat16 over float16: same exponent range as float32, so no overflow. Gradient accumulation simulates a large batch; gradient checkpointing trades ~30% speed for large memory savings.
• Distribution is data / tensor / pipeline parallel, usually combined — and optimiser state is roughly twice the weights with Adam, which is why it gets sharded.
• Fine-tune at a 10–100× lower learning rate or you cause catastrophic forgetting. Diffusion beat GANs because a denoising objective is stable where an adversarial game is not.

Self-test: What two properties of images do convolutions build in? · Why can't an RNN use a whole GPU on one sequence? · Why does a residual connection stop gradients vanishing? · Why do transformers use layer norm rather than batch norm? · What does bfloat16 protect against that float16 does not? · What does training loss falling while validation loss rises mean?