Search papers, labs, and topics across Lattice.
This study investigates the performance of the orthogonal optimizer Muon compared to Adam by analyzing the loss landscapes of Transformers through spectral probing. The authors find that while Muon outperforms Adam, it underutilizes the bulk of the loss landscape, leading to the introduction of Spectral-Aware Muon (SAMuon), which optimally allocates step sizes across different directions. Both SAMuon and its simplified variant outperform Adam and Muon across various model scales and batch sizes, requiring significantly fewer training tokens to achieve equivalent validation loss.
Muon's advantage over Adam is rooted in a nuanced spectral profile that reveals underutilization of the loss landscape's bulk, leading to the development of a more efficient optimizer, SAMuon.
Orthogonal optimisers such as Muon can substantially accelerate large language model pretraining relative to Adam, yet the mechanism remains incompletely understood. We investigate this through an out-of-sample spectral probing analysis of Transformer loss landscapes. At checkpoints along real training trajectories, we decompose each momentum buffer into its singular directions and estimate the loss-optimal step size along each direction on held-out data. The resulting spectral profile is anisotropic yet stable across batches and training stages, and consistent across the optimisers and model scales: a volatile head operating at the Edge-of-Stability supports a much smaller step size than the tolerant bulk, which permits substantially larger steps. This profile provides a unified spectral allocation account of why Muon outperforms Adam, which outperforms SGD. It also exposes a limitation of Muon's uniform scaling: it still underutilises the bulk. Guided by this finding, we introduce Spectral-Aware Muon (SAMuon), which holds the head at the Muon scale and amplifies the bulk using a static spectral prior. We provide two variants: the complete SAMuon follows the measured profile using a low-rank randomised SVD and the simplified SAMuon-lite uses a two-level approximation via rank-one power iteration. Neither method adds persistent optimiser state or notable extra FLOPs beyond Muon at scale, and the idealised exact-whitening versions of both retain Muon's asymptotic convergence rate under standard assumptions. Across"modded-nanogpt"models from 124M to 1B parameters, both variants outperform tuned AdamW and Muon (Scion implementation) baselines in all evaluated model-scale and batch-size configurations. SAMuon requires 13.3% to 24.0% fewer training tokens to reach the same validation loss as Muon, while SAMuon-lite retains most of this gain with near-zero wall-clock overhead.
