Muon

A neural network’s weight matrices turn input vectors into output vectors. When we update a matrix, we change what it does to all those inputs. We’d like to decrease the training loss while keeping that change small. But there is more than one way to measure a small change to a matrix.

Gradient descent uses the sum of squared changes to the entries. Muon uses a different geometry: it limits how much the update can change the output for any unit input. We can derive both directions from the same local optimization problem. The difference comes from the norm we put on the update.

How should a matrix’s action on its inputs shape the update we take?

We’ll see what the two norms allow, derive the resulting updates, and use polynomial approximation to compute Muon’s direction efficiently.

Choosing a Small Update

Let \(\mathbf W\in\mathbb R^{d\times k}\) be a layer’s weight matrix. For a column-vector input \(\mathbf x\in\mathbb R^k\), the layer produces \(\mathbf W\mathbf x\in\mathbb R^d\). An update \(\boldsymbol\Delta\in\mathbb R^{d\times k}\) changes that output by

\[ (\mathbf W+\boldsymbol\Delta)\mathbf x-\mathbf W\mathbf x=\boldsymbol\Delta\mathbf x. \]

So we can measure the update either through its entries or through what it does to inputs. The Frobenius norm adds up the squared entry changes:

\[ \|\boldsymbol\Delta\|_F^2=\sum_{i=1}^d\sum_{j=1}^k\Delta_{ij}^2. \]

This also equals the sum of the squared singular values. It gives the update one shared budget across its input directions. Spending more in one direction leaves less for the others.

The spectral norm instead measures the largest output change over unit inputs:

\[ \|\boldsymbol\Delta\|_2=\max_{\|\mathbf x\|_2=1}\|\boldsymbol\Delta\mathbf x\|_2=\sigma_1(\boldsymbol\Delta). \]

A spectral budget \(\|\boldsymbol\Delta\|_2\leq\rho\) guarantees that the output changes by at most \(\rho\|\mathbf x\|_2\) for every input. Several perpendicular directions can each change by this amount without exceeding the budget. The Frobenius norm would add their squared changes together.

Now use the loss to choose among the permitted updates. Write \(\mathbf G=\nabla_{\mathbf W}\mathcal L\) for the gradient. The first-order approximation is

\[ \mathcal L(\mathbf W+\boldsymbol\Delta)\approx\mathcal L(\mathbf W)+\langle\mathbf G,\boldsymbol\Delta\rangle_F,\qquad \langle\mathbf G,\boldsymbol\Delta\rangle_F=\sum_{i=1}^d\sum_{j=1}^kG_{ij}\Delta_{ij}. \]

Fix a step budget \(\rho>0\). Our problem is to make the predicted loss change as negative as possible:

\[ \min_{\|\boldsymbol\Delta\|\leq\rho}\ \langle\mathbf G,\boldsymbol\Delta\rangle_F. \]

This is steepest descent, with the norm specifying which steps count as small. Before solving it for a general matrix, look at the diagonal case below. For \(\boldsymbol\Delta=\operatorname{diag}(\delta_1,\delta_2)\), the Frobenius constraint is \(\delta_1^2+\delta_2^2\leq\rho^2\), a disk. The spectral constraint requires both \(|\delta_1|\leq\rho\) and \(|\delta_2|\leq\rho\), giving a square. Both panels use the same gradient and the same loss approximation.

For a diagonal update with entries δ₁ and δ₂, a unit Frobenius budget permits the disk δ₁² + δ₂² ≤ 1, while a unit spectral budget permits the square max(|δ₁|, |δ₂|) ≤ 1. With gradient diag(1, s), the best Frobenius step is proportional to (−1, −s). The best spectral step is (−1, −1) for s > 0. Dragging the shared trial step compares its feasibility and predicted loss under both budgets.

The lines mark equal predicted loss. Moving a line in the descent direction lowers that loss until it reaches the edge of the permitted region. On the disk, the best step points opposite the gradient. On the square, it reaches a corner as long as both gradient entries are positive. Even a small positive second entry gives us a reason to use the full allowance in that direction.

For the Frobenius constraint, Cauchy–Schwarz makes this argument precise. If \(\mathbf G\ne\mathbf0\), then

with equality at \(\boldsymbol\Delta=-\rho\mathbf G/\|\mathbf G\|_F\). This is the gradient-descent direction. Ordinary gradient descent takes \(-\alpha\mathbf G\), so its Frobenius step length is \(\rho=\alpha\|\mathbf G\|_F\). Comparing fixed budgets isolates the direction from the learning rate.

