Muon优化器教程：原理与替换AdamW实践

一、教程目标：你将完成什么

在训练大规模神经网络时，选择哪种优化器直接影响收敛速度、稳定性和算力成本。AdamW 长期是主流默认选择。近两年，以 Muon 为代表的一类矩阵正交化优化器而受到关注。有研究在 160 亿参数的混合专家（MoE）模型 上用 Muon 完成训练，实验表明：在计算最优配置下，Muon 只需约 52% 的训练计算量，就能达到 AdamW 的同等表现。

本教程以“引入 Muon、替换 AdamW”为主线，帮助你完成以下任务：

1. 理解原理：弄清 Muon 优化思想的核心，即为什么对矩阵型参数“先做正交化，再做更新”更有效；
2. 实现代码：写出带 Newton-Schulz 正交化算子的优化器，并做好参数分组与学习率换算；
3. 设计实验：设计小规模对照实验，评估 Muon 在自己的模型与数据上是否有效；
4. 迁移扩缩：把 Muon 接入已有的 AdamW 训练流程，并迁移到更大规模；
5. 验证排查：掌握收敛正确性验证、性能监控、常见问题排查与成本优化方法。

适合读者

• 训练框架与平台开发者：需要实现或集成新的优化器；
• 算法工程师：负责大模型预训练或继续训练，希望降低达到目标精度所需的算力；
• 技术负责人：需要评估 Muon 的收益、风险与工程成本，再决定是否推广。

教程约定：本文以通用的张量并行、数据并行训练框架为背景。实现 Muon 时，推荐使用深度学习编译器提供的矩阵乘法融合能力。这类编译器能自动把多个矩阵乘与逐元素运算融合为少量高性能算子。代码为通用伪代码，不绑定任何具体框架接口。

二、从 AdamW 到 Muon：Muon 体现了什么新优化思想

2.1 AdamW 的工作方式

AdamW 对每个参数独立维护一阶动量（梯度滑动平均）和二阶动量（梯度平方滑动平均）。更新时，它用二阶动量对每个元素做归一化，使单个参数的有效步长大致稳定在“学习率 × 符号方向”的量级。权重衰减则与梯度更新解耦，直接按一定比例收缩参数。

这种做法实现简单、对噪声梯度鲁棒。但深入看矩阵型参数时会发现一个问题：Adam 调整的是“每个元素”，而神经网络真正关心的是“整张权重矩阵作用到激活上所产生的整体变化”。

设隐藏层权重为 $W\in\mathbb{R}^{n\times m}$，其更新为 $\Delta W$，输入激活的均方幅度为 $R$，则该层输出变化的幅度可近似估计为：

$|\Delta W x|<em>2 \approx R\sqrt{\sum</em>{i,j}(\Delta W_{ij})^2}=R|\Delta W|_F$

也就是说，这一层对训练的影响由更新矩阵的 Frobenius 范数决定，而 Adam 把这个范数“交给”了梯度的统计特性。不同形状、不同层、不同训练阶段的梯度谱差异很大，更新范数也随之漂移，需要靠单独调整各层学习率来弥补。权重矩阵的形状又随模型宽度和深度变化，这会让超参数在模型规模变化时的迁移变得困难。

2.2 Muon 的核心思想：把更新方向“正交化”

Muon 的全称是 MomentUm Orthogonalized by Newton-Schulz。步骤分为两步：

1. 维护动量：与普通动量法一样，维护梯度矩阵的滑动平均 $M$；
2. 正交化后更新：对动量矩阵 $M$ 做奇异值分解 $M=USV^\top$（$U$、$V$ 为正交矩阵，$S$ 为对角奇异值矩阵），取其正交部分 $UV^\top$ 作为更新方向：
   $$W_{t+1}=W_t-\eta\,UV^\top$$

正交化有两层直观含义：

