Adam was published in 2014, and a decade later it, or its decoupled-weight-decay sibling AdamW, still trains most frontier models. That is remarkable for an algorithm designed when a large network had tens of millions of parameters. It also has a cost nobody worried about at the time. For every parameter, Adam keeps two extra numbers, so a 405-billion-parameter model needs about 3.2 terabytes just for Adam's state in float32, before counting the weights, the gradients, or a single activation.

Two things changed recently. Large runs exposed the parts of Adam that do not scale gracefully (its $\epsilon$, its $\beta_2$, how weight decay interacts with the learning rate), and a new optimizer, Muon, which replaces Adam's per-coordinate scaling with an orthogonalized matrix update, set speed records and was then scaled to a 16-billion-parameter model with about half the compute of AdamW. This section covers what Adam costs, how to run AdamW correctly at scale, the memory-saving alternatives (Adafactor, Lion), the matrix-preconditioning family (Shampoo, SOAP), and Muon, whose core, the Newton-Schulz iteration, we derive from scratch.

We build on §0.6, which derived Adam with its bias correction, implemented AdamW's decoupled weight decay, showed that weights forget on a timescale of $1/(\eta\lambda)$ steps, and explained why large runs use $\beta_2 = 0.95$. Here we take those as known and ask what changes at scale, and we use the number formats of §8.5 to count the bytes. Throughout, $\theta$ is the parameters, $\mathbf{g}_t$ the gradient at step $t$, $\eta$ the learning rate, $\lambda$ the weight decay, $\mathbf{m}_t$ and $\mathbf{v}_t$ Adam's first and second moments with decay rates $\beta_1$ and $\beta_2$, and $\epsilon$ Adam's small constant.

What Adam costs in memory

§0.6 counted 16 bytes per parameter with everything in FP32 (weights, gradients, and two moments, 4 bytes each). Mixed precision changes the split but not the total, because the 16-bit working copy of the weights comes with an FP32 master copy (§8.5 explains why each copy has the precision it has). ZeRO's accounting (Rajbhandari et al. 2020) is the one everyone quotes.

In plain words, Adam's two moments are half of the total, and the optimizer side (master weights plus moments) is three quarters. For a 7B model that is 112 GB, more than one 80 GB GPU, and for LLaMA 3 405B it is about 6.5 TB, the memory of 81 H100s before a single activation is stored. This is why distributed training shards optimizer state first (§10.3), and why every optimizer in this section is judged partly by its state.

import numpy as np

def optimizer_state_numbers(shapes, optimizer):
    raise NotImplementedError

AdamW at scale

AdamW's update, from §0.6, is

with the decay applied directly to the weights rather than added to the gradient (Loshchilov & Hutter 2019). Three of its details behave differently at scale than in a textbook, and each one has caused real failures.

Weight decay sets an equilibrium, and a timescale

Why decay weights at all in a model that sees each token once and cannot overfit in the classical sense? Not mainly as a regularizer. With normalization layers, the function a layer computes barely depends on the scale of its weights (§5.2), so what weight decay really controls is the effective learning rate, the size of each update relative to the size of the weights.

Here is the derivation. Adam's normalized update $\mathbf{u}_t = \hat{\mathbf{m}}_t/(\sqrt{\hat{\mathbf{v}}_t}+\epsilon)$ has a roughly constant size, say $\lVert\mathbf{u}\rVert^2 \approx P\rho^2$ for a tensor with $P$ entries whose update has root-mean-square $\rho$ (typically 0.2 to 0.4 for Adam in practice). In high dimensions the update is nearly orthogonal to the current weights, so

Setting $\lVert\theta_t\rVert = \lVert\theta_{t-1}\rVert$ and using $(1-\eta\lambda)^2 \approx 1 - 2\eta\lambda$ for small $\eta\lambda$ gives the equilibrium

In plain words, the weights grow or shrink until each step changes them by a relative amount of about $\sqrt{2\eta\lambda}$. They get there on the timescale of $1/(\eta\lambda)$ steps that §0.6 found by unrolling the update, the same timescale on which they forget their initialization, about 33,000 steps for $\eta = 3\times10^{-4}$ and $\lambda = 0.1$. Wang and Aitchison (2024) argue that this timescale, as a fraction of the run, is the quantity to hold fixed when changing model or dataset size, which tells you how $\lambda$ should move when $\eta$ does.

