最近在训练大模型和深层网络时Muon 这个优化器频繁出现在社区讨论里。不少实验显示在相同的参数量、相同的 token 数下Muon 的收敛曲线明显比 AdamW 更平滑最终 loss 也更低。与此同时“Spectral Allocation”这个词开始和 Muon 绑定出现很多人第一次看到时会误以为它只是某种学习率调整策略其实它更像是一种针对矩阵参数特性的更新方向分配机制。这篇文章会从一个比较容易理解的角度拆解 Muon 与 Adam 的本质区别解释为什么 Muon 在矩阵参数较多的网络结构中能胜过 Adam并给出可以跑起来的 PyTorch 实现和一组改进方向。适合正在做模型训练调参的开发者也适合想深入理解优化器原理的初学者。1. 背景从 SGD 到 Adam再到 Muon1.1 优化器到底在优化什么神经网络的训练过程本质上是在不断求解一个高维非凸优化问题。每一步我们都需要根据当前参数计算梯度并根据梯度方向更新参数。优化器负责回答两个问题更新方向是什么以及每一步走多远。SGD 是最朴素的选择它直接用梯度方向作为更新方向用学习率控制步长。它的问题在于梯度在不同维度上的尺度差异很大如果直接使用统一学习率某些维度会收敛很慢某些维度可能会震荡。Adam 的核心改进就是利用梯度的一阶矩和二阶矩估计对每个参数坐标做归一化相当于给每个维度一个独立的步长。这种方式在大多数情况下非常有效尤其是训练 Transformer 这类深层模型。但是Adam 有一个隐藏的局限它的归一化是针对“单个参数元素”的而神经网络中的参数往往是以矩阵形式组织的例如线性层权重、卷积核、注意力投影矩阵。这些矩阵本身具有谱结构奇异值分布、条件数等Adam 完全没有利用这种结构。1.2 什么是 Spectral AllocationSpectral Allocation 可以直译为“谱分配”。它关注的是优化器在更新一个矩阵参数时不同奇异方向上的更新量应该如何分配。每一轮梯度更新都可以看作是在参数空间中移动一个向量。这个向量可以分解到矩阵的奇异向量方向上。Adam 的做法是按元素归一化相当于对每个坐标施加相同的缩放不会考虑奇异结构。Muon 的做法则不同它会先对动量缓冲做正交化处理使更新方向更符合矩阵的参数几何相当于把更新能量重新分配到主要奇异方向上。换句话说Adam 是“每个参数坐标独立处理”Muon 是“按矩阵整体谱结构协调更新”。Spectral Allocation 指的就是 Muon 在做正交化时对矩阵各奇异方向进行的这种能量重新分配。1.3 Muon 适用于哪些场景Muon 并不是要完全取代 Adam。它在以下场景中更有优势模型主体由大量矩阵参数构成例如 Transformer、MLP-Mixer 等。存在明显的长尾奇异值分布例如大矩阵经过预训练或初始化后主奇异方向占主导。训练非常大 batch 的模型需要更稳定的更新方向。参数量极大希望减少优化器状态占用同时保持收敛质量。从目前社区反馈和公开实验来看Muon 在大模型预训练阶段的表现尤其亮眼这也是它最近热度上升的直接原因。2. 为什么 Muon 会优于 Adam核心原理对比2.1 Adam 的坐标级更新机制先看 Adam 的典型更新公式m_t beta1 * m_{t-1} (1 - beta1) * g_t v_t beta2 * v_{t-1} (1 - beta2) * g_t^2 m_hat m_t / (1 - beta1^t) v_hat v_t / (1 - beta2^t) theta_t theta_{t-1} - lr * m_hat / (sqrt(v_hat) eps)其中g_t是梯度m_t是一阶矩v_t是二阶矩。Adam 会对每个标量元素分别计算m_hat和v_hat然后做逐元素除法。这种做法的优点是实现简单、鲁棒性强尤其适合稀疏特征和不同量级的参数。缺点是它假设每个参数坐标是独立的忽略参数之间的相关性。对于矩阵参数这意味着 Adam 的更新方向可能和矩阵的自然梯度方向不一致在谱结构较强的模型中会导致收敛变慢或振荡。AdamW 在 Adam 基础上将权重衰减从梯度中解耦解决了一部分正则化问题但它并没有改变逐元素归一化的本质。因此当我们说“Muon 优于 AdamW”时核心不是简单的权重衰减差异而是更新方向的几何结构差异。2.2 Muon 是如何工作的Muon 的思路可以拆成几个步骤计算梯度G。用动量累积梯度得到动量矩阵U。对U做正交化处理使其趋向于一个正交矩阵。用U作为更新方向乘上学习率后更新参数。关键是第 3 步。正交化并不等于 PCA也不是 SVD 分解后直接替换而是通过 Newton-Schur 迭代来近似极分解中的正交因子。这样既能保留动量信息又能让更新方向落在 Stiefel 流形的切空间附近。这种处理相当于给矩阵参数一个“全局预条件”让更新方向在谱结构上更加合理。2.3 谱结构差异为什么这对收敛很重要考虑一个参数矩阵W它的 SVD 分解为W U Σ V^T。其中Σ对角线上的奇异值可能相差很大。在训练过程中梯度G在不同奇异方向上的分量也往往差异很大。Adam 对每个元素做归一化等价于给所有方向一个接近一致的学习率。这会使得主奇异方向上的梯度被过度放大而小奇异方向上的梯度被过度抑制。结果就是更新方向偏离了原本利于收敛的轨迹。Muon 正交化之后更新矩阵U的奇异值接近 1。也就是说各个方向的更新幅度被归一化到同一个量级但从矩阵几何角度来看这等于把更新能量平均分配到了所有方向反而避免了 Adam 中“大方向过大、小方向过小”的问题。这里要说明一点Muon 并不是在所有任务上都一定优于 Adam。它的优势更体现在矩阵参数占主导、模型深度较大、训练步数较多的场景。在这些场景下正确的谱分配比逐元素自适应更重要。2.4 直观比喻可以把优化过程想象成在山区行走。Adam 的做法是每踏出一步都根据脚下的坡度分别调整前后左右方向的步子大小。这在地形复杂时很有用但缺点是你不会去看远处的山脊走向。Muon 的做法是每走几步就拿出指南针修正方向保证整体前进方向更接近山谷的走向。虽然每一步的计算更复杂但长期来看路线更直不容易走偏。这个比喻并不完全严谨但能解释为什么 Muon 的收敛曲线通常更平滑也为什么不那么容易在中途“卡住”。3. 环境准备与实验设计在动手实现 Muon 之前先把实验环境准备好。这里的版本只是参考实际使用时以你本机环境为准。3.1 基础环境操作系统Ubuntu 20.04 / 22.04Windows WSL2 也可以。Python3.9 或 3.10。PyTorch1.13 或 2.x 都可以推荐 2.1 以上。CUDA11.8 或 12.x不强制 GPU纯 CPU 也能跑通逻辑。建议创建一个独立的虚拟环境conda create -n muon-test python3.10 -y conda activate muon-test pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install numpy matplotlib如果不想装 GPU 版直接pip install torch也可以本文示例的数据量很小CPU 就能跑。3.2 实验设计思路为了对比 Muon 和 Adam我们构造一个简单的二分类任务使用一个小型 MLP。模型只有两层线性层中间用 GELU 激活。因为 Muon 对矩阵参数有效线性层的权重矩阵正好能体现差异。实验脚本会分别用 AdamW 和 Muon 训练同一个模型并记录训练 loss。最后对比两条 loss 曲线。需要注意的是为了公平对比两者必须使用相同的参数初始化、相同的数据顺序和相同的学习率搜索策略。在实际工程中Muon 的最佳学习率往往和 AdamW 不同所以对比时可以各取一个接近最优的学习率而不是强行用同一个数值。4. Muon 优化器核心拆解与可运行实现4.1 算法伪代码我们这里实现的是一个常见版本的 Muon流程如下输入参数集 P梯度 G动量缓冲 M 超参学习率 lr动量系数 mu正交化迭代步数 ns_steps for each参数 p in P: g 获取对应梯度 m M[p] m mu * m g if p 是二维矩阵并且行数 列数: u m mu * g # nesterov 形式 u newton_schur(u, ns_steps) # 正交化 else: u m mu * g # 保持原始动量方向 p p - lr * u不同实现中nesterov 的细节会有差异但总体思路一致。如果你的模型包含偏置项、归一化层参数或 embedding 参数通常不对它们做正交化只使用原始动量更新。4.2 Newton-Schur 正交化迭代正交化是 Muon 最核心的步骤。对于一个形状为(m, n)且m n的矩阵Q我们希望迭代后满足Q Q^T ≈ I。常见的迭代公式是Q 1.5 * Q - 0.5 * Q Q.T Q这个公式来自极分解的 Newton 迭代。当Q的行数小于等于列数时Q Q^T收敛到单位矩阵Q收敛到接近正交行矩阵。迭代次数通常取 5 到 10 次即可。为什么不是直接做 SVD因为 SVD 每步计算成本太高不适合作为优化器内部操作。Newton-Schur 迭代只需要矩阵乘法在 GPU 上非常高效并且梯度可以正常回传。4.3 PyTorch 实现 Muon下面给出一份可以直接使用的 PyTorch 实现。这份实现参考了社区常见写法适合理解原理也可以直接集成到自己的训练循环中。# 文件路径muon_optimizer.py import torch def newton_schur(Q: torch.Tensor, steps: int 5) - torch.Tensor: 使用 Newton-Schur 迭代将 Q 的行方向趋于正交。 要求 Q 的 shape 满足行数 列数。 for _ in range(steps): QTQ Q Q.T factor 1.5 * Q - 0.5 * Q QTQ Q factor return Q class Muon(torch.optim.Optimizer): Muon 优化器。 对于二维矩阵参数且行数 列数时对动量更新做正交化。 其他参数偏置、标量、embedding 等仅使用动量更新。 def __init__( self, params, lr: float 0.02, momentum: float 0.95, nesterov: bool True, ns_steps: int 5, ): defaults dict(lrlr, momentummomentum, nesterovnesterov, ns_stepsns_steps) super().__init__(params, defaults) torch.no_grad() def step(self, closureNone): loss None if closure is not None: with torch.enable_grad(): loss closure() for group in self.param_groups: lr group[lr] momentum group[momentum] nesterov group[nesterov] ns_steps group[ns_steps] for p in group[params]: if p.grad is None: continue grad p.grad.data if grad.is_sparse: raise RuntimeError(Muon 暂时不支持稀疏梯度) state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p) buf state[momentum_buffer] buf.mul_(momentum).add_(grad) if nesterov: update buf momentum * grad else: update buf.clone() # 对二维矩阵且行数 列数的参数做正交化 if p.dim() 2 and p.shape[0] p.shape[1]: update newton_schur(update, ns_steps) p.add_(update, alpha-lr) return loss这里的newton_schur会原地修改Q但在 PyTorch 的运算中我们创建了新张量。需要注意如果在update上直接原地操作可能会破坏计算图。上述实现中Q实际上是update buf momentum * grad产生的新张量因此可以安全迭代。4.4 完整训练对比脚本下面用一个非常小的 MLP 模型来对比 Muon 和 AdamW。为了方便阅读我把训练逻辑封装成函数。# 文件路径train_compare.py import torch import torch.nn as nn from torch.optim import AdamW from torch.utils.data import DataLoader, TensorDataset from muon_optimizer import Muon def build_model(): return nn.Sequential( nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32), nn.GELU(), nn.Linear(32, 2), ) def make_data(num_samples4096, batch_size128): torch.manual_seed(0) X torch.randn(num_samples, 32) y (X[:, 0] * X[:, 1] X[:, 2] 0).long() dataset TensorDataset(X, y) loader DataLoader(dataset, batch_sizebatch_size, shuffleTrue) return loader def train(model, optimizer, loader, epochs20): model.train() criterion nn.CrossEntropyLoss() losses [] for epoch in range(epochs): total_loss 0.0 for X, y in loader: optimizer.zero_grad() out model(X) loss criterion(out, y) loss.backward() optimizer.step() total_loss loss.item() * X.size(0) avg_loss total_loss / len(loader.dataset) losses.append(avg_loss) print(fepoch {epoch 1:02d}, loss: {avg_loss:.4f}) return losses if __name__ __main__: loader make_data() model_adam build_model() optimizer_adam AdamW(model_adam.parameters(), lr1e-3, weight_decay0.01) print( AdamW ) losses_adam train(model_adam, optimizer_adam, loader) model_muon build_model() optimizer_muon Muon(model_muon.parameters(), lr0.02, momentum0.95, ns_steps5) print( Muon ) losses_muon train(model_muon, optimizer_muon, loader)注意这个示例只是为了演示代码结构实际对比中 MLP 太小Muon 的优势不会很明显。要看到更明显的差异建议在 Transformer 的一到两层线性层上测试或者使用更大的矩阵参数。4.5 运行方式与预期结果保存muon_optimizer.py和train_compare.py后在终端运行python train_compare.py输出大概是每个 epoch 的 loss。通常 AdamW 在前几个 epoch 下降很快后期会出现一些震荡Muon 如果学习率设置合理整体下降会更平稳。如果第一次跑不出明显差异先检查两点是否所有矩阵参数都满足p.shape[0] p.shape[1]。对于常见的nn.Linear(in_features, out_features)权重形状是(out_features, in_features)如果out_features大于in_features不满足条件就不会做正交化。Muon 的学习率是否偏大或偏小。一般从0.01到0.05之间尝试。4.6 Muon 与 Adam 在内存和计算上的区别Adam 需要维护两个状态一阶动量m和二阶动量v。Muon 只需要维护一个动量缓冲但需要在更新时做矩阵乘法迭代。对于一个(m, n)的矩阵参数Adam 的状态占用是2 * m * n个浮点数Muon 的状态占用是m * n个浮点数。在大模型场景下优化器状态占比很大Muon 可以明显节省显存。计算量方面Newton-Schur 迭代的核心是Q Q.T Q复杂度大约是 O(m^2 n)。如果矩阵很扁即 m 远小于 n这个计算开销并不大。这也是为什么 Muon 通常只对行数小于等于列数的矩阵做正交化。5. 如何改进 Muon工程实验与调优方向Muon 不是银弹。直接套用 Muon 可能遇到收敛不稳、loss 不下降等问题。下面整理一些实际工程中行之有效的改进方向。5.1 混合使用 AdamW 和 MuonMuon 适合矩阵参数但偏置、LayerNorm 参数、Embedding 参数并不一定适合正交化。常见做法是对权重矩阵使用 Muon。对偏置、LayerNorm 参数、Embedding 使用 AdamW。这样可以避免对非矩阵参数做无意义的正交化同时保留 AdamW 在精细参数上的稳定性。示例配置from torch.optim import AdamW from muon_optimizer import Muon def configure_optimizers(model): muon_params [] adamw_params [] for name, param in model.named_parameters(): if not param.requires_grad: continue if param.dim() 2 or (param.dim() 2 and param.shape[0] param.shape[1]): # 例如 Linear 权重、Conv 权重的前两维 muon_params.append(param) else: adamw_params.append(param) optim_groups [ {params: muon_params, lr: 0.02}, {params: adamw_params, lr: 1e-3, weight_decay: 0.01}, ] muon Muon(muon_params, lr0.02) adamw AdamW(adamw_params, lr1e-3, weight_decay0.01) return muon, adamw这里我故意写成了两个优化器对象。如果希望更统一也可以把所有参数放在同一个优化器中通过param_group区分。但分开两个优化器在代码上更清晰。5.2 对 Embedding 单独处理在大模型训练中Embedding 矩阵往往也很大。对 Embedding 做正交化容易破坏嵌入空间的语义结构。社区实践中通常对 Embedding 使用 AdamW而不是 Muon。如果你的模型词表很大Embedding 的更新非常稀疏更推荐使用 Adafactor 或 AdamW 这类自适应优化器。Muon 的正交化会改变整个嵌入矩阵方向对稀疏更新并不友好。5.3 调整正交化迭代步数Newton-Schur 迭代步数ns_steps是一个重要超参数。步数越多更新方向越接近正交但计算量也越大。ns_steps1时正交化效果很弱稳定性最好适合小模型。ns_steps5是常见默认值。ns_steps10更接近完全正交但可能损失一些小方向上的信息。建议先用 5观察 loss 是否稳定。如果 loss 突然升高尝试减小步数如果收敛变慢尝试增大步数。5.4 配合权重衰减Muon 的原始设计通常不包含权重衰减因为正交化更新本身具有一定的正则化效果。但如果你希望显式控制权重衰减可以借鉴 AdamW 的解耦思路在更新后单独对参数做衰减# 在每个 step 中加上 p.mul_(1 - lr * weight_decay)不过要注意权重衰减系数不能设置太高否则正交化带来的几何优势会被抵消。建议从0.01开始尝试。5.5 学习率预热与调度Muon 对学习率比较敏感。直接使用较大学习率可能在初期造成震荡。推荐使用 warmup 加 cosine decaydef get_lr(step, total_steps, warmup_steps, peak_lr): if step warmup_steps: return peak_lr * step / warmup_steps progress (step - warmup_steps) / (total_steps - warmup_steps) return peak_lr * 0.5 * (1 torch.cos(torch.pi * progress))训练初期使用较小的学习率可以让动量缓冲充分估计方向之后再提高学习率会稳定很多。5.6 与梯度裁剪结合Muon 的正交化可以部分抑制梯度爆炸但并不意味着不需要梯度裁剪。特别是在训练 Transformer 时梯度范数可能突然变大。建议设置最大梯度范数例如 1.0torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)裁剪顺序要在 optimizer.step() 之前。5.7 改进正交化的版本带缩放的正交化一些改进版本会在正交化后乘上一个缩放系数让更新方向不至于完全损失梯度幅度。例如scale torch.norm(update, dim-1, keepdimTrue) / torch.norm(orthogonalized, dim-1, keepdimTrue) update orthogonalized * scale这种做法相当于保留每个方向上的能量但仍然让方向正交。如果你发现 Muon 收敛太慢可以尝试这个思路。6. 常见问题与排查思路6.1 训练直接出现 NaN问题现象可能原因解决思路loss 变为 NaN或参数出现 NaN学习率过大降低学习率建议从 0.01 开始只有特定层出现 NaN正交化迭代步数过多减少 ns_steps混合精度训练时出现 NaN更新方向与 loss scaling 不匹配在 Muon 前先检查梯度是否有限或使用torch.cuda.amp.GradScaler定位方法在optimizer.step()前检查每个参数梯度的范数for name, param in model.named_parameters(): if param.grad is not None and not torch.isfinite(param.grad).all(): print(f梯度异常{name})6.2 收敛速度比 AdamW 还慢通常是因为正交化的方向过于“平均”导致主梯度方向信息被稀释。此时可以降低正交化步数到 2 或 3。对更新方向乘以一个缩放系数。只对最大矩阵参数使用 Muon其余使用 AdamW。6.3 内存占用没有明显下降Muon 相比 Adam 少维护一个二阶矩状态但如果你同时保存了多个优化器状态或者为了对比实验复制了多个模型内存占用依然会很高。使用前先检查显存分布nvidia-smi如果显存不足可以尝试torch.utils.checkpoint减少激活内存或者使用更小的 batch size。6.4 分布式训练中状态同步问题Muon 的动量缓冲和普通优化器一致本身就支持DDP和FSDP。但在 DDP 中梯度同步后每个进程的newton_schur计算应该得到相同结果因为输入相同。如果出现不同 rank 结果不一致通常是浮点累加顺序导致可以在update计算前添加torch.distributed.barrier()但大多数情况下不需要。6.5 与 DeepSpeed 和 FSDP 的兼容性Muon 的自定义参数分组方式可能会与 DeepSpeed 的优化器封装冲突。如果你使用deepspeed.initialize需要把 Muon 配置为自定义优化器并在配置文件中关掉默认优化器。更简单的办法是先用 PyTorch 原生的 FSDP 测试等跑通后再考虑迁移到 DeepSpeed。7. 最佳实践与工程建议7.1 不要把 Muon 用于所有参数很多初学者会写这样的代码optimizer Muon(model.parameters())这种做法通常效果不好。因为偏置、归一化参数、Embedding 并不适合正交化。更好的策略是对二维权重矩阵使用 Muon。对一维参数和 Embedding 使用 AdamW。对于卷积核可以先按(out_channels, in_channels, k_h, k_w)展平成二维后再做正交化但这一步需要谨慎直接展平可能破坏局部空间结构。7.2 记录详细的实验日志Muon 的超参数比较敏感包括学习率、momentum、ns_steps、是否使用 nesterov。建议使用wandb或tensorboard记录训练 loss 和验证 loss。梯度范数。更新方向与梯度方向的内积。正交化前后更新矩阵的奇异值分布。这些信息能帮你判断 Muon 是否正常工作。7.3 初始化条件不可忽视Muon 对初始化的要求并不比 Adam 更低。如果初始化时矩阵的奇异值分布很不平衡前几步正交化迭代可能会让更新方向剧烈变化。建议初始化时使用较小的标准差例如0.02或者在训练前先做几步 warmup。7.4 在更大模型上验证优势单机小规模实验中Muon 的优势可能不明显因为小模型的谱结构比较简单。如果你希望验证 Muon 是否适合你的业务建议至少在一个 10 层以上的 Transformer 或 5 层以上的 MLP 上测试并和 AdamW 做完整的 lr 扫描。7.5 保持代码可回滚Muon 仍在快速演进不同实现可能存在差异。建议在代码中锁定版本并保留一个 AdamW 的训练路径。当 Muon 出现奇怪问题时可以快速切换回 AdamW 继续实验。8. 总结Muon 的核心价值在于它从矩阵谱结构出发对更新方向做了正交化处理实现了 Adam 无法做到的“全局谱分配”。Adam 适合大多数通用场景但当你面对大量矩阵参数时Muon 往往能带来更平稳的收敛和更好的最终效果。本文从原理、实现、改进到排查给出了一个比较完整的入门路径。你可以先在小型 MLP 上跑通代码再逐步迁移到 Transformer 结构重点对比 loss 曲线和最终精度。不要直接跳到生产环境大规模训练先在小规模任务上积累对超参数的直觉。如果你准备在自己的模型上尝试 Muon建议把混合优化器策略、学习率 warmup、梯度裁剪这三件事一起做。它们能大大提高 Muon 的稳定性。后续还可以进一步阅读极分解、Stiefel 流形优化和自然梯度相关的资料这些知识会帮助你更灵活地调整 Muon 的行为。
Muon优化器解析：从Spectral Allocation到PyTorch实现与调优
最近在训练大模型和深层网络时Muon 这个优化器频繁出现在社区讨论里。不少实验显示在相同的参数量、相同的 token 数下Muon 的收敛曲线明显比 AdamW 更平滑最终 loss 也更低。与此同时“Spectral Allocation”这个词开始和 Muon 绑定出现很多人第一次看到时会误以为它只是某种学习率调整策略其实它更像是一种针对矩阵参数特性的更新方向分配机制。这篇文章会从一个比较容易理解的角度拆解 Muon 与 Adam 的本质区别解释为什么 Muon 在矩阵参数较多的网络结构中能胜过 Adam并给出可以跑起来的 PyTorch 实现和一组改进方向。适合正在做模型训练调参的开发者也适合想深入理解优化器原理的初学者。1. 背景从 SGD 到 Adam再到 Muon1.1 优化器到底在优化什么神经网络的训练过程本质上是在不断求解一个高维非凸优化问题。每一步我们都需要根据当前参数计算梯度并根据梯度方向更新参数。优化器负责回答两个问题更新方向是什么以及每一步走多远。SGD 是最朴素的选择它直接用梯度方向作为更新方向用学习率控制步长。它的问题在于梯度在不同维度上的尺度差异很大如果直接使用统一学习率某些维度会收敛很慢某些维度可能会震荡。Adam 的核心改进就是利用梯度的一阶矩和二阶矩估计对每个参数坐标做归一化相当于给每个维度一个独立的步长。这种方式在大多数情况下非常有效尤其是训练 Transformer 这类深层模型。但是Adam 有一个隐藏的局限它的归一化是针对“单个参数元素”的而神经网络中的参数往往是以矩阵形式组织的例如线性层权重、卷积核、注意力投影矩阵。这些矩阵本身具有谱结构奇异值分布、条件数等Adam 完全没有利用这种结构。1.2 什么是 Spectral AllocationSpectral Allocation 可以直译为“谱分配”。它关注的是优化器在更新一个矩阵参数时不同奇异方向上的更新量应该如何分配。每一轮梯度更新都可以看作是在参数空间中移动一个向量。这个向量可以分解到矩阵的奇异向量方向上。Adam 的做法是按元素归一化相当于对每个坐标施加相同的缩放不会考虑奇异结构。Muon 的做法则不同它会先对动量缓冲做正交化处理使更新方向更符合矩阵的参数几何相当于把更新能量重新分配到主要奇异方向上。换句话说Adam 是“每个参数坐标独立处理”Muon 是“按矩阵整体谱结构协调更新”。Spectral Allocation 指的就是 Muon 在做正交化时对矩阵各奇异方向进行的这种能量重新分配。1.3 Muon 适用于哪些场景Muon 并不是要完全取代 Adam。它在以下场景中更有优势模型主体由大量矩阵参数构成例如 Transformer、MLP-Mixer 等。存在明显的长尾奇异值分布例如大矩阵经过预训练或初始化后主奇异方向占主导。训练非常大 batch 的模型需要更稳定的更新方向。参数量极大希望减少优化器状态占用同时保持收敛质量。从目前社区反馈和公开实验来看Muon 在大模型预训练阶段的表现尤其亮眼这也是它最近热度上升的直接原因。2. 为什么 Muon 会优于 Adam核心原理对比2.1 Adam 的坐标级更新机制先看 Adam 的典型更新公式m_t beta1 * m_{t-1} (1 - beta1) * g_t v_t beta2 * v_{t-1} (1 - beta2) * g_t^2 m_hat m_t / (1 - beta1^t) v_hat v_t / (1 - beta2^t) theta_t theta_{t-1} - lr * m_hat / (sqrt(v_hat) eps)其中g_t是梯度m_t是一阶矩v_t是二阶矩。Adam 会对每个标量元素分别计算m_hat和v_hat然后做逐元素除法。这种做法的优点是实现简单、鲁棒性强尤其适合稀疏特征和不同量级的参数。缺点是它假设每个参数坐标是独立的忽略参数之间的相关性。对于矩阵参数这意味着 Adam 的更新方向可能和矩阵的自然梯度方向不一致在谱结构较强的模型中会导致收敛变慢或振荡。AdamW 在 Adam 基础上将权重衰减从梯度中解耦解决了一部分正则化问题但它并没有改变逐元素归一化的本质。因此当我们说“Muon 优于 AdamW”时核心不是简单的权重衰减差异而是更新方向的几何结构差异。2.2 Muon 是如何工作的Muon 的思路可以拆成几个步骤计算梯度G。用动量累积梯度得到动量矩阵U。对U做正交化处理使其趋向于一个正交矩阵。用U作为更新方向乘上学习率后更新参数。关键是第 3 步。正交化并不等于 PCA也不是 SVD 分解后直接替换而是通过 Newton-Schur 迭代来近似极分解中的正交因子。这样既能保留动量信息又能让更新方向落在 Stiefel 流形的切空间附近。这种处理相当于给矩阵参数一个“全局预条件”让更新方向在谱结构上更加合理。2.3 谱结构差异为什么这对收敛很重要考虑一个参数矩阵W它的 SVD 分解为W U Σ V^T。其中Σ对角线上的奇异值可能相差很大。在训练过程中梯度G在不同奇异方向上的分量也往往差异很大。Adam 对每个元素做归一化等价于给所有方向一个接近一致的学习率。这会使得主奇异方向上的梯度被过度放大而小奇异方向上的梯度被过度抑制。结果就是更新方向偏离了原本利于收敛的轨迹。Muon 正交化之后更新矩阵U的奇异值接近 1。也就是说各个方向的更新幅度被归一化到同一个量级但从矩阵几何角度来看这等于把更新能量平均分配到了所有方向反而避免了 Adam 中“大方向过大、小方向过小”的问题。这里要说明一点Muon 并不是在所有任务上都一定优于 Adam。它的优势更体现在矩阵参数占主导、模型深度较大、训练步数较多的场景。在这些场景下正确的谱分配比逐元素自适应更重要。2.4 直观比喻可以把优化过程想象成在山区行走。Adam 的做法是每踏出一步都根据脚下的坡度分别调整前后左右方向的步子大小。这在地形复杂时很有用但缺点是你不会去看远处的山脊走向。Muon 的做法是每走几步就拿出指南针修正方向保证整体前进方向更接近山谷的走向。虽然每一步的计算更复杂但长期来看路线更直不容易走偏。这个比喻并不完全严谨但能解释为什么 Muon 的收敛曲线通常更平滑也为什么不那么容易在中途“卡住”。3. 环境准备与实验设计在动手实现 Muon 之前先把实验环境准备好。这里的版本只是参考实际使用时以你本机环境为准。3.1 基础环境操作系统Ubuntu 20.04 / 22.04Windows WSL2 也可以。Python3.9 或 3.10。PyTorch1.13 或 2.x 都可以推荐 2.1 以上。CUDA11.8 或 12.x不强制 GPU纯 CPU 也能跑通逻辑。建议创建一个独立的虚拟环境conda create -n muon-test python3.10 -y conda activate muon-test pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install numpy matplotlib如果不想装 GPU 版直接pip install torch也可以本文示例的数据量很小CPU 就能跑。3.2 实验设计思路为了对比 Muon 和 Adam我们构造一个简单的二分类任务使用一个小型 MLP。模型只有两层线性层中间用 GELU 激活。因为 Muon 对矩阵参数有效线性层的权重矩阵正好能体现差异。实验脚本会分别用 AdamW 和 Muon 训练同一个模型并记录训练 loss。最后对比两条 loss 曲线。需要注意的是为了公平对比两者必须使用相同的参数初始化、相同的数据顺序和相同的学习率搜索策略。在实际工程中Muon 的最佳学习率往往和 AdamW 不同所以对比时可以各取一个接近最优的学习率而不是强行用同一个数值。4. Muon 优化器核心拆解与可运行实现4.1 算法伪代码我们这里实现的是一个常见版本的 Muon流程如下输入参数集 P梯度 G动量缓冲 M 超参学习率 lr动量系数 mu正交化迭代步数 ns_steps for each参数 p in P: g 获取对应梯度 m M[p] m mu * m g if p 是二维矩阵并且行数 列数: u m mu * g # nesterov 形式 u newton_schur(u, ns_steps) # 正交化 else: u m mu * g # 保持原始动量方向 p p - lr * u不同实现中nesterov 的细节会有差异但总体思路一致。如果你的模型包含偏置项、归一化层参数或 embedding 参数通常不对它们做正交化只使用原始动量更新。4.2 Newton-Schur 正交化迭代正交化是 Muon 最核心的步骤。对于一个形状为(m, n)且m n的矩阵Q我们希望迭代后满足Q Q^T ≈ I。常见的迭代公式是Q 1.5 * Q - 0.5 * Q Q.T Q这个公式来自极分解的 Newton 迭代。当Q的行数小于等于列数时Q Q^T收敛到单位矩阵Q收敛到接近正交行矩阵。迭代次数通常取 5 到 10 次即可。为什么不是直接做 SVD因为 SVD 每步计算成本太高不适合作为优化器内部操作。Newton-Schur 迭代只需要矩阵乘法在 GPU 上非常高效并且梯度可以正常回传。4.3 PyTorch 实现 Muon下面给出一份可以直接使用的 PyTorch 实现。这份实现参考了社区常见写法适合理解原理也可以直接集成到自己的训练循环中。# 文件路径muon_optimizer.py import torch def newton_schur(Q: torch.Tensor, steps: int 5) - torch.Tensor: 使用 Newton-Schur 迭代将 Q 的行方向趋于正交。 要求 Q 的 shape 满足行数 列数。 for _ in range(steps): QTQ Q Q.T factor 1.5 * Q - 0.5 * Q QTQ Q factor return Q class Muon(torch.optim.Optimizer): Muon 优化器。 对于二维矩阵参数且行数 列数时对动量更新做正交化。 其他参数偏置、标量、embedding 等仅使用动量更新。 def __init__( self, params, lr: float 0.02, momentum: float 0.95, nesterov: bool True, ns_steps: int 5, ): defaults dict(lrlr, momentummomentum, nesterovnesterov, ns_stepsns_steps) super().__init__(params, defaults) torch.no_grad() def step(self, closureNone): loss None if closure is not None: with torch.enable_grad(): loss closure() for group in self.param_groups: lr group[lr] momentum group[momentum] nesterov group[nesterov] ns_steps group[ns_steps] for p in group[params]: if p.grad is None: continue grad p.grad.data if grad.is_sparse: raise RuntimeError(Muon 暂时不支持稀疏梯度) state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p) buf state[momentum_buffer] buf.mul_(momentum).add_(grad) if nesterov: update buf momentum * grad else: update buf.clone() # 对二维矩阵且行数 列数的参数做正交化 if p.dim() 2 and p.shape[0] p.shape[1]: update newton_schur(update, ns_steps) p.add_(update, alpha-lr) return loss这里的newton_schur会原地修改Q但在 PyTorch 的运算中我们创建了新张量。需要注意如果在update上直接原地操作可能会破坏计算图。上述实现中Q实际上是update buf momentum * grad产生的新张量因此可以安全迭代。4.4 完整训练对比脚本下面用一个非常小的 MLP 模型来对比 Muon 和 AdamW。为了方便阅读我把训练逻辑封装成函数。# 文件路径train_compare.py import torch import torch.nn as nn from torch.optim import AdamW from torch.utils.data import DataLoader, TensorDataset from muon_optimizer import Muon def build_model(): return nn.Sequential( nn.Linear(32, 64), nn.GELU(), nn.Linear(64, 32), nn.GELU(), nn.Linear(32, 2), ) def make_data(num_samples4096, batch_size128): torch.manual_seed(0) X torch.randn(num_samples, 32) y (X[:, 0] * X[:, 1] X[:, 2] 0).long() dataset TensorDataset(X, y) loader DataLoader(dataset, batch_sizebatch_size, shuffleTrue) return loader def train(model, optimizer, loader, epochs20): model.train() criterion nn.CrossEntropyLoss() losses [] for epoch in range(epochs): total_loss 0.0 for X, y in loader: optimizer.zero_grad() out model(X) loss criterion(out, y) loss.backward() optimizer.step() total_loss loss.item() * X.size(0) avg_loss total_loss / len(loader.dataset) losses.append(avg_loss) print(fepoch {epoch 1:02d}, loss: {avg_loss:.4f}) return losses if __name__ __main__: loader make_data() model_adam build_model() optimizer_adam AdamW(model_adam.parameters(), lr1e-3, weight_decay0.01) print( AdamW ) losses_adam train(model_adam, optimizer_adam, loader) model_muon build_model() optimizer_muon Muon(model_muon.parameters(), lr0.02, momentum0.95, ns_steps5) print( Muon ) losses_muon train(model_muon, optimizer_muon, loader)注意这个示例只是为了演示代码结构实际对比中 MLP 太小Muon 的优势不会很明显。要看到更明显的差异建议在 Transformer 的一到两层线性层上测试或者使用更大的矩阵参数。4.5 运行方式与预期结果保存muon_optimizer.py和train_compare.py后在终端运行python train_compare.py输出大概是每个 epoch 的 loss。通常 AdamW 在前几个 epoch 下降很快后期会出现一些震荡Muon 如果学习率设置合理整体下降会更平稳。如果第一次跑不出明显差异先检查两点是否所有矩阵参数都满足p.shape[0] p.shape[1]。对于常见的nn.Linear(in_features, out_features)权重形状是(out_features, in_features)如果out_features大于in_features不满足条件就不会做正交化。Muon 的学习率是否偏大或偏小。一般从0.01到0.05之间尝试。4.6 Muon 与 Adam 在内存和计算上的区别Adam 需要维护两个状态一阶动量m和二阶动量v。Muon 只需要维护一个动量缓冲但需要在更新时做矩阵乘法迭代。对于一个(m, n)的矩阵参数Adam 的状态占用是2 * m * n个浮点数Muon 的状态占用是m * n个浮点数。在大模型场景下优化器状态占比很大Muon 可以明显节省显存。计算量方面Newton-Schur 迭代的核心是Q Q.T Q复杂度大约是 O(m^2 n)。如果矩阵很扁即 m 远小于 n这个计算开销并不大。这也是为什么 Muon 通常只对行数小于等于列数的矩阵做正交化。5. 如何改进 Muon工程实验与调优方向Muon 不是银弹。直接套用 Muon 可能遇到收敛不稳、loss 不下降等问题。下面整理一些实际工程中行之有效的改进方向。5.1 混合使用 AdamW 和 MuonMuon 适合矩阵参数但偏置、LayerNorm 参数、Embedding 参数并不一定适合正交化。常见做法是对权重矩阵使用 Muon。对偏置、LayerNorm 参数、Embedding 使用 AdamW。这样可以避免对非矩阵参数做无意义的正交化同时保留 AdamW 在精细参数上的稳定性。示例配置from torch.optim import AdamW from muon_optimizer import Muon def configure_optimizers(model): muon_params [] adamw_params [] for name, param in model.named_parameters(): if not param.requires_grad: continue if param.dim() 2 or (param.dim() 2 and param.shape[0] param.shape[1]): # 例如 Linear 权重、Conv 权重的前两维 muon_params.append(param) else: adamw_params.append(param) optim_groups [ {params: muon_params, lr: 0.02}, {params: adamw_params, lr: 1e-3, weight_decay: 0.01}, ] muon Muon(muon_params, lr0.02) adamw AdamW(adamw_params, lr1e-3, weight_decay0.01) return muon, adamw这里我故意写成了两个优化器对象。如果希望更统一也可以把所有参数放在同一个优化器中通过param_group区分。但分开两个优化器在代码上更清晰。5.2 对 Embedding 单独处理在大模型训练中Embedding 矩阵往往也很大。对 Embedding 做正交化容易破坏嵌入空间的语义结构。社区实践中通常对 Embedding 使用 AdamW而不是 Muon。如果你的模型词表很大Embedding 的更新非常稀疏更推荐使用 Adafactor 或 AdamW 这类自适应优化器。Muon 的正交化会改变整个嵌入矩阵方向对稀疏更新并不友好。5.3 调整正交化迭代步数Newton-Schur 迭代步数ns_steps是一个重要超参数。步数越多更新方向越接近正交但计算量也越大。ns_steps1时正交化效果很弱稳定性最好适合小模型。ns_steps5是常见默认值。ns_steps10更接近完全正交但可能损失一些小方向上的信息。建议先用 5观察 loss 是否稳定。如果 loss 突然升高尝试减小步数如果收敛变慢尝试增大步数。5.4 配合权重衰减Muon 的原始设计通常不包含权重衰减因为正交化更新本身具有一定的正则化效果。但如果你希望显式控制权重衰减可以借鉴 AdamW 的解耦思路在更新后单独对参数做衰减# 在每个 step 中加上 p.mul_(1 - lr * weight_decay)不过要注意权重衰减系数不能设置太高否则正交化带来的几何优势会被抵消。建议从0.01开始尝试。5.5 学习率预热与调度Muon 对学习率比较敏感。直接使用较大学习率可能在初期造成震荡。推荐使用 warmup 加 cosine decaydef get_lr(step, total_steps, warmup_steps, peak_lr): if step warmup_steps: return peak_lr * step / warmup_steps progress (step - warmup_steps) / (total_steps - warmup_steps) return peak_lr * 0.5 * (1 torch.cos(torch.pi * progress))训练初期使用较小的学习率可以让动量缓冲充分估计方向之后再提高学习率会稳定很多。5.6 与梯度裁剪结合Muon 的正交化可以部分抑制梯度爆炸但并不意味着不需要梯度裁剪。特别是在训练 Transformer 时梯度范数可能突然变大。建议设置最大梯度范数例如 1.0torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)裁剪顺序要在 optimizer.step() 之前。5.7 改进正交化的版本带缩放的正交化一些改进版本会在正交化后乘上一个缩放系数让更新方向不至于完全损失梯度幅度。例如scale torch.norm(update, dim-1, keepdimTrue) / torch.norm(orthogonalized, dim-1, keepdimTrue) update orthogonalized * scale这种做法相当于保留每个方向上的能量但仍然让方向正交。如果你发现 Muon 收敛太慢可以尝试这个思路。6. 常见问题与排查思路6.1 训练直接出现 NaN问题现象可能原因解决思路loss 变为 NaN或参数出现 NaN学习率过大降低学习率建议从 0.01 开始只有特定层出现 NaN正交化迭代步数过多减少 ns_steps混合精度训练时出现 NaN更新方向与 loss scaling 不匹配在 Muon 前先检查梯度是否有限或使用torch.cuda.amp.GradScaler定位方法在optimizer.step()前检查每个参数梯度的范数for name, param in model.named_parameters(): if param.grad is not None and not torch.isfinite(param.grad).all(): print(f梯度异常{name})6.2 收敛速度比 AdamW 还慢通常是因为正交化的方向过于“平均”导致主梯度方向信息被稀释。此时可以降低正交化步数到 2 或 3。对更新方向乘以一个缩放系数。只对最大矩阵参数使用 Muon其余使用 AdamW。6.3 内存占用没有明显下降Muon 相比 Adam 少维护一个二阶矩状态但如果你同时保存了多个优化器状态或者为了对比实验复制了多个模型内存占用依然会很高。使用前先检查显存分布nvidia-smi如果显存不足可以尝试torch.utils.checkpoint减少激活内存或者使用更小的 batch size。6.4 分布式训练中状态同步问题Muon 的动量缓冲和普通优化器一致本身就支持DDP和FSDP。但在 DDP 中梯度同步后每个进程的newton_schur计算应该得到相同结果因为输入相同。如果出现不同 rank 结果不一致通常是浮点累加顺序导致可以在update计算前添加torch.distributed.barrier()但大多数情况下不需要。6.5 与 DeepSpeed 和 FSDP 的兼容性Muon 的自定义参数分组方式可能会与 DeepSpeed 的优化器封装冲突。如果你使用deepspeed.initialize需要把 Muon 配置为自定义优化器并在配置文件中关掉默认优化器。更简单的办法是先用 PyTorch 原生的 FSDP 测试等跑通后再考虑迁移到 DeepSpeed。7. 最佳实践与工程建议7.1 不要把 Muon 用于所有参数很多初学者会写这样的代码optimizer Muon(model.parameters())这种做法通常效果不好。因为偏置、归一化参数、Embedding 并不适合正交化。更好的策略是对二维权重矩阵使用 Muon。对一维参数和 Embedding 使用 AdamW。对于卷积核可以先按(out_channels, in_channels, k_h, k_w)展平成二维后再做正交化但这一步需要谨慎直接展平可能破坏局部空间结构。7.2 记录详细的实验日志Muon 的超参数比较敏感包括学习率、momentum、ns_steps、是否使用 nesterov。建议使用wandb或tensorboard记录训练 loss 和验证 loss。梯度范数。更新方向与梯度方向的内积。正交化前后更新矩阵的奇异值分布。这些信息能帮你判断 Muon 是否正常工作。7.3 初始化条件不可忽视Muon 对初始化的要求并不比 Adam 更低。如果初始化时矩阵的奇异值分布很不平衡前几步正交化迭代可能会让更新方向剧烈变化。建议初始化时使用较小的标准差例如0.02或者在训练前先做几步 warmup。7.4 在更大模型上验证优势单机小规模实验中Muon 的优势可能不明显因为小模型的谱结构比较简单。如果你希望验证 Muon 是否适合你的业务建议至少在一个 10 层以上的 Transformer 或 5 层以上的 MLP 上测试并和 AdamW 做完整的 lr 扫描。7.5 保持代码可回滚Muon 仍在快速演进不同实现可能存在差异。建议在代码中锁定版本并保留一个 AdamW 的训练路径。当 Muon 出现奇怪问题时可以快速切换回 AdamW 继续实验。8. 总结Muon 的核心价值在于它从矩阵谱结构出发对更新方向做了正交化处理实现了 Adam 无法做到的“全局谱分配”。Adam 适合大多数通用场景但当你面对大量矩阵参数时Muon 往往能带来更平稳的收敛和更好的最终效果。本文从原理、实现、改进到排查给出了一个比较完整的入门路径。你可以先在小型 MLP 上跑通代码再逐步迁移到 Transformer 结构重点对比 loss 曲线和最终精度。不要直接跳到生产环境大规模训练先在小规模任务上积累对超参数的直觉。如果你准备在自己的模型上尝试 Muon建议把混合优化器策略、学习率 warmup、梯度裁剪这三件事一起做。它们能大大提高 Muon 的稳定性。后续还可以进一步阅读极分解、Stiefel 流形优化和自然梯度相关的资料这些知识会帮助你更灵活地调整 Muon 的行为。