• 各方向同等推进：在 $M$ 的每个奇异方向上，更新步长都被拉齐到大致相同的大小（奇异值 $s_i$ 被替换为 1），既不会被梯度中少数“大方向”主导，也不会让“小方向”进展过慢。
• 只定方向、不定幅度：更新只携带方向信息，幅度由学习率 $\eta$ 统一决定。由此，$n\times m$ 权重矩阵的更新 Frobenius 范数恒为 $\sqrt{\min(n,m)}$。结合 2.1 节的近似，只要学习率随矩阵短边按 $1/\sqrt{\min(n,m)}$ 缩放，不同形状的矩阵在同一学习率下就能产生量级一致的层输出变化。

这与主流的“逐元素自适应”思路不同：Muon 把“更新幅度”的调节权从优化器内部转移到学习率上。这样大幅降低了对梯度统计量和精细调参的依赖，超参数也更容易在不同模型规模之间迁移。

补充说明：以 Shampoo 为代表的方法也带有“对梯度结构做预处理”的思想，Muon 可以理解为用更低的开销实现了类似的正交化效果。

2.3 为什么用 Newton-Schulz 迭代近似

直接对每张权重矩阵做 SVD，开销大，在 GPU 上的吞吐远低于普通矩阵乘法，低精度下的稳定性也不理想。

Newton-Schulz 迭代只做矩阵乘法和逐元素运算，属于“后端友好”的计算，很适合张量核心。它还能借助编译器融合为少量高效算子。常用的五次迭代形式为：

$X_{k+1}=a_1X_k+a_3(X_kX_k^\top)X_k+a_5(X_kX_k^\top)^2X_k$

其中 $a_1=3.4445,\ a_3=-4.7750,\ a_5=2.0315$。迭代前先用 Frobenius 范数归一化输入，使所有奇异值不超过 1。这几组系数经过调优，能在约 5 步内把非零奇异值映射到 0.7～1.3 的区间。因此 5 步迭代的精度通常已足够训练使用，实际训练中极少能看出它与精确正交化的区别。

对形状为 $n\times m$ 的矩阵，若 $n>m$，先转置成“瘦高”矩阵再做迭代，使主计算规模降为 $m^2\times n$。每步的算术运算量约为 $6m^2n$，5 步约 $30m^2n$，相当于一次 $m\times m$ 规模矩阵乘法的几十倍。对 Transformer 中“宽而矮”的典型权重矩阵（$m\ll n$），这部分开销通常只占该层前向加反向运算量的几个百分点。

2.4 收益、成本与适用边界

已知收益（来自上述公开实验）：

• 计算最优配置下约 52% 的训练计算量即可达到同等表现，等效训练效率约提升一倍；
• 该团队训练的 160 亿参数 MoE 模型中，激活参数量约为总参数量的 1/5 量级；
• 在该配置下，Muon 的峰值显存低于 AdamW，通信开销略高。

需要客观看待的地方：

1. 不能外推为“端到端快一倍”。约 2 倍的加速来自两个来源叠加：一是每个 Muon 参数所需的训练计算量减少；二是达到同等质量所需的迭代次数减少。前者只覆盖矩阵型参数，后者依赖模型结构和数据。端到端收益要在自己的模型上实测。
2. 学习率更敏感。正交化后更新幅度由学习率单独控制，对学习率越界的容错更小，需要重新调参并观察训练曲线。
3. 可见性变化。缺少逐元素二阶信息，高方差参数缺少“刹车”，依靠全局监控来发现异常。
4. 推理部署不变。Muon 只改变训练方式，模型的网络结构和推理接口不变。已有 checkpoint 评估新优化器时，可直接复用现有推理链路。

三、适用场景与前置准备

3.1 适用场景

场景	适用性	说明
从零预训练的稠密 Transformer	适合	已有公开实验覆盖同类结构，收益最明确
混合专家模型预训练	适合	共享专家、路由等矩阵参数均可使用
继续预训练、中期训练	可尝试	需从 AdamW 权重平滑切换，并建立对照
全参数微调	可评估	矩阵参数规模有限时，收益不确定
小模型、短训练	需实测	正交化开销占比较高，可能得不偿失

优先使用 Muon 的参数：隐藏层中二维的线性层权重，包括注意力中的 $Q/K/V/O$ 投影、前馈网络的上投影与下投影。

建议继续使用 AdamW 的参数：

• 嵌入层与输出层：每行对应一个 token 的更新，正交化会让不同 token 的嵌入相互干扰；这两类参数的梯度稀疏度高、方差大，对逐元素自适应更敏感；
• 一维参数：归一化层的缩放系数、偏置等。正交矩阵对一维向量没有定义，Muon 在理论上不适用。

这些参数在 Transformer 中通常只占很小的比例，用 AdamW 处理不会对整体效率造成明显拖累。

3.2 前置准备

1. 环境：支持低精度矩阵运算的加速卡、可靠的集合通信环境、主流深度学习框架及深度学习编译器；
2. 代码：可修改优化器与训练循环的训练框架，支持检查点保存、恢复与梯度累积；
3. 对照基线：一套能跑通、超参数合理的 AdamW 训练配置，作为比较基准；
4. 实验资源：一批固定、可复现的训练数据，以及一台与预想生产规模“同代、同互联”的机器。用它跑小规模对照，结论才能更好地外推。

如果训练任务运行在百度智能云等云平台：计算资源可选用 GPU 类云服务器或托管训练平台；检查点可存放于对象存储；集群互联带宽、GPU 利用率、显存占用可借助监控告警能力持续观察。本文的实现与这些环境无关，本地集群或混合部署同样适用。

四、实施步骤

整体流程如下：

搭建对照环境 → 实现正交化算子 → 实现优化器更新 → 划分参数组并换算学习率 → 小规模对照实验 → 接入已有训练流程并扩规模 → 验证与排查

步骤 1：搭建可复现的对照环境

做什么：固定数据采样顺序、随机种子和框架版本，用 AdamW 跑一组小规模基线训练。

为什么做：Muon 的收益来自“更少的计算达到同等表现”。没有可靠的对照，就无法区分收益来自优化器本身，还是来自数据顺序等随机因素。

注意：基线超参数应调至合理水平，可以借鉴所选用模型结构公开的预训练配置作为起点。否则对照不公平，容易高估或低估 Muon。

步骤 2：实现 Newton-Schulz 正交化算子

做什么：实现“归一化、按需转置、五次迭代”的正交化函数，并做数值正确性校验。

为什么做：正交化是 Muon 的核心，其精度和开销直接决定训练是否稳定、是否值得。

实现要点：

• 归一化用全局 Frobenius 范数，避免奇异值大于 1 导致迭代溢出；
• 低精度计算前先上浮到 32 位浮点计算范数，避免下溢；
• 长宽比适中的矩阵收益最大，极端瘦长的矩阵可尝试分块处理（见常见问题第 6 条）；
• 可通过深度学习编译器导出计算图，使多个矩阵乘与逐元素运算融合为少量算子。

通用伪代码：

A1, A3, A5 = 3.4445, -4.7750, 2.0315          # 五次迭代系数
def zeropower_via_newtonschulz(G, steps=5, eps=1e-7):
    # G: 二维梯度或动量张量，支持在分布式场景下跨进程通信
    X = G.astype(float32)
    X = X / (X.norm() + eps)                   # Frobenius 归一化：奇异值 <= 1
    trans = X.shape[0] > X.shape[1]
    if trans:                                  # 保证迭代中 X @ X^T 的规模较小
        X = X.transpose(-2, -1)
    for _ in range(steps):
        A = X @ X.transpose(-2, -1)            # 第 1 次矩阵乘
        B = A @ A                              # 第 2 次
        X = A1 * X + (A3 * A + A5 * B) @ X     # 第 3 次（编译器可融合）
    return X.transpose(-2, -1) if trans else X

正确性校验：

