Muon
An Optimizer from First Principles
Named after the subatomic particle, Muon replaces AdamW's element-wise heuristics with a single, geometrically principled operation: matrix orthogonalization. This page rebuilds the entire idea from scratch — every choice explained, every alternative considered.
Deep Dive — 12 Sections
A neural network starts as random noise — millions of numbers with no useful structure. Training is the process of adjusting those numbers so the network does something useful. But there are millions (or billions) of numbers to adjust simultaneously. How?
The Gradient: Which Way is Downhill?
For any given set of weights, we can compute a loss — a single number measuring how wrong the network is. The gradient tells us, for each weight, which direction would decrease the loss. It's a vector pointing “downhill” in weight-space.
The simplest optimizer — Stochastic Gradient Descent (SGD) — just follows the gradient:
Why Not Just Use SGD?
SGD treats every direction equally. But loss landscapes are not equally curved in every direction. Imagine a narrow valley — steep walls on the sides, gentle slope along the bottom:
- —SGD bounces between the steep walls while barely making progress along the valley floor
- —The learning rate must be tiny to avoid exploding on the steep walls
- —This makes progress along the gentle direction agonizingly slow
The fundamental tension: different directions need different step sizes. This is what motivates every optimizer since SGD.
The Question That Led to Muon
Adam solved the “different directions need different step sizes” problem with per-element scaling. But what if there's a better way — one that respects the matrix structure of neural network weights, instead of treating them as flat bags of numbers?
Each major optimizer added one key idea. Understanding this lineage makes Muon's contribution crystal clear.
Added: Running average of gradients (momentum). Instead of following the current gradient alone, average it with recent gradients. This smooths out noise and builds up speed in consistent directions.
Still missing: All directions still get the same step size
Added: Per-element adaptive scaling. Track both the mean (m) and variance (v) of each element's gradient. Divide by √v to normalize — elements with large gradients get scaled down, elements with small gradients get scaled up.
Still missing: Treats each weight as independent number. Ignores matrix structure entirely. Needs 2 buffers (2× memory).
Added: Decoupled weight decay. Instead of adding L2 penalty to the loss (which gets scaled by Adam's adaptive rates), apply weight decay directly to weights. Small but important fix for regularization.
Still missing: Still per-element. Still ignores that weights form matrices. Still 2× memory. This is the reigning champion that Muon aims to replace.
Key insight: Neural network weights are matrices — linear transformations between vector spaces. Instead of per-element normalization (Adam), orthogonalize the entire gradient matrix. This equalizes all singular values to 1 — uniformly scaled in every direction, respecting the matrix geometry.
Optimizer Trajectories on an Elongated Loss Surface
L(x, y) = 50x² + y² — a 50:1 condition ratio that punishes naive gradient descent
SGD follows the raw gradient. The x-gradient is 50x larger than the y-gradient, so it oscillates wildly across the narrow valley while creeping along it — most energy is wasted fighting curvature.
Adam adapts per-coordinate: it divides each gradient component by its running RMS, partly equalizing the two directions. This damps the x-oscillations, but the ratio of first to second moments still biases the path.
Muon orthogonalizes the gradient (via Newton-Schulz iteration in practice), which erases all scale differences between directions. The optimizer sees an effectively isotropic landscape and takes the most direct path to the minimum.
This is the conceptual leap that makes Muon possible. Every other optimizer treats weights as a flat list of numbers. Muon says: “Wait — these numbers form a matrix, and that matrix has geometric meaning.”
How Adam Sees Weights
A 4096×4096 weight matrix = 16,777,216 independent numbers. Each gets its own running mean and variance. No element knows about any other.
Each cell is optimized independently. The matrix structure is invisible.
How Muon Sees Weights
A 4096×4096 weight matrix = a linear transformation from R4096 → R4096. It has singular values describing how much it stretches space in each direction.
The whole matrix is one geometric object. Optimization should respect its structure.
Why Does This Matter? A Concrete Example
Consider a gradient matrix with SVD: G = UΣVT, where Σ = diag(10, 1, 0.1).
Per-element normalization scrambles the singular structure. The σ=10 direction gets attention, σ=0.1 gets ignored. The resulting update matrix may not even preserve the original directions.
Orthogonalization: O = U·I·VT. Same directions (U, V), but all singular values → 1. Every direction gets equal update magnitude. The network explores the full space uniformly.
What About Non-Matrix Parameters?
Muon only applies to weight matrices (the hidden layers). Parameters that aren't matrices — like layer norm scales, biases, and embeddings — still use AdamW. In practice this is fine: matrix parameters are 95%+ of a transformer's total parameters.
Orthogonalization is the mathematical heart of Muon. Given a gradient momentum matrix M, we want to find the closest orthogonal matrix O — one where all singular values are exactly 1. This is called the polar decomposition.
SVD Decomposition — M = UΣVT
A matrix transforms the unit circle into an ellipse. The singular values (σ1, σ2) are the semi-axis lengths — they measure how much the matrix stretches space along each principal direction.
What if we don't orthogonalize? When σ1 ≫ σ2, the dominant direction receives disproportionately large updates. Over many steps, this causes the optimizer to repeatedly reinforce the same direction while starving others — gradients become rank-deficient and training slows. The condition number (σ1/σ2 = 4.0) measures this imbalance.
What Orthogonalization Does
Given M with SVD: M = UΣVT, where Σ = diag(σ₁, σ₂, ..., σₙ):
It preserves the directions (U, V stay the same) but equalizes the magnitudes (all σᵢ → 1).
What If We Chose Differently?
- ?Just normalize (M/‖M‖)? This preserves the relative singular values. A dominant direction stays dominant. Doesn't solve the problem.
- ?Clip singular values? Better, but you'd need to compute SVD (expensive) and choose a clipping threshold (new hyperparameter).
- ?Per-row/column normalization? Doesn't respect the global singular structure. Two rows could conspire to still create a dominant direction.
Polar decomposition is the unique operation that preserves directions while making the matrix “as close to orthogonal as possible.”
The Intuition: Audio Equalization
Think of an audio mixer. The gradient is like a music signal where bass is at volume 10 and treble is at volume 0.1. Adam adjusts the volumes per-frequency but imperfectly. Muon does a perfect equalization — every frequency gets set to exactly the same volume. You hear the full spectrum equally. In optimization terms: the network explores all learnable directions with equal intensity.
Computing the exact polar decomposition requires SVD — an O(n³) operation that's notoriously slow on GPUs. Muon needs a fast approximation. Enter the Newton-Schulz iteration: a polynomial that converges to the orthogonal factor using only matrix multiplications.
Newton-Schulz Convergence
Watch singular values converge to 1 through iterative polynomial updates
The Iteration
Each iteration is 2-3 matrix multiplications. On a GPU, matrix multiply is the fastest operation available — this is perfectly hardware-aligned.
Why These Specific Coefficients?
The polynomial f(x) = ax + bx³ + cx⁵ must have x=1 as a super-attracting fixed point: f(1)=1 and f'(1)=0. This ensures singular values near 1 stay put, and far-away values race toward 1.
What If We Used More Iterations?
N=10 gives a slightly more accurate result, but it doesn't improve training performance — because bfloat16 only has ~3 decimal digits of precision anyway. 5 iterations already saturate bf16 accuracy. More iterations just waste compute without improving the result where it matters.
Here's a beautiful theoretical perspective: both Adam and Muon can be understood as steepest descent under different norm constraints. The choice of norm determines which directions the optimizer amplifies.
Steepest Descent Under Different Norms
Adam and Muon are both doing steepest descent — just under different norms. The choice of norm determines which directions get amplified. Drag the angle slider to rotate the gradient and watch how each optimizer picks a different update direction.
Why spectral norm matters: Weight matrices are linear operators — they transform vectors, not individual numbers. The spectral norm measures what matters: how much the operator stretches inputs. By doing steepest descent under this norm, Muon produces updates that respect the geometric structure of the weight space, leading to more efficient optimization.
SGD: Frobenius Norm
“Find the update with ‖ΔW‖F ≤ 1 that decreases loss the most.”
The Frobenius norm treats the matrix as a flat vector. The unit ball is a sphere. Steepest descent = gradient direction.
Adam: Max-of-Max Norm
“Find the update where no single element exceeds 1 that decreases loss the most.”
Element-wise max norm. The unit ball is a hypercube. Steepest descent = sign of gradient (like sign-SGD).
Muon: Spectral Norm
“Find the update where ‖ΔW‖spectral ≤ 1 that decreases loss the most.”
The spectral norm = largest singular value. This is the natural operator norm for matrices. Steepest descent = orthogonalized gradient.
Why Spectral Norm is Natural for Weights
Neural network weights are linear operators — they map input vectors to output vectors. The spectral norm measures the maximum “stretch” of this operator. This is the natural notion of “size” for a linear map, just as absolute value is natural for scalars.
Using element-wise norms (Frobenius, Max) is like measuring the length of a matrix by counting how many digits its entries have — technically a number, but geometrically meaningless. The spectral norm respects what the matrix does.
The original Muon (from Keller Jordan's blog post) worked brilliantly on small models. But when the Moonshot AI team tried to scale it to billions of parameters and hundreds of billions of tokens, it broke. Here's why and how they fixed it.
Weight Decay at Scale
Why Muon needs weight decay to survive bf16 training
Weight RMS Over Training
Just right
Weights stay bounded. Stable training even in bf16.
Why does this matter? Muon's orthogonalized updates push weight matrices toward large singular values over time. Without decay, weight norms grow without bound — eventually exceeding the representable range of bfloat16 (~3.4 x 10^38 but effective precision degrades much earlier).
Weight decay applies W = (1 - lambda) * W each step, creating a steady pull toward zero that balances the outward push of gradient updates. At lambda ~0.1, this equilibrium keeps weight RMS in a stable, bf16-safe range.
The Problem: Unbounded Weight Growth
Muon's orthogonal update O always has singular values = 1. This means:
- 1.Each step adds a fixed-magnitude perturbation to W
- 2.Nothing ever shrinks W — updates only add, never subtract from magnitude
- 3.Over millions of steps: ‖W‖ grows as ~√t (random walk)
- 4.Eventually exceeds bfloat16 max (~65504) → NaN → training collapses
At small scale (≤1B params, ≤10B tokens), this growth is too slow to cause problems. At large scale, it's fatal.
The Fix: Weight Decay (λ=0.1)
Add a term that gently shrinks weights toward zero each step:
The λW term creates a restoring force. Growth from the update is balanced by shrinkage from decay, creating a stable equilibrium for weight norms.
Bonus: weight decay also acts as regularization (just like in AdamW), actually improving final loss compared to vanilla Muon — even on small models where overflow isn't a risk.
Why Didn't the Original Paper Include Weight Decay?
Keller Jordan's original Muon was tested on NanoGPT speedruns — small models (≤1B), short training (≤10B tokens). At that scale, weight growth is negligible. The decay fix was only discovered when Moonshot AI attempted to train at 100B+ token scale and observed the instability firsthand.
The second critical fix for scaling. Orthogonalization produces updates with RMS = 1/√max(A,B) for an [A×B] matrix. This means different-shaped layers get wildly different update magnitudes — not because the gradient says so, but purely because of their shape.
Per-Parameter Update Scaling
Normalizing Muon updates across different matrix shapes
Wildly unequal updates
Raw Muon update RMS = 1/sqrt(max(m,n)). Large matrices like FFN layers get tiny updates (0.0095) while small head projections get massive ones (0.088) — a 9x difference that destabilizes training.
The scaling fix: Muon's orthogonal projection produces updates whose RMS is exactly 1/sqrt(max(m,n)) for an [m x n] matrix. Multiplying by 0.2 * sqrt(max(m,n)) cancels the shape dependence, giving every parameter a uniform effective learning rate comparable to AdamW's natural scale.
The Math: Why RMS Depends on Shape
For an orthogonal matrix O of shape [A, B] (assume A ≥ B for full rank):
This is a mathematical consequence of orthogonalization, not a bug. But it means a [4096×128] matrix gets 5× larger RMS updates than a [4096×128256] matrix — for no good optimization reason.
Baseline
Global scale by hidden dim. Doesn't account for varying shapes.
Update Norm
Force all updates to exactly 0.2 RMS. Works but adds compute.
Adjusted LR ✓
Per-matrix scaling. Elegant, efficient, matches target RMS.
The Key Insight: Reuse AdamW's Hyperparameters
By scaling Muon's updates to match AdamW's typical RMS (~0.2-0.4), you can directly reuse AdamW's learning rate and weight decay — no hyperparameter search needed. This is why the paper claims Muon works “out-of-the-box” at any scale.
Training at scale means distributing across hundreds or thousands of GPUs. The standard technique — ZeRO-1 — partitions optimizer state across GPUs. But Muon has a unique challenge: Newton-Schulz needs the full gradient matrix, not just a partition.
Distributed Muon Algorithm
Average gradients across DP group. Each GPU gets its 1/N partition.
M_local = μ·M_local + G_local. Each GPU updates only its partition.
Reconstruct complete momentum matrix M from all partitions. This is the extra step Muon needs.
Run 5 iterations on the full matrix to get orthogonal update O.
Each GPU extracts its 1/N partition of O. Discard the rest.
Apply update to local partition, then reconstruct full weights.
Why Can't We Just Orthogonalize Each Partition?
Orthogonalization is a global operation — it depends on the relationship between ALL rows of the matrix. If GPU 1 has rows 1-1024 and GPU 2 has rows 1025-2048, each can't independently orthogonalize its rows because the result depends on the interaction between them. This is why the DP Gather step is unavoidable.
Scaling laws quantify how loss decreases as you increase compute. The paper ran rigorous experiments with models from 399M to 1.5B parameters to compare Muon and AdamW in the compute-optimal regime.
Muon vs AdamW Scaling Laws
Scaling laws describe how loss decreases as you spend more compute. Better optimizers shift the curve down — achieving the same loss with less compute. Drag the marker to explore different compute budgets.
What scaling laws tell us: In deep learning, loss follows a power law as compute increases: L = A * C^B. A lower curve means the optimizer extracts more performance per FLOP. Muon consistently outperforms AdamW across model sizes from 399M to 1.5B parameters, achieving roughly 2x compute efficiency — the same loss at about half the training cost. This gap persists (and may widen) as models scale, making optimizer choice increasingly important.
Muon Scaling Law
Lower constant = better efficiency at every compute budget.
AdamW Scaling Law
Higher constant means more compute needed for same loss.
What Does “2× Efficiency” Mean Concretely?
To prove Muon works at frontier scale, the team trained Moonlight — a 3B-activated / 16B-total parameter Mixture-of-Experts model on 5.7 trillion tokens.
Moonlight vs. Comparable Models
| Benchmark | Moonlight | Llama3.2-3B | Qwen2.5-3B | DSV2-Lite |
|---|---|---|---|---|
| Training Tokens | 5.7T | 9T | 18T | 5.7T |
| Optimizer | Muon | AdamW | Unknown | AdamW |
| MMLU | 70.0 | 54.7 | 65.6 | 58.3 |
| HumanEval | 48.1 | 28.0 | 42.1 | 29.9 |
| GSM8K | 77.4 | 34.0 | 79.1 | 41.1 |
| MATH | 45.3 | 8.5 | 42.6 | 17.1 |
| MBPP | 63.8 | 48.7 | 57.1 | 43.2 |
Moonlight vs Moonlight-A
Moonlight-A is the same architecture trained with AdamW instead of Muon. At 1.2T tokens, Moonlight (Muon) already outperforms Moonlight-A on most benchmarks — demonstrating the optimizer advantage directly, controlling for everything else.
SVD Entropy Confirmation
Over 90% of Moonlight's weight matrices show higher SVD entropy than Moonlight-A's. This confirms the theory: Muon's orthogonal updates produce more balanced weight matrices that utilize diverse optimization directions, not just the few dominant ones.
Open Question: The SFT Mismatch
A Muon-pretrained model fine-tuned with Muon outperforms one fine-tuned with AdamW. But an AdamW-pretrained model does not benefit from Muon fine-tuning. The optimizers seem to create different weight structures that don't transfer well. This remains an active research question.
Muon Optimizer — From First Principles
Question 1 of 10Why does Muon orthogonalize the momentum matrix instead of using it directly like SGD?