Source-linked AI summary
MUON+: Towards More Effective Muon via One Additional Normalization Step for LLM Pre-training
Ruijie Zhang, Yequan Zhao, Ziyue Liu, Zhengyang Wang, Yupeng Su, Liyan Tan, Zheng Zhang
TL;DR
Large-scale pre-training is computationally costly, and practical polar iterations in Muon can amplify update imbalance rather than eliminate it. The paper introduces MUON+, which adds one normalization step after orthogonalization, and reports consistent improvements over Muon across GPT and LLaMA experiments, with up to 37.1% pre-training speedup.
Problem
Pre-training costs are prohibitively high at extreme model and dataset scales, while practical polar iterations can amplify Muon’s column/row update imbalance.
Source: §1 · pp. 1–2Method
MUON+ adds a single normalization step after orthogonalization to mitigate update imbalance, without adding optimizer state.
Source: §1 · p. 2Results
Takeaways & Limitations
Across the evaluated GPT and LLaMA settings, the added normalization improves pre-training performance with zero memory overhead.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Muon has recently emerged as a strong optimizer for large language model pre-training, orthogonalizing the momentum matrix via Newton--Schulz polar iterations. A natural intuition is that polar iterations, by flattening the singular spectrum to all ones, should also eliminate column- and row-wise norm imbalance in the update. We show that this is not true in practice: practical polar steps can substantially amplify the imbalance. We term this the post-polar imbalanced update problem, and prove that such imbalance tightens the second-order term in a blockwise descent analysis, weakening Muon's per-step descent guarantee. Motivated by this analysis, we propose Muon+, a one-line fix that inserts a single normalization step after polar orthogonalization. Muon+ adds no optimizer state. Across pre-training experiments on GPT and LLaMA models from 60M to 7B parameters, spanning both compute-optimal budgets and extended token-to-parameter ratios up to approximately 200, Muon+ consistently outperforms Muon in terms of training and validation perplexity, leading to significant overall pre-training speedup.
1 Introduction
Muon orthogonalizes gradient momentum to address rank collapse, but practical polar iterations can amplify update imbalance. MUON+ adds a normalization step and consistently outperforms Muon across the evaluated GPT and LLaMA pre-training settings.
1 Introduction
- The high computational cost of pre-training at extreme model and dataset scales has motivated research into more efficient optimizers.Adam and AdamW remain dominant, alongside proposals aimed at reducing computing or memory costs.Source: §1 · p. 1
- Muon orthogonalizes gradient momentum using Newton–Schulz iterations to flatten the singular spectrum and mitigate update rank collapse.Unlike Adam- or SGD-based optimizers, Muon operates on a matrix; Ortho(M) is the closest semi-orthogonal matrix function to M.Source: §1 · pp. 1–2
- MUON+ consistently outperforms Muon across GPT and LLaMA pre-training from 60M to 7B parameters, reducing validation perplexity by up to 2.02 and speeding training by up to 37.1%.It also improves average downstream zero-shot accuracy by up to 1.3% across seven tasks, with zero additional optimizer states.Source: Figure 1 · p. 2, §1 · p. 2
- Practical Newton–Schulz settings can amplify column/row update imbalance, which the authors identify as the post-polar imbalanced update problem.Their proposed MUON+ fix adds one normalization step after orthogonalization, with zero memory overhead and negligible computing overhead.Source: §1 · p. 2
2 Method
Section 2 defines update imbalance through nonuniform row and column norms, shows that practical polar iterations can worsen it, analyzes its effect on local descent guarantees, and introduces MUON+ as a normalization fix.
2.1 Update Imbalance Caused by Muon’s Polar Iterations
- Muon’s polar step can amplify row- or column-norm imbalance in practice; five Newton–Schulz steps still amplify variance, and non-square layers retain irreducible variance in one direction.The imbalance is measured through the variance of row or column norms; NorMuon and PolarExpress also fail to eliminate it under realistic pretraining dynamics.Source: §2.1 · pp. 3, 5
- Variance contributions depend on the Newton–Schulz derivative profile: spectral pairs with |g′(ξij)| > 1 increase variance, while those with |g′(ξij)| < 1 reduce it.The expected variance shift of one polar step can be positive, although variance decreases asymptotically as iterations tend to infinity.Source: §2.1 · p. 4
2.2 Imbalanced Updates Hurt Convergence in Muon
- MUON+ remains stable at larger learning rates, while Muon shows unstable training and degraded validation loss.The comparison is reported for stable step sizes.Source: Figure 5 · p. 5
- Under local smoothness and a near-block-diagonal Hessian, update imbalance weakens Muon’s one-step loss-decrease guarantee.With positive gradient–Hessian correlation preserved after the polar step, larger column-gradient imbalance raises block-weighted curvature and reduces the largest stable step size.Source: §2.2 · p. 5
2.3 An Easy-to-Deploy Solution: MUON+
- MUON+ adds a single normalization after polar iterations to reduce update imbalance.The method considers column-wise, row-wise, and sequentially composed normalizations.Source: §2.3 · p. 6
3 Experiments
The experiments compare MUON+ with Muon across compute-optimal and extended-training regimes, training efficiency, and ablations of normalization direction and polar methods.
3.1 Pre-training in Compute-Optimal Settings
- MUON+ consistently outperforms baselines across evaluated GPT and LLaMA scales in compute-optimal pre-training.The evaluation spans GPT models and LLaMA models from 60M to 7B parameters.Source: Table 1 · p. 6, §3.1 · pp. 6–7, Table 2 · p. 7
3.2 Pre-Training with High Token-to-Parameter Ratios
- At a token-to-parameter ratio of approximately 200, MUON+ beats Muon on both overtrained models with a stable training gap, and wins most downstream tasks with higher average accuracy.The overtraining runs use 72 billion FineWeb tokens; downstream results are for the GPT-Base overtraining checkpoint.Source: Table 3 · p. 7, §3.2 · pp. 7–8, Table 4 · p. 8
3.3 Training Efficiency
- MUON+ speeds pre-training by up to 37.1% to reach Muon’s final training loss, with nearly the same per-step runtime and memory cost.The comparison uses the same data scale as the compute-optimal experiments.Source: Table 5 · p. 8, §3.3 · p. 8
3.4 Ablation Study
- Across evaluated settings, bidirectional normalization outperforms single-direction normalization, and normalization improves optimization over baseline Muon as model scale increases.The direction ablation compares none, col, row, col_row, and row_col variants.Source: Table 6 · p. 8, §3.4 · p. 8
- MUON+ outperforms Muon across the tested polar methods, and adds gains even when the orthogonal update is computed exactly with SVD.The tested iterative methods are You, Jordan, and PolarExpress, each with five iterations.Source: §3.4 · p. 8, Table 7 · p. 9, Table 8 · p. 9
4 Related Work and Comparison with MUON+
Related work improves Muon through polar-function design, preconditioning, or theoretical analysis. MUON+ is complementary to polar-function methods, outperforms the described competing preconditioners, and offers a parameter-space account of update imbalance.
Optimizing the polar function
- Polar-function improvements are complementary to MUON+, and combining them with MUON+ yields additional performance improvements.The passage names You, POLAREXPRESS, and TURBO-MUON as examples.Source: §4 · p. 9
Pre-conditioning techniques
- MUON+ consistently outperforms the described preconditioning competitors NorMuon, AdaMuon, and Mano, without requiring additional optimizer states.NorMuon and AdaMuon require additional optimizer states; MUON+ does not.Source: §4 · p. 9
Theoretical foundations of Muon
- Prior theoretical views explain Muon through spectral, non-Euclidean, or isotropic-curvature formulations, while this work analyzes parameter-space imbalance and its effect on convergence.The passage also connects the analysis to why preconditioning methods may be effective.Source: §4 · p. 9
5 Conclusion
The paper concludes that Muon’s polar iterations can amplify parameter-space update imbalance, reducing its largest stable step size. MUON+ adds one normalization step after orthogonalization and improves pre-training performance across GPT and LLaMA experiments with zero memory overhead and significant overall speedup.
5 Conclusion
- Across GPT and LLaMA pre-training from 60M to 7B parameters, MUON+ consistently improved pre-training performance, with zero memory overhead and significant overall training speedup.Source: §5 · p. 9
A Blockwise Descent Analysis for Muon
The appendix develops a blockwise descent analysis for Muon under local Hessian assumptions and examines whether the polar step preserves relative block magnitudes. It reports that pre-polar rows that are relatively large tend to remain relatively large afterward.
A Blockwise Descent Analysis for Muon
- The analysis partitions the vectorized matrix into nonoverlapping blocks, with columns or rows as the two canonical choices.Under column blocks, each block is a column of W; under row blocks, each block is a row.Source: A Blockwise Descent Analysis for Muon · p. 13
A.1 Block-Hessian and weighted curvature
- The block-Hessian analysis defines each block's principal Hessian submatrix and an update-weighted block curvature.The update is described as orthogonalized, with a semi-orthogonal matrix.Source: §A.1 · p. 13
A.2 Local assumptions
- The local analysis assumes a Hessian-Lipschitz neighborhood and a Hessian that is near block diagonal under the chosen partition.The stated assumptions are local to the iterate.Source: §A.2 · p. 13
A.3 Technical lemma
- A technical lemma uses the fundamental theorem of calculus and the local Hessian-Lipschitz assumption to bound a remainder.The proof combines terms to establish the lemma.Source: §A.3 · p. 14
A.4 One-step descent bound: formal statement and proof
- The one-step descent bound is stated for a Muon-style orthogonalized update under local Hessian-Lipschitz and near-block-diagonal Hessian assumptions.The analysis also notes the near-block-diagonal slack for the selected partition.Source: §A.2 · p. 13, §A.4 · pp. 14–15
A.5 Positive Correlation Before and After the Polar Step
- Figure 7 shows consistently positive, generally substantial Spearman correlations between pre-polar and post-polar row norms across training.The measurements use row norms and average across layers.Source: §A.5 · p. 16, Figure 7 · p. 16
B.1 Proof of Lemma 2.1
- The proof of Lemma 2.1 expresses row energies using the singular spectrum and evaluates their variance under a Haar-distributed orthogonal factor.It also uses that total row energy is deterministic by orthogonality.Source: §B.1 · p. 17
B.2 Proof of Corollary 2.2
The proof links the Newton–Schulz map's behavior near zero to an increase in expected row-energy variance after transformation. The appendix also includes pseudocode for the Muon+ update.
B.2 Proof of Corollary 2.2
- For Muon's Newton–Schulz map, the proof considers singular values in [0, 1] and a map designed to move them toward 1.The map is written as φ(x) = ax + bx^3 + cx^5.Source: §B.2 · p. 18
- For λ1, . . . , λn in [0, δ] that are not all equal, the proof establishes a strictly positive expected variance difference: E[Var(s(Q)) − Var(s(M))] > 0.When a > 1, |g′(λ)| > 1 on a neighborhood of 0.Source: §B.2 · pp. 18–19
C Pseudocode of Muon+
- The appendix provides Python pseudocode for the Muon+ update.Source: C Pseudocode of Muon+ · p. 20
D.1 Jordan Coefficients [15]
The appendix presents imbalance-region figures for three polar-iteration coefficient schemes and lists architecture-configuration tables for GPT and LLaMA-style models.
D.1 Jordan Coefficients [15]
- The Jordan-coefficient subsection gives the coefficients used in and presents regions of imbalance amplification and reduction.Figure 8 depicts those regions.Source: §D.1 · p. 21, Figure 8 · p. 21
D.2 You Coefficients [6]
- The You scheme optimizes coefficients at each iteration and presents per-step regions of imbalance amplification and reduction.Figure 9 depicts the regions.Source: §D.2 · p. 21, Figure 9 · p. 21
D.3 PolarExpress Coefficients [2]
- PolarExpress [2] uses iteration-dependent coefficients, with subsequent coefficients numerically equal to (1.875, −1.25, 0.375), and has an imbalance-region figure.Figure 10 shows regions of imbalance amplification and reduction.Source: §D.3 · p. 22, Figure 10 · p. 22
E.1 Model Configurations
- The appendix includes architecture-configuration tables for GPT and LLaMA-style models.These are Tables 10 and 11, respectively.Source: Table 10 · p. 23, Table 11 · p. 23
E.2 Training Configurations
The training-configuration tables report model-specific settings for GPT and LLaMA FineWeb runs. They also specify settings for extended runs that train 72 billion tokens per configuration.
E.2.1 GPT Models
- GPT-Small/Base/Large/Huge FineWeb runs use Jordan orthogonalization, with sequence lengths of 2048/4096/8192 for Small/Base/Large and a constant-then-linear-decay learning-rate schedule.The learning rate is constant for the first 40% of training steps, then decays linearly to zero.Source: Table 12 · p. 23
- Each overtraining GPT-Base setting trains 72 billion FineWeb tokens with sequence length 4096.The settings are reported in Table 13.Source: Table 13 · p. 24
E.2.2 LLaMA Models
- LLaMA FineWeb runs use Jordan orthogonalization, with sequence lengths of 1024 for 60M/130M models and 4096 for 350M/1B models.Norm(col_row) and Norm(row_col) yield nearly identical performance.Source: Table 14 · p. 25
- Each overtraining LLaMA-350M setting trains 72 billion FineWeb tokens with sequence length 4096.The settings are reported in Table 15.Source: Table 15 · p. 25
F Detailed Sweep Experiments
The detailed sweeps compare best validation perplexity across norm settings and learning rates for LLaMA and GPT models. They also examine five-seed results and controls for learning-rate shifts and Frobenius-norm changes.
F Detailed Sweep Experiments
- Tables 16 and 17 compare best validation perplexity by norm setting and learning rate for LLaMA and GPT models, respectively.Figure 11 visualizes the LLaMA-350M sweep.Source: Table 16 · p. 26, Table 17 · p. 27, Figure 11 · p. 27
- Table 18 reports mean validation perplexity and standard deviation over five seeds for GPT-Base and LLaMA-350M.Lower perplexity is better; the caption says hyperparameters follow Tables 16 and 17.Source: Table 18 · p. 27
- Learning-rate sweeps show MUON+'s optimal learning rate is essentially unchanged from Muon's, suggesting its improvement is not due to an effective-rate shift.The authors also match the normalized update's Frobenius norm to its pre-normalization value, with results reported in Table 19.Source: F Detailed Sweep Experiments · pp. 27–28, Table 19 · p. 28
G Supplementary results
The supplementary results report per-step runtime for normalization variants and visualize the imbalance trend.
G.1 Per-step runtime
- All normalization variants introduce only negligible runtime overhead compared with Muon.Runtime is measured on an H100 GPU with mixed precision (bf16), batch size 4, and sequence length 4096, in ms/step.Source: §G.1 · p. 29, Table 20 · p. 29
G.2 Imbalance trend
- Figure 12 is a detailed visualization of the imbalance trend.Source: Figure 12 · p. 30