• 在 CPU 或加速卡上构造已知矩阵 $M=U\,\mathrm{diag}(s)\,V^\top$，比较输出 $X$ 与 $UV^\top$；
• 计算 $X$ 的奇异值，应全部落在约 0.7～1.3 区间；
• 再验证一次性质：$|(X^\top X-I)\,x|$ 对随机向量 $x$ 应很小。

注意：迭代次数增加会提高精度，但线性增加开销。5 步是公开实现中的常用值，可在小规模上对比 3、5、8 步的训练曲线后再决定。

步骤 3：实现带权重衰减的 Muon 更新

做什么：在优化器的更新函数中为每个参数维护动量缓冲区，并依次完成动量更新、（可选的）Nesterov 修正、正交化、参数更新与解耦权重衰减。

为什么做：

• 直接对瞬时梯度做正交化，会丢掉动量的降噪作用；先累积动量再正交化，才能兼顾稳定与方向质量。
• 正交化后的更新幅度与梯度无关，AdamW 式的梯度裁剪对 Muon 参数基本不起作用，权重衰减强度必须按 Muon 的更新幅度重新评估。

通用伪代码：

# 伪代码示意：以“参数状态”为单位，逐个参数执行
buf = zeros_like(W)                       # 动量缓冲区（每个参数一份）
buf = mu * buf + G                        # 动量累积，mu 常取 0.95
if nesterov:                              # 可选：向前看一步
    buf = G + mu * buf
U = zeropower_via_newtonschulz(buf)       # 正交化
W = W - lr * U                            # 学习率已按矩阵形状换算，见步骤 4
if wd > 0:
    W = W - lr * wd * W                   # 解耦权重衰减

注意：

• 动量系数常取 0.95，高于 AdamW 常用的 0.9，以更好地平滑梯度噪声；
• 是否使用 Nesterov 属于实现细节，应通过对照实验确定；
• 权重衰减需按 Muon 的更新幅度重新设定，不能直接沿用 AdamW 的数值；
• 若优化器在加速卡上以低精度反向传播（例如 16 位浮点）维护参数副本，应确认正交化使用足够精度的副本，避免长时间累积的数值偏差。

步骤 4：划分参数组并换算学习率

做什么：按参数名或模块类型把参数分为两个优化器组，并为 Muon 组设置按形状换算的学习率。

为什么做：如果对所有参数统一使用 Muon，嵌入层与输出层会出现额外扰动；如果不做学习率换算，形状差异大的矩阵会得到强度差异悬殊的更新。

通用伪代码：

def param_group_of(name, shape):
    is_mat = len(shape) == 2
    is_embed_or_head = any(k in name for k in ("embed", "head", "lm_head", "output"))
    if is_mat and not is_embed_or_head:
        scale = 0.2 * (max(shape) ** 0.5)     # 随长边按平方根放大，系数取常用经验值
        return {"base_lr": 0.02, "lr": 0.02 * scale, "wd": 1.0}
    return {"base_lr": 4e-4, "lr": 4e-4, "wd": 0.1}

如何理解换算式：若把 Muon 参数的基准学习率（短边为 1 时的值）记为 $\eta_0$，实际学习率取 $\eta_0\cdot c\cdot\sqrt{\max(n,m)}$，则 $n\times m$ 矩阵的更新 Frobenius 范数为 $\eta_0 c\sqrt{\min\cdot\max}=\eta_0 c\sqrt{nm}$，恰好与参数量的平方根成正比，而不是与矩阵形状的具体组合有关。这使同一个学习率可以覆盖同一层中形状各异的矩阵。

注意：

• 代码中的数值仅用于演示“换算方式”，实际取值必须以对照实验为准；
• 常见的做法是先在小模型上扫 Muon 组的学习率，再固定系数 $c$，放大到目标模型。

步骤 5：设计并执行小规模对照实验

做什么：在固定算力预算下，让 AdamW 与 Muon 从相同初始化出发，训练同一份数据。绘制“训练损失—计算量”曲线，比较达到同等损失时双方消耗的浮点运算量之比。随后再在相同训练计算量下比较最终损失。

为什么做：

