Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

本文提出了名为 Mousse 的新型优化器，它通过结合 Shampoo 的曲率感知预条件技术与 Muon 的谱优化方法，在黎曼流形上实现了各向异性的信任区域约束，从而在几乎不增加计算开销的情况下显著提升了语言模型的训练效率。

要点

这篇论文介绍了一种名为 Mousse（慕斯）的新算法，它是用来训练大型人工智能模型（比如大语言模型）的“加速器”。

为了让你轻松理解，我们可以把训练 AI 模型想象成在崎岖不平的山地里寻找最低点（山谷）的过程。

1. 背景：之前的登山者（Muon）

在 Mousse 出现之前，有一种很流行的登山方法叫 Muon。

• 它的做法：Muon 就像是一个严格守纪律的登山队。它规定，无论你在哪个方向走，每一步的“最大跨度”必须完全一样（就像给每个人发了一根长度固定的登山杖）。
• 优点：这种纪律性让队伍走得很稳，不会乱跑，训练速度很快。
• 缺点：现实中的山（AI 的损失函数地形）是非常复杂的。有的地方是陡峭的悬崖（高曲率），有的地方是平坦的草地（低曲率）。
  • Muon 的“一刀切”策略有个大问题：在陡峭的悬崖边，它不敢迈大步（怕摔死）；在平坦的草地上，它又不敢迈大步（因为规则限制了最大步长）。结果就是，该快的时候快不起来，该稳的时候又太保守。

2. 主角登场：Mousse（慕斯）

Mousse 的名字来源于一种甜点（慕斯），寓意它像慕斯一样，既保留了 Muon 的“结构”（像慕斯的骨架），又加入了“顺滑”的曲率感知（像慕斯的口感）。

• 核心创新：Mousse 发现，Muon 的问题在于它是在“错误的地图”上走路的。它没有考虑到山地的曲率（哪里陡、哪里平）。
• Mousse 的做法：
  1. 先“熨平”地形：Mousse 会先利用一种叫 Shampoo 的技术，给地图做一个“预处理”。想象一下，它把陡峭的悬崖“压平”，把平坦的草地“垫高”，让原本崎岖的山地瞬间变成平坦的球面。
  2. 再执行 Muon 的规则：在这个被“熨平”后的新地图上，Mousse 再让 Muon 那个“严格守纪律”的登山队出发。
  3. 最后还原：走完后，再把步长还原回真实的地形。

通俗比喻：
想象你在玩一个迷宫游戏。

• Muon 是拿着一个固定长度的尺子在走，不管前面是墙还是空地，它都只能走尺子那么长。
• Mousse 则是先拿一个魔法透镜（Shampoo 预处理）看迷宫。透过这个透镜，所有的墙都变矮了，所有的坑都填平了，迷宫变得像操场一样平坦。
• 然后，Mousse 依然拿着那把固定长度的尺子（Muon 的纪律）在平坦的操场上走。
• 因为操场是平的，尺子能发挥最大效率；等走完了，再透过透镜看回原来的迷宫，你会发现你其实已经走了很远很远的路，而且避开了所有悬崖。

3. 为什么 Mousse 更厉害？

论文通过实验证明，Mousse 比 Muon 强在两点：

1. 跑得更快（省步数）：
   在训练 8 亿参数（800M）的模型时，Mousse 比 Muon 少用了约 12% 的步数就能达到同样的效果。

   • 比喻：如果 Muon 需要走 100 步下山，Mousse 只需要走 88 步就能到山脚，而且走得更稳。

2. 不费力气（低开销）：
   以前有些高级算法（比如 SOAP）虽然也能感知地形，但它们太“重”了，计算量大，像背着大石头爬山。

   • Mousse 很聪明，它只增加了3%的计算时间，却换来了巨大的性能提升。它就像给登山队配了一副轻便的护目镜，而不是让他们背个沉重的背包。

4. 关键技巧（如何让它稳定）

