Muon

Every optimizer we have built, from gradient descent through momentum to Adam, treats the parameters as one flat vector: a weight matrix \(\mathbf{W}\in\mathbb{R}^{d\times k}\) becomes a list of \(dk\) numbers before any update touches it. Last lecture accepted that reshaping and asked which entries we may move: LoRA froze \(\mathbf{W}\) and trained a rank-\(r\) correction. Muon keeps every entry trainable and examines the effect of treating a weight matrix as a linear map rather than a flat vector.

A weight matrix is a linear map, but ordinary gradient descent measures its step after flattening the matrix into a vector. The resulting algorithm depends on that choice of norm.

Steepest Descent Depends on the Norm

Recall the justification for the gradient direction from the Optimization lecture: among all steps of a fixed size, the negative gradient decreases the loss fastest. For a matrix parameter, the choice of norm enters the steepest-descent calculation directly. Write \(\mathbf{G}\) for the gradient of the loss with respect to the weight matrix at the current weights: \[ \mathbf{G} = \nabla_\mathbf{W}\mathcal{L}(\mathbf{W}) \in \mathbb{R}^{d\times k} . \] It has one entry for each entry of \(\mathbf{W}\), so the first-order change in the loss caused by a step \(\mathbf{\Delta}\in\mathbb{R}^{d\times k}\) is the ordinary entrywise dot product of the two matrices: \[ \mathcal{L}(\mathbf{W}+\mathbf{\Delta}) \approx \mathcal{L}(\mathbf{W}) + \langle\mathbf{G},\mathbf{\Delta}\rangle , \quad\text{where}\quad \langle\mathbf{G},\mathbf{\Delta}\rangle = \sum_{i=1}^d\sum_{j=1}^k [\mathbf{G}]_{i,j}\,[\mathbf{\Delta}]_{i,j} . \] Steepest descent is now a small optimization problem of its own, with the learning rate \(\alpha>0\) as a budget on the step’s size: \[ \mathbf{\Delta} = \arg\min_{\|\mathbf{\Delta}\|\,\leq\,\alpha} \ \langle\mathbf{G},\mathbf{\Delta}\rangle . \] Everything depends on which norm \(\|\cdot\|\) sits in that constraint, and nothing above chooses it for us.

Flattening the matrix chooses the Frobenius norm, the Euclidean length of the \(dk\) entries laid out in a row: \[ \|\mathbf{M}\|_F = \sqrt{\sum_{i=1}^d\sum_{j=1}^k [\mathbf{M}]_{i,j}^2} = \sqrt{\sum_i\sigma_i^2} , \] where \(\sigma_1\geq\sigma_2\geq\cdots\) are the singular values of \(\mathbf{M}\) and the second equality is the identity from the Low-rank Approximation lecture. Under this budget, the Cauchy–Schwarz inequality applied to the two matrices read as vectors of length \(dk\) bounds how negative the inner product can get: \[ \langle\mathbf{G},\mathbf{\Delta}\rangle \ \geq\ -\|\mathbf{G}\|_F\,\|\mathbf{\Delta}\|_F \ \geq\ -\alpha\,\|\mathbf{G}\|_F , \] with equality exactly when \(\mathbf{\Delta}\) points opposite \(\mathbf{G}\) and has norm \(\alpha\): \[ \mathbf{\Delta}^{\mathrm{frob}} = -\alpha\,\frac{\mathbf{G}}{\|\mathbf{G}\|_F} . \] That is plain gradient descent, with the learning rate absorbing the normalization, so every optimizer we have written down so far was solving this problem without saying so.

A linear map has a size of its own, independent of how its entries are stored: the largest factor by which it stretches a vector. That is the spectral norm from the Depth-enablers lecture, which the Decompositions lecture identified with the top singular value: \[ \|\mathbf{M}\|_2 = \max_{\mathbf{x}\neq\mathbf{0}}\frac{\|\mathbf{M}\mathbf{x}\|_2}{\|\mathbf{x}\|_2} = \sigma_1 . \] The Frobenius norm adds up the squares of every singular value; the spectral norm keeps only the largest. Spreading a step’s total movement evenly over fifty directions instead of one leaves its Frobenius norm unchanged but makes each singular value \(\sqrt{50}\) times smaller than the single value it replaced, so under the spectral norm the spread-out step costs about seven times less.

