NorMuon: Rendere Muon più efficiente e scalabile

NorMuon: Making Muon more efficient and scalable

October 7, 2025

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

cs.AI

Abstract

La scelta dell'ottimizzatore influisce significativamente sull'efficienza dell'addestramento e sui costi computazionali dei grandi modelli linguistici (LLM). Recentemente, l'ottimizzatore Muon ha dimostrato risultati promettenti ortogonalizzando gli aggiornamenti dei parametri, migliorando la geometria dell'ottimizzazione attraverso un migliore condizionamento. Nonostante l'emergere di Muon come candidato successore di Adam, il potenziale per sfruttare congiuntamente i loro punti di forza non è stato sistematicamente esplorato. In questo lavoro, colmiamo questa lacuna proponendo NorMuon (Neuron-wise Normalized Muon), un ottimizzatore che combina sinergicamente l'ortogonalizzazione con tassi di apprendimento adattivi a livello di neurone. La nostra analisi rivela che, sebbene Muon riduca efficacemente i numeri di condizionamento, gli aggiornamenti risultanti presentano norme dei neuroni altamente non uniformi, causando il dominio di alcuni neuroni nel processo di ottimizzazione. NorMuon affronta questo squilibrio mantenendo statistiche del momento del secondo ordine per ogni neurone e applicando una normalizzazione per riga dopo l'ortogonalizzazione, garantendo un utilizzo bilanciato dei parametri preservando i benefici del condizionamento di Muon. Per consentire un'implementazione pratica su larga scala, sviluppiamo un'implementazione distribuita efficiente nel framework FSDP2 che distribuisce strategicamente i calcoli di ortogonalizzazione tra i dispositivi. Esperimenti su più scale di modelli dimostrano che NorMuon supera costantemente sia Adam che Muon, raggiungendo un'efficienza di addestramento migliore del 21,74% rispetto ad Adam e un miglioramento dell'11,31% rispetto a Muon in un contesto di pre-addestramento da 1,1 miliardi di parametri, mantenendo un'impronta di memoria comparabile a quella di Muon. I nostri risultati suggeriscono che l'ortogonalizzazione e i tassi di apprendimento adattivi sono approcci complementari piuttosto che concorrenti, aprendo nuove strade per la progettazione di ottimizzatori nell'apprendimento profondo su larga scala.

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.