JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

The Optimization Loop: Optimizer, Schedule, Precision

Inside every training step is the same four-beat rhythm. Meet AdamW, the warmup-decay schedule, mixed precision, and the loss spikes that haunt engineers at 3am.

The four-beat rhythm of a step

However giant the model, every training step is the same loop. (1) Feed a batch of token sequences through the network. (2) Compute the cross-entropy loss against the true next tokens. (3) Run backpropagation to get a gradient — the direction in weight-space that would reduce the loss. (4) Let the optimizer nudge every weight a small step along that direction. Repeat hundreds of thousands of times, and the loss curve grinds downward.

The four-beat training step: feed a batch forward, compute the loss, backpropagate the gradients, then update the weights.

A loop diagram: a batch flows through the model to a prediction, a loss is computed, and an update feeds back to the model.

Three choices govern whether this loop is stable and fast: which optimizer you use, how you schedule the step size over time, and what numerical precision you compute in. Get them right and the run is boring (the best kind). Get them wrong and you waste a fortune watching the loss explode.

AdamW: the default optimizer

Almost every LLM is trained with AdamW. Plain gradient descent steps every weight by the same global rule, which is fragile when some weights need big updates and others need tiny ones. Adam fixes this by giving each weight its own adaptive step size, estimated from a running average of its recent gradients (the first moment) and their variance (the second moment). The result is far more robust on the messy loss landscapes of deep transformers.

An optimizer rolls downhill on the loss surface — AdamW adapts the step size per weight instead of using one global rule.

Interactive gradient descent: a point rolls down a curved loss surface toward the minimum.

The W in AdamW is decoupled weight decay: a gentle pull of every weight toward zero, applied separately from the gradient. It acts as regularization, keeps weights from drifting too large, and noticeably improves stability at scale. Doing decay the older, coupled way subtly breaks Adam's math — the decoupled version is now the standard.

The learning-rate schedule: warmup then decay

The single most important hyperparameter is the learning rate — the size of each step. But you don't keep it constant; you follow a warmup-then-decay schedule. For the first few thousand steps you warm up, ramping the learning rate from near zero up to its peak. Then you decay it slowly back down, usually along a cosine curve, over the rest of training.

Why both halves? Early on, the freshly initialized weights are fragile and the gradients are wild; a big step would knock the model into a bad region, so you warm up gently. Late in training, near a good minimum, big steps would bounce you around uselessly, so you decay to take ever-finer steps and settle in. The schedule is the difference between a clean descent and a diverging mess.

\eta_t=\eta_{\min}+\tfrac{1}{2}\,(\eta_{\max}-\eta_{\min})\left(1+\cos\!\left(\pi\,\frac{t-T_{\text{warm}}}{T-T_{\text{warm}}}\right)\right)

After a linear warmup to the peak, a cosine schedule decays the learning rate smoothly toward its floor.

step      learning rate
0         0.0          (start)
2000      3.0e-4       (peak, end of warmup)
100000    1.5e-4       (cosine decay, halfway)
200000    3.0e-5       (near end, ~10% of peak)
# warmup: linear 0 -> peak; then cosine peak -> floor
A warmup-then-cosine-decay schedule. The peak and warmup length are tuned per model.

Mixed precision: half the bits, twice the speed

Doing every calculation in full 32-bit floating point is accurate but slow and memory-hungry. Mixed-precision training does the heavy matrix multiplies in a 16-bit format (like bfloat16) while keeping a master copy of the weights and the loss accumulation in higher precision. You get roughly double the throughput and half the memory, for almost no loss in quality — which is why every large run uses it.

The catch is range. Older 16-bit formats could not represent very small gradient values and silently rounded them to zero, stalling learning. The bfloat16 format keeps the wide exponent range of 32-bit (trading away some fractional precision) and largely solves this, which is why it became the workhorse format for LLM pretraining.

Loss spikes: when a giant run wobbles

Train for weeks and you will eventually see it: the loss, descending smoothly, suddenly jumps upward — a loss spike. Sometimes it recovers on its own; sometimes it diverges and the run is ruined. Spikes come from a rare batch of pathological data, a numerical overflow, or a too-aggressive learning rate hitting a sharp part of the landscape. At billion-parameter scale they are a fact of life, not a bug you can fully eliminate.

  1. Gradient clipping: cap the gradient's overall size so one wild batch can't blow up the weights.
  2. Frequent checkpoints: save the full model state often, so a divergence costs hours, not weeks.
  3. Rewind and skip: on a bad spike, roll back to the last good checkpoint and skip the offending batch of data.
  4. Lower the peak / extend warmup: a slightly smaller learning rate trades a little speed for a lot of stability.
\mathbf{g}\leftarrow \mathbf{g}\cdot\min\!\left(1,\ \frac{c}{\lVert \mathbf{g}\rVert}\right)

Gradient clipping rescales any gradient whose norm exceeds the threshold c, taming the spikes that derail a long run.