Muon$^p$: 分数谱幂的Muon优化器

# Muonp：具有分数谱次数的 Muon

来源：https://arxiv.org/html/2606.13867

###### 摘要

Muon 是一种日益广泛使用的优化器，它将梯度 `G=USV^⊤` 替换为其极因子 `UV^⊤`，从而压平了奇异谱。然而，完全压平丢弃了可能对自适应至关重要的奇异值信息。我们引入了 Muonp，这是一种 Muon 风格的优化器，它使用分数谱次数更新 `US^p V^⊤`（其中有理数 `p∈(0,1)`），在 Muon 和梯度下降之间进行插值。为了使其实用，我们证明了任何固定的单变量多项式迭代都无法计算分数谱次数，并进一步推导出低次奇数双变量递推关系，这些递推关系仅使用矩阵乘法即可近似 `US^p V^⊤`，从而保留了 Muon 仅涉及矩阵乘法的结构和计算复杂度。我们证明 Muonp 在 `q=1+1/p` 的 Schatten q-范数下最大化损失的线性改进。实验上，Muonp 在微调场景中尤其有效：在十亿级模型上，Muonp 提高了验证困惑度和下游任务性能。我们进一步通过谱几何的视角分析了 Muonp 不太适用的情形。我们的结果揭示了保留奇异谱可以在何时带来显著收益的重要见解，并介绍了一种实现这些收益的原则性方法。我们的代码已在 https://github.com/princeton-pli/muon-p 公开。

## 1 引言

