Optimizers and schedulers: upgrades to the step
Beat four of the loop — "nudge parameters against the gradient" — has forty years of engineering behind it. You need the two ideas that matter and the two names you'll actually type.
Idea 1: momentum
Run the editor. The loss surface is a narrow valley — gentle along one axis, steep along the other (real networks are full of these). Plain descent (momentum 0) zigzags: the steep direction bounces it side to side while the gentle direction inches along. Momentum keeps a running velocity: zigzag components cancel across steps, the consistent direction compounds. Watch the output: momentum 0.5 lands two orders of magnitude closer to the optimum than plain descent — and momentum 0.9 overshoots and rings, ending farther away than either. Momentum is a knob with a sweet spot, not a free lunch, and the demo shows both sides.
Idea 2: per-parameter step sizes
Different parameters see wildly different gradient scales (an embedding for a rare token vs a bias touched every step). Adaptive optimizers track each parameter's recent gradient magnitudes and scale its steps individually — big steps for quiet parameters, careful steps for loud ones.
Adam combines both ideas — momentum plus per-parameter scaling — and is the sensible default across deep learning (its weight-decay-corrected variant AdamW is standard for transformers). Plain SGD (+momentum) remains common in vision and as the thing Adam is compared against.
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
One note worth stating plainly: optimizer choice is a default-then-diagnose decision, not a leaderboard to agonize over. Adam with a sane lr gets you to "training works"; everything after that is curves.
Schedulers: the lr as a plan, not a constant
Last lesson's dilemma — big steps early would help, small steps
late would help — has an obvious resolution: change the lr over
training. A scheduler does it on a plan: warm up from tiny to
peak over the first steps (protects the fragile random-init
phase), then decay — often along a cosine — toward near-zero for
fine-grained settling. In torch, schedulers wrap the optimizer and
advance with scheduler.step().
Where AI specifically gets this wrong
- Exotic optimizers on day one. Generated code sometimes reaches for something fashionable when the actual problem is the lr or the data. Adam + lr sweep first.
- Scheduler misconfigured against epochs vs steps. A cosine schedule meant for 10,000 steps, stepped once per epoch, never leaves warmup. Print the lr each epoch — one line, catches it.
- Comparing optimizers at one fixed lr. Each optimizer's best lr differs; a single-lr bake-off measures the lr, not the optimizer.