NorMuon让优化器更高效,通过平衡神经元更新提升大模型训练速度。
NorMuon: Making Muon more efficient and scalable
- 在Muon基础上引入神经元级自适应学习率,结合正交化与归一化。
- 11亿参数训练中比Muon快11.31%,比Adam快21.74%。
- 适合追求高效率的大规模语言模型训练团队使用。
优化器选择显著影响大语言模型的训练效率与计算成本。近期,基于参数更新正交化的Muon优化器通过改善优化几何结构,展现出良好性能,成为Adam的潜在替代方案。然而,其与现有优势方法的协同潜力尚未系统探索。本文提出NorMuon(神经元级归一化Muon),通过在正交化后引入逐行归一化并保留每个神经元的二阶动量统计,实现参数更新的均衡利用,同时保持良好的条件数特性。为支持大规模部署,我们基于FSDP2框架设计了高效的分布式实现,智能分配正交化计算任务。在多个模型规模上的实验表明,NorMuon持续优于Adam与Muon:在11亿参数预训练设置下,相较Adam提升21.74%训练效率,相较Muon提升11.31%,且内存开销与Muon相当。研究揭示正交化与自适应学习率具有互补性,为大模型优化器设计开辟新方向。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。