Abstract
Muon and related matrix-sign optimizers are increasingly used to pre-train large language models, yet their impact on the internal geometry of individual weight matrices remains poorly understood. This preliminary report proposes a unified framework based on a single idealizing assumption: the exact scale invariance of the loss under weight rescaling, which approximately holds in normalization-heavy networks.
Under this assumption, plain SGD imposes a built-in $1/||W||$ brake on its update size, while Muon's matrix-sign step removes this brake, leading both Frobenius and spectral norms to drift outward faster (at rates of $t^{1/2}$ versus $t^{1/4}$). We further observe that the spectral-norm perturbation contains a non-negative second-order term. This indicates that a lightweight "spectral cap"—which projects out only the first-order growth of the single top singular direction from each update—can control the output covariance $W K_X W^T$ without freezing training: the weights continue to learn through non-top directions, top-direction rotation, and top switching.
We relate this cap to the min-entropy ($H_{\infty}$) of the singular-value spectrum. We then examine three systems trained with Muon: a nanoGPT feed-forward projection, a 64-expert mixture-of-experts router, and the query/key projections of a bf16 FlashAttention block. In each case, the cap increases isotropy and, at the margins—such as a router collapsing to a single expert and the near-divergence of one attention head—prevents concrete failures while leaving validation loss essentially unchanged.
We emphasize that the scale-invariance assumption is strong and that these small-scale results are preliminary; comments are welcome.
Blogger's Review: This research provides a fresh perspective on the application of Muon optimizers in large language models, particularly regarding the control of weight matrix geometry. By introducing the concept of a spectral cap, the authors not only demonstrate how to enhance isotropy but also offer significant insights for future optimization algorithm designs. It is worth noting that the assumptions in the study require further validation to ensure effectiveness in broader application scenarios.