how you could have invented muon from memory

AdamW is a comfortable default. Then someone remembers that a matrix is a transformation, and suddenly there is a suspiciously simple optimizer to explain.

Let's try to invent it. We have a model, some data, and enough calculus and linear algebra to be dangerous. How should we change the weights?

We'll start with the obvious answer, gradient descent. Then we'll notice that a matrix does more than store a list of numbers. Taking that seriously will lead us to Muon's update. The aim is to make each choice follow from the previous one, so the formula becomes something we can reconstruct.

The optional explanations are folded away. Open whichever bits you need; there is no entrance exam.

§A function, some data, and a loss

A model is a mathematical function $f_\theta:\mathcal X\to\mathcal Y$, with adjustable parameters $\theta$. We choose those parameters to approximate a relationship we do not know. In regression, call that relationship $f_*$; after fitting, our estimate is $\hat f=f_{\hat\theta}\approx f_*$. The hat means "estimated from data".

This covers a line $f_\theta(x)=ax+b$, a vision network, or an LLM. For an LLM, the output is a probability distribution over possible next tokens. We are estimating a conditional distribution, rather than assuming every input has one inevitable answer.

A loss turns "how well are we doing?" into a number to minimize. It doesn't have to be a distance in the formal mathematical sense. Here, it measures what we intuitively call distance or difference: how far the model's predictions are from what we want.

For regression, our training data consists of pairs $(x_i,y_i)$. The input is $x_i$, the desired output is $y_i$, and the model predicts $f_\theta(x_i)$. For a language model, $x_t$ is the observed token at position $t$, and $x_{<t}$ is the preceding context. Two common losses are

The first is mean squared error over $N$ input/output pairs. The second is token cross-entropy over $T$ predicted tokens: it penalizes assigning low probability to the token that actually occurred. A lower loss means improvement on this objective; choosing a useful objective is our responsibility.

Why does low token probability give a large loss?

For a target distribution $y$ and predicted probabilities $p$, cross-entropy is $-\sum_j y_j\log p_j$. A one-hot target puts all its mass on the observed token, so only $-\log p(\text{observed token})$ remains.

Giving the correct token probability $0.8$ costs about $0.223$ nats. Giving it $0.2$ costs about $1.609$. The negative logarithm makes confident mistakes expensive. Averaging this quantity over tokens gives the loss above.

Training moves $\theta$ through a space with one dimension per parameter. We want a useful direction and a sensible distance. We get local information about the loss, not a map showing the best route to the bottom.

§First, take a step downhill

So far, $\theta$ has meant all the model's adjustable parameters. To work out one update, focus on the weight matrix of a single layer, $W\in\mathbb R^{m\times n}$, and hold the other parameters fixed. We will write the loss as $L(W)$ to show which parameters we are changing. It is the same training objective, viewed as a smaller optimization problem.

Once we have a rule for this matrix, we can apply it to other weight matrices too. Biases and other parameter shapes may use a different update rule. For now, assume $L(W)$ is differentiable: small weight changes have a first-order effect that we can predict from the local slopes.

An optimizer adds a step $D$, with the same shape as $W$:

Each $D_{ij}$ says how much to change the corresponding weight $W_{ij}$. We keep uppercase $W$ and $D$ for the whole matrices throughout this derivation. In the small examples below, lowercase $w_1,w_2$ denote individual weights.

The gradient matrix $G$ collects the partial derivatives, $G_{ij}=\partial L/\partial W_{ij}$. Each entry measures the local change in loss per tiny change in that weight, with the others held fixed. A positive slope says increasing the weight raises the loss locally; a negative slope says the opposite.

The gradient gives us slopes. To turn them into an actual step, choose a positive scale $\alpha$, called the learning rate, and define

The minus sign reverses the slopes; $\alpha$ controls how far we move along them. This is gradient descent. Whether it lowers the actual loss still depends on how large a step we take.

Derivatives and partial derivatives, without the ceremony

For $f(w)=w^2$, changing $w=1$ by $h$ changes the output by $2h+h^2$. Divide by $h$ to get the average slope $2+h$. As $h$ approaches zero, this approaches the derivative $f'(1)=2$.

Play animationRestart← BackNext →Shrink the interval and watch the slope settle.

A derivative measures slope. A second derivative measures how that slope changes, which is curvature. If every partial derivative is zero, we have a stationary point. It could be a minimum, a maximum, or a saddle. Zero does not come with a "best possible answer" certificate.

Work out a two-weight loss and its gradient

Let two predictions be $\hat y_1=w_1$ and $\hat y_2=2w_2$, both with target zero. Half the sum of squared errors gives

$$\begin{aligned} L(w_1,w_2) &=\tfrac12\bigl[(\hat y_1-0)^2+(\hat y_2-0)^2\bigr]\\[10pt] &=\tfrac12\bigl[w_1^2+(2w_2)^2\bigr]\\[10pt] &=\tfrac12(w_1^2+4w_2^2). \end{aligned}$$

The $1/2$ tidies the derivatives. The $4$ comes from squaring the second prediction's factor of $2$. In this two-weight example, collect the derivatives in a vector $g$:

$$\begin{aligned} g&=\left(\frac{\partial L}{\partial w_1},\frac{\partial L}{\partial w_2}\right)\\[12pt] &=\left(\tfrac12\cdot2w_1,\tfrac12\cdot8w_2\right)\\[10pt] &=(w_1,4w_2). \end{aligned}$$

