Muown: Row-Norm Control for Muon Optimization

TL;DR

Muown is a novel optimizer for language model pre-training that improves upon Muon by explicitly controlling row-norms, leading to better perplexity and training stability across various model sizes.

Contribution

We introduce Muown, a new optimizer that explicitly manages row-norms in Muon, achieving optimal convergence rates and reducing spectral norm drift during training.

Findings

Muown outperforms Muon, SOAP, AdamW, and Lion in perplexity across multiple model sizes.

Muown widens the plateau of near-optimal learning rates, enhancing training robustness.

Muown reduces sensitivity to weight decay and avoids spectral norm drift with minimal overhead.

Abstract

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the $\ell_\infty$ geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.