Optimizers

You have a loss function with a few billion inputs. You can afford to evaluate its gradient, on a random slice of the data, a few hundred thousand times. Each evaluation is noisy, the landscape is stretched by factors of a million in some directions and flat in others, and you only get one attempt. How do you move?

That is the whole optimizer problem. Everything in this chapter, from plain gradient descent to Muon, is a different answer to three sub-questions: which direction do we step, how far, and how do we stop the noise and the bad conditioning from wrecking us.

We will build the answers up one term at a time. Each new optimizer adds exactly one idea to the previous one, and every idea is there because something broke without it.

The problem: a bad landscape seen through a noisy keyhole

Picture the loss as a landscape over the parameters. Three things make it hard.

Bad conditioning. In some directions the loss curves sharply (a narrow canyon), in others it is almost flat (a long valley floor). The ratio of the biggest curvature to the smallest is the condition number. For a transformer it is enormous: embedding rows for rare tokens see tiny gradients, attention logits see huge ones, and the same learning rate has to serve both.

Noise. We never see the true gradient, only a minibatch estimate. Its variance depends on batch size and on which parameters we are looking at.

Scale. Billions of parameters means we cannot store anything like a Hessian, and even a second copy of the parameters costs real memory. Every optimizer is a memory budget in disguise.

Gradient descent and SGD

Start from the simplest possible rule. The gradient points uphill, so step downhill by a fixed fraction of it.

Here $\theta$ is the parameter vector, $\eta$ (eta) is the learning rate, and $g_t$ is the gradient at step $t$. That is all: direction is the negative gradient, distance is $\eta$ times its length.

That last sentence is the whole conditioning problem in one line. If one direction has curvature $c_{\max}$ and another has $c_{\min}$, the learning rate must be below $2/c_{\max}$ to be stable, but the flat direction then converges at a rate of $1 - \eta\, c_{\min} \approx 1 - 2c_{\min}/c_{\max}$. With a condition number of $10^4$ you need thousands of steps to move along the flat axis. Look again at the pink path in Figure 1.

Minibatch noise

We cannot afford the full-dataset gradient, so we average over a minibatch of $B$ examples. The result is an unbiased but noisy estimate: its mean is the true gradient, and its variance shrinks like $1/B$.

This matters in two ways. First, the noise sets a floor: with a fixed learning rate, SGD does not settle at the minimum but jitters around it in a cloud whose radius scales with $\eta \sigma$, where $\sigma$ is the gradient noise. Decaying $\eta$ shrinks the cloud, which is one reason learning-rate schedules exist. Second, noise is not evenly spread. Rare tokens get a gradient in only a few batches, so their embedding rows see a signal that is mostly zero with occasional spikes. Keep that picture; it is why Adam wins on transformers.

Momentum: the heavy ball

What breaks in Figure 1 is that consecutive gradients disagree along the steep axis (they flip sign) and agree along the flat axis (they keep pointing the same way). Averaging consecutive gradients would cancel the first and reinforce the second.

So keep a running velocity that accumulates gradients, and step along the velocity instead of the raw gradient.

$\beta$ (typically 0.9) is how much of the old velocity survives each step. What just happened: $v_t$ is a weighted sum of all past gradients, $g_t + \beta g_{t-1} + \beta^2 g_{t-2} + \dots$. Gradients that keep agreeing add up to at most $1/(1-\beta)$ times a single gradient (10× for $\beta = 0.9$), while gradients that alternate cancel.

Because of that $1/(1-\beta)$ factor, people sometimes write the update as $v_t = \beta v_{t-1} + (1-\beta) g_t$, an exponential moving average with the same fixed point but with the ten-fold gain folded into $\eta$. PyTorch's SGD uses the first form; Adam (below) uses the second. When you compare learning rates across papers, check which one they mean.

Nesterov momentum, briefly

