企业官网建设流程全解析

最近在训练大模型和深层网络时，Muon 这个优化器频繁出现在社区讨论里。不少实验显示，在相同的参数量、相同的 token 数下，Muon 的收敛曲线明显比 AdamW 更平滑，最终 loss 也更低。与此同时，“Spectral Allocation”这个词开始和 Muon 绑定出现，很多人第一次看到时会误以为它只是某种学习率调整策略，其实它更像是一种针对矩阵参数特性的更新方向分配机制。

这篇文章会从一个比较容易理解的角度，拆解 Muon 与 Adam 的本质区别，解释为什么 Muon 在矩阵参数较多的网络结构中能胜过 Adam，并给出可以跑起来的 PyTorch 实现和一组改进方向。适合正在做模型训练调参的开发者，也适合想深入理解优化器原理的初学者。

1. 背景：从 SGD 到 Adam，再到 Muon

1.1 优化器到底在优化什么

神经网络的训练过程，本质上是在不断求解一个高维非凸优化问题。每一步，我们都需要根据当前参数计算梯度，并根据梯度方向更新参数。优化器负责回答两个问题：更新方向是什么，以及每一步走多远。

SGD 是最朴素的选择，它直接用梯度方向作为更新方向，用学习率控制步长。它的问题在于梯度在不同维度上的尺度差异很大，如果直接使用统一学习率，某些维度会收敛很慢，某些维度可能会震荡。Adam 的核心改进，就是利用梯度的一阶矩和二阶矩估计，对每个参数坐标做归一化，相当于给每个维度一个独立的步长。这种方式在大多数情况下非常有效，尤其是训练 Transformer 这类深层模型。

但是，Adam 有一个隐藏的局限：它的归一化是针对“单个参数元素”的，而神经网络中的参数往往是以矩阵形式组织的，例如线性层权重、卷积核、注意力投影矩阵。这些矩阵本身具有谱结构（奇异值分布、条件数等），Adam 完全没有利用这种结构。

1.2 什么是 Spectral Allocation

Spectral 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 的思路可以拆成几个步骤：

1. 计算梯度G。
2. 用动量累积梯度，得到动量矩阵U。
3. 对U做正交化处理，使其趋向于一个正交矩阵。
4. 用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.04，Windows WSL2 也可以。
• Python：3.9 或 3.10。
• PyTorch：1.13 或 2.x 都可以，推荐 2.1 以上。
• CUDA：11.8 或 12.x，不强制 GPU，纯 CPU 也能跑通逻辑。

建议创建一个独立的虚拟环境：

conda create -n muon-test python=3.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(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps) super().__init__(params, defaults) @torch.no_grad() def step(self, closure=None): 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_samples=4096, batch_size=128): 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_size=batch_size, shuffle=True) return loader def train(model, optimizer, loader, epochs=20): 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(f"epoch {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(), lr=1e-3, weight_decay=0.01) print("===== AdamW =====") losses_adam = train(model_adam, optimizer_adam, loader) model_muon = build_model() optimizer_muon = Muon(model_muon.parameters(), lr=0.02, momentum=0.95, ns_steps=5) 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 和 Muon

Muon 适合矩阵参数，但偏置、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, lr=0.02) adamw = AdamW(adamw_params, lr=1e-3, weight_decay=0.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_steps=1时，正交化效果很弱，稳定性最好，适合小模型。
• ns_steps=5是常见默认值。
• ns_steps=10更接近完全正交，但可能损失一些小方向上的信息。

建议先用 5，观察 loss 是否稳定。如果 loss 突然升高，尝试减小步数；如果收敛变慢，尝试增大步数。

5.4 配合权重衰减

Muon 的原始设计通常不包含权重衰减，因为正交化更新本身具有一定的正则化效果。但如果你希望显式控制权重衰减，可以借鉴 AdamW 的解耦思路，在更新后单独对参数做衰减：

# 在每个 step 中加上 p.mul_(1 - lr * weight_decay)

不过要注意，权重衰减系数不能设置太高，否则正交化带来的几何优势会被抵消。建议从0.01开始尝试。

5.5 学习率预热与调度

Muon 对学习率比较敏感。直接使用较大学习率可能在初期造成震荡。推荐使用 warmup 加 cosine decay：

def 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.0：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

裁剪顺序要在 optimizer.step() 之前。

5.7 改进正交化的版本：带缩放的正交化

一些改进版本会在正交化后乘上一个缩放系数，让更新方向不至于完全损失梯度幅度。例如：

scale = torch.norm(update, dim=-1, keepdim=True) / torch.norm(orthogonalized, dim=-1, keepdim=True) update = orthogonalized * scale

这种做法相当于保留每个方向上的能量，但仍然让方向正交。如果你发现 Muon 收敛太慢，可以尝试这个思路。

6. 常见问题与排查思路

6.1 训练直接出现 NaN

定位方法：在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 的行为。