What an Optimizer Actually Does

One question, three knobs

Training is a loop: compute a gradient, change the weights, repeat. The optimizer owns the second half of that sentence, and every optimizer ever shipped is answering the same question — given everything I have seen about the gradient so far, and a fixed budget of memory, which direction do I move, and how far?

Strip the names away and each one resolves into three decisions. How to smooth the gradient in time, because a minibatch gradient is a noisy sample and averaging trades responsiveness for variance. How to rescale it in space, because the loss surface is not equally steep in every direction and one step size cannot be right for every coordinate. And which norm defines a step of size one, because "steepest descent" is meaningless until you say what distance means.

All of them fit one template:

SGD sets m = g and P = I. Adam sets m to an exponential moving average and P to a diagonal matrix of per-coordinate gradient scales. Muon keeps the same m but makes P−1 the map that replaces the momentum matrix with the nearest semi-orthogonal one. That is genuinely the whole taxonomy. The rest is bookkeeping about what you can afford to store.

Here is what those choices look like when you let them run. Same landscape, same starting point, same wall-clock budget of steps — six pens, six answers.

Two things are worth noticing before we take any of it apart. First, the ranking is not fixed — switch to the Rosenbrock valley and the adaptive methods lose their advantage, because the problem there is curvature that rotates, which no diagonal preconditioner can see. Second, every method that beats plain SGD does so by storing something. That storage is the subject of half this article.

The Newton budget

Start from the only honest first principle available: near the current weights, the loss looks like a quadratic.

Minimise the right-hand side over Δ and you get the Newton step, Δ = −H−1g. It is the correct answer. It is also, for a 7-billion-parameter model, a 49-quintillion-entry matrix that you would then need to invert. Nobody is doing that.

So every practical optimizer is a compression of H⁻¹ under a memory budget. That single sentence organises the entire field, and it is the thing to reach for when a new optimizer shows up: ask what structure it assumes about the curvature, and what that assumption costs per parameter.

The ladder runs from exact to free:

Read the rungs as a budget, not a quality ranking. A full Hessian costs n floats per parameter — impossible. Block-diagonal per layer is still nℓ floats per parameter within a layer, also impossible for a 4096×4096 matrix. Kronecker factorisation — assume the curvature of a weight matrix W ∈ ℝm×n factors as L ⊗ R — drops the cost to (m² + n²)/(mn) floats per parameter, which for real transformer shapes lands between 1 and 4. That is Shampoo. A diagonal approximation is exactly 1 float per parameter: Adagrad, RMSProp, Adam. A scalar is free: that is plain SGD with a tuned learning rate.

Picking a geometry

There is a second first-principles route to the same place, and it is the one that actually explains Muon. Forget curvature for a moment and define the step as a small, constrained greedy move:

This is a linear objective over a convex set, so the answer always sits on the boundary — and which boundary point depends entirely on the shape of the ball. Change the norm and you change the optimizer. Nothing else moves.

The Euclidean ball is a sphere, so the tangent point is −r·g/‖g‖: that is SGD. The max-norm ball is a cube, so the answer is always a corner, −r·sign(g): that is sign descent, and Adam in its ε→0 limit is very close to it. The ℓ1 ball is a cross-polytope, so the answer is a single vertex — move only the largest coordinate — which is greedy coordinate descent.

Now do it for a weight matrix. The natural norm on a linear map is the spectral norm, ‖W‖2 = σmax, because that is what bounds how much the layer's output can change when its input changes. Write the gradient's SVD as G = UΣV⊤. In singular-value coordinates the spectral ball is a box, so the same corner logic applies — push every singular value to the boundary:

This is also why the two compose so awkwardly and why practical Muon runs are hybrids: an embedding table or a bias vector has no meaningful singular-value structure, so those tensors stay on AdamW while the 2-D hidden weights go to Muon.

Smoothing in time

Momentum is the oldest of the three knobs and the least mysterious, but two things about it are routinely misread: what it does to your effective learning rate, and why it helps at all.

Two conventions, one factor of (1 − β)

