MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

Da Chang; Qiankun Shi; Lvgang Zhang; Yu Li; Ruijie Zhang; Yao Lu; Yongxiang Liu; Ganzhao Yuan

cs.LGMar 30, 2026

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 ·

Abstract

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 $\widetilde{\mathcal 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\in\{\mathrm{RC},\mathrm{R},\mathrm{C}\}

; default

s=\mathrm{R}

) 1: Input: Initial parameters

\mathbf{X}_{1}\in\mathbb{R}^{m\times n}

, learning rates

\eta_{t}>0

, momentum coefficients

\{\beta_{t}\}_{t\geq 1}\subset[0,1)

, weight decay

\lambda_{t}\geq 0

, mode

s\in\{\mathrm{RC},\mathrm{R},\mathrm{C}\}

, stability term

\varepsilon\geq 0

, Nesterov flag, scaling constant

a=0.2\sqrt{\max(m,n)}

. 2: Initialize

\mathbf{M}_{0}=\mathbf{0}

3: for

t=1

T

do 4:

\mathbf{G}_{t}=\nabla f(\mathbf{X}_{t};\xi_{t})

\mathbf{M}_{t}=\beta_{t}\mathbf{M}_{t-1}+(1-\beta_{t})\mathbf{G}_{t}

6: if Nesterov then 7:

\widetilde{\mathbf{M}}_{t}=\beta_{t+1}\mathbf{M}_{t}+(1-\beta_{t+1})\mathbf{G}_{t}

8: else 9:

\widetilde{\mathbf{M}}_{t}=\mathbf{M}_{t}

10: end if 11:

\hat{\mathbf{M}}_{t}=\textbf{DiagPre}(\widetilde{\mathbf{M}}_{t},s,\varepsilon)

12:

\mathbf{O}_{t}=\mathrm{NS5}(\hat{\mathbf{M}}_{t})

13:

\mathbf{X}_{t+1}=(1-\lambda_{t}\eta_{t})\mathbf{X}_{t}-a\eta_{t}\mathbf{O}_{t}

14: end for

Algorithm 2 DiagPre 1: Input: Momentum

\widetilde{\mathbf{M}}_{t}

, mode

s\in\{\mathrm{RC},\mathrm{R},\mathrm{C}\}

, stability term

\varepsilon\geq 0

2: if

s\in\{\mathrm{RC},\mathrm{R}\}

then 3:

\mathbf{D}_{r,t}=\operatorname{diag}(\operatorname{rowsum}(\widetilde{\mathbf{M}}_{t}\odot\widetilde{\mathbf{M}}_{t})+\varepsilon)

4: end if 5: if

s\in\{\mathrm{RC},\mathrm{C}\}

then 6:

\mathbf{D}_{c,t}=\operatorname{diag}(\operatorname{colsum}(\widetilde{\mathbf{M}}_{t}\odot\widetilde{\mathbf{M}}_{t})+\varepsilon)

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

\hat{\mathbf{M}}_{t}

Remark 2.1 . Unless otherwise stated, we use

s=\mathrm{R}

;

\mathrm{RC}

and

\mathrm{C}

are retained for analysis and ablation. For varying

\beta_{t}

, the Nesterov line follows the EMA-buffer convention; for

\beta_{t}\equiv\beta

, it reduces to the Nesterov form Yuan et al. (2024) ; Chang et al. (2025) .

Table 1

Panel of Figure 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.

Panel of Figure 2 — 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.

Panel of Figure 3 — 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
Optimizer	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 $\downarrow$ 0.56	12.52	0.26	15.07 $\downarrow$ 0.14	28.18	0.63	12.86 $\downarrow$ 0.20	34.81	2.78

Table 1: Best validation perplexity, peak memory usage, and wall-clock training time on LLaMA2/C4.

Panel of Figure 4 — 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

Panel of Figure 5 — Figure 5: Learning-rate sweeps for LLaMA2-130M and LLaMA2-350M trained on C4 for 2.6B and 7.5B tokens, respectively.

Panel of Figure 6 — 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 .

Panel of Figure 7 — 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 $\|\mathrm{NS}_{5}(S(\mathbf{M}))-\mathrm{Orth}(S(\mathbf{M}))\|_{F}$ , while the bottom panel reports $\|\mathrm{Orth}(S(\mathbf{M}))-\mathrm{Orth}(\mathbf{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.

Panel of Figure 8 — 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
$(\beta_{1},\beta_{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
$(\beta_{1},\beta_{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
$(\beta_{1},\beta_{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
$(\beta_{1},\beta_{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

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 321 papers

AMO: Operator-level Adaptive Muon Orthogonalization
May 18, 2026Xinlin Zhuang, Panyi Ouyang, Yichen Li +7Language Model Pretraining Newton-Schulz Iteration
The Chinese University of Hong Kong · Shopee · MBZUAI +3
MUON+: Towards More Effective Muon via One Additional Normalization Step for LLM Pre-training
Feb 25, 2026Ruijie Zhang, Yequan Zhao, Ziyue Liu +4Muon Optimizer Deep Learning Optimization
University of California at Santa Barbara
Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates
Dec 4, 2025Thibaut Boissin, Thomas Massena, Franck Mamalet +1Newton-Schulz Iteration Orthogonal Optimization
IRIT-MISFIT · Institut de Recherche Technologique Saint-Exupery, France · DTIPG - SNCF, IRIT-MISFIT +2

May 18, 2026J/K move · Enter open · S save

AMO: Operator-level Adaptive Muon Orthogonalization

Xinlin Zhuang, Panyi Ouyang, Yichen Li, Jiangming Shi+6

The Chinese University of Hong Kong · Shopee · MBZUAI+3

Language Model Pretraining Newton-Schulz Iteration Muon Optimizer Second-Order Optimization Adaptive Gradient Methods

Open paper AskPDF arXiv

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.

Abstract

Neural Network Optimization

Explore similar work

AMO: Operator-level Adaptive Muon Orthogonalization

MUON+: Towards More Effective Muon via One Additional Normalization Step for LLM Pre-training

Turbo-Muon: Almost-Orthogonal Pre-Conditioning for Fast Muon Updates