MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration
Authors: Da Chang, Qiankun Shi, Lvgang Zhang, Yu Li, Ruijie Zhang, Yao Lu, Yongxiang Liu, Ganzhao Yuan
Organizations: Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · Pengcheng Laboratory · University of Chinese Academy of Sciences · Sun Yat-sen University · Southern University of Science and Technology · George Washington University · Shenzhen University of Advanced Technology
Orthogonalized-update optimizers such as Muon improve training of matrix-valued parameters, but existing extensions typically either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it. We introduce {\method}, a lightweight family of pre-orthogonalization equilibration schemes for Muon with three forms: two-sided row/column normalization (RC), row normalization (R), and column normalization (C). By rebalancing the momentum matrix before finite-step Newton--Schulz orthogonalization, {\method} improves the geometry seen by orthogonalization. We show that finite-step orthogonalization is governed by the input spectrum, especially stable rank and condition number, and that row/column normalization acts as a zeroth-order surrogate for whitening. For hidden matrix weights, R is the default variant. Theoretically, {\method} (R) retains the standard O(T−1/4) Muon-type nonconvex stationarity guarantee with decoupled weight decay and a horizon-free diminishing learning-rate schedule, and extends it to finite-step NS5 up to an explicit inexactness constant. In LLaMA2 pretraining on C4, {\method} (R) consistently outperforms Muon on 130M, 350M, and 1B models, with faster convergence and lower validation perplexity. The code is available at the MuonEq codebase.
Figures & tables
Algorithm 1 MuonEq (mode s∈{RC,R,C} ; default s=R ) 1: Input: Initial parameters X1∈Rm×n , learning rates ηt>0 , momentum coefficients {βt}t≥1⊂[0,1) , weight decay λt≥0 , mode s∈{RC,R,C} , stability term ε≥0 , Nesterov flag, scaling constant a=0.2max(m,n) . 2: Initialize M0=0 3: for t=1 to T do 4: Gt=∇f(Xt;ξt) 5: Mt=βtMt−1+(1−βt)Gt 6: if Nesterov then 7: Mt=βt+1Mt+(1−βt+1)Gt 8: else 9: Mt=Mt 10: end if 11: M^t=DiagPre(Mt,s,ε) 12: Ot=NS5(M^t) 13: Xt+1=(1−λtηt)Xt−aηtOt 14: end for
Algorithm 2 DiagPre 1: Input: Momentum Mt , mode s∈{RC,R,C} , stability term ε≥0 2: if s∈{RC,R} then 3: Dr,t=diag(rowsum(Mt⊙Mt)+ε) 4: end if 5: if s∈{RC,C} then 6: Dc,t=diag(colsum(Mt⊙Mt)+ε) 7: end if 8: \hat{\mathbf{M}}_{t}=\left\{\begin{array}[]{ll}\mathbf{D}_{r,t}^{-1/2}\widetilde{\mathbf{M}}_{t}\mathbf{D}_{c,t}^{-1/2},&\hbox{s=\mathrm{RC};}\\ \mathbf{D}_{r,t}^{-1/2}\widetilde{\mathbf{M}}_{t},&\hbox{s=\mathrm{R};}\\ \widetilde{\mathbf{M}}_{t}\mathbf{D}_{c,t}^{-1/2},&\hbox{s=\mathrm{C}.}\end{array}\right. 9: return M^t Remark 2.1 . Unless otherwise stated, we use s=R ; RC and C are retained for analysis and ablation. For varying βt , the Nesterov line follows the EMA-buffer convention; for βt≡β , it reduces to the Nesterov form Yuan et al. (2024) ; Chang et al. (2025) .
Table 1
Figure 1: Random Gaussian matrices with controlled shapes and spectral spreads. Top: finite-step relative Frobenius error to the exact polar factor across Newton–Schulz steps. Bottom: raw and post-normalization condition numbers. Two-sided row/column normalization yields the smallest error and the most consistent spectral compression.
Figure 2: Finite-step orthogonalization error across Newton–Schulz steps at 1%, 10%, 50%, and 100% of training. In each column, the top panel shows the module-wise median and the bottom panel shows the module-wise mean; shaded bands denote the 25%–75% range. Two-sided row/column normalization decays fastest.
Figure 3: Comparison of LLaMA2 130M, 350M, and 1B from three perspectives. In each column, the top panel corresponds to the 130M model, the middle panel corresponds to the 350M model, and the bottom panel corresponds to the 1B model. From left to right, the panels show training loss versus tokens, validation loss versus tokens, and validation loss versus wall-clock time.
Optimizer
130M
350M
1B
Perplexity
Memory
Time
Perplexity
Memory
Time
Perplexity
Memory
Time
AdamW Loshchilov and Hutter (2019)
21.35
12.94
0.21
16.25
29.43
0.53
14.71
39.56
2.47
MARS-AdamW Yuan et al. (2024)
20.03
12.99
0.23
15.62
29.47
0.57
13.22
39.57
2.50
Muon Jordan et al. () ; Liu et al. (2025a)
20.52
12.52
0.25
15.21
28.18
0.61
13.06
34.81
2.76
Muon-Nes Jordan et al. () ; Liu et al. (2025a) ; Chang et al. (2025)
20.18
12.52
0.25
15.19
28.18
0.61
13.24
34.81
2.76
MuonEq
19.96 ↓ 0.56
12.52
0.26
15.07 ↓ 0.14
28.18
0.63
12.86 ↓ 0.20
34.81
2.78
Table 1: Best validation perplexity, peak memory usage, and wall-clock training time on LLaMA2/C4.
Figure 4: (a,b) Validation perplexity and step time vs. Newton–Schulz iterations K on GPT2-small/FineWeb (10.5B tokens); (c,d) validation perplexity heatmaps over learning rate and momentum at K=5 (2.6B tokens).
Table 7
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Learning-rate sweeps for LLaMA2-130M and LLaMA2-350M trained on C4 for 2.6B and 7.5B tokens, respectively.
Figure 6: Singular-value entropy and stable rank of Muon momentum matrices under different normalization schemes during LLaMA2-130M training on C4 for 2.6B tokens. RowColNorm consistently improves both metrics over Direct , ColNorm , and RowNorm .
Figure 7: Layerwise Muon NS5 bias decomposition under different normalization schemes during LLaMA2-130M training on C4 for 2.6B tokens. The top panel reports ∥NS5(S(M))−Orth(S(M))∥F , while the bottom panel reports ∥Orth(S(M))−Orth(M)∥F , both evaluated across layers at the same training step. These plots reveal how the choice of normalization shifts the error budget between the finite-step approximation error term and the preconditioning bias term.
Figure 8: Validation perplexity heatmaps on GPT2-small/FineWeb (2.6B tokens) over learning rate and momentum at K=4 and K=5 .
Hyper-parameter
AdamW
MARS-AdamW
Muon
Muon-Nes
MuonEq
MuonEq-Nes
Max Learning Rate
5e-3
3e-3
3e-3
3e-3
3e-3
3e-3
Warmup Steps
500
Batch Size
128
Maximum Length
4096
Weight Decay
0.1
(β1,β2)
(0.9,0.95)
Appendix
Table 4: Hyperparameters used for training LLaMA2-130M on C4
Hyper-parameter
AdamW
MARS-AdamW
Muon
Muon-Nes
MuonEq
MuonEq-Nes
Max Learning Rate
1.5e-3
2e-3
1.5e-3
2e-3
2e-3
2e-3
Warmup Steps
500
Batch Size
256
Maximum Length
4096
Weight Decay
0.1
(β1,β2)
(0.9,0.95)
Appendix
Table 5: Hyperparameters used for training LLaMA2-350M on C4
Hyper-parameter
AdamW
MARS-AdamW
Muon
Muon-Nes
MuonEq
MuonEq-Nes
Max Learning Rate
{5e-4, 8e-4}
Warmup Steps
1000
Batch Size
512
Maximum Length
4096
Weight Decay
0.1
(β1,β2)
(0.9,0.95)
Appendix
Table 6: Hyperparameters used for training LLaMA2-1B on C4
Hyper-parameter
SGD
AdamW
Muon
Muon-Nes
MuonEq-Nes (RC)
MuonEq-Nes (C)
MuonEq-Nes (R)
Max Learning Rate
5e-2
1e-3
5e-2
5e-2
5e-2
5e-2
5e-2
Warmup Steps
0
Epochs
100
Batch Size
128
(β1,β2)
(0.9,0.999)
(0.9,0.95)
Muon-Momentum
✗
✗
0.9
Appendix
Table 7: Hyperparameters used for training ResNet-18 on CIFAR10
Hyper-parameter
AdamW
AdaMuon
Muon+ (colrow)
Mousse
Muon-Nes
MuonEq-Nes (RC)
MuonEq-Nes (C)
MuonEq-Nes (R)
Max Learning Rate
2e-3
5e-3
5e-3
2e-3
2e-3
2e-3
2e-3
2e-3
Warmup Steps
1000
Total Steps
20000
Batch Size
128
Maximum Length
4096
Weight Decay
0.1
Appendix
Table 8: Hyperparameters used for training GPT2-small on FineWeb
Muon has recently emerged as a competitive alternative to AdamW for large-scale pre-training, with orthogonalization via Newton-Schulz (NS) iteration as its core operation. Standard Muon applies a uniform NS schedule to all parameter matrices, overlooking possible differences in orthogonalization difficulty and its impact on performance. Through a systematic empirical study, we show that this per-matrix heterogeneity is pervasive and strongly associated with matrix geometry, which evolves dynamically across operator types, training stages, and network depths. Therefore, uniform NS schedules can lead to uneven orthogonalization quality across the model. Motivated by these findings, we propose Operator-level Adaptive Muon Orthogonalization (AMO), an observe-then-commit method that measures weight geometry by operator type early in training and then uses these signals to allocate the NS budget for the remainder of training. AMO delivers consistent improvements over uniform-schedule Muon across standard, prolonged, and continual pre-training, surpassing the strongest baseline by +0.76 on Llama3.1-1.4B and +0.51 on Qwen3-1.7B in average downstream performance of 12 evaluation tasks, with gains persisting at Llama3.1-4B scale.
Xinlin Zhuang, Panyi Ouyang, Yichen Li +7
The Chinese University of Hong Kong · Shopee · MBZUAI +3
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.
Orthogonality-based optimizers, such as Muon, have recently shown strong performance across large-scale training and community-driven efficiency challenges. However, these methods rely on a costly gradient orthogonalization step. Even efficient iterative approximations such as Newton-Schulz remain expensive, typically requiring dozens of matrix multiplications to converge. We introduce a pre-conditioning procedure that improves the initialization of the Newton--Schulz iterations while incurring negligible overhead. Furthermore, our pre-conditioning reduces the initial polar error and enables the removal of one Newton-Schulz iteration (out of the five iterations usually used in practice). The resulting implementation significantly reduces Muon's overhead. At the end-to-end training level, we observe consistent runtime improvements across speed-run and standard benchmarks, including ∼3% reductions in training time on multiple fast training benchmarks, while matching reference performance on both language and vision tasks. Crucially, these improvements require no hyperparameter tuning and can be adopted as a simple drop-in replacement. Beyond empirical gains, we provide theoretical insight into the geometry of the update and its potential robustness against feature collapse. Our code is publicly available on github, in optax and huggingface kernels.
Thibaut Boissin, Thomas Massena, Franck Mamalet +1
IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT +2