Muon+: Towards Better Muon via One Additional Normalization Step

📄 arXiv: 2602.21545v1 📥 PDF

作者: Ruijie Zhang, Yequan Zhao, Ziyue Liu, Zhengyang Wang, Zheng Zhang

分类: cs.LG

发布日期: 2026-02-25

🔗 代码/项目: GITHUB

💡 一句话要点

Muon+：通过额外的归一化步骤提升Muon优化器性能

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

关键词: 优化器 大语言模型 预训练 梯度正交化 归一化 深度学习

📋 核心要点

  1. 现有Muon优化器在训练大型语言模型时，虽然通过梯度正交化取得了一定效果，但仍有提升空间。
  2. Muon+的核心在于在Muon的正交化步骤后，增加一个额外的归一化步骤，以进一步稳定训练过程。
  3. 实验结果表明，Muon+在不同规模和架构的模型上，均能持续提升训练和验证困惑度，尤其是在高T2P比率下。

📝 摘要（中文）

Muon优化器通过梯度（或动量）正交化，在大语言模型的预训练中表现出良好的性能。本文提出了一种简单而有效的Muon增强方法，称为Muon+，它在正交化之后引入了一个额外的归一化步骤。通过对各种模型规模和架构进行广泛的预训练实验，证明了Muon+的有效性。我们的评估包括参数量从130M到774M的GPT风格模型，以及参数量从60M到1B的LLaMA风格模型。我们全面评估了Muon+在计算最优训练机制中的有效性，并将token-to-parameter (T2P) 比率扩展到工业级的约200。实验结果表明，与Muon相比，Muon+在训练和验证困惑度上都提供了持续的提升。代码已开源。

🔬 方法详解

问题定义：Muon优化器旨在通过梯度或动量正交化来改善大型语言模型的训练过程。然而，即使进行了正交化，训练过程仍然可能不稳定，导致收敛速度慢或最终性能不佳。现有的Muon优化器可能无法充分利用计算资源，尤其是在高token-to-parameter (T2P) 比率下，训练效率有待提高。

核心思路：Muon+的核心思路是在Muon优化器的正交化步骤之后，添加一个额外的归一化步骤。这个额外的归一化步骤旨在进一步稳定梯度，防止梯度爆炸或消失，从而改善训练的收敛性和最终性能。通过更有效地利用计算资源，尤其是在高T2P比率下，提高训练效率。

技术框架：Muon+的整体框架与Muon类似，主要包括梯度计算、动量更新、梯度正交化等步骤。关键区别在于，在梯度正交化之后，Muon+增加了一个额外的归一化步骤。这个归一化步骤对正交化后的梯度进行处理，使其具有更稳定的尺度。然后，使用归一化后的梯度更新模型参数。

关键创新：Muon+的最重要的技术创新点在于引入了额外的归一化步骤。与Muon相比，Muon+在正交化后对梯度进行了进一步的约束，从而提高了训练的稳定性。这种额外的归一化步骤可以被视为一种正则化技术，有助于防止过拟合，并提高模型的泛化能力。

关键设计：Muon+的关键设计在于选择合适的归一化方法。论文中可能采用了如Layer Normalization或Weight Normalization等方法。具体的归一化参数（例如，均值和方差的计算方式）以及归一化层的放置位置，都是需要仔细调整的关键技术细节。此外，与其他超参数（如学习率、动量系数等）的配合也至关重要。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Muon+在不同规模的GPT和LLaMA风格模型上，均能持续提升训练和验证困惑度。在高T2P比率（约200）下，Muon+的优势更加明显，表明其在高计算效率场景下的有效性。与Muon相比，Muon+能够更快地收敛，并达到更低的困惑度，从而证明了其优越的性能。

🎯 应用场景

Muon+优化器可广泛应用于各种大型语言模型的预训练任务，尤其是在计算资源有限或需要高T2P比率的场景下。它可以提高模型的训练效率和最终性能，从而降低训练成本，并提升模型的生成质量和泛化能力。该方法在自然语言处理、机器翻译、文本生成等领域具有广泛的应用前景。

📄 摘要（原文）

The Muon optimizer has demonstrated promising performance in pre-training large language models through gradient (or momentum) orthogonalization. In this work, we propose a simple yet effective enhancement to Muon, namely Muon+, which introduces an additional normalization step after orthogonalization. We demonstrate the effectiveness of Muon+ through extensive pre-training experiments across a wide range of model scales and architectures. Our evaluation includes GPT-style models ranging from 130M to 774M parameters and LLaMA-style models ranging from 60M to 1B parameters. We comprehensively evaluate the effectiveness of Muon+ in the compute-optimal training regime and further extend the token-to-parameter (T2P) ratio to an industrial level of $\approx 200$. Experimental results show that Muon+ provides a consistent boost on training and validation perplexity over Muon. We provide our code here: https://github.com/K1seki221/MuonPlus.