$\text{Muon}^p$：分数阶谱幂更新机制，全面提升十亿级大模型微调性能
Muon$^p$: Muon with Fractional Spectral Powers
在深度学习尤其是大语言模型的训练中，优化器的设计直接决定了模型学习的效率与最终的性能上限。长期以来，AdamW 等基于标量动量和方差调整的优化器占据了主导地位。然而，随着模型规模的爆炸式增长，研究人员开始将目光投向利用权重和梯度矩阵结构的二阶或伪二阶优化器。
ArXiv URL：https://arxiv.org/abs/2606.13867v1
近期，Muon 优化器凭借其在大型语言模型预训练中的卓越表现脱颖而出。Muon 的核心思想是通过极分解（Polar Decomposition）将梯度矩阵的奇异值谱完全展平，使得参数更新在各个奇异向量方向上保持均匀。这种机制有效防止了少数主导方向在训练初期垄断更新，从而在预训练阶段实现了显著的扩展性。
然而，普林斯顿大学（Princeton University）的研究团队敏锐地指出了这种“全谱展平”机制的局限性：在微调或特定任务适应阶段，完全抛弃奇异值所包含的幅度信息，反而会阻碍模型对关键特征的捕捉。为了解决这一两难困境，他们提出了一种全新的优化器——$\text{Muon}^p$。这不仅是 Muon 的一个优雅推广，更引入了“分数阶谱幂”（Fractional Spectral Powers）的概念，通过构建双变量多项式迭代算法，在保留 Muon 极致计算效率的同时，实现了从标准梯度下降到完全谱展平之间的连续插值。
谱展平的困境与 $\text{Muon}^p$ 的理论破局
理解 $\text{Muon}^p$ 的起点，在于剖析神经网络优化中的几何意义。假设在某一步更新中，损失函数关于某层权重的梯度矩阵为 $G$。我们可以对其进行奇异值分解（SVD），记作 $G = U S V^\top$，其中 $S$ 是包含奇异值的对角矩阵，代表了梯度在不同正交方向上的“强度”。
传统的梯度下降法直接使用 $G$ 进行更新，这意味着更新步长完全正比于奇异值的大小。而 Muon 优化器则采取了一种极端的策略：它剥离了 $S$，直接使用正交因子 $U V^\top$ 进行参数更新。从几何上看，这相当于把所有非零奇异值强制设为 1，即完全展平了奇异值谱（Flattening the singular spectrum）。这种做法在预训练初期非常有效，因为它强迫模型均匀地探索所有可能的特征方向。
但在微调（Finetuning）阶段，情况发生了根本性变化。预训练模型已经具备了良好的特征表示，微调本质上是在一个低维的、结构化的子空间内进行特定任务的适应。此时，较大的奇异值往往指向了对当前任务至关重要的残差方向，如果强行展平谱，就会丢失这些宝贵的适应性信号。
普林斯顿团队给出的答案是 $\text{Muon}^p$。它并没有完全丢弃 $S$，也没有完全保留 $S$，而是使用 $U S^p V^\top$ 作为更新方向，其中 $p$ 是介于 0 和 1 之间的有理数。
这一构造并非拍脑门得出的启发式设计，而是有着极其坚实的数学基础。研究证明，$\text{Muon}^p$ 实际上是在特定的矩阵范数下，使得损失函数局部下降最快的方向。具体而言，如果在算子范数下寻找最速下降方向，我们会得到 Muon 的更新规则；而 $\text{Muon}^p$ 则最大化了在 Schatten $q$-范数（其中 $q = 1 + \frac{1}{p}$）约束下的线性损失改进。这一理论结果完美地将 $\text{Muon}^p$ 置身于规范化最速下降（Norm-based steepest descent）的理论框架之中。
计算效率的奇迹：从单变量到双变量多项式
理论上 $U S^p V^\top$ 非常完美，但在工程实现上却面临着巨大的障碍。
要精确计算 $U S^p V^\top$，最直接的方法是进行完整的奇异值分解。然而，对于现代大模型动辄数千乘数千的权重矩阵来说，SVD 的计算复杂度极高，且在 GPU 上的并行效率极差，这在要求毫秒级响应的优化器单步更新中是完全不可接受的。
最初的 Muon 优化器之所以实用，是因为它巧妙地利用了牛顿-舒尔茨（Newton-Schulz）迭代法。这种方法只需通过简单的矩阵乘法，迭代计算一个固定的单变量奇多项式 $f(S)$，就能快速逼近 $U V^\top$。
自然地，研究人员试图寻找一个单变量多项式迭代来逼近分数阶谱幂。但普林斯顿团队在数学上给出了一个否定性的结论：对于任何不等于 1 的有理数 $p \in (0,1)$，不存在任何单变量多项式 $f$ 能够通过迭代 $\lim_{n \to \infty} U f(S_n) V^\top$ 来收敛到 $U S^p V^\top$。 原因在于，单变量多项式迭代本质上是对奇异值的无记忆映射，它无法将任意初始奇异值 $y$ 精确地吸引到其分数次幂 $y^p$ 上，除非它“记住”了最初的矩阵。
为了打破这一计算瓶颈，作者们创造性地引入了 低次奇双变量多项式递归（Low-degree odd bivariate recurrences）。
他们构造了一个双变量多项式 $f(x, y)$。如果初始矩阵是 $G = U S V^\top$，迭代过程变为计算 $G_n = U S_n V^\top$，其中 $S_0 = S$，而 $S_{n+1} = f(S_n, S)$。这里的 $y$ 就扮演了“记住初始矩阵”的角色。
但这又引发了下一个问题：我们只有矩阵 $G$ 和 $G_n$，如何在不进行 SVD 分解出 $U$、$S$、$V$ 的情况下，仅仅利用全矩阵级别的乘法来计算出 $U f(S_n, S) V^\top$ 呢？
研究团队证明，只要多项式 $f(x, y)$ 满足“奇函数”的性质，即 $f(-x, -y) = -f(x, y)$，那么通过提取 $G$ 和 $G_n$，并交替乘上 $(G_n^\top G_n)$ 和 $(G^\top G)$，所有的正交矩阵 $U$ 和 $V$ 因子都会在内部乘法中由于 $U^\top U = I$ 和 $V^\top V = I$ 的性质被完美抵消。最终，更新步骤可以完全化简为一系列标准的矩阵乘积。
例如，对于 $p = 1/3$，作者推导出的核心迭代形式为：
\[f(x,y) = x + c(y - x^3)\]对应到矩阵运算，不需要任何奇异值分解，只需对当前矩阵和初始梯度矩阵进行几次常数级别的矩阵乘法。这使得 $\text{Muon}^p$ 在算法层面上完全保留了 Muon 仅依赖矩阵乘法（Matrix-multiplication-only）的优良结构。在实际的 GPU 运行测试中，其单步优化时间与标准 Muon 处于同一量级，具备了在十亿、百亿参数规模上落地应用的基础。
实验验证：微调场景下的全面超越
为了验证 $\text{Muon}^p$ 的实际威力，研究团队在 Llama3-1B 等十亿级大模型上进行了大规模的微调实验，涵盖了从数学推理、代码编写到常识问答等多种下游任务。实验默认采用 $p=1/3$。
在数学推理与代码任务中，$\text{Muon}^p$ 的优势展现得淋漓尽致。
上图展示了在 Numina 数据集上微调 Llama3-1B 时，不同学习率下的超参数扫描热力图。可以清晰地看到，无论学习率如何变动，$\text{Muon}^p$ 的表现区间始终优于标准 Muon。最终，在 GSM8K、Numina 上的 pass@8 验证准确率，以及 OpenCodeInstruct 上的损失指标，$\text{Muon}^p$ 均实现了显著的超越。
在语言建模和常识推理任务上，使用 3 亿 Token 的 Fineweb 和 6.25 亿 Token 的 Pile of Law 数据进行微调时，$\text{Muon}^p$ 同样带来了更低的验证困惑度（Perplexity）和更高的下游评估指标（如 Winogrande、ARC 等）。值得注意的是，由于 Llama3-1B 原始模型是用 AdamW 预训练的，直接用 Muon 微调会产生优化器不匹配的问题。为了排除这一干扰，作者甚至专门用标准 Muon 重新从头预训练了一个 Llama3-1B 模型，随后在其上对比微调效果。结果依然坚挺：无论底层模型是用什么优化器预训练的，$\text{Muon}^p$ 在微调阶段的降维打击能力都稳定存在。
深度剖析：谱几何视角下的动力学机制
$\text{Muon}^p$ 为什么有效？为何它的优势主要体现在微调阶段，而在从零开始的预训练阶段却往往不如标准 Muon？研究团队通过谱几何（Spectral Geometry）的视角，给出了极为深刻的物理洞察。
如上图所示，在 135M 的 SmolLM 基座模型上，预训练阶段（左图）标准 Muon 的损失下降明显快于 $\text{Muon}^p$；而在微调阶段（右图），局势发生了逆转，$\text{Muon}^p$ 占据了主导地位。
这种差异的核心在于任务空间的本征维度。在预训练初期，模型的目标是尽可能广泛地探索并发现数据中的通用特征模式。此时，如果任由梯度按自然幅度更新，少数具有高频率或高增益的主导方向会迅速垄断参数的变化，导致那些代表联想记忆或长尾模式的奇异向量得不到充分学习。Muon 强行展平频谱，强制模型在所有方向上进行各向同性（Isotropic）的均匀更新，恰好契合了预训练“广撒网”的需求。
相反，在微调阶段，模型并不需要重新建立世界观。大量的研究表明，预训练模型针对下游任务的适应往往发生在一个极低维的子空间内。梯度矩阵中那些较大的奇异值，精确地标定了哪些特征方向是当前微调任务最迫切需要调整的。此时如果继续使用 $p=0$ 的 Muon，等于抹杀了这些极其宝贵的方向信号。$\text{Muon}^p$ 使用 $p \in (0,1)$，在抑制噪声的同时，温和地保留了重要主奇异向量的优先权。
为了验证这不是多项式逼近误差导致的假象，作者直接使用精确的 SVD 计算了 $UV^\top$ 和 $US^{1/3}V^\top$ 进行比对。
精确 SVD 的实验（上图）完美复现了前面的结论：预训练确实从无差别的全谱更新中获益，而特定领域的适应性微调必须依赖奇异值的幅度信息。
此外，引入分数幂 $p$，本质上等同于为每个奇异向量方向隐式地分配了一个自适应学习率。
上图展示了在微调中不同 $p$ 值的效果验证。可以看到，最优解既不是完全展平的 $p=0$（Muon），也不是完全保持原状的 $p=1$（普通梯度），而是落在 $p=1/3$ 到 $p=1/5$ 这样的中间地带。这证明了分数阶介于两者之间的插值特性，不仅具有极高的鲁棒性，还能切实逼近微调的最优几何流形。
更有趣的是，$\text{Muon}^p$ 与标准 Muon 之间存在极其平滑的过渡能力。在实际训练中，我们可以直接实施一种“课程学习”（Curriculum Learning）。
上图展示了在训练后期，不需要像切换到 AdamW 那样丢弃并重置复杂的动量状态，仅仅将优化器从 Muon 切换为 $\text{Muon}^p$（只需更改底层调用的逼近多项式），就能在最后 500 步中触发一次极其可观的损失骤降。这种无缝衔接的特性，赋予了它极强的工程灵活性。
最后，作者追踪了训练过程中模型权重的有效秩（Effective Rank）。
有效秩的变化曲线强有力地印证了理论猜想：由于 $\text{Muon}^p$ 相较于 Muon 更加压制微小奇异向量，集中力量在主导方向上进行更新，它在优化过程中实质上驱动模型在一个维度更低的子空间内进行演化。这对于旨在避免破坏预训练泛化能力的微调过程而言，是一种天然且极其优美的正则化。
总结
普林斯顿大学提出的 $\text{Muon}^p$，不仅仅是对现有优化算法的一次修补，更是通过深厚的代数几何与多项式逼近理论，为大语言模型的优化开辟了全新的设计空间。通过将计算效率极高的双变量奇多项式递归引入优化器底层，$\text{Muon}^p$ 成功找回了在谱展平过程中遗失的宝贵幅度信号。
这项工作清晰地揭示了这样一个深刻的规律：好的优化器不应仅仅追求梯度的快速下降，更应契合模型当前所处的几何流形。预训练需要各向同性的广阔探索，而微调则需要低维空间的精准聚焦。$\text{Muon}^p$ 提供了一个原则性的工具，让开发者能够在算力开销几乎不增加的情况下，优雅地掌控这股谱维度的力量。对于当前致力于挖掘十亿乃至百亿参数开源模型潜力的研究者与工业界团队而言，这无疑是一个极具参考价值与落地潜力的优化新利器。