Deep learning
How multi-layer networks are built and trained. Backpropagation as the chain rule, keeping gradients healthy with initialisation and normalisation, and the optimiser, schedule and regularisation choices.
Overview
Deep learning is machine learning with models built from many layers of simple differentiable operations, trained end to end by gradient descent. The layers can be fully connected, convolutional, recurrent or attention-based, but the training machinery is shared: a forward pass computes the loss, automatic differentiation computes the gradient of that loss with respect to every parameter, and an optimiser updates the parameters. Everything in transformers, vision and speech sits on this base.
This hub matters to a working engineer because most training failures are failures of this machinery rather than of the architecture. A loss that goes to NaN at step 300, a model that trains to 60 percent and stops, a validation curve that turns up after epoch 4: each has a short list of likely causes, and the list is the content of this page.
Key ideas
A network is a composition of simple functions
A multilayer perceptron applies a matrix multiply, adds a bias, passes the result through a non-linearity such as ReLU, and repeats. Without the non-linearity the stack would collapse to one matrix. With it, a network of modest width can approximate any continuous function in principle, but in practice depth is what makes approximation cheap: each layer builds features out of the previous layer’s features. The parameter count is the sum of the weight matrices, so a layer from 1,024 to 1,024 units costs about 1 M parameters, and a transformer block is a few such layers with some structure between them.
Backpropagation is the chain rule, organised
The loss depends on the last layer’s output, which depends on the layer before, and so on back to the input. Backpropagation applies the chain rule from the loss backwards, reusing each layer’s gradient to compute the one below, so the cost of a full gradient is about twice the cost of a forward pass regardless of how many parameters there are. PyTorch and JAX record the forward computation as a graph and replay it in reverse. The price is memory: every intermediate activation is kept until the backward pass has used it, which is why batch size and sequence length, not parameter count alone, decide whether a model fits on a GPU. Gradient checkpointing trades recomputation for that memory.
Keeping the signal healthy
A gradient that passes through fifty layers is a product of fifty factors. If each is a little above one the gradient explodes, if a little below it vanishes, and either way the lower layers stop learning. Three habits keep the factors near one. Initialise weights with the right scale for the activation, which is the Kaiming scheme for ReLU and the Xavier scheme for tanh. Normalise activations inside the network: batch normalisation uses minibatch statistics and suits convolutional nets, layer normalisation uses per-example statistics and suits transformers, and RMSNorm drops the mean subtraction for speed. Add residual connections so that each block learns a correction to its input rather than a whole new representation, which is why networks past about twenty layers became trainable after 2015. Gradient clipping is the safety net for the cases these do not cover.
Optimisers, schedules and regularisation
SGD with momentum is the reference optimiser and still the best choice for many convolutional nets. Adam adapts the step per parameter from running estimates of the gradient’s mean and variance, tolerates a wider range of learning rates and is the default for transformers. AdamW separates weight decay from the gradient update, which is the form you should use. The learning rate matters more than the optimiser: sweep it over a few powers of ten before touching anything else. Warm up from a small value for the first few hundred steps, then decay, usually with a cosine schedule to a tenth or a hundredth of the peak.
Regularisation fights variance. Weight decay shrinks parameters, dropout zeroes random activations during training and is switched off at evaluation, data augmentation makes the dataset look larger, and early stopping on the validation loss is the cheapest of all. Mixed precision (bfloat16 compute with float32 master weights) is not regularisation but belongs in the same checklist, since it halves memory and roughly doubles throughput on modern GPUs and Apple Silicon at almost no cost in accuracy.
Where to start
Read The training loop first, which writes gradient descent by hand, then Your first neural network in PyTorch, which hands the gradient to autograd. The PyTorch training loop reference is the page to keep open while writing code, and tensor shapes and broadcasting explains the error messages. The MNIST CNN project is the first full training run with a schedule and a validation curve.
Tools and platforms
| Tool | What for | Licence | Platform |
|---|---|---|---|
| PyTorch | Autograd, modules, optimisers, the default research and production framework | BSD-3 | PyTorch |
| JAX with Flax and Optax | Functional transforms, TPU support, composable optimisers | Apache-2.0 | JAX |
| Keras 3 | High-level API over PyTorch, JAX or TensorFlow | Apache-2.0 | |
| PyTorch Lightning | Training loop boilerplate, checkpoints, multi-GPU | Apache-2.0 | Open source |
| MLX | Array framework for Apple Silicon with unified memory | MIT | Apple |
| TensorBoard | Loss curves, histograms, embeddings | Apache-2.0 | |
| Weights and Biases | Hosted experiment tracking and sweeps | Proprietary service, MIT client | Open source |
Common mistakes
- Forgetting
optimizer.zero_grad(), so gradients accumulate across steps and the loss wanders. - Evaluating with the model still in training mode. Dropout stays on and batch norm uses batch statistics, so the numbers are wrong and noisy.
- Picking a learning rate without a sweep. The symptom of too high is loss spikes or NaN, of too low is a flat curve that looks like a data problem.
- Applying weight decay to biases and normalisation parameters. Most reference implementations exclude them.
- Starving the GPU with a slow data loader, then concluding the model is slow. Profile one step before changing anything.
- Overfitting a toy dataset and declaring the architecture works. Overfitting one batch is a useful smoke test, but it says nothing about generalisation.
Going deeper
- Goodfellow, Bengio and Courville, Deep Learning, 2016, free online.
- He et al., Delving Deep into Rectifiers, 2015, the Kaiming initialisation paper, and Deep Residual Learning for Image Recognition, 2015.
- Ioffe and Szegedy, Batch Normalization, 2015, and Ba, Kiros and Hinton, Layer Normalization, 2016.
- Loshchilov and Hutter, Decoupled Weight Decay Regularization, 2017, the AdamW paper.
- The PyTorch autograd mechanics and automatic mixed precision notes.