Muon, Part 1: A Matrix Has More Than Coordinates

September 2026 · 6 min read + interactive lab

Jump to the loss-contour trajectory demo ↓

Why treating a 2D weight matrix as a flat list of numbers starves most learning directions—and how Muon fixes it by setting every active singular value of the update to $1$.

Optimizers like SGD, RMSprop, Adam, and Lion treat a weight matrix $W\in\mathbb R^{m\times n}$ as a flat list of $mn$ independent numbers, updating each entry $W_{ij}$ on its own. Muon (MomentUm Orthogonalized by Newton–schulz) instead treats $W$ as a 2D transformation: it accumulates a momentum matrix $M_t$, then equalizes its singular values so every active direction of the layer takes a step of the exact same size.

1. Wait—isn’t the gradient a vector? Why are $G_t$ and momentum $M_t$ matrices?

In Chapters 1–6 (SGD through Adam), we stacked every parameter of the network into one long 1D column vector $w\in\mathbb R^d$, so the gradient $g_t$ and momentum $m_t$ were also written as 1D vectors. That notation worked because coordinate-wise optimizers update each weight $W_{ij}$ completely on its own, so the shape of the parameter tensor never mattered.

Inside an actual neural network, however, a linear layer $y=Wx$ connecting $n$ inputs to $m$ outputs stores its weights in a 2D grid (matrix) $W\in\mathbb R^{m\times n}$. How do the layer’s gradient $G_t$ and momentum $M_t$ become $m\times n$ matrices?

• 1. Every weight $W_{ij}$ has its own derivative $\partial\ell/\partial W_{ij}$: Since output $i$ is $y_i=\sum_{j=1}^n W_{ij}x_j$, let $x\in\mathbb R^{n\times 1}$ be the input column vector and $\delta=\partial\ell/\partial y\in\mathbb R^{m\times 1}$ be the backpropagated output-gradient column vector. By the chain rule, the derivative for weight $(i,j)$ is $\frac{\partial\ell}{\partial W_{ij}}=\delta_i x_j$.
• 2. 2D in shape ($m\times n$), 1D in rank (rank $1$) for one example: By the matrix-multiplication rule $(m\times\mathbf{1})\times(\mathbf{1}\times n)=(m\times n)$, multiplying the column vector $\delta$ by the row vector $x^\top$ produces a 2D grid with $m$ rows and $n$ columns:
  $$G=\frac{\partial\ell}{\partial W}=\delta x^\top=\begin{bmatrix}\delta_1\\\vdots\\\delta_m\end{bmatrix}_{m\times 1}\begin{bmatrix}x_1&\cdots&x_n\end{bmatrix}_{1\times n}=\begin{bmatrix}\delta_1 x_1&\cdots&\delta_1 x_n\\\vdots&\ddots&\vdots\\\delta_m x_1&\cdots&\delta_m x_n\end{bmatrix}_{m\times n}.$$
  Notice, however, that for a single example, every row of $\delta x^\top$ is just a scalar multiple $\delta_i$ of the same row vector $x^\top$—so a single example’s gradient matrix has rank $1$ (it only spans a single 1D direction, with $\sigma_1>0$ and $\sigma_2=\cdots=0$).
• 3. Batching and momentum $M_t$ build up multiple directions (rank $r$): Averaging over $B$ different examples in a mini-batch ($G_t=\frac{1}{B}\sum_{b=1}^B\delta_b x_b^\top$) and accumulating across training steps in the momentum buffer ($M_t=\beta M_{t-1}+G_t$) adds many different rank-$1$ matrices together. That sum is still an $(m\times n)$ 2D matrix, but now has rank $r\le\min(m,n)$ with many active directions:
  $$G_t=\frac{1}{B}\sum_{b=1}^B\delta_b x_b^\top\in\mathbb R^{m\times n},\qquad M_t=\beta M_{t-1}+G_t\in\mathbb R^{m\times n}.$$

Why a momentum matrix has skewed singular values

Once we view the momentum buffer $M\in\mathbb R^{m\times n}$ as a 2D matrix (of rank $r\le\min(m,n)$), we can decompose it via the Singular Value Decomposition (SVD) into a sum of perpendicular rank-$1$ directions:

Why is $v_i\in\mathbb R^n$ called an input direction and $u_i\in\mathbb R^m$ called an output direction? Look at what happens when a single term $\sigma_i u_i v_i^\top$ multiplies a layer input vector $x\in\mathbb R^n$ (just as $W$ multiplies $x$ in $y=Wx$):

And because the vectors $v_1,\dots,v_r$ are mutually perpendicular unit vectors ($v_i^\top v_i=1$ and $v_j^\top v_i=0$ for $j\ne i$), feeding $x=v_i$ into the full matrix $M$ zeros out every term except the $i$-th one:

So each rank-$1$ term $\sigma_i u_i v_i^\top$ acts as a one-channel pipe from the layer’s input space to its output space (matching the shape of $\delta_b x_b^\top$ above):

• Input direction $v_i\in\mathbb R^n$ ($V^\top V=I$): a unit vector in the $n$-dimensional input space that the matrix reads from via the dot product $v_i^\top x$ (sitting in the same right-hand $1\times n$ slot as the input $x_b^\top$).
• Output direction $u_i\in\mathbb R^m$ ($U^\top U=I$): a unit vector in the $m$-dimensional output space that the matrix writes to (sitting in the same left-hand $m\times 1$ slot as the output gradient $\delta_b$).
• Singular value $\sigma_i>0$: the stretch factor (gain) measuring how strongly $M$ amplifies input direction $v_i$ into output direction $u_i$ ($\|Mv_i\|_2=\sigma_i$).

In real networks, batch activations $x_b$ and output gradients $\delta_b$ concentrate heavily along a few dominant directions. As a result, the top singular value $\sigma_1$ of $M$ is huge while the remaining singular values are tiny ($\sigma_1\gg\sigma_2\gg\cdots\gg\sigma_r$). If you update with raw momentum $\Delta W=-\eta M$, the optimizer takes a giant step along the single loudest direction $\sigma_1 u_1 v_1^\top$ while barely moving the quieter directions that still carry useful learning signal.

2. Muon’s fix: set every active singular value to $1$

Muon keeps the input and output directions $(U,V)$ of the momentum matrix intact, but replaces every nonzero singular value $\sigma_i$ with $1$. In linear algebra, this matrix $Q=UV^\top$ is called the orthogonal polar factor of $M$:

Think of $Q=UV^\top$ as the matrix version of the $\operatorname{sign}(\cdot)$ function: just as $\operatorname{sign}(g)$ normalizes every nonzero scalar to $\pm 1$, Muon normalizes every nonzero singular value $\sigma_i$ to $1$. When $M$ has full rank ($r=\min(m,n)$), the resulting update matrix $Q$ has perpendicular unit columns or rows:

3. Why $UV^\top$? Steepest descent with a cap on layer stretch

Why is $Q=UV^\top$ the best update direction for a linear layer? We can derive it in three simple steps without any matrix-trace tricks:

Step 1: How much does a weight update $D=\Delta W$ change the loss?
For a 1D vector, nudging parameters by $\Delta w$ changes the loss by the dot product $g^\top\Delta w=\sum_k g_k\Delta w_k$ (slope $\times$ step). For a 2D weight matrix $W\in\mathbb R^{m\times n}$, nudging every entry $W_{pq}$ by $D_{pq}$ changes the loss by the exact same entry-by-entry sum of (slope $\times$ step), denoted $\langle M,D\rangle_F$:

To decrease the loss as much as possible, we want to make this sum $\langle M,D\rangle_F$ as negative as possible.

Step 2: Cap how much the layer’s output can jump ($\|D\|_2\le\eta$).
When we change the weights by $D$, a unit input vector $\|x\|_2=1$ sees its layer output change by $\Delta y=Dx$. To keep the layer’s activations stable, we cap the worst-case output jump across all unit inputs using the spectral norm:

Step 3: Plug in the SVD $M=\sum_{i=1}^r\sigma_i u_i v_i^\top$.
For a single rank-$1$ term $u_i v_i^\top$, its entry-by-entry dot product with $D$ is simply $u_i^\top Dv_i$—the scalar measuring how much $D$ maps input direction $v_i$ into output direction $u_i$. Substituting the SVD of $M$ splits the loss change into $r$ independent channels:

Because $u_i$ and $v_i$ are unit vectors, the stretch cap $\|D\|_2\le\eta$ bounds each channel’s step between $-\eta$ and $+\eta$:

Since every singular value is positive ($\sigma_i>0$), how do we make $\sum_{i=1}^r\sigma_i(u_i^\top Dv_i)$ as negative as possible? Push every single active channel to its full negative limit $u_i^\top Dv_i=-\eta$:

In short, the three main optimizer families just differ in what they equalize:

4. Lab: matrix normalization versus entry-wise signs

Build a $2\times 2$ matrix $M=U\operatorname{diag}(s_1,s_2)V^\top$ below to compare how the raw matrix $M$, Muon’s polar factor $Q=UV^\top$, and the entry-wise sign matrix $\operatorname{sign}(M)$ transform a unit circle.

• Vary $(s_1,s_2)$ while keeping both positive: the middle panel ($UV^\top$) never changes at all, because $UV^\top$ depends only on the directions $U,V$, not on the magnitudes $(s_1,s_2)$.
• Why does $\operatorname{sign}(M)=\begin{bmatrix}1&-1\\1&-1\end{bmatrix}$ have singular values $2, 0$ yet still show a $y$-component? Singular values are not the $x$-axis and $y$-axis lengths—they are the length and width of the output ellipse along its tilted singular directions $u_1$ and $u_2$. Multiplying any input $x=\begin{bmatrix}x_1\\x_2\end{bmatrix}$ by $\begin{bmatrix}1&-1\\1&-1\end{bmatrix}$ gives $\begin{bmatrix}x_1-x_2\\x_1-x_2\end{bmatrix}=(x_1-x_2)\begin{bmatrix}1\\1\end{bmatrix}$: every output lands on the $45^\circ$ diagonal line $y=x$ (which has both $x$ and $y$ components!). Here $\sigma_1=2$ is the stretch along that $45^\circ$ diagonal $u_1=\tfrac{1}{\sqrt 2}\begin{bmatrix}1\\1\end{bmatrix}$, while $\sigma_2=0$ means the circle has zero thickness along the perpendicular $-45^\circ$ direction $u_2=\tfrac{1}{\sqrt 2}\begin{bmatrix}-1\\1\end{bmatrix}$. Notice the flaw of entry-wise $\operatorname{sign}(M)$: even though $M$ had two active directions ($s_1=3, s_2=0.5$), taking entry-wise signs made both rows identical and collapsed the 2D circle into a 1D line!
• Rotate the angles: notice how $UV^\top$ rotates smoothly with the basis, whereas $\operatorname{sign}(M)$ snaps discontinuously whenever a matrix entry crosses zero.
• Set $s_2=0$ (rank-deficient): the partial polar factor sets positive singular values to $1$ and leaves zero singular values at $0$—it never invents noise along empty directions.

5. Lab: watch Muon descend through the loss contours

This lab runs Muon on the same eight-example quadratic bowl as the SGD demo, embedded in a $2\times 2$ weight matrix $W=\left[\begin{smallmatrix}w_1&w_2\\0&0\end{smallmatrix}\right]$. Each step updates the momentum matrix, normalizes by the Frobenius norm, and runs $5$ Newton–Schulz iterations to orthogonalize the update before stepping $W$.

6. Complete Muon algorithm, shape scaling, and memory

Computing an exact SVD on every weight matrix is slow on GPUs, so Muon approximates $UV^\top$ using $5$ fast matrix-multiply steps of a Newton–Schulz polynomial (derived in Part 2). Putting momentum $B_t$, optional Nesterov lookahead $\widetilde B_t$, Newton–Schulz orthogonalization, and decoupled weight decay $\lambda$ together gives the complete update for a 2D weight matrix $W_t\in\mathbb R^{m\times n}$:

Two practical details explain how Step 3 and Step 4 work:

• Why divide by $\|\widetilde B_t\|_F+\epsilon$ before Newton–Schulz?
  Newton–Schulz (derived in Part 2) pushes singular values toward $1$ by repeatedly plugging each singular value $s$ into an odd polynomial such as $s_{k+1}=1.5s_k-0.5s_k^3$ while leaving the directions $u_i,v_i$ untouched. That polynomial only converges to $1$ if the starting singular values lie inside $(0,1]$:
  • If $\widetilde B_t$ has a singular value $\sigma_1>1$ (e.g., $\sigma_1=3$) and you skip normalization, the $-0.5s_k^3$ term dominates and explodes to $\pm\infty$ (`NaN`) in just a few steps:
    $$\text{Direct }\operatorname{NS}(\widetilde B_t)\text{ on }\sigma_1=3:\quad 3\;\xrightarrow{\text{step 1}}\;1.5(3)-0.5(3^3)=-9\;\xrightarrow{\text{step 2}}\;+351\;\xrightarrow{\text{step 3}}\;-2.16\times 10^7\;\longrightarrow\;\text{NaN!}$$
  • If $\widetilde B_t$ has tiny singular values (e.g., $\sigma_1=10^{-4}$ late in training) and you skip normalization, $s_k^3\approx 0$, so 5 steps only multiply $\sigma_1$ by a small factor ($\sim 1.5^5\times 10^{-4}\approx 0.00076\ll 1$), failing to reach $1$.
  Dividing by $\|\widetilde B_t\|_F+\epsilon=\sqrt{\sum_j\sigma_j^2}+\epsilon$ fixes both issues at once: it rescales every singular value into $(0,1]$ without changing $(u_i,v_i)$, so 5 Newton–Schulz steps safely push all of them to $\approx 1$, giving $Q_t\approx UV^\top$.
• Why multiply by $\sqrt{\max(1,m/n)}$? An orthogonal matrix $Q\in\mathbb R^{m\times n}$ has $\min(m,n)$ singular values equal to $1$, so the sum of squares of all $mn$ entries is $\|Q\|_F^2=\min(m,n)$. That means the typical (RMS) size of each entry $Q_{ij}$ is $\sqrt{\min(m,n)/(mn)}=1/\sqrt{\max(m,n)}$. When a matrix has many more rows than columns ($m\gg n$), its individual entries get smaller; multiplying by $\sqrt{\max(1,m/n)}$ keeps the per-parameter step size consistent across tall layers.

# One Muon step for a 2D hidden weight matrix W of shape (m, n)
B = beta * B + grad
B_tilde = grad + beta * B if nesterov else B
X0 = B_tilde / (B_tilde.norm() + 1e-7)     # Shrink all singular values into (0, 1]
Q = newton_schulz_5(X0)                    # Push all active singular values -> 1 (Q ≈ U @ V.T)
scale = max(1.0, rows / cols) ** 0.5       # Aspect-ratio scaling for tall matrices
W *= 1 - lr * weight_decay                 # Decoupled weight decay
W -= lr * scale * Q

Optimizer state memory: Unlike AdamW (which stores two buffers $m_t$ and $v_t$, costing $+8\text{ bytes/parameter}$ in FP32), Muon stores only one buffer per parameter—the momentum matrix $B_t$ ($+4\text{ bytes/parameter}$ in FP32, matching SGD with Momentum)—because $Q_t$ is computed on the fly and immediately discarded.

7. Parameter routing: which layers use Muon vs. AdamW?

Muon is built specifically for dense 2D hidden-layer weight matrices. In a Transformer or MLP, non-matrix and vocabulary parameters should still be optimized with AdamW:

8. Questions worth answering

Sources

Keller Jordan, Muon: An optimizer for hidden layers in neural networks; Muon reference implementation; Bernstein and Newhouse, Old Optimizer, New Norm: An Anthology.