2024年底，独立研究员Keller Jordan提出了一种新型神经网络优化器——Muon（全称MomentUm Orthogonalized by Newton‑Schulz）。这一优化器专为神经网络中的二维权重矩阵（hidden layers的weight matrices）设计，通过将Nesterov动量与Newton‑Schulz正交化迭代相结合，实现了对权重矩阵几何结构的精确优化。2025‑2026年间，Muon迅速在AI社区流行，尤其在安德烈·卡帕西（Andrej Karpathy）的nanoGPT、nanochat、autoresearch等极简项目中被广泛采用，成为当前“速度纪录”（speedrun）训练方案的核心组件。PyTorch 2.11版本已将其正式纳入torch.optim.Muon，标志着这一优化器进入主流框架。
Muon的诞生背景是传统优化器（如AdamW）在处理大型Transformer模型时存在的结构性缺陷。AdamW等优化器将权重矩阵视为一堆独立的标量进行元素级更新，忽略了矩阵作为几何对象的内在结构，导致梯度在不同方向上的相互干扰较大。Muon则将权重矩阵作为一个整体进行优化，通过动量更新后立即进行正交化，使更新方向保持更好的几何性质，从而提升训练稳定性和收敛速度。
Muon的PyTorch实现代码与核心机制解析
Muon的官方实现非常简洁高效。Keller Jordan的GitHub仓库（https://github.com/KellerJordan/Muon）提供了完整代码，PyTorch官方也已集成。核心思想是“先动量更新，再正交化”。
以下是简化后的核心实现（基于官方代码，PyTorch风格）：
Python
import torch
from torch.optim import Optimizer
def newtonschulz5(G, steps=5, eps=1e-7):
"""Newton-Schulz迭代实现矩阵正交化"""
assert G.ndim == 2
a, b, c = (3.4445, -4.7750, 2.0315) # Keller Jordan优选多项式系数
X = G.bfloat16()
X = X / (X.norm() + eps)
if G.size(0) > G.size(1):
X = X.T
for _ in range(steps):
A = X @ X.T
B = b * A + c * A @ A
X = a * X + B @ X
if G.size(0) > G.size(1):
X = X.T
return X
class Muon(Optimizer):
def __init__(self, params, lr=0.1, momentum=0.95, nesterov=True, steps=5):
defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, steps=steps)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
g = p.grad
if g.ndim != 2: # 只对2D权重矩阵应用Muon
continue
state = self.state[p]
if len(state) == 0:
state['momentum'] = torch.zeros_like(g)
# Nesterov动量更新
state['momentum'].mul_(group['momentum']).add_(g, alpha=1 - group['momentum'])
g = state['momentum'] if group['nesterov'] else g
# Newton-Schulz正交化
g_ortho = newtonschulz5(g, steps=group['steps'])
# 应用更新
p.data.add_(g_ortho, alpha=-group['lr'])
return loss
关键解析：
- Newton‑Schulz迭代：通过5次迭代快速将更新矩阵逼近正交矩阵（operator norm接近1），避免梯度爆炸或消失。
- 混合使用：实际项目中通常采用Muon + AdamW混合策略——对二维权重矩阵（attention和FFN层）使用Muon，对bias、embedding等1D参数使用AdamW。
- 超参数：学习率通常设为0.1左右（远高于AdamW的0.001），动量0.95，Newton‑Schulz步数5步即可收敛。
Muon与AdamW的详细对比
AdamW是目前LLM训练的事实标准优化器，其核心是元素级自适应学习率（基于梯度一阶和二阶矩）。Muon则是一种矩阵感知（matrix‑aware）优化器，两者在以下维度形成鲜明对比：
Figure 1. Optimizer comparison by sample efficiency. [reproducible logs]
1、优化对象：
- AdamW：将权重视为一堆独立标量，元素级更新。
- Muon：将2D权重矩阵视为几何整体，先动量更新再正交化。
2、收敛速度：
- 实验显示，Muon在nanoGPT、CIFAR‑10 speedrun等任务中可将训练速度提升1.35x~2x。
- 在LLM训练中，Muon + AdamW混合方案往往能以更少步数达到相同loss，部分场景下节省约50%计算成本。
3、内存占用：
- Muon仅维护一个动量状态，内存开销比AdamW低约33%。
- 混合方案在保持性能的同时，进一步降低峰值内存。
4、稳定性与泛化：
- Muon通过正交化减少了谱偏差（spectral bias），使更新更均匀，尤其在重尾分布（heavy‑tailed）任务中表现更好。
- AdamW在长训练中容易出现梯度尖峰，而Muon的更新更平滑。
5、适用范围：
- AdamW通用性强，适合所有参数。
- Muon专为2D权重设计，与AdamW混合使用效果最佳（embedding/bias用AdamW，hidden weights用Muon）。
在具体项目中的使用方式
1、Karpathy的autoresearch项目：
- 在train.py中，Muon被用作默认优化器之一。
- 代理可自主修改优化器配置（Muon或AdamW），在固定5分钟预算内运行实验，通过val_bpb指标判断改进效果。
- 典型代码片段：
Python
from muon import Muon optimizer = Muon(model.parameters(), lr=0.1, momentum=0.95)
2、nanoGPT / nanochat：
- 在Karpathy的极简GPT训练项目中，Muon + AdamW混合优化器帮助打破多项速度纪录。
- 通常对QKV和FFN权重矩阵单独应用Muon，对embedding和bias使用AdamW。
3、大规模LLM训练：
- Moonshot、Kimi等团队已在千亿参数模型中验证Muon的可扩展性。
- 混合方案已成为当前最优实践：Muon负责隐藏层几何优化，AdamW负责其余参数的自适应调整。
实际效果与社区反馈
在nanoGPT speedrun中，切换到Muon后训练速度提升约35%。在CIFAR‑10等基准上，Muon也展现出显著优势。社区普遍认为，Muon是自Adam以来优化器领域最重要的创新之一，尤其适合Transformer架构的LLM训练。
总结与展望
Muon优化器以其矩阵感知设计和高效正交化机制，为LLM训练提供了全新选择。它不是完全取代AdamW，而是与之形成高效互补。通过在Karpathy等极简项目中的落地，Muon已证明其在实际训练中的价值。随着PyTorch官方集成和更多大规模验证，Muon有望成为下一代AI训练的标准组件之一。
从元素级优化到矩阵级几何优化，Muon代表了优化器从“通用”向“结构感知”的演进方向。在AI训练成本居高不下的今天，这一创新不仅加速了研究迭代，也为更高效、更绿色的AI基础设施奠定了基础。
▶️ 喜欢这篇内容？创作真的很辛苦！
欢迎关注我们的 YouTube 频道：
AI Rosetta
你的点赞、订阅和分享，是对我们最大的支持！🙏
更多AI新闻请访问：
https://www.airosetta.com