LLMModule 12 · Video 08

Muon and MuonClip: orthogonalised updates for LLM training

Muon treats each weight matrix as a matrix, not a bag of numbers. It takes the momentum update, keeps its directions and throws away their uneven sizes, using a few cheap matrix multiplications. This video goes from Adam to Muon to MuonClip, the version Moonshot AI used to pre-train Kimi K2 without loss spikes. This guide explains each step from the original write-ups and lets you watch the Newton–Schulz iteration work.

Video length17:52

Published13 October 2025

LevelIntermediate: matrix multiplication and the SVD help; Video 07 first

Builds onVideo 07 · AdamW, 30u30 Day 13 · attention, Module 05 · MLA

Video by Jia-Bin Huang · Watch on YouTube

Start here (no maths needed)

Muon is a newer rule for training AI models. Instead of nudging each number on its own, it looks at a whole grid of numbers at once and makes sure every direction of change gets a fair share of the step.

Think of it like this

Imagine a choir where a couple of loud singers drown everyone else out. The conductor wants every voice heard, so she asks the loud ones to sing softer and the quiet ones to sing louder, until everyone sings at the same volume. They still sing their own notes; only the volume changes. Muon does this to the update for each layer: the "notes" are the directions of change, the "volume" is how big each one is, and Muon sets every volume to about the same level. A trick called Newton–Schulz does the evening-out with a few quick rounds of multiplication, like the conductor giving the same instruction five times.

How it works

Compute the gradient for a layer's weight grid and blend it with recent gradients (momentum), like Adam does.
Notice that this update is lopsided: a few directions are huge and most are tiny.
Run five rounds of a simple multiply-and-add recipe that pushes every direction's size towards 1.
Use that evened-out grid as the step.
In very big models, also watch the attention scores, and gently shrink the query and key weights of any head whose scores get too big (MuonClip).

Why robots care

Big robot models are trained the same way as big language models. If an optimiser can reach the same quality with less computing, or without the training run crashing halfway, that saves money and time for anyone training a robot foundation model. Right now this connection is mostly a promise: Muon's published wins are on language models.

Try it

Stretch a rubber band between your fingers in one direction only and look at how lopsided its shape is. Now stretch it equally in all directions into a circle. Muon turns the lopsided ellipse into a circle, but keeps it pointing the same way.
In the demo below, leave Coefficients on Muon's tuned set and slide NS steps from 0 to 5: watch the tiny singular values shoot up. Then switch to the classic cubic recipe and see how many more steps it needs. Raise Condition number to make the update more lopsided.

Words to know

Weight matrix: The grid of numbers inside one layer that turns its input into its output.
Singular values: How strongly a matrix stretches things along each of its main directions.
Orthogonalise: Keep a matrix's directions but set all its stretches to 1.
Newton–Schulz iteration: A short repeated recipe of matrix multiplications that orthogonalises a matrix approximately, with no SVD.
Attention logit: The raw score between a query and a key before the softmax turns scores into weights.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Watch

Chapter guide

Time	Chapter	What to look for
00:00	Why Muon?	The motivation: an optimiser reported to beat well-tuned AdamW on speed records, and the question of whether that holds at scale.
00:36	Reviewing Adam	Adam's per-parameter scaling. Note that it treats every entry of a matrix independently; that's the assumption Muon drops. Video 07 covers Adam fully.
02:13	Linear layer	Why a weight matrix should be judged by how much it can change the layer's output, which leads to the spectral norm and to singular values.
04:24	Solving orthogonalization with SVD	Replacing the update G = USVᵀ with UVᵀ: same directions, all singular values set to 1. Then why a full SVD each step is too slow.
06:28	Newton–Schulz iteration: odd polynomial matrix	The key identity: an odd matrix polynomial acts on the singular values alone, so you never need the SVD itself.
08:11	Newton–Schulz iteration: example	How repeated application pushes values in (0, 1] towards 1, and how the coefficients trade speed near zero against precision near one. The demo below plays the same game.
10:35	The Muon optimizer	The full algorithm: momentum, then Newton–Schulz, then the step; which parameters use Muon and which stay on AdamW.
11:49	The exploding attention logit crisis	The instability that appears when scaling Muon up: max attention logits growing into the hundreds or thousands, and why existing fixes didn't fit.
15:13	MuonClip: extending QK-clip to Multi-head Latent Attention (MLA)	Rescaling query and key weights per head after each step, and how this is split across MLA's head-specific and shared components.
17:24	Results of MuonClip	Max-logit curves and the loss curve for Kimi K2's pre-training.

