Warmup Guide — Autograd & Training Dynamics

Zero-to-senior primer for Phase 03. We start from "what is a gradient and why does a model need one" and end at the full training loop: the computation graph, reverse-mode automatic differentiation (the engine behind loss.backward()), the softmax+cross-entropy (p − y) gradient, the optimizer family SGD → momentum → Adam → AdamW, learning-rate schedules, gradient clipping and loss spikes, numerical precision and mixed precision, and gradient accumulation. Then we build all of it as a tiny scalar autograd engine and watch a real net learn XOR. When you finish, loss.backward() is no longer a black box — it is the thing you wrote.

Table of Contents

• Chapter 1: Why Gradients — Learning as Descent
• Chapter 2: The Computation Graph
• Chapter 3: The Chain Rule and Reverse-Mode Autodiff
• Chapter 4: Backprop on the Graph — One Pass, All Gradients
• Chapter 5: Softmax + Cross-Entropy and the Clean (p − y) Gradient
• Chapter 6: Optimizers — SGD → Momentum → Adam → AdamW
• Chapter 7: Learning-Rate Schedules — Warmup + Cosine
• Chapter 8: Gradient Clipping and Loss Spikes
• Chapter 9: Numerical Precision, Mixed Precision, Gradient Accumulation
• Chapter 10: This Is PyTorch Autograd
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Why Gradients — Learning as Descent

From zero. A neural network is a function f(x; θ) with a pile of tunable numbers θ (the parameters). "Training" means: pick θ so the model's outputs are good. We make "good" precise with a loss L(θ) — a single number that is small when the model is right and large when it's wrong. Training is then a minimization problem: find θ that makes L small.

The one tool we have. θ has billions of dimensions, so we can't search it. But we can ask a local question: if I nudge each parameter a tiny bit, does the loss go up or down, and how fast? That is exactly the gradient ∇L, the vector of partial derivatives ∂L/∂θ_i. The gradient points in the direction of steepest increase, so to decrease the loss we step the other way:

$$ \theta \leftarrow \theta - \eta \cdot \nabla L(\theta). $$

This is gradient descent, and η (the learning rate) is how big a step we take. The entire field of training is (a) computing ∇L efficiently and (b) choosing a good step from it. Chapters 2–4 are (a); chapters 6–8 are (b).

Why this framing matters. Once you see training as "compute the gradient, take a step, repeat," everything else slots in: backprop is how we compute the gradient cheaply; Adam is a smarter step; the LR schedule is how the step size changes over time; clipping is what to do when the gradient is huge. Hold this picture and nothing in the phase is mysterious.

Common misconception. "The model figures out the weights itself." No — gradient descent moves the weights, and backprop supplies the gradients it descends on. There is no magic; there is calculus and a loop. The senior move is to know the loop well enough to debug it.

Chapter 2: The Computation Graph

What it is. Any expression you compute is a directed graph: leaves are inputs/parameters, internal nodes are operations, and edges point from an operation to the values it produced. Take

$$ L = (w \cdot x + b - y)^2. $$

As a graph: w and x feed a *; that and b feed a +; subtract y; square the result. Each node holds a value (the forward result) and remembers which nodes it came from.

   w ──┐
       ├─(*)── u ──┐
   x ──┘           ├─(+)── p ──┐
   b ──────────────┘           ├─(−)── e ──(square)── L
   y ──────────────────────────┘

Why build the graph at all? Because the graph is the program, and the derivative of a program is computed by walking that same graph backward (Chapter 3). In our lab, every arithmetic operation on a Value creates a new node that stores its parents (_prev) and a closure (_backward) describing its local derivative. The graph is built as a side effect of the forward pass — you just compute L normally, and the tape records itself.