The derivation also explains two practical choices. PyTorch's AdamW multiplies the decay by the learning rate, $\theta \leftarrow \theta(1 - \eta\lambda)$, so when the schedule lowers $\eta$, decay weakens too, and the two are coupled when you tune them. Wortsman et al. (2023) found that independent weight decay, $\theta \leftarrow \theta(1 - \lambda')$ with $\lambda'$ not multiplied by $\eta$, made the final loss less sensitive to the learning rate in their small-scale stability study. And it explains what not to decay. Gains of normalization layers and biases have no free scale to control, so decaying them only drags them away from useful values. OLMo 2 also stopped decaying the embedding matrix, as one of a set of changes made for training stability (Team OLMo 2024).

Both runs end at the predicted RMS of about 0.031, one growing from a tiny initialization and one shrinking from a large one, and both get most of the way there within a few multiples of the timescale $1/(\eta\lambda) \approx 3{,}333$ steps. The equilibrium does not care where you started, only about $\eta$, $\lambda$, and the size of the updates. (We used a learning rate ten times larger than a typical LLM's so the cell runs quickly. At $\eta = 3\times10^{-4}$ the timescale is 33,000 steps.)

Beta2 and the sharper failure at scale

§0.6 explained why large runs use $\beta_2 = 0.95$ (GPT-3, LLaMA, OLMo, DeepSeek) rather than Adam's default 0.999. The second moment then averages over about 20 steps instead of 1,000, so when gradients suddenly grow, the denominator $\sqrt{\hat{\mathbf{v}}}$ catches up in tens of steps rather than hundreds, and the update cannot stay far above its usual size for long. At scale there is a second failure that a short memory alone does not prevent.

Molybog et al. (2023) described it in runs from 7B to 546B parameters. Some layers, typically early ones, have gradients smaller than $\epsilon$ for long stretches, so their updates are tiny and their gradients become strongly correlated from step to step. When a rare batch pushes those gradients above $\epsilon$, the updates jump to roughly sign-sized values all at once, in a coordinated direction, and the loss spikes. Their list of remedies reads like a stability checklist (§8.8). Skip the batch, lower the learning rate, lower $\epsilon$, reduce the batch size, or adjust $\beta_1$ and $\beta_2$.

Why epsilon matters at scale

$\epsilon$ exists to avoid dividing by zero. §0.6 noted one reason to make it larger, as Llama 2 did with $10^{-5}$. It puts a floor under the denominator, so a parameter whose gradient is tiny and noisy cannot take full-size steps on noise. At scale the opposite danger takes over. The default, $10^{-8}$, was chosen when per-coordinate gradients were much larger than that. In a large model trained with a large batch, the root-mean-square gradient of a single weight can fall to around $10^{-8}$, and then $\epsilon$ is no longer negligible. Per coordinate, the update is

so when $\sqrt{\hat v} \approx \epsilon$, every update is halved, and for smaller gradients Adam turns into SGD with learning rate $\eta/\epsilon$, whose steps are tiny because the gradients themselves are, and which no longer adapts to each coordinate's scale. Wortsman et al. (2023) observed exactly this, with gradient RMS approaching $10^{-8}$ at scale and updates shrinking. Lowering $\epsilon$ to $10^{-15}$ improved their 4.8B model, while raising it to $10^{-6}$ made it diverge. OLMo 2 lowered $\epsilon$ from $10^{-5}$ to $10^{-8}$ for the same reason. The right $\epsilon$ depends on the gradient scale, which shrinks with batch size and model size, so a value inherited from a smaller run is worth checking against the logged gradient RMS.

You implemented AdamW in Exercise 0.6d, and §0.6's PyTorch block puts matrices and 1-D tensors in separate parameter groups so only the matrices are decayed. At scale the same code carries three extra decisions, all from this section. Decide whether the decay is multiplied by the learning rate (PyTorch's default) or independent of it. Add the embedding to the no-decay group if you follow OLMo 2. And set $\epsilon$ from the gradient RMS you actually observe, not from a default.

Adafactor, a second moment in factored form

Adafactor (Shazeer & Stern 2018) attacks Adam's memory directly. For a weight matrix $\mathbf{W} \in \R^{n\times m}$ with gradient $\mathbf{G}$, Adam stores a full $n\times m$ matrix $\mathbf{V}$ of second moments. Adafactor stores only its row sums $\mathbf{r} \in \R^n$ and column sums $\mathbf{c} \in \R^m$, each kept as an exponential moving average of the corresponding sums of $\mathbf{G}^2$ (plus a tiny $\epsilon_1 = 10^{-30}$), and reconstructs

In plain words, guess each entry's second moment from how large its row is and how large its column is, which costs $n + m$ numbers instead of $nm$. For a $4096\times 4096$ matrix, that is 8,192 numbers instead of 16.8 million.

Why this particular rank-one guess? Because it is the best one under the right notion of error.

Adafactor made three more changes, each to save memory or add stability. It drops the first moment altogether ($\beta_1 = 0$), saving another full copy. It clips each update to a root-mean-square of at most $d = 1$, $\hat{\mathbf{U}} = \mathbf{U}/\max(1, \operatorname{RMS}(\mathbf{U})/d)$, which stops the large steps that a stale second moment would otherwise produce (the same failure as $\beta_2 = 0.999$ above). And it uses an increasing decay rate, $\hat\beta_{2,t} = 1 - t^{-0.8}$, so the average covers more steps as training goes on, and a relative step size, the learning rate times the RMS of the parameter itself. T5 was trained with Adafactor, and PaLM used Adafactor without factorization, keeping its $\beta_2$ schedule and parameter-scaled steps (Chowdhery et al. 2022).

import numpy as np

def adafactor_step(G, r, c, beta2, eps1=1e-30, d=1.0):
    raise NotImplementedError

Lion, an optimizer found by search

Lion (Chen et al. 2023) was discovered by a symbolic search over programs, not designed by hand, and it is strikingly simple. With $\beta_1 = 0.9$ and $\beta_2 = 0.99$,

In plain words, take the sign of a blend of the momentum and the current gradient, so every coordinate moves by exactly $\eta$, then update the momentum with a slower average. Two different betas mean that the direction used now leans more on the current gradient than the momentum that is carried forward does.

Lion keeps only one state tensor instead of Adam's two, which saves 4 of the 16 bytes per parameter of mixed-precision training. Because every coordinate of a sign update has magnitude 1, while Adam's update has an RMS around 0.2 to 0.4, Lion needs a learning rate 3 to 10 times smaller than AdamW's and a weight decay 3 to 10 times larger to keep $\eta\lambda$ similar. The paper reported up to 5 times less pretraining compute for vision Transformers on JFT and solid results on language modeling. Sign-based updates are also cheap to communicate across GPUs, since one bit per coordinate carries the direction.

import numpy as np

def lion_step(theta, g, m, lr, beta1=0.9, beta2=0.99, wd=0.0):
    raise NotImplementedError

Shampoo and SOAP, preconditioning with matrices

Adam rescales each coordinate separately, which is a diagonal preconditioner. It ignores correlations between coordinates. The ideal of full-matrix AdaGrad, which would precondition with the inverse square root of the full $P\times P$ matrix of gradient outer products, is hopeless for $P$ in the billions. Shampoo (Gupta et al. 2018) exploits the matrix structure of each layer instead. For a weight $\mathbf{W} \in \R^{m\times n}$ with gradient $\mathbf{G}_t$, it accumulates a left and a right statistic and preconditions on both sides,

In plain words, whiten the gradient along its rows with one small matrix and along its columns with another. The Kronecker product $\mathbf{L}^{1/2}\otimes\mathbf{R}^{1/2}$ approximates the full-matrix statistic with $m^2 + n^2$ numbers instead of $m^2n^2$.

There is a beautiful special case that connects Shampoo to the next optimizer. Drop the accumulation and use only the current gradient, with SVD $\mathbf{G} = \mathbf{U}\boldsymbol\Sigma\mathbf{V}^\top$. Then $\mathbf{L} = \mathbf{U}\boldsymbol\Sigma^2\mathbf{U}^\top$, $\mathbf{R} = \mathbf{V}\boldsymbol\Sigma^2\mathbf{V}^\top$, and

