Muon优化器如何以谱分配视角超越Adam：原理、实现与改进-平芜编程栈

这次我们来看一个跟训练大模型直接相关的优化器问题：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 在大模型上几乎成了默认配置。

再看带动量的 SGD：

momentum_buffer = momentum * momentum_buffer + g param -= lr * momentum_buffer

SGD 的更新方向完全由梯度主导，没有逐坐标归一化，因此对大尺度差异敏感，需要精细调学习率。

从这三个优化器可以提炼出一条主线：优化器到底怎么做“归一化”。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, k=5) W_{t+1} = W_t - η · O_t

β 通常取 0.95。也可以加 Nesterov 加速，用M_t * β + G_t作为实际正交化的输入。

关键在 Newton-Schulz 迭代。它不需要做 SVD，只靠矩阵乘法就能把输入矩阵投影到正交矩阵附近。

def zeropower_via_newtonschulz(G, steps=5, eps=1e-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 的一种应用：只有矩阵参数才谈得上谱，向量参数就用坐标归一化。

下面用一个表直接对比三种优化器的谱分配方式：

对比维度	SGD	Adam / AdamW	Muon
更新依据	原始梯度	逐元素归一化梯度	动量矩阵正交化
谱方向处理	按奇异值权重分配	坐标白化，谱方向被混合	均匀分配
一维参数	同样适用	自适应学习率	退回 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, steps=5) 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 参数：继续用 AdamW；
attention 的 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, steps=5, eps=1e-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, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.01): defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps, betas=betas, eps=eps, weight_decay=weight_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, steps=group["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, alpha=1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - 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(), lr=0.02, momentum=0.95, weight_decay=0.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_in=64, d_hidden=256, d_out=16): 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_size=32, n_batches=200, seed=0): 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(), lr=lr) 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, topk=8): """取出某个矩阵参数对应的更新量，返回前 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_matrices=False) return s[:topk].detach().cpu().numpy()

对同一模型分别用 AdamW 和 Muon 跑一步，输出奇异值。通常可以看到：AdamW 的更新矩阵奇异值可能前两个特别大、后面快速衰减；Muon 的更新矩阵所有奇异值都接近 1。这就是“谱分配”差异的直接证据。