Muon Optimizer: Audit Orthogonal Updates

Muon optimizer implementations can produce finite updates while routing the wrong parameters or scaling rectangular matrices incorrectly. This tutorial audits one complete step—momentum, Newton–Schulz orthogonalization, shape-aware scaling, weight decay, and AdamW fallback—on small matrices before any training claim.

Muon optimizer parameter routing is the first invariant

A Muon optimizer step begins before any matrix polynomial runs. It begins with a manifest that says which named parameter goes to which optimizer. Hidden weight matrices are candidates for Muon; embeddings, the output head, biases, normalization gains, and scalar parameters stay on an AdamW path. Shape alone is not enough because an embedding table is also a rank-two tensor. Name, role, rank, and an explicit exception rule belong in the receipt.

That separation is consistent with the original Muon write-up, which presents Muon for two-dimensional hidden weights while retaining AdamW for other parameters. The safe implementation turns that convention into data: each record carries a stable parameter name, semantic role, shape, route, and reason. A duplicate route or an unknown role is an error rather than an invitation to guess.

Routing deserves its own test because a misrouted embedding can still yield finite numbers. The training loop runs, loss may move, and no stack trace identifies the conceptual bug. This is why the audit lab includes an embedding-misroute mutant. It proves that the validator checks semantic role instead of merely accepting every matrix-valued parameter. This explicit AdamW parameter routing rule also keeps excluded matrix-valued parameters visible rather than silently mixing them into the Muon path.

The manifest also makes review practical. A teammate can compare the expected count and total elements on each route without reading optimizer internals. The same discipline appears in the multi-token prediction loss audit: declare membership first, then calculate. In both cases, silent inclusion is more dangerous than a loud arithmetic failure.

Write the one-step ledger before the update loop

For one routed matrix, freeze the parameter W, gradient G, prior momentum M, learning rate, momentum coefficient, Newton–Schulz coefficients, iteration count, update scale, and weight-decay coefficient. Then expose every intermediate in order. A minimal Muon optimizer ledger records the momentum state, the Nesterov combination, the working orientation, the normalized matrix, each polynomial iterate, the restored orientation, the shape-aware scale, the decoupled decay term, and the final parameter.

The ordering matters. Momentum is stateful, while the orthogonalization transform acts on the proposed update. Shape scaling is applied to that transformed direction. Decoupled weight decay acts on the parameter, not by being mixed into the gradient before the matrix transform. If those stages are collapsed into one expression, a reviewer cannot tell whether a plausible output came from the intended contract.

Use tiny matrices whose values fit on a page. The lab defaults to a 2×3 hidden matrix and constrains all dimensions to 2 through 4. Each input cell is finite and capped in magnitude. That bound is not a claim about production weights; it keeps the local computation, CSV rows, and independent oracle small enough to inspect.

The final receipt contains inputs and outputs, but it also contains stage hashes and scalar summaries: Frobenius norm, maximum absolute entry, shape, and finite-state checks. A matching final matrix without matching stages is not accepted. This mirrors the accounting principle in gradient accumulation: equivalent-looking totals do not excuse a different denominator or operation order.

Treat Newton–Schulz orthogonalization as an approximation

The Muon optimizer uses a short Newton–Schulz-style polynomial to push the singular values of an update toward a common band. The five-step default is useful precisely because it is an inexpensive approximation. It must not be described as an exact polar factor, and the audit must not require a zero orthogonality residual after five iterations.

Start by scaling the working matrix so its norm is safely inside the polynomial's operating region. For a wide matrix X, each iteration forms A = XXᵀ and then updates X with a declared polynomial in A multiplied by X. A tall matrix is transposed first, processed in the wide orientation, and transposed back. The coefficient tuple and exact iteration count are receipt fields rather than hidden constants.

The trace should show more than a final pass badge. After each Newton–Schulz orthogonalization step, report the singular-value spread, the finite-state flag, and an orientation-appropriate residual. The PyTorch Muon documentation is the API reference for the production optimizer, but the local lab deliberately implements a bounded educational ledger instead of importing a framework.

An independent oracle computes the closest orthogonal or semi-orthogonal factor for the tiny matrix with a separate Jacobi eigensolver. It reports distance between the polynomial result and that reference plus both residuals. Passing means the audited polynomial matches its own direct reimplementation and improves the declared measures within tolerance. It does not mean the approximation equals the oracle bit for bit.

Rectangular geometry decides the multiplication side

A square example hides the most important Muon optimizer implementation branch. If rows are no greater than columns, the wide working matrix can use XXᵀ on the left. If the original update is tall, transposing it produces that same wide contract. Omitting the transpose or forming XᵀX while still multiplying on the left either creates a dimension error or silently changes the transform in code that happens to broadcast or reshape.

The ledger therefore stores original shape, working shape, whether a transpose occurred, Gram shape, and multiplication side. Those five fields let an independent checker reconstruct the geometry without trusting the implementation's matrix helpers. The missing-tall-transpose and wrong-multiplication-side mutants each change one contract field and must be rejected for a distinct reason.

Orientation also determines the residual you can interpret. A wide semi-orthogonal matrix is judged with XXᵀ near identity; a tall result is judged with XᵀX near identity after restoring its original shape. Asking a rectangular matrix to satisfy both is mathematically impossible. The article's figure and lab label the relevant identity explicitly so “orthogonal” does not become a vague green badge.

This distinction is especially useful when models change width. A route that once held square projections may later include expansion matrices. Shape-aware tests catch that transition. The same practice applies to FP8 calibration range: make the numerical domain part of the release contract, because a method that behaves well in one range or shape is not automatically valid in another.