Deriving the Muon Update

The SVD lets us carry the diagonal picture over to any matrix. Write the compact SVD of the nonzero gradient as

\[ \mathbf G=\mathbf U\boldsymbol\Sigma\mathbf V^\top=\sum_{i=1}^r\sigma_i\mathbf u_i\mathbf v_i^\top, \]

where \(r\) is its rank, \(\sigma_1\geq\cdots\geq\sigma_r>0\), \(\mathbf U\in\mathbb R^{d\times r}\), \(\mathbf V\in\mathbb R^{k\times r}\), and \(\boldsymbol\Sigma\in\mathbb R^{r\times r}\). The columns of \(\mathbf U\) and \(\mathbf V\) are orthonormal. In direction \(\mathbf v_i\), the gradient matrix produces \(\mathbf G\mathbf v_i=\sigma_i\mathbf u_i\): an input direction, an output direction, and a stretch.

Consider an update in the opposite directions, with nonnegative amounts \(a_1,\ldots,a_r\):

\[ \boldsymbol\Delta=-\sum_{i=1}^r a_i\mathbf u_i\mathbf v_i^\top,\qquad \langle\mathbf G,\boldsymbol\Delta\rangle_F=-\sum_{i=1}^r\sigma_i a_i. \]

Increasing any \(a_i\) decreases the predicted loss because \(\sigma_i>0\). Under a Frobenius budget, the amounts must share \(\sum_{i=1}^r a_i^2\leq\rho^2\). The gradient direction allocates them in proportion to \(\sigma_i\). Under a spectral budget, each amount only needs to satisfy \(a_i\leq\rho\). We can set all of them to \(\rho\).

This suggests \(\boldsymbol\Delta=-\rho\mathbf U\mathbf V^\top\). We still need to check that an update using other directions cannot do better.

Claim: Under the spectral constraint \(\|\boldsymbol\Delta\|_2\leq\rho\), a steepest update is

\[ \boldsymbol\Delta=-\rho\mathbf U\mathbf V^\top,\qquad \langle\mathbf G,\boldsymbol\Delta\rangle_F=-\rho\sum_{i=1}^r\sigma_i. \]

Proof

Expand the inner product using the SVD. For any permitted update,

\[ \langle\mathbf G,\boldsymbol\Delta\rangle_F=\sum_{i=1}^r\sigma_i\mathbf u_i^\top\boldsymbol\Delta\mathbf v_i\geq-\rho\sum_{i=1}^r\sigma_i. \]

Each term obeys the bound

\[ \mathbf u_i^\top\boldsymbol\Delta\mathbf v_i\geq-\|\mathbf u_i\|_2\|\boldsymbol\Delta\mathbf v_i\|_2\geq-\rho, \]

since both singular vectors have length one. For the proposed update, \(\boldsymbol\Delta\mathbf v_i=-\rho\mathbf u_i\), so every term attains its lower bound. Its positive singular values are all \(\rho\), so it satisfies the constraint as well.

The matrix \(\mathbf Q=\mathbf U\mathbf V^\top\) is the polar factor on these active directions. It retains the singular vectors and replaces each positive singular value by one. It maps each \(\mathbf v_i\) to \(\mathbf u_i\) without the original stretch \(\sigma_i\). On directions absent from the compact SVD it is zero. If the gradient itself is zero, take \(\mathbf Q=\mathbf0\).

The figure returns to our diagonal two-dimensional example. Both updates act on the same unit inputs. The arrows reverse because these are descent updates, and the outputs shown are changes \(\boldsymbol\Delta\mathbf x\), not the layer’s full outputs. The gradient update gives less movement to the weaker singular direction, while the polar update uses the full allowance in both directions when their singular values are positive. The highlighted input is shared between the panels, so we can follow what each update does to exactly the same vector.

Both panels use gradient diag(1, s) and a unit norm budget. The Frobenius-normalized gradient step maps the unit input circle to an ellipse. The polar step maps it to a unit circle while s is positive. Dragging the same input point shows its output under both updates. At s = 0, both chosen updates leave the missing direction unchanged.

There is a sharp distinction near zero. However small the second singular value becomes, the polar update gives that direction full movement as long as it is positive. At exactly zero, our chosen update leaves the direction alone. This jump will matter when we try to approximate the update.

A spectral budget permits more total squared movement than a Frobenius budget of the same numerical size. That explains the larger allowed decrease in our linear approximation. It does not establish which update will train a nonlinear network better. The mathematical distinction is how movement is allocated across directions (Old Optimizer, New Norm).

