Muon: Why the Shape of a Weight Matrix Matters

September 2026 · 4 min read

Before you start: Momentum and singular values; the chapter introduces the required matrix norms.

A weight matrix transforms vectors between feature spaces. Flattening it into independent coordinates discards that organization when choosing an update. Muon uses the matrix structure of a momentum-like direction: it approximately equalizes its nonzero singular values before applying a parameter step.

Read the operation through singular values

Write a matrix direction $M=U\Sigma V^\top$. For a full-rank rectangular matrix, its thin polar factor is $UV^\top$: singular directions remain, while singular values become 1. At rank deficiency, the canonical partial polar factor acts on nonzero singular directions and leaves null directions zero; an SVD completion can choose additional null-space directions. Numerical implementations need a defined convention.

For M=diag(3,1), the polar direction is diag(1,1). The first singular direction no longer receives three times the update magnitude of the second. This differs from dividing by the Frobenius norm, which preserves the 3:1 ratio. It also differs from elementwise sign, which depends on coordinate entries rather than singular directions.

A steepest-direction problem

The operator norm bounds the largest singular value; the nuclear norm sums singular values. With this unit constraint, aligning D’s singular vectors with G and using singular values 1 maximizes the first-order decrease from a step −ηD. This is the matrix analogue of choosing a unit Euclidean gradient direction, under a different geometry.

Practical Muon typically transforms momentum rather than the exact current gradient, so the simple steepest-descent derivation motivates the direction but is not a complete convergence proof for every implementation. Momentum, approximation, shape scaling, weight decay, and stochastic gradients all matter.

Avoid an SVD at every large training step

Newton–Schulz-style polynomial iterations approximate the polar transformation with matrix multiplications. Normalize the input to an appropriate range, iterate a chosen polynomial, and transpose rectangular matrices when useful for computational cost. The common Muon implementations use tuned finite-step approximations; do not assume every polynomial converges to the exact polar factor under arbitrary scaling.

import numpy as np

def exact_partial_polar(matrix, tolerance=1e-10):
    # Small reference implementation for checking a fast approximation.
    u, singular_values, vt = np.linalg.svd(matrix, full_matrices=False)
    keep = (singular_values > tolerance).astype(float)
    return (u * keep) @ vt

assert np.allclose(exact_partial_polar(np.diag([3., 1.])),
                   np.eye(2))

Use this small SVD reference to inspect singular values and approximation error. For a zero matrix, a zero update is a natural convention. Test rectangular and low-rank inputs too; a square full-rank example alone can hide shape or transpose errors.

Which parameters receive matrix updates?

Hidden-layer weight matrices are natural candidates. Biases, normalization scales, embeddings, and output heads may use a separate optimizer or a distinct recipe. Convolutional kernels require a declared reshape. Partitioning a large matrix into blocks changes the transformation, so distributed sharding is also an algorithmic choice, not merely a storage detail.

Shape-dependent scaling can match an update’s typical magnitude across differently shaped layers. Compare the actual parameter-update norm when changing a recipe. Using identical nominal learning rates across optimizers does not ensure comparable update sizes, and it is not a substitute for a fair tuning budget.

Measure both optimization and overhead

Record validation loss versus tokens, steps, and wall-clock time; matrix-transform overhead; peak memory; parameter groups; precision; and approximation settings. Compare across seeds and model sizes. An optimizer that reduces steps but adds expensive local matrix work may or may not win overall. The next chapter adds communication to this accounting.

Sources and further reading

• Keller Jordan: Muon implementation — Matrix-shaped updates, Newton–Schulz approximation, and parameter grouping.
• Mishra, Trivedi & Kumar: SignMuon, arXiv v1 — The specific majority-vote matrix optimizer discussed here; similarly named methods can use different algorithms.