9  Optimization Algorithms

Every model in this book has been trained by a line we have not yet examined. Since fitting our first models in Chapter 2 we have constructed an optimizer — SGD at first, Adam soon after — handed it the parameters and a learning rate, and let the training loop invoke it once per minibatch, trusting the loss to fall. It nearly always did. This chapter opens that box. Inside is not one idea but three decisions, made anew by every algorithm: a descent direction — which way counts as “down”, a question whose answer turns out to depend on which norm measures the size of a step; a step size over time — how boldly to move, and how that boldness should change over a run; and a way of living with noise — since almost every gradient is estimated from a minibatch rather than computed exactly. The twelve sections of this chapter are these three decisions unfolded, roughly in the order history made them.

Two properties of the loss surface make the decisions consequential. The first is curvature: a deep network’s loss rises steeply along some directions of parameter space and barely at all along others, and a single step size must serve both — too bold and the steep directions oscillate out of control, too timid and the flat directions never arrive. The second is noise: the exact gradient costs a full pass over the dataset, so any method that scales settles for a minibatch estimate whose variance is ours to choose. Section 9.1 maps this terrain. The five sections after it climb the classical ladder, each rung repairing a failure of the one below: gradient descent and the ideal of preconditioning (Section 9.2), stochastic gradients and why learning rates must decay (Section 9.3), minibatching and what the hardware has to say about it (Section 9.4), momentum against curvature (Section 9.5), and per-coordinate scaling from AdaGrad through RMSProp to Adam (Section 9.6) — where we also build the tiny transformer language model on which the rest of the chapter runs its experiments.

The ladder’s top rung dates from 2014, and for years textbooks stopped there. Practice did not. The second half of the chapter is the layer a practitioner actually configures today: decoupled weight decay, the rule of what not to decay, and the memory arithmetic of optimizer state (Section 9.7); warmup, cosine decay, and the warmup–stable–decay schedules of large-model training (Section 9.8); the realization that steepest descent depends on the norm — under one norm it recovers sign descent and, in essence, Adam, and under the spectral norm it yields Muon, the first credible challenger to Adam’s decade (Section 9.9); how large a batch can grow before more parallelism stops buying anything (Section 9.10); how to tune a small model and transfer the result to one too expensive to tune (Section 9.11); and the craft of running real training — recipes, gradient clipping, weight averaging, and how to sweep (Section 9.12).

A word on what this chapter is not. It states results, demonstrates phenomena, and gives intuition; it does not prove. The descent lemma, the condition-number law, momentum’s \(\sqrt{\kappa}\) acceleration, the Robbins–Monro conditions, Adam’s bias correction, and the convex analysis underneath all of them live in the optimization chapter of the mathematical appendix (Chapter 25), which is written as this chapter’s proof volume. Where a section here says “one can show”, the appendix shows it. The two chapters can be read in either order; the experiments here give the theorems something to predict.

The modern layer is younger than it looks, and parts of it are still moving. Decoupled weight decay was published in 2017 and became a universal default only years later; warmup–stable–decay schedules entered large-model practice around 2024; Muon went from a speed-run leaderboard in late 2024 to reported trillion-parameter production runs within about a year — for optimizers, unusually fast. The field has also grown a benchmarking discipline: matched-tuning comparisons and public benchmarks now routinely shrink headline claims, and more than one celebrated optimizer has failed to replicate under them. Where the evidence is unsettled we say so in place, and the chapter’s comparisons follow one rule throughout: tuned against tuned, never a tuned challenger against a default baseline.

Resources and Further Reading

The references below follow the chapter’s arc — the convex foundations, the classical ladder, the modern layer, and the craft of tuning. All are freely accessible online except where noted. The optimization chapter of the mathematical appendix (Chapter 25) keeps its own resource list for the theory side — convex-optimization texts and courses with proofs — and we do not repeat those entries here.

Books

• Convex Optimization — Boyd & Vandenberghe — free PDF; the standard reference behind the vocabulary this chapter uses informally — conditioning, convergence rates, duality, projections — and the right place to see the analyses that Section 9.2 and Section 9.3 state without proof done properly.
• Numerical Optimization — Nocedal & Wright — the classical machinery in full: line search, trust regions, and the quasi-Newton methods that Section 9.2 names as the ideal deep learning cannot afford (paywalled, noted; widely held in university libraries).

Courses and video lectures

• Stanford CS336: Language Modeling from Scratch — Assignment 1 — free; the graded version of this chapter’s exercises: implement AdamW exactly as Section 9.7 does, account for optimizer-state memory byte by byte, and run the learning-rate and batch-size sweeps that Section 9.12 turns into method; the accompanying lectures are on YouTube.

Foundational and current papers

• Old Optimizer, New Norm: An Anthology — Bernstein & Newhouse (2024) — free; the unification that organizes Section 9.9: SGD, sign descent/Adam, and Shampoo are each steepest descent under a different norm, which turns the chapter’s zoo of methods into one question asked three ways.
• An Empirical Model of Large-Batch Training — McCandlish et al. (2018) — free; defines the gradient-noise scale and the critical batch size, the two quantities measured at the center of Section 9.10, and predicts when doubling the batch stops halving the steps.
• Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective — Wen et al. (2024) — free; the modern upgrade of the ill-conditioned valley of Section 9.1: a river-valley landscape in which the stable phase travels along the river and the decay phase descends its bank, explaining the WSD loss cliff of Section 9.8.
• Fantastic Pretraining Optimizers and Where to Find Them — Stanford (2025) — free; re-benchmarks ten optimizers under matched tuning and watches most claimed speedups over AdamW shrink — the fair-comparison discipline that Section 9.9 and Section 9.12 adopt as a rule.
• Benchmarking Neural Network Training Algorithms — Dahl et al. (2023) — free; the MLCommons AlgoPerf benchmark (code and results): why optimizer verdicts depend on the comparison protocol, the evidence standard behind the caveats of Section 9.9 and Section 9.12.

Tutorials, notes, and interactive

• Why Momentum Really Works — Gabriel Goh, Distill (2017) — free, interactive; the damping and acceleration story of Section 9.5 with sliders for \(\eta\) and \(\beta\) — the fastest way to internalize why the critical \(\beta\) exists and what ringing looks like.
• An Overview of Gradient Descent Optimization Algorithms — Sebastian Ruder (2016) — free; the field guide to the classical ladder of Section 9.3 through Section 9.6, and a historical marker: everything a practitioner needed in 2016, and a measure of how much the modern layer has added since.
• Deep Learning Tuning Playbook — Godbole et al., Google Research — free; the scientific/nuisance/fixed-hyperparameter methodology and budget-tiered sweeps that Section 9.12 teaches, from the team that ran them at production scale.
• Muon: An Optimizer for Hidden Layers in Neural Networks — Keller Jordan (2024) — free; the original post: design decisions, Newton–Schulz coefficients, and ablations behind the optimizer that Section 9.9 builds from scratch.
• modded-nanogpt — Keller Jordan et al. — free; the speedrun repository where Muon first proved itself, with every record documented and reproducible — the evidence culture that Section 9.9 holds up as a model.
• Deriving Muon — Jeremy Bernstein — free; a compact derivation of Muon from the steepest-descent-under-a-norm principle, the note-form companion to the derivation in Section 9.9.
• The Practitioner’s Guide to the Maximal Update Parameterization — EleutherAI — free; muP implemented step by step with the coordinate-check experiments of Section 9.11, including the failure modes a first implementation actually hits.