Matrix optimizers: K-FAC, Shampoo and Muon

Before this lesson: lesson 06 (Adam; its first step is $\eta\operatorname{sign}(g)$), lesson 10 (Newton, the generalized Gauss–Newton matrix and the Fisher), linear algebra 08 (SVD, polar decomposition), linear algebra 11 (spectral and nuclear norms, dual norms), linear algebra 16 (Kronecker product and vec). Time: about 80 minutes, plus the notebook.

---

Why this matters #

Adam scales each coordinate separately: its preconditioner is diagonal. Full curvature (Newton) is far too expensive for a network. Matrix optimizers are the practical middle ground, and in 2024–2025 they moved into practice:

• Distributed Shampoo won the external-tuning track of the AlgoPerf training-speed benchmark;
• Muon set records in the NanoGPT speedrun, and Moonshot AI trained Kimi K2 (1 trillion parameters) with a Muon variant;
• K-FAC-type curvature is used for influence functions of large language models and for physics-informed networks.

The theory behind Muon is short and useful beyond it: an optimizer is a choice of norm, and the update is the steepest step in that norm.

1. Predict first #

2. The picture #

Optimizers as preconditioners #

Most optimizers have the form $\Delta w = -\eta\, P^{-1} g$ for some matrix $P$:

Method	$P$	Cost for $d$ parameters
gradient descent	$I$	$O(d)$
Adam (lesson 06)	a diagonal matrix (one scale per coordinate)	$O(d)$
Newton (lesson 10)	the full Hessian (or GGN, Fisher)	$O(d^2)$ memory, $O(d^3)$ time
K-FAC, Shampoo	per layer, a Kronecker product of two small matrices	$O(m^2 + n^2)$ memory for an $m \times n$ weight

For a weight matrix $W \in \R^{m \times n}$ there are $d = mn$ parameters, and a full $P$ has $(mn)^2$ entries. A Kronecker product $R \otimes L$ of an $n \times n$ matrix $R$ and an $m \times m$ matrix $L$ is also $mn \times mn$, but it is stored with only $m^2 + n^2$ numbers, and it acts on the gradient matrix by multiplication from the left and the right: $(R \otimes L)\operatorname{vec}(G) = \operatorname{vec}(L\,G\,R^\top)$ (linear algebra 16).

Convention for this lesson. $\operatorname{vec}$ stacks the columns of a matrix (as in Martens and Grosse). Then $\operatorname{vec}(L\,X\,R^\top) = (R \otimes L)\operatorname{vec}(X)$ and $\operatorname{vec}(\delta\, a^\top) = a \otimes \delta$. With row stacking, the two factors swap places.

Steepest descent depends on how you measure a step #

For a small step $\Delta$, the loss changes by about $\langle g, \Delta\rangle$. The steepest step of size 1 is the $\Delta$ with $\|\Delta\| = 1$ that makes $\langle g, \Delta\rangle$ most negative. The answer depends on the norm. Take $g = (3, 1)$:

• Euclidean norm $\ell_2$: $\Delta = -g/\|g\|_2 = -(0.949, 0.316)$. This is gradient descent.
• Max norm $\ell_\infty$ (each coordinate may move by at most 1): $\Delta = -(1, 1) = -\operatorname{sign}(g)$. This is sign descent. Adam's first step is $-\eta\operatorname{sign}(g)$, and Adam with its moving averages switched off is sign descent.
• $\ell_1$ norm (the total movement is at most 1): $\Delta = -(1, 0)$. Move only the coordinate with the largest $|g_i|$: coordinate descent.

Interactive figure: loading…

For a matrix: the spectral norm and $UV^\top$ #

A weight matrix $W$ maps inputs to outputs. A natural size for a change $\Delta W$ is the spectral norm $\|\Delta W\|_2 = \max_{x \ne 0}\|\Delta W x\|/\|x\| = \sigma_{\max}(\Delta W)$, the largest singular value. It bounds how far the change can move any output: $\|\Delta W x\| \le \|\Delta W\|_2\,\|x\|$.

Section 3 proves that the steepest step under this norm is $-UV^\top$, where $G = U\Sigma V^\top$ is the SVD of the gradient: keep the singular vectors of $G$, and set all its singular values to 1. Gradient descent moves along each singular direction $u_i v_i^\top$ in proportion to $\sigma_i$; the orthogonalized step gives every direction the same size.

The simplest case: a diagonal gradient $G = \operatorname{diag}(5, -2)$. Its SVD has $U = I$, $\Sigma = \operatorname{diag}(5, 2)$ and $V = \operatorname{diag}(1, -1)$. So $UV^\top = \operatorname{diag}(1, -1) = \operatorname{sign}(G)$: for a diagonal gradient, orthogonalization is sign descent. In general it is "sign descent on the singular values".

Interactive figure: loading…

3. Definitions and main results #

K-FAC: Kronecker-factored curvature #

Take a linear layer $s = W a$ with input $a \in \R^n$ and $W \in \R^{m \times n}$ (a bias can be handled by appending a 1 to $a$). Backpropagation gives $\delta = \partial \ell/\partial s$ for the loss $\ell$ of one example, and the weight gradient is the outer product $\nabla_W \ell = \delta\, a^\top$.

One example: an exact Kronecker product. With column stacking, $\operatorname{vec}(\delta a^\top) = a \otimes \delta$, so

Many examples: the approximation. The Fisher block of the layer is the average of this matrix over examples, with the targets $y$ sampled from the model's own predictive distribution:

