Muown: Row-Norm Control for Muon Optimization
- URL: http://arxiv.org/abs/2605.10797v1
- Date: Mon, 11 May 2026 16:26:06 GMT
- Title: Muown: Row-Norm Control for Muon Optimization
- Abstract summary: Muon is a strong competitor to AdamW for language model pre-training.<n>Muown improves perplexity over Muon, SOAP, AdamW, and Lion.
- Score: 41.541871767555484
- License: http://creativecommons.org/licenses/by/4.0/
- 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.
Related papers
- Spectral Allocation: Why Muon Outperforms Adam, and How to Improve Muon [12.170977868230622]
Orthogonal optimisers such as Muon can substantially accelerate large language model pretraining relative to Adam.<n>We investigate this through an out-of-sample spectral probing analysis of Transformer loss landscapes.<n>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.
arXiv Detail & Related papers (2026-08-26T16:41:12Z) - MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning [14.373662883104808]
We propose MALT to reduce the sensitivity of Muon to curvature geometry of the loss landscape.<n>MALT uses two-sided diagonal preconditioners with low memory and overhead to capture the curvature geometry of the loss landscape.
arXiv Detail & Related papers (2026-08-05T17:26:47Z) - Scale Weight Decay and Train Better [0.0]
We propose to scale weight decay by a fraction of the peak learning rate $/_max$.<n>We prove that this preserves the stationarity guarantees of the corresponding unregularized methods for both a gradient descent and a non-Euclidean Muon.<n>Using a steady-state analysis, we explain why under standard weight decay the weight norm shrinks steadily as training proceeds, whereas under scaled weight decay it settles to roughly constant value.
arXiv Detail & Related papers (2026-07-26T17:55:40Z) - Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors [110.10775872372047]
We propose emphMagnitude-Direction-Direction (MD) Decoupling.<n>It factorizes each weight into a fixed-norm direction on a hypersphere and learnable per-row and per-column magnitude gains, updated at separate learning rates.<n>Across both Adam and Muon, MD Decoupling improves on well-tuned baselines, transfers the optimal LR across model width without retuning, and continues to help at scale on large Mixture-of-Experts models.
arXiv Detail & Related papers (2026-06-24T15:40:26Z) - Why Muon Outperforms Adam: A Curvature Perspective [49.85900116602357]
Muon improves training efficiency over Adam in large language-model training by about two times.<n>Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective.
arXiv Detail & Related papers (2026-06-03T09:40:30Z) - Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training [56.323119575322146]
Extra-Merge consistently outperforms standard merging baselines in experiments across GPT-2 and LLaMA families.<n>It yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon citepjordan2024muon.
arXiv Detail & Related papers (2026-05-26T02:48:34Z) - Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory [77.27772368491698]
Spectrals such as Muon have recently shown strong empirical performance in large-scale language model training.<n>We study this question through the linear associative memory problem.<n>We show that the storage capacity of Muon significantly exceeds that of SGD.
arXiv Detail & Related papers (2026-03-27T16:13:18Z) - Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning [72.62839712454196]
Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
arXiv Detail & Related papers (2026-03-10T14:03:49Z) - NuMuon: Nuclear-Norm-Constrained Muon for Compressible LLM Training [50.27276603708547]
We show that despite imposing full-rank updates, Muon-trained models exhibit pronounced low-rank structure in their weight matrices and are readily compressible under standard pipelines.<n>We propose NuMuon, which augments Muon with a nuclear-norm constraint on the update direction, further constraining the learned weights toward low-rank structure.
arXiv Detail & Related papers (2026-03-04T00:10:14Z) - Muon in Associative Memory Learning: Training Dynamics and Scaling Laws [23.350512542598803]
We study Muon in a linear associative memory model with softmax retrieval and a hierarchical frequency spectrum over query-answer pairs.<n>We show that Muon mitigates this imbalance, leading to faster and more uniform progress.
arXiv Detail & Related papers (2026-02-05T14:49:40Z) - ECO: Quantized Training without Full-Precision Master Weights [58.97082407934466]
Error-Compensating (ECO) eliminates master weights by applying updates directly to quantized parameters.<n>We show that ECO converges to a constant-radius neighborhood of the optimum, while naive master-weight removal can incur an error that is inversely proportional to the learning rate.
arXiv Detail & Related papers (2026-01-29T18:35:01Z) - Variance-Adaptive Muon: Accelerating LLM Pretraining with NSR-Modulated and Variance-Scaled Momentum [19.385264518362472]
Large Language Models (LLMs) achieve competitive performance across diverse natural language processing (NLP) tasks.<n>We propose two variants of Muon, Muon-NSR and Muon-VS, which apply variance-adaptive normalization to momentum.<n>Experiments on GPT-2 and LLaMA pretraining demonstrate that our proposed methods accelerate convergence and consistently achieve lower validation loss than both competitive, well-tuned AdamW and Muon baselines.
arXiv Detail & Related papers (2026-01-21T02:41:56Z) - Cautious Weight Decay [23.198565281737896]
Cautious Weight Decay (CWD) is a one-line, agnostic modification that applies weight decay only to parameter whose signs align with update.<n>CWD is a drop-in change for coordinates such as AdamW, Lion and Muon.<n>For language model pre-training and ImageNet classification, CWD consistently improves final loss and accuracy.
arXiv Detail & Related papers (2025-10-14T11:32:55Z) - Low-rank Orthogonalization for Large-scale Matrix Optimization with Applications to Foundation Model Training [3.1922198632169327]
Recently, the Muon citejordanmuon has gained significant attention for its strong performance in foundation model training.<n>We propose low-rank matrix-signed gradient descent and a low-rank variant of Muon.
arXiv Detail & Related papers (2025-09-15T14:28:53Z) - Convergence Bound and Critical Batch Size of Muon Optimizer [1.2289361708127877]
We provide convergence proofs for Muon across four practical settings.<n>We show that the addition of weight decay yields strictly tighter theoretical bounds.<n>We derive the critical batch size for Muon that minimizes the computational cost of training.
arXiv Detail & Related papers (2025-07-02T11:03:13Z) - Muon is Scalable for LLM Training [50.68746986439438]
We introduce Moonlight, a Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon.<n>Our model improves the current frontier, achieving better performance with much fewer training FLOPs compared to prior models.<n>We open-source our distributed Muon implementation that is memory optimal and communication efficient.
arXiv Detail & Related papers (2025-02-24T09:12:29Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.