A real network chains layers together. An early weight affects later activations, then the prediction, then the loss. Backpropagation applies the chain rule through those layers to compute the same kind of partial derivatives.

Choosing $D$ gets harder when the loss has curved valleys, flat regions, saddles, or several minima. A downhill step can settle in a shallow basin even when a deeper one exists elsewhere.

The surface below has two minima. The first run settles in the higher one. Try the other basin to see a different result. You can also drag the graph to look around and zoom with the buttons.

The exact loss behind the picture

The surface uses

$$\begin{aligned} L_{\mathrm{draw}}(x,y)&=(x^2-1)^2\\[6pt] &\quad+0.3(x-1)^2+8y^2. \end{aligned}$$

Its gradient is $(4x(x^2-1)+0.6(x-1),\,16y)$. The global minimum is $(1,0)$: all three nonnegative terms vanish there. There is also a strict local minimum at $x=(-1-\sqrt{0.4})/2\approx-0.816$, $y=0$, separated from the deeper basin by a saddle.

The animation recomputes the gradient at each step, using $\alpha=0.09$. The same function drives the surface and the path. Moving between iterates is only a visual aid; the optimization steps are discrete. This is a two-weight illustration, not a claim that shallow local minima are the main difficulty in every large network.

§Turn many slopes into one score

We have a candidate step $D$, but its entries only tell us how much to change the weights. They do not tell us whether those changes help. For that, we need to know how the loss responds. That is the information in $G$.

Hold every weight except $W_{ij}$ fixed. Recall that $G_{ij}=\partial L/\partial W_{ij}$ is the local slope: loss change per unit change in this weight, in the limit of a tiny move. Near our current weights, we approximate the curve by a straight line. On a straight line, change in height is slope times horizontal movement. Here that becomes

For example, if $G_{ij}=3$ and we increase that weight by $D_{ij}=0.01$, the predicted loss change is $3\times0.01=0.03$. Moving it by $-0.01$ instead predicts $-0.03$, a decrease. A negative gradient entry reverses which direction helps. The same weight change can help or hurt depending on its slope.

When several weights change together, differentiability tells us that their first-order effects add. This does not assume the weights act independently: the slopes depend on the whole current matrix, and the linear approximation leaves a remainder. We will make that approximation explicit in the next section. For now, our combined score is

This is the ordinary dot product of the entries, keeping them arranged as matrices. Its name is the Frobenius inner product. For any two matrices $A$ and $B$ of the same shape, we write

Here $A^\top$ is the transpose, which swaps rows and columns. $\operatorname{tr}$ denotes the trace, the sum of a square matrix's diagonal entries. The whole expression returns one number.

What does the trace notation mean?

For a square matrix $C$, the trace is the sum of its diagonal entries:

$$\operatorname{tr}(C)=\sum_j C_{jj}.$$

For example,

$$\begin{aligned} C&=\begin{bmatrix}3&0\\2&-2\end{bmatrix},\\[10pt] \operatorname{tr}(C)&=3+(-2)=1. \end{aligned}$$

The entry $2$ below the diagonal contributes nothing. Trace takes a square matrix and returns one scalar.

Why does it appear in our score? The $j$th diagonal entry of $A^\top B$ is the dot product of column $j$ of $A$ with column $j$ of $B$:

$$\begin{aligned} (A^\top B)_{jj}&=\sum_i A_{ij}B_{ij},\\[10pt] \operatorname{tr}(A^\top B)&=\sum_j\sum_i A_{ij}B_{ij}. \end{aligned}$$

Adding those diagonal entries counts every matching pair once. That gives the same inner product we started with.

We will also need a way to measure the size of a step. Using the same entries, define the Frobenius norm:

It is the Euclidean length we would get by listing all the entries in a vector. This length and our score are related: the difference between the squared lengths of $A+B$ and $A-B$ isolates their inner product.

This is the real polarization identity.

Where does the polarization formula come from?

Start with the norm's definition. Squaring it removes the square root, and each entry of $A+B$ is $A_{ij}+B_{ij}$. Apply the ordinary identity $(a+b)^2=a^2+2ab+b^2$ to each entry:

$$\begin{aligned} \|A+B\|_F^2&=\sum_{i,j}(A_{ij}+B_{ij})^2\\[10pt] &=\sum_{i,j}A_{ij}^2+\sum_{i,j}B_{ij}^2\\[6pt] &\quad+2\sum_{i,j}A_{ij}B_{ij}\\[10pt] &=\|A\|_F^2+\|B\|_F^2\\[6pt] &\quad+2\langle A,B\rangle_F. \end{aligned}$$

The last line simply recognizes the two squared norms and the inner product from their definitions. Replacing $B$ by $-B$ changes the sign of the cross term:

$$\begin{aligned} \|A-B\|_F^2&=\|A\|_F^2+\|B\|_F^2\\[6pt] &\quad-2\langle A,B\rangle_F. \end{aligned}$$

Subtract the second expansion from the first. The squared norms cancel, and the cross terms leave $4\langle A,B\rangle_F$. Dividing by four gives the formula above.

§Predict whether a step will lower the loss

Our sum of slopes times changes looks like a useful score. Now we need to connect it to the quantity we actually care about, the loss after taking the step.

