Muon
The optimizer Muon [JJ+, ’24] uses the matrix shape of a neural network’s weights when computing an update. It starts with momentum, then changes its singular values before applying the update. To understand this step, we will compare ways of measuring a matrix update’s size, then derive and approximate the best update under each constraint.
Update Budgets
Let \(\mathbf W\in\mathbb R^{d\times k}\) be a weight matrix and \(\mathbf G=\nabla_{\mathbf W}\mathcal{L}(\mathbf W)\in\mathbb R^{d\times k}\) its gradient. For a small update \(\boldsymbol\Delta\in\mathbb R^{d\times k}\), \[ \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}^k G_{ij}\Delta_{ij}. \] Steepest descent asks which update makes this linear approximation smallest under a size budget: \[ \underset{\|\boldsymbol\Delta\|\leq\alpha}{\operatorname{minimize}} \ \langle\mathbf G,\boldsymbol\Delta\rangle_F. \] The choice of norm determines which updates satisfy the constraint.
For any \(\mathbf M\in\mathbb R^{d\times k}\), let \(\sigma_1(\mathbf M),\ldots,\sigma_{\min(d,k)}(\mathbf M)\) be its singular values in decreasing order, including zeros. Flattening the matrix into \(dk\) entries chooses the Frobenius norm \[ \|\mathbf M\|_F=\sqrt{\sum_{i=1}^d\sum_{j=1}^k M_{ij}^2} =\sqrt{\sum_{i=1}^{\min(d,k)}\sigma_i(\mathbf M)^2}. \]
For a nonzero gradient and a positive budget, the familiar Cauchy–Schwarz argument gives \[ \langle\mathbf G,\boldsymbol\Delta\rangle_F \geq-\|\mathbf G\|_F\|\boldsymbol\Delta\|_F \geq-\alpha\|\mathbf G\|_F. \] Equality is attained by \(\boldsymbol\Delta_{\rm F}=-\alpha\mathbf G/\|\mathbf G\|_F\). Thus treating the entries as a vector makes the steepest step proportional to the gradient.
Under a Frobenius budget, movement in one singular direction leaves less squared length for the others. The optimal step therefore preserves the gradient’s relative singular values.
A matrix also has a size as a linear map. Its spectral norm is its largest singular value, \[ \|\mathbf M\|_2=\max_{\mathbf x\ne0} \frac{\|\mathbf M\mathbf x\|_2}{\|\mathbf x\|_2}=\sigma_1(\mathbf M). \] Changing the layer by \(\boldsymbol\Delta\) changes its output by at most \[ \|\boldsymbol\Delta\mathbf x\|_2 \leq \|\boldsymbol\Delta\|_2\|\mathbf x\|_2. \] Thus \(\|\boldsymbol\Delta\|_2\leq\alpha\) guarantees that no input direction moves by more than \(\alpha\) times its length. The Frobenius norm instead adds movement across all singular directions.
If a step has \(r\) equal nonzero singular values \(c\), its spectral norm is \(c\) while its Frobenius norm is \(c\sqrt r\). A spectral budget bounds the largest change to any unit input. A Frobenius budget bounds the square root of the total squared change across an orthonormal input basis.
The Polar Factor
Write a compact singular value decomposition \[ \mathbf G=\mathbf U\boldsymbol\Sigma\mathbf V^\top =\sum_{i=1}^r\sigma_i\mathbf u_i\mathbf v_i^\top, \] where \(r=\operatorname{rank}(\mathbf G)>0\). Here \(\sigma_i=\sigma_i(\mathbf G)>0\) for \(i=1,\ldots,r\), and the columns \(\mathbf u_i\in\mathbb R^d\) of \(\mathbf U\in\mathbb R^{d\times r}\) and \(\mathbf v_i\in\mathbb R^k\) of \(\mathbf V\in\mathbb R^{k\times r}\) are orthonormal within each family.
Claim: If \(\mathbf G\ne\mathbf0\) and \(\alpha>0\), under \(\|\boldsymbol\Delta\|_2\leq\alpha\) a steepest update is the negative polar factor \[ \boldsymbol\Delta_{\rm spec}=-\alpha\mathbf U\mathbf V^\top. \] The minimum linearized change is \[ \min_{\|\boldsymbol\Delta\|_2\leq\alpha} \langle\mathbf G,\boldsymbol\Delta\rangle_F =-\alpha\sum_{i=1}^r\sigma_i. \]
Proof
First expand the linearized change using the SVD and the trace definition of the Frobenius inner product: \[ \begin{aligned} \langle\mathbf G,\boldsymbol\Delta\rangle_F &=\operatorname{tr}(\mathbf G^\top\boldsymbol\Delta)\\ &=\sum_{i=1}^r\sigma_i \mathbf u_i^\top\boldsymbol\Delta\mathbf v_i. \end{aligned} \] For unit singular vectors and a feasible \(\boldsymbol\Delta\), \[ \mathbf u_i^\top\boldsymbol\Delta\mathbf v_i \geq-\|\mathbf u_i\|_2\|\boldsymbol\Delta\|_2\|\mathbf v_i\|_2 \geq-\alpha. \] Since every \(\sigma_i>0\), summing gives \(\langle\mathbf G,\boldsymbol\Delta\rangle_F \geq-\alpha\sum_{i=1}^r\sigma_i\). The proposed update has spectral norm \(\alpha\) because \(\mathbf U\mathbf V^\top\) has nonzero singular values equal to one. It also satisfies \((-\alpha\mathbf U\mathbf V^\top)\mathbf v_i=-\alpha\mathbf u_i\), so every summand attains its lower bound. Therefore the negative polar factor minimizes the linearized loss under the spectral-norm constraint.
When \(\mathbf G\) is rank deficient, the minimizer need not be unique because an update may also act on directions absent from the gradient. The claim identifies one canonical minimizer.
The matrix \(\mathbf U\mathbf V^\top\) is the polar factor of \(\mathbf G\). It retains the singular directions and replaces every nonzero singular value by one.
For a rectangular or rank-deficient gradient, this produces a partial isometry rather than a square orthogonal matrix.
A unit Frobenius budget normalizes the gradient by its Frobenius norm. A unit spectral budget permits a polar step with every active singular value equal to one.
The linked panels compare singular values with the image of a unit circle under the update. Frobenius and spectral constraints impose different update budgets, so their optimal decreases answer different optimization questions.
Newton–Schulz Iteration
Computing an SVD of every momentum matrix would be expensive. Newton–Schulz iterations approximate its polar factor using matrix multiplications.
First rescale a nonzero momentum-smoothed gradient \(\mathbf M\) so that its singular values lie at most \(1\), for example \[ \mathbf X^{(0)}=\frac{\mathbf M}{\|\mathbf M\|_F}. \] If \(s_{\min}>0\) is its smallest nonzero singular value, then every active singular value lies in \([s_{\min},1]\). A small \(s_{\min}\) means the matrix is poorly conditioned and the weakest active direction has farther to travel. To see the mechanism, apply the classical Newton–Schulz step \[ \mathbf X^{(t+1)} =\frac32\mathbf X^{(t)}-\frac12\mathbf X^{(t)} \big(\mathbf X^{(t)}\big)^\top\mathbf X^{(t)}. \]
Claim: If every singular value of \(\mathbf X^{(0)}\) lies in \([0,1]\), the Newton–Schulz iteration preserves its singular vectors, sends each singular value through \[ p(x)=\frac{3x-x^3}{2}. \] Every initially positive singular value converges to one, while every zero singular value stays zero. Thus the iterates converge to the polar factor in exact arithmetic.
Proof
Let \(\mathbf X^{(t)}=\mathbf U\boldsymbol\Sigma\mathbf V^\top\) be a compact SVD. Orthonormality gives \[ \begin{aligned} \mathbf X^{(t)}(\mathbf X^{(t)})^\top\mathbf X^{(t)} &=\mathbf U\boldsymbol\Sigma\mathbf V^\top \mathbf V\boldsymbol\Sigma\mathbf U^\top \mathbf U\boldsymbol\Sigma\mathbf V^\top\\ &=\mathbf U\boldsymbol\Sigma^3\mathbf V^\top. \end{aligned} \] Substituting into the iteration yields \[ \mathbf X^{(t+1)} =\mathbf U\left(\frac{3\boldsymbol\Sigma-\boldsymbol\Sigma^3}{2}\right) \mathbf V^\top. \] The singular vectors are unchanged, and each singular value follows \(p\). For \(0<x<1\), \[ p(x)-x=\frac{x(1-x^2)}2>0, \qquad 1-p(x)=\frac{(1-x)^2(x+2)}2\geq0. \] An active singular value below one therefore increases but never exceeds one. One already equal to one remains fixed because \(p(1)=1\). The increasing sequence has a limit \(L\in(0,1]\), and continuity gives \(L=p(L)\). Expanding this equation, \[ 2L=3L-L^3 \quad\Longleftrightarrow\quad L(L^2-1)=0. \] The only solution in \((0,1]\) is \(L=1\). Since \(p(0)=0\), zero singular values remain zero. The limiting matrix is consequently \(\mathbf U\mathbf V^\top\), the polar factor on the active singular subspaces.
The singular directions remain fixed while the positive singular values approach one. An exactly zero value remains zero.
Each singular value follows p(x) = (3x − x³)/2. Values in (0, 1] approach one and zero remains zero.
A small positive starting value takes longer to approach one, which explains the dependence on conditioning.
The classical cubic makes every positive singular value converge to one, but practical Muon uses a tuned quintic polynomial and only a few steps to obtain an approximate polar direction [JJ+, ’24]. The practical map increases small singular values more quickly but does not make the finite-step output exactly orthogonal.
Muon first averages recent gradients into a momentum matrix and then approximately orthogonalizes it. Applying the nonlinear map after averaging allows raw magnitudes to cancel or reinforce before their relative singular strengths are discarded. Scalar and vector parameters use an ordinary optimizer because they are not matrix maps.
The polar direction is optimal for the linearized loss under a spectral budget. The nonlinear loss still requires a suitable learning rate. Treating a matrix as a map changes the promise that defines an admissible step. Newton–Schulz makes that matrix-aware direction practical to approximate.