For a constant gradient the accumulator converges to g/(1−β) and the EMA converges to g. With β = 0.9 that is a 10× difference in effective learning rate between two lines of code that both call themselves momentum. Porting an LR between frameworks without checking which one you are in is a classic way to lose a day.

Why it works: κ becomes √κ

On a quadratic with condition number κ, plain gradient descent contracts the error by (κ−1)/(κ+1) per step — for κ = 10⁴ that is 0.9998, which is a disaster. Heavy-ball momentum with the optimal β contracts by (√κ−1)/(√κ+1). The condition number is square-rooted. That is not a heuristic speedup, it is the accelerated rate, and it is the entire reason momentum is not optional.

Below β* the iterate is over-damped and crawls; above it, under-damped and it rings. The default 0.9 is an implicit bet that your problem's conditioning is around 10³.

The window's effective length is 1/(1−β) gradients, with mean age β/(1−β) steps. At β = 0.9 you are averaging roughly the last ten minibatches; at β = 0.99, the last hundred. That number should be compared against how fast your loss landscape is actually moving — a long window on a run that is still changing regime gives you a smoothed direction that points where the model used to want to go.

Nesterov

Nesterov's variant evaluates the gradient at the point momentum is about to carry you to, rather than where you are: gt = ∇f(θt − ηβmt−1). Practically it acts as a first-order correction that damps overshoot — the update gets an extra β(mt − mt−1) term that pulls back when the momentum is growing. It buys a real but modest constant on convex problems and a smaller one on deep nets, which is why it survives as a flag rather than a default.

Scaling in space

The second knob assumes the curvature is diagonal and estimates it from the only thing you have: the gradients themselves. The lineage is three papers deep and each one fixes exactly one flaw in the last.

Adagrad's v is a running sum, so 1/√v decays monotonically like 1/√t — an annealing schedule welded into the optimizer. That is provably right for convex problems and fatal for deep ones, where the learning rate hits zero long before the model is done. RMSProp swaps the sum for an exponential moving average, which turns "all of history" into "the last 1/(1−β₂) steps." Adam adds momentum on the numerator and fixes the initialisation bias on both.

What bias correction is actually correcting

Both moments start at zero, so both are biased toward zero early. The usual gloss — "it stops the first steps from being too small" — has the sign backwards. Work it out at t = 1 with the defaults:

The numerator is suppressed by 10× but the denominator by ~32×, so the ratio is inflated by a factor of √(1−β₂)⁻¹·(1−β₁) ≈ 3.2. Without correction Adam takes its largest steps at initialisation, when the gradient estimate is worst. Bias correction is a built-in warmup, and it is one reason short-warmup recipes are viable at all.

Weight decay, coupled and not

The difference between Adam and AdamW is four characters of code and a real change in what the regulariser does.

In the first form the penalty is fed through the adaptive denominator, so the decay a weight actually receives is ηλθ / (√v̂ + ε). A parameter with a long history of large gradients has a large √v̂ and therefore gets less regularisation — precisely backwards, since those are the parameters doing the most work and most likely to grow. Decoupling makes the shrinkage uniform, which is what "weight decay" meant in the first place.

Decay sets an equilibrium weight norm

Worth deriving once, because it tells you what λ actually controls. Adam's update has RMS ≈ 1 and is roughly uncorrelated with θ, so per coordinate θt+1 = (1 − ηλ)θt − η ut with E[u²] ≈ 1. Taking second moments at steady state:

λ does not "shrink weights toward zero." It picks a scale, jointly with the learning rate, and the network sits at that scale. With η = 3·10⁻⁴ and λ = 0.1 that is an RMS around 0.04. It also means that halving your peak LR halves η/λ and changes the operating point of every normalised layer — which is why LR and weight decay have to be retuned together, and why some recipes parameterise λ relative to peak LR so the equilibrium stays put.

Spending less state

Adam's 8 bytes per parameter of state is, for a large model, the single largest line item you can attack without touching the model. Four different attacks, all in production somewhere.

Adafactor — factor the second moment

For a weight matrix W ∈ ℝm×n, the second-moment matrix V is mn numbers. Adafactor's observation is that V is well approximated by a rank-1 matrix, and the best rank-1 approximation under I-divergence is just the outer product of its row and column sums:

For a 4096×11008 MLP matrix that is 45 M floats collapsing to 15 K — a 3000× reduction on that tensor, and effectively zero bytes per parameter. Adafactor also drops the first moment by default (β₁ = 0) and compensates with update clipping on the RMS. It trained T5 and it is still the right answer when memory is the binding constraint, but it does lose something: the rank-1 assumption is a real assumption, and runs are noticeably twitchier without momentum.

8-bit optimizers — quantise the state

Store m and v as 8-bit integers instead of fp32, dequantising on the fly inside the step. Naive quantisation fails because a single outlier sets the scale for the whole tensor; the fix is block-wise quantisation — split the state into blocks of 2048, normalise each block by its own absolute max, and use a non-linear (dynamic-exponent) code so small magnitudes keep resolution. State drops from 8 bytes to 2 bytes per parameter with no measurable loss on standard recipes. The failure mode to know about is that v is more sensitive than m, because errors there pass through a reciprocal square root.

Lion — throw away the second moment entirely

One state tensor, 4 bytes per parameter. Lion leans all the way into the sign-descent reading of Adam: if the update was going to have magnitude ≈ 1 everywhere anyway, skip the machinery that produces the 1. Two consequences follow directly and both are mandatory in practice — the update norm is now exactly √(number of params) regardless of gradient scale, so the LR must drop by roughly 3–10×, and since decay competes against a fixed-size update, λ must rise by roughly the same factor to keep η/λ in the same place.

Adam-mini — share the denominator

Rather than one v per parameter, partition parameters into blocks that plausibly share curvature (per attention head, per output neuron) and keep one v per block. State drops close to 4 bytes per parameter. It is the diagonal approximation, coarsened one more notch down the ladder from §2.

Matrices, not vectors

Everything so far treats a weight tensor as a bag of independent scalars. That throws away the one structural fact you definitely have: it is a linear map, and its rows and columns are not interchangeable with each other.

Shampoo

Assume the curvature of W ∈ ℝm×n factors as a Kronecker product, and maintain the two factors separately:

The quarter powers are not arbitrary: they are what makes the two-sided product equal (L ⊗ R)−1/2 applied to vec(G), the Kronecker analogue of Adagrad's inverse square root. The cost is m² + n² floats of accumulator plus the same again for the cached roots, and an eigendecomposition that is O(m³ + n³) — which is why every real implementation amortises the root computation over 50–100 steps and runs it off the critical path. SOAP is the sharpest version of this idea: keep Shampoo's eigenbasis, update it rarely, and just run ordinary Adam in that rotated basis. You get the rotation from Shampoo and the cheap per-coordinate scaling from Adam.

Muon

Muon is the geometry argument from §3, made executable. Take the momentum matrix, replace it with the nearest semi-orthogonal matrix, step. The only question is how to orthogonalise a 4096×11008 matrix on every step without an SVD, and the answer is a Newton–Schulz iteration: a fixed odd polynomial applied to the matrix, which acts on the singular values without ever computing them.

Because X and A share singular vectors, the matrix iteration is exactly the scalar map σ ← aσ + bσ³ + cσ⁵ applied independently to every singular value. Watch it run:

The coefficients are the interesting part. A textbook Newton–Schulz iteration converges to exactly 1; these do not — their fixed point wobbles in a band roughly [0.7, 1.3]. That is deliberate. They are tuned to maximise the slope at σ = 0, so that the tiny singular values — the directions the raw gradient has almost no signal in — get lifted into the band within five steps instead of fifty. An optimizer does not need an orthogonal matrix; it needs a matrix whose singular values are all the same order of magnitude. Accepting ±30% error buys a 10× reduction in iterations and lets the whole thing run in bf16.

What Muon costs

Per m×n matrix with n = min(m,n), one iteration is three matmuls: XX⊤ (2n²m), A² (2n³), BX (2n²m). Five iterations, divided by the mn parameters:

The overhead scales with model width and inversely with tokens per step, and at any realistic large-batch configuration it disappears. What does not disappear is the communication, which is the real reason Muon is harder than it looks — see §10.

The family around it

• Scion / modular norms push the same argument further: pick a norm per layer type — spectral for hidden matrices, RMS→∞ for embeddings, sign for biases — and take a norm-constrained step everywhere, so nothing falls back to AdamW.
• Dion attacks the distributed problem: orthonormalise in a low-rank subspace so the update can be computed on sharded parameters without gathering the full matrix on any one rank.
• MARS layers variance reduction on top, correcting the momentum with a scaled gradient difference between consecutive steps.
• Weight decay is not optional at scale. The original Muon had none; large-scale reproductions added decoupled decay and per-parameter update-RMS matching, and reported that both were required past a few billion parameters.

The memory ledger

Yes — you keep them resident. Optimizer state is not a transient buffer; it is one tensor the same shape as the parameter, allocated before step 1 and read and written on every step until the run ends. It is also the part of the memory bill that people consistently underestimate, because unlike activations it does not show up in a "batch size 1" sanity check.

The full mixed-precision accounting, per parameter:

Shampoo's entry is the one that needs a caveat, because it is the only row where the cost depends on tensor shape rather than parameter count. Storage is (m² + n²)/(mn) floats per parameter: for a square 4096×4096 that is 2 floats, but for a 4096×11008 MLP matrix the n² term dominates and it climbs to 3.1. Tall-skinny tensors are where Kronecker methods get expensive.

Sharding, in one formula

With N ranks and K bytes of optimizer state plus master weights per parameter (K = 12 for AdamW: 4 master + 4 m + 4 v), the steady-state per-rank cost is:

For 7B on 8 ranks that is 112 GB replicated, 38.5 GB at ZeRO-1, 26.3 GB at ZeRO-2, and 14 GB at ZeRO-3. The important and slightly non-obvious fact: ZeRO-1 is free. A gradient all-reduce is internally a reduce-scatter followed by an all-gather, and the reduce-scatter phase already leaves each rank holding the fully-summed slice it owns. Sharding optimizer state costs exactly zero extra bytes on the wire — you were computing that partition anyway.

The bill nobody budgets for

Your checkpoint contains master weights plus optimizer state — 12 bytes per parameter, or 84 GB for a 7B model, seven times the size of the deployable bf16 weights. That is what makes checkpointing slow, and it is why the resharding format matters: if you save flat per-rank shards, you cannot resume the run on a different world size. And if you resume without the optimizer state, the moments restart at zero and you get a visible loss bump from the bias-correction transient.

Where a step actually costs you

The step is bandwidth-bound, not compute-bound

Count the traffic for one fused AdamW step, per parameter: read the bf16 gradient (2 B), the fp32 master weight (4 B), m (4 B) and v (4 B); write back master, m, v (12 B) and the bf16 weight (2 B). Roughly 28 bytes moved for about 12 FLOPs.

So the step time is just bytes ÷ bandwidth: 7B × 28 B = 196 GB, at 3.35 TB/s ≈ 59 ms unsharded, or ~7 ms per rank at ZeRO-3 on 8 GPUs. Against a 2–8 second training step that is noise. Three things turn it back into a problem.

1. Kernel launch overhead

A 7B transformer has a few hundred parameter tensors, many of them small. A naive Python loop issues five or six tiny kernels per tensor per step, and at ~5 µs of launch overhead each you can spend more time launching than computing. foreach (multi-tensor apply) batches tensors into a handful of kernels; fused goes further and does the whole update in one. On small models this is the difference between a 2 ms step and a 40 ms step.

2. Offload

Moving optimizer state to host memory frees 12 B/param of HBM and costs you the PCIe link. For 7B that is 84 GB across a Gen5 ×16 link at ~64 GB/s — about 1.3 s each way, plus a CPU-side update that is an order of magnitude slower than the GPU's. Offload is a way to make a run possible, never a way to make it fast.

3. Communication — and why Muon is harder than it looks

