深入探讨 Muon 及其扩展

Muon 优化器最近因其强大的经验性能和对矩阵形状参数使用的正交化更新而受到广泛关注，但其背后的机制及其与自适应优化器如 Adam 的关系仍不充分理解。本文通过统一的谱方法来回应这些问题。具体而言，我们将 Muon 视为形式 $U \boldsymbol{\Sigma}^{p} V'$ 的 $p=0$ 端点，考虑额外的变体，其中 $p=1/2$, $p=1/4$, 和 $p=1$。这些变换分别应用于动量 SGD 中的一阶矩更新以及 Adam 中的根均方（RMS）归一化梯度更新。为实现高效计算，我们发展了一种避免显式奇异值分解的耦合牛顿迭代。通过控制实验，我们发现 RMS 归一化更新比一阶矩更新更稳定。此外，尽管谱压缩在一阶矩更新下提供了强大的稳定性优势，Muon 更新（$p=0$）并不一致地优于 Adam。这些结果表明，Muon 最佳地理解为一种谱归一化形式，但并非通用性更优的优化方法。我们的源代码将发布于 https://github.com/Ocram7/BeyondMuon。