NorMuon: Making Muon more efficient and scalable

📄 arXiv: 2510.05491v1 📥 PDF

作者: Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, Tuo Zhao

分类: cs.LG, cs.CL

发布日期: 2025-10-07

💡 一句话要点

提出NorMuon优化器，结合正交化与神经元自适应学习率，提升大模型训练效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 优化器 大规模语言模型 正交化 自适应学习率 神经元归一化 分布式训练 FSDP2

📋 核心要点

  1. 现有优化器在大模型训练中存在效率瓶颈，Muon虽然通过正交化改善了优化几何，但神经元更新不均衡。
  2. NorMuon结合正交化与神经元自适应学习率，通过神经元归一化解决更新不均衡问题，提升参数利用率。
  3. 实验表明，NorMuon在多个模型规模上优于Adam和Muon，在11亿参数预训练中效率分别提升21.74%和11.31%。

📝 摘要（中文）

优化器的选择对大型语言模型（LLM）的训练效率和计算成本有显著影响。Muon优化器通过正交化参数更新，改善优化几何结构，展现出良好的效果。尽管Muon有潜力成为Adam的继任者，但尚未系统地探索如何联合利用它们的优势。本文提出了NorMuon（神经元归一化Muon），该优化器协同结合了正交化和神经元级别的自适应学习率。分析表明，Muon虽然有效地降低了条件数，但由此产生的更新表现出高度不均匀的神经元范数，导致某些神经元主导优化过程。NorMuon通过维护每个神经元的二阶动量统计量，并在正交化后应用行向归一化来解决这种不平衡，从而在保持Muon的条件优势的同时，确保参数的均衡利用。为了实现大规模的实际部署，我们开发了FSDP2框架下的高效分布式实现，该实现策略性地将正交化计算分布在各个设备上。在多个模型规模上的实验表明，NorMuon始终优于Adam和Muon，在11亿参数的预训练设置下，比Adam提高了21.74%的训练效率，比Muon提高了11.31%，同时保持了与Muon相当的内存占用。我们的研究结果表明，正交化和自适应学习率是互补而非竞争的方法，为大规模深度学习中的优化器设计开辟了新的途径。

🔬 方法详解

问题定义：在大规模语言模型训练中，优化器的选择至关重要。Muon优化器通过正交化参数更新来改善优化几何，但其更新后的神经元范数高度不均匀，导致部分神经元主导优化过程，限制了整体训练效率。现有方法未能有效解决这种神经元更新不均衡的问题。

核心思路：NorMuon的核心思路是结合Muon的正交化优势和神经元级别的自适应学习率。通过在正交化后对每个神经元的更新进行归一化，平衡不同神经元的更新幅度，从而更有效地利用模型参数，加速收敛。这种设计旨在弥补Muon在神经元更新均衡性方面的不足。

技术框架：NorMuon的整体框架基于Muon优化器，并在其基础上增加了神经元归一化模块。主要流程包括：1) 计算梯度；2) 应用Muon的正交化更新；3) 维护每个神经元的二阶动量统计量；4) 对每个神经元的更新进行行向归一化；5) 更新模型参数。该框架利用FSDP2进行分布式训练，将正交化计算分布到多个设备上。

关键创新：NorMuon的关键创新在于将正交化和神经元级别的自适应学习率相结合。与Muon相比，NorMuon通过神经元归一化解决了更新不均衡的问题，提高了参数利用率。与Adam相比，NorMuon利用正交化改善了优化几何，从而加速了收敛。这种结合是现有优化器设计中未充分探索的。

关键设计：NorMuon的关键设计包括：1) 使用二阶动量统计量来估计每个神经元的更新尺度；2) 在正交化更新后，对每个神经元的更新进行行向归一化，确保所有神经元的更新幅度大致相同；3) 利用FSDP2框架进行分布式训练，并通过策略性地分配正交化计算来提高效率。具体的归一化方法和动量参数的选择需要根据具体任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，NorMuon在多个模型规模上均优于Adam和Muon。在11亿参数的预训练设置下，NorMuon比Adam提高了21.74%的训练效率，比Muon提高了11.31%，同时保持了与Muon相当的内存占用。这些结果表明，NorMuon能够显著提升大模型的训练效率。

🎯 应用场景

NorMuon优化器可广泛应用于大规模深度学习模型的训练，尤其是在计算资源有限的情况下。其高效性和可扩展性使其成为训练大型语言模型、视觉模型以及其他复杂模型的理想选择。该研究为优化器设计提供了新的思路，有望推动人工智能领域的发展。

📄 摘要（原文）

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.