Two flavors. Frameworks build this graph either statically (define the whole graph once, then run it many times — TensorFlow 1.x, JAX's jit) or dynamically (rebuild it every forward pass, so control flow is plain Python — PyTorch eager, and our Value engine). Dynamic graphs are what make PyTorch feel like normal code; the price is rebuilding the tape each step.

Common misconception. "The graph stores the math symbolically and differentiates it like a CAS (computer algebra system)." It does not. Autodiff is not symbolic differentiation and not finite differences. It stores numeric local derivatives at each node and composes them — exact (to floating point) and cheap. That distinction is a favorite interview probe.

Chapter 3: The Chain Rule and Reverse-Mode Autodiff

The chain rule. If L depends on u and u depends on w, then

$$ \frac{\partial L}{\partial w} = \frac{\partial L}{\partial u} \cdot \frac{\partial u}{\partial w}. $$

In words: the sensitivity of the loss to w is the sensitivity to u times the local sensitivity of u to w. Every node only needs to know its local derivative (∂u/∂w for the op it performed); the chain rule stitches the locals into the global gradient.

Two ways to apply it. Suppose the computation goes inputs → ... → loss.

• Forward-mode computes, for one chosen input, how every downstream value changes — it pushes derivatives forward alongside the values. To get gradients w.r.t. all n parameters you'd run it n times. Cheap when you have few inputs, many outputs.
• Reverse-mode computes, for one chosen output (the loss, a scalar!), how it changes w.r.t. every input — it pulls derivatives backward from the output. One pass gives all n gradients. Cheap when you have many inputs, one output — which is exactly a neural network: billions of parameters, one scalar loss.

This is the whole reason backprop exists. We have N parameters and 1 loss, so reverse-mode is the right algorithm: a single backward pass yields ∂L/∂θ_i for every parameter, at roughly the cost of one extra forward pass. (Recall the 6ND from Phase 00 — the "× 3" is forward + the ~2× of this backward pass.) Forward-mode would cost N passes. That asymmetry is why every framework does reverse-mode and why this is the single most important algorithm in the curriculum.

The local derivatives you need. For each op, the local rule the node stores:

Each row is one _backward closure in the lab. That's the entire engine.

Common misconception. "Reverse-mode is always better." It's better for our shape (many in, one out). If you had one input and many outputs (rare in ML, common in sensitivity analysis), forward-mode wins. Knowing why you pick reverse-mode — the input/output count — is the senior answer, not "because PyTorch does."

Chapter 4: Backprop on the Graph — One Pass, All Gradients

The algorithm. Backpropagation is reverse-mode autodiff run on the computation graph:

1. Forward pass — compute every node's value, building the graph (each node remembers its parents and its _backward closure).
2. Seed the output node's gradient to 1.0 — because ∂L/∂L = 1.
3. Reverse topological order — process nodes so that a node is handled only after every node that uses it. Call each node's _backward, which uses the chain rule to add this node's contribution into its parents' .grad.

Why topological order, exactly? A node's gradient is the sum over all paths from it to the loss. If node x feeds two later nodes, x.grad isn't complete until both of those have pushed their contributions. Processing in reverse-topo order guarantees that by the time we read a node's grad to push it upstream, every downstream consumer has already added to it. Get the order wrong and you read a half-finished gradient.

forward:   build graph, fill .data           ──►
                                              loss
backward:  seed loss.grad = 1.0              ◄──
           reverse-topo: each node does
           parent.grad += local_deriv * node.grad

Gradients accumulate — this matters twice. Inside one backward pass, a shared node (the diamond: x used in two places) must sum its incoming gradients, so every _backward uses +=, never =. And across training steps, grads keep accumulating in .grad until you reset them — which is precisely why PyTorch makes you call optimizer.zero_grad() (or set grads to 0.0 at the top of the loop, as our train does) every step. Forget it and you sum this step's gradient onto last step's: the classic silent training bug.

A concrete hand-check (the lab's first test). For z = x*y + x at x=2, y=3: dz/dx = y + 1 = 4 (one path through the *, one through the bare + x), dz/dy = x = 2. Run .backward() and you should get exactly x.grad == 4, y.grad == 2. If you can derive that by hand and the engine agrees, your chain rule is right.

The soul test: finite differences. The independent ground truth is the definition of a derivative:

$$ \frac{\partial L}{\partial x} \approx \frac{L(x+h) - L(x-h)}{2h} \quad (\text{central difference}). $$

It's slow (O(params) re-evaluations) and a little inexact, but it needs no calculus — so it's the perfect referee. If your analytic backprop matches central finite differences on a gnarly composite like tanh(x)² · e^x + log(x+3), your engine is correct. This is the test that proves an autograd implementation; every framework's test suite has it (torch.autograd.gradcheck).

Common misconception. "Backprop is its own algorithm separate from autodiff." Backprop is reverse-mode autodiff specialized to neural-net graphs. The 1986 "backprop" paper popularized it for nets; the underlying reverse-mode idea is older and fully general.

Chapter 5: Softmax + Cross-Entropy and the Clean (p − y) Gradient

Turning logits into probabilities. A classifier/LM outputs raw scores (logits) z. Softmax maps them to a probability distribution:

$$ p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}. $$

The stability trick (non-negotiable). e^{1000} overflows to inf. Subtract the max logit m = max_j z_j first:

$$ p_i = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}. $$

This is mathematically identical (the e^{-m} cancels top and bottom) but now every exponent is ≤ 0, so nothing overflows. Our softmax does this, and a test feeds it [1000, 1001, 1002] to prove it. Never write a softmax without the max-subtraction — it's a guaranteed interview ding.

The loss. Cross-entropy measures how far the predicted distribution p is from the true one-hot target y (a 1 at the correct class, 0 elsewhere):

$$ \text{CE} = -\sum_i y_i \log p_i = -\log p_{\text{target}}. $$

For a one-hot target it collapses to "negative log-probability of the right answer." Build it on the stable softmax so log never receives 0.

The beautiful gradient. Differentiate CE through softmax w.r.t. the logits and almost everything cancels:

$$ \boxed{;\frac{\partial \text{CE}}{\partial z_i} = p_i - y_i;} $$

— predicted probability minus the one-hot target. If the model is confident and right, p ≈ y and the gradient is near zero (nothing to fix); if it's confidently wrong, the gradient is large. This (p − y) form is why the softmax+CE pair is the default classification loss everywhere: the backward is a single subtraction. Our lab verifies it numerically — the analytic (p − y) must match finite differences, and the gradient on the target logit is p_target − 1.

Why frameworks fuse them. F.cross_entropy in PyTorch takes raw logits, not probabilities, and computes softmax + log + the loss together using log-sum-exp — both for stability and to use the clean fused gradient directly. If you ever softmax-then-CE manually, you've reimplemented it less safely.

Common misconception. "Cross-entropy needs probabilities as input." The math does, but the production op takes logits and does the softmax internally — feeding it already-softmaxed values double-softmaxes and silently wrecks training. State this; interviewers love it.

Chapter 6: Optimizers — SGD → Momentum → Adam → AdamW

We have the gradient. How do we turn it into a step? This is a short evolutionary ladder; know each rung and what the next one fixed.

1. SGD. The literal descent rule:

$$ \theta \leftarrow \theta - \eta , g, \qquad g = \nabla L. $$