The mean of a product is replaced by the product of the means. This is exact when $aa^\top$ and $\delta\delta^\top$ are independent; in general it is an approximation. Martens and Grosse (2015), who introduced K-FAC (Kronecker-factored approximate curvature), call it a major approximation that still keeps the coarse structure of the Fisher. (They write $G$ for the second factor; we write $D$, because $G$ is the gradient matrix in this lesson.) If you use the data labels instead of sampled targets, you get the "empirical Fisher", which loses the link to the Gauss–Newton matrix of lesson 10.

Why the step is cheap. Two Kronecker identities, $(A \otimes D)^{-1} = A^{-1} \otimes D^{-1}$ and $(A \otimes D)\operatorname{vec}(X) = \operatorname{vec}(D X A)$ for symmetric $A$, give

You invert an $n \times n$ and an $m \times m$ matrix, $O(n^3 + m^3)$ operations, instead of one matrix of size $mn$, which costs $O(m^3 n^3)$. In practice K-FAC uses one block per layer (it ignores the curvature between layers), moving averages of $A$ and $D$, damping of each factor ($A + \gamma I$, $D + \gamma I$), and recomputes the inverses only every few steps.

Interactive figure: loading…

Steepest descent under a norm #

Let $\|\cdot\|$ be a norm. Its dual norm is

The dual of $\ell_2$ is $\ell_2$; the dual of $\ell_\infty$ is $\ell_1$; the dual of $\ell_1$ is $\ell_\infty$ (linear algebra 11). For matrices with $\langle G, O\rangle = \operatorname{tr}(G^\top O)$, the dual of the Frobenius norm is the Frobenius norm, and the dual of the spectral norm is the nuclear norm $\|G\|_{\rm nuc} = \sum_i \sigma_i$ (Theorem 2 below proves it).

Assume the loss satisfies a quadratic upper bound in this norm, $f(w + \Delta) \le f(w) + \langle g, \Delta\rangle + \tfrac{L}{2}\|\Delta\|^2$, like the descent lemma of lesson 02 but with any norm. Minimizing the right side gives the step.

Proposition 1 (steepest descent; Bernstein and Newhouse, 2024). For $L > 0$,

Proof.

1. Write $\Delta = -c\,t$ with a length $c \ge 0$ and a direction $\|t\| = 1$. The objective is $-c\,\langle g, t\rangle + \tfrac{L}{2}c^2$.
2. For every fixed $c$, the best direction maximizes $\langle g, t\rangle$; the maximum value is $\|g\|_*$.
3. Minimize $-c\,\|g\|_* + \tfrac{L}{2}c^2$ over $c$: the derivative is zero at $c = \|g\|_*/L$. $\blacksquare$

So the direction is the unit vector most aligned with $g$, and the length is the dual norm divided by $L$ (Bernstein and Newhouse call $L$ the "sharpness"). Each norm gives an optimizer:

Space	Norm	Dual norm	Steepest direction $-t$ with	Optimizer
$\R^d$	$\ell_2$	$\ell_2$	$t = g/\|g\|_2$	gradient descent
$\R^d$	$\ell_\infty$	$\ell_1$	$t = \operatorname{sign}(g)$	sign descent (Adam without moving averages)
$\R^{m\times n}$	Frobenius	Frobenius	$t = G/\|G\|_F$	gradient descent
$\R^{m\times n}$	spectral	nuclear	$t = UV^\top$	Muon; Shampoo without accumulation

Theorem 2: the spectral steepest direction is the polar factor #

Theorem 2. Let $G \in \R^{m\times n}$ have rank $r$ and reduced SVD $G = U\Sigma V^\top = \sum_{i=1}^r \sigma_i u_i v_i^\top$. Then

and the steepest step is $\arg\min_{\Delta W}\big[\langle G, \Delta W\rangle + \tfrac{L}{2}\|\Delta W\|_2^2\big] = -\dfrac{\|G\|_{\rm nuc}}{L}\,UV^\top$.

Why this is true (intuition). $\langle G, O\rangle = \sum_i \sigma_i\,(u_i^\top O v_i)$: each singular direction contributes $\sigma_i$ times "how much $O$ sends $v_i$ to $u_i$". A matrix with spectral norm at most 1 can send each $v_i$ to at most a unit vector, so the best it can do is $O v_i = u_i$ for every $i$. That is $UV^\top$.

Proof.

1. $\langle G, O\rangle = \operatorname{tr}(G^\top O) = \operatorname{tr}\big(\sum_i \sigma_i v_i u_i^\top O\big) = \sum_i \sigma_i\, u_i^\top O v_i$ (the trace of $v_i u_i^\top O$ equals $u_i^\top O v_i$).
2. By Cauchy–Schwarz, $u_i^\top O v_i \le \|u_i\|\,\|O v_i\| \le \|O\|_2 \le 1$, because $u_i$ and $v_i$ are unit vectors. So $\langle G, O\rangle \le \sum_i \sigma_i$.
3. $O = UV^\top$ has all nonzero singular values equal to 1, so $\|O\|_2 = 1$. And $u_i^\top UV^\top v_i = e_i^\top e_i = 1$ for every $i$, so it attains the bound. This also shows that the spectral and nuclear norms are dual.
4. Proposition 1 with the spectral norm gives the step $-(\|G\|_{\rm nuc}/L)\,UV^\top$. $\blacksquare$