To choose between the norms, consider what a change to \(\mathbf{W}\) does to the layer that contains it. The layer maps an input \(\mathbf{x}\in\mathbb{R}^k\) to \(\mathbf{W}\mathbf{x}\in\mathbb{R}^d\), so changing the weights by \(\mathbf{\Delta}\) moves that output by \[ \big\|(\mathbf{W}+\mathbf{\Delta})\mathbf{x} - \mathbf{W}\mathbf{x}\big\|_2 = \|\mathbf{\Delta}\mathbf{x}\|_2 \leq \|\mathbf{\Delta}\|_2\,\|\mathbf{x}\|_2 , \] where the inequality is the definition of the spectral norm. A spectral-norm budget guarantees that no input’s output moves by more than \(\alpha\) times that input’s length. A Frobenius budget can be small in aggregate while changing one direction substantially. That is the concern the Depth-enablers lecture raised about products of layer Jacobians, asked now about the update rather than the weights.

The Optimal Update Is the Orthogonalized Gradient

Claim: Let \(\mathbf{G} = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top = \sum_{i=1}^r \sigma_i\mathbf{u}_i\mathbf{v}_i^\top\) be the SVD of the gradient, with rank \(r\), orthonormal columns in \(\mathbf{U}\in\mathbb{R}^{d\times r}\) and \(\mathbf{V}\in\mathbb{R}^{k\times r}\), and the nonzero singular values on the diagonal of \(\mathbf{\Sigma}\in\mathbb{R}^{r\times r}\). Then the steepest-descent step under a spectral-norm budget of \(\alpha\) is \[ \mathbf{\Delta}^{\mathrm{spec}} = -\alpha\,\mathbf{U}\mathbf{V}^\top , \] and it decreases the linearized loss by \(\alpha\sum_{i=1}^r\sigma_i\).

Proof of Claim

Substitute the outer-product form of \(\mathbf{G}\) into the inner product and pull the sum out: \[ \begin{align*} \langle\mathbf{G},\mathbf{\Delta}\rangle &= \Big\langle \sum_{i=1}^r \sigma_i\mathbf{u}_i\mathbf{v}_i^\top,\ \mathbf{\Delta}\Big\rangle \\&= \sum_{i=1}^r \sigma_i\,\big\langle \mathbf{u}_i\mathbf{v}_i^\top,\ \mathbf{\Delta}\big\rangle \\&= \sum_{i=1}^r \sigma_i\ \mathbf{u}_i^\top\mathbf{\Delta}\mathbf{v}_i , \end{align*} \] where the second equality is linearity of the inner product and the third writes out the entrywise definition, \(\langle\mathbf{a}\mathbf{b}^\top,\mathbf{\Delta}\rangle = \sum_{i,j}a_ib_j[\mathbf{\Delta}]_{i,j} = \mathbf{a}^\top\mathbf{\Delta}\mathbf{b}\). Now bound one term at a time. Since \(\mathbf{u}_i\) and \(\mathbf{v}_i\) are unit vectors, \[ \big|\mathbf{u}_i^\top\mathbf{\Delta}\mathbf{v}_i\big| \leq \|\mathbf{u}_i\|_2\,\|\mathbf{\Delta}\mathbf{v}_i\|_2 \leq \|\mathbf{\Delta}\|_2 \leq \alpha , \] by Cauchy–Schwarz, then the definition of the spectral norm, then the budget. Every \(\sigma_i\) is nonnegative, so the whole sum cannot fall below \(-\alpha\sum_i\sigma_i\). It remains to exhibit a step that attains that value, and \(\mathbf{\Delta} = -\alpha\mathbf{U}\mathbf{V}^\top\) does. It is admissible: \(\mathbf{U}\mathbf{V}^\top\) sends each \(\mathbf{v}_i\) to the unit vector \(\mathbf{u}_i\) and sends everything orthogonal to all the \(\mathbf{v}_i\) to \(\mathbf{0}\), so it stretches no vector by more than \(1\) and \(\|\mathbf{\Delta}\|_2 = \alpha\). And it attains the bound, orthonormality gives \(\mathbf{u}_i^\top(\mathbf{U}\mathbf{V}^\top)\mathbf{v}_i = \mathbf{u}_i^\top\mathbf{u}_i = 1\) for every \(i\), so \(\langle\mathbf{G},\mathbf{\Delta}\rangle = -\alpha\sum_i\sigma_i\).

