Hello, I'm Shrijith Venkatramana, and I'm building LiveReview — a blast-radius aware AI code review built for your business-critical systems. Star us to help devs discover the project, give it a try, and share your feedback to help improve the product.
Most LLM developers know the AdamW update by heart:
take the gradient, keep moving averages, normalize the update, change the weights.
But there is a question hiding underneath all of this:
What exactly is a 4096 x 4096 weight matrix?
AdamW mostly treats it as 16 million scalar coordinates.
Muon treats it as a matrix.
That distinction is the whole story.
Muon is an optimizer for the hidden matrix parameters of neural networks. Its core operation takes the momentum update, looks at its singular directions, throws away the singular-value magnitudes, and keeps the directions. It then approximates this operation cheaply with a few Newton-Schulz iterations.
The result is an optimizer that has produced faster training in small-model competitions, scaled to multi-billion-parameter language models, and is now part of the mainstream PyTorch optimization stack.
The interesting part is not merely that "Muon beats AdamW."
The interesting part is why someone thought to optimize a neural network matrix as a matrix in the first place.
1. The story starts with a matrix
Consider a transformer linear layer:
y = W x
where W might be a 4096 x 4096 matrix.
During backpropagation we obtain a gradient:
G = dL/dW
AdamW maintains statistics for each scalar element of G.
Very roughly:
m_t = beta1 * m_(t-1) + (1-beta1) * G_t
v_t = beta2 * v_(t-1) + (1-beta2) * G_t^2
update = m_t / sqrt(v_t)
The important detail is G_t^2: the second-moment estimate is elementwise.
This gives Adam a coordinate-wise view of the parameter space.
Muon starts from a different observation:
W is not merely a bag of numbers. It represents a linear transformation.
For a matrix, one of the natural ways to understand that transformation is through its singular value decomposition:
G = U Sigma V^T
Here:
U and V = directions
Sigma = magnitudes along those directions
So there are two different questions:
- Which directions does the gradient want to move?
- How large should each of those directions be?
Muon makes a very particular choice:
preserve the directions, but approximately equalize their magnitudes.
That is what orthogonalizing the update means.
The idea appeared publicly in October 2024, when Keller Jordan and collaborators were working on the NanoGPT speedrunning competition. On October 15, 2024, a Muon-based run set a new training-speed record, improving the previous result by about 35%. The project then became a collaboration involving people such as Jeremy Bernstein, Laker Newhouse, Yuchen Jin, Vlado Boza, Jiacheng You, and Franz Cesista. Jordan's account is unusually concrete about the engineering: Boza found that treating Q, K, and V separately worked better; Jin pushed experiments to larger models and supplied much of the H100 compute; Bernstein, You, and Cesista reduced the cost of the matrix orthogonalization itself.
That history matters because Muon did not emerge from a giant benchmark suite first.
It emerged from people trying to make a tiny training program go faster.
2. The key idea: keep the directions, flatten the singular values
Suppose an update matrix has the form
G = U diag(20, 3, 0.2) V^T
The gradient has three principal matrix directions.
One direction has magnitude 20.
Another has magnitude 3.
Another has magnitude 0.2.
Muon constructs an approximation of
U diag(1, 1, 1) V^T
up to the appropriate scaling for the matrix shape.
The singular vectors remain.
The singular values disappear.
For a square matrix, this produces an ordinary orthogonal matrix. For a rectangular matrix, it produces a semi-orthogonal matrix.
Another way to write the operation is:
Ortho(G) ~= U V^T
This has a useful geometric interpretation.
Imagine that your update is saying:
"Move strongly in direction A,
somewhat in direction B,
and barely at all in direction C."
Muon says:
"Those are the important directions.
Let's give them roughly equal opportunity to affect the layer."
The original Muon write-up notes that transformer updates often have high condition numbers: a few directions can dominate the update while other directions have much smaller singular values. The authors proposed that orthogonalization may help those lower-magnitude directions contribute more. That explanation is an empirical hypothesis rather than the complete theoretical justification for Muon.
This also explains why Muon is fundamentally different from simply changing Adam's hyperparameters.
Adam changes how each coordinate is scaled.
Muon changes the matrix geometry of the update.
There is an analogy to dimensionality reduction, but in reverse.
PCA asks:
Which directions contain most of the variation?
Muon asks:
What if the update contains a few dominant directions,
but I want the matrix update to retain all of its principal directions
at comparable scale?
3. But an SVD every training step would be ridiculous
There is an obvious problem.
If we literally want:
G = U Sigma V^T
G -> U V^T
we could compute an SVD.
For a huge transformer, doing a full SVD for every large weight matrix at every optimizer step would be an unattractive idea.
Muon's practical insight is:
we do not need to compute the SVD explicitly.
Instead, use Newton-Schulz iteration.
Start by normalizing the matrix:
X = G / ||G||_F
Then repeatedly apply a matrix polynomial.
The production Muon implementation uses:
X_next = a X
+ b (X X^T) X
+ c (X X^T)^2 X
with coefficients approximately:
a = 3.4445
b = -4.7750
c = 2.0315
and typically five iterations.
Why does this work?
Take the SVD:
X = U Sigma V^T
The polynomial operation preserves the singular vectors:
p(X) = U p(Sigma) V^T
So instead of manipulating the whole matrix conceptually, we can think about what the polynomial does to each singular value.
The quintic mapping is:
p(s) = a s + b s^3 + c s^5
Repeatedly apply it.
The goal is for the singular values to converge toward 1.
So:
U Sigma V^T
|
v
U p(Sigma) V^T
|
v
U p(p(Sigma)) V^T
|
v
U I V^T
The fascinating implementation detail is that we never explicitly calculate U, Sigma, or V.
We only perform matrix multiplications.
This is where numerical linear algebra meets GPU engineering.
The early Muon work considered several ways of doing the orthogonalization. SVD was too slow. Other Newton-style methods had numerical problems in lower precision. Newton-Schulz could be run efficiently in bfloat16, which made it much more suitable for modern accelerators. The coefficients themselves were tuned experimentally; Jordan describes researchers using Desmos to explore polynomial shapes during the NanoGPT speedrun.
That is a useful general lesson for ML engineers:
an algorithm that is mathematically expensive may become practical when you find a formulation that maps onto the hardware's favorite operations.
4. The deeper idea: Muon is really about choosing a geometry
There is a deeper way to understand all of this.
Suppose a linear layer is:
y = W x
and we change the weights by:
W -> W + dW
The resulting change in the output is:
dy = dW x
So we can ask:
How large should a weight update be if I care about controlling the change it causes to the layer's output?
Suppose we measure vectors using RMS:
||x||_RMS = sqrt((1/d) sum_i x_i^2)
Then a matrix has an operator norm describing its maximum RMS-to-RMS amplification:
||W||_(RMS->RMS)
Now imagine the optimization problem:
minimize <G, dW>
subject to ||dW||_(RMS->RMS) <= eta
In plain English:
Choose the update that gives the largest first-order
decrease in loss, while limiting how much the layer
can change its outputs.
The solution involves the orthogonalized gradient:
dW ~= -eta * scale * U V^T
So the U V^T operation is not merely an arbitrary trick.
It arises from asking what "the biggest useful update" means under a matrix norm that is tied to the behavior of a linear layer.
This is part of Jeremy Bernstein and Laker Newhouse's broader work on modular duality. Their 2025 ICML paper develops a framework in which different neural-network modules can be assigned different geometries, with GPU-friendly dualization procedures for layers such as Linear and Conv2D. Newton-Schulz appears naturally in that construction.
This perspective also connects Muon to Shampoo.
Without its accumulation mechanism, the Shampoo update can be algebraically reduced to an orthogonalized gradient:
G
|
v
U Sigma V^T
|
v
U V^T
So Muon can be viewed as a particularly cheap, momentum-based way of getting this matrix-aware behavior.
That is one reason the optimizer is intellectually interesting.
It is less about inventing another collection of moving averages and more about asking:
What metric should a neural-network layer use for optimization?
5. Getting from a 2024 speedrun to real LLM training
The early results were promising, but there was a serious problem:
Would this thing actually scale?
Moonshot AI addressed that question in the 2025 paper Muon is Scalable for LLM Training.
They identified two practical issues that mattered at larger scale:
1. Weight decay
2. Correct scaling of the Muon update
The second point is particularly important.
Muon produces an orthogonalized matrix whose RMS behavior depends on the dimensions of the matrix.
A 1024 x 1024 matrix and a 8192 x 8192 matrix cannot simply receive the identical raw update scale and be expected to behave identically.
Moonshot introduced an update scaling rule designed to make Muon's update RMS comparable to AdamW's. Their experiments reported roughly 2x computational efficiency at compute-optimal training, with comparable performance reached using roughly 52% of the training FLOPs of the AdamW counterparts in their scaling experiments.
They also trained Moonlight, a 3B/16B mixture-of-experts model, on 5.7 trillion tokens using Muon.
This is where the distinction between "interesting optimizer paper" and "useful engineering technique" becomes important.
A 2x efficiency result is economically meaningful only when the comparison is properly controlled.
For example, if a training run costs:
$1,000,000
and the compute requirement genuinely falls by 48%, the idealized savings are:
$1,000,000 * 0.48 = $480,000
But actual GPU spend is not a pure FLOP meter.
You also have:
GPU utilization
communication
checkpointing
data loading
optimizer implementation
network topology
engineering time
failed runs
So "48% fewer FLOPs" should be read as an opportunity for lower cost, rather than as a promise of a 48% lower cloud bill.
There is also a useful memory difference.
Adam-like optimizers commonly maintain two moment tensors:
m
v
Muon's core optimizer state contains one momentum buffer.
Ignoring parameter replicas, master weights, sharding, and datatype choices:
AdamW optimizer state: 2 x parameter bytes
Muon optimizer state: 1 x parameter bytes
At 100B parameters, if those state tensors were stored in fp32:
100B * 4 bytes = 400 GB
AdamW moments:
2 * 400 GB = 800 GB
Muon momentum:
1 * 400 GB = 400 GB
That difference becomes relevant when optimizer state is one of the constraints determining how many GPUs a training job needs.
The computational cost of Newton-Schulz is also less frightening than the name suggests.
For an n x m matrix, with m <= n, the Muon write-up derives an extra cost of roughly:
6 T n m^2
FLOPs for T Newton-Schulz steps.
The corresponding forward-plus-backward cost for the linear layer scales roughly like:
6 n m B
where B is the number of tokens processed by the layer in the batch.
The ratio is therefore approximately:
overhead ~= T m / B
Take a hypothetical training setup:
model width m = 4096
tokens per batch B = 4,000,000
Newton-Schulz steps T = 5
Then:
overhead ~= 5 * 4096 / 4,000,000
~= 0.00512
~= 0.51%
That is the operational trick.
The expensive-looking matrix computation is amortized across millions of tokens.
For the actual NanoGPT speedrun configuration discussed by Jordan, the corresponding estimate was about 0.7%.
As of 2026, the idea has also moved into mainstream infrastructure. Current PyTorch documentation exposes torch.optim.Muon, including different update-scaling modes, and the DeepSpeed team added Muon support in June 2026.
6. What does a developer actually do with Muon?
The first mistake would be:
"Replace AdamW everywhere with Muon."
That is not how the original method is intended to be used.
Muon is primarily for 2D hidden weight matrices.
Parameters such as:
embeddings
biases
LayerNorm parameters
other 1D parameters
input layers
output heads
remain on a conventional optimizer such as AdamW.
The original experiments also found that Q, K, and V were better treated as separate matrices rather than as one fused QKV matrix.
Conceptually, your optimizer setup looks like:
hidden Linear weights -> Muon
embeddings -> AdamW
normalization -> AdamW
biases -> AdamW
LM head -> AdamW
With current PyTorch, a schematic setup looks like this:
muon_opt = torch.optim.Muon(
muon_params,
lr=3e-4,
weight_decay=0.01,
momentum=0.95,
nesterov=True,
adjust_lr_fn="match_rms_adamw",
)
adamw_opt = torch.optim.AdamW(
adamw_params,
lr=3e-4,
weight_decay=0.01,
)
The important part is not the exact numbers above.
The important part is constructing muon_params deliberately rather than doing:
[p for p in model.parameters() if p.ndim == 2]
because a tensor being 2D does not automatically mean that its optimization geometry should be Muon's.
Also, do not blindly copy early Muon examples that use a fixed lr=0.02.
There have been several generations of scaling conventions. Keller Jordan's original implementation, Moonshot's RMS-matching variant, and Bernstein's theoretical scaling rule use different parameter-shape-dependent factors. Current PyTorch exposes all three approaches.
For a developer evaluating Muon, I would therefore treat the optimizer as part of the training configuration, not as a one-line AdamW replacement.
A reasonable experiment is:
same architecture
same data
same tokens
same batch size
same hardware
same evaluation checkpoints
compare:
AdamW
Muon + AdamW hybrid
Then measure:
validation loss vs tokens
validation loss vs FLOPs
validation loss vs GPU-hours
peak memory
step time
The last three matter because an optimizer can improve sample efficiency while making each step slower, or reduce FLOPs without reducing wall-clock time on a communication-bound cluster.
7. The larger lesson
Muon is interesting because it exposes a broader idea that is easy to miss when working with large neural networks:
the optimizer contains assumptions about what a parameter means.
Adam implicitly says:
parameters are coordinates
gradients are coordinates
normalize coordinates independently
Muon says:
these parameters form matrices
matrices represent linear operators
linear operators have meaningful singular directions
optimize those operators using a matrix-aware geometry
That shift is bigger than the particular Newton-Schulz polynomial.
The Newton-Schulz iteration is the implementation technique.
The deeper idea is choosing a geometry that matches the structure of the object being optimized.
That may also explain why optimizer research sometimes looks disconnected from ordinary software engineering. An optimizer sounds like a small implementation detail:
optimizer.step()
Yet changing that one line changes the effective geometry of a trillion-parameter computation.
Muon began as an October 2024 experiment in a community speedrun, accumulated contributions from researchers and engineers working on the math and GPU implementation, scaled into the Moonlight training run, and subsequently became available in major training infrastructure. The interesting question now is less "Is Muon the replacement for AdamW?" and more:
How many other parts of deep-learning systems are still being optimized using mathematical abstractions chosen for convenience rather than for the structure of the object itself?
What other neural-network components do you think deserve their own optimization geometry?
Your team's attention is limited, and the deluge of AI-generated code is making it harder to keep production reliable and secure without slowing you down.
I'm building LiveReview, a blast-radius aware AI code review built for your business-critical systems.
Instead of presenting every diff with equal emphasis, LiveReview scores each change by blast radius — how far its impact reaches through your call graph — so you can focus attention where it actually matters.
Spend code review effort where business risk is highest — not spread evenly across every diff.
⭐ Star it on GitHub:
HexmosTech
/
LiveReview
Blast-Radius Aware AI Code Review for Business-Critical Systems
LiveReview: Blast-Radius Aware AI Code Review for Business-Critical Systems
LiveReview is an AI code reviewer that scores every hunk of a diff by blast radius: how far a change reaches through your call graph, how much persistent state it touches, and how well-tested it is. A 3-line change to a shared auth check can outrank a 300-line UI tweak. Your team's attention goes to the highest-risk code first, not spread evenly across every diff.
blast-radius-demo.mp4
LiveReview's Blast Radius & Review Priority scoring, live in the diff viewer.
The exact math, not a black box
Visualize blast radius at a glance
Every factor that feeds the score
Here's the goal:
- A 3-line fix in a function used by 40 other files, that also writes to a database, should score high.
- A 300-line UI change in one file, fully covered by…
Click below to try LiveReview with your codebase:
Top comments (0)