How Muon Lost Its Geometry

TL;DR

Ligo Biosciences’ Arda Göreci traces how Muon’s per-matrix scaling factor — the term that should preserve maximal-update parameterization (µP) for rectangular layers — diverged from the geometrically correct $\sqrt{\text{fan-out}/\text{fan-in}}$ as the optimizer moved from NanoGPT speedruns to trillion-parameter LLMs. Two breaking variants now ship as defaults: Keller Jordan’s $\max(1, \sqrt{\text{fan-out}/\text{fan-in}})$ (clipped) and Moonshot’s $\sqrt{\max(\text{fan-in}, \text{fan-out})}$ (only equivalent when all matrices share a fan-in). Both are exposed as user-selectable options in torch.optim.Muon (PyTorch 2.9) and optax.contrib.muon, neither matches the correct version, and the post argues MuonClip and similar patches are symptoms of this broken scaling rule rather than independent stability improvements. Microsoft’s Dion implementation is called out as the one popular variant that gets the scaling factor right. A CIFAR-10 LR-transfer experiment confirms the correct factor gives the lowest test loss with no LR drift across widths; the Moonshot factor systematically shifts the optimum to lower LRs as width grows. The motivating context is biomolecular structure prediction (AlphaFold3-style architectures), where rectangular projections are the rule rather than the exception.

Key claims

Muon orthogonalizes the update via Newton-Schulz so the spectral norm of the update is 1, but the quantity that bounds the change in activation RMS is the RMS→RMS operator norm — for a rectangular matrix these differ by a factor of $\sqrt{\text{fan-in}/\text{fan-out}}$ , so the geometrically-correct scaling factor on the update is $\sqrt{\text{fan-out}/\text{fan-in}}$ [§“What Muon is trying to do”].
Keller Jordan’s speedrun implementation applies $\max(1, \sqrt{\text{fan-out}/\text{fan-in}})$ , which clips the reciprocal correction whenever fan-out < fan-in — wrong on down-projections specifically [Table “Where implementations diverge”].
Moonshot AI’s Moonlight implementation uses $\sqrt{\max(\text{fan-in}, \text{fan-out})}$ , which is equivalent to Keller’s only if fan-in is the same for all matrices in the network — an assumption that breaks immediately for Multi-head Latent Attention’s compress/decompress projections [§“The first large-scale stress test: Moonlight”].
The MLA down-projection in DeepSeek-V3-style architectures gets updated about 5× larger than µP prescribes under Moonshot’s factor, which the post argues is the underlying cause of the attention-logit explosions that MuonClip was introduced to patch [§“The first large-scale stress test: Moonlight”].
A CIFAR-10 MLP-bottleneck LR-transfer sweep across 5 widths and 5 Muon variants (Adam baseline, µP Adam, correct $\sqrt{\text{fan-out}/\text{fan-in}}$ Muon, Keller’s $\max(1, \cdot)$ Muon, Moonshot’s $\sqrt{\max}$ Muon) shows: the correct factor gives the lowest absolute test loss at every width with no LR drift; Moonshot’s drifts the optimum to lower LRs as width grows; Keller’s drifts less but still drifts [§“A CIFAR-10 learning-rate transfer test”].
Microsoft’s Dion implementation applies the correct $\sqrt{\text{fan-out}/\text{fan-in}}$ factor — the one popular Muon variant that ships µP-correct scaling [Table “Where implementations diverge”].
torch.optim.Muon (PyTorch 2.9) exposes the ambiguity as a user choice between Keller-style and Moonshot-style scaling; neither option is the µP-correct version, and the documentation frames this as flexibility rather than as a missing default [§“Ambiguity built into the major frameworks”].
For AlphaFold3-style architectures with many rectangular projections, the per-layer LR error under Keller’s factor ranges from 0% (square or up-projection-dominated layers) to +553% (atom → xyz update, $128 \to 3$ ); under Moonshot’s factor it ranges from −71% to +167% [§“When everything gets the wrong learning rate”, inline table].