In plain words, Shampoo without memory replaces the gradient by its orthogonal part, keeping its singular vectors and setting every singular value to 1. Hold on to that, because it is exactly what Muon computes.

Shampoo's practical problems are cost and numerics. Inverse fourth roots of large matrices need an eigendecomposition, which is slow and unstable in low precision, so implementations recompute them only every few dozen or hundred steps and run them in float32 on the side. A distributed implementation of Shampoo won the external-tuning track of the AlgoPerf training-algorithms benchmark in 2024. SOAP (Vyas et al. 2024) simplified the idea further. Rotate the gradient into the eigenbasis of Shampoo's $\mathbf{L}$ and $\mathbf{R}$, run ordinary AdamW there, and rotate back, refreshing the eigenbasis only every 10 steps or so. At 360M and 660M parameters with a 2-million-token batch, SOAP needed over 40 percent fewer iterations and over 35 percent less wall-clock time than AdamW.

Muon, orthogonalized updates

Muon (Jordan et al. 2024, "MomentUm Orthogonalized by Newton-Schulz") applies to each 2-D weight matrix of the hidden layers. It keeps a momentum buffer, as SGD with momentum would, and replaces the update by its orthogonal part $\mathbf{U}\mathbf{V}^\top$ before applying it. Embeddings, the output head, and all 1-D parameters (gains, biases) are trained with AdamW as usual.

Why orthogonalize

There are two answers, one empirical and one geometric.

The empirical one comes from looking at updates. Jordan observed that the updates SGD-momentum and Adam produce for Transformer weight matrices have very high condition numbers. They are close to low rank, dominated by a few directions, so a few directions in weight space get almost all of the learning while many others, possibly important, get almost none. Orthogonalizing sets every singular value to 1 and gives the rare directions the same step size as the dominant ones.

The geometric answer explains why this is a principled choice. Gradient descent is steepest descent under the Euclidean norm, the direction that decreases the loss fastest per unit of step length. Change the norm and you change the algorithm. For a linear layer, the natural norm on an update $\Delta\mathbf{W}$ is the spectral norm $\lVert\Delta\mathbf{W}\rVert_{\text{op}}$, its largest singular value, because $\lVert\Delta\mathbf{W}\mathbf{x}\rVert \le \lVert\Delta\mathbf{W}\rVert_{\text{op}}\lVert\mathbf{x}\rVert$ bounds how much the layer's output can change for any input (Bernstein & Newhouse 2024). So ask for the step that decreases the linearized loss the most while changing no output by more than $\eta$ times its input,

where $\langle\mathbf{G}, \Delta\mathbf{W}\rangle = \operatorname{tr}(\mathbf{G}^\top\Delta\mathbf{W})$ is the first-order change in the loss. Write $\mathbf{G} = \sum_i \sigma_i\mathbf{u}_i\mathbf{v}_i^\top$. Then

because $|\mathbf{u}^\top\mathbf{A}\mathbf{v}| \le \lVert\mathbf{A}\rVert_{\text{op}}$ for unit vectors $\mathbf{u}$ and $\mathbf{v}$. The bound is reached by $\Delta\mathbf{W} = -\eta\sum_i\mathbf{u}_i\mathbf{v}_i^\top = -\eta\,\mathbf{U}\mathbf{V}^\top$, since then $\mathbf{u}_i^\top\Delta\mathbf{W}\mathbf{v}_i = -\eta$ for every $i$, and $\lVert\mathbf{U}\mathbf{V}^\top\rVert_{\text{op}} = 1$. In plain words, the steepest step under the spectral norm is the gradient with all its singular values set to the same size. Muon applies this to the momentum rather than the raw gradient.

Newton-Schulz, orthogonalizing with matrix products

Computing $\mathbf{U}\mathbf{V}^\top$ with an SVD every step is slow on a GPU. Newton-Schulz iterations compute it with nothing but matrix multiplications, which GPUs do best, and in bfloat16.

