9  Optimization Algorithms

Every model in this book has used an optimizer to update its parameters once per minibatch. This chapter examines the principles behind those updates. An optimization algorithm makes three related choices: the descent direction, the step size and its variation over time, and the treatment of noise in minibatch gradients. The appropriate descent direction depends on how the size of an update is measured; the step size must account for curvature and training time; and the gradient variance depends on batching and averaging.

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 slowly along others. A single step size must serve both, so a large step causes oscillation in steep directions while a small step makes slow progress in flat directions. The second is noise: an exact gradient costs a full pass over the dataset, so scalable methods use minibatch estimates whose variance depends on the batch size. Section 9.1 introduces these properties. The next five sections develop gradient descent and preconditioning (Section 9.2), stochastic gradients and learning-rate decay (Section 9.3), the computational effects of minibatching (Section 9.4), momentum for ill-conditioned objectives (Section 9.5), and per-coordinate scaling from AdaGrad through RMSProp to Adam (Section 9.6). The Adam section also builds the tiny transformer language model used in later experiments.

The second half removes simplifying assumptions from this progression. AdamW separates shrinkage from adaptive preconditioning, and learning-rate schedules vary the step over time. Muon changes the geometry used to choose matrix updates. Batch-size experiments connect gradient variance to parallel compute, while scaling and practice address hyperparameter transfer, clipping, weight averaging, and matched comparisons.

This chapter emphasizes computations and controlled experiments. The mathematical appendix (Chapter 26) proves the descent lemma, the condition-number law, momentum’s \(\sqrt{\kappa}\) acceleration, the Robbins–Monro conditions, and Adam’s bias correction under explicit assumptions. The two treatments can be read in either order.

Several methods in the second half are recent, and their relative performance remains protocol-dependent. The comparisons therefore tune each optimizer under the same budget and distinguish evidence from small testbeds, public benchmarks, and reported production runs.

Resources and Further Reading

The references below cover convex foundations, classical methods, current optimizers, and tuning. All are freely accessible online except where noted. The optimization chapter of the mathematical appendix (Chapter 26) 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 carried out in full.
• Numerical Optimization — Nocedal & Wright — a comprehensive treatment of line search, trust regions, and the quasi-Newton methods discussed in Section 9.2 (paywalled; 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 compares these methods through a shared geometric question.
• 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 finds that many reported speedups over AdamW shrink; it motivates the matched-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; an interactive treatment of damping and acceleration with sliders for \(\eta\) and \(\beta\), illustrating the critical value of \(\beta\) and oscillatory trajectories.
• An Overview of Gradient Descent Optimization Algorithms — Sebastian Ruder (2016) — free; a survey of the classical progression from Section 9.3 through Section 9.6 and a useful record of common practice in 2016.
• 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 in which Muon was first demonstrated, with documented and reproducible records; its reporting practice informs the evidence standard 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.