NorMuon: Tornando o Muon mais eficiente e escalável

NorMuon: Making Muon more efficient and scalable

October 7, 2025

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

cs.AI

Resumo

A escolha do otimizador impacta significativamente a eficiência do treinamento e os custos computacionais de grandes modelos de linguagem (LLMs). Recentemente, o otimizador Muon demonstrou resultados promissores ao ortogonalizar as atualizações de parâmetros, melhorando a geometria de otimização por meio de um melhor condicionamento. Apesar do Muon emergir como um candidato a sucessor do Adam, o potencial de aproveitar conjuntamente seus pontos fortes não foi sistematicamente explorado. Neste trabalho, preenchemos essa lacuna propondo o NorMuon (Muon Normalizado por Neurônio), um otimizador que combina sinergicamente a ortogonalização com taxas de aprendizado adaptativas em nível de neurônio. Nossa análise revela que, embora o Muon reduza efetivamente os números de condição, as atualizações resultantes exibem normas de neurônios altamente não uniformes, fazendo com que certos neurônios dominem o processo de otimização. O NorMuon aborda esse desequilíbrio mantendo estatísticas de momento de segunda ordem para cada neurônio e aplicando normalização por linha após a ortogonalização, garantindo uma utilização equilibrada dos parâmetros enquanto preserva os benefícios de condicionamento do Muon. Para permitir a implantação prática em escala, desenvolvemos uma implementação distribuída eficiente sob o framework FSDP2 que distribui estrategicamente os cálculos de ortogonalização entre dispositivos. Experimentos em múltiplas escalas de modelos demonstram que o NorMuon supera consistentemente tanto o Adam quanto o Muon, alcançando 21,74% de melhor eficiência de treinamento em relação ao Adam e 11,31% de melhoria sobre o Muon em um cenário de pré-treinamento de 1,1 B, enquanto mantém uma pegada de memória comparável ao Muon. Nossos achados sugerem que a ortogonalização e as taxas de aprendizado adaptativas são abordagens complementares, e não concorrentes, abrindo novos caminhos para o design de otimizadores em aprendizado profundo em larga 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.