Optimizers: SGD to Adam to what's next

In my posts on what a neural network actually is and backpropagation, I ended training with one sentence: compute the gradient, step the other way, repeat. That sentence hides a whole field. The gradient tells you which way is downhill. It does not tell you how far to step, whether to trust this one gradient or the last hundred, or whether every weight should move by the same amount. The rule that answers those questions is the optimizer.

Optimizers are on my list of things I’m looking into right now, partly because for almost a decade the answer was boring (use Adam) and in the last two years it stopped being boring. So this post walks the whole line, one fix at a time, and ends with what people are actually training big models with in the middle of 2026.

Gradient descent on one weight

Start with the smallest possible example. One weight w, and a loss that is lowest at w = 3:

loss(w)     = (w - 3)^2
gradient(w) = 2 * (w - 3)
update:       w = w - lr * gradient(w)

lr is the learning rate, the step size. Start at w = 0 with lr = 0.1:

step 0: w = 0.00   gradient = -6.0    -> w = 0 + 0.6    = 0.60
step 1: w = 0.60   gradient = -4.8    -> w = 0.6 + 0.48 = 1.08
step 2: w = 1.08   gradient = -3.84   -> w = 1.464

Each step closes 20 percent of the remaining gap to 3. It gets there, slowly. Now try lr = 1.0. The first step jumps from 0 to 6, the next jumps from 6 back to 0, and it bounces forever. At lr = 1.1 each bounce is bigger than the last and the weight flies off to infinity.

That one example already has the main tension of the whole field. Small steps are safe and slow. Big steps are fast until they blow up. And the “right” step size depended on the curvature of this specific loss, the 2 in 2 * (w - 3). A real network has billions of weights, and the curvature is different in every direction.

The “stochastic” in SGD

The real loss is an average over the entire training set. Computing the exact gradient means running every example through the network before taking a single step, which for a language model is trillions of tokens. Most large-model training uses minibatches instead.

Instead you grab a small random batch, compute the gradient on just that, and step. That gradient is noisy. It points roughly downhill on average, but any one batch can be off. This is stochastic gradient descent, SGD. The math behind “a noisy estimate of the right direction still gets you there if the steps shrink properly” goes back to Robbins and Monro in 1951, decades before anyone was training neural networks.¹

The noise turns out to be fine, and maybe even useful. Cheaper noisy steps can make better use of a fixed compute budget. The thing is, noise plus uneven curvature makes a specific problem much worse.

The ravine problem and momentum

Picture a long narrow valley. It is steep across and nearly flat along the bottom. The gradient mostly points across the valley, because that’s where the slope is. So gradient descent bounces between the walls and barely moves along the floor, which is the direction you actually wanted to go.

If you lower the learning rate to stop the bouncing, progress along the floor gets even slower. If you raise it, the bouncing gets worse and eventually diverges. One learning rate can’t serve both directions.

The fix is from Boris Polyak in 1964, and it’s usually called the heavy ball method, or just momentum.² Instead of stepping along the current gradient, you keep a running velocity that remembers past gradients:

v = beta * v + gradient
w = w - lr * v

With beta = 0.9, the velocity is roughly a sum of the last ten or so gradients, weighted toward the recent ones. Think about what that does in the ravine. The across-the-valley components flip sign every step (left wall, right wall, left wall), so they mostly cancel in the sum. The along-the-floor component points the same way every step, so it adds up. You get a ball that builds speed in the consistent direction and damps out the bouncing.

There’s a variant called Nesterov momentum that looks ahead before computing the gradient, and it tends to behave a little better. That detail comes back later, because Muon uses it.

SGD with momentum trained most of the famous image networks. It works really well when someone has carefully tuned the learning rate and its schedule. That “when” is the catch.

Adam: a learning rate per weight

Momentum fixes direction. It doesn’t fix scale. Some weights in a network get huge gradients and some get tiny ones. Embeddings for rare words barely ever see a gradient at all. One global learning rate is too big for some of them and too small for others at the same time.

Adam, from Diederik Kingma and Jimmy Ba (on arXiv in December 2014, published at ICLR 2015), keeps two running averages for every single weight.³ One is momentum, the average gradient. The other is the average of the squared gradient, which tracks how big that weight’s gradients usually are. Then it divides one by the square root of the other:

m = b1 * m + (1 - b1) * g        # average gradient (momentum)
v = b2 * v + (1 - b2) * g * g    # average squared gradient (scale)
m_hat = m / (1 - b1**t)          # bias correction, since m and v start at 0
v_hat = v / (1 - b2**t)
w = w - lr * m_hat / (sqrt(v_hat) + eps)

The defaults in the paper are lr = 0.001, b1 = 0.9, b2 = 0.999, eps = 1e-8, and honestly a huge number of people never changed them.

