Worked lesson: gradients must accumulate
Let x = 3, a = x × x and y = a + a. The forward pass gives a = 9 and y = 18.
a is used twice, and each use contributes 1 to dy/da, so its accumulated gradient is 2.
The two inputs to x × x then each contribute x × 2 = 6 to dy/dx. The final gradient is 12.
Assigning a gradient with = instead of accumulating with += silently loses a path.
def objective(x):
a = x * x
return a + a
x, epsilon = 3.0, 1e-5
numeric_gradient = (objective(x + epsilon) - objective(x - epsilon)) / (2 * epsilon)
assert abs(numeric_gradient - 12.0) < 1e-6This numerical gradient is a small, independent check on your backward pass, for debugging rather than training. It needs many forward evaluations and can be unstable near non-differentiable points or with a badly chosen step size.
Try it yourself. Change the graph to y = a × a, keeping a = x × x. At x = 3, predict
the value and derivative before running the numerical check. Which reused nodes need accumulation?
Hint and explained answer
The objective is x⁴, so its value is 81 and its derivative is 4x³ = 108. Each use of a
contributes 9, giving dy/da = 18. Each use of x then contributes 3 × 18 = 54, adding up to
108. Test this graph as well as plain addition and multiplication. A backward pass that only works
on a simple chain is incomplete.
Core concepts
Backpropagation, by building it. Write a scalar autograd engine: a Value class with the
arithmetic operations, a topological sort for the backward pass, and correct gradient accumulation
when a node is used more than once. That last detail is where most people's bug lives, and finding
it is the moment the chain rule stops being a formula.
PyTorch as a language. Tensors, Dataset and DataLoader, and the explicit training loop you
should be able to implement and explain. Know why model.train() and model.eval() differ, why
zero_grad() exists, what torch.no_grad() saves you, and why both the model and the data need to
be on the same device.
Start with this explicit loop and inspect each state change:
for epoch in range(num_epochs):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()For validation, call model.eval() and run the forward pass inside torch.no_grad() or
torch.inference_mode(). Evaluation mode changes dropout and batch-normalisation behaviour.
Turning off autograd avoids building a backward graph. Neither does the other's job. When
debugging, check mode and gradient state separately rather than guessing. After
validation, switch back to training mode at the start of each epoch.
PyTorch's optimisation tutorial
shows separate training and evaluation loops.
Check the loss contract too. For ordinary multiclass classification, pass unnormalised logits to
CrossEntropyLoss, with class-index targets of the expected shape and dtype. Applying softmax first
changes the objective. Use valid probability targets only when the task calls for them. The
loss documentation
also explains ignored targets and reduction. For epoch loss, sum the loss numerator and its
intended denominator across batches, then divide. A plain average of batch means can overweight a
short last batch or a sequence with few supervised tokens.
Training dynamics, and how to break them. This is the part that separates people who can debug training from people who change the architecture and hope. Instrument your runs with loss curves for both splits, gradient norms per layer, and activation histograms.
Five habits fix most problems:
- Try to overfit ten clean samples with regularisation and augmentation turned off. An unexpected failure is a clue, though noisy labels, a constrained model or the loss can limit the fit. Check shapes, target encoding and gradient flow before adding complexity.
- Compare against a trivial baseline.
- Sweep the learning rate across orders of magnitude and plot it.
- Turn off every extra (augmentation, scheduler, dropout), get a working baseline, then add one thing at a time.
- Look at the data. Print ten samples as the model sees them, after transforms. A surprising share of "it will not learn" turns out to be inverted labels or a normalisation applied twice.
When it still will not learn, work down this checklist. Change one thing per run and log it.
- Overfit one small batch. If a few examples cannot reach near-zero training loss, suspect the loss, targets or update step before the architecture.
- Data, labels and shapes. Assert shapes and dtypes where the loss is computed, and check label
values against the output layer.
(B, 1)outputs against(B,)targets broadcast to(B, B)in a regression loss, as in the Chapter 2 exercise. - Loss at initialisation. With near-uniform initial predictions, multiclass cross-entropy starts
near
ln(C), about 2.30 for ten classes. Much higher suggests overconfident initialisation. An output bias matching the class prior or target mean is a cheap fix. - Learning rate, then warmup and schedule. Find a working constant rate by sweeping orders of magnitude. Only then add a short warmup for early instability, and a decay schedule.
- AdamW and weight decay.
Adamwithweight_decayadds an L2 term to the gradient by default.AdamWdecouples the decay from the adaptive update. Check which one you use, and whether biases and normalisation weights are excluded from decay. - Gradient norms and clipping. Log the global norm each step (
clip_grad_norm_returns it). Clip to cap occasional spikes, callingscaler.unscale_(optimizer)first if you use a gradient scaler. Clipping on most steps acts as a smaller learning rate. - Vanishing or exploding gradients. Early-layer gradient norms orders of magnitude below later ones, with a flat loss, suggest vanishing. Growing norms, loss spikes or NaNs suggest exploding. Check initialisation for the activation, normalisation, residual paths and the learning rate.
Karpathy's recipe for training neural networks and Google's Deep Learning Tuning Playbook expand these steps. Both have notes in resources/.
Transfer learning. When suitable pretrained weights exist, adapt them and compare against a small from-scratch baseline. For tiny datasets, freeze the backbone and train a new head. Unfreeze more as you get more data, and try lower learning rates for earlier layers. CNNs make weight sharing and local inductive bias concrete. Vision transformers and self-supervised encoders are useful challengers when their data, memory and latency needs fit. Avoid claiming one family wins every small-data or edge task. Augment training data only, deduplicate before splitting, and hold out whole subjects or capture sessions when deployment needs that generalisation.
Freezing parameters with requires_grad_(False) does not freeze batch-normalisation buffers. Decide
whether to keep the pretrained running statistics or adapt them on training data. To keep them,
return the relevant backbone modules to evaluation mode after each outer model.train() call,
leaving the new head in training mode. Record the choice, and keep validation data out of
statistics updates.
PyTorch's BatchNorm documentation
describes running statistics. Evaluation mode and gradient freezing are separate controls.
Attention, built by hand. Follow Karpathy's build-a-GPT lecture, typing rather than watching, then build a byte-pair tokeniser. Inspect how it splits whitespace, numbers and several languages. Tokenisation affects the input and token cost, but is only one possible cause of reasoning failures. Add tests for causal masking, padding, position handling and encode/decode round trips. Split a language corpus by source document before making token windows, so neighbouring or duplicated text cannot leak into validation. Learn BPE merges and any fitted preprocessing from training text only, then keep them fixed for validation and the final test.
The scale has a useful derivation. Suppose the d query and key components are independent, with
zero mean and unit variance. Then each product has variance 1 and the dot product, their sum, has
variance d. Dividing by sqrt(d) restores variance 1 under those assumptions. This helps stop dot
products growing with head dimension and over-concentrating the softmax. Learned components need not
meet those assumptions exactly, so this explains the scale rather than proving exact normalisation
for every trained head. Apply a causal mask before softmax, so a masked future token gets no
attention weight, however large its unmasked score.
To test causal masking, change a future token and check earlier logits stay the same within numerical tolerance, in evaluation mode. Shift targets correctly for next-token prediction, checking whether the model already shifts them internally. Exclude padding from the loss too. Masking it in attention alone leaves its training targets in place.
Read Attention Is All You Need after this, not before. It is a 2017 machine-translation paper describing an encoder-decoder model. Compare its attention mechanism with your decoder-only model. There is no useful percentage for how much of its design a current model shares. Read cold as a first resource, it tends to discourage people. Read after you have built the thing, it is easier to connect to the code. Its notes are in resources/.
Efficient training, after correctness. Profile an eager baseline, then change one thing at a time. Automatic mixed precision can cut memory and runtime on supported hardware. Choose the dtype for the device, check loss and gradients stay finite, and use a gradient scaler when your precision and setup need one (PyTorch AMP recipe).
Gradient accumulation changes the effective batch size. Scale the loss correctly and step the optimiser at the intended interval. With variable-length sequences, divide the summed token loss by all non-ignored prediction targets in the accumulation window, including a short final window. Averaging each microbatch's mean equally changes token weighting. Check your pinned trainer's handling so you do not divide twice. The Transformers accumulation guide documents this denominator. Accumulation does not reproduce a large batch's batch-normalisation statistics. Activation checkpointing trades recomputation for activation memory. These are different controls, not interchangeable fixes.
Try torch.compile once the model is correct. Record first-run compilation separately from
steady-state time, because changing shapes and graph breaks can erase the benefit. Report peak
memory, examples or tokens per second, and the same validation metric as the eager model. A slower
or less stable result is still a useful experiment
(official compile tutorial).
Where recent research fits. The material below was checked on 9 October 2026. It is orientation and optional depth within the 90 hours. Reproducing a frontier model is extra work.
| Source and evidence | What to learn | Bounded exercise |
|---|---|---|
| CS336 Spring 2026, official course | Architecture/resource accounting, data filtering and evaluation precede mid/post-training | Read the schedule, then trace one tiny model from corpus to checkpoint; reserve full assignments for Chapter 7 |
| FlashAttention-4, peer-reviewed MLSys 2026 paper | Attention performance depends on memory movement, kernel scheduling and hardware | Explain why the reported hardware result does not transfer automatically to your laptop; use a supported attention backend and profile it |
| DeepSeek-R1, peer-reviewed Nature paper published 17 September 2025 | Reinforcement learning with verifiable rewards is a separate stage from pretraining and supervised adaptation | Describe a reliable task verifier, then one way a model could exploit it; reasoning benchmarks alone show no general retention or pricing benefit |
For adaptation, tell full fine-tuning apart from LoRA, which trains low-rank weight updates, and QLoRA, which trains adapters on a quantised base model. Neither guarantees lower serving latency or unchanged task quality. Supervised fine-tuning learns from demonstrations and preference optimisation from comparisons. RL with verifiable rewards needs a real verifier. Learn these distinctions here. Run one bounded adaptation experiment only when it addresses a measured failure in your project: optionally in Chapter 5's week 31 fine-tuning decision, or as Chapter 7 depth work. The primary papers and current implementation guides are in the resource list.