3.9 Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo¶

The optimizer is the engine that turns gradients into weight updates. Choose it well and a 70-billion-parameter model converges smoothly on a fixed token budget; choose it poorly and you either waste GPU-months or watch the loss spike into NaNs. For most of the deep-learning era one optimizer — Adam, and its weight-decay-fixed cousin AdamW — has been the default for training transformers, and for good reason. But Adam carries a hidden tax: it stores two extra full-precision tensors per parameter, and at scale that memory cost rivals the model weights themselves. That tension — fast, robust convergence versus memory and compute overhead — is the throughline of this chapter, and it is exactly what newer optimizers like Adafactor, Lion, Shampoo, and Muon attack from different angles.

We will build up from the gradient-descent first principles you saw in Calculus, Optimization & Convexity, derive Adam and its bias correction carefully, implement AdamW from scratch, account for optimizer-state memory the way a systems engineer must, and then tour the modern menagerie: factored second moments (Adafactor), sign-based updates (Lion), full preconditioners (Shampoo), and the newest newsmaker, Muon, which orthogonalizes the update of every weight matrix. By the end you should be able to pick an optimizer for a given memory and stability budget, and defend that choice in an interview.

From Gradient Descent to Momentum¶

The starting point is plain stochastic gradient descent (SGD). Given a loss \(L(\theta)\) and a minibatch estimate \(g_t = \nabla_\theta L_{\mathcal{B}_t}(\theta_{t-1})\) of its gradient, SGD takes the step

where \(\eta\) is the learning rate. This is the cheapest possible optimizer: zero extra state, one tensor of gradients, one fused multiply-subtract. Its weakness is that the raw gradient is a noisy, badly-scaled descent direction. In a ravine — a loss surface that is steep across one axis and nearly flat along another, which is the typical geometry of deep nets — SGD zig-zags across the steep walls while crawling along the flat valley floor. The condition number of the Hessian, the ratio of its largest to smallest eigenvalue, controls how bad this is; for transformers it can be enormous.

Momentum fixes the zig-zag by accumulating an exponentially-weighted average of past gradients, a “velocity” \(v_t\), and stepping along that instead:

Here \(\mu \in [0,1)\) (typically 0.9) is the momentum coefficient. The intuition is physical: the velocity behaves like a heavy ball rolling downhill. Oscillating components of the gradient cancel across steps, while the consistent down-valley component accumulates. The geometric-series identity tells us that a steady gradient \(g\) produces a terminal velocity of \(g/(1-\mu)\), so momentum with \(\mu=0.9\) effectively multiplies the step length along consistent directions by \(10\times\). That is also why you usually drop the learning rate when you increase momentum.

A subtle but important variant is Nesterov momentum, which evaluates the gradient at the look-ahead point \(\theta_{t-1} - \eta\mu v_{t-1}\) rather than at \(\theta_{t-1}\). The correction term gives a slightly more responsive update and a better convergence rate on convex problems. SGD with Nesterov momentum and a well-tuned learning-rate schedule remains the gold standard for training convolutional vision models, and it generalizes beautifully. So why do we not train transformers with it?

import torch

def sgd_momentum_step(params, grads, velocities, lr=0.1, mu=0.9, nesterov=False):
    """One step of (Nesterov) momentum SGD, in-place. Pure PyTorch tensors."""
    for p, g, v in zip(params, grads, velocities):
        v.mul_(mu).add_(g)              # v <- mu*v + g
        if nesterov:
            update = g.add(v, alpha=mu)  # g + mu*v  (look-ahead)
        else:
            update = v
        p.add_(update, alpha=-lr)        # theta <- theta - lr * update

The answer is scale heterogeneity. In a transformer the gradient magnitudes differ wildly across parameters — embedding rows that fire rarely versus LayerNorm gains versus attention projections — and a single global learning rate cannot serve all of them. We need a per-parameter adaptive learning rate. That is what Adam provides.

Adam and AdamW: Derivation, Bias Correction, Decoupled Decay¶

Adam (Kingma & Ba, Adam: A Method for Stochastic Optimization, 2015) combines two ideas: momentum on the gradient (the first moment), and a per-coordinate rescaling by the running magnitude of the gradient (the second moment, an idea inherited from RMSProp and AdaGrad). It maintains two exponential moving averages:

where \(g_t^2\) is elementwise. Typical defaults are \(\beta_1 = 0.9\), \(\beta_2 = 0.999\) (for LLMs \(\beta_2 = 0.95\) is common — more on that below). \(m_t\) estimates \(\mathbb{E}[g]\) and \(v_t\) estimates \(\mathbb{E}[g^2]\).

Why bias correction is necessary¶

Both averages are initialized at zero, which biases them toward zero in early steps. Unroll \(v_t\) assuming a stationary gradient distribution:

Taking expectations and assuming \(\mathbb{E}[g_i^2]\) is approximately constant \(\approx \mathbb{E}[g_t^2]\),

The factor \((1-\beta_2^{\,t})\) is the bias: at \(t=1\) with \(\beta_2=0.999\), \(v_1\) is only \(0.1\%\) of the true second moment, so \(\sqrt{v_1}\) is \(\approx\sqrt{1000}\approx 31\times\) too small. The first moment is biased too (\(m_1 = (1-\beta_1)g = 0.1\,g\)), and the two biases partially cancel — an update with neither correction is inflated by \((1-\beta_1)/\sqrt{1-\beta_2} = 0.1/0.0316 \approx 3.2\times\) at \(t=1\), still enough to wreck a freshly initialized model. Crucially the two biases decay at different rates (\(\beta_1^t\) vs \(\beta_2^t\)), so the mismatch persists for hundreds of steps and must be removed from each moment separately:

The final update normalizes the first moment by the root second moment, with a small \(\epsilon\) (e.g. \(10^{-8}\)) for numerical safety:

The deep insight is in the ratio \(\hat{m}_t / \sqrt{\hat{v}_t}\). For a coordinate with a consistent gradient, \(\hat{m}_t \approx \sqrt{\hat{v}_t}\), so the update magnitude is \(\approx \eta\) — Adam takes a step of roughly unit size in each coordinate, regardless of the gradient’s absolute scale. This is the scale-invariance that makes Adam so robust: you can change the loss scaling (or use mixed precision, see Mixed Precision, bf16 & FP8 Training) and Adam’s effective step is unchanged. It is also why Adam tolerates the wildly heterogeneous gradient scales of a transformer that defeat plain SGD.

Decoupled weight decay: the W in AdamW¶

L2 regularization adds \(\tfrac{\lambda}{2}\lVert\theta\rVert^2\) to the loss, which contributes \(\lambda\theta\) to the gradient. In plain SGD this is identical to weight decay (shrinking \(\theta\) toward zero each step). But in Adam they are not equivalent: the \(\lambda\theta\) term flows through the \(m_t\) and \(v_t\) machinery and gets divided by \(\sqrt{\hat v_t}\), so parameters with large gradient magnitude (large \(v_t\)) get less decay — exactly backwards from what you want. Loshchilov & Hutter (Decoupled Weight Decay Regularization, 2019) showed this hurts, and proposed AdamW, which applies the decay directly to the weights, decoupled from the adaptive term:

Equivalently \(\theta_t = (1-\eta\lambda)\,\theta_{t-1} - \eta\,\hat m_t/(\sqrt{\hat v_t}+\epsilon)\): a clean multiplicative shrink toward zero plus the adaptive step. AdamW is the de-facto standard for pretraining every modern LLM. A crucial practical detail: do not decay 1-D parameters — biases, LayerNorm/RMSNorm gains, and usually embeddings — only the 2-D weight matrices. Decaying norm gains pulls them toward zero and destabilizes training.

Implementing AdamW from scratch¶

Here is a complete, correct, heavily-commented AdamW that you could drop into a real training loop. It mirrors the math above and matches torch.optim.AdamW numerically.

import torch
from torch.optim import Optimizer

