For about a decade, the optimizer was the part of a training run that people argued about the least. Adam came out in 2014 and AdamW fixed how it handles weight decay in 2017. From GPT-3 to Llama 3, most large models were trained with more or less the same update rule: two moving averages, a square root and a division. Over the same period, attention, normalisation and the feed-forward block were all redesigned several times, while the optimizer stayed mostly the same.
In my previous article on mHC (Manifold-Constrained Hyper-Connections), I mentioned that Adam is being challenged by Muon, and production models are already using it. DeepSeek-V4 trains most of its modules with Muon, alongside mHC, and Moonshot’s Kimi K2, a 1-trillion-parameter mixture of experts with 32B active parameters, was pre-trained on 15.5 trillion tokens with a Muon variant with no loss spikes.
Keller Jordan published Muon as a blog post in December 2024, rather than as a paper, after it had set speed records on two hobbyist benchmarks: training a CIFAR-10 classifier and training GPT-2 small on FineWeb (the NanoGPT speedrun). The idea is to take the momentum we would normally use with SGD, replace it with the nearest orthogonal matrix, and step in that direction. Muon is short for MomentUm Orthogonalized by Newton–Schulz.
To give a sense of the gains, on CIFAR-10 Muon cut the record from 3.3 to 2.6 A100-seconds, and on the NanoGPT speedrun it made training 1.35× faster. Later, Moonshot AI’s scaling-law experiments put Muon at about 2× the compute efficiency of AdamW, and they used it to train Moonlight, a 16B mixture of experts with 3B active parameters, on 5.7T tokens.
In this article, we go through Muon step by step, the same way we did in the mHC article. We start with what an optimizer does to a weight matrix and what Adam misses, then look at the Muon update rule and why an orthogonal step is a better step (this turns out to be a question about norms). After that we cover Newton–Schulz, the iteration that makes Muon cheap, then the code, and finally some limitation and consideration before using Muon.
Every training step has the same form: W_t = W_{t−1} − η · update. Here η (the Greek letter eta, pronounced "AY-tuh") is the learning rate and W is the weight matrix. Optimizers only differ in how they turn the gradient G into that update. One of the oldest options that still works well is SGD with momentum, which keeps a running sum of past gradients and steps along it.
Adam adds a second running average, this time of the squared gradient, and divides by its square root.
Notice the lower-case letters in Equation 2. Adam works on one number at a time. For example, a 768 × 768 attention projection (Weight Matrix) in GPT-2 small is just 589,824 separate scalars to Adam. Each one gets its own step size, and none of them knows which row or column it belongs to.
But a weight matrix is really a linear map: a 768-dimensional vector goes in and another one comes out. What an update does to that map depends on the matrix as a whole, and the tool we use to describe that is the singular value decomposition (SVD).
Keller Jordan’s starting point was an empirical observation. For the 2D weights of a transformer, the updates produced by SGD-momentum and Adam have a very high condition number, which means a few singular values dominate and the rest are tiny.
For example, take five directions of a momentum matrix with singular values 0.9, 0.3, 0.1, 0.01 and 0.001. A step along M moves the weights 900 times further along the first direction than along the last. So the update keeps pushing the same few directions, and a direction that is rare but useful hardly moves at all.
Muon is Equation 1 with one extra step between the momentum and the weight update.
If we compare this with Equation 3, the SVD writes the momentum as U Σ Vᵀ. Muon keeps U and V, which hold the directions, and drops Σ, which holds the magnitudes, so every singular value becomes 1. The result, U Vᵀ, is the orthogonal matrix closest to the momentum. It points the same way as the momentum, but it moves every direction by the same amount.
With exact orthogonalization, the same five directions go in at 0.9, 0.3, 0.1, 0.01 and 0.001 and all come out at 1. The dominant direction gets cut back a little, while the rare direction at 0.001 gets scaled up about a thousand times. So the update now moves every direction of the matrix, where before most of it went into two or three. (Muon’s five Newton–Schulz steps land close to 1 rather than exactly on it, as we’ll see in the Newton–Schulz section.)
Adam also evens out magnitudes, but it does so per coordinate, in whatever basis the weights happen to be stored in. Muon does it per direction, in the matrix’s own basis. That is the main difference between the two, and in the next section we’ll see why the matrix’s own basis is the better choice.
A few details from the reference implementation complete the rule:
💡 Recap Note: Momentum is an exponential weighted average
Muon’s first step is ordinary momentum, so it helps to see what momentum computes. An exponential weighted average (the same smoothing idea used in time-series forecasting) blends the running average with the newest value:
With β = 0.95, this becomes V_t = 0.95·V_{t−1} + 0.05·a_t. Most of the weight stays on the history, so noisy mini-batch gradients get smoothed out and the optimizer is less likely to zigzag or stall. If we replace a_t with the gradient, and call the average M_t and the coefficient μ as in Equation 1, we get SGD with momentum:
Muon keeps this momentum buffer but changes what happens next. SGD applies the smoothed gradient directly, and Adam rescales each parameter by its own gradient history. Muon instead orthogonalizes the whole momentum matrix, as in Equation 4, which pushes all singular values toward 1 so that a few dominant directions can’t control the update.
Keller’s write-up (Equation 1) drops the (1 − μ) factor, while the reference code keeps it. Either works, because orthogonalization ignores the overall scale of M_t. Orthogonalization needs a matrix, so Muon is applied only to 2D hidden-layer weights. Biases, embeddings and output heads are usually still trained with AdamW.
Dropping the singular values might look like losing information, but it makes more sense once we write down what an optimizer step is trying to do. Every first-order step answers the same question: which small change to W lowers the loss the most? To define "small" we need to pick a norm, and different norms give different answers.
Bernstein and Newhouse, in Old Optimizer, New Norm, show that familiar optimizers are this one problem with different norms, once their moving averages are switched off.
Why does the last row of Table 1 come out as U Vᵀ? If we write the gradient with its SVD, the loss change splits into one term per direction: ⟨G, ΔW⟩ = Σᵢ σᵢ · uᵢᵀ ΔW vᵢ. Each term is the gradient's strength in that direction, σᵢ, times how far the step moves along it, uᵢᵀ ΔW vᵢ. A spectral-norm budget of η caps that move at η for every direction at the same time. The budget isn't shared between directions, so using all of it on one direction still leaves the full budget for the others. The best step therefore moves every direction by the full η against the gradient, and that step is ΔW = −η U Vᵀ. Strong and weak directions get the same step, whatever their σᵢ.
The Frobenius norm in the first row works differently. There the budget is shared across directions, so the best step spends more of it where σᵢ is large and keeps the gradient’s uneven shape. That gives normalised SGD, and it’s the reason the rare directions get so little of the update.
The spectral norm is a better budget for a hidden layer because of what the layer does. It takes an activation vector x and returns W x, so after a step the output changes by ΔW x. The spectral norm ‖ΔW‖ is the largest factor by which ΔW can stretch any input, which gives us the bound ‖ΔW x‖ ≤ ‖ΔW‖ · ‖x‖. By capping it, we make sure no token's output moves by more than a fixed amount in one step, whatever its activation looks like.
Adam’s ℓ_∞ norm caps each entry instead, and that tells us little about how far an output vector can move. For example, if every entry of a 768 × 768 step sits at the cap η and the signs line up, as they do when one direction dominates the gradient, the step can stretch an input by up to 768η.
This is also why the rare directions matter. Under the spectral budget, a direction with σ = 0.001 costs as much of the step as a direction with σ = 0.9, so the optimizer has no reason to neglect it. This matches Keller Jordan’s original intuition that orthogonalization boosts the rare directions, which are small in the update but still important for learning.
The direct way to get U Vᵀ is to run an SVD on every weight matrix at every step. An SVD is iterative, runs slowly on GPUs and doesn't work well in bfloat16, so Muon doesn't compute it. Instead, it uses a property of odd matrix polynomials.
If X = U S Vᵀ, then XXᵀX = U S³ Vᵀ, and the whole right-hand side becomes U φ(S) Vᵀ with φ(σ) = aσ + bσ³ + cσ⁵. The directions U and V pass through unchanged, and each singular value simply goes through a scalar polynomial. If we apply the step five times, each σ goes through φ five times, and we only ever use matrix multiplications, which is exactly what GPUs are good at.
Before the first step, Muon divides the momentum by its Frobenius norm. The Frobenius norm is never smaller than the largest singular value, so after this every σ sits in (0, 1], which is the range φ is designed for.
The speed comes from the choice of coefficients. The textbook Newton–Schulz polynomial, 1.5σ − 0.5σ³, converges to exactly 1, but slowly from below: a direction at 0.01 only reaches 0.076 after five steps. Keller Jordan tuned (a, b, c) to make the slope at zero as large as possible (3.4445 per step), while still keeping every value within roughly 0.7 to 1.3. As a result, small values grow quickly. The trade-off is that φ no longer settles exactly at 1, so the values overshoot and oscillate inside a band.
For the five example directions, the condition number drops from 900 to about 2.3. Every direction that starts between 0.01 and 1 ends up between 0.68 and 1.13, whereas the textbook polynomial would leave the two rare directions almost where they started. Keller reports that an error of up to about 0.3 around 1 doesn’t hurt the loss curve, and the speedrun results support that.
DeepSeek-V4 tightens this band. It runs ten steps in two stages: eight with Keller’s coefficients for fast growth, then two with (2, −1.5, 0.5) to settle the singular values precisely at 1.
The cost is small. Each step is three matrix multiplications in bfloat16, and Keller bounds Muon’s FLOP overhead at T·m / B, where T is the number of steps, m the model width and B the batch size in tokens. For the NanoGPT speedrun (m = 768, B = 524,288) that's 0.7%, and for Llama 405B (m = 16,384, B = 16M) it's 0.5%.
Since a matrix acts on space, the easiest way to picture a Muon step is geometrically. Any 2 × 2 matrix maps the unit circle to an ellipse. The ellipse’s semi-axes are the matrix’s singular values, and they point along the columns of U. Let's take a 2 × 2 momentum with singular values 0.9 and 0.1, which is small enough to draw, and pass it through the same normalisation and five Newton–Schulz steps.
The momentum squeezes the circle into an ellipse nine times longer than it is wide, so a step along M moves one direction and barely touches the other. After Newton–Schulz, the two stretches are 0.70 and 0.75. The arrows still point in the same directions as before. Only their lengths have changed, and the rare direction has gone from 0.1 to 0.75.
We can’t draw an ellipse at full size, but we can show the matrix itself. Figure 4 builds a 5 × 5 momentum with the same five singular values (0.9, 0.3, 0.1, 0.01 and 0.001) and follows it all the way to the weight update.
In M, the middle column holds 75% of the squared Frobenius norm and the largest entry in every row, so a step along M would mostly change how the layer reads one input coordinate. In O, the mass is spread out: the largest column holds 38% of the squared norm, and every column carries some of it. The spectrum below the heatmaps shows the same thing. The five singular values go from 0.9, 0.3, 0.1, 0.01 and 0.001 to values between 0.49 and 1.13, and the condition number drops from 900 to 2.3.
These values are a little different from the ones in Figure 2 because Muon first divides M by its Frobenius norm, 0.954, so the five directions enter Newton–Schulz at 0.94, 0.31, 0.10, 0.010 and 0.001. The weights then take the step W ← W − η · O, which moves every direction of W by between 0.49η and 1.13η.
Muon only optimizes hidden layers, the 2D matrices that map one activation vector to another. Everything else stays on AdamW, so in practice a Muon run always uses two optimizers.
There’s a theoretical reason for the embedding rule. The spectral-norm argument above assumes a layer that multiplies dense activation vectors, but an embedding only ever receives one-hot inputs. So what matters there is how far each token’s row moves, which is a per-row limit rather than a spectral one. The output head is the one case the theory doesn’t explain, and Keller reports that AdamW simply works better there.
In practice, the split is just a filter on p.ndim and on module names. DeepSeek-V4 follows the same pattern: Muon for most modules, and AdamW for the embedding, the prediction head, all RMSNorm weights, and the static biases and gating factors of the mHC modules from the previous article. Those are vectors and scalars, so there's no matrix for Muon to work with.
The whole optimizer fits on one screen**.** Our version follows the reference implementation, trimmed down to a single GPU so that every line maps back to an equation. We start with newton_schulz5, which is Equation 6 in a loop: it casts the update to bfloat16, divides it by its Frobenius norm, and runs the five steps on the wide side of the matrix.
import torchdef newton_schulz5(G, steps=5, eps=1e-7): """Approximate U V^T for G = U S V^T using only matrix multiplications.""" a, b, c = 3.4445, -4.7750, 2.0315 X = G.bfloat16() tall = X.size(0) > X.size(1) if tall: # work on the wide side, so X @ X.T is the small Gram matrix X = X.T X = X / (X.norm() + eps) # Frobenius norm >= largest singular value, so every sigma <= 1 for _ in range(steps): A = X @ X.T # rows x rows B = b * A + c * A @ A X = a * X + B @ X # aX + b(XX^T)X + c(XX^T)^2 X if tall: X = X.T return X.to(G.dtype)
The transpose is there for cost. For a 3072 × 768 MLP matrix, X @ X.T on the tall side would be 3072 × 3072, while on the wide side it's 768 × 768, sixteen times fewer entries. Grouping the polynomial as B @ X keeps both Gram products, A and A @ A, at that smaller size.
The Muon class wraps newton_schulz5 with momentum, shape scaling and decoupled weight decay, in the same order as Figure 1. Each parameter gets one momentum_buffer, which is the only optimizer state Muon keeps. Since p.grad is read as one whole matrix, this version assumes every parameter lives on a single device.
class Muon(torch.optim.Optimizer): """Muon for 2D hidden weights only. Embeddings, head, gains and biases go to AdamW.""" def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.0): defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, weight_decay=weight_decay) super().__init__(params, defaults) @torch.no_grad() def step(self): for group in self.param_groups: mu, lr = group["momentum"], group["lr"] for p in group["params"]: if p.grad is None: continue g = p.grad state = self.state[p] if "momentum_buffer" not in state: state["momentum_buffer"] = torch.zeros_like(g) buf = state["momentum_buffer"] buf.lerp_(g, 1 - mu) # M = mu*M + (1-mu)*G update = g.lerp(buf, mu) if group["nesterov"] else buf O = newton_schulz5(update, steps=group["ns_steps"]) O *= max(1, p.size(0) / p.size(1)) ** 0.5 # shape scaling p.mul_(1 - lr * group["weight_decay"]) # decoupled weight decay p.add_(O, alpha=-lr) # W = W - lr * O
As in the recap note, the momentum here is an exponential average, μM + (1−μ)G, while Equation 1 uses a plain sum. This doesn't change the update, because Newton–Schulz divides by the Frobenius norm first and any constant factor disappears. The Nesterov line, g.lerp(buf, mu), is the blend of fresh gradient and momentum from the update rule.
The last piece is the parameter split from Table 2, with two optimizers stepped together. is_hidden sends every 2D parameter to Muon unless its name contains embed or lm_head, and everything else, including all 1D gains and biases, goes to AdamW. Parameter names differ between model implementations, so it's worth printing the names from model.named_parameters() once and checking which group each embedding lands in before you train.
is_hidden = lambda n, p: p.ndim >= 2 and "embed" not in n and "lm_head" not in nhidden = [p for n, p in model.named_parameters() if is_hidden(n, p)]others = [p for n, p in model.named_parameters() if not is_hidden(n, p)]opt_muon = Muon(hidden, lr=0.02, momentum=0.95)opt_adamw = torch.optim.AdamW(others, lr=3e-4, betas=(0.9, 0.95))loss.backward()opt_muon.step(); opt_adamw.step()opt_muon.zero_grad(); opt_adamw.zero_grad()
For example, if we run newton_schulz5 on a 768 × 768 matrix with singular values 0.9, 0.3, 0.1, 0.01 and 0.001 (and the other 763 at about 0.011), the bfloat16 result lands within 0.005 of the scalar polynomial in Figure 2. A random 3072 × 768 matrix comes out with every singular value between 0.68 and 1.15, close to the band we saw in the Newton–Schulz section.
Most of Muon’s results come from pre-training transformers from scratch, and that is where it works best, since a 2× saving in compute is worth the extra engineering there. Outside that setting, there are some trade-offs worth knowing about.
If you’re pre-training a transformer from scratch, Muon is worth the extra engineering. It treats each hidden weight as a linear map instead of a list of numbers and takes the steepest step under the spectral norm, U Vᵀ. Five Newton–Schulz steps keep the cost of that step under 1% of the training FLOPs at typical LLM batch sizes. A good place to start is the split in Table 2, checked against your model's parameter names. For long runs, add the weight decay and update-RMS and keep an eye on the attention logits, since those two fixes and QK-Clip are what made Muon stable at a trillion parameters. If you're fine-tuning a checkpoint that was pre-trained with AdamW, stick with AdamW.
The idea behind Muon is close to the one behind mHC. mHC keeps the residual mixing matrix on the doubly stochastic manifold so its gain can’t build up across layers. Muon keeps the update close to an orthogonal matrix so no direction can dominate a step and no output can move too far. Both methods constrain a matrix that would otherwise be free in order to make training more stable, and DeepSeek-V4 uses both. Fine-tuning is still an open question, and until Muon-pretrained checkpoints become common, most of us will meet Muon at pre-training scale first.
Muon: The Orthogonalized Optimizer Explained Through Equations, Intuition, Code and Results was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.