Training Loops & Optimization
The engine room of learning
A training loop is forward → loss → backward → update, repeated in epochs over batches. Optimizers (SGD, Adam), batch size, LR schedules and gradient clipping turn that loop from fragile to production-usable.
What you will learn
- Describe one step of the training loop end to end
- Compare SGD, Adam and modern variants
- Choose batch size and LR schedule with evidence
What is this?
Training loops iterate: sample a batch, run the forward pass, compute loss, backpropate gradients, apply an optimizer update, repeat. Epochs are full passes over data; the optimizer, batch size, learning-rate schedule and clipping turn this mechanical loop into a stable, convergent process. Almost every 'my fine-tune diverged' story is a loop-parameter story.
Why do you need it?
- Batch size trades gradient noise for throughput and memory.
- Adam's per-parameter scaling makes it the default for deep/fine-tuned networks.
- Warm-up and decay schedules stabilize training at scale.
- Clipping bounds worst-case gradient steps, protecting against one bad batch.
Mental model
An exercise routine with reps and sets
The training loop
Each rep (batch) improves a little; each set (epoch) reviews the whole routine. Rest days (drops in LR), pacing (schedule) and a spotter (clipping) keep the athlete improving without injury.
How it works
- Sampler loads a batch; loss computed from predictions; backward fills gradients.
- The optimizer adjusts weights: SGD uses raw gradient steps; Adam scales per-parameter by past gradient statistics.
- An LR scheduler warms up then decays, often cosine-style, over the run.
- Gradient clipping caps the update norm, preventing single-batch blowups.
- Checkpoints save each epoch so a crash resumes instead of restarting.
One training step.
In production
- Fine-tuning on GPUs uses AdamW + linear warmup + cosine decay, with gradient clipping at 1.0.
- A distributed job uses gradient accumulation (same optimizer visibility with small memory batches).
- Resumable checkpoints + validation snapshot every N steps drive long stable runs.
Best practices
- Log loss, LR and gradient norm every step batch; inspect weekly, not yearly.
- Use Adam(W) by default; plain SGD with momentum only when you've measured it better.
- Warm up LR over early steps for deep nets.
- Make runs resumable from checkpoints from day one.
Common mistakes
- Very large batch without LR scaling — a common diverging recipe.
- No gradient clipping on a model that occasionally spikes.
- Judging the run from a single step's loss noise.
- Calling 'done' on training loss before checking the validation curve.
Common issues
Divergence to NaNLearning rate too high or unstable layer scales
Lower LR, clip gradients, warm up, check input feature scale.
Slow-but-steady lossLR too low or a plateau
Try LR warm-up, restart with higher peak LR, or cyclic decay.
Checkpoint resume changes resultsRNG/sampler state not saved
Save RNG and sampler epoch/offset in the checkpoint.
Debugging tips
- Loss doesn't move on the first batch — wiring is wrong; test with one-token/one-class sanity.
- Spiky validation + flat training typically = LR too high near a plateau.
- GPU memory plateau — lower batch size and use gradient accumulation.
Interview checkpoint
Hands-on lab
- Implement forward, loss, gradient descent by hand
- Watch the learning-rate effect on convergence
- Verify batching and epochs change nothing logically
Hand-Roll a Training Step
30 min · intermediate
Quick quiz
What does gradient clipping do?
Adam is a reasonable default optimizer for deep networks.