Towards a Principled Muon under $μ\mathsf{P}$: Ensuring Spectral Conditions throughout Training
John Zhao
TL;DR
The paper addresses ensuring the μP spectral conditions for matrix-based optimizers like Muon during long-horizon LLM training. It introduces Muon++ as a projection-based variant that enforces spectral constraints on updates, enabling μP-compatible scaling without heavy weight normalization and enabling adaptive, data-aware conditioning. The key contributions include a concrete constrained optimization formulation, a projection-based solution that operates via the SVD of projected gradients, and practical considerations for ultra-wide models along with an initial step toward data-dependent spectral adaptation. This approach aims to bridge μP theory and practical deployment of matrix-based optimizers, improving stable hyperparameter transfer and scalable training of large models.
Abstract
The $μ$-parameterization ($μ$P) provides a principled foundation for large language model (LLM) training by prescribing width-independent learning dynamics, which in turn enables predictable scaling behavior and robust hyperparameter transfer across model sizes. A central requirement of $μ$P is the satisfaction of certain spectral conditions on weight matrices, which ensure consistent feature learning and optimization behavior as model width grows. While these conditions are well understood in theory, guaranteeing their validity in practical training for matrix-based optimizers such as Muon is still under studied. Existing works that study Muon under $μ$P exhibit important limitations: they either do not ensure that the spectral conditions hold throughout the entire training horizon, or require repeated spectral normalization (or Newton-Schulz iterations) applied to both weights and updates, leading to significant computational overhead and reduced practicality. In this work, we show how to reliably guarantee the spectral conditions required by $μ$P for Muon during the entire training process. Our key insight is that for moderately large models, maintaining spectral control at the level of optimizer updates alone is sufficient to preserve $μ$P-compatible scaling, eliminating the need for explicit spectral normalization of the weights. Based on this principle, we develop a variant of Muon, namely Muon++, that satisfies spectral condition throughout the training process. Our results bridge the gap between the theoretical promises of $μ$P and the practical deployment of matrix-based optimizers in long-horizon training. We also take the first step towards an adaptive spectral condition by incorporating data-dependent effects, making it better suited for long-horizon LLM training.