class AdamW(Optimizer):
    """From-scratch AdamW (Loshchilov & Hutter, 2019).

    Stores two state tensors per parameter: exp_avg (m) and exp_avg_sq (v).
    This is the memory tax we account for in the next section.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.95),
                 eps=1e-8, weight_decay=0.1):
        # betas: (beta1, beta2). For LLMs beta2=0.95 is the common choice.
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for group in self.param_groups:
            lr, (b1, b2) = group["lr"], group["betas"]
            eps, wd = group["eps"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                if g.is_sparse:
                    raise RuntimeError("AdamW does not support sparse grads")

                state = self.state[p]
                if len(state) == 0:                    # lazy init
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(p)      # m
                    state["exp_avg_sq"] = torch.zeros_like(p)   # v
                m, v = state["exp_avg"], state["exp_avg_sq"]
                state["step"] += 1
                t = state["step"]

                # --- Decoupled weight decay: multiplicative shrink toward 0 ---
                # Applied to the *weights*, NOT folded into the gradient.
                if wd != 0:
                    p.mul_(1.0 - lr * wd)

                # --- Update biased first and second moment estimates ---
                m.mul_(b1).add_(g, alpha=1.0 - b1)          # m = b1*m + (1-b1)*g
                v.mul_(b2).addcmul_(g, g, value=1.0 - b2)   # v = b2*v + (1-b2)*g^2

                # --- Bias correction ---
                bias_c1 = 1.0 - b1 ** t
                bias_c2 = 1.0 - b2 ** t
                # Fold bias_c2 into the denominator; step_size folds bias_c1.
                denom = (v.sqrt() / (bias_c2 ** 0.5)).add_(eps)
                step_size = lr / bias_c1

                # theta <- theta - step_size * m / denom
                p.addcdiv_(m, denom, value=-step_size)
        return loss

A few implementation notes that separate a toy from a production optimizer. The @torch.no_grad() decorator is mandatory — the update itself must not build an autograd graph. Operations are in-place (mul_, add_, addcmul_, addcdiv_) to avoid allocating new tensors every step; for a 7B model a single full-size temporary is 14 GB in bf16. Real implementations go further with fused or foreach kernels (torch.optim.AdamW(..., fused=True)) that batch the elementwise math across all parameters into one CUDA launch, which is a large throughput win when you have thousands of parameter tensors.

Parameter groups and the libraries you actually call¶

In a real run you never hand model.parameters() to the optimizer as one blob — the decay/no-decay split from above has to be expressed as parameter groups, and the optimizer itself comes from a library. This is the whole “optimizer layer” of a training script:

import torch

model = torch.nn.Sequential(              # stand-in for your transformer
    torch.nn.Embedding(1000, 64), torch.nn.LayerNorm(64), torch.nn.Linear(64, 64)
)

def param_groups(model, weight_decay=0.1):
    """Two groups: 2-D matmul weights decay, 1-D params (norms, biases) do not."""
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        # nanoGPT/GPT-3 convention: everything with ndim >= 2 decays, including
        # the embedding and LM head. Some recipes exclude those by name — be
        # explicit either way, and never let norm gains or biases into `decay`.
        (decay if p.ndim >= 2 else no_decay).append(p)
    return [{"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0}]

# The 99% case: PyTorch's fused AdamW. `fused=True` runs the whole elementwise
# update as one multi-tensor CUDA kernel (params must be CUDA + floating point);
# the default `foreach=True` path is the CPU/other-backend fallback.
opt = torch.optim.AdamW(param_groups(model), lr=3e-4, betas=(0.9, 0.95),
                        eps=1e-8, fused=torch.cuda.is_available())

# Memory-bound instead? bitsandbytes is a drop-in replacement, same signature:
#   import bitsandbytes as bnb
#   opt = bnb.optim.AdamW8bit(param_groups(model), lr=3e-4, betas=(0.9, 0.95))
# (bitsandbytes' GlobalOptimManager lets you force the embedding table back to
#  32-bit optimizer state, which is the usual stability precaution.)

The same layer looks different inside the big training frameworks, and it is worth knowing which knob is which: DeepSpeed swaps in deepspeed.ops.adam.FusedAdam automatically, and DeepSpeedCPUAdam when you enable ZeRO-Offload so the optimizer step runs on CPU; Megatron-LM shards optimizer state ZeRO-1-style behind --use-distributed-optimizer; PyTorch FSDP shards it as a consequence of sharding the parameters themselves (see Megatron-LM, DeepSpeed & Parallelism in Practice). In every case the math is the AdamW you just wrote — only the sharding and the kernel change.

Large-batch training: LARS, LAMB and the trust ratio¶

When you scale the global batch to tens of thousands of sequences to keep more GPUs busy, a single global learning rate becomes the bottleneck again — this time across layers. The ratio \(\lVert\Delta\theta\rVert / \lVert\theta\rVert\) varies by orders of magnitude between layers, and the layer with the largest ratio diverges first. LARS (You et al., 2017) and its Adam-based successor LAMB (You et al., Large Batch Optimization for Deep Learning: Training BERT in 76 Minutes, 2020) fix this with a per-layer trust ratio. Writing \(u_t\) for the layer’s raw update — for LAMB, the AdamW update \(\hat m_t/(\sqrt{\hat v_t}+\epsilon) + \lambda\theta_{t-1}\) — the step becomes

with the ratio clamped to 1 when either norm is zero. Every layer now moves by a fixed fraction of its own weight norm per step, which is what let BERT pretrain at batch size 32k. For LLM pretraining today, tuned AdamW with warmup and muP-style width scaling (see Learning Rate Schedules, Warmup, Batch Size & Hyperparameters) has largely displaced LAMB — but the diagnostic is permanent: log \(\lVert\Delta\theta\rVert/\lVert\theta\rVert\) per layer during training. A common rule of thumb is that healthy training sits on the order of \(10^{-3}\); a layer sitting orders of magnitude above the rest is where your next loss spike will come from.

The Memory Cost of Optimizer States¶

Here is the systems reality that motivates everything in the rest of the chapter. Consider a model with \(P\) parameters. With AdamW you store, at minimum:

That last row appears because mixed-precision training keeps a high-precision master copy of the weights so that tiny updates are not lost to bf16 rounding (see Mixed Precision, bf16 & FP8 Training). The headline number from the ZeRO/DeepSpeed analysis is 16 bytes per parameter for the model+optimizer state (2+2+4+4+4 = 16), of which the optimizer alone (m, v, master weights) is 12 bytes — three times the 4 bytes of bf16 weights and gradients together.

Two orthogonal strategies attack this. The first is systems-level: ZeRO/FSDP shards the 12 bytes of optimizer state across \(N\) data-parallel ranks, so each rank holds only \(12P/N\) bytes — no change to the math, pure distribution. The second is algorithmic: redesign the optimizer to store less state per parameter. Adafactor, Lion, and the sign-based family live here, and that is where we turn next. The two strategies compose: you can shard a Lion optimizer too.

Adafactor: Factoring the Second Moment¶

Adafactor (Shazeer & Stern, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost, 2018) was born from exactly the memory pressure above, originally for training large T5 models on TPUs. Its key observation: the costliest state is the second moment \(v\), a full \(P\)-element tensor. For a weight matrix \(W \in \mathbb{R}^{n\times m}\), instead of storing the full \(n\times m\) matrix of second moments, Adafactor stores only a row vector \(R\in\mathbb{R}^{n}\) and a column vector \(C\in\mathbb{R}^{m}\) and reconstructs a rank-1 approximation:

This is the best rank-1 (in a generalized-KL sense) factorization of the true second-moment matrix. The memory for the second moment drops from \(O(nm)\) to \(O(n+m)\) — sublinear in the parameter count. For a \(4096\times 4096\) matrix that is \(16.7\)M numbers versus \(8192\), a \(2000\times\) reduction for that tensor’s second-moment state.

Adafactor pairs this with two more memory moves: it can drop the first moment entirely (\(\beta_1 = 0\), no momentum), and it uses relative step sizes scaled by the RMS of the parameters so it needs no external learning-rate tuning. It also adds update clipping by RMS norm for stability. The factored update for a matrix, in pseudocode:

import torch

def adafactor_matrix_step(W, G, R, C, t, lr, beta2=0.999, eps1=1e-30, eps2=1e-3):
    """One Adafactor step for a 2D weight W with grad G.
    R: row accumulator (n,), C: col accumulator (m,). No first moment here.
    """
    n, m = W.shape
    g2 = G * G + eps1                       # squared grad, floored

    # Decayed running averages of row sums and column sums of g^2
    beta2_t = 1.0 - t ** (-0.8)             # Adafactor's time-dependent decay
    R.mul_(beta2_t).add_(g2.mean(dim=1), alpha=1 - beta2_t)   # (n,)
    C.mul_(beta2_t).add_(g2.mean(dim=0), alpha=1 - beta2_t)   # (m,)

    # Rank-1 reconstruction of the second-moment estimate V_hat (n,m)
    R_factor = (R / R.mean()).rsqrt().unsqueeze(1)   # (n,1)
    C_factor = C.rsqrt().unsqueeze(0)                # (1,m)
    update = G * R_factor * C_factor                 # G / sqrt(V_hat)

    # RMS-clip the update (Adafactor's stability trick)
    rms = update.pow(2).mean().sqrt()
    update = update / max(1.0, (rms / 1.0).item())

    # Relative step size scaled by parameter RMS
    param_rms = W.pow(2).mean().sqrt().clamp_min(eps2)
    W.add_(update, alpha=-lr * param_rms.item())

The trade-off is real: Adafactor’s factored second moment and missing momentum make it slightly noisier and sometimes less stable than AdamW, and it can need more babysitting (warmup, the eps2 floor). But it cut optimizer memory roughly in half-to-two-thirds and made T5-scale training feasible. It remains popular for fine-tuning large models on memory-constrained hardware, and its rank-1 factorization idea echoes in later work. A related modern option, 8-bit Adam (Dettmers et al.), takes the orthogonal route of keeping full \(m\) and \(v\) but quantizing them to 8 bits with block-wise scaling, cutting their footprint \(4\times\) with almost no quality loss — and composes with everything else.

Lion: Learning the Sign of the Update¶

Lion (Chen et al., Symbolic Discovery of Optimization Algorithms, 2023) was discovered by a program-search procedure over optimizer programs, and the winner is startlingly simple. “Lion” stands for Evolved Sign Momentum. It keeps a single momentum buffer (so only 4 extra bytes/param versus Adam’s 8) and the update direction is the sign of an interpolated momentum:

Note the two different interpolations: the update uses \(\beta_1\) (e.g. 0.9) while the momentum buffer is updated with \(\beta_2\) (e.g. 0.99). The \(\operatorname{sign}\) is the heart of it: every parameter moves by exactly \(\pm\eta\) (plus decay), independent of gradient magnitude. This is an extreme form of the scale-invariance Adam approximates — Lion makes it exact and uniform.

import torch

def lion_step(p, g, m, lr=1e-4, beta1=0.9, beta2=0.99, wd=0.0):
    """One Lion update, in-place. Only ONE state tensor m per parameter."""
    # Update direction uses an interpolation with beta1...
    c = m.mul(beta1).add(g, alpha=1.0 - beta1)   # beta1*m + (1-beta1)*g (temp)
    update = c.sign()                            # +/-1 per coordinate
    if wd != 0:
        p.mul_(1.0 - lr * wd)                    # decoupled weight decay
    p.add_(update, alpha=-lr)                    # theta <- theta - lr*sign(c)
    # ...but the stored momentum uses beta2 (note: g, not c)
    m.mul_(beta2).add_(g, alpha=1.0 - beta2)

Because the update magnitude is uniformly \(\eta\), Lion’s effective step is larger and more uniform than AdamW’s, so the recommended learning rate is roughly \(3\)–\(10\times\) smaller than Adam’s, and the weight decay correspondingly \(3\)–\(10\times\) larger (to keep \(\eta\lambda\) in a sane range). When tuned, Lion matches or beats AdamW on many vision and language pretraining tasks while using half the optimizer memory and slightly less compute (no square root, no second buffer). Its weaknesses: the sign update injects more gradient noise, so it tends to need larger batch sizes to behave, and it can be touchier near the end of training. Still, Lion is the cleanest demonstration that you do not need a per-coordinate magnitude estimate at all — a good sign plus momentum is often enough.

Second-Order Methods: Shampoo and Preconditioning¶

Everything so far uses only diagonal curvature information — Adam’s \(\sqrt{v}\) is a diagonal preconditioner, one scalar per coordinate. The ideal update, from a second-order Taylor expansion, is Newton’s method:

where \(H\) is the Hessian (or, in practice, the Fisher / empirical second-moment matrix). \(H^{-1}\) rotates and rescales the gradient to undo the loss surface’s anisotropy, fixing the ill-conditioning that diagonal methods can only partially address. The catch is fatal at scale: for \(P\) parameters \(H\) is \(P\times P\), which is \(\sim 5\times 10^{19}\) entries for a 7B model. Storing, let alone inverting, the full Hessian is impossible.

Shampoo (Gupta, Koren & Singer, Shampoo: Preconditioned Stochastic Tensor Optimization, 2018) makes second-order methods tractable for matrix-shaped parameters by using a Kronecker-factored preconditioner. For a weight matrix \(W\in\mathbb{R}^{n\times m}\) with gradient \(G\), it maintains two much smaller matrices:

and preconditions on both sides with their inverse fourth roots:

This approximates the full \(nm\times nm\) preconditioner by the Kronecker product \(L^{1/2}\otimes R^{1/2}\), capturing curvature between rows and between columns separately. The memory is \(O(n^2 + m^2)\) instead of \(O(n^2 m^2)\), and Shampoo provably converges faster per step on many problems. Its costs are the periodic computation of matrix inverse-roots (eigendecompositions of \(n\times n\) and \(m\times m\) matrices, done every \(K\) steps to amortize) and the \(L,R\) accumulators themselves. A distributed implementation (Distributed Shampoo, Anil et al.) won the 2024 AlgoPerf optimization benchmark, demonstrating that well-engineered second-order methods can beat AdamW on wall-clock time, not just step count. The practical barriers — kernel complexity, inverse-root numerics, and integrating with parameter sharding — have kept it out of most mainstream LLM recipes, but it directly inspired the optimizer we cover last.

Muon: Orthogonalizing the Update¶

Muon (Jordan et al., 2024) is the most prominent of the recent entrants and has driven much of the optimizer excitement since 2024, graduating from record-setting nanoGPT speedruns to frontier-scale production training (Moonshot’s Kimi K2 and Zhipu’s GLM-4.5 both train on Muon variants). The name stands for MomentUm Orthogonalized by Newton-Schulz. Its premise is geometric: for a 2-D weight matrix, the momentum update \(M_t\) is typically dominated by a few large singular directions — it is effectively low-rank, so most of the update’s “energy” pushes along a handful of directions and starves the rest. Muon fixes this by replacing the momentum update with its orthogonalization: the nearest semi-orthogonal matrix, which has all singular values equal to 1.

Formally, if \(M_t = U\Sigma V^\top\) is the singular value decomposition (SVD) of the momentum, Muon’s update is

Setting every singular value to 1 makes the update spectrally uniform — it pushes equally along every singular direction of the momentum, rather than letting the top few dominate. This is the matrix-valued analogue of Lion’s per-coordinate sign: Lion normalizes each scalar entry to \(\pm 1\); Muon normalizes each singular value to 1. Both are aggressive ways of discarding magnitude and keeping direction, but Muon respects the matrix structure of the weight.

The genius is how it computes \(UV^\top\) without an SVD (which is slow and hard to do well in bf16 on GPUs). It uses a Newton-Schulz iteration: a fixed sequence of matrix-multiply-only steps that drives the singular values toward 1. Starting from a spectrally-normalized \(X_0 = M/\lVert M\rVert_F\), it repeats a cubic polynomial in \(X\):

with carefully chosen coefficients \((a,b,c)\approx(3.4445, -4.7750, 2.0315)\) tuned so that roughly 5 iterations push all singular values close to 1. Crucially these are just matmuls — they run at full GPU throughput in bf16, so the orthogonalization adds only modest overhead.

import torch

@torch.no_grad()
def newton_schulz5(G, steps=5, eps=1e-7):
    """Compute an approximate UV^T (orthogonalization) of G via Newton-Schulz.
    Matmul-only; runs in bf16. Coefficients tuned to converge singular values->1.
    """
    assert G.ndim == 2
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16()
    transposed = False
    if X.size(0) > X.size(1):           # iterate on the smaller dimension
        X = X.T
        transposed = True
    X = X / (X.norm() + eps)            # spectral pre-normalization
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    if transposed:
        X = X.T
    return X

@torch.no_grad()
def muon_step(W, G, momentum_buf, lr=0.02, mu=0.95, ns_steps=5):
    """One Muon update for a 2D weight W. Only ONE state buffer (momentum)."""
    momentum_buf.mul_(mu).add_(G)               # standard heavy-ball momentum
    # Reference implementations default to Nesterov here, i.e. orthogonalize
    # (G + mu * momentum_buf) instead of the buffer itself.
    update = newton_schulz5(momentum_buf, steps=ns_steps)  # orthogonalize
    # Scale by sqrt(max(rows,cols)) so RMS of update ~ 1, matching AdamW's scale
    scale = (max(W.shape) ** 0.5)
    W.add_(update, alpha=-lr * scale)

Muon’s properties make it a compelling AdamW replacement for the bulk of an LLM’s parameters:

• Memory. Like momentum SGD and Lion, it stores one buffer per parameter (momentum), not two — half of Adam’s optimizer-state memory. There is no second-moment tensor at all.
• Only for 2-D weights. Orthogonalization is defined for matrices. The standard recipe is hybrid: use Muon for the 2-D hidden weight matrices (attention and MLP projections) and a small AdamW for the 1-D parameters and the input embedding / output head, which are not matrix-multiplied in the same sense and behave better under Adam.
• Scale matching. The \(\sqrt{\max(n,m)}\) factor makes Muon’s update RMS comparable to AdamW’s, so learning-rate intuition transfers and you can reuse much of an AdamW schedule.
• Reported gains. On small-scale benchmarks (nanoGPT speedruns) and, more recently, at MoE scale (Moonshot’s 16B/3B-active Moonlight on 5.7T tokens), Muon reaches a target loss in meaningfully fewer steps/tokens than tuned AdamW — roughly \(2\times\) compute efficiency at the compute-optimal frontier — while using less memory.

The mechanism connects cleanly to Shampoo: orthogonalizing \(M = U\Sigma V^\top\) to \(UV^\top\) is exactly applying the preconditioner \((MM^\top)^{-1/2}M\), a “whitening” of the update — the same spectral idea as Shampoo’s inverse-root preconditioning, but computed cheaply with matmuls and applied to the momentum rather than accumulated second moments. Muon can be read as a streamlined, GPU-friendly descendant of the Shampoo line.

Wiring the hybrid recipe¶

“Hybrid” means two optimizer objects over two disjoint parameter sets, both stepped every iteration. That is all there is to it:

import torch

model = torch.nn.Sequential(              # stand-in for your transformer
    torch.nn.Embedding(1000, 64), torch.nn.LayerNorm(64), torch.nn.Linear(64, 64)
)

# Muon for the 2-D hidden weights; AdamW for embeddings, LM head, norms, biases.
hidden, others = [], []
for name, p in model.named_parameters():
    is_hidden_matrix = p.ndim == 2 and not any(k in name
                                               for k in ("embed", "lm_head"))
    (hidden if is_hidden_matrix else others).append(p)

adamw = torch.optim.AdamW(others, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.0)
bufs = {p: torch.zeros_like(p) for p in hidden}   # one momentum buffer per matrix

def hybrid_step(lr_muon=0.02):
    for p in hidden:                              # muon_step defined above
        muon_step(p.data, p.grad, bufs[p], lr=lr_muon)
    adamw.step()
    adamw.zero_grad(set_to_none=True)
    for p in hidden:
        p.grad = None

One systems caveat that bites as soon as you leave a single GPU: Newton-Schulz needs the whole 2-D matrix on one device, but ZeRO/FSDP shard parameters and gradients flat across ranks. A distributed Muon therefore has to gather each matrix’s gradient (or replicate the orthogonalization) before it can iterate — that engineering, on top of ZeRO-1-style optimizer sharding, is exactly what Moonshot’s open-source Moonlight implementation contributes. Reference implementations live in KellerJordan/modded-nanogpt (single-node) and MoonshotAI/Moonlight (distributed). We use this hybrid — Muon on hidden matrices, AdamW on the rest, with a QK-clip for attention stability — to train the capstone model in Optimizer & Schedule: Muon + MuonClip and Warmup-Stable-Decay, where every constant is pinned to a real 100M-parameter run.

Choosing an Optimizer: A Practical Comparison¶

The table below summarizes the state per parameter (the memory tax) and the character of each method. “Buffers/param” counts the optimizer-specific state tensors beyond weights and gradients.

A pragmatic decision procedure for a new pretraining run:

1. Default to AdamW with \(\beta=(0.9, 0.95)\), weight decay \(0.1\) on 2-D weights only, gradient clipping at norm 1.0. It is the most documented and forgiving choice and the safe baseline for an interview answer.
2. If optimizer memory is the binding constraint and you cannot shard further, reach for 8-bit AdamW (smallest behavioral change) or Adafactor (most aggressive, more babysitting).
3. If you want speed and have appetite for tuning, try Muon for the matrices + AdamW for embeddings/head; it is the most exciting current option and halves optimizer-state memory on the bulk of parameters.
4. Always re-sweep learning rate, warmup, and weight decay when you change family — the hyperparameters do not transfer.

Further reading¶

• Kingma & Ba, Adam: A Method for Stochastic Optimization (2015) — the original Adam derivation and bias correction.
• Loshchilov & Hutter, Decoupled Weight Decay Regularization (2019) — AdamW and why decoupling matters.
• Shazeer & Stern, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost (2018) — factored second moments.
• Chen et al., Symbolic Discovery of Optimization Algorithms (2023) — the Lion optimizer.
• You et al., Large Batch Optimization for Deep Learning: Training BERT in 76 Minutes (2020) — LARS/LAMB and the layer-wise trust ratio.
• Gupta, Koren & Singer, Shampoo: Preconditioned Stochastic Tensor Optimization (2018); Anil et al., Scalable Second Order Optimization for Deep Learning (Distributed Shampoo).
• Jordan et al., Muon (2024) — orthogonalized momentum via Newton-Schulz; see also the nanoGPT speedrun writeups.
• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2020) — the optimizer-state memory analysis and sharding.
• Dettmers et al., 8-bit Optimizers via Block-wise Quantization (2022) — quantized Adam states.

Exercises¶

1. (Weight decay.) A colleague configures AdamW with a single parameter group and applies weight decay \(\lambda = 0.1\) to every parameter in a transformer, including RMSNorm gains and biases. (a) Explain why L2 regularization (adding \(\tfrac{\lambda}{2}\lVert\theta\rVert^2\) to the loss) and decoupled weight decay are not equivalent under Adam, even though they are identical under plain SGD. (b) Why is decaying the 1-D RMSNorm gains specifically harmful? What should the correct setup look like?

2. (Momentum, by hand.) Consider heavy-ball momentum \(v_t = \mu v_{t-1} + g_t\), \(\theta_t = \theta_{t-1} - \eta v_t\), driven by a constant gradient \(g\) starting from \(v_0 = 0\). (a) Derive the terminal (steady-state) velocity and evaluate it for \(\mu = 0.9\) and \(\mu = 0.98\). (b) If you raise \(\mu\) from \(0.9\) to \(0.98\) and want to keep the same steady-state step length \(\eta v_\infty\), by what factor must you change \(\eta\)? © This is why momentum and learning rate are coupled. State the coupling in one sentence.

3. (Bias correction, by hand.) The chapter notes that with \(\beta_2 = 0.999\), the uncorrected second-moment estimate at \(t=1\) is only \(0.1\%\) of the true value, so the \(v\)-side of the raw step is about \(\sqrt{1000}\approx 31\times\) too large (partly offset by the first moment’s own bias). (a) Redo this for the LLM-standard \(\beta_2 = 0.95\): compute the bias factor \((1-\beta_2^{\,t})\) at \(t=1\) and \(t=3\), and the resulting “too-large” factor \(1/\sqrt{1-\beta_2^{\,t}}\) on the step. Comment on why \(\beta_2=0.95\) needs far less protection than \(\beta_2=0.999\). (b) Now show that with full bias correction, a constant gradient \(g\) produces a bias-corrected Adam step of exactly \(\eta\,\mathrm{sign}(g)\) at every \(t\) (take \(\epsilon\to 0\)), independent of \(\beta_1,\beta_2,t\). This is the scale-invariance property.

4. (Optimizer-memory accounting.) Consider a \(30\times 10^9\)-parameter dense model trained in bf16 mixed precision with AdamW (\(m\) and \(v\) in fp32, plus an fp32 master copy of the weights). (a) Give the memory for weights, gradients, and each of the three optimizer tensors, and the total bytes/param. (b) How much memory is the optimizer state alone, and how many 80 GB H100s would that state occupy (unsharded)? © You switch the 2-D matrices to an optimizer that keeps a single momentum buffer instead of \(m\) and \(v\) (Lion or Muon). Assuming essentially all \(30\)B params are 2-D, how much optimizer-state memory do you save, and what is the new H100 count for the optimizer state?

5. (Implementation.) The chapter gives Lion as a functional lion_step. Turn it into a proper torch.optim.Optimizer subclass, mirroring the from-scratch AdamW in the chapter: lazy per-parameter state initialization, one momentum buffer per parameter, decoupled weight decay, the \(\beta_1\)-interpolated sign update, and the \(\beta_2\)-updated momentum buffer. Keep it in-place and under @torch.no_grad().

import torch
from torch.optim import Optimizer

class Lion(Optimizer):
    """From-scratch Lion (Chen et al., 2023). One state tensor per param."""
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0):
        # Lion's LR is typically ~3-10x smaller than AdamW's, decay larger.
        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for group in self.param_groups:
            lr, (b1, b2) = group["lr"], group["betas"]
            wd = group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                state = self.state[p]
                if len(state) == 0:                       # lazy init
                    state["exp_avg"] = torch.zeros_like(p)  # m (only buffer)
                m = state["exp_avg"]

                # --- Decoupled weight decay (multiplicative shrink) ---
                if wd != 0:
                    p.mul_(1.0 - lr * wd)

                # --- Update direction: sign of the beta1 interpolation ---
                c = m.mul(b1).add(g, alpha=1.0 - b1)      # temp, not stored
                p.add_(c.sign(), alpha=-lr)               # theta -= lr*sign(c)

                # --- Momentum buffer uses beta2 and the raw grad g ---
                m.mul_(b2).add_(g, alpha=1.0 - b2)
        return loss

6. (Shampoo memory and the Muon connection.) Shampoo maintains \(L\in\mathbb{R}^{n\times n}\) and \(R\in\mathbb{R}^{m\times m}\) for a weight matrix \(W\in\mathbb{R}^{n\times m}\); AdamW maintains \(m,v\in\mathbb{R}^{n\times m}\). (a) Count the optimizer-state numbers for each on a square \(4096\times4096\) matrix and on a rectangular \(4096\times1024\) matrix. Is Shampoo always cheaper than Adam? What does this tell you about why people use Shampoo? (b) The chapter says orthogonalizing the momentum \(M = U\Sigma V^\top\) to \(UV^\top\) equals applying the preconditioner \((MM^\top)^{-1/2}M\). Verify this identity and explain in one line how it connects Muon to Shampoo’s spectral idea.