Nesterov's variant evaluates the gradient at the look-ahead point $\theta_t - \eta\beta v_{t-1}$ (where the velocity is about to take us) instead of at $\theta_t$. If the ball is about to overshoot, the look-ahead gradient already points back, so the correction arrives one step earlier. In practice it is a mild but consistent improvement over classical momentum, and it is the default inside Muon.

Adaptive methods: one learning rate per parameter

Momentum fixes the zig-zag but not the scale mismatch. If parameter A sees gradients of size $10$ and parameter B sees gradients of size $0.01$, a single $\eta$ is either far too big for A or hopelessly small for B. The fix is obvious once said: give each parameter its own learning rate, chosen automatically from the size of the gradients it has been getting.

AdaGrad

AdaGrad (Duchi et al., 2011) keeps, per parameter, the sum of all squared gradients so far, and divides the step by its square root.

All operations are elementwise. A parameter that has seen big gradients gets a small effective learning rate; one that has seen tiny gradients gets a large one. The catch: $s_t$ only grows, so every learning rate decays to zero and training stalls. Good for convex problems, bad for a million-step transformer run.

RMSProp

The fix is to forget: replace the sum with an exponential moving average, so the normalizer tracks recent gradient size instead of the whole history.

With $\beta_2 = 0.9$, $s_t$ is roughly the mean of the last ten squared gradients. Divide by its root and the update has a size of about $\eta$ in every coordinate regardless of gradient scale. RMSProp was proposed by Tieleman and Hinton in a 2012 lecture and never got a paper; it is essentially AdaGrad with a leaky memory.

Adam: two moments and a bias correction

Adam (Kingma and Ba, 2014) combines the two ideas: a momentum-style average of the gradient (the first moment) divided by an RMSProp-style average of its square (the second moment). One more wrinkle makes it work at step one.

First, the two moving averages. Both start at zero.

Now the problem. At $t = 1$ with $\beta_1 = 0.9$, $m_1 = 0.1\, g_1$: the average is ten times too small because it was initialized at zero. Same for $v_1 = 0.001\, g_1^2$, a thousand times too small. Their ratio is badly off. So we divide each by the total weight it has accumulated so far, which is $1 - \beta^t$.

At $t=1$ this divides $m_1$ by $0.1$ and $v_1$ by $0.001$, exactly undoing the shrinkage. As $t$ grows the correction fades to 1. Finally the step:

What just happened: the direction is a smoothed gradient, the scale is normalized per parameter, and $\epsilon$ (typically $10^{-8}$) stops a parameter with zero gradient history from dividing by zero. For a gradient that is consistent in sign, $\hat m / \sqrt{\hat v} \approx \pm 1$, so every parameter moves by about $\eta$ per step. Adam is nearly a sign-of-gradient method with a smoothed sign.

What ε actually does

$\epsilon$ looks like numerical hygiene, but it also sets a floor on the normalizer. If $\sqrt{\hat v}$ is comparable to $\epsilon$, the update degrades toward plain SGD with learning rate $\eta/\epsilon$. For tiny-gradient parameters in mixed precision this can matter: some large-model runs use $\epsilon = 10^{-6}$ or larger to damp noise, others keep $10^{-8}$. It is a real hyperparameter, just a boring one.

Why Adam beats SGD on transformers

On image classifiers, SGD with momentum matches or beats Adam. On transformers, SGD is not even close. Why?

The short answer is scale mismatch plus heavy-tailed noise. A transformer's parameters live at very different gradient scales: an embedding row for a token that appears once per ten thousand batches sees a gradient that is zero almost always and large when it fires; a LayerNorm gain sees a dense, smooth gradient every step; attention logit weights sit somewhere in between. SGD's single learning rate cannot serve all of them. Adam's per-parameter normalizer makes the rare embedding row move as far per step as the busy LayerNorm gain.

