Muon

The problem, from first principles

Adam: full update rule scales each parameter's step by its OWN historical gradient magnitude, but treats every entry of a weight matrix independently, entry-by-entry, with no notion that the matrix as a WHOLE has a shape and a structure (rows, columns, singular directions). For 2D weight matrices specifically (the vast majority of a transformer's parameters), this leaves real structure on the table: some directions in weight-space matter far more than others for how the layer actually transforms its input, and Adam's per-entry scaling doesn't distinguish between them.

Muon (short for "MomentUm Orthogonalized by Newton-schulz") takes a different, matrix-aware approach: instead of adapting each entry independently, it takes the momentum-accumulated update (a full matrix) and ORTHOGONALIZES it, rescaling it so every one of its singular values moves toward 1, before using it to step. This has the effect of treating every "direction" the update wants to move in roughly equally, rather than letting a few dominant directions (large singular values) drown out the rest, the same "normalize before combining" spirit 04-feature-scaling's standardization applies to input features, applied here to an optimizer's own update matrix instead.

From theory to code

Theory uses a fast, fixed-coefficient iterative approximation (a specific quintic polynomial map, applied a handful of times) to push a matrix's singular values toward 1 cheaply, using only matrix multiplications, no actual SVD ever computed. This is deliberately NOT run to full convergence, a few steps get most singular values substantially closer to 1 than they started, which is all Muon's update actually needs.

Implement newton_schulz_orthogonalize(G, steps=5, eps=1e-7) first, then muon_step(param, grad, momentum_buf, lr, momentum=0.95) on top of it.

Constraints

• newton_schulz_orthogonalize must handle tall, wide, AND square matrices, always returning a result the same shape as G.
• Normalize G by its Frobenius norm (np.linalg.norm(G)) before iterating.
• muon_step accumulates momentum first (SGD + Momentum's own formula), THEN orthogonalizes the resulting momentum buffer, never the raw gradient directly.

Hints

Open one at a time. Each gives away a little more than the last.