AI Engineer Path
Phase 4 · Deep Learning4 Jan – 10 Jan

Week 16

Optimising a neural network

Make training fast, stable and well-generalising.

Why this week matters

Layer norm deserves special attention: it is inside every transformer block you will ever touch. Batch norm is the one you'll meet in CNNs.

Done when

You've run an optimiser sweep on one model and can explain the differences you observed.

Concepts

5 lessons · tick each one once you could explain it

Mini-batch SGD estimates the gradient on a small batch: noisy but cheap, and the noise even helps escape poor regions. Momentum keeps a running average of past gradients, so steps build speed along consistent directions and damp oscillations across a narrow valley.

RMSprop divides each parameter's step by a running average of its squared gradients: parameters with large, noisy gradients take smaller steps. Adam combines momentum (first moment) with RMSprop-style scaling (second moment), plus bias correction for the early steps.

AdamW decouples weight decay from the adaptive update. It is the default for transformers. Watch all four race across the same awkward valley in the simulation.

mt=β1mt−1+(1−β1)gt,  vt=β2vt−1+(1−β2)gt2,  θt=θt−1−αm^tv^t+ϵm_t = \beta_1 m_{t-1} + (1-\beta_1)g_t,\ \ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2,\ \ \theta_t = \theta_{t-1} - \alpha\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}

Going deeper

Adam's effective step size is roughly the learning rate regardless of gradient scale, which makes it forgiving. Its weakness was weight decay interacting badly with adaptive scaling, fixed by AdamW's decoupled decay.

Newer optimisers (Lion, Sophia, Muon) target memory or speed at LLM scale; AdamW remains the safe default. Optimiser state (two values per parameter for Adam) is a major part of training memory.

Best resources for this lesson

The learning rate is usually the most important hyperparameter. Too high diverges; too low wastes compute and can get stuck.

Warmup ramps the rate up from near zero over the first steps, which stabilises Adam's early statistics and is standard for transformers. Decay then lowers it: step decay, exponential, or cosine decay to a small floor. One-cycle policies go up then down within a run.

The LR range test (increase the rate each batch and plot loss) finds a sensible maximum quickly.

αt=αmin⁡+12(αmax⁡−αmin⁡)(1+cos⁡πtT)\alpha_t = \alpha_{\min} + \tfrac12(\alpha_{\max}-\alpha_{\min})\Big(1 + \cos\frac{\pi t}{T}\Big)

Going deeper

Warmup exists because Adam's second-moment estimates are unreliable in the first steps, and large early updates can wreck a freshly initialised transformer.

Warmup-stable-decay schedules (hold the LR flat, decay sharply at the end) let you extend training without restarting, increasingly common for LLM pretraining.

Best resources for this lesson

Dropout zeroes each activation with probability p during training (scaling the rest by 1/(1−p)), so the network can't rely on any single unit: like training an ensemble of thinned networks. It is switched off at inference, which is why model.eval() matters.

Early stopping monitors validation loss and keeps the checkpoint from the best epoch, stopping after it fails to improve for a 'patience' window. It is free regularisation.

Weight decay (L2) is the third standard regulariser.

Going deeper

At inference, dropout is turned off and activations are scaled so their expected value matches training ('inverted dropout' does the scaling during training instead).

Large pretrained transformers often use little or no dropout during pretraining because they see each example once. It matters more when fine-tuning on small datasets.

Common pitfalls

  • Forgetting model.eval() at inference: dropout stays on and predictions become random.

Best resources for this lesson

Batch normalisation normalises each feature across the examples in a mini-batch, then applies a learned scale and shift. It speeds training in CNNs but depends on batch statistics: it behaves differently at train and test time (running averages) and breaks with small batches or variable-length sequences.

Layer normalisation normalises across the features of a single example. No batch dependence and identical behaviour at train and test time, which makes it ideal for sequences. Every transformer block applies it (usually 'pre-norm', before attention and before the MLP). RMSNorm, used in LLaMA, drops the mean subtraction.

Both keep activations in a well-behaved range so gradients flow and higher learning rates are safe.

LN(x)=γ⊙x−μfeaturesσfeatures2+ϵ+β\text{LN}(x) = \gamma \odot \frac{x - \mu_{\text{features}}}{\sqrt{\sigma^2_{\text{features}} + \epsilon}} + \beta

Going deeper

