NorMuon: Повышение эффективности и масштабируемости Muon

NorMuon: Making Muon more efficient and scalable

October 7, 2025

Авторы: Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, Tuo Zhao

cs.AI

Аннотация

Выбор оптимизатора существенно влияет на эффективность обучения и вычислительные затраты больших языковых моделей (LLM). Недавно оптимизатор Muon продемонстрировал многообещающие результаты за счет ортогонализации обновлений параметров, улучшая геометрию оптимизации через лучшее обусловливание. Несмотря на появление Muon как потенциального преемника Adam, возможность совместного использования их преимуществ систематически не исследовалась. В данной работе мы устраняем этот пробел, предлагая NorMuon (Neuron-wise Normalized Muon) — оптимизатор, который синергетически сочетает ортогонализацию с адаптивными скоростями обучения на уровне нейронов. Наш анализ показывает, что хотя Muon эффективно снижает числа обусловленности, результирующие обновления демонстрируют высокую неравномерность норм нейронов, что приводит к доминированию некоторых нейронов в процессе оптимизации. NorMuon устраняет этот дисбаланс, поддерживая статистику второго момента для каждого нейрона и применяя построчную нормализацию после ортогонализации, обеспечивая сбалансированное использование параметров при сохранении преимуществ обусловленности Muon. Для практического масштабирования мы разработали эффективную распределенную реализацию в рамках FSDP2, стратегически распределяя вычисления ортогонализации между устройствами. Эксперименты на моделях различных масштабов демонстрируют, что NorMuon стабильно превосходит как Adam, так и Muon, достигая на 21,74% лучшей эффективности обучения по сравнению с Adam и на 11,31% улучшения по сравнению с Muon в настройке предварительного обучения на 1,1 млрд параметров, при этом сохраняя сопоставимый с Muon объем памяти. Наши результаты свидетельствуют о том, что ортогонализация и адаптивные скорости обучения являются взаимодополняющими, а не конкурирующими подходами, открывая новые направления в разработке оптимизаторов для крупномасштабного глубокого обучения.

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.