Training Loops & Optimization
Backpropagation, the training loop step by step, and how SGD, Adam, and AdamW actually differ — including the bug AdamW specifically fixed
Training Loops & Optimization
TL;DR
Every training step follows the same four moves: compute predictions (forward pass), measure the loss, compute gradients via backpropagation (loss.backward()), and update parameters against those gradients (optimizer.step()). Which optimizer does that last step, and how large a step (the learning rate), are the decisions that most determine whether training actually works.
| Property | Value |
|---|---|
| Level | Intermediate |
| Reading time | ~20 minutes |
| Prerequisites | Neural Network Fundamentals |
| You will understand | The full training loop, backpropagation conceptually, and how to choose an optimizer and learning rate schedule |
The Training Loop, Line by Line
for batch in dataloader:
optimizer.zero_grad() # 1. clear old gradients
predictions = model(batch.inputs) # 2. forward pass
loss = loss_fn(predictions, batch.labels) # 3. compute the loss
loss.backward() # 4. backward pass — compute gradients
optimizer.step() # 5. update parametersoptimizer.zero_grad() is not a formality. PyTorch accumulates gradients into .grad by default every time .backward() is called — it does not clear them automatically. Skip zero_grad() and each step's gradients get added on top of the previous step's, silently corrupting training in a way that often doesn't crash, just quietly trains worse. This is one of the most common real bugs in a hand-written training loop.
Backpropagation: Where the Gradient Comes From
loss.backward() computes the gradient of the loss with respect to every parameter in the model, using the chain rule from calculus, automatically — this is what "autodiff" means.
A tiny worked example makes the mechanism concrete. Suppose:
z = w * x (a linear operation)
y = relu(z) (a nonlinearity)
loss = (y - target)^2To find how much loss changes per unit change in w, the chain rule multiplies the local derivatives along the path from w to loss:
d(loss)/dw = d(loss)/dy * dy/dz * dz/dwEach factor is a local derivative — easy to compute in isolation — and the chain rule is exactly the rule for combining them into the global effect of w on the final loss. PyTorch's autodiff engine builds a graph of every operation during the forward pass, then walks it backward applying exactly this rule at every node, for every parameter, automatically.
x = torch.tensor(2.0)
w = torch.tensor(0.5, requires_grad=True)
target = torch.tensor(3.0)
z = w * x
y = torch.relu(z)
loss = (y - target) ** 2
loss.backward()
print(w.grad) # d(loss)/dw, computed automatically via the chain ruleOptimizers: What Actually Happens on .step()
SGD, Momentum, and Adam/AdamW
SGD (Stochastic Gradient Descent)
The base case: param -= learning_rate * param.grad. Simple, but every parameter moves at the same fixed rate regardless of how noisy or how steep its gradient has been recently.
SGD with momentum
Accumulates a running average of past gradients and moves in that averaged direction instead of the raw current gradient — smooths out noisy updates and helps the optimizer keep moving through flat or bumpy regions of the loss surface instead of stalling.
Adam / AdamW
RecommendedTracks a per-parameter adaptive learning rate, based on running estimates of both the gradient's mean and its variance — parameters with consistently large gradients get smaller effective steps, and vice versa. This adaptivity is why Adam-family optimizers are the default for training transformers, largely without needing careful per-parameter tuning.
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)The Bug AdamW Specifically Fixed
The original Adam combined its adaptive updates with L2 regularization by adding a weight_decay * param term directly into the gradient before Adam's adaptive scaling was applied. That's a real bug, not just a stylistic choice: because Adam divides updates by a running estimate of gradient magnitude, folding weight decay into the gradient means the effective amount of decay applied to each parameter gets warped by that same adaptive scaling — parameters with large gradients get proportionally less decay than intended, and vice versa.
AdamW fixes this by applying weight decay directly to the parameters, decoupled from the gradient-based adaptive update entirely:
# Conceptually, AdamW's update separates into two independent pieces:
# 1. The adaptive Adam update, from the gradient alone
param -= lr * adam_adaptive_step(grad)
# 2. Weight decay, applied directly and uniformly — not run through Adam's scaling
param -= lr * weight_decay * paramThis is why virtually every modern model is trained with AdamW, not Adam. It's not a minor variant — it's the same optimizer with a specific, well-understood bug in how it interacted with regularization fixed, which is why the field converged on it as the default almost universally once the fix was published.
Learning Rate Schedules
The learning rate rarely stays constant through training:
A standard warmup + cosine decay schedule
Warmup exists because early gradients are unreliable. At the very start of training, the model's parameters are essentially random, so the first gradients computed can be unusually large or noisy — taking a full-sized optimizer step on that noisy signal risks destabilizing training before it even gets going. Ramping the learning rate up gradually gives the model a few steps to settle into a more reasonable region before taking full-sized steps.
Gradient Clipping
If a gradient's magnitude spikes — from an unstable batch, a numerical edge case, or just bad luck — the resulting parameter update can be enormous, throwing training into a much worse region or producing NaN values outright. Gradient clipping caps the gradient's overall norm before the optimizer step:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()If the combined gradient norm across all parameters exceeds max_norm, every gradient is scaled down proportionally so the norm lands exactly at the cap — direction is preserved, only magnitude is reduced.
Batch Size and Its Trade-off
Each batch's gradient is an estimate of the true gradient over the whole dataset — a small batch gives a noisier estimate, a large batch a smoother one, at proportionally higher compute cost per step.
| Batch size | Gradient estimate | Trade-off |
|---|---|---|
| Small | Noisier | Cheaper per step, but noise can act as a mild regularizer, and can also make training less stable at a fixed learning rate |
| Large | Smoother, closer to the true gradient | More expensive per step, generally tolerates (and often needs) a higher learning rate to make comparable progress per epoch |
Learning rate and batch size aren't independent choices — a learning rate tuned for one batch size usually needs adjusting if the batch size changes significantly. This is why scaling up training (more GPUs, bigger batches) is rarely a pure "just add more compute" change without revisiting the learning rate too.
Concept Checks
Check yourself
What actually goes wrong if optimizer.zero_grad() is accidentally omitted from a training loop?
Gradients accumulate across steps instead of being freshly computed each time, because PyTorch adds newly computed gradients into .grad rather than replacing them by default. The optimizer then updates parameters based on a sum of gradients from multiple steps rather than the current step alone, which corrupts training in a way that usually doesn't crash — it just trains noticeably worse without an obvious error message.
Why did AdamW replace Adam as the default for training transformers, rather than being just a stylistic preference?
Because Adam's original way of combining weight decay with its adaptive update had a real bug: folding weight decay into the gradient before Adam's adaptive scaling meant the effective decay applied to each parameter got distorted by that same per-parameter scaling. AdamW fixes this by applying weight decay directly to the parameters, decoupled from the adaptive gradient update, which is a correctness fix rather than a preference.
Why does a learning rate warmup phase matter more early in training than later?
Because the model's parameters start essentially random, so the very first gradients computed can be unusually large or noisy — taking a full-sized step on that unreliable signal risks destabilizing training before it's found any reasonable region of parameter space. Ramping the learning rate up gradually gives the model a chance to settle before full-sized updates begin.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
| The training loop | Forward → loss → backward → optimizer step → zero_grad, repeated |
zero_grad() | Required — gradients accumulate by default unless explicitly cleared |
| Backpropagation | The chain rule, applied automatically across the whole computation graph |
| SGD / momentum / Adam | Increasingly adaptive ways of turning a gradient into a parameter update |
| AdamW | Decouples weight decay from Adam's adaptive scaling — fixes a real bug in Adam+L2 |
| Warmup | Ramps the learning rate up gradually, since early gradients are unreliable |
| Gradient clipping | Caps the gradient norm to prevent destabilizing updates |
| Batch size / learning rate coupling | A learning rate tuned for one batch size usually needs revisiting if batch size changes |
Next
With training mechanics covered, the next page builds the architecture underneath every modern language model: The Transformer Architecture.
Neural Network Fundamentals
nn.Module, the layers that make up a network, and why a nonlinearity between them is what makes depth mean anything
The Transformer Architecture
Attention built up from scratch — Query/Key/Value, the scaling factor, multi-head attention, positional encoding, and a full transformer block