The phenomenon of grokking, where a neural network abruptly generalises long after achieving near-perfect training accuracy, has become a useful testbed for understanding optimisation dynamics. When Tveit and colleagues reported that the Muon optimizer reaches the grokking threshold on modular arithmetic substantially faster than AdamW, the natural follow-up question was: which part of Muon is doing the work? Muon combines two distinct mechanisms, an orthogonalized momentum update via Newton-Schulz iteration, and a spectral-norm-based scaling factor. Prior work credited both without isolating either. This paper runs the ablation that was missing.
The practical stakes are modest at the scale studied here, a one-layer transformer on modular addition mod 97, but the methodological lessons generalise. If we cannot correctly attribute a speedup to its source in a regime where experiments take seconds to minutes, we are unlikely to get the attribution right when they take days.
What the Paper Actually Does
The authors decompose Muon into its two components and test them independently across multi-seed and multi-learning-rate sweeps. The four configurations are: full Muon, orthogonalize-only (Newton-Schulz without spectral scaling), spectral-only (scaling without orthogonalization), and AdamW as the baseline. Two metrics are tracked throughout: first-crossing time, the step at which validation accuracy first exceeds 0.95, and stable-grok time, the step at which it exceeds 0.95 and stays there for the remainder of training.
The distinction between these metrics is not a minor bookkeeping choice. At higher learning rates, AdamW can cross the threshold and oscillate back below it dozens of times, making first-crossing a misleading quantity. The paper pre-registers both metrics and reports both, which is the right call. The two primaries, orthogonalize-only versus AdamW and reduced Newton-Schulz iteration count versus canonical Muon, both survive Holm-Bonferroni correction. Secondary comparisons are clearly flagged as exploratory.
Key Findings
- Orthogonalization is the active ingredient. Orthogonalize-only matches full Muon in steps to stable grokking. Spectral-only is statistically indistinguishable from AdamW and shows higher variance. This verdict holds across the learning rate sweep from 3e-4 to 3e-3.
- The grokked solution is genuinely lower-norm, not just less perturbed. Orthogonalizing optimizers reach generalisation at roughly 3x lower spectral norm in the hidden weight matrices. Critically, the authors control for embedding movement and find the lower norm reflects a qualitatively different solution, not simply a smaller update magnitude. This is a mechanistic claim, and it is carefully hedged: the Fourier circuit analysis uses a coarse proxy and the authors acknowledge it.
- Newton-Schulz iteration count governs a speed-stability tradeoff. Reducing from five iterations to one speeds up first-crossing but makes the grokked solution fragile, with transient collapses in validation accuracy that scale with learning rate. Five iterations is the rate-robust default; one iteration is only safe at small learning rates. The headline that fewer iterations is faster inverts under the stable-grok metric at higher learning rates.
- Spectral scaling can be dropped at no measured cost. This is a clean negative result and a useful one. The spectral scaling factor is a per-matrix scalar, not a spectral-norm projection in the stricter sense, and it appears to be inert in this setting.
Methodological Contributions
The paper's most transferable contribution may be the argument about metrics. Grokking speed is typically reported as first-crossing time, which is both intuitive and gameable by any optimizer that produces noisy trajectories. An optimizer that crosses the threshold at step 500 and then oscillates below 0.9 for the next 2000 steps is not actually better than one that crosses cleanly at step 700. The stable-grok metric captures this, and the paper shows concretely that the ranking of their own proposed variant inverts depending on which metric you use.
This is worth taking seriously beyond grokking. Evaluation metrics that reward transient rather than sustained behaviour are a recurring problem in empirical ML, and the authors demonstrate the hazard with a real example from their own results rather than a hypothetical.
The multi-seed, multi-learning-rate design is also worth noting. Single-seed results on grokking are unreliable because the phenomenon has high stochastic variance, and learning-rate sensitivity means that a result true at one rate may not hold at another. Running the sweep and reporting significance tests with multiple comparison correction is the minimum standard for this kind of claim, and it is not universally applied in the grokking literature.
Limitations and Open Questions
The authors are candid about scope. One architecture, one scale, one task family. The computational cost of Newton-Schulz iteration is not free at scale, and a finding that orthogonalization is the active ingredient at 128-dimensional hidden layers does not immediately tell us whether the tradeoff holds for large language model training runs where the iteration cost becomes non-trivial. The spectral norm analysis relies on a coarse proxy for the Fourier circuit structure identified by Nanda and colleagues, and the mechanistic claims in Section 4 should be read as suggestive rather than definitive.
There are also questions the paper raises without answering. If orthogonalization drives lower-norm solutions, what is the geometric mechanism? Is the Newton-Schulz iteration acting as an implicit regularizer, or is it changing the effective loss landscape in a more fundamental way? The connection to the broader literature on implicit regularization in gradient descent feels underexplored, though that may be a deliberate scope decision.
The code is released at https://github.com/louiswang524/muon-grokking-frontier, which makes the multi-seed sweep results verifiable. That matters here because the central claims rest on statistical comparisons rather than a single dramatic curve.
For anyone working on optimizer design or using grokking as a diagnostic tool, the takeaway is practical: orthogonalization is doing the work in Muon, spectral scaling is free to remove, and any speed claim about grokking should be checked against a stability-aware metric before being trusted. The paper is a good example of how to do a focused ablation study properly, even if the scope limits how far the conclusions can travel.