NorMuon: Steigerung der Effizienz und Skalierbarkeit von Muon

NorMuon: Making Muon more efficient and scalable

October 7, 2025

papers.authors: Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, Tuo Zhao

cs.AI

papers.abstract

Die Wahl des Optimierers hat einen erheblichen Einfluss auf die Trainings effizienz und die Rechenkosten von großen Sprachmodellen (LLMs). Kürzlich hat der Muon-Optimierer vielversprechende Ergebnisse gezeigt, indem er Parameteraktualisierungen orthogonalisiert und die Optimierungsgeometrie durch eine bessere Konditionierung verbessert. Obwohl Muon als potenzieller Nachfolger von Adam aufgetreten ist, wurde das Potenzial, ihre Stärken gemeinsam zu nutzen, bisher nicht systematisch untersucht. In dieser Arbeit schließen wir diese Lücke, indem wir NorMuon (Neuron-wise Normalized Muon) vorschlagen, einen Optimierer, der Orthogonalisierung mit adaptiven Lernraten auf Neuronenniveau synergetisch kombiniert. Unsere Analyse zeigt, dass Muon zwar effektiv Konditionszahlen reduziert, die resultierenden Aktualisierungen jedoch stark nicht-uniforme Neuronennormen aufweisen, was dazu führt, dass bestimmte Neuronen den Optimierungsprozess dominieren. NorMuon behebt dieses Ungleichgewicht, indem es zweite Momente für jedes Neuron beibehält und nach der Orthogonalisierung eine zeilenweise Normalisierung anwendet, wodurch eine ausgewogene Parameterausnutzung gewährleistet wird, während die Konditionierungsvorteile von Muon erhalten bleiben. Um eine praktische Implementierung in großem Maßstab zu ermöglichen, entwickeln wir eine effiziente verteilte Implementierung unter dem FSDP2-Framework, das Orthogonalisierungsberechnungen strategisch über Geräte verteilt. Experimente über mehrere Modellskalen hinweg zeigen, dass NorMuon durchweg sowohl Adam als auch Muon übertrifft und eine 21,74 % bessere Trainings effizienz als Adam und eine 11,31 % Verbesserung gegenüber Muon im 1,1 B-Pretraining-Setting erzielt, während es einen vergleichbaren Speicherbedarf wie Muon beibehält. Unsere Ergebnisse deuten darauf hin, dass Orthogonalisierung und adaptive Lernraten eher komplementär als konkurrierende Ansätze sind und neue Wege für das Design von Optimierern im groß angelegten Deep Learning eröffnen.

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.