Muon+: Towards Better Muon via One Additional Normalization Step

📄 arXiv: 2602.21545 📥 PDF

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

分类: cs.LG

发布日期: 2026-02-28

💡 一句话要点

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

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

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

📋 核心要点

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

📝 摘要（中文）

Muon优化器通过梯度（或动量）正交化，在大语言模型的预训练中表现出良好的性能。本文提出了一种简单而有效的Muon增强方法，称为Muon+，它在正交化之后引入了一个额外的归一化步骤。通过在各种模型规模和架构上的大量预训练实验，证明了Muon+的有效性。评估包括从1.3亿到7.74亿参数的GPT风格模型，以及从6000万到10亿参数的LLaMA风格模型。我们全面评估了Muon+在计算最优训练机制中的有效性，并将token-to-parameter (T2P) 比率扩展到工业级的约200。实验结果表明，Muon+在训练和验证困惑度方面均优于Muon。

🔬 方法详解

问题定义：Muon优化器旨在通过梯度正交化来改善大型语言模型的训练过程。然而，即使经过正交化，梯度可能仍然存在尺度不一致的问题，这会影响训练的稳定性和收敛速度。现有Muon优化器可能无法充分利用计算资源，尤其是在token-to-parameter (T2P) 比例较高的情况下，训练效果可能受限。

核心思路：Muon+的核心思路是在梯度正交化之后，增加一个额外的归一化步骤。这个归一化步骤旨在确保梯度在正交化后具有更一致的尺度，从而提高训练的稳定性和效率。通过对正交化后的梯度进行归一化，可以更好地控制梯度的大小，避免梯度爆炸或消失的问题。

技术框架：Muon+的整体框架与Muon相似，主要包括以下几个步骤：1. 计算梯度；2. 计算动量（可选）；3. 对梯度（或动量）进行正交化；4. 对正交化后的梯度进行归一化；5. 使用归一化后的梯度更新模型参数。Muon+的关键区别在于第4步，即额外的归一化步骤。

关键创新：Muon+最重要的技术创新点是在正交化之后增加了一个额外的归一化步骤。与Muon相比，Muon+通过归一化进一步优化了梯度流，使得训练过程更加稳定和高效。这种额外的归一化步骤可以被视为对Muon的补充，旨在解决正交化后梯度尺度不一致的问题。

关键设计：论文中没有明确指出归一化步骤的具体实现方式，但常见的归一化方法包括L2归一化、Layer Normalization等。具体选择哪种归一化方法可能需要根据具体的模型架构和数据集进行调整。此外，论文还强调了在高T2P比率下的有效性，这意味着Muon+可能更适合于大规模的预训练任务。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Muon+在各种模型规模（从1.3亿到10亿参数）和架构（GPT和LLaMA风格）上均优于Muon。尤其是在高token-to-parameter比率（约200）下，Muon+在训练和验证困惑度方面均取得了显著提升，证明了其在大规模预训练任务中的有效性。

🎯 应用场景

Muon+优化器可广泛应用于大型语言模型的预训练，尤其是在计算资源有限的情况下。通过提高训练效率和稳定性，Muon+可以帮助研究人员和工程师更快地训练出性能更优越的模型。此外，Muon+在高token-to-parameter比率下的有效性，使其特别适用于工业界的大规模预训练任务，有助于降低训练成本，加速模型迭代。

📄 摘要（原文）

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:this https URL.