STAGE 0 · 2.5.4 · 3 NODES

The optimizer: AdamW

Backprop hands you a gradient — the downhill direction for every weight. The optimizer decides how big a step to take in that direction. This is AdamW, the optimizer almost every large model is trained with: it gives each weight its own step size, and a schedule paces the whole run.

THE 3 STEPS

01Momentum smooths the descent 02Per-weight step sizes 03The learning-rate schedule

NODE 01 / 03

Plain steps, or smart steps?

You already saw a model roll downhill with one step size. Watch what happens on a stretched-out valley when SGD races Adam — same start, same 16 steps.

SGD takes the raw gradient step. In a narrow valley the step must stay small or it bounces off the steep walls — so it crawls along the flat floor.
Adam adds momentum. It averages recent gradients, so it builds speed down the valley instead of zig-zagging across it.
Same step budget, different progress. After 16 steps Adam is at the bottom; SGD is still working its way there.

GOOD TO KNOW · THE NAME

Why “AdamW”

Adam = Adaptive Moment estimation — it tracks the running moments (averages) of each weight's gradient. The W adds weight decay: a gentle pull of every weight toward zero on each step, which curbs overfitting. AdamW is the default optimizer for training large language models.

SGD vs ADAM · SAME LOSS SURFACE

Flip between the two — SGD bounces across the steep walls; Adam glides down the middle.

Plain SGD — step = learning rate × gradient, nothing else. Simple, but one step size has to serve every weight, so it stalls in stretched valleys.

AFTER 16 STEPS
OPTIMIZERSGD
DIST TO MIN0.32
STATUSstill descending
How close each optimizer got to the lowest-loss point in the same number of steps.

NODE 02 / 03

Every weight gets its own step size

Here's the trick that makes Adam adaptive. It watches how big and how jumpy each weight's gradients have been, and shrinks the step for the noisy ones.

One global step size is unfair. A weight with huge, jerky gradients and a weight with tiny, steady ones can't share the same step.
Adam divides each weight's step by its own recent gradient size — the root-mean-square, a plain average of how large the gradients have been. Big, noisy gradients → smaller, safer step.
The result: every weight moves at a sane, similar pace. Flip to Adam and both weights take a controlled step, whatever their raw gradient.

RAW STEP vs ADAM-SCALED STEP

Switch RAW to Adam-scaled — the noisy weight's giant step shrinks to match the steady one.

Raw step (SGD) — the step is just learning-rate × gradient, so the noisy weight lurches while the steady one barely moves.

NODE 03 / 03

Pacing the whole run: the schedule

The learning rate isn't one fixed number — it changes over training. Almost every run warms it up, holds it high, then decays it toward zero.

Warm up first. Start the rate near zero and ramp it up over the first few thousand steps, so early wild gradients don't blow up the weights.
Then decay. Lower the rate as training goes on, so the model settles into a good minimum instead of bouncing around it.
Two common shapes. Cosine glides smoothly down to a floor; WSD holds the peak flat, then drops fast at the very end.

LEARNING RATE OVER TRAINING

Drag the warmup and switch the decay shape — see how the curve paces the run.

WARMUP LENGTH6%

Cosine decay — after warmup, the rate follows a smooth half-cosine down to a small floor. The most common choice.

SCHEDULE
WARMUP6k / 100k
PEAK LR3e-4
DECAYS TO3e-5
A typical 100k-step run: how long the warmup lasts, the peak rate, and where the rate ends up.

EXPLAIN IT BACK

Backprop already gives the exact downhill direction for every weight. So why isn't a single, fixed learning rate enough?

NEXT: SUPERVISED FINE-TUNING

You can now train a model end to end: data → tokens → transformer → loss → backprop → optimizer. Next we stop building the base model and start shaping it — supervised fine-tuning, where a pretrained model learns to follow instructions from example answers.

← Back: backpropagation Continue to fine-tuning →

Language: English