Muon applies the polar transformation to the momentum-smoothed gradient \(\mathbf M\) from SGD and Adam. Recent gradients can reinforce or cancel before we change their singular values. Using the compact SVD of \(\mathbf M\), the idealized update is

\[ \mathbf M=\mathbf U\boldsymbol\Sigma\mathbf V^\top,\qquad \mathbf W\leftarrow\mathbf W-\alpha\mathbf U\mathbf V^\top. \]

The learning rate \(\alpha>0\) sets the spectral step size. Computing an SVD at every update would be expensive, so we next need a way to approximate this operation with matrix multiplication.

A Matrix Sign Problem

First suppose \(\mathbf M\) is symmetric. In an eigenvector direction it multiplies by a real eigenvalue \(\lambda\). Its polar factor should preserve the direction and any reversal, while removing the magnitude. That means replacing a positive eigenvalue by \(1\), a negative one by \(-1\), and zero by zero. This is the same jump at zero we saw in the figure. The scalar operation is the sign function:

\[ \operatorname{sign}(\lambda)=\begin{cases}1,&\lambda>0,\\0,&\lambda=0,\\-1,&\lambda<0.\end{cases} \]

Applied to a matrix, sign acts on its eigenvalues, not its entries. We want a computational way to make that change without first finding the eigenvectors.

The same idea works for a rectangular matrix. Put \(\mathbf M\in\mathbb R^{d\times k}\) in the off-diagonal blocks of a symmetric matrix \(\mathbf H\in\mathbb R^{(d+k)\times(d+k)}\):

\[ \mathbf H=\begin{bmatrix}\mathbf0&\mathbf M\\\mathbf M^\top&\mathbf0\end{bmatrix},\qquad \operatorname{sign}(\mathbf H)=\begin{bmatrix}\mathbf0&\mathbf Q\\\mathbf Q^\top&\mathbf0\end{bmatrix}. \]

The nonzero eigenvalues of \(\mathbf H\) are the pairs \(\sigma_i,-\sigma_i\). Replacing them by \(1,-1\) removes the singular-value magnitudes from the off-diagonal blocks, giving precisely the polar factor we need. We use this larger matrix to understand the problem. Our implementation will work directly with \(\mathbf M\).

Why does the block-matrix identity hold?

A singular pair satisfies \(\mathbf M\mathbf v_i=\sigma_i\mathbf u_i\) and \(\mathbf M^\top\mathbf u_i=\sigma_i\mathbf v_i\). Consequently,

\[ \mathbf H\frac1{\sqrt2}\begin{bmatrix}\mathbf u_i\\\mathbf v_i\end{bmatrix}=\sigma_i\frac1{\sqrt2}\begin{bmatrix}\mathbf u_i\\\mathbf v_i\end{bmatrix},\qquad \mathbf H\frac1{\sqrt2}\begin{bmatrix}\mathbf u_i\\-\mathbf v_i\end{bmatrix}=-\sigma_i\frac1{\sqrt2}\begin{bmatrix}\mathbf u_i\\-\mathbf v_i\end{bmatrix}. \]

These eigenvectors have length one. In the eigendecomposition of \(\operatorname{sign}(\mathbf H)\), their outer products receive weights \(1\) and \(-1\). Subtracting cancels the diagonal blocks and leaves \(\mathbf u_i\mathbf v_i^\top\) in the upper-right block. Summing over the positive singular values gives \(\mathbf U\mathbf V^\top=\mathbf Q\). The remaining eigenvalues are zero and contribute nothing.

Polynomials give us a way to change eigenvalues using only multiplication and addition. If \(\mathbf H\mathbf v=\lambda\mathbf v\), then \(\mathbf H^j\mathbf v=\lambda^j\mathbf v\) for every nonnegative integer \(j\). Thus for any polynomial \(p\),

\[ p(\mathbf H)\mathbf v=p(\lambda)\mathbf v. \]

A polynomial that approximates sign on the relevant eigenvalues approximates the matrix operation we want. The approximation problem has become one-dimensional, even though the update matrix may have many entries.

Computing the Update

Start by normalizing a nonzero momentum matrix:

\[ \mathbf X^{(0)}=\frac{\mathbf M}{\|\mathbf M\|_F}. \]

Its singular values now lie in \([0,1]\), and positive rescaling does not change its polar factor. For \(\mathbf M=\mathbf0\), return the zero update directly.

