Muon and AdamW under a fixed compute budget

An optimizer changes the cost of a step

A comparison between optimizers often begins with training loss against update count. That is useful for studying optimization behavior, but it leaves out an important systems question: what did each update cost?

Muon uses matrix operations to approximately orthogonalize momentum updates for selected weight matrices. AdamW applies coordinate-wise updates with adaptive moment estimates. This difference changes both the optimization rule and the work performed by the accelerator. The Muon reference implementation is a useful starting point for the algorithm; my focus here is the cost of implementing and evaluating variants in a Transformer training loop.

In my JAX TPU training code, I added optimizer targets, matrix splitting, iteration controls, and component timing. These controls let me ask whether the expensive part of a training step is the model or the mechanism that updates its parameters.

Why model utilization can miss the bottleneck

A model-FLOPs utilization estimate divides an estimate of the model’s arithmetic rate by the device’s nominal peak. If the numerator excludes optimizer work while elapsed time includes it, an expensive optimizer lowers that number. The estimate can be internally consistent and still tell an incomplete story.

The measured step time includes more than the neural network:

This is a conceptual accounting identity, not a promise that independently timed components will sum exactly: a compiler can fuse work, and a runtime can overlap it. I use component timings to form a hypothesis, then check the full step.

Wide feed-forward matrices are a particularly useful place to look. Repeated matrix operations in an approximate orthogonalization can be substantial, even when the forward pass is efficient. More iterations, different matrix shapes, and different ways of grouping weights can change this cost. A higher model-only utilization number does not by itself tell me which optimizer reaches a target quality sooner.

Two public profiles that should not be ranked

The repository records the following short-window profiles:

These are configuration observations in the README, not a controlled Muon-versus-AdamW result. The model and batch size both differ. It would be incorrect to conclude from this table that AdamW is faster for the same model, or that either optimizer produces better models.

The 453 positions/s measurement also describes Go training positions, not language-model tokens. Translating a systems technique between domains requires translating the workload as well as the code.

A faster update may be a different optimizer

My implementation exposes several ways to investigate optimizer cost:

• Apply Muon only to attention matrices, leaving other parameters on AdamW.
• Split wide matrices into smaller blocks before the approximate orthogonalization.
• Change the number of orthogonalization iterations.
• Change precision and parameter-tree layout independently where possible.

These knobs do not all mean the same thing. Splitting a matrix changes which directions are coupled by an update. Applying Muon to fewer weights changes the optimization rule for the remaining parameters. Reducing iterations changes the approximation. These are algorithmic ablations as well as implementation experiments.

That distinction matters when an optimization looks promising. A drop in update time is only the first result. I still need to know whether the change preserves the quality reached with the available data and compute. I would not describe these options as equivalent, lossless implementations of full-matrix Muon.

The comparison I would actually run

For a quality-under-budget comparison, I would first freeze the architecture, data split, global batch, effective sample order, numerical precision, and evaluation protocol. I would then compare a small set of explicitly named candidates: AdamW, the selected full-matrix Muon baseline, attention-only Muon, and one matrix-splitting variant.

Each candidate needs an equal, declared hyperparameter-search budget. Forcing every optimizer to use the same learning rate is not automatically fair. Allowing one candidate many more tuning attempts is not fair either. I would report both the tuning policy and the selected settings.

The outcomes I care about are:

1. Validation quality versus training positions, which describes sample efficiency.
2. Validation quality versus elapsed accelerator time, which describes time efficiency.
3. Peak memory and optimizer-state storage, which determine feasible model and batch sizes.
4. Downstream task evaluation on a fixed, independently selected protocol.

A second experiment could allow each optimizer to choose its best feasible batch size. That answers a deployment question, but it changes the comparison and should be reported separately. Fixed-model attribution and best-system selection are both useful experiments; they need different controls.

What would change my mind?

If an attention-only configuration loses enough sample efficiency that it reaches a quality target later, its cheaper update is not a win under that time budget. If a full-matrix variant wins only after much longer training, the conclusion depends on the intended run length. A result that reverses across model widths is a reason to investigate the scaling of update cost, not to select only the favorable width.

My current public notes establish a working implementation and throughput observations. They do not establish a quality-matched winner. The next useful result would connect optimizer cost to a learning curve under a controlled budget.

Sources

• Muon reference implementation.
• JAX training loop and optimizer implementation.
• Public configuration notes and limitations.
• Batch-size sweep and Muon configuration sweep.

The proposed quality comparison above is an experiment plan, not a claim that those runs have already been completed.

Code snapshot reviewed: 90c12af. Later revisions may change implementation details.