English

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

Machine Learning 2026-05-12 v2 Machine Learning

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 O~(T1/4)\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 \href{https://github.com/MaeChd/muon-eq}{MuonEq codebase}.

Cite

@article{arxiv.2603.28254,
  title  = {MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration},
  author = {Da Chang and Qiankun Shi and Lvgang Zhang and Yu Li and Ruijie Zhang and Yao Lu and Yongxiang Liu and Ganzhao Yuan},
  journal= {arXiv preprint arXiv:2603.28254},
  year   = {2026}
}
R2 v1 2026-07-01T11:43:50.702Z