The key fact is that odd matrix polynomials act on singular values independently. If $\mathbf{X} = \mathbf{U}\boldsymbol\Sigma\mathbf{V}^\top$, then $\mathbf{X}\mathbf{X}^\top\mathbf{X} = \mathbf{U}\boldsymbol\Sigma\mathbf{V}^\top\mathbf{V}\boldsymbol\Sigma\mathbf{U}^\top\mathbf{U}\boldsymbol\Sigma\mathbf{V}^\top = \mathbf{U}\boldsymbol\Sigma^3\mathbf{V}^\top$, and in general $(\mathbf{X}\mathbf{X}^\top)^k\mathbf{X} = \mathbf{U}\boldsymbol\Sigma^{2k+1}\mathbf{V}^\top$. So the iteration

keeps $\mathbf{U}$ and $\mathbf{V}$ fixed and applies the scalar polynomial $\varphi(\sigma) = a\sigma + b\sigma^3 + c\sigma^5$ to every singular value at once. In plain words, to orthogonalize a matrix we only need a scalar polynomial that, iterated, sends every number in $(0, 1]$ to 1.

The classic choice is the cubic $\varphi(\sigma) = \tfrac32\sigma - \tfrac12\sigma^3$, the Newton-Schulz iteration for the polar decomposition. Its properties are easy to check.

• 1 is a fixed point, $\varphi(1) = \tfrac32 - \tfrac12 = 1$.
• Convergence near 1 is quadratic. Writing $\sigma = 1 - \delta$, $\varphi(1 - \delta) = 1 - \tfrac32\delta^2 + \tfrac12\delta^3$, so the error goes from $\delta$ to about $1.5\,\delta^2$.
• Every $\sigma$ in $(0, 1)$ increases toward 1, since $\varphi(\sigma)/\sigma = \tfrac32 - \tfrac12\sigma^2 > 1$ there, and $\varphi$ maps $(0, 1)$ into $(0, 1)$. Values in $(1, \sqrt3)$ are mapped into $(0, 1)$ and then converge too, while values above $\sqrt3$ go negative and diverge.

To keep every singular value in the safe range, first divide by the Frobenius norm, $\mathbf{X}_0 = \mathbf{G}/\lVert\mathbf{G}\rVert_F$. Since $\lVert\mathbf{G}\rVert_F^2 = \sum_i\sigma_i^2 \ge \sigma_{\max}^2$, all singular values of $\mathbf{X}_0$ are at most 1.

The cubic has one weakness. Near zero it only multiplies by $\varphi'(0) = 1.5$ per step, and normalizing by the Frobenius norm makes small singular values very small. A singular value of $10^{-3}$ needs about 20 cubic steps to reach 1. Jordan's insight was that Muon does not need exactly $\mathbf{U}\mathbf{V}^\top$. He chose a quintic with a large slope at zero, $(a, b, c) = (3.4445, -4.7750, 2.0315)$, so $\varphi'(0) = 3.4445$, and accepted that it does not converge to 1 but oscillates in a band around it. After 5 steps, every initial singular value from $10^{-2}$ to 1 lands between about 0.68 and 1.13, so the result is $\mathbf{U}\mathbf{S}'\mathbf{V}^\top$ with $\mathbf{S}'$ close to the identity, which works as well in practice.

This matrix has condition number 1,000, and after normalization its singular values span 0.0002 to 0.54. Five cubic steps barely move the small ones (only 21 percent reach the band from 0.68 to 1.2), and 25 cubic steps converge to $\mathbf{U}\mathbf{V}^\top$ exactly. Five quintic steps put 77 percent of the singular values in the band and lift the smallest from 0.0002 to about 0.1, a 500-fold boost, at a fifth of the cost of the 25 cubic steps. Explore the scalar maps below.

import numpy as np

def newton_schulz(G, steps=5, coeffs=(3.4445, -4.7750, 2.0315), eps=1e-7):
    raise NotImplementedError

The full algorithm, and making it scale

Putting it together, Muon's step for a matrix $\mathbf{W} \in \R^{A\times B}$ with gradient $\mathbf{G}$ is

with momentum $\mu = 0.95$, Nesterov's look-ahead $\mathbf{N}$, five Newton-Schulz steps, and a shape-dependent scale $s$. Two scales are in use. Jordan's original sets $s = \sqrt{\max(1, A/B)}$. The Moonshot team, scaling Muon up, found two changes crucial (Liu et al. 2025). They added decoupled weight decay, which plain Muon lacked, and chose $s$ so the update has the same root-mean-square as a typical AdamW update. Since $\mathbf{U}\mathbf{V}^\top$ for an $A\times B$ matrix has $\min(A, B)$ singular values equal to 1, its squared Frobenius norm is $\min(A, B)$ and