Zhang et al. (2020) measured the noise in transformer gradients and found it heavy-tailed: occasional enormous gradient samples dominate the average. Under such noise, methods that normalize (and, they argue, gradient clipping, which does something similar) converge where plain SGD does not. Kunstner et al. (2023) went a step further: they showed that the sign-descent behaviour of Adam (recall that $\hat m / \sqrt{\hat v} \approx \pm 1$) explains most of the gap, and that even sign descent with momentum closes much of it, particularly at large batch sizes where the noise is tamed.

AdamW: decoupled weight decay

Weight decay is the oldest regularizer there is: pull every parameter gently toward zero each step, so weights stay small and the model does not overfit to whatever quirks the data has. With SGD there are two equivalent ways to write it. You can add an $L_2$ penalty $\frac{\lambda}{2}\|\theta\|^2$ to the loss, which adds $\lambda\theta$ to the gradient, or you can shrink the weights directly after the gradient step. With SGD these give identical updates: $\theta \leftarrow \theta - \eta(g + \lambda\theta) = (1-\eta\lambda)\theta - \eta g$.

Under Adam they are not the same, and this is the entire content of the AdamW paper (Loshchilov and Hutter, 2019). If you add $\lambda\theta$ to the gradient, it goes through the $1/\sqrt{\hat v}$ normalizer like everything else. A weight with large historical gradients gets its decay term divided by a large number, so it is barely regularized at all; a weight with tiny gradients gets a huge effective decay. The regularization strength depends on the gradient history, which is not what anyone wanted.

AdamW's fix: keep the decay out of the gradient and apply it separately, after the normalized step.

Now every weight shrinks by the same fraction $\eta\lambda$ per step, regardless of its gradients. Note that $\lambda$ multiplies $\eta$, so when the learning-rate schedule decays, so does the decay; some implementations keep them separate. PyTorch's AdamW uses this coupled-to-$\eta$ form.

Learning rate: warmup and schedules

Adam's per-parameter scaling assumes $\hat v$ is a decent estimate of the gradient's typical size. In the first few hundred steps it is not: it has seen a handful of samples from a network whose gradients are about to change wildly as the loss drops. The remedy is warmup: ramp $\eta$ linearly from zero over the first 1–2% of training, so the early, poorly-normalized steps are small.

After warmup, decay. Cosine decay to about 10% of the peak is the common LLM default; some runs use linear decay to zero, and "warmup-stable-decay" (a long flat plateau followed by a short sharp decay) makes it easier to extend a run. The reason to decay at all is the noise cloud from earlier: a smaller $\eta$ shrinks the jitter around the minimum, and a large fraction of the final loss improvement arrives during the decay phase.

Gradient clipping

Occasionally a minibatch produces a gradient hundreds of times larger than usual: a bad batch, an attention logit spike, a numerical hiccup. Adam's $v$ takes ~1,000 steps to absorb such an outlier, so the normalizer is briefly far too small and the model takes a huge step. That is a loss spike, and it can take thousands of steps to recover from.

The standard defence is global-norm clipping (Pascanu et al., 2013): if the norm of the whole gradient vector exceeds a threshold $c$, scale the whole vector down to have norm $c$.

Direction is preserved, only the length is capped. Almost every LLM run uses $c = 1.0$. It is cheap insurance, and Zhang et al. (2020) argue it is also part of why adaptive-style methods handle heavy-tailed noise. Watch the fraction of steps that get clipped: if it climbs from a few percent to most steps, something upstream is wrong (usually the learning rate).

Large-batch training: LARS and LAMB

When you train on thousands of GPUs you want batches of millions of tokens, and past a certain critical batch size (McCandlish et al., 2018) bigger batches stop buying proportionally faster convergence. The problem: to use a huge batch efficiently you must raise the learning rate, and at high learning rates some layers become unstable before others.

