Spectral Cap Improves Muon Optimizer Stability in LLM Training

This preliminary research explores the internal geometry of weight matrices when training large language models (LLMs) with Muon and similar matrix-sign optimizers. The authors observe that unlike plain SGD, Muon's matrix-sign step removes a built-in brake on update size, causing Frobenius and spectral norms to drift outwards faster. This can lead to issues like a router collapsing to a single expert or attention heads diverging. To address this, the paper proposes a lightweight "spectral cap." This cap projects out only the first-order growth of the single top singular direction from each update, effectively controlling the output covariance without freezing training. The weight matrix continues to learn through other directions and rotations. Case studies on nanoGPT feed-forward projections, mixture-of-experts routers, and FlashAttention blocks demonstrate that this cap increases isotropy and prevents concrete failures at the margins, all while leaving validation loss essentially unchanged. The authors emphasize the preliminary nature of these small-scale results and the strong assumption of exact scale invariance.