在深度学习训练里Adam 长期以来是默认选项但它并不是所有参数的最佳更新方式。最近关于 Muon 优化器以及光谱分配Spectral Allocation的讨论把训练优化器的问题重新拉回到一个老问题当我们更新一个权重矩阵时更新量到底应该按哪些方向分配。这篇文章会讲清楚 Muon 为什么能在不少任务上优于 Adam也会给出 Muon 的完整实现路径和改进方法适合已经熟悉 PyTorch、想把自己训练流程从 AdamW 迁移到 Muon 的算法工程师和学生。读完以后你能独立完成一次最小规模的 Muon 实验也知道改哪些超参数、看哪些指标、踩到哪些坑。1. 光谱分配到底在解决什么从 Adam 的更新方式说起在讨论 Muon 之前有必要先理解 Adam 为什么“好用”又为什么在权重矩阵上不一定最优。只有把 Adam 的更新机制拆开才能理解光谱分配这个概念的真正含义。1.1 Adam 的逐坐标归一化优点与隐患Adam 的核心更新公式可以写成m_t beta1 * m_{t-1} (1 - beta1) * g_t v_t beta2 * v_{t-1} (1 - beta2) * g_t^2 theta_t theta_{t-1} - lr * m_t_hat / (sqrt(v_t_hat) eps)其中m_t是一阶动量v_t是二阶动量也就是梯度每个元素平方的滑动平均。Adam 的关键设计是对参数的每一个坐标做一次归一化。梯度大的坐标会被sqrt(v)压小梯度一直很小的坐标会被放大一些。这样每个坐标的更新幅度大致在同一数量级因此 Adam 对学习率不那么敏感也适合稀疏梯度。但“逐坐标归一化”也有隐患。它假设每个坐标是独立的信息通道可权重矩阵的坐标之间往往存在强烈相关性。一个矩阵梯度G的各个元素并不是独立变量它的真正结构体现在奇异值谱上。Adam 只看元素的二阶矩忽略了行与列之间的相关性于是容易被噪声大的坐标带偏。1.2 权重矩阵的更新也许应该看“方向”而不是“坐标”把权重矩阵看成m x n的线性变换。它的梯度G做奇异值分解后可以写成G U Σ V^T其中U和V是正交矩阵Σ diag(σ_1, σ_2, ..., σ_k)是奇异值。如果把梯度按奇异值方向拆开每个方向对最终更新量的贡献取决于σ_i的大小。SGD 的做法很直接梯度大的方向更新也大也就是“更新预算按奇异值比例分配”。而真实训练中权重矩阵的梯度谱往往长尾分布前面几个大奇异值方向信息量大后面的小奇异值方向更接近噪声。Adam 的做法是按元素归一化它会把小方差坐标对应的方向放大到和大方向同等的更新量级。这个特性在遇到大量噪声坐标时会浪费更新预算。1.3 光谱分配把更新量按奇异方向做预算分配光谱分配Spectral Allocation不是一个像 Adam 那样有严格公式的优化器我更愿意把它理解成一种设计优化器的视角对于权重矩阵参数更新预算应该按奇异值方向来分配而不是按坐标来分配。不同优化器在这条光谱上处于不同位置优化器分配方式特点SGD按奇异值大小成比例分配主方向更新大噪声方向更新小Adam按坐标二阶矩归一化每个坐标同量级容易放大噪声方向Muon正交化后均匀分配每个谱方向等量更新幅度收敛自然梯度类按曲率倒数分配理论上更准但计算成本高Muon 选的是“均匀分配”这条路径对梯度矩阵做正交化得到一个近似正交的更新方向再用学习率控制步长。这个选择看起来反直觉但在大模型预训练任务上表现得很好这也是它值得单独分析的原因。2. Muon 优化器原理拆解Muon 在公开的论文和技术报告中被描述为“适用于神经网络隐藏层参数的优化器”。它在概念上并不复杂维护动量对动量矩阵做正交化然后沿正交化后的方向更新参数。2.1 核心对梯度做正交化而不是归一化Adam 对梯度做的是元素级归一化Muon 对梯度做的是矩阵级正交化。所谓正交化就是给定矩阵A找出一个“最接近A且各列正交”的矩阵O。从数学上看A的极分解可以写成A O P其中O是有正交列的方向矩阵P是半正定矩阵。Muon 想要的方向是O。对A做奇异值分解后O U V^T。也就是说Muon 把梯度矩阵的幅度Σ去掉了只留下方向信息。这样更新的范数稳定不会因为某个 batch 的梯度范数突然变大而剧烈抖动。2.2 Newton-Schulz 迭代不求 SVD也能逼近正交化直接对每个梯度做 SVD 显然不现实尤其在分布式训练中。Muon 使用 Newton-Schulz 迭代来逼近A (A^T A)^{-1/2}这个“正交因子”。它只需要矩阵乘法和转置非常适合 GPU 上的张量运算。一个常用的迭代形式是X A / ||A||_2 每次迭代: X a * X b * (X X.T) X c * (X X.T) (X X.T) X这其实是多项式近似。不同实现会选择不同的迭代次数和系数比较常见的是 5 次迭代。迭代次数太少正交化不充分迭代次数太多收益递减。在实际实现里还会处理行数小于列数的转置情况避免计算形状不匹配的问题。2.3 Muon 更新步骤的最小公式表达把整个流程写下来Muon 每个 step 做的事是G 当前梯度 M momentum * M (1 - momentum) * G O NewtonSchulz(M) # 正交化 W W - lr * O # 更新权重 W W * (1 - lr * weight_decay) # 解耦权重衰减可选的 AdamW 式写法这里有几个细节值得注意。第一正交化作用在动量M上不是直接作用在原始梯度G上。第二weight_decay是解耦的不进入动量也不参与正交化和 AdamW 的“解耦权重衰减”一致。第三正交化后的矩阵范数量级接近 1所以 Muon 的lr和 Adam 的lr不在同一尺度不能直接沿用。2.4 与 Adam、AdamW、SGD 的差异对比AdamW 相对 Adam 的改进是把权重衰减从梯度里拆出来单独乘以学习率。理解这一点后再看 Muon会发现 Muon 的权重衰减设计也遵循同样的思路W W - lr * lambda * W而不是G lambda * W。对比项SGDAdamAdamWMuon更新方向原始梯度元素级归一化梯度元素级归一化梯度矩阵正交化梯度是否用二阶矩否是是不用二阶矩权重衰减耦合或解耦耦合解耦解耦适用参数类型全参数全参数全参数推荐 2D 权重矩阵更新范数随梯度变化近似按坐标稳定近似按坐标稳定稳定在 lr 量级谱方向分配按奇异值比例按坐标方差按坐标方差各方向均等这张表并不是说 Muon 在所有场景都优于 Adam。实际上1D 参数如偏置、LayerNorm 的 scale 并不适合做正交化这类参数保留在 Adam 里更合理。这也是后续“Muon Adam 混合优化”这一改进策略的基础。3. 为什么 Muon 能比 Adam 更好三个层面的原因“Muon 优于 Adam”并不是一个普适结论但在大语言模型预训练这类任务上它确实表现出明显优势。原因可以从噪声放大、矩阵谱结构、曲率感知三个角度分析。3.1 Adam 会把噪声坐标放大Muon 会抑制这种放大在训练后期或者小 batch 场景下梯度中包含大量噪声。Adam 对每个坐标独立做归一化会将原本很小平方向上的噪声直接放大到与主方向相同的量级。也就是说Adam 的更新方向里相当一部分预算被分配给了噪声方向。Muon 不按坐标缩放而是先对整个梯度矩阵做谱归一化再通过正交化去除幅度信息。正交化过程等价于把更新方向投影到“最接近梯度的正交矩阵”上噪声方向贡献会被显著压缩。对质量差的 batchMuon 的更新方向更稳定。3.2 嵌入矩阵和隐藏层的谱结构更适合各向同性更新在大模型训练中Embedding 矩阵和 MLP 权重矩阵是体量最大的参数模块。这些矩阵有一个共同特点参数之间存在强相关性梯度矩阵的方向分布比坐标分布更有意义。当网络希望把一个词向量往某个方向推时决定这个词向量语义的是它与其它向量的相对夹角而不是每个元素单独的值。Muon 的正交化等价于在新的更新步中尽量保持权重的“旋转结构”这与 embedding 的语义学习方式是契合的。这也是 Muon 论文和相关实现中会把 Embedding 层也划入 Muon 参数组的原因。3.3 曲率不一致时正交化更新更接近自然梯度思想自然梯度法的核心思想是在参数空间中更新不应该按欧氏距离走而应该按“对模型输出影响相同”的方向走。对权重矩阵而言最粗糙的近似就是让每个谱方向对输出的影响尽量均衡。Adam 的逐坐标归一化没有考虑坐标之间的相关性。Muon 的正交化在某种程度上考虑了矩阵左乘右乘的结构右乘正交矩阵不改变权重的行空间关系左乘正交矩阵不改变列空间关系。这让 Muon 更新的方向对权重的正交变换具有不变性在损失景观几何各向异性明显时比 Adam 更接近我们想要的“按影响分配步长”。3.4 实证场景什么任务更容易看到 Muon 优势在公开的大模型训练实践中Muon 曾被用于十亿到数十亿参数规模模型的预训练对比对象往往是 AdamW。这类实践通常报告了更低的预训练 loss或者相同的 loss 下更少的总步数。保守地说Muon 在以下条件下更容易体现出优势模型参数以 2D 权重矩阵和 Embedding 为主训练步数较多batch 较大使用 bf16 混合精度当前模型的 AdamW 训练流程已经调得比较稳loss 曲线进入平台期。对小型分类模型、图像模型或参数极少的任务Muon 的增益可能不明显甚至因为正交化开销导致训练变慢。因此不建议直接拿 Muon 替换所有项目的 Adam。4. 在 PyTorch 中实现并替换优化器理论清楚以后接下来落地。下面给出一个最小可运行的 Muon 实现并通过参数组拆分把 Muon 和 Adam 结合起来。4.1 环境准备建议环境如下Python 3.9 以上PyTorch 2.0 以上支持torch.linalg.matrix_norm单卡或小规模多卡测试均可在代码里主要用到了torch.no_grad()、torch.linalg.matrix_norm和基本的张量运算不需要额外安装第三方优化器库。如果已有项目用 transformers 或自己写的 training loop只需要替换optimizer部分。4.2 最小 Muon 实现import torch from torch.optim import Optimizer def orthogonalize(G, steps5): 对矩阵 G 做 Newton-Schulz 正交化返回近似正交的更新方向。 m, n G.shape if min(m, n) 1: return G # 短边转置保证主循环里 m n数值上更稳定 if m n: return orthogonalize(G.T, steps).T # 谱范数归一化保证迭代收敛 G G / (torch.linalg.matrix_norm(G, ord2) 1e-8) # 公开实现中常见的多项式系数实际项目可做小范围调整 a, b, c 3.4445, -4.7750, 2.0315 X G for _ in range(steps): A X X.T B b * A c * A A X a * X B X return X class Muon(Optimizer): def __init__(self, params, lr2e-2, momentum0.95, weight_decay0.0, ns_steps5): defaults dict(lrlr, momentummomentum, weight_decayweight_decay, ns_stepsns_steps) super().__init__(params, defaults) torch.no_grad() def step(self, closureNone): loss closure() if closure is not None else None for group in self.param_groups: lr group[lr] beta group[momentum] wd group[weight_decay] steps group[ns_steps] for p in group[params]: if p.grad is None: continue g p.grad state self.state[p] if m not in state: state[m] torch.zeros_like(g) m state[m] m.mul_(beta).add_(g, alpha1 - beta) # AdamW 式解耦权重衰减不影响动量 if wd ! 0: p.mul_(1 - lr * wd) # 2D 且两个维度都大于 1 的参数做正交化 if p.ndim 2 and min(p.shape) 1: o orthogonalize(m, steps) else: o m p.add_(o, alpha-lr) return loss这段代码的关键点有三个。第一orthogonalize里先做谱范数归一化如果跳过这一步梯度范数很大时 Newton-Schulz 迭代可能不收敛。第二正交化只作用在 2D 权重矩阵上因为对一维向量做正交化基本只是缩放或符号翻转没有实际收益。第三Momentum用的是零中心动量没有做 bias correction这是 Muon 常见实现与 PyTorch Adam 的一个差异。4.3 参数组拆分哪些参数用 Muon哪些用 Adam实际项目中不要把所有参数都丢给 Muon。推荐做法是2D 权重矩阵用 Muon1D 参数和 embedding 之外的偏置、norm 参数用 Adam。用法如下def split_params(named_params): muon_params [] adam_params [] for name, param in named_params: if param.ndim 2 and min(param.shape) 1: muon_params.append(param) else: adam_params.append(param) return muon_params, adam_params muon_params, adam_params split_params(model.named_parameters()) optimizer Muon(muon_params, lr2e-2, weight_decay0.01) optimizer.add_param_group({ params: adam_params, lr: 1e-3, weight_decay: 0.0, })这样做是因为 2D 权重矩阵是 Muon 发挥优势的主场而偏置和缩放参数数量少、对学习率敏感度低保留在 Adam 里更安全。如果项目里有自定义的 2D 参数但语义上不是“线性变换矩阵”例如位置编码的 2D 表示需要先评估再决定是否纳入 Muon。4.4 学习率与权重衰减的初始配置Muon 的正交化更新方向范数约等于 1所以它的学习率与 Adam 不能直接类比。常见的做法是以2e-2或1e-2作为 Muon 组的起点Adam 组保持1e-3左右。如果你的模型比较小也可以从5e-3尝试。在公开的 Muon 训练配方中学习率通常会给得比 AdamW 更大但必须结合自己的任务重新标定。参数建议初始值说明Muon lr1e-2 到 2e-2更新范数稳定可先用大一点的学习率Muon momentum0.95与 Adam 的 beta1 类似Muon weight_decay0.01 到 0.1解耦式按模型大小调整Adam lr1e-3 左右用于偏置、norm 等 1D 参数ns_steps5迭代次数测试时可调整到 3 或 6这里特别强调不要直接沿用之前 AdamW 搜索出来的学习率。AdamW 的 3e-4 对 Muon 来说通常会偏小导致收敛变慢如果直接套用 SGD 常用的 1e-1又会因为更新太大而发散。4.5 小规模验证方法建议先用一个 1 亿参数以内的语言模型或小型 MLP 跑一次 100 到 500 步的训练记录loss、grad_norm和update_norm。判断标准不是“最终准确率”而是三件事Muon 组的更新范数是否稳定在lr量级loss 曲线是否比 AdamW 基线下降更快或者更低是否有参数组的 grad norm 突然异常。total_grad_norm 0.0 for group in optimizer.param_groups: for p in group[params]: if p.grad is not None: total_grad_norm p.grad.norm().item() ** 2 total_grad_norm total_grad_norm ** 0.5 print(fstep{step} loss{loss.item():.4f} grad_norm{total_grad_norm:.4f})如果grad_norm在一个 step 内跳变超过 10 倍优先怀疑梯度裁剪策略或某个 batch 的数据异常。如果grad_norm正常但 loss 不下降优先怀疑 Muon 的lr太小或正交化迭代次数不足。5. 如何进一步改进 MuonMuon 本身是一个不错的基础优化器但要真正落地在生产训练里还需要结合工程实践做几类改进。5.1 混合优化器Muon Adam效果常常最好前文已经写了参数组拆分这里再强调一遍其动机Muon 擅长更新“有旋转结构”的矩阵参数Adam 擅长逐坐标自适应地更新容易受限的参数。对一个大模型来说最合理的不是“全 Muon 或全 Adam”而是“矩阵归 Muon向量归 Adam”。在一个典型的大模型里这样做Embedding 权重矩阵 - Muon Self-Attention 的 QKV/O 权重 - Muon FFN 的两个权重矩阵 - Muon 偏置 b、LayerNorm 的 gamma - Adam 最终分类头 - Muon 或 Adam 均可需单独试验分类头如果与 embedding 共享权重需要和 embedding 使用同一优化器避免两套更新策略相互干扰。如果分类头是独立矩阵可以先统计它的梯度谱再决定归属。5.2 分层学习率与自动谱尺度虽然 Muon 的正交化本身已经隐式做了谱归一化但不同层级的学习率需求还是不一样。底层 embedding 和顶层分类头通常需要更小或更大的学习率。常见做法是在参数组上单独设置lroptimizer Muon( [ {params: embedding_params, lr: 1e-2}, {params: ffn_params, lr: 2e-2}, ], momentum0.95, weight_decay0.01, )更进一步的思路是“自动谱尺度”在正交化之后再乘上一个与谱分布相关的系数让更新量不再是完全均等而是稍微偏向信息量大的谱方向。这种思路介于 Muon 和 SGD 之间目前还没有统一的公开标准实现可以把它当成实验方向不要直接上生产。5.3 权重衰减策略Muon 的权重衰减必须用解耦形式。如果沿用旧的 L2 正则化写法即在梯度的位置加上wd * W再送入动量那么正交化会把权重衰减的方向也一起改变权重“衰减”就变成了“按正交化方向收缩”行为不可控。推荐做法W W - lr * wd * W # 解耦衰减不进入动量 W W - lr * O # 正交化更新权重衰减值的选择建议和 AdamW 基线保持一致例如从 0.01 开始。在grad_norm稳定的前提下如果 loss 下降变慢可以尝试调小权重衰减到 0.001。5.4 混合精度、梯度裁剪与分布式注意事项Muon 在 bf16 混合精度下表现稳定但需要注意正交化中的矩阵乘法精度。一个常见做法是在 bf16 下完成动量累加但在计算X X.T时临时提升到 fp32避免累加误差影响正交性。PyTorch 的 autocast 不一定自动处理这一步必要时在orthogonalize内部手动转精度。梯度裁剪仍然需要保留。虽然 Muon 正交化会削弱梯度的绝对幅度但动量 buffer 可能仍会出现极端值。一般使用torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)与 AdamW 流程相同。分布式训练里Muon 的真正优势是它不需要为每个参数维护二阶矩v_t只是多了一步正交化计算。正交化只涉及当前进程的局部矩阵乘不引入额外通信。相比 AdamW 每个 step 都要额外同步二阶矩Muon 在通信量上反而更省。5.5 可探索的改进方向如果你的实验已经跑通基础 Muon可以尝试下面几个方向调整 Newton-Schulz 迭代次数观察 3 次、5 次、8 次的 loss 差异在训练不同阶段切换优化器例如前半段 AdamW后半段 Muon使用 WSD 或余弦学习率调度配合 Muon 的稳定更新范数对正交化后的O再做一次行范数归一化适合某些需要严格控制更新步长的场景。这些方向没有标准答案但每一个都需要单独做对照实验不要同时改多个变量。6. 常见问题与排查从 Adam 切到 Muon最常见的不是“跑不起来”而是“看起来在训练但效果不对”。下面按现象给出排查路径。问题现象常见原因检查方式处理建议训练直接发散loss 变 NaNMuon lr 过大或未做谱范数归一化打印正交化前 G 的谱范数调低 lr 或检查orthogonalize是否除以谱范数loss 能下降但明显慢于 AdamWMuon lr 过小观察 Muon 组 update_norm 是否只有 lr 的 0.1 倍提高 Muon lr 到 2e-2 附近某个梯度值突然异常大momentum 累积了异常 batch查看 grad_norm 历史曲线增加梯度裁剪检查数据管线是否混入坏样本正交化速度太慢迭代次数过多或矩阵过大profile 每个 step 的耗时减少 ns_steps 到 3或在 fp32 和 bf16 之间取平衡1D 参数也被正交化参数组划分不对打印各参数组名称按ndim 2把参数划到 Adam权重衰减看起来没生效用了耦合 L2 而非解耦检查weight_decay是否进入动量改为p.mul_(1 - lr * wd)的解耦形式下面再展开几个容易混淆的细节。第一个坑对 1D 向量做正交化。向量只有两个方向正交化结果基本就是正负号信息量很小。如果你把 LayerNorm 的 scale 参数也丢给 Muon不仅没有收益还可能因为更新幅度过大导致 loss 震荡。第二个坑Newton-Schulz 迭代次数设成 1 或 2。正交化不充分时X还没有收敛到正交矩阵更新方向仍然带有一部分幅度信息这和 Adam 的区别就变小了。至少从 5 次开始再逐步调整。第三个坑直接套用 Adam 的 weight decay 数值。AdamW 里常见的 0.01 可能在 Muon 上偏大尤其当模型较大时。如果训练初期 loss 下降正常中后期grad_norm越来越小优先怀疑 weight decay 过大压制了有效的更新量。7. 落地建议与可复用清单最后给出一个从 AdamW 切换到 Muon 的落地流程以及一份可以直接用于评审或实验记录的检查清单。7.1 从 Adam 切换到 Muon 的决策流程不要因为 Muon 效果好就全量替换。更稳妥的顺序是固定训练数据和模型结构先复现一份 AdamW 基线。在小型模型上实现 Muon 参数组拆分。用小学习率跑 100 步确认 loss 能下降、grad_norm 稳定。逐步调大 Muon 学习率找到平稳区间。与 AdamW 做相同训练步数的对照实验比较 loss 曲线和验证指标。确认收益后再迁移到大模型。这个流程的好处是每一步都可以通过 loss 和 grad_norm 判断优化器本身是否正常工作而不是糊里糊涂把问题混进模型改动里。7.2 发布前检查清单在正式训练任务里使用 Muon 之前建议逐项确认只有 2D 权重矩阵使用 Muon1D 参数和 norm 参数使用 Adam权重衰减是解耦形式没有进入动量orthogonalize内部做了谱范数归一化Newton-Schulz 迭代次数不少于 3 次Muon 参数组的lr是单独标定的不是直接沿用 AdamW梯度裁剪保留max_norm与基线一致混合精度下矩阵乘法使用 fp32 或经过精度控制记录 Muon 组的update_norm用于判断学习率是否合适保留一份 AdamW 基线代码方便随时对照。7.3 快速判断收益的方法如果你不想完整调参只想快速知道 Muon 是否适合自己当前的模型可以做一个“同 loss 步数对比”固定训练步数为 200 步分别用 AdamW 和 Muon 训练记录到达相同 loss 的步数差。如果 Muon 能在八成步数内达到 AdamW 的 loss就值得继续调参如果连这个都达不到先不要扩大规模。Muon 的价值不在于替代 Adam而在于提供了一种比“逐坐标归一化”更合理的矩阵参数更新视角。光谱分配的思路也提醒我们优化器设计不只是调 beta 和 lr更是在回答一个问题更新预算到底应该按什么原则分配给参数空间里的各个方向。把这个问题想清楚以后用不用 Muon 其实都只是选择问题而不再是黑盒调参。