LARS (You et al., 2017) and LAMB (You et al., 2019) address this by normalizing the update per layer: the step for a weight matrix is scaled so that its norm is a fixed fraction of the weight's own norm, $\|\Delta W\| \propto \|W\|$, in what is called a trust ratio. LAMB applies this on top of Adam. It made batch sizes of 32k–64k viable for BERT pretraining, cutting wall-clock time from days to about an hour on a large TPU pod. For typical LLM pretraining today, batches are chosen below the critical size and plain AdamW is used; LAMB shows up mostly in extreme scale-out settings.

Memory: Adafactor, 8-bit Adam, GaLore

Here is a number people forget: for a model with $N$ parameters, Adam stores $2N$ extra floats ($m$ and $v$). In fp32 that is 8 bytes per parameter, on top of the 2-byte bf16 weight and its fp32 master copy. For a 70B model the optimizer state alone is 560 GB. Optimizer memory, not weights, is often what forces you onto more GPUs.

Adafactor: factor the second moment

Shazeer and Stern (2018) noticed that $v$ for a weight matrix $W \in \mathbb{R}^{m \times n}$ is itself an $m \times n$ matrix, and that it can be approximated by a rank-1 outer product of its row sums and column sums.

Storing $R$ and $C$ costs $m + n$ numbers instead of $mn$. For a $4096 \times 4096$ matrix that is 8k numbers instead of 16M, a 2,000× saving on the second moment. Adafactor also drops the first moment by default (or keeps it in low precision) and uses a relative step size ($\eta$ scaled by the RMS of the weights). It was the optimizer for T5 and PaLM.

8-bit Adam: quantize the state

Dettmers et al. (2022) keep the full $m$ and $v$ but store each in 8 bits using block-wise quantization: every block of 2,048 values gets its own scale, and values are mapped to a non-uniform (dynamic) 8-bit grid. States are dequantized to fp32 for the update and re-quantized after. Optimizer memory drops 4× with, in their experiments, no loss in quality. This is what bitsandbytes gives you.

GaLore: project the gradient

GaLore (Zhao et al., 2024) takes a different angle. The gradient of a weight matrix tends to be low-rank, so project it onto a rank-$r$ subspace (via an SVD refreshed every few hundred steps), run Adam in that small space, and project the update back. The optimizer state becomes $O((m+n)r)$ instead of $O(mn)$. The authors report pretraining a 7B model on a single 24 GB consumer GPU. Unlike LoRA, the full weights are still trained, only the optimizer state is compressed.

Lion: keep only the sign

Chen et al. (2023) used an automated search over update rules and found a strikingly simple winner: interpolate the gradient with momentum, take the sign, step.

Only one state tensor ($m$), so half of Adam's memory. Since every coordinate moves by exactly $\pm\eta$, Lion needs a learning rate 3–10× smaller than Adam and a correspondingly larger $\lambda$ to get the same effective decay. The sign makes explicit what Adam does approximately (recall $\hat m/\sqrt{\hat v} \approx \pm 1$), which is a neat confirmation of the Kunstner et al. (2023) story.

Shampoo and SOAP: preconditioning with Kronecker factors

Adam's normalizer is diagonal: each parameter is scaled independently. But the parameters of a weight matrix are not independent; the rows interact through the input and the columns through the output. A full second-order preconditioner for an $m \times n$ matrix would be an $mn \times mn$ matrix, which is hopeless. Shampoo (Gupta et al., 2018) approximates it as a Kronecker product of two small matrices.

What just happened: $L$ captures how gradient rows correlate, $R$ how columns correlate, and the inverse fourth roots on both sides whiten the gradient in both directions at once. For a $4096 \times 4096$ layer you store two $4096 \times 4096$ matrices rather than one $16\text{M} \times 16\text{M}$ one, and you need a matrix root, which is computed every few hundred steps and amortized.

A distributed Shampoo implementation won the external-tuning track of the AlgoPerf benchmark in 2024, and SOAP (Vyas et al., 2024) sharpened the idea: run Adam in the eigenbasis of Shampoo's preconditioners, refreshing that basis only occasionally. SOAP reports fewer steps and less wall-clock than AdamW on 360M–660M models, at the cost of more memory and a couple of extra hyperparameters.

