想象一下,你正试图将一块沉重且形状怪异的巨石推上山坡。这块巨石代表了训练计算机大脑(神经网络)时的“损失”(误差)。你的目标是尽快将巨石推至谷底(完美解)。
长期以来,实现这一目标的标准方法是SGD(随机梯度下降)。将 SGD 想象为一群推巨石的人。他们观察坡度,朝着看起来最陡峭的方向推。然而,这块巨石形状怪异:它有一侧极其光滑且陡峭,而另一侧则较为平缓。
旧方法(SGD)的问题:
由于巨石的一侧过于光滑,团队必须非常轻柔地推。如果推得太用力,他们就会失控地滑向错误的一侧,导致整个操作崩溃。这意味着他们必须采取微小而谨慎的步伐,使得旅程非常缓慢。
新英雄:Muon
这篇论文介绍了一种名为Muon的新方法。Muon 不仅仅是在推,它更像是一位神奇的机械师,在每次推动之前都能重塑巨石。
以下是 Muon 的工作原理,基于论文的核心概念:
1. 谱平坦化(Spectral Flattening): “整平”技巧
论文将 Muon 的制胜法宝称为**“谱平坦化”**。
- 类比:想象巨石有一个“尖刺”侧(非常陡峭且危险的方向)和几个“平坦”侧。在数学术语中,这些被称为“奇异值”。尖刺侧是问题所在,因为它迫使团队走得很慢。
- Muon 的做法:Muon 使用一种数学工具(牛顿 - 舒尔茨迭代)来“整平”那个尖刺侧。它并不改变团队需要前往的方向;它只是让陡峭的坡度变得不那么陡,让平缓的坡度变得稍微陡一些。它将一块参差不齐、凹凸不平的岩石变成了一颗光滑圆润的球体。
- 结果:因为岩石现在变得圆润且平衡,团队可以推得更用力而不会打滑。他们可以迈出巨大的步伐(更大的学习率),而这些步伐若用旧方法会导致立即崩溃。
2. 为何更快:“平均”与“最坏”
论文证明了两点关于其为何有效的主要事实:
速度限制:
- SGD受限于巨石最坏的部分(最陡峭的尖刺)。如果尖刺比其余部分陡峭 100 倍,团队必须以 1/100 的速度行走。
- Muon受限于平均陡峭程度。通过整平尖刺,Muon 忽略了“最坏情况”,转而关注平均值。这使得团队能够以更接近其真实潜力的速度行走。
收敛(到达底部):
- 论文表明,Muon 不仅仅是迈更大的步子;它迈的是更好的步子。它充当了“预条件器”,这是一种 fancy 的说法,意指它调整地形,使每一步都能更高效地将巨石推向目标。
- 即使你强迫旧团队(SGD)和新团队(Muon)以完全相同的速度行走,Muon 仍然能更快地到达谷底,因为它的路径更直接,且更少摇摆。
3. 实验:证明其有效性
作者在标准的计算机视觉任务(从 CIFAR-10 数据集中识别猫、狗等图像)上测试了这一点。
- “跳水”测试:他们尝试以极高的速度推巨石。旧团队(SGD)立即从悬崖上跌落(发散),训练失败。新团队(Muon)则保持平稳行走并达到了目标。
- “赛跑”测试:当他们强迫两个团队以相同的中等速度行走时,Muon 仍然赢得了比赛,比旧团队提前了几个“圈”(epoch)达到了高准确率分数。
4. 关于“归一化”的新见解
论文还发现了常用工具“批归一化”(Batch Normalization)为何有效的原因。
- 发现:作者意识到,批归一化之所以有效,是因为它自然地“整平”了输入数据,类似于 Muon 对巨石所做的操作。
- 新想法:他们提出了一条设计这些工具的新规则:不仅要平衡数值,还要平衡数据的“形状”。如果你能确保输入数据没有某个“超级陡峭”的方向,你就可以用更大的步伐、更快地训练计算机大脑。他们测试了一种名为"FrobNorm"的新方法,该方法正是如此,结果发现它允许在极高的速度下进行训练,而其他方法在此速度下会失败。
总结
简而言之,Muon是一种更智能的 AI 训练方式。它在每一步之前,将问题中混乱、不平坦的地形“整平”。这防止了训练被困在陡峭、危险的山坡上,使 AI 能够学得更快、迈出更大的步伐,并比以前的方法更可靠地到达解决方案。
技术摘要:谱平坦化是 Muon 所需的一切
问题陈述
训练现代神经网络涉及在高度非凸的景观上进行优化,该景观的特征包括平坦区域、尖锐方向和鞍点。优化器的选择对稳定性和效率至关重要。虽然像 Adam 和 AdamW 这样的坐标自适应方法提高了鲁棒性,但它们并未完全解决矩阵值参数所带来的几何挑战。最近,Muon 作为一种高性能替代方案出现,它在更新前对动量缓冲区进行正交化,从而允许使用更大的学习率并实现更快的收敛。然而,从理论上解释 Muon 为何在最大稳定步长和收敛率方面优于标准梯度下降(SGD)的机制仍未得到充分探索。具体而言,现有理论尚未定量解释正交化如何改变学习率瓶颈或优化问题的有效条件数。
方法论
本文通过将 Muon 的更新规则表述为一种谱平坦化形式,提供了对其的理论分析。作者提出了两个主要的理论结果,并通过实验进行了验证:
理论框架:
- Muon 定义: Muon 通过计算梯度 Gt,然后应用牛顿 - 舒尔茨(Newton–Schulz)迭代来近似极因子 Ot≈(GtGt⊤)−1/2Gt=UV⊤,从而更新权重 W。该操作保留了梯度的奇异向量,同时将奇异值替换为 1。
- 海森矩阵近似: 分析利用了一种克鲁内克分解的曲率模型(K-FAC),将海森矩阵近似为 H≈X⊤X⊗GtGt⊤,其中 X 代表输入激活值,Gt 代表梯度。
- 学习率分析: 作者推导了 SGD 和 Muon 的精确单步下降阈值。他们比较了两种优化器的最大稳定学习率(ηmax),表明 SGD 受限于梯度的最大奇异值(σmax),而 Muon 的稳定性则由平均奇异值(m1∑σi)决定。
- 收敛性分析: Muon 被重新表述为一种带有预条件子 P=In⊗(GtGt⊤)−1/2 的预条件梯度方法。在 Polyak–Łojasiewicz (PL) 条件和相对平滑性下,作者分析了收敛因子。他们证明,当梯度协方差是各向异性时,Muon 的有效条件比(α~/β~)严格优于 SGD 的条件比(α/β),具体改进程度与梯度的谱展宽相关。
实验设计:
- 设置: 实验在 CIFAR-10 上使用 CifarNet 进行。为了隔离优化器对矩阵值参数的影响,采用了一种双优化器策略:卷积核由目标优化器(Muon 或 SGD)更新,而偏置和其他参数使用固定的参考优化器。
- 稳定性测试: 在不使用批量归一化(BN)的情况下进行学习率扫描,以观察内在稳定性。
- 收敛性测试: 两种优化器在相同的学习率下(以及分别在各自的最佳速率下)运行,以测量收敛速度和每步收敛比(rt)。
- 归一化原理: 提出并测试了一种新的归一化方案,即FrobNorm(通过 Frobenius 范数进行归一化),以验证理论主张:降低输入谱尺度(λmax(X⊤X))能够实现更高的学习率。
主要贡献
最大学习率推导: 本文证明了 Muon 的最大稳定步长与梯度的平均奇异值成比例,而非最大奇异值。
- 对于 SGD:ηmax∝λmax(H)2≈λmax(X⊤X)σmax(G)2。
- 对于 Muon:ηmax∝λmax(H)2m∑σi≈λmax(X⊤X)σmax(G)2mσmax(G)∑σi。
- 这表明谱平坦化防止了单个主导奇异方向决定步长约束,使 Muon 能够容忍显著更大的学习率。
收敛率加速: 作者表明,Muon 充当了一种预条件子,改善了有效条件比。
- 在 K-FAC 下,Muon 的收敛因子为 1−β~α~,其中 β~α~=βαλmax(GtGt⊤)λmin(GtGt⊤)。
- 由于 λmaxλmin<1,项 … 小于 1,但推导表明在预条件系统的上下文中 β~α~>βα,从而在梯度协方差病态时导致严格更快的线性收敛率。
新的归一化原理: 分析揭示,最大稳定学习率与 λmax(X⊤X) 成反比。作者提出,归一化层应旨在降低此谱尺度。他们通过 FrobNorm 验证了这一点,表明它使得在标准训练失败的高学习率下(例如 0.25)进行训练成为可能,特别是对于 SGD。
实验结果
- 稳定性: 在没有批量归一化的情况下,SGD 在 Muon 保持稳定并继续降低损失的学习率下,在前几次迭代内即发散。Muon 表现出参数和梯度范数的受控增长,而 SGD 则表现出灾难性的爆炸。
- 收敛速度: 在相同的学习率下(例如 η=0.05),Muon 比 SGD 提前几个 epoch 达到精度里程碑(70%–91%)。经验收敛比 rt 对于 Muon 始终更低,证实了更快线性速率的理论预测。
- 归一化影响: FrobNorm 的应用使得两种优化器都能在高学习率下稳定训练(CifarNet 为 0.25,GPT-2 风格的 Transformer 为 0.2),验证了控制层输入谱尺度是实现更大步长的通用原理这一假设。
意义与主张
本文声称提供了 Muon 经验成功的原则性几何解释。它提出,Muon 的优势不仅仅是一种启发式方法,而是谱平坦化的直接结果,该结果:
- 将学习率约束与最大奇异值解耦,允许更大的步长。
- 充当单侧预条件子,改善优化景观的有效条件数,从而实现更快的收敛。
作者承认了局限性,指出他们的理论分析假设了确定性全批量设置和精确的极因子,而实际实现使用的是随机梯度、有限牛顿 - 舒尔茨迭代和动量。他们表示,将这些结果扩展到随机设置和更大规模的架构仍然是未来工作的重要方向。本文 concludes 这些结果有助于将 Muon 的经验性能转化为更精确的理论理解,可能指导未来优化器和归一化策略的设计。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。