When $G$ has full rank the maximizer is unique. $UV^\top$ is the orthogonal factor $Q$ of the polar decomposition $G = QP$, with $Q = UV^\top$ and $P = V\Sigma V^\top \succeq 0$ (linear algebra 08). It is also the closest semi-orthogonal matrix to $G$ in the Frobenius norm (Bernstein and Newhouse, Proposition 4). Muon uses the direction $UV^\top$ with a tuned learning rate in place of the length $\|G\|_{\rm nuc}/L$.

Why the spectral norm? Bernstein and Newhouse (Proposition 6) give the simplest case: a linear predictor $y = Wx$ with the square loss $\tfrac{1}{2n}\sum_i \tfrac{1}{d_{\rm out}}\|y_i - Wx_i\|^2$ and inputs with $\|x_i\|_2 = \sqrt{d_{\rm in}}$ satisfies

Minimizing this upper bound is exactly a spectral steepest-descent step. There is also an empirical reason (Jordan, 2024): the gradient and momentum matrices of transformer layers are almost low rank, with a few dominant directions. Orthogonalization raises the weak, "rare" directions, which may still matter for learning.

Newton–Schulz: $UV^\top$ with matrix products only #

An SVD is slow on GPUs. The Newton–Schulz iteration uses only matrix multiplications, which GPUs do very fast, also in low precision.

Lemma 3. Let $X = USV^\top$ with orthonormal columns in $U$ and $V$ and a diagonal $S \ge 0$. For coefficients $(a, b, c)$, the step

equals $U\,p(S)\,V^\top$ with the odd polynomial $p(\sigma) = a\sigma + b\sigma^3 + c\sigma^5$ applied to each singular value.

Proof.

1. $XX^\top = USV^\top VSU^\top = US^2U^\top$, because $V^\top V = I$.
2. $(XX^\top)X = US^2U^\top USV^\top = US^3V^\top$, because $U^\top U = I$. In the same way, $(XX^\top)^2X = US^5V^\top$.
3. Add the three terms: $X^+ = U(aS + bS^3 + cS^5)V^\top$. $\blacksquare$

So the singular vectors never change, and each singular value follows the scalar map $p$. After $k$ steps, $X_k = U\,p^{\circ k}(S)\,V^\top$, where $p^{\circ k}$ is $p$ applied $k$ times. If $p^{\circ k}(\sigma) \to 1$ for every singular value that occurs, then $X_k \to UV^\top$. The algorithm:

1. Normalize: $X_0 = G/(\|G\|_F + \epsilon)$. Since $\sigma_{\max} \le \|G\|_F$, all singular values are now in $[0, 1]$.
2. Transpose tall matrices: if $m > n$, work with $X^\top$, so that $XX^\top$ is the smaller Gram matrix.
3. Iterate with the chosen coefficients.

The cubic (classical) iteration, $(a, b, c) = (1.5, -0.5, 0)$, so $p(\sigma) = 1.5\sigma - 0.5\sigma^3$:

• Fixed points: $p(\sigma) = \sigma \iff 0.5\,\sigma(1 - \sigma^2) = 0 \iff \sigma \in \{0, 1, -1\}$.
• $p'(\sigma) = 1.5 - 1.5\sigma^2$, so $p'(1) = 0$. Near 1 the convergence is quadratic: $p(1 - e) = 1 - 1.5e^2 + 0.5e^3$.
• $p'(0) = 1.5$: a small singular value grows only by about 1.5 per step. From $\sigma = 0.01$ it needs 14 steps to pass 0.99.
• It converges to 1 for every $\sigma \in (0, \sqrt3)$; the normalization guarantees this.

Muon's quintic (Jordan, 2024), $(a, b, c) = (3.4445, -4.7750, 2.0315)$:

• The slope at 0 is $a = 3.4445$, so small singular values grow fast. From $0.01$, five steps give $0.034, 0.118, 0.400, 1.093, 0.699$.
• Jordan chose $a$ as large as possible while keeping the iterates in about $[0.7, 1.3]$.
• The price: 1 is not a fixed point, since $p(1) = 3.4445 - 4.7750 + 2.0315 = 0.701$. The iterates keep moving inside a band instead of converging. After 5 steps, normalized singular values that start in $[0.01, 1]$ end in about $[0.68, 1.13]$.

So Muon's output is $US'V^\top$ with $S' \approx I$, not exactly $UV^\top$. Jordan reports that this does not hurt training. The 5 steps run in bfloat16.

Interactive figure: loading…

Cost. For $m \le n$, one step costs about $4m^2n + 2m^3$ operations. Jordan bounds Muon's extra cost in transformer training by $Tm/B$ ($T$ Newton–Schulz steps, usually 5; model width $m$; batch size $B$ in tokens): about 0.7% of the training compute for the NanoGPT speedrun and 0.5% at the scale of Llama 3 405B.

Muon #

For each 2D weight $W \in \R^{m \times n}$ of a hidden layer ($m$ = fan-out rows, $n$ = fan-in columns), with gradient $G_t$, momentum coefficient $\beta$ (0.95 by default), learning rate $\eta$ and weight decay $\lambda$, torch.optim.Muon does:

1. momentum: $B_t = \beta B_{t-1} + G_t$ (starting from $B_0 = 0$);
2. Nesterov: $\tilde B_t = G_t + \beta B_t$ (without Nesterov, $\tilde B_t = B_t$);
3. $O_t = \mathrm{NS}_5(\tilde B_t)$: five quintic Newton–Schulz steps;
4. $W_t = (1 - \eta\lambda)\,W_{t-1} - \eta\, s\, O_t$, with a shape factor $s$: $s = \sqrt{\max(1, m/n)}$ (Jordan's original choice, PyTorch's default) or $s = 0.2\sqrt{\max(m, n)}$ (Moonshot, adjust_lr_fn="match_rms_adamw").

Some implementations write the momentum as a moving average, $B_t = \beta B_{t-1} + (1 - \beta)G_t$. This only changes the size of the buffer, and Newton–Schulz divides by the Frobenius norm, so $O_t$ is the same. Muon is "Shampoo without accumulation" (next subsection) applied to the momentum.

Update size. A full-rank semi-orthogonal $m \times n$ matrix has $\|O\|_F^2 = \min(m, n)$, so its root-mean-square (RMS) entry is $\sqrt{\min(m, n)/(mn)} = 1/\sqrt{\max(m, n)}$. AdamW's update RMS is usually about 0.2 to 0.4 (Liu et al., 2025). The factor $0.2\sqrt{\max(m, n)}$ makes Muon's update RMS equal to 0.2, so the learning rate and weight decay tuned for AdamW can be reused.

Which parameters. Muon is for the 2D weight matrices of hidden layers (a convolution kernel is flattened to 2D over its last three dimensions). Embeddings, the output layer, biases and normalization gains use AdamW. Jordan found AdamW better for the embedding and output layers even though they are 2D; biases and gains are vectors, so "orthogonalizing" them makes no sense (for a vector, $UV^\top$ is just $g/\|g\|$).

Shampoo #

Algorithm (Gupta, Koren and Singer, ICML 2018). For $W \in \R^{m\times n}$, start from $L_0 = \epsilon I_m$, $R_0 = \epsilon I_n$, and at each step

It stores $m^2 + n^2$ numbers instead of the $(mn)^2$ of full-matrix AdaGrad, and the inverse fourth roots cost $O(m^3 + n^3)$ (an eigendecomposition of each factor).

Why the power $-1/4$. Full-matrix AdaGrad preconditions $\operatorname{vec}(G_t)$ with $\big(\sum_s \operatorname{vec}(G_s)\operatorname{vec}(G_s)^\top\big)^{1/2}$. Gupta et al. bound the sum inside by a multiple of a Kronecker product of $L_t^{1/2}$ and $R_t^{1/2}$ (their Lemma 8). Its square root has the factors $L_t^{1/4}$ and $R_t^{1/4}$, one on each side of $G_t$.

Proposition 4 (Shampoo without accumulation is $UV^\top$). With $L = GG^\top$ and $R = G^\top G$, and the inverse roots taken on the range of $G$ (the limit $\epsilon \to 0$):

Proof.

1. $GG^\top = U\Sigma^2U^\top$, so $(GG^\top)^{-1/4} = U\Sigma^{-1/2}U^\top$ on the range of $U$.
2. In the same way, $(G^\top G)^{-1/4} = V\Sigma^{-1/2}V^\top$.
3. The product is $U\Sigma^{-1/2}\,(U^\top U)\,\Sigma\,(V^\top V)\,\Sigma^{-1/2}V^\top = U\Sigma^{-1/2}\Sigma\Sigma^{-1/2}V^\top = UV^\top$. $\blacksquare$

After $k$ identical gradients, $L = kGG^\top$ and $R = kG^\top G$, so the update is $k^{-1/4}\,k^{-1/4}\,UV^\top = UV^\top/\sqrt{k}$: the $1/\sqrt{t}$ step decay of AdaGrad. Practical versions (Distributed Shampoo; Shi et al., 2023) use moving averages for $L$ and $R$, recompute the roots only every few steps, and use "grafting": Shampoo gives the direction, and a simpler method such as Adam gives the step size of each layer.

Shampoo and K-FAC. For one example $\nabla_W = \delta a^\top$, so $\nabla_W\nabla_W^\top = \|a\|^2\,\delta\delta^\top$ and $\nabla_W^\top\nabla_W = \|\delta\|^2\,aa^\top$: K-FAC's factors, each weighted by the squared norm of the other side. Shampoo needs no extra backward pass with sampled targets, and it uses the power $-1/4$ instead of $-1$.

SOAP: Adam in Shampoo's eigenbasis #

Vyas et al. (2024) show that Shampoo with the power $1/2$ (that is, $L^{-1/2}GR^{-1/2}$ with averaged $L$ and $R$) is equivalent to Adafactor (an Adam variant whose second moment is factored into row and column statistics), run in the eigenbasis of $L$ and $R$. SOAP ("ShampoO with Adam in the Preconditioner's eigenbasis") runs full Adam in that basis:

1. Let $Q_L$ and $Q_R$ be the eigenvectors of $L$ and $R$ (moving averages of $GG^\top$ and $G^\top G$).
2. Rotate the gradient, $G' = Q_L^\top G Q_R$, and run Adam's moment updates on $G'$ (the first moment is kept in the original space and rotated the same way).
3. Rotate the Adam step back: $W \leftarrow W - \eta\,Q_L\,(\text{Adam step})\,Q_R^\top$. Refresh $Q_L$ and $Q_R$ every $f$ steps.