The step keeps the gradient’s singular vectors and discards its singular values: every direction the gradient uses at all gets pushed exactly as hard as every other, and the directions it ignores, where \(\sigma_i = 0\), are not pushed at all. This matrix is the polar factor of \(\mathbf{G}\), and computing it is called orthogonalizing the gradient. Compare the two decreases we have derived, \(\alpha\sum_i\sigma_i\) for the orthogonalized step against \(\alpha\sqrt{\sum_i\sigma_i^2}\) for the Frobenius one: the orthogonalized step collects something from every direction, while the raw gradient’s own length is dominated by its top few.

In the plot, the teal bars are one matrix’s singular values, rescaled so the largest is \(1\), and the black dashed line is where orthogonalization sends all of them. The shaded band is the interval \([\ell,1]\) that the rescaled singular values occupy, with \(\ell\) the smallest of them; the band is wide here because this matrix has condition number \(\kappa = \sigma_1/\sigma_r = 1/\ell \approx 38\), the Optimization lecture’s measure of how stretched a problem is. Collapsing that band onto the dashed line is all orthogonalization does, and it is all Muon does.

A nearly rank-one gradient and a perfectly flat one with the same singular vectors receive the identical update. Orthogonalization discards the gradient’s weighting of the directions. A budget on the worst-case stretch is already fully spent by the top direction, so holding the others back makes the step no cheaper and the loss decrease smaller. The Optimization lecture addressed poor conditioning by rescaling features. Here the update is reshaped so that its singular directions have equal magnitude.

Computing \(\mathbf{U}\) and \(\mathbf{V}\) with an SVD at every training step is impractical.

Cheap Orthogonalization via Newton–Schulz

The case for avoiding the SVD is not asymptotic, since an SVD and a handful of same-shape matrix multiplications both cost \(O(dk\min(d,k))\); it is that matrix multiplication is the one operation accelerators run at full throughput and at the low precision training already uses, while an SVD parallelizes badly and wants high precision. So we want the polar factor expressed in matrix multiplications and nothing else.

The Newton–Schulz iteration does that. Rescale the gradient so that all of its singular values are at most \(1\), which dividing by the Frobenius norm accomplishes since \(\sigma_1\leq\|\mathbf{G}\|_F\): \[ \mathbf{X}^{(0)} = \frac{\mathbf{G}}{\|\mathbf{G}\|_F} . \] Write \(\ell>0\) for the smallest nonzero singular value of \(\mathbf{X}^{(0)}\), as in the spectrum figure, so that every singular value of \(\mathbf{X}^{(0)}\) is either \(0\) or lies in the interval \([\ell,1]\). Then repeat the cubic step \[ \mathbf{X}^{(t+1)} = \frac32\mathbf{X}^{(t)} - \frac12\,\mathbf{X}^{(t)}\big(\mathbf{X}^{(t)}\big)^\top\mathbf{X}^{(t)} , \] which is two matrix multiplications and a weighted sum, nothing else.

The following fact, proved in class, shows how a cubic in \(\mathbf{X}\) acts on its singular values.