Under FSDP the per-step wire traffic is a reduce-scatter of gradients and an all-gather of parameters, roughly 2 bytes per parameter each way in bf16. Muon adds one more thing: Newton–Schulz needs the whole matrix. FSDP2 shards parameters along dimension 0, so each rank holds a horizontal stripe of every weight matrix — and a stripe cannot be orthogonalised. The standard implementation all-gathers the momentum for each 2-D tensor, runs the iteration on one rank (matrices assigned round-robin), then broadcasts the result back.

On a single node this is invisible. Across an Ethernet-class fabric it is not, and it lands on the critical path because the optimizer step cannot start until every gradient has arrived. This asymmetry — cheap inside a node, expensive across one — is the whole reason the distributed-Muon variants exist: keep the orthogonalisation sharded, or approximate it in a low-rank subspace, so no full matrix ever has to cross the slow fabric.

Nuances that bite

• ε is not a rounding guard. m̂/(√v̂ + ε) and m̂/√(v̂ + ε) are different optimizers for small gradients. In the first form, once √v̂ ≪ ε the coordinate silently reverts to SGD with LR η/ε. Defaults of 1e-8 are fine at fp32 gradient scales; runs with heavily normalised losses often need 1e-15 to avoid accidentally freezing low-gradient parameters.
• β₂ staleness causes loss spikes. The window is 1/(1−β₂) steps — 1000 at the default 0.999. If the loss regime shifts faster than that (a data-mix change, a curriculum boundary, a sudden burst of long documents), v is stale and too small, the update is too large, and you get a spike. Lowering β₂ to 0.95 shortens the memory; the alternative used at scale is simply to detect the spike and skip the batch.
• Gradient clipping happens before the optimizer, and interacts with it. A global-norm clip rescales all gradients by one scalar, which Adam's denominator then partially undoes — it mostly protects v from ingesting an outlier, not the current step. That is the right mental model for why clipping still helps a method that already normalises per coordinate.
• Accumulate gradients in fp32 past a handful of microbatches. bf16 has 8 mantissa bits; summing 32 microbatch gradients into a bf16 buffer loses small contributions to rounding entirely. The cost is 2 extra bytes per parameter and it is almost always worth it.
• Master weights exist because of one addition. θ ← θ − ηu with η·u around 1e-7 and θ around 1e-2 is a no-op in bf16 — the update falls off the bottom of the mantissa. Keeping θ in fp32 fixes it. The alternative is stochastic rounding, which makes the update land probabilistically and lets you drop 4 bytes per parameter; it works, but it is a genuine change to the training dynamics, not a free win.
• Under gradient accumulation, the moments see accumulated gradients. β₁ and β₂ windows are measured in optimizer steps, not microbatches. Doubling accumulation while halving the step count halves your effective momentum window in wall-clock terms.
• Learning rates do not transfer across widths for free. Under µP, Adam's LR for hidden matrices should scale as 1/fan_in. Muon's spectral normalisation gets much closer to width-invariance for free, which is a large part of its practical appeal — one fewer thing to re-sweep when you scale the model.
• A weight EMA is another full copy. If your eval uses an EMA of the weights, budget 2–4 more bytes per parameter and remember it also has to be checkpointed and resharded.
• Muon touches only 2-D parameters. Embeddings, the LM head, LayerNorm gains, and biases stay on AdamW — so a Muon run has two optimizers, two learning rates, and two weight decays to tune, and the memory saving applies only to the hidden weights (typically 70–85% of parameters).
• Check whether your `foreach` path silently allocates. Multi-tensor apply on a fresh set of intermediate buffers can add a transient copy of the state. On a run sitting at 95% HBM, that is the OOM you cannot reproduce at smaller batch size.

Choosing one

A short decision procedure that is right more often than not.

And the standing advice regardless of which one you pick: measure the update-to-weight RMS ratio per layer. One scalar per tensor, logged every few hundred steps. Healthy training keeps RMS(Δθ)/RMS(θ) in a band around 10⁻³. If a layer drifts an order of magnitude off, you have found your problem before the loss curve shows it — and that diagnostic works identically for every optimizer in this article, because it measures the only thing they all agree they are doing.

References

Grouped by topic. Every entry links to its primary source.