优化算法在塑造神经网络的学习方式中发挥着核心作用，影响着训练的效率以及最终模型的质量 [23 (https://arxiv.org/html/2606.13867#bib.bib23), 36 (https://arxiv.org/html/2606.13867#bib.bib36), 22 (https://arxiv.org/html/2606.13867#bib.bib22)]。Muon 已成为一种有前景的深度网络优化器，因为其谱压平更新促进了跨奇异方向的更平衡学习，防止主导模式垄断训练，并已被证明在 LLM 预训练中有效缩放 [22 (https://arxiv.org/html/2606.13867#bib.bib22), 32 (https://arxiv.org/html/2606.13867#bib.bib32), 54 (https://arxiv.org/html/2606.13867#bib.bib54), 44 (https://arxiv.org/html/2606.13867#bib.bib44), 51 (https://arxiv.org/html/2606.13867#bib.bib51)]。Muon 通过将梯度 `G=USV^⊤` 替换为其极因子 `UV^⊤` 来实现这一点，这会压平奇异谱并使更新在奇异方向上更均匀 [22 (https://arxiv.org/html/2606.13867#bib.bib22)]。但完全压平也丢弃了奇异值信息，这引出了一个自然问题：我们能否保留 Muon 在矩阵层面的谱优势，同时保留一些对自适应重要的幅度结构？

我们用 `Muon^p` 来回答这个问题，这是一种 Muon 风格的优化器，对于有理数 `p∈(0,1)`，使用 `US^p V^⊤` 进行更新。这一族具有清晰的几何解释：`US^p V^⊤` 是在 `q=1+1/p` 的 Schatten q-范数下的最速下降方向（定理 2.1 (https://arxiv.org/html/2606.13867#S2.Thmtheorem1)）。我们识别并解决了推导 Muonp 并使其实现实用的关键挑战。与 Muon 的极因子不同，分数谱次数无法通过迭代一个固定的单变量多项式获得；任意有理数次幂需要一个保留原始矩阵的递推关系（引理 2.2 (https://arxiv.org/html/2606.13867#S2.Thmtheorem2)）。因此，我们推导了低次奇数双变量递推关系，保留了 Muon 仅涉及矩阵乘法的结构（定理 2.3 (https://arxiv.org/html/2606.13867#S2.Thmtheorem3)），从而得到对 Muon 实现仅需一行代码更改的实用更新（算法 1 (https://arxiv.org/html/2606.13867#alg1)）。虽然先前的工作 [42 (https://arxiv.org/html/2606.13867#bib.bib42)] 已经提出了针对特定 p 值的分数次幂，但我们从第一性原理出发引入了一个系统的公式，并明确推导了适用于*任意*分数次幂 p 的 Muonp。实验上，Muonp 在微调中最为有益。在各种模型上，Muonp 改善了数学和代码推理、验证困惑度以及下游任务性能。我们进一步分析了 Muonp 何时不如 Muon 适用的情形。综上所述，这些结果将 Muonp 定位为 Muon 的一个原则性推广，同时表明正确的谱几何是依赖于场景的。

##### 贡献。我们的主要贡献是：

- • 我们引入了 Muonp，一种通用的 Muon 风格优化器，计算*任意*有理数 `p∈(0,1)` 的分数谱次数更新 `US^p V^⊤`，而无需显式的 SVD。
- • 我们证明没有固定的单变量多项式迭代可以计算这些分数次幂，从而引出了新颖的双变量多项式递推关系。
- • Muonp 在数学推理、编程、困惑度和下游任务性能方面展现出相对于 Muon 的微调收益。
- • 我们通过谱学习率自适应、学习课程和有效秩动态的视角，检验了 Muonp 的优势和劣势。

## 2 Muonp：具有分数谱次数的 Muon

##### 动机：神经网络的优化器根据损失函数的梯度 G 计算权重的更新。最简单的优化器选择与 G 线性成正比的权重更新。AdamW 优化器通过权重中梯度的运行平均值来更新每个权重，并除以该梯度平方的运行平均值。Muon 利用了梯度矩阵和权重都可以解释为矩阵这一事实。如果 `G=USV^⊤` 是梯度矩阵的奇异值分解，那么 Muon 通过添加 `UV^⊤` 来更新权重，因此与最简单的优化器相比，权重在小奇异值方向上的更新更多，在大奇异值方向上的更新更少。对于某些应用，我们认为 Muon 过于极端：我们不应该完全消除奇异值中包含的信息，而应该在 Muon 和最简单的优化器之间进行折衷。Muonp 通过为选定的 p（介于 0 和 1 之间）添加 `US^p V^⊤` 来更新权重，从而实现这一点。

##### Muonp 在 Schatten q-范数下最大化损失改进。Muonp 的一个理论依据来自基于范数的最速下降：Muon 更新规则在权重更新的算子范数大小给定下最大化损失的线性改进 [6 (https://arxiv.org/html/2606.13867#bib.bib6)]。Muonp 在权重更新的 Schatten q-范数大小给定（其中 `q=1+1/p`）下最大化损失的线性改进。定理 2.1 (https://arxiv.org/html/2606.13867#S2.Thmtheorem1) 正式陈述了这一结果，证明见附录 A.2 (https://arxiv.org/html/2606.13867#A1.SS2)。使用 Schatten 范数进行最速下降的思想也在 [10 (https://arxiv.org/html/2606.13867#bib.bib10)] 和 [7 (https://arxiv.org/html/2606.13867#bib.bib7)] 中被考虑。

###### 定理 2.1。对于任意实数 `p∈(0,1)`，令 `q=1+1/p`。设 G 是一个 `n×n` 矩阵，奇异值分解为 `USV^⊤`。设 `⟨,⟩` 表示矩阵的逐元素点积，`||·||_q` 表示 Schatten q-范数。那么，在 `||M||_q` 有界的 n×n 矩阵 M 中，`⟨G,M⟩` 的最大值由 `US^p V^⊤` 的标量倍数达到。

##### 计算问题：Muonp 包含两个组成部分：更新规则 `US^p V^⊤` 的选择，以及近似 `US^p V^⊤` 的算法。要使 Muonp 高效，我们需要一个快速算法，在已知 `USV^⊤` 的情况下计算 `US^p V^⊤`。通过执行奇异值分解从 `USV^⊤` 计算 `US^p V^⊤` 的朴素算法太慢。在 Muon 中计算 `UV^⊤` 也存在同样的困难。Muon 通过使用迭代算法在知道 `G=USV^⊤` 的情况下快速近似 `UV^⊤` 来规避这个问题。我们证明了他们的迭代方法的直接迁移无法用于计算 `US^p V^⊤`，但 Muonp 使用一种新的变体方法，可以计算任意有理数 `p∈(0,1)` 的 `US^p V^⊤`。本节的其余部分致力于解释这种新的迭代方法。

##### Muon 公式回顾：对于矩阵 G 及其奇异值分解 `G=USV^⊤`（其中 S 是一个具有正项的对角矩阵），Muon 优化器使用 Newton-Schulz 算法来近似 `UV^⊤`，而无需计算奇异值分解。该算法基于两个观察：

1. 1. 我们有 `GG^⊤ G = US^3 V^⊤` 和 `GG^⊤ GG^⊤ GG^⊤ G = US^5 V^⊤`。对于任何奇多项式 `f(x) = Σ_{i=0}^d a_{2i+1} x^{2i+1}`，这些以及类似的恒等式允许我们仅使用矩阵乘法（而非奇异值分解）将 `Uf(S)V^⊤` 计算为 `Σ_{i=0}^d a_{2i+1} G (G^⊤ G)^i`。
2. 2. 存在奇多项式 f，使得迭代 `f(x), f(f(x)), f(f(f(x)))` 对起始范围 `(0,1]` 内的所有 x 收敛到 1。因此，固定任意这样的奇多项式 f，在将 G 归一化使其最大奇异值最多为 1 后，Muon 迭代地应用操作 `USV^⊤ → Uf(S)V^⊤` 来近似 `UV^⊤`，近似质量随迭代次数增加而提高。精确多项式 f 的选择是为了通过平衡达到收敛所需的迭代次数与每次迭代的计算时间（取决于多项式的次数），使计算足够好近似的时间尽可能短。

##### 我们的算法（引言）：Muonp 应用了该策略的一种推广，为介于 0 和 1 之间的任意有理数 p 计算 `US^p V^⊤`。（类似的方法可以应用于负指数或大于 1 的指数。）不可能通过迭代任何形如 `USV^⊤ → Uf(S)V^⊤` 的操作来实现这一点：

###### 引理 2.2（迭代单变量多项式不计算有理数次幂）。对于任何不等于 1 的实数 `p∈(0,1)`，不存在任何单变量多项式 f，使得对于所有奇异值最多为 1 的可逆矩阵 `G=USV^⊤`，如果我们设 `S_0=S` 并对所有 `n≥0` 设 `S_{n+1}=f(S_n)`，那么有 `lim_{n→∞} Uf(S_n)V^⊤ = US^p V^⊤`。

对于特征值而非奇异值，类似于引理 2.2 (https://arxiv.org/html/2606.13867#S2.Thmtheorem2) 的结论已在数值线性代数文献中被观察到 [13 (https://arxiv.org/html/2606.13867#bib.bib13), 16 (https://arxiv.org/html/2606.13867#bib.bib16)]。相反，为了收敛到 `US^p V^⊤`，迭代算法必须记住初始矩阵。因此，我们将考虑一个双变量多项式 `f(x,y)`。如果 `G=USV^⊤` 是我们的初始矩阵，Muonp 迭代地计算矩阵 `G_n = U S_n V^⊤`，其中 `S_0=S` 且 `S_{n+1}=f(S_n,S)`。此外，`S_n` 的选择使得 `G_n` 收敛到 `US^p V^⊤`。这可以对 `(0,1)` 中的任意*有理*数 p 实现。为了证明这是可能的，我们需要两个事实（其证明推迟到附录）：

##### 第一个关键事实（计算多项式）：对于 `f(x,y)` 是一个两个变量的多项式，在意义 `f(-x,-y) = -f(x,y)` 下是奇的，给定矩阵 `G_n = U S_n V^⊤` 和 `G = U S V^⊤`，我们可以使用矩阵乘法和转置来计算 `U f(S_n, S) V^⊤`。更精确地说，我们有如下定理：

###### 定理 2.3（计算奇数双变量多项式）。
1. 1. 设 `f(x,y)` 是一个具有实系数的两个变量的奇多项式。我们可以将 f 写成如下形式：
   `f(x,y) = Σ_{i,j=0}^d a_{2i+1,2j} x^{2i+1} y^{2j} + Σ_{i,j=0}^d a_{2i,2j+1} x^{2i} y^{2j+1}` (2.1)
   对于某个自然数 d 以及实数元组 `(a_{2i+1,2j})_{i,j=0}^d`, `(a_{2i,2j+1})_{i,j=0}^d`。
2. 2. 对于由公式 (2.1 (https://arxiv.org/html/2606.13867#S2.E1)) 给出的多项式 f，对于 `G=USV^⊤` 和 `G_n=U S_n V^⊤`，我们有：
   `U f(S_n, S) V^⊤ = Σ_{i,j} a_{2i+1,2j} G_n (G_n^⊤ G_n)^i (G^⊤ G)^j + Σ_{i,j} a_{2i,2j+1} G (G_n^⊤ G_n)^i (G^⊤ G)^j`.

##### 第二个关键事实（多项式的存在性）：存在奇多项式 `f(x,y)`，使得对于任意 `y∈(0,1]`，设 `x_0=y` 并对所有 `n≥0` 设 `x_{n+1}=f(x_n, y)`，我们有 `lim_{n→∞} x_n = y^p`。更精确地说，我们有如下定理：

###### 定理 2.4（奇数双变量多项式计算有理数次幂）。设 p 是 `(0,1)` 中的一个有理数。那么存在一个奇多项式 `f(x,y)`，使得对于任意 `y∈(0,1]`，如果我们设 `x_0=y` 并对所有 `n≥0` 设 `x_{n+1}=f(x_n, y)`，那么有 `lim_{n→∞} x_n = y^p`。

##### 我们的算法（结论）：对于 Muonp，我们选择如定理 2.4 (https://arxiv.org/html/2606.13867#S2.Thmtheorem4) 所述的多项式 f。给定一个输入矩阵 `G=USV^⊤`，我们设 `S_0=S`，对所有 `n≥0` 设 `S_{n+1}=f(S_n, S)`，以及对所有 `n≥0` 设 `G_n = U S_n V^⊤`。我们应用定理 2.3 (https://arxiv.org/html/2606.13867#S2.Thmtheorem3) 从 `G_n` 和 G 计算 `G_{n+1}`。根据定理 2.4 (https://arxiv.org/html/2606.13867#S2.Thmtheorem4)，我们有 `lim_{n→∞} S_n = S^p`，因此 `lim_{n→∞} G_n = U S^p V^⊤`。因此，在固定次数 N 次迭代后停止，我们计算 `US^p V^⊤` 的一个近似 `G_N`。当 `p=0` 且 `f(x,y)` 不依赖于 y 时，我们得到原始 Muon 算法作为 Muonp 的一个特例。与 Muon 一样，Muonp 在 bfloat16 实现中是数值稳定的，允许在现代硬件上快速实现。

算法 1 对于 `p=1/3` 的 Muonp 梯度更新。通过根据方程 (3.1 (https://arxiv.org/html/2606.13867#S3.E1)) 更改第 4 行，可以轻松适应 p 的其他值，包括 Muon 的 `p=0` 情况。
1: 输入：梯度 G，总步数 N。
2: 初始化 `G_0 ← G`
3: for n ← 0 到 N-1 do
4:    `G_{n+1} = G_n + c (G - G_n G_n^⊤ G_n)`
5: end for
6: 返回 Muonp 梯度更新 `G_N`.

##### 定理 2.3 (https://arxiv.org/html/2606.13867#S2.Thmtheorem3) 和定理 2.4 (https://arxiv.org/html/2606.13867#S2.Thmtheorem4) 证明中的关键思想。定理 2.3 (https://arxiv.org/html/2606.13867#S2.Thmtheorem3) 的证明通过应用定义，然后利用正交矩阵的定义消去所有不需要的因子 `U, U^⊤, V, V^⊤` 来完成。定理 2.4 (https://arxiv.org/html/2606.13867#S2.Thmtheorem4) 的证明需要构造