ML from Scratch · Lecture 5 of 5
Neural Networks III: Training That Actually Works
Losses, mini-batches, momentum and Adam, sensible initialisation, normalisation and regularisation. The practical toolkit that turns backprop into a model that generalises.
Loss functions
- Regression: MSE $\frac1n\sum(y-\hat y)^2$ (Gaussian noise); MAE or Huber for robustness to outliers.
- Binary classification: sigmoid + binary cross-entropy.
- Multi-class: softmax $p_k=e^{s_k}/\sum_je^{s_j}$ + categorical cross-entropy $-\log p_{y}$.
Cross-entropy is the negative log-likelihood. Minimising it is maximum likelihood, the same principle behind logistic regression.
Stochastic gradient descent
Computing the full gradient over millions of examples per step is wasteful. Use a random mini-batch $B$ (32–1024 examples):
The noisy gradient is unbiased, cheap, and its noise even helps escape sharp minima. The learning rate η is the most important hyperparameter. Too high diverges, too low crawls.
Better optimisers
Momentum accumulates a velocity so consistent directions speed up and oscillations cancel:
Adam keeps running means of the gradient and its square and takes per-parameter steps scaled by the gradient's typical size:
with bias corrections $\hat m=m/(1-\beta_1^t)$, $\hat v=v/(1-\beta_2^t)$. Defaults $\beta_1=0.9$, $\beta_2=0.999$, η = 1e-3 work surprisingly often. AdamW decouples weight decay from the adaptive step and is today's default for transformers.
The loss surface here is a curved, narrow valley, typical of real networks. Plain gradient descent zig-zags across the valley walls. Momentum builds speed along the floor. Adam rescales each direction and heads more directly to the minimum. Push the learning rate up and watch plain GD start to bounce.
Initialisation
Weights must break symmetry (not all equal) and keep activations' variance stable from layer to layer:
The factor 2 for ReLU compensates for it zeroing half of its inputs.
Normalisation and residuals
- Batch normalisation: standardise each feature over the mini-batch, then rescale with learnable γ, β. It smooths the loss landscape and allows higher learning rates.
- Layer normalisation: standardise across features per example. It's the default in transformers and RNNs.
- Residual connections $h_{l+1}=h_l+F(h_l)$: gradients flow through the identity path, which makes 100+ layer networks trainable.
Regularisation
- Weight decay (L2): add $\frac\lambda2\|\theta\|^2$; prefers small, smooth weights.
- Dropout: during training zero each hidden unit with probability $p$ and scale the rest by $1/(1-p)$. This trains an implicit ensemble of thinned networks, bagging in disguise (Lecture 2).
- Early stopping: keep the weights from the epoch with the best validation loss.
- Data augmentation: label-preserving transformations (crops, flips, noise) multiply the data.
A training recipe
- Overfit a tiny batch first. If you can't get its loss to ~0, there's a bug.
- Start with AdamW, a sensible η (1e-3), He/Xavier init and normalisation layers.
- Use a learning-rate schedule: warm-up, then cosine or step decay.
- Watch train vs validation loss. A gap means more regularisation or data. Both high means a bigger model or more training.
- Clip gradients (norm ≈ 1) for recurrent and very deep models.
Key takeaways
- Cross-entropy = maximum likelihood; mini-batch SGD gives cheap, unbiased, helpful-noise gradients.
- Momentum smooths oscillation; Adam/AdamW adapt per-parameter step sizes and are strong defaults.
- He/Xavier init, normalisation, residuals, weight decay, dropout and early stopping make deep nets train and generalise.
Exercises
- Show that dropout with rescaling 1/(1−p) keeps the expected activation unchanged.
- Implement Adam for the XOR network from Lecture 4 and compare epochs-to-converge with plain GD.
- Derive He initialisation by requiring Var(z_{l+1}) = Var(z_l) for ReLU layers.