Here’s the concrete way to see what the division does. Say two weights have gradients of 100 and 0.01 on the first step. With plain SGD at lr = 0.001, the first weight moves by 0.1 and the second by 0.00001. Ten thousand times apart. With Adam, on step one m_hat equals g and sqrt(v_hat) equals |g|, so the update is basically lr * sign(g). Both weights move by about 0.001. Adam cares about which way each weight wants to go and how consistently, much more than about the raw size of its gradient.

That’s why Adam is so forgiving. You don’t have to hand-tune a learning rate for every layer, because it normalizes each weight’s step on its own. It became the default almost everywhere. The original Transformer paper trained with Adam plus a warmup schedule,⁴ which I touched on when I read that paper in 2026, and many later language models have used Adam or AdamW.

AdamW: the weight decay bug

Weight decay means shrinking every weight a little toward zero each step, so the network doesn’t lean on huge weights. In plain SGD you can get it by adding a penalty 0.5 * wd * w^2 to the loss. The gradient of that penalty is wd * w, so it shows up in the gradient and shrinks the weight. For SGD, “add a penalty to the loss” and “shrink the weights directly” are mathematically the same thing.

Ilya Loshchilov and Frank Hutter pointed out that for Adam they are not the same.⁵ If you put the penalty into the gradient, Adam divides it by sqrt(v_hat) along with everything else. So weights with big, frequent gradients get barely any decay, and weights with small gradients get a lot. That’s backwards from what you wanted, and the effective regularization ends up tangled with the learning rate.

Their fix is almost embarrassingly small. Take the decay out of the gradient and apply it directly:

w = w - lr * (m_hat / (sqrt(v_hat) + eps) + wd * w)

That’s AdamW, “decoupled weight decay.” It was posted in November 2017 and published at ICLR 2019, and it became the standard for training transformers. To give one public example, Meta’s Llama 2 paper lists AdamW with b1 = 0.9, b2 = 0.95, and weight decay 0.1.⁶ AdamW became a common starting point for language-model training.

What I find interesting about AdamW is that the whole contribution is noticing that two things everyone assumed were equal weren’t. No new machinery. Just reading the update rule carefully.

What Adam ignores

Here’s the limitation that sets up everything after Adam. Adam treats every weight as its own separate number. A layer’s weight matrix with 4096 by 4096 entries is, to Adam, 16.7 million unrelated scalars, each with its own little m and v. It never looks at the matrix as a matrix.

But a weight matrix is a linear map. It takes a vector in and sends a vector out. The gradient of that matrix also has structure. In practice the gradient for a hidden layer tends to be dominated by a few directions. If you take its singular value decomposition, a handful of singular values are big and the rest are tiny. So a plain update pushes hard along a few directions the network already “agrees” on, and barely moves along all the rest. Adam’s per-entry scaling doesn’t fix this, because the imbalance lives across rows and columns jointly, not in individual entries.

Both of the newer optimizers I want to talk about start from that observation.

Muon: orthogonalize the update

Keller Jordan first used Muon in the NanoGPT speedrun in October 2024 and wrote it up with collaborators in a December 2024 blog post.⁷ The speedrun is a public competition to train a small GPT-2-sized model to a target loss in as little time as possible. When Muon was adopted there, it improved the training speed by about 35 percent, and it stayed the optimizer behind the records that followed.

The idea, for each 2D weight matrix in a hidden layer:

1. Compute the momentum of the gradient, like SGD with Nesterov momentum.
2. Replace that momentum matrix with its closest orthogonal matrix.
3. Step along that.

“Closest orthogonal matrix” sounds abstract, so here’s the concrete version. Write the momentum matrix as G = U S V^T, where S holds the singular values. Orthogonalizing means throwing away S and keeping U V^T. Every direction the gradient points in gets a step of the same size. The few dominant directions stop drowning out the rare ones.

Computing an SVD every step on every matrix would be slow. Muon instead uses a Newton-Schulz iteration, which only needs matrix multiplies, and those are exactly what GPUs are good at. Here’s the core of it, following the reference implementation:

def newton_schulz5(G, steps=5):
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G / (G.norm() + 1e-7)       # scale so all singular values are <= 1
    flip = G.shape[0] > G.shape[1]
    if flip:
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X           # pushes every singular value toward ~1
    return X.T if flip else X

Each pass applies a fixed polynomial to every singular value at once without ever computing them. Those three odd-looking coefficients were tuned so that five passes push small singular values up fast. It doesn’t land exactly on 1. It lands in a rough band around it, and the post argues that’s good enough in practice and much cheaper than doing it exactly.

A few practical details matter. Muon is only for the 2D hidden-layer matrices. The embedding layer, the output head, and all the scalar and vector parameters like biases and norm gains still use AdamW. The post estimates the extra compute at under 1 percent of a training run.⁷ PyTorch now ships torch.optim.Muon, and its docs show exactly this pattern of splitting parameters by p.ndim == 2 and sending the rest to AdamW.⁸