Method

The argument is a mix of textual derivation and a small empirical confirmation. The derivation follows Jeremy Bernstein’s RMS→RMS operator-norm framing: a linear layer’s output change after a weight update is bounded by $\|\Delta W\|_{\text{RMS} \to \text{RMS}}$ , which for a Newton-Schulz-orthogonalized update equals $\sqrt{\text{fan-in}/\text{fan-out}}$ rather than 1. To get $\|\Delta y\|_\text{RMS} \le \eta$ — the µP-feature-learning condition — the update needs the reciprocal correction $\sqrt{\text{fan-out}/\text{fan-in}}$ as a per-matrix multiplier. The post then walks through how popular implementations either clip this correction (Keller) or replace it with a non-equivalent surrogate (Moonshot’s $\sqrt{\max}$ ), and how torch.optim.Muon and Optax fossilize both as user options.

The empirical confirmation is a CIFAR-10 MLP with a width bottleneck in the middle of six layers (chosen specifically to force rectangular matrices), trained at five widths with five optimizer settings: Adam (default), Adam under µP LR, Muon with $\sqrt{\text{fan-out}/\text{fan-in}}$ , Muon with $\max(1, \sqrt{\text{fan-out}/\text{fan-in}})$ , and Muon with $\sqrt{\max(\text{fan-in}, \text{fan-out})}$ . The figure ([§“A CIFAR-10 learning-rate transfer test”]) plots test loss against learning rate at each width.

Results

Correct $\sqrt{\text{fan-out}/\text{fan-in}}$ Muon: lowest absolute test loss of any setting; LR optimum stays at roughly the same value as width changes [§“A CIFAR-10 learning-rate transfer test”].
Moonshot’s $\sqrt{\max}$ factor: optimum drifts steadily to lower LRs as width grows (i.e. would require re-tuning at every scale) [§“A CIFAR-10 learning-rate transfer test”].
Keller’s $\max(1, \cdot)$ factor: drifts less than Moonshot’s but more than the correct version.
AlphaFold3-style worked example: under Keller’s scaling, +553% LR error on the atom → xyz update ( $128 \to 3$ ) and +300% on the MSA pair-weight bias ( $128 \to 8$ ); under Moonshot’s, +167% and +63% respectively [§“When everything gets the wrong learning rate”].
MLA-specific argument: down-projection updated ≈5× too large vs µP prescription under Moonshot’s factor in a Moonlight-style setup [§“The first large-scale stress test: Moonlight”].
The fix is a one-line change in PyTorch and Optax to add the correct scaling option as a default [§“Conclusion”].

Why it’s interesting

This blog supplies the missing causal story behind two longstanding wiki threads. First, it directly substantiates the “Muon is half-aligned with µP” conflict line on Hyperparameter scaling laws and Training stability at scale: Controlled LLM Training on Spectral Sphere argues µP’s motivation still stands and that closing the loop on weight + update spectra restores LR transfer; this post argues a narrower version of the same — that even the update-side half is wrong as implemented, and a one-line fix to the scaling factor restores the LR-transfer property at the optimizer level without needing SSO-style retraction. Second, it reframes QK-Clip: Taking Muon Further on the Scaleup Journey‘s MuonClip and the QK-Clip lineage as a symptom rather than a fix — the MLA-specific instability that QK-Clip was introduced to patch is exactly what Moonshot’s $\sqrt{\max}$ factor would predict from a 5× over-updated down-projection. Third, it complicates the framework-default question opened by Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819): PyTorch 2.9 shipped Muon with both wrong scaling options as user-selectable defaults, which is exactly the “the bug fossilized as API surface” outcome that paper’s open question warned about. The Tilde Research line (Aurora: A Leverage-Aware Optimizer for Rectangular Matrices, Compositional Muon — partner-whitened optimizer that controls composed functional updates in attention) is independently working on Muon’s rectangular-matrix and compositional gaps; this post argues a simpler, additive fix exists.