NorMuon: Muon efficiënter en schaalbaarder maken

NorMuon: Making Muon more efficient and scalable

October 7, 2025

Auteurs: Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, Tuo Zhao

cs.AI

Samenvatting

De keuze van de optimizer heeft een aanzienlijke invloed op de trainings efficiëntie en de rekenkosten van grote taalmodellen (LLM's). Recentelijk heeft de Muon-optimizer veelbelovende resultaten laten zien door parameterupdates te orthogonaliseren, waardoor de optimalisatiegeometrie wordt verbeterd door betere conditionering. Ondanks dat Muon naar voren is gekomen als een mogelijke opvolger van Adam, is het potentieel om hun sterke punten gezamenlijk te benutten nog niet systematisch onderzocht. In dit werk overbruggen we deze kloof door NorMuon (Neuron-wise Normalized Muon) voor te stellen, een optimizer die orthogonalisatie synergetisch combineert met adaptieve leerpercentages op neuronniveau. Onze analyse laat zien dat Muon weliswaar effectief conditionele getallen reduceert, maar dat de resulterende updates sterk niet-uniforme neuronnormen vertonen, waardoor bepaalde neuronen het optimalisatieproces domineren. NorMuon lost deze onbalans op door tweede-orde momentumstatistieken voor elk neuron te behouden en rijgewijze normalisatie toe te passen na orthogonalisatie, waardoor een gebalanceerde parameterbenutting wordt gegarandeerd terwijl de conditioneringsvoordelen van Muon behouden blijven. Om praktische implementatie op grote schaal mogelijk te maken, ontwikkelen we een efficiënte gedistribueerde implementatie binnen het FSDP2-framework die orthogonalisatieberekeningen strategisch verdeelt over apparaten. Experimenten over meerdere modelschalen tonen aan dat NorMuon consistent beter presteert dan zowel Adam als Muon, met een 21,74% betere trainings efficiëntie dan Adam en een 11,31% verbetering ten opzichte van Muon in een 1,1 B pretrainingsinstelling, terwijl het een vergelijkbaar geheugenverbruik behoudt als Muon. Onze bevindingen suggereren dat orthogonalisatie en adaptieve leerpercentages complementair zijn in plaats van concurrerende benaderingen, wat nieuwe mogelijkheden opent voor het ontwerpen van optimizers in grootschalige deep learning.

English

The choice of optimizer significantly impacts the training efficiency and computational costs of large language models (LLMs). Recently, the Muon optimizer has demonstrated promising results by orthogonalizing parameter updates, improving optimization geometry through better conditioning. Despite Muon's emergence as a candidate successor to Adam, the potential for jointly leveraging their strengths has not been systematically explored. In this work, we bridge this gap by proposing NorMuon (Neuron-wise Normalized Muon), an optimizer that synergistically combines orthogonalization with neuron-level adaptive learning rates. Our analysis reveals that while Muon effectively reduces condition numbers, the resulting updates exhibit highly non-uniform neuron norms, causing certain neurons to dominate the optimization process. NorMuon addresses this imbalance by maintaining second-order momentum statistics for each neuron and applying row-wise normalization after orthogonalization, ensuring balanced parameter utilization while preserving Muon's conditioning benefits. To enable practical deployment at scale, we develop an efficient distributed implementation under the FSDP2 framework that strategically distributes orthogonalization computations across devices. Experiments across multiple model scales demonstrate that NorMuon consistently outperforms both Adam and Muon, achieving 21.74% better training efficiency than Adam and 11.31% improvement over Muon on 1.1 B pretraining setting, while maintaining a comparable memory footprint to Muon. Our findings suggest that orthogonalization and adaptive learning rates are complementary rather than competing approaches, opening new avenues for optimizer design in large-scale deep learning.