Does it scale?

A speedrun on a small model is one thing. The real question was whether Muon holds up on big models, and that’s where Moonshot AI came in. Their February 2025 paper, “Muon is Scalable for LLM Training,” found two changes that mattered: add weight decay (the same lesson AdamW learned), and rescale each matrix’s update so its overall size matches what AdamW would produce.⁹ That second trick means you can reuse learning rates and weight decay already tuned for AdamW, which removes a big practical barrier. Their scaling-law experiments put Muon at about 2x the compute efficiency of AdamW, and they trained Moonlight, a mixture-of-experts model with about 16B total and 3B active parameters, on 5.7 trillion tokens with it.

Then they went much bigger. Kimi K2, released in July 2025, is a mixture-of-experts model with 1 trillion total parameters and 32 billion active, pretrained on 15.5 trillion tokens with MuonClip, Muon plus a fix called QK-clip that keeps the attention logits from blowing up.¹⁰ They report zero loss spikes over the whole run, which for a model that size is a big deal. Loss spikes are one of the main things that wreck large training runs.

It isn’t only Moonshot. Zhipu’s GLM-4.5 report (355B total, 32B active) says they used Muon for every parameter except word embeddings, biases and RMSNorm weights, with five Newton-Schulz steps and momentum 0.95.¹¹

Shampoo and SOAP: the older, heavier relative

Muon has an older cousin. Shampoo, from Vineet Gupta, Tomer Koren and Yoram Singer in 2018, also treats a weight matrix as a matrix.¹² Instead of one scale per entry like Adam, it keeps two matrices per layer: one built from G G^T (how the rows of the gradient move together) and one from G^T G (how the columns move together). The update multiplies the gradient on the left and right by inverse fourth roots of those:

L = L + G @ G.T
R = R + G.T @ G
update = L^(-1/4) @ G @ R^(-1/4)

Here’s a fact I like. If you skip the running sum and use only the current gradient, then with G = U S V^T you get L^(-1/4) = U S^(-1/2) U^T and R^(-1/4) = V S^(-1/2) V^T, and the product collapses to U V^T. That’s exactly Muon’s orthogonalized update. In this idealized, unregularized calculation, the update matches Muon’s exact polar factor. Practical implementations add momentum, numerical safeguards, and approximations. Shampoo with memory is more informed but has to store and invert those extra matrices, which is expensive.

Shampoo got a big credibility boost from AlgoPerf, a benchmark that compares training algorithms on wall-clock time across many different workloads. In the first competition, a distributed implementation of Shampoo won the external tuning track, beating Adam-style methods even when measured in real time rather than steps.¹³ That implementation comes from a team at Meta, who report at most about a 10 percent slowdown per step compared to Adam-style methods.¹⁴

SOAP (Vyas and colleagues, 2024) is a refinement.¹⁵ Its observation is that Shampoo is closely related to running a memory-efficient version of Adam in a rotated coordinate system, specifically the eigenbasis of Shampoo’s preconditioner matrices. So SOAP literally does that: rotate into that basis, run Adam there, rotate back. It only adds one hyperparameter over Adam, how often to recompute the basis. On 360M and 660M parameter language models in a large-batch setting, they report over 40 percent fewer iterations and over 35 percent less wall-clock time than AdamW.

What people actually use, mid-2026

This is the part the hype usually skips, so I tried to only say what I could find documented.

• AdamW is still the default. Most models with public training details that I’ve seen still use it, and it’s what you should reach for if you don’t have a reason not to.
• Muon is real at frontier scale. Kimi K2 at a trillion parameters and GLM-4.5 at 355B are both publicly documented as trained with Muon or a Muon variant, and it’s in PyTorch itself. This stopped being a speedrun trick.
• Shampoo and SOAP win benchmarks and have a strong research case, but I couldn’t find a publicly documented frontier-scale language model trained with either as its main optimizer. That doesn’t mean nobody does it. Labs don’t always publish their optimizer. It means I can’t point to one.

My read on why Muon spread faster than Shampoo, even though Shampoo is older and arguably more principled: Muon has no extra state beyond momentum, it’s a few lines of code, and after Moonshot’s rescaling trick it drops in with your existing AdamW settings. Engineering friction decides a lot of these things.

What I’d want to try

It’s exciting that there is still room to improve something as basic as how a model updates its weights. If a different update gets us to the same result with less training, that could change which experiments people can afford to run.

I’d like to compare these methods on a small transformer with the same training budget and enough care put into tuning each one. Watching the updates themselves might be as interesting as watching the loss. Does changing their shape help throughout training, or only at certain points? That seems like a more useful place to start than assuming the newest optimizer should replace every older one.

Sources