为了让这个“魔法透镜”不失效，作者还加了两个小补丁：

• 迹归一化 (Trace Normalization)：就像给地图做“标准化”，防止某些层的数据太大或太小，导致透镜变形。
• 谱调节 (Spectral Tempering)：就像调节透镜的“焦距”。如果焦距太猛（过度矫正），反而会把路看歪；Mousse 发现把焦距调得柔和一点（参数设为 0.125 而不是 0.25），效果反而最好。

总结

Mousse 就是给原本就很优秀的 Muon 算法，加上了**“地形感知眼镜”。
它不再盲目地用同一种步长去走所有路，而是先看清哪里陡、哪里平，把路“修平”后再走。这使得 AI 模型训练得更快、更稳、更省资源**。对于想要训练大模型的公司来说，这意味着能省下大量的时间和电费。

技术摘要

论文概览

标题：Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning
核心主题：提出了一种名为 Mousse 的新型优化器，旨在解决当前流行的谱优化方法（如 Muon）在深度神经网络中因假设各向同性（Isotropic）优化景观而导致的次优问题。Mousse 通过结合 Shampoo 的结构化曲率估计与 Muon 的谱约束，实现了在保持计算高效的同时，显著提升大语言模型（LLM）的训练收敛速度和最终性能。

---

1. 问题背景 (Problem)

• 现有方法的局限性：
  • Muon 优化器：近期在大规模预训练中表现出色，它通过将更新步长限制在 Stiefel 流形上（利用 Newton-Schulz 迭代进行正交化），实现了谱优化。然而，Muon 隐含地假设优化景观是各向同性的（即所有特征方向的曲率相同），强制所有特征方向上的更新范数均匀。
  • 实际挑战：深度神经网络的损失景观实际上是高度各向异性且病态的（ill-conditioned），其曲率谱呈现重尾分布（heavy-tailed）。
  • 后果：Muon 的“平等主义”约束在曲率高的方向上可能放大不稳定性，而在平坦方向上又限制了必要的进展，导致次优的收敛效果。
• 现有改进的不足：虽然已有工作尝试改进 Muon，但尚未有效解决如何将 Muon 的各向同性约束与神经网络高度各向异性的曲率景观对齐这一核心几何矛盾。

---

2. 方法论 (Methodology)

Mousse (Muon Optimization Utilizing Shampoo's Structural Estimation) 的核心思想是：在“白化”的坐标系统中执行谱优化。

2.1 几何框架重构

• 白化变换：Mousse 首先利用 Shampoo 的 Kronecker 分解曲率统计量（$L$ 和 $R$ 矩阵）对梯度进行预处理（白化）。这相当于将局部优化景观“球化”（sphering），使其在该坐标系下近似各向同性。
• 约束优化问题：
  • 传统 Muon 求解的是：$\min \langle G, U \rangle$ s.t. $\|U\|_{op} \le 1$（在原始欧氏空间）。
  • Mousse 求解的是：$\min \langle G, U \rangle$ s.t. $\|P U Q\|_{op} \le 1$，其中 $P=L^{1/4}, Q=R^{1/4}$ 是白化因子。
• 数学推导：
  1. 定义白化后的梯度 $\tilde{G} = P^{-1} G Q^{-1}$。
  2. 在白化坐标系中，问题转化为标准的 Muon 形式：$\min \text{Tr}(\tilde{G}^T Y)$ s.t. $\|Y\|_{op} \le 1$。
  3. 解为 $Y = -\text{msign}(\tilde{G})$（通过 Newton-Schulz 迭代近似）。
  4. 最终更新步长 $\Delta W = -L^{-1/4} \text{msign}(L^{-1/4} G R^{-1/4}) R^{-1/4}$。

2.2 关键工程技巧 (Engineering Insights)

为了稳定二阶谱优化，论文提出了两项关键技术：

1. 迹归一化 (Trace Normalization)：
   • 问题：不同层的 $L$ 和 $R$ 矩阵绝对数值差异巨大，导致统一的阻尼项 $\epsilon$ 难以生效。
   • 方案：在特征分解前，将协方差矩阵归一化，使其迹（Trace）等于维度（即平均特征值为 1）。这确保了阻尼项在所有模块中具有相对一致的效应。
2. 谱调温 (Spectral Tempering)：
   • 问题：标准的 Shampoo 指数 $\alpha=0.25$ 对于 Mousse 过于激进，会导致平坦方向上的梯度被过度放大。
   • 方案：发现使用更温和的指数 $\alpha=0.125$ 效果更佳。这平衡了谱优化的各向同性稳定性与二阶方法的各向异性加速能力。
3. 梯度嫁接 (Gradient Grafting)：
   • 为了维持更新步长的幅度稳定，防止 RMS 范数随训练衰减，Mousse 采用了梯度嫁接技术，将更新幅度与优化方向解耦。
4. 单侧预条件 (Single-Sided Preconditioning)：
   • 实验发现仅使用左侧预条件器（$L$，对应激活统计量）即可达到与双侧几乎相同的性能，且能显著降低计算和内存开销。

---

3. 主要贡献 (Key Contributions)

1. 统一的几何框架：从理论上证明了 Mousse 是各向异性几何下双范数最大化问题的最优解。通过引入 Kronecker 分解的曲率统计量，成功弥合了谱方法（Muon）与二阶预条件器（Shampoo）之间的鸿沟。
2. 稳健的工程实践：系统分析了谱约束与重尾曲率估计结合时的稳定性挑战，提出了迹归一化、谱调温等关键技巧，为大尺度二阶谱优化提供了实用指南。
3. 帕累托最优的效率：在 160M 到 800M 参数量的语言模型上进行了广泛实验。Mousse 在显著降低训练步数的同时，仅增加了极小的计算开销，确立了新的预训练效率基准。

---

4. 实验结果 (Results)

• 收敛速度与最终性能：
  • 在 800M 模型上，Mousse 相比 Muon 减少了约 12% 的训练步数即可达到相同的验证损失。
  • 在相同的训练步数下，Mousse 的最终验证损失比 Muon 低约 0.012，且优于 SOAP 和 AdamW。
  • 在不同模型规模（160M - 800M）和不同学习率下，Mousse 均表现出一致的优势。
• 计算效率与内存：
  • 训练速度：Mousse 的训练速度几乎与高效的 Muon 相同（仅增加约 3% 的墙钟时间），远快于 SOAP。
  • 内存占用：由于不需要像 Adam 类优化器那样维护完整的二阶动量状态（$v$），且采用单侧预条件技术，Mousse 的峰值内存占用约为 SOAP 的 88%，与 Muon 相当（约 1.05 倍 Muon）。
• 可扩展性：随着模型规模增大，Mousse 的优势依然保持，验证了其在大模型预训练中的可扩展性。

---

5. 意义与影响 (Significance)

• 理论突破：揭示了谱优化方法（如 Muon）在原始参数空间直接应用的几何缺陷，并提出了通过白化坐标变换来修正这一缺陷的理论框架。
• 工业界价值：为大规模语言模型的预训练提供了一种**“低成本、高收益”**的优化方案。它保留了 Muon 的内存效率和收敛稳定性，同时通过引入曲率感知显著提升了样本效率（Sample Efficiency）。
• 未来方向：Mousse 的架构天然兼容 Shampoo 领域的最新进展（如 Kronecker-Sum 预条件器），为未来进一步优化大模型训练提供了灵活的接口。

总结：Mousse 通过“先白化，后谱约束”的策略，成功解决了 Muon 在病态优化景观中的几何失配问题，在不显著增加计算成本的前提下，实现了比当前 SOTA 优化器（Muon, SOAP, AdamW）更优的训练效率和最终模型性能。