这次我们来看一个跟训练大模型直接相关的优化器问题Muon 为什么能打赢 Adam以及怎么从“谱分配Spectral Allocation”的角度继续改进 Muon。Muon 是月之暗面Moonshot AI在训练 Moonlight 系列模型时公开的一类优化器。它没有引入复杂的二阶 Hessian 信息而是用动量加牛顿-舒尔茨Newton-Schulz迭代把更新矩阵正交化再做一次全局学习率缩放。从公开资料看Muon 在相同模型结构下比 Adam 收敛更快Moonlight-16B 用相对更少的训练 token 就达到了不错的推理和数学表现。这让 Muon 成为继 LAMB、Lion、Sophia 之后又一个值得替换 AdamW 的优化器候选。这篇文章不打算停留在“Muon 很厉害”这个层面而是回答三个问题从谱空间看Muon 相比 Adam 到底改了什么自己怎么写一个可用的 Muon并接到现有训练代码里基于“谱分配”的思路后续还能从哪些方向改进 Muon内容按“算法原理 - 谱视角分析 - PyTorch 实现 - 实验验证 - 改进方向 - 排错与最佳实践”展开。适合正在做预训练、大模型微调或者单纯想研究优化器的读者。全文给出可复现代码和检查清单对显存没有特殊要求。但要提前说明Muon 的收益在矩阵参数为主的大模型上最明显小规模任务上不一定能看出差距。1. 核心能力速览能力项说明项目/算法Muon 优化器Spectral Allocation 视角提出背景月之暗面 Moonlight 系列模型训练2025 年公开核心机制动量 Newton-Schulz 正交化 全局学习率相比 Adam 的核心差异对二维及以上参数按矩阵谱结构归一化而不是按元素归一化适用参数二维及以上矩阵参数一维参数仍走 AdamW硬件要求只增加矩阵乘法开销常规 PyTorch 训练环境即可显存占用比 AdamW 少一个二阶矩缓冲区但要保留动量总体与 AdamW 同量级实际以本机测试为准启动方式代码内替换优化器即可无独立服务是否支持 API不涉及这是一段训练/优化代码是否支持批量任务不涉及从规格上看Muon 不是一个“部署工具”而是一个直接写进训练脚本的优化器。它最大的门槛不是环境而是理解它对矩阵参数做了什么。2. 先厘清Adam、AdamW、SGD 的差别讨论 Muon 之前先把基线优化器的区别说清楚。很多文章把 Adam 和 AdamW 混着写实际它们两个在权重衰减上完全不同。Adam 的更新公式如下# Adam 伪代码 m beta1 * m (1 - beta1) * g # 一阶动量 v beta2 * v (1 - beta2) * g * g # 二阶动量 param - lr * m / (sqrt(v) eps) # 归一化更新Adam 的问题在于它把 L2 正则直接写进了梯度里。对一般任务问题不大但预训练大模型时L2 正则与 Adam 的自适应学习率耦合会导致权重大小和损失项之间互相干扰泛化变差。AdamW 的做法是解耦权重衰减先把参数按固定比例缩小再走 Adam 更新。# AdamW 伪代码weight decay 与梯度解耦 param - lr * weight_decay * param m beta1 * m (1 - beta1) * g v beta2 * v (1 - beta2) * g * g param - lr * m / (sqrt(v) eps)这看起来只是移动了一行代码但 AdamW 在大模型上几乎成了默认配置。再看带动量的 SGDmomentum_buffer momentum * momentum_buffer g param - lr * momentum_bufferSGD 的更新方向完全由梯度主导没有逐坐标归一化因此对大尺度差异敏感需要精细调学习率。从这三个优化器可以提炼出一条主线优化器到底怎么做“归一化”。Adam 系列做的是逐元素归一化Muon 做的是矩阵级正交化。下面从谱的角度展开。3. 谱视角为什么“按元素归一化”不够考虑一层网络的权重矩阵 W ∈ R^{m×n}梯度矩阵 G 也可以看作同一形状的矩阵。对它做奇异值分解G U Σ V^T其中 Σ 的对角线是奇异值 σ1 ≥ σ2 ≥ ... ≥ σmin(m,n)。这些奇异值反映梯度在不同“谱方向”上的强度。深层网络中权重矩阵梯度的奇异值通常非常不均匀少数方向贡献了绝大部分梯度能量其余方向接近噪声。SGD 直接把 G 当作更新更新能量与奇异值成正比。大奇异值方向会主导训练这也是 SGD 在大模型上容易震荡的深层原因。Adam 按元素归一化它把每个坐标的梯度除以自己的 RMS。这个操作等于在“坐标基”下做对角缩放但矩阵的左奇异方向、右奇异方向会被打乱。换句话说Adam 做的不是谱空间的白化而是坐标空间的白化。Muon 做正交化把动量矩阵投影到“正交矩阵集合”附近使所有非零奇异值都被拉向 1。更新在每个谱方向上能量接近相等。这是真正的“谱白化”。这就是“Spectral Allocation”的核心同一个学习率预算你要怎么分配给不同的谱方向SGD 是“按奇异值大小分配”Adam 是“没有明确谱目标的坐标分配”Muon 是“均匀分配”。Adam 在很多任务上效果好的原因是逐坐标归一化让每个参数都有了自己的“等效学习率”解决了梯度尺度差异问题。但它仍然默认“参数坐标之间互相独立”。对一维向量这个假设没问题对真正的权重矩阵它丢掉了矩阵的谱结构。Muon 正是从这一点上做改进。4. Muon 算法拆解Muon 的完整更新流程分三步动量、正交化、全局缩放。M_t β M_{t-1} (1-β) G_t O_t NewtonSchulz(M_t, k5) W_{t1} W_t - η · O_tβ 通常取 0.95。也可以加 Nesterov 加速用M_t * β G_t作为实际正交化的输入。关键在 Newton-Schulz 迭代。它不需要做 SVD只靠矩阵乘法就能把输入矩阵投影到正交矩阵附近。def zeropower_via_newtonschulz(G, steps5, eps1e-7): Newton-Schulz 迭代把矩阵投影到正交矩阵附近。 输入 G 可以是任意二维矩阵返回与 G 同形状的近似正交矩阵。 # 归一化保证迭代收敛半径 G G / (torch.linalg.matrix_norm(G) eps) 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这里系数 a、b、c 是预先算好的五次多项式迭代系数。每次迭代做两次矩阵乘法5 步就是 10 次矩阵乘法。对 m、n 不超过几千的矩阵来说开销远小于 SVD。为什么不直接用 SVD原因有三个SVD 的分层计算在 GPU 上不稳定尤其混合精度下容易出现 NaN。SVD 无法高效地嵌入到反向传播后的优化步骤里矩阵分解本身会破坏显存和计算图。Newton-Schulz 只涉及矩阵乘CUDA 核友好几十行就能实现。对于非方阵迭代过程中X X.T得到 m×m 矩阵。如果 m 远大于 n可以先转置让 Gram 矩阵更小。最终结果仍然是一个“近似半正交”的矩阵。对卷积权重这类四维参数需要先 reshape 成二维矩阵再处理。5. Spectral Allocation为什么 Muon 比 Adam 收敛更快从公开资料看Muon 之所以在大模型训练上优于 Adam可以从四个层面解释。5.1 避免少数奇异方向主导Adam 的逐元素归一化解决了“梯度绝对尺度”问题但没有解决“谱方向尺度不均”问题。当梯度矩阵存在一个很大的主导奇异值时Adam 在坐标空间的归一化会让更新在主导方向上仍然偏大形成震荡。Muon 的正交化把更新矩阵的最大奇异值拉回 1从根源上避免了大方向过冲。5.2 改善矩阵参数的条件数损失函数在参数空间中的曲率如果很不均匀Adam 的逐元素估计相当于一个对角预条件子。但权重矩阵的参数之间存在强烈的“左右奇异方向耦合”用对角近似这类耦合会被忽略。Muon 的矩阵级正交化相当于在奇异值空间做统一缩放比对角近似更贴合矩阵参数的真实几何。5.3 保留左右奇异方向的耦合关系SVD 中左右奇异向量分别对应输出通道和输入通道的组合方向。Adam 按元素缩放会把这种组合关系拆散Muon 是在矩阵乘法层面做的投影保留了 U 和 V 的结构信息。Transformer 的 QKV 投影和 FFN 中间层都是矩阵乘法对这类结构自然更友好。5.4 动量与正交化的组合效应动量平滑梯度噪声正交化去掉尺度差异。两者结合之后优化轨迹更像是在“方向空间”里做稳定行走。实践中的一个直观表现是同样的学习率下Muon 的 loss 曲线通常更平滑不容易出现 Adam 早期那种剧烈抖动。需要强调Muon 不会对一维参数做正交化。向量的“谱”没有明确的左右空间结构正交化反而可能破坏维度间的尺度信息。所以原版实现把 bias、LayerNorm 的 scale/shift、embedding 之类的一维参数交给 AdamW 处理矩阵参数走 Muon。这也是 Spectral Allocation 的一种应用只有矩阵参数才谈得上谱向量参数就用坐标归一化。下面用一个表直接对比三种优化器的谱分配方式对比维度SGDAdam / AdamWMuon更新依据原始梯度逐元素归一化梯度动量矩阵正交化谱方向处理按奇异值权重分配坐标白化谱方向被混合均匀分配一维参数同样适用自适应学习率退回 AdamW额外显存一份动量两份动量缓冲一份动量 少量临时量理论依据一阶方向对角近似二阶矩矩阵谱归一化大模型表现稳定但收敛慢默认基线公开数据收敛更快6. 怎么进一步改进 Muon标题里的“How to Improve Muon”下面给几个可以落地的改进方向。这些方向不保证每个都有效但都可以在训练脚本里做成开关一次只改一个变量用固定种子做对照实验。6.1 谱插值在 SGD 与 Muon 之间滑动标准 Muon 把所有奇异值拉向 1但训练早期可能希望保留部分原始尺度让参数移动更快。可以引入一个 α 参数控制“正交化强度”alpha 0.8 # 可设为 schedule从 0.6 升到 1.0 mom_norm mom / (torch.linalg.matrix_norm(mom) 1e-7) ortho zeropower_via_newtonschulz(mom, steps5) update (1 - alpha) * mom_norm alpha * ortho param.data.add_(update, alpha-lr)α0 时接近带动量的 SGDα1 时是标准 Muon。训练早期用小 α 让模型快速探索后期用大 α 稳定收敛。这可能比固定为 1 更灵活。6.2 自适应谱缩放Adaptive Spectral Scaling正交化把所有奇异值变为 1但有些方向可能确实应该走大步。改进思路是用低成本方法估计动量矩阵前 k 个奇异方向在这些方向上额外做缩放把 Adam 的“逐坐标二阶矩”升级成“逐谱方向二阶矩”。具体方案每 N 步用 power iteration 估计矩阵 MM^T 的 top-k 特征值得到近似奇异值 r_i。对正交化结果在这 k 个方向上乘上系数 f(r_i)。系数函数可以设计为1 / sqrt(r_i c)这类形式模仿 Adam 的二阶矩归一化。这样的自适应谱缩放比完整 SVD 便宜得多又能保留谱分配的思想。6.3 谱截断与低秩分配当动量矩阵的奇异值衰减很快时有效秩很小后半段谱方向基本是噪声。可以只保留前 r 个奇异方向做更新其余方向用小尺度更新。这样既降低对噪声方向的敏感度又减少计算量。严格做 SVD 代价高改用 randomized SVD 或者在上一步 power iteration 基础上扩展即可。6.4 与 AdamW 的混合参数分组不是所有矩阵参数都适合均匀谱分配。embedding 和输出头的梯度谱分布与 FFN 很不一样通常已经自带较规范的结构。一个稳妥做法是embedding、lm_head、bias、norm 参数继续用 AdamWattention 的 q/k/v/o 矩阵、FFN 中间矩阵用 Muon。在超参搜索时还可以给不同参数组设置独立学习率。这是“按模块做谱分配”比全局统一替换更容易取得正向收益。6.5 学习率与谱结构联动Muon 的更新矩阵谱范数近似为 1学习率直接决定“步长”。可以周期性地计算动量矩阵的谱分布离散程度比如用奇异值方差做一个简单控制器谱越分散说明方向差异大学习率适当调小谱越集中说明更新方向一致学习率可以放大。这个控制器不需要精确 SVD用 power iteration 估计最大和最小奇异值比值即可。比固定 schedule 更贴近实际训练状态但需要额外写监控代码。6.6 与 weight decay、LayerScale 的耦合原版 Muon 采用解耦权重衰减和 AdamW 一致。如果网络里用了 LayerScale、输出头缩放之类的技巧需要注意正交化会改变参数更新的绝对幅度LayerScale 的初始值可能需要重新调。一个改进方向是让正交化前的输入缩放与层缩放保持一致避免一层被“压扁”另一层被“放大”。7. PyTorch 实现与验证方法下面给出一个可直接运行的 Muon 实现。它把二维及以上参数走 Muon一维参数走 AdamW。import torch from torch.optim import Optimizer def zeropower_via_newtonschulz(G, steps5, eps1e-7): Newton-Schulz 迭代把矩阵投影到正交矩阵附近。 G 为二维矩阵返回与 G 同形状的近似正交矩阵。 G G / (torch.linalg.matrix_norm(G) eps) 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): Muon 优化器 - 二维及以上参数动量 Newton-Schulz 正交化 - 一维参数AdamW 风格更新 def __init__(self, params, lr0.02, momentum0.95, nesterovTrue, ns_steps5, betas(0.9, 0.95), eps1e-8, weight_decay0.01): defaults dict(lrlr, momentummomentum, nesterovnesterov, ns_stepsns_steps, betasbetas, epseps, weight_decayweight_decay) super().__init__(params, defaults) torch.no_grad() def step(self): for group in self.param_groups: lr group[lr] for p in group[params]: if p.grad is None: continue grad p.grad.data state self.state[p] if group[weight_decay] 0: p.data.mul_(1 - lr * group[weight_decay]) if p.ndim 2: if mom not in state: state[mom] torch.zeros_like(grad) mom state[mom] mom.mul_(group[momentum]).add_(grad) if group[nesterov]: update mom * group[momentum] grad else: update mom update zeropower_via_newtonschulz( update, stepsgroup[ns_steps] ) p.data.add_(update, alpha-lr) else: if exp_avg not in state: state[exp_avg] torch.zeros_like(grad) state[exp_avg_sq] torch.zeros_like(grad) state[step] 0 exp_avg state[exp_avg] exp_avg_sq state[exp_avg_sq] beta1, beta2 group[betas] state[step] 1 exp_avg.mul_(beta1).add_(grad, alpha1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value1 - beta2) bias_corr1 1 - beta1 ** state[step] bias_corr2 1 - beta2 ** state[step] denom (exp_avg_sq.sqrt() / (bias_corr2 ** 0.5)).add_(group[eps]) p.data.addcdiv_(exp_avg, denom, value-lr / bias_corr1)使用方式和普通 PyTorch 优化器一致model MyTransformer() optimizer Muon(model.parameters(), lr0.02, momentum0.95, weight_decay0.01) for x, y in dataloader: optimizer.zero_grad() loss model.compute_loss(x, y) loss.backward() optimizer.step()7.1 一个可复现的对照实验为了快速验证 Muon 和 AdamW 的差别可以先用一个两层 MLP 跑回归任务。注意这个任务非常小Muon 的优势不一定能体现出来但能验证代码正确性。import torch import torch.nn as nn class MLP(nn.Module): def __init__(self, d_in64, d_hidden256, d_out16): super().__init__() self.fc1 nn.Linear(d_in, d_hidden) self.fc2 nn.Linear(d_hidden, d_hidden) self.fc3 nn.Linear(d_hidden, d_out) self.act nn.GELU() def forward(self, x): h self.act(self.fc1(x)) h self.act(self.fc2(h)) return self.fc3(h) def make_data(batch_size32, n_batches200, seed0): torch.manual_seed(seed) for _ in range(n_batches): x torch.randn(batch_size, 64) w torch.randn(64, 16) / 8.0 y x w 0.05 * torch.randn(batch_size, 16) yield x, y def train_compare(opt_class, lr, label): torch.manual_seed(42) model MLP() optimizer opt_class(model.parameters(), lrlr) for step in range(101): model.train() for x, y in make_data(): optimizer.zero_grad() loss nn.functional.mse_loss(model(x), y) loss.backward() optimizer.step() if step % 20 0: print(f[{label}] step{step:3d} loss{loss.item():.6f}) if __name__ __main__: train_compare(torch.optim.AdamW, 1e-3, AdamW) train_compare(Muon, 0.005, Muon )Muon 默认学习率和 AdamW 不是一个量级需要单独搜索。小任务上建议从 0.002 到 0.02 之间试。7.2 从谱上观察区别除了看 loss还可以直接看两个优化器更新矩阵的奇异值分布。def update_spectrum(optimizer, p, topk8): 取出某个矩阵参数对应的更新量返回前 topk 个奇异值 state optimizer.state[p] if mom in state: update state[mom] elif exp_avg in state: update state[exp_avg] else: update p.grad.data update_mat update.reshape(update.size(0), -1).float() u, s, v torch.linalg.svd(update_mat, full_matricesFalse) return s[:topk].detach().cpu().numpy()对同一模型分别用 AdamW 和 Muon 跑一步输出奇异值。通常可以看到AdamW 的更新矩阵奇异值可能前两个特别大、后面快速衰减Muon 的更新矩阵所有奇异值都接近 1。这就是“谱分配”差异的直接证据。8. 资源占用与性能观察Muon 的显存变化需要分两部分看。第一它不做逐坐标二阶矩估计所以比 AdamW 少一个v_t缓冲区。对单个矩阵参数AdamW 要存 exp_avg 和 exp_avg_sq 两份 bufferMuon 只需要一份动量 buffer。从参数状态的角度Muon 更省显存。第二Newton-Schulz 迭代会产生临时矩阵。和 AdamW 一样临时变量在算子内部