• 两种优化器的单步开销不同，只有统一横轴为计算量，比较才公平；
• 公开实验结论出自特定的模型规模与数据配比，需要验证它在自己场景中的表现；
• “同等质量”需要明确指标：是同等训练损失，还是同等下游评测分数，结论可能不同。

建议的实验规模：

1. 优先使用尺寸接近目标模型的“代理模型”，例如深度相近而宽度减半，以保持宽深比接近；
2. 计算量不足时退而用较小模型，并记录代理模型与目标模型的差异；
3. 每个配置至少设置两个随机种子，避免单点结论。

注意：计算量应把正交化的矩阵乘和集合通信折算进去，按“模型浮点运算量 + 优化器开销”统计，而不是只数迭代次数。

步骤 6：接入已有训练流程，并逐步扩大规模

做什么：在训练框架中以“参数组 + 对应优化器”的方式接入 Muon；先单机跑通，再过渡到张量并行、数据并行及二者组合。

为什么做：已有流程中包含学习率调度、梯度裁剪、检查点、混合精度、监控等大量基础设施。用适配层引入 Muon，改动最小，也最不容易出错。

为什么需要分布式适配：

• 数据并行：各进程上的参数一致，正交化本地完成即可。分布式优化器会改变参数副本的分布方式，需要让正交化在“完整参数”而非分片上执行。
• 张量并行：权重矩阵沿某个维度切分到多个进程，每个进程只持有分片，无法在本地得到完整动量矩阵。因此需要先聚合出完整动量、再执行正交化、再切回各自的分片。聚合可以在迭代期间异步预取，以隐藏通信开销。
• 检查点：Muon 的动量缓冲区需要随检查点一并保存和恢复，否则续训时动量清零，会出现短暂的训练抖动。

注意：若采用“复制参数、按批次分片、反向后再归约”的数据并行实现，需要在归约后、正交化前对动量做同步，避免各进程的正交化结果出现偏差。

五、关键配置项说明

配置项	作用	常见取值	调整风险
Muon 动量系数	控制动量平滑程度	约 0.95	偏大易导致对梯度变化反应迟钝
Nesterov 修正	用“预测位置”的动量做更新	开/关，需实验确定	可能放大学习率过高时的不稳定
正交化迭代步数	精度与开销的权衡	5	过少则近似变差，过多则开销上升
Muon 参数基准学习率	决定整体更新幅度	需通过扫描确定，常见量级为 $10^{-2}$	偏大易引起损失尖峰甚至发散
形状换算系数	统一不同形状矩阵的更新强度	常用经验值 0.2 量级	偏离合适区间会改变各层之间的平衡
Muon 组权重衰减	约束参数幅度	需单独调优	沿用过 AdamW 值可能导致过度收缩
AdamW 组学习率	嵌入、输出层及一维参数更新	通常显著低于 Muon 组	偏大易导致输出层梯度异常
正交化计算精度	数值稳定与吞吐	低精度存储、高精度计算	低精度存储可能在长训练中引入偏差
梯度裁剪阈值	限制全局梯度范数	沿用基线	对 Muon 参数几乎无约束作用

关于低精度存储的说明：公开的分布式实现中，动量缓冲区与正交化迭代通常以低精度存储。这是优化训练吞吐与显存的常见做法，但需要在长训练中持续观察损失与评测指标是否出现缓慢劣化。

六、结果验证：如何判断接入成功

6.1 算子级验证

• 对随机矩阵，输出奇异值应集中在 1 附近；
• 与“完整 SVD 取正交部分”的参考结果对比，逐元素相对误差应足够小；
• 同一张量在单卡与多卡聚合路径下结果一致（考虑低精度误差）。

6.2 训练级验证

• 收敛曲线：与 AdamW 基线对照，在同等计算量下观察损失差距是否符合预期；
• 更新幅度：统计 Muon 参数每步更新的 Frobenius 范数，应接近“学习率 × √min(n, m)”；若偏离较大，说明实现有误；
• 稳定性：监控损失尖峰频率、梯度范数分布和激活数值范围。