A differentiable function is locally approximated by its value plus a linear change. Applied to the weights, that linear change is exactly $\langle G,D\rangle_F$. This is the first-order Taylor expansion:

The first term is our current loss. The second is the predicted change. A more negative inner product means a larger predicted decrease. The approximation is local, so the step size still matters.

Why does the local slope predict a change in loss?

Taylor approximates a function near a point using derivatives there. For a sufficiently smooth scalar function,

$$\begin{aligned} f(x+h)&=f(x)+f'(x)h\\[8pt] &\quad+\frac{f''(x)}{2!}h^2+\cdots\\[8pt] &\quad+\frac{f^{(N)}(x)}{N!}h^N+R_N(h). \end{aligned}$$

The remainder $R_N$ records the error. If the next derivative is continuous and bounded in magnitude by $B$ along the interval, $|R_N(h)|\leq B|h|^{N+1}/(N+1)!$. Keeping only the value and slope gives the first-order approximation.

For $f(x)=x^2$ at $x=1$, the tangent predicts $1+2h$ and the exact value is $1+2h+h^2$. Watch the terms appear with their geometric meaning.

Play animationRestart← BackNext →Value, slope, and the part the tangent misses.

For matrices, restrict the loss to a line, $\phi(t)=L(W+tD)$. The chain rule gives

$$\phi'(0)=\sum_{i,j}\frac{\partial L}{\partial W_{ij}}D_{ij} =\langle G,D\rangle_F.$$

That is where Frobenius enters: it packages the sum of slopes times changes. Differentiability gives the precise statement

$$\begin{aligned} L(W+D)&=L(W)+\langle G,D\rangle_F\\[6pt] &\quad+R(D). \end{aligned}$$

The remainder becomes small relative to the step:

$$\frac{R(D)}{\|D\|_F}\longrightarrow0\quad\text{as }D\to0.$$

A quadratic error bound needs extra smoothness, such as a continuous Hessian nearby. For our earlier quadratic, the remainder is exactly $\tfrac12(d_1^2+4d_2^2)$. An oversized step can let this positive remainder overwhelm the predicted decrease.

§Decide how far we are willing to move

We can now compare steps by their predicted change in loss. But the prediction only uses local slopes. We choose a size limit so we stay close enough for it to be useful. A huge downhill step can overshoot, raise the actual loss, and send training bouncing around. The tangent line will happily recommend it anyway.

Call the limit $\eta>0$, our step budget. We choose both its size and how to measure it; a limit alone does not guarantee convergence. Measuring size with the Frobenius norm gives

Read this as: among all steps within the budget, choose the one with the most negative predicted loss change. The gradient $G$ is fixed at our current weights; the step $D$ is what we get to choose.

To see the geometry, imagine listing the matrix entries in vectors. For nonzero $G$ and $D$, call the angle between those vectors $\varphi$:

The quantity $\|D\|_F\cos\varphi$ is the signed projection of the step onto the gradient direction. Multiplying by $\|G\|_F$ turns that displacement into a predicted loss change. Moving with the gradient gives a positive score; moving against it gives a negative one. A perpendicular step scores zero, which predicts no first-order improvement. We want the most negative score our budget allows.

Cauchy–Schwarz says $|\langle G,D\rangle_F|\leq\|G\|_F\|D\|_F$. Taking the negative side of that bound and applying our budget,

For $G\neq0$, choose $D_F=-\eta G/\|G\|_F$. Substitution gives $\langle G,D_F\rangle_F=-\eta\|G\|_F$, exactly the lower bound. There is our minus sign and our best direction.

Ordinary SGD uses $-\alpha G$: the same direction, with step length $\alpha\|G\|_F$ instead of a fixed $\eta$.

Our budget so far measures the combined size of the changes to the matrix entries. There is another useful question: how much will this update change what the layer outputs?

To see the difference, feed the same input vector $x$ into the layer before and after the update. Here $x$ means the activations entering this layer. For the linear part of the layer, call the output $z$:

The old output $Wx$ cancels. An unchanged bias would cancel too. What remains, $Dx$, is exactly the change caused by updating the weights. Even in one dimension, this distinction matters: changing a weight from 2 to 2.1 changes the output for input 3 from 6 to 6.3. The weight changed by 0.1; the output changed by $0.1\times3=0.3$.

So we can choose a budget based on the update's effect on the layer's outputs. We need it to cover different inputs, not just one example. That leads to the next question: which input direction does $D$ change the most?

§Measure the largest output change

Keep the same loss prediction, but measure the update by its largest output change per unit of input length.

We must fix the input length before comparing directions. Otherwise, simply doubling $x$ doubles $Dx$, even though the update $D$ has not changed. Try every input with length 1 and measure the resulting change $Dx$. The largest length we find is the spectral norm, also called the Euclidean operator norm:

Here $\|x\|_2$ is Euclidean vector length (the size of the vector, $\sqrt{x_1^2+\cdots+x_n^2}$). Fixing it at 1 lets us compare directions fairly.

A spectral budget $\|D\|_{\mathrm{op}}\leq\eta$ guarantees $\|Dx\|_2\leq\eta\|x\|_2$ for every input. It controls this layer's worst output change per unit input. It does not, by itself, control the whole network's loss curvature.

For a diagonal update, the two budgets are easy to see:

Set $\eta=1$. Each point in the animation is a candidate update $(d_1,d_2)$. Shading marks allowed steps; the outline is the budget limit. Frobenius gives a disk, and the spectral norm gives a square.