Simple, memory-free, and with a good schedule it generalizes beautifully (it's still used to train many vision models). Weakness: one global η for all parameters, and it zig-zags in ravines.

2. Momentum. Accumulate an exponentially-weighted average of past gradients (a "velocity") and step along it:

$$ v \leftarrow \mu v + g, \qquad \theta \leftarrow \theta - \eta v. $$

This damps oscillation across a ravine and accelerates along its floor — like a heavy ball rolling downhill. μ ≈ 0.9 is typical.

3. Adam (Adaptive Moment estimation). Keep two running averages: the 1st moment m (mean of gradients, like momentum) and the 2nd moment v (mean of squared gradients, a per-parameter scale):

$$ m \leftarrow \beta_1 m + (1-\beta_1) g, \qquad v \leftarrow \beta_2 v + (1-\beta_2) g^2. $$

Because m, v start at 0 they're biased toward zero early, so bias-correct:

$$ \hat m = \frac{m}{1-\beta_1^{,t}}, \qquad \hat v = \frac{v}{1-\beta_2^{,t}}, $$

then take a per-parameter step scaled by the inverse root of the 2nd moment:

$$ \theta \leftarrow \theta - \eta , \frac{\hat m}{\sqrt{\hat v} + \epsilon}. $$

The √v̂ denominator gives each parameter its own effective learning rate — big-gradient parameters get smaller steps, small-gradient ones get larger — which is why Adam "just works" on transformers with almost no tuning. t is the 1-based step count (for bias correction); defaults β = (0.9, 0.999), ε = 1e-8.

4. AdamW — decoupled weight decay. Weight decay shrinks weights toward zero to regularize. The "obvious" way (plain Adam with L2) adds λθ to the gradient — but then Adam's adaptive √v̂ denominator rescales that decay differently per parameter, so the regularization you asked for isn't the regularization you get. AdamW decouples it: do the Adam step and, separately, shrink the weight directly:

$$ \theta \leftarrow \theta - \eta \lambda \theta \quad (\text{decoupled decay}), \qquad \theta \leftarrow \theta - \eta \frac{\hat m}{\sqrt{\hat v}+\epsilon}. $$

This small change measurably improves generalization and is the default optimizer for training transformers. Our adamw_step implements exactly this — and a test proves that with zero gradient the weight still shrinks (decay acting alone).

Common misconception. "AdamW is just Adam with L2 weight decay." No — the whole point is that L2-in-the-gradient and decoupled decay are not equivalent under an adaptive optimizer. The "W" is a genuine algorithmic fix, which is exactly why interviewers ask the difference.

Chapter 7: Learning-Rate Schedules — Warmup + Cosine

Why the LR shouldn't be constant. Early in training the weights are random and gradients are wild; a big step can send the loss to NaN. Late in training you want small steps to settle into a minimum rather than bounce around it. So the canonical schedule has two phases:

• Linear warmup — ramp η from ~0 up to base_lr over the first few hundred/thousand steps. This is especially important for Adam, whose 2nd-moment estimate is noisy and unreliable at the very start (t small) — a warmup hides that fragile early period.
• Cosine decay — after warmup, anneal η from base_lr down toward 0 following a half-cosine:

$$ \eta(t) = \eta_{\max} \cdot \tfrac12\Big(1 + \cos\big(\pi \cdot \text{progress}\big)\Big), \quad \text{progress} \in [0,1]. $$

 lr
  base ┤        ____
       │     __/    \__
       │    /          \__
       │   / warmup        \___ cosine → ~0
       └──┴──────────────────────► step

Cosine is smooth (no abrupt drops to destabilize Adam's moments) and spends a lot of time at low LR near the end, which empirically finds better minima. Our lr_schedule implements both branches; the tests check the warmup ramp, the peak at base_lr, and that it decays to ~0.

Common misconception. "Warmup is a minor trick." Skip warmup on a large transformer with Adam and you can diverge in the first few hundred steps. It's load-bearing, not decorative.

Chapter 8: Gradient Clipping and Loss Spikes

The problem. Occasionally a batch produces a huge gradient — a hard example, a numerical edge, a rare token. One oversized step can launch the weights into a bad region and spike the loss (you watch it shoot from 2.0 to 9.0 on a dashboard). On recurrent/transformer training this is common enough to need a standing defense.

Global-norm clipping. Treat all parameter gradients as one long vector, compute its L2 norm, and if it exceeds a threshold, scale the whole vector down so its norm equals the threshold:

$$ \text{norm} = \sqrt{\textstyle\sum_i g_i^2}, \qquad \text{if norm} > \tau:; g_i \leftarrow g_i \cdot \frac{\tau}{\text{norm}}. $$

Scaling the whole vector preserves its direction (so you still descend correctly) but caps its length (so no single step is catastrophic). Our clip_grad_norm does exactly this and returns the pre-clip norm — which is itself a useful signal: log it, and a sudden spike in the gradient norm is your early warning that a loss spike is coming.

The debugging ladder for a loss spike. When loss explodes mid-run, check, in order: (1) is the LR too high after warmup? (2) is gradient clipping on, and is the pre-clip norm spiking? (3) is this fp16 overflow producing inf/NaN (Chapter 9)? (4) is it a single pathological batch? The engineer who logs gradient norm and LR finds the cause in minutes.

Common misconception. "Clipping changes the optimum / hurts the model." Global-norm clipping preserves gradient direction; it only bounds step length on the rare giant-gradient step. It's a stability tool, not a regularizer that distorts where you converge.

Chapter 9: Numerical Precision, Mixed Precision, Gradient Accumulation

This chapter is conceptual (the lab is fp32 stdlib floats), but the ideas are senior table-stakes.

The precision menu.

Why range matters more than precision for training. fp16's 5 exponent bits mean small gradients underflow to 0 (they vanish) and large activations overflow to inf. bf16 keeps fp32's exponent, so it almost never overflows/underflows — you trade mantissa precision (usually fine, since gradients are noisy anyway) for range. This is why modern accelerators train in bf16 and why it "just works" where fp16 needs babysitting.

Mixed precision. Do the heavy matmuls in fp16/bf16 (fast, half the bandwidth — recall Phase 00's roofline: decode is bandwidth-bound) but keep a master copy of the weights and the optimizer state in fp32, and accumulate sensitive reductions in fp32. You get most of the speed/memory win without the accuracy loss.

Loss scaling (the fp16 fix). Because fp16 gradients underflow to zero, multiply the loss by a big constant S before backward (which scales every gradient by S, lifting them out of the underflow zone), then divide the gradients by S before the optimizer step. Dynamic loss scaling auto-tunes S: raise it when things are fine, halve it when an inf/NaN appears. bf16's range usually makes loss scaling unnecessary — another reason it's preferred.

Gradient accumulation. Want an effective batch of 256 but only 32 fits in memory? Run 8 micro-batches, summing (accumulating) their gradients without stepping, then take one optimizer step and zero the grads. Mathematically it's a 256-batch update; in memory it's a 32-batch footprint. This is the same grad-accumulation behavior (+=) you implemented in backward — now exploited deliberately. The cost is time (more forward/backward passes per step), not memory.

Common misconception. "fp16 and bf16 are interchangeable 16-bit formats." They are not — bf16 trades precision for range and is far more forgiving for training; fp16 needs loss scaling to avoid gradient underflow. Naming which one and why is the senior answer.

Chapter 10: This Is PyTorch Autograd

Everything you built maps one-to-one onto the real framework — same algorithm, just vectorized.

The one thing to internalize: a torch.Tensor with requires_grad=True is a Value whose .data is a multi-dimensional array instead of one float, and whose _backward operates on whole tensors instead of scalars. The reason GPUs help is not a different algorithm — it's that one node is now a matmul over millions of numbers that the hardware does in parallel. You built the real thing; production just made the nodes bigger.

Lab Walkthrough Guidance

The lab (lab-01-autograd-engine) turns this guide into code. Suggested order (matches the file):

1. Value core ops — __add__, __mul__, __pow__, exp/log/tanh/relu (Ch 2–3). For each, set out.data and an out._backward closure that does parent.grad += local · out.grad. Use += everywhere.
2. Derived ops — __neg__/__sub__/__truediv__ and the r-variants, all in terms of the core ops (no new _backward needed).
3. backward() (Ch 4) — build the reverse topological order with a visited set, seed self.grad = 1.0, walk reversed. Hand-check z = x*y + x first, then make the finite-difference soul test pass — that single green test means your chain rule is correct.
4. softmax / cross_entropy (Ch 5) — subtract the max; build CE on the stable softmax; verify the gradient is p − y.
5. Neuron/Layer/MLP — deterministic init from random.Random(seed); the final layer is linear (logits). Check the parameter count.
6. clip_grad_norm, sgd_step, adamw_step, lr_schedule (Ch 6–8) — clip returns the pre-clip norm; AdamW does decoupled decay and the bias-corrected Adam step.
7. train (the soul training test) — zero grads, sum CE over the 4 XOR examples, one .backward(), optional clip, optimizer step, record loss. It must strictly decrease and be deterministic for a fixed seed.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution (31 tests).
• You can derive dz/dx, dz/dy for z = x*y + x and match .backward().
• The finite-difference test is green — your analytic gradients match central differences on a gnarly composite (the proof your engine is correct).
• You can explain why a reused node accumulates its gradient (and why that forces zero_grad).
• You can write a stable softmax (max-subtraction) and derive the (p − y) CE gradient from memory.
• You can state, without notes, the SGD → momentum → Adam → AdamW ladder and what decoupled weight decay fixes.
• Your XOR training loss is strictly decreasing and deterministic for a fixed seed.

Interview Q&A

• "What does loss.backward() actually do?" — Walks the computation graph in reverse topological order, applying the chain rule: each node adds its local-derivative × upstream-grad into its parents' .grad. One pass fills every parameter's gradient because the loss is a single scalar (reverse-mode).
• "Why does reverse-mode give all gradients in one pass, and when is forward-mode better?" — Reverse-mode is O(1) passes for many-inputs-one-output (nets: N params, 1 loss). Forward-mode is O(inputs) but O(1) in outputs, so it wins when inputs are few and outputs many (the opposite shape).
• "Derive the softmax + cross-entropy gradient." — p_i = softmax(z)_i, CE = −log p_target; ∂CE/∂z_i = p_i − y_i (probabilities minus one-hot). Near-zero when confident-and-right, large when confident-and-wrong.
• "Why subtract the max in softmax?" — e^{large} overflows; subtracting max makes every exponent ≤ 0 with identical math (the factor cancels). Without it, large logits give inf/NaN.
• "SGD vs Adam vs AdamW?" — SGD: one global LR. Adam: per-parameter LR via the 2nd-moment estimate, plus momentum and bias correction. AdamW: Adam with decoupled weight decay (shrink the weight directly rather than adding L2 to the gradient, which the adaptive denominator would distort). AdamW is the transformer default.
• "Why bias-correct in Adam?" — m, v start at 0, so early estimates are biased toward zero; dividing by 1 − β^t corrects it, which matters most in the first steps (and is why warmup helps).
• "Why warmup + cosine?" — Warmup avoids early divergence while Adam's moment estimates are noisy; cosine anneals the step smoothly so late training settles into a minimum.
• "Your loss spiked at step 4000 — what do you check?" — LR too high post-warmup; gradient norm spiking (clip it, log the pre-clip norm); fp16 overflow → NaN (loss scaling / bf16); a pathological batch. Global-norm clipping caps the giant step without changing direction.
• "fp16 vs bf16?" — Same size; bf16 keeps fp32's exponent range (rarely overflows/underflows, modern training default), fp16 has more mantissa but a tiny range so it needs loss scaling to stop gradient underflow.
• "What is gradient accumulation and why use it?" — Sum gradients over several micro-batches before one optimizer step to simulate a larger batch within a small memory budget; it's the same += accumulation, used deliberately. Costs time, not memory.
• "Why must you zero_grad() every step?" — .grad accumulates by design (so shared nodes and micro-batches sum correctly); without zeroing, this step's gradient piles onto the last step's.
• "How do you know your autograd is correct?" — gradcheck: compare analytic gradients against central finite differences. If they match on a nontrivial composite, the engine is right.

References

• Rumelhart, Hinton & Williams, Learning representations by back-propagating errors (1986) — the backprop paper.
• Baydin et al., Automatic Differentiation in Machine Learning: a Survey (2018) — forward vs reverse mode, the definitive overview.
• Andrej Karpathy, micrograd (github.com/karpathy/micrograd) and Neural Networks: Zero to Hero — the spiritual source of this lab; watch "The spelled-out intro to neural networks and backpropagation."
• Kingma & Ba, Adam: A Method for Stochastic Optimization (2015).
• Loshchilov & Hutter, Decoupled Weight Decay Regularization ("AdamW", 2019).
• Micikevicius et al., Mixed Precision Training (2018) — fp16, master weights, loss scaling.
• PyTorch docs — Autograd mechanics, torch.optim.AdamW, torch.nn.utils.clip_grad_norm_, torch.cuda.amp (automatic mixed precision), and torch.autograd.gradcheck.
• Goodfellow, Bengio & Courville, Deep Learning (2016), ch. 6.5 (backprop) and ch. 8 (optimization).