Reproduce “Towards a Principled Muon under $μ\mathsf{P}$: Ensuring Spectral Conditions throughout Training” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
John Zhao
This paper introduces Muon++, an optimizer designed to ensure Maximal-update Parameterization (μP) spectral conditions throughout large language model training. It achieves this by controlling the spectral properties of optimizer updates, demonstrating a computationally affordable method for stable and efficient hyperparameter transfer and proposing adaptive spectral conditions based on weight correlations.
Key insights
- For moderately large models, maintaining spectral control solely at the level of optimizer updates is sufficient to preserve μP-compatible scaling behavior, eliminating the need for computationally expensive explicit weight normalization.
- The spectral constraint on weight matrices can be automatically satisfied if the update is orthogonal to the top singular vectors of the weight matrix, provided the learning rate is within a specific admissible range.
- Static, data-independent spectral conditions may be insufficient for very long training horizons; evolving data-dependent correlations within weight matrices necessitate adaptive scaling adjustments.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.