The preconditioning frequency $f$ is the only hyperparameter beyond Adam's. The intuition: Adam's diagonal scaling works best in a basis where the curvature is nearly diagonal, and Shampoo's eigenbasis is such a basis. For language models with 360M and 660M parameters at a large batch size, Vyas et al. report over 40% fewer iterations and over 35% less wall-clock time than AdamW, and about 20% better than Shampoo on both.

The evidence, as of 2025–2026 #

These are current results; they may change as methods and baselines improve.

• AlgoPerf (Kasimbeg et al., 2025): the winner of the external-tuning ruleset used Distributed Shampoo. Non-diagonal preconditioning beat Adam-type methods even in wall-clock time. (The self-tuning ruleset was won by Schedule-Free AdamW.)
• Speedruns (Jordan, 2024): Muon improved the NanoGPT speedrun record (3.28 validation loss on FineWeb) by a factor of 1.35, and the CIFAR-10 record (94% accuracy) from 3.3 to 2.6 A100-seconds.
• Moonlight (Liu et al., "Muon is Scalable for LLM Training", 2025): two changes were needed at scale, weight decay and the per-matrix update scale $0.2\sqrt{\max(m, n)}$. Their scaling-law runs showed about $2\times$ the compute efficiency of AdamW under compute-optimal training. Moonlight is a mixture-of-experts model with 3B active and 16B total parameters, trained on 5.7T tokens.
• Kimi K2 (2025): a mixture-of-experts model with 1T total and 32B active parameters, trained on 15.5T tokens with MuonClip (Muon plus "QK-Clip", which rescales the query and key weights of an attention head when its largest attention logit grows too large), with zero loss spikes.
• A careful comparison (Wen, Hall, Ma and Liang, "Fantastic Pretraining Optimizers and Where to Find Them", 2025): with equally careful tuning of every optimizer, and comparison at the end of training, the fastest methods were matrix-based (Muon, SOAP), but their speedup over AdamW fell from about $1.4\times$ at 0.1B parameters to about $1.1\times$ at 1.2B. Many earlier claims compared against a weakly tuned AdamW.

Answers to section 1. (1) Both become about 1: exactly 1 for $UV^\top$, and about 0.70 each for Muon's five quintic steps (the normalized values $0.99995$ and $0.0100$ end at $0.696$ and $0.699$). The weak direction now gets the same step as the strong one: this is the steepest step under the spectral norm, and for nearly low-rank updates it raises the rare directions. (2) The Fisher block of a layer with $mn = 10^6$ weights has $(10^6)^2 = 10^{12}$ entries. K-FAC's two factors have $1000^2 + 1000^2 = 2\cdot10^6$. (3) $(GG^\top)^{-1/4}G(G^\top G)^{-1/4} = UV^\top$ (Proposition 4): the same direction as Muon.

4. Worked example #

Take the gradient $G = \begin{pmatrix} 3 & 0 \\ 4 & 5 \end{pmatrix}$.

SVD. $G^\top G = \begin{pmatrix} 25 & 20 \\ 20 & 25 \end{pmatrix}$ has eigenvalues $45$ and $5$, so $\sigma_1 = 3\sqrt5 \approx 6.708$ and $\sigma_2 = \sqrt5 \approx 2.236$ (ratio 3). The right singular vectors are $v_1 = (1, 1)/\sqrt2$ and $v_2 = (1, -1)/\sqrt2$. Then $u_i = Gv_i/\sigma_i$:

Polar factor.

This is a rotation by $\arctan(1/2) \approx 26.57°$. With $P = V\Sigma V^\top = \sqrt5\begin{pmatrix} 2 & 1 \\ 1 & 2 \end{pmatrix}$ you can check $QP = G$.

Theorem 2 in numbers. $\|G\|_{\rm nuc} = 3\sqrt5 + \sqrt5 = 4\sqrt5 \approx 8.944$, and $\langle G, UV^\top\rangle = (3\cdot2 + 0\cdot(-1) + 4\cdot1 + 5\cdot2)/\sqrt5 = 20/\sqrt5 \approx 8.944$. Two other matrices with spectral norm 1 do worse:

• $G/\sigma_1$ gives $\|G\|_F^2/\sigma_1 = 50/(3\sqrt5) \approx 7.454$;
• $\operatorname{sign}(G)/\|\operatorname{sign}(G)\|_2$, with $\operatorname{sign}(G) = \begin{pmatrix} 1 & 0 \\ 1 & 1\end{pmatrix}$ and $\|\operatorname{sign}(G)\|_2 = (1 + \sqrt5)/2 \approx 1.618$, gives $12/1.618 \approx 7.416$.

Gradient descent keeps the ratio $3 : 1$ between the two singular directions; $UV^\top$ makes it $1 : 1$.

Newton–Schulz. $\|G\|_F = \sqrt{50} \approx 7.071$, so the normalized singular values are $(3/\sqrt{10},\ 1/\sqrt{10}) \approx (0.949, 0.316)$.

• Cubic: $p(0.949) = 0.996$ and $p(0.316) = 0.459$; then $(1.000, 0.640)$, $(1.000, 0.829)$, $(1.000, 0.958)$, $(1.000, 0.997)$. The largest error is $10^{-5}$ after 6 steps and about $10^{-10}$ after 7. The large value converges at once (quadratic convergence); the small one first grows by about 1.5 per step.
• Muon's quintic: $(0.949, 0.316) \to (0.752, 0.945) \to (1.048, 0.757)$, then $(0.682, 1.041) \to (1.134, 0.682) \to (0.753, 1.134)$. After 5 steps the update is $U\operatorname{diag}(0.753, 1.134)V^\top$. It is not orthogonal, and the strong direction now has the smaller value. But the ratio of the singular values went from 3 to 1.5.