Claim: Let \(q(x) = c_1x + c_3x^3 + c_5x^5 + \cdots\) be a polynomial with only odd powers, and define its matrix version by replacing each \(x^{2m+1}\) with \(\mathbf{X}(\mathbf{X}^\top\mathbf{X})^m\): \[ q(\mathbf{X}) = c_1\mathbf{X} + c_3\,\mathbf{X}\mathbf{X}^\top\mathbf{X} + c_5\,\mathbf{X}\mathbf{X}^\top\mathbf{X}\mathbf{X}^\top\mathbf{X} + \cdots . \] If \(\mathbf{X} = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top\), then \(q(\mathbf{X}) = \mathbf{U}\,q(\mathbf{\Sigma})\,\mathbf{V}^\top\): the singular vectors come through untouched and \(q\) is applied to each singular value on its own.

Proof of Claim

Every term is \(\mathbf{X}(\mathbf{X}^\top\mathbf{X})^m\) for some \(m\geq0\) (the \(m=0\) term is \(\mathbf{X}\) itself), so it is enough to handle one such term. Substitute the SVD and collapse the interior with orthonormality of the singular vectors: \[ \begin{align*} \mathbf{X}^\top\mathbf{X} &= \mathbf{V}\mathbf{\Sigma}\mathbf{U}^\top\mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top = \mathbf{V}\mathbf{\Sigma}^2\mathbf{V}^\top \\ (\mathbf{X}^\top\mathbf{X})^m &= \mathbf{V}\mathbf{\Sigma}^{2m}\mathbf{V}^\top \\ \mathbf{X}(\mathbf{X}^\top\mathbf{X})^m &= \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top\mathbf{V}\mathbf{\Sigma}^{2m}\mathbf{V}^\top = \mathbf{U}\mathbf{\Sigma}^{2m+1}\mathbf{V}^\top , \end{align*} \] using \(\mathbf{U}^\top\mathbf{U} = \mathbf{I}\) on the first line, then \(\mathbf{V}^\top\mathbf{V} = \mathbf{I}\) repeatedly to multiply out the \(m\)th power on the second, and once more on the third. Adding the terms back up with their coefficients, \[ q(\mathbf{X}) = \mathbf{U}\Big(\sum_m c_{2m+1}\mathbf{\Sigma}^{2m+1}\Big)\mathbf{V}^\top = \mathbf{U}\,q(\mathbf{\Sigma})\,\mathbf{V}^\top , \] and \(\mathbf{\Sigma}\) is diagonal, so \(q(\mathbf{\Sigma})\) is the diagonal matrix whose \(i\)th entry is \(q(\sigma_i)\).

The Newton–Schulz step is this claim applied to one particular cubic: \[ p(x) = \frac{3x-x^3}{2} . \] The iteration on matrices is an iteration on numbers, run in parallel on every singular value at once, so we never compute \(\mathbf{U}\) or \(\mathbf{V}\) and never need to. All we need is a polynomial that sends every number in \([\ell,1]\) to \(1\), and sends \(0\) to \(0\) so that unused directions stay unused.

Fixed Points and Convergence

We first analyze where \(p\) sends a scalar. The fixed points solve \(x = p(x)\): \[ x = \frac{3x-x^3}{2} \iff 2x = 3x-x^3 \iff x^3-x = 0 \iff x(x-1)(x+1) = 0 , \] so \(p\) fixes \(x=-1\), \(x=0\), and \(x=1\), and nothing else. Two of those are exactly the values we want held in place: a singular value already at \(1\) stays put, and a direction the gradient never uses stays at \(0\). The question is what happens to everything in between. Start at a small singular value \(x^{(0)}=0.1\) and iterate twice: \[ x^{(1)} = \frac{3(0.1)-(0.1)^3}{2} = \frac{0.3-0.001}{2} = 0.1495, \qquad x^{(2)} = \frac{3(0.1495)-(0.1495)^3}{2} \approx 0.2226 . \] Each step multiplied the number by about \(3/2\), which is what the cubic does while \(x^3\) is negligible next to \(3x\), since \(p(x)\approx\frac32x\) near zero. As \(x\) approaches \(1\), the cubic term cancels the excess exactly and pins the value there: small singular values climb slowly at first and only snap into place once they are close to \(1\), which is why a badly conditioned gradient needs more Newton–Schulz iterations.

