Deriving Muon | Chris Ge

Goal: keep the RMS norms of each activation small (bounded by 1)
For determining \(\Delta W\), the change to the weight matrix, choose it to maximize the effect on the loss while maintaining that the output doesn’t change too much
how much does it change? scales with the operator norm of \(W\)
- hence we choose the update \(\Delta W\) to minimize \(\langle G, \Delta W \rangle\) (make loss go down) subject to \(\lVert \Delta W \rVert_{\text{RMS} \to \text{RMS}} \le \eta\), where \(G\) is the gradient

Two (three) ways to view the closed form of this constrained optimization

If the gradient’s SVD is \(G = U \Sigma V^T\), then \(\Delta W = \eta \, U V^T\): “orthogonalization” because it’s just setting all the singular values to \(\eta\).
\(\Delta W = \eta \, G (G^T G)^{-1/2}\). This is useful for seeing that \(\Delta W\) depends only on the gradient (and \(\eta\)).
The third, less important interpretation is that \(\Delta W\) is the closest orthogonal matrix (as measured in Frobenius norm) to the gradient. This is just the formulation given in the second blog post https://kellerjordan.github.io/posts/muon/.

Discussion / Comparison to SGD, external to the paper

How does this compare to a variational formulation of SGD? The insight is that SGD treats weights as a vector (we can flatten a matrix into a vector if needed), whereas Muon treated weights as a matrix (connects to Why Muon?)
can write SGD as \(\min_{\Delta v} \langle g, \Delta v \rangle\) subject to \(\lVert \Delta v \rVert_2^2 \le \eta^2 \lVert g \rVert_2^2\) (proof that this gives you back SGD as closed form by Cauchy–Schwarz)

Why use Muon? It converges much faster than AdamW

in part because it’s basing the updates off the whole matrix gradient, and using matrix properties (i.e. SVD) of the matrix, while SGD/AdamW updates weights component by component, as vectors
this results in the AdamW updates being close to low rank, i.e. for some principal components, there’s low singular values so little learning in those directions, while Muon made all singular values equal

The rest of Muon is just this neat trick (Newton–Schulz) to efficiently compute the \(UV^T\) term. Apparently if you just apply an odd polynomial \(p\) (only odd-degree terms) to \(U \Sigma V^T\), then \(p(U \Sigma V^T) = U \, p(\Sigma) \, V^T\), so we just choose \(p\) such that \(p(p(\cdots p(\Sigma) \cdots)) \approx I\) and iteratively apply \(p\) to \(U \Sigma V^T\)