๐ค AI Summary
This work addresses the discrepancy between theory and practice in the Muon optimizer for large language model training, which exhibits rapid early convergence but suffers from instability and oscillation in later stages. The authors introduce a novel โvalleyโ theoretical framework grounded in optimization trajectory analysis, integrating a mixed spike matrix sensing model with spectral river methodology to elucidate how Muon rapidly descends along dominant information directions yet frequently overshoots. Building on these insights, they propose a two-stage optimization strategy: employing Muon initially to accelerate convergence and subsequently switching to a gradient descentโtype optimizer to stabilize training. Both theoretical analysis and preliminary experiments demonstrate that this approach effectively enhances training stability and final model performance in large language models.
๐ Abstract
Recently, Muon has gained substantial attention as an appealing alternative to Adam-like optimizers, with many works highlighting its advantages through spectral normalization and improved conditioning. Yet this positive theoretical narrative contrasts with its empirical performance in large language model (LLM) training, where Muon's gains over Adam/AdamW are often mixed, schedule-sensitive, and not uniformly superior. To address this gap, we develop a trajectory-level theory characterizing both the strengths and limitations of Muon. We introduce a mixed-spiked matrix sensing model whose sensing operator decomposes into signal, spike, and bulk components, capturing a mixture of anisotropic structure and long-tail information reminiscent of LLM training. On top of it, we adopted a river-valley perspective in which we view the landscape as composed of a river direction flowing to the desired solution and hill directions encoding nuisance or task-irrelevant information. In the momentum-free setting, we show that Muon moves faster along the information-bearing river direction during early optimization, but can converge much more slowly near the river bottom than gradient descent. We then extend the river-valley perspective to general nonconvex objectives with momentum by studying points on the spectral river. There, while Muon converges faster early on, its orthogonalized update removes residual scale information, making it prone to overshooting and oscillation near the target solution. Together, these results suggest that our characterizations extend beyond spiked matrix sensing and motivate switching to GD-like refinement optimizers in the final phase, rather than relying only on a fixed learning-rate schedule for Muon. We also provide preliminary evidence supporting this two-stage approach in language model training experiments.