Muon: updating a matrix

Checked: 2026-10-10 · Prerequisite: momentum and AdamW · For a first pass, read the first 2 sections and the comparison table.

A linear layer has a weight matrix and a gradient matrix. AdamW adapts updates using each coordinate's history; Muon preserves matrix structure and changes the geometry of the update. It changes parameter updates, not the loss or the attention computation.

Start with a 2 × 2 example

Suppose the momentum update is

$$M=\begin{bmatrix}6&0\\0&2\end{bmatrix}.$$

Both directions receive an update, but the first has 3 times the magnitude. Dividing the matrix by its norm shrinks both together; the ratio stays 3:1.

The idealized Muon operation differs. Write the singular value decomposition as $U\Sigma V^\top$, then replace each nonzero singular value with 1. Keeping only the corresponding singular vectors $U_r,V_r$ gives $U_rV_r^\top$. In this full-rank example the result is the identity, with a 1:1 ratio. Left and right singular vectors are preserved; the entries are not all set to 1. The original author's explanation develops this orthogonalization view.

This explains the operation, not a guarantee of better learning. A small singular value might describe a useful direction or mostly noise. Equalizing magnitudes does not make the gradient more accurate for free.

For $\operatorname{diag}(2,0)$, our reference instead gives $\operatorname{diag}(1,0)$, not the identity: the zero direction stays zero. Multiplying all the returned $U,V$ vectors can include arbitrary completions associated with zero singular values. The code below uses a threshold to decide which singular values to retain.

What happens in one update?

loss → gradient matrix → momentum → approximate orthogonalization → scaling → weight update
                                                                                 ↑
                                                                       decoupled weight decay

A simplified version without Nesterov is

$$M_t=\mu M_{t-1}+G_t,$$ $$O_t\approx\operatorname{Polar}(M_t),\qquad W_t=(1-\eta_t\lambda)W_{t-1}-\eta_t s(m,n)O_t.$$

Here $s(m,n)$ depends only on the matrix row and column counts. $\mu$ is momentum, $\lambda$ is weight decay, and $s$ is the implementation's scaling factor. Actual code may use Nesterov or an EMA momentum convention; mixing equations from different implementations does not establish step-by-step equivalence.

The PyTorch Muon documentation lists these choices separately, including learning-rate adjustments. Record the library version and configuration rather than treating one page's defaults as a universal definition.

What does Newton–Schulz approximate?

Running SVD every step is expensive. Muon instead uses matrix-multiplication-based Newton–Schulz iterations on the normalized update. Starting with $X_0=M/(\lVert M\rVert_F+\epsilon)$, one family of iterations is

$$A_k=X_kX_k^\top,\qquad X_{k+1}=aX_k+(bA_k+cA_k^2)X_k.$$

The operation acts on the update, not on the weight matrix $W$. A rectangular matrix also cannot satisfy both identity conditions: full column rank can give $O^\top O=I$, while full row rank can give $OO^\top=I$.

An important detail in the pinned original implementation: the quintic coefficients (3.4445, -4.775, 2.0315) prioritize useful behavior within a small iteration budget. Repeating them does not guarantee exact convergence to the ideal polar factor. Substitute a singular value of 1:

$$3.4445-4.775+2.0315=0.701.$$

One is not even a fixed point of that polynomial. “More iterations must converge exactly to 1” therefore cannot be the explanation. Approximate orthogonalization is not exact orthogonalization at arbitrary iteration counts.

import math

diagonal = [6.0, 2.0]
scale = math.sqrt(sum(value * value for value in diagonal)) + 1e-7
transformed = [value / scale for value in diagonal]
for step in range(5):
    transformed = [3.4445 * value - 4.775 * value**3 + 2.0315 * value**5
                   for value in transformed]
assert math.isclose(transformed[0], 0.7530, abs_tol=1e-4)
assert math.isclose(transformed[1], 1.1337, abs_tol=1e-4)

Check the geometry with SVD

This is an educational exact reference, not an efficient Muon implementation. It omits momentum and weight decay, requires PyTorch, and accepts real matrices only. For rank-deficient inputs it preserves the null space rather than inventing directions for zero singular values.

import torch


def polar_reference(matrix):
    try:
        inferred = torch.as_tensor(matrix)
    except RuntimeError:
        inferred = torch.as_tensor(matrix, dtype=torch.float64)
    if torch.is_complex(inferred):
        raise ValueError("expected a real matrix, not complex values")
    matrix = torch.as_tensor(matrix, dtype=torch.float64)
    if matrix.ndim != 2 or matrix.numel() == 0:
        raise ValueError("expected a nonempty real matrix")
    if not torch.isfinite(matrix).all():
        raise ValueError("matrix must be finite")
    scale = matrix.abs().max()
    if scale == 0:
        return torch.zeros_like(matrix)
    left, singular, right = torch.linalg.svd(matrix / scale, full_matrices=False)
    threshold = max(matrix.shape) * torch.finfo(matrix.dtype).eps * singular.max()
    retained = (singular > threshold).to(matrix.dtype)
    return (left * retained) @ right


direction = polar_reference([[6.0, 0.0], [0.0, 2.0]])
assert torch.allclose(direction, torch.eye(2, dtype=torch.float64))
assert torch.count_nonzero(polar_reference(torch.zeros(2, 3))) == 0

Try [[1, 1], [0, 1]]: the result is not the elementwise sign. Near-zero singular values depend on precision and thresholds; this reference's threshold convention is not a compatibility guarantee for a production optimizer.

Why isn't this just changing an optimizer name?

First, scaling. For a full-rank $m\times n$ matrix, the idealized $O$ has $\min(m,n)$ singular values equal to 1, so

$$\operatorname{RMS}(O)=\sqrt{\frac{\min(m,n)}{mn}} =\frac{1}{\sqrt{\max(m,n)}}.$$

With the same learning rate, larger matrices receive smaller per-entry RMS updates. The scaling paper uses $0.2\sqrt{\max(m,n)}$ to match an empirical AdamW update scale. That constant is not a theorem; the original shape adjustment and RMS matching are different configurations.

For example, the original option in PyTorch v2.9.0 uses $\sqrt{\max(1,m/n)}$. For ideal full-rank directions, before multiplying by the learning rate (reference values, not a claim that finite-step NS is exactly orthogonal):

As of October 2026, the stable PyTorch 2.14 documentation adds a third option, spectral_unclamped: the factor is $\sqrt{m/n}$, no longer clamped below at 1. The $8\times2$ row stays at 0.7071, while $2\times8$ becomes $0.3536\times0.5\approx0.1768$. Leaving the argument unset still means original.

In the v2.9.0 implementation, weight decay uses the base learning rate, not the shape-adjusted rate. With $\eta=0.1,\lambda=0.2$, the decay term is $0.98W$; doubling the update direction does not change it to $0.96W$.

Second, parameter groups. The classical hybrid uses Muon for hidden matrices and AdamW for embeddings, output layers, biases, and normalization parameters. Embeddings are themselves matrices, so ndim == 2 is not a sufficient selector. Deduplicate tied embedding/output weights by object identity to avoid two updates to the same parameter. The original repository's usage guide distinguishes these groups.

Finally, distribution. Orthogonalizing matrix shards separately generally differs from orthogonalizing the full matrix before slicing. If the framework flattens or shards parameters, establish how logical matrices are reconstructed and what communication is needed; do not copy an elementwise AdamW sharding scheme unchanged.

import math

momentum = torch.tensor([[1.0, 1.0], [0.0, 1.0]], dtype=torch.float64)
whole_update = polar_reference(momentum)
row_updates = torch.cat([polar_reference(row) for row in momentum.split(1, dim=0)])
expected_whole = torch.tensor([[2.0, 1.0], [-1.0, 2.0]], dtype=torch.float64) / math.sqrt(5)
torch.testing.assert_close(whole_update, expected_whole)
assert math.isclose(torch.dot(row_updates[0], row_updates[1]).item(), 1 / math.sqrt(2))
assert not torch.allclose(row_updates, whole_update)

What is a meaningful AdamW comparison?

A practical starting experiment fixes a small model and data order, gives both optimizers the same number of learning-rate/weight-decay trials, and reports repeated-run validation curves, step time, total time, and peak memory. This is an experiment design; this site has not run it to establish that Muon beats AdamW.

What needs qualification in 2026?

• “Embeddings must always use AdamW” is too broad. October 2026's AF-Muon studies specialized tied-embedding updates. It is a new design, not permission to apply classical Muon unchanged. Only its abstract and problem setting have been checked here; its experiments have not been reproduced.
• “Switch to Muon and scale up directly” is too optimistic. The Kimi K2 report §2.1 (July 2025) found that in a run with 9B activated and 53B total parameters, plain Muon drove the maximum attention logit past 1,000 quickly. Its MuonClip checks each head's maximum logit after every update and, above a threshold τ (100 for K2), shrinks that head's query / key projection weights. K2 used it to pre-train on 15.5T tokens and reports no loss spike. By 2026 Muon is in production-scale training of this kind: the DeepSeek-V4 model card also states that it uses Muon (embeddings and some other parameters stay on AdamW; see the V4 note). At scale, monitor new failure signals such as logit growth.
• “Muon guarantees convergence” is also too broad. An August 2026 convergence analysis examines counterexamples and error conditions in specific stochastic optimization problems. That establishes neither universal failure in LLM training nor success without assumptions.

Understand the update geometry, then ask which operation a newer variant changes. That is more useful than memorizing claims that an optimizer is obsolete. Continue with training memory and distributed training to place the costs in a complete training run.