Muon: orthogonalize the momentum

Muon (Jordan et al., 2024) is the optimizer that made people re-examine whether Adam is the final word for transformers. It is simple to state, cheap, and, as reported in a 2025 paper by Liu et al. and in the Kimi K2 technical report, has been used at trillion-parameter scale with a modified clipping scheme.

The problem Muon solves

Look at the gradient (or momentum) of a single weight matrix in a transformer. Take its SVD: $M = U \Sigma V^\top$. Empirically its singular values are wildly uneven: a few directions carry almost all the energy, and the rest are tiny. An SGD step along $M$ moves the weight almost entirely along those few dominant directions. Adam's elementwise scaling does not fix this, because the dominant directions are spread across all elements, not concentrated in a few coordinates.

The rare directions are exactly the ones we want to learn from: they are the "rare features" of the matrix update, analogous to rare tokens for embeddings. Muon's idea is to give every direction the same step: replace $M$ by the closest orthogonal matrix, $UV^\top$, which has every singular value equal to 1.

Why not just compute the SVD?

An SVD of a $4096 \times 4096$ matrix every step for every layer is too slow and does not run well in bf16 on GPUs. The trick is that we do not need the SVD itself, only $UV^\top$, and there is a classical iteration for that: Newton–Schulz. Applying an odd polynomial $p(X) = aX + b(XX^\top)X + c(XX^\top)^2 X$ to a matrix acts on its singular values alone: each $\sigma_i$ becomes $p(\sigma_i) = a\sigma_i + b\sigma_i^3 + c\sigma_i^5$ while $U$ and $V$ stay fixed. Choose the coefficients so that repeated application pushes every $\sigma$ in $(0, 1]$ toward 1, and you get $UV^\top$ from matrix multiplications only.

The Muon recipe, with the coefficients from the reference implementation:

The Frobenius normalization first guarantees every singular value is at most 1, which is the range where the polynomial is contractive toward 1. Five iterations are enough. These particular coefficients were tuned for speed rather than exact convergence: they push tiny singular values up very fast (the slope at zero is $a = 3.44$), at the price of the values settling into a band around 1 rather than at exactly 1. For an optimizer that is fine; what matters is that no direction is starved.

The Muon algorithm

Put it together. For each 2-D weight matrix $W$ with gradient $G$:

with Nesterov-style momentum ($\mu = 0.95$) by default, and, in most implementations, a scale factor on $O_t$ such as $\sqrt{\max(1, m/n)}$ so that the update's RMS does not depend on the matrix's aspect ratio. Weight decay is applied decoupled, as in AdamW.

What Muon applies to, and what it does not

Orthogonalization is a matrix operation, so Muon only makes sense for parameters that are matrices acting as linear maps: the attention projections and MLP weights. Embeddings, the output head, biases, and LayerNorm/RMSNorm gains are handled by AdamW in every practical setup. The reasons are partly principled (an embedding table is a lookup, not a linear map on activations; its rows should be scaled per token, which is what Adam does) and partly empirical (the output head trained with Muon was reported to be worse).

Convolution kernels can be reshaped to matrices, and attention's fused QKV weight is usually split into its three matrices before orthogonalizing, so that each gets its own equalized update.

import torch

