NorMuon : Rendre Muon plus efficace et évolutif

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

Le choix de l'optimiseur a un impact significatif sur l'efficacité de l'entraînement et les coûts computationnels des grands modèles de langage (LLM). Récemment, l'optimiseur Muon a démontré des résultats prometteurs en orthogonalisant les mises à jour des paramètres, améliorant ainsi la géométrie de l'optimisation grâce à un meilleur conditionnement. Bien que Muon émerge comme un candidat successeur potentiel à Adam, le potentiel d'une exploitation conjointe de leurs forces n'a pas été systématiquement exploré. Dans ce travail, nous comblons cette lacune en proposant NorMuon (Muon Normalisé par Neurone), un optimiseur qui combine de manière synergique l'orthogonalisation avec des taux d'apprentissage adaptatifs au niveau des neurones. Notre analyse révèle que si Muon réduit efficacement les nombres de conditionnement, les mises à jour résultantes présentent des normes de neurones très non uniformes, ce qui entraîne la domination de certains neurones dans le processus d'optimisation. NorMuon corrige ce déséquilibre en maintenant des statistiques de moment du second ordre pour chaque neurone et en appliquant une normalisation par ligne après l'orthogonalisation, assurant ainsi une utilisation équilibrée des paramètres tout en préservant les avantages de conditionnement de Muon. Pour permettre un déploiement pratique à grande échelle, nous développons une implémentation distribuée efficace sous le cadre FSDP2 qui répartit stratégiquement les calculs d'orthogonalisation entre les dispositifs. Les expériences menées sur plusieurs échelles de modèles démontrent que NorMuon surpasse systématiquement à la fois Adam et Muon, atteignant une efficacité d'entraînement 21,74 % supérieure à celle d'Adam et une amélioration de 11,31 % par rapport à Muon dans un contexte de pré-entraînement de 1,1 milliard de paramètres, tout en maintenant une empreinte mémoire comparable à celle de Muon. Nos résultats suggèrent que l'orthogonalisation et les taux d'apprentissage adaptatifs sont des approches complémentaires plutôt que concurrentes, ouvrant de nouvelles voies pour la conception d'optimiseurs dans l'apprentissage profond à grande échelle.

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.