2. The ideas, explained

Adam, seen as steepest descent

Bernstein and Newhouse's "Old Optimizer, New Norm: An Anthology" (2024) re-reads popular optimisers as steepest descent under a chosen norm. Switch off Adam's moving averages and its update becomes g/|g| elementwise, which is sign descent: the steepest-descent direction when step size is measured with the max norm (ℓ∞) over all parameters. That norm ignores the fact that some parameters form a matrix that acts on vectors. The paper's suggestion is to give each tensor a norm suited to its role in the network, and for the hidden linear layers that norm is spectral.

The linear-layer view: why the spectral norm

A linear layer computes y = Wx. Changing the weights by ΔW changes the output by ΔW·x, and the worst case over inputs is bounded by the spectral norm (the largest singular value):

‖ΔW · x‖ ≤ ‖ΔW‖₂ · ‖x‖           ‖ΔW‖₂ = largest singular value of ΔW

steepest descent under the spectral norm:
  minimise ⟨G, ΔW⟩  subject to  ‖ΔW‖₂ ≤ η
  with G = U S Vᵀ (SVD)   →   ΔW = −η · U Vᵀ

The answer keeps the gradient's singular vectors and replaces every singular value with 1. Every direction moves by the same amount, so the step changes the output by a controlled amount whatever the input. Jeremy Bernstein's "Deriving Muon" post works this derivation out carefully, with the norm scaled to the layer's input and output dimensions. Keller Jordan's write-up adds an empirical motivation: momentum and Adam updates for transformer matrices typically have a very high condition number, close to low-rank, and orthogonalising boosts the "rare directions" that are small but still useful.

Orthogonalisation with the SVD, and why not

The operation is the polar factor of G, the nearest semi-orthogonal matrix in Frobenius norm:

Ortho(G) = argmin_O ‖O − G‖_F   s.t.  OᵀO = I or OOᵀ = I
         = U Vᵀ                  where G = U S Vᵀ

Computing an SVD for every weight matrix at every step is far too slow on GPUs, which are built for matrix multiplication, not for iterative decompositions. Shampoo-style inverse matrix roots can do it but, as Jordan notes, need fp32 to stay stable. The fix is an iteration made only of matrix multiplies, stable in bf16.

Newton–Schulz: an odd polynomial on singular values

The trick behind the 06:28 chapter: for G = USVᵀ, any odd polynomial of the form aG + b(GGᵀ)G + c(GGᵀ)²G equals U·φ(S)·Vᵀ with φ(x) = ax + bx³ + cx⁵ applied to each singular value. Matrix multiplications therefore reshape the singular values without touching the singular vectors, and you never compute U, S or V.

X₀ = G / ‖G‖_F                    so every singular value is in (0, 1]
repeat N times:
  A = X Xᵀ
  X = a·X + (b·A + c·A·A)·X        singular values: σ → φ(σ) = aσ + bσ³ + cσ⁵

classic cubic:      (a, b, c) = (1.5, −0.5, 0)         converges to exactly 1, slowly near 0
simple quintic:     (2, −1.5, 0.5)                      φ(1) = 1, φ′(1) = 0
Muon's tuned set:   (3.4445, −4.7750, 2.0315)           large slope a near 0; ends in a band near 1

The slope at zero is a, so a large a lifts tiny singular values quickly: roughly a× per step while they're small. Keller Jordan tuned the coefficients to maximise a while keeping the long-run values roughly within 0.7–1.3, after finding that a band that loose did not hurt training. With those coefficients, five steps are enough in his experiments. Cost: for an n×m matrix (m ≤ n), each step costs about 2(2nm² + m³) FLOPs, and the write-up bounds the total overhead relative to the forward and backward pass by Tm/B, with T the number of steps, m the model width and B the batch in tokens. Its worked examples come to 0.7% for the NanoGPT speedrun and 0.5% for a Llama-405B-sized run.

Interactive: Newton–Schulz pushing singular values towards 1

Left: the curve φᴺ(x), what N steps do to a starting singular value x, with the 0.7–1.3 band shaded and the line y = 1 dashed. Right: 16 singular values of a lopsided update after Frobenius normalisation, sorted, before (grey outlines) and after N steps (blue bars). Their spread is set by the condition number; values are drawn from a fixed seed. Because the iteration acts on each singular value separately, these bars are exactly what the matrix iteration would produce.

