Muon: An optimizer for hidden layers in neural networks

引言

在深度学习领域，优化算法对模型训练效率和性能起着关键作用。从经典的随机梯度下降 (SGD) 及其动量法，到自适应优化方法 Adam/AdamW 等，一系列优化器大大加速了神经网络的收敛。然而，随着模型规模和数据量的爆炸式增长，研究者们不断探索更高效的新型优化策略，以进一步缩短训练时间和提高性能
Muon (Momentum Orthogonalized by Newton-Schulz) 便是近期引起广泛关注的创新成果之一，代表一种利用牛顿–舒尔茨迭代对梯度动量进行正交化处理的优化算法。Muon优化器专门针对神经网络隐藏层的二维权重参数设计，其核心思想是在每次参数更新前，对梯度更新矩阵进行近似正交化处理，从而改善优化动力学。凭借这一独特机制，Muon 在实际训练中取得了惊人的加速效果：例如，它被用于 NanoGPT 和 CIFAR-10 等任务的训练，并刷新了这些任务的速度纪录
一句话总结 Muon 的原理：利用牛顿–舒尔茨迭代高效地逼近梯度矩阵的正交化形式，并将其用于更新神经网络中的二维权重参数，从而加速训练收敛

Muon 相关链接

代码实现：https://github.com/KellerJordan/Muon
博客介绍：
- https://kellerjordan.github.io/posts/muon/
- https://jeremybernste.in/writing/deriving-muon

背景

Muon 优化器的设计初衷是针对神经网络中的线性变换层（即权重为矩阵的层，例如全连接层、卷积层等）提供一种专门优化方案。传统优化器（如 SGD 或 Adam）对所有参数一视同仁地应用统一的更新规则，而 Muon 采取了一种模块化的视角：根据层类型的不同采用不同的优化策略，以充分利用每类层结构的特性。这种“为不同网络模块定制优化器”的理念正是深度学习优化最新趋势的一部分。Muon 优化器聚焦于二维权重矩阵参数（例如全连接层的权重矩阵或卷积核张量展平后的矩阵），通过对这些矩阵形式的梯度更新进行特殊处理，以提升优化效果
在这里插入图片描述

Muon 优化器具体实现

Muon 是针对神经网络隐藏层二维参数的优化器。其算法定义如下：

梯度计算与动量累积：在每次参数更新迭代 $t$ ，首先计算当前参数 $\theta_{t-1}$ 的梯度 $G_t = \nabla_{\theta}\mathcal{L}t(\theta{t-1})$ 。然后将其累积到动量矩阵 $B_t$ 中，即 $B_t \leftarrow \mu B_{t-1} + G_t$ ，其中 $\mu$ 是动量超参数（类似于SGD动量法中的动量因子）。初始化时 $B_0 = 0$ 。
以下伪代码（Algorithm 2 Muon）展示了Muon优化算法的主要循环步骤，其中 $B_t$ 表示动量梯度累积， $\mu$ 是动量系数， $\eta$ 是学习率。在每次迭代中，先计算梯度 $G_t$ 并累积到 $B_t$ 中，然后对 $B_t$ 执行Newton-Schulz迭代（即步骤5的 $O_t \leftarrow \mathrm{NewtonSchulz5}(B_t)$ ），最后用得到的正交化更新 $O_t$ 对参数进行梯度下降更新（步骤6）

其中，‘NewtonSchulz5’被定义为以下 Newton-Schulz 矩阵迭代

# Pytorch 代码
def newtonschulz5(G, steps=5, eps=1e-7):assert G.ndim == 2a, b, c = (3.4445, -4.7750, 2.0315)X = G.bfloat16()X /= (X.norm() + eps)if G.size(0) > G.size(1):X = X.Tfor _ in range(steps):A = X @ X.TB = b * A + c * A @ AX = a * X + B @ Xif G.size(0) > G.size(1):X = X.Treturn X

牛顿–舒尔茨迭代介绍

这一步是 Muon 的核心创新，即对累积梯度 $B_t$ 应用 Newton-Schulz 迭代以计算其近似正交化矩阵 $O_t$ 。Newton-Schulz 迭代是一种用于近似计算矩阵函数（例如矩阵平方根或逆）的算法。这里它被用于逼近 $B_t$ 的逆平方根（或相关变换），以达到正交化的效果。Muon使用固定步数的Newton-Schulz迭代来处理 $B_t$ ，典型地迭代 5 次就能取得足够精度。在实现中，作者选择在半精度（bfloat16）下运行该迭代以提高效率，并发现这种迭代在低精度下仍然数值稳定。相比直接进行奇异值分解(SVD)，Newton-Schulz迭代可以充分利用GPU的并行计算能力，大幅降低计算开销。例如，作者最初实现的Newton-Schulz迭代每次更新对形状为 $n\times m$ 的参数需约 $6n m^2$ 次浮点运算，随后经过改进将复杂度降低到 $4n m^2 + 2 m^3$ （在实际中 $\le n$ ，因此复杂度进一步受限于参数矩阵较小的一维），显著提升了运算效率。通过这样处理，Muon在每步更新中增加的计算和内存开销相对可控，却换来了更有效的梯度方向调整。

Muon 原理分析