The useful matrix identity is

\[ \mathbf X\mathbf X^\top\mathbf X=\mathbf U\boldsymbol\Sigma^3\mathbf V^\top\qquad\text{when }\mathbf X=\mathbf U\boldsymbol\Sigma\mathbf V^\top. \]

The middle factors \(\mathbf V^\top\mathbf V\) and \(\mathbf U^\top\mathbf U\) cancel to identities. We have cubed each stretch without computing its singular vectors! Therefore the cubic \(p_3(x)=(3x-x^3)/2\) can be applied through

\[ \mathbf X^{(t+1)}=\frac32\mathbf X^{(t)}-\frac12\mathbf X^{(t)}(\mathbf X^{(t)})^\top\mathbf X^{(t)}. \]

This is the classical Newton–Schulz iteration. Every positive singular value follows the scalar iteration \(x\leftarrow p_3(x)\). For \(0<x<1\),

\[ p_3(x)-x=\frac{x(1-x^2)}2>0,\qquad 1-p_3(x)=\frac{(1-x)^2(x+2)}2\geq0. \]

It increases without crossing one. The sequence therefore has a limit, and the only positive fixed point in \([0,1]\) is one. Zero stays zero. Since the polynomial is odd, the corresponding negative eigenvalues in the normalized block matrix approach \(-1\). In exact arithmetic, repeated steps recover the polar factor.

Convergence alone does not tell us how quickly we get a useful update. Near zero, \(p_3(x)\approx\tfrac32x\), so a tiny positive singular value grows slowly. The original Muon implementation uses a tuned quintic instead:

\[ p_5(x)=ax+bx^3+cx^5,\qquad (a,b,c)=(3.4445,-4.7750,2.0315). \]

Its larger slope at zero moves small values faster. It also overshoots one, and \(p_5(1)\ne1\), so it is not an iteration that converges exactly to the sign function. It is designed to give a useful approximate update after a few steps. The figure compares both polynomials on the same normalized starting spectrum and with the same number of matrix products. A cubic step takes two products and a quintic step takes three, so the cubic gets more iterations within that budget.

With the same matrix-product budget for both methods, the left panel compares repeated cubic and Muon quintic polynomials with sign on negative and positive eigenvalues. The right panel applies the same compositions to a shared Frobenius-normalized spectrum. The cubic converges to one slowly from small positive values. The quintic moves them faster but overshoots and does not converge exactly to one. Zero stays zero under both. The figure reports approximation error and two versus three matrix products per iteration.

On the left, each curve shows the result of repeatedly applying a polynomial to an initial value within the selected computation budget. On the right, we apply those same compositions to the matrix’s singular values. Small positive values take longer to reach the target, and an exact zero remains zero. The quintic gains speed near zero in exchange for uneven final stretches. With more computation, the cubic can become more accurate. Which approximation is better depends on the starting spectrum and how much accuracy we need.

The figure reports the largest error over the active directions. If \(x_i^{(t)}\) is the scalar iterate in direction \(i\), the shared singular vectors give

\[ \|\mathbf X^{(t)}-\mathbf Q\|_2=\max_{1\leq i\leq r}|x_i^{(t)}-1|. \]

Zero directions contribute no error because both maps leave them at zero. Thus this scalar comparison measures the largest error in the approximate update’s action on a unit input.

To evaluate the quintic efficiently, form \(\mathbf A^{(t)}=\mathbf X^{(t)}(\mathbf X^{(t)})^\top\) and compute

\[ \mathbf X^{(t+1)}=\left(a\mathbf I_d+b\mathbf A^{(t)}+c(\mathbf A^{(t)})^2\right)\mathbf X^{(t)}. \]

This requires three matrix products per iteration, compared with two for the cubic. Equal iteration counts therefore do not mean equal work. For a tall matrix, transpose first so \(\mathbf A^{(t)}\) uses the smaller dimension, then transpose the result back. These product counts help compare the algorithms, though actual time also depends on matrix shapes and hardware.

A better polynomial can reach a useful approximation with less computation, but we must choose what accuracy is worth paying for. The sign function jumps at zero, so no continuous polynomial can approximate it uniformly to arbitrary accuracy near that jump. Very small nonzero singular values remain the difficult case. The demo shows how this affects the classical iteration.

Muon chooses its direction by limiting the largest output change over input directions. Polynomial iteration makes that direction practical to approximate during training. The norm determines what update we want, and approximation determines how much computation we spend getting it.