Shampoo. $GG^\top = \begin{pmatrix} 9 & 12 \\ 12 & 41 \end{pmatrix}$ also has eigenvalues 45 and 5, and $(GG^\top)^{-1/4}G(G^\top G)^{-1/4} = UV^\top$ exactly. After 4 identical gradients the update is $UV^\top/2$.

Update RMS. For an MLP weight with $m = 3072$ rows and $n = 768$ columns, $UV^\top$ has RMS $1/\sqrt{3072} \approx 0.018$. Jordan's factor $\sqrt{3072/768} = 2$ gives $0.036$; Moonshot's factor $0.2\sqrt{3072} \approx 11.1$ gives $0.2$.

K-FAC in one dimension. With one input and one output, the Fisher block is the number $\mathbb{E}[a^2\delta^2]$. For two examples $(a, \delta) = (1, 2)$ and $(2, 1)$, the exact value is $\tfrac12(1\cdot4 + 4\cdot1) = 4$, but K-FAC gives $\mathbb{E}[a^2]\,\mathbb{E}[\delta^2] = 2.5 \cdot 2.5 = 6.25$. Here large $a$ comes with small $\delta$, so $a^2$ and $\delta^2$ are negatively correlated, and the product of the means is too large.

5. In code #

The notebook asks you to implement Newton–Schulz, Shampoo's inverse roots and K-FAC yourself. Here NumPy and SciPy check the hand results: scipy.linalg.polar gives $UV^\top$, and fractional_matrix_power gives the inverse fourth roots of Shampoo.

import numpy as np
from scipy.linalg import polar, fractional_matrix_power

G = np.array([[3.0, 0.0], [4.0, 5.0]])
U, s, Vt = np.linalg.svd(G)
Q, P = polar(G)                                       # G = Q P with Q = U Vᵀ, P = V Σ Vᵀ
print(s.round(4), np.allclose(Q, U @ Vt))             # [6.7082 2.2361] True
print(np.sqrt(5) * Q)                                 # [[ 2. -1.] [ 1.  2.]]

# Theorem 2: among matrices with spectral norm 1, U Vᵀ has the largest <G, O> = ‖G‖_nuc
for name, O in [("UV^T", Q), ("G/sigma_max", G / s[0]),
                ("sign(G)", np.sign(G) / np.linalg.norm(np.sign(G), 2))]:
    print(f"{name:12s} <G,O> = {np.sum(G * O):.4f}")
print(f"nuclear norm {s.sum():.4f}")                  # 8.9443; the three rows: 8.9443, 7.4536, 7.4164

def newton_schulz(G, coeffs, steps):
    a, b, c = coeffs
    X = G / np.linalg.norm(G)                         # Frobenius norm: all singular values ≤ 1
    tall = X.shape[0] > X.shape[1]
    if tall:
        X = X.T                                       # work with the smaller Gram matrix X Xᵀ
    for _ in range(steps):
        A = X @ X.T
        X = a * X + (b * A + c * A @ A) @ X
    return X.T if tall else X

cubic, quintic = (1.5, -0.5, 0.0), (3.4445, -4.7750, 2.0315)
print(np.abs(newton_schulz(G, cubic, 7) - Q).max())                         # about 1e-10
print(np.linalg.svd(newton_schulz(G, quintic, 5), compute_uv=False).round(3))  # [1.134 0.753]

# Shampoo without accumulation gives U Vᵀ; after k equal gradients it gives U Vᵀ/√k
for k in [1, 4]:
    S = fractional_matrix_power(k * G @ G.T, -0.25) @ G @ fractional_matrix_power(k * G.T @ G, -0.25)
    print(k, np.allclose(S, Q / np.sqrt(k)))          # 1 True, then 4 True

# K-FAC for one example: vec(δ aᵀ) = a ⊗ δ (columns stacked)
rng = np.random.default_rng(0)
a, delta = rng.standard_normal(3), rng.standard_normal(2)
v = np.outer(delta, a).flatten(order="F")
print(np.allclose(np.outer(v, v), np.kron(np.outer(a, a), np.outer(delta, delta))))   # True
# the cheap step: (A ⊗ D)^{-1} vec(M) = vec(D^{-1} M A^{-1})
A_f, D_f, M = np.diag([1.0, 2.0, 3.0]), np.array([[2.0, 1.0], [1.0, 2.0]]), rng.standard_normal((2, 3))
lhs = np.linalg.solve(np.kron(A_f, D_f), M.flatten(order="F"))
print(np.allclose(lhs, (np.linalg.inv(D_f) @ M @ np.linalg.inv(A_f)).flatten(order="F")))  # True
# two examples in 1D: mean of products vs product of means
a2, d2 = np.array([1.0, 2.0]), np.array([2.0, 1.0])
print(np.mean(a2**2 * d2**2), np.mean(a2**2) * np.mean(d2**2))   # 4.0 6.25

# a 64 x 256 gradient whose singular values span a factor of 100
U2, _ = np.linalg.qr(rng.standard_normal((64, 64)))
V2, _ = np.linalg.qr(rng.standard_normal((256, 64)))
G2 = U2 @ np.diag(np.geomspace(1.0, 0.01, 64)) @ V2.T
O2 = newton_schulz(G2, quintic, 5)
sv = np.linalg.svd(O2, compute_uv=False)
print(sv.min().round(2), sv.max().round(2))           # 0.68 1.2: a band around 1
update = 0.2 * np.sqrt(max(64, 256)) * O2             # Moonshot's shape factor
print(round(np.sqrt(np.mean(update**2)), 3))          # 0.184, close to AdamW's typical 0.2