Muon 在每次迭代时首先使用常规方法计算梯度并累积动量，然后对动量梯度矩阵进行**正交化(orthogonalization)**处理，再用于更新权重。所谓对梯度矩阵正交化，是指找到一个与原梯度矩阵最接近的“半正交矩阵”作为更新，其中“半正交”意味着矩阵的行向量或列向量是正交的（满足 $O^T O = I$ 或 $O O^T = I$ ）。等价地，这相当于将梯度矩阵进行奇异值分解 $G=USV^T$ ，然后用 $UV^T$ （即将奇异值矩阵 $S$ 置为单位矩阵）替换原梯度——这一结果被称为“正交化梯度”。Muon 的更新步骤正是用这样的正交化梯度替代原始梯度动量进行参数调整。

那么，引入正交化的直观意义何在？ 简单来说，它可以丰富更新的方向性。研究者通过对 Transformer 等网络的梯度观察发现，传统优化器产生的梯度更新矩阵往往条件数很高，接近于秩亏（也就是被少数几个主导方向所支配）。换言之，许多神经元对应的更新方向非常相似，梯度矩阵接近低秩。这意味着一些“罕见方向”（对应较小奇异值的方向）在更新中作用很弱，但这些方向可能对进一步降低损失仍然重要。通过正交化处理，Muon 实质上放大了这些罕见方向在更新中的作用，因为正交化将所有奇异值归一，无论原本大小，从而赋予每个独立方向以均等的尺度权重。这种调整有助于避免训练过程中某些重要方向被忽略，提高了搜索参数空间的效率。从经验上看，正是这种对梯度更新方向分量的重新均衡，使得 Muon 相较于传统优化方法展现出更快的收敛和更高的效率。

Muon 实践注意事项

在具体实现中，还有一些实用技巧和细节：首先，Muon 目前仅针对二维权重张量（包括将卷积核展平成矩阵的情况）应用。对于标量或向量参数（例如偏置项、LayerNorm 中的缩放参数等），以及输入层和输出层的权重，作者建议仍采用常规优化器 (如 AdamW) 进行训练。这部分原因在于：嵌入层（输入层）参数的优化动态确实有别于其他层，需要特殊对待；输出层是否需要区别对待则是基于经验的优化选择。实际实验表明，如果对 Transformer 模型的词嵌入层和最后分类头仍使用 AdamW 优化，而对中间的线性层使用 Muon 优化，能够取得最佳性能。其次，在动量的具体形式上，Muon 采用Nesterov动量略优于传统动量，因此作者在公开实现中将 Nesterov 作为默认方案。最后，对于 Transformer 中特有的多头注意力，研究者发现将查询 Q、键 K、值 V 各自的权重矩阵分别应用 Muon 优化效果更佳，而不是将它们合并为一个大矩阵一起正交化。这一拆分处理使每个子矩阵的正交化更精确，有助于训练稳定。上述这些实现层面的经验总结，体现了 Muon 在不同网络结构中的适配性和灵活性。

实验

Muon 取得以下优化结果
- 将 CIFAR-10 训练到 94% 准确率的速度记录从 3.3 A100-秒提高到 2.6 A100-秒
- 在 FineWeb 上将训练到 3.28 验证损失（一项被称为 NanoGPT 速度竞赛的竞争任务）的速度记录提高了 1.35 倍。
- 在扩展到 7.74 亿和 15 亿参数时，持续显示训练速度的改进。
- 在 10 个 8xH100 小时内，将一个 15 亿参数的 Transformer 训练到 GPT-2 XL 在 HellaSwag 上的性能水平。使用 AdamW 实现相同结果需要 13.3 小时。
下图展示了在NanoGPT语言模型任务中，不同优化器随时间推进的验证集损失下降曲线（横轴为在8×H100 GPU上消耗的实际时间，纵轴为验证损失）比较。可以看到，Muon（紫色曲线）相对于标准的Adam（蓝色）收敛更快。与此同时，图中标注的每步迭代耗时显示，Muon每步约142毫秒，与Adam的139毫秒几乎相当。这意味着Muon在提供更快收敛的同时，并未引入明显的单步时间开销

Muon 后续工作进展

kimi 团队对 Muon 做了进一步优化：https://ar5iv.labs.arxiv.org/html/2502.16982
- 由于Muon需要对每个二维权重矩阵执行Newton-Schulz迭代，如果在分布式环境下一个矩阵分散在不同设备上，如何高效地完成迭代成为一个问题。kimi 这里开发了内存优化且通信高效的分布式Muon实现。
- 最初的Muon实验主要集中在十亿参数量级以内的模型上（如1.5B参数的Transformer）。一个悬而未决的问题是：Muon能否在数百亿乃至上万亿参数、万亿级别token数据的大规模训练中保持优势并稳定运行？kimi 这里的结果展示了通过在 Muon 中引入权重衰减(weight decay)以及按参数尺度调整更新幅度等改进，可以使 Muon 在无需特殊超参数调优的情况下，直接应用于超大规模模型的训练

总结

Muon 激发了人们对于优化器创新的热情——长久以来，Adam/AdamW 几乎是大多数任务的不二之选，而 Muon 证明了通过深入挖掘梯度结构特性，我们仍能找到显著优于现有方法的新途径。
当然，Muon 目前也并非完美无缺。它对模型结构有特定假设（主要作用于稠密的线性层），在一些场景下的通用性还有待进一步验证。此外，对于更大规模模型、更长训练阶段（如微调、强化学习等）的效果也需要更多实验来检验