def newton_schulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Approximate U V^T for G = U S V^T using the quintic Newton-Schulz iteration."""
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.to(torch.bfloat16) if G.is_cuda else G.float()
    transposed = X.size(0) > X.size(1)
    if transposed:                      # work on the "wide" orientation: X X^T is the smaller product
        X = X.T
    X = X / (X.norm() + 1e-7)           # Frobenius-normalize so all singular values are <= 1
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X               # a X + b (X X^T) X + c (X X^T)^2 X
    return (X.T if transposed else X).to(G.dtype)


class Muon(torch.optim.Optimizer):
    """Momentum + orthogonalized update. Use only for 2-D weights; give the rest to AdamW."""
    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, weight_decay=0.0):
        super().__init__(params, dict(lr=lr, momentum=momentum, nesterov=nesterov, weight_decay=weight_decay))

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None or p.ndim != 2:
                    continue
                st = self.state[p]
                if "buf" not in st:
                    st["buf"] = torch.zeros_like(p.grad)
                buf = st["buf"]
                buf.mul_(group["momentum"]).add_(p.grad)
                g = p.grad.add(buf, alpha=group["momentum"]) if group["nesterov"] else buf
                upd = newton_schulz5(g)
                upd *= max(1.0, p.size(0) / p.size(1)) ** 0.5   # keep update RMS independent of shape
                if group["weight_decay"]:
                    p.mul_(1 - group["lr"] * group["weight_decay"])
                p.add_(upd, alpha=-group["lr"])

Note the three practical choices in the code: work on the orientation where $XX^\top$ is the smaller matrix, run the iteration in bf16 (the coefficients were chosen to tolerate it), and apply the aspect-ratio scale. The companion code/lumen/optimizers.py has this alongside from-scratch SGD, momentum, Adam and AdamW, all with the same interface as torch.optim so you can swap them into code/lumen/train.py.

Comparison table

Memory is counted in extra numbers per parameter (the weight itself and its gradient not included). "Hyperparameters" lists the ones you realistically tune.

Practical LLM settings

These are typical values from published pretraining recipes (GPT-3, Brown et al. 2020; LLaMA, Touvron et al. 2023) and open reproductions, not laws. The one robust rule is that peak learning rate falls as the model grows.

Practice

Further reading

• Kingma & Ba (2014). Adam: A Method for Stochastic Optimization. The original; Section 3 has the bias-correction derivation.
• Loshchilov & Hutter (2019). Decoupled Weight Decay Regularization. Why AdamW exists, with the figure that convinced everyone.
• Zhang et al. (2020). Why are Adaptive Methods Good for Attention Models? Heavy-tailed gradient noise in transformers and the role of clipping.
• Kunstner et al. (2023). Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be. The sign-descent explanation.
• Pascanu et al. (2013). On the difficulty of training recurrent neural networks. Where gradient-norm clipping was formalized.
• Shazeer & Stern (2018). Adafactor: Adaptive Learning Rates with Sublinear Memory Cost. Factored second moments.
• Dettmers et al. (2022). 8-bit Optimizers via Block-wise Quantization. Quantized optimizer state.
• Zhao et al. (2024). GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection.
• Chen et al. (2023). Symbolic Discovery of Optimization Algorithms. The Lion optimizer.
• You et al. (2019). Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. LAMB.
• McCandlish et al. (2018). An Empirical Model of Large-Batch Training. The critical batch size.
• Gupta et al. (2018). Shampoo: Preconditioned Stochastic Tensor Optimization. Kronecker-factored preconditioning.
• Vyas et al. (2024). SOAP: Improving and Stabilizing Shampoo using Adam.
• Jordan et al. (2024). Muon: An optimizer for hidden layers in neural networks. The blog post that introduced Muon, with the Newton–Schulz coefficients and the reference implementation.
• Bernstein & Newhouse (2024). Old Optimizer, New Norm: An Anthology. Adam, Shampoo and Muon as steepest descent under different norms.
• Liu et al. (2025). Muon is Scalable for LLM Training. Reported scaling of Muon to a 16B MoE, with decoupled decay and RMS matching.
• Moonshot AI (2025). Kimi K2: Open Agentic Intelligence. Reports MuonClip and QK-clip for a 1T-parameter MoE.
• Brown et al. (2020). Language Models are Few-Shot Learners. Appendix has the GPT-3 optimizer settings by model size.
• Touvron et al. (2023). LLaMA: Open and Efficient Foundation Language Models. A clean, widely-copied AdamW recipe.