The update RMS is 0.184 and not exactly 0.2, because the singular values are only near 1. Real Muon runs the iteration in bfloat16, which moves the last digits a little.

6. Where it shows up #

• Large-model pretraining. PyTorch ships torch.optim.Muon with three shape factors (adjust_lr_fn = "original", "match_rms_adamw" or "spectral_unclamped"). A typical setup has two optimizers: Muon for the hidden matrices and AdamW for everything else. When you compare it with AdamW, tune the AdamW baseline just as carefully (Wen et al., 2025).
• Distributed training. Muon keeps one state matrix per weight (the momentum) instead of AdamW's two, so it needs half of AdamW's optimizer-state memory. With sharded optimizer states (ZeRO, FSDP), each matrix must be gathered before Newton–Schulz; Moonshot's distributed Muon does this.
• Curvature for analysis. Influence functions for large language models use K-FAC-type Hessian approximations (EK-FAC; Grosse et al., 2023). Laplace approximations of the posterior of a network often use K-FAC factors.
• ML for science. For physics-informed networks, Dangel, Müller and Zeinhofer (NeurIPS 2024) extend K-FAC to losses with differential operators. Transformer foundation models for weather or seismic data have the same parameter types as language models, so the same Muon rules apply. Small full-batch physics problems are still served better by curvature methods (L-BFGS, Gauss–Newton; lesson 10).

7. Common mistakes #

• Muon on every parameter. Embeddings, the output layer, biases and gains belong to AdamW.
• No normalization before Newton–Schulz. For a singular value above $\sqrt3$, the cubic iteration does not converge to 1, and the quintic is tuned for $[0, 1]$. Divide by the Frobenius norm first.
• Expecting an exact $UV^\top$ from Muon. Five quintic steps give singular values in a band around 1, on purpose.
• Reusing AdamW's learning rate with Jordan's shape factor. The update RMS is then $\sqrt{\max(1, m/n)}/\sqrt{\max(m, n)}$ (0.036 for a $3072 \times 768$ weight), far below AdamW's 0.2. Use the $0.2\sqrt{\max(m, n)}$ factor, or retune the learning rate.
• Mixing up the Kronecker order. With column stacking, $\operatorname{vec}(\delta a^\top) = a \otimes \delta$ and the K-FAC step is $D^{-1}\nabla_W A^{-1}$. With row stacking, the factors swap. Fix one convention and test it against np.kron on a small case.
• K-FAC with the data labels (the empirical Fisher) or without damping.
• Calling Shampoo "K-FAC". Both use two Kronecker factors, but with different statistics (gradient matrices vs activations and output gradients) and different powers ($-1/4$ vs $-1$).

8. Check yourself #

1. What is the steepest-descent step under the $\ell_\infty$ norm for $g = (3, -1, 0.5)$ and $L = 1$?
   Answer

   The dual norm is $\ell_1$: $\|g\|_1 = 4.5$, and the unit direction is $\operatorname{sign}(g) = (1, -1, 1)$. By Proposition 1, $\Delta = -4.5\,(1, -1, 1)$.
2. $G = \operatorname{diag}(5, -2)$. What are $UV^\top$ and the Shampoo update without accumulation?
   Answer

   Both are $\operatorname{diag}(1, -1) = \operatorname{sign}(G)$. For example, $GG^\top = \operatorname{diag}(25, 4)$, so $(GG^\top)^{-1/4}G(G^\top G)^{-1/4} = \operatorname{diag}(5^{-1/2}\cdot5\cdot5^{-1/2},\ 2^{-1/2}\cdot(-2)\cdot2^{-1/2}) = \operatorname{diag}(1, -1)$.
3. A single example passes through a linear layer, so $G = \delta a^\top$ with nonzero vectors $\delta$ and $a$. What is $UV^\top$, and what does Muon do here?
   Answer

   The reduced SVD has one singular value $\sigma = \|\delta\|\|a\|$, with $u = \delta/\|\delta\|$ and $v = a/\|a\|$. So $UV^\top = \delta a^\top/(\|\delta\|\|a\|) = G/\|G\|_F$. For a rank-1 gradient, orthogonalization is just normalized gradient descent. (Because $G$ is not full rank, the steepest direction is not unique; $UV^\top$ from the reduced SVD is one solution.)
4. Shampoo receives the same gradient $G$ at every step. What is the update after $k$ steps, and which older method has the same step decay?
   Answer

   $L = kGG^\top$ and $R = kG^\top G$ (ignoring $\epsilon$), so the update is $k^{-1/4}\cdot k^{-1/4}\,UV^\top = UV^\top/\sqrt{k}$. AdaGrad also gives steps that decay like $1/\sqrt{t}$ for a constant gradient.
5. Why does the cubic iteration need many steps for a singular value of $0.01$, how does Muon's quintic fix this, and what does it give up?
   Answer

   The cubic map has slope $p'(0) = 1.5$, so small values grow only 1.5 times per step (14 steps from 0.01 to 0.99). The quintic has slope 3.4445 at 0 and reaches the band around 1 in 5 steps. The price: $p(1) = 0.701$, so 1 is not a fixed point, and the singular values stay in a band of about $[0.68, 1.13]$ instead of converging to 1.