Coefficients: NS steps: Condition number:

Smallest σ

–

Largest σ

–

σ in 0.7–1.3

–

Spread (max ÷ min)

–

Try this: at the defaults (tuned coefficients, 5 steps, condition number 1,000) the smallest value rises from about 0.001 to about 0.36 and 14 of 16 values sit in the band, while the largest overshoots to about 1.2. Switch to the simple quintic: after 5 steps the smallest is only about 0.02, and it takes 10 steps to reach about 0.67. The classic cubic is slower still. Notice the tuned set never settles at exactly 1: it oscillates within roughly 0.68–1.2 here, which is the deliberate trade Jordan describes.

The Muon optimizer

Keller Jordan introduced Muon (MomentUm Orthogonalized by Newton-Schulz) in late 2024 for the NanoGPT and CIFAR-10 speedrunning competitions. For each 2D hidden-layer weight:

M ← μ · M + G                       SGD momentum (μ is the momentum coefficient)
O = NewtonSchulz5(M)                ≈ U Vᵀ of the momentum
W ← W − η · O

Practical rules from the write-up: Muon is only for 2D hidden-layer matrices (convolution filters can be flattened). Scalars, vectors, the embedding and the output head stay on AdamW, which Jordan reports matters for performance even though those layers are 2D. Memory is the same as SGD with momentum: one buffer per parameter instead of Adam's two. Momentum goes before the orthogonalisation; the write-up notes that an earlier method which orthogonalised first and averaged afterwards did worse.

Reported results in the write-up include a CIFAR-10 speed record cut from 3.3 to 2.6 A100-seconds, a 1.35× improvement on the NanoGPT speedrun record, and a 1.5B-parameter model reaching GPT-2 XL level on HellaSwag in 10 8×H100-hours against 13.3 for AdamW.

Scaling it up (Moonlight). Liu et al. (Moonshot AI, 2025), "Muon is Scalable for LLM Training", identified two changes needed at scale: add weight decay, and rescale each matrix's update so its RMS matches what AdamW would produce, so learning rates and weight decay carry over without retuning. The rescaling in Kimi K2's algorithm is O · 0.2 · √max(n, m). Their scaling-law experiments report about 2× compute efficiency over AdamW, and they trained Moonlight, a 3B-active / 16B-total MoE model, on 5.7T tokens with it.

W ← W − η · ( 0.2 · √max(n, m) · NewtonSchulz(M)  +  λ · W )

A sober comparison. Wen, Hall, Ma and Liang's "Fantastic Pretraining Optimizers and Where to Find Them" (2025) tuned ten optimisers carefully across 0.1B–1.2B models. They found the fastest ones, Muon and Soap, both precondition with matrices, but the speedup over well-tuned AdamW shrank with scale, from about 1.4× at 0.1B to about 1.1× at 1.2B. Treat "2×" claims as setting-dependent.

The exploding-logit problem

The Kimi K2 report (Kimi Team, 2025) describes the instability met when scaling Muon. Define the max logit of head h over a batch:

S_max^h = (1/√d) · max over samples, i, j of  q_i^h · k_j^h

In a 9B-active / 53B-total MoE trained with vanilla Muon, max logits quickly passed 1,000, which the report links to loss spikes and occasional divergence; it says this happened more often with Muon than with AdamW in their experiments. Two known fixes didn't fit. Logit soft-capping clips the scores, but the query–key dot products can keep growing underneath. QK-Norm normalises queries and keys, but with Multi-head Latent Attention the keys are never fully materialised at inference, so there is nothing to normalise. (Module 05 explains MLA.)

MuonClip: QK-clip, per head, adapted to MLA

QK-Clip leaves the forward and backward pass alone and instead rescales weights after the optimiser step, whenever a head's max logit exceeded a threshold τ:

γ_h = min(1, τ / S_max^h)                    per head, τ = 100 for K2
MHA:   W_q^h ← γ_h^α · W_q^h,  W_k^h ← γ_h^(1−α) · W_k^h      α = 0.5
MLA:   head-specific q^C, k^C   × √γ_h
       head-specific rotary q^R × γ_h
       shared rotary k^R        untouched (it is shared across heads)

Because the logit is a product of a query and a key, scaling both by √γ scales the logit by γ, bringing that head back to τ. Clipping per head matters because only a few heads blow up; a global γ would needlessly shrink healthy heads. In MLA the rotary key part is shared by all heads, so it is left alone and the head's own rotary query part takes the full γ. MuonClip is then Moonlight's Muon (weight decay and RMS matching) plus QK-Clip.

The report's results: for K2's 15.5T-token pre-training with τ = 100, max logits were capped at 100 early on and decayed to a normal range after roughly 30% of training, with no adjustment to τ, and the per-step loss curve showed no spikes. An appendix reports small-scale ablations in which QK-Clip did not hurt model quality.

3. Where the course goes deeper

Topic	Course page
Adam, AdamW and the optimiser family tree	Video 07 · Optimisers to AdamW
Attention, queries, keys and the softmax	30u30 Day 13 · Attention is all you need
Multi-head Latent Attention and the KV cache	Module 05 · Trillion-scale architecture (length); Module 08 · Stage 08 · Inference-aware design
Training instability and large MoE runs	Module 05 · Trillion-scale architecture; Module 08 · Stage 03 · Beyond transformers
Optimiser memory under ZeRO sharding	Module 08 · Stage 06 · Distributed training
Matrix maths and autograd from scratch	Module 08 · Stage 01 · Tensor maths and autograd

Why it matters for robotics

The link is indirect. Muon's published evidence is on language-model pre-training. I'm not aware of a robot-policy paper showing it beats AdamW, so treat it as an LLM-training tool that robot foundation models may inherit.
Robot foundation models are transformers. VLAs are built on language and vision backbones and trained with the same recipes, so a cheaper or more stable optimiser for their hidden matrices would carry over to their pre-training.
Stability is worth more than speed for expensive runs. QK-Clip's appeal is a long run with no loss spikes. A robot lab with one shot at a large pre-training run cares about that.
Memory. One optimiser buffer instead of two frees about a quarter of model-state memory in mixed precision, which helps on smaller training clusters.
Fine-tuning on a robot dataset is a different regime. Small-data fine-tunes are often done with LoRA and AdamW; nothing in these sources tells you Muon helps there.

Self-check

What does Muon do to the singular values and singular vectors of the momentum matrix?

It keeps the singular vectors and pushes every singular value towards 1, approximately replacing USVᵀ with UVᵀ.

Why can Newton–Schulz orthogonalise a matrix without ever computing its SVD?

An odd matrix polynomial aX + b(XXᵀ)X + c(XXᵀ)²X equals U·φ(S)·Vᵀ, so it reshapes singular values elementwise while leaving U and V unchanged. Choosing φ so that repeated application sends (0, 1] towards 1 orthogonalises using only matrix multiplications.

Why is G divided by its Frobenius norm first?

The iteration only converges for singular values in (0, 1]. The Frobenius norm is at least the largest singular value, so dividing by it puts every singular value in range, and orthogonalisation ignores overall scale anyway.

Which parameters should not use Muon, according to Keller Jordan's write-up?

Scalars and vectors (biases, norm gains), plus the input embedding and the output head, which stay on AdamW.

Why couldn't Kimi K2 use QK-Norm to stop exploding logits?

K2 uses Multi-head Latent Attention, where the keys are not fully materialised at inference, so there's no explicit key vector to normalise. QK-Clip instead rescales the query and key weights after each step.

A head has max logit 400 and τ = 100. What happens to its MLA weights under QK-Clip?

γ = 100/400 = 0.25. The head-specific q^C and k^C weights are scaled by √0.25 = 0.5, the head-specific rotary query weights by 0.25, and the shared rotary key is untouched.

References

From the video's description:

Muon: An optimizer for hidden layers in neural networks (Keller Jordan, blog, December 2024)
Deriving Muon (Jeremy Bernstein, blog)
Old Optimizer, New Norm: An Anthology (Bernstein and Newhouse, 2024)
Muon is Scalable for LLM Training (Liu et al., Moonshot AI, 2025)
MuonClip: Kimi K2: Open Agentic Intelligence (Kimi Team, 2025)
Fantastic Pretraining Optimizers and Where to Find Them (Wen, Hall, Ma and Liang, 2025). The description's link for this entry points to the Anthology paper above (arXiv 2409.20325); this is the paper with this title.
Course link: Module 05: Kimi K3 trains with a per-head variant of Muon

The description also links the author's AdamW video (Video 07 here) and notes the video was made with Manim.