6.3 性能级验证

• 单步耗时：记录加入 Muon 前后的单步耗时，分解正交化、通信和参数更新各部分占比；
• 资源占用：对比峰值显存，以及加速卡矩阵运算单元的利用率；
• 扩展效率：改变张量并行度时，观察每卡吞吐的变化。

6.4 质量级验证

• 检查点转换后，推理结果应与训练中的前向计算一致；
• 在留出集上比较困惑度、下游任务指标等，确认达到预期。

七、常见问题与排查

现象	常见原因	排查思路
训练早期出现损失尖峰或发散	Muon 学习率偏大；从 AdamW 中间权重切换时梯度尺度突变；Nesterov 放大波动	降低 Muon 学习率；切换前后加入短暂回退；关闭 Nesterov 后对比
吞吐下降超出预期	正交化未与张量并行/序列并行结合；未借助编译器融合；小矩阵上算子启动开销大	分解单步耗时定位瓶颈；按形状批量执行正交化；确认编译优化已生效
张量并行下损失异常	分片上直接做正交化，或聚合后未正确切回	校验各进程更新后的参数是否一致；用算子级验证交叉检查
显存峰值上升	临时张量和低精度副本未复用	复用缓冲区；逐参数处理；检查缓存分配策略
续训后损失短暂上升	动量未随检查点保存	确认优化器状态完整写入和恢复
梯度裁剪不起作用	正交化后更新幅度与梯度无关	这是预期现象，应把裁剪理解为对 AdamW 组的保护
某些层（如最后一层）行为异常	权重极度瘦长，正交化近似精度不足	调整并行切分方式；比较不同迭代步数的效果
小模型上未见收益	正交化开销占比高；对照基线本身很强	计算损失曲线交点；评估收益能否覆盖额外实现成本

八、优化建议

• 性能：让正交化与反向计算重叠，聚合通信与计算重叠；对形状相近的矩阵批量执行正交化；优先使用融合了矩阵乘与逐元素运算的实现。
• 稳定性：从较小学习率起步并延长预热；保留梯度范数监控；把损失尖峰的自动回退机制覆盖到 Muon 组。
• 可维护性：用适配层隔离 Muon 实现，训练主循环只依赖“参数组 + 优化器”接口；为算子补充单元测试与数值回归测试；在配置文件中集中管理所有超参数并记录版本。
• 成本：在计算最优点附近停止训练，避免额外迭代；利用 Muon 峰值显存较低的特点，评估更小显存规格的卡或更高并行度；在百度智能云上训练时，可结合监控指标做弹性伸缩，并在夜间低谷时段跑对照实验；用日志服务统一采集训练指标，便于复盘。
• 安全：检查点可能含模型权重等敏感信息。保存到对象存储时，应开启身份认证与访问控制、服务端加密和访问审计；跨环境传输时校验完整性。

九、总结与后续方向

Muon 体现的是“把优化器的职责聚焦于方向，把幅度的调节权交给学习率”这一新思路。它以较低开销近似对矩阵参数做正交化，使更新在各个奇异方向上更均衡，对梯度统计的依赖也更小。公开实验表明，在 160 亿参数 MoE 模型上，Muon 可在计算最优配置下以约 52% 的计算量达到 AdamW 的同等表现，说明它在部分大模型场景中具备实用价值。

落地时可以记住三句话：

1. 先对照：用“损失—计算量”曲线确认收益；
2. 再分组：矩阵参数用 Muon，嵌入、输出层和一维参数用 AdamW；
3. 后扩规模：把分布式适配和通信重叠解决在扩大规模之前。

后续可继续关注的方向包括：正交化优化器与二阶方法思想的融合、面向 MoE 专家权重的专门适配、并行切分与正交化更紧密的结合，以及在大规模训练中更完整的稳定性经验。对多数团队而言，合理的路径是：先在内部小规模任务上验证，再逐步推广到更大规模的预训练。