6. A linear layer has $n = 1024$ inputs and $m = 4096$ outputs. Compare the cost of inverting its exact Fisher block with K-FAC's two inversions.
   Answer

   The exact block has size $mn = 4{,}194{,}304$, so about $(mn)^3 \approx 7.4\times10^{19}$ operations. K-FAC inverts a $1024\times1024$ and a $4096\times4096$ matrix: $1024^3 + 4096^3 \approx 7.0\times10^{10}$, about $10^9$ times less.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Define the dual norm of a norm $\|\cdot\|$.

   Show answer

2. What is the steepest-descent direction under the spectral norm for a gradient $G = U\Sigma V^\top$?

   Show answer

3. What is the steepest unit step for $g = (3, 1)$ under $\ell_\infty$?

   Show answer

4. What is the purpose of the Newton–Schulz iteration in Muon?

   Show answer

5. Describe Muon in one sentence.

   Show answer

6. Write the Shampoo update.

   Show answer

7. What is Shampoo without accumulation equal to?

   Show answer

8. Describe SOAP in one sentence.

   Show answer

9. What is the K-FAC approximation of the Fisher block of a linear layer, and what is its step?

   Show answer

10. Why is K-FAC exact for a single example?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "What is Muon, and why orthogonalize?" — Momentum SGD where the update of each 2D hidden matrix is orthogonalized (about $UV^\top$) by five Newton–Schulz steps; AdamW handles the rest. It is steepest descent under the spectral norm (Theorem 2), so the rare directions of a nearly low-rank update get a real step.
• "How are Muon and Shampoo related?" — Shampoo without accumulation is $(GG^\top)^{-1/4}G(G^\top G)^{-1/4} = UV^\top$; Muon applies this to the momentum and computes it with Newton–Schulz.
• "What does K-FAC approximate, and why is it cheap?" — The layer's Fisher block $\mathbb{E}[aa^\top \otimes \delta\delta^\top]$ by $A \otimes D$; the step $D^{-1}\nabla_W A^{-1}$ needs two small inversions instead of one huge one.
• "How do you switch an AdamW recipe to Muon?" — Muon for 2D hidden weights only, the $0.2\sqrt{\max(m, n)}$ scale to reuse AdamW's learning rate and weight decay, keep the weight decay, watch the attention logits.
• "Does Muon beat AdamW?" — Speedruns, Moonlight (about $2\times$ compute efficiency) and Kimi K2 say yes; a careful benchmark with equally tuned baselines finds about $1.4\times$ at 0.1B parameters, shrinking to $1.1\times$ at 1.2B (as of 2025–2026).

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Check numerically that $UV^\top$ maximizes $\langle G, O\rangle$ over spectral-norm-1 matrices (random search against the nuclear norm).
2. Implement the Newton–Schulz iteration (cubic and quintic) and compare with the SVD polar factor.
3. Compute matrix inverse fourth roots by eigendecomposition and check that Shampoo without accumulation gives $UV^\top$.
4. K-FAC for one linear layer: compute $A$, $D$ and the exact Fisher block from samples, and the relative error of $A \otimes D$.
5. Muon vs Adam vs gradient descent on a small matrix regression problem with an ill-conditioned input.

11. Watch and read #

This Simple Optimizer Is Revolutionizing How We Train AI | MuonJia-Bin Huang · 18 min

Orthogonalization with the SVD, the Newton–Schulz iteration as an odd matrix polynomial, Muon, and MuonClip for attention logits.

Jeremy Bernstein - Depths of First Order OptimizationJeremy Bernstein · Cohere · 48 min

Steepest descent under a norm, and how the spectral norm leads to Shampoo and Muon (section 3).

2nd-order Optimization for Neural Network TrainingJames Martens · Microsoft Research · 87 min

Hessian-free optimization and K-FAC, by the author of both methods.

Nikhil Vyas - SOAP: Improving and Stabilizing Shampoo using AdamNikhil Vyas · Cohere · 42 min

Why Shampoo is Adafactor in a rotated basis, and how SOAP runs Adam in that basis.

• Bernstein and Newhouse, Old Optimizer, New Norm: An Anthology (2024): Proposition 1 (steepest descent), Propositions 4–6 (Shampoo and the spectral norm), Table 1.
• Jeremy Bernstein, Deriving Muon (blog).
• Keller Jordan, Muon: An optimizer for hidden layers in neural networks (2024): the Newton–Schulz coefficients, the cost analysis, the results.
• Martens and Grosse, Optimizing Neural Networks with Kronecker-factored Approximate Curvature (2015): §3 (the approximation), §5 (statistics with sampled targets), §6.3 (damping).
• Gupta, Koren and Singer, Shampoo: Preconditioned Stochastic Tensor Optimization (2018), and Shi et al., Distributed Shampoo (2023).
• Vyas et al., SOAP: Improving and Stabilizing Shampoo using Adam (2024).
• Liu et al., Muon is Scalable for LLM Training (2025), and Kimi Team, Kimi K2: Open Agentic Intelligence (2025).
• Kasimbeg et al., Accelerating Neural Network Training: An Analysis of the AlgoPerf Competition (2025), and Wen et al., Fantastic Pretraining Optimizers and Where to Find Them (2025).
• PyTorch, torch.optim.Muon (the algorithm box and the three shape factors).