NorMuon: Haciendo a Muon más eficiente y escalable

Resumen

La elección del optimizador impacta significativamente la eficiencia del entrenamiento y los costos computacionales de los grandes modelos de lenguaje (LLMs, por sus siglas en inglés). Recientemente, el optimizador Muon ha demostrado resultados prometedores al ortogonalizar las actualizaciones de parámetros, mejorando la geometría de optimización mediante un mejor condicionamiento. A pesar de que Muon ha surgido como un candidato sucesor de Adam, el potencial para aprovechar conjuntamente sus fortalezas no ha sido explorado sistemáticamente. En este trabajo, cerramos esta brecha proponiendo NorMuon (Muon Normalizado por Neurona), un optimizador que combina sinérgicamente la ortogonalización con tasas de aprendizaje adaptativas a nivel de neurona. Nuestro análisis revela que, aunque Muon reduce efectivamente los números de condición, las actualizaciones resultantes exhiben normas de neurona altamente no uniformes, lo que hace que ciertas neuronas dominen el proceso de optimización. NorMuon aborda este desequilibrio manteniendo estadísticas de momento de segundo orden para cada neurona y aplicando normalización por filas después de la ortogonalización, asegurando un uso equilibrado de los parámetros mientras preserva los beneficios de condicionamiento de Muon. Para permitir una implementación práctica a gran escala, desarrollamos una implementación distribuida eficiente bajo el marco FSDP2 que distribuye estratégicamente los cálculos de ortogonalización entre dispositivos. Los experimentos en múltiples escalas de modelos demuestran que NorMuon supera consistentemente tanto a Adam como a Muon, logrando un 21.74% mejor eficiencia de entrenamiento que Adam y un 11.31% de mejora sobre Muon en un escenario de preentrenamiento de 1.1 B, manteniendo una huella de memoria comparable a la de Muon. Nuestros hallazgos sugieren que la ortogonalización y las tasas de aprendizaje adaptativas son enfoques complementarios en lugar de competidores, abriendo nuevas vías para el diseño de optimizadores en el aprendizaje profundo a gran escala.

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.

NorMuon: Haciendo a Muon más eficiente y escalable

NorMuon: Making Muon more efficient and scalable

Resumen

Support