gives an RMS of 0.2, matching AdamW's typical 0.2 to 0.4. That lets Muon reuse AdamW's learning rate and weight decay, and the full update is $\mathbf{W} \leftarrow \mathbf{W} - \eta\big(0.2\sqrt{\max(A, B)}\,\mathbf{O} + \lambda\mathbf{W}\big)$. With these changes, Muon reached the same loss as AdamW with about 52 percent of the training FLOPs in their compute-optimal scaling experiments, and they trained Moonlight, a mixture-of-experts model with 3B active and 16B total parameters, on 5.7 trillion tokens with it.

What does Muon cost? Each Newton-Schulz step on an $m\times n$ matrix ($m \le n$) costs about $4m^2n + 2m^3$ FLOPs, at most $6m^2n$, so five steps cost at most $30m^2n$. The matrix's share of a training step with a batch of $B$ tokens (here $B$ is the batch size again, not the matrix dimension above) is $6mnB$ FLOPs (forward and backward, §5.6), so the overhead is at most

For a model of width $m = 4096$ trained with a batch of 4 million tokens, that is 0.5 percent. This is the formula $Tm/B$ (with $T = 5$ steps) from Jordan's write-up. The real costs are elsewhere. Newton-Schulz needs the whole matrix on one device, which clashes with sharded optimizers that split each matrix across GPUs (§10.3), and it adds latency even when its FLOPs are small. Moonshot's distributed implementation gathers each matrix, orthogonalizes it, and scatters the result.

import numpy as np

def muon_step(W, G, M, lr, momentum=0.95, wd=0.1, nesterov=True):
    raise NotImplementedError

Choosing an optimizer

The honest summary for a practitioner is that AdamW with $\beta_2 = 0.95$, $\lambda = 0.1$ on matrices only, and a small enough $\epsilon$ remains the safe default. Muon is the strongest new contender, with published wins at the 16B scale and adoption in at least one frontier-scale open model. Comparisons between optimizers are notoriously easy to get wrong (an untuned baseline, a different learning-rate schedule, a short run that favors one method), so treat any reported speedup as a claim about the tuning budget as well as the algorithm. Hyperparameters that transfer across model widths are their own topic (§9.5).

In PyTorch

PyTorch now ships Muon, so the real way is a few lines. Matrices of hidden layers go to Muon, everything else to AdamW, and we check one step of torch.optim.Muon against our NumPy Newton-Schulz.

The two agree to about the precision of bfloat16, which is the format PyTorch runs Newton-Schulz in. That the iteration is stable in bfloat16 at all is one of its practical virtues, since an SVD in bfloat16 would be useless.

Read more

• Loshchilov & Hutter (2019). Decoupled Weight Decay Regularization. AdamW, and why L2 and weight decay differ for adaptive methods.
• Shazeer & Stern (2018). Adafactor, Adaptive Learning Rates with Sublinear Memory Cost. The factored second moment, update clipping, and relative step sizes.
• Chen et al. (2023). Symbolic Discovery of Optimization Algorithms. Lion, and how it was found by program search.
• Gupta et al. (2018). Shampoo, Preconditioned Stochastic Tensor Optimization, and Vyas et al. (2024), SOAP. Matrix preconditioning, and Adam in Shampoo's eigenbasis.
• Jordan et al. (2024). Muon, An optimizer for hidden layers in neural networks. The original write-up, with the quintic coefficients and the overhead analysis.
• Liu et al. (2025). Muon is Scalable for LLM Training. Weight decay, RMS matching, distributed Muon, and Moonlight.
• Bernstein & Newhouse (2024). Old Optimizer, New Norm, An Anthology. Steepest descent under different norms, connecting Adam, Shampoo, and orthogonalized updates.
• Wortsman et al. (2023). Small-scale proxies for large-scale Transformer training instabilities. Independent weight decay, the $\epsilon$ problem, and more, studied at small scale.

Watch next

▶

▶

▶