Pre-norm (LayerNorm before each sub-layer) trains more stably at depth than the original post-norm transformer, which is why GPT-2 onward uses it. RMSNorm drops the mean-centring for speed with similar results.

BatchNorm's train/inference mismatch (batch statistics vs running averages) is a classic source of 'works in training, fails in production' bugs.

Best resources for this lesson

Where this comes back

  • Week 19LayerNorm wraps attention and the MLP in every transformer block.

Backprop multiplies many factors together. If they are mostly below 1, gradients vanish and early layers stop learning; above 1, they explode and training diverges.

Xavier/Glorot initialisation (for tanh) and He/Kaiming initialisation (for ReLU) scale random weights by the layer's fan-in so activation variance stays roughly constant through depth. Residual connections and normalisation layers help further.

Gradient clipping caps the gradient norm (commonly 1.0) before the update, a cheap guard against occasional spikes that is standard in transformer and RNN training.

W∼N ⁣(0, 2nin)  (He),g←g⋅min⁡ ⁣(1,c∥g∥)W \sim \mathcal{N}\!\Big(0,\ \frac{2}{n_{\text{in}}}\Big)\ \ (\text{He}),\qquad g \leftarrow g\cdot\min\!\Big(1, \frac{c}{\|g\|}\Big)

Going deeper

Residual networks often scale the last layer of each block by 1/√(2·layers) at initialisation so the residual stream's variance doesn't grow with depth, a trick used in GPT-2.

Monitor per-layer gradient norms and update-to-weight ratios (around 1e-3 is healthy); Karpathy's makemore lectures show exactly how.

Best resources for this lesson

Where this comes back

  • Week 18Vanishing gradients through time are why plain RNNs fail on long sequences.

Practice

Hands-on work that makes the lessons stick. Warm-ups take minutes; stretch goals are optional.

  1. Core

    Optimiser sweep

    Train one model with SGD, SGD+momentum, RMSprop, Adam and AdamW at 3 learning rates each. Plot all loss curves and explain the differences.

  2. Warm-up

    Break it, then fix it

    Initialise a 10-layer network with too-large and too-small weights and watch activations explode or vanish; fix it with He initialisation and LayerNorm.

  3. Stretch

    LR range test

    Implement an LR range test and pick a learning rate for a one-cycle schedule from the plot.

This week, day by day

Dates follow your pace from Settings. Open the notebook icon to log hours and notes.

  1. Day 106Monday4 Jan2 h planned

    SGD, momentum, RMSprop

  2. Day 107Tuesday5 Jan2 h planned

    Adam / AdamW and learning-rate schedules

  3. Day 108Wednesday6 Jan2 h planned

    Dropout and early stopping

  4. Day 109Thursday7 Jan2 h planned

    Batch norm vs layer norm (layer norm is load-bearing in transformers)

  5. Day 110Friday8 Jan2 h planned

    Weight init, vanishing/exploding gradients, gradient clipping

  6. Day 111Saturday9 Jan3 h planned

    Experiment: one model, sweep optimisers, log and compare results

  7. Day 112Sunday10 JanReview

    Review the week, finish anything unfinished, rest

Watch

100 Days of Deep LearningPrimary

CampusX · playlist

Neural Networks / Deep Learning

StatQuest · playlist

Improving Deep Neural Networks (DLS Course 2)

Andrew Ng · DeepLearning.AI · playlist

Optimisers, regularisation and batch norm, explained carefully.

Batch Normalization, Explained

CodeEmporium

NYU Deep Learning (LeCun & Canziani)

Alfredo Canziani · playlist

CMU 11-785 Introduction to Deep Learning

Carnegie Mellon University · playlist

Read and use

Interview prep

Questions this week's material gets asked as. Answer out loud first, then open the outline.

Explain Adam and why AdamW is preferred.
  • First moment (momentum) and second moment (per-parameter scaling)
  • Bias correction for early steps
  • AdamW decouples weight decay from the adaptive update
BatchNorm vs LayerNorm: which do transformers use and why?
  • LayerNorm normalises per example across features
  • No batch dependence; same at train and inference
  • Works for variable-length sequences and small batches

Check yourself

Five questions. The done-when test above is the real bar; this is a quick self-check.

  1. 1.What does momentum do?

  2. 2.Why do transformers use layer norm instead of batch norm?

  3. 3.Predictions change every time you run inference. Likely cause?

  4. 4.Training loss suddenly spikes to NaN. A standard guard is…

  5. 5.AdamW differs from Adam by…