MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration
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 \href{https://github.com/MaeChd/muon-eq}{MuonEq codebase}.
Code (0)
등록된 구현이 없습니다.
Similar Papers 제목 키워드 기반
Muon$^2$: Boosting Muon via Adaptive Second-Moment Preconditioning
Muon has emerged as a promising optimizer for large-scale foundation model pre-training by exploiting the matrix structure of neural network updates through iterative orthogonalization. However, the orthogonalization qua…
Analysis of Reaction Network Systems Using Tropical Geometry
We discuss a novel analysis method for reaction network systems with polynomial or rational rate functions. This method is based on computing tropical equilibrations defined by the equality of at least two dominant monom…
Qualitative dynamics of chemical reaction networks: an investigation using partial tropical equilibrations
We discuss a method to describe the qualitative dynamics of chemical reaction networks in terms of symbolic dynamics. The method, that can be applied to mass-action reaction networks with separated timescales, uses solut…
Towards Continual Learning Desiderata via HSIC-Bottleneck Orthogonalization and Equiangular Embedding
Deep neural networks are susceptible to catastrophic forgetting when trained on sequential tasks. Various continual learning (CL) methods often rely on exemplar buffers or/and network expansion for balancing model stabil…
Continual LearningThe Invisible Hand of Laplace: the Role of Market Structure in Price Convergence and Oscillation
A fundamental question about a market is under what conditions, and then how rapidly, does price signaling cause price equilibration. Qualitatively, this ought to depend on how well-connected the market is. We address th…