Write the score as $c(d_1,d_2)=g_1d_1+g_2d_2$. The blue arrow points along $-g=(-g_1,-g_2)$, where this score decreases. Its drawn length is just for readability. The orange dot is a candidate update $D$, and the dashed line passes through all updates with the same score as that dot.

Start at the origin: no weight change, score zero. Now slide the dashed line in the arrow's direction. More negative scores mean larger predicted decreases. Stop when the line is about to lose contact with the allowed region. The final orange dot is the best allowed update. The arrow tells us which way is better; the boundary decides how far we can go.

Keep our loss prediction in view:

We changed which steps are allowed. The highlighted score is still the Frobenius inner product.

§Choose a step in each independent direction

The diagonal examples show what the new budget allows. We now need the best step for a general matrix, whose useful directions may not line up with the coordinate axes:

The useful fact is that we can choose perpendicular input and output coordinates in which the matrix only stretches directions. This is the singular value decomposition, or SVD. It gives unit input directions $v_i$, unit output directions $u_i$, and stretch factors $\sigma_i$, called singular values. In these coordinates, $Gv_i=\sigma_i u_i$.

We will use the SVD theorem without proving it here. It says that every real matrix, square or rectangular, can be written as $U\Sigma V^\top$. You can accept that result for this derivation, or follow the construction in Gilbert Strang's SVD lecture. What matters for our optimizer is what the factors do.

For an $m\times n$ matrix, the full factorization has $U\in\mathbb R^{m\times m}$, $\Sigma\in\mathbb R^{m\times n}$, and $V\in\mathbb R^{n\times n}$. The middle factor has the same rectangular shape as the original matrix, with nonnegative entries on its diagonal and zeros everywhere else. $U$ and $V$ are orthogonal, so they preserve lengths.

Read $U\Sigma V^\top$ right to left: $V^\top$ changes input coordinates, $\Sigma$ stretches each coordinate, and $U$ changes to output coordinates. Those coordinate changes can include reflections. The original matrix need not look diagonal at all.

For $G=\begin{bmatrix}3&0\\0&1\end{bmatrix}$, the useful directions are already the axes. It triples one and leaves the other unchanged, so $U=V=I$ and $\Sigma=G$. A less conveniently aligned matrix needs the coordinate changes too.

The animation gives us the directions; now we can write the score in those coordinates. The compact SVD keeps only the $r$ directions with positive stretches.

For nonzero $G$, its compact SVD is

Here $r$ counts the independent active directions, and $\Sigma_r$ stores the positive stretches on its diagonal. The columns of $U_r$ and $V_r$ are the output and input directions.

The largest $\sigma_i$ is $\|G\|_{\mathrm{op}}$. Combining all their squared lengths gives $\|G\|_F=\sqrt{\sum_i\sigma_i^2}$.

Unlike eigenvalue diagonalization, SVD uses separate input and output coordinates. That is why the same construction works for rectangular matrices too.

Now the score separates into those directions:

Read one term before adding them all. The vector $Dv_i$ is the output change caused by the update when the input is $v_i$. The scalar $u_i^\top Dv_i$ is its signed component along the unit output direction $u_i$. Positive means along $u_i$; negative means against it. The positive number $\sigma_i$ weights that component in our score.

Our spectral budget limits the output change for every unit input. Each $v_i$ is one of those inputs, so

A signed component cannot have greater magnitude than the vector itself. A vector of length at most $\eta$ therefore cannot project farther than $-\eta$ in the opposite direction. Written as inequalities,

The most negative value occurs when $Dv_i=-\eta u_i$: use the full allowed length and point exactly against $u_i$. Multiplying each bound by $\sigma_i>0$ preserves its direction. Adding gives

That is a lower bound. To prove it is the best attainable score, we still need one matrix that sends every active $v_i$ to $-\eta u_i$ without breaking the budget. Choose

To check this, multiply the proposed update by one input direction $v_i$. Let $e_i$ be the coordinate vector with a 1 in position $i$ and zeros elsewhere. The columns of $V_r$ have unit length and are perpendicular to one another, so their dot products with $v_i$ give

Multiplying a matrix by $e_i$ selects its $i$th column, so $U_re_i=u_i$. Substitute those two facts, one at a time:

Thus this single update matrix sends every active input direction to the opposite output direction with length $\eta$. Every term in the score reaches its lower bound.

It also respects the budget for combinations of inputs. $V_r^\top$ extracts the coordinates in the active input directions without increasing length; $U_r$ maps those coordinates into orthonormal output directions without changing their length. Multiplication by $\eta$ scales that length by $\eta$. So $\|D_*x\|_2\leq\eta\|x\|_2$ for every $x$, with equality on each active $v_i$.

We have both pieces: no allowed update can score lower, and this update achieves the bound. After factoring out the minus sign and step size, its effect is to replace the positive stretches (the singular values stored in $\Sigma_r$) by ones:

Call $U_rV_r^\top$ the polar factor, written $\operatorname{polar}(G)$. It keeps the active input and output directions and replaces their positive stretches by 1. We have derived the matrix transformation at the center of Muon. Next we need to decide which gradient estimate to feed into it, and how to compute it cheaply.

§Give the direction some memory

So far we have treated $G$ as the direction to trust. In training, we usually compute it from a minibatch, a small sample of the data. Different samples give different gradients. Before applying the polar transformation, combine recent gradients into a steadier estimate $M_t$:

Start with $M_0=0$ and choose $0\leq\beta<1$. With $\beta=0.8$, keep 80% of the old estimate and add 20% of the incoming gradient. The animation builds that sum.

Now put that memory on a loss surface. The dark path follows each noisy gradient; the blue path averages them before stepping. Both start at the same weights, use the same learning rate, and receive the same sequence of batch perturbations. Each computes its gradient at its own current position.

Watch the dark path bounce across the valley while the blue path filters those changing directions. Then remove the noise and try longer memory. The blue path can pass the minimum because its stored direction takes time to reverse. Momentum reduces jitter in this example; it does not guarantee that we avoid overshooting. The learning rate and memory length both matter. Why Momentum Really Works explores this behavior in detail.

Use this momentum estimate in our update:

Momentum followed by orthogonalization gives the idealized core of Muon. Its step optimizes the local score built from $M_t$, so it need not decrease the current minibatch's loss.

How much history is in that average, and why this order?

Expanding the recurrence gives $M_t=(1-\beta)\sum_{s=1}^t\beta^{t-s}G_s$. Older gradients fade geometrically. The weights sum to $1-\beta^t$, so starting at zero pulls early estimates toward zero.

Average first, then orthogonalize. These operations do not commute. In one dimension, averaging $10$ and $-1$ gives a positive number whose sign is $+1$. Averaging their signs gives zero. The sign function is the scalar version of the polar transformation.

So far, our optimizer has a clean formula and an SVD bill. Let's deal with the bill.

§Change the stretches without finding the directions

Computing a full decomposition at every training step would be expensive. But look at what we actually need from the momentum matrix:

The step from $\Sigma_r$ to $I_r$ is the important choice here. Let's work through why we made it, starting with two numbers instead of the whole factorization.

Suppose the momentum matrix is diagonal, with stretches (singular values on the diagonal of $\Sigma_r$) equal to 10 and 1. Its input and output directions are just the two coordinate axes. Set the allowed output-change budget to $\eta=1$. We are choosing an update, so we may choose its length in each of those directions.

Recall our score: multiply each slope by its proposed change, then add. Here we use the momentum values 10 and 1 as averaged slopes. A more negative total is better according to that score. Consider these two choices:

The first choice preserves the momentum's 10-to-1 ratio. The second moves farther along the weaker direction. Both obey the same budget, and the second has the more negative score. Keeping that ratio would give up an improvement we are allowed to take.

Why doesn't moving in both directions break the budget? Take any unit input $x=(x_1,x_2)$, so $x_1^2+x_2^2=1$. The second update sends it to $(-x_1,-x_2)$, whose length is still 1. The input cannot have length 1 along both axes at once. That is why the spectral budget permits an update length of 1 along each axis. A Frobenius budget would instead measure the entries together as $\sqrt{1^2+1^2}=\sqrt2$ and reject this update when $\eta=1$. These are different definitions of a small update.

Now connect those numbers to the general matrix. Recall that the SVD describes each input direction $v_i$ and its corresponding output direction $u_i$:

Here $\sigma_i>0$ is the original stretch. It tells us what $M$ does; it does not require our update $D$ to use that same length. Introduce $a_i\geq0$ as the update length we choose in direction $i$. To move against the momentum, we ask for

This is our proposed update rule, not another property of the SVD. It means: when the input points along $v_i$, the change in the layer's output points opposite $u_i$, with length $a_i$. One matrix that does this in every active direction is

As in the multiplication we checked earlier, $V_r^\top$ selects a coordinate, the diagonal matrix scales it by $a_i$, and $U_r$ turns it into $u_i$. For this choice of $D$, the largest output stretch is $\max_i a_i$. The budget therefore says each $a_i$ may be at most $\eta$.

The score also comes from earlier: we used $\langle G,D\rangle_F$ to predict the loss change. After introducing momentum, we use $\langle M,D\rangle_F$ to choose a step from the averaged slopes. This is a momentum-based score, not an exact prediction from the current gradient. Expanding it in the singular directions gives the same calculation as before:

The last equality uses $u_i^\top u_i=1$. Each positive $\sigma_i$ multiplies a negative term, so increasing $a_i$ makes the score smaller. We can increase every $a_i$ to $\eta$ without exceeding the spectral budget. The earlier lower-bound argument shows this choice is optimal even among updates outside this proposed form.

All chosen lengths are now equal. Factoring out their common size $\eta$ leaves ones on the diagonal:

That is why we replace the positive stretches (the singular values stored in $\Sigma_r$) by 1 in $Q$. We keep the directions and choose new update lengths to fit our budget. The learning rate supplies the common size afterward. Only the middle factor becomes an identity; $Q$ still maps the input directions to their corresponding output directions and generally is not the identity matrix.

Newton–Schulz offers a way to approach this target using repeated matrix multiplications. Before deriving its update, we need two tools: a multiplication that changes the stretches without changing their directions, and a scalar rule that brings each stretch toward 1 and keeps it there.

Start with the multiplication. Let $X$ be our current approximation. The following identity changes its stretches without computing the directions explicitly.

If $X=U_r\Sigma_rV_r^\top$, then

Multiplying by the transpose and then itself cubes every singular value while retaining the singular directions. That means

We can apply $p(s)=as+bs^3$ to all singular values using matrix multiplications alone. We never have to compute $U_r$ or $V_r$. That leaves a scalar design problem: choose $p$ so repeating it brings positive stretches (singular values stored in $\Sigma_r$) to 1.

Show the factors cancelling

The compact SVD has $V_r^\top V_r=I_r$ and $U_r^\top U_r=I_r$. Expand the product:

$$\begin{aligned} &XX^\top X\\[8pt] &=(U_r\Sigma_rV_r^\top)(V_r\Sigma_rU_r^\top) (U_r\Sigma_rV_r^\top)\\[10pt] &=U_r\Sigma_r\underbrace{V_r^\top V_r}_{I_r} \Sigma_r\underbrace{U_r^\top U_r}_{I_r}\Sigma_rV_r^\top\\[10pt] &=U_r\Sigma_r^3V_r^\top. \end{aligned}$$

The same argument produces higher odd powers of $\Sigma_r$. If a polynomial produces a negative value, that sign can be absorbed into a singular vector; the actual singular value is its magnitude. Our classical cubic below stays nonnegative on the interval we use.

§Make one the destination, and keep it there

Now we can work on one stretch at a time. If all the positive diagonal entries approach 1, then

That is why 1 is the destination. Reaching it in the scalar calculation gives us the polar direction in the matrix calculation. We still need a rule for getting there.

Take one positive stretch $s_0$ and repeatedly feed the result of a function back into that same function:

In general, $s_{k+1}=p(s_k)$. Here $k$ counts these inner iterations, not training steps. We want the numbers to approach 1. Once a number gets there, another iteration should leave it there, so we require $p(1)=1$. A value that a function leaves unchanged is called a fixed point.

Staying at the destination is only half the job. Nearby values also need to move toward it. For example, $p(s)=\tfrac12s+\tfrac12$ sends $0.9$ to $0.95$, then $0.975$, halving the distance to 1 each time. But $p(s)=2s-1$ sends $0.9$ to $0.8$, then $0.6$. Both fix 1; only the first attracts nearby values. A fixed point can be a place you approach or a place you get kicked away from.

Watch a value go into the function, then come back as its next input

A fixed point is where $y=p(s)$ meets $y=s$. Orange goes vertically to evaluate the function, then horizontally to the diagonal to reuse the output. The number also moves into the formula above the graph, so the geometry and substitution tell the same story.

Play animationRestart← BackNext →Plug in the number. Reuse the result.

We want the attracting behavior, with fast error reduction near the target. Write the current stretch as $s=1+e$, where $e$ is its error. The local approximation we used for the loss now tells us how the next error behaves. Assuming $p(1)=1$,

So $p'(1)$ multiplies the leading part of the error. A slope with magnitude below 1 shrinks nearby errors; choosing $p'(1)=0$ removes that linear term altogether. For the cubic $p(s)=as+bs^3$ that our matrix multiplication can compute, the two requirements are

Substituting the cubic and its derivative gives

Subtract the first equation from the second: $2b=-1$. So $b=-1/2$ and $a=3/2$. Putting these coefficients back into the scalar rule and its matrix counterpart gives the classical Newton–Schulz iteration:

Normalize a nonzero momentum matrix with $X_0=M/\|M\|_F$. Since $\|M\|_{\mathrm{op}}\leq\|M\|_F$, its singular values start in $[0,1]$. Repeated cubic steps bring each positive one toward 1 and keep zero at zero. Consequently, the middle factor approaches $I_r$, and $X_k$ approaches $Q=U_rV_r^\top$. The scalar fixed point gives us the matrix direction we wanted. Zero stretches stay zero, so this also preserves the active subspace when the original matrix has missing directions.

Why do the iterates actually reach one?

For $0<s<1$,

$$\begin{aligned} p(s)-s&=\tfrac12s(1-s^2)>0,\\[10pt] 1-p(s)&=\tfrac12(1-s)^2(s+2)\geq0. \end{aligned}$$

The value increases without crossing 1. Its limit must be a fixed point; the only positive one in this interval is 1. The squared error factor gives quadratic convergence near 1. Very small starting values can still need many iterations to get close.

More generally, $|p'(1)|<1$ gives local attraction. Setting it to zero improves the local error reduction; it does not, on its own, prove that arbitrary starting values converge. That is why we checked the interval.

§Get a useful approximation in five steps

We now have an iteration that converges, but a tiny initial stretch can take many steps to get close to 1. Every training step contains this inner matrix calculation. We need to stop it somewhere, or we will spend the training budget polishing one update. The original Muon experiments found five inner steps sufficient for the models tested. Five is a compute-budget choice, not a convergence guarantee. They use this tuned quintic:

Its slope at zero is $3.4445$, versus the cubic's $1.5$, so small values initially grow faster. There is a catch: $q(1)=0.701$. One is not even a fixed point.

The bargain is fast early progress toward unit-sized stretches, with some error left over. Running this same quintic forever does not finish the job; it can keep oscillating. Our two cubic conditions do not derive these tuned coefficients. The Newton–Schulz notes explain this choice and the resulting oscillation.

Play animationRestart← BackNext →Show five stepsTry start at 1Substitute, calculate, and reuse each result.

A quintic $q(s)=as+bs^3+cs^5$ has three coefficients. We could impose $q(1)=1$ and $q'(1)=0$, then use the remaining freedom to shape the interval. Muon's original coefficients make a different tradeoff. Polynomial design includes choosing the conditions, not only solving for the coefficients.

Show the matrix code and the practical optimizer details

Here is the five-step direction calculation in float32. The reference implementation uses bfloat16 for speed.

import torch

@torch.no_grad()
def muon_direction(M, steps=5, eps=1e-7):
    """Approximate polar direction for one finite 2D matrix."""
    assert M.ndim == 2
    X = M.float()
    transposed = X.shape[0] > X.shape[1]
    if transposed:
        X = X.T
    X = X / (X.norm() + eps)  # Frobenius norm; zero stays zero.
    a, b, c = 3.4445, -4.7750, 2.0315
    for _ in range(steps):
        A = X @ X.T
        X = a * X + (b * A + c * (A @ A)) @ X
    return X.T if transposed else X

Transposing a tall matrix makes $XX^\top$ smaller and cheaper to multiply; the return restores its shape. The quintic uses the same SVD algebra, adding a $\Sigma_r^5$ term.

The reference optimizer also uses a Nesterov-style blend $H_t=(1-\beta)G_t+\beta M_t$ by default, applies matrix-shape scaling, and supports decoupled weight decay. The finite quintic output need not have spectral norm exactly 1.

The reference setup uses Muon for hidden weight matrices and another optimizer, such as AdamW, for embeddings, the output layer, and scalar or vector parameters. The function above computes the direction, not the full training recipe.

§Put the pieces together: one Muon step

We have enough parts. Time to assemble the optimizer.

At training step $t$, backpropagation gives the gradient $G_t$. Blend it into the momentum, normalize that matrix, and run five polynomial steps. Then scale the result by the learning rate and subtract it from the weights:

Run the recurrence for $k=0,\ldots,4$, with $a=3.4445$, $b=-4.7750$, and $c=2.0315$. The tiny $\varepsilon$ prevents division by zero. The optional decay $\lambda$ shrinks the old weights separately from the gradient calculation. Set it to zero to recover the update we derived.

This is a minimal Muon step. Practical variants can blend in the current gradient again with Nesterov momentum, scale the update for the matrix shape, or use different polynomial schedules. The reference implementation makes those choices explicit. The central operation stays the same: turn the momentum matrix into an approximate polar direction.

Here $Q=UV^\top$ is the exact polar factor of the momentum matrix, computed as a reference for the illustration. Muon computes $X_5$ instead. The solid arrows approach the dotted ones without landing exactly on them, because five quintic steps leave residual stretch.

Momentum chooses which matrix we trust. Normalization puts its stretches in range. The polynomial reshapes them, and the minus sign turns the result into a weight update. The updated weights themselves need not be orthogonal.

§Same gradients, different optimizers

We have assembled a complete Muon step. To see what its matrix transformation changes, give the same gradient sequence to three update rules.

SGD uses the gradient directly. Adam keeps moving averages of both gradients and their elementwise squares, then divides each coordinate by its own estimated scale. AdamW adds weight decay separately. Muon rescales the momentum matrix along its singular directions.

Let $R_t$ denote the gradient-driven part before the learning rate and decay:

An entry is one number in the matrix, such as the number in row 1, column 2. The symbol $\oslash$ means dividing each number by the number in the same position in another matrix. AdamW keeps a separate running scale for each position. For example, if two momentum entries are 4 and 1, and their estimated scales are 4 and 1, their contributions become $4/4=1$ and $1/1=1$. Those scales come from past squared gradients; this is an example of the arithmetic, not a fixed rule that all entries become 1.

That is entrywise scaling: rescaling each matrix entry separately. Muon instead uses the matrix's singular directions, which can involve several entries moving together. The next comparison asks whether that distinction changes the path through a loss function.

The first gradient has four positive entries. SGD keeps their relative sizes. AdamW's first bias-corrected gradient contribution is nearly the sign of each entry, so it is nearly rank one here. Muon's result keeps both matrix directions active. That difference comes from the rules, before we have trained anything.

The useful conclusion is that Muon gives weaker singular directions room to contribute instead of letting the strongest one dominate the update. That is what the rounder shape shows. It does not mean every step is smaller or gentler. These pictures normalize away the overall size; use the printed norms to compare size. The learning rate still decides how far we move, and a large rate can overshoot.

§Now give them a loss to minimize

The replay showed how the same gradients produce different steps. During training, those steps also change where the next gradient is evaluated. To see that feedback, give all three optimizers an actual loss and let each follow its own path.

Here is a quadratic loss with a narrow valley. In the two matrix directions shown below, it is

The bottom is at $(0,0)$. The loss curves forty times more sharply along $a$ than along $b$. All three methods start at $(1.5,1.5)$, then recompute their gradients at their own new weights. Each curve is an actual optimization run.

Why compare two ways of writing this same valley? AdamW gives each matrix entry its own scale. We want to see whether its behavior changes when one useful movement is spread across several entries instead of stored in just one. Muon's rule uses singular directions, so this is a concrete way to compare what the two rules pay attention to.

Here $a$ measures movement across the steep side of the valley, and $b$ measures movement along its shallow side. The optimizer updates the entries of $W$, not the letters $a$ and $b$ directly. We can encode those two coordinates in $W$ in either of these ways.

Separate entries. Put $a$ in the top-left entry and $b$ in the bottom-right entry. Hold $a=1$ and increase only $b$ from 0 to 0.2:

Only the bottom-right number changes. AdamW's scale for that entry is based on gradients for the shallow movement alone.

Mixed entries. Store the same coordinates using $W_{11}=W_{22}=(a+b)/2$ and $W_{12}=W_{21}=(a-b)/2$. Make exactly the same change, $a=1$ and $b:0\to0.2$:

Now the shallow movement needs all four numbers to change: add 0.1 on the diagonal and subtract 0.1 off it. Increasing $a$ would instead add to all four entries. Each position therefore participates in both the steep and shallow movements. Its gradient history mixes their contributions, and AdamW's separate scale for that position cannot assign an independent scale to each movement.

The button Use separate entries for a and b switches from the second encoding to the first. It changes the formula for $W$ and reruns all three optimizers. It does not change $L(a,b)$, so the bowl itself stays the same. Watch whether the optimizer follows the same path through that bowl. The formula and example matrix above the plots show the selected encoding; the matrices below show the actual weights during the run.

Both encodings start from $a=b=1.5$, which gives $W=1.5I$ in either case. That is why the starting weights look identical even though subsequent updates can differ. With the initial mixed encoding and learning rates, SGD bounces across the steep direction, AdamW makes slow progress along the shallow direction, and Muon gets farther down the valley in 80 steps. Use the switch to test how much of that result depends on the encoding.

Now press Use separate entries for a and b. The valley has the same shape, but its directions line up with the individual entries AdamW rescales. AdamW becomes much faster; SGD and Muon follow the same paths as before. Or keep the mixed basis and raise AdamW's learning rate to $0.32$. It catches up and finishes below Muon. AdamW has declined our invitation to be deprecated.

Muon responds to the matrix's coupled directions, so rotating this example does not change its path in valley coordinates. That is a precise sense in which its update follows the matrix geometry. Calling it "more natural" is less useful than seeing which changes of coordinates leave its behavior intact.

There is only one minimum in this quadratic bowl. None of these paths demonstrates escaping a local minimum. At a point with zero gradient and no stored momentum, Muon also has no direction to follow. Momentum can carry an optimizer through a flat or turning region, but replacing singular values by one does not by itself guarantee escape.

The test answers the question we started with. Here, changing how the two movements are distributed across matrix entries changes AdamW's path, while SGD and Muon keep the same paths in the valley. AdamW scales those entries separately; Muon rescales the singular directions that combine them. The learning-rate control also shows that this difference does not establish a universal winner. We count training steps here; a Muon step does more work, so these curves do not measure wall-clock speed.

The matrix loss, gradients, and exact settings

The full problem has a $2\times2$ weight matrix:

$$\begin{aligned} L(W)&=\tfrac12\operatorname{tr}(W^\top H W),\\[10pt] G&=HW,\\[10pt] H&=B\begin{bmatrix}40&0\\0&1\end{bmatrix}B^\top. \end{aligned}$$

The mixed basis uses $B=\tfrac1{\sqrt2}\begin{bmatrix}1&-1\\1&1\end{bmatrix}$; the aligned basis uses $B=I$. All methods start at $W_0=1.5I$. For these two choices of $B$, their iterates stay in the slice $W=B\operatorname{diag}(a,b)B^\top$. Substitution gives the two-dimensional loss plotted above. The code updates all four matrix entries, then reads off $a$ and $b$.

SGD has no momentum and starts with learning rate $0.04$. AdamW starts at $0.08$ or $0.32$, with $\beta_1=0.8$, $\beta_2=0.9$ and $\varepsilon=10^{-8}$. Muon starts at $0.08$, with $\beta=0.8$, five quintic iterations and $\varepsilon=10^{-7}$, without Nesterov or a shape multiplier. Every method has zero weight decay, so AdamW is equivalent to Adam in this example. Momentum histories start at zero; AdamW applies bias correction.

Each method scales its initial rate by the same schedule over 80 training steps:

$$\eta_t=\eta_0\left[0.1+0.45\left(1+\cos\frac{\pi t}{80}\right)\right].$$

Here $t=0,\ldots,79$. These settings make the geometric effect visible. The basis and learning-rate controls let us test how fragile any claim about the winner would be. The five inner polynomial iterations belong to each Muon training step; they are a different count from the 80 steps on the horizontal axis.

That is the optimizer we set out to invent. Use the local loss prediction to choose a direction, decide what "small" means for a matrix, remember recent gradients, and approximate the resulting polar transform cheaply. The remaining question is how well those choices work on the model we actually want to train.

§Where these choices came from

We reconstructed the update by choosing a loss prediction, an output-change budget, a gradient average, and a short polynomial calculation. That is a route to understanding the method. For its historical development and experimental results:

• Keller Jordan's original Muon writeup describes the optimizer and its early experiments. Start here for the original account.
• Modular Duality in Deep Learning, by Bernstein and Newhouse, develops the broader norm-based framework and matrix iterations behind this route.
• Muon is Scalable for LLM Training, by Liu and collaborators, studies weight decay, update scaling, distributed implementation, and the Moonlight models.
• Practical Efficiency of Muon for Pretraining, by Essential AI, examines compute versus wall-clock time, large batches, and hyperparameter transfer with μP.
• The Polar Express, by Amsel and collaborators, designs iteration-dependent polynomials for faster polar approximation and addresses low-precision computation.

If I had to rebuild the argument, I'd start with $Dx$: how much should one update be allowed to change the layer's output?