AI summaryⓘ
The authors studied Muon, a new optimizer designed to improve learning by adjusting weight updates differently than usual. They tested Muon against the popular AdamW optimizer on a specific state-space model called Mamba-2 130M, changing only which parts of the model's weights Muon was applied to. They found that applying Muon only to the output layer worked better than applying it to the input layer or both. This improvement mainly helped the model learn more efficiently with fewer tokens, and the effect was consistent across different datasets and longer training times. The authors also observed that Muon's benefits were not due to how well-conditioned the model parts were, since improving input projection conditioning did not lead to gains.
Muon optimizerAdamWNewton-Schulz iterationspectral normstate-space modelsweight matricestoken efficiencyconditioningMamba-2 130Moutput projection
Authors
Arslan Battalov, Karim Kramin, Alexander Markotenko, Sofia Sinitsina
Abstract
Muon is a recent optimizer that orthogonalizes the update to each weight matrix with a Newton-Schulz iteration, which performs steepest descent under the spectral norm. Almost all the evidence for it comes from Transformer models, and its behavior on state-space models is largely unreported. We compare Muon with AdamW on Mamba-2 130M under a controlled protocol that varies only which weight groups are trained with Muon. The benefit is localized. Muon on the output projection alone beats Muon on the input projection or on both. The advantage is mainly one of token efficiency. It holds on two corpora and two token budgets, and persists when training continues well past the compute-optimal point. Conditioning does not explain the gain. Muon lowers the condition number of whichever projection it trains, but the better-conditioned input projection is not the one that helps.