Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
Addressing the limitation of one-sided row normalization in existing Muon variants, the authors introduce MeqMuon to balance both row and column update magnitudes adaptively across varying matrix imbalance patterns. The optimizer eliminates AdamW second-moment tracking entirely, substantially lowering optimizer-state memory overhead during large-scale training. Across LLM pretraining evaluations, MeqMuon achieves superior convergence rates and final loss compared to AdamW, standard Muon, and related normalization baselines.
One-sided matrix normalization leaves performance on the table: equilibrating both rows and columns accelerates LLM pretraining convergence while completely shedding AdamW's second-moment memory footprint.
The success of large language models (LLMs) has been accompanied by continued growth in model size and pretraining costs. Muon offers high accuracy and training efficiency in LLM pretraining. Recent work introduces row-wise normalization into Muon to balance update magnitudes and improve pretraining performance. However, row-wise normalization alone cannot accommodate different imbalance patterns in update matrices. In this paper, we propose an improved Muon optimizer, called \underline{m}atrix-\underline{eq}uilibrating Muon~(MeqMuon), for LLM pretraining. MeqMuon balances both row and column magnitudes through normalization that can be automatically tailored to different imbalance patterns without manual intervention. Moreover, MeqMuon eliminates the need to store AdamW's second-moment estimates, reducing optimizer-state memory usage. Empirical results demonstrate that MeqMuon achieves better convergence performance than AdamW, Muon, and other baselines in LLM pretraining.