Training loops, optimizers and model debugging

Why it matters

Training failures are usually quiet: the loss falls a little, metrics look plausible, and the model is still wrong because of a masking bug, misaligned labels or a learning rate off by a factor of ten. Interviewers ask you to read a loss curve or justify optimizer choices because it reveals whether you have trained models yourself and whether you debug methodically.

Core ideas

Loop anatomy. Each step: set train mode, fetch a batch, forward, compute the loss, zero gradients, backward, optionally clip, optimizer.step(), scheduler.step(). Evaluation uses eval mode (dropout off, batch norm on running statistics) and no gradient tracking. Forgetting to zero gradients accumulates them across steps; forgetting eval mode makes validation noisy and wrong. Gradient accumulation deliberately sums gradients over several micro-batches before one optimizer step to emulate a larger batch; scale the loss accordingly.

Optimizers.

• SGD with momentum keeps a velocity v <- mu * v + g and steps along it, damping oscillation across steep directions and accelerating along consistent ones.
• Adam keeps exponential moving averages of the gradient and of the squared gradient, corrects both for their zero initialization, and divides each parameter's step by the square root of its second-moment estimate. That is two extra tensors per parameter, so optimizer state is about twice the model size at the same precision.
• AdamW applies weight decay directly to the weights instead of adding an L2 term to the gradient, where Adam's per-parameter scaling would weaken the decay for parameters with large gradients. It is the usual choice for transformers, often with biases and normalization weights excluded from decay.

Learning rate. The most sensitive hyperparameter: sweep it on a log scale and watch the early loss. Warmup ramps the rate up over the first steps, because early second-moment estimates are noisy and the network is far from a stable regime, so full-size steps can diverge or spike. Afterwards, decay (cosine or linear) lets the model settle. Changing the batch size usually means retuning the learning rate.

Gradient hygiene. Clip by global norm (rescale all gradients together when their joint norm exceeds a threshold) to survive rare spikes without distorting every update. Log the gradient norm: sudden jumps often precede divergence, and near-zero norms in early layers indicate vanishing gradients.

Checkpointing. Save model weights, optimizer state, scheduler state, the step counter, RNG states and the data-loader position, written atomically (temporary file, then rename). Resuming from weights alone silently resets Adam's moments and the schedule, which shows up as a loss spike. Keep the best checkpoint by validation metric as well as the latest.

Debugging recipe.

1. Check the initial loss: for C balanced classes it should be near ln(C), about 2.3 for 10 classes. Far off means a bug in the loss, labels or initialization.
2. Overfit one small batch to near-zero loss. If you cannot, the problem is in the model, loss or optimizer wiring, not in generalization.
3. Look at real inputs and labels exactly as the model receives them after the data pipeline: augmentation, normalization, label alignment, padding and masks.
4. Check shapes and broadcasting: subtracting a (n, 1) tensor from an (n,) tensor silently yields (n, n).
5. Compare against a trivial baseline, change one thing at a time, and log every run's configuration and seed.

Reproducibility. Seed Python, NumPy and the framework; even then some GPU kernels are nondeterministic, so measure seed-to-seed variance before attributing a small difference to a change.

Common misconceptions

• "Adam does not need learning-rate tuning." It is less sensitive than SGD, but the learning rate and schedule still dominate results.
• "Adam with L2 is the same as AdamW." The L2 gradient gets rescaled by Adam's adaptive denominator; decoupled decay does not.
• "Clipping fixes a learning rate that is too high." It caps rare spikes; if clipping triggers on most steps, the step size or the model is wrong.
• "If the training loss falls, the pipeline is correct." Networks can fit shuffled labels or leaked features; validate the data and compare with baselines.
• "Saving the weights is a checkpoint." Without optimizer, scheduler and RNG state, resuming is not the same run.
• "One run shows the improvement." Differences smaller than seed-to-seed variation are not evidence.

How to approach it

Write a training loop from scratch at least once, including checkpointing, and verify exact resumption by comparing loss curves with an uninterrupted run. Inject bugs one at a time (missing zero_grad, eval mode forgotten, labels shifted by one, learning rate times 100) and learn their symptoms. Senior-level discussions add loss spikes at scale, batch-size scaling and bisecting a training regression.