Problem 18 runs the design in reverse. Starting only from the class of odd cubics and requirements for safety on \([0,1]\), it asks whether the coefficient \(3/2\) is conventional or forced by maximizing progress on the smallest singular values. Practical implementations do better by changing the polynomial as they go: at each step, pick the odd polynomial that lifts the current worst point \(\ell_t\) as high as a bounded-degree polynomial can. That is the Polar Express of Amsel, Persson, Musco, and Gower, the orthogonalization routine Muon implementations use in place of a fixed cubic.

Matrix functions

The last two sections applied a scalar function to a matrix’s singular values and left its singular vectors alone, the move this whole unit has been making. The Low-rank Approximation lecture’s truncated SVD is \(f(\sigma)=\sigma\) for the top \(k\) singular values and \(f(\sigma)=0\) for the rest. The pseudoinverse, from the Optimization lecture, is \(f(\sigma)=1/\sigma\) for every \(\sigma>0\) and \(f(0)=0\). Today’s orthogonalization is \(f(\sigma)=1\) for every \(\sigma>0\) and \(f(0)=0\), which is \(\mathrm{sign}\) on the nonnegative numbers. The same construction gives three operations through different choices of \(f\): \[ f\Big(\sum_i \sigma_i\mathbf{u}_i\mathbf{v}_i^\top\Big) = \sum_i f(\sigma_i)\,\mathbf{u}_i\mathbf{v}_i^\top . \] The singular vectors remain fixed while \(f\) transforms the singular values: truncation sets small values to zero, the pseudoinverse inverts nonzero values, and orthogonalization maps every nonzero value to \(1\). What is left is to package it as an optimizer we can actually run.

Muon in Practice

Muon puts the orthogonalized step together with momentum from the Gradient Descent lecture. One step, for one weight matrix \(\mathbf{W}\):

1. Compute the minibatch gradient \(\mathbf{G}^{(t)}\) and fold it into a momentum buffer, \(\mathbf{M}^{(t)} = \beta\mathbf{M}^{(t-1)} + \mathbf{G}^{(t)}\), with \(\beta\) near \(0.9\) as usual.
2. Orthogonalize the momentum rather than the raw gradient: run a few Newton–Schulz steps starting from \(\mathbf{M}^{(t)}/\|\mathbf{M}^{(t)}\|_F\) to get an approximate polar factor \(\mathbf{O}^{(t)}\).
3. Take the step \(\mathbf{W}^{(t+1)} = \mathbf{W}^{(t)} - \alpha\,\mathbf{O}^{(t)}\).

Five Newton–Schulz steps is the usual choice, so the added cost is about ten matrix multiplications per weight matrix per step. It is the momentum that gets orthogonalized because momentum is the direction we actually intend to move; averaging first also keeps a single noisy minibatch from deciding which directions count as used.

Muon applies only to parameters that are genuinely matrices; biases and gains, which are not linear maps, are left to Adam. Since \(\|\mathbf{O}^{(t)}\|_2\approx1\) by construction, every Muon step has spectral norm about \(\alpha\) whatever the gradient’s size, so the learning rate sets the step size directly instead of only scaling it.

A deep network’s end-to-end behavior depends on a product of many weight matrices, whose conditioning degrades when individual matrices develop a few dominant singular directions. An update whose own singular values are all \(1\) moves every direction it touches by the same amount, so it does not concentrate a weight matrix’s spectrum the way a raw gradient step does.

The class demo trains the same small network on MNIST two ways, Adam against Muon, changing nothing but the optimizer. The orthogonalized step drives the training loss down considerably faster, while test accuracy moves less: optimizing faster and generalizing better turn out to be separate questions.

Muon follows the polar factor of a matrix gradient rather than the raw gradient. The next unit studies a different use of matrix structure: convolution and attention constrain the architecture before optimization begins.