Search papers, labs, and topics across Lattice.
This study evaluates the performance of the Muon optimizer compared to AdamW in the context of state-space models, specifically using the Mamba-2 130M architecture. The findings reveal that Muon's benefits are localized to the output projection, demonstrating superior token efficiency over both the input projection and a combined approach. Notably, the advantage persists across different corpora and token budgets, indicating that conditioning does not account for the observed improvements in performance.
Muon optimizer's unique advantage lies in its ability to enhance token efficiency specifically when applied to the output projection, challenging assumptions about conditioning in state-space models.
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.
