Optimizers: SGD to AdamW

Problem

Gradient descent says: move every parameter a small step against its gradient, θ ← θ − η∇L(θ). Two things make that harder than it sounds. First, the true gradient is an average over the whole training set, and computing it for every step would be far too slow, so every step uses a noisy estimate from a small batch. Second, the loss surface is badly shaped. In some directions it is steep, in others nearly flat, and one learning rate η has to serve all of them: small enough not to blow up in the steep directions, which makes it crawl in the flat ones.

An optimizer is a rule for turning the stream of noisy gradients into steps that make fast, stable progress anyway. The rules on this page differ in what they remember between steps, and that memory (the optimizer's state) costs real memory: for Adam, two extra numbers per parameter.

Intuition

Picture a long, narrow valley: gentle along its floor, steep across it. Plain gradient descent points mostly across the valley, because that's where the slope is, so it bounces from wall to wall and creeps along the floor. Raise the learning rate to creep faster and the bouncing grows until it diverges.

• Momentum keeps a running sum of past gradients, like a heavy ball. The across-the-valley components alternate in sign and cancel; the along-the-floor components agree and add up, so the ball builds speed where it should.
• RMSProp divides each parameter's gradient by a running average of its own size. The steep direction, with big gradients, gets a small step; the flat direction gets a big one. Every parameter gets its own effective learning rate.
• Adam does both: a running average of the gradient (momentum) divided by a running average of its square (RMSProp), plus a correction for the averages starting at zero. Its steps are about η in size whatever the gradient's scale.
• AdamW is Adam with weight decay applied directly to the weights rather than folded into the gradient, where Adam's division would distort it.

Each dot is one step. The y direction is 50× steeper, so SGD's learning rate is capped by it and the x direction barely moves. Momentum and Adam both fix that, in different ways. Adam's steps are all about η = 0.3 long, which is too big for the last stretch, so it overshoots and circles the minimum instead of settling.

Mechanics

SGD and minibatches

The training loss is an average over N examples, L(θ) = (1/N) Σ ℓ_i(θ). Stochastic gradient descent replaces its gradient with the average over a random minibatch B of B examples:

g = (1/B) Σ_i∈B ∇ℓ_i(θ),   θ ← θ − η g

g is an unbiased estimate of the full gradient: right on average, with noise whose standard deviation falls like 1/√B. Quadrupling the batch halves the noise but costs four times the compute per step, which is why moderate batches win on total compute. The noise isn't purely harmful either; it helps escape saddle points and is thought to favour flatter minima that generalize better. Batch size and learning rate interact: bigger batches allow bigger η, roughly in proportion for SGD, up to a limit.

On a quadratic, the step along a direction with curvature a multiplies the error by (1 − ηa). That is stable only while |1 − ηa| < 1, that is η < 2/a, so the steepest direction sets the largest usable learning rate, and the flattest direction then sets how slowly everything converges. In the valley above, a = 50 across it caps η at 0.04; at η = 0.036 the x error shrinks by only 0.964 per step. The ratio of the largest curvature to the smallest, the condition number, is the core difficulty every rule below attacks.

Momentum and Nesterov

Momentum keeps a velocity buffer b and steps along it:

b ← μb + g,   θ ← θ − η b

With μ = 0.9, a gradient that keeps the same sign is added up to an effective 1/(1 − μ) = 10 times; one that alternates mostly cancels. That is exactly the valley's problem: the floor direction is consistent, the wall-to-wall direction alternates. PyTorch's SGD stores b as momentum_buffer, initializes it to the first gradient, and (with the default dampening=0) uses exactly this form; some texts write b ← μb + (1 − μ)g instead, which only rescales η.

Nesterov momentum steps along g + μb instead of b: it applies the momentum "look-ahead" to the current gradient as well, which reacts sooner when the velocity starts carrying the parameters uphill. It is a small change with a better worst-case theory for convex problems. Note what it does to the step on the fresh gradient: η(1 + μ) instead of η, so it needs a smaller learning rate. In the race, Nesterov at η = 0.03 diverges, while at η = 0.02 it reaches 0.0943.

RMSProp

RMSProp keeps a running average of each parameter's squared gradient and divides by its root:

s ← αs + (1 − α)g²,   θ ← θ − η g / (√s + ε)

Everything is element-wise, so each parameter gets its own scale. If a parameter's gradients are consistently around size 50, √s ≈ 50 and the step is about η; if they're around 1, the step is also about η. That flattens the valley's condition number for any problem whose steep directions line up with the parameter axes. PyTorch's RMSprop defaults to α = 0.99, ε = 10⁻⁸ and stores s as square_avg. Its ancestor AdaGrad summed g² over all time instead of averaging, so its steps shrank forever; the moving average is what lets RMSProp keep learning.

Adam: moment estimates and bias correction

Adam keeps both averages, called the first and second moment estimates:

m ← β₁m + (1 − β₁)g,   v ← β₂v + (1 − β₂)g²

m̂ = m / (1 − β₁^t),   v̂ = v / (1 − β₂^t),   θ ← θ − η m̂ / (√v̂ + ε)

with defaults η = 10⁻³, β₁ = 0.9, β₂ = 0.999, ε = 10⁻⁸. m is momentum and v is RMSProp's average, so the step m̂/√v̂ is roughly "the gradient's sign, smoothed", and its size is at most about η per parameter.

The hats are the bias correction. Both averages start at 0, so early on they are too small: after one step m = (1 − β₁)g = 0.1g and v = 0.001g². Dividing by 1 − β^t rescales them to unbiased estimates of the gradient's mean and mean square. Without the correction, the first step would be η · 0.1g / √(0.001g²) = 3.162η, and because v's average is much slower than m's, the overshoot gets worse before it gets better, peaking near t = 12.

For a constant gradient, the corrected step is exactly η from step 1. The uncorrected step starts 3.16× too large, reaches 6.57× at t = 12, and is still 1.26× at t = 1000. Early steps are exactly when the parameters are worst, which is why learning-rate warmup is often used on top.

Because the step is normalized, Adam is far less sensitive to the gradient's scale than SGD, and one learning rate works across layers whose gradients differ by orders of magnitude. That, plus robustness to noisy and sparse gradients, is why it became the default for transformers. It is not a free win: on some problems, especially image classification with well-tuned schedules, SGD with momentum still generalizes slightly better.

AdamW: decoupled weight decay

Weight decay shrinks every weight toward zero each step. For SGD, adding an L2 penalty (λ/2)‖θ‖² to the loss does exactly that: the gradient gains a term λθ, and the step −η(g + λθ) is the usual step plus θ ← (1 − ηλ)θ. The two ideas coincide, and SGD(weight_decay=λ) implements decay by adding λθ to the gradient.

Adam breaks the equivalence. Adam(weight_decay=λ) also adds λθ to the gradient, and then divides by √v̂. A weight whose gradients are large gets its decay divided down to almost nothing; a weight with tiny gradients gets its decay amplified. Loshchilov and Hutter's AdamW decouples the two: it shrinks the weight directly, θ ← (1 − ηλ)θ, and then takes an ordinary Adam step on the unmodified gradient. In PyTorch that is the single line param.mul_(1 - lr * weight_decay), and AdamW is Adam with decoupled_weight_decay=True and a default λ of 10⁻² instead of 0.

With L2 inside Adam, the decay term λθ is divided by the gradient's running size. When the gradients are ±0.01, that division makes the decay act about 100 times harder and the weight is gone (0.2295 after 100 steps). AdamW's two curves lie on top of each other: 0.3614 after 1000 steps either way, matching (1 − ηλ)¹⁰⁰⁰ = 0.3677.

Muon: orthogonalized momentum for weight matrices

Every optimizer so far treats a weight matrix as a bag of independent numbers. Muon treats it as a matrix. Its gradient G (or the momentum built from it) can have a few large singular values and many tiny ones, and then a plain step moves the layer a lot along a handful of directions and barely at all along the rest. Muon replaces the update with its nearest semi-orthogonal matrix: write the momentum as M = USV^T and step along UV^T, which has the same singular directions but every singular value set to 1. Every direction the gradient points in gets an equal share of the step.

An exact SVD per step would be too slow, so Muon runs a few Newton–Schulz iterations instead. Divide M by its Frobenius norm, which puts every singular value in [0, 1], then repeat X ← aX + (bXX^T + c(XX^T)²)X. Each iteration only multiplies matrices, and it acts on each singular value s independently as the odd polynomial φ(s) = as + bs³ + cs⁵. PyTorch's coefficients, (a, b, c) = (3.4445, -4.775, 2.0315), are chosen to make the slope at zero, a, as large as possible, so small singular values grow quickly, at the cost of exactness: φ(1) = 0.701, so the iteration never settles at exactly 1. The pinned docstring says the result is US′V^T with S′ roughly uniform on (0.5, 1.5), and that this does not hurt model performance relative to UV^T. The matrix runs in bfloat16, and a tall matrix is transposed first so XX^T is the smaller Gram matrix.

The example gradient [3, 1, 0, 0.5], [1, 2, 1, 0], [0, 0.5, 0.2, 0.1] has a condition number of 23.43. Five steps with PyTorch's coefficients bring its singular values to [1.121, 0.9991, 0.6878]: not exactly 1, which is why the result still differs from the exact UV^T by 0.3348 in Frobenius norm, but no direction is left out. Computed in float64 by muon_newton_schulz, a stdlib copy of _zeropower_via_newtonschulz.

The rest of torch.optim.Muon is short. Momentum is one buffer per matrix, updated in place as buf.lerp_(grad, 1 - momentum): an exponential average, even though the docstring's formula writes B_t = μB_t−1 + g_t. With nesterov=True (the default) the orthogonalized matrix is grad.lerp(buf, momentum). Weight decay is decoupled as in AdamW, with the unadjusted rate, and then the step is scaled by a shape-dependent factor, because an orthogonal matrix's entries shrink as the matrix grows. For an A×B weight:

"match_rms_adamw" is the rule from Moonshot's Muon is Scalable for LLM Training: it scales the update so its entries have about the same RMS (0.2) as a typical AdamW step, which lets Muon reuse AdamW's learning rate and weight decay. Muon only accepts 2-D parameters and raises on anything else; the docstring says biases, embeddings and other non-2-D parameters should go to a standard optimizer such as AdamW, and its example builds the two side by side. Its state is one momentum buffer, half of Adam's.

At scale, Muon has one known failure. Kimi K2's report says that training a 53B-parameter MoE with plain Muon let the largest attention logit, the maximum of q·k/√d over a batch, climb past 1000 in some heads, which risks loss spikes. Their fix, QK-Clip, watches each head's maximum logit S_max and, whenever it passes a threshold τ, rescales that head's query and key weights after the optimizer step by √γ each, with γ = min(1, τ/S_max), so the product q·k shrinks by exactly γ. Heads under τ are untouched, and nothing in the forward pass changes. With τ = 100:

Muon with QK-Clip, plus weight decay and the RMS-matching scale, is what the report calls MuonClip; K2 used τ = 100 and pre-trained a 1T-parameter MoE on 15.5T tokens with no loss spikes. For multi-head latent attention the report scales only the per-head components and leaves the shared rotary key alone. PyTorch's Muon does not include QK-Clip at v2.15.0, since the clip needs the attention logits, which the optimizer never sees.

Optimizer state memory

Every bit of memory an optimizer keeps is per parameter. For a model with 7 billion parameters in float32:

The weights themselves are another 28 GB and the gradients another 28 GB, so Adam training needs 16 bytes per parameter, 112 GB at 7B, before a single activation is stored. Mixed-precision training usually keeps the same total: bfloat16 weights and gradients (4 bytes) plus a float32 master copy and Adam's two moments (12 bytes). This is why optimizer state is the first thing sharded across GPUs in large-model training (ZeRO and FSDP), and why memory-lean variants exist: Adafactor factorizes v into row and column averages, and 8-bit Adam quantizes both moments. PyTorch also keeps a tiny step counter per parameter.

In PyTorch

Every optimizer subclasses torch.optim.Optimizer. It is built from an iterable of parameters (or of parameter groups, dicts that each carry their own lr, weight_decay and so on), keeps self.param_groups and a self.state dict keyed by parameter, and exposes three calls a training loop uses:

• zero_grad() clears every parameter's .grad, by default by setting it to None, because backward accumulates;
• step() reads each .grad, updates the state and changes the parameters in place, under torch.no_grad();
• state_dict() / load_state_dict() save and restore the state, so a resumed run continues with the same moments rather than restarting Adam's warm-up.

Adam.step gathers the tensors of each group and calls a functional adam, which dispatches to one of three implementations: _single_tensor_adam, a plain loop over parameters that is easiest to read; a foreach version that batches the same operations over lists of tensors; and a fused kernel that does the whole update in one GPU launch. The default on GPU is foreach. They compute the same update.

Worked example

The race in numbers

All six runs on the valley above, for 30 steps from loss 33:

Implementation

import torch

theta = torch.tensor([0.5, -0.3], requires_grad=True)
opt = torch.optim.Adam([theta], lr=1e-3, betas=(0.9, 0.999), eps=1e-8)
for g in ([0.02, -4.0], [0.01, -3.0], [0.03, -5.0]):
    opt.zero_grad()
    theta.grad = torch.tensor(g)        # stand-in for loss.backward()
    opt.step()
theta                                   # tensor([ 0.4971, -0.2970])
opt.state[theta].keys()                 # step, exp_avg, exp_avg_sq

# the usual setup: AdamW, with no decay on biases and normalization weights
decay = [p for n, p in model.named_parameters() if p.ndim >= 2]
no_decay = [p for n, p in model.named_parameters() if p.ndim < 2]
opt = torch.optim.AdamW([
    {"params": decay, "weight_decay": 0.1},
    {"params": no_decay, "weight_decay": 0.0},
], lr=3e-4, betas=(0.9, 0.95))

for x, y in loader:
    opt.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    opt.step()

torch.save(opt.state_dict(), "opt.pt")  # m and v are part of the training state

Every rule on this page, in plain Python:

import math

def step(kind, params, grads, state, lr, mu=0.9, alpha=0.99,
         betas=(0.9, 0.999), eps=1e-8, wd=0.0):
    t = state["t"] = state.get("t", 0) + 1
    if kind == "adamw":
        params = [p * (1 - lr * wd) for p in params]
    else:
        grads = [g + wd * p for g, p in zip(grads, params)]
    if kind == "sgd":
        return [p - lr * g for p, g in zip(params, grads)]
    if kind in ("momentum", "nesterov"):
        buf = state["buf"] = grads if t == 1 else [mu * b + g for b, g in zip(state["buf"], grads)]
        d = buf if kind == "momentum" else [g + mu * b for g, b in zip(grads, buf)]
        return [p - lr * x for p, x in zip(params, d)]
    if kind == "rmsprop":
        s = state["s"] = [alpha * s + (1 - alpha) * g * g
                          for s, g in zip(state.get("s", [0.0] * len(grads)), grads)]
        return [p - lr * g / (math.sqrt(v) + eps) for p, g, v in zip(params, grads, s)]
    b1, b2 = betas
    m = state["m"] = [b1 * m + (1 - b1) * g for m, g in zip(state.get("m", [0.0] * len(grads)), grads)]
    v = state["v"] = [b2 * v + (1 - b2) * g * g for v, g in zip(state.get("v", [0.0] * len(grads)), grads)]
    return [p - lr * (mi / (1 - b1**t)) / (math.sqrt(vi / (1 - b2**t)) + eps)
            for p, mi, vi in zip(params, m, v)]

theta, state = [0.5, -0.3], {}
for g in ([0.02, -4.0], [0.01, -3.0], [0.03, -5.0]):
    theta = step("adam", theta, g, state, lr=1e-3)
theta                                   # [0.497125..., -0.297029...]

Tradeoffs

• Adaptive versus plain. Adam and AdamW need little tuning and handle badly scaled and sparse gradients; they're the default for transformers. SGD with momentum uses half the state memory and, with a carefully tuned schedule, often generalizes a little better on vision tasks.
• Memory. Adam doubles the parameter memory again, 8 extra bytes per float32 parameter; at billions of parameters that drives sharding, 8-bit states or factored variants like Adafactor.
• Hyperparameters still matter. The learning rate dominates everything else. For large models β₂ = 0.95 is common, because 0.999 averages over about a thousand steps and reacts too slowly when gradients spike; ε matters in low precision, where 10⁻⁸ underflows in float16.
• L2 or decoupled decay. With SGD they're the same. With Adam use AdamW, whose decay strength doesn't depend on the gradient scale, and usually exclude biases and normalization weights from decay.
• Implementation. foreach and fused versions compute the same update with far fewer kernel launches; the single-tensor loop exists for clarity and for unusual tensor types.

Check yourself

Connections

References

• Robbins and Monro. A Stochastic Approximation Method. Annals of Mathematical Statistics, 1951.
• Polyak. Some Methods of Speeding Up the Convergence of Iteration Methods. USSR Computational Mathematics and Mathematical Physics, 1964 (momentum).
• Nesterov. A Method for Solving the Convex Programming Problem with Convergence Rate O(1/k²). Soviet Mathematics Doklady, 1983.
• Duchi, Hazan and Singer. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. JMLR, 2011 (AdaGrad).
• Tieleman and Hinton. Lecture 6.5, RMSProp. Coursera: Neural Networks for Machine Learning, 2012.
• Sutskever, Martens, Dahl and Hinton. On the Importance of Initialization and Momentum in Deep Learning. ICML, 2013.
• Kingma and Ba. Adam: A Method for Stochastic Optimization. ICLR, 2015 (arXiv 1412.6980).
• Loshchilov and Hutter. Decoupled Weight Decay Regularization. ICLR, 2019 (arXiv 1711.05101).
• Reddi, Kale and Kumar. On the Convergence of Adam and Beyond. ICLR, 2018 (arXiv 1904.09237; AMSGrad).
• Rajbhandari, Rasley, Ruwase and He. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC, 2020 (arXiv 1910.02054).
• Keller Jordan. Muon: An Optimizer for Hidden Layers in Neural Networks. Blog post, 2024 (kellerjordan.github.io/posts/muon).
• Moonshot AI. Muon is Scalable for LLM Training. 2025 (arXiv 2502.16982; the match_rms_adamw scale).
• Jeremy Bernstein. Deriving Muon. 2025 (jeremybernste.in/writing/deriving-muon).
• Kimi Team. Kimi K2: Open Agentic Intelligence. 2025 (arXiv 2507.20534; QK-Clip and MuonClip).
• PyTorch v2.15.0: torch/optim/sgd.py, torch/optim/adam.py, torch/optim/adamw.py, torch/optim/_muon.py and torch/optim/optimizer.py.