Apply update scale and decoupled decay in order

After the polynomial transform, the Muon optimizer applies a shape-aware update scale. The exact policy belongs in configuration and evidence because rectangular matrices can otherwise receive systematically different step magnitudes. The lab uses one declared rule, records its formula and value, and includes an omitted-scale mutant. The purpose is not to recommend a universal coefficient; it is to make the chosen update scale impossible to lose during refactoring.

Decoupled weight decay is a second ordering constraint. The parameter contribution is computed from the pre-step parameter and applied outside momentum and Newton–Schulz orthogonalization. Coupling decay into the gradient changes the matrix that the polynomial sees. That can alter direction, not just magnitude, while still producing a finite update.

For each cell, the receipt separates transformed update, scaled update, decay contribution, and final delta. The sum must reconcile to the emitted parameter matrix within a tight numeric tolerance. The coupled-decay mutant intentionally injects λW into the gradient before momentum. A validator that inspects only the final shape would miss it; the staged ledger identifies both the wrong ordering and the changed stage hash.

There is a useful product lesson here. Optimizer adoption is not one binary choice. Parameter routing, approximation coefficients, scale policy, decay policy, learning rate, and fallback optimizer are separate decisions with separate owners. Keeping those fields visible makes experiments comparable and rollbackable even when the high-level label remains “Muon optimizer.”

Use an independent small-matrix polar oracle

A Muon optimizer test that calls the same helper twice is not independent. The audit lab therefore keeps the production-style polynomial path separate from a small-matrix reference. The reference forms the appropriate Gram matrix, diagonalizes it with bounded Jacobi rotations, applies an inverse square-root with an explicit epsilon, and reconstructs a polar factor. It is slower and more general than needed for the main path, which is exactly why it is valuable as an oracle.

The oracle returns three comparisons: distance to the reference factor, singular-value spread, and the orientation-appropriate identity residual. Each is reported, not overinterpreted. Five polynomial steps may leave a nonzero residual. A small distance on one synthetic matrix is not evidence of training quality, convergence, or accelerator speed.

The scalable Muon study supplies important research context for large language model training. This local receipt does not reproduce its experiments. It checks whether one bounded arithmetic step conforms to the declared route, orientation, polynomial, scale, and decay contract. That truth boundary stays next to every export.

Independence also extends to serialization. JSON preserves the full structured receipt, including matrices and mutation outcomes. CSV contains one row per stage and scalar summary so it can be diffed without executing JavaScript. Both exports include the same input hash and configuration fields. A verifier recomputes the hash from canonical input data and rejects an export whose provenance was edited after the calculation.

Make five negative fixtures first-class evidence

Happy-path output proves very little when several wrong programs produce matrices of the right size. The Muon optimizer lab runs five named mutants beside the main receipt. EMBEDDING_MISROUTE sends an embedding matrix to Muon. MISSING_TALL_TRANSPOSE processes a tall update in the wrong orientation. WRONG_MULTIPLICATION_SIDE builds the wrong Gram contract. OMITTED_SCALE removes the declared shape factor. COUPLED_DECAY changes operation order.

Each fixture has a specific expectation and diagnostic code. The test fails if a mutant survives, but it also fails if a different guard happens to reject it first. That precision keeps one broad “invalid” condition from masking a missing invariant. The exported receipt includes the mutated field, expected code, observed code, and pass state.

Mutation tests should remain bounded and deterministic. They reuse the same frozen parameter, gradient, momentum, and configuration as the control. No random matrix is generated, no timing result is collected, and no network request occurs. A repeated run must produce the same canonical JSON and CSV bytes except for the user-initiated download wrapper, which carries no timestamp.

The browser test also probes malicious and malformed inputs: non-finite cells, dimensions outside 2–4, duplicate parameter names, unsupported semantic roles, and excessive decimal text. Every case must become visible validation text, never HTML. Keyboard users can edit the matrix, run the receipt, execute fixtures, and invoke exports without relying on color or pointer gestures.

Make the adoption decision from receipts, not labels

A sensible Muon optimizer rollout begins with a small inventory. Count eligible hidden matrices, excluded matrices, total elements, and exceptions. Run the bounded ledger on representative wide, tall, and square shapes. Compare the production implementation against the independent oracle and all five mutants. Then add model-level experiments as a separate evidence layer.

The go/no-go table should have explicit owners. An optimizer engineer owns routing and arithmetic parity. A training lead owns hyperparameters and convergence experiments. An infrastructure owner measures memory and throughput on the actual accelerator stack. A model owner evaluates downstream quality. None of those later decisions can be inferred from this lab, but all benefit from its precise starting state.

Keep an AdamW fallback for every excluded or unsupported parameter and for emergency rollback. Record configuration beside checkpoints so a run can be reconstructed. If parameter names or shapes change, fail the manifest comparison before training rather than silently accepting a new route. Long-context work such as context parallelism follows the same principle: system topology belongs in the experiment record, not in tribal memory.

The deliverable is deliberately modest: one inspectable update, one independent approximation comparison, five killed mutants, and deterministic exports. That is enough to establish implementation integrity before expensive evaluation begins. It is not a recommendation that every model use Muon, nor a claim about speed, convergence, or quality.

Runnable local artifact — The lab audits arithmetic and parameter routing on bounded 2×2 through 4×4 matrices only; it makes no convergence, speedup, model-quality, or hyperparameter claim.

Enter one bounded matrix record, trace momentum and Nesterov state, transpose tall updates, execute five declared polynomial steps, apply shape scaling and decoupled decay, compare a separate small-matrix polar oracle, kill five mutations, and export canonical JSON or CSV.