Muon$^p$：具有分数光谱功率的 Muon

Muon 是一种越来越广泛使用的优化器，它用极因子 $UV^\top$ 代替梯度 $G=USV^\top$，从而平坦化奇异光谱。然而，完全扁平化会丢弃可能对适应很重要的奇异值信息。我们引入 Muon$^p$，一种 Muon 风格的优化器，它使用分数谱功率更新 $US^pV^\top$ 来实现有理 $p\in(0,1)$，在 Muon 和梯度下降之间进行插值。为了使其实用，我们证明分数谱功率不能通过任何固定的单变量多项式迭代来计算，并且进一步导出仅使用矩阵乘法近似 $US^pV^\top$ 的低次奇二元递归，保留 Muon 的仅矩阵乘法结构和计算复杂性。我们表明，对于 $q=1+\frac{1}{p}$，Muon$^p$ 在 Schatten $q$-范数下最大化损失的线性改进。根据经验，Muon$^p$ 对于微调特别有效：在十亿级模型上，Muon$^p$ 提高了验证复杂性和下游任务性能。我们通过光谱几何的角度进一步分析什么时候 Muon$^p$ 不太合适。我们的结果揭示了关于何时保留奇异光谱可以带